fix(runtime): abandon a concurrency slot won after Shutdown cancelled the tails

concurrent_test.go

@@ -4,8 +4,10 @@ import (
 	"context"
 	"log/slog"
 	"net/http"
+	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -284,3 +286,67 @@ func TestConcurrentConstructionValidation(t *testing.T) {
 		t.Fatal("expected a negative max concurrent bound to fail")
 	}
 }
+
+// TestConcurrentSlotFreedDuringShutdownIsAbandoned pins that a waiter whose
+// slot frees while Shutdown cancels the tails abandons the wait. The tails
+// cancel in random order, so the test repeats the scenario.
+func TestConcurrentSlotFreedDuringShutdownIsAbandoned(t *testing.T) {
+	t.Parallel()
+	if runtime.GOMAXPROCS(0) < 2 {
+		t.Skip("the first handler must run while Shutdown cancels the tails")
+	}
+
+	for i := range 10 {
+		state := newFakeState()
+		adapter := newFakeAdapter("fake")
+		var logs syncBuffer
+		bot := newConcurrentRuntime(t, state, adapter, &logs, func(o *chat.RuntimeOptions) {
+			o.MaxConcurrent = 1
+			o.DetachTimeout = time.Hour
+		})
+
+		children := make([]context.CancelFunc, 0, 20000)
+		firstStarted := make(chan struct{})
+		var waiterRan atomic.Bool
+		bot.OnNewMention(func(ctx context.Context, ev *chat.MessageEvent) error {
+			if ev.Event.ID != "first" {
+				waiterRan.Store(true)
+				return nil
+			}
+			// Shutdown cancels these after it wakes this handler, so the
+			// waiter's context stays live while this handler frees the slot.
+			for range cap(children) {
+				_, cancel := context.WithCancel(ctx)
+				children = append(children, cancel)
+			}
+			close(firstStarted)
+			<-ctx.Done()
+			return ctx.Err()
+		})
+
+		if status := postEvent(t, bot, "fake", mentionEvent("first", "fake:v1:thread-1")); status != http.StatusOK {
+			t.Fatalf("first status = %d", status)
+		}
+		select {
+		case <-firstStarted:
+		case <-time.After(5 * time.Second):
+			t.Fatal("first handler did not start")
+		}
+		if status := postEvent(t, bot, "fake", mentionEvent("waiter", "fake:v1:thread-2")); status != http.StatusOK {
+			t.Fatalf("waiter status = %d", status)
+		}
+		if err := bot.Shutdown(context.Background()); err != nil {
+			t.Fatalf("shutdown error: %v", err)
+		}
+		for _, cancel := range children {
+			cancel()
+		}
+
+		if waiterRan.Load() {
+			t.Fatalf("iteration %d: waiter ran its handler during Shutdown; logs:\n%s", i, logs.String())
+		}
+		if !strings.Contains(logs.String(), "chat concurrent slot wait abandoned") {
+			t.Fatalf("iteration %d: abandoned slot wait not surfaced via observation; logs:\n%s", i, logs.String())
+		}
+	}
+}

runtime.go

@@ -1086,7 +1086,7 @@ func (c *Chat) startDetachedTail(work preludeWork) {
 		defer tailCancel()
 
 		if work.noLock {
-			if !c.acquireConcurrencySlot(tailCtx, work.event) {
+			if !c.acquireDetachedConcurrencySlot(tailCtx, work.event) {
 				c.safeEnd(work.span, OutcomeIgnored, RouteAttr(work.route))
 				return
 			}
@@ -1215,6 +1215,23 @@ func (c *Chat) acquireConcurrencySlot(ctx context.Context, event *Event) bool {
 	}
 }
 
+// acquireDetachedConcurrencySlot is acquireConcurrencySlot for a detached
+// tail. Shutdown cancels the tails one at a time, so a tail can win a slot
+// that a cancelled tail freed before its own context is cancelled. Such a
+// tail releases the slot and abandons the wait.
+func (c *Chat) acquireDetachedConcurrencySlot(ctx context.Context, event *Event) bool {
+	if !c.acquireConcurrencySlot(ctx, event) {
+		return false
+	}
+	err := c.baseCtx.Err()
+	if err == nil {
+		return true
+	}
+	c.releaseConcurrencySlot()
+	c.logger.Info("chat concurrent slot wait abandoned", "adapter", event.Adapter, "event_id", event.ID, "thread_id", event.ThreadID, "error", err)
+	return false
+}
+
 func (c *Chat) releaseConcurrencySlot() {
 	<-c.concurrencySlots
 }