Merge branch 'thehyve/fix-crunch-documentation' Fix a typo in Crunch Dispatch install...
[arvados.git] / services / keep-balance / balance_run_test.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package main
6
7 import (
8         "encoding/json"
9         "fmt"
10         "io"
11         "io/ioutil"
12         "log"
13         "net/http"
14         "net/http/httptest"
15         "strings"
16         "sync"
17         "time"
18
19         "git.curoverse.com/arvados.git/sdk/go/arvados"
20
21         check "gopkg.in/check.v1"
22 )
23
24 var _ = check.Suite(&runSuite{})
25
26 type reqTracker struct {
27         reqs []http.Request
28         sync.Mutex
29 }
30
31 func (rt *reqTracker) Count() int {
32         rt.Lock()
33         defer rt.Unlock()
34         return len(rt.reqs)
35 }
36
37 func (rt *reqTracker) Add(req *http.Request) int {
38         rt.Lock()
39         defer rt.Unlock()
40         rt.reqs = append(rt.reqs, *req)
41         return len(rt.reqs)
42 }
43
44 var stubServices = []arvados.KeepService{
45         {
46                 UUID:           "zzzzz-bi6l4-000000000000000",
47                 ServiceHost:    "keep0.zzzzz.arvadosapi.com",
48                 ServicePort:    25107,
49                 ServiceSSLFlag: false,
50                 ServiceType:    "disk",
51         },
52         {
53                 UUID:           "zzzzz-bi6l4-000000000000001",
54                 ServiceHost:    "keep1.zzzzz.arvadosapi.com",
55                 ServicePort:    25107,
56                 ServiceSSLFlag: false,
57                 ServiceType:    "disk",
58         },
59         {
60                 UUID:           "zzzzz-bi6l4-000000000000002",
61                 ServiceHost:    "keep2.zzzzz.arvadosapi.com",
62                 ServicePort:    25107,
63                 ServiceSSLFlag: false,
64                 ServiceType:    "disk",
65         },
66         {
67                 UUID:           "zzzzz-bi6l4-000000000000003",
68                 ServiceHost:    "keep3.zzzzz.arvadosapi.com",
69                 ServicePort:    25107,
70                 ServiceSSLFlag: false,
71                 ServiceType:    "disk",
72         },
73         {
74                 UUID:           "zzzzz-bi6l4-h0a0xwut9qa6g3a",
75                 ServiceHost:    "keep.zzzzz.arvadosapi.com",
76                 ServicePort:    25333,
77                 ServiceSSLFlag: true,
78                 ServiceType:    "proxy",
79         },
80 }
81
82 var stubMounts = map[string][]arvados.KeepMount{
83         "keep0.zzzzz.arvadosapi.com:25107": {{
84                 UUID:     "zzzzz-ivpuk-000000000000000",
85                 DeviceID: "keep0-vol0",
86         }},
87         "keep1.zzzzz.arvadosapi.com:25107": {{
88                 UUID:     "zzzzz-ivpuk-100000000000000",
89                 DeviceID: "keep1-vol0",
90         }},
91         "keep2.zzzzz.arvadosapi.com:25107": {{
92                 UUID:     "zzzzz-ivpuk-200000000000000",
93                 DeviceID: "keep2-vol0",
94         }},
95         "keep3.zzzzz.arvadosapi.com:25107": {{
96                 UUID:     "zzzzz-ivpuk-300000000000000",
97                 DeviceID: "keep3-vol0",
98         }},
99 }
100
101 // stubServer is an HTTP transport that intercepts and processes all
102 // requests using its own handlers.
103 type stubServer struct {
104         mux      *http.ServeMux
105         srv      *httptest.Server
106         mutex    sync.Mutex
107         Requests reqTracker
108         logf     func(string, ...interface{})
109 }
110
111 // Start initializes the stub server and returns an *http.Client that
112 // uses the stub server to handle all requests.
113 //
114 // A stubServer that has been started should eventually be shut down
115 // with Close().
116 func (s *stubServer) Start() *http.Client {
117         // Set up a config.Client that forwards all requests to s.mux
118         // via s.srv. Test cases will attach handlers to s.mux to get
119         // the desired responses.
120         s.mux = http.NewServeMux()
121         s.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
122                 s.mutex.Lock()
123                 s.Requests.Add(r)
124                 s.mutex.Unlock()
125                 w.Header().Set("Content-Type", "application/json")
126                 s.mux.ServeHTTP(w, r)
127         }))
128         return &http.Client{Transport: s}
129 }
130
131 func (s *stubServer) RoundTrip(req *http.Request) (*http.Response, error) {
132         w := httptest.NewRecorder()
133         s.mux.ServeHTTP(w, req)
134         return &http.Response{
135                 StatusCode: w.Code,
136                 Status:     fmt.Sprintf("%d %s", w.Code, http.StatusText(w.Code)),
137                 Header:     w.HeaderMap,
138                 Body:       ioutil.NopCloser(w.Body)}, nil
139 }
140
141 // Close releases resources used by the server.
142 func (s *stubServer) Close() {
143         s.srv.Close()
144 }
145
146 func (s *stubServer) serveStatic(path, data string) *reqTracker {
147         rt := &reqTracker{}
148         s.mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
149                 rt.Add(r)
150                 if r.Body != nil {
151                         ioutil.ReadAll(r.Body)
152                         r.Body.Close()
153                 }
154                 io.WriteString(w, data)
155         })
156         return rt
157 }
158
159 func (s *stubServer) serveCurrentUserAdmin() *reqTracker {
160         return s.serveStatic("/arvados/v1/users/current",
161                 `{"uuid":"zzzzz-tpzed-000000000000000","is_admin":true,"is_active":true}`)
162 }
163
164 func (s *stubServer) serveCurrentUserNotAdmin() *reqTracker {
165         return s.serveStatic("/arvados/v1/users/current",
166                 `{"uuid":"zzzzz-tpzed-000000000000000","is_admin":false,"is_active":true}`)
167 }
168
169 func (s *stubServer) serveDiscoveryDoc() *reqTracker {
170         return s.serveStatic("/discovery/v1/apis/arvados/v1/rest",
171                 `{"defaultCollectionReplication":2}`)
172 }
173
174 func (s *stubServer) serveZeroCollections() *reqTracker {
175         return s.serveStatic("/arvados/v1/collections",
176                 `{"items":[],"items_available":0}`)
177 }
178
179 func (s *stubServer) serveFooBarFileCollections() *reqTracker {
180         rt := &reqTracker{}
181         s.mux.HandleFunc("/arvados/v1/collections", func(w http.ResponseWriter, r *http.Request) {
182                 r.ParseForm()
183                 rt.Add(r)
184                 if strings.Contains(r.Form.Get("filters"), `modified_at`) {
185                         io.WriteString(w, `{"items_available":0,"items":[]}`)
186                 } else {
187                         io.WriteString(w, `{"items_available":2,"items":[
188                                 {"uuid":"zzzzz-4zz18-ehbhgtheo8909or","portable_data_hash":"fa7aeb5140e2848d39b416daeef4ffc5+45","manifest_text":". 37b51d194a7513e45b56f6524f2d51f2+3 0:3:bar\n","modified_at":"2014-02-03T17:22:54Z"},
189                                 {"uuid":"zzzzz-4zz18-znfnqtbbv4spc3w","portable_data_hash":"1f4b0bc7583c2a7f9102c395f4ffc5e3+45","manifest_text":". acbd18db4cc2f85cedef654fccc4a4d8+3 0:3:foo\n","modified_at":"2014-02-03T17:22:54Z"}]}`)
190                 }
191         })
192         return rt
193 }
194
195 func (s *stubServer) serveCollectionsButSkipOne() *reqTracker {
196         rt := &reqTracker{}
197         s.mux.HandleFunc("/arvados/v1/collections", func(w http.ResponseWriter, r *http.Request) {
198                 r.ParseForm()
199                 rt.Add(r)
200                 if strings.Contains(r.Form.Get("filters"), `"modified_at","\u003c="`) {
201                         io.WriteString(w, `{"items_available":3,"items":[]}`)
202                 } else if strings.Contains(r.Form.Get("filters"), `"modified_at","\u003e`) {
203                         io.WriteString(w, `{"items_available":0,"items":[]}`)
204                 } else if strings.Contains(r.Form.Get("filters"), `"modified_at","="`) && strings.Contains(r.Form.Get("filters"), `"uuid","\u003e"`) {
205                         io.WriteString(w, `{"items_available":0,"items":[]}`)
206                 } else {
207                         io.WriteString(w, `{"items_available":2,"items":[
208                                 {"uuid":"zzzzz-4zz18-ehbhgtheo8909or","portable_data_hash":"fa7aeb5140e2848d39b416daeef4ffc5+45","manifest_text":". 37b51d194a7513e45b56f6524f2d51f2+3 0:3:bar\n","modified_at":"2014-02-03T17:22:54Z"},
209                                 {"uuid":"zzzzz-4zz18-znfnqtbbv4spc3w","portable_data_hash":"1f4b0bc7583c2a7f9102c395f4ffc5e3+45","manifest_text":". acbd18db4cc2f85cedef654fccc4a4d8+3 0:3:foo\n","modified_at":"2014-02-03T17:22:54Z"}]}`)
210                 }
211         })
212         return rt
213 }
214
215 func (s *stubServer) serveZeroKeepServices() *reqTracker {
216         return s.serveJSON("/arvados/v1/keep_services", arvados.KeepServiceList{})
217 }
218
219 func (s *stubServer) serveKeepServices(svcs []arvados.KeepService) *reqTracker {
220         return s.serveJSON("/arvados/v1/keep_services", arvados.KeepServiceList{
221                 ItemsAvailable: len(svcs),
222                 Items:          svcs,
223         })
224 }
225
226 func (s *stubServer) serveJSON(path string, resp interface{}) *reqTracker {
227         rt := &reqTracker{}
228         s.mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
229                 rt.Add(r)
230                 json.NewEncoder(w).Encode(resp)
231         })
232         return rt
233 }
234
235 func (s *stubServer) serveKeepstoreMounts() *reqTracker {
236         rt := &reqTracker{}
237         s.mux.HandleFunc("/mounts", func(w http.ResponseWriter, r *http.Request) {
238                 rt.Add(r)
239                 json.NewEncoder(w).Encode(stubMounts[r.Host])
240         })
241         return rt
242 }
243
244 func (s *stubServer) serveKeepstoreIndexFoo4Bar1() *reqTracker {
245         rt := &reqTracker{}
246         s.mux.HandleFunc("/index/", func(w http.ResponseWriter, r *http.Request) {
247                 count := rt.Add(r)
248                 if r.Host == "keep0.zzzzz.arvadosapi.com:25107" {
249                         io.WriteString(w, "37b51d194a7513e45b56f6524f2d51f2+3 12345678\n")
250                 }
251                 fmt.Fprintf(w, "acbd18db4cc2f85cedef654fccc4a4d8+3 %d\n\n", 12345678+count)
252         })
253         for _, mounts := range stubMounts {
254                 for i, mnt := range mounts {
255                         i := i
256                         s.mux.HandleFunc(fmt.Sprintf("/mounts/%s/blocks", mnt.UUID), func(w http.ResponseWriter, r *http.Request) {
257                                 count := rt.Add(r)
258                                 if i == 0 && r.Host == "keep0.zzzzz.arvadosapi.com:25107" {
259                                         io.WriteString(w, "37b51d194a7513e45b56f6524f2d51f2+3 12345678\n")
260                                 }
261                                 if i == 0 {
262                                         fmt.Fprintf(w, "acbd18db4cc2f85cedef654fccc4a4d8+3 %d\n", 12345678+count)
263                                 }
264                                 fmt.Fprintf(w, "\n")
265                         })
266                 }
267         }
268         return rt
269 }
270
271 func (s *stubServer) serveKeepstoreTrash() *reqTracker {
272         return s.serveStatic("/trash", `{}`)
273 }
274
275 func (s *stubServer) serveKeepstorePull() *reqTracker {
276         return s.serveStatic("/pull", `{}`)
277 }
278
279 type runSuite struct {
280         stub   stubServer
281         config Config
282 }
283
284 // make a log.Logger that writes to the current test's c.Log().
285 func (s *runSuite) logger(c *check.C) *log.Logger {
286         r, w := io.Pipe()
287         go func() {
288                 buf := make([]byte, 10000)
289                 for {
290                         n, err := r.Read(buf)
291                         if n > 0 {
292                                 if buf[n-1] == '\n' {
293                                         n--
294                                 }
295                                 c.Log(string(buf[:n]))
296                         }
297                         if err != nil {
298                                 break
299                         }
300                 }
301         }()
302         return log.New(w, "", log.LstdFlags)
303 }
304
305 func (s *runSuite) SetUpTest(c *check.C) {
306         s.config = Config{
307                 Client: arvados.Client{
308                         AuthToken: "xyzzy",
309                         APIHost:   "zzzzz.arvadosapi.com",
310                         Client:    s.stub.Start()},
311                 KeepServiceTypes: []string{"disk"}}
312         s.stub.serveDiscoveryDoc()
313         s.stub.logf = c.Logf
314 }
315
316 func (s *runSuite) TearDownTest(c *check.C) {
317         s.stub.Close()
318 }
319
320 func (s *runSuite) TestRefuseZeroCollections(c *check.C) {
321         opts := RunOptions{
322                 CommitPulls: true,
323                 CommitTrash: true,
324                 Logger:      s.logger(c),
325         }
326         s.stub.serveCurrentUserAdmin()
327         s.stub.serveZeroCollections()
328         s.stub.serveKeepServices(stubServices)
329         s.stub.serveKeepstoreMounts()
330         s.stub.serveKeepstoreIndexFoo4Bar1()
331         trashReqs := s.stub.serveKeepstoreTrash()
332         pullReqs := s.stub.serveKeepstorePull()
333         _, err := (&Balancer{}).Run(s.config, opts)
334         c.Check(err, check.ErrorMatches, "received zero collections")
335         c.Check(trashReqs.Count(), check.Equals, 4)
336         c.Check(pullReqs.Count(), check.Equals, 0)
337 }
338
339 func (s *runSuite) TestServiceTypes(c *check.C) {
340         opts := RunOptions{
341                 CommitPulls: true,
342                 CommitTrash: true,
343                 Logger:      s.logger(c),
344         }
345         s.config.KeepServiceTypes = []string{"unlisted-type"}
346         s.stub.serveCurrentUserAdmin()
347         s.stub.serveFooBarFileCollections()
348         s.stub.serveKeepServices(stubServices)
349         s.stub.serveKeepstoreMounts()
350         indexReqs := s.stub.serveKeepstoreIndexFoo4Bar1()
351         trashReqs := s.stub.serveKeepstoreTrash()
352         _, err := (&Balancer{}).Run(s.config, opts)
353         c.Check(err, check.IsNil)
354         c.Check(indexReqs.Count(), check.Equals, 0)
355         c.Check(trashReqs.Count(), check.Equals, 0)
356 }
357
358 func (s *runSuite) TestRefuseNonAdmin(c *check.C) {
359         opts := RunOptions{
360                 CommitPulls: true,
361                 CommitTrash: true,
362                 Logger:      s.logger(c),
363         }
364         s.stub.serveCurrentUserNotAdmin()
365         s.stub.serveZeroCollections()
366         s.stub.serveKeepServices(stubServices)
367         s.stub.serveKeepstoreMounts()
368         trashReqs := s.stub.serveKeepstoreTrash()
369         pullReqs := s.stub.serveKeepstorePull()
370         _, err := (&Balancer{}).Run(s.config, opts)
371         c.Check(err, check.ErrorMatches, "current user .* is not .* admin user")
372         c.Check(trashReqs.Count(), check.Equals, 0)
373         c.Check(pullReqs.Count(), check.Equals, 0)
374 }
375
376 func (s *runSuite) TestDetectSkippedCollections(c *check.C) {
377         opts := RunOptions{
378                 CommitPulls: true,
379                 CommitTrash: true,
380                 Logger:      s.logger(c),
381         }
382         s.stub.serveCurrentUserAdmin()
383         s.stub.serveCollectionsButSkipOne()
384         s.stub.serveKeepServices(stubServices)
385         s.stub.serveKeepstoreMounts()
386         s.stub.serveKeepstoreIndexFoo4Bar1()
387         trashReqs := s.stub.serveKeepstoreTrash()
388         pullReqs := s.stub.serveKeepstorePull()
389         _, err := (&Balancer{}).Run(s.config, opts)
390         c.Check(err, check.ErrorMatches, `Retrieved 2 collections with modtime <= .* but server now reports there are 3 collections.*`)
391         c.Check(trashReqs.Count(), check.Equals, 4)
392         c.Check(pullReqs.Count(), check.Equals, 0)
393 }
394
395 func (s *runSuite) TestDryRun(c *check.C) {
396         opts := RunOptions{
397                 CommitPulls: false,
398                 CommitTrash: false,
399                 Logger:      s.logger(c),
400         }
401         s.stub.serveCurrentUserAdmin()
402         collReqs := s.stub.serveFooBarFileCollections()
403         s.stub.serveKeepServices(stubServices)
404         s.stub.serveKeepstoreMounts()
405         s.stub.serveKeepstoreIndexFoo4Bar1()
406         trashReqs := s.stub.serveKeepstoreTrash()
407         pullReqs := s.stub.serveKeepstorePull()
408         var bal Balancer
409         _, err := bal.Run(s.config, opts)
410         c.Check(err, check.IsNil)
411         for _, req := range collReqs.reqs {
412                 c.Check(req.Form.Get("include_trash"), check.Equals, "true")
413         }
414         c.Check(trashReqs.Count(), check.Equals, 0)
415         c.Check(pullReqs.Count(), check.Equals, 0)
416         c.Check(bal.stats.pulls, check.Not(check.Equals), 0)
417         c.Check(bal.stats.underrep.replicas, check.Not(check.Equals), 0)
418         c.Check(bal.stats.overrep.replicas, check.Not(check.Equals), 0)
419 }
420
421 func (s *runSuite) TestCommit(c *check.C) {
422         opts := RunOptions{
423                 CommitPulls: true,
424                 CommitTrash: true,
425                 Logger:      s.logger(c),
426                 Dumper:      s.logger(c),
427         }
428         s.stub.serveCurrentUserAdmin()
429         s.stub.serveFooBarFileCollections()
430         s.stub.serveKeepServices(stubServices)
431         s.stub.serveKeepstoreMounts()
432         s.stub.serveKeepstoreIndexFoo4Bar1()
433         trashReqs := s.stub.serveKeepstoreTrash()
434         pullReqs := s.stub.serveKeepstorePull()
435         var bal Balancer
436         _, err := bal.Run(s.config, opts)
437         c.Check(err, check.IsNil)
438         c.Check(trashReqs.Count(), check.Equals, 8)
439         c.Check(pullReqs.Count(), check.Equals, 4)
440         // "foo" block is overreplicated by 2
441         c.Check(bal.stats.trashes, check.Equals, 2)
442         // "bar" block is underreplicated by 1, and its only copy is
443         // in a poor rendezvous position
444         c.Check(bal.stats.pulls, check.Equals, 2)
445 }
446
447 func (s *runSuite) TestRunForever(c *check.C) {
448         opts := RunOptions{
449                 CommitPulls: true,
450                 CommitTrash: true,
451                 Logger:      s.logger(c),
452                 Dumper:      s.logger(c),
453         }
454         s.stub.serveCurrentUserAdmin()
455         s.stub.serveFooBarFileCollections()
456         s.stub.serveKeepServices(stubServices)
457         s.stub.serveKeepstoreMounts()
458         s.stub.serveKeepstoreIndexFoo4Bar1()
459         trashReqs := s.stub.serveKeepstoreTrash()
460         pullReqs := s.stub.serveKeepstorePull()
461
462         stop := make(chan interface{})
463         s.config.RunPeriod = arvados.Duration(time.Millisecond)
464         go RunForever(s.config, opts, stop)
465
466         // Each run should send 4 pull lists + 4 trash lists. The
467         // first run should also send 4 empty trash lists at
468         // startup. We should complete all four runs in much less than
469         // a second.
470         for t0 := time.Now(); pullReqs.Count() < 16 && time.Since(t0) < 10*time.Second; {
471                 time.Sleep(time.Millisecond)
472         }
473         stop <- true
474         c.Check(pullReqs.Count() >= 16, check.Equals, true)
475         c.Check(trashReqs.Count(), check.Equals, pullReqs.Count()+4)
476 }