fix(storage): honor --enable-rapid-writes flag during object overwrites for pirlo buckets (#5004)

Author: vipnydav

internal/storage/bucket_handle.go

@@ -193,17 +193,29 @@ func (bh *bucketHandle) getObjectHandleWithPreconditionsSet(req *gcs.CreateObjec
 	return obj
 }
 
+func (bh *bucketHandle) getStorageClassForCreateObject(defaultStorageClass string) string {
+	// When rapid writes are enabled on an RCU bucket, objects should be explicitly
+	// created with the RAPID storage class in the zonal cache rather than
+	// defaulting to the bucket's default storage class.
+	// If rapid writes are disabled, we explicitly clear the storage class to prevent
+	// inheriting the RAPID storage class from an overwritten object and write it to
+	// the bucket's default storage class.
+	switch bh.BucketType().Pirlo {
+	case gcs.PirloStateRapidWritesEnabled:
+		return storageClassRapid
+	case gcs.PirloStateRapidWritesDisabled:
+		return ""
+	default:
+		return defaultStorageClass
+	}
+}
+
 func (bh *bucketHandle) CreateObject(ctx context.Context, req *gcs.CreateObjectRequest) (o *gcs.Object, err error) {
 	defer func() {
 		err = gcs.GetGCSError(err)
 	}()
 
-	// When rapid writes are enabled on an RCU bucket, objects should be explicitly
-	// created with the RAPID storage class in the zonal cache rather than
-	// defaulting to the bucket's default storage class
-	if bh.BucketType().Pirlo == gcs.PirloStateRapidWritesEnabled {
-		req.StorageClass = storageClassRapid
-	}
+	req.StorageClass = bh.getStorageClassForCreateObject(req.StorageClass)
 
 	obj := bh.getObjectHandleWithPreconditionsSet(req)
 
@@ -247,12 +259,7 @@ func (bh *bucketHandle) CreateObject(ctx context.Context, req *gcs.CreateObjectR
 }
 
 func (bh *bucketHandle) CreateObjectChunkWriter(ctx context.Context, req *gcs.CreateObjectRequest, chunkSize int, callBack func(bytesUploadedSoFar int64)) (gcs.Writer, error) {
-	// When rapid writes are enabled on an RCU bucket, objects should be explicitly
-	// created with the RAPID storage class in the zonal cache rather than
-	// defaulting to the bucket's default storage class.
-	if bh.BucketType().Pirlo == gcs.PirloStateRapidWritesEnabled {
-		req.StorageClass = storageClassRapid
-	}
+	req.StorageClass = bh.getStorageClassForCreateObject(req.StorageClass)
 
 	obj := bh.getObjectHandleWithPreconditionsSet(req)
 

internal/storage/bucket_handle_test.go

@@ -688,32 +688,45 @@ func (testSuite *BucketHandleTest) TestBucketHandle_StorageClassOverrides() {
 	bucketScenarios := []struct {
 		name                        string
 		bucketType                  gcs.BucketType
+		initialStorageClass         string
 		expectedWriterStorageClass  string
 		expectedCreatedStorageClass string
 		canTestCreateObject         bool
 	}{
 		{
 			name:                        "StandardBucket",
 			bucketType:                  gcs.BucketType{Pirlo: gcs.PirloStateNone},
+			initialStorageClass:         "",
 			expectedWriterStorageClass:  "",
 			expectedCreatedStorageClass: "STANDARD",
 			canTestCreateObject:         true,
 		},
+		{
+			name:                        "StandardBucket_InheritsStorageClass",
+			bucketType:                  gcs.BucketType{Pirlo: gcs.PirloStateNone},
+			initialStorageClass:         "COLDLINE",
+			expectedWriterStorageClass:  "COLDLINE",
+			expectedCreatedStorageClass: "COLDLINE",
+			canTestCreateObject:         true,
+		},
 		{
 			name:                       "ZonalBucket",
 			bucketType:                 gcs.BucketType{Zonal: true},
+			initialStorageClass:        "",
 			expectedWriterStorageClass: "",    // Zonal buckets do not use the RAPID storage class.
 			canTestCreateObject:        false, // Fails on HTTP append.
 		},
 		{
 			name:                       "PirloBucket_RapidEnabled",
 			bucketType:                 gcs.BucketType{Pirlo: gcs.PirloStateRapidWritesEnabled},
+			initialStorageClass:        "STANDARD", // Simulate inheriting from standard source
 			expectedWriterStorageClass: storageClassRapid,
 			canTestCreateObject:        false, // Fails on HTTP append.
 		},
 		{
 			name:                        "PirloBucket_RapidDisabled",
 			bucketType:                  gcs.BucketType{Pirlo: gcs.PirloStateRapidWritesDisabled},
+			initialStorageClass:         storageClassRapid, // Simulate inheriting from rapid source
 			expectedWriterStorageClass:  "",
 			expectedCreatedStorageClass: "STANDARD",
 			canTestCreateObject:         true,
@@ -725,7 +738,7 @@ func (testSuite *BucketHandleTest) TestBucketHandle_StorageClassOverrides() {
 			createBucketHandle(testSuite, &controlpb.StorageLayout{})
 			testSuite.bucketHandle.bucketType = &scenario.bucketType
 			testSuite.bucketHandle.writeConfig = &cfg.WriteConfig{}
-			req := &gcs.CreateObjectRequest{Name: "test_object_2"}
+			req := &gcs.CreateObjectRequest{Name: "test_object_2", StorageClass: scenario.initialStorageClass}
 
 			w, err := testSuite.bucketHandle.CreateObjectChunkWriter(context.Background(), req, 1024, nil)
 
@@ -740,7 +753,7 @@ func (testSuite *BucketHandleTest) TestBucketHandle_StorageClassOverrides() {
 				createBucketHandle(testSuite, &controlpb.StorageLayout{})
 				testSuite.bucketHandle.bucketType = &scenario.bucketType
 				testSuite.bucketHandle.writeConfig = &cfg.WriteConfig{}
-				req := &gcs.CreateObjectRequest{Name: "test_object_1", Contents: strings.NewReader("data")}
+				req := &gcs.CreateObjectRequest{Name: "test_object_1", Contents: strings.NewReader("data"), StorageClass: scenario.initialStorageClass}
 
 				o, err := testSuite.bucketHandle.CreateObject(context.Background(), req)