package trace import ( "time" "github.com/yaoapp/yao/trace/types" ) // State management using channel-based serialization (no locks needed) // All state mutations go through a single worker goroutine // managerState holds all mutable state (accessed only by state worker) type managerState struct { rootNode *types.TraceNode currentNodes []*types.TraceNode spaces map[string]*types.TraceSpace traceStatus types.TraceStatus completed bool updates []*types.TraceUpdate subscribers map[string]chan *types.TraceUpdate } // State command interface - all commands are processed serially type stateCommand interface { execute(s *managerState) } // Commands with response channels for synchronous operations // --- Root Node Commands --- type cmdSetRoot struct { node *types.TraceNode } func (c *cmdSetRoot) execute(s *managerState) { s.rootNode = c.node } type cmdGetRoot struct { resp chan *types.TraceNode } func (c *cmdGetRoot) execute(s *managerState) { c.resp <- s.rootNode } // --- Current Nodes Commands --- type cmdSetCurrentNodes struct { nodes []*types.TraceNode } func (c *cmdSetCurrentNodes) execute(s *managerState) { s.currentNodes = c.nodes } type cmdGetCurrentNodes struct { resp chan []*types.TraceNode } func (c *cmdGetCurrentNodes) execute(s *managerState) { // Return a copy to prevent external mutation nodes := make([]*types.TraceNode, len(s.currentNodes)) copy(nodes, s.currentNodes) c.resp <- nodes } type cmdUpdateRootAndCurrent struct { root *types.TraceNode current []*types.TraceNode } func (c *cmdUpdateRootAndCurrent) execute(s *managerState) { s.rootNode = c.root s.currentNodes = c.current } // --- Space Commands --- type cmdGetSpace struct { id string resp chan *types.TraceSpace } func (c *cmdGetSpace) execute(s *managerState) { c.resp <- s.spaces[c.id] } type cmdSetSpace struct { id string space *types.TraceSpace } func (c *cmdSetSpace) execute(s *managerState) { s.spaces[c.id] = c.space } type cmdDeleteSpace struct { id string } func (c *cmdDeleteSpace) execute(s *managerState) { delete(s.spaces, c.id) } type cmdGetAllSpaces struct { resp chan []*types.TraceSpace } func (c *cmdGetAllSpaces) execute(s *managerState) { spaces := make([]*types.TraceSpace, 0, len(s.spaces)) for _, space := range s.spaces { spaces = append(spaces, space) } c.resp <- spaces } // --- Trace Status Commands --- type cmdSetTraceStatus struct { status types.TraceStatus } func (c *cmdSetTraceStatus) execute(s *managerState) { s.traceStatus = c.status } type cmdGetTraceStatus struct { resp chan types.TraceStatus } func (c *cmdGetTraceStatus) execute(s *managerState) { c.resp <- s.traceStatus } // --- Completion Commands --- type cmdMarkCompleted struct { resp chan bool // Returns true if marked, false if already completed } func (c *cmdMarkCompleted) execute(s *managerState) { if s.completed { c.resp <- false } else { s.completed = true c.resp <- true } } type cmdIsCompleted struct { resp chan bool } func (c *cmdIsCompleted) execute(s *managerState) { c.resp <- s.completed } // --- Update Commands --- type cmdAddUpdate struct { update *types.TraceUpdate } func (c *cmdAddUpdate) execute(s *managerState) { s.updates = append(s.updates, c.update) } type cmdGetUpdates struct { since int64 resp chan []*types.TraceUpdate } func (c *cmdGetUpdates) execute(s *managerState) { filtered := make([]*types.TraceUpdate, 0) for _, update := range s.updates { if update.Timestamp >= c.since { filtered = append(filtered, update) } } c.resp <- filtered } type cmdSetUpdates struct { updates []*types.TraceUpdate } func (c *cmdSetUpdates) execute(s *managerState) { s.updates = c.updates } // --- Subscriber Commands --- type cmdAddSubscriber struct { id string ch chan *types.TraceUpdate } func (c *cmdAddSubscriber) execute(s *managerState) { s.subscribers[c.id] = c.ch } type cmdRemoveSubscriber struct { id string } func (c *cmdRemoveSubscriber) execute(s *managerState) { delete(s.subscribers, c.id) } type cmdGetSubscribers struct { resp chan map[string]chan *types.TraceUpdate } func (c *cmdGetSubscribers) execute(s *managerState) { // Return a copy of the map subs := make(map[string]chan *types.TraceUpdate, len(s.subscribers)) for id, ch := range s.subscribers { subs[id] = ch } c.resp <- subs } // --- Broadcast Command (special - sends to all subscribers) --- type cmdBroadcast struct { update *types.TraceUpdate } func (c *cmdBroadcast) execute(s *managerState) { // Send to all subscribers (non-blocking, with panic recovery) for _, ch := range s.subscribers { func(channel chan *types.TraceUpdate) { defer func() { // Recover from panic if channel is closed if r := recover(); r != nil { // Channel was closed, ignore } }() select { case channel <- c.update: default: // Subscriber is slow, skip (non-blocking) } }(ch) } } // --- Space KV Commands (for concurrent safety) --- // These ensure all operations on a space are serialized through state worker type cmdSpaceKVOp struct { spaceID string fn func() error resp chan error } func (c *cmdSpaceKVOp) execute(s *managerState) { // Execute the operation (typically a driver call) // The function is provided by caller and executed serially here err := c.fn() c.resp <- err } // State worker - processes all commands serially in a single goroutine func (m *manager) startStateWorker() { // Initialize state state := &managerState{ rootNode: nil, currentNodes: []*types.TraceNode{}, spaces: make(map[string]*types.TraceSpace), traceStatus: types.TraceStatusPending, completed: false, updates: make([]*types.TraceUpdate, 0, 100), subscribers: make(map[string]chan *types.TraceUpdate), } // Process commands until context is cancelled or trace is completed for { select { case cmd, ok := <-m.stateCmdChan: if !ok { // Channel closed return } cmd.execute(state) // Exit after processing completion if state.completed { // Drain remaining commands with timeout drainTimer := time.NewTimer(100 * time.Millisecond) defer drainTimer.Stop() drainLoop: for { select { case cmd := <-m.stateCmdChan: cmd.execute(state) case <-drainTimer.C: break drainLoop } } return } case <-m.ctx.Done(): // Context cancelled - continue processing for a short time to handle cancellation // Then exit to prevent deadlock time.Sleep(10 * time.Millisecond) return } } } // Helper methods for manager to send commands func (m *manager) stateSetRoot(node *types.TraceNode) { m.stateCmdChan <- &cmdSetRoot{node: node} } func (m *manager) stateGetRoot() *types.TraceNode { resp := make(chan *types.TraceNode, 1) m.stateCmdChan <- &cmdGetRoot{resp: resp} return <-resp } func (m *manager) stateSetCurrentNodes(nodes []*types.TraceNode) { m.stateCmdChan <- &cmdSetCurrentNodes{nodes: nodes} } func (m *manager) stateGetCurrentNodes() []*types.TraceNode { resp := make(chan []*types.TraceNode, 1) m.stateCmdChan <- &cmdGetCurrentNodes{resp: resp} return <-resp } func (m *manager) stateUpdateRootAndCurrent(root *types.TraceNode, current []*types.TraceNode) { m.stateCmdChan <- &cmdUpdateRootAndCurrent{root: root, current: current} } func (m *manager) stateGetSpace(id string) (*types.TraceSpace, bool) { resp := make(chan *types.TraceSpace, 1) m.stateCmdChan <- &cmdGetSpace{id: id, resp: resp} space := <-resp return space, space != nil } func (m *manager) stateSetSpace(id string, space *types.TraceSpace) { m.stateCmdChan <- &cmdSetSpace{id: id, space: space} } func (m *manager) stateDeleteSpace(id string) { m.stateCmdChan <- &cmdDeleteSpace{id: id} } func (m *manager) stateGetAllSpaces() []*types.TraceSpace { resp := make(chan []*types.TraceSpace, 1) m.stateCmdChan <- &cmdGetAllSpaces{resp: resp} return <-resp } func (m *manager) stateSetTraceStatus(status types.TraceStatus) { m.stateCmdChan <- &cmdSetTraceStatus{status: status} } func (m *manager) stateGetTraceStatus() types.TraceStatus { resp := make(chan types.TraceStatus, 1) m.stateCmdChan <- &cmdGetTraceStatus{resp: resp} return <-resp } func (m *manager) stateMarkCompleted() bool { resp := make(chan bool, 1) m.stateCmdChan <- &cmdMarkCompleted{resp: resp} return <-resp } func (m *manager) stateIsCompleted() bool { resp := make(chan bool, 1) m.stateCmdChan <- &cmdIsCompleted{resp: resp} return <-resp } func (m *manager) stateAddUpdate(update *types.TraceUpdate) { m.stateCmdChan <- &cmdAddUpdate{update: update} } func (m *manager) stateGetUpdates(since int64) []*types.TraceUpdate { resp := make(chan []*types.TraceUpdate, 1) m.stateCmdChan <- &cmdGetUpdates{since: since, resp: resp} return <-resp } func (m *manager) stateSetUpdates(updates []*types.TraceUpdate) { m.stateCmdChan <- &cmdSetUpdates{updates: updates} } func (m *manager) stateAddSubscriber(id string, ch chan *types.TraceUpdate) { m.stateCmdChan <- &cmdAddSubscriber{id: id, ch: ch} } func (m *manager) stateRemoveSubscriber(id string) { m.stateCmdChan <- &cmdRemoveSubscriber{id: id} } func (m *manager) stateGetSubscribers() map[string]chan *types.TraceUpdate { resp := make(chan map[string]chan *types.TraceUpdate, 1) m.stateCmdChan <- &cmdGetSubscribers{resp: resp} return <-resp } func (m *manager) stateBroadcast(update *types.TraceUpdate) { m.stateCmdChan <- &cmdBroadcast{update: update} } // stateExecuteSpaceOp executes a space operation serially through state worker func (m *manager) stateExecuteSpaceOp(spaceID string, fn func() error) error { resp := make(chan error, 1) m.stateCmdChan <- &cmdSpaceKVOp{spaceID: spaceID, fn: fn, resp: resp} return <-resp }