core/http: add graceful stop serve

internal/controlplane/server.go

@@ -33,6 +33,7 @@ import (
 	"github.com/pomerium/pomerium/pkg/envoy/files"
 	pom_grpc "github.com/pomerium/pomerium/pkg/grpc"
 	"github.com/pomerium/pomerium/pkg/grpcutil"
+	"github.com/pomerium/pomerium/pkg/httputil"
 )
 
 // A Service can be mounted on the control plane.
@@ -195,29 +196,13 @@ func (srv *Server) Run(ctx context.Context) error {
 		{"metrics", srv.MetricsListener, srv.MetricsRouter},
 	} {
 		entry := entry
-		hsrv := (&http.Server{
-			BaseContext: func(li net.Listener) context.Context {
-				return ctx
-			},
-			Handler: entry.Handler,
-		})
 
 		// start the HTTP server
 		eg.Go(func() error {
 			log.Info(ctx).
 				Str("addr", entry.Listener.Addr().String()).
 				Msgf("starting control-plane %s server", entry.Name)
-			return hsrv.Serve(entry.Listener)
-		})
-
-		// gracefully stop the HTTP server on context cancellation
-		eg.Go(func() error {
-			<-ctx.Done()
-
-			ctx, cleanup := context.WithTimeout(ctx, time.Second*5)
-			defer cleanup()
-
-			return hsrv.Shutdown(ctx)
+			return httputil.ServeWithGracefulStop(ctx, entry.Handler, entry.Listener, time.Second*5)
 		})
 	}
 

pkg/contextutil/contextutil.go

@@ -64,3 +64,25 @@ func (mc *mergedCtx) Value(key interface{}) interface{} {
 	}
 	return mc.doneCtx.Value(key)
 }
+
+type onlyValues struct {
+	context.Context
+}
+
+// OnlyValues returns a derived context that removes deadlines and cancellation,
+// but keeps values.
+func OnlyValues(ctx context.Context) context.Context {
+	return onlyValues{ctx}
+}
+
+func (o onlyValues) Deadline() (time.Time, bool) {
+	return time.Time{}, false
+}
+
+func (o onlyValues) Done() <-chan struct{} {
+	return nil
+}
+
+func (o onlyValues) Err() error {
+	return nil
+}

pkg/httputil/serve.go

@@ -0,0 +1,48 @@
+// Package httputil contains additional functionality for working with http.
+package httputil
+
+import (
+	"context"
+	"errors"
+	"net"
+	"net/http"
+	"time"
+
+	"github.com/pomerium/pomerium/pkg/contextutil"
+)
+
+// ServeWithGracefulStop serves the HTTP listener until ctx.Done(), and then gracefully stops and waits for gracefulTimeout
+// before definitively stopping.
+func ServeWithGracefulStop(ctx context.Context, handler http.Handler, li net.Listener, gracefulTimeout time.Duration) error {
+	// create a context that will be used for the http requests
+	// it will only be cancelled when baseCancel is called but will
+	// preserve the values from ctx
+	baseCtx, baseCancel := context.WithCancelCause(contextutil.OnlyValues(ctx))
+
+	srv := http.Server{
+		Handler: handler,
+		BaseContext: func(l net.Listener) context.Context {
+			return baseCtx
+		},
+	}
+
+	go func() {
+		<-ctx.Done()
+
+		// create a context that will cancel after the graceful timeout
+		timeoutCtx, clearTimeout := context.WithTimeout(context.Background(), gracefulTimeout)
+		defer clearTimeout()
+
+		// shut the http server down
+		_ = srv.Shutdown(timeoutCtx)
+
+		// cancel the base context used for http requests
+		baseCancel(ctx.Err())
+	}()
+
+	err := srv.Serve(li)
+	if errors.Is(err, http.ErrServerClosed) {
+		err = nil
+	}
+	return err
+}

pkg/httputil/serve_test.go

@@ -0,0 +1,98 @@
+package httputil_test
+
+import (
+	"context"
+	"io"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/pomerium/pomerium/pkg/httputil"
+)
+
+func TestServeWithGracefulStop(t *testing.T) {
+	t.Parallel()
+
+	t.Run("immediate", func(t *testing.T) {
+		t.Parallel()
+
+		li, err := net.Listen("tcp4", "127.0.0.1:0")
+		require.NoError(t, err)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		})
+
+		now := time.Now()
+		err = httputil.ServeWithGracefulStop(ctx, h, li, time.Millisecond*100)
+		elapsed := time.Since(now)
+		assert.Nil(t, err)
+		assert.Less(t, elapsed, time.Millisecond*100, "should complete immediately")
+	})
+	t.Run("graceful", func(t *testing.T) {
+		t.Parallel()
+
+		li, err := net.Listen("tcp4", "127.0.0.1:0")
+		require.NoError(t, err)
+
+		h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			switch r.URL.Path {
+			case "/":
+				w.WriteHeader(http.StatusNoContent)
+			case "/wait":
+				w.WriteHeader(http.StatusOK)
+				w.Write([]byte("\n"))
+				w.(http.Flusher).Flush()
+				select {
+				case <-r.Context().Done():
+				case <-make(chan struct{}):
+				}
+			default:
+				http.NotFound(w, r)
+			}
+		})
+
+		now := time.Now()
+		ctx, cancel := context.WithCancel(context.Background())
+		eg, ctx := errgroup.WithContext(ctx)
+		eg.Go(func() error {
+			return httputil.ServeWithGracefulStop(ctx, h, li, time.Millisecond*100)
+		})
+		eg.Go(func() error {
+			// poll until the server is ready
+			for {
+				res, err := http.Get("http://" + li.Addr().String() + "/")
+				if err != nil {
+					continue
+				}
+				res.Body.Close()
+
+				break
+			}
+
+			// issue a stream request that will last indefinitely
+			res, err := http.Get("http://" + li.Addr().String() + "/wait")
+			if err != nil {
+				return err
+			}
+
+			cancel()
+
+			// wait until the request completes (should stop after the graceful timeout)
+			io.ReadAll(res.Body)
+			res.Body.Close()
+
+			return nil
+		})
+		eg.Wait()
+		elapsed := time.Since(now)
+		assert.Greater(t, elapsed, time.Millisecond*100, "should complete after 100ms")
+	})
+}