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).
type candidateNode struct {
	entry    *edgenode.NodeEntry
	capacity int
}

type groupPolicy struct {
	maxQueue     int
	queueTimeout time.Duration
}

type inflightRec struct {
	groupKey string
	nodeID   string
}

type admitResult struct {
	node *edgenode.NodeEntry
	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.
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
}

func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, candidates []candidateNode) *candidateNode {
	for i := range candidates {
		c := &candidates[i]
		if group.inflight[c.entry.NodeID] < c.capacity {
			return c
		}
	}
	return nil
}

// admit selects an available node for the given model group, or queues the
// request until a slot opens. Blocks until a node is assigned, the queue
// timeout expires, or ctx is cancelled.
func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy) (*edgenode.NodeEntry, error) {
	m.mu.Lock()

	group := m.getOrCreateGroupLocked(groupKey, policy)
	if group.adapter == "" {
		group.adapter = adapter
	}
	if group.target == "" {
		group.target = target
	}

	candidate := m.findAvailableNodeLocked(group, candidates)
	if candidate != nil {
		group.inflight[candidate.entry.NodeID]++
		m.mu.Unlock()
		return candidate.entry, 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.node, 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.node != nil {
				m.releaseSlot(groupKey, res.node.NodeID)
			}
		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.node != nil {
				m.releaseSlot(groupKey, res.node.NodeID)
			}
		default:
		}
		return nil, ctx.Err()
	}
}

// trackInflight records a dispatched run so release events can find it.
func (m *modelQueueManager) trackInflight(groupKey, runID, nodeID string) {
	m.mu.Lock()
	m.inflightByRun[runID] = inflightRec{groupKey: groupKey, nodeID: nodeID}
	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)
	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
		}

		if group.inflight[nodeID] > 0 {
			group.inflight[nodeID] = 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.
func (m *modelQueueManager) releaseSlot(groupKey, nodeID string) {
	m.mu.Lock()
	m.releaseSlotLocked(groupKey, nodeID)
	m.mu.Unlock()
}

func (m *modelQueueManager) releaseSlotLocked(groupKey, nodeID string) {
	group, ok := m.groups[groupKey]
	if !ok {
		return
	}
	if group.inflight[nodeID] > 0 {
		group.inflight[nodeID]--
	}
	m.tryDispatchLocked(group)
}

// 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:]
		group.inflight[candidate.entry.NodeID]++

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

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),
		})
	}

	return snaps
}

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
}

func resolveSnapshotAdapterName(rec *edgenode.NodeRecord, adapterType, target string) (string, bool) {
	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
	}
}