X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/0f74b22fb1b64f983debb78d7946eb26728fcc52..d158890921e8821782420711513a8b28ce0a1280:/tools/keep-exercise/keep-exercise.go

diff --git a/tools/keep-exercise/keep-exercise.go b/tools/keep-exercise/keep-exercise.go
index a94c01e55b..4131d752b8 100644
--- a/tools/keep-exercise/keep-exercise.go
+++ b/tools/keep-exercise/keep-exercise.go
@@ -36,6 +36,8 @@ var (
 	VaryThread    = flag.Bool("vary-thread", false, "use -wthreads different data blocks")
 	Replicas      = flag.Int("replicas", 1, "replication level for writing")
 	StatsInterval = flag.Duration("stats-interval", time.Second, "time interval between IO stats reports, or 0 to disable")
+	ServiceURL    = flag.String("url", "", "specify scheme://host of a single keep service to exercise (instead of using all advertised services like normal clients)")
+	ServiceUUID   = flag.String("uuid", "", "specify UUID of a single advertised keep service to exercise")
 )
 
 func main() {
@@ -45,13 +47,15 @@ func main() {
 	if err != nil {
 		log.Fatal(err)
 	}
-	kc, err := keepclient.MakeKeepClient(&arv)
+	kc, err := keepclient.MakeKeepClient(arv)
 	if err != nil {
 		log.Fatal(err)
 	}
 	kc.Want_replicas = *Replicas
 	kc.Client.Timeout = 10 * time.Minute
 
+	overrideServices(kc)
+
 	nextBuf := make(chan []byte, *WriteThreads)
 	nextLocator := make(chan string, *ReadThreads+*WriteThreads)
 
@@ -109,6 +113,7 @@ func makeBufs(nextBuf chan []byte, threadID int) {
 	}
 	for {
 		if *VaryRequest {
+			buf = make([]byte, *BlockSize)
 			if _, err := io.ReadFull(rand.Reader, buf); err != nil {
 				log.Fatal(err)
 			}
@@ -155,3 +160,23 @@ func doReads(kc *keepclient.KeepClient, nextLocator chan string) {
 		bytesInChan <- uint64(n)
 	}
 }
+
+func overrideServices(kc *keepclient.KeepClient) {
+	roots := make(map[string]string)
+	if *ServiceURL != "" {
+		roots["zzzzz-bi6l4-000000000000000"] = *ServiceURL
+	} else if *ServiceUUID != "" {
+		for uuid, url := range kc.GatewayRoots() {
+			if uuid == *ServiceUUID {
+				roots[uuid] = url
+				break
+			}
+		}
+		if len(roots) == 0 {
+			log.Fatalf("Service %q was not in list advertised by API %+q", *ServiceUUID, kc.GatewayRoots())
+		}
+	} else {
+		return
+	}
+	kc.SetServiceRoots(roots, roots, roots)
+}