br/storage: enable async prefetch data

br/pkg/lightning/backend/external/bench_test.go

@@ -37,10 +37,8 @@ func openTestingStorage(t *testing.T) storage.ExternalStorage {
 	if *testingStorageURI == "" {
 		t.Skip("testingStorageURI is not set")
 	}
-	b, err := storage.ParseBackend(*testingStorageURI, nil)
-	intest.Assert(err == nil)
-	s, err := storage.New(context.Background(), b, nil)
-	intest.Assert(err == nil)
+	s, err := storage.NewFromURL(context.Background(), *testingStorageURI, nil)
+	intest.AssertNoError(err)
 	return s
 }
 
@@ -111,7 +109,7 @@ func writePlainFile(s *writeTestSuite) {
 	offset := 0
 	flush := func(w storage.ExternalFileWriter) {
 		n, err := w.Write(ctx, buf[:offset])
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 		intest.Assert(offset == n)
 		offset = 0
 	}
@@ -120,7 +118,7 @@ func writePlainFile(s *writeTestSuite) {
 		s.beforeCreateWriter()
 	}
 	writer, err := s.store.Create(ctx, "test/plain_file", nil)
-	intest.Assert(err == nil)
+	intest.AssertNoError(err)
 	key, val, _ := s.source.next()
 	for key != nil {
 		if offset+len(key)+len(val) > len(buf) {
@@ -135,7 +133,7 @@ func writePlainFile(s *writeTestSuite) {
 		s.beforeWriterClose()
 	}
 	err = writer.Close(ctx)
-	intest.Assert(err == nil)
+	intest.AssertNoError(err)
 	if s.afterWriterClose != nil {
 		s.afterWriterClose()
 	}
@@ -153,14 +151,14 @@ func writeExternalFile(s *writeTestSuite) {
 	key, val, h := s.source.next()
 	for key != nil {
 		err := writer.WriteRow(ctx, key, val, h)
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 		key, val, h = s.source.next()
 	}
 	if s.beforeWriterClose != nil {
 		s.beforeWriterClose()
 	}
 	err := writer.Close(ctx)
-	intest.Assert(err == nil)
+	intest.AssertNoError(err)
 	if s.afterWriterClose != nil {
 		s.afterWriterClose()
 	}
@@ -180,17 +178,17 @@ func TestCompareWriter(t *testing.T) {
 	beforeTest := func() {
 		fileIdx++
 		file, err = os.Create(fmt.Sprintf("cpu-profile-%d.prof", fileIdx))
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 		err = pprof.StartCPUProfile(file)
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 		now = time.Now()
 	}
 	beforeClose := func() {
 		file, err = os.Create(fmt.Sprintf("heap-profile-%d.prof", fileIdx))
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 		// check heap profile to see the memory usage is expected
 		err = pprof.WriteHeapProfile(file)
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 	}
 	afterClose := func() {
 		elapsed = time.Since(now)
@@ -228,16 +226,16 @@ type readTestSuite struct {
 
 func readFileSequential(s *readTestSuite) {
 	ctx := context.Background()
-	files, _, err := GetAllFileNames(ctx, s.store, "evenly_distributed")
-	intest.Assert(err == nil)
+	files, _, err := GetAllFileNames(ctx, s.store, "/evenly_distributed")
+	intest.AssertNoError(err)
 
 	buf := make([]byte, s.memoryLimit)
 	if s.beforeCreateReader != nil {
 		s.beforeCreateReader()
 	}
 	for i, file := range files {
 		reader, err := s.store.Open(ctx, file, nil)
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 		_, err = reader.Read(buf)
 		for err == nil {
 			_, err = reader.Read(buf)
@@ -249,7 +247,7 @@ func readFileSequential(s *readTestSuite) {
 			}
 		}
 		err = reader.Close()
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 	}
 	if s.afterReaderClose != nil {
 		s.afterReaderClose()
@@ -258,8 +256,8 @@ func readFileSequential(s *readTestSuite) {
 
 func readFileConcurrently(s *readTestSuite) {
 	ctx := context.Background()
-	files, _, err := GetAllFileNames(ctx, s.store, "evenly_distributed")
-	intest.Assert(err == nil)
+	files, _, err := GetAllFileNames(ctx, s.store, "/evenly_distributed")
+	intest.AssertNoError(err)
 
 	conc := min(s.concurrency, len(files))
 	var eg errgroup.Group
@@ -273,7 +271,7 @@ func readFileConcurrently(s *readTestSuite) {
 		eg.Go(func() error {
 			buf := make([]byte, s.memoryLimit/conc)
 			reader, err := s.store.Open(ctx, file, nil)
-			intest.Assert(err == nil)
+			intest.AssertNoError(err)
 			_, err = reader.Read(buf)
 			for err == nil {
 				_, err = reader.Read(buf)
@@ -285,12 +283,12 @@ func readFileConcurrently(s *readTestSuite) {
 				}
 			})
 			err = reader.Close()
-			intest.Assert(err == nil)
+			intest.AssertNoError(err)
 			return nil
 		})
 	}
 	err = eg.Wait()
-	intest.Assert(err == nil)
+	intest.AssertNoError(err)
 	if s.afterReaderClose != nil {
 		s.afterReaderClose()
 	}
@@ -303,12 +301,12 @@ func createEvenlyDistributedFiles(
 	store := openTestingStorage(t)
 	ctx := context.Background()
 
-	files, statFiles, err := GetAllFileNames(ctx, store, "evenly_distributed")
-	intest.Assert(err == nil)
+	files, statFiles, err := GetAllFileNames(ctx, store, "/evenly_distributed")
+	intest.AssertNoError(err)
 	err = store.DeleteFiles(ctx, files)
-	intest.Assert(err == nil)
+	intest.AssertNoError(err)
 	err = store.DeleteFiles(ctx, statFiles)
-	intest.Assert(err == nil)
+	intest.AssertNoError(err)
 
 	value := make([]byte, 100)
 	kvCnt := 0
@@ -317,7 +315,7 @@ func createEvenlyDistributedFiles(
 			SetMemorySizeLimit(uint64(float64(fileSize) * 1.1))
 		writer := builder.Build(
 			store,
-			"evenly_distributed",
+			"/evenly_distributed",
 			fmt.Sprintf("%d", i),
 		)
 
@@ -326,21 +324,21 @@ func createEvenlyDistributedFiles(
 		for totalSize < fileSize {
 			key := fmt.Sprintf("key_%09d", keyIdx)
 			err := writer.WriteRow(ctx, []byte(key), value, nil)
-			intest.Assert(err == nil)
+			intest.AssertNoError(err)
 			keyIdx += fileCount
 			totalSize += len(key) + len(value)
 			kvCnt++
 		}
 		err := writer.Close(ctx)
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 	}
 	return store, kvCnt
 }
 
 func readMergeIter(s *readTestSuite) {
 	ctx := context.Background()
-	files, _, err := GetAllFileNames(ctx, s.store, "evenly_distributed")
-	intest.Assert(err == nil)
+	files, _, err := GetAllFileNames(ctx, s.store, "/evenly_distributed")
+	intest.AssertNoError(err)
 
 	if s.beforeCreateReader != nil {
 		s.beforeCreateReader()
@@ -349,7 +347,7 @@ func readMergeIter(s *readTestSuite) {
 	readBufSize := s.memoryLimit / len(files)
 	zeroOffsets := make([]uint64, len(files))
 	iter, err := NewMergeKVIter(ctx, files, zeroOffsets, s.store, readBufSize, false)
-	intest.Assert(err == nil)
+	intest.AssertNoError(err)
 
 	kvCnt := 0
 	for iter.Next() {
@@ -362,7 +360,7 @@ func readMergeIter(s *readTestSuite) {
 	}
 	intest.Assert(kvCnt == s.totalKVCnt)
 	err = iter.Close()
-	intest.Assert(err == nil)
+	intest.AssertNoError(err)
 	if s.afterReaderClose != nil {
 		s.afterReaderClose()
 	}
@@ -383,17 +381,17 @@ func TestCompareReader(t *testing.T) {
 	beforeTest := func() {
 		fileIdx++
 		file, err = os.Create(fmt.Sprintf("cpu-profile-%d.prof", fileIdx))
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 		err = pprof.StartCPUProfile(file)
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 		now = time.Now()
 	}
 	beforeClose := func() {
 		file, err = os.Create(fmt.Sprintf("heap-profile-%d.prof", fileIdx))
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 		// check heap profile to see the memory usage is expected
 		err = pprof.WriteHeapProfile(file)
-		intest.Assert(err == nil)
+		intest.AssertNoError(err)
 	}
 	afterClose := func() {
 		elapsed = time.Since(now)
@@ -409,23 +407,24 @@ func TestCompareReader(t *testing.T) {
 		beforeReaderClose:  beforeClose,
 		afterReaderClose:   afterClose,
 	}
-	readFileSequential(suite)
+
+	readMergeIter(suite)
 	t.Logf(
-		"sequential read speed for %d bytes: %.2f MB/s",
+		"merge iter read speed for %d bytes: %.2f MB/s",
 		fileSize*fileCnt,
 		float64(fileSize*fileCnt)/elapsed.Seconds()/1024/1024,
 	)
 
-	readFileConcurrently(suite)
+	readFileSequential(suite)
 	t.Logf(
-		"concurrent read speed for %d bytes: %.2f MB/s",
+		"sequential read speed for %d bytes: %.2f MB/s",
 		fileSize*fileCnt,
 		float64(fileSize*fileCnt)/elapsed.Seconds()/1024/1024,
 	)
 
-	readMergeIter(suite)
+	readFileConcurrently(suite)
 	t.Logf(
-		"merge iter read speed for %d bytes: %.2f MB/s",
+		"concurrent read speed for %d bytes: %.2f MB/s",
 		fileSize*fileCnt,
 		float64(fileSize*fileCnt)/elapsed.Seconds()/1024/1024,
 	)

br/pkg/lightning/backend/external/byte_reader.go

@@ -68,8 +68,9 @@ func openStoreReaderAndSeek(
 	store storage.ExternalStorage,
 	name string,
 	initFileOffset uint64,
+	prefetchSize int,
 ) (storage.ExternalFileReader, error) {
-	storageReader, err := store.Open(ctx, name, nil)
+	storageReader, err := store.Open(ctx, name, &storage.ReaderOption{PrefetchSize: prefetchSize})
 	if err != nil {
 		return nil, err
 	}

br/pkg/lightning/backend/external/kv_reader.go

@@ -36,11 +36,12 @@ func newKVReader(
 	initFileOffset uint64,
 	bufSize int,
 ) (*kvReader, error) {
-	sr, err := openStoreReaderAndSeek(ctx, store, name, initFileOffset)
+	oneThird := bufSize / 3
+	sr, err := openStoreReaderAndSeek(ctx, store, name, initFileOffset, oneThird*2)
 	if err != nil {
 		return nil, err
 	}
-	br, err := newByteReader(ctx, sr, bufSize)
+	br, err := newByteReader(ctx, sr, oneThird)
 	if err != nil {
 		return nil, err
 	}

br/pkg/lightning/backend/external/stat_reader.go

@@ -26,7 +26,7 @@ type statsReader struct {
 }
 
 func newStatsReader(ctx context.Context, store storage.ExternalStorage, name string, bufSize int) (*statsReader, error) {
-	sr, err := openStoreReaderAndSeek(ctx, store, name, 0)
+	sr, err := openStoreReaderAndSeek(ctx, store, name, 0, 0)
 	if err != nil {
 		return nil, err
 	}

br/pkg/storage/BUILD.bazel

@@ -28,6 +28,7 @@ go_library(
         "//br/pkg/logutil",
         "//pkg/sessionctx/variable",
         "//pkg/util/intest",
+        "//pkg/util/prefetch",
         "@com_github_aliyun_alibaba_cloud_sdk_go//sdk/auth/credentials",
         "@com_github_aliyun_alibaba_cloud_sdk_go//sdk/auth/credentials/providers",
         "@com_github_aws_aws_sdk_go//aws",

br/pkg/storage/s3.go

@@ -34,6 +34,7 @@ import (
 	"github.com/pingcap/log"
 	berrors "github.com/pingcap/tidb/br/pkg/errors"
 	"github.com/pingcap/tidb/br/pkg/logutil"
+	"github.com/pingcap/tidb/pkg/util/prefetch"
 	"github.com/spf13/pflag"
 	"go.uber.org/zap"
 )
@@ -735,6 +736,9 @@ func (rs *S3Storage) Open(ctx context.Context, path string, o *ReaderOption) (Ex
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
+	if o != nil && o.PrefetchSize > 0 {
+		reader = prefetch.NewReader(reader, o.PrefetchSize)
+	}
 	return &s3ObjectReader{
 		storage:   rs,
 		name:      path,
@@ -874,6 +878,9 @@ func (r *s3ObjectReader) Read(p []byte) (n int, err error) {
 	if maxCnt > int64(len(p)) {
 		maxCnt = int64(len(p))
 	}
+	if maxCnt == 0 {
+		return 0, io.EOF
+	}
 	n, err = r.reader.Read(p[:maxCnt])
 	// TODO: maybe we should use !errors.Is(err, io.EOF) here to avoid error lint, but currently, pingcap/errors
 	// doesn't implement this method yet.

br/pkg/storage/storage.go

@@ -97,6 +97,8 @@ type ReaderOption struct {
 	StartOffset *int64
 	// EndOffset is exclusive. And it's incompatible with Seek.
 	EndOffset *int64
+	// PrefetchSize will switch to NewPrefetchReader if value is positive.
+	PrefetchSize int
 }
 
 // ExternalStorage represents a kind of file system storage.
@@ -195,6 +197,25 @@ func NewWithDefaultOpt(ctx context.Context, backend *backuppb.StorageBackend) (E
 	return New(ctx, backend, &opts)
 }
 
+// NewFromURL creates an ExternalStorage from URL.
+func NewFromURL(ctx context.Context, uri string, opts *ExternalStorageOptions) (ExternalStorage, error) {
+	if len(uri) == 0 {
+		return nil, errors.Annotate(berrors.ErrStorageInvalidConfig, "empty store is not allowed")
+	}
+	u, err := ParseRawURL(uri)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	if u.Scheme == "memstore" {
+		return NewMemStorage(), nil
+	}
+	b, err := parseBackend(u, uri, nil)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	return New(ctx, b, opts)
+}
+
 // New creates an ExternalStorage with options.
 func New(ctx context.Context, backend *backuppb.StorageBackend, opts *ExternalStorageOptions) (ExternalStorage, error) {
 	if opts == nil {

br/pkg/storage/storage_test.go

@@ -3,6 +3,7 @@
 package storage_test
 
 import (
+	"context"
 	"net/http"
 	"testing"
 
@@ -24,3 +25,10 @@ func TestDefaultHttpClient(t *testing.T) {
 	require.Equal(t, int(concurrency), transport.MaxIdleConnsPerHost)
 	require.Equal(t, int(concurrency), transport.MaxIdleConns)
 }
+
+func TestNewMemStorage(t *testing.T) {
+	url := "memstore://"
+	s, err := storage.NewFromURL(context.Background(), url, nil)
+	require.NoError(t, err)
+	require.IsType(t, (*storage.MemStorage)(nil), s)
+}