*: support concurrent write for S3 writer (pingcap#45723) (pingcap#49185)

ti-chi-bot · web-flow · commit 79bcffb362b7 · 2023-12-06T09:22:21.000Z
ref pingcap#45719, close pingcap#48607
diff --git a/br/pkg/mock/storage/storage.go b/br/pkg/mock/storage/storage.go
diff --git a/br/pkg/storage/azblob.go b/br/pkg/storage/azblob.go
@@ -415,7 +415,7 @@ func (s *AzureBlobStorage) URI() string {
 }
 
 // Create implements the StorageWriter interface.
-func (s *AzureBlobStorage) Create(_ context.Context, name string) (ExternalFileWriter, error) {
+func (s *AzureBlobStorage) Create(_ context.Context, name string, _ *WriterOption) (ExternalFileWriter, error) {
 	client := s.containerClient.NewBlockBlobClient(s.withPrefix(name))
 	uploader := &azblobUploader{
 		blobClient: client,
diff --git a/br/pkg/storage/compress.go b/br/pkg/storage/compress.go
@@ -24,15 +24,15 @@ func WithCompression(inner ExternalStorage, compressionType CompressType) Extern
 	return &withCompression{ExternalStorage: inner, compressType: compressionType}
 }
 
-func (w *withCompression) Create(ctx context.Context, name string) (ExternalFileWriter, error) {
+func (w *withCompression) Create(ctx context.Context, name string, _ *WriterOption) (ExternalFileWriter, error) {
 	var (
 		writer ExternalFileWriter
 		err    error
 	)
 	if s3Storage, ok := w.ExternalStorage.(*S3Storage); ok {
 		writer, err = s3Storage.CreateUploader(ctx, name)
 	} else {
-		writer, err = w.ExternalStorage.Create(ctx, name)
+		writer, err = w.ExternalStorage.Create(ctx, name, nil)
 	}
 	if err != nil {
 		return nil, errors.Trace(err)
diff --git a/br/pkg/storage/gcs.go b/br/pkg/storage/gcs.go
@@ -246,7 +246,7 @@ func (s *GCSStorage) URI() string {
 }
 
 // Create implements ExternalStorage interface.
-func (s *GCSStorage) Create(ctx context.Context, name string) (ExternalFileWriter, error) {
+func (s *GCSStorage) Create(ctx context.Context, name string, _ *WriterOption) (ExternalFileWriter, error) {
 	object := s.objectName(name)
 	wc := s.bucket.Object(object).NewWriter(ctx)
 	wc.StorageClass = s.gcs.StorageClass
diff --git a/br/pkg/storage/hdfs.go b/br/pkg/storage/hdfs.go
@@ -123,7 +123,7 @@ func (s *HDFSStorage) URI() string {
 }
 
 // Create opens a file writer by path. path is relative path to storage base path
-func (*HDFSStorage) Create(_ context.Context, _ string) (ExternalFileWriter, error) {
+func (*HDFSStorage) Create(_ context.Context, _ string, _ *WriterOption) (ExternalFileWriter, error) {
 	return nil, errors.Annotatef(berrors.ErrUnsupportedOperation, "currently HDFS backend only support rawkv backup")
 }
 
diff --git a/br/pkg/storage/local.go b/br/pkg/storage/local.go
@@ -131,7 +131,7 @@ func (l *LocalStorage) Open(_ context.Context, path string) (ExternalFileReader,
 }
 
 // Create implements ExternalStorage interface.
-func (l *LocalStorage) Create(_ context.Context, name string) (ExternalFileWriter, error) {
+func (l *LocalStorage) Create(_ context.Context, name string, _ *WriterOption) (ExternalFileWriter, error) {
 	file, err := os.Create(filepath.Join(l.base, name))
 	if err != nil {
 		return nil, errors.Trace(err)
diff --git a/br/pkg/storage/local_test.go b/br/pkg/storage/local_test.go
@@ -25,7 +25,7 @@ func TestDeleteFile(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, false, ret)
 
-	_, err = store.Create(context.Background(), name)
+	_, err = store.Create(context.Background(), name, nil)
 	require.NoError(t, err)
 
 	ret, err = store.FileExists(context.Background(), name)
diff --git a/br/pkg/storage/memstore.go b/br/pkg/storage/memstore.go
@@ -219,7 +219,7 @@ func (*MemStorage) URI() string {
 // Create creates a file and returning a writer to write data into.
 // When the writer is closed, the data is stored in the file.
 // It implements the `ExternalStorage` interface
-func (s *MemStorage) Create(ctx context.Context, name string) (ExternalFileWriter, error) {
+func (s *MemStorage) Create(ctx context.Context, name string, _ *WriterOption) (ExternalFileWriter, error) {
 	select {
 	case <-ctx.Done():
 		return nil, ctx.Err()
diff --git a/br/pkg/storage/memstore_test.go b/br/pkg/storage/memstore_test.go
@@ -50,7 +50,7 @@ func TestMemStoreBasic(t *testing.T) {
 	require.NotNil(t, err)
 
 	// create a writer to write
-	w, err := store.Create(ctx, "/hello.txt")
+	w, err := store.Create(ctx, "/hello.txt", nil)
 	require.Nil(t, err)
 	_, err = w.Write(ctx, []byte("hello world 3"))
 	require.Nil(t, err)
diff --git a/br/pkg/storage/noop.go b/br/pkg/storage/noop.go
@@ -43,7 +43,7 @@ func (*noopStorage) URI() string {
 }
 
 // Create implements ExternalStorage interface.
-func (*noopStorage) Create(_ context.Context, _ string) (ExternalFileWriter, error) {
+func (*noopStorage) Create(_ context.Context, _ string, _ *WriterOption) (ExternalFileWriter, error) {
 	return &noopWriter{}, nil
 }
 
diff --git a/br/pkg/storage/s3.go b/br/pkg/storage/s3.go
@@ -12,6 +12,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	alicred "github.com/aliyun/alibaba-cloud-sdk-go/sdk/auth/credentials"
@@ -912,11 +913,59 @@ func (rs *S3Storage) CreateUploader(ctx context.Context, name string) (ExternalF
 	}, nil
 }
 
-// Create creates multi upload request.
-func (rs *S3Storage) Create(ctx context.Context, name string) (ExternalFileWriter, error) {
-	uploader, err := rs.CreateUploader(ctx, name)
+type s3ObjectWriter struct {
+	wd  *io.PipeWriter
+	wg  *sync.WaitGroup
+	err error
+}
+
+// Write implement the io.Writer interface.
+func (s *s3ObjectWriter) Write(_ context.Context, p []byte) (int, error) {
+	return s.wd.Write(p)
+}
+
+// Close implement the io.Closer interface.
+func (s *s3ObjectWriter) Close(_ context.Context) error {
+	err := s.wd.Close()
 	if err != nil {
-		return nil, err
+		return err
+	}
+	s.wg.Wait()
+	return s.err
+}
+
+// Create creates multi upload request.
+func (rs *S3Storage) Create(ctx context.Context, name string, option *WriterOption) (ExternalFileWriter, error) {
+	var uploader ExternalFileWriter
+	var err error
+	if option == nil || option.Concurrency <= 1 {
+		uploader, err = rs.CreateUploader(ctx, name)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		up := s3manager.NewUploaderWithClient(rs.svc, func(u *s3manager.Uploader) {
+			u.Concurrency = option.Concurrency
+			u.BufferProvider = s3manager.NewBufferedReadSeekerWriteToPool(option.Concurrency * 8 * 1024 * 1024)
+		})
+		rd, wd := io.Pipe()
+		upParams := &s3manager.UploadInput{
+			Bucket: aws.String(rs.options.Bucket),
+			Key:    aws.String(rs.options.Prefix + name),
+			Body:   rd,
+		}
+		s3Writer := &s3ObjectWriter{wd: wd, wg: &sync.WaitGroup{}}
+		s3Writer.wg.Add(1)
+		go func() {
+			_, err := up.UploadWithContext(ctx, upParams)
+			err1 := rd.Close()
+			if err != nil {
+				log.Warn("upload to s3 failed", zap.String("filename", name), zap.Error(err), zap.Error(err1))
+			}
+			s3Writer.err = err
+			s3Writer.wg.Done()
+		}()
+		uploader = s3Writer
 	}
 	uploaderWriter := newBufferedWriter(uploader, hardcodedS3ChunkSize, NoCompression)
 	return uploaderWriter, nil
diff --git a/br/pkg/storage/storage.go b/br/pkg/storage/storage.go
@@ -78,6 +78,10 @@ type Writer interface {
 	Close(ctx context.Context) error
 }
 
+type WriterOption struct {
+	Concurrency int
+}
+
 // ExternalStorage represents a kind of file system storage.
 type ExternalStorage interface {
 	// WriteFile writes a complete file to storage, similar to os.WriteFile, but WriteFile should be atomic
@@ -101,8 +105,8 @@ type ExternalStorage interface {
 	// URI returns the base path as a URI
 	URI() string
 
-	// Create opens a file writer by path. path is relative path to storage base path
-	Create(ctx context.Context, path string) (ExternalFileWriter, error)
+	// Create opens a file writer by path. path is relative path to storage base path. Currently only s3 implemented WriterOption
+	Create(ctx context.Context, path string, option *WriterOption) (ExternalFileWriter, error)
 	// Rename file name from oldFileName to newFileName
 	Rename(ctx context.Context, oldFileName, newFileName string) error
 }
diff --git a/br/pkg/storage/writer_test.go b/br/pkg/storage/writer_test.go
@@ -29,7 +29,7 @@ func TestExternalFileWriter(t *testing.T) {
 		storage, err := Create(ctx, backend, true)
 		require.NoError(t, err)
 		fileName := strings.ReplaceAll(test.name, " ", "-") + ".txt"
-		writer, err := storage.Create(ctx, fileName)
+		writer, err := storage.Create(ctx, fileName, nil)
 		require.NoError(t, err)
 		for _, str := range test.content {
 			p := []byte(str)
@@ -105,7 +105,7 @@ func TestCompressReaderWriter(t *testing.T) {
 		storage = WithCompression(storage, test.compressType)
 		suffix := createSuffixString(test.compressType)
 		fileName := strings.ReplaceAll(test.name, " ", "-") + suffix
-		writer, err := storage.Create(ctx, fileName)
+		writer, err := storage.Create(ctx, fileName, nil)
 		require.NoError(t, err)
 		for _, str := range test.content {
 			p := []byte(str)
diff --git a/br/pkg/task/stream_test.go b/br/pkg/task/stream_test.go
@@ -155,7 +155,7 @@ func fakeCheckpointFiles(
 		filename := fmt.Sprintf("%v.ts", info.storeID)
 		buff := make([]byte, 8)
 		binary.LittleEndian.PutUint64(buff, info.global_checkpoint)
-		if _, err := s.Create(ctx, filename); err != nil {
+		if _, err := s.Create(ctx, filename, nil); err != nil {
 			return errors.Trace(err)
 		}
 		if err := s.WriteFile(ctx, filename, buff); err != nil {
diff --git a/dumpling/export/writer_util.go b/dumpling/export/writer_util.go
@@ -454,7 +454,7 @@ func writeBytes(tctx *tcontext.Context, writer storage.ExternalFileWriter, p []b
 func buildFileWriter(tctx *tcontext.Context, s storage.ExternalStorage, fileName string, compressType storage.CompressType) (storage.ExternalFileWriter, func(ctx context.Context) error, error) {
 	fileName += compressFileSuffix(compressType)
 	fullPath := s.URI() + "/" + fileName
-	writer, err := storage.WithCompression(s, compressType).Create(tctx, fileName)
+	writer, err := storage.WithCompression(s, compressType).Create(tctx, fileName, nil)
 	if err != nil {
 		tctx.L().Warn("fail to open file",
 			zap.String("path", fullPath),
@@ -487,7 +487,7 @@ func buildInterceptFileWriter(pCtx *tcontext.Context, s storage.ExternalStorage,
 	initRoutine := func() error {
 		// use separated context pCtx here to make sure context used in ExternalFile won't be canceled before close,
 		// which will cause a context canceled error when closing gcs's Writer
-		w, err := storage.WithCompression(s, compressType).Create(pCtx, fileName)
+		w, err := storage.WithCompression(s, compressType).Create(pCtx, fileName, nil)
 		if err != nil {
 			pCtx.L().Warn("fail to open file",
 				zap.String("path", fullPath),

Original file line number	Diff line number	Diff line change
`@@ -415,7 +415,7 @@ func (s *AzureBlobStorage) URI() string {`
`415`	`415`	`}`
`416`	`416`
`417`	`417`	`// Create implements the StorageWriter interface.`
`418`		`-func (s *AzureBlobStorage) Create(_ context.Context, name string) (ExternalFileWriter, error) {`
	`418`	`+func (s AzureBlobStorage) Create(_ context.Context, name string, _ WriterOption) (ExternalFileWriter, error) {`
`419`	`419`	`client := s.containerClient.NewBlockBlobClient(s.withPrefix(name))`
`420`	`420`	`uploader := &azblobUploader{`
`421`	`421`	`blobClient: client,`
Original file line number	Diff line number	Diff line change
`@@ -24,15 +24,15 @@ func WithCompression(inner ExternalStorage, compressionType CompressType) Extern`
`24`	`24`	`return &withCompression{ExternalStorage: inner, compressType: compressionType}`
`25`	`25`	`}`
`26`	`26`
`27`		`-func (w *withCompression) Create(ctx context.Context, name string) (ExternalFileWriter, error) {`
	`27`	`+func (w withCompression) Create(ctx context.Context, name string, _ WriterOption) (ExternalFileWriter, error) {`
`28`	`28`	`var (`
`29`	`29`	`writer ExternalFileWriter`
`30`	`30`	`err error`
`31`	`31`	`)`
`32`	`32`	`if s3Storage, ok := w.ExternalStorage.(*S3Storage); ok {`
`33`	`33`	`writer, err = s3Storage.CreateUploader(ctx, name)`
`34`	`34`	`} else {`
`35`		`- writer, err = w.ExternalStorage.Create(ctx, name)`
	`35`	`+ writer, err = w.ExternalStorage.Create(ctx, name, nil)`
`36`	`36`	`}`
`37`	`37`	`if err != nil {`
`38`	`38`	`return nil, errors.Trace(err)`
Original file line number	Diff line number	Diff line change
`@@ -246,7 +246,7 @@ func (s *GCSStorage) URI() string {`
`246`	`246`	`}`
`247`	`247`
`248`	`248`	`// Create implements ExternalStorage interface.`
`249`		`-func (s *GCSStorage) Create(ctx context.Context, name string) (ExternalFileWriter, error) {`
	`249`	`+func (s GCSStorage) Create(ctx context.Context, name string, _ WriterOption) (ExternalFileWriter, error) {`
`250`	`250`	`object := s.objectName(name)`
`251`	`251`	`wc := s.bucket.Object(object).NewWriter(ctx)`
`252`	`252`	`wc.StorageClass = s.gcs.StorageClass`
Original file line number	Diff line number	Diff line change
`@@ -123,7 +123,7 @@ func (s *HDFSStorage) URI() string {`
`123`	`123`	`}`
`124`	`124`
`125`	`125`	`// Create opens a file writer by path. path is relative path to storage base path`
`126`		`-func (*HDFSStorage) Create(_ context.Context, _ string) (ExternalFileWriter, error) {`
	`126`	`+func (HDFSStorage) Create(_ context.Context, _ string, _ WriterOption) (ExternalFileWriter, error) {`
`127`	`127`	`return nil, errors.Annotatef(berrors.ErrUnsupportedOperation, "currently HDFS backend only support rawkv backup")`
`128`	`128`	`}`
`129`	`129`
Original file line number	Diff line number	Diff line change
`@@ -131,7 +131,7 @@ func (l *LocalStorage) Open(_ context.Context, path string) (ExternalFileReader,`
`131`	`131`	`}`
`132`	`132`
`133`	`133`	`// Create implements ExternalStorage interface.`
`134`		`-func (l *LocalStorage) Create(_ context.Context, name string) (ExternalFileWriter, error) {`
	`134`	`+func (l LocalStorage) Create(_ context.Context, name string, _ WriterOption) (ExternalFileWriter, error) {`
`135`	`135`	`file, err := os.Create(filepath.Join(l.base, name))`
`136`	`136`	`if err != nil {`
`137`	`137`	`return nil, errors.Trace(err)`
Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ func (*noopStorage) URI() string {`
`43`	`43`	`}`
`44`	`44`
`45`	`45`	`// Create implements ExternalStorage interface.`
`46`		`-func (*noopStorage) Create(_ context.Context, _ string) (ExternalFileWriter, error) {`
	`46`	`+func (noopStorage) Create(_ context.Context, _ string, _ WriterOption) (ExternalFileWriter, error) {`
`47`	`47`	`return &noopWriter{}, nil`
`48`	`48`	`}`
`49`	`49`