br: backup/restore statistic data in the partition dimension

DEPS.bzl

@@ -5849,13 +5849,13 @@ def go_deps():
         name = "com_github_pingcap_kvproto",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/pingcap/kvproto",
-        sha256 = "5cf7fd02dc19b5a84b74fb0f5b066e147be92a964b6a03604436e9ee9a2ce752",
-        strip_prefix = "github.com/pingcap/[email protected]20231222062942-c0c73f41d0b2",
+        sha256 = "69c0e5916c604b358dc854873cb855e053c19da1cab0163fa48797eb119537f9",
+        strip_prefix = "github.com/pingcap/[email protected]20231226064240-4f28b82c7860",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20231222062942-c0c73f41d0b2.zip",
-            "http://ats.apps.svc/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20231222062942-c0c73f41d0b2.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20231222062942-c0c73f41d0b2.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20231222062942-c0c73f41d0b2.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20231226064240-4f28b82c7860.zip",
+            "http://ats.apps.svc/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20231226064240-4f28b82c7860.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20231226064240-4f28b82c7860.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/kvproto/com_github_pingcap_kvproto-v0.0.0-20231226064240-4f28b82c7860.zip",
         ],
     )
     go_repository(

br/cmd/br/debug.go

@@ -80,7 +80,7 @@ func newCheckSumCommand() *cobra.Command {
 			}
 
 			reader := metautil.NewMetaReader(backupMeta, s, &cfg.CipherInfo)
-			dbs, err := utils.LoadBackupTables(ctx, reader)
+			dbs, err := metautil.LoadBackupTables(ctx, reader)
 			if err != nil {
 				return errors.Trace(err)
 			}
@@ -173,7 +173,7 @@ func newBackupMetaValidateCommand() *cobra.Command {
 				return errors.Trace(err)
 			}
 			reader := metautil.NewMetaReader(backupMeta, s, &cfg.CipherInfo)
-			dbs, err := utils.LoadBackupTables(ctx, reader)
+			dbs, err := metautil.LoadBackupTables(ctx, reader)
 			if err != nil {
 				log.Error("load tables failed", zap.Error(err))
 				return errors.Trace(err)

br/pkg/backup/schema.go

@@ -40,6 +40,7 @@ type schemaInfo struct {
 	totalKvs   uint64
 	totalBytes uint64
 	stats      *util.JSONTable
+	statsIndex []*backuppb.StatsFileIndex
 }
 
 type iterFuncTp func(kv.Storage, func(*model.DBInfo, *model.TableInfo)) error
@@ -147,8 +148,10 @@ func (ss *Schemas) BackupSchemas(
 					}
 				}
 				if statsHandle != nil {
-					if err := schema.dumpStatsToJSON(statsHandle, backupTS); err != nil {
+					statsWriter := metaWriter.NewStatsWriter()
+					if err := schema.dumpStatsToJSON(ctx, statsWriter, statsHandle, backupTS); err != nil {
 						logger.Error("dump table stats failed", logutil.ShortError(err))
+						return errors.Trace(err)
 					}
 				}
 			}
@@ -209,15 +212,19 @@ func (s *schemaInfo) calculateChecksum(
 	return nil
 }
 
-func (s *schemaInfo) dumpStatsToJSON(statsHandle *handle.Handle, backupTS uint64) error {
+func (s *schemaInfo) dumpStatsToJSON(ctx context.Context, statsWriter *metautil.StatsWriter, statsHandle *handle.Handle, backupTS uint64) error {
 	log.Info("dump stats to json", zap.Stringer("db", s.dbInfo.Name), zap.Stringer("table", s.tableInfo.Name))
-	jsonTable, err := statsHandle.DumpStatsToJSONBySnapshot(
-		s.dbInfo.Name.String(), s.tableInfo, backupTS, true)
-	if err != nil {
+	if err := statsHandle.PersistStatsBySnapshot(
+		ctx, s.dbInfo.Name.String(), s.tableInfo, backupTS, statsWriter.BackupStats,
+	); err != nil {
 		return errors.Trace(err)
 	}
 
-	s.stats = jsonTable
+	statsFileIndexes, err := statsWriter.BackupStatsDone(ctx)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	s.statsIndex = statsFileIndexes
 	return nil
 }
 
@@ -249,5 +256,6 @@ func (s *schemaInfo) encodeToSchema() (*backuppb.Schema, error) {
 		TotalKvs:   s.totalKvs,
 		TotalBytes: s.totalBytes,
 		Stats:      statsBytes,
+		StatsIndex: s.statsIndex,
 	}, nil
 }

br/pkg/backup/schema_test.go

@@ -243,7 +243,7 @@ func TestBuildBackupRangeAndSchemaWithBrokenStats(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, schemas, 1)
 	// the stats should be empty, but other than that everything should be backed up.
-	require.Nil(t, schemas[0].Stats)
+	require.Nil(t, schemas[0].StatsFileIndexes)
 	require.NotZerof(t, schemas[0].Crc64Xor, "%v", schemas[0])
 	require.NotZerof(t, schemas[0].TotalKvs, "%v", schemas[0])
 	require.NotZerof(t, schemas[0].TotalBytes, "%v", schemas[0])
@@ -270,7 +270,7 @@ func TestBuildBackupRangeAndSchemaWithBrokenStats(t *testing.T) {
 	schemas2 := GetSchemasFromMeta(t, es2)
 	require.Len(t, schemas2, 1)
 	// the stats should now be filled, and other than that the result should be equivalent to the first backup.
-	require.NotNil(t, schemas2[0].Stats)
+	require.True(t, len(schemas2[0].StatsFileIndexes[0].InlineData) > 0 || len(schemas2[0].StatsFileIndexes[0].Name) > 0)
 	require.Equal(t, schemas[0].Crc64Xor, schemas2[0].Crc64Xor)
 	require.Equal(t, schemas[0].TotalKvs, schemas2[0].TotalKvs)
 	require.Equal(t, schemas[0].TotalBytes, schemas2[0].TotalBytes)

br/pkg/metautil/BUILD.bazel

@@ -2,15 +2,22 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "metautil",
-    srcs = ["metafile.go"],
+    srcs = [
+        "load.go",
+        "metafile.go",
+        "statsfile.go",
+    ],
     importpath = "github.com/pingcap/tidb/br/pkg/metautil",
     visibility = ["//visibility:public"],
     deps = [
         "//br/pkg/errors",
         "//br/pkg/logutil",
         "//br/pkg/storage",
         "//br/pkg/summary",
+        "//br/pkg/utils",
         "//pkg/parser/model",
+        "//pkg/statistics/handle",
+        "//pkg/statistics/handle/types",
         "//pkg/statistics/handle/util",
         "//pkg/tablecodec",
         "//pkg/util/encrypt",
@@ -21,6 +28,7 @@ go_library(
         "@com_github_pingcap_kvproto//pkg/brpb",
         "@com_github_pingcap_kvproto//pkg/encryptionpb",
         "@com_github_pingcap_log//:log",
+        "@org_golang_x_sync//errgroup",
         "@org_uber_go_zap//:zap",
     ],
 )
@@ -29,18 +37,29 @@ go_test(
     name = "metautil_test",
     timeout = "short",
     srcs = [
+        "load_test.go",
         "main_test.go",
         "metafile_test.go",
+        "statsfile_test.go",
     ],
     embed = [":metautil"],
     flaky = True,
-    shard_count = 6,
+    shard_count = 9,
     deps = [
         "//br/pkg/mock/storage",
+        "//br/pkg/storage",
+        "//br/pkg/utils",
+        "//pkg/parser/model",
+        "//pkg/statistics/handle/types",
+        "//pkg/statistics/handle/util",
+        "//pkg/tablecodec",
         "//pkg/testkit/testsetup",
+        "@com_github_golang_protobuf//proto",
         "@com_github_pingcap_kvproto//pkg/brpb",
         "@com_github_pingcap_kvproto//pkg/encryptionpb",
+        "@com_github_pingcap_tipb//go-tipb",
         "@com_github_stretchr_testify//require",
+        "@org_golang_x_sync//errgroup",
         "@org_uber_go_goleak//:goleak",
         "@org_uber_go_mock//gomock",
     ],

br/pkg/metautil/load.go

@@ -0,0 +1,75 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package metautil
+
+import (
+	"context"
+
+	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/pkg/parser/model"
+)
+
+// Database wraps the schema and tables of a database.
+type Database struct {
+	Info   *model.DBInfo
+	Tables []*Table
+}
+
+// GetTable returns a table of the database by name.
+func (db *Database) GetTable(name string) *Table {
+	for _, table := range db.Tables {
+		if table.Info.Name.String() == name {
+			return table
+		}
+	}
+	return nil
+}
+
+// LoadBackupTables loads schemas from BackupMeta.
+func LoadBackupTables(ctx context.Context, reader *MetaReader) (map[string]*Database, error) {
+	ch := make(chan *Table)
+	errCh := make(chan error)
+	go func() {
+		if err := reader.ReadSchemasFiles(ctx, ch); err != nil {
+			errCh <- errors.Trace(err)
+		}
+		close(ch)
+	}()
+
+	databases := make(map[string]*Database)
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case err := <-errCh:
+			return nil, errors.Trace(err)
+		case table, ok := <-ch:
+			if !ok {
+				close(errCh)
+				return databases, nil
+			}
+			dbName := table.DB.Name.String()
+			db, ok := databases[dbName]
+			if !ok {
+				db = &Database{
+					Info:   table.DB,
+					Tables: make([]*Table, 0),
+				}
+				databases[dbName] = db
+			}
+			db.Tables = append(db.Tables, table)
+		}
+	}
+}

br/pkg/utils/schema_test.go → br/pkg/metautil/load_test.go

@@ -1,6 +1,17 @@
-// Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
-
-package utils
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package metautil
 
 import (
 	"context"
@@ -11,7 +22,6 @@ import (
 	"github.com/golang/protobuf/proto"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/kvproto/pkg/encryptionpb"
-	"github.com/pingcap/tidb/br/pkg/metautil"
 	"github.com/pingcap/tidb/br/pkg/storage"
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/statistics/handle/util"
@@ -84,12 +94,12 @@ func TestLoadBackupMeta(t *testing.T) {
 	require.NoError(t, err)
 
 	ctx := context.Background()
-	err = store.WriteFile(ctx, metautil.MetaFile, data)
+	err = store.WriteFile(ctx, MetaFile, data)
 	require.NoError(t, err)
 
 	dbs, err := LoadBackupTables(
 		ctx,
-		metautil.NewMetaReader(
+		NewMetaReader(
 			meta,
 			store,
 			&backuppb.CipherInfo{
@@ -179,12 +189,12 @@ func TestLoadBackupMetaPartionTable(t *testing.T) {
 	require.NoError(t, err)
 
 	ctx := context.Background()
-	err = store.WriteFile(ctx, metautil.MetaFile, data)
+	err = store.WriteFile(ctx, MetaFile, data)
 	require.NoError(t, err)
 
 	dbs, err := LoadBackupTables(
 		ctx,
-		metautil.NewMetaReader(
+		NewMetaReader(
 			meta,
 			store,
 			&backuppb.CipherInfo{
@@ -265,12 +275,12 @@ func BenchmarkLoadBackupMeta64(b *testing.B) {
 		require.NoError(b, err)
 
 		ctx := context.Background()
-		err = store.WriteFile(ctx, metautil.MetaFile, data)
+		err = store.WriteFile(ctx, MetaFile, data)
 		require.NoError(b, err)
 
 		dbs, err := LoadBackupTables(
 			ctx,
-			metautil.NewMetaReader(
+			NewMetaReader(
 				meta,
 				store,
 				&backuppb.CipherInfo{
@@ -297,12 +307,12 @@ func BenchmarkLoadBackupMeta1024(b *testing.B) {
 		require.NoError(b, err)
 
 		ctx := context.Background()
-		err = store.WriteFile(ctx, metautil.MetaFile, data)
+		err = store.WriteFile(ctx, MetaFile, data)
 		require.NoError(b, err)
 
 		dbs, err := LoadBackupTables(
 			ctx,
-			metautil.NewMetaReader(
+			NewMetaReader(
 				meta,
 				store,
 				&backuppb.CipherInfo{
@@ -329,12 +339,12 @@ func BenchmarkLoadBackupMeta10240(b *testing.B) {
 		require.NoError(b, err)
 
 		ctx := context.Background()
-		err = store.WriteFile(ctx, metautil.MetaFile, data)
+		err = store.WriteFile(ctx, MetaFile, data)
 		require.NoError(b, err)
 
 		dbs, err := LoadBackupTables(
 			ctx,
-			metautil.NewMetaReader(
+			NewMetaReader(
 				meta,
 				store,
 				&backuppb.CipherInfo{