1// SPDX-License-Identifier: MIT
2
3package registryproxy
4
5import (
6 "context"
7 "io"
8 "net/http"
9 "net/http/httptest"
10 "strings"
11 "sync"
12 "testing"
13 "time"
14)
15
16// At most MaxDownloads cold blobs download at once; one more waits for
17// SlotWait and then gets 503 (dockerd retries or pulls from Docker Hub), a
18// joiner of a running download takes no slot, and a slot frees once a
19// download ends.
20func TestColdDownloadCap(t *testing.T) {
21 u := newStallingUpstream(t)
22 const limit = 3
23 const wait = 50 * time.Millisecond
24 p := proxyWith(t, u.srv, Config{AbandonGrace: 10 * time.Millisecond, MaxDownloads: limit, SlotWait: wait})
25 ctx, cancel := context.WithCancel(context.Background())
26 var wg sync.WaitGroup
27 for i := range limit {
28 wg.Go(func() {
29 p.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, blobPath(i), nil).WithContext(ctx))
30 })
31 }
32 waitFor(t, p.s, "the cap to fill", func() bool { return p.s.downloads == limit })
33 start := time.Now()
34 w := request(t, p, http.MethodGet, blobPath(limit))
35 if w.Code != http.StatusServiceUnavailable || w.Header().Get("Retry-After") == "" {
36 t.Fatalf("over the cap: %d", w.Code)
37 }
38 if d := time.Since(start); d < wait {
39 t.Fatalf("503 after %v, before the slot wait", d)
40 }
41 // Joining a running download takes no slot: the joiner is a waiter of the
42 // flight, and the number of downloads does not move.
43 key := kindBlob.key(sum("blob-0")[7:])
44 go p.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, blobPath(0), nil).WithContext(ctx))
45 waitFor(t, p.s, "the joiner", func() bool { return waiters(p.s, key) == 2 })
46 p.s.mu.Lock()
47 downloads := p.s.downloads
48 p.s.mu.Unlock()
49 if downloads != limit {
50 t.Fatalf("%d downloads with a joiner, want %d", downloads, limit)
51 }
52 cancel()
53 wg.Wait()
54 waitFor(t, p.s, "the slots to free", func() bool { return p.s.downloads == 0 })
55 if u.started.Load() != limit {
56 t.Fatalf("%d downloads started", u.started.Load())
57 }
58}
59
60// A sandbox that saturates its quota does not hold back another sandbox: its
61// download starts at once while the global cap still has room.
62func TestQuotaIsolatesChannels(t *testing.T) {
63 u := newStallingUpstream(t)
64 p := proxyWith(t, u.srv, Config{SlotWait: time.Hour, AbandonGrace: 10 * time.Millisecond})
65 a, b := NewQuota(3), NewQuota(3)
66 ctxA, cancelA := context.WithCancel(WithQuota(context.Background(), a))
67 defer cancelA()
68 for i := range 3 {
69 serveAsync(p, ctxA, blobPath(i), nil)
70 }
71 waitFor(t, p.s, "A's quota to fill", func() bool { return a.used.Load() == 3 })
72 // A's fourth waits for its own quota and starts nothing.
73 serveAsync(p, ctxA, blobPath(3), nil)
74 ctxB, cancelB := context.WithCancel(WithQuota(context.Background(), b))
75 defer cancelB()
76 start := time.Now()
77 serveAsync(p, ctxB, blobPath(10), nil)
78 waitFor(t, p.s, "B's download to start", func() bool { return b.used.Load() == 1 })
79 if d := time.Since(start); d > time.Second {
80 t.Fatalf("B started after %v", d)
81 }
82 // The quotas are taken before the downloads reach upstream: wait for the
83 // 4 of A and B there, then check under the store lock that A's fourth
84 // took no slot.
85 waitFor(t, p.s, "4 downloads upstream", func() bool { return u.started.Load() == 4 })
86 p.s.mu.Lock()
87 downloads := p.s.downloads
88 p.s.mu.Unlock()
89 if downloads != 4 {
90 t.Fatalf("%d downloads, want A's 3 and B's 1", downloads)
91 }
92 cancelA()
93 cancelB()
94 waitFor(t, p.s, "every slot to free", func() bool { return p.s.downloads == 0 })
95 if a.used.Load() != 0 || b.used.Load() != 0 {
96 t.Fatalf("quota not released: %d %d", a.used.Load(), b.used.Load())
97 }
98}
99
100// A request over its quota waits, kept alive, and starts its download as soon
101// as one of its channel's ends, with no 503.
102func TestQuotaWaiterGetsFreedSlot(t *testing.T) {
103 body := strings.Repeat("layer", 100)
104 release := make(chan struct{})
105 var once sync.Once
106 unblock := func() { once.Do(func() { close(release) }) }
107 t.Cleanup(unblock)
108 srv := upstreamWith(t, func(w http.ResponseWriter, r *http.Request) {
109 if strings.HasSuffix(r.URL.Path, sum(body)) {
110 _, _ = io.WriteString(w, body)
111 return
112 }
113 w.Header().Set("Content-Length", "1000")
114 _, _ = io.WriteString(w, "x")
115 w.(http.Flusher).Flush()
116 select {
117 case <-release:
118 case <-r.Context().Done():
119 }
120 })
121 p := proxyWith(t, srv, Config{SlotWait: time.Hour, KeepAliveInterval: 5 * time.Millisecond})
122 q := NewQuota(1)
123 ctx := WithQuota(context.Background(), q)
124 first := serveAsync(p, ctx, blobPath(0), nil)
125 waitFor(t, p.s, "the quota to fill", func() bool { return q.used.Load() == 1 })
126 w := &keepAliveRecorder{ResponseRecorder: httptest.NewRecorder()}
127 second := serveAsync(p, ctx, "/v2/library/alpine/blobs/"+sum(body), w)
128 waitKeptAlive(t, w) // waiting for a slot
129 select {
130 case code := <-second:
131 t.Fatalf("did not wait for a slot: %d", code)
132 default:
133 }
134 unblock() // the first download ends (digest mismatch), freeing the slot
135 if code := <-first; code != http.StatusBadGateway {
136 t.Fatalf("first: %d", code)
137 }
138 if code := <-second; code != http.StatusOK || w.Body.String() != body {
139 t.Fatalf("second: %d", code)
140 }
141}
142
143// After SlotWait without a slot, the request gets 503.
144func TestQuotaWaitTimesOut(t *testing.T) {
145 u := newStallingUpstream(t)
146 const wait = 50 * time.Millisecond
147 p := proxyWith(t, u.srv, Config{SlotWait: wait, AbandonGrace: 10 * time.Millisecond})
148 q := NewQuota(1)
149 ctx, cancel := context.WithCancel(WithQuota(context.Background(), q))
150 defer cancel()
151 serveAsync(p, ctx, blobPath(0), nil)
152 waitFor(t, p.s, "the quota to fill", func() bool { return q.used.Load() == 1 })
153 start := time.Now()
154 code := <-serveAsync(p, ctx, blobPath(1), nil)
155 if code != http.StatusServiceUnavailable {
156 t.Fatalf("over the quota: %d", code)
157 }
158 if d := time.Since(start); d < wait {
159 t.Fatalf("503 after %v", d)
160 }
161 if u.started.Load() != 1 {
162 t.Fatalf("%d downloads started", u.started.Load())
163 }
164}
165
166// A request that disconnects while it waits for a slot leaves at once, with no
167// slot and no 503.
168func TestQuotaWaiterDisconnects(t *testing.T) {
169 u := newStallingUpstream(t)
170 const wait = time.Hour
171 p := proxyWith(t, u.srv, Config{SlotWait: wait, AbandonGrace: 10 * time.Millisecond, KeepAliveInterval: 5 * time.Millisecond})
172 q := NewQuota(1)
173 ctx, cancel := context.WithCancel(WithQuota(context.Background(), q))
174 defer cancel()
175 serveAsync(p, ctx, blobPath(0), nil)
176 waitFor(t, p.s, "the quota to fill", func() bool { return q.used.Load() == 1 })
177 wctx, wcancel := context.WithCancel(WithQuota(context.Background(), q))
178 w := &keepAliveRecorder{ResponseRecorder: httptest.NewRecorder()}
179 waiting := serveAsync(p, wctx, blobPath(1), w)
180 waitKeptAlive(t, w) // really waiting for a slot
181 start := time.Now()
182 wcancel()
183 code := <-waiting
184 if d := time.Since(start); d > time.Second {
185 t.Fatalf("left after %v", d)
186 }
187 if code == http.StatusServiceUnavailable {
188 t.Fatal("a disconnected waiter got 503")
189 }
190 if u.started.Load() != 1 || q.used.Load() != 1 {
191 t.Fatalf("%d downloads started, quota %d", u.started.Load(), q.used.Load())
192 }
193}
194
195// Joining a download in progress takes none of the joiner's quota.
196func TestQuotaJoinersExempt(t *testing.T) {
197 u := newStallingUpstream(t)
198 p := proxyWith(t, u.srv, Config{SlotWait: 20 * time.Millisecond, AbandonGrace: 10 * time.Millisecond})
199 a, b := NewQuota(1), NewQuota(1)
200 ctxA, cancelA := context.WithCancel(WithQuota(context.Background(), a))
201 defer cancelA()
202 serveAsync(p, ctxA, blobPath(0), nil)
203 waitFor(t, p.s, "A's download", func() bool { return a.used.Load() == 1 })
204 ctxB, cancelB := context.WithCancel(WithQuota(context.Background(), b))
205 defer cancelB()
206 serveAsync(p, ctxB, blobPath(0), nil) // joins A's
207 key := kindBlob.key(sum("blob-0")[7:])
208 waitFor(t, p.s, "B to join", func() bool { return waiters(p.s, key) == 2 })
209 if b.used.Load() != 0 {
210 t.Fatal("a joiner took quota")
211 }
212 // B's quota is still free for a download of its own.
213 serveAsync(p, ctxB, blobPath(1), nil)
214 waitFor(t, p.s, "B's own download", func() bool { return b.used.Load() == 1 })
215}