core: fix graceful stop (#4866)

core: fix graceful stop (#4865)

* core/grpc: fix graceful stop

* core/http: add graceful stop serve

Co-authored-by: Caleb Doxsey <cdoxsey@pomerium.com>

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")
+	})
+}