chore(storage): add endpoint option to CLI [benchmarking] (#8269)

googleapis · Jul 14, 2023 · 715715c · 715715c
1 parent d9bb34f
commit 715715c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/storage/internal/benchmarks/benchmark_result.go b/storage/internal/benchmarks/benchmark_result.go
@@ -218,6 +218,11 @@ func (br *benchmarkResult) cloudMonitoring() []byte {
 		sb.WriteString(",")
 	}
 
+	if len(opts.endpoint) > 0 {
+		sb.WriteString(makeStringQuoted("endpoint", opts.endpoint))
+		sb.WriteString(",")
+	}
+
 	sb.WriteString(makeStringQuoted("code_version", codeVersion))
 	sb.WriteString(",")
 	sb.WriteString(makeStringQuoted("go_version", goVersion))

diff --git a/storage/internal/benchmarks/client_pool.go b/storage/internal/benchmarks/client_pool.go
@@ -105,6 +105,7 @@ func initializeClientPools(ctx context.Context, opts *benchmarkOptions) func() {
 					readBufferSize:  opts.readBufferSize,
 					useJSON:         false,
 					setGCSFuseOpts:  opts.useGCSFuseConfig,
+					endpoint:        opts.endpoint,
 				})
 			},
 			opts.numClients,
@@ -122,6 +123,7 @@ func initializeClientPools(ctx context.Context, opts *benchmarkOptions) func() {
 					readBufferSize:  opts.readBufferSize,
 					useJSON:         true,
 					setGCSFuseOpts:  opts.useGCSFuseConfig,
+					endpoint:        opts.endpoint,
 				})
 			},
 			opts.numClients,
@@ -136,6 +138,7 @@ func initializeClientPools(ctx context.Context, opts *benchmarkOptions) func() {
 					writeBufferSize:    opts.writeBufferSize,
 					readBufferSize:     opts.readBufferSize,
 					connectionPoolSize: opts.connPoolSize,
+					endpoint:           opts.endpoint,
 				})
 			},
 			opts.numClients,
@@ -177,6 +180,7 @@ var clientMu sync.Mutex
 // Client config
 type clientConfig struct {
 	writeBufferSize, readBufferSize int
+	endpoint                        string
 	useJSON                         bool // only applicable to HTTP Clients
 	setGCSFuseOpts                  bool // only applicable to HTTP Clients
 	connectionPoolSize              int  // only applicable to GRPC Clients
@@ -185,6 +189,10 @@ type clientConfig struct {
 func initializeHTTPClient(ctx context.Context, config clientConfig) (*storage.Client, error) {
 	opts := []option.ClientOption{}
 
+	if len(config.endpoint) > 0 {
+		opts = append(opts, option.WithEndpoint(config.endpoint))
+	}
+
 	if config.writeBufferSize != useDefault || config.readBufferSize != useDefault || config.setGCSFuseOpts {
 		// We need to modify the underlying HTTP client
 		base := http.DefaultTransport.(*http.Transport).Clone()
@@ -236,6 +244,10 @@ func initializeHTTPClient(ctx context.Context, config clientConfig) (*storage.Cl
 func initializeGRPCClient(ctx context.Context, config clientConfig) (*storage.Client, error) {
 	opts := []option.ClientOption{option.WithGRPCConnectionPool(config.connectionPoolSize)}
 
+	if len(config.endpoint) > 0 {
+		opts = append(opts, option.WithEndpoint(config.endpoint))
+	}
+
 	if config.writeBufferSize != useDefault {
 		opts = append(opts, option.WithGRPCDialOption(grpc.WithWriteBufferSize(config.writeBufferSize)))
 	}

diff --git a/storage/internal/benchmarks/main.go b/storage/internal/benchmarks/main.go
@@ -83,6 +83,7 @@ type benchmarkOptions struct {
 	numObjectsPerDirectory int
 
 	useGCSFuseConfig bool
+	endpoint         string
 }
 
 func (b *benchmarkOptions) validate() error {
@@ -155,6 +156,7 @@ func parseFlags() {
 	flag.Int64Var(&opts.maxReadOffset, "maximum_read_offset", 0, "maximum read offset in bytes")
 
 	flag.BoolVar(&opts.useGCSFuseConfig, "gcs_fuse", false, "use GCSFuse configs on HTTP client creation")
+	flag.StringVar(&opts.endpoint, "endpoint", "", "endpoint to set on Storage Client")
 
 	flag.IntVar(&opts.readBufferSize, "read_buffer_size", useDefault, "read buffer size in bytes")
 	flag.IntVar(&opts.writeBufferSize, "write_buffer_size", useDefault, "write buffer size in bytes")