crossbar/internal/limiter/limiter.go

// Package limiter hands out at most `parallel` concurrent slots per (host, model) and lets at
// most `queue_max` requests wait in a FIFO. A request that finds the queue full is refused at
// once so the caller can retry elsewhere; a waiting request can cancel and leave without leaking
// a slot or a queue place.
package limiter

import (
	"context"
	"errors"
	"sync"
	"time"
)

// ErrQueueFull is returned by Acquire when the queue is already at queue_max; the caller may try a
// different host.
var ErrQueueFull = errors.New("queue full")

// defaults is what an unconfigured (host, model) behaves as: one slot, no waiting room.
const (
	defaultParallel = 1
	defaultQueueMax = 0
)

// pair holds the live state for one (host, model): how many slots exist, how many are taken, and
// the FIFO of waiters. All fields are guarded by Limiter.mu.
type pair struct {
	parallel   int
	queueMax   int
	inflight   int
	waiters    []chan struct{}
	configured bool
}

// Limiter tracks one pair per (host, model). Safe for concurrent use.
type Limiter struct {
	mu    sync.Mutex
	pairs map[pairKey]*pair
}

type pairKey struct {
	host  string
	model string
}

// New returns an empty Limiter.
func New() *Limiter {
	return &Limiter{pairs: make(map[pairKey]*pair)}
}

// Configure sets the slot and queue limits for one (host, model). It may be called before any
// request or after one has created the pair with the defaults; either way the limits apply.
func (l *Limiter) Configure(host, model string, parallel, queueMax int) {
	l.mu.Lock()
	defer l.mu.Unlock()
	p := l.pairLocked(host, model)
	p.parallel = parallel
	p.queueMax = queueMax
	p.configured = true
}

// pairLocked returns the pair for (host, model), creating it with the unconfigured defaults if it
// does not exist yet. The caller holds l.mu.
func (l *Limiter) pairLocked(host, model string) *pair {
	k := pairKey{host, model}
	p := l.pairs[k]
	if p == nil {
		p = &pair{parallel: defaultParallel, queueMax: defaultQueueMax}
		l.pairs[k] = p
	}
	return p
}

// Acquire blocks until a slot is held. It returns a release that gives the slot back exactly once
// (a second call is a no-op), how long the caller spent in the queue, and an error: ErrQueueFull
// when the queue is already full (returned immediately, without waiting), or ctx.Err() when the
// context ends while waiting.
func (l *Limiter) Acquire(ctx context.Context, host, model string) (release func(), waited time.Duration, err error) {
	l.mu.Lock()
	start := time.Now()
	p := l.pairLocked(host, model)
	if p.inflight < p.parallel {
		p.inflight++
		l.mu.Unlock()
		return l.release(p), time.Since(start), nil
	}
	if len(p.waiters) >= p.queueMax {
		l.mu.Unlock()
		return nil, time.Since(start), ErrQueueFull
	}
	waiter := make(chan struct{})
	p.waiters = append(p.waiters, waiter)
	l.mu.Unlock()

	select {
	case <-ctx.Done():
		// The slot may have been handed to us the instant the context ended; release it either
		// way so neither a slot nor a queue place leaks.
		l.mu.Lock()
		if !p.dropWaiter(waiter) {
			l.mu.Unlock()
			l.release(p)
			return nil, time.Since(start), ctx.Err()
		}
		l.mu.Unlock()
		return nil, time.Since(start), ctx.Err()
	case <-waiter:
		return l.release(p), time.Since(start), nil
	}
}

// Track counts one request against (host, model) without waiting and without refusing: in flight
// may exceed parallel. The returned release is idempotent.
func (l *Limiter) Track(host, model string) func() {
	l.mu.Lock()
	p := l.pairLocked(host, model)
	p.inflight++
	l.mu.Unlock()
	return l.release(p)
}

// release returns the function the caller holds for a slot: it hands the slot to the next waiter
// only while there is room (in flight at or below parallel), otherwise it counts the slot back. It is
// safe to call through the sync.Once that Acquire wrapped it in. The same release serves Track, whose
// tracked load can push in flight past parallel, so a release there cannot free a slot that exists.
func (l *Limiter) release(p *pair) func() {
	var once sync.Once
	return func() {
		once.Do(func() {
			l.mu.Lock()
			defer l.mu.Unlock()
			if len(p.waiters) > 0 && p.inflight <= p.parallel {
				next := p.waiters[0]
				p.waiters = p.waiters[1:]
				close(next)
				return
			}
			p.inflight--
		})
	}
}

// dropWaiter removes w from the middle of the queue. It reports whether w was there; false means
// the slot was already handed to w (its channel closed) and the caller must give it back.
func (p *pair) dropWaiter(w chan struct{}) bool {
	for i, cw := range p.waiters {
		if cw == w {
			p.waiters = append(p.waiters[:i], p.waiters[i+1:]...)
			return true
		}
	}
	return false
}

// InFlight returns the number of held slots for (host, model).
func (l *Limiter) InFlight(host, model string) int {
	l.mu.Lock()
	defer l.mu.Unlock()
	p := l.pairs[pairKey{host, model}]
	if p == nil {
		return 0
	}
	return p.inflight
}

// Queued returns the number of requests waiting for (host, model).
func (l *Limiter) Queued(host, model string) int {
	l.mu.Lock()
	defer l.mu.Unlock()
	p := l.pairs[pairKey{host, model}]
	if p == nil {
		return 0
	}
	return len(p.waiters)
}

// FreeSlots sums the unused slots over the host's configured models, never counting below zero for
// one. An unknown host has none.
func (l *Limiter) FreeSlots(host string) int {
	l.mu.Lock()
	defer l.mu.Unlock()
	total := 0
	for k, p := range l.pairs {
		if k.host != host || !p.configured {
			continue
		}
		if free := p.parallel - p.inflight; free > 0 {
			total += free
		}
	}
	return total
}