diff --git a/pkg/objstore/azblob.go b/pkg/objstore/azblob.go index f9186248a43e5..9bfc3fe9e05b9 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.PartSize > 0 { + chunkSize = int(wo.PartSize) + } + if wo != nil && wo.Concurrency > 1 { + 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,66 @@ type azblobUploader struct { cpkScope *blob.CPKScopeInfo cpkInfo *blob.CPKInfo + + // eg/egCtx drive concurrent upload; nil means Write stages blocks + // 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 error +} + +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 errors.Annotate(err, "Failed to upload block to azure blob") + } + return nil } func (u *azblobUploader) Write(ctx context.Context, data []byte) (int, error) { + if u.eg != nil { + if err := u.egCtx.Err(); err != nil { + return 0, err + } + } generatedUUID, err := uuid.NewUUID() if err != nil { return 0, errors.Annotate(err, "Fail to generate uuid") } blockID := base64.StdEncoding.EncodeToString([]byte(generatedUUID.String())) + u.blockIDList = append(u.blockIDList, blockID) - _, 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") + 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) + 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 { + err = u.eg.Wait() + } + if 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 options := &blockblob.CommitBlockListOptions{ CPKScopeInfo: u.cpkScope, @@ -920,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) } 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) +} 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