Repository: deislabs/ratify (go)

Issue:
Ratify should gracefully shutdown during a rolling upgrade
   Discussed the runtime update during the community meeting. Currently the standards is a rolling upgrade, where new instances of ratify will be swapped in.. note there is a action item to confirm if the old pods do a "graceful exist" @akashsinghal

_Originally posted by @susanshi in https://github.com/deislabs/ratify/issues/932#issuecomment-1643055960_
            

Investigation notes from Akash,

     Following up on our conversation earlier this week regarding request dropping concerns during Ratify upgrade scenarios. I looked into our current implementation and it seems that we are not handling this currently. The SIGTERM signal is received by the ratify process when K8s does a rolling upgrade and terminates the pod. Currently SIGTERM just terminates the process without waiting for the server to properly Shutdown. I believe we would need to update the server to explicitly listen and block for the SIGTERM and then execute a proper server shutdown command with a time buffer (time to allow for requests in flight to complete). There's guidance online I've found that we can follow for Ratify to gracefully shutdown

Pull request description:
feat: add graceful shutdown for http server
# Description

## What this PR does / why we need it:
- Adds support for `Shutdown` command to invoked on SIGTERM signal or interrupt OS command.
- Adds channel to wait on shutdown process to complete (6 second context timeout)

## Which issue(s) this PR fixes *(optional, using `fixes #<issue number>(, fixes #<issue_number>, ...)` format, will close the issue(s) when the PR gets merged)*:
Fixes #942 

## Type of change

Please delete options that are not relevant.

- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] Helm Chart Change (any edit/addition/update that is necessary for changes merged to the `main` branch)
- [ ] This change requires a documentation update

# How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Please also list any relevant details for your test configuration

- [ ] Test A
- [ ] Test B

# Checklist:

- [ ] Does the affected code have corresponding tests?
- [ ] Are the changes documented, not just with inline documentation, but also with conceptual documentation such as an overview of a new feature, or task-based documentation like a tutorial? Consider if this change should be announced on your project blog.
- [ ] Does this introduce breaking changes that would require an announcement or bumping the major version?
- [ ] Do all new files have appropriate license header?

# Post Merge Requirements
- [ ] MAINTAINERS: manually trigger the "Publish Package" workflow after merging any PR that indicates `Helm Chart Change`


Patch:
diff --git a/httpserver/server.go b/httpserver/server.go
index a63ef83..8f76364 100644
--- a/httpserver/server.go
+++ b/httpserver/server.go
@@ -23,8 +23,10 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/deislabs/ratify/config"
@@ -152,14 +154,9 @@ func (server *Server) Run(certRotatorReady chan struct{}) error {
 			MinVersion:         tls.VersionTLS13,
 		}
 
-		if err := svr.ServeTLS(lsnr, certFile, keyFile); err != nil {
-			logrus.Errorf("failed to start server: %v", err)
-			return err
-		}
-		return nil
+		return startServerWithGracefulShutdown(true, svr, lsnr, certFile, keyFile)
 	}
-	logrus.Info("starting server without TLS")
-	return svr.Serve(lsnr)
+	return startServerWithGracefulShutdown(false, svr, lsnr, "", "")
 }
 
 func (server *Server) register(method, path string, handler ContextHandler) {
@@ -190,3 +187,36 @@ type ServerAddrNotFoundError struct{}
 func (err ServerAddrNotFoundError) Error() string {
 	return "The http server address configuration is not set. Skipping server creation"
 }
+
+// startServerWithGracefulShutdown starts the server and waits for SIGINT or SIGTERM to shutdown the server gracefully
+func startServerWithGracefulShutdown(isTLSEnabled bool, svr *http.Server, lsnr net.Listener, certFile string, keyFile string) error {
+	connectionsClosed := make(chan struct{})
+	// wait for SIGINT or SIGTERM to shutdown the server gracefully
+	go func() {
+		sigint := make(chan os.Signal, 1)
+		signal.Notify(sigint, os.Interrupt, syscall.SIGTERM)
+		<-sigint
+		logrus.Info("shutting down ratify server...")
+		ctx, cancel := context.WithTimeout(context.Background(), 6*time.Second) // wait 6 seconds for existing connections to finish
+		defer cancel()
+		if err := svr.Shutdown(ctx); err != nil {
+			logrus.Errorf("failed to shutdown ratify server: %v", err)
+		}
+		close(connectionsClosed)
+	}()
+	if isTLSEnabled {
+		if err := svr.ServeTLS(lsnr, certFile, keyFile); err != nil && err != http.ErrServerClosed { // http.ErrServerClosed is immediately returned when server.Shutdown() is called
+			logrus.Errorf("failed to start server: %v", err)
+			return err
+		}
+	} else {
+		if err := svr.Serve(lsnr); err != nil && err != http.ErrServerClosed { // http.ErrServerClosed is immediately returned when server.Shutdown() is called
+			logrus.Errorf("failed to start server: %v", err)
+			return err
+		}
+	}
+	<-connectionsClosed // wait for server to shutdown
+	logrus.Info("ratify server shutdown complete")
+	time.Sleep(5 * time.Second) // wait for controllers/manager to shutdown
+	return nil
+}


Test patch:
diff --git a/httpserver/server_test.go b/httpserver/server_test.go
index 65f4f97..430baa9 100644
--- a/httpserver/server_test.go
+++ b/httpserver/server_test.go
@@ -20,8 +20,11 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"syscall"
 	"testing"
 	"time"
 
@@ -444,3 +447,52 @@ func TestServer_Verify_ParseReference_Failure(t *testing.T) {
 		}
 	})
 }
+
+// TestServe_serverGracefulShutdown tests the case where the server is shutdown gracefully
+func TestServer_serverGracefulShutdown(t *testing.T) {
+	// create a server that sleeps for 5 seconds before responding
+	ts := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(5 * time.Second)
+		fmt.Fprintln(w, "request succeeded")
+	}))
+
+	// start the server
+	go func() {
+		_ = startServerWithGracefulShutdown(false, ts.Config, ts.Listener, "", "")
+	}()
+
+	// wait a second for server to come online
+	time.Sleep(1 * time.Second)
+
+	// create a client that makes a request to the server and validates response
+	client := &http.Client{Transport: &http.Transport{}}
+	go func() {
+		url := "http://" + ts.Listener.Addr().String()
+		res, err := client.Get(url)
+		if err != nil {
+			fmt.Printf("Expected no error but got %v", err)
+		}
+
+		body, err := io.ReadAll(res.Body)
+		res.Body.Close()
+		if err != nil {
+			fmt.Printf("Expected no error but got %v", err)
+		}
+		if string(body) != "request succeeded\n" {
+			fmt.Printf("Expected response body to be 'request succeeded' but got %s", string(body))
+		}
+	}()
+	// wait a second for client to make request and reach server
+	time.Sleep(1 * time.Second)
+	// send SIGTERM to server
+	proc, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		t.Fatalf("Expected no error but got %v", err)
+	}
+	err = proc.Signal(syscall.SIGTERM)
+	if err != nil {
+		t.Fatalf("Expected no error but got %v", err)
+	}
+	// wait some time to see shutdown logs
+	time.Sleep(5 * time.Second)
+}