ING-1399: Add long running request case for graceful shutdown

Author: Westwooo

gateway/gateway.go

@@ -409,6 +409,8 @@ func (g *Gateway) Run(ctx context.Context) error {
 			Password:        config.Password,
 		})
 
+		ctx, cancel := context.WithCancel(ctx)
+
 		config.Logger.Info("initializing protostellar system")
 		gatewaySys, err := system.NewSystem(&system.SystemOptions{
 			Logger:      config.Logger.Named("gateway-system"),
@@ -432,6 +434,7 @@ func (g *Gateway) Run(ctx context.Context) error {
 			},
 			AlphaEndpoints: config.AlphaEndpoints,
 			Debug:          config.Debug,
+			Cancel:         cancel,
 		})
 		if err != nil {
 			config.Logger.Error("error creating legacy proxy")

gateway/system/system.go

@@ -59,13 +59,17 @@ type SystemOptions struct {
 	DapiTlsConfig  *tls.Config
 	AlphaEndpoints bool
 	Debug          bool
+
+	Cancel context.CancelFunc
 }
 
 type System struct {
 	logger *zap.Logger
 
 	dataServer *grpc.Server
 	dapiServer *http.Server
+
+	cancelFunc context.CancelFunc
 }
 
 func NewSystem(opts *SystemOptions) (*System, error) {
@@ -242,6 +246,7 @@ func NewSystem(opts *SystemOptions) (*System, error) {
 		logger:     opts.Logger,
 		dataServer: dataSrv,
 		dapiServer: dapiSrv,
+		cancelFunc: opts.Cancel,
 	}
 
 	return s, nil
@@ -299,9 +304,12 @@ func (s *System) Shutdown() {
 			defer wg.Done()
 			s.dapiServer.SetKeepAlivesEnabled(false)
 			time.Sleep(time.Second * 5)
-			_ = s.dapiServer.Shutdown(context.Background())
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+			defer cancel()
+			_ = s.dapiServer.Shutdown(ctx)
 		}()
 	}
 
 	wg.Wait()
+	s.cancelFunc()
 }

gateway/test/dapi_graceful_shutdown_test.go

@@ -1,6 +1,7 @@
 package test
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"errors"
@@ -44,7 +45,7 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 		GrpcCertificate: *gwCert,
 		DapiCertificate: *gwCert,
 		NumInstances:    1,
-
+		ProxyServices:   []string{"query"},
 		StartupCallback: func(m *gateway.StartupInfo) {
 			gwStartInfoCh <- m
 		},
@@ -94,6 +95,8 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 	var eofs int
 	var unexpectedErr error
 	wg.Add(1)
+
+	// Continually run basic requests against the gateway and continue during shutdown.
 	go func() {
 		defer wg.Done()
 
@@ -104,7 +107,7 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 				case errors.Is(err, io.EOF):
 					// There is a small window between the flushing the request bytes to the socket and the http handler receiving
 					// the bytes where the server sees the connection as idle and will close it. We record these errors and
-					// include them when checking all requests recieved responses.
+					// include them when checking all requests received responses.
 					eofs++
 				case errors.Is(err, syscall.ECONNREFUSED):
 					// This is what we expect to see once the listeners have closed
@@ -113,7 +116,7 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 					// wroteRequest hook has not triggered. Such racy errors are unavoidable when running directly against an
 					// http server.
 				default:
-					// Any errors not mentinoned above are not expected and should cause a failure to be investigated.
+					// Any errors not mentioned above are not expected and should cause a failure to be investigated.
 					unexpectedErr = err
 					return
 				}
@@ -126,9 +129,37 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 		}
 	}()
 
+	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"}`)
+		reqBody := bytes.NewReader(statement)
+
+		req, err := http.NewRequestWithContext(context.Background(), "POST", fmt.Sprintf("https://%s/_p/query/query/service", dapiAddr), reqBody)
+		assert.NoError(s.T(), err, "failed to create request with client trace")
+
+		req.SetBasicAuth(testConfig.CbUser, testConfig.CbPass)
+		req.Header.Set("Content-Type", "application/json")
+
+		slowReqSent <- struct{}{}
+		_, err = dapiCli.Do(req)
+
+		// After the configured period shutdown will reuturn and we cancel the context. This triggers the go routine in Serve()
+		// that stops the http server forcibly, causing EOFs on existing connections doing work.
+		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
+
 	gw.Shutdown()
 
 	wg.Wait()