yao/trace/subscription.go
Max 47a35b7e63 Refactor trace manager for improved state management and concurrency
- Replaced the previous node and space management with a channel-based state management system in `manager.go`, enhancing concurrency handling.
- Implemented methods for loading and saving trace updates to disk in `local/driver.go` and `store/driver.go`, allowing for persistent state across sessions.
- Updated subscription handling in `subscription.go` to streamline the process of broadcasting updates to active subscribers.
- Enhanced node and trace status management, including cancellation and completion states, to provide better control over trace execution.
- Adjusted related tests to ensure compatibility with the new state management approach.
2025-11-18 18:08:20 +08:00

75 lines
2 KiB
Go

package trace
import (
"time"
gonanoid "github.com/matoous/go-nanoid/v2"
"github.com/yaoapp/yao/trace/types"
)
// Subscribe creates a new subscription for trace updates (real-time from now)
func (m *manager) Subscribe() (<-chan *types.TraceUpdate, error) {
return m.subscribe(time.Now().Unix())
}
// SubscribeFrom creates a subscription starting from a specific timestamp
func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error) {
return m.subscribe(since)
}
// subscribe is the internal implementation for subscriptions
func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) {
// Generate unique subscriber ID
subID, _ := gonanoid.Generate("0123456789abcdefghijklmnopqrstuvwxyz", 12)
// Create update channel
updateCh := make(chan *types.TraceUpdate, 100)
// Register subscriber
m.stateAddSubscriber(subID, updateCh)
// Start replay and stream goroutine (will auto-cleanup on completion)
go m.replayAndStream(subID, updateCh, since)
return updateCh, nil
}
// replayAndStream replays historical updates and streams new ones
func (m *manager) replayAndStream(subID string, ch chan *types.TraceUpdate, since int64) {
// Auto-cleanup on exit - MUST remove from map before closing channel
defer func() {
// Remove from subscribers map first to prevent new broadcasts
m.stateRemoveSubscriber(subID)
// Close channel (any in-flight broadcasts will be caught by recover)
close(ch)
}()
// Get historical updates
updates := m.stateGetUpdates(since)
// Replay historical updates
for _, update := range updates {
select {
case ch <- update:
case <-m.ctx.Done():
return
}
}
// Continue streaming new updates
// The channel will receive updates via broadcast from addUpdate
// Monitor completion to know when to exit
ticker := time.NewTicker(100 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-ticker.C:
if m.stateIsCompleted() {
return
}
case <-m.ctx.Done():
return
}
}
}