Implement graceful shutdown of HTTP server

README.md

@@ -137,10 +137,8 @@ Environmental variables:
 | `read_timeout`         | Yes          | HTTP timeout for reading the payload from the client caller (in seconds) |
 | `write_timeout`        | Yes          | HTTP timeout for writing a response body from your function (in seconds)  |
 | `exec_timeout`         | Yes          | Exec timeout for process exec'd for each incoming request (in seconds). Disabled if set to 0. |
-| `port`                 | Yes          | Specify an alternative TCP port fo testing |
+| `port`                 | Yes          | Specify an alternative TCP port for testing |
 | `write_debug`          | No           | Write all output, error messages, and additional information to the logs. Default is false. |
 | `content_type`         | Yes          | Force a specific Content-Type response for all responses - only in forking/serializing modes. |
-| `suppress_lock`        | No           | The watchdog will attempt to write a lockfile to /tmp/ for swarm healthchecks - set this to true to disable behaviour. |
+| `suppress_lock`        | Yes           | The watchdog will attempt to write a lockfile to /tmp/ for swarm healthchecks - set this to true to disable behaviour. |
 | `upstream_url`         | Yes          | `http` mode only - where to forward requests i.e. 127.0.0.1:5000 |
-
-> Note: the .lock file is implemented for health-checking, but cannot be disabled yet. You must create this file in /tmp/.

config/config.go

@@ -18,6 +18,7 @@ type WatchdogConfig struct {
 	ContentType      string
 	InjectCGIHeaders bool
 	OperationalMode  int
+	SuppressLock     bool
 }
 
 // Process returns a string for the process and a slice for the arguments from the FunctionProcess.
@@ -59,6 +60,7 @@ func New(env []string) (WatchdogConfig, error) {
 		ExecTimeout:      getDuration(envMap, "exec_timeout", time.Second*10),
 		OperationalMode:  ModeStreaming,
 		ContentType:      contentType,
+		SuppressLock:     getBool(envMap, "suppress_lock", false),
 	}
 
 	if val := envMap["mode"]; len(val) > 0 {
@@ -103,3 +105,13 @@ func getInt(env map[string]string, key string, defaultValue int) int {
 
 	return result
 }
+
+func getBool(env map[string]string, key string, defaultValue bool) bool {
+	result := defaultValue
+	if val, exists := env[key]; exists {
+		if val == "true" {
+			result = true
+		}
+	}
+	return result
+}

main.go

@@ -1,19 +1,26 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"log"
 	"net/http"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
+	"syscall"
+	"time"
 
 	"github.com/openfaas-incubator/of-watchdog/config"
 	"github.com/openfaas-incubator/of-watchdog/executor"
 )
 
+var acceptingConnections bool
+
 func main() {
+	acceptingConnections = false
 	watchdogConfig, configErr := config.New(os.Environ())
 	if configErr != nil {
 		fmt.Fprintf(os.Stderr, configErr.Error())
@@ -36,12 +43,22 @@ func main() {
 
 	log.Printf("OperationalMode: %s\n", config.WatchdogMode(watchdogConfig.OperationalMode))
 
-	if err := lock(); err != nil {
-		log.Panic(err.Error())
+	if !watchdogConfig.SuppressLock {
+		path, lockErr := lock()
+
+		if lockErr != nil {
+			log.Panicf("Cannot write %s. To disable lock-file set env suppress_lock=true.\n Error: %s.\n", path, lockErr.Error())
+		} else {
+			acceptingConnections = true
+		}
+	} else {
+		log.Println("Warning: \"suppress_lock\" is enabled. No automated health-checks will be in place for your function.")
+		acceptingConnections = true
 	}
 
 	http.HandleFunc("/", requestHandler)
-	log.Fatal(s.ListenAndServe())
+	// log.Fatal(s.ListenAndServe())
+	listenUntilShutdown(watchdogConfig.ExecTimeout, s)
 }
 
 func buildRequestHandler(watchdogConfig config.WatchdogConfig) http.HandlerFunc {
@@ -68,11 +85,12 @@ func buildRequestHandler(watchdogConfig config.WatchdogConfig) http.HandlerFunc
 	return requestHandler
 }
 
-func lock() error {
+func lock() (string, error) {
 	lockFile := filepath.Join(os.TempDir(), ".lock")
 	log.Printf("Writing lock file at: %s", lockFile)
-	return ioutil.WriteFile(lockFile, nil, 0600)
-
+	writeErr := ioutil.WriteFile(lockFile, nil, 0600)
+	acceptingConnections = true
+	return lockFile, writeErr
 }
 
 func makeAfterBurnRequestHandler(watchdogConfig config.WatchdogConfig) func(http.ResponseWriter, *http.Request) {
@@ -223,3 +241,33 @@ func makeHTTPRequestHandler(watchdogConfig config.WatchdogConfig) func(http.Resp
 
 	}
 }
+
+func listenUntilShutdown(shutdownTimeout time.Duration, s *http.Server) {
+	idleConnsClosed := make(chan struct{})
+	go func() {
+		sig := make(chan os.Signal, 1)
+		signal.Notify(sig, syscall.SIGTERM)
+
+		<-sig
+
+		log.Printf("SIGTERM received.. shutting down server")
+
+		acceptingConnections = false
+
+		if err := s.Shutdown(context.Background()); err != nil {
+			// Error from closing listeners, or context timeout:
+			log.Printf("Error in Shutdown: %v", err)
+		}
+
+		<-time.Tick(shutdownTimeout)
+
+		close(idleConnsClosed)
+	}()
+
+	if err := s.ListenAndServe(); err != http.ErrServerClosed {
+		log.Printf("Error ListenAndServe: %v", err)
+		close(idleConnsClosed)
+	}
+
+	<-idleConnsClosed
+}