1// SPDX-License-Identifier: MIT
2
3package registryproxy
4
5import (
6 "context"
7 "fmt"
8 "io"
9 "net/http"
10 "net/http/httptest"
11 "os"
12 "path/filepath"
13 "strings"
14 "sync"
15 "sync/atomic"
16 "testing"
17 "time"
18
19 "github.com/bnema/zerowrap"
20)
21
22// Concurrent requests for one digest share a single fetch.
23func TestSingleFlight(t *testing.T) {
24 s := mustOpen(t, storeCfg(t.TempDir(), 1<<20))
25 var calls atomic.Int32
26 var wg sync.WaitGroup
27 start := make(chan struct{})
28 for range 20 {
29 wg.Go(func() {
30 <-start
31 f, _, _, err := s.ensure(context.Background(), kindBlob, sum("content"), func(_ context.Context, w io.Writer) (string, error) {
32 calls.Add(1)
33 _, err := io.WriteString(w, "content")
34 return "", err
35 }, keepalive{})
36 if err != nil {
37 t.Error(err)
38 return
39 }
40 _ = f.Close()
41 })
42 }
43 close(start)
44 wg.Wait()
45 if calls.Load() != 1 {
46 t.Fatalf("fetches %d", calls.Load())
47 }
48}
49
50// A guest that requests many digests and disconnects leaves no download
51// behind: once a flight has no waiter for the grace period, it is cancelled
52// upstream and its temporary file removed.
53func TestAbandonedDownloadsCancelled(t *testing.T) {
54 u := newStallingUpstream(t)
55 const n = 20
56 p := proxyWith(t, u.srv, Config{AbandonGrace: 20 * time.Millisecond, MaxDownloads: n})
57 for i := range n {
58 serveCancelled(t, p, u, blobPath(i), int32(i+1))
59 }
60 waitFor(t, p.s, "every flight to end", func() bool { return len(p.s.flights) == 0 && p.s.downloads == 0 && p.s.pending == 0 })
61 for end := time.Now().Add(5 * time.Second); u.cancelled.Load() != n; {
62 if time.Now().After(end) {
63 t.Fatalf("upstream saw %d cancellations of %d", u.cancelled.Load(), n)
64 }
65 time.Sleep(time.Millisecond)
66 }
67 waitFor(t, p.s, "tmp/ to be empty", func() bool {
68 ds, _ := os.ReadDir(filepath.Join(p.s.dir, "tmp"))
69 return len(ds) == 0
70 })
71}
72
73// A request that joins within the grace period keeps the download.
74func TestRejoinWithinGraceKeepsDownload(t *testing.T) {
75 body := strings.Repeat("layer", 4000)
76 release := make(chan struct{})
77 var gets, cancelled atomic.Int32
78 srv := upstreamWith(t, func(w http.ResponseWriter, r *http.Request) {
79 gets.Add(1)
80 w.Header().Set("Content-Length", fmt.Sprint(len(body)))
81 _, _ = io.WriteString(w, body[:10])
82 w.(http.Flusher).Flush()
83 select {
84 case <-release:
85 _, _ = io.WriteString(w, body[10:])
86 case <-r.Context().Done():
87 cancelled.Add(1)
88 }
89 })
90 var once sync.Once
91 t.Cleanup(func() { once.Do(func() { close(release) }) })
92 p := proxyWith(t, srv, Config{AbandonGrace: time.Hour})
93 path := "/v2/library/alpine/blobs/" + sum(body)
94 ctx, cancel := context.WithCancel(context.Background())
95 first := make(chan struct{})
96 go func() {
97 defer close(first)
98 p.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, path, nil).WithContext(ctx))
99 }()
100 key := kindBlob.key(sum(body)[7:])
101 waitFor(t, p.s, "the first waiter", func() bool { return waiters(p.s, key) == 1 })
102 cancel()
103 <-first
104 w := httptest.NewRecorder()
105 second := make(chan struct{})
106 go func() { defer close(second); p.ServeHTTP(w, httptest.NewRequest(http.MethodGet, path, nil)) }()
107 waitFor(t, p.s, "the second waiter", func() bool { return waiters(p.s, key) == 1 })
108 once.Do(func() { close(release) })
109 <-second
110 if w.Code != http.StatusOK || w.Body.String() != body || gets.Load() != 1 || cancelled.Load() != 0 {
111 t.Fatalf("code %d, %d fetches, %d cancelled", w.Code, gets.Load(), cancelled.Load())
112 }
113}
114
115// The shared download does not belong to the client that started it. That
116// client disconnects mid-download; a second client, joining within the grace
117// period, still gets the blob and the blob is committed, from one upstream
118// fetch.
119func TestBlobFlightSurvivesInitiatorDisconnect(t *testing.T) {
120 body := strings.Repeat("layer", 4000)
121 half, release := make(chan struct{}), make(chan struct{})
122 var releaseOnce sync.Once
123 unblock := func() { releaseOnce.Do(func() { close(release) }) }
124 var gets atomic.Int32
125 srv := upstreamWith(t, func(w http.ResponseWriter, r *http.Request) {
126 gets.Add(1)
127 w.Header().Set("Content-Length", fmt.Sprint(len(body)))
128 _, _ = io.WriteString(w, body[:len(body)/2])
129 w.(http.Flusher).Flush()
130 close(half)
131 <-release
132 _, _ = io.WriteString(w, body[len(body)/2:])
133 })
134 t.Cleanup(unblock) // before srv.Close, which waits for the handler
135 p := proxyWith(t, srv, Config{})
136 path := "/v2/library/alpine/blobs/" + sum(body)
137
138 ctx, cancel := context.WithCancel(context.Background())
139 first := httptest.NewRecorder()
140 done := make(chan struct{})
141 go func() {
142 defer close(done)
143 p.ServeHTTP(first, httptest.NewRequest(http.MethodGet, path, nil).WithContext(ctx))
144 }()
145 <-half
146 cancel()
147 <-done
148 if strings.Contains(first.Body.String(), "layer") {
149 t.Fatal("unverified bytes sent to the initiator")
150 }
151 key := kindBlob.key(sum(body)[7:])
152 p.s.mu.Lock()
153 left := waiters(p.s, key)
154 p.s.mu.Unlock()
155 if left != 0 {
156 t.Fatalf("waiters after the disconnect: %d", left)
157 }
158
159 second := httptest.NewRecorder()
160 done = make(chan struct{})
161 go func() {
162 defer close(done)
163 p.ServeHTTP(second, httptest.NewRequest(http.MethodGet, path, nil))
164 }()
165 waitFor(t, p.s, "the second client to join", func() bool { return waiters(p.s, key) == 1 })
166 unblock()
167 <-done
168 if second.Code != http.StatusOK || second.Body.String() != body {
169 t.Fatalf("second client: %d", second.Code)
170 }
171 if _, err := os.Stat(filepath.Join(p.s.dir, key)); err != nil {
172 t.Fatalf("blob not committed: %v", err)
173 }
174 if gets.Load() != 1 {
175 t.Fatalf("%d upstream fetches", gets.Load())
176 }
177}
178
179// A client waiting for a progressing download is kept alive, so the cache
180// channel's progress deadline follows the download.
181func TestWaiterKeepAlive(t *testing.T) {
182 body := strings.Repeat("layer", 4000)
183 srv := upstreamWith(t, slowBlob(body, 10, 30*time.Millisecond))
184 p := proxyWith(t, srv, Config{KeepAliveInterval: 10 * time.Millisecond})
185 w := &keepAliveRecorder{ResponseRecorder: httptest.NewRecorder()}
186 p.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v2/library/alpine/blobs/"+sum(body), nil))
187 if w.Code != http.StatusOK || w.calls.Load() == 0 {
188 t.Fatalf("code %d, %d keepalives", w.Code, w.calls.Load())
189 }
190}
191
192// A panic in a download fails its flight and frees its slots; it does not
193// leave the waiters hanging.
194func TestPanickingFetchFailsFlight(t *testing.T) {
195 c := storeCfg(t.TempDir(), 1<<20)
196 c.maxDownloads = 1
197 s := mustOpen(t, c)
198 q := NewQuota(1)
199 ctx, cancel := context.WithTimeout(WithQuota(context.Background(), q), 5*time.Second)
200 defer cancel()
201 _, _, _, err := s.ensure(ctx, kindBlob, sum("x"), func(context.Context, io.Writer) (string, error) {
202 panic("boom")
203 }, keepalive{})
204 if err == nil || !strings.Contains(err.Error(), "panicked") {
205 t.Fatalf("err %v", err)
206 }
207 s.mu.Lock()
208 defer s.mu.Unlock()
209 if s.downloads != 0 || len(s.flights) != 0 || q.used.Load() != 0 || s.pending != 0 {
210 t.Fatalf("leaked: downloads=%d flights=%d quota=%d pending=%d", s.downloads, len(s.flights), q.used.Load(), s.pending)
211 }
212 if ds, _ := os.ReadDir(filepath.Join(s.dir, "tmp")); len(ds) != 0 {
213 t.Fatal("temp file kept")
214 }
215}
216
217// A request that joins a manifest flight logs hit=false, with its job.
218func TestJoinerLogsMiss(t *testing.T) {
219 release := make(chan struct{})
220 srv := upstreamWith(t, func(w http.ResponseWriter, r *http.Request) {
221 w.Header().Set("Content-Type", "application/vnd.oci.image.index.v1+json")
222 w.Header().Set("Docker-Content-Digest", sum(manifestBody))
223 if r.Method == http.MethodGet {
224 <-release
225 _, _ = io.WriteString(w, manifestBody)
226 }
227 })
228 var logs syncBuffer
229 p := proxyWith(t, srv, Config{Log: zerowrap.New(zerowrap.Config{Level: "info", Format: "json", Output: &logs})})
230 var wg sync.WaitGroup
231 for i := range 2 {
232 wg.Go(func() {
233 ctx := WithJob(context.Background(), fmt.Sprint("job-", i))
234 p.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/v2/library/alpine/manifests/3", nil).WithContext(ctx))
235 })
236 }
237 key := kindManifest.key(sum(manifestBody)[7:])
238 waitFor(t, p.s, "two waiters", func() bool { return waiters(p.s, key) == 2 })
239 close(release)
240 wg.Wait()
241 out := logs.String()
242 if strings.Count(out, `"hit":false`) != 2 || strings.Contains(out, `"hit":true`) {
243 t.Fatalf("logs: %s", out)
244 }
245 if !strings.Contains(out, `"job_id":"job-0"`) || !strings.Contains(out, `"job_id":"job-1"`) {
246 t.Fatalf("no job_id on the registry lines: %s", out)
247 }
248}