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/grpcutil/serve_test.go

@@ -0,0 +1,100 @@
+package grpcutil_test
+
+import (
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/health"
+	"google.golang.org/grpc/health/grpc_health_v1"
+
+	"github.com/pomerium/pomerium/pkg/grpcutil"
+)
+
+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)
+
+		srv := grpc.NewServer()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		now := time.Now()
+		err = grpcutil.ServeWithGracefulStop(ctx, srv, 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)
+
+		srv := grpc.NewServer()
+		hsrv := health.NewServer()
+		grpc_health_v1.RegisterHealthServer(srv, hsrv)
+		hsrv.SetServingStatus("test", grpc_health_v1.HealthCheckResponse_SERVING)
+
+		now := time.Now()
+		ctx, cancel := context.WithCancel(context.Background())
+		eg, ctx := errgroup.WithContext(ctx)
+		eg.Go(func() error {
+			return grpcutil.ServeWithGracefulStop(ctx, srv, li, time.Millisecond*100)
+		})
+		eg.Go(func() error {
+			var cc *grpc.ClientConn
+			for {
+				var err error
+				cc, err = grpc.Dial(li.Addr().String(),
+					grpc.WithTransportCredentials(insecure.NewCredentials()))
+				if err != nil {
+					continue
+				}
+
+				break
+			}
+
+			c := grpc_health_v1.NewHealthClient(cc)
+
+			// wait till the server is ready
+			for {
+				_, err := c.Check(ctx, &grpc_health_v1.HealthCheckRequest{
+					Service: "test",
+				})
+				if err != nil {
+					return err
+				}
+
+				break
+			}
+
+			// start streaming to hold open the server during graceful stop
+			_, err = c.Watch(context.Background(), &grpc_health_v1.HealthCheckRequest{
+				Service: "test",
+			})
+			if err != nil {
+				return err
+			}
+
+			cancel()
+
+			return nil
+		})
+		eg.Wait()
+		elapsed := time.Since(now)
+		assert.Greater(t, elapsed, time.Millisecond*100, "should complete after 100ms")
+	})
+}