diff --git a/storage/integration_test.go b/storage/integration_test.go index c40011446fed..abf72ceaa83c 100644 --- a/storage/integration_test.go +++ b/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 diff --git a/storage/pcu.go b/storage/pcu.go index 51b15acd93d7..5b504064d8e4 100644 --- a/storage/pcu.go +++ b/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) diff --git a/storage/pcu_test.go b/storage/pcu_test.go index 5e3a20517974..ffc30e7cdea6 100644 --- a/storage/pcu_test.go +++ b/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() }) diff --git a/storage/writer.go b/storage/writer.go index 62f53871cd81..b0fc7d5bbd86 100644 --- a/storage/writer.go +++ b/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)