feat: add graceful shutdown for proxy server (#776)

Signed-off-by: Anish Ramasekar <anish.ramasekar@gmail.com>

cmd/proxy/main.go

@@ -1,8 +1,12 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"monis.app/mlog"
 	"sigs.k8s.io/controller-runtime/pkg/manager/signals"
@@ -48,13 +52,31 @@ func mainErr() error {
 		return nil
 	}
 
+	ctx := withShutdownSignal(context.Background())
+
 	p, err := proxy.NewProxy(proxyPort, mlog.New().WithName("proxy"))
 	if err != nil {
 		return fmt.Errorf("setup: failed to create proxy: %w", err)
 	}
-	if err := p.Run(); err != nil {
+	if err := p.Run(ctx); err != nil {
 		return fmt.Errorf("setup: failed to run proxy: %w", err)
 	}
 
 	return nil
 }
+
+// withShutdownSignal returns a copy of the parent context that will close if
+// the process receives termination signals.
+func withShutdownSignal(ctx context.Context) context.Context {
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, syscall.SIGTERM, syscall.SIGINT, os.Interrupt)
+
+	nctx, cancel := context.WithCancel(ctx)
+
+	go func() {
+		<-signalChan
+		mlog.Info("received shutdown signal")
+		cancel()
+	}()
+	return nctx
+}

pkg/proxy/proxy.go

@@ -40,7 +40,7 @@ var (
 )
 
 type Proxy interface {
-	Run() error
+	Run(ctx context.Context) error
 }
 
 type proxy struct {
@@ -88,7 +88,7 @@ func NewProxy(port int, logger mlog.Logger) (Proxy, error) {
 }
 
 // Run runs the proxy server
-func (p *proxy) Run() error {
+func (p *proxy) Run(ctx context.Context) error {
 	rtr := mux.NewRouter()
 	rtr.PathPrefix(tokenPathPrefix).HandlerFunc(p.msiHandler)
 	rtr.PathPrefix(readyzPathPrefix).HandlerFunc(p.readyzHandler)
@@ -100,7 +100,20 @@ func (p *proxy) Run() error {
 		ReadHeaderTimeout: 5 * time.Second,
 		Handler:           rtr,
 	}
-	return server.ListenAndServe()
+
+	go func() {
+		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			panic(err)
+		}
+	}()
+
+	<-ctx.Done()
+
+	p.logger.Info("shutting down the proxy server")
+	// shutdown the server gracefully with a 5 second timeout
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	return server.Shutdown(shutdownCtx)
 }
 
 func (p *proxy) msiHandler(w http.ResponseWriter, r *http.Request) {