Add one-shot mode for sealed execution

Sealed execution allows only one request to execute before the
process exits, causing a restart, for a fresh process.

With SlicerVM, this means a hardware per request barrier.

Signed-off-by: Alex Ellis (OpenFaaS Ltd) <alexellis2@gmail.com>

Author: alexellis

README.md

@@ -205,6 +205,7 @@ Environmental variables:
 | `log_call_id`                    | In HTTP mode, when printing a response code, content-length and timing, include the X-Call-Id header at the end of the line in brackets i.e. `[079d9ff9-d7b7-4e37-b195-5ad520e6f797]` or `[none]` when it's empty. Default: `false` |
 | `max_inflight`                   |  Limit the maximum number of requests in flight, and return a HTTP status 429 when exceeded           |
 | `mode`                           |  The mode which of-watchdog operates in, Default `streaming` [see doc](#3-streaming-fork-modestreaming---default). Options are [http](#1-http-modehttp), [serialising fork](#2-serializing-fork-modeserializing), [streaming fork](#3-streaming-fork-modestreaming---default), [static](#4-static-modestatic) |
+| `one_shot`                       |  When set to `true`, accept the first genuine invoke request, then immediately begin graceful shutdown and reject subsequent invoke requests. Readiness and health endpoints do not trigger this mode. |
 | `port`                           |  Specify an alternative TCP port for testing. Default: `8080`            |
 | `prefix_logs`                    |  When set to `true` the watchdog will add a prefix of "Date Time" + "stderr/stdout" to every line read from the function process. Default `true`             |
 | `read_timeout`                   |  HTTP timeout for reading the payload from the client caller (in seconds)          |
@@ -214,11 +215,12 @@ Environmental variables:
 | `upstream_url`                   |  Alias for `http_upstream_url`                                                          |
 | `write_timeout`                  |  HTTP timeout for writing a response body from your function (in seconds)          |
 
+`shutdown_after_first_request` remains accepted as a deprecated alias for `one_shot`.
+
 Unsupported options from the [Classic Watchdog](https://github.com/openfaas/classic-watchdog):
 
 | Option               | Usage                                                                                         |
 | -------------------- | --------------------------------------------------------------------------------------------- |
 | `write_debug`        | In the classic watchdog, this prints the response body out to the console |
 | `read_debug`         | In the classic watchdog, this prints the request body out to the console |
 | `combined_output`    | In the classic watchdog, this returns STDOUT and STDERR in the function's HTTP response, when off it only returns STDOUT and prints STDERR to the logs of the watchdog |
-

config/config.go

@@ -69,6 +69,10 @@ type WatchdogConfig struct {
 	// HTTP mode.
 	LogCallId bool
 
+	// OneShot marks the watchdog as unhealthy and begins graceful shutdown
+	// when the first genuine invoke request starts.
+	OneShot bool
+
 	// Handler is the HTTP handler to use in "inproc" mode
 	Handler http.HandlerFunc
 }
@@ -178,6 +182,7 @@ func New(env []string) (WatchdogConfig, error) {
 		LogBufferSize:       logBufferSize,
 		ReadyEndpoint:       envMap["ready_path"],
 		LogCallId:           logCallId,
+		OneShot:             getBool(envMap, "one_shot") || getBool(envMap, "shutdown_after_first_request"),
 	}
 
 	if val := envMap["mode"]; len(val) > 0 {

config/config_test.go

@@ -389,3 +389,61 @@ func Test_NonParsableString_parseIntOrDurationValue(t *testing.T) {
 		t.Error(fmt.Sprintf("want: %q got: %q", want, got))
 	}
 }
+
+func TestNewParsesOneShot(t *testing.T) {
+	cfg, err := New([]string{
+		"fprocess=/bin/cat",
+		"mode=streaming",
+		"one_shot=true",
+	})
+	if err != nil {
+		t.Fatalf("expected config to parse, got error: %v", err)
+	}
+
+	if !cfg.OneShot {
+		t.Fatalf("expected OneShot to be true")
+	}
+}
+
+func TestNewParsesDeprecatedShutdownAfterFirstRequestAlias(t *testing.T) {
+	cfg, err := New([]string{
+		"fprocess=/bin/cat",
+		"mode=streaming",
+		"shutdown_after_first_request=true",
+	})
+	if err != nil {
+		t.Fatalf("expected config to parse, got error: %v", err)
+	}
+
+	if !cfg.OneShot {
+		t.Fatalf("expected deprecated alias to enable OneShot")
+	}
+}
+
+func TestNewParsesOneShotNumericValues(t *testing.T) {
+	cfg, err := New([]string{
+		"fprocess=/bin/cat",
+		"mode=streaming",
+		"one_shot=1",
+	})
+	if err != nil {
+		t.Fatalf("expected config to parse, got error: %v", err)
+	}
+
+	if !cfg.OneShot {
+		t.Fatalf("expected one_shot=1 to enable OneShot")
+	}
+
+	cfg, err = New([]string{
+		"fprocess=/bin/cat",
+		"mode=streaming",
+		"one_shot=0",
+	})
+	if err != nil {
+		t.Fatalf("expected config to parse, got error: %v", err)
+	}
+
+	if cfg.OneShot {
+		t.Fatalf("expected one_shot=0 to disable OneShot")
+	}
+}

pkg/watchdog.go

@@ -2,6 +2,7 @@ package pkg
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -12,6 +13,7 @@ import (
 	"os/signal"
 	"path/filepath"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"syscall"
 	"time"
@@ -50,8 +52,25 @@ func (w *Watchdog) Start(ctx context.Context) error {
 	baseFunctionHandler := buildRequestHandler(w.config, w.config.PrefixLogs)
 	requestHandler := baseFunctionHandler
 
+	drainCh := make(chan string, 1)
+	var drainOnce sync.Once
+	startDrain := func(reason string) {
+		drainOnce.Do(func() {
+			if err := markUnhealthy(); err != nil {
+				log.Printf("Unable to mark server as unhealthy: %s\n", err.Error())
+			}
+
+			log.Printf("Scheduling graceful shutdown: %s\n", reason)
+			drainCh <- reason
+		})
+	}
+
+	if w.config.OneShot {
+		requestHandler = makeOneShotHandler(requestHandler, w.config.ReadyEndpoint, startDrain)
+	}
+
 	if w.config.JWTAuthentication {
-		handler, err := makeJWTAuthHandler(w.config, baseFunctionHandler)
+		handler, err := makeJWTAuthHandler(w.config, requestHandler)
 		if err != nil {
 			return fmt.Errorf("error creating JWTAuthMiddleware: %w", err)
 		}
@@ -111,6 +130,7 @@ func (w *Watchdog) Start(ctx context.Context) error {
 		w.config.HealthcheckInterval,
 		w.config.HTTPWriteTimeout,
 		w.config.SuppressLock,
+		drainCh,
 		&httpMetrics)
 
 	return nil
@@ -122,32 +142,46 @@ func markUnhealthy() error {
 	path := filepath.Join(os.TempDir(), ".lock")
 	log.Printf("Removing lock-file : %s\n", path)
 	removeErr := os.Remove(path)
+	if errors.Is(removeErr, os.ErrNotExist) {
+		return nil
+	}
 	return removeErr
 }
 
-func listenUntilShutdown(s *http.Server, shutdownCtx context.Context, healthcheckInterval time.Duration, writeTimeout time.Duration, suppressLock bool, httpMetrics *metrics.Http) error {
+func listenUntilShutdown(s *http.Server, shutdownCtx context.Context, healthcheckInterval time.Duration, writeTimeout time.Duration, suppressLock bool, drain <-chan string, httpMetrics *metrics.Http) error {
 
 	idleConnsClosed := make(chan struct{})
 	go func() {
 		sig := make(chan os.Signal, 1)
 		signal.Notify(sig, syscall.SIGTERM)
 
 		reason := ""
+		drainDelay := healthcheckInterval
 
 		select {
 		case <-sig:
 			reason = "SIGTERM"
 		case <-shutdownCtx.Done():
 			reason = "Context cancelled"
+		case reason = <-drain:
+			drainDelay = 0
 		}
 
-		log.Printf("%s: no new connections in %s\n", reason, healthcheckInterval.String())
+		if drainDelay > 0 {
+			log.Printf("%s: no new connections in %s\n", reason, drainDelay.String())
+		} else {
+			log.Printf("%s: no new connections allowed immediately\n", reason)
+		}
 
 		if err := markUnhealthy(); err != nil {
 			log.Printf("Unable to mark server as unhealthy: %s\n", err.Error())
 		}
 
-		<-time.Tick(healthcheckInterval)
+		if drainDelay > 0 {
+			timer := time.NewTimer(drainDelay)
+			defer timer.Stop()
+			<-timer.C
+		}
 
 		connections := int64(testutil.ToFloat64(httpMetrics.InFlight))
 		log.Printf("No new connections allowed, draining: %d requests\n", connections)
@@ -193,6 +227,34 @@ func listenUntilShutdown(s *http.Server, shutdownCtx context.Context, healthchec
 	return nil
 }
 
+func makeOneShotHandler(next http.Handler, readyEndpoint string, startDrain func(string)) http.Handler {
+	var served int32
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL != nil {
+			switch r.URL.Path {
+			case "/_/health", "/_/ready":
+				next.ServeHTTP(w, r)
+				return
+			case readyEndpoint:
+				if readyEndpoint != "" {
+					next.ServeHTTP(w, r)
+					return
+				}
+			}
+		}
+
+		if !atomic.CompareAndSwapInt32(&served, 0, 1) {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			w.Write([]byte("watchdog is draining after serving a request"))
+			return
+		}
+
+		startDrain("one_shot")
+		next.ServeHTTP(w, r)
+	})
+}
+
 func buildRequestHandler(cfg config.WatchdogConfig, prefixLogs bool) http.Handler {
 	var requestHandler http.HandlerFunc
 

pkg/watchdog_test.go

@@ -0,0 +1,77 @@
+package pkg
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+)
+
+func TestMakeOneShotHandlerDrainsAndRejectsSubsequentRequests(t *testing.T) {
+	var calls int32
+	var drains int32
+
+	handler := makeOneShotHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&calls, 1)
+		w.WriteHeader(http.StatusAccepted)
+		w.Write([]byte("ok"))
+	}), "", func(reason string) {
+		atomic.AddInt32(&drains, 1)
+	})
+
+	firstReq := httptest.NewRequest(http.MethodPost, "/", nil)
+	firstRes := httptest.NewRecorder()
+	handler.ServeHTTP(firstRes, firstReq)
+
+	if firstRes.Code != http.StatusAccepted {
+		t.Fatalf("expected first request to pass through, got status %d", firstRes.Code)
+	}
+
+	secondReq := httptest.NewRequest(http.MethodPost, "/", nil)
+	secondRes := httptest.NewRecorder()
+	handler.ServeHTTP(secondRes, secondReq)
+
+	if secondRes.Code != http.StatusServiceUnavailable {
+		t.Fatalf("expected second request to be rejected, got status %d", secondRes.Code)
+	}
+
+	if got := atomic.LoadInt32(&calls); got != 1 {
+		t.Fatalf("expected next handler to be called once, got %d", got)
+	}
+
+	if got := atomic.LoadInt32(&drains); got != 1 {
+		t.Fatalf("expected drain to be scheduled once, got %d", got)
+	}
+}
+
+func TestMakeOneShotHandlerIgnoresReadyEndpoint(t *testing.T) {
+	var calls int32
+	var drains int32
+
+	handler := makeOneShotHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt32(&calls, 1)
+		w.WriteHeader(http.StatusOK)
+	}), "/ready", func(reason string) {
+		atomic.AddInt32(&drains, 1)
+	})
+
+	readyReq := httptest.NewRequest(http.MethodGet, "/ready", nil)
+	readyRes := httptest.NewRecorder()
+	handler.ServeHTTP(readyRes, readyReq)
+
+	firstReq := httptest.NewRequest(http.MethodPost, "/", nil)
+	firstRes := httptest.NewRecorder()
+	handler.ServeHTTP(firstRes, firstReq)
+
+	if firstRes.Code != http.StatusOK {
+		t.Fatalf("expected first real request to pass through, got status %d", firstRes.Code)
+	}
+
+	if got := atomic.LoadInt32(&calls); got != 2 {
+		t.Fatalf("expected handler to be called for readiness and first invoke, got %d", got)
+	}
+
+	if got := atomic.LoadInt32(&drains); got != 1 {
+		t.Fatalf("expected drain to be scheduled once for the real request, got %d", got)
+	}
+}

vendor/modules.txt

@@ -10,10 +10,6 @@ github.com/docker/go-units
 # github.com/golang-jwt/jwt/v5 v5.3.0
 ## explicit; go 1.21
 github.com/golang-jwt/jwt/v5
-# github.com/google/uuid v1.6.0
-## explicit
-# github.com/gorilla/mux v1.8.1
-## explicit; go 1.20
 # github.com/kylelemons/godebug v1.1.0
 ## explicit; go 1.11
 github.com/kylelemons/godebug/diff