From d463996ce9dd24710763f32d77b3347b474a2ed5 Mon Sep 17 00:00:00 2001
From: Alexander Neumann <alexander@bumpern.de>
Date: Fri, 26 Apr 2024 19:21:32 +0200
Subject: [PATCH] serve: Properly shut down the server

---
 cmd/restic/cmd_serve.go | 41 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/cmd/restic/cmd_serve.go b/cmd/restic/cmd_serve.go
index d00ce201e..d101ec2c6 100644
--- a/cmd/restic/cmd_serve.go
+++ b/cmd/restic/cmd_serve.go
@@ -2,7 +2,10 @@ package main
 
 import (
 	"context"
+	"fmt"
+	"net"
 	"net/http"
+	"time"
 
 	"github.com/spf13/cobra"
 
@@ -35,6 +38,8 @@ func init() {
 	cmdFlags.StringVarP(&serveOptions.Listen, "listen", "l", "localhost:3080", "set the listen host name and `address`")
 }
 
+const serverShutdownTimeout = 30 * time.Second
+
 func runWebServer(ctx context.Context, opts ServeOptions, gopts GlobalOptions, args []string) error {
 	if len(args) > 0 {
 		return errors.Fatal("this command does not accept additional arguments")
@@ -57,10 +62,42 @@ func runWebServer(ctx context.Context, opts ServeOptions, gopts GlobalOptions, a
 		return err
 	}
 
-	srv := server.New(repo, snapshotLister, TimeFormat)
+	srv := http.Server{
+		BaseContext: func(l net.Listener) context.Context {
+			// just return the global context
+			return ctx
+		},
+		Handler: server.New(repo, snapshotLister, TimeFormat),
+	}
+
+	listener, err := net.Listen("tcp", opts.Listen)
+	if err != nil {
+		return fmt.Errorf("start listener: %v", err)
+	}
+
+	// wait until context is cancelled, then close listener
+	go func() {
+		<-ctx.Done()
+		Printf("gracefully shutting down server\n")
+
+		ctxTimeout, cancel := context.WithTimeout(context.Background(), serverShutdownTimeout)
+		defer cancel()
+
+		_ = srv.Shutdown(ctxTimeout)
+	}()
 
 	Printf("Now serving the repository at http://%s\n", opts.Listen)
 	Printf("When finished, quit with Ctrl-c here.\n")
 
-	return http.ListenAndServe(opts.Listen, srv)
+	err = srv.Serve(listener)
+
+	if errors.Is(err, http.ErrServerClosed) {
+		err = nil
+	}
+
+	if err != nil {
+		return fmt.Errorf("serve: %v", err)
+	}
+
+	return nil
 }