fix: Ensure all connections persist until queue-proxy drain

Fixes: Websockets (and some HTTP) closing abruptly when queue-proxy
undergoes drain.

Due to hijacked connections in net/http not being respected when
server.Shutdown is called, any active websocket connections currently
end as soon as the queue-proxy calls .Shutdown. See
gorilla/websocket#448 and golang/go#17721 for details. This patch fixes
this issue by introducing an atomic counter of active requests, which
increments as a request comes in and decrements as a request handler
terminates. During drain, this counter must reach zero or adhere to the
revision timeout, in order to call .Shutdown.

Further, this prevents pre-mature closing of connections in the user
container due to misconfigured SIGTERM handling, by delaying the SIGTERM
send until the queue-proxy has verified it has fully drained.

Author: elijah-rou

.gitignore

@@ -9,3 +9,6 @@
 # Temporary output of build tools
 bazel-*
 *.out
+
+# Repomix outputs
+repomix*.xml

pkg/activator/net/throttler.go

@@ -100,7 +100,13 @@ func (p *podTracker) Capacity() int {
 	if p.b == nil {
 		return 1
 	}
-	return p.b.Capacity()
+	capacity := p.b.Capacity()
+	// Safe conversion: breaker capacity is always reasonable for int
+	// Check for overflow before conversion
+	if capacity > 0x7FFFFFFF {
+		return 0x7FFFFFFF // Return max int32 value
+	}
+	return int(capacity)
 }
 
 func (p *podTracker) UpdateConcurrency(c int) {
@@ -118,7 +124,7 @@ func (p *podTracker) Reserve(ctx context.Context) (func(), bool) {
 }
 
 type breaker interface {
-	Capacity() int
+	Capacity() uint64
 	Maybe(ctx context.Context, thunk func()) error
 	UpdateConcurrency(int)
 	Reserve(ctx context.Context) (func(), bool)
@@ -721,8 +727,13 @@ func newInfiniteBreaker(logger *zap.SugaredLogger) *infiniteBreaker {
 }
 
 // Capacity returns the current capacity of the breaker
-func (ib *infiniteBreaker) Capacity() int {
-	return int(ib.concurrency.Load())
+func (ib *infiniteBreaker) Capacity() uint64 {
+	concurrency := ib.concurrency.Load()
+	// Safe conversion: concurrency is int32 and we check for non-negative
+	if concurrency >= 0 {
+		return uint64(concurrency)
+	}
+	return 0
 }
 
 func zeroOrOne(x int) int32 {

pkg/activator/net/throttler_test.go

@@ -226,7 +226,7 @@ func TestThrottlerUpdateCapacity(t *testing.T) {
 				rt.breaker = newInfiniteBreaker(logger)
 			}
 			rt.updateCapacity(tt.capacity)
-			if got := rt.breaker.Capacity(); got != tt.want {
+			if got := rt.breaker.Capacity(); got != uint64(tt.want) {
 				t.Errorf("Capacity = %d, want: %d", got, tt.want)
 			}
 			if tt.checkAssignedPod {
@@ -560,7 +560,7 @@ func TestThrottlerSuccesses(t *testing.T) {
 				rt.mux.RLock()
 				defer rt.mux.RUnlock()
 				if *cc != 0 {
-					return rt.activatorIndex.Load() != -1 && rt.breaker.Capacity() == wantCapacity &&
+					return rt.activatorIndex.Load() != -1 && rt.breaker.Capacity() == uint64(wantCapacity) &&
 						sortedTrackers(rt.assignedTrackers), nil
 				}
 				// If CC=0 then verify number of backends, rather the capacity of breaker.
@@ -638,7 +638,7 @@ func TestPodAssignmentFinite(t *testing.T) {
 	if got, want := trackerDestSet(rt.assignedTrackers), sets.New("ip0", "ip4"); !got.Equal(want) {
 		t.Errorf("Assigned trackers = %v, want: %v, diff: %s", got, want, cmp.Diff(want, got))
 	}
-	if got, want := rt.breaker.Capacity(), 2*42; got != want {
+	if got, want := rt.breaker.Capacity(), uint64(2*42); got != want {
 		t.Errorf("TotalCapacity = %d, want: %d", got, want)
 	}
 	if got, want := rt.assignedTrackers[0].Capacity(), 42; got != want {
@@ -657,7 +657,7 @@ func TestPodAssignmentFinite(t *testing.T) {
 	if got, want := len(rt.assignedTrackers), 0; got != want {
 		t.Errorf("NumAssignedTrackers = %d, want: %d", got, want)
 	}
-	if got, want := rt.breaker.Capacity(), 0; got != want {
+	if got, want := rt.breaker.Capacity(), uint64(0); got != want {
 		t.Errorf("TotalCapacity = %d, want: %d", got, want)
 	}
 }
@@ -687,7 +687,7 @@ func TestPodAssignmentInfinite(t *testing.T) {
 	if got, want := len(rt.assignedTrackers), 3; got != want {
 		t.Errorf("NumAssigned trackers = %d, want: %d", got, want)
 	}
-	if got, want := rt.breaker.Capacity(), 1; got != want {
+	if got, want := rt.breaker.Capacity(), uint64(1); got != want {
 		t.Errorf("TotalCapacity = %d, want: %d", got, want)
 	}
 	if got, want := rt.assignedTrackers[0].Capacity(), 1; got != want {
@@ -703,7 +703,7 @@ func TestPodAssignmentInfinite(t *testing.T) {
 	if got, want := len(rt.assignedTrackers), 0; got != want {
 		t.Errorf("NumAssignedTrackers = %d, want: %d", got, want)
 	}
-	if got, want := rt.breaker.Capacity(), 0; got != want {
+	if got, want := rt.breaker.Capacity(), uint64(0); got != want {
 		t.Errorf("TotalCapacity = %d, want: %d", got, want)
 	}
 }
@@ -935,7 +935,7 @@ func TestInfiniteBreaker(t *testing.T) {
 	}
 
 	// Verify initial condition.
-	if got, want := b.Capacity(), 0; got != want {
+	if got, want := b.Capacity(), uint64(0); got != want {
 		t.Errorf("Cap=%d, want: %d", got, want)
 	}
 	if _, ok := b.Reserve(context.Background()); ok != true {
@@ -949,7 +949,7 @@ func TestInfiniteBreaker(t *testing.T) {
 	}
 
 	b.UpdateConcurrency(1)
-	if got, want := b.Capacity(), 1; got != want {
+	if got, want := b.Capacity(), uint64(1); got != want {
 		t.Errorf("Cap=%d, want: %d", got, want)
 	}
 
@@ -976,7 +976,7 @@ func TestInfiniteBreaker(t *testing.T) {
 	if err := b.Maybe(ctx, nil); err == nil {
 		t.Error("Should have failed, but didn't")
 	}
-	if got, want := b.Capacity(), 0; got != want {
+	if got, want := b.Capacity(), uint64(0); got != want {
 		t.Errorf("Cap=%d, want: %d", got, want)
 	}
 
@@ -1212,7 +1212,7 @@ func TestAssignSlice(t *testing.T) {
 			t.Errorf("Got=%v, want: %v; diff: %s", got, want,
 				cmp.Diff(want, got, opt))
 		}
-		if got, want := got[0].b.Capacity(), 0; got != want {
+		if got, want := got[0].b.Capacity(), uint64(0); got != want {
 			t.Errorf("Capacity for the tail pod = %d, want: %d", got, want)
 		}
 	})

pkg/autoscaler/metrics/stat.pb.go

@@ -43,7 +43,7 @@ type BreakerParams struct {
 // executions in excess of the concurrency limit. Function call attempts
 // beyond the limit of the queue are failed immediately.
 type Breaker struct {
-	inFlight   atomic.Int64
+	pending    atomic.Int64
 	totalSlots int64
 	sem        *semaphore
 
@@ -83,10 +83,10 @@ func NewBreaker(params BreakerParams) *Breaker {
 func (b *Breaker) tryAcquirePending() bool {
 	// This is an atomic version of:
 	//
-	// if inFlight == totalSlots {
+	// if pending == totalSlots {
 	//   return false
 	// } else {
-	//   inFlight++
+	//   pending++
 	//   return true
 	// }
 	//
@@ -96,19 +96,20 @@ func (b *Breaker) tryAcquirePending() bool {
 	// (it fails if we're raced to it) or if we don't fulfill the condition
 	// anymore.
 	for {
-		cur := b.inFlight.Load()
+		cur := b.pending.Load()
+		// 10000 + containerConcurrency = totalSlots
 		if cur == b.totalSlots {
 			return false
 		}
-		if b.inFlight.CompareAndSwap(cur, cur+1) {
+		if b.pending.CompareAndSwap(cur, cur+1) {
 			return true
 		}
 	}
 }
 
 // releasePending releases a slot on the pending "queue".
 func (b *Breaker) releasePending() {
-	b.inFlight.Add(-1)
+	b.pending.Add(-1)
 }
 
 // Reserve reserves an execution slot in the breaker, to permit
@@ -154,9 +155,9 @@ func (b *Breaker) Maybe(ctx context.Context, thunk func()) error {
 	return nil
 }
 
-// InFlight returns the number of requests currently in flight in this breaker.
-func (b *Breaker) InFlight() int {
-	return int(b.inFlight.Load())
+// Pending returns the number of requests currently pending to this breaker.
+func (b *Breaker) Pending() int {
+	return int(b.pending.Load())
 }
 
 // UpdateConcurrency updates the maximum number of in-flight requests.
@@ -165,10 +166,15 @@ func (b *Breaker) UpdateConcurrency(size int) {
 }
 
 // Capacity returns the number of allowed in-flight requests on this breaker.
-func (b *Breaker) Capacity() int {
+func (b *Breaker) Capacity() uint64 {
 	return b.sem.Capacity()
 }
 
+// InFlight returns the number of requests currently in-flight on this breaker.
+func (b *Breaker) InFlight() uint64 {
+	return b.sem.InFlight()
+}
+
 // newSemaphore creates a semaphore with the desired initial capacity.
 func newSemaphore(maxCapacity, initialCapacity int) *semaphore {
 	queue := make(chan struct{}, maxCapacity)
@@ -288,9 +294,15 @@ func (s *semaphore) updateCapacity(size int) {
 }
 
 // Capacity is the capacity of the semaphore.
-func (s *semaphore) Capacity() int {
+func (s *semaphore) Capacity() uint64 {
 	capacity, _ := unpack(s.state.Load())
-	return int(capacity) //nolint:gosec // TODO(dprotaso) - capacity should be uint64
+	return capacity
+}
+
+// InFlight is the number of the inflight requests of the semaphore.
+func (s *semaphore) InFlight() uint64 {
+	_, inFlight := unpack(s.state.Load())
+	return inFlight
 }
 
 // unpack takes an uint64 and returns two uint32 (as uint64) comprised of the leftmost

pkg/queue/breaker.go

@@ -212,12 +212,12 @@ func TestBreakerUpdateConcurrency(t *testing.T) {
 	params := BreakerParams{QueueDepth: 1, MaxConcurrency: 1, InitialCapacity: 0}
 	b := NewBreaker(params)
 	b.UpdateConcurrency(1)
-	if got, want := b.Capacity(), 1; got != want {
+	if got, want := b.Capacity(), uint64(1); got != want {
 		t.Errorf("Capacity() = %d, want: %d", got, want)
 	}
 
 	b.UpdateConcurrency(0)
-	if got, want := b.Capacity(), 0; got != want {
+	if got, want := b.Capacity(), uint64(0); got != want {
 		t.Errorf("Capacity() = %d, want: %d", got, want)
 	}
 }
@@ -294,12 +294,12 @@ func TestSemaphoreRelease(t *testing.T) {
 func TestSemaphoreUpdateCapacity(t *testing.T) {
 	const initialCapacity = 1
 	sem := newSemaphore(3, initialCapacity)
-	if got, want := sem.Capacity(), 1; got != want {
+	if got, want := sem.Capacity(), uint64(1); got != want {
 		t.Errorf("Capacity = %d, want: %d", got, want)
 	}
 	sem.acquire(context.Background())
 	sem.updateCapacity(initialCapacity + 2)
-	if got, want := sem.Capacity(), 3; got != want {
+	if got, want := sem.Capacity(), uint64(3); got != want {
 		t.Errorf("Capacity = %d, want: %d", got, want)
 	}
 }

pkg/queue/breaker_test.go

@@ -85,7 +85,7 @@ func (h *appRequestMetricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Requ
 	startTime := h.clock.Now()
 
 	if h.breaker != nil {
-		h.queueLen.Record(r.Context(), int64(h.breaker.InFlight()))
+		h.queueLen.Record(r.Context(), int64(h.breaker.Pending()))
 	}
 	defer func() {
 		// Filter probe requests for revision metrics.

pkg/queue/request_metric.go

@@ -18,8 +18,11 @@ package sharedmain
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"net/http"
+	"strings"
+	"sync/atomic"
 	"time"
 
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
@@ -30,6 +33,7 @@ import (
 	netheader "knative.dev/networking/pkg/http/header"
 	netproxy "knative.dev/networking/pkg/http/proxy"
 	netstats "knative.dev/networking/pkg/http/stats"
+	"knative.dev/pkg/network"
 	pkghandler "knative.dev/pkg/network/handlers"
 	"knative.dev/serving/pkg/activator"
 	pkghttp "knative.dev/serving/pkg/http"
@@ -46,6 +50,7 @@ func mainHandler(
 	logger *zap.SugaredLogger,
 	mp metric.MeterProvider,
 	tp trace.TracerProvider,
+	pendingRequests *atomic.Int32,
 ) (http.Handler, *pkghandler.Drainer) {
 	target := net.JoinHostPort("127.0.0.1", env.UserPort)
 	tracer := tp.Tracer("knative.dev/serving/pkg/queue")
@@ -73,6 +78,7 @@ func mainHandler(
 
 	composedHandler = requestAppMetricsHandler(logger, composedHandler, breaker, mp)
 	composedHandler = queue.ProxyHandler(tracer, breaker, stats, composedHandler)
+
 	composedHandler = queue.ForwardedShimHandler(composedHandler)
 	composedHandler = handler.NewTimeoutHandler(composedHandler, "request timeout", func(r *http.Request) (time.Duration, time.Duration, time.Duration) {
 		return timeout, responseStartTimeout, idleTimeout
@@ -81,6 +87,8 @@ func mainHandler(
 	composedHandler = queue.NewRouteTagHandler(composedHandler)
 	composedHandler = withFullDuplex(composedHandler, env.EnableHTTPFullDuplex, logger)
 
+	composedHandler = withRequestCounter(composedHandler, pendingRequests)
+
 	drainer := &pkghandler.Drainer{
 		QuietPeriod: drainSleepDuration,
 		// Add Activator probe header to the drainer so it can handle probes directly from activator
@@ -105,11 +113,10 @@ func mainHandler(
 			return !netheader.IsProbe(r)
 		}),
 	)
-
 	return composedHandler, drainer
 }
 
-func adminHandler(ctx context.Context, logger *zap.SugaredLogger, drainer *pkghandler.Drainer) http.Handler {
+func adminHandler(ctx context.Context, logger *zap.SugaredLogger, drainer *pkghandler.Drainer, pendingRequests *atomic.Int32) http.Handler {
 	mux := http.NewServeMux()
 	mux.HandleFunc(queue.RequestQueueDrainPath, func(w http.ResponseWriter, r *http.Request) {
 		logger.Info("Attached drain handler from user-container", r)
@@ -130,6 +137,17 @@ func adminHandler(ctx context.Context, logger *zap.SugaredLogger, drainer *pkgha
 		w.WriteHeader(http.StatusOK)
 	})
 
+	// New endpoint that returns 200 only when all requests are drained
+	mux.HandleFunc("/drain-complete", func(w http.ResponseWriter, r *http.Request) {
+		if pendingRequests.Load() <= 0 {
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte("drained"))
+		} else {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			fmt.Fprintf(w, "pending requests: %d", pendingRequests.Load())
+		}
+	})
+
 	return mux
 }
 
@@ -145,3 +163,13 @@ func withFullDuplex(h http.Handler, enableFullDuplex bool, logger *zap.SugaredLo
 		h.ServeHTTP(w, r)
 	})
 }
+
+func withRequestCounter(h http.Handler, pendingRequests *atomic.Int32) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get(network.ProbeHeaderName) != network.ProbeHeaderValue && !strings.HasPrefix(r.Header.Get("User-Agent"), "kube-probe/") {
+			pendingRequests.Add(1)
+			defer pendingRequests.Add(-1)
+		}
+		h.ServeHTTP(w, r)
+	})
+}