fix(gateway): send on closed chan - (*Handle).webRequestQ

gateway/handle.go

@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	jsoniter "github.com/json-iterator/go"
@@ -94,6 +95,7 @@ type Handle struct {
 	// other state
 
 	backendConfigInitialised bool
+	inFlightRequestsCount    *atomic.Uint64
 
 	trackCounterMu    sync.Mutex // protects trackSuccessCount and trackFailureCount
 	trackSuccessCount int

gateway/handle_lifecycle.go

@@ -7,6 +7,7 @@ import (
 	"net/http/httputil"
 	"net/url"
 	"strconv"
+	"sync/atomic"
 	"time"
 
 	"github.com/rudderlabs/rudder-schemas/go/stream"
@@ -68,6 +69,7 @@ func (gw *Handle) Setup(
 	gw.versionHandler = versionHandler
 	gw.rsourcesService = rsourcesService
 	gw.sourcehandle = sourcehandle
+	gw.inFlightRequestsCount = new(atomic.Uint64)
 
 	// Port where GW is running
 	gw.conf.webPort = config.GetIntVar(8080, 1, "Gateway.webPort")
@@ -383,6 +385,18 @@ func (gw *Handle) StartWebHandler(ctx context.Context) error {
 		gw.logger.Child("rsources_failed_keys"),
 	)
 	srvMux.Use(
+		func(h http.Handler) http.Handler {
+			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				_ = gw.inFlightRequestsCount.Add(1)
+				h.ServeHTTP(w, r)
+				for {
+					prev := gw.inFlightRequestsCount.Load()
+					if gw.inFlightRequestsCount.CompareAndSwap(prev, prev-1) {
+						break
+					}
+				}
+			})
+		},
 		chiware.StatMiddleware(ctx, stats.Default, component),
 		middleware.LimitConcurrentRequests(gw.conf.maxConcurrentRequests),
 		middleware.UncompressMiddleware,
@@ -482,6 +496,10 @@ func (gw *Handle) Shutdown() error {
 		return err
 	}
 
+	for gw.inFlightRequestsCount.Load() != 0 {
+		time.Sleep(250 * time.Millisecond)
+	}
+
 	// UserWebRequestWorkers
 	for _, worker := range gw.userWebRequestWorkers {
 		close(worker.webRequestQ)