Feature/graceful shutdown

README.md

@@ -11,27 +11,32 @@ only HTTP triggers are supported, but additional triggers may be supported in th
 
 ### How it works
 Lambda Function in AWS:
-1. The main function calls the lambda.Start function, passing in the Handler function.
-2. AWS Lambda invokes the Handler function when an HTTP request is received.
+1. The main function calls the `lambda.Start` function, passing in the `Handler` function.
+2. AWS Lambda invokes the `Handler` function when an HTTP request is received.
+3. Alternatively, the main function can call `lambda.StartWithOptions` to enable graceful shutdown with an optional shutdown function(s).
+
 ![diagram1](docs/images/diagram1.png)
 
 Testing a Lambda Function outside of AWS:
-1. The main function calls the httpadapter.Start function, passing in the Handler function and the port number to listen on.
-2. The httpadapter listens for incoming HTTP requests on the specified port and invokes the Handler function when a request is received.
+1. The main function calls the `httpadapter.Start` function, passing in the `Handler` function and the port number to listen on.
+2. The httpadapter listens for incoming HTTP requests on the specified port and invokes the `Handler` function when a request is received.
+3. Alternatively, the main function can call `httpadapter.StartWithOptions` to enable graceful shutdown with an optional shutdown function(s).
+
 ![diagram2](docs/images/diagram2.png)
 
 
 ## Features
-* Supports APIGatewayV2HTTP, APIGatewayProxy, and ALBTargetGroup events.
-* Supports handler functions with and without context.Context parameters.
+* Supports `APIGatewayV2HTTP`, `APIGatewayProxy`, and `ALBTargetGroup` events.
+* Supports handler functions with and without `context.Context` parameters.
 * Supports both values and pointers for handler function request events.
 * Supports both values and pointers for handler function response events.
 
 ## Installation
 ```go get -u github.com/Evernorth/aws-lambda-go-adapter```
 
 ## Usage
-```
+### Start
+```go
 package main
 
 import (
@@ -71,6 +76,59 @@ func main() {
 	}
 }
 ```
+### StartWithOptions WithEnableSIGTERM
+This is useful for application shutdown logic needed for a graceful shutdown. The httpadapter will mimic the AWS Lambda behavior of allowing ~500ms for shutdown before it is terminated. 
+```go
+package main
+
+import (
+	"context"
+	"github.com/Evernorth/aws-lambda-go-adapter/httpadapter"
+	"github.com/Evernorth/aws-lambda-go-adapter/pkg/util"
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/aws/aws-lambda-go/lambda"
+	"net/http"
+)
+
+// Handler is the Lambda handler function.  It returns a 200 status code with a "Hello, World!" message for GET requests,
+// and a 405 status code for all other requests.
+func Handler(ctx context.Context, request events.APIGatewayV2HTTPRequest) (events.APIGatewayV2HTTPResponse, error) {
+
+	if request.RequestContext.HTTP.Method == http.MethodGet {
+		return events.APIGatewayV2HTTPResponse{
+			StatusCode: http.StatusOK,
+			Body:       "Hello, World!",
+			Headers: map[string]string{
+				"Content-Type": "text/html",
+			},
+		}, nil
+	}
+
+	return events.APIGatewayV2HTTPResponse{
+		StatusCode: http.StatusMethodNotAllowed,
+		Body:       "Method not allowed",
+	}, nil
+}
+
+func shutdown() {
+    // Perform any necessary cleanup here, such as closing database connections or releasing resources.
+    // This function is called during graceful shutdown.
+}
+
+func main() {
+	if util.IsLambdaRuntime() { 
+	    // Start AWS Lambda function with SIGTERM support for graceful shutdown. 
+	    // App shutdown logic has ~500ms to complete before it's terminated with a Lambda SIGKILL 
+	    lambda.StartWithOptions(Handler, 
+	        lambda.WithEnableSIGTERM(shutdown))
+	} else {
+	    // Start HTTP server with SIGTERM support for graceful shutdown.
+	    // App shutdown logic has ~500ms to complete before it's terminated with a "SIGKILL" panic
+	    httpadapter.StartWithOptions(8080, Handler,
+	        httpadapter.WithEnableSIGTERM(shutdown))
+	}
+}
+```
 
 ## Dependencies
 See the [go.mod](go.mod) file.

httpadapter/httpadapter.go

@@ -37,24 +37,66 @@ func SetLogger(l *slog.Logger) {
 	logger = l
 }
 
+type adapterOptions struct {
+	//baseContext                      context.Context
+	enableSIGTERM bool
+	sigtermFuncs  []func()
+}
+
+type Option func(*adapterOptions)
+
 // Start starts the HTTP server on the specified port and listens for incoming requests.  When a request is received,
 // it is converted to the appropriate Lambda event type and passed to the handler function.  The response from the
 // handler function is then converted to an HTTP response and returned to the client.
+// See StartWithOptions for other start options.
 func Start(port int, handler interface{}) {
+	StartWithOptions(port, handler)
+}
+
+// StartWithOptions starts the HTTP server on the specified port and listens for incoming requests.  When a request is received,
+// it is converted to the appropriate Lambda event type and passed to the handler function.  The response from the
+// handler function is then converted to an HTTP response and returned to the client.
+func StartWithOptions(port int, handler interface{}, options ...Option) {
+
 	reflectHandler(handler)
 
+	opts := &adapterOptions{}
+	for _, option := range options {
+		option(opts)
+	}
+
 	switch delegateHandlerType {
 	case apigwV2HandlerType:
-		listenAndServe(port, handleRequestForApigwV2)
+		listenAndServe(port, handleRequestForApigwV2, opts)
 	case albHandlerType:
-		listenAndServe(port, handleRequestForAlb)
+		listenAndServe(port, handleRequestForAlb, opts)
 	case apigwHandlerType:
-		listenAndServe(port, handleRequestForApigw)
+		listenAndServe(port, handleRequestForApigw, opts)
 	default:
 		panic("unsupported handler type")
 	}
 }
 
+// WithEnableSIGTERM enables SIGTERM behavior with the HTTP server for graceful shutdown.
+// Optionally, an array of shutdown functions to run on SIGTERM may be provided.
+// After HTTP server graceful shutdown, the provided functions are invoked within a ~500ms timeout. If the shutdown functions
+// do not complete within the timeout limit, a "SIGKILL" panic will occur. This enables testing in a local environment,
+// mimicking AWS Lambda's SIGTERM and SIGKILL behavior.
+//
+// Usage:
+//
+//	httpadapter.StartWithOptions(8080, Handler,
+//		lambda.WithEnableSIGTERM(func() {
+//			log.Print("Cleaning up application components...")
+//		})
+//	)
+func WithEnableSIGTERM(sigtermFuncs ...func()) Option {
+	return Option(func(opts *adapterOptions) {
+		opts.sigtermFuncs = append(opts.sigtermFuncs, sigtermFuncs...)
+		opts.enableSIGTERM = true
+	})
+}
+
 // reflectHandler reflects the handler function to determine the input and output event types, and whether a context is
 // required.  It panics if the function signature is not supported.
 func reflectHandler(handler interface{}) {

httpadapter/listener.go

@@ -1,21 +1,86 @@
 package httpadapter
 
 import (
+	"context"
 	"log/slog"
 	"net/http"
+	"os/signal"
 	"strconv"
+	"syscall"
+	"time"
 )
 
-func listenAndServe(port int, handler func(httpResponseWriter http.ResponseWriter, httpRequest *http.Request)) {
+func listenAndServe(port int, handler func(httpResponseWriter http.ResponseWriter, httpRequest *http.Request), opts *adapterOptions) {
 	logger.Info("Starting http listener.",
 		slog.Int("port", port))
 
 	http.HandleFunc("/", handler)
 	portStr := ":" + strconv.Itoa(port)
-	err := http.ListenAndServe(portStr, nil)
-	if err != nil {
-		logger.Error("Could not start http listener.",
-			slog.Any("err", err))
+	server := &http.Server{
+		Addr:    portStr,
+		Handler: nil, // Use the typical defaultServerMux with the handler set by http.HandleFunc
+	}
+
+	if opts.enableSIGTERM {
+		sigtermListenAndServe(server, opts)
+	} else {
+		defaultListenAndServe(server)
+	}
+}
+
+func defaultListenAndServe(server *http.Server) {
+	// start the server and ignore the error if it is http.ErrServerClosed from Shutdown
+	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		logger.Error("Could not start http listener.", slog.Any("err", err))
 		panic(err)
 	}
 }
+
+func sigtermListenAndServe(server *http.Server, opts *adapterOptions) {
+	// Create a context that cancels on SIGTERM or SIGINT
+	sigCtx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
+	defer stop()
+
+	// Start the server in a goroutine
+	go func() {
+		defaultListenAndServe(server)
+	}()
+
+	// Wait for the termination signal
+	<-sigCtx.Done()
+	logger.Info("SIGTERM: Shutting server down gracefully...")
+
+	// Create a context with a timeout for a limited clean server shutdown (5s should be plenty)
+	// Shutdown usually completes much faster, of course.
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	if err := server.Shutdown(shutdownCtx); err != nil {
+		logger.Error("SIGTERM: Graceful server shutdown failed", slog.Any("err", err))
+	}
+	logger.Info("SIGTERM: Server shutdown gracefully")
+
+	// Perform cleanup with a 500ms limit like AWS Lambda SIGTERM behavior. This will help ensure
+	// that cleanup functions can be locally tested with the same behavior as in AWS Lambda.
+	if len(opts.sigtermFuncs) > 0 {
+		logger.Info("SIGTERM: Running shutdown functions (mimicking the 500ms AWS Lambda limit)")
+		timedCtx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+		defer cancel()
+		done := make(chan struct{})
+		go func() {
+			for _, f := range opts.sigtermFuncs {
+				f()
+			}
+			close(done)
+		}()
+		select {
+		case <-done:
+			logger.Info("SIGTERM: Shutdown functions completed")
+		case <-timedCtx.Done():
+			//Panic is more attention-grabbing than an error log.
+			panic("SIGKILL: Shutdown functions did not complete within the AWS Lambda limit")
+		}
+	}
+
+	logger.Info("SIGTERM: Completed, exiting...")
+}

httpadapter/listener_test.go

@@ -0,0 +1,84 @@
+package httpadapter
+
+import (
+	"net/http"
+	"os"
+	"testing"
+	"time"
+)
+
+func TestSigtermListenAndServe(t *testing.T) {
+
+	opts := &adapterOptions{
+		enableSIGTERM: true,
+	}
+	sigtermListenAndServeTest(t, opts, false)
+}
+
+func TestSigtermListenAndServe_WithFuncs(t *testing.T) {
+
+	opts := &adapterOptions{
+		enableSIGTERM: true,
+		sigtermFuncs: []func(){
+			func() {
+				t.Log("Running cleanup function 1")
+			},
+			func() {
+				t.Log("Running cleanup function 2")
+			},
+		},
+	}
+	sigtermListenAndServeTest(t, opts, false)
+}
+
+func TestSigtermListenAndServe_ExpectFuncTimeout(t *testing.T) {
+	opts := &adapterOptions{
+		enableSIGTERM: true,
+		sigtermFuncs: []func(){
+			func() {
+				t.Log("Long running cleanup function")
+				time.Sleep(502 * time.Millisecond) // Simulate a long-running function (>500ms)
+			},
+		},
+	}
+	sigtermListenAndServeTest(t, opts, true)
+}
+
+func sigtermListenAndServeTest(t *testing.T, opts *adapterOptions, expectPanic bool) {
+	// Create a dummy HTTP server that does nothing, just to test the SIGTERM handling
+	server := &http.Server{
+		Addr: ":0", // Use a port no one should care about
+	}
+
+	// Create a channel to signal when the server has started
+	done := make(chan struct{})
+	go func() {
+		// Use defer to ensure that we recover from a timeout panic to assert
+		// behavior of AWS Lambda where cleanup functions take too long
+		defer func() {
+			if r := recover(); r == nil && expectPanic {
+				t.Fatal("Expected panic due to cleanup function timeout, but no panic occurred")
+			}
+		}()
+
+		sigtermListenAndServe(server, opts)
+		close(done)
+	}()
+
+	// Wait for the server to start
+	time.Sleep(20 * time.Millisecond)
+
+	// Send SIGINT to self
+	p, _ := os.FindProcess(os.Getpid())
+	_ = p.Signal(os.Interrupt)
+
+	// Wait for the server to exit or timeout
+	select {
+	case <-done:
+		t.Log("server successfully exited on SIGINT")
+	case <-time.After(1 * time.Second):
+		if !expectPanic {
+			t.Fatal("server did not exit on SIGINT")
+		}
+	}
+}