some more fix

Signed-off-by: Wenqi Mou <[email protected]>

Author: Tristan1900

br/pkg/registry/registration.go

@@ -148,6 +148,11 @@ type RegistrationInfo struct {
 	Cmd string
 }
 
+type RegistrationInfoWithID struct {
+	RegistrationInfo
+	restoreID uint64
+}
+
 // Registry manages registrations of restore tasks
 type Registry struct {
 	se glue.Session
@@ -368,13 +373,13 @@ func (r *Registry) PauseTask(ctx context.Context, restoreID uint64) error {
 }
 
 // GetRegistrationsByMaxID returns all registrations with IDs smaller than maxID
-func (r *Registry) GetRegistrationsByMaxID(ctx context.Context, maxID uint64) ([]RegistrationInfo, error) {
+func (r *Registry) GetRegistrationsByMaxID(ctx context.Context, maxID uint64) ([]RegistrationInfoWithID, error) {
 	if err := r.createTableIfNotExist(ctx); err != nil {
 		return nil, errors.Trace(err)
 	}
 
 	selectSQL := fmt.Sprintf(selectRegistrationsByMaxIDSQLTemplate, RegistrationDBName, RegistrationTableName)
-	registrations := make([]RegistrationInfo, 0)
+	registrations := make([]RegistrationInfoWithID, 0)
 
 	execCtx := r.se.GetSessionCtx().GetRestrictedSQLExecutor()
 	rows, _, errSQL := execCtx.ExecRestrictedSQL(
@@ -408,7 +413,12 @@ func (r *Registry) GetRegistrationsByMaxID(ctx context.Context, maxID uint64) ([
 			Cmd:               cmd,
 		}
 
-		registrations = append(registrations, info)
+		infoWithID := RegistrationInfoWithID{
+			info,
+			row.GetUint64(0),
+		}
+
+		registrations = append(registrations, infoWithID)
 	}
 
 	return registrations, nil
@@ -460,16 +470,17 @@ func (r *Registry) CheckTablesWithRegisteredTasks(
 func (r *Registry) checkForTableConflicts(
 	tracker *utils.PiTRIdTracker,
 	tables []*metautil.Table,
-	regInfo RegistrationInfo,
+	regInfo RegistrationInfoWithID,
 	f filter.Filter,
-	restoreID uint64,
+	curRestoreID uint64,
 ) error {
 	// function to handle conflict when found
 	handleConflict := func(dbName, tableName string) error {
 		log.Warn("table already covered by another restore task",
+			zap.Uint64("existing_restore_id", regInfo.restoreID),
+			zap.Uint64("current_restore_id", curRestoreID),
 			zap.String("database", dbName),
 			zap.String("table", tableName),
-			zap.Uint64("current_restore_id", restoreID),
 			zap.Strings("filter_strings", regInfo.FilterStrings),
 			zap.Uint64("start_ts", regInfo.StartTS),
 			zap.Uint64("restored_ts", regInfo.RestoredTS),
@@ -478,8 +489,8 @@ func (r *Registry) checkForTableConflicts(
 			zap.String("cmd", regInfo.Cmd))
 		return errors.Annotatef(berrors.ErrTablesAlreadyExisted,
 			"table %s.%s cannot be restored by current task with ID %d "+
-				"because it is already being restored by task (time range: %d->%d, cmd: %s)",
-			dbName, tableName, restoreID, regInfo.StartTS, regInfo.RestoredTS, regInfo.Cmd)
+				"because it is already being restored by task (restoreId: %d, time range: %d->%d, cmd: %s)",
+			dbName, tableName, curRestoreID, regInfo.restoreID, regInfo.StartTS, regInfo.RestoredTS, regInfo.Cmd)
 	}
 
 	// Use PiTRTableTracker if available for PiTR task

br/pkg/task/restore.go

@@ -842,20 +842,18 @@ func RunRestore(c context.Context, g glue.Glue, cmdName string, cfg *RestoreConf
 		// if checkpoint is not persisted, let's just unregister the task since we don't need it
 		if hasCheckpointPersisted(c, cfg) {
 			log.Info("pausing restore task from registry",
-				zap.Uint64("restoreId", cfg.RestoreID), zap.Error(err))
+				zap.Uint64("restoreId", cfg.RestoreID), zap.Error(restoreError))
 			if err := restoreRegistry.PauseTask(c, cfg.RestoreID); err != nil {
 				log.Error("failed to pause restore task from registry",
 					zap.Uint64("restoreId", cfg.RestoreID), zap.Error(err))
 			}
 		} else {
 			log.Info("unregistering restore task from registry",
-				zap.Uint64("restoreId", cfg.RestoreID), zap.Error(err))
+				zap.Uint64("restoreId", cfg.RestoreID), zap.Error(restoreError))
 			if err := restoreRegistry.Unregister(c, cfg.RestoreID); err != nil {
 				log.Error("failed to unregister restore task from registry",
 					zap.Uint64("restoreId", cfg.RestoreID), zap.Error(err))
 			}
-			// clean up checkpoint just in case there are some empty checkpoint db/table created but no data
-			cleanUpCheckpoints(c, cfg, cmdName)
 		}
 
 		return errors.Trace(restoreError)
@@ -874,32 +872,29 @@ func RunRestore(c context.Context, g glue.Glue, cmdName string, cfg *RestoreConf
 
 func cleanUpCheckpoints(ctx context.Context, cfg *RestoreConfig, cmdName string) {
 	if cfg.UseCheckpoint {
-		if IsStreamRestore(cmdName) {
-			log.Info("start to remove checkpoint data for PITR restore")
+		log.Info("start to remove checkpoint data for PITR restore")
+		if cfg.logCheckpointMetaManager != nil {
 			err := cfg.logCheckpointMetaManager.RemoveCheckpointData(ctx)
 			if err != nil {
 				log.Warn("failed to remove checkpoint data for log restore", zap.Error(err))
 			}
-			err = cfg.sstCheckpointMetaManager.RemoveCheckpointData(ctx)
+		}
+		if cfg.sstCheckpointMetaManager != nil {
+			err := cfg.sstCheckpointMetaManager.RemoveCheckpointData(ctx)
 			if err != nil {
 				log.Warn("failed to remove checkpoint data for compacted restore", zap.Error(err))
 			}
-			// Skip removing snapshot checkpoint data if this is a pure log restore
-			// (i.e. restoring only from log backup without a base snapshot backup),
-			// since snapshotCheckpointMetaManager would be nil in that case
-			if cfg.snapshotCheckpointMetaManager != nil {
-				err = cfg.snapshotCheckpointMetaManager.RemoveCheckpointData(ctx)
-				if err != nil {
-					log.Warn("failed to remove checkpoint data for snapshot restore", zap.Error(err))
-				}
-			}
-		} else {
+		}
+		// Skip removing snapshot checkpoint data if this is a pure log restore
+		// (i.e. restoring only from log backup without a base snapshot backup),
+		// since snapshotCheckpointMetaManager would be nil in that case
+		if cfg.snapshotCheckpointMetaManager != nil {
 			err := cfg.snapshotCheckpointMetaManager.RemoveCheckpointData(ctx)
 			if err != nil {
 				log.Warn("failed to remove checkpoint data for snapshot restore", zap.Error(err))
 			}
 		}
-		log.Info("all the checkpoint data removed.")
+		log.Info("all checkpoint data removed.")
 	} else {
 		log.Info("checkpoint not enabled, skip to remove checkpoint data")
 	}

br/tests/br_pitr/run.sh

@@ -117,6 +117,8 @@ latest_log_db=$(run_sql "select table_schema from information_schema.tables wher
 latest_sst_db=$(run_sql "select table_schema from information_schema.tables where table_schema like '__TiDB_BR_Temporary_Custom_SST_Restore_Checkpoint%' order by table_schema desc limit 1;" | tail -n 1 | awk '{print $2}')
 run_sql "DROP DATABASE IF EXISTS \`$latest_log_db\`;"
 run_sql "DROP DATABASE IF EXISTS \`$latest_sst_db\`;"
+run_sql "DROP DATABASE IF EXISTS __TiDB_BR_Temporary_Restore_Registration_DB;"
+
 run_br --pd $PD_ADDR restore point -s "local://$TEST_DIR/$PREFIX/log" --full-backup-storage "local://$TEST_DIR/$PREFIX/full" > $res_file 2>&1 || ( cat $res_file && exit 1 )
 
 check_result

tests/realtikvtest/brietest/brie_test.go

@@ -135,8 +135,6 @@ func TestCancel(t *testing.T) {
 // cleanupRegistry drops the registry tables and database
 func cleanupRegistry(tk *testkit.TestKit) {
 	registryDB := "__TiDB_BR_Temporary_Restore_Registration_DB"
-	registryTable := "restore_registry"
-	tk.MustExec(fmt.Sprintf("DROP TABLE IF EXISTS %s.%s", registryDB, registryTable))
 	tk.MustExec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", registryDB))
 }
 
@@ -207,8 +205,7 @@ func TestExistedTables(t *testing.T) {
 		require.NoError(t, err)
 
 		_, err = session.ResultSetToStringSlice(context.Background(), tk.Session(), res)
-		// due to previous restore didn't succeed with checkpoint enabled
-		require.ErrorContains(t, err, "it is already being restored by task")
+		require.ErrorContains(t, err, "table already exists")
 	}()
 	select {
 	case <-time.After(20 * time.Second):