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