Add 'tools/arvbox/' from commit 'd3d368758db1f4a9fa5b89f77b5ee61d68ef5b72'
[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         "time"
25
26         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
27         "git.curoverse.com/arvados.git/sdk/go/keepclient"
28 )
29
30 // Command line config knobs
31 var (
32         BlockSize     = flag.Int("block-size", keepclient.BLOCKSIZE, "bytes per read/write op")
33         ReadThreads   = flag.Int("rthreads", 1, "number of concurrent readers")
34         WriteThreads  = flag.Int("wthreads", 1, "number of concurrent writers")
35         VaryRequest   = flag.Bool("vary-request", false, "vary the data for each request: consumes disk space, exercises write behavior")
36         VaryThread    = flag.Bool("vary-thread", false, "use -wthreads different data blocks")
37         Replicas      = flag.Int("replicas", 1, "replication level for writing")
38         StatsInterval = flag.Duration("stats-interval", time.Second, "time interval between IO stats reports, or 0 to disable")
39         ServiceURL    = flag.String("url", "", "specify scheme://host of a single keep service to exercise (instead of using all advertised services like normal clients)")
40         ServiceUUID   = flag.String("uuid", "", "specify UUID of a single advertised keep service to exercise")
41 )
42
43 func main() {
44         flag.Parse()
45
46         arv, err := arvadosclient.MakeArvadosClient()
47         if err != nil {
48                 log.Fatal(err)
49         }
50         kc, err := keepclient.MakeKeepClient(&arv)
51         if err != nil {
52                 log.Fatal(err)
53         }
54         kc.Want_replicas = *Replicas
55         kc.Client.Timeout = 10 * time.Minute
56
57         overrideServices(kc)
58
59         nextBuf := make(chan []byte, *WriteThreads)
60         nextLocator := make(chan string, *ReadThreads+*WriteThreads)
61
62         go countBeans(nextLocator)
63         for i := 0; i < *WriteThreads; i++ {
64                 go makeBufs(nextBuf, i)
65                 go doWrites(kc, nextBuf, nextLocator)
66         }
67         for i := 0; i < *ReadThreads; i++ {
68                 go doReads(kc, nextLocator)
69         }
70         <-make(chan struct{})
71 }
72
73 // Send 1234 to bytesInChan when we receive 1234 bytes from keepstore.
74 var bytesInChan = make(chan uint64)
75 var bytesOutChan = make(chan uint64)
76
77 // Send struct{}{} to errorsChan when an error happens.
78 var errorsChan = make(chan struct{})
79
80 func countBeans(nextLocator chan string) {
81         t0 := time.Now()
82         var tickChan <-chan time.Time
83         if *StatsInterval > 0 {
84                 tickChan = time.NewTicker(*StatsInterval).C
85         }
86         var bytesIn uint64
87         var bytesOut uint64
88         var errors uint64
89         for {
90                 select {
91                 case <-tickChan:
92                         elapsed := time.Since(t0)
93                         log.Printf("%v elapsed: read %v bytes (%.1f MiB/s), wrote %v bytes (%.1f MiB/s), errors %d",
94                                 elapsed,
95                                 bytesIn, (float64(bytesIn) / elapsed.Seconds() / 1048576),
96                                 bytesOut, (float64(bytesOut) / elapsed.Seconds() / 1048576),
97                                 errors,
98                         )
99                 case i := <-bytesInChan:
100                         bytesIn += i
101                 case o := <-bytesOutChan:
102                         bytesOut += o
103                 case <-errorsChan:
104                         errors++
105                 }
106         }
107 }
108
109 func makeBufs(nextBuf chan []byte, threadID int) {
110         buf := make([]byte, *BlockSize)
111         if *VaryThread {
112                 binary.PutVarint(buf, int64(threadID))
113         }
114         for {
115                 if *VaryRequest {
116                         buf = make([]byte, *BlockSize)
117                         if _, err := io.ReadFull(rand.Reader, buf); err != nil {
118                                 log.Fatal(err)
119                         }
120                 }
121                 nextBuf <- buf
122         }
123 }
124
125 func doWrites(kc *keepclient.KeepClient, nextBuf chan []byte, nextLocator chan string) {
126         for buf := range nextBuf {
127                 locator, _, err := kc.PutB(buf)
128                 if err != nil {
129                         log.Print(err)
130                         errorsChan <- struct{}{}
131                         continue
132                 }
133                 bytesOutChan <- uint64(len(buf))
134                 for cap(nextLocator) > len(nextLocator)+*WriteThreads {
135                         // Give the readers something to do, unless
136                         // they have lots queued up already.
137                         nextLocator <- locator
138                 }
139         }
140 }
141
142 func doReads(kc *keepclient.KeepClient, nextLocator chan string) {
143         for locator := range nextLocator {
144                 rdr, size, url, err := kc.Get(locator)
145                 if err != nil {
146                         log.Print(err)
147                         errorsChan <- struct{}{}
148                         continue
149                 }
150                 n, err := io.Copy(ioutil.Discard, rdr)
151                 rdr.Close()
152                 if n != size || err != nil {
153                         log.Printf("Got %d bytes (expected %d) from %s: %v", n, size, url, err)
154                         errorsChan <- struct{}{}
155                         continue
156                         // Note we don't count the bytes received in
157                         // partial/corrupt responses: we are measuring
158                         // throughput, not resource consumption.
159                 }
160                 bytesInChan <- uint64(n)
161         }
162 }
163
164 func overrideServices(kc *keepclient.KeepClient) {
165         roots := make(map[string]string)
166         if *ServiceURL != "" {
167                 roots["zzzzz-bi6l4-000000000000000"] = *ServiceURL
168         } else if *ServiceUUID != "" {
169                 for uuid, url := range kc.GatewayRoots() {
170                         if uuid == *ServiceUUID {
171                                 roots[uuid] = url
172                                 break
173                         }
174                 }
175                 if len(roots) == 0 {
176                         log.Fatalf("Service %q was not in list advertised by API %+q", *ServiceUUID, kc.GatewayRoots())
177                 }
178         } else {
179                 return
180         }
181         kc.SetServiceRoots(roots, roots, roots)
182 }