From e88354f7fdd7d453552cbbf6e1e7c8f512b85a42 Mon Sep 17 00:00:00 2001
From: "Han Verstraete (OpenFaaS Ltd)" <han@openfaas.com>
Date: Fri, 8 Dec 2023 23:22:28 +0100
Subject: [PATCH] Support shutdown of bootstrap server

Signed-off-by: Han Verstraete (OpenFaaS Ltd) <han@openfaas.com>
---
 serve.go | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/serve.go b/serve.go
index 67ce6db..535986e 100644
--- a/serve.go
+++ b/serve.go
@@ -4,6 +4,8 @@
 package bootstrap
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"log"
 	"net/http"
@@ -31,7 +33,7 @@ func Router() *mux.Router {
 }
 
 // Serve load your handlers into the correct OpenFaaS route spec. This function is blocking.
-func Serve(handlers *types.FaaSHandlers, config *types.FaaSConfig) {
+func Serve(ctx context.Context, handlers *types.FaaSHandlers, config *types.FaaSConfig) {
 
 	if config.EnableBasicAuth {
 		reader := auth.ReadBasicAuthFromDisk{
@@ -118,5 +120,16 @@ func Serve(handlers *types.FaaSHandlers, config *types.FaaSConfig) {
 		Handler:        r,
 	}
 
-	log.Fatal(s.ListenAndServe())
+	// Start server in a goroutine
+	go func() {
+		if err := s.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
+			log.Fatal(err)
+		}
+	}()
+
+	// Shutdown server when context is done.
+	<-ctx.Done()
+	if err := s.Shutdown(context.Background()); err != nil {
+		log.Printf("Failed to shut down provider gracefully: %s", err)
+	}
 }