Add a maximum timeout for shutting down

The shutdown operation now has a maximum duration to wait for active connections to complete. Signed-off-by: Alex Ellis (OpenFaaS Ltd) <[email protected]>
openfaas · Nov 2, 2021 · a693532 · a693532
1 parent b77a0ec
commit a693532
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 12 deletions.
diff --git a/template/golang-http/main.go b/template/golang-http/main.go
@@ -27,6 +27,7 @@ const defaultTimeout = 10 * time.Second
 func main() {
 	readTimeout := parseIntOrDurationValue(os.Getenv("read_timeout"), defaultTimeout)
 	writeTimeout := parseIntOrDurationValue(os.Getenv("write_timeout"), defaultTimeout)
+	healthInterval := parseIntOrDurationValue(os.Getenv("healthcheck_interval"), writeTimeout)
 
 	s := &http.Server{
 		Addr:           fmt.Sprintf(":%d", 8082),
@@ -36,28 +37,28 @@ func main() {
 	}
 
 	http.HandleFunc("/", makeRequestHandler())
-	listenUntilShutdown(s, writeTimeout)
+	listenUntilShutdown(s, healthInterval, writeTimeout)
 }
 
-func listenUntilShutdown(s *http.Server, shutdownTimeout time.Duration) {
+func listenUntilShutdown(s *http.Server, shutdownTimeout time.Duration, writeTimeout time.Duration) {
 	idleConnsClosed := make(chan struct{})
 	go func() {
 		sig := make(chan os.Signal, 1)
 		signal.Notify(sig, syscall.SIGTERM)
 
 		<-sig
 
-		log.Printf("[entrypoint] SIGTERM received.. shutting down server in %s\n", shutdownTimeout.String())
-
+		log.Printf("[entrypoint] SIGTERM: no connections in: %s", shutdownTimeout.String())
 		<-time.Tick(shutdownTimeout)
 
-		if err := s.Shutdown(context.Background()); err != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), writeTimeout)
+		defer cancel()
+
+		if err := s.Shutdown(ctx); err != nil {
 			log.Printf("[entrypoint] Error in Shutdown: %v", err)
 		}
 
-		log.Printf("[entrypoint] No new connections allowed. Exiting in: %s\n", shutdownTimeout.String())
-
-		<-time.Tick(shutdownTimeout)
+		log.Printf("[entrypoint] Exiting.")
 
 		close(idleConnsClosed)
 	}()

diff --git a/template/golang-middleware/main.go b/template/golang-middleware/main.go
@@ -35,10 +35,10 @@ func main() {
 
 	http.HandleFunc("/", function.Handle)
 
-	listenUntilShutdown(s, healthInterval)
+	listenUntilShutdown(s, healthInterval, writeTimeout)
 }
 
-func listenUntilShutdown(s *http.Server, shutdownTimeout time.Duration) {
+func listenUntilShutdown(s *http.Server, shutdownTimeout time.Duration, writeTimeout time.Duration) {
 	idleConnsClosed := make(chan struct{})
 	go func() {
 		sig := make(chan os.Signal, 1)
@@ -49,11 +49,14 @@ func listenUntilShutdown(s *http.Server, shutdownTimeout time.Duration) {
 		log.Printf("[entrypoint] SIGTERM: no connections in: %s", shutdownTimeout.String())
 		<-time.Tick(shutdownTimeout)
 
-		if err := s.Shutdown(context.Background()); err != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), writeTimeout)
+		defer cancel()
+
+		if err := s.Shutdown(ctx); err != nil {
 			log.Printf("[entrypoint] Error in Shutdown: %v", err)
 		}
 
-		log.Printf("[entrypoint] Exiting")
+		log.Printf("[entrypoint] Exiting.")
 
 		close(idleConnsClosed)
 	}()