api: enforce Connect stream lifecycle policies

Bound stream idle time and lifetime, track active RPC cancellation, bound terminated unary response writes, and cancel long-lived handlers during shutdown so every admission slot is eventually released.

Signed-off-by: Siavash Safi <siavash@cloudflare.com>

Author: siavashs

api/api.go

@@ -85,6 +85,12 @@ type Options struct {
 	// ConnectUnaryTimeout is the timeout for Connect unary RPCs. Zero inherits
 	// Timeout. A negative value disables the Connect unary timeout.
 	ConnectUnaryTimeout time.Duration
+	// ConnectStreamIdleTimeout limits the time between messages on a Connect
+	// stream. Non-positive values disable the timeout.
+	ConnectStreamIdleTimeout time.Duration
+	// ConnectStreamLifetime limits the total lifetime of a Connect stream.
+	// Non-positive values disable the timeout.
+	ConnectStreamLifetime time.Duration
 	// ConnectReadMaxBytes limits each incoming Connect protobuf message.
 	// Non-positive values do not set a limit.
 	ConnectReadMaxBytes int
@@ -149,6 +155,8 @@ func (o Options) resolve() effectiveOptions {
 			UnaryConcurrency:    unaryConcurrency,
 			StreamConcurrency:   streamConcurrency,
 			UnaryTimeout:        unaryTimeout,
+			StreamIdleTimeout:   max(o.ConnectStreamIdleTimeout, 0),
+			StreamLifetime:      max(o.ConnectStreamLifetime, 0),
 			ReadMaxBytes:        max(o.ConnectReadMaxBytes, 0),
 			SendMaxBytes:        max(o.ConnectSendMaxBytes, 0),
 			MaxRequestBodyBytes: max(o.ConnectMaxRequestBodyBytes, 0),
@@ -320,6 +328,13 @@ func (api *API) Update(cfg *config.Config, setAlertStatus func(ctx context.Conte
 	}
 }
 
+// Shutdown rejects new Connect RPCs and cancels active RPCs.
+func (api *API) Shutdown() {
+	if api.connect != nil {
+		api.connect.Shutdown()
+	}
+}
+
 func (api *API) limitHandler(h http.Handler) http.Handler {
 	limited := api.concurrencyLimitHandler(h)
 	if api.timeout <= 0 {

api/api_test.go

@@ -87,6 +87,8 @@ func TestOptionsResolve(t *testing.T) {
 	effective := (Options{
 		Concurrency:                3,
 		Timeout:                    time.Minute,
+		ConnectStreamIdleTimeout:   -time.Second,
+		ConnectStreamLifetime:      -time.Second,
 		ConnectReadMaxBytes:        -1,
 		ConnectSendMaxBytes:        -1,
 		ConnectMaxRequestBodyBytes: -1,
@@ -96,6 +98,8 @@ func TestOptionsResolve(t *testing.T) {
 	require.Equal(t, 3, effective.connect.UnaryConcurrency)
 	require.Equal(t, 3, effective.connect.StreamConcurrency)
 	require.Equal(t, time.Minute, effective.connect.UnaryTimeout)
+	require.Zero(t, effective.connect.StreamIdleTimeout)
+	require.Zero(t, effective.connect.StreamLifetime)
 	require.Zero(t, effective.connect.ReadMaxBytes)
 	require.Zero(t, effective.connect.SendMaxBytes)
 	require.Zero(t, effective.connect.MaxRequestBodyBytes)