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

package service

import (
	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
}

// dropLeasesForNodeLocked forgets every lease held on a node without decrementing
// per-lease counters; the caller zeroes that node's resource state wholesale.
// Must be called with m.mu held.
func (m *modelQueueManager) dropLeasesForNodeLocked(nodeID string) {
	for id, lease := range m.leases {
		if lease.nodeID != nodeID {
			continue
		}
		delete(m.leases, id)
		if lease.runID != "" {
			delete(m.leaseByRun, lease.runID)
		}
		lease.state = leaseStateReleased
	}
}

// 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()
}

// 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()

	m.dropLeasesForNodeLocked(nodeID)

	for key, res := range m.resources {
		if key.nodeID == nodeID {
			res.inFlight = 0
			res.longInFlight = 0
			res.orphan = true
		}
	}

	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".
		// Long-context in-flight counters are keyed identically, so clear them
		// alongside the normal counters.
		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)
					delete(group.longInflight, slot)
					continue
				}
			} else {
				// Legacy slot: exact nodeID match.
				if slot == nodeID {
					delete(group.inflight, slot)
					delete(group.longInflight, slot)
				}
			}
		}

	}

	// One global pump after every group has had the disconnected node stripped
	// from its candidate lists, 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()
}

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