package service import ( "context" "errors" "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 !generationEligible(c.generation, res) { 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 } // generationEligible reports whether a candidate's connection generation may // still acquire res. An untracked candidate (generation 0 — legacy/direct // dispatch and hand-built fixtures) is always eligible. A tracked candidate is // eligible when it is at least as new as the generation that last owned the // resource; on an orphaned (disconnected) resource it must be strictly newer, // because only a genuine reconnect — never the disconnected generation itself — // may re-acquire the slot. func generationEligible(generation uint64, res *providerResourceState) bool { if generation == 0 || res == nil { return true } if res.orphan { return generation > res.generation } return generation >= res.generation } // 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 } // Generation fence: a candidate resolved before the owning connection // disconnected or was superseded by a reconnect must not reserve. Only a // strictly newer generation may re-acquire an orphaned (disconnected) // resource, so the disconnected generation cannot resurrect its own slot. if !generationEligible(candidate.generation, res) { 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 if candidate.generation > res.generation { res.generation = candidate.generation } 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. // // The return value is tri-state: // - resolveOk: usable candidates returned; the pump continues to // findAvailableNodeLocked (which may skip the item if all candidates // are at capacity — temporary block, not terminal). // - resolveNoCandidates: the live candidate universe is empty. The pump // removes the item from the queue and delivers errProviderUnavailable. // - resolveResolverError: the resolver returned an error distinct from // errProviderUnavailable. The item remains queued so a later live-state // re-evaluation can recover from a catalog/config resolver fault. // // Orphaned candidates (disconnected nodes whose resources were cleared by // releaseNode) are filtered out. If all candidates are orphaned, the result // is resolveNoCandidates — the pump has no live target to dispatch to. // Must be called with m.mu held. func (m *modelQueueManager) resolveQueuedCandidatesLocked(item *queueItem) ([]candidateNode, resolveOutcome, error) { if item.resolveCandidates == nil { return item.candidates, resolveOk, nil } candidates, err := item.resolveCandidates() if err != nil { if errors.Is(err, errProviderUnavailable) { return nil, resolveNoCandidates, nil } return nil, resolveResolverError, err } if len(candidates) == 0 { return nil, resolveNoCandidates, nil } // 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) } if len(filtered) == 0 { // Every candidate resolved was orphaned — there is no live target // to dispatch to. Treat as no-live-candidate terminal. return nil, resolveNoCandidates, nil } return filtered, resolveOk, nil } // 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. // // Resolver outcomes are tri-state: resolveOk continues to selection // (findAvailableNodeLocked); resolveNoCandidates is terminal and delivers typed // unavailable; resolveResolverError leaves the item queued for a later live-state // re-evaluation. 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, outcome, _ := m.resolveQueuedCandidatesLocked(ref.item) switch outcome { case resolveOk: candidate := m.findAvailableNodeLocked(ref.group, candidates, ref.item.long) if candidate == nil { // All candidates exist but are at capacity/long-capacity or // filtered out by live-state checks. The item stays queued — // this is the temporary-block path, not terminal. 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 case resolveNoCandidates: // No live candidates — terminal unavailable. Remove the item and // deliver a typed signal so the caller does not wait on a dead end. m.removeQueuedItemLocked(ref.group, ref.item) select { case ref.item.waitCh <- admitResult{err: fmt.Errorf("model group %q: %w", ref.group.key, errProviderUnavailable)}: default: } return true case resolveResolverError: // Resolver returned a non-terminal error. Keep the item queued so a // corrected catalog/config snapshot can dispatch it on a later pump. continue } } 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, generation: candidate.generation, 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 }