1// SPDX-License-Identifier: MIT
2
3package cachesvc
4
5import (
6 "bytes"
7 "context"
8 "crypto/sha256"
9 "fmt"
10 "io"
11 "net/http"
12 "net/http/httptest"
13 "strings"
14 "sync"
15 "sync/atomic"
16 "testing"
17 "time"
18
19 "github.com/bnema/zerowrap"
20
21 "bnema.dev/ermite-runner/internal/registryproxy"
22)
23
24// The registry requests of one channel in progress at once are capped: one
25// over the cap gets 503 at once, the cache API of that channel still answers,
26// and another channel has its own slots.
27func TestRegistryPerChannelCap(t *testing.T) {
28 const limit = 2
29 var inside atomic.Int32
30 release := make(chan struct{})
31 svc, err := New(Config{Dir: t.TempDir(), RegistryPerChannel: limit, Registry: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
32 inside.Add(1)
33 <-release
34 _, _ = io.WriteString(w, "ok")
35 })})
36 if err != nil {
37 t.Fatal(err)
38 }
39 t.Cleanup(func() { _ = svc.CloseAll() })
40 var once sync.Once
41 unblock := func() { once.Do(func() { close(release) }) }
42 t.Cleanup(unblock)
43 a, b := channel(t, svc, "sb-a", "job-a"), channel(t, svc, "sb-b", "job-b")
44 get := func(c *http.Client) int {
45 resp, err := c.Get("http://cache" + RegistryPath + "library/alpine/blobs/x")
46 if err != nil {
47 t.Error(err)
48 return 0
49 }
50 _, _ = io.Copy(io.Discard, resp.Body)
51 _ = resp.Body.Close()
52 return resp.StatusCode
53 }
54 var wg sync.WaitGroup
55 codes := make(chan int, limit)
56 for range limit {
57 wg.Go(func() { codes <- get(a) })
58 }
59 waitUntil(t, func() bool { return inside.Load() == limit })
60 if code := get(a); code != http.StatusServiceUnavailable {
61 t.Fatalf("over the cap: %d", code)
62 }
63 if code := lookupStatus(t, a, ""); code != http.StatusForbidden {
64 t.Fatalf("cache API behind the registry cap: %d", code)
65 }
66 other := make(chan int, 1)
67 go func() { other <- get(b) }()
68 waitUntil(t, func() bool { return inside.Load() == limit+1 })
69 unblock()
70 wg.Wait()
71 close(codes)
72 for code := range codes {
73 if code != http.StatusOK {
74 t.Fatalf("within the cap: %d", code)
75 }
76 }
77 if code := <-other; code != http.StatusOK {
78 t.Fatalf("other channel: %d", code)
79 }
80 if code := get(a); code != http.StatusOK {
81 t.Fatalf("slots not released: %d", code)
82 }
83}
84
85// stallingHub fakes Docker Hub for a real registryproxy: a manifest, and blobs
86// whose download stalls until the test ends. It counts the blob downloads
87// started.
88func stallingHub(t *testing.T) (*httptest.Server, *atomic.Int32, string) {
89 t.Helper()
90 const manifest = `{"schemaVersion":2}`
91 digest := fmt.Sprintf("sha256:%x", sha256.Sum256([]byte(manifest)))
92 var started atomic.Int32
93 done := make(chan struct{})
94 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
95 switch {
96 case r.URL.Path == "/token":
97 _, _ = io.WriteString(w, `{"token":"public"}`)
98 case strings.Contains(r.URL.Path, "/manifests/"):
99 w.Header().Set("Content-Type", "application/vnd.oci.image.index.v1+json")
100 w.Header().Set("Docker-Content-Digest", digest)
101 _, _ = io.WriteString(w, manifest)
102 default:
103 started.Add(1)
104 w.Header().Set("Content-Length", "1000000")
105 _, _ = io.WriteString(w, "partial")
106 w.(http.Flusher).Flush()
107 select {
108 case <-r.Context().Done():
109 case <-done:
110 }
111 }
112 }))
113 t.Cleanup(srv.Close)
114 t.Cleanup(func() { close(done) }) // before srv.Close, which waits for the handlers
115 return srv, &started, digest
116}
117
118// Each channel has its own quota of cold downloads: a channel that has
119// started RegistryDownloadsPerChannel of them waits for its own to end, while
120// another channel starts its download at once.
121func TestRegistryQuotaPerChannel(t *testing.T) {
122 hub, started, _ := stallingHub(t)
123 proxy, err := registryproxy.New(registryproxy.Config{Dir: t.TempDir(), Max: 1 << 30, Upstream: hub.URL, AuthURL: hub.URL + "/token",
124 SlotWait: time.Hour, AbandonGrace: 10 * time.Millisecond})
125 if err != nil {
126 t.Fatal(err)
127 }
128 svc, err := New(Config{Dir: t.TempDir(), Registry: proxy})
129 if err != nil {
130 t.Fatal(err)
131 }
132 t.Cleanup(func() { _ = svc.CloseAll() })
133 ctx, cancel := context.WithCancel(context.Background())
134 defer cancel()
135 var wg sync.WaitGroup
136 defer wg.Wait()
137 defer cancel()
138 pull := func(c *http.Client, i int) {
139 req, _ := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("http://cache%slibrary/alpine/blobs/sha256:%x", RegistryPath, sha256.Sum256(fmt.Appendf(nil, "blob-%d", i))), nil)
140 wg.Go(func() {
141 if resp, err := c.Do(req); err == nil {
142 _ = resp.Body.Close()
143 }
144 })
145 }
146 a, b := channel(t, svc, "sb-a", "job-a"), channel(t, svc, "sb-b", "job-b")
147 for i := range RegistryDownloadsPerChannel + 1 {
148 pull(a, i)
149 }
150 waitUntil(t, func() bool { return started.Load() == RegistryDownloadsPerChannel })
151 time.Sleep(50 * time.Millisecond) // the one over A's quota would have started by now
152 if n := started.Load(); n != RegistryDownloadsPerChannel {
153 t.Fatalf("channel A started %d downloads, its quota is %d", n, RegistryDownloadsPerChannel)
154 }
155 pull(b, 100)
156 waitUntil(t, func() bool { return started.Load() == RegistryDownloadsPerChannel+1 })
157}
158
159// The registry lines of a channel carry its job.
160func TestRegistryLogsJob(t *testing.T) {
161 hub, _, digest := stallingHub(t)
162 var logs bytes.Buffer
163 var mu sync.Mutex
164 proxy, err := registryproxy.New(registryproxy.Config{Dir: t.TempDir(), Max: 1 << 20, Upstream: hub.URL, AuthURL: hub.URL + "/token",
165 Log: zerowrap.New(zerowrap.Config{Level: "info", Format: "json", Output: writerFunc(func(p []byte) (int, error) {
166 mu.Lock()
167 defer mu.Unlock()
168 return logs.Write(p)
169 })})})
170 if err != nil {
171 t.Fatal(err)
172 }
173 svc, err := New(Config{Dir: t.TempDir(), Registry: proxy})
174 if err != nil {
175 t.Fatal(err)
176 }
177 t.Cleanup(func() { _ = svc.CloseAll() })
178 resp, err := channel(t, svc, "sb-a", "job-a").Get("http://cache" + RegistryPath + "library/alpine/manifests/" + digest)
179 if err != nil {
180 t.Fatal(err)
181 }
182 _ = resp.Body.Close()
183 if resp.StatusCode != http.StatusOK {
184 t.Fatal(resp.StatusCode)
185 }
186 mu.Lock()
187 defer mu.Unlock()
188 if !strings.Contains(logs.String(), `"job_id":"job-a"`) || !strings.Contains(logs.String(), `"message":"registry manifest"`) {
189 t.Fatalf("logs: %s", logs.String())
190 }
191}
192
193type writerFunc func(p []byte) (int, error)
194
195func (f writerFunc) Write(p []byte) (int, error) { return f(p) }