Repository: aws/amazon-eks-pod-identity-webhook

Issue:
In-Flight Requests to Webhook Terminated Abruptly on Pod Shutdown
**What happened**:

We observed intermittent AWS IAM permission errors affecting a small number of pods (approximately once every few weeks). Investigation revealed that these pods were missing environment variables normally injected by the pod-identity-webhook (PIW). Kubernetes audit logs confirmed that the API server’s request to PIW resulted in no mutation:
```json
{"configuration":"pod-identity-webhook","webhook":"pod-identity-webhook.amazonaws.com","mutated":false}
``` 
The failure coincided with the PIW pod handling the request shutting down. Due to `failurePolicy: Ignore` in the mutating webhook configuration, the pod creation proceeded without mutation, resulting in the missing environment variables.

**What you expected to happen**:

PIW should perform a graceful shutdown to allow in-flight requests to complete. In this case, the webhook request should have been served before the PIW process terminated.

**How to reproduce it (as minimally and precisely as possible)**:

The root cause appears to be a race condition in the shutdown logic, as seen here:
https://github.com/aws/amazon-eks-pod-identity-webhook/blob/c8426f2d7700a5f86fe26df416134ba5dff31ac5/main.go#L306-L334

Key contributing factors:
1. PIW server is spun off in a goroutine.
2. Metrics server is running in the main goroutine and prevents main from exiting as long as it is running.
3. The code has logic for graceful shutdown via `ShutdownFromContext`. However, as soon as metrics server shuts down, `main` finishes and the whole program terminates without waiting for graceful shutdown of the PIW server.

**Anything else we need to know?**:

I’d be happy to contribute a fix to address this issue and improve the shutdown behavior.

**Environment**:
- AWS Region: us-east-1
- EKS Platform version: Kops managed clusters instead of EKS.
- Kubernetes version: v1.30.8
- Webhook Version: v0.6.1


Discussion:

Solution patch:
diff --git a/main.go b/main.go
index cabde09..3ba0755 100644
--- a/main.go
+++ b/main.go
@@ -16,9 +16,11 @@
 package main
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"crypto/x509/pkix"
+	"errors"
 	goflag "flag"
 	"fmt"
 	"math/rand"
@@ -26,7 +28,6 @@ import (
 	"os"
 	"strings"
 	"time"
-	"context"
 
 	"github.com/aws/amazon-eks-pod-identity-webhook/pkg"
 	"github.com/aws/amazon-eks-pod-identity-webhook/pkg/cache"
@@ -309,26 +310,45 @@ func main() {
 		Handler:   mux,
 		TLSConfig: tlsConfig,
 	}
-
-	handler.ShutdownFromContext(signalHandlerCtx, server, time.Duration(10)*time.Second)
+	serverDone, serverShutdownError := handler.ShutdownFromContext(signalHandlerCtx, server, time.Duration(10)*time.Second)
 
 	metricsServer := &http.Server{
 		Addr:    metricsAddr,
 		Handler: metricsMux,
 	}
+	metricsServerDone, metricsServerShutdownError := handler.ShutdownFromContext(signalHandlerCtx, metricsServer, time.Duration(10)*time.Second)
 
-	handler.ShutdownFromContext(signalHandlerCtx, metricsServer, time.Duration(10)*time.Second)
-
+	// Start webhook server
 	go func() {
 		klog.Infof("Listening on %s", addr)
-		if err := server.ListenAndServeTLS("", ""); err != http.ErrServerClosed {
+		if err := server.ListenAndServeTLS("", ""); !errors.Is(err, http.ErrServerClosed) {
+			klog.Fatalf("Error listening: %q", err)
+		}
+	}()
+
+	// Start metrics server
+	go func() {
+		klog.Infof("Listening on %s for metrics", metricsAddr)
+		if err := metricsServer.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
 			klog.Fatalf("Error listening: %q", err)
 		}
 	}()
 
-	klog.Infof("Listening on %s for metrics", metricsAddr)
-	if err := metricsServer.ListenAndServe(); err != http.ErrServerClosed {
-		klog.Fatalf("Error listening: %q", err)
+	klog.Info("Waiting for webhook and metrics servers to shutdown...")
+
+	// Wait for webhook server shutdown
+	<-serverDone
+	if err := <-serverShutdownError; err != nil {
+		klog.Errorf("Webhook server shutdown error: %v", err)
+	} else {
+		klog.Infof("Webhook server shutdown gracefully")
+	}
+
+	// Wait for metrics server shutdown
+	<-metricsServerDone
+	if err := <-metricsServerShutdownError; err != nil {
+		klog.Errorf("Metrics server shutdown error: %v", err)
+	} else {
+		klog.Infof("Metrics server shutdown gracefully")
 	}
-	klog.Info("Graceflully closed")
 }
diff --git a/pkg/handler/shutdown.go b/pkg/handler/shutdown.go
index e370582..ed1bd65 100644
--- a/pkg/handler/shutdown.go
+++ b/pkg/handler/shutdown.go
@@ -17,24 +17,55 @@ package handler
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"net/http"
 	"time"
 
 	"k8s.io/klog/v2"
 )
 
-func ShutdownFromContext(ctx context.Context, server *http.Server, timeout time.Duration) {
+// ShutdownFromContext gracefully shuts down an HTTP server when the provided context is cancelled.
+// It waits for ctx to be cancelled, then attempts a graceful shutdown with the given timeout.
+// If graceful shutdown fails, it will attempt to force close the server.
+//
+// Parameters:
+//   - ctx: Context to watch for cancellation signal
+//   - server: HTTP server to shutdown
+//   - timeout: Maximum time to wait for graceful shutdown
+//
+// Returns:
+//   - <-chan struct{}: Closed when shutdown process completes
+//   - <-chan error: Receives shutdown error if shutdown wasn't graceful (nil on success)
+func ShutdownFromContext(ctx context.Context, server *http.Server, timeout time.Duration) (<-chan struct{}, <-chan error) {
+	errCh := make(chan error, 1)
+	doneCh := make(chan struct{})
+
 	go func() {
+		defer close(doneCh)
+		defer close(errCh)
+
 		<-ctx.Done()
+		klog.Info("Context cancelled, beginning graceful shutdown")
 
 		shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
 		defer cancel()
 
 		if err := server.Shutdown(shutdownCtx); err != nil {
-			klog.Errorf("Error shutting server down: %v", err)
-			if err := server.Close(); err != nil {
-				klog.Fatalf("Error closing server: %v", err)
+			// Try to force close the server because the server didn't shutdown gracefully within the timeout
+			if closeErr := server.Close(); closeErr != nil {
+				errCh <- errors.Join(
+					fmt.Errorf("graceful server shutdown failed: %w", err),
+					fmt.Errorf("server force close failed: %w", closeErr),
+				)
+			} else {
+				errCh <- fmt.Errorf("graceful server shutdown failed, server has been force closed: %w", err)
 			}
+		} else {
+			// Graceful shutdown succeeded
+			errCh <- nil
 		}
 	}()
+
+	return doneCh, errCh
 }



Test patch:
diff --git a/pkg/handler/shutdown_test.go b/pkg/handler/shutdown_test.go
new file mode 100644
index 0000000..e6c690a
--- /dev/null
+++ b/pkg/handler/shutdown_test.go
@@ -0,0 +1,103 @@
+/*
+  Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  A copy of the License is located at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+  or in the "license" file accompanying this file. This file is distributed
+  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+  express or implied. See the License for the specific language governing
+  permissions and limitations under the License.
+*/
+
+package handler
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+)
+
+func TestShutdownFromContext_SuccessfulGracefulShutdown(t *testing.T) {
+	// Create a test server
+	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	// Create a context that we can cancel
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// Call ShutdownFromContext
+	doneCh, errCh := ShutdownFromContext(ctx, server.Config, 2*time.Second)
+
+	// Start the server
+	server.Start()
+	defer server.Close()
+
+	// Cancel the context to trigger shutdown
+	cancel()
+
+	// Wait for shutdown to complete
+	select {
+	case <-doneCh:
+		// Shutdown completed
+	case <-time.After(5 * time.Second):
+		t.Fatal("Shutdown did not complete within expected time")
+	}
+
+	// Check that no error occurred (graceful shutdown succeeded).
+	if err := <-errCh; err != nil {
+		t.Errorf("Expected no error for graceful shutdown, got: %v", err)
+	}
+}
+
+func TestShutdownFromContext_GracefulShutdownTimeout(t *testing.T) {
+	// Create a test server with a slow handler that will delay shutdown
+	testServer := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Simulate a slow handler that takes time to complete
+		time.Sleep(200 * time.Millisecond)
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	// Start the server
+	testServer.Start()
+	defer testServer.Close()
+
+	// Start a request to the slow handler to create an active connection during shutdown
+	go func() {
+		client := &http.Client{}
+		client.Get(testServer.URL)
+	}()
+
+	// Give the request time to start
+	time.Sleep(10 * time.Millisecond)
+
+	// Create a context that we can cancel
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// Call ShutdownFromContext with a very short timeout to force timeout
+	doneCh, errCh := ShutdownFromContext(ctx, testServer.Config, 1*time.Nanosecond)
+
+	// Cancel the context to trigger shutdown
+	cancel()
+
+	// Wait for shutdown to complete
+	select {
+	case <-doneCh:
+		// Shutdown completed
+	case <-time.After(1 * time.Second):
+		t.Fatal("Shutdown did not complete within expected time")
+	}
+
+	// Check that we got an error (timeout should always cause an error)
+	if err := <-errCh; err == nil {
+		t.Error("Expected error for graceful shutdown timeout, but got nil")
+	} else {
+		t.Logf("Received expected error with very short timeout: %v", err)
+	}
+}