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

@@ -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
 		}
+		func() {
+			defer func() { recover() }()
 			select {
 			case entry.ch <- ev:
 			default:
-			// Subscriber chan full, skip (non-blocking)
 			}
+		}()
 	}
 }
 
-// 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.

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

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
 			}

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)
+// Subscribe creates a new subscription for trace updates (replays all historical events from the beginning).
-func (m *manager) Subscribe() (<-chan *types.TraceUpdate, error) {
+// 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
+// SubscribeFrom creates a subscription starting from a specific timestamp.
-func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error) {
+// 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
 }

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

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)
 	}

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

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)
 	}

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)
+			}
+		})
+	}
+}

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

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 subscribes to trace updates (replay history + real-time).
-	Subscribe() (<-chan *TraceUpdate, error)
+	// Returns the update channel and a cancel function. The caller MUST call
-	// SubscribeFrom subscribes from a specific timestamp (for resume)
+	// cancel when done (e.g., client disconnect) to release the goroutine.
-	SubscribeFrom(since int64) (<-chan *TraceUpdate, error)
+	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