internal/packlimit/packlimit_test.go

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