X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/494de455828983fe827c96cf425fac6492618308..9a26fa413d0890eac3a2578876508efc2cd32374:/tools/keep-exercise/keep-exercise.go

diff --git a/tools/keep-exercise/keep-exercise.go b/tools/keep-exercise/keep-exercise.go
index a94c01e55b..6c8a866291 100644
--- a/tools/keep-exercise/keep-exercise.go
+++ b/tools/keep-exercise/keep-exercise.go
@@ -1,3 +1,7 @@
+// Copyright (C) The Arvados Authors. All rights reserved.
+//
+// SPDX-License-Identifier: AGPL-3.0
+
 // Testing tool for Keep services.
 //
 // keepexercise helps measure throughput and test reliability under
@@ -21,6 +25,7 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"net/http"
 	"time"
 
 	"git.curoverse.com/arvados.git/sdk/go/arvadosclient"
@@ -36,6 +41,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,18 +52,26 @@ 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
 
-	nextBuf := make(chan []byte, *WriteThreads)
+	transport := *(http.DefaultTransport.(*http.Transport))
+	transport.TLSClientConfig = arvadosclient.MakeTLSConfig(arv.ApiInsecure)
+	kc.HTTPClient = &http.Client{
+		Timeout:   10 * time.Minute,
+		Transport: &transport,
+	}
+
+	overrideServices(kc)
+
 	nextLocator := make(chan string, *ReadThreads+*WriteThreads)
 
 	go countBeans(nextLocator)
 	for i := 0; i < *WriteThreads; i++ {
+		nextBuf := make(chan []byte, 1)
 		go makeBufs(nextBuf, i)
 		go doWrites(kc, nextBuf, nextLocator)
 	}
@@ -102,22 +117,28 @@ func countBeans(nextLocator chan string) {
 	}
 }
 
-func makeBufs(nextBuf chan []byte, threadID int) {
+func makeBufs(nextBuf chan<- []byte, threadID int) {
 	buf := make([]byte, *BlockSize)
 	if *VaryThread {
 		binary.PutVarint(buf, int64(threadID))
 	}
+	randSize := 524288
+	if randSize > *BlockSize {
+		randSize = *BlockSize
+	}
 	for {
 		if *VaryRequest {
-			if _, err := io.ReadFull(rand.Reader, buf); err != nil {
+			rnd := make([]byte, randSize)
+			if _, err := io.ReadFull(rand.Reader, rnd); err != nil {
 				log.Fatal(err)
 			}
+			buf = append(rnd, buf[randSize:]...)
 		}
 		nextBuf <- buf
 	}
 }
 
-func doWrites(kc *keepclient.KeepClient, nextBuf chan []byte, nextLocator chan string) {
+func doWrites(kc *keepclient.KeepClient, nextBuf <-chan []byte, nextLocator chan<- string) {
 	for buf := range nextBuf {
 		locator, _, err := kc.PutB(buf)
 		if err != nil {
@@ -134,7 +155,7 @@ func doWrites(kc *keepclient.KeepClient, nextBuf chan []byte, nextLocator chan s
 	}
 }
 
-func doReads(kc *keepclient.KeepClient, nextLocator chan string) {
+func doReads(kc *keepclient.KeepClient, nextLocator <-chan string) {
 	for locator := range nextLocator {
 		rdr, size, url, err := kc.Get(locator)
 		if err != nil {
@@ -155,3 +176,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)
+}