6c8a866291b10ca0efcff8e7d9f50ca6d727703b
[arvados.git] / tools / keep-exercise / keep-exercise.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 // Testing tool for Keep services.
6 //
7 // keepexercise helps measure throughput and test reliability under
8 // various usage patterns.
9 //
10 // By default, it reads and writes blocks containing 2^26 NUL
11 // bytes. This generates network traffic without consuming much disk
12 // space.
13 //
14 // For a more realistic test, enable -vary-request. Warning: this will
15 // fill your storage volumes with random data if you leave it running,
16 // which can cost you money or leave you with too little room for
17 // useful data.
18 //
19 package main
20
21 import (
22         "crypto/rand"
23         "encoding/binary"
24         "flag"
25         "io"
26         "io/ioutil"
27         "log"
28         "net/http"
29         "time"
30
31         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
32         "git.curoverse.com/arvados.git/sdk/go/keepclient"
33 )
34
35 // Command line config knobs
36 var (
37         BlockSize     = flag.Int("block-size", keepclient.BLOCKSIZE, "bytes per read/write op")
38         ReadThreads   = flag.Int("rthreads", 1, "number of concurrent readers")
39         WriteThreads  = flag.Int("wthreads", 1, "number of concurrent writers")
40         VaryRequest   = flag.Bool("vary-request", false, "vary the data for each request: consumes disk space, exercises write behavior")
41         VaryThread    = flag.Bool("vary-thread", false, "use -wthreads different data blocks")
42         Replicas      = flag.Int("replicas", 1, "replication level for writing")
43         StatsInterval = flag.Duration("stats-interval", time.Second, "time interval between IO stats reports, or 0 to disable")
44         ServiceURL    = flag.String("url", "", "specify scheme://host of a single keep service to exercise (instead of using all advertised services like normal clients)")
45         ServiceUUID   = flag.String("uuid", "", "specify UUID of a single advertised keep service to exercise")
46 )
47
48 func main() {
49         flag.Parse()
50
51         arv, err := arvadosclient.MakeArvadosClient()
52         if err != nil {
53                 log.Fatal(err)
54         }
55         kc, err := keepclient.MakeKeepClient(arv)
56         if err != nil {
57                 log.Fatal(err)
58         }
59         kc.Want_replicas = *Replicas
60
61         transport := *(http.DefaultTransport.(*http.Transport))
62         transport.TLSClientConfig = arvadosclient.MakeTLSConfig(arv.ApiInsecure)
63         kc.HTTPClient = &http.Client{
64                 Timeout:   10 * time.Minute,
65                 Transport: &transport,
66         }
67
68         overrideServices(kc)
69
70         nextLocator := make(chan string, *ReadThreads+*WriteThreads)
71
72         go countBeans(nextLocator)
73         for i := 0; i < *WriteThreads; i++ {
74                 nextBuf := make(chan []byte, 1)
75                 go makeBufs(nextBuf, i)
76                 go doWrites(kc, nextBuf, nextLocator)
77         }
78         for i := 0; i < *ReadThreads; i++ {
79                 go doReads(kc, nextLocator)
80         }
81         <-make(chan struct{})
82 }
83
84 // Send 1234 to bytesInChan when we receive 1234 bytes from keepstore.
85 var bytesInChan = make(chan uint64)
86 var bytesOutChan = make(chan uint64)
87
88 // Send struct{}{} to errorsChan when an error happens.
89 var errorsChan = make(chan struct{})
90
91 func countBeans(nextLocator chan string) {
92         t0 := time.Now()
93         var tickChan <-chan time.Time
94         if *StatsInterval > 0 {
95                 tickChan = time.NewTicker(*StatsInterval).C
96         }
97         var bytesIn uint64
98         var bytesOut uint64
99         var errors uint64
100         for {
101                 select {
102                 case <-tickChan:
103                         elapsed := time.Since(t0)
104                         log.Printf("%v elapsed: read %v bytes (%.1f MiB/s), wrote %v bytes (%.1f MiB/s), errors %d",
105                                 elapsed,
106                                 bytesIn, (float64(bytesIn) / elapsed.Seconds() / 1048576),
107                                 bytesOut, (float64(bytesOut) / elapsed.Seconds() / 1048576),
108                                 errors,
109                         )
110                 case i := <-bytesInChan:
111                         bytesIn += i
112                 case o := <-bytesOutChan:
113                         bytesOut += o
114                 case <-errorsChan:
115                         errors++
116                 }
117         }
118 }
119
120 func makeBufs(nextBuf chan<- []byte, threadID int) {
121         buf := make([]byte, *BlockSize)
122         if *VaryThread {
123                 binary.PutVarint(buf, int64(threadID))
124         }
125         randSize := 524288
126         if randSize > *BlockSize {
127                 randSize = *BlockSize
128         }
129         for {
130                 if *VaryRequest {
131                         rnd := make([]byte, randSize)
132                         if _, err := io.ReadFull(rand.Reader, rnd); err != nil {
133                                 log.Fatal(err)
134                         }
135                         buf = append(rnd, buf[randSize:]...)
136                 }
137                 nextBuf <- buf
138         }
139 }
140
141 func doWrites(kc *keepclient.KeepClient, nextBuf <-chan []byte, nextLocator chan<- string) {
142         for buf := range nextBuf {
143                 locator, _, err := kc.PutB(buf)
144                 if err != nil {
145                         log.Print(err)
146                         errorsChan <- struct{}{}
147                         continue
148                 }
149                 bytesOutChan <- uint64(len(buf))
150                 for cap(nextLocator) > len(nextLocator)+*WriteThreads {
151                         // Give the readers something to do, unless
152                         // they have lots queued up already.
153                         nextLocator <- locator
154                 }
155         }
156 }
157
158 func doReads(kc *keepclient.KeepClient, nextLocator <-chan string) {
159         for locator := range nextLocator {
160                 rdr, size, url, err := kc.Get(locator)
161                 if err != nil {
162                         log.Print(err)
163                         errorsChan <- struct{}{}
164                         continue
165                 }
166                 n, err := io.Copy(ioutil.Discard, rdr)
167                 rdr.Close()
168                 if n != size || err != nil {
169                         log.Printf("Got %d bytes (expected %d) from %s: %v", n, size, url, err)
170                         errorsChan <- struct{}{}
171                         continue
172                         // Note we don't count the bytes received in
173                         // partial/corrupt responses: we are measuring
174                         // throughput, not resource consumption.
175                 }
176                 bytesInChan <- uint64(n)
177         }
178 }
179
180 func overrideServices(kc *keepclient.KeepClient) {
181         roots := make(map[string]string)
182         if *ServiceURL != "" {
183                 roots["zzzzz-bi6l4-000000000000000"] = *ServiceURL
184         } else if *ServiceUUID != "" {
185                 for uuid, url := range kc.GatewayRoots() {
186                         if uuid == *ServiceUUID {
187                                 roots[uuid] = url
188                                 break
189                         }
190                 }
191                 if len(roots) == 0 {
192                         log.Fatalf("Service %q was not in list advertised by API %+q", *ServiceUUID, kc.GatewayRoots())
193                 }
194         } else {
195                 return
196         }
197         kc.SetServiceRoots(roots, roots, roots)
198 }