Skip to content

chore(storage): support PartSizeHint and optimize cleanup #19930

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cpriti-os wants to merge 4 commits into
base: main
Choose a base branch
from
+71 −53
Open

chore(storage): support PartSizeHint and optimize cleanup #19930

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3fece61
chore(storage): support PartSizeHint and optimize parallel upload com…
cpriti-os Jun 9, 2026
90ee391
chore(storage): fix goimports formatting in pcu.go
cpriti-os Jun 9, 2026
bfd5d8c
improve comments
cpriti-os Jun 10, 2026
6bd94d9
Merge branch 'main' into pcu-improve
cpriti-os Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 14 additions & 14 deletions storage/integration_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8904,26 +8904,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,
},
}

Expand Down Expand Up @@ -9015,8 +9015,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
Expand Down Expand Up @@ -9070,7 +9070,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 {
Expand Down Expand Up @@ -9100,7 +9100,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

Expand Down
50 changes: 31 additions & 19 deletions storage/pcu.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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/"
Expand All @@ -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.
Comment thread
cpriti-os marked this conversation as resolved.
// 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)).
Expand All @@ -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
Expand Down Expand Up @@ -125,6 +123,7 @@ type pcuState struct {
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)
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -218,7 +217,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

Expand Down Expand Up @@ -354,7 +353,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 {
Expand All @@ -373,7 +372,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
}
Expand Down Expand Up @@ -441,9 +440,16 @@ 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.
// Cleanup is only attempted if the upload failed and the final compose did not succeed.
// We do it in the background to not block returning.
Comment thread
cpriti-os marked this conversation as resolved.
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)
}
}()
Comment on lines 448 to 456

@gemini-code-assist gemini-code-assist Bot Jun 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

When the parent context is canceled, the background cleanup goroutine go s.doCleanupFn(s) will run using the canceled context s.ctx. This causes all subsequent GCS API calls (such as deleting temporary objects) to fail immediately with context.Canceled, resulting in leaked temporary objects in the bucket.

To prevent this, use context.WithoutCancel to ensure the cleanup operation is not prematurely terminated when the parent context is canceled.

Suggested change
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)
}
}()
defer func() {
s.mu.Lock()
hasErr := s.firstErr != nil
finalSucceeded := s.finalComposeSucceeded
s.mu.Unlock()
if hasErr && !finalSucceeded {
s.mu.Lock()
s.ctx = context.WithoutCancel(s.ctx)
s.mu.Unlock()
go s.doCleanupFn(s)
}
}()
References
  1. When a cleanup function must be attempted even if its parent context is canceled, use context.WithoutCancel to ensure the operation is not prematurely terminated.


s.mu.Lock()
Expand Down Expand Up @@ -524,6 +530,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 {
Expand Down Expand Up @@ -554,12 +561,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)
Expand Down
28 changes: 16 additions & 12 deletions storage/pcu_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
},
Expand Down Expand Up @@ -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) {
Expand All @@ -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)
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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++
Expand Down Expand Up @@ -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()
})
Expand Down
2 changes: 1 addition & 1 deletion storage/writer.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,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)
Expand Down
Loading