br: add restoreid to idmap table (#61278)

close #61277

Author: Tristan1900

br/pkg/restore/log_client/client.go

@@ -201,6 +201,7 @@ type LogClient struct {
 	currentTS uint64
 
 	upstreamClusterID uint64
+	restoreID         uint64
 
 	// the query to insert rows into table `gc_delete_range`, lack of ts.
 	deleteRangeQuery          []*stream.PreDelRangeQuery

br/pkg/restore/log_client/export_test.go

@@ -99,6 +99,7 @@ func TEST_NewLogClient(clusterID, startTS, restoreTS, upstreamClusterID uint64,
 		dom:               dom,
 		unsafeSession:     se,
 		upstreamClusterID: upstreamClusterID,
+		restoreID:         0,
 		LogFileManager: &LogFileManager{
 			startTS:   startTS,
 			restoreTS: restoreTS,

br/pkg/restore/log_client/id_map.go

@@ -107,14 +107,14 @@ func (rc *LogClient) saveIDMap2Table(ctx context.Context, dbMaps []*backuppb.Pit
 		return errors.Trace(err)
 	}
 	// clean the dirty id map at first
-	err = rc.unsafeSession.ExecuteInternal(ctx, "DELETE FROM mysql.tidb_pitr_id_map WHERE restored_ts = %? and upstream_cluster_id = %?;", rc.restoreTS, rc.upstreamClusterID)
+	err = rc.unsafeSession.ExecuteInternal(ctx, "DELETE FROM mysql.tidb_pitr_id_map WHERE restored_ts = %? and upstream_cluster_id = %? and restore_id = %?;", rc.restoreTS, rc.upstreamClusterID, rc.restoreID)
 	if err != nil {
 		return errors.Trace(err)
 	}
-	replacePitrIDMapSQL := "REPLACE INTO mysql.tidb_pitr_id_map (restored_ts, upstream_cluster_id, segment_id, id_map) VALUES (%?, %?, %?, %?);"
+	replacePitrIDMapSQL := "REPLACE INTO mysql.tidb_pitr_id_map (restore_id, restored_ts, upstream_cluster_id, segment_id, id_map) VALUES (%?, %?, %?, %?, %?);"
 	for startIdx, segmentId := 0, 0; startIdx < len(data); segmentId += 1 {
 		endIdx := min(startIdx+PITRIdMapBlockSize, len(data))
-		err := rc.unsafeSession.ExecuteInternal(ctx, replacePitrIDMapSQL, rc.restoreTS, rc.upstreamClusterID, segmentId, data[startIdx:endIdx])
+		err := rc.unsafeSession.ExecuteInternal(ctx, replacePitrIDMapSQL, rc.restoreID, rc.restoreTS, rc.upstreamClusterID, segmentId, data[startIdx:endIdx])
 		if err != nil {
 			return errors.Trace(err)
 		}
@@ -173,12 +173,13 @@ func (rc *LogClient) loadSchemasMapFromTable(
 	ctx context.Context,
 	restoredTS uint64,
 ) ([]*backuppb.PitrDBMap, error) {
-	getPitrIDMapSQL := "SELECT segment_id, id_map FROM mysql.tidb_pitr_id_map WHERE restored_ts = %? and upstream_cluster_id = %? ORDER BY segment_id;"
+	getPitrIDMapSQL := "SELECT segment_id, id_map FROM mysql.tidb_pitr_id_map WHERE restore_id = %? and restored_ts = %? and upstream_cluster_id = %? ORDER BY segment_id;"
 	execCtx := rc.unsafeSession.GetSessionCtx().GetRestrictedSQLExecutor()
 	rows, _, errSQL := execCtx.ExecRestrictedSQL(
 		kv.WithInternalSourceType(ctx, kv.InternalTxnBR),
 		nil,
 		getPitrIDMapSQL,
+		rc.restoreID,
 		restoredTS,
 		rc.upstreamClusterID,
 	)

br/pkg/restore/snap_client/systable_restore_test.go

@@ -121,7 +121,7 @@ func TestCheckSysTableCompatibility(t *testing.T) {
 //
 // The above variables are in the file br/pkg/restore/systable_restore.go
 func TestMonitorTheSystemTableIncremental(t *testing.T) {
-	require.Equal(t, int64(247), session.CurrentBootstrapVersion)
+	require.Equal(t, int64(248), session.CurrentBootstrapVersion)
 }
 
 func TestIsStatsTemporaryTable(t *testing.T) {

pkg/session/BUILD.bazel

@@ -133,7 +133,6 @@ go_library(
         "@com_github_tikv_client_go_v2//txnkv/transaction",
         "@com_github_tikv_client_go_v2//util",
         "@io_etcd_go_etcd_client_v3//:client",
-        "@io_etcd_go_etcd_client_v3//concurrency",
         "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
         "@org_uber_go_zap//zapcore",

pkg/session/bootstrap.go

@@ -34,7 +34,6 @@ import (
 	"github.com/pingcap/tidb/pkg/bindinfo"
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/domain"
-	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/infoschema"
 	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
@@ -62,7 +61,6 @@ import (
 	"github.com/pingcap/tidb/pkg/util/sqlescape"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
 	"github.com/pingcap/tidb/pkg/util/timeutil"
-	"go.etcd.io/etcd/client/v3/concurrency"
 	"go.uber.org/zap"
 )
 
@@ -721,13 +719,16 @@ const (
 		KEY (status));`
 
 	// CreatePITRIDMap is a table that records the id map from upstream to downstream for PITR.
+	// set restore id default to 0 to make it compatible for old BR tool to restore to a new TiDB, such case should be
+	// rare though.
 	CreatePITRIDMap = `CREATE TABLE IF NOT EXISTS mysql.tidb_pitr_id_map (
+		restore_id BIGINT NOT NULL DEFAULT 0,
 		restored_ts BIGINT NOT NULL,
 		upstream_cluster_id BIGINT NOT NULL,
 		segment_id BIGINT NOT NULL,
 		id_map BLOB(524288) NOT NULL,
 		update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-		PRIMARY KEY (restored_ts, upstream_cluster_id, segment_id));`
+		PRIMARY KEY (restore_id, restored_ts, upstream_cluster_id, segment_id));`
 
 	// DropMySQLIndexUsageTable removes the table `mysql.schema_index_usage`
 	DropMySQLIndexUsageTable = "DROP TABLE IF EXISTS mysql.schema_index_usage"
@@ -1284,11 +1285,15 @@ const (
 	// version 247
 	// Add last_stats_histograms_version to mysql.stats_meta.
 	version247 = 247
+
+	// version 248
+	// Update mysql.tidb_pitr_id_map to add restore_id as a primary key field
+	version248 = 248
 )
 
 // currentBootstrapVersion is defined as a variable, so we can modify its value for testing.
 // please make sure this is the largest version
-var currentBootstrapVersion int64 = version247
+var currentBootstrapVersion int64 = version248
 
 // DDL owner key's expired time is ManagerSessionTTL seconds, we should wait the time and give more time to have a chance to finish it.
 var internalSQLTimeout = owner.ManagerSessionTTL + 15
@@ -1471,6 +1476,7 @@ var (
 		upgradeToVer245,
 		upgradeToVer246,
 		upgradeToVer247,
+		upgradeToVer248,
 	}
 )
 
@@ -1626,31 +1632,6 @@ func upgrade(s sessiontypes.Session) {
 	}
 }
 
-// checkOwnerVersion is used to wait the DDL owner to be elected in the cluster and check it is the same version as this TiDB.
-func checkOwnerVersion(ctx context.Context, dom *domain.Domain) (bool, error) {
-	ticker := time.NewTicker(100 * time.Millisecond)
-	defer ticker.Stop()
-	logutil.BgLogger().Info("Waiting for the DDL owner to be elected in the cluster")
-	for {
-		select {
-		case <-ctx.Done():
-			return false, ctx.Err()
-		case <-ticker.C:
-			ownerID, err := dom.DDL().OwnerManager().GetOwnerID(ctx)
-			if err == concurrency.ErrElectionNoLeader {
-				continue
-			}
-			info, err := infosync.GetAllServerInfo(ctx)
-			if err != nil {
-				return false, err
-			}
-			if s, ok := info[ownerID]; ok {
-				return s.Version == mysql.ServerVersion, nil
-			}
-		}
-	}
-}
-
 // upgradeToVer2 updates to version 2.
 func upgradeToVer2(s sessiontypes.Session, ver int64) {
 	if ver >= version2 {
@@ -3477,6 +3458,15 @@ func upgradeToVer247(s sessiontypes.Session, ver int64) {
 	doReentrantDDL(s, "ALTER TABLE mysql.stats_meta ADD COLUMN last_stats_histograms_version bigint unsigned DEFAULT NULL", infoschema.ErrColumnExists)
 }
 
+func upgradeToVer248(s sessiontypes.Session, ver int64) {
+	if ver >= version248 {
+		return
+	}
+	doReentrantDDL(s, "ALTER TABLE mysql.tidb_pitr_id_map ADD COLUMN restore_id BIGINT NOT NULL DEFAULT 0", infoschema.ErrColumnExists)
+	doReentrantDDL(s, "ALTER TABLE mysql.tidb_pitr_id_map DROP PRIMARY KEY")
+	doReentrantDDL(s, "ALTER TABLE mysql.tidb_pitr_id_map ADD PRIMARY KEY(restore_id, restored_ts, upstream_cluster_id, segment_id)")
+}
+
 // initGlobalVariableIfNotExists initialize a global variable with specific val if it does not exist.
 func initGlobalVariableIfNotExists(s sessiontypes.Session, name string, val any) {
 	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnBootstrap)