Add and implement zonal buckets config option.

Signed-off-by: NickGoog <[email protected]>

Author: NickGoog

README.md

@@ -373,6 +373,7 @@ config:
   bucket: ""
   service_account: ""
   use_grpc: false
+  use_zonal_buckets: false
   grpc_conn_pool_size: 0
   http_config:
     idle_conn_timeout: 0s

providers/gcs/gcs.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/storage"
+	"cloud.google.com/go/storage/experimental"
 	"github.com/go-kit/log"
 	"github.com/pkg/errors"
 	"github.com/prometheus/common/version"
@@ -42,6 +43,10 @@ type Config struct {
 	Bucket         string `yaml:"bucket"`
 	ServiceAccount string `yaml:"service_account"`
 	UseGRPC        bool   `yaml:"use_grpc"`
+	// Both upload and download paths will use zonal bucket gRPC APIs by default.
+	// Zonal buckets are currently allowlisted; please contact your Google account
+	// manager if interested.
+	UseZonalBuckets bool `yaml:"use_zonal_buckets"`
 	// GRPCConnPoolSize controls the size of the gRPC connection pool and should only be used
 	// when direct path is not enabled.
 	// See https://pkg.go.dev/cloud.google.com/go/storage#hdr-Experimental_gRPC_API for more details
@@ -164,8 +169,14 @@ func newBucket(ctx context.Context, logger log.Logger, gc Config, opts []option.
 		opts = append(opts,
 			option.WithGRPCConnectionPool(gc.GRPCConnPoolSize),
 		)
+		if gc.UseZonalBuckets {
+			opts = append(opts, experimental.WithZonalBucketAPIs())
+		}
 		gcsClient, err = storage.NewGRPCClient(ctx, opts...)
 	} else {
+		if gc.UseZonalBuckets {
+			return nil, errors.New("Zonal buckets require `use_grpc: true`")
+		}
 		gcsClient, err = storage.NewClient(ctx, opts...)
 	}
 	if err != nil {

providers/gcs/gcs_test.go

@@ -74,6 +74,54 @@ func TestNewBucketWithConfig_ShouldCreateGRPC(t *testing.T) {
 	testutil.Assert(t, bkt != nil, "expected bucket to be created")
 }
 
+func TestNewBucketWithConfig_ShouldCreateZonalBucket(t *testing.T) {
+	cfg := Config{
+		Bucket:          "test-bucket",
+		ServiceAccount:  "",
+		UseGRPC:         true,
+		UseZonalBuckets: true,
+	}
+
+	// Note, the emulator does not currently have any specific zonal bucket logic.
+	svr, err := gcsemu.NewServer("127.0.0.1:0", gcsemu.Options{})
+	testutil.Ok(t, err)
+	err = os.Setenv("STORAGE_EMULATOR_HOST", svr.Addr)
+	testutil.Ok(t, err)
+	err = os.Setenv("GCS_EMULATOR_HOST", svr.Addr)
+	testutil.Ok(t, err)
+	err = os.Setenv("STORAGE_EMULATOR_HOST_GRPC", svr.Addr)
+	testutil.Ok(t, err)
+
+	bkt, err := NewBucketWithConfig(context.Background(), log.NewNopLogger(), cfg, "test-bucket", nil)
+	testutil.Ok(t, err)
+
+	// Check if the bucket is created.
+	testutil.Assert(t, bkt != nil, "expected bucket to be created")
+}
+
+func TestNewBucketWithConfig_ShouldNotCreateZonalBucketWithoutGRPC(t *testing.T) {
+	cfg := Config{
+		Bucket:          "test-bucket",
+		ServiceAccount:  "",
+		UseGRPC:         false,
+		UseZonalBuckets: true,
+	}
+
+	// Note, the emulator does not currently have any specific zonal bucket logic.
+	svr, err := gcsemu.NewServer("127.0.0.1:0", gcsemu.Options{})
+	testutil.Ok(t, err)
+	err = os.Setenv("STORAGE_EMULATOR_HOST", svr.Addr)
+	testutil.Ok(t, err)
+	err = os.Setenv("GCS_EMULATOR_HOST", svr.Addr)
+	testutil.Ok(t, err)
+	err = os.Setenv("STORAGE_EMULATOR_HOST_GRPC", svr.Addr)
+	testutil.Ok(t, err)
+
+	bkt, err := NewBucketWithConfig(context.Background(), log.NewNopLogger(), cfg, "test-bucket", nil)
+	testutil.Assert(t, bkt == nil, "did not expect bucket to be created")
+	testutil.Equals(t, err.Error(), "Zonal buckets require `use_grpc: true`")
+}
+
 func TestParseConfig_ChunkSize(t *testing.T) {
 	for _, tc := range []struct {
 		name       string