iop/apps/edge/internal/service/model_queue_release.go

package service

import (
	edgeevents "iop/apps/edge/internal/events"
	eventpkg "iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

// startEventWatcher subscribes to all run and node events, releasing in-flight
// slots when runs terminate or nodes disconnect. Returns a stop function.
func (m *modelQueueManager) startEventWatcher(bus *edgeevents.Bus) func() {
	runCh, unsubRun := bus.SubscribeAllRuns(256)
	nodeCh, unsubNode := bus.SubscribeAllNodes(64)
	done := make(chan struct{})
	go func() {
		defer unsubRun()
		defer unsubNode()
		for {
			select {
			case e, ok := <-runCh:
				if !ok {
					return
				}
				if isTerminalRunEvent(e) {
					m.releaseRun(e.GetRunId(), e.GetType())
				}
			case e, ok := <-nodeCh:
				if !ok {
					return
				}
				if e.GetType() == eventpkg.TypeNodeDisconnected {
					m.releaseNode(e.GetNodeId(), "disconnected")
				}
			case <-done:
				return
			}
		}
	}()
	return func() { close(done) }
}

func isTerminalRunEvent(e *iop.RunEvent) bool {
	t := e.GetType()
	return t == "complete" || t == "error" || t == "cancelled"
}

// releaseRun releases the in-flight slot for a terminated run and dispatches
// the next queued item if one is waiting.
func (m *modelQueueManager) releaseRun(runID, reason string) {
	m.mu.Lock()
	rec, ok := m.inflightByRun[runID]
	if !ok {
		m.mu.Unlock()
		return
	}
	delete(m.inflightByRun, runID)
	m.releaseSlotLocked(rec.groupKey, rec.nodeID, rec.providerID, rec.long)
	m.mu.Unlock()
}

// releaseNode resets all in-flight slots on a disconnected node, removes it
// from all queued items' candidate lists, and tries to re-dispatch to any
// remaining live candidates.
func (m *modelQueueManager) releaseNode(nodeID, reason string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	for runID, rec := range m.inflightByRun {
		if rec.nodeID == nodeID {
			delete(m.inflightByRun, runID)
		}
	}

	for _, group := range m.groups {
		// Remove the disconnected node from all queued items so a future
		// tryDispatch cannot pick it as a stale available candidate.
		for _, item := range group.queue {
			filtered := make([]candidateNode, 0, len(item.candidates))
			for _, c := range item.candidates {
				if c.entry.NodeID != nodeID {
					filtered = append(filtered, c)
				}
			}
			item.candidates = filtered
		}

		// Clear all slot keys that start with this nodeID.
		// For provider-pool: "nodeID:providerID". For legacy: "nodeID".
		// Long-context in-flight counters are keyed identically, so clear them
		// alongside the normal counters.
		for slot := range group.inflight {
			if colonIdx := findLastColon(slot); colonIdx > 0 {
				candidateNodeID := slot[:colonIdx]
				// Only clear if the node part matches.
				if candidateNodeID == nodeID {
					delete(group.inflight, slot)
					delete(group.longInflight, slot)
					continue
				}
			} else {
				// Legacy slot: exact nodeID match.
				if slot == nodeID {
					delete(group.inflight, slot)
					delete(group.longInflight, slot)
				}
			}
		}

		// tryDispatchLocked will handle remaining queued items, but we need
		// to call it if we cleared any inflight.
		// Since we already deleted entries, tryDispatchLocked will find available slots.
		// We need to re-check: if any slots were cleared and there are queued items.
		if len(group.queue) > 0 {
			m.tryDispatchLocked(group)
		}
	}
}

// releaseSlot decrements the in-flight count and tries to dispatch the next
// queued item. Used when admit-path I/O fails after the slot was reserved.
// For provider-pool dispatches (nodeID:providerID slot), pass providerID.
// For legacy dispatches, pass empty providerID so nodeID is used directly.
func (m *modelQueueManager) releaseSlot(groupKey, nodeID string, providerID ...string) {
	var pid string
	if len(providerID) > 0 {
		pid = providerID[0]
	}
	m.releaseSlotWithLong(groupKey, nodeID, pid, false)
}

// releaseSlotWithLong releases a reserved slot, additionally freeing a long-context
// slot when longSlot is true. Used by admit-path failure/race handling where the
// caller knows whether a long slot was reserved.
func (m *modelQueueManager) releaseSlotWithLong(groupKey, nodeID, providerID string, longSlot bool) {
	m.mu.Lock()
	m.releaseSlotLocked(groupKey, nodeID, providerID, longSlot)
	m.mu.Unlock()
}

func (m *modelQueueManager) releaseSlotLocked(groupKey, nodeID string, providerID string, longSlot bool) {
	group, ok := m.groups[groupKey]
	if !ok {
		return
	}
	// Use provider-aware slot key for provider-pool dispatches.
	slot := nodeID
	if providerID != "" {
		slot = nodeID + ":" + providerID
	}
	if group.inflight[slot] > 0 {
		group.inflight[slot]--
	}
	if longSlot && group.longInflight[slot] > 0 {
		group.longInflight[slot]--
	}
	m.tryDispatchLocked(group)
}