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

package service

import (
	"context"
	"fmt"
	"sort"
	"time"

	"iop/packages/go/config"
)

// normalizeGroupPolicy applies queue defaults while preserving an explicit
// zero queue timeout as "no timeout".
func normalizeGroupPolicy(policy groupPolicy) groupPolicy {
	if policy.maxQueue <= 0 {
		policy.maxQueue = defaultGroupMaxQueue
	}
	if policy.queueTimeout > 0 {
		policy.queueTimeoutSet = true
	}
	if !policy.queueTimeoutSet {
		policy.queueTimeout = defaultQueueTimeout
		policy.queueTimeoutSet = true
	}
	return policy
}

// getOrCreateGroupLocked returns an existing group or creates one with the
// given policy (applying defaults for unset 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 {
		policy = normalizeGroupPolicy(policy)
		g = &modelQueueGroup{
			key:          key,
			policy:       policy,
			inflight:     make(map[string]int),
			longInflight: 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 configured 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
		g.policy.queueTimeoutSet = true
	} else if policy.queueTimeoutSet {
		g.policy.queueTimeout = policy.queueTimeout
		g.policy.queueTimeoutSet = true
	}
}

// findAvailableNodeLocked returns the available candidate with the lowest
// absolute in_flight count. Full candidates are skipped, so a provider that is
// at capacity is deferred behind providers that still have room. When in_flight
// is equal, the candidate with the lower priority value wins. When both
// in_flight and priority are equal, all tied candidates remain in best[] so
// that the rotation logic in the caller picks the next slot after the last
// selected one.
func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, candidates []candidateNode, long bool) *candidateNode {
	var best []candidateNode
	bestInflight := 0
	bestPriority := 0
	for i := range candidates {
		c, live := m.liveCandidateLocked(&candidates[i])
		if !live {
			continue
		}
		if c.capacity <= 0 {
			continue
		}
		var inflight int
		if c.providerID != "" {
			res, eligible := m.resourceForCandidateLocked(&c)
			if !eligible || res == nil {
				continue
			}
			if !res.canReserve(long) {
				continue
			}
			inflight = res.inFlight
		} else {
			slot := c.slotKey()
			inflight = group.inflight[slot]
			if inflight >= c.capacity {
				continue
			}
			// Long-context gate: a long request is excluded from a provider whose
			// long slots are full. Providers with long_context_capacity == 0 declare
			// no dedicated long-slot limit and are not gated here. Normal requests
			// (long == false) never consult long slot state.
			if long && c.longContextCapacity > 0 && group.longInflight[slot] >= c.longContextCapacity {
				continue
			}
		}
		switch {
		case len(best) == 0 || inflight < bestInflight:
			best = []candidateNode{c}
			bestInflight = inflight
			bestPriority = c.priority
		case inflight == bestInflight:
			if c.priority < bestPriority {
				best = []candidateNode{c}
				bestPriority = c.priority
			} else if c.priority == bestPriority {
				best = append(best, c)
			}
		}
	}
	if len(best) == 0 {
		return nil
	}
	sort.Slice(best, func(i, j int) bool {
		return candidateLess(&best[i], &best[j])
	})
	idx := 0
	if group.lastSelectedSlot != "" {
		for i := range best {
			if best[i].slotKey() == group.lastSelectedSlot {
				idx = (i + 1) % len(best)
				break
			}
		}
	}
	return &best[idx]
}

// candidateLess provides a deterministic ordering for equal-inflight/priority rotation:
// providerID first, then nodeID.
func candidateLess(a, b *candidateNode) bool {
	if b == nil {
		return true
	}
	if a.providerID != b.providerID {
		return a.providerID < b.providerID
	}
	return a.entry.NodeID < b.entry.NodeID
}

// liveCandidateLocked re-resolves a candidate against the current node store and
// reports whether it may still be dispatched to. What the queued candidate owns
// is its identity — node, provider, served target — and that never changes; the
// dispatch inputs derived from provider config (capacity, long-context capacity,
// priority, adapter, execution path) are re-read here, because an item admitted
// from the queue may have been waiting across any number of config refreshes.
// A provider that disappeared from its node, or that was disabled after the item
// was enqueued, is not eligible at all.
//
// Legacy candidates (no provider id) and fixtures without a store keep their
// enqueue-time fields: there is no live provider config to re-read for them.
// Must be called with m.mu held.
func (m *modelQueueManager) liveCandidateLocked(c *candidateNode) (candidateNode, bool) {
	live := *c
	if c.providerID == "" || m.store == nil || c.entry == nil {
		return live, true
	}
	rec, ok := m.store.FindByID(c.entry.NodeID)
	if !ok || rec == nil {
		return live, false
	}
	for _, prov := range rec.Providers {
		if prov.ID != c.providerID {
			continue
		}
		if !config.ProviderEnabled(prov) {
			return live, false
		}
		applyProviderDispatchFields(&live, prov)
		return live, true
	}
	return live, false
}

// reserveCandidateLocked re-checks live eligibility and, if it still holds,
// reserves the slot and mints the lease that owns its release. Selection and
// reservation therefore share one critical section: a release, refresh, or
// disconnect that lands between a scan and its reservation cannot leave the
// scheduler holding a slot on a provider that no longer has room for it. Returns
// the lease id, or false when the candidate became ineligible. Must be called
// with m.mu held.
func (m *modelQueueManager) reserveCandidateLocked(group *modelQueueGroup, candidate *candidateNode, long bool) (uint64, bool) {
	live, eligible := m.liveCandidateLocked(candidate)
	if !eligible || live.capacity <= 0 {
		return 0, false
	}

	slot := candidate.slotKey()
	if candidate.providerID != "" {
		res, ok := m.resourceForCandidateLocked(candidate)
		if !ok || res == nil {
			return 0, false
		}
		if !res.canReserve(long) {
			return 0, false
		}
		// Clear orphan flag on re-admit: a disconnected node may reconnect and
		// the new lease must not be blocked by the stale disconnect marker.
		res.orphan = false
		candidate.capacity = res.capacity
		candidate.longContextCapacity = res.longCapacity
		candidate.priority = live.priority
		candidate.adapter = live.adapter
		candidate.providerType = live.providerType
		candidate.executionPath = live.executionPath
		res.reserve(long)
	} else {
		if group.inflight[slot] >= live.capacity {
			return 0, false
		}
		if long && live.longContextCapacity > 0 && group.longInflight[slot] >= live.longContextCapacity {
			return 0, false
		}
		group.inflight[slot]++
		if long && live.longContextCapacity > 0 {
			group.longInflight[slot]++
		}
	}

	return m.newLeaseLocked(group.key, candidate, long), true
}

// enqueueItemLocked appends an item to its group queue and stamps it with the
// next global arrival number. Every pending item enters the scheduler here, so
// the sequence is assigned exactly once. The optional resolveCandidates closure
// lets the pump recompute the candidate universe from live state at dispatch
// time instead of trusting the enqueue-time snapshot; when nil the manager
// falls back to the candidates slice so legacy callers keep working. Must be
// called with m.mu held.
func (m *modelQueueManager) enqueueItemLocked(group *modelQueueGroup, item *queueItem, resolveCandidates func() ([]candidateNode, error)) {
	m.enqueueSeq++
	item.enqueueSeq = m.enqueueSeq
	item.resolveCandidates = resolveCandidates
	group.queue = append(group.queue, item)
}

// pendingRef locates a queued item together with the group that stores it.
type pendingRef struct {
	group *modelQueueGroup
	item  *queueItem
}

// pendingItemsLocked returns every queued item across all model groups in global
// arrival order. Group queues stay the one storage for pending items and the
// global order is derived from them, so there is a single place where an item is
// inserted and a single place where one is removed — no second index can drift
// out of step with the queues on a cancel, timeout, or disconnect path. Items
// that share a sequence (fixtures that build queues directly) fall back to group
// key and in-group position, which preserves group-local FIFO. Must be called
// with m.mu held.
func (m *modelQueueManager) pendingItemsLocked() []pendingRef {
	groupKeys := make([]string, 0, len(m.groups))
	for key, group := range m.groups {
		if len(group.queue) > 0 {
			groupKeys = append(groupKeys, key)
		}
	}
	if len(groupKeys) == 0 {
		return nil
	}
	sort.Strings(groupKeys)

	pending := make([]pendingRef, 0, len(groupKeys))
	for _, key := range groupKeys {
		group := m.groups[key]
		for _, item := range group.queue {
			pending = append(pending, pendingRef{group: group, item: item})
		}
	}
	sort.SliceStable(pending, func(i, j int) bool {
		return pending[i].item.enqueueSeq < pending[j].item.enqueueSeq
	})
	return pending
}

// pumpAllLocked re-evaluates every pending item in every model group against
// current provider state. It exists because provider resources are shared across
// model groups: a lease released by one group frees a slot that the earliest
// waiter may be sitting in a different group's queue for, and pumping only the
// releasing group would leave that waiter asleep until its own timeout.
//
// Each round dispatches at most one item and then rescans, because a dispatch
// consumes capacity that changes what the remaining items may do. Must be called
// with m.mu held.
func (m *modelQueueManager) pumpAllLocked() {
	for m.pumpOnceLocked() {
	}
}

// resolveQueuedCandidatesLocked recomputes the candidate universe for item
// under the manager lock. When the item owns a live resolver it calls that
// resolver to rebuild the slice from current store/catalog/registry state;
// otherwise it returns the enqueue-time snapshot. A resolver error or an empty
// result leaves the item queued — stale snapshot fallback is deliberately not
// used so a provider that was disabled at enqueue time cannot be resurrected
// by a stale list. Orphaned candidates (disconnected nodes whose resources
// were cleared by releaseNode) are filtered out so the pump never dispatches
// to a node that is no longer connected. Returns the candidate slice and
// whether any candidates are available. Must be called with m.mu held.
func (m *modelQueueManager) resolveQueuedCandidatesLocked(item *queueItem) ([]candidateNode, bool) {
	if item.resolveCandidates == nil {
		return item.candidates, true
	}
	candidates, err := item.resolveCandidates()
	if err != nil || len(candidates) == 0 {
		return nil, false
	}
	// Filter out orphaned candidates (disconnected nodes whose resources were
	// cleared by releaseNode). The resolver rebuilds the universe from live
	// state, but it has no view of the manager's orphan tracking, so we drop
	// those entries here before the pump scans them.
	filtered := make([]candidateNode, 0, len(candidates))
	for i := range candidates {
		c := &candidates[i]
		if c.entry == nil {
			filtered = append(filtered, *c)
			continue
		}
		key := providerResourceKey{nodeID: c.entry.NodeID, providerID: c.providerID}
		if r, ok := m.resources[key]; ok && r.orphan {
			continue
		}
		filtered = append(filtered, *c)
	}
	return filtered, true
}

// pumpOnceLocked expires timed-out items and dispatches the earliest globally
// dispatchable one, reporting whether it changed anything. Items that cannot be
// dispatched under current state are skipped rather than blocking the scan, so a
// long request whose long slots are full never holds up a normal request behind
// it — in its own group or in any other (SDD D05/D06). Among items that *can*
// dispatch, the global enqueue sequence decides, which is what keeps FIFO across
// groups and prevents a busy group from starving a quiet one.
func (m *modelQueueManager) pumpOnceLocked() bool {
	pending := m.pendingItemsLocked()

	now := time.Now()
	for _, ref := range pending {
		if ref.item.deadline.IsZero() || !now.After(ref.item.deadline) {
			continue
		}
		m.removeQueuedItemLocked(ref.group, ref.item)
		select {
		case ref.item.waitCh <- admitResult{err: fmt.Errorf("model group %q: %w", ref.group.key, errQueueTimeout)}:
		default:
		}
		return true
	}

	for _, ref := range pending {
		candidates, ok := m.resolveQueuedCandidatesLocked(ref.item)
		if !ok {
			continue
		}
		candidate := m.findAvailableNodeLocked(ref.group, candidates, ref.item.long)
		if candidate == nil {
			continue
		}
		leaseID, ok := m.reserveCandidateLocked(ref.group, candidate, ref.item.long)
		if !ok {
			// Eligibility vanished between selection and reservation; leave the
			// item queued and try the next one.
			continue
		}
		candidate.leaseID = leaseID
		m.removeQueuedItemLocked(ref.group, ref.item)

		select {
		case ref.item.waitCh <- admitResult{candidate: candidate}:
			ref.group.lastSelectedSlot = candidate.slotKey()
		default:
			// Caller already timed out or cancelled; release the reserved slot.
			// No pump here: we are already inside the dispatch loop.
			m.releaseLeaseLocked(leaseID)
		}
		return true
	}

	return false
}

// 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.
// The optional long flag marks the request as long-context so it is gated on
// provider long slots and reserves a long slot on dispatch. The optional
// resolveCandidates closure lets the pump recompute the candidate universe from
// live state at dispatch time; when nil the enqueue-time snapshot is used.
// The optional providerPool flag scopes admission under the manager's global
// policy when true.
func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy, resolveCandidates func() ([]candidateNode, error), long bool, providerPool bool) (*candidateNode, error) {
	candidate, _, err := m.admitWithReason(ctx, groupKey, adapter, target, candidates, policy, resolveCandidates, long, providerPool)
	return candidate, err
}

func (m *modelQueueManager) admitWithReason(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy, resolveCandidates func() ([]candidateNode, error), long bool, providerPool bool) (*candidateNode, string, 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, long)
	if candidate != nil {
		// The lease record is created in the same critical section that reserved
		// the resource, so a reserved slot is never observable without the lease
		// identity that owns its release.
		leaseID, ok := m.reserveCandidateLocked(group, candidate, long)
		if !ok {
			m.mu.Unlock()
			return nil, "", fmt.Errorf("model group %q: candidate ineligible", groupKey)
		}
		candidate.leaseID = leaseID
		group.lastSelectedSlot = candidate.slotKey()
		m.mu.Unlock()
		return candidate, "dispatched", nil
	}

	if providerPool {
		// Provider-pool admission uses the manager's root policy for the
		// global cap. Legacy adapter pending items are never counted against it.
		// Zero max_queue normalizes to the canonical default.
		maxQ := m.providerPoolPolicy.maxQueue
		if maxQ <= 0 {
			maxQ = defaultGroupMaxQueue
		}
		if m.pendingProviderPoolCountLocked() >= maxQ {
			m.mu.Unlock()
			return nil, "", fmt.Errorf("provider pool: %w", errQueueFull)
		}
	} else if len(group.queue) >= group.policy.maxQueue {
		m.mu.Unlock()
		return nil, "", fmt.Errorf("model group %q: %w", groupKey, errQueueFull)
	}

	// Determine queue reason.
	reason := "capacity_full"
	if long {
		longBlocked := false
		for _, c := range candidates {
			if c.capacity <= 0 {
				continue
			}
			var inflight, longInflight int
			var capacity, longCapacity int
			if c.providerID != "" {
				res, eligible := m.resourceForCandidateLocked(&c)
				if !eligible || res == nil {
					continue
				}
				inflight = res.inFlight
				longInflight = res.longInFlight
				capacity = res.capacity
				longCapacity = res.longCapacity
			} else {
				slot := c.slotKey()
				inflight = group.inflight[slot]
				longInflight = group.longInflight[slot]
				capacity = c.capacity
				longCapacity = c.longContextCapacity
			}
			if inflight < capacity {
				if longCapacity > 0 && longInflight >= longCapacity {
					longBlocked = true
					break
				}
			}
		}
		if longBlocked {
			reason = "long_context_capacity_full"
		}
	}

	now := time.Now()
	var deadline time.Time
	if providerPool {
		if m.providerPoolPolicy.queueTimeoutSet {
			deadline = deadlineFrom(now, m.providerPoolPolicy.queueTimeout)
		} else if group.policy.queueTimeout > 0 {
			deadline = deadlineFrom(now, group.policy.queueTimeout)
		}
	} else if group.policy.queueTimeout > 0 {
		deadline = deadlineFrom(now, group.policy.queueTimeout)
	}
	item := &queueItem{
		candidates:        candidates,
		resolveCandidates: resolveCandidates,
		waitCh:            make(chan admitResult, 1),
		enqueuedAt:        now,
		deadline:          deadline,
		policyChange:      make(chan struct{}, 1),
		long:              long,
		providerPool:      providerPool,
		reason:            reason,
	}
	m.enqueueItemLocked(group, item, resolveCandidates)
	m.mu.Unlock()

	return m.waitForAdmission(ctx, groupKey, item)
}

// waitForAdmission blocks until the request is admitted, its queue timeout
// fires, or ctx is cancelled. A zero deadline disables only the timer channel;
// the same loop keeps observing policy changes so both initial-zero and
// disable/re-enable transitions can arm a future deadline. Timer delivery is a
// re-check signal only: membership and the current deadline are re-read under
// m.mu before an item can be timed out.
func (m *modelQueueManager) waitForAdmission(ctx context.Context, groupKey string, item *queueItem) (*candidateNode, string, error) {
	var timer *time.Timer
	var timerC <-chan time.Time
	resetTimer := func(deadline time.Time) {
		if timer != nil {
			if !timer.Stop() {
				select {
				case <-timer.C:
				default:
				}
			}
		}
		if deadline.IsZero() {
			timerC = nil
			return
		}
		delay := time.Until(deadline)
		if delay < 0 {
			delay = 0
		}
		if timer == nil {
			timer = time.NewTimer(delay)
		} else {
			timer.Reset(delay)
		}
		timerC = timer.C
	}
	defer func() {
		if timer != nil {
			timer.Stop()
		}
	}()

	m.mu.Lock()
	deadline := item.deadline
	m.mu.Unlock()
	resetTimer(deadline)

	for {
		select {
		case res := <-item.waitCh:
			return res.candidate, item.reason, res.err
		case <-ctx.Done():
			m.mu.Lock()
			m.removeItemLocked(groupKey, item)
			m.mu.Unlock()
			select {
			case res := <-item.waitCh:
				if res.candidate != nil {
					m.releaseLease(res.candidate.leaseID, "queue-abandoned")
				}
			default:
			}
			return nil, "", ctx.Err()
		case <-timerC:
			m.mu.Lock()
			deadline = item.deadline
			queued := m.isQueuedItemLocked(groupKey, item)
			if queued && !deadline.IsZero() && !time.Now().Before(deadline) {
				m.removeItemLocked(groupKey, item)
				m.mu.Unlock()
				select {
				case res := <-item.waitCh:
					if res.candidate != nil {
						m.releaseLease(res.candidate.leaseID, "queue-abandoned")
					}
				default:
				}
				return nil, "", fmt.Errorf("model group %q: %w", groupKey, errQueueTimeout)
			}
			m.mu.Unlock()
			if !queued {
				// A dispatcher or the queue pump already removed the item and
				// published its result. Disable the stale timer and let waitCh
				// deliver the winning outcome.
				resetTimer(time.Time{})
				continue
			}
			resetTimer(deadline)
		case <-item.policyChange:
			m.mu.Lock()
			deadline = item.deadline
			m.mu.Unlock()
			resetTimer(deadline)
		}
	}
}

// deadlineFrom returns a zero time when timeout is zero (meaning no internal
// queue deadline — the waiter only waits for ctx cancellation) and otherwise
// returns time.Now().Add(timeout).
func deadlineFrom(now time.Time, timeout time.Duration) time.Time {
	if timeout <= 0 {
		return time.Time{}
	}
	return now.Add(timeout)
}

// newLeaseLocked mints a lease for a slot that was just reserved for candidate.
// Must be called with m.mu held, in the same critical section as the reserve.
// The long flag is narrowed to what was actually reserved: a long slot is only
// taken when the request is long and the provider declares a long-context limit,
// so the release path frees exactly what admission took.
func (m *modelQueueManager) newLeaseLocked(groupKey string, candidate *candidateNode, long bool) uint64 {
	m.leaseSeq++
	id := m.leaseSeq
	m.leases[id] = &providerLease{
		id:         id,
		groupKey:   groupKey,
		nodeID:     candidate.entry.NodeID,
		providerID: candidate.providerID,
		long:       long && candidate.longContextCapacity > 0,
		state:      leaseStateReserved,
	}
	return id
}

// trackLease binds a dispatched run id to its lease so run terminal events and
// node disconnects can resolve the lease that owns the slot. A lease that was
// already released (a terminal cause won the race before the run id was bound)
// stays released; binding it again would resurrect a freed slot.
func (m *modelQueueManager) trackLease(leaseID uint64, runID string) {
	if leaseID == 0 || runID == "" {
		return
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	lease, ok := m.leases[leaseID]
	if !ok || lease.state == leaseStateReleased {
		return
	}
	lease.state = leaseStateTracked
	lease.runID = runID
	m.leaseByRun[runID] = leaseID
}

// removeQueuedItemLocked drops one pending item from its group queue. Every
// removal path — dispatch, timeout, cancel — goes through here, so the queues
// remain the single authority on what is still pending. Must be called with
// m.mu held.
func (m *modelQueueManager) removeQueuedItemLocked(group *modelQueueGroup, item *queueItem) {
	if group == nil {
		return
	}
	for i, it := range group.queue {
		if it == item {
			group.queue = append(group.queue[:i], group.queue[i+1:]...)
			return
		}
	}
}

func (m *modelQueueManager) removeItemLocked(groupKey string, item *queueItem) {
	m.removeQueuedItemLocked(m.groups[groupKey], item)
}

func (m *modelQueueManager) isQueuedItemLocked(groupKey string, item *queueItem) bool {
	group := m.groups[groupKey]
	if group == nil {
		return false
	}
	for _, queued := range group.queue {
		if queued == item {
			return true
		}
	}
	return false
}