Graceful shutdown

cmd/servercmd.go

@@ -2,10 +2,16 @@ package cmd
 
 import (
 	"context"
+	"errors"
 	"fmt"
-	"log"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
 
 	"github.com/spf13/cobra"
+	"go.ntppool.org/common/logger"
 	"go.ntppool.org/data-api/server"
 	"golang.org/x/sync/errgroup"
 )
@@ -25,25 +31,36 @@ func (cli *CLI) serverCmd() *cobra.Command {
 }
 
 func (cli *CLI) serverCLI(cmd *cobra.Command, args []string) error {
+	log := logger.Setup()
 
 	// cfg := cli.Config
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer cancel()
+
 	g, ctx := errgroup.WithContext(ctx)
 
+	srv, err := server.NewServer(ctx, cfgFile)
+	if err != nil {
+		return fmt.Errorf("srv setup: %s", err)
+	}
+
 	g.Go(func() error {
-		srv, err := server.NewServer(ctx, cfgFile)
-		if err != nil {
-			return fmt.Errorf("srv setup: %s", err)
-		}
 		return srv.Run()
 	})
 
-	err := g.Wait()
-	if err != nil {
-		log.Printf("server error: %s", err)
+	g.Go(func() error {
+		<-ctx.Done()
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		return srv.Shutdown(shutdownCtx)
+	})
+
+	err = g.Wait()
+	if err != nil && !errors.Is(err, http.ErrServerClosed) {
+		log.Error("server error", "err", err)
 	}
 
-	cancel()
-	return err
-
+	// don't tell cobra something went wrong as it'll just
+	// print usage information
+	return nil
 }