diff --git a/server/internal/runstream/broker.go b/server/internal/runstream/broker.go index 9071c0b9..91d0f799 100644 --- a/server/internal/runstream/broker.go +++ b/server/internal/runstream/broker.go @@ -10,7 +10,7 @@ import ( const DefaultBufferSize = 64 type Broker struct { - mu sync.Mutex + mu sync.Mutex // serializes subscriber sends and closes as well as run state bufferSize int runs map[string]*runState } @@ -33,21 +33,16 @@ func (b *Broker) Publish(runID string, ev connector.PromptEvent) { return } b.mu.Lock() + defer b.mu.Unlock() st := b.stateLocked(runID) if st.closed { - b.mu.Unlock() return } st.events = append(st.events, ev) if len(st.events) > b.bufferSize { st.events = append([]connector.PromptEvent(nil), st.events[len(st.events)-b.bufferSize:]...) } - subs := make([]chan connector.PromptEvent, 0, len(st.subscribers)) for ch := range st.subscribers { - subs = append(subs, ch) - } - b.mu.Unlock() - for _, ch := range subs { select { case ch <- ev: default: @@ -63,30 +58,19 @@ func (b *Broker) Subscribe(ctx context.Context, runID string) <-chan connector.P } b.mu.Lock() st := b.stateLocked(runID) - replay := append([]connector.PromptEvent(nil), st.events...) - closed := st.closed - if !closed { - st.subscribers[out] = struct{}{} + for _, ev := range st.events { + out <- ev } + if st.closed { + close(out) + b.mu.Unlock() + return out + } + st.subscribers[out] = struct{}{} b.mu.Unlock() go func() { - defer func() { - if !closed { - b.unsubscribe(runID, out) - } - }() - for _, ev := range replay { - select { - case <-ctx.Done(): - return - case out <- ev: - } - } - if closed { - close(out) - return - } <-ctx.Done() + b.unsubscribe(runID, out) }() return out } @@ -96,19 +80,14 @@ func (b *Broker) Finish(runID string) { return } b.mu.Lock() + defer b.mu.Unlock() st, ok := b.runs[runID] if !ok || st.closed { - b.mu.Unlock() return } st.closed = true - subs := make([]chan connector.PromptEvent, 0, len(st.subscribers)) for ch := range st.subscribers { - subs = append(subs, ch) delete(st.subscribers, ch) - } - b.mu.Unlock() - for _, ch := range subs { close(ch) } } diff --git a/server/internal/runstream/broker_test.go b/server/internal/runstream/broker_test.go index fa092330..ca38e711 100644 --- a/server/internal/runstream/broker_test.go +++ b/server/internal/runstream/broker_test.go @@ -2,6 +2,7 @@ package runstream import ( "context" + "runtime" "testing" "time" @@ -42,6 +43,36 @@ func TestCancelCleansUp(t *testing.T) { } } +func TestPublishConcurrentWithCancelDoesNotPanic(t *testing.T) { + b := NewBroker(1) + ctx, cancel := context.WithCancel(context.Background()) + const runID = "run-publish-cancel-race" + const subscriberCount = 4096 + for range subscriberCount { + _ = b.Subscribe(ctx, runID) + } + + start := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer close(done) + <-start + b.Publish(runID, connector.PromptEvent{Type: connector.EventDelta, Delta: "x"}) + }() + close(start) + runtime.Gosched() + cancel() + <-done + + deadline := time.Now().Add(5 * time.Second) + for b.SubscriberCount(runID) != 0 && time.Now().Before(deadline) { + runtime.Gosched() + } + if got := b.SubscriberCount(runID); got != 0 { + t.Fatalf("subscribers after cancel = %d, want 0", got) + } +} + func TestMultiSubscriber(t *testing.T) { b := NewBroker(64) ctx, cancel := context.WithCancel(context.Background())