feat: add handling for SIGCHLD to cleanup child processes (#2043)

Author: jorng

cmd/proxy/main.go

@@ -63,22 +63,6 @@ func main() {
 		Handler:           handler,
 		ReadHeaderTimeout: 2 * time.Second,
 	}
-	idleConnsClosed := make(chan struct{})
-
-	go func() {
-		sigint := make(chan os.Signal, 1)
-		signal.Notify(sigint, shutdown.GetSignals()...)
-		s := <-sigint
-		logger.WithField("signal", s).Infof("Received signal, shutting down server")
-
-		// We received an interrupt signal, shut down.
-		ctx, cancel := context.WithTimeout(context.Background(), time.Second*time.Duration(conf.ShutdownTimeout))
-		defer cancel()
-		if err := srv.Shutdown(ctx); err != nil {
-			logger.WithError(err).Fatal("Could not shut down server")
-		}
-		close(idleConnsClosed)
-	}()
 
 	if conf.EnablePprof {
 		go func() {
@@ -111,15 +95,31 @@ func main() {
 		}
 	}
 
-	if conf.TLSCertFile != "" && conf.TLSKeyFile != "" {
-		err = srv.ServeTLS(ln, conf.TLSCertFile, conf.TLSKeyFile)
-	} else {
-		err = srv.Serve(ln)
-	}
+	signalCtx, signalStop := signal.NotifyContext(context.Background(), shutdown.GetSignals()...)
+	reaper := shutdown.ChildProcReaper(signalCtx, logger.Logger)
 
-	if !errors.Is(err, http.ErrServerClosed) {
-		logger.WithError(err).Fatal("Could not start server")
-	}
+	go func() {
+		defer signalStop()
+		if conf.TLSCertFile != "" && conf.TLSKeyFile != "" {
+			err = srv.ServeTLS(ln, conf.TLSCertFile, conf.TLSKeyFile)
+		} else {
+			err = srv.Serve(ln)
+		}
 
-	<-idleConnsClosed
+		if !errors.Is(err, http.ErrServerClosed) {
+			logger.WithError(err).Fatal("Could not start server")
+		}
+	}()
+
+	// Wait for shutdown signal, then cleanup before exit.
+	<-signalCtx.Done()
+	logger.Infof("Shutting down server")
+
+	// We received an interrupt signal, shut down.
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second*time.Duration(conf.ShutdownTimeout))
+	defer cancel()
+	if err := srv.Shutdown(shutdownCtx); err != nil {
+		logger.WithError(err).Fatal("Could not shut down server")
+	}
+	<-reaper.Done()
 }

internal/shutdown/signals.go

@@ -3,8 +3,13 @@
 package shutdown
 
 import (
+	"context"
+	"errors"
 	"os"
+	"os/signal"
 	"syscall"
+
+	"github.com/sirupsen/logrus"
 )
 
 // GetSignals returns the appropriate signals to catch for a clean shutdown, dependent on the OS.
@@ -13,3 +18,42 @@ import (
 func GetSignals() []os.Signal {
 	return []os.Signal{os.Interrupt, syscall.SIGTERM}
 }
+
+// ChildProcReaper spawns a goroutine to listen for SIGCHLD signals to cleanup
+// zombie child processes. The returned context will be canceled once all child
+// processes have been cleaned up, and should be waited on before exiting.
+//
+// This only applies to Unix platforms, and returns an already canceled context
+// on Windows.
+func ChildProcReaper(ctx context.Context, logger logrus.FieldLogger) context.Context {
+	sigChld := make(chan os.Signal, 1)
+	signal.Notify(sigChld, syscall.SIGCHLD)
+	done, cancel := context.WithCancel(context.WithoutCancel(ctx))
+	go func() {
+		defer cancel()
+		for {
+			select {
+			case <-ctx.Done():
+				reap(logger)
+				return
+			case <-sigChld:
+				reap(logger)
+			}
+		}
+	}()
+	return done
+}
+
+func reap(logger logrus.FieldLogger) {
+	for {
+		var wstatus syscall.WaitStatus
+		pid, err := syscall.Wait4(-1, &wstatus, syscall.WNOHANG, nil)
+		if err != nil && !errors.Is(err, syscall.ECHILD) {
+			logger.Errorf("failed to reap child process: %v", err)
+			continue
+		} else if pid <= 0 {
+			return
+		}
+		logger.Infof("reaped child process %v, exit status: %v", pid, wstatus.ExitStatus())
+	}
+}

internal/shutdown/signals_notunix.go

@@ -2,9 +2,26 @@
 
 package shutdown
 
-import "os"
+import (
+	"context"
+	"os"
+
+	"github.com/sirupsen/logrus"
+)
 
 // GetSignals returns the appropriate signals to catch for a clean shutdown, dependent on the OS.
 func GetSignals() []os.Signal {
 	return []os.Signal{os.Interrupt}
 }
+
+// ChildProcReaper spawns a goroutine to listen for SIGCHLD signals to cleanup
+// zombie child processes. The returned context will be canceled once all child
+// processes have been cleaned up, and should be waited on before exiting.
+//
+// This only applies to Unix platforms, and returns an already canceled context
+// on Windows.
+func ChildProcReaper(ctx context.Context, logger logrus.FieldLogger) context.Context {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	return ctx
+}