internal/packlimit/packlimit_test.go
315 lines · 7594 bytes
1package packlimit
2
3import (
4 "bytes"
5 "errors"
6 "log/slog"
7 "math"
8 "strings"
9 "testing"
10 "time"
11)
12
13func TestGlobalCap(t *testing.T) {
14 l := New(2, 0, 0, time.Second)
15 r1, err1 := l.Acquire(nil, "a")
16 r2, err2 := l.Acquire(nil, "b")
17 if err1 != nil || err2 != nil {
18 t.Fatal(err1, err2)
19 }
20 if _, err := l.Acquire(nil, "c"); !errors.Is(err, ErrBusy) {
21 t.Fatalf("third with no queue: %v", err)
22 }
23 r1()
24 r1() // a second release is a no-op
25 r3, err := l.Acquire(nil, "c")
26 if err != nil {
27 t.Fatal(err)
28 }
29 if _, err := l.Acquire(nil, "d"); !errors.Is(err, ErrBusy) {
30 t.Fatalf("double release freed two slots: %v", err)
31 }
32 r2()
33 r3()
34}
35
36func TestPerPrincipalCap(t *testing.T) {
37 l := New(4, 1, 4, 50*time.Millisecond)
38 ra, err := l.Acquire(nil, "a")
39 if err != nil {
40 t.Fatal(err)
41 }
42 defer ra()
43 if _, err := l.Acquire(nil, "a"); !errors.Is(err, ErrBusy) {
44 t.Fatalf("second for a: %v", err)
45 }
46 rb, err := l.Acquire(nil, "b")
47 if err != nil {
48 t.Fatalf("b blocked by a: %v", err)
49 }
50 rb()
51}
52
53func TestWaiterGetsReleasedSlot(t *testing.T) {
54 l := New(1, 0, 1, 5*time.Second)
55 r1, _ := l.Acquire(nil, "a")
56 got := make(chan error, 1)
57 go func() {
58 r, err := l.Acquire(nil, "b")
59 if err == nil {
60 r()
61 }
62 got <- err
63 }()
64 waitQueued(t, l, 1)
65 r1()
66 select {
67 case err := <-got:
68 if err != nil {
69 t.Fatal(err)
70 }
71 case <-time.After(2 * time.Second):
72 t.Fatal("waiter never got the slot")
73 }
74}
75
76func TestQueueIsBounded(t *testing.T) {
77 l := New(1, 0, 1, 5*time.Second)
78 r1, _ := l.Acquire(nil, "a")
79 defer r1()
80 go l.Acquire(nil, "b")
81 waitQueued(t, l, 1)
82 if _, err := l.Acquire(nil, "c"); !errors.Is(err, ErrBusy) {
83 t.Fatalf("queue over its bound: %v", err)
84 }
85}
86
87// A principal cannot fill the queue on its own.
88func TestPrincipalQueueIsBounded(t *testing.T) {
89 l := New(1, 1, 8, 5*time.Second)
90 r1, _ := l.Acquire(nil, "x")
91 defer r1()
92 go l.Acquire(nil, "a")
93 waitQueued(t, l, 1)
94 if _, err := l.Acquire(nil, "a"); !errors.Is(err, ErrBusy) {
95 t.Fatalf("second waiter for a: %v", err)
96 }
97}
98
99func TestClientGoneWhileQueued(t *testing.T) {
100 l := New(1, 0, 1, 5*time.Second)
101 r1, _ := l.Acquire(nil, "a")
102 done := make(chan struct{})
103 close(done)
104 if _, err := l.Acquire(done, "b"); !errors.Is(err, ErrGone) {
105 t.Fatalf("got %v, want ErrGone", err)
106 }
107 assertQueueEmpty(t, l)
108 r1()
109 assertHeldEmpty(t, l)
110}
111
112func TestWaitRunsOut(t *testing.T) {
113 l := New(1, 0, 1, 20*time.Millisecond)
114 r1, _ := l.Acquire(nil, "a")
115 if _, err := l.Acquire(nil, "b"); !errors.Is(err, ErrBusy) {
116 t.Fatalf("got %v, want ErrBusy", err)
117 }
118 assertQueueEmpty(t, l)
119 r1()
120 assertHeldEmpty(t, l)
121}
122
123func assertQueueEmpty(t *testing.T, l *Limiter) {
124 t.Helper()
125 l.mu.Lock()
126 defer l.mu.Unlock()
127 if l.queued != 0 || len(l.waiting) != 0 {
128 t.Fatalf("queue not cleaned up: queued=%d waiting=%v", l.queued, l.waiting)
129 }
130}
131
132func assertHeldEmpty(t *testing.T, l *Limiter) {
133 t.Helper()
134 l.mu.Lock()
135 defer l.mu.Unlock()
136 if len(l.held) != 0 {
137 t.Fatalf("held not cleaned up: %v", l.held)
138 }
139}
140
141func TestNilLimiterNeverWaits(t *testing.T) {
142 var l *Limiter
143 if l = New(0, 1, 1, time.Second); l != nil {
144 t.Fatal("max 0 should mean no limit")
145 }
146 r, err := l.Acquire(nil, "a")
147 if err != nil {
148 t.Fatal(err)
149 }
150 r()
151}
152
153// A waiter whose done channel closes just as a slot frees up must not be
154// granted the slot: it has to see ErrGone, and the slot must go to
155// someone else instead of leaking to an abandoned caller.
156func TestGivenUpWaiterNeverGetsSlot(t *testing.T) {
157 l := New(1, 0, 1, 10*time.Second)
158 _, _ = l.Acquire(nil, "a") // holds the only slot
159
160 done := make(chan struct{})
161 got := make(chan error, 1)
162 go func() {
163 r, err := l.Acquire(done, "b")
164 if err == nil {
165 r()
166 }
167 got <- err
168 }()
169 waitQueued(t, l, 1)
170
171 // Close done and free a's slot in the same critical section, so
172 // changed and done become ready to b's select at the same instant
173 // — the exact race the done-check-before-fits ordering in Acquire
174 // has to win, whichever the select picks.
175 l.mu.Lock()
176 close(done)
177 l.running--
178 delete(l.held, "a")
179 close(l.changed)
180 l.changed = make(chan struct{})
181 l.mu.Unlock()
182
183 select {
184 case err := <-got:
185 if !errors.Is(err, ErrGone) {
186 t.Fatalf("got %v, want ErrGone", err)
187 }
188 case <-time.After(2 * time.Second):
189 t.Fatal("b never returned")
190 }
191
192 // The slot must still be free for someone else: b must not hold it.
193 r2, err := l.Acquire(nil, "c")
194 if err != nil {
195 t.Fatalf("slot leaked to the abandoned waiter: %v", err)
196 }
197 r2()
198}
199
200func waitQueued(t *testing.T, l *Limiter, n int) {
201 t.Helper()
202 deadline := time.Now().Add(2 * time.Second)
203 for time.Now().Before(deadline) {
204 l.mu.Lock()
205 q := l.queued
206 l.mu.Unlock()
207 if q == n {
208 return
209 }
210 time.Sleep(time.Millisecond)
211 }
212 t.Fatalf("queue never reached %d", n)
213}
214
215// config maps pack_queue = -1 to math.MaxInt: waiters are not turned
216// away for want of queue room.
217func TestUnboundedQueue(t *testing.T) {
218 l := New(1, 0, math.MaxInt, 5*time.Second)
219 r1, _ := l.Acquire(nil, "a")
220 done := make(chan struct{})
221 defer close(done)
222 for i := 0; i < 64; i++ {
223 go l.Acquire(done, "b")
224 }
225 waitQueued(t, l, 64)
226 r1()
227}
228
229// Anonymous clients cannot hold every slot: with max 3 and "ip:" capped
230// at 2, a third anonymous request queues and an account still gets in.
231func TestClassCap(t *testing.T) {
232 l := New(3, 0, 4, 5*time.Second)
233 l.CapClass("ip:", 2)
234 r1, err1 := l.Acquire(nil, "ip:192.0.2.1")
235 r2, err2 := l.Acquire(nil, "ip:192.0.2.2")
236 if err1 != nil || err2 != nil {
237 t.Fatal(err1, err2)
238 }
239 got := make(chan error, 1)
240 go func() {
241 r, err := l.Acquire(nil, "ip:192.0.2.3")
242 if err == nil {
243 r()
244 }
245 got <- err
246 }()
247 waitQueued(t, l, 1)
248 ru, err := l.Acquire(nil, "user:1")
249 if err != nil {
250 t.Fatalf("account refused the free slot: %v", err)
251 }
252 select {
253 case err := <-got:
254 t.Fatalf("third anonymous request did not queue: %v", err)
255 default:
256 }
257 r1()
258 select {
259 case err := <-got:
260 if err != nil {
261 t.Fatal(err)
262 }
263 case <-time.After(2 * time.Second):
264 t.Fatal("queued anonymous request never got the freed slot")
265 }
266 r2()
267 ru()
268 assertHeldEmpty(t, l)
269 if l.classHeld != 0 {
270 t.Fatalf("classHeld = %d after every release", l.classHeld)
271 }
272}
273
274func TestAddrPrincipal(t *testing.T) {
275 for in, want := range map[string]string{
276 "192.0.2.7": "ip:192.0.2.7",
277 "::ffff:192.0.2.7": "ip:192.0.2.7",
278 "2001:db8:1:2:3:4:5:6": "ip:2001:db8:1:2::/64",
279 "2001:db8:1:2:ffff::1": "ip:2001:db8:1:2::/64",
280 "fe80::1%en0": "ip:fe80::/64",
281 "not-an-address": "ip:not-an-address",
282 } {
283 if got := AddrPrincipal(in); got != want {
284 t.Errorf("AddrPrincipal(%q) = %q, want %q", in, got, want)
285 }
286 }
287}
288
289// A refusal is logged once a minute per transport, with the principal's
290// class and never its address.
291func TestRefusedLogsOncePerTransport(t *testing.T) {
292 var buf bytes.Buffer
293 old := slog.Default()
294 slog.SetDefault(slog.New(slog.NewTextHandler(&buf, nil)))
295 t.Cleanup(func() { slog.SetDefault(old) })
296
297 l := New(1, 0, 0, time.Second)
298 l.Refused("http", "ip:192.0.2.7", ErrBusy)
299 l.Refused("http", "ip:192.0.2.8", ErrBusy)
300 l.Refused("ssh", "user:4", ErrGone)
301 out := buf.String()
302 if n := strings.Count(out, "\n"); n != 2 {
303 t.Fatalf("%d lines, want 2:\n%s", n, out)
304 }
305 for _, want := range []string{"transport=http class=ip reason=busy", "transport=ssh class=user reason=gone"} {
306 if !strings.Contains(out, want) {
307 t.Errorf("missing %q in:\n%s", want, out)
308 }
309 }
310 if strings.Contains(out, "192.0.2") || strings.Contains(out, "user:4") {
311 t.Fatalf("principal logged:\n%s", out)
312 }
313 var none *Limiter
314 none.Refused("git", "ip:x", ErrBusy)
315}