download schema log files concurrently

Signed-off-by: Jianjun Liao <jianjun.liao@outlook.com>

Author: Leavrth

br/pkg/restore/log_client/client.go

@@ -83,6 +83,7 @@ import (
 
 const MetaKVBatchSize = 64 * 1024 * 1024
 const maxSplitKeysOnce = 10240
+const maxReadMetaKVFilesConcurrency uint = 128
 
 // rawKVBatchCount specifies the count of entries that the rawkv client puts into TiKV.
 const rawKVBatchCount = 64
@@ -1121,15 +1122,15 @@ func SeparateAndSortFilesByCF(files []*backuppb.DataFileInfo) ([]*backuppb.DataF
 	// The error of transactions of meta could happen if restore write CF events successfully,
 	// but failed to restore default CF events.
 	for _, f := range files {
-		if f.Cf == consts.WriteCF {
-			filesInWriteCF = append(filesInWriteCF, f)
-			continue
-		}
-		if f.Type == backuppb.FileType_Delete {
-			log.Warn("internal error: detected delete file of meta key, skip it", zap.Any("file", f))
+		if !shouldReadMetaKVFile(f) {
+			if f.Type == backuppb.FileType_Delete {
+				log.Warn("internal error: detected delete file of meta key, skip it", zap.Any("file", f))
+			}
 			continue
 		}
-		if f.Cf == consts.DefaultCF {
+		if f.Cf == consts.WriteCF {
+			filesInWriteCF = append(filesInWriteCF, f)
+		} else {
 			filesInDefaultCF = append(filesInDefaultCF, f)
 		}
 	}
@@ -1140,6 +1141,26 @@ func SeparateAndSortFilesByCF(files []*backuppb.DataFileInfo) ([]*backuppb.DataF
 	return filesInDefaultCF, filesInWriteCF
 }
 
+func shouldReadMetaKVFile(file *backuppb.DataFileInfo) bool {
+	if file.Cf == consts.WriteCF {
+		return true
+	}
+	if file.Type == backuppb.FileType_Delete {
+		return false
+	}
+	return file.Cf == consts.DefaultCF
+}
+
+func countReadableMetaKVFiles(files []*backuppb.DataFileInfo) int {
+	count := 0
+	for _, file := range files {
+		if shouldReadMetaKVFile(file) {
+			count++
+		}
+	}
+	return count
+}
+
 // LoadAndProcessMetaKVFilesInBatch restores meta kv files to TiKV in strict TS order. It does so in batch and after
 // success it triggers an update so every TiDB node can pick up the restored content.
 func LoadAndProcessMetaKVFilesInBatch(
@@ -1290,14 +1311,28 @@ func (rc *LogClient) filterAndSortKvEntriesFromFiles(
 	}
 
 	// read all entries from files.
-	for _, f := range files {
-		es, filteredOutEs, err := rc.ReadFilteredEntriesFromFiles(ctx, f, filterTS)
-		if err != nil {
+	if len(files) > 0 {
+		eg, egCtx := errgroup.WithContext(ctx)
+		workerPool := tidbutil.NewWorkerPool(min(uint(len(files)), maxReadMetaKVFilesConcurrency), "read meta kv files")
+		var entriesLock sync.Mutex
+		for _, f := range files {
+			file := f
+			workerPool.ApplyOnErrorGroup(eg, func() error {
+				es, filteredOutEs, err := rc.ReadFilteredEntriesFromFiles(egCtx, file, filterTS)
+				if err != nil {
+					return errors.Trace(err)
+				}
+
+				entriesLock.Lock()
+				curKvEntries = append(curKvEntries, es...)
+				filteredOutKvEntries = append(filteredOutKvEntries, filteredOutEs...)
+				entriesLock.Unlock()
+				return nil
+			})
+		}
+		if err := eg.Wait(); err != nil {
 			return nil, nil, errors.Trace(err)
 		}
-
-		curKvEntries = append(curKvEntries, es...)
-		filteredOutKvEntries = append(filteredOutKvEntries, filteredOutEs...)
 	}
 
 	// sort these entries.

br/pkg/restore/log_client/client_test.go

@@ -772,6 +772,34 @@ func TestSortMetaKVFiles(t *testing.T) {
 	require.Equal(t, files[2].Path, "f3")
 	require.Equal(t, files[3].Path, "f4")
 	require.Equal(t, files[4].Path, "f5")
+
+	files = []*backuppb.DataFileInfo{
+		{
+			Path: "write-put",
+			Cf:   consts.WriteCF,
+			Type: backuppb.FileType_Put,
+		},
+		{
+			Path: "write-delete",
+			Cf:   consts.WriteCF,
+			Type: backuppb.FileType_Delete,
+		},
+		{
+			Path: "default-put",
+			Cf:   consts.DefaultCF,
+			Type: backuppb.FileType_Put,
+		},
+		{
+			Path: "default-delete",
+			Cf:   consts.DefaultCF,
+			Type: backuppb.FileType_Delete,
+		},
+	}
+	defaultFiles, writeFiles := logclient.SeparateAndSortFilesByCF(files)
+	require.Len(t, defaultFiles, 1)
+	require.Equal(t, "default-put", defaultFiles[0].Path)
+	require.Len(t, writeFiles, 2)
+	require.Equal(t, 3, logclient.TEST_CountReadableMetaKVFiles(files))
 }
 
 func toLogDataFileInfoIter(logIter iter.TryNextor[*backuppb.DataFileInfo]) logclient.LogIter {

br/pkg/restore/log_client/export_test.go

@@ -16,6 +16,7 @@ package logclient
 
 import (
 	"context"
+	"sync/atomic"
 
 	"github.com/pingcap/errors"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
@@ -121,21 +122,48 @@ func TEST_NewLogFileManager(startTS, restoreTS, shiftStartTS uint64, helper stre
 	}
 }
 
+func TEST_CountReadableMetaKVFiles(files []*backuppb.DataFileInfo) int {
+	return countReadableMetaKVFiles(files)
+}
+
 type FakeStreamMetadataHelper struct {
 	streamMetadataHelper
 
-	Data []byte
+	Data      []byte
+	ReadGate  <-chan struct{}
+	active    atomic.Int32
+	maxActive atomic.Int32
+}
+
+func (helper *FakeStreamMetadataHelper) ActiveReadCount() int32 {
+	return helper.active.Load()
+}
+
+func (helper *FakeStreamMetadataHelper) MaxActiveReadCount() int32 {
+	return helper.maxActive.Load()
 }
 
 func (helper *FakeStreamMetadataHelper) ReadFile(
 	ctx context.Context,
 	path string,
 	offset uint64,
 	length uint64,
+	rawLength uint64,
 	compressionType backuppb.CompressionType,
 	storage storeapi.Storage,
 	encryptionInfo *encryptionpb.FileEncryptionInfo,
 ) ([]byte, error) {
+	active := helper.active.Add(1)
+	for {
+		maxActive := helper.maxActive.Load()
+		if active <= maxActive || helper.maxActive.CompareAndSwap(maxActive, active) {
+			break
+		}
+	}
+	defer helper.active.Add(-1)
+	if helper.ReadGate != nil {
+		<-helper.ReadGate
+	}
 	return helper.Data[offset : offset+length], nil
 }
 

br/pkg/restore/log_client/log_file_manager.go

@@ -82,6 +82,7 @@ type streamMetadataHelper interface {
 		path string,
 		offset uint64,
 		length uint64,
+		rawLength uint64,
 		compressionType backuppb.CompressionType,
 		storage storeapi.Storage,
 		encryptionInfo *encryptionpb.FileEncryptionInfo,
@@ -336,7 +337,7 @@ func (lm *LogFileManager) collectDDLFilesAndPrepareCache(
 
 	dataFileInfos := make([]*backuppb.DataFileInfo, 0)
 	for _, g := range fs.Item {
-		lm.helper.InitCacheEntry(g.Path, len(g.FileMetas))
+		lm.helper.InitCacheEntry(g.Path, countReadableMetaKVFiles(g.FileMetas))
 		dataFileInfos = append(dataFileInfos, g.FileMetas...)
 	}
 
@@ -460,7 +461,7 @@ func (lm *LogFileManager) ReadFilteredEntriesFromFiles(
 	kvEntries := make([]*KvEntryWithTS, 0)
 	filteredOutKvEntries := make([]*KvEntryWithTS, 0)
 
-	buff, err := lm.helper.ReadFile(ctx, file.Path, file.RangeOffset, file.RangeLength, file.CompressionType,
+	buff, err := lm.helper.ReadFile(ctx, file.Path, file.RangeOffset, file.RangeLength, file.Length, file.CompressionType,
 		lm.storage, file.FileEncryptionInfo)
 	if err != nil {
 		return nil, nil, errors.Trace(err)

br/pkg/restore/log_client/log_file_manager_test.go

@@ -16,6 +16,7 @@ import (
 	"strings"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/pingcap/errors"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
@@ -612,6 +613,7 @@ func generateKvData() ([]byte, logclient.Log) {
 		Sha256:      sha256[:],
 		RangeOffset: rangeOffset,
 		RangeLength: rangeLength,
+		Length:      rangeLength,
 	}
 }
 
@@ -646,4 +648,26 @@ func TestReadAllEntries(t *testing.T) {
 			encodekvEntryWithTS("mDDL", 65),
 		}, nextKvEntries)
 	}
+	{
+		readGate := make(chan struct{})
+		helper := &logclient.FakeStreamMetadataHelper{Data: data, ReadGate: readGate}
+		fm := logclient.TEST_NewLogFileManager(35, 75, 25, helper)
+		file.Cf = consts.DefaultCF
+		errCh := make(chan error, 4)
+		for range 4 {
+			go func() {
+				_, _, err := fm.ReadFilteredEntriesFromFiles(ctx, file, 50)
+				errCh <- err
+			}()
+		}
+		require.Eventually(t, func() bool {
+			return helper.ActiveReadCount() == 4
+		}, time.Second, 10*time.Millisecond)
+		require.Equal(t, int32(4), helper.MaxActiveReadCount())
+		close(readGate)
+		for range 4 {
+			require.NoError(t, <-errCh)
+		}
+		require.Equal(t, int32(4), helper.MaxActiveReadCount())
+	}
 }

br/pkg/stream/stream_mgr.go

@@ -19,6 +19,7 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"strings"
+	"sync"
 
 	"github.com/klauspost/compress/zstd"
 	"github.com/pingcap/errors"
@@ -156,6 +157,7 @@ func BuildObserveMetaRange() *kv.KeyRange {
 }
 
 type ContentRef struct {
+	mu       sync.Mutex
 	init_ref int
 	ref      int
 	data     []byte
@@ -164,6 +166,7 @@ type ContentRef struct {
 // MetadataHelper make restore/truncate compatible with metadataV1 and metadataV2.
 type MetadataHelper struct {
 	cache             map[string]*ContentRef
+	cacheMu           sync.RWMutex
 	decoder           *zstd.Decoder
 	encryptionManager *encryption.Manager
 }
@@ -194,19 +197,21 @@ func (m *MetadataHelper) InitCacheEntry(path string, ref int) {
 	if ref <= 0 {
 		return
 	}
+	m.cacheMu.Lock()
+	defer m.cacheMu.Unlock()
 	m.cache[path] = &ContentRef{
 		init_ref: ref,
 		ref:      ref,
 		data:     nil,
 	}
 }
 
-func (m *MetadataHelper) decodeCompressedData(data []byte, compressionType backuppb.CompressionType) ([]byte, error) {
+func (m *MetadataHelper) decodeCompressedData(data []byte, rawLength uint64, compressionType backuppb.CompressionType) ([]byte, error) {
 	switch compressionType {
 	case backuppb.CompressionType_UNKNOWN:
 		return data, nil
 	case backuppb.CompressionType_ZSTD:
-		return m.decoder.DecodeAll(data, nil)
+		return m.decoder.DecodeAll(data, make([]byte, 0, rawLength))
 	}
 	return nil, errors.Errorf(
 		"failed to decode compressed data: compression type is unimplemented. type id is %d", compressionType)
@@ -246,12 +251,15 @@ func (m *MetadataHelper) ReadFile(
 	path string,
 	offset uint64,
 	length uint64,
+	rawLength uint64,
 	compressionType backuppb.CompressionType,
 	storage storeapi.Storage,
 	encryptionInfo *encryptionpb.FileEncryptionInfo,
 ) ([]byte, error) {
 	var err error
+	m.cacheMu.RLock()
 	cref, exist := m.cache[path]
+	m.cacheMu.RUnlock()
 	if !exist {
 		// Only files from metaV2 are cached,
 		// so the file should be from metaV1.
@@ -268,29 +276,33 @@ func (m *MetadataHelper) ReadFile(
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
-		return m.decodeCompressedData(decryptedData, compressionType)
+		return m.decodeCompressedData(decryptedData, rawLength, compressionType)
 	}
 
+	cref.mu.Lock()
 	cref.ref -= 1
 
 	if len(cref.data) == 0 {
 		cref.data, err = storage.ReadFile(ctx, path)
 		if err != nil {
+			cref.mu.Unlock()
 			return nil, errors.Trace(err)
 		}
 	}
-	// decrypt if needed
-	decryptedData, err := m.verifyChecksumAndDecryptIfNeeded(ctx, cref.data[offset:offset+length], encryptionInfo)
-	if err != nil {
-		return nil, errors.Trace(err)
-	}
-	buf, err := m.decodeCompressedData(decryptedData, compressionType)
-
+	data := cref.data[offset : offset+length]
 	if cref.ref <= 0 {
 		// need reset reference information.
 		cref.data = nil
 		cref.ref = cref.init_ref
 	}
+	cref.mu.Unlock()
+
+	// decrypt if needed
+	decryptedData, err := m.verifyChecksumAndDecryptIfNeeded(ctx, data, encryptionInfo)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	buf, err := m.decodeCompressedData(decryptedData, rawLength, compressionType)
 
 	return buf, errors.Trace(err)
 }