chore(storage): add warmup option [benchmarks]

storage/internal/benchmarks/directory_benchmark.go

@@ -25,6 +25,9 @@ import (
 	"time"
 
 	"cloud.google.com/go/storage"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/api/iterator"
 )
@@ -172,6 +175,13 @@ func (r *directoryBenchmark) cleanup() error {
 }
 
 func (r *directoryBenchmark) uploadDirectory(ctx context.Context, numWorkers int) (elapsedTime time.Duration, err error) {
+	var span trace.Span
+	ctx, span = otel.GetTracerProvider().Tracer(tracerName).Start(ctx, "uploadDirectory")
+	span.SetAttributes(
+		attribute.KeyValue{"num_workers", attribute.IntValue(numWorkers)},
+	)
+	defer span.End()
+
 	benchGroup, ctx := errgroup.WithContext(ctx)
 	benchGroup.SetLimit(numWorkers)
 
@@ -220,6 +230,13 @@ func (r *directoryBenchmark) uploadDirectory(ctx context.Context, numWorkers int
 }
 
 func (r *directoryBenchmark) downloadDirectory(ctx context.Context, numWorkers int) (elapsedTime time.Duration, err error) {
+	var span trace.Span
+	ctx, span = otel.GetTracerProvider().Tracer(tracerName).Start(ctx, "downloadDirectory")
+	span.SetAttributes(
+		attribute.KeyValue{"num_workers", attribute.IntValue(numWorkers)},
+	)
+	defer span.End()
+
 	benchGroup, ctx := errgroup.WithContext(ctx)
 	benchGroup.SetLimit(numWorkers)
 
@@ -299,6 +316,13 @@ func (r *directoryBenchmark) downloadDirectory(ctx context.Context, numWorkers i
 }
 
 func (r *directoryBenchmark) run(ctx context.Context) error {
+	var span trace.Span
+	ctx, span = otel.GetTracerProvider().Tracer(tracerName).Start(ctx, "directoryBenchmark")
+	span.SetAttributes(
+		attribute.KeyValue{"api", attribute.StringValue(string(r.opts.api))},
+		attribute.KeyValue{"object_size", attribute.Int64Value(r.opts.objectSize)})
+	defer span.End()
+
 	// Upload
 	err := runOneOp(ctx, r.writeResult, func() (time.Duration, error) {
 		return r.uploadDirectory(ctx, r.numWorkers)

storage/internal/benchmarks/main.go

@@ -96,6 +96,7 @@ type benchmarkOptions struct {
 
 	enableTracing   bool
 	traceSampleRate float64
+	warmup          time.Duration
 }
 
 func (b *benchmarkOptions) validate() error {
@@ -194,6 +195,8 @@ func parseFlags() {
 	flag.IntVar(&opts.workload, "workload", 1, "which workload to run")
 	flag.IntVar(&opts.numObjectsPerDirectory, "directory_num_objects", 1000, "total number of objects in directory")
 
+	flag.DurationVar(&opts.warmup, "warmup", 0, "time to warmup benchmarks; w1r3 benchmarks will be run for this duration without recording any results")
+
 	flag.Parse()
 
 	if len(projectID) < 1 {
@@ -283,6 +286,9 @@ func main() {
 		log.Fatalf("populateDependencyVersions: %v", err)
 	}
 
+	if err := warmupW1R3(ctx, opts); err != nil {
+		log.Fatal(err)
+	}
 	recordResultGroup, _ := errgroup.WithContext(ctx)
 	startRecordingResults(w, recordResultGroup, opts.outType)
 

storage/internal/benchmarks/w1r3.go

@@ -149,7 +149,7 @@ func (r *w1r3) cleanup() error {
 
 func (r *w1r3) run(ctx context.Context) error {
 	// Use the same client for write and reads as the api is the same
-	client := getClient(ctx, r.writeResult.params.api)
+	client := getClient(ctx, r.readResults[0].params.api)
 
 	var span trace.Span
 	ctx, span = otel.GetTracerProvider().Tracer(tracerName).Start(ctx, "w1r3")

storage/internal/benchmarks/warmup.go

@@ -0,0 +1,72 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"golang.org/x/sync/errgroup"
+)
+
+func warmupW1R3(ctx context.Context, opts *benchmarkOptions) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	discardBenchmarkResults(ctx)
+
+	warmupGroup, ctx := errgroup.WithContext(ctx)
+	warmupGroup.SetLimit(opts.numWorkers)
+
+	for deadline := time.Now().Add(opts.warmup); time.Now().Before(deadline); {
+		warmupGroup.Go(func() error {
+			benchmark := &w1r3{opts: opts, bucketName: opts.bucket}
+
+			if err := benchmark.setup(ctx); err != nil {
+				return fmt.Errorf("warmup setup failed: %v", err)
+			}
+			if err := benchmark.run(ctx); err != nil {
+				return fmt.Errorf("warmup run failed: %v", err)
+			}
+			if err := benchmark.cleanup(); err != nil {
+				return fmt.Errorf("warmup cleanup failed: %v", err)
+			}
+			return nil
+		})
+	}
+
+	return warmupGroup.Wait()
+}
+
+// discardBenchmarkResults consumes benchmark results until the provided context
+// is cancelled
+func discardBenchmarkResults(ctx context.Context) {
+	results = make(chan benchmarkResult)
+
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				close(results)
+				return
+			case _, ok := <-results:
+				if !ok {
+					return
+				}
+			}
+		}
+	}()
+}