adding retry option for control client

GoogleCloudPlatform · May 3, 2024 · d262ca9 · d262ca9
1 parent 23832da
commit d262ca9
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 43 deletions.
diff --git a/internal/storage/storage_handle.go b/internal/storage/storage_handle.go
@@ -18,7 +18,6 @@ import (
 	"fmt"
 	"net/http"
 	"os"
-	"time"
 
 	"cloud.google.com/go/storage"
 	control "cloud.google.com/go/storage/control/apiv2"
@@ -29,7 +28,6 @@ import (
 	"golang.org/x/net/context"
 	option "google.golang.org/api/option"
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
 
 	// Side effect to run grpc client with direct-path on gcp machine.
@@ -83,38 +81,22 @@ func createClientOptionForGRPCClient(clientConfig *storageutil.StorageClientConf
 	return
 }
 
-func StorageControlClientRetriesForCreateAndDeleteFolder(clientConfig *storageutil.StorageClientConfig) *control.StorageControlCallOptions {
+func storageControlClientRetryOptions(clientConfig *storageutil.StorageClientConfig) []gax.CallOption {
+	return []gax.CallOption{ gax.WithRetry(func() gax.Retryer {
+		return gax.OnErrorFunc(gax.Backoff{
+			Max:        clientConfig.MaxRetrySleep,
+			Multiplier: clientConfig.RetryMultiplier,
+		}, storageutil.ShouldRetry)
+	})}
+}
+
+func storageControlClientRetries(clientConfig *storageutil.StorageClientConfig) *control.StorageControlCallOptions {
 	return &control.StorageControlCallOptions{
-		CreateFolder: []gax.CallOption{
-			gax.WithRetry(func() gax.Retryer {
-				return gax.Retryer(
-					storage.WithBackoff(
-						gax.Backoff{
-							Max:        clientConfig.MaxRetrySleep,
-							Multiplier: clientConfig.RetryMultiplier,
-						}),
-					storage.WithErrorFunc(storageutil.ShouldRetry),
-					storage.WithPolicy(storage.RetryAlways)),
-			}
-		},
-		DeleteFolder: []gax.CallOption{
-			gax.WithTimeout(60000 * time.Millisecond),
-		},
-		GetFolder: []gax.CallOption{
-			gax.WithRetry(func() gax.Retryer {
-				return gax.OnCodes([]codes.Code{
-					codes.ResourceExhausted,
-					codes.Unavailable,
-					codes.DeadlineExceeded,
-					codes.Internal,
-					codes.Unknown,
-				}, gax.Backoff{
-					Initial:    1000 * time.Millisecond,
-					Max:        60000 * time.Millisecond,
-					Multiplier: 2.00,
-				})
-			}),
-		},
+		CreateFolder:     storageControlClientRetryOptions(clientConfig),
+		DeleteFolder:     storageControlClientRetryOptions(clientConfig),
+		GetFolder:        storageControlClientRetryOptions(clientConfig),
+		RenameFolder:     storageControlClientRetryOptions(clientConfig),
+		GetStorageLayout: storageControlClientRetryOptions(clientConfig),
 	}
 }
 
@@ -129,21 +111,13 @@ func createGRPCControlClientHandle(ctx context.Context, clientConfig *storageuti
 		return nil, fmt.Errorf("Error in getting clientOpts for gRPC client: %w", err)
 	}
 
-	retryer := gax.OnCodes([]codes.Code{
-		codes.Unavailable,      // Retry on server unavailable errors
-		codes.DeadlineExceeded, // Retry on deadline exceeded errors
-	}, gax.Backoff{
-		Initial:    100 * time.Millisecond,
-		Max:        1000 * time.Millisecond,
-		Multiplier: 1.3,
-	})
-
 	sc, err = control.NewStorageControlClient(ctx, clientOpts...)
 	if err != nil {
 		err = fmt.Errorf("NewStorageControlClient: %w", err)
 	}
 
-	sc = control.StorageControlClient{CallOptions: control.StorageControlCallOptions{CreateFolder:}}
+	// Set retries for control client.
+	sc = &control.StorageControlClient{CallOptions: storageControlClientRetries(clientConfig)}
 
 	// Unset the environment variable, since it's used only while creation of grpc client.
 	if err := os.Unsetenv("GOOGLE_CLOUD_ENABLE_DIRECT_PATH_XDS"); err != nil {

diff --git a/internal/storage/storage_handle_test.go b/internal/storage/storage_handle_test.go
@@ -296,3 +296,23 @@ func (testSuite *StorageHandleTest) TestCreateStorageHandleWithEnableHNSTrue() {
 	assert.Nil(testSuite.T(), err)
 	assert.NotNil(testSuite.T(),sh)
 }
+
+func (testSuite *StorageHandleTest) TestStorageControlClientRetryOptions() {
+	sc := storageutil.GetDefaultStorageClientConfig()
+
+	gaxOpts := storageControlClientRetryOptions(&sc)
+
+	assert.NotNil(testSuite.T(), gaxOpts)
+}
+
+func (testSuite *StorageHandleTest) TestStorageControlClientRetries() {
+	sc := storageutil.GetDefaultStorageClientConfig()
+
+  storageControlOpts := storageControlClientRetries(&sc)
+
+	assert.NotNil(testSuite.T(), storageControlOpts.CreateFolder)
+	assert.NotNil(testSuite.T(), storageControlOpts.DeleteFolder)
+	assert.NotNil(testSuite.T(), storageControlOpts.RenameFolder)
+	assert.NotNil(testSuite.T(), storageControlOpts.GetFolder)
+	assert.NotNil(testSuite.T(), storageControlOpts.GetStorageLayout)
+}