From 5fac9166301b0668456dbae74b968fb183048e1a Mon Sep 17 00:00:00 2001 From: Prince Kumar Date: Sat, 2 Nov 2024 12:20:23 +0000 Subject: [PATCH] feat: adding support of read-stall-retry --- main.go | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/main.go b/main.go index 33edd58..517920f 100644 --- a/main.go +++ b/main.go @@ -8,7 +8,6 @@ import ( "io" "log" "net/http" - // Register the pprof endpoints under the web server root at /debug/pprof _ "net/http/pprof" "os" @@ -21,6 +20,7 @@ import ( "cloud.google.com/go/profiler" "cloud.google.com/go/storage" + "cloud.google.com/go/storage/experimental" "github.com/googleapis/gax-go/v2" "go.opencensus.io/stats" "golang.org/x/oauth2" @@ -68,6 +68,9 @@ var ( projectID = flag.String("project_id", "", "project ID to run profiler with; only required when running outside of GCP.") version = flag.String("version", "original", "version to run profiler with") + // Enable read stall retry. + enableReadStallRetry = flag.Bool("enable-read-stall-retry", false, "Enable read stall retry") + eG errgroup.Group ) @@ -113,7 +116,15 @@ func CreateHTTPClient(ctx context.Context, isHTTP2 bool) (client *storage.Client UserAgent: "prince", } - return storage.NewClient(ctx, option.WithHTTPClient(httpClient)) + if *enableReadStallRetry { + return storage.NewClient(ctx, option.WithHTTPClient(httpClient), + experimental.WithReadStallTimeout(&experimental.ReadStallTimeoutConfig{ + Min: time.Second, + TargetPercentile: 0.99, + })) + } else { + return storage.NewClient(ctx, option.WithHTTPClient(httpClient)) + } } // CreateGrpcClient creates grpc client.