Improvements to Server lifecycle.

- NewServer doesn't allocate any resources besides the server object
  itself. So eg context.WithCancel is delayed until s.Serve().
- Add a demonstration of graceful shutdown on signals to the cgi
  example.

1 files changed, 30 insertions, 12 deletions

diff --git a/gemini/serve.go b/gemini/serve.go
index 8fd6b57..abf127e 100644
--- a/gemini/serve.go
+++ b/gemini/serve.go
@@ -25,11 +25,8 @@ func NewServer(
 	address string,
 	handler Handler,
 ) (*Server, error) {
-	ctx, cancel := context.WithCancel(ctx)
-
 	listener, err := net.Listen(network, address)
 	if err != nil {
-		cancel()
 		return nil, err
 	}
 
@@ -37,29 +34,34 @@ func NewServer(
 		ctx:      ctx,
 		network:  network,
 		address:  address,
-		cancel:   cancel,
 		wg:       &sync.WaitGroup{},
 		listener: tls.NewListener(listener, tlsConfig),
 		handler:  handler,
 	}
-	go s.propagateCancel()
 
 	return s, nil
 }
 
-func (s *Server) Close() {
-	s.cancel()
-	s.wg.Wait()
-}
-
-func (s *Server) Serve() {
+// Serve starts the server and blocks until it is closed.
+//
+// This function will allocate resources which are not cleaned up until
+// Close() is called.
+func (s *Server) Serve() error {
 	s.wg.Add(1)
 	defer s.wg.Done()
 
+	s.ctx, s.cancel = context.WithCancel(s.ctx)
+
+	s.wg.Add(1)
+	go s.propagateCancel()
+
 	for {
 		conn, err := s.listener.Accept()
 		if err != nil {
-			return
+			if s.closed() {
+				err = nil
+			}
+			return err
 		}
 
 		s.wg.Add(1)
@@ -67,19 +69,33 @@ func (s *Server) Serve() {
 	}
 }
 
+// Close begins a graceful shutdown of the server.
+//
+// It cancels the server's context which interrupts all concurrently running
+// request handlers, if they support it. It then blocks until all resources
+// have been cleaned up and all request handlers have completed.
+func (s *Server) Close() {
+	s.cancel()
+	s.wg.Wait()
+}
+
+// Network returns the network type on which the server is running.
 func (s *Server) Network() string {
 	return s.network
 }
 
+// Address returns the address on which the server is listening.
 func (s *Server) Address() string {
 	return s.address
 }
 
+// Hostname returns just the hostname portion of the listen address.
 func (s *Server) Hostname() string {
 	host, _, _ := net.SplitHostPort(s.address)
 	return host
 }
 
+// Port returns the port on which the server is listening.
 func (s *Server) Port() string {
 	_, portStr, _ := net.SplitHostPort(s.address)
 	return portStr
@@ -110,6 +126,8 @@ func (s *Server) handleConn(conn net.Conn) {
 
 func (s *Server) propagateCancel() {
 	go func() {
+		defer s.wg.Done()
+
 		<-s.ctx.Done()
 		_ = s.listener.Close()
 	}()