chore(storage): support PartSizeHint and optimize cleanup

storage/integration_test.go

@@ -9323,26 +9323,26 @@ func TestIntegration_ParallelUpload(t *testing.T) {
 			{
 				name:     "small object",
 				content:  bytes.Repeat([]byte("a"), 1<<20), // 1 MiB
-				config:   ParallelUploadConfig{PartSize: 5 << 20, MaxConcurrency: 2},
+				config:   ParallelUploadConfig{PartSizeHint: 8 << 20, MaxConcurrency: 2},
 				expected: 1 << 20,
 			},
 			{
 				name:     "exact part size",
-				content:  bytes.Repeat([]byte("b"), 5<<20), // 5 MiB
-				config:   ParallelUploadConfig{PartSize: 5 << 20, MaxConcurrency: 2},
-				expected: 5 << 20,
+				content:  bytes.Repeat([]byte("b"), 8<<20), // 8 MiB
+				config:   ParallelUploadConfig{PartSizeHint: 8 << 20, MaxConcurrency: 2},
+				expected: 8 << 20,
 			},
 			{
 				name:     "large object",
-				content:  bytes.Repeat([]byte("c"), 12<<20), // 12 MiB
-				config:   ParallelUploadConfig{PartSize: 5 << 20, MaxConcurrency: 4},
-				expected: 12 << 20,
+				content:  bytes.Repeat([]byte("c"), 18<<20), // 18 MiB
+				config:   ParallelUploadConfig{PartSizeHint: 8 << 20, MaxConcurrency: 4},
+				expected: 18 << 20,
 			},
 			{
 				name:     "recursive composition logic",
-				content:  bytes.Repeat([]byte("d"), 165<<20), // 165 MiB = 33 parts of 5 MiB each (tests recursive compose)
-				config:   ParallelUploadConfig{PartSize: 5 << 20, MaxConcurrency: 16},
-				expected: 165 << 20,
+				content:  bytes.Repeat([]byte("d"), 264<<20), // 264 MiB = 33 parts of 8 MiB each (tests recursive compose)
+				config:   ParallelUploadConfig{PartSizeHint: 8 << 20, MaxConcurrency: 16},
+				expected: 264 << 20,
 			},
 		}
 
@@ -9434,8 +9434,8 @@ func TestIntegration_ParallelUploadConcurrency(t *testing.T) {
 	multiTransportTest(ctx, t, func(t *testing.T, ctx context.Context, bucket string, _ string, client *Client) {
 		h := testHelper{t}
 
-		content := bytes.Repeat([]byte("z"), 15<<20) // 15 MiB
-		config := ParallelUploadConfig{PartSize: 5 << 20, MaxConcurrency: 4}
+		content := bytes.Repeat([]byte("z"), 24<<20) // 24 MiB
+		config := ParallelUploadConfig{PartSizeHint: 8 << 20, MaxConcurrency: 4}
 
 		var wg sync.WaitGroup
 		numUploads := 5
@@ -9489,7 +9489,7 @@ func TestIntegration_ParallelUpload_ChecksumValidation(t *testing.T) {
 	multiTransportTest(ctx, t, func(t *testing.T, ctx context.Context, bucket string, _ string, client *Client) {
 		h := testHelper{t}
 
-		content := bytes.Repeat([]byte("z"), 12<<20) // 12 MiB
+		content := bytes.Repeat([]byte("z"), 18<<20) // 18 MiB
 		correctCRC := crc32.Checksum(content, crc32cTable)
 
 		testCases := []struct {
@@ -9519,7 +9519,7 @@ func TestIntegration_ParallelUpload_ChecksumValidation(t *testing.T) {
 
 				w := obj.NewWriter(ctx)
 				w.EnableParallelUpload = true
-				w.ParallelUploadConfig = ParallelUploadConfig{PartSize: 5 << 20, MaxConcurrency: 2}
+				w.ParallelUploadConfig = ParallelUploadConfig{PartSizeHint: 8 << 20, MaxConcurrency: 2}
 				w.SendCRC32C = true
 				w.CRC32C = tc.crc32c
 

storage/pcu.go

@@ -30,7 +30,7 @@ import (
 
 const (
 	defaultPartSize      = 16 * 1024 * 1024 // 16 MiB
-	minPartSize          = 5 * 1024 * 1024  // 5 MiB
+	minPartSize          = 8 * 1024 * 1024  // 8 MiB
 	baseWorkers          = 4
 	maxWorkers           = 16
 	tmpObjectPrefix      = "gcs-go-sdk-pu-tmp/"
@@ -47,15 +47,13 @@ const (
 //
 // **Note:** This feature is currently experimental and its API surface may change
 // in future releases. It is not yet recommended for production use.
-//
-// TODO(b/521239530): Add option to delete source parts after compose operation and remove cleanup logic.
 type ParallelUploadConfig struct {
 
-	// PartSize is the size of each part to be uploaded in parallel.
-	// Defaults to 16MiB. If a value less than 5MiB is provided, it will be
-	// automatically increased to 5MiB. The value is automatically rounded up
-	// to the nearest multiple of 256KiB.
-	PartSize int
+	// PartSizeHint is the preferred size of each part to be uploaded in parallel.
+	// This is used as a hint; the actual part size used may be different (e.g. server-determined).
+	// Defaults to 16MiB. If a value less than 8MiB is provided, it will be
+	// automatically increased to 8MiB.
+	PartSizeHint int
 
 	// MaxConcurrency is the number of goroutines to use for uploading parts in parallel.
 	// Defaults to a dynamic value based on the number of CPUs (min(4 + NumCPU/2, 16)).
@@ -64,10 +62,10 @@ type ParallelUploadConfig struct {
 
 // defaults fills in values for the configuration options.
 func (c *ParallelUploadConfig) defaults() {
-	if c.PartSize == 0 {
-		c.PartSize = defaultPartSize
-	} else if c.PartSize < minPartSize {
-		c.PartSize = minPartSize
+	if c.PartSizeHint == 0 {
+		c.PartSizeHint = defaultPartSize
+	} else if c.PartSizeHint < minPartSize {
+		c.PartSizeHint = minPartSize
 	}
 	// Use a heuristic for the number of workers: start with 4, add 1 for
 	// every 2 CPUs, but don't exceed a cap of 16. This provides a
@@ -118,13 +116,14 @@ type pcuState struct {
 	bytesBuffered   int64
 	buffersAlloc    int
 
-	bufferCh    chan []byte
-	uploadCh    chan uploadTask
-	resultCh    chan uploadResult
-	workerWG    sync.WaitGroup
-	collectorWG sync.WaitGroup
-	started     bool
-	closeOnce   sync.Once
+	bufferCh              chan []byte
+	uploadCh              chan uploadTask
+	resultCh              chan uploadResult
+	workerWG              sync.WaitGroup
+	collectorWG           sync.WaitGroup
+	started               bool
+	closeOnce             sync.Once
+	finalComposeSucceeded bool
 
 	// Function to upload a part; can be overridden for testing.
 	uploadPartFn func(s *pcuState, task uploadTask) (*ObjectHandle, *ObjectAttrs, error)
@@ -176,8 +175,8 @@ func (w *Writer) initPCU(ctx context.Context) error {
 	cfg := &w.ParallelUploadConfig
 	cfg.defaults()
 
-	// Ensure PartSize is a multiple of googleapi.MinUploadChunkSize.
-	cfg.PartSize = gRPCChunkSize(cfg.PartSize)
+	// Ensure PartSizeHint is a multiple of googleapi.MinUploadChunkSize.
+	cfg.PartSizeHint = gRPCChunkSize(cfg.PartSizeHint)
 
 	s := newPCUSettings(cfg.MaxConcurrency)
 
@@ -219,7 +218,7 @@ func (w *Writer) initPCU(ctx context.Context) error {
 	go state.resultCollector()
 
 	// Handle to get the first buffer.
-	state.currentBuffer = make([]byte, cfg.PartSize)
+	state.currentBuffer = make([]byte, cfg.PartSizeHint)
 	state.buffersAlloc = 1
 	state.bytesBuffered = 0
 
@@ -355,7 +354,7 @@ func (s *pcuState) write(p []byte) (int, error) {
 				s.bytesBuffered = 0
 			default:
 				if s.buffersAlloc < s.settings.bufferPoolSize {
-					s.currentBuffer = make([]byte, s.config.PartSize)
+					s.currentBuffer = make([]byte, s.config.PartSizeHint)
 					s.buffersAlloc++
 					s.bytesBuffered = 0
 				} else {
@@ -374,7 +373,7 @@ func (s *pcuState) write(p []byte) (int, error) {
 		p = p[n:]
 
 		// If the buffer is full, dispatch it to a worker.
-		if s.bytesBuffered == int64(s.config.PartSize) {
+		if s.bytesBuffered == int64(s.config.PartSizeHint) {
 			if err := s.flushCurrentBuffer(); err != nil {
 				return total - len(p), err
 			}
@@ -442,9 +441,18 @@ func (s *pcuState) close() error {
 		close(s.resultCh)
 		s.collectorWG.Wait()
 
-		// Cleanup is always attempted. We do it in the background to not block returning.
+		// Manual cleanup is only attempted if the upload failed or the final compose did not succeed.
+		// If the upload and composition succeed, the GCS compose operations will delete the source
+		// parts automatically (via DeleteSourceObjects = true).
+		// We do it in the background to not block returning.
 		defer func() {
-			go s.doCleanupFn(s)
+			s.mu.Lock()
+			hasErr := s.firstErr != nil
+			finalSucceeded := s.finalComposeSucceeded
+			s.mu.Unlock()
+			if hasErr && !finalSucceeded {
+				go s.doCleanupFn(s)
+			}
 		}()
 
 		s.mu.Lock()
@@ -525,6 +533,7 @@ func (s *pcuState) composeParts() error {
 
 				interHandle := s.w.o.c.Bucket(s.w.o.bucket).Object(compName)
 				composer := interHandle.ComposerFrom(finalComps[start:end]...)
+				composer.DeleteSourceObjects = true
 
 				_, err := s.composeFn(s.ctx, composer)
 				if err != nil {
@@ -555,12 +564,17 @@ func (s *pcuState) composeParts() error {
 	composer.ObjectAttrs = s.w.ObjectAttrs
 	composer.KMSKeyName = s.w.ObjectAttrs.KMSKeyName
 	composer.SendCRC32C = s.w.SendCRC32C
+	composer.DeleteSourceObjects = true
 
 	attrs, err := s.composeFn(s.ctx, composer)
 	if err != nil {
 		return err
 	}
 
+	s.mu.Lock()
+	s.finalComposeSucceeded = true
+	s.mu.Unlock()
+
 	// Perform client-side CRC32C validation if a user-provided checksum was specified.
 	if s.w.SendCRC32C && s.w.CRC32C != 0 && attrs.CRC32C != s.w.CRC32C {
 		return fmt.Errorf("storage: object was uploaded, but its CRC32C (%d) does not match the expected CRC32C (%d)", attrs.CRC32C, s.w.CRC32C)

storage/pcu_test.go

@@ -71,28 +71,28 @@ func TestParallelUploadConfig_defaults(t *testing.T) {
 			name: "all defaults",
 			in:   &ParallelUploadConfig{},
 			want: &ParallelUploadConfig{
-				PartSize:       defaultPartSize,
+				PartSizeHint:   defaultPartSize,
 				MaxConcurrency: expectedWorkers,
 			},
 		},
 		{
 			name: "user-provided values are respected",
 			in: &ParallelUploadConfig{
-				PartSize:       10 * 1024 * 1024, // 10 MiB
+				PartSizeHint:   10 * 1024 * 1024, // 10 MiB
 				MaxConcurrency: 10,
 			},
 			want: &ParallelUploadConfig{
-				PartSize:       10 * 1024 * 1024,
+				PartSizeHint:   10 * 1024 * 1024,
 				MaxConcurrency: 10,
 			},
 		},
 		{
-			name: "PartSize below minimum is adjusted",
+			name: "PartSizeHint below minimum is adjusted",
 			in: &ParallelUploadConfig{
-				PartSize: 1024 * 1024, // 1 MiB, below the 5 MiB minimum.
+				PartSizeHint: 1024 * 1024, // 1 MiB, below the 8 MiB minimum.
 			},
 			want: &ParallelUploadConfig{
-				PartSize:       minPartSize,
+				PartSizeHint:   minPartSize,
 				MaxConcurrency: expectedWorkers,
 			},
 		},
@@ -424,7 +424,7 @@ func TestPCUWorker_WriteContextCancellation(t *testing.T) {
 				c:      &Client{},
 			},
 		},
-		config:   &ParallelUploadConfig{PartSize: 10, MaxConcurrency: 1},
+		config:   &ParallelUploadConfig{PartSizeHint: 10, MaxConcurrency: 1},
 		settings: &pcuSettings{bufferPoolSize: 2},
 		partMap:  make(map[int]*ObjectHandle),
 		doCleanupFn: func(s *pcuState) {
@@ -448,7 +448,7 @@ func TestPCUWorker_WriteContextCancellation(t *testing.T) {
 		state.bufferCh <- make([]byte, 10)
 	}
 
-	// Trigger a flush by writing exact PartSize.
+	// Trigger a flush by writing exact part size.
 	n, err := state.write([]byte("0123456789"))
 	if err != nil {
 		t.Fatalf("state.write failed: %v", err)
@@ -554,7 +554,7 @@ func TestPCUState_Write(t *testing.T) {
 				started:  true,
 				bufferCh: make(chan []byte, 5),
 				uploadCh: make(chan uploadTask, 5),
-				config:   &ParallelUploadConfig{PartSize: PartSize},
+				config:   &ParallelUploadConfig{PartSizeHint: PartSize},
 			}
 
 			// Pre-fill buffer pool.
@@ -805,6 +805,10 @@ func TestPCUState_ComposeParts(t *testing.T) {
 				mu.Lock()
 				defer mu.Unlock()
 
+				if !c.DeleteSourceObjects {
+					t.Errorf("expected DeleteSourceObjects to be true, got false")
+				}
+
 				// If the destination isn't the final object, it's an intermediate.
 				if c.dst.object != "final-dest" {
 					intermediateCount++
@@ -1082,11 +1086,11 @@ func TestPCUState_Close(t *testing.T) {
 				t.Errorf("expectCompose %v, but composeCalled was %v", tc.expectCompose, composeCalled)
 			}
 
-			// Wait for background cleanup to execute
+			// Wait for background cleanup to execute if failure is expected.
 			time.Sleep(10 * time.Millisecond)
 			mu.Lock()
-			if !cleanupCalled {
-				t.Errorf("cleanup logic was not executed")
+			if cleanupCalled != tc.expectError {
+				t.Errorf("cleanupCalled = %v; want %v", cleanupCalled, tc.expectError)
 			}
 			mu.Unlock()
 		})

storage/writer.go

@@ -180,7 +180,7 @@ type Writer struct {
 	//
 	// For parallel uploads, progress is reported when each part is successfully uploaded.
 	// Therefore, the progress may be delayed relative to the standard upload,
-	// and jump in increments of PartSize (e.g. 16MiB).
+	// and jump in increments of the part size used (e.g. 16MiB).
 	//
 	// ProgressFunc should return quickly without blocking.
 	ProgressFunc func(int64)