feat(api): add connection-aware readiness probe for SSE

When SSE connections reach the configured threshold (default 95% of
--sse-max-connections), the readiness probe returns 503. This removes
the pod from Kubernetes Endpoints and GCLB NEG, preventing new
connections from being routed to saturated pods while keeping existing
SSE connections alive.

The Envoy sidecar circuit breaker is ineffective for SSE over HTTP/2
because multiplexed streams are not counted as active TCP connections
or requests. This readiness check bypasses Envoy and reads the
Dispatcher's connection count directly.

Closes #2786

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: t-kikuc

pkg/api/cmd/server.go

@@ -143,6 +143,7 @@ type server struct {
 	cacheInvalidationTopic    *string
 	sseHeartbeatInterval      *time.Duration
 	sseMaxConnections         *int
+	sseReadinessThreshold    *float64
 }
 
 func RegisterCommand(r cli.CommandRegistry, p cli.ParentCommand) cli.Command {
@@ -339,6 +340,9 @@ func RegisterCommand(r cli.CommandRegistry, p cli.ParentCommand) cli.Command {
 		sseMaxConnections: cmd.Flag("sse-max-connections",
 			"Maximum number of concurrent SSE connections per pod.",
 		).Default("10000").Int(),
+		sseReadinessThreshold: cmd.Flag("sse-readiness-threshold",
+			"Fraction of sse-max-connections at which the readiness probe starts failing (0.0-1.0).",
+		).Default("0.95").Float64(),
 	}
 	r.RegisterCommand(server)
 	return server
@@ -679,9 +683,15 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
 	healthCheckCtx, healthCheckCancel := context.WithCancel(context.Background())
 	defer healthCheckCancel()
 
+	sseReadinessLimit := int(float64(*s.sseMaxConnections) * *s.sseReadinessThreshold)
+	sseReadinessCheck := func() bool {
+		current, _ := streamDispatcher.ConnectionStatus()
+		return current < sseReadinessLimit
+	}
 	healthChecker := health.NewGrpcChecker(
 		health.WithTimeout(5*time.Second),
 		health.WithCheck("metrics", metrics.Check),
+		health.WithReadinessCheck("sse-connections", sseReadinessCheck),
 	)
 	go healthChecker.Run(healthCheckCtx)
 
@@ -729,6 +739,7 @@ func (s *server) Run(ctx context.Context, metrics metrics.Metrics, logger *zap.L
 		api.Version, api.Service,
 		health.WithTimeout(5*time.Second),
 		health.WithCheck("metrics", metrics.Check),
+		health.WithReadinessCheck("sse-connections", sseReadinessCheck),
 	)
 	go restHealthChecker.Run(healthCheckCtx)
 

pkg/api/stream/dispatcher.go

@@ -69,6 +69,13 @@ func NewDispatcher(maxConns int, fetchFeatures FeaturesFetcher, logger *zap.Logg
 	}
 }
 
+// ConnectionStatus returns the current and maximum number of SSE connections.
+func (d *Dispatcher) ConnectionStatus() (current, max int) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	return d.totalConns, d.maxConns
+}
+
 // Shutdown signals all active SSE handlers to exit immediately.
 func (d *Dispatcher) Shutdown() {
 	d.shutdownOnce.Do(func() { close(d.shutdownCh) })

pkg/api/stream/dispatcher_test.go

@@ -135,6 +135,32 @@ func TestDispatcherRegisterMaxConns(t *testing.T) {
 	}
 }
 
+func TestDispatcherConnectionStatus(t *testing.T) {
+	t.Parallel()
+	d := NewDispatcher(100, nil, zap.NewNop())
+
+	cur, max := d.ConnectionStatus()
+	assert.Equal(t, 0, cur)
+	assert.Equal(t, 100, max)
+
+	_, dereg1, err := d.register("env-1", "tag-A", "src")
+	require.NoError(t, err)
+	_, dereg2, err := d.register("env-1", "tag-B", "src")
+	require.NoError(t, err)
+
+	cur, max = d.ConnectionStatus()
+	assert.Equal(t, 2, cur)
+	assert.Equal(t, 100, max)
+
+	dereg1()
+	cur, _ = d.ConnectionStatus()
+	assert.Equal(t, 1, cur)
+
+	dereg2()
+	cur, _ = d.ConnectionStatus()
+	assert.Equal(t, 0, cur)
+}
+
 func TestDispatcherRegisterSlotFreedByDeregister(t *testing.T) {
 	t.Parallel()
 	maxConns := 1

pkg/health/health.go

@@ -45,13 +45,18 @@ func (s Status) String() string {
 
 type check func(context.Context) Status
 
+// ReadinessCheck is called synchronously on each readiness probe.
+// It returns true when the pod should accept new traffic.
+type ReadinessCheck func() bool
+
 type checker struct {
 	status  uint32
 	stopped uint32 // 0 = running, 1 = stopped
 
-	interval time.Duration
-	timeout  time.Duration
-	checks   map[string]check
+	interval        time.Duration
+	timeout         time.Duration
+	checks          map[string]check
+	readinessChecks map[string]ReadinessCheck
 }
 
 type option func(*checker)
@@ -77,12 +82,19 @@ func WithTimeout(timeout time.Duration) option {
 	}
 }
 
+func WithReadinessCheck(name string, rc ReadinessCheck) option {
+	return func(c *checker) {
+		c.readinessChecks[name] = rc
+	}
+}
+
 func newChecker(opts ...option) *checker {
 	checker := &checker{
-		status:   uint32(Unhealthy),
-		interval: 10 * time.Second,
-		timeout:  5 * time.Second,
-		checks:   make(map[string]check),
+		status:          uint32(Unhealthy),
+		interval:        10 * time.Second,
+		timeout:         5 * time.Second,
+		checks:          make(map[string]check),
+		readinessChecks: make(map[string]ReadinessCheck),
 	}
 	for _, o := range opts {
 		o(checker)
@@ -135,11 +147,16 @@ func (hc *checker) check(ctx context.Context) {
 }
 
 func (hc *checker) ServeReadyHTTP(resp http.ResponseWriter, req *http.Request) {
-	// Readiness check: return 503 if health checks fail
 	if hc.getStatus() == Unhealthy {
 		resp.WriteHeader(http.StatusServiceUnavailable)
 		return
 	}
+	for _, rc := range hc.readinessChecks {
+		if !rc() {
+			resp.WriteHeader(http.StatusServiceUnavailable)
+			return
+		}
+	}
 	resp.WriteHeader(http.StatusOK)
 }
 

pkg/health/health_test.go

@@ -137,6 +137,67 @@ func TestHTTPReadyUnhealthy(t *testing.T) {
 	}
 }
 
+func TestReadinessCheckBlocksReady(t *testing.T) {
+	t.Parallel()
+	ready := true
+	checker := NewRestChecker(version, service,
+		WithReadinessCheck("capacity", func() bool { return ready }),
+	)
+	checker.check(context.Background())
+
+	req := httptest.NewRequest("GET", fmt.Sprintf("%s%s%s", version, service, readyPath), nil)
+	resp := httptest.NewRecorder()
+	checker.ServeReadyHTTP(resp, req)
+	if resp.Code != http.StatusOK {
+		t.Errorf("Expected 200 when readiness check passes, got %d", resp.Code)
+	}
+
+	ready = false
+	resp = httptest.NewRecorder()
+	checker.ServeReadyHTTP(resp, req)
+	if resp.Code != http.StatusServiceUnavailable {
+		t.Errorf("Expected 503 when readiness check fails, got %d", resp.Code)
+	}
+}
+
+func TestReadinessCheckDoesNotAffectLiveness(t *testing.T) {
+	t.Parallel()
+	checker := NewRestChecker(version, service,
+		WithReadinessCheck("capacity", func() bool { return false }),
+	)
+	checker.check(context.Background())
+
+	req := httptest.NewRequest("GET", getTargetPath(t), nil)
+	resp := httptest.NewRecorder()
+	checker.ServeLiveHTTP(resp, req)
+	if resp.Code != http.StatusOK {
+		t.Errorf("Liveness should not be affected by readiness check, got %d", resp.Code)
+	}
+}
+
+func TestGRPCReadinessCheckBlocksReady(t *testing.T) {
+	t.Parallel()
+	ready := true
+	checker := NewGrpcChecker(
+		WithReadinessCheck("capacity", func() bool { return ready }),
+	)
+	checker.check(context.Background())
+
+	req := httptest.NewRequest("GET", readyPath, nil)
+	resp := httptest.NewRecorder()
+	checker.ServeHTTP(resp, req)
+	if resp.Code != http.StatusOK {
+		t.Errorf("Expected 200 when readiness check passes, got %d", resp.Code)
+	}
+
+	ready = false
+	resp = httptest.NewRecorder()
+	checker.ServeHTTP(resp, req)
+	if resp.Code != http.StatusServiceUnavailable {
+		t.Errorf("Expected 503 when readiness check fails, got %d", resp.Code)
+	}
+}
+
 func TestHTTPHealthAffectedByStop(t *testing.T) {
 	t.Parallel()
 	patterns := []struct {