chore(storage): add configurable per op timeout, print on timeout [be…

…nchmarking] (#8191)
googleapis · Jul 7, 2023 · ed73916 · ed73916
1 parent e1f52db
commit ed73916
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 3 deletions.
diff --git a/storage/internal/benchmarks/README.md b/storage/internal/benchmarks/README.md
@@ -31,6 +31,7 @@ This runs 1000 iterations on 512kib to 2Gib files in the background, sending out
 | -connection_pool_size | GRPC connection pool size | any positive integer | 4 |
 | -force_garbage_collection | whether to force garbage collection <br> before every write or read benchmark |  `true` or `false` (present/not present) | `false` |
 | -timeout | timeout (maximum time running benchmarks) <br> the program may run for longer while it finishes running processes | any [time.Duration](https://pkg.go.dev/time#Duration) | `1h` |
+| -timeout_per_op | timeout on a single upload or download | any [time.Duration](https://pkg.go.dev/time#Duration) | `5m` |
 | -workload | `1` will run a w1r3 (write 1 read 3) benchmark <br> `6` will run a benchmark uploading and downloading (once each) <br> a single directory with `-directory_num_objects` number of files (no subdirectories)  | `1` or `6` | `1` |
 | -directory_num_objects | total number of objects in a directory (directory will only contain files, <br> no subdirectories); only applies to workload 6 | any positive integer | `1000` |
 

diff --git a/storage/internal/benchmarks/directory_benchmark.go b/storage/internal/benchmarks/directory_benchmark.go
@@ -205,6 +205,7 @@ func (r *directoryBenchmark) uploadDirectory(ctx context.Context, numWorkers int
 				object:              objectName,
 				useDefaultChunkSize: opts.minChunkSize == useDefault || opts.maxChunkSize == useDefault,
 				objectPath:          filePath,
+				timeout:             r.opts.timeoutPerOp,
 			})
 			return err
 		})
@@ -278,6 +279,7 @@ func (r *directoryBenchmark) downloadDirectory(ctx context.Context, numWorkers i
 				downloadToDirectory: r.downloadDirectoryPath,
 				rangeStart:          rangeStart,
 				rangeLength:         rangeLength,
+				timeout:             r.opts.timeoutPerOp,
 			})
 			return err
 		})

diff --git a/storage/internal/benchmarks/download_benchmark.go b/storage/internal/benchmarks/download_benchmark.go
@@ -33,6 +33,7 @@ type downloadOpts struct {
 	rangeStart          int64
 	rangeLength         int64
 	downloadToDirectory string
+	timeout             time.Duration
 }
 
 func downloadBenchmark(ctx context.Context, dopts downloadOpts) (elapsedTime time.Duration, rerr error) {
@@ -45,7 +46,7 @@ func downloadBenchmark(ctx context.Context, dopts downloadOpts) (elapsedTime tim
 	defer func() { elapsedTime = time.Since(start) }()
 
 	// Set additional timeout
-	ctx, cancel := context.WithTimeout(ctx, time.Minute*2)
+	ctx, cancel := context.WithTimeout(ctx, dopts.timeout)
 	defer cancel()
 
 	o := dopts.client.Bucket(dopts.bucket).Object(dopts.object)

diff --git a/storage/internal/benchmarks/main.go b/storage/internal/benchmarks/main.go
@@ -73,7 +73,8 @@ type benchmarkOptions struct {
 	forceGC      bool
 	connPoolSize int
 
-	timeout time.Duration
+	timeout      time.Duration
+	timeoutPerOp time.Duration
 
 	continueOnFail bool
 
@@ -166,6 +167,7 @@ func parseFlags() {
 	flag.BoolVar(&opts.forceGC, "force_garbage_collection", false, "force garbage collection at the beginning of each upload")
 
 	flag.DurationVar(&opts.timeout, "timeout", time.Hour, "timeout")
+	flag.DurationVar(&opts.timeoutPerOp, "timeout_per_op", time.Minute*5, "timeout per upload/download")
 	flag.StringVar(&outputFile, "o", "", "file to output results to - if empty, will output to stdout")
 
 	flag.BoolVar(&opts.continueOnFail, "continue_on_fail", false, "continue even if a run fails")
@@ -209,6 +211,12 @@ func main() {
 	ctx, cancel := context.WithDeadline(context.Background(), start.Add(opts.timeout))
 	defer cancel()
 
+	// Print a message once deadline is exceeded
+	go func() {
+		<-ctx.Done()
+		log.Printf("total configured timeout exceeded")
+	}()
+
 	// Create bucket if necessary
 	if len(opts.bucket) < 1 {
 		opts.bucket = randomName(bucketPrefix)

diff --git a/storage/internal/benchmarks/upload_benchmark.go b/storage/internal/benchmarks/upload_benchmark.go
@@ -38,6 +38,7 @@ type uploadOpts struct {
 	object              string
 	useDefaultChunkSize bool
 	objectPath          string
+	timeout             time.Duration
 }
 
 func uploadBenchmark(ctx context.Context, uopts uploadOpts) (elapsedTime time.Duration, rerr error) {
@@ -50,7 +51,7 @@ func uploadBenchmark(ctx context.Context, uopts uploadOpts) (elapsedTime time.Du
 	defer func() { elapsedTime = time.Since(start) }()
 
 	// Set additional timeout
-	ctx, cancel := context.WithTimeout(ctx, time.Minute*2)
+	ctx, cancel := context.WithTimeout(ctx, uopts.timeout)
 	defer cancel()
 
 	// Open file

diff --git a/storage/internal/benchmarks/w1r3.go b/storage/internal/benchmarks/w1r3.go
@@ -157,6 +157,7 @@ func (r *w1r3) run(ctx context.Context) error {
 			object:              r.objectName,
 			useDefaultChunkSize: opts.minChunkSize == useDefault || opts.maxChunkSize == useDefault,
 			objectPath:          r.objectPath,
+			timeout:             r.opts.timeoutPerOp,
 		})
 	})
 
@@ -186,6 +187,7 @@ func (r *w1r3) run(ctx context.Context) error {
 				rangeStart:          rangeStart,
 				rangeLength:         rangeLength,
 				downloadToDirectory: r.directoryPath,
+				timeout:             r.opts.timeoutPerOp,
 			})
 		})
 		if err != nil {