From 148242fc80b26a716f9b3bf1cca68db93cea1153 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 23 Jul 2026 11:56:08 +0800 Subject: [PATCH 1/9] objstore: add concurrent Azure block upload via errgroup AzureBlobStorage.Create ignored WriterOption, so azblobUploader staged blocks serially regardless of the requested concurrency, while S3 and GCS already had concurrent writers. It now stages blocks concurrently with an errgroup bounded by SetLimit(Concurrency), chunked by PartSize. blockID is generated in the serial Write and appended in file order, so the block list needs no locking; the derived context cancels in-flight StageBlock calls as soon as one fails, and Close skips CommitBlockList on error (staged blocks are garbage-collected by Azure). Callers passing no option keep the original synchronous path. Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/objstore/BUILD.bazel | 2 +- pkg/objstore/azblob.go | 73 +++++++++++++++++++++++++++++++++---- pkg/objstore/azblob_test.go | 30 +++++++++++++++ 3 files changed, 96 insertions(+), 9 deletions(-) diff --git a/pkg/objstore/BUILD.bazel b/pkg/objstore/BUILD.bazel index 413bdc86362a0..ee0a8c7a80ff1 100644 --- a/pkg/objstore/BUILD.bazel +++ b/pkg/objstore/BUILD.bazel @@ -86,7 +86,7 @@ go_test( ], embed = [":objstore"], flaky = True, - shard_count = 45, + shard_count = 46, deps = [ "//pkg/objstore/compressedio", "//pkg/objstore/objectio", diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index f9186248a43e5..98aab979adac2 100644 --- a/pkg/objstore/azblob.go +++ b/pkg/objstore/azblob.go @@ -48,6 +48,7 @@ import ( "github.com/pingcap/tidb/pkg/objstore/storeapi" "github.com/spf13/pflag" "go.uber.org/zap" + "golang.org/x/sync/errgroup" ) const ( @@ -717,7 +718,7 @@ func (s *AzureBlobStorage) URI() string { const azblobChunkSize = 64 * 1024 * 1024 // Create implements the StorageWriter interface. -func (s *AzureBlobStorage) Create(_ context.Context, name string, _ *storeapi.WriterOption) (objectio.Writer, error) { +func (s *AzureBlobStorage) Create(ctx context.Context, name string, wo *storeapi.WriterOption) (objectio.Writer, error) { client := s.containerClient.NewBlockBlobClient(s.withPrefix(name)) uploader := &azblobUploader{ blobClient: client, @@ -730,7 +731,18 @@ func (s *AzureBlobStorage) Create(_ context.Context, name string, _ *storeapi.Wr cpkInfo: s.cpkInfo, } - uploaderWriter := objectio.NewBufferedWriter(uploader, azblobChunkSize, compressedio.NoCompression, nil) + chunkSize := azblobChunkSize + if wo != nil && wo.Concurrency > 1 { + if wo.PartSize > 0 { + chunkSize = int(wo.PartSize) + } + eg, egCtx := errgroup.WithContext(ctx) + eg.SetLimit(wo.Concurrency) + uploader.eg = eg + uploader.egCtx = egCtx + } + + uploaderWriter := objectio.NewBufferedWriter(uploader, chunkSize, compressedio.NoCompression, nil) return uploaderWriter, nil } @@ -889,28 +901,73 @@ type azblobUploader struct { cpkScope *blob.CPKScopeInfo cpkInfo *blob.CPKInfo + + // eg is set only when concurrent upload is requested through WriterOption. + // When it is nil, Write stages blocks synchronously, which stays the default + // path for callers that pass no option. egCtx cancels the in-flight stagers + // as soon as one of them fails. + eg *errgroup.Group + egCtx context.Context } -func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { +func newBlockID() (string, error) { generatedUUID, err := uuid.NewUUID() if err != nil { - return 0, errors.Annotate(err, "Fail to generate uuid") + return "", errors.Annotate(err, "Fail to generate uuid") } - blockID := base64.StdEncoding.EncodeToString([]byte(generatedUUID.String())) + // A block ID must be a fixed-length string within a single commit; the + // canonical UUID is always 36 chars, so its base64 form is always 48. + return base64.StdEncoding.EncodeToString([]byte(generatedUUID.String())), nil +} - _, err = u.blobClient.StageBlock(ctx, blockID, newNopCloser(bytes.NewReader(data)), &blockblob.StageBlockOptions{ +func (u *azblobUploader) stageBlock(ctx context.Context, blockID string, data []byte) error { + _, err := u.blobClient.StageBlock(ctx, blockID, newNopCloser(bytes.NewReader(data)), &blockblob.StageBlockOptions{ CPKScopeInfo: u.cpkScope, CPKInfo: u.cpkInfo, }) if err != nil { - return 0, errors.Annotate(err, "Failed to upload block to azure blob") + return errors.Annotate(err, "Failed to upload block to azure blob") } - u.blockIDList = append(u.blockIDList, blockID) + return nil +} +func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { + blockID, err := newBlockID() + if err != nil { + return 0, err + } + + if u.eg == nil { + if err := u.stageBlock(ctx, blockID, data); err != nil { + return 0, err + } + u.blockIDList = append(u.blockIDList, blockID) + return len(data), nil + } + + // Write is called serially by the buffered writer, so appending here keeps + // blockIDList in file order; the background stagers never touch it. + u.blockIDList = append(u.blockIDList, blockID) + // The buffered writer reuses its backing array once Write returns, so copy + // before staging on a background goroutine. SetLimit bounds the in-flight + // uploads, so Go blocks here once the pool is full. + buf := make([]byte, len(data)) + copy(buf, data) + u.eg.Go(func() error { + return u.stageBlock(u.egCtx, blockID, buf) + }) return len(data), nil } func (u *azblobUploader) Close(ctx context.Context) error { + if u.eg != nil { + if err := u.eg.Wait(); err != nil { + // Staged but uncommitted blocks are garbage-collected by Azure, so + // skip CommitBlockList and surface the failure instead. + return err + } + } + // the encryption scope and the access tier can not be both in the HTTP headers options := &blockblob.CommitBlockListOptions{ CPKScopeInfo: u.cpkScope, diff --git a/pkg/objstore/azblob_test.go b/pkg/objstore/azblob_test.go index ec166abf896f2..d7000eee3e963 100644 --- a/pkg/objstore/azblob_test.go +++ b/pkg/objstore/azblob_test.go @@ -513,3 +513,33 @@ func TestCopyObject(t *testing.T) { require.NoError(t, err) require.Equal(t, srcReader, reader) } + +func TestAzblobConcurrentUpload(t *testing.T) { + ctx := context.Background() + options := &backuppb.AzureBlobStorage{Bucket: "test", Prefix: "concurrent/"} + builder := &sharedKeyAzuriteClientBuilder{} + skip, err := createContainer(ctx, builder, options.Bucket) + if skip || err != nil { + t.Skip("azurite is not running, skip test") + } + + azblobStorage, err := newAzureBlobStorageWithClientBuilder(ctx, options, builder) + require.NoError(t, err) + + data := make([]byte, 20*1024*1024) + _, err = rand.Read(data) + require.NoError(t, err) + + w, err := azblobStorage.Create(ctx, "concurrent.bin", &storeapi.WriterOption{ + Concurrency: 4, + PartSize: 4 * 1024 * 1024, + }) + require.NoError(t, err) + _, err = w.Write(ctx, data) + require.NoError(t, err) + require.NoError(t, w.Close(ctx)) + + got, err := azblobStorage.ReadFile(ctx, "concurrent.bin") + require.NoError(t, err) + require.Equal(t, data, got) +} From 0fcb3f8b7ece3a8b54a549fdf12e69d95c3381a7 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 23 Jul 2026 14:25:45 +0800 Subject: [PATCH 2/9] objstore: trim azblob uploader comments Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/objstore/azblob.go | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index 98aab979adac2..d7963bde8fba0 100644 --- a/pkg/objstore/azblob.go +++ b/pkg/objstore/azblob.go @@ -902,10 +902,8 @@ type azblobUploader struct { cpkScope *blob.CPKScopeInfo cpkInfo *blob.CPKInfo - // eg is set only when concurrent upload is requested through WriterOption. - // When it is nil, Write stages blocks synchronously, which stays the default - // path for callers that pass no option. egCtx cancels the in-flight stagers - // as soon as one of them fails. + // eg/egCtx are non-nil only for concurrent upload; otherwise Write stages + // blocks synchronously. egCtx cancels in-flight stagers on the first failure. eg *errgroup.Group egCtx context.Context } @@ -915,8 +913,6 @@ func newBlockID() (string, error) { if err != nil { return "", errors.Annotate(err, "Fail to generate uuid") } - // A block ID must be a fixed-length string within a single commit; the - // canonical UUID is always 36 chars, so its base64 form is always 48. return base64.StdEncoding.EncodeToString([]byte(generatedUUID.String())), nil } From f1f2f5efd74b00c50e3e44b5a59d2422599d0648 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 23 Jul 2026 14:31:53 +0800 Subject: [PATCH 3/9] objstore: drop verbose comments in azblob Write Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/objstore/azblob.go | 5 ----- 1 file changed, 5 deletions(-) diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index d7963bde8fba0..e5cedfa4a8222 100644 --- a/pkg/objstore/azblob.go +++ b/pkg/objstore/azblob.go @@ -941,12 +941,7 @@ func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { return len(data), nil } - // Write is called serially by the buffered writer, so appending here keeps - // blockIDList in file order; the background stagers never touch it. u.blockIDList = append(u.blockIDList, blockID) - // The buffered writer reuses its backing array once Write returns, so copy - // before staging on a background goroutine. SetLimit bounds the in-flight - // uploads, so Go blocks here once the pool is full. buf := make([]byte, len(data)) copy(buf, data) u.eg.Go(func() error { From 2b09de73b1da3640ff2d819ff0d3c58f9b286234 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 23 Jul 2026 14:34:32 +0800 Subject: [PATCH 4/9] objstore: inline blockID generation into azblob Write Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/objstore/azblob.go | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index e5cedfa4a8222..4862b184dd80c 100644 --- a/pkg/objstore/azblob.go +++ b/pkg/objstore/azblob.go @@ -908,14 +908,6 @@ type azblobUploader struct { egCtx context.Context } -func newBlockID() (string, error) { - generatedUUID, err := uuid.NewUUID() - if err != nil { - return "", errors.Annotate(err, "Fail to generate uuid") - } - return base64.StdEncoding.EncodeToString([]byte(generatedUUID.String())), nil -} - func (u *azblobUploader) stageBlock(ctx context.Context, blockID string, data []byte) error { _, err := u.blobClient.StageBlock(ctx, blockID, newNopCloser(bytes.NewReader(data)), &blockblob.StageBlockOptions{ CPKScopeInfo: u.cpkScope, @@ -928,10 +920,11 @@ func (u *azblobUploader) stageBlock(ctx context.Context, blockID string, data [] } func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { - blockID, err := newBlockID() + generatedUUID, err := uuid.NewUUID() if err != nil { - return 0, err + return 0, errors.Annotate(err, "Fail to generate uuid") } + blockID := base64.StdEncoding.EncodeToString([]byte(generatedUUID.String())) if u.eg == nil { if err := u.stageBlock(ctx, blockID, data); err != nil { From f1ae5c7dd0617a2d82b0880393cabf50388a7fe0 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 23 Jul 2026 14:41:46 +0800 Subject: [PATCH 5/9] objstore: note Azure uncommitted-block GC window in the comment Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/objstore/azblob.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index 4862b184dd80c..b44ddfee40e42 100644 --- a/pkg/objstore/azblob.go +++ b/pkg/objstore/azblob.go @@ -946,7 +946,7 @@ func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { func (u *azblobUploader) Close(ctx context.Context) error { if u.eg != nil { if err := u.eg.Wait(); err != nil { - // Staged but uncommitted blocks are garbage-collected by Azure, so + // Uncommitted blocks are garbage-collected by Azure after a week, so // skip CommitBlockList and surface the failure instead. return err } From be601387e8561027cbe85d16c13e6c206f1d0359 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 23 Jul 2026 14:53:54 +0800 Subject: [PATCH 6/9] objstore: make azblob sync upload all-or-nothing too Record the first stage failure in a shared atomic.Error inside stageBlock, so both the synchronous and concurrent paths skip CommitBlockList on any failure (previously the sync path could commit a truncated blob from the blocks staged before the failure). Close now checks the shared error uniformly. Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/objstore/azblob.go | 42 +++++++++++++++++++++++------------------- 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index b44ddfee40e42..8b91f81c73b88 100644 --- a/pkg/objstore/azblob.go +++ b/pkg/objstore/azblob.go @@ -47,6 +47,7 @@ import ( "github.com/pingcap/tidb/pkg/objstore/objectio" "github.com/pingcap/tidb/pkg/objstore/storeapi" "github.com/spf13/pflag" + "go.uber.org/atomic" "go.uber.org/zap" "golang.org/x/sync/errgroup" ) @@ -902,10 +903,11 @@ type azblobUploader struct { cpkScope *blob.CPKScopeInfo cpkInfo *blob.CPKInfo - // eg/egCtx are non-nil only for concurrent upload; otherwise Write stages - // blocks synchronously. egCtx cancels in-flight stagers on the first failure. + // eg/egCtx drive concurrent upload; nil means Write stages blocks + // synchronously. err records the first stage failure on either path. eg *errgroup.Group egCtx context.Context + err atomic.Error } func (u *azblobUploader) stageBlock(ctx context.Context, blockID string, data []byte) error { @@ -914,7 +916,9 @@ func (u *azblobUploader) stageBlock(ctx context.Context, blockID string, data [] CPKInfo: u.cpkInfo, }) if err != nil { - return errors.Annotate(err, "Failed to upload block to azure blob") + err = errors.Annotate(err, "Failed to upload block to azure blob") + u.err.CompareAndSwap(nil, err) + return err } return nil } @@ -925,31 +929,31 @@ func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { return 0, errors.Annotate(err, "Fail to generate uuid") } blockID := base64.StdEncoding.EncodeToString([]byte(generatedUUID.String())) + u.blockIDList = append(u.blockIDList, blockID) - if u.eg == nil { - if err := u.stageBlock(ctx, blockID, data); err != nil { - return 0, err - } - u.blockIDList = append(u.blockIDList, blockID) + if u.eg != nil { + buf := make([]byte, len(data)) + copy(buf, data) + u.eg.Go(func() error { + return u.stageBlock(u.egCtx, blockID, buf) + }) return len(data), nil } - u.blockIDList = append(u.blockIDList, blockID) - buf := make([]byte, len(data)) - copy(buf, data) - u.eg.Go(func() error { - return u.stageBlock(u.egCtx, blockID, buf) - }) + if err := u.stageBlock(ctx, blockID, data); err != nil { + return 0, err + } return len(data), nil } func (u *azblobUploader) Close(ctx context.Context) error { if u.eg != nil { - if err := u.eg.Wait(); err != nil { - // Uncommitted blocks are garbage-collected by Azure after a week, so - // skip CommitBlockList and surface the failure instead. - return err - } + _ = u.eg.Wait() + } + if err := u.err.Load(); err != nil { + // Uncommitted blocks are garbage-collected by Azure after a week, so + // skip CommitBlockList and surface the failure instead. + return err } // the encryption scope and the access tier can not be both in the HTTP headers From 03821d9a4b46ebf8737e7adc3b0085191a42e237 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 23 Jul 2026 16:41:43 +0800 Subject: [PATCH 7/9] objstore: honor Azure PartSize regardless of concurrency; fix Create doc Address review feedback: - azblob Create now applies WriterOption.PartSize whenever it is set, matching the s3like backend, instead of only when Concurrency > 1. - Storage.Create doc no longer claims only s3 implements WriterOption; it now points readers to each backend's Create. Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/objstore/azblob.go | 6 +++--- pkg/objstore/storeapi/storage.go | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index 8b91f81c73b88..260a6025325ea 100644 --- a/pkg/objstore/azblob.go +++ b/pkg/objstore/azblob.go @@ -733,10 +733,10 @@ func (s *AzureBlobStorage) Create(ctx context.Context, name string, wo *storeapi } chunkSize := azblobChunkSize + if wo != nil && wo.PartSize > 0 { + chunkSize = int(wo.PartSize) + } if wo != nil && wo.Concurrency > 1 { - if wo.PartSize > 0 { - chunkSize = int(wo.PartSize) - } eg, egCtx := errgroup.WithContext(ctx) eg.SetLimit(wo.Concurrency) uploader.eg = eg diff --git a/pkg/objstore/storeapi/storage.go b/pkg/objstore/storeapi/storage.go index ca48bb63ad71d..0f1b51f17080c 100644 --- a/pkg/objstore/storeapi/storage.go +++ b/pkg/objstore/storeapi/storage.go @@ -164,8 +164,8 @@ type Storage interface { URI() string // Create opens a file writer by path. path is relative path to storage base - // path. The old file under same path will be overwritten. Currently only s3 - // implemented WriterOption. + // path. The old file under same path will be overwritten. Not all backends + // honor WriterOption; see each backend's Create implementation. Create(ctx context.Context, path string, option *WriterOption) (objectio.Writer, error) // Rename file name from oldFileName to newFileName Rename(ctx context.Context, oldFileName, newFileName string) error From 72358764b6b713bef6f60ce7825800fc3cd567a1 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 23 Jul 2026 17:01:55 +0800 Subject: [PATCH 8/9] objstore: fail fast in azblob Write once a block upload has failed Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/objstore/azblob.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index 260a6025325ea..cfe69c938a633 100644 --- a/pkg/objstore/azblob.go +++ b/pkg/objstore/azblob.go @@ -924,6 +924,9 @@ func (u *azblobUploader) stageBlock(ctx context.Context, blockID string, data [] } func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { + if err := u.err.Load(); err != nil { + return 0, err + } generatedUUID, err := uuid.NewUUID() if err != nil { return 0, errors.Annotate(err, "Fail to generate uuid") From ff6b68c02279f5127a53a303344fb9267eef7ffc Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 24 Jul 2026 10:37:53 +0800 Subject: [PATCH 9/9] objstore: simplify azblob concurrent error handling Address review feedback: - stageBlock is now a pure function; the first-error recording no longer happens as a side effect inside it. - The concurrent path's error comes from errgroup's Wait(); err only records the synchronous path's failure, so it can be a plain error instead of an atomic one. Close reads whichever applies. - Write's fail-fast now checks egCtx (cancelled by errgroup on the first failure) instead of the shared error. Co-Authored-By: Claude Opus 4.8 (1M context) --- pkg/objstore/azblob.go | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index cfe69c938a633..9bfc3fe9e05b9 100644 --- a/pkg/objstore/azblob.go +++ b/pkg/objstore/azblob.go @@ -47,7 +47,6 @@ import ( "github.com/pingcap/tidb/pkg/objstore/objectio" "github.com/pingcap/tidb/pkg/objstore/storeapi" "github.com/spf13/pflag" - "go.uber.org/atomic" "go.uber.org/zap" "golang.org/x/sync/errgroup" ) @@ -904,10 +903,11 @@ type azblobUploader struct { cpkInfo *blob.CPKInfo // eg/egCtx drive concurrent upload; nil means Write stages blocks - // synchronously. err records the first stage failure on either path. + // synchronously. err holds the synchronous path's stage failure; the + // concurrent path's error comes from eg.Wait(). eg *errgroup.Group egCtx context.Context - err atomic.Error + err error } func (u *azblobUploader) stageBlock(ctx context.Context, blockID string, data []byte) error { @@ -916,16 +916,16 @@ func (u *azblobUploader) stageBlock(ctx context.Context, blockID string, data [] CPKInfo: u.cpkInfo, }) if err != nil { - err = errors.Annotate(err, "Failed to upload block to azure blob") - u.err.CompareAndSwap(nil, err) - return err + return errors.Annotate(err, "Failed to upload block to azure blob") } return nil } func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { - if err := u.err.Load(); err != nil { - return 0, err + if u.eg != nil { + if err := u.egCtx.Err(); err != nil { + return 0, err + } } generatedUUID, err := uuid.NewUUID() if err != nil { @@ -944,16 +944,18 @@ func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { } if err := u.stageBlock(ctx, blockID, data); err != nil { + u.err = err return 0, err } return len(data), nil } func (u *azblobUploader) Close(ctx context.Context) error { + err := u.err if u.eg != nil { - _ = u.eg.Wait() + err = u.eg.Wait() } - if err := u.err.Load(); err != nil { + if err != nil { // Uncommitted blocks are garbage-collected by Azure after a week, so // skip CommitBlockList and surface the failure instead. return err @@ -968,6 +970,6 @@ func (u *azblobUploader) Close(ctx context.Context) error { if len(u.accessTier) > 0 { options.Tier = &u.accessTier } - _, err := u.blobClient.CommitBlockList(ctx, u.blockIDList, options) + _, err = u.blobClient.CommitBlockList(ctx, u.blockIDList, options) return errors.Trace(err) }