Gracefully shutdown HTTP, gRPC, and Prom servers

cmd/app/grpc.go

@@ -21,8 +21,10 @@ import (
 	"fmt"
 	"net"
 	"os"
+	"os/signal"
 	"runtime"
 	"sync"
+	"syscall"
 
 	"github.com/fsnotify/fsnotify"
 	"github.com/goadesign/goa/grpc/middleware"
@@ -197,7 +199,7 @@ func (g *grpcServer) setupPrometheus(reg *prometheus.Registry) {
 	grpc_prometheus.Register(g.Server)
 }
 
-func (g *grpcServer) startTCPListener() {
+func (g *grpcServer) startTCPListener(wg *sync.WaitGroup) {
 	// lis is closed by g.Server.Serve() upon exit
 	lis, err := net.Listen("tcp", g.grpcServerEndpoint)
 	if err != nil {
@@ -206,13 +208,30 @@ func (g *grpcServer) startTCPListener() {
 
 	g.grpcServerEndpoint = lis.Addr().String()
 	log.Logger.Infof("listening on grpc at %s", g.grpcServerEndpoint)
+
+	idleConnsClosed := make(chan struct{})
+	go func() {
+		sigint := make(chan os.Signal, 1)
+		signal.Notify(sigint, syscall.SIGINT, syscall.SIGTERM)
+		<-sigint
+
+		// received an interrupt signal, shut down
+		g.Server.GracefulStop()
+		close(idleConnsClosed)
+		log.Logger.Info("stopped grpc server")
+	}()
+
+	wg.Add(1)
 	go func() {
 		if g.tlsCertWatcher != nil {
 			defer g.tlsCertWatcher.Close()
 		}
 		if err := g.Server.Serve(lis); err != nil {
 			log.Logger.Errorf("error shutting down grpcServer: %w", err)
 		}
+		<-idleConnsClosed
+		wg.Done()
+		log.Logger.Info("grpc server shutdown")
 	}()
 }
 

cmd/app/http.go

@@ -21,8 +21,12 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"os"
+	"os/signal"
 	"strconv"
 	"strings"
+	"sync"
+	"syscall"
 	"time"
 
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
@@ -102,12 +106,32 @@ func createHTTPServer(ctx context.Context, serverEndpoint string, grpcServer, le
 	return httpServer{&api, serverEndpoint}
 }
 
-func (h httpServer) startListener() {
+func (h httpServer) startListener(wg *sync.WaitGroup) {
 	log.Logger.Infof("listening on http at %s", h.httpServerEndpoint)
+
+	idleConnsClosed := make(chan struct{})
+	go func() {
+		sigint := make(chan os.Signal, 1)
+		signal.Notify(sigint, syscall.SIGINT, syscall.SIGTERM)
+		<-sigint
+
+		// received an interrupt signal, shut down
+		if err := h.Shutdown(context.Background()); err != nil {
+			// error from closing listeners, or context timeout
+			log.Logger.Errorf("HTTP server Shutdown: %v", err)
+		}
+		close(idleConnsClosed)
+		log.Logger.Info("stopped http server")
+	}()
+
+	wg.Add(1)
 	go func() {
 		if err := h.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
-			log.Logger.Fatal(err)
+			log.Logger.Error(err)
 		}
+		<-idleConnsClosed
+		wg.Done()
+		log.Logger.Info("http server shutdown")
 	}()
 }
 

cmd/app/http_test.go

@@ -27,6 +27,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/sigstore/fulcio/pkg/ca"
@@ -50,7 +51,8 @@ func setupHTTPServer(t *testing.T) (httpServer, string) {
 	if err != nil {
 		t.Error(err)
 	}
-	grpcServer.startTCPListener()
+	var wg sync.WaitGroup
+	grpcServer.startTCPListener(&wg)
 	conn, err := grpc.Dial(grpcServer.grpcServerEndpoint, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	defer func() {
 		if conn != nil {
@@ -95,7 +97,9 @@ func setupHTTPServerWithGRPCTLS(t *testing.T) (httpServer, string) {
 	if err != nil {
 		t.Error(err)
 	}
-	grpcServer.startTCPListener()
+
+	var wg sync.WaitGroup
+	grpcServer.startTCPListener(&wg)
 	conn, err := grpc.Dial(grpcServer.grpcServerEndpoint, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	defer func() {
 		if conn != nil {

cmd/app/serve.go

@@ -19,13 +19,17 @@ import (
 	"bytes"
 	"context"
 	"crypto/x509"
+	"errors"
 	"flag"
 	"fmt"
 	"net"
 	"net/http"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
+	"sync"
+	"syscall"
 	"time"
 
 	"chainguard.dev/go-grpc-kit/pkg/duplex"
@@ -286,6 +290,9 @@ func runServeCmd(cmd *cobra.Command, args []string) { //nolint: revive
 		return
 	}
 
+	// waiting for http and grpc servers to shutdown gracefully
+	var wg sync.WaitGroup
+
 	httpServerEndpoint := fmt.Sprintf("%v:%v", viper.GetString("http-host"), viper.GetString("http-port"))
 
 	reg := prometheus.NewRegistry()
@@ -295,7 +302,7 @@ func runServeCmd(cmd *cobra.Command, args []string) { //nolint: revive
 		log.Logger.Fatal(err)
 	}
 	grpcServer.setupPrometheus(reg)
-	grpcServer.startTCPListener()
+	grpcServer.startTCPListener(&wg)
 
 	legacyGRPCServer, err := createLegacyGRPCServer(cfg, grpcServer.caService)
 	if err != nil {
@@ -304,15 +311,37 @@ func runServeCmd(cmd *cobra.Command, args []string) { //nolint: revive
 	legacyGRPCServer.startUnixListener()
 
 	httpServer := createHTTPServer(ctx, httpServerEndpoint, grpcServer, legacyGRPCServer)
-	httpServer.startListener()
+	httpServer.startListener(&wg)
 
 	readHeaderTimeout := viper.GetDuration("read-header-timeout")
 	prom := http.Server{
 		Addr:              fmt.Sprintf(":%v", viper.GetString("metrics-port")),
 		Handler:           promhttp.Handler(),
 		ReadHeaderTimeout: readHeaderTimeout,
 	}
-	log.Logger.Error(prom.ListenAndServe())
+
+	idleConnsClosed := make(chan struct{})
+	go func() {
+		sigint := make(chan os.Signal, 1)
+		signal.Notify(sigint, syscall.SIGINT, syscall.SIGTERM)
+		<-sigint
+
+		// received an interrupt signal, shut down
+		if err := prom.Shutdown(context.Background()); err != nil {
+			// error from closing listeners, or context timeout
+			log.Logger.Errorf("HTTP server Shutdown: %v", err)
+		}
+		close(idleConnsClosed)
+		log.Logger.Info("stopped prom server")
+	}()
+	if err := prom.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+		log.Logger.Error(err)
+	}
+	<-idleConnsClosed
+	log.Logger.Info("prom server shutdown")
+
+	// wait for http and grpc servers to shutdown
+	wg.Wait()
 }
 
 func checkServeCmdConfigFile() error {