quota.go RawHistoryFolder

1// SPDX-License-Identifier: MIT
2
3package registryproxy
4
5import (
6	"context"
7	"sync/atomic"
8)
9
10// Quota bounds the cold blob downloads started on behalf of one client, such as
11// the channel of one sandbox, so that no client can take every download slot of
12// the Proxy. A download counts against the quota of the request that started it
13// until it really ends, even when only other clients still wait for it;
14// requests that join a download in progress take none.
15type Quota struct {
16	max  int64
17	used atomic.Int64
18}
19
20// NewQuota returns a quota of n concurrent downloads; n must be positive.
21func NewQuota(n int) *Quota { return &Quota{max: int64(max(n, 1))} }
22
23// WithQuota attaches q to the context of a request to the Proxy. A request
24// without a quota is bound by the Proxy's global cap only.
25func WithQuota(ctx context.Context, q *Quota) context.Context {
26	return context.WithValue(ctx, quotaKey{}, q)
27}
28
29type quotaKey struct{}
30
31func quotaFrom(ctx context.Context) *Quota {
32	q, _ := ctx.Value(quotaKey{}).(*Quota)
33	return q
34}
35
36// WithJob attaches the id of the job a request serves to its context: the
37// Proxy logs it as job_id on its registry lines.
38func WithJob(ctx context.Context, job string) context.Context {
39	return context.WithValue(ctx, jobKey{}, job)
40}
41
42type jobKey struct{}
43
44func jobFrom(ctx context.Context) string {
45	job, _ := ctx.Value(jobKey{}).(string)
46	return job
47}
48
49func (q *Quota) tryAcquire() bool {
50	for {
51		n := q.used.Load()
52		if n >= q.max {
53			return false
54		}
55		if q.used.CompareAndSwap(n, n+1) {
56			return true
57		}
58	}
59}
60
61func (q *Quota) release() { q.used.Add(-1) }