75 lines
2.6 KiB
Go
75 lines
2.6 KiB
Go
package service
|
|
|
|
// queueReservation owns one provider-pool queue admission. Admission reserves a
|
|
// slot on the selected node/provider — and a long-context slot when the request
|
|
// is long and the provider declares a long-context limit — before any request
|
|
// is built or sent. Every dispatch path that fails after admission hands the
|
|
// reservation back through release, so the slot is freed exactly once.
|
|
//
|
|
// The reservation has two lifetimes. Before track, only the admitted slot
|
|
// exists and release frees it directly. After track, the run is registered
|
|
// inflight and release must go through the run registration instead, so the
|
|
// event watcher and the release path cannot both free the same slot.
|
|
//
|
|
// A dispatch that succeeds calls handOff: the run event watcher (normalized) or
|
|
// the tunnel terminal frame / close (tunnel) owns the release from that point.
|
|
type queueReservation struct {
|
|
queue *modelQueueManager
|
|
groupKey string
|
|
nodeID string
|
|
providerID string
|
|
long bool
|
|
|
|
runID string
|
|
released bool
|
|
}
|
|
|
|
// newQueueReservation records the slot admitWithReason just reserved for the
|
|
// selected candidate.
|
|
func newQueueReservation(queue *modelQueueManager, groupKey string, selected *candidateNode, long bool) *queueReservation {
|
|
return &queueReservation{
|
|
queue: queue,
|
|
groupKey: groupKey,
|
|
nodeID: selected.entry.NodeID,
|
|
providerID: selected.providerID,
|
|
// A long slot is reserved only for long requests on providers that
|
|
// declare a long-context limit; the release/track path must match that
|
|
// reservation.
|
|
long: long && selected.longContextCapacity > 0,
|
|
}
|
|
}
|
|
|
|
// track registers the run as inflight and moves the reservation onto the run
|
|
// id. It runs before the send so a terminal event arriving early still finds
|
|
// the registration to release.
|
|
func (res *queueReservation) track(runID string) {
|
|
if res == nil || res.queue == nil {
|
|
return
|
|
}
|
|
res.runID = runID
|
|
res.queue.trackInflight(res.groupKey, runID, res.nodeID, res.providerID, res.long)
|
|
}
|
|
|
|
// release frees the reservation exactly once. Repeated calls, and calls after
|
|
// handOff, are no-ops.
|
|
func (res *queueReservation) release(reason string) {
|
|
if res == nil || res.queue == nil || res.released {
|
|
return
|
|
}
|
|
res.released = true
|
|
if res.runID != "" {
|
|
res.queue.releaseRun(res.runID, reason)
|
|
return
|
|
}
|
|
res.queue.releaseSlotWithLong(res.groupKey, res.nodeID, res.providerID, res.long)
|
|
}
|
|
|
|
// handOff marks the reservation as owned by the dispatched run or tunnel
|
|
// lifecycle, so no later release on this object can free a slot that now
|
|
// belongs to a live request.
|
|
func (res *queueReservation) handOff() {
|
|
if res == nil {
|
|
return
|
|
}
|
|
res.released = true
|
|
}
|