internal/packlimit/packlimit_test.go

  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}
316
317// Two limiters log apart: a named one says what it limits, and its
318// once-a-minute window does not silence the other's.
319func TestRefusedNamesTheLimit(t *testing.T) {
320	var buf bytes.Buffer
321	old := slog.Default()
322	slog.SetDefault(slog.New(slog.NewTextHandler(&buf, nil)))
323	t.Cleanup(func() { slog.SetDefault(old) })
324
325	packs := New(1, 0, 0, time.Second)
326	pushes := New(1, 0, 0, time.Second)
327	pushes.Name("push")
328	packs.Refused("ssh", "user:4", ErrBusy)
329	pushes.Refused("ssh", "key:9", ErrBusy)
330	out := buf.String()
331	if !strings.Contains(out, `"pack limit: request turned away`) || !strings.Contains(out, `"push limit: request turned away`) {
332		t.Fatalf("want one line per limit:\n%s", out)
333	}
334	if !strings.Contains(out, "class=key") || strings.Contains(out, "key:9") {
335		t.Fatalf("deploy key principal logged or class missing:\n%s", out)
336	}
337	var none *Limiter
338	none.Name("push")
339}
340
341// With a waiting cap above per, one principal runs per and queues up to
342// perQueue, and is refused past that.
343func TestPerPrincipalQueueCap(t *testing.T) {
344	l := New(4, 1, 16, 5*time.Second)
345	l.CapQueue(4)
346	r, err := l.Acquire(nil, "a")
347	if err != nil {
348		t.Fatal(err)
349	}
350	for i := 1; i <= 4; i++ {
351		go func() {
352			if r, err := l.Acquire(nil, "a"); err == nil {
353				r()
354			}
355		}()
356		waitQueued(t, l, i)
357	}
358	if _, err := l.Acquire(nil, "a"); !errors.Is(err, ErrBusy) {
359		t.Fatalf("sixth for a: %v", err)
360	}
361	rb, err := l.Acquire(nil, "b")
362	if err != nil {
363		t.Fatalf("b blocked by a's queue: %v", err)
364	}
365	rb()
366	r()
367}