From 1e4e1224e2f9364eedb10fe9e329d7777a3fdaa2 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 24 Feb 2026 00:22:08 +0800 Subject: [PATCH] Enhance subscription management by adding cancel functions and improving channel handling - Update the Subscribe and SubscribeFrom methods to return a cancel function, ensuring proper resource cleanup when subscriptions are no longer needed. - Modify unsubscribe logic to close channels safely, preventing potential panics from sending on closed channels. - Enhance test cases to utilize the new cancel functionality, ensuring robust handling of subscriptions in various scenarios. --- event/sub.go | 41 +++++-- event/sub_test.go | 6 +- openapi/trace/events.go | 7 +- trace/subscription.go | 32 ++++-- trace/trace_autocomplete_test.go | 3 +- trace/trace_bench_test.go | 3 +- trace/trace_concurrent_test.go | 11 +- trace/trace_mem_test.go | 6 +- trace/trace_subscription_leak_test.go | 157 ++++++++++++++++++++++++++ trace/trace_subscription_test.go | 15 ++- trace/types/interfaces.go | 11 +- 11 files changed, 252 insertions(+), 40 deletions(-) create mode 100644 trace/trace_subscription_leak_test.go diff --git a/event/sub.go b/event/sub.go index 089ad94d..9e6815c1 100644 --- a/event/sub.go +++ b/event/sub.go @@ -54,14 +54,25 @@ func (sm *subManager) subscribe(pattern string, ch chan<- *types.Event, opts ... return id } -// unsubscribe removes a subscriber by ID. +// unsubscribe removes a subscriber by ID and closes its channel +// so that any goroutine blocked on `range ch` will unblock and exit. func (sm *subManager) unsubscribe(id string) { sm.mu.Lock() - defer sm.mu.Unlock() + entry, ok := sm.entries[id] delete(sm.entries, id) + sm.mu.Unlock() + + if ok && entry.ch != nil { + func() { + defer func() { recover() }() + close(entry.ch) + }() + } } // notify sends an event to all matching subscribers (non-blocking). +// Recovers from send-on-closed-channel panics that may occur if +// unsubscribe closes a channel concurrently. func (sm *subManager) notify(ev *types.Event) { sm.mu.RLock() defer sm.mu.RUnlock() @@ -73,19 +84,31 @@ func (sm *subManager) notify(ev *types.Event) { if entry.filter != nil && !entry.filter(ev) { continue } - select { - case entry.ch <- ev: - default: - // Subscriber chan full, skip (non-blocking) - } + func() { + defer func() { recover() }() + select { + case entry.ch <- ev: + default: + } + }() } } -// clear removes all subscribers. Used during Stop. +// clear removes all subscribers and closes their channels. Used during Stop. func (sm *subManager) clear() { sm.mu.Lock() - defer sm.mu.Unlock() + old := sm.entries sm.entries = make(map[string]*subEntry) + sm.mu.Unlock() + + for _, entry := range old { + if entry.ch != nil { + func() { + defer func() { recover() }() + close(entry.ch) + }() + } + } } // Subscribe dynamically subscribes to events matching the given pattern. diff --git a/event/sub_test.go b/event/sub_test.go index abdde4f9..b51d93ba 100644 --- a/event/sub_test.go +++ b/event/sub_test.go @@ -164,6 +164,7 @@ func TestSubscribe_StopClearsSubscribers(t *testing.T) { } // drainChan reads up to n events from ch within timeout. +// Stops early if the channel is closed. func drainChan(ch chan *types.Event, n int, timeout time.Duration) []*types.Event { var result []*types.Event timer := time.NewTimer(timeout) @@ -171,7 +172,10 @@ func drainChan(ch chan *types.Event, n int, timeout time.Duration) []*types.Even for range n { select { - case ev := <-ch: + case ev, ok := <-ch: + if !ok { + return result + } result = append(result, ev) case <-timer.C: return result diff --git a/openapi/trace/events.go b/openapi/trace/events.go index a869e2aa..a4f74e59 100644 --- a/openapi/trace/events.go +++ b/openapi/trace/events.go @@ -58,13 +58,14 @@ func handleStreamMode(c *gin.Context, manager types.Manager, info *types.TraceIn c.Header("X-Accel-Buffering", "no") // Subscribe to trace updates - updates, err := manager.Subscribe() + updates, cancel, err := manager.Subscribe() if err != nil { // Send error as SSE event fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"Failed to subscribe: %s\"}\n\n", err.Error()) c.Writer.Flush() return } + defer cancel() // Stream events ctx := c.Request.Context() @@ -73,22 +74,18 @@ func handleStreamMode(c *gin.Context, manager types.Manager, info *types.TraceIn for { select { case <-clientGone: - // Client disconnected return case update, ok := <-updates: if !ok { - // Channel closed return } - // Format and send SSE event err := sendSSEEvent(c.Writer, *update) if err != nil { return } - // Check if trace is complete if update.Type == types.UpdateTypeComplete { return } diff --git a/trace/subscription.go b/trace/subscription.go index ac76cf12..3d1f53f8 100644 --- a/trace/subscription.go +++ b/trace/subscription.go @@ -2,6 +2,7 @@ package trace import ( "fmt" + "sync" "github.com/yaoapp/yao/event" eventTypes "github.com/yaoapp/yao/event/types" @@ -12,13 +13,16 @@ func dedupKey(u *types.TraceUpdate) string { return fmt.Sprintf("%s:%s:%d", u.Type, u.NodeID, u.Timestamp) } -// Subscribe creates a new subscription for trace updates (replays all historical events from the beginning) -func (m *manager) Subscribe() (<-chan *types.TraceUpdate, error) { +// Subscribe creates a new subscription for trace updates (replays all historical events from the beginning). +// Returns the update channel and a cancel function. The caller MUST call +// cancel when done (e.g., client disconnect) to release the goroutine. +func (m *manager) Subscribe() (<-chan *types.TraceUpdate, func(), error) { return m.subscribe(0) } -// SubscribeFrom creates a subscription starting from a specific timestamp -func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error) { +// SubscribeFrom creates a subscription starting from a specific timestamp. +// Returns the update channel and a cancel function. +func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, func(), error) { return m.subscribe(since) } @@ -26,12 +30,14 @@ func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error) // updates, then streams live events via the event service's Subscriber. // The subscriber is registered BEFORE reading historical state to prevent // missing events that occur between the state snapshot and subscriber setup. -func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) { +// +// The returned cancel function triggers event.Unsubscribe which closes +// liveCh, causing the goroutine to exit via `for range liveCh`. +func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, func(), error) { bufferSize := 1000 out := make(chan *types.TraceUpdate, bufferSize) - // Register live subscriber FIRST to avoid missing events between snapshot and subscribe. liveCh := make(chan *eventTypes.Event, bufferSize) traceID := m.traceID subID := event.Subscribe("trace.*", liveCh, event.Filter(func(ev *eventTypes.Event) bool { @@ -42,19 +48,23 @@ func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) { return update.TraceID == traceID })) - // THEN snapshot historical updates (may overlap with live events). historical := m.stateGetUpdates(since) - // Build a set of historical event identifiers for dedup. - // Key: "type:nodeID:timestamp" is unique enough for trace events. histSeen := make(map[string]struct{}, len(historical)) for _, u := range historical { histSeen[dedupKey(u)] = struct{}{} } + var cancelOnce sync.Once + cancel := func() { + cancelOnce.Do(func() { + event.Unsubscribe(subID) + }) + } + go func() { defer close(out) - defer event.Unsubscribe(subID) + defer cancel() for _, update := range historical { out <- update @@ -77,5 +87,5 @@ func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) { } }() - return out, nil + return out, cancel, nil } diff --git a/trace/trace_autocomplete_test.go b/trace/trace_autocomplete_test.go index f7d08855..d746d8e1 100644 --- a/trace/trace_autocomplete_test.go +++ b/trace/trace_autocomplete_test.go @@ -331,8 +331,9 @@ func TestAutoCompleteParentEvents(t *testing.T) { defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...) // Subscribe to updates - updates, err := manager.Subscribe() + updates, cancel, err := manager.Subscribe() assert.NoError(t, err) + defer cancel() // Collect updates in background var receivedUpdates []*types.TraceUpdate diff --git a/trace/trace_bench_test.go b/trace/trace_bench_test.go index 319fc8b3..afb74d94 100644 --- a/trace/trace_bench_test.go +++ b/trace/trace_bench_test.go @@ -266,7 +266,7 @@ func BenchmarkSubscription(b *testing.B) { } // Subscribe - updates, err := manager.Subscribe() + updates, cancel, err := manager.Subscribe() if err != nil { b.Fatalf("Failed to subscribe: %s", err.Error()) } @@ -301,6 +301,7 @@ func BenchmarkSubscription(b *testing.B) { } } + cancel() trace.Release(traceID) trace.Remove(ctx, trace.Local, traceID) } diff --git a/trace/trace_concurrent_test.go b/trace/trace_concurrent_test.go index 97e3c994..1c96547f 100644 --- a/trace/trace_concurrent_test.go +++ b/trace/trace_concurrent_test.go @@ -179,6 +179,7 @@ func TestConcurrentSubscribers(t *testing.T) { var wg sync.WaitGroup numSubscribers := 5 subscribers := make([]<-chan *types.TraceUpdate, numSubscribers) + cancels := make([]func(), numSubscribers) var mu sync.Mutex for i := 0; i < numSubscribers; i++ { @@ -186,14 +187,22 @@ func TestConcurrentSubscribers(t *testing.T) { go func(idx int) { defer wg.Done() - sub, err := manager.Subscribe() + sub, cancelSub, err := manager.Subscribe() assert.NoError(t, err) mu.Lock() subscribers[idx] = sub + cancels[idx] = cancelSub mu.Unlock() }(i) } + defer func() { + for _, c := range cancels { + if c != nil { + c() + } + } + }() wg.Wait() // Verify all subscriptions were created diff --git a/trace/trace_mem_test.go b/trace/trace_mem_test.go index f4f227a5..77e38604 100644 --- a/trace/trace_mem_test.go +++ b/trace/trace_mem_test.go @@ -259,10 +259,11 @@ func TestMemoryLeakComplexScenarios(t *testing.T) { { name: "WithSubscription", execute: func(m types.Manager) error { - updates, err := m.Subscribe() + updates, cancel, err := m.Subscribe() if err != nil { return err } + defer cancel() // Drain updates in background with timeout done := make(chan bool) @@ -563,7 +564,7 @@ func TestGoroutineLeak(t *testing.T) { } // Subscribe (creates goroutines) - updates, err := manager.Subscribe() + updates, cancel, err := manager.Subscribe() if err != nil { t.Errorf("Subscribe failed at iteration %d: %s", i, err.Error()) } @@ -598,6 +599,7 @@ func TestGoroutineLeak(t *testing.T) { } } + cancel() trace.Release(traceID) trace.Remove(ctx, trace.Local, traceID) } diff --git a/trace/trace_subscription_leak_test.go b/trace/trace_subscription_leak_test.go new file mode 100644 index 00000000..adcde0fc --- /dev/null +++ b/trace/trace_subscription_leak_test.go @@ -0,0 +1,157 @@ +package trace_test + +import ( + "context" + "runtime" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/event" + "github.com/yaoapp/yao/trace" + "github.com/yaoapp/yao/trace/types" +) + +// stableGoroutines waits for runtime to settle and returns goroutine count. +func stableGoroutines() int { + for i := 0; i < 5; i++ { + runtime.GC() + runtime.Gosched() + time.Sleep(10 * time.Millisecond) + } + return runtime.NumGoroutine() +} + +// TestLeak_SubscriptionClientDisconnect reproduces the goroutine leak that +// occurs when an SSE client subscribes to a trace and then disconnects +// without the trace ever completing (no UpdateTypeComplete sent). +// +// The subscription goroutine in subscription.go blocks on +// `for ev := range liveCh` and never exits because: +// 1. liveCh is never closed (event.Unsubscribe only deletes the map entry) +// 2. The goroutine only returns on UpdateTypeComplete +// 3. No context/cancellation mechanism exists +// +// This simulates the real-world scenario: SSE handler returns on client +// disconnect, but the subscription goroutine keeps running forever. +func TestLeak_SubscriptionClientDisconnect(t *testing.T) { + drivers := trace.GetTestDrivers() + + for _, d := range drivers { + t.Run(d.Name, func(t *testing.T) { + ctx := context.Background() + + before := stableGoroutines() + + const numClients = 10 + + for i := 0; i < numClients; i++ { + traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...) + assert.NoError(t, err) + + // Client subscribes (like SSE handler calling manager.Subscribe()) + updates, cancel, err := manager.Subscribe() + assert.NoError(t, err) + assert.NotNil(t, updates) + + // Simulate some trace activity + _, err = manager.Add("step", types.TraceNodeOption{Label: "Processing"}) + assert.NoError(t, err) + + // Read a couple of events (like the SSE handler would) + timeout := time.After(500 * time.Millisecond) + drain: + for { + select { + case _, ok := <-updates: + if !ok { + break drain + } + case <-timeout: + break drain + } + } + + // Client disconnects: SSE handler calls cancel (deferred). + // This triggers event.Unsubscribe which closes liveCh, + // allowing the subscription goroutine to exit. + cancel() + + trace.Release(traceID) + } + + // Wait for goroutines to settle + time.Sleep(1 * time.Second) + after := stableGoroutines() + + leaked := after - before + t.Logf("goroutines: before=%d after=%d leaked=%d (over %d simulated client disconnects)", before, after, leaked, numClients) + + // Each Subscribe() spawns a goroutine that should eventually exit. + // If it doesn't, we'll see roughly numClients leaked goroutines. + if leaked >= numClients { + t.Errorf("goroutine leak detected: %d goroutines leaked after %d client disconnects. "+ + "Subscription goroutines are not cleaned up when clients disconnect without trace completion.", + leaked, numClients) + } + }) + } +} + +// TestLeak_SubscriptionEventServiceStop reproduces the goroutine leak when +// event.Stop() is called (e.g., during shutdown) while subscriptions are active. +// +// event.Stop() calls smgr.clear() which deletes all subscriber entries but +// does NOT close their channels, leaving goroutines blocked on `range liveCh`. +func TestLeak_SubscriptionEventServiceStop(t *testing.T) { + drivers := trace.GetTestDrivers() + + for _, d := range drivers { + t.Run(d.Name, func(t *testing.T) { + ctx := context.Background() + + before := stableGoroutines() + + const numSubs = 5 + + traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...) + assert.NoError(t, err) + + // Create multiple subscriptions (simulating multiple SSE clients) + for i := 0; i < numSubs; i++ { + _, _, err := manager.Subscribe() + assert.NoError(t, err) + } + + // Simulate some activity + _, err = manager.Add("work", types.TraceNodeOption{Label: "Working"}) + assert.NoError(t, err) + + time.Sleep(200 * time.Millisecond) + + // Stop event service (like during server shutdown) + err = event.Stop(ctx) + assert.NoError(t, err) + + // Restart event service for other tests + err = event.Start() + if err != nil && err != event.ErrAlreadyStart { + t.Fatalf("Failed to restart event service: %v", err) + } + + trace.Release(traceID) + + time.Sleep(1 * time.Second) + after := stableGoroutines() + + leaked := after - before + t.Logf("goroutines: before=%d after=%d leaked=%d (over %d subscriptions + event.Stop)", before, after, leaked, numSubs) + + if leaked >= numSubs { + t.Errorf("goroutine leak detected: %d goroutines leaked after event.Stop() with %d active subscriptions. "+ + "smgr.clear() does not close subscriber channels.", + leaked, numSubs) + } + }) + } +} diff --git a/trace/trace_subscription_test.go b/trace/trace_subscription_test.go index 77d4b81d..5a1283c4 100644 --- a/trace/trace_subscription_test.go +++ b/trace/trace_subscription_test.go @@ -24,9 +24,10 @@ func TestSubscription(t *testing.T) { defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...) // Subscribe to updates - updates, err := manager.Subscribe() + updates, cancel, err := manager.Subscribe() assert.NoError(t, err) assert.NotNil(t, updates) + defer cancel() // Collect updates in background var receivedUpdates []*types.TraceUpdate @@ -146,9 +147,10 @@ func TestSubscribeFrom(t *testing.T) { // Real scenario: User refreshes page and resumes from last known timestamp // This should replay events from resumeTimestamp onwards - updates, err := manager.SubscribeFrom(resumeTimestamp) + updates, cancelSub, err := manager.SubscribeFrom(resumeTimestamp) assert.NoError(t, err) assert.NotNil(t, updates) + defer cancelSub() // Collect updates var receivedUpdates []*types.TraceUpdate @@ -233,14 +235,17 @@ func TestMultipleSubscribers(t *testing.T) { defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...) // Create multiple subscribers - sub1, err := manager.Subscribe() + sub1, cancel1, err := manager.Subscribe() assert.NoError(t, err) + defer cancel1() - sub2, err := manager.Subscribe() + sub2, cancel2, err := manager.Subscribe() assert.NoError(t, err) + defer cancel2() - sub3, err := manager.Subscribe() + sub3, cancel3, err := manager.Subscribe() assert.NoError(t, err) + defer cancel3() // Collect updates from all subscribers var wg sync.WaitGroup diff --git a/trace/types/interfaces.go b/trace/types/interfaces.go index 125c0a16..f4d7438c 100644 --- a/trace/types/interfaces.go +++ b/trace/types/interfaces.go @@ -48,10 +48,13 @@ type Manager interface { MarkComplete() error // Subscription Operations - // Subscribe subscribes to trace updates (replay history + real-time) - Subscribe() (<-chan *TraceUpdate, error) - // SubscribeFrom subscribes from a specific timestamp (for resume) - SubscribeFrom(since int64) (<-chan *TraceUpdate, error) + // Subscribe subscribes to trace updates (replay history + real-time). + // Returns the update channel and a cancel function. The caller MUST call + // cancel when done (e.g., client disconnect) to release the goroutine. + Subscribe() (<-chan *TraceUpdate, func(), error) + // SubscribeFrom subscribes from a specific timestamp (for resume). + // Returns the update channel and a cancel function. + SubscribeFrom(since int64) (<-chan *TraceUpdate, func(), error) // IsComplete checks if the trace is completed IsComplete() bool