internal/packlimit/packlimit.go

  1// Package packlimit bounds concurrent git pack generation. upload-pack
  2// and upload-archive over SSH, smart HTTP and git:// draw on one
  3// budget: a global cap, a cap per principal (an account, or a client
  4// address on the anonymous transports), and a bounded queue whose
  5// waiters give up after a fixed wait or when the client goes away.
  6// Waiters are not served in order; a new arrival can take a freed slot
  7// ahead of them, and the wait bounds how long any one of them waits.
  8package packlimit
  9
 10import (
 11	"errors"
 12	"net/netip"
 13	"strings"
 14	"sync"
 15	"time"
 16)
 17
 18var (
 19	ErrBusy = errors.New("the server is busy generating packs for other clients; try again in a minute")
 20	ErrGone = errors.New("client went away while queued")
 21)
 22
 23type Limiter struct {
 24	max, per, queue int
 25	wait            time.Duration
 26
 27	// Principals starting with class may hold at most classCap slots
 28	// between them; classCap 0 is no class cap.
 29	class    string
 30	classCap int
 31
 32	mu        sync.Mutex
 33	running   int
 34	classHeld int
 35	queued    int
 36	held      map[string]int // running, per principal
 37	waiting   map[string]int // queued, per principal
 38	changed   chan struct{}  // closed and replaced on every release
 39}
 40
 41// New returns a limiter, or nil — no limit — when max is not positive.
 42func New(max, per, queue int, wait time.Duration) *Limiter {
 43	if max <= 0 {
 44		return nil
 45	}
 46	return &Limiter{max: max, per: per, queue: queue, wait: wait,
 47		held: map[string]int{}, waiting: map[string]int{}, changed: make(chan struct{})}
 48}
 49
 50// CapClass caps the slots that principals starting with prefix may hold
 51// between them. Call it before the limiter is in use.
 52func (l *Limiter) CapClass(prefix string, n int) {
 53	if l == nil {
 54		return
 55	}
 56	l.class, l.classCap = prefix, n
 57}
 58
 59// AddrPrincipal is the principal for an unauthenticated client at addr:
 60// an IPv4 address as is, an IPv6 address by its /64, since one host
 61// commonly holds a whole /64. An address that does not parse is used
 62// as given.
 63func AddrPrincipal(addr string) string {
 64	a, err := netip.ParseAddr(addr)
 65	if err != nil {
 66		return "ip:" + addr
 67	}
 68	a = a.WithZone("").Unmap()
 69	if a.Is4() {
 70		return "ip:" + a.String()
 71	}
 72	return "ip:" + netip.PrefixFrom(a, 64).Masked().String()
 73}
 74
 75// Acquire takes a slot for principal, queueing when none is free.
 76// done, when it closes, ends the wait. Once Acquire returns a nil
 77// error, the caller holds the slot and must call release — once git
 78// has exited — regardless of what its own context has done since:
 79// done closing after that point does not release the slot on the
 80// caller's behalf.
 81func (l *Limiter) Acquire(done <-chan struct{}, principal string) (release func(), err error) {
 82	if l == nil {
 83		return func() {}, nil
 84	}
 85	l.mu.Lock()
 86	if l.fits(principal) {
 87		l.take(principal)
 88		l.mu.Unlock()
 89		return l.releaser(principal), nil
 90	}
 91	if l.queued >= l.queue || (l.per > 0 && l.waiting[principal] >= l.per) {
 92		l.mu.Unlock()
 93		return nil, ErrBusy
 94	}
 95	l.queued++
 96	l.waiting[principal]++
 97	l.mu.Unlock()
 98	defer func() {
 99		l.mu.Lock()
100		l.queued--
101		if l.waiting[principal]--; l.waiting[principal] == 0 {
102			delete(l.waiting, principal)
103		}
104		l.mu.Unlock()
105	}()
106
107	timer := time.NewTimer(l.wait)
108	defer timer.Stop()
109	for {
110		l.mu.Lock()
111		// changed and done can both be ready at once — a slot can
112		// free up at the same moment the caller gives up. select
113		// among the wake sources would then pick between them at
114		// random, so re-check done first, under the lock, on every
115		// pass: this makes the limiter prefer ErrGone whenever both
116		// are ready, instead of leaving it to chance which one a
117		// given pass observes.
118		select {
119		case <-done:
120			l.mu.Unlock()
121			return nil, ErrGone
122		default:
123		}
124		if l.fits(principal) {
125			l.take(principal)
126			l.mu.Unlock()
127			return l.releaser(principal), nil
128		}
129		changed := l.changed
130		l.mu.Unlock()
131		select {
132		case <-changed:
133		case <-timer.C:
134			return nil, ErrBusy
135		case <-done:
136			return nil, ErrGone
137		}
138	}
139}
140
141func (l *Limiter) fits(principal string) bool {
142	if l.inClass(principal) && l.classHeld >= l.classCap {
143		return false
144	}
145	return l.running < l.max && (l.per <= 0 || l.held[principal] < l.per)
146}
147
148func (l *Limiter) inClass(principal string) bool {
149	return l.classCap > 0 && strings.HasPrefix(principal, l.class)
150}
151
152func (l *Limiter) take(principal string) {
153	l.running++
154	l.held[principal]++
155	if l.inClass(principal) {
156		l.classHeld++
157	}
158}
159
160func (l *Limiter) releaser(principal string) func() {
161	var once sync.Once
162	return func() {
163		once.Do(func() {
164			l.mu.Lock()
165			defer l.mu.Unlock()
166			l.running--
167			if l.inClass(principal) {
168				l.classHeld--
169			}
170			if l.held[principal]--; l.held[principal] == 0 {
171				delete(l.held, principal)
172			}
173			close(l.changed)
174			l.changed = make(chan struct{})
175		})
176	}
177}