- 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.
75 lines
2 KiB
Go
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
|
|
}
|
|
}
|
|
}
|