server: implement graceful shutdown

2 files changed, 74 insertions, 8 deletions

diff --git a/cmd/server/main.go b/cmd/server/main.go
index 73e2fd3..426e2ff 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -5,6 +5,8 @@ import (
 	"log"
 	"log/slog"
 	"os"
+	"os/signal"
+	"sync"
 
 	"website/internal/server"
 
@@ -29,5 +31,34 @@ func main() {
 		log.Panicf("parsing runtime configuration: %v", err)
 	}
 
-	server.Start(&runtimeConfig)
+	c := make(chan os.Signal, 2)
+	signal.Notify(c, os.Interrupt)
+	sv, err := server.New(&runtimeConfig)
+	if err != nil {
+		log.Fatalf("error setting up server: %v", err)
+	}
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		sig := <-c
+		log.Printf("signal captured: %v", sig)
+		<-sv.Stop()
+		slog.Debug("server stopped")
+	}()
+
+	sErr := make(chan error)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		sErr <- sv.Start()
+	}()
+	log.Printf("server listening on %s", sv.Addr)
+
+	err = <-sErr
+	if err != nil {
+		// Error starting or closing listener:
+		log.Fatalf("error: %v", err)
+	}
+	wg.Wait()
 }
diff --git a/internal/server/server.go b/internal/server/server.go
index 9c43fdc..1404f46 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"log"
@@ -39,6 +40,10 @@ type HTTPError struct {
 	Code    int
 }
 
+type Server struct {
+	*http.Server
+}
+
 func canonicalisePath(path string) (cPath string, differs bool) {
 	cPath = path
 	if strings.HasSuffix(path, "/index.html") {
@@ -106,7 +111,7 @@ func fixupMIMETypes() {
 	}
 }
 
-func Start(runtimeConfig *Config) error {
+func New(runtimeConfig *Config) (*Server, error) {
 	fixupMIMETypes()
 
 	var err error
@@ -141,6 +146,7 @@ func Start(runtimeConfig *Config) error {
 		Repanic: true,
 	})
 
+	top := http.NewServeMux()
 	mux := http.NewServeMux()
 	slog.Debug("binding main handler to", "host", runtimeConfig.BaseURL.Hostname()+"/")
 	mux.Handle(runtimeConfig.BaseURL.Hostname()+"/", webHandler(serveFile))
@@ -156,7 +162,7 @@ func Start(runtimeConfig *Config) error {
 	} else {
 		logWriter = os.Stdout
 	}
-	http.Handle("/",
+	top.Handle("/",
 		sentryHandler.Handle(
 			wrapHandlerWithLogging(mux, wrappedHandlerOptions{
 				defaultHostname: runtimeConfig.BaseURL.Hostname(),
@@ -165,14 +171,43 @@ func Start(runtimeConfig *Config) error {
 		),
 	)
 	// no logging, no sentry
-	http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
+	top.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)
 	})
 
 	listenAddress := net.JoinHostPort(runtimeConfig.ListenAddress, runtimeConfig.Port)
+	return &Server{
+		&http.Server{
+			Addr:    listenAddress,
+			Handler: top,
+		},
+	}, nil
+}
+
+func (s *Server) Start() error {
+	if err := s.ListenAndServe(); err != http.ErrServerClosed {
+		return err
+	}
+	return nil
+}
+
+func (s *Server) Stop() chan struct{} {
+	slog.Debug("stop called")
+
+	idleConnsClosed := make(chan struct{})
+
+	go func() {
+		slog.Debug("shutting down server")
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		err := s.Server.Shutdown(ctx)
+		slog.Debug("server shut down")
+		if err != nil {
+			// Error from closing listeners, or context timeout:
+			log.Printf("HTTP server Shutdown: %v", err)
+		}
+		close(idleConnsClosed)
+	}()
 
-	done := make(chan bool)
-	go http.ListenAndServe(listenAddress, nil)
-	log.Printf("server listening on %s", runtimeConfig.BaseURL.String())
-	<-done
+	return idleConnsClosed
 }