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

package service

import (
	"context"
	"fmt"
	"sync"
	"time"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	eventpkg "iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

const (
	defaultNodeCapacity  = 1
	defaultGroupMaxQueue = 16
	defaultQueueTimeout  = 30 * time.Second
)

var (
	errQueueFull    = fmt.Errorf("queue is full")
	errQueueTimeout = fmt.Errorf("queue timeout")
)

// candidateNode pairs a registry entry with the per-request capacity derived
// from the node's adapter config (or Runtime.Concurrency as fallback).
// For provider-pool candidates, providerID, adapter, and servedTarget carry
// the globally-unique provider id, the dispatch adapter key, and the concrete
// model name to dispatch.
type candidateNode struct {
	entry        *edgenode.NodeEntry
	capacity     int
	providerID   string // non-empty for provider-pool candidates
	adapter      string // non-empty for provider-pool candidates; dispatch adapter key
	servedTarget string // concrete served model name; used for target rewrite
}

// slotKey returns a unique slot key for inflight accounting.
// For provider-pool candidates (providerID non-empty) it uses "nodeID:providerID"
// so that multiple providers on the same node are tracked independently.
// For legacy candidates it falls back to nodeID only.
func (c *candidateNode) slotKey() string {
	if c.providerID != "" {
		return c.entry.NodeID + ":" + c.providerID
	}
	return c.entry.NodeID
}

type groupPolicy struct {
	maxQueue     int
	queueTimeout time.Duration
}

type inflightRec struct {
	groupKey   string
	nodeID     string
	providerID string // non-empty for provider-pool dispatches
}

type admitResult struct {
	candidate *candidateNode
	err       error
}

type queueItem struct {
	candidates []candidateNode
	waitCh     chan admitResult
	deadline   time.Time
}

type modelQueueGroup struct {
	key      string
	policy   groupPolicy
	queue    []*queueItem
	inflight map[string]int // nodeID → current in-flight count
	adapter  string
	target   string
}

type modelQueueManager struct {
	mu            sync.Mutex
	groups        map[string]*modelQueueGroup
	inflightByRun map[string]inflightRec // runID → {groupKey, nodeID}
	store         *edgenode.NodeStore
}

func newModelQueueManager(store *edgenode.NodeStore) *modelQueueManager {
	return &modelQueueManager{
		groups:        make(map[string]*modelQueueGroup),
		inflightByRun: make(map[string]inflightRec),
		store:         store,
	}
}

func (m *modelQueueManager) setStore(store *edgenode.NodeStore) {
	m.mu.Lock()
	m.store = store
	m.mu.Unlock()
}

// 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"
}

// getOrCreateGroupLocked returns an existing group or creates one with the
// given policy (applying defaults for zero fields). Policy is only applied
// at creation; existing groups retain their current policy until an explicit
// updateGroupPolicyLocked call (used for config-refresh policy propagation).
func (m *modelQueueManager) getOrCreateGroupLocked(key string, policy groupPolicy) *modelQueueGroup {
	g, ok := m.groups[key]
	if !ok {
		if policy.maxQueue <= 0 {
			policy.maxQueue = defaultGroupMaxQueue
		}
		if policy.queueTimeout <= 0 {
			policy.queueTimeout = defaultQueueTimeout
		}
		g = &modelQueueGroup{
			key:      key,
			policy:   policy,
			inflight: make(map[string]int),
		}
		m.groups[key] = g
	}
	return g
}

// updateGroupPolicyLocked refreshes an existing group's queue policy so that
// provider max_queue / queue_timeout_ms changes applied via config refresh take
// effect on subsequent admissions. Only non-zero policy fields override the
// current values; in-flight counts and queued items are preserved untouched.
func (m *modelQueueManager) updateGroupPolicyLocked(g *modelQueueGroup, policy groupPolicy) {
	if g == nil {
		return
	}
	if policy.maxQueue > 0 {
		g.policy.maxQueue = policy.maxQueue
	}
	if policy.queueTimeout > 0 {
		g.policy.queueTimeout = policy.queueTimeout
	}
}

// findAvailableNodeLocked returns the available candidate with the lowest
// in_flight/capacity load ratio. Ties are broken deterministically by
// providerID then nodeID to produce a stable ordering across calls.
// Uses provider-aware slot keys so that multiple providers on the same node
// are tracked independently.
func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, candidates []candidateNode) *candidateNode {
	var best *candidateNode
	bestRatio := 2.0 // sentinel: any valid ratio is in [0, 1)
	for i := range candidates {
		c := &candidates[i]
		slot := c.slotKey()
		inflight := group.inflight[slot]
		if inflight >= c.capacity {
			continue
		}
		ratio := float64(inflight) / float64(c.capacity)
		if ratio < bestRatio || (ratio == bestRatio && candidateLess(c, best)) {
			best = c
			bestRatio = ratio
		}
	}
	return best
}

// candidateLess provides a deterministic tie-break ordering for candidates
// with equal load ratios: providerID first, then nodeID.
func candidateLess(a, b *candidateNode) bool {
	if a.providerID != b.providerID {
		return a.providerID < b.providerID
	}
	return a.entry.NodeID < b.entry.NodeID
}

// admit selects an available candidate for the given model group, or queues
// the request until a slot opens. Blocks until a candidate is assigned, the
// queue timeout expires, or ctx is cancelled. Returns the selected candidateNode
// so callers can rewrite the dispatch target for provider-pool requests.
func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy) (*candidateNode, error) {
	m.mu.Lock()

	group := m.getOrCreateGroupLocked(groupKey, policy)
	// Refresh the policy on an already-existing group so capacity/queue changes
	// applied via config refresh (and recomputed per-request from the live store)
	// take effect on this and subsequent admissions.
	m.updateGroupPolicyLocked(group, policy)
	if group.adapter == "" {
		group.adapter = adapter
	}
	if group.target == "" {
		group.target = target
	}

	candidate := m.findAvailableNodeLocked(group, candidates)
	if candidate != nil {
		slot := candidate.slotKey()
		group.inflight[slot]++
		m.mu.Unlock()
		return candidate, nil
	}

	if len(group.queue) >= group.policy.maxQueue {
		m.mu.Unlock()
		return nil, fmt.Errorf("model group %q: %w", groupKey, errQueueFull)
	}

	item := &queueItem{
		candidates: candidates,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(group.policy.queueTimeout),
	}
	group.queue = append(group.queue, item)
	timeout := group.policy.queueTimeout
	m.mu.Unlock()

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	select {
	case res := <-item.waitCh:
		return res.candidate, res.err
	case <-timer.C:
		m.mu.Lock()
		m.removeItemLocked(groupKey, item)
		m.mu.Unlock()
		// Handle race: dispatch may have sent to waitCh just before timeout fired.
		select {
		case res := <-item.waitCh:
			if res.candidate != nil {
				// Use public releaseSlot (which acquires m.mu) because we are
				// outside the lock here.  This avoids shared-state mutation
				// without mutex that the original releaseSlotByCandidate call
				// would cause.
				slot := res.candidate.slotKey()
				var nodeID, providerID string
				if colonIdx := findLastColon(slot); colonIdx > 0 {
					nodeID = slot[:colonIdx]
					providerID = slot[colonIdx+1:]
				} else {
					nodeID = slot
				}
				m.releaseSlot(groupKey, nodeID, providerID)
			}
		default:
		}
		return nil, fmt.Errorf("model group %q: %w", groupKey, errQueueTimeout)
	case <-ctx.Done():
		m.mu.Lock()
		m.removeItemLocked(groupKey, item)
		m.mu.Unlock()
		// Handle race: dispatch may have sent to waitCh just before cancellation.
		select {
		case res := <-item.waitCh:
			if res.candidate != nil {
				// Same fix as timeout path: use public releaseSlot with
				// explicit mutex to avoid concurrent map/slice mutation.
				slot := res.candidate.slotKey()
				var nodeID, providerID string
				if colonIdx := findLastColon(slot); colonIdx > 0 {
					nodeID = slot[:colonIdx]
					providerID = slot[colonIdx+1:]
				} else {
					nodeID = slot
				}
				m.releaseSlot(groupKey, nodeID, providerID)
			}
		default:
		}
		return nil, ctx.Err()
	}
}

// trackInflight records a dispatched run so release events can find it.
// providerID is non-empty for provider-pool dispatches and is used to
// attribute in-flight counts to the correct provider snapshot.
func (m *modelQueueManager) trackInflight(groupKey, runID, nodeID, providerID string) {
	m.mu.Lock()
	m.inflightByRun[runID] = inflightRec{groupKey: groupKey, nodeID: nodeID, providerID: providerID}
	m.mu.Unlock()
}

// 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)
	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".
		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)
					continue
				}
			} else {
				// Legacy slot: exact nodeID match.
				if slot == nodeID {
					delete(group.inflight, 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) {
	m.mu.Lock()
	var pid string
	if len(providerID) > 0 {
		pid = providerID[0]
	}
	m.releaseSlotLocked(groupKey, nodeID, pid)
	m.mu.Unlock()
}

func (m *modelQueueManager) releaseSlotLocked(groupKey, nodeID string, providerID string) {
	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]--
	}
	m.tryDispatchLocked(group)
}

// releaseSlotByCandidate releases the slot identified by the candidate directly.
// IMPORTANT: This function MUST be called with m.mu held because it calls
// releaseSlotLocked which modifies shared state (group.inflight, tryDispatchLocked).
// Use releaseSlot instead when releasing from an unlocked context.
func (m *modelQueueManager) releaseSlotByCandidate(groupKey string, candidate *candidateNode) {
	slot := candidate.slotKey()
	// For provider-pool candidates, slot is "nodeID:providerID".
	// We need to split it back for the locked release.
	var nodeID, providerID string
	if colonIdx := findLastColon(slot); colonIdx > 0 {
		nodeID = slot[:colonIdx]
		providerID = slot[colonIdx+1:]
	} else {
		nodeID = slot
	}
	m.releaseSlotLocked(groupKey, nodeID, providerID)
}

// findLastColon returns the index of the last ':' in s, or -1 if not found.
func findLastColon(s string) int {
	for i := len(s) - 1; i >= 0; i-- {
		if s[i] == ':' {
			return i
		}
	}
	return -1
}

// tryDispatchLocked attempts to dispatch the head of the queue to an available
// node. Must be called with m.mu held.
func (m *modelQueueManager) tryDispatchLocked(group *modelQueueGroup) {
	for len(group.queue) > 0 {
		head := group.queue[0]

		if time.Now().After(head.deadline) {
			group.queue = group.queue[1:]
			select {
			case head.waitCh <- admitResult{err: fmt.Errorf("model group %q: %w", group.key, errQueueTimeout)}:
			default:
			}
			continue
		}

		candidate := m.findAvailableNodeLocked(group, head.candidates)
		if candidate == nil {
			break
		}

		group.queue = group.queue[1:]
		slot := candidate.slotKey()
		group.inflight[slot]++

		select {
		case head.waitCh <- admitResult{candidate: candidate}:
			return
		default:
			// Caller already timed out or cancelled; release the reserved slot and try next.
			group.inflight[slot]--
		}
	}
}

func (m *modelQueueManager) removeItemLocked(groupKey string, item *queueItem) {
	group, ok := m.groups[groupKey]
	if !ok {
		return
	}
	for i, it := range group.queue {
		if it == item {
			group.queue = append(group.queue[:i], group.queue[i+1:]...)
			return
		}
	}
}

func (m *modelQueueManager) getSnapshotForNode(nodeID string, rec *edgenode.NodeRecord) []*iop.ProviderSnapshot {
	m.mu.Lock()
	defer m.mu.Unlock()

	var snaps []*iop.ProviderSnapshot

	concurrencyFallback := 1
	if rec.Runtime.Concurrency > 0 {
		concurrencyFallback = rec.Runtime.Concurrency
	}

	// 1. CLI
	if rec.Adapters.CLI.Enabled {
		capVal := concurrencyFallback
		inflight, queued := m.getStatsForAdapterLocked(nodeID, rec, "cli")
		snaps = append(snaps, &iop.ProviderSnapshot{
			Adapter:  "cli",
			Status:   "available",
			Capacity: int32(capVal),
			InFlight: int32(inflight),
			Queued:   int32(queued),
		})
	}

	// 2. Ollama
	for _, inst := range rec.Adapters.OllamaInstances {
		if !inst.Enabled {
			continue
		}
		name := inst.Name
		if name == "" {
			name = "ollama"
		}
		capVal := inst.Capacity
		if capVal <= 0 {
			capVal = concurrencyFallback
		}
		inflight, queued := m.getStatsForAdapterLocked(nodeID, rec, name)
		snaps = append(snaps, &iop.ProviderSnapshot{
			Adapter:  name,
			Status:   "available",
			Capacity: int32(capVal),
			InFlight: int32(inflight),
			Queued:   int32(queued),
		})
	}

	// 3. vLLM
	for _, inst := range rec.Adapters.VllmInstances {
		if !inst.Enabled {
			continue
		}
		name := inst.Name
		if name == "" {
			name = "vllm"
		}
		capVal := inst.Capacity
		if capVal <= 0 {
			capVal = concurrencyFallback
		}
		inflight, queued := m.getStatsForAdapterLocked(nodeID, rec, name)
		snaps = append(snaps, &iop.ProviderSnapshot{
			Adapter:  name,
			Status:   "available",
			Capacity: int32(capVal),
			InFlight: int32(inflight),
			Queued:   int32(queued),
		})
	}

	// 4. OpenAI Compat
	for _, inst := range rec.Adapters.OpenAICompatInstances {
		if !inst.Enabled {
			continue
		}
		name := inst.Name
		if name == "" {
			name = "openai_compat"
		}
		capVal := inst.Capacity
		if capVal <= 0 {
			capVal = concurrencyFallback
		}
		inflight, queued := m.getStatsForAdapterLocked(nodeID, rec, name)
		snaps = append(snaps, &iop.ProviderSnapshot{
			Adapter:  name,
			Status:   "available",
			Capacity: int32(capVal),
			InFlight: int32(inflight),
			Queued:   int32(queued),
		})
	}

	// 5. Provider catalog from nodes[].providers[] config (provider pool schema).
	//    Stats are resolved via providerID from inflightByRun and queue candidates
	//    so that provider-pool dispatch (keyed by model alias, not adapter/target)
	//    is correctly attributed to the right provider snapshot.
	for _, prov := range rec.Providers {
		if prov.ID == "" {
			continue
		}
		capVal := prov.Capacity
		if capVal <= 0 {
			capVal = concurrencyFallback
		}
		inflight, queued := m.getStatsForProviderLocked(nodeID, prov.ID)

		servedModels := make([]string, len(prov.Models))
		copy(servedModels, prov.Models)

		lifecycleCaps := make([]string, len(prov.LifecycleCapabilities))
		copy(lifecycleCaps, prov.LifecycleCapabilities)

		var loadRatio float32
		if capVal > 0 {
			loadRatio = float32(inflight) / float32(capVal)
		}

		snaps = append(snaps, &iop.ProviderSnapshot{
			Adapter:               prov.Adapter,
			Status:                "available",
			Capacity:              int32(capVal),
			InFlight:              int32(inflight),
			Queued:                int32(queued),
			Id:                    prov.ID,
			Type:                  prov.Type,
			Category:              string(prov.Category),
			ServedModels:          servedModels,
			Health:                prov.Health,
			LoadRatio:             loadRatio,
			LifecycleCapabilities: lifecycleCaps,
		})
	}

	return snaps
}

// getStatsForProviderLocked returns in-flight and queued counts for a
// provider-pool provider identified by (nodeID, providerID). Must be called
// with m.mu held.
func (m *modelQueueManager) getStatsForProviderLocked(nodeID, providerID string) (inFlight, queued int) {
	for _, rec := range m.inflightByRun {
		if rec.nodeID == nodeID && rec.providerID == providerID {
			inFlight++
		}
	}
	for _, group := range m.groups {
		for _, item := range group.queue {
			for _, c := range item.candidates {
				if c.entry.NodeID == nodeID && c.providerID == providerID {
					queued++
					break
				}
			}
		}
	}
	return
}

func (m *modelQueueManager) getStatsForAdapterLocked(nodeID string, rec *edgenode.NodeRecord, adapterName string) (inFlight, queued int) {
	for _, group := range m.groups {
		canonical, ok := resolveSnapshotAdapterName(rec, group.adapter, group.target)
		if ok && canonical == adapterName {
			if val, ok := group.inflight[nodeID]; ok {
				inFlight += val
			}
			for _, item := range group.queue {
				for _, c := range item.candidates {
					if c.entry.NodeID == nodeID {
						queued++
						break
					}
				}
			}
		}
	}
	return
}

// resolveSnapshotAdapterName returns the adapter/instance key used to match
// queue state for a provider snapshot.
// When adapterType is empty it falls back to the provider id (target) so that
// queue stats are still resolved instead of being dropped on an empty key.
func resolveSnapshotAdapterName(rec *edgenode.NodeRecord, adapterType, target string) (string, bool) {
	// Fallback to provider id when adapter is not set.
	if adapterType == "" {
		adapterType = target
	}
	if rec == nil {
		return adapterType, true
	}

	// 1. Exact instance Name match (highest priority).
	for _, inst := range rec.Adapters.OllamaInstances {
		if inst.Name == adapterType {
			return adapterType, true
		}
	}
	for _, inst := range rec.Adapters.VllmInstances {
		if inst.Name == adapterType {
			return adapterType, true
		}
	}
	for _, inst := range rec.Adapters.OpenAICompatInstances {
		if inst.Name == adapterType {
			return adapterType, true
		}
	}

	// 2. Type-name route.
	switch adapterType {
	case "ollama":
		var enabled []string
		for _, inst := range rec.Adapters.OllamaInstances {
			if inst.Enabled {
				name := inst.Name
				if name == "" {
					name = "ollama"
				}
				enabled = append(enabled, name)
			}
		}
		switch len(enabled) {
		case 0:
			return "ollama", true
		case 1:
			return enabled[0], true
		default:
			return "", false
		}

	case "vllm":
		var enabled []string
		for _, inst := range rec.Adapters.VllmInstances {
			if inst.Enabled {
				name := inst.Name
				if name == "" {
					name = "vllm"
				}
				enabled = append(enabled, name)
			}
		}
		switch len(enabled) {
		case 0:
			return "vllm", true
		case 1:
			return enabled[0], true
		default:
			return "", false
		}

	case "openai_compat":
		var enabled []string
		for _, inst := range rec.Adapters.OpenAICompatInstances {
			if inst.Enabled {
				name := inst.Name
				if name == "" {
					name = "openai_compat"
				}
				enabled = append(enabled, name)
			}
		}
		switch len(enabled) {
		case 0:
			return "openai_compat", true
		case 1:
			return enabled[0], true
		default:
			return "", false
		}

	case "cli":
		return "cli", true

	default:
		return adapterType, true
	}
}