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

Author: Westwooo

gateway/system/system.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net"
 	"net/http"
 	"sync"
 	"time"
@@ -66,6 +67,8 @@ type System struct {
 
 	dataServer *grpc.Server
 	dapiServer *http.Server
+
+	cancelFunc context.CancelFunc
 }
 
 func NewSystem(opts *SystemOptions) (*System, error) {
@@ -231,17 +234,21 @@ func NewSystem(opts *SystemOptions) (*System, error) {
 	}
 	httpHandler = c.Handler(httpHandler)
 
+	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
 	dapiSrv := &http.Server{
 		ReadHeaderTimeout: time.Second * 5,
 		IdleTimeout:       time.Second * 60,
 		Handler:           httpHandler,
 		TLSConfig:         opts.DapiTlsConfig,
+		BaseContext:       func(net.Listener) context.Context { return shutdownCtx },
 	}
 
 	s := &System{
 		logger:     opts.Logger,
 		dataServer: dataSrv,
 		dapiServer: dapiSrv,
+
+		cancelFunc: shutdownCancel,
 	}
 
 	return s, nil
@@ -299,7 +306,10 @@ 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)
+			s.cancelFunc()
 		}()
 	}
 

gateway/test/dapi_graceful_shutdown_test.go

@@ -1,6 +1,7 @@
 package test
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"errors"
@@ -44,6 +45,8 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 		GrpcCertificate: *gwCert,
 		DapiCertificate: *gwCert,
 		NumInstances:    1,
+		ProxyServices:   []string{"query"},
+		Debug:           true,
 
 		StartupCallback: func(m *gateway.StartupInfo) {
 			gwStartInfoCh <- m
@@ -94,6 +97,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()
 
@@ -126,9 +131,41 @@ 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{}{}
+		resp, err := dapiCli.Do(req)
+		assert.NoError(s.T(), err)
+		assert.Equal(s.T(), http.StatusBadGateway, resp.StatusCode, "expected to recieve a 502 response from the gateway")
+
+		defer func() { _ = resp.Body.Close() }()
+		body, err := io.ReadAll(resp.Body)
+		assert.NoError(s.T(), err)
+
+		assert.Equal(s.T(), string(body), "failed to receive response: context canceled")
+	}()
+
 	// 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()