internal/packlimit/packlimit.go

227 lines · 6350 bytes

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