1e4e1224e2
- 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.
91 lines
2.4 KiB
Go
91 lines
2.4 KiB
Go
package trace
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/yaoapp/yao/event"
|
|
eventTypes "github.com/yaoapp/yao/event/types"
|
|
"github.com/yaoapp/yao/trace/types"
|
|
)
|
|
|
|
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).
|
|
// 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.
|
|
// Returns the update channel and a cancel function.
|
|
func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, func(), error) {
|
|
return m.subscribe(since)
|
|
}
|
|
|
|
// subscribe creates a subscription channel that first replays historical
|
|
// 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.
|
|
//
|
|
// 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)
|
|
|
|
liveCh := make(chan *eventTypes.Event, bufferSize)
|
|
traceID := m.traceID
|
|
subID := event.Subscribe("trace.*", liveCh, event.Filter(func(ev *eventTypes.Event) bool {
|
|
update, ok := ev.Payload.(*types.TraceUpdate)
|
|
if !ok {
|
|
return false
|
|
}
|
|
return update.TraceID == traceID
|
|
}))
|
|
|
|
historical := m.stateGetUpdates(since)
|
|
|
|
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 cancel()
|
|
|
|
for _, update := range historical {
|
|
out <- update
|
|
}
|
|
|
|
for ev := range liveCh {
|
|
update, ok := ev.Payload.(*types.TraceUpdate)
|
|
if !ok {
|
|
continue
|
|
}
|
|
key := dedupKey(update)
|
|
if _, dup := histSeen[key]; dup {
|
|
delete(histSeen, key)
|
|
continue
|
|
}
|
|
out <- update
|
|
if update.Type == types.UpdateTypeComplete {
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
return out, cancel, nil
|
|
}
|