llmgw/internal/router/admission/queue.go

package admission

import "time"

// waiter is one queued request.
type waiter struct {
	tenant   string
	model    string
	ready    chan struct{}
	enqueued time.Time
	deadline time.Time
	// admitted guards the race between the dispatcher admitting a waiter and
	// the waiter giving up on its own.
	admitted bool
	// err is why the waiter was turned away while queued — displaced by a
	// higher band — set before ready is closed.
	err error
}

// bandQueue holds one band's waiters, grouped by tenant.
//
// Grouping is what makes fairness possible: a single tenant sending a thousand
// requests occupies one slot in the rotation, not a thousand, so a traffic
// spike from one caller cannot consume the band.
type bandQueue struct {
	band Band

	byTenant map[string][]*waiter
	// rotation is the round-robin order of tenants with pending work.
	rotation []string
	next     int
	count    int
}

func newBandQueue(b Band) *bandQueue {
	return &bandQueue{band: b, byTenant: map[string][]*waiter{}}
}

func (q *bandQueue) len() int { return q.count }

func (q *bandQueue) push(w *waiter) {
	if _, ok := q.byTenant[w.tenant]; !ok {
		q.byTenant[w.tenant] = nil
		q.rotation = append(q.rotation, w.tenant)
	}
	q.byTenant[w.tenant] = append(q.byTenant[w.tenant], w)
	q.count++
}

// peek returns the waiter the next pop would return, without removing it.
func (q *bandQueue) peek() *waiter {
	if q.count == 0 {
		return nil
	}
	idx, tenant := q.nextTenant()
	if idx < 0 {
		return nil
	}
	_ = tenant
	return q.byTenant[q.rotation[idx]][0]
}

// pop removes and returns the next waiter in tenant rotation.
func (q *bandQueue) pop() *waiter {
	if q.count == 0 {
		return nil
	}
	idx, _ := q.nextTenant()
	if idx < 0 {
		return nil
	}

	tenant := q.rotation[idx]
	ws := q.byTenant[tenant]
	w := ws[0]
	q.byTenant[tenant] = ws[1:]
	q.count--

	if len(q.byTenant[tenant]) == 0 {
		delete(q.byTenant, tenant)
		q.rotation = append(q.rotation[:idx], q.rotation[idx+1:]...)
		if idx < q.next {
			q.next--
		}
	} else {
		q.next = idx + 1
	}
	if len(q.rotation) > 0 {
		q.next %= len(q.rotation)
	} else {
		q.next = 0
	}
	return w
}

// nextTenant finds the next rotation entry that still has waiters.
func (q *bandQueue) nextTenant() (int, string) {
	n := len(q.rotation)
	for i := range n {
		idx := (q.next + i) % n
		if len(q.byTenant[q.rotation[idx]]) > 0 {
			return idx, q.rotation[idx]
		}
	}
	return -1, ""
}

// remove takes a specific waiter out, for cancellation and TTL expiry.
func (q *bandQueue) remove(target *waiter) bool {
	ws, ok := q.byTenant[target.tenant]
	if !ok {
		return false
	}
	for i, w := range ws {
		if w != target {
			continue
		}
		q.byTenant[target.tenant] = append(ws[:i:i], ws[i+1:]...)
		q.count--
		if len(q.byTenant[target.tenant]) == 0 {
			delete(q.byTenant, target.tenant)
			q.dropTenant(target.tenant)
		}
		return true
	}
	return false
}

func (q *bandQueue) dropTenant(tenant string) {
	for i, t := range q.rotation {
		if t != tenant {
			continue
		}
		q.rotation = append(q.rotation[:i], q.rotation[i+1:]...)
		if i < q.next {
			q.next--
		}
		break
	}
	if len(q.rotation) > 0 {
		q.next %= len(q.rotation)
	} else {
		q.next = 0
	}
}

// dropExpired removes waiters past their deadline, calling onDrop for each.
// Their Acquire call is already unblocking on its own timer; this keeps the
// queue from carrying corpses that would occupy capacity.
func (q *bandQueue) dropExpired(now time.Time, onDrop func()) {
	if q.count == 0 {
		return
	}
	for tenant, ws := range q.byTenant {
		kept := ws[:0]
		for _, w := range ws {
			if now.After(w.deadline) && !w.admitted {
				q.count--
				onDrop()
				continue
			}
			kept = append(kept, w)
		}
		if len(kept) == 0 {
			delete(q.byTenant, tenant)
			q.dropTenant(tenant)
			continue
		}
		q.byTenant[tenant] = kept
	}
}

// newest is the most recently queued waiter, across every tenant: the one
// displacement takes, since it has invested the least waiting.
func (q *bandQueue) newest() *waiter {
	var out *waiter
	for _, ws := range q.byTenant {
		for _, w := range ws {
			if out == nil || w.enqueued.After(out.enqueued) {
				out = w
			}
		}
	}
	return out
}