1// SPDX-License-Identifier: MIT
2
3package cachesvc
4
5import (
6 "bytes"
7 "encoding/json"
8 "fmt"
9 "io"
10 "net/http"
11 "net/url"
12 "sync"
13 "testing"
14 "time"
15)
16
17const (
18 bwRate = 200_000 // bytes per second of a limited channel
19 bwSize = 100_000 // bytes per transfer: 0.5 s at bwRate
20)
21
22// bwLocation looks up key k for the job's token on channel c and returns the
23// archive location.
24func bwLocation(t *testing.T, c *http.Client, token string) string {
25 t.Helper()
26 req, _ := http.NewRequest("GET", "http://cache"+APIPath+"cache?keys=k&version="+v1, nil)
27 req.Header.Set("Authorization", "Bearer "+token)
28 resp, err := c.Do(req)
29 if err != nil {
30 t.Fatal(err)
31 }
32 defer func() { _ = resp.Body.Close() }()
33 var e struct{ ArchiveLocation string }
34 if err := json.NewDecoder(resp.Body).Decode(&e); err != nil || resp.StatusCode != 200 {
35 t.Fatalf("lookup: %d %v", resp.StatusCode, err)
36 }
37 return e.ArchiveLocation
38}
39
40// hangGuard is the generous upper bound of a transfer: it only catches a hang,
41// never a slow machine.
42const hangGuard = 8 * time.Second
43
44// limiterOf returns the limiter installed on the channel id (nil: unlimited).
45func limiterOf(t *testing.T, svc *Service, id string) *bandwidth {
46 t.Helper()
47 svc.mu.Lock()
48 defer svc.mu.Unlock()
49 l := svc.channels[id]
50 if l == nil {
51 t.Fatalf("no channel %s", id)
52 }
53 return l.bw
54}
55
56// bucketUsed reports whether any byte was ever taken from b.
57func bucketUsed(b *bandwidth) bool {
58 b.mu.Lock()
59 defer b.mu.Unlock()
60 return !b.free.IsZero()
61}
62
63// bwDownload fetches the archive and returns how long the whole transfer took.
64func bwDownload(t *testing.T, c *http.Client, loc string) time.Duration {
65 t.Helper()
66 start := time.Now()
67 code, body := archiveStatus(t, c, "GET", loc)
68 if code != 200 || len(body) != bwSize {
69 t.Errorf("download: %d, %d bytes", code, len(body))
70 }
71 return time.Since(start)
72}
73
74// bwRig has one stored archive of bwSize bytes (job-a, key k).
75func bwRig(t *testing.T) (*rig, string) {
76 t.Helper()
77 r := newRig(t, 1<<22)
78 j := r.job("job-a", repoA, main_, main_)
79 // One chunk (save would send 4-byte ones).
80 code, id := j.reserve("k", v1, bwSize)
81 if code != http.StatusCreated {
82 t.Fatalf("reserve: %d", code)
83 }
84 if code := j.chunk(id, bytes.Repeat([]byte("x"), bwSize), 0); code != http.StatusNoContent {
85 t.Fatalf("chunk: %d", code)
86 }
87 if code := j.commit(id, bwSize); code != http.StatusNoContent {
88 t.Fatalf("commit: %d", code)
89 }
90 return r, j.token
91}
92
93// A channel with a rate serves a download at about that rate: neither much
94// faster (the limit holds) nor stalled (the transfer keeps going to the end).
95func TestBandwidthLimitsADownload(t *testing.T) {
96 r, tok := bwRig(t)
97 c := channel(t, r.svc, "er0000000a", "job-a", withRate(bwRate))
98 loc := bwLocation(t, c, tok)
99 got := bwDownload(t, c, loc)
100 want := time.Second * bwSize / bwRate
101 if got < want*8/10 || got > want+hangGuard {
102 t.Fatalf("download took %v at %d B/s, want at least %v", got, bwRate, want*8/10)
103 }
104}
105
106// The request bodies are limited too.
107func TestBandwidthLimitsAnUpload(t *testing.T) {
108 r := newRig(t, 1<<22)
109 c := channel(t, r.svc, "er0000000a", "job-a", withRate(bwRate))
110 j := r.job("job-a", repoA, main_, main_)
111 code, id := j.reserve("k", v1, bwSize)
112 if code != http.StatusCreated {
113 t.Fatalf("reserve: %d", code)
114 }
115 req, _ := http.NewRequest("PATCH", "http://cache"+APIPath+fmt.Sprintf("caches/%d", id), bytes.NewReader(bytes.Repeat([]byte("x"), bwSize)))
116 req.Header.Set("Authorization", "Bearer "+j.token)
117 req.Header.Set("Content-Range", fmt.Sprintf("bytes 0-%d/*", bwSize-1))
118 start := time.Now()
119 resp, err := c.Do(req)
120 if err != nil {
121 t.Fatal(err)
122 }
123 _ = resp.Body.Close()
124 got := time.Since(start)
125 if resp.StatusCode != http.StatusNoContent {
126 t.Fatalf("upload: %d", resp.StatusCode)
127 }
128 want := time.Second * bwSize / bwRate
129 if got < want*8/10 || got > want+hangGuard {
130 t.Fatalf("upload took %v at %d B/s, want at least %v", got, bwRate, want*8/10)
131 }
132 if code := j.commit(id, bwSize); code != http.StatusNoContent {
133 t.Fatalf("commit: %d", code)
134 }
135}
136
137// The limit is the sandbox's, not the connection's: two simultaneous downloads
138// on one channel take as long as the two in a row at that rate would, and a
139// channel of another sandbox, with no limit, is not slowed by them.
140func TestBandwidthIsSharedByTheConnectionsOfASandboxOnly(t *testing.T) {
141 r, tok := bwRig(t)
142 a := channel(t, r.svc, "er0000000a", "job-a", withRate(bwRate))
143 b := channel(t, r.svc, "er0000000b", "job-a") // same job id, so the same URL; no limit
144 loc := bwLocation(t, b, tok)
145
146 if limiterOf(t, r.svc, "er0000000a") == nil || limiterOf(t, r.svc, "er0000000b") != nil {
147 t.Fatal("A must be limited and B must have no limiter")
148 }
149
150 start := time.Now()
151 var wg sync.WaitGroup
152 for range 2 {
153 wg.Add(1)
154 go func() { defer wg.Done(); bwDownload(t, a, loc) }()
155 }
156 bwDownload(t, b, loc) // B has no limiter, so it never waits on A's bucket
157 wg.Wait()
158 total := time.Since(start)
159
160 want := time.Second * 2 * bwSize / bwRate // 1 s for the two; 0.5 s if each had the rate
161 if total < want*8/10 || total > want+hangGuard {
162 t.Fatalf("two downloads took %v in all, want at least %v", total, want*8/10)
163 }
164}
165
166// Two channels with a rate each do not draw on one bucket: the limit is per
167// Serve, so two sandboxes at once take the time of one.
168func TestBandwidthIsPerChannel(t *testing.T) {
169 r, tok := bwRig(t)
170 a := channel(t, r.svc, "er0000000a", "job-a", withRate(bwRate))
171 b := channel(t, r.svc, "er0000000b", "job-a", withRate(bwRate))
172 loc := bwLocation(t, a, tok)
173 la, lb := limiterOf(t, r.svc, "er0000000a"), limiterOf(t, r.svc, "er0000000b")
174 if la == nil || lb == nil || la == lb {
175 t.Fatal("each channel must have its own limiter")
176 }
177 // Consuming from A leaves B's bucket untouched.
178 bwDownload(t, a, loc)
179 if !bucketUsed(la) || bucketUsed(lb) {
180 t.Fatalf("after a download on A only: A used=%v, B used=%v", bucketUsed(la), bucketUsed(lb))
181 }
182 bwDownload(t, b, loc)
183 if !bucketUsed(lb) {
184 t.Fatal("B's own download did not draw on B's bucket")
185 }
186}
187
188// Without a rate nothing is limited or wrapped, and a number of Mbit/s is
189// bytes per second the way the VM interface counts it.
190func TestBandwidthUnsetIsUnlimited(t *testing.T) {
191 if newBandwidth(mbitToBytes(0)) != nil || newBandwidth(mbitToBytes(-1)) != nil {
192 t.Fatal("0 must mean unlimited")
193 }
194 if got := mbitToBytes(1000); got != 125_000_000 {
195 t.Fatalf("1000 Mbit/s = %d B/s", got)
196 }
197 r, tok := bwRig(t)
198 c := channel(t, r.svc, "er0000000a", "job-a", WithBandwidth(0))
199 loc := bwLocation(t, c, tok)
200 if limiterOf(t, r.svc, "er0000000a") != nil {
201 t.Fatal("WithBandwidth(0) installed a limiter")
202 }
203 if got := bwDownload(t, c, loc); got > hangGuard {
204 t.Fatalf("unlimited download took %v", got)
205 }
206}
207
208// A throttled transfer is cut when its connection is, not left sleeping.
209func TestBandwidthSleepEndsWithTheConnection(t *testing.T) {
210 r, tok := bwRig(t)
211 c := channel(t, r.svc, "er0000000a", "job-a", withRate(5_000)) // 20 s for the archive
212 loc := bwLocation(t, c, tok)
213 u, err := url.Parse(loc)
214 if err != nil {
215 t.Fatal(err)
216 }
217 req, _ := http.NewRequest("GET", "http://cache"+u.RequestURI(), nil)
218 resp, err := c.Do(req)
219 if err != nil {
220 t.Fatal(err)
221 }
222 buf := make([]byte, 1000)
223 if _, err := io.ReadFull(resp.Body, buf); err != nil {
224 t.Fatalf("the throttled transfer does not make progress: %v", err)
225 }
226 _ = resp.Body.Close()
227 done := make(chan struct{})
228 go func() { r.svc.Stop("er0000000a"); close(done) }()
229 select {
230 case <-done:
231 case <-time.After(5 * time.Second):
232 t.Fatal("Stop waits for a throttled transfer to finish")
233 }
234}