fix: reserve chat hook dispatch capacity for running turns (#27656)

## Context Follow-up fix from live UAT of the merged chat lifecycle hooks stack (#27430). Its companion UAT fix (#27655) has merged, so this targets `main` directly. ## Why? UAT measured a burst of 1,500 concurrent chat creations against a consumer with 1.2s latency. 255 were admitted and 1,245 got `502 hook_dispatch_failed (over_capacity)`, which is correct fail-closed behavior. The collateral wasn't: the same burst failed 24 `stop` dispatches, parking chats that had already been admitted and had already executed tools. One 256-slot semaphore served every event, so new-work admission could take every slot and kill turns in flight. Callers now classify each dispatch as admission or generation, and admission draws from a 192-slot gate held *before* the shared pool. At least 64 shared slots stay reachable only by dispatches for work a chat already admitted. The dispatcher is per `coderd` replica, so these limits are per replica, not deployment-wide, and the docs say so. **The caller classifies, not the event type.** Event type isn't a reliable proxy in either direction: a subagent spawn dispatches `user_prompt_submit` from inside a running turn, and the edit path dispatches `session_start` at admission time. `CapacityClassUnset` is rejected in `Dispatch`, so a new call site fails closed rather than silently inheriting a share. **Acquisition order is load-bearing.** Admission takes its own gate first. Taking a shared slot first would let admissions queued on the gate occupy the very capacity the reserve protects. `acquireCapacity` is the only path that takes either pool, so the order can't be bypassed. ## What this does not guarantee Nothing bounds how many turns generate concurrently, so the 192/64 split is a judgement call, not a derived ceiling. This stops an *admission* burst from consuming every slot; it does not make the remainder sufficient. A large enough generation load can still exhaust the reserve and error a running chat. The docs say so explicitly rather than promising a guarantee the code doesn't deliver. Generation can now take all 256 slots, so generation traffic starves admission harder than before. That's the intended priority: rejecting a new prompt is recoverable, ending a turn that already ran tools is not. ## Testing Red-green proved both new tests. Removing the release-on-failure path fails `RefusedSharedAcquireReleasesAdmission` deterministically; removing the expired-deadline check fails `ExpiredDeadlineRefusesFreeSlot` in 18/30 runs. That deadline check fixes a real race found in review. `acquire` previously shared one `time.Timer` across both acquires. Because `select` picks a ready case at random, an admission dispatch could take a slot after its capacity deadline had passed. Measured over 300 trials: 135 late acquisitions, worst overshoot 2.1ms. `acquire` now takes an absolute deadline and refuses an expired one before selecting, which measures 0/300. Go: `coderd/x/agenthooks/...` and `coderd/x/chatd/...`, plus `-race -count=3` on the dispatcher. > Mux opened this PR on Mike's behalf.
2026-09-24 15:04:27 +08:00 · 2026-08-03 19:04:54 +02:00
parent 7bd9f5ec93
commit fc24c27dfd
10 changed files with 224 additions and 43 deletions
@@ -28,11 +28,22 @@ import (

 const (
 	maxConcurrentDispatches = 256
-	maxResponseBodyBytes    = 1_048_576
-	maxModelContextBytes    = 16_384
-	capacityWaitLimit       = 250 * time.Millisecond
-	retryBackoff            = 250 * time.Millisecond
-	clockSkewLeeway         = 30 * time.Second
+	// Keep capacity reachable by work that a chat already admitted.
+	maxAdmissionDispatches = 192
+	maxResponseBodyBytes   = 1_048_576
+	maxModelContextBytes   = 16_384
+	capacityWaitLimit      = 250 * time.Millisecond
+	retryBackoff           = 250 * time.Millisecond
+	clockSkewLeeway        = 30 * time.Second
+)
+
+// CapacityClass selects which share of dispatch capacity an event draws from.
+type CapacityClass int
+
+const (
+	CapacityClassUnset CapacityClass = iota
+	CapacityClassAdmission
+	CapacityClassGeneration
 )

 // Result classifies the terminal outcome of a dispatch attempt.
@@ -53,7 +64,8 @@ const (
 type Event struct {
 	Type agenthooks.EventType
 	agenthooks.ChatRef
-	Data any
+	Data     any
+	Capacity CapacityClass
 }

 // Error preserves the attempt ID and failure class.
@@ -90,6 +102,7 @@ type Dispatcher struct {
 	deploymentID string
 	userAgent    string
 	semaphore    chan struct{}
+	admission    chan struct{}
 	metrics      *metrics
 }

@@ -167,6 +180,7 @@ func New(
 		deploymentID: deploymentID,
 		userAgent:    "coderd-agenthooks/" + coderVersion,
 		semaphore:    make(chan struct{}, maxConcurrentDispatches),
+		admission:    make(chan struct{}, maxAdmissionDispatches),
 		metrics:      newMetrics(reg),
 	}
 }
@@ -187,25 +201,21 @@ func (d *Dispatcher) Dispatch(ctx context.Context, event Event) (agenthooks.Resp
 		return agenthooks.Response{}, uuid.Nil, xerrors.Errorf("chat hook URL rejected: %w", d.hookURLErr)
 	}

+	switch event.Capacity {
+	case CapacityClassAdmission, CapacityClassGeneration:
+	default:
+		return agenthooks.Response{}, uuid.Nil, xerrors.Errorf("dispatch event %q has no capacity class", event.Type)
+	}
+
 	startedAt := time.Now()
 	dispatchID := uuid.New()
-	wait := max(min(d.timeout, capacityWaitLimit), 0)
-	capacityTimer := time.NewTimer(wait)
-	defer capacityTimer.Stop()
+	capacityDeadline := startedAt.Add(max(min(d.timeout, capacityWaitLimit), 0))

-	// The capacity wait runs against its own timer rather than a dispatch
-	// deadline, so a timeout shorter than capacityWaitLimit cannot make the
-	// over-capacity and caller-cancellation cases race.
-	select {
-	case d.semaphore <- struct{}{}:
-		defer func() { <-d.semaphore }()
-	case <-ctx.Done():
-		outcome := dispatchOutcome{result: ResultTimeout, err: ctx.Err()}
-		return agenthooks.Response{}, dispatchID, d.finish(ctx, event, dispatchID, startedAt, outcome)
-	case <-capacityTimer.C:
-		outcome := dispatchOutcome{result: ResultOverCapacity, err: context.DeadlineExceeded}
-		return agenthooks.Response{}, dispatchID, d.finish(ctx, event, dispatchID, startedAt, outcome)
+	release, refused, ok := d.acquireCapacity(ctx, event.Capacity, capacityDeadline)
+	if !ok {
+		return agenthooks.Response{}, dispatchID, d.finish(ctx, event, dispatchID, startedAt, refused)
 	}
+	defer release()

 	// Both post attempts share whatever remains of the configured timeout so
 	// that waiting for capacity cannot extend the dispatch past it.
@@ -219,6 +229,55 @@ func (d *Dispatcher) Dispatch(ctx context.Context, event Event) (agenthooks.Resp
 	return outcome.response, dispatchID, nil
 }

+// Admission acquires its gate before the shared pool so queued admission
+// dispatches cannot occupy the reserved capacity.
+func (d *Dispatcher) acquireCapacity(
+	ctx context.Context,
+	capacity CapacityClass,
+	deadline time.Time,
+) (release func(), outcome dispatchOutcome, ok bool) {
+	if capacity == CapacityClassAdmission {
+		releaseAdmission, refused, admitted := acquire(ctx, d.admission, deadline)
+		if !admitted {
+			return nil, refused, false
+		}
+		releaseShared, refused, acquired := acquire(ctx, d.semaphore, deadline)
+		if !acquired {
+			releaseAdmission()
+			return nil, refused, false
+		}
+		return func() {
+			releaseShared()
+			releaseAdmission()
+		}, dispatchOutcome{}, true
+	}
+	return acquire(ctx, d.semaphore, deadline)
+}
+
+// Check the deadline before select because select randomly chooses among ready
+// cases, including a free slot and an expired timer.
+func acquire(
+	ctx context.Context,
+	pool chan struct{},
+	deadline time.Time,
+) (release func(), outcome dispatchOutcome, ok bool) {
+	overCapacity := dispatchOutcome{result: ResultOverCapacity, err: context.DeadlineExceeded}
+	remaining := time.Until(deadline)
+	if remaining <= 0 {
+		return nil, overCapacity, false
+	}
+	timer := time.NewTimer(remaining)
+	defer timer.Stop()
+	select {
+	case pool <- struct{}{}:
+		return func() { <-pool }, dispatchOutcome{}, true
+	case <-ctx.Done():
+		return nil, dispatchOutcome{result: ResultTimeout, err: ctx.Err()}, false
+	case <-timer.C:
+		return nil, overCapacity, false
+	}
+}
+
 func (d *Dispatcher) finish(
 	ctx context.Context,
 	event Event,
@@ -671,7 +671,8 @@ func newTestEvent(t *testing.T, eventType agenthooks.EventType, data any) Event
 			ChatID:  uuid.New(),
 			OwnerID: uuid.New(),
 		},
-		Data: data,
+		Data:     data,
+		Capacity: CapacityClassGeneration,
 	}
 }

@@ -699,3 +700,111 @@ type roundTripFunc func(*http.Request) (*http.Response, error)
 func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
 	return f(req)
 }
+
+func TestDispatcherCapacityClassRequired(t *testing.T) {
+	t.Parallel()
+
+	event := newTestEvent(t, agenthooks.EventStop, agenthooks.StopData{})
+	event.Capacity = CapacityClassUnset
+	dispatcher := newTestDispatcher(t, nil, "https://unused.test", time.Second)
+
+	_, dispatchID, err := dispatcher.Dispatch(testutil.Context(t, testutil.WaitShort), event)
+	require.ErrorContains(t, err, "no capacity class")
+	require.Equal(t, uuid.Nil, dispatchID)
+}
+
+func TestDispatcherAdmissionReserve(t *testing.T) {
+	t.Parallel()
+
+	fill := func(t *testing.T, pool chan struct{}, count int) {
+		t.Helper()
+		for range count {
+			pool <- struct{}{}
+		}
+		t.Cleanup(func() {
+			for range count {
+				<-pool
+			}
+		})
+	}
+
+	t.Run("AdmissionReleasesBothPools", func(t *testing.T) {
+		t.Parallel()
+
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			_, err := w.Write([]byte(`{}`))
+			assert.NoError(t, err)
+		}))
+		t.Cleanup(server.Close)
+
+		dispatcher := New(
+			testutil.Logger(t), server.Client(), server.URL, testSecret, testutil.WaitShort,
+			testDeploymentID, testVersion, prometheus.NewRegistry(),
+		)
+		event := newTestEvent(t, agenthooks.EventUserPromptSubmit, agenthooks.UserPromptSubmitData{Prompt: "hi"})
+		event.Capacity = CapacityClassAdmission
+
+		_, _, err := dispatcher.Dispatch(testutil.Context(t, testutil.WaitLong), event)
+		require.NoError(t, err)
+		require.Empty(t, dispatcher.admission)
+		require.Empty(t, dispatcher.semaphore)
+	})
+
+	t.Run("SaturatedAdmissionRefusesAdmission", func(t *testing.T) {
+		t.Parallel()
+
+		dispatcher := newTestDispatcher(t, nil, "https://unused.test", 10*time.Millisecond)
+		fill(t, dispatcher.admission, maxAdmissionDispatches)
+
+		event := newTestEvent(t, agenthooks.EventUserPromptSubmit, agenthooks.UserPromptSubmitData{Prompt: "hi"})
+		event.Capacity = CapacityClassAdmission
+		_, _, err := dispatcher.Dispatch(testutil.Context(t, testutil.WaitLong), event)
+		assertDispatchErrorClass(t, err, ResultOverCapacity)
+		require.Empty(t, dispatcher.semaphore, "a refused admission must not hold a shared slot")
+	})
+
+	t.Run("ExpiredDeadlineRefusesFreeSlot", func(t *testing.T) {
+		t.Parallel()
+
+		pool := make(chan struct{}, 1)
+		release, outcome, ok := acquire(testutil.Context(t, testutil.WaitShort), pool, time.Now().Add(-time.Millisecond))
+		require.False(t, ok)
+		require.Nil(t, release)
+		require.Equal(t, ResultOverCapacity, outcome.result)
+		require.Empty(t, pool)
+	})
+
+	t.Run("RefusedSharedAcquireReleasesAdmission", func(t *testing.T) {
+		t.Parallel()
+
+		dispatcher := newTestDispatcher(t, nil, "https://unused.test", 10*time.Millisecond)
+		fill(t, dispatcher.semaphore, maxConcurrentDispatches)
+
+		event := newTestEvent(t, agenthooks.EventUserPromptSubmit, agenthooks.UserPromptSubmitData{Prompt: "hi"})
+		event.Capacity = CapacityClassAdmission
+		_, _, err := dispatcher.Dispatch(testutil.Context(t, testutil.WaitLong), event)
+		assertDispatchErrorClass(t, err, ResultOverCapacity)
+		require.Empty(t, dispatcher.admission, "an admission refused by the shared pool must release its gate token")
+	})
+
+	t.Run("SaturatedAdmissionStillServesGeneration", func(t *testing.T) {
+		t.Parallel()
+
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			_, err := w.Write([]byte(`{}`))
+			assert.NoError(t, err)
+		}))
+		t.Cleanup(server.Close)
+
+		dispatcher := New(
+			testutil.Logger(t), server.Client(), server.URL, testSecret, testutil.WaitShort,
+			testDeploymentID, testVersion, prometheus.NewRegistry(),
+		)
+		fill(t, dispatcher.admission, maxAdmissionDispatches)
+		fill(t, dispatcher.semaphore, maxAdmissionDispatches)
+
+		event := newTestEvent(t, agenthooks.EventStop, agenthooks.StopData{})
+		_, _, err := dispatcher.Dispatch(testutil.Context(t, testutil.WaitLong), event)
+		require.NoError(t, err)
+	})
+}