ING-1399: Fix flakiness in graceful shutdown tests

gateway/gateway.go

@@ -19,6 +19,7 @@ import (
 	"github.com/couchbase/stellar-gateway/gateway/dapiimpl"
 	"github.com/couchbase/stellar-gateway/gateway/dapiimpl/proxy"
 	"github.com/couchbase/stellar-gateway/gateway/dataimpl"
+	"github.com/couchbase/stellar-gateway/gateway/hooks"
 	"github.com/couchbase/stellar-gateway/gateway/ratelimiting"
 	"github.com/couchbase/stellar-gateway/gateway/system"
 	"github.com/couchbase/stellar-gateway/pkg/metrics"
@@ -42,6 +43,7 @@ type StartupInfo struct {
 	ServerGroup   string
 	AdvertiseAddr string
 	ServicePorts  ServicePorts
+	HooksManager  *hooks.HooksManager
 }
 
 type Config struct {
@@ -489,6 +491,7 @@ func (g *Gateway) Run(ctx context.Context) error {
 					PS:   boundPsPort,
 					DAPI: boundDapiPort,
 				},
+				HooksManager: gatewaySys.HooksManager(),
 			})
 		}
 

gateway/hooks/hooksmanager.go

@@ -2,6 +2,7 @@ package hooks
 
 import (
 	"errors"
+	"net/http"
 	"sync"
 
 	"github.com/couchbase/goprotostellar/genproto/internal_hooks_v1"
@@ -67,3 +68,7 @@ func (m *HooksManager) Server() internal_hooks_v1.HooksServiceServer {
 func (m *HooksManager) UnaryInterceptor() grpc.UnaryServerInterceptor {
 	return makeGrpcUnaryInterceptor(m, m.logger)
 }
+
+func (m *HooksManager) HTTPMiddleware() func(http.Handler) http.Handler {
+	return makeHTTPMiddleware(m, m.logger)
+}

gateway/hooks/httpinterceptor.go

@@ -0,0 +1,44 @@
+package hooks
+
+import (
+	"net/http"
+
+	"go.uber.org/zap"
+)
+
+func makeHTTPMiddleware(manager *HooksManager, log *zap.Logger) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			hooksID := r.Header.Get("X-Hooks-ID")
+			if hooksID == "" {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			hooksContext := manager.GetHooksContext(hooksID)
+			if hooksContext == nil {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			barrierID := r.Header.Get("X-Barrier-ID")
+			if barrierID == "" {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			log.Info("http request waiting on barrier",
+				zap.String("hooks-id", hooksID),
+				zap.String("barrier-id", barrierID))
+
+			barrier := hooksContext.GetBarrier(barrierID)
+			barrier.Wait(r.Context(), hooksID+":"+barrierID, nil)
+
+			log.Info("http request done waiting on barrier",
+				zap.String("hooks-id", hooksID),
+				zap.String("barrier-id", barrierID))
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}

gateway/system/system.go

@@ -70,6 +70,7 @@ type System struct {
 	dataServer *grpc.Server
 	dapiServer *http.Server
 
+	hooksManager    *hooks.HooksManager
 	shutdownTimeout time.Duration
 }
 
@@ -231,6 +232,9 @@ func NewSystem(opts *SystemOptions) (*System, error) {
 	})
 
 	var httpHandler http.Handler = mux
+	if opts.Debug {
+		httpHandler = hooksManager.HTTPMiddleware()(httpHandler)
+	}
 	if opts.RateLimiter != nil {
 		httpHandler = opts.RateLimiter.HttpMiddleware(httpHandler)
 	}
@@ -252,12 +256,17 @@ func NewSystem(opts *SystemOptions) (*System, error) {
 		logger:          opts.Logger,
 		dataServer:      dataSrv,
 		dapiServer:      dapiSrv,
+		hooksManager:    hooksManager,
 		shutdownTimeout: opts.ShutdownTimeout,
 	}
 
 	return s, nil
 }
 
+func (s *System) HooksManager() *hooks.HooksManager {
+	return s.hooksManager
+}
+
 func (s *System) Serve(ctx context.Context, l *Listeners) error {
 	var wg sync.WaitGroup