ING-1429: Add shutdown timeout to grpc server

Author: Westwooo

gateway/system/system.go

@@ -46,6 +46,7 @@ import (
 )
 
 const maxMsgSize = 25 * 1024 * 1024 // 25MiB
+const shutdownTimeout = 30 * time.Second
 
 type SystemOptions struct {
 	Logger *zap.Logger
@@ -301,7 +302,18 @@ func (s *System) Shutdown() {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			s.dataServer.GracefulStop()
+
+			// GracefulStop has no timeout mechanism so we need to take this approach
+			done := make(chan struct{})
+			go func() {
+				s.dataServer.GracefulStop()
+				close(done)
+			}()
+
+			select {
+			case <-done:
+			case <-time.After(shutdownTimeout):
+			}
 		}()
 	}
 

…eway/test/dapi_graceful_shutdown_test.go gateway/test/graceful_shutdown_test.gogateway/test/dapi_graceful_shutdown_test.go renamed to gateway/test/graceful_shutdown_test.go gateway/test/dapi_graceful_shutdown_test.go renamed to gateway/test/graceful_shutdown_test.go

@@ -14,12 +14,17 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/couchbase/goprotostellar/genproto/query_v1"
 	"github.com/couchbase/stellar-gateway/gateway"
 	"github.com/couchbase/stellar-gateway/testutils"
 	"github.com/couchbase/stellar-gateway/utils/selfsignedcert"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/status"
 )
 
 func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
@@ -56,8 +61,12 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 		s.T().Fatalf("failed to initialize graceful-shutdown-gateway: %s", err)
 	}
 
+	var wg sync.WaitGroup
+
 	gwClosedCh := make(chan struct{})
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		err := gw.Run(context.Background())
 		if err != nil {
 			s.T().Errorf("graceful-shutdown-gateway run failed: %s", err)
@@ -69,6 +78,13 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 
 	startInfo := <-gwStartInfoCh
 
+	connAddr := fmt.Sprintf("%s:%d", "127.0.0.1", startInfo.ServicePorts.PS)
+	conn, err := grpc.NewClient(connAddr,
+		grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{InsecureSkipVerify: true})),
+		grpc.WithUserAgent("test-client"),
+	)
+	assert.NoError(s.T(), err, "failed to create grpc client connection to gateway")
+
 	dapiAddr := fmt.Sprintf("%s:%d", "127.0.0.1", startInfo.ServicePorts.DAPI)
 	dapiCli := &http.Client{
 		Transport: &http.Transport{
@@ -92,7 +108,6 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 	assert.NoError(s.T(), err)
 
 	respCloseChan := make(chan (bool), 10000)
-	var wg sync.WaitGroup
 
 	var eofs int
 	var unexpectedErr error
@@ -133,13 +148,36 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 
 	wg.Add(1)
 
+	slowQuery := "SELECT SUM(t + t2) AS total FROM ARRAY_RANGE(0,8000) AS t, ARRAY_RANGE(0,4000) AS t2"
+	slowGrpcReqSent := make(chan any)
+	go func() {
+		defer wg.Done()
+		queryClient := query_v1.NewQueryServiceClient(conn)
+
+		slowGrpcReqSent <- struct{}{}
+		client, err := queryClient.Query(context.Background(), &query_v1.QueryRequest{
+			BucketName: &s.bucketName,
+			Statement:  slowQuery,
+		}, grpc.PerRPCCredentials(s.basicRpcCreds))
+		assert.NoError(s.T(), err, "failed to send grpc query request")
+
+		_, err = client.Recv()
+		assert.Error(s.T(), err, "expected error receiving from grpc stream after shutdown")
+
+		st, ok := status.FromError(err)
+		assert.True(s.T(), ok)
+		assert.Equal(s.T(), codes.Unavailable, st.Code())
+		assert.Contains(s.T(), st.Message(), "EOF")
+	}()
+
+	wg.Add(1)
 	slowReqSent := make(chan any)
 	// Send a long running request (this one takes approx 2 mins) and ensure that the gateway forcibly stops this request
 	// during shutdown.
 	go func() {
 		defer wg.Done()
 
-		statement := []byte(`{"statement": "SELECT SUM(t + t2) AS total FROM ARRAY_RANGE(0,8000) AS t, ARRAY_RANGE(0,4000) AS t2"}`)
+		statement := []byte(fmt.Sprintf(`{"statement": "%s"}`, slowQuery))
 		reqBody := bytes.NewReader(statement)
 
 		req, err := http.NewRequestWithContext(context.Background(), "POST", fmt.Sprintf("https://%s/_p/query/query/service", dapiAddr), reqBody)
@@ -156,11 +194,12 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 		assert.ErrorIs(s.T(), err, io.EOF)
 	}()
 
-	// Allow some requests to run against the gateway before shutting down
-	time.Sleep(time.Second)
-
 	// Wait for the slow running request to be sent before starting shutdown
 	<-slowReqSent
+	<-slowGrpcReqSent
+
+	// Allow the server to start processing the requests
+	time.Sleep(time.Second * 5)
 
 	gw.Shutdown()