git.iko.soy/crush
on main help sign in
fix(server): close a queued-prompt cancel race
d7a814cChristian Rocha
internal/agent/agent.go+26−9
compare with main
unifiedsplit
@@ -359,6 +359,30 @@ func (a *sessionAgent) enqueueCall(call SessionAgentCall) {
359 a.messageQueue.Set(call.SessionID, existing)
359 a.messageQueue.Set(call.SessionID, existing)
360}
360}
361
361
362+// drainUncanceledQueue copies and clears the session's message queue and
363+// returns the queued calls not covered by a pending cancel. The queue
364+// copy/delete and the cancel-mark check happen under the per-session
365+// dispatch mutex so the filtering is atomic against a concurrent Cancel:
366+// canceledBySeq requires the caller to hold that mutex, and evaluating it
367+// here (rather than after unlocking) prevents a cancel recorded between
368+// the drain and the check from being observed inconsistently. The
369+// returned survivors are processed by the caller without the lock held.
370+func (a *sessionAgent) drainUncanceledQueue(sessionID string) []SessionAgentCall {
371+ dispatchLock := a.sessionMu(sessionID)
372+ dispatchLock.Lock()
373+ defer dispatchLock.Unlock()
374+ queuedCalls, _ := a.messageQueue.Get(sessionID)
375+ a.messageQueue.Del(sessionID)
376+ survivors := queuedCalls[:0]
377+ for _, queued := range queuedCalls {
378+ if a.canceledBySeq(sessionID, queued.acceptSeq) {
379+ continue
380+ }
381+ survivors = append(survivors, queued)
382+ }
383+ return survivors
384+}
385+
362// clearPendingCancel removes any pending-cancel mark for sessionID. It
386// clearPendingCancel removes any pending-cancel mark for sessionID. It
363// takes the per-session dispatch lock so it is ordered against Cancel
387// takes the per-session dispatch lock so it is ordered against Cancel
364// and the dispatch handoff.
388// and the dispatch handoff.
@@ -717,15 +741,8 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *
717 // part of this step. Coverage is per-call by accept sequence
741 // part of this step. Coverage is per-call by accept sequence
718 // so a follow-up queued after the cancel (higher seq) is
742 // so a follow-up queued after the cancel (higher seq) is
719 // still folded in.
743 // still folded in.
720− dispatchLock := a.sessionMu(call.SessionID)
744+ survivors := a.drainUncanceledQueue(call.SessionID)
721− dispatchLock.Lock()
722− queuedCalls, _ := a.messageQueue.Get(call.SessionID)
723− a.messageQueue.Del(call.SessionID)
724− dispatchLock.Unlock()
725− for _, queued := range queuedCalls {
745+ for _, queued := range survivors {
726− if a.canceledBySeq(call.SessionID, queued.acceptSeq) {
727− continue
728− }
729 userMessage, createErr := a.createUserMessage(callContext, queued)
746 userMessage, createErr := a.createUserMessage(callContext, queued)
730 if createErr != nil {
747 if createErr != nil {
731 return callContext, prepared, createErr
748 return callContext, prepared, createErr
internal/agent/run_complete_test.go+56
@@ -58,6 +58,62 @@ func TestSessionAgentRun_QueueStripsOnComplete(t *testing.T) {
58 "RunComplete still correlates with the originating SendMessage")
58 "RunComplete still correlates with the originating SendMessage")
59}
59}
60
60
61+// TestDrainUncanceledQueue_FiltersUnderDispatchLock verifies that the
62+// queue drain evaluates the per-session cancel mark while holding the
63+// dispatch mutex (canceledBySeq's documented precondition). Queued calls
64+// at or below the cancel high-water mark are dropped, calls queued after
65+// the cancel (higher seq) survive, untracked enqueues (seq == 0) are
66+// dropped whenever any mark is present, and the queue is cleared.
67+func TestDrainUncanceledQueue_FiltersUnderDispatchLock(t *testing.T) {
68+ t.Parallel()
69+
70+ env := testEnv(t)
71+ a := NewSessionAgent(SessionAgentOptions{
72+ Sessions: env.sessions,
73+ Messages: env.messages,
74+ }).(*sessionAgent)
75+
76+ const sessionID = "drain-session"
77+ a.messageQueue.Set(sessionID, []SessionAgentCall{
78+ {SessionID: sessionID, Prompt: "below", acceptSeq: 1},
79+ {SessionID: sessionID, Prompt: "at-mark", acceptSeq: 2},
80+ {SessionID: sessionID, Prompt: "after", acceptSeq: 3},
81+ {SessionID: sessionID, Prompt: "untracked", acceptSeq: 0},
82+ })
83+ // Cancel high-water mark at seq 2: seq <= 2 and seq == 0 are covered.
84+ a.cancelMark.Set(sessionID, 2)
85+
86+ survivors := a.drainUncanceledQueue(sessionID)
87+
88+ require.Len(t, survivors, 1,
89+ "only the follow-up queued after the cancel (seq > mark) must survive")
90+ require.Equal(t, "after", survivors[0].Prompt)
91+
92+ _, ok := a.messageQueue.Get(sessionID)
93+ require.False(t, ok, "drain must clear the session message queue")
94+}
95+
96+// TestDrainUncanceledQueue_NoMarkKeepsAll verifies that with no cancel
97+// mark recorded, every queued call survives the drain.
98+func TestDrainUncanceledQueue_NoMarkKeepsAll(t *testing.T) {
99+ t.Parallel()
100+
101+ env := testEnv(t)
102+ a := NewSessionAgent(SessionAgentOptions{
103+ Sessions: env.sessions,
104+ Messages: env.messages,
105+ }).(*sessionAgent)
106+
107+ const sessionID = "drain-nomark"
108+ a.messageQueue.Set(sessionID, []SessionAgentCall{
109+ {SessionID: sessionID, Prompt: "a", acceptSeq: 0},
110+ {SessionID: sessionID, Prompt: "b", acceptSeq: 5},
111+ })
112+
113+ survivors := a.drainUncanceledQueue(sessionID)
114+ require.Len(t, survivors, 2, "no cancel mark means all queued calls survive")
115+}
116+
61// TestRunCompletePublisher_MustDeliverOverTakesPublish exercises the
117// TestRunCompletePublisher_MustDeliverOverTakesPublish exercises the
62// pubsub.Publisher interface change end-to-end: a Broker is the only
118// pubsub.Publisher interface change end-to-end: a Broker is the only
63// concrete Publisher implementation and must satisfy both Publish and
119// concrete Publisher implementation and must satisfy both Publish and