From 197d8e4cb0170356dd20755efcf1d336c4c38583 Mon Sep 17 00:00:00 2001
From: tjpcc <tjp@ctrl-c.club>
Date: Wed, 11 Jan 2023 10:33:44 -0700
Subject: 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.
---
 examples/cgi/main.go | 14 ++++++++++++--
 gemini/serve.go      | 42 ++++++++++++++++++++++++++++++------------
 2 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/examples/cgi/main.go b/examples/cgi/main.go
index d1df220..3dfec29 100644
--- a/examples/cgi/main.go
+++ b/examples/cgi/main.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"log"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"tildegit.org/tjp/gus/contrib/cgi"
 	guslog "tildegit.org/tjp/gus/contrib/log"
@@ -26,12 +28,20 @@ func main() {
 	// add stdout logging to the request handler
 	handler := guslog.Requests(os.Stdout, nil)(cgiHandler)
 
+	// set up signals to trigger graceful shutdown
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGHUP)
+	defer stop()
+
 	// run the server
-	server, err := gemini.NewServer(context.Background(), tlsconf, "tcp4", ":1965", handler)
+	server, err := gemini.NewServer(ctx, tlsconf, "tcp4", ":1965", handler)
 	if err != nil {
 		log.Fatal(err)
 	}
-	server.Serve()
+	defer server.Close()
+
+	if err := server.Serve(); err != nil {
+		log.Fatal(err)
+	}
 }
 
 func envConfig() (string, string) {
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()
 	}()
-- 
cgit v1.2.3