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

package service

import (
	"fmt"

	iop "iop/proto/gen/iop"
)

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

// releaseLease frees the resources held by one lease, exactly once, whatever the
// cause: a send failure, a normalized terminal run event, a tunnel END/ERROR
// frame, a tunnel close, or a node disconnect. All of them converge here, and the
// lease map — mutated only under m.mu — decides which caller performs the single
// transition into released. A lease id that is already gone is a no-op, so a
// losing racer can never decrement a counter that now belongs to another request.
func (m *modelQueueManager) releaseLease(leaseID uint64, reason string) {
	if leaseID == 0 {
		return
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.leases[leaseID]; !ok {
		return
	}
	if !m.releaseLeaseLocked(leaseID) {
		return
	}
	// Pump every group, not just the one that held the lease: the freed slot
	// belongs to a provider resource that other model groups queue against, and
	// the earliest waiter for it may be in any of them.
	m.pumpAllLocked()
}

// releaseLeaseLocked performs the reserved/tracked → released transition and
// frees the slot the lease holds, without pumping the queue. Returns true when
// this call is the one that performed the transition. Must be called with m.mu
// held.
func (m *modelQueueManager) releaseLeaseLocked(leaseID uint64) bool {
	lease, ok := m.leases[leaseID]
	if !ok {
		return false
	}
	delete(m.leases, leaseID)
	if lease.runID != "" {
		delete(m.leaseByRun, lease.runID)
	}
	lease.state = leaseStateReleased
	m.decrementSlotLocked(lease.groupKey, lease.nodeID, lease.providerID, lease.long)
	return true
}

// fenceNodeGenerationLocked fences the disconnected connection identified by
// (nodeID, generation): it settles leases through the exactly-once release path
// so each provider resource counter is returned per lease, and marks matching
// provider resources orphan. Under a generation-scoped fence a strictly newer
// reconnect's leases and resources are left untouched; under a whole-node fence
// (generation 0) it also zeroes any residual counters seeded by fixtures
// without a backing lease. Returns whether any lease was settled so the caller
// can tell an authoritative fence from a stale no-op. Must be called with m.mu
// held.
func (m *modelQueueManager) fenceNodeGenerationLocked(nodeID string, generation uint64) bool {
	fenceAll := generation == 0
	settledLease := m.settleLeasesForNodeLocked(nodeID, generation)

	for key, res := range m.resources {
		if key.nodeID != nodeID {
			continue
		}
		if !fenceAll && res.generation > generation {
			continue
		}
		if fenceAll {
			res.inFlight = 0
			res.longInFlight = 0
		}
		res.orphan = true
	}
	return settledLease
}

// resolveAndPumpAllLocked rebuilds every queued item's candidate universe
// through its live resolver (or filters orphaned and fenced-node snapshot
// candidates when no resolver is attached) and then runs the global pump.
// Terminal no-candidate items are settled immediately so the pump only deals
// with dispatchable or temporarily-blocked waiters. excludeNodeID, excludeGen
// and fenceAll describe the disconnect context so snapshot items that lack a
// live resolver are also purged of the fenced node — the same signal the old
// candidate-snapshot removal loop used — while provider-pool items with a
// resolver get a fresh resolution against current state. This is the
// authoritative disconnect settlement pass: surviving-provider fallback,
// terminal unavailable, and normal dispatch all resolve in the same
// deterministic global pass, independent of event bus delivery. Must be called
// with m.mu held.
func (m *modelQueueManager) resolveAndPumpAllLocked(excludeNodeID string, excludeGen uint64, fenceAll bool) {
	for _, group := range m.groups {
		for _, item := range group.queue {
			if item.resolveCandidates != nil {
				candidates, outcome, _ := m.resolveQueuedCandidatesLocked(item)
				if outcome == resolveNoCandidates {
					// Terminal no-candidate: remove immediately, before the pump,
					// so the dispatch pass only sees items that can actually run.
					m.removeQueuedItemLocked(group, item)
					select {
					case item.waitCh <- admitResult{err: fmt.Errorf("model group %q: %w", group.key, errProviderUnavailable)}:
					default:
					}
					continue
				}
				if outcome == resolveResolverError {
					// Leave the item queued for a later pump; the resolver fault
					// is recoverable and must not block the dispatch pass.
					continue
				}
				if candidates != nil {
					item.candidates = candidates
				}
				continue
			}
			// Snapshot-candidate path: filter orphaned resources AND the fenced
			// node itself so a pump dispatch cannot hand them to an offline
			// connection. A strictly newer reconnect on the same node is
			// preserved (generation-scoped fence).
			filtered := make([]candidateNode, 0, len(item.candidates))
			for _, c := range item.candidates {
				if c.entry == nil {
					filtered = append(filtered, c)
					continue
				}
				if excludeNodeID != "" && (fenceAll || c.generation <= excludeGen) && c.entry.NodeID == excludeNodeID {
					continue
				}
				key := providerResourceKey{nodeID: c.entry.NodeID, providerID: c.providerID}
				if r, ok := m.resources[key]; ok && r.orphan {
					continue
				}
				filtered = append(filtered, c)
			}
			item.candidates = filtered
		}
	}
	// One global pump after every group's candidates are refreshed, so a waiter
	// can only be handed a provider that is still live, and a waiter in any
	// group can pick up the remaining ones.
	m.pumpAllLocked()
}

// releaseNode fences the disconnected connection identified by (nodeID,
// generation): it drops that connection's leases, zeroes the matching provider
// resources to offline, rebuilds every queued item's candidate universe through
// its live resolver, and runs one global pump so surviving-provider fallback,
// terminal unavailable, and normal dispatch all resolve in the same pass. A
// generation of 0 fences the node unconditionally (untracked/legacy disconnect).
// A specific generation fences only resources whose current owner generation is
// that one or older: a stale callback for an already-superseded connection
// cannot zero the slot a live reconnect now holds, so its live resource and
// events are left unchanged.
func (m *modelQueueManager) releaseNode(nodeID string, generation uint64, reason string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	fenceAll := generation == 0

	// Fence the generation: settle leases and mark matching resources orphan.
	settledLease := m.fenceNodeGenerationLocked(nodeID, generation)

	// A generation-scoped callback that fenced neither a lease nor a resource is
	// a stale close for an already-superseded connection: the live owner's state
	// and queued candidate lists are authoritative, so leave them untouched.
	if !fenceAll && !settledLease {
		return
	}

	// A whole-node fence (generation 0) clears any residual legacy/provider slot
	// counters this node still holds — legacy (nodeID-only) fixtures and any
	// counter a lease did not account for. Provider-pool counters live on the
	// resource state already marked orphan above.
	if fenceAll {
		for _, group := range m.groups {
			for slot := range group.inflight {
				if colonIdx := findLastColon(slot); colonIdx > 0 {
					if slot[:colonIdx] == nodeID {
						delete(group.inflight, slot)
						delete(group.longInflight, slot)
					}
				} else if slot == nodeID {
					delete(group.inflight, slot)
					delete(group.longInflight, slot)
				}
			}
		}
	}

	// Rebuild every queued item's candidate universe through the live resolver
	// and run one global pump so that surviving-provider fallback, terminal
	// unavailable, and normal dispatch all resolve in the same deterministic
	// pass. This replaces the previous candidate-snapshot removal loop: items
	// with a live resolver get re-evaluated against the current store/catalog/
	// registry state (including orphan filtering), and items without a live
	// resolver have the fenced node and orphaned resources purged from their
	// candidate lists, so the pump sees only live targets and can settle
	// terminal no-candidate items directly.
	m.resolveAndPumpAllLocked(nodeID, generation, fenceAll)
}

// activateNode is the reconnect counterpart to releaseNode: it restores the
// provider resources owned by an accepted connection identified by (nodeID,
// generation) to an available state for that generation, then rebuilds every
// queued item's candidate universe through its live resolver and runs one global
// pump. A disconnect marked those resources orphan and stranded the waiters
// queued against them; the reconnect alone re-activates the resources and
// re-dispatches the waiters in global enqueue order, with no new request, config
// refresh, or lease release to trigger the pump.
//
// No disconnect context is passed to the rebuild pass (excludeNodeID empty,
// fenceAll false), so nothing is purged: the reconnected node's candidates are
// made live and dispatched, and every other group's waiters are re-evaluated
// against current state in the same deterministic pass.
//
// isCurrentOwner is the ownership linearization gate. It is evaluated under m.mu,
// in the same critical section that mutates the resources and pumps, and it is
// serialized against releaseNode (which also holds m.mu). A ready callback whose
// connection a disconnect settled — unregistered its entry and released its
// leases — between the transport handshake and this activation therefore observes
// the disconnect and is a no-op: the check and the activation cannot straddle the
// disconnect the way a check performed before m.mu was acquired could. The
// callback itself reads the registry (m.mu → registry-lock order, matching the
// pump's live resolver); the registry lock is never held across m.mu, so no lock
// inversion is introduced. A nil callback is treated as always-current for
// untracked/legacy callers and fixtures.
func (m *modelQueueManager) activateNode(nodeID string, generation uint64, isCurrentOwner func() bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if isCurrentOwner != nil && !isCurrentOwner() {
		// A stale/superseded generation: the disconnect that removed this
		// connection already settled its resources under this same lock, and
		// re-activating them here would resurrect a slot the live owner (or none)
		// now holds.
		return
	}
	m.activateNodeGenerationLocked(nodeID, generation)
	m.resolveAndPumpAllLocked("", 0, false)
}

// settleLeasesForNodeLocked returns the resources held by the leases a disconnect
// fences on a node, routing each through the exactly-once releaseLeaseLocked
// decrement path so per-lease provider counters are returned — never zeroed
// wholesale. A generation of 0 fences every lease on the node (untracked/legacy
// disconnect); a specific generation fences that owner and any older one, leaving
// a strictly newer reconnect's leases — and the shared resource counters they
// still hold — untouched. It reports whether it settled any lease so the caller
// can tell an authoritative fence from a stale no-op. Must be called with m.mu
// held.
func (m *modelQueueManager) settleLeasesForNodeLocked(nodeID string, generation uint64) bool {
	fenceAll := generation == 0
	// Collect first: releaseLeaseLocked mutates m.leases, so it cannot run inside
	// the range over the same map.
	matching := make([]uint64, 0, len(m.leases))
	for id, lease := range m.leases {
		if lease.nodeID != nodeID {
			continue
		}
		if !fenceAll && lease.generation > generation {
			continue
		}
		matching = append(matching, id)
	}
	for _, id := range matching {
		m.releaseLeaseLocked(id)
	}
	return len(matching) > 0
}

// releaseRun releases the lease owning a terminated run and dispatches the next
// queued item if one is waiting. Safe to call for an unknown or already-released
// run: the lease index resolves nothing and the call is a no-op.
func (m *modelQueueManager) releaseRun(runID, reason string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	leaseID, ok := m.leaseByRun[runID]
	if !ok {
		return
	}
	if m.leases[leaseID] == nil {
		delete(m.leaseByRun, runID)
		return
	}
	if !m.releaseLeaseLocked(leaseID) {
		return
	}
	m.pumpAllLocked()
}

// 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) {
	if group := m.decrementSlotLocked(groupKey, nodeID, providerID, longSlot); group == nil {
		return
	}
	m.pumpAllLocked()
}

// decrementSlotLocked frees the resources a reservation holds without pumping the
// queue, returning the group so the caller can decide whether to dispatch. Must
// be called with m.mu held.
func (m *modelQueueManager) decrementSlotLocked(groupKey, nodeID, providerID string, longSlot bool) *modelQueueGroup {
	group, ok := m.groups[groupKey]
	if !ok {
		return nil
	}
	// Use provider-aware slot key for provider-pool dispatches.
	slot := nodeID
	if providerID != "" {
		slot = nodeID + ":" + providerID
		key := providerResourceKey{nodeID: nodeID, providerID: providerID}
		if res, ok := m.resources[key]; ok {
			res.release(longSlot)
			if res.orphan && res.inFlight == 0 && res.longInFlight == 0 {
				delete(m.resources, key)
			}
		} else {
			if group.inflight[slot] > 0 {
				group.inflight[slot]--
			}
			if longSlot && group.longInflight[slot] > 0 {
				group.longInflight[slot]--
			}
		}
	} else {
		if group.inflight[slot] > 0 {
			group.inflight[slot]--
		}
		if longSlot && group.longInflight[slot] > 0 {
			group.longInflight[slot]--
		}
	}
	return group
}