This is an automated cherry-pick of pingcap#61610

Signed-off-by: ti-chi-bot <[email protected]>

Author: Tristan1900

br/pkg/registry/BUILD.bazel

@@ -15,6 +15,7 @@ go_library(
         "//br/pkg/utils",
         "//pkg/domain",
         "//pkg/kv",
+        "//pkg/util/chunk",
         "//pkg/util/sqlexec",
         "//pkg/util/table-filter",
         "@com_github_pingcap_errors//:errors",

br/pkg/registry/registration.go

@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/br/pkg/utils"
 	"github.com/pingcap/tidb/pkg/domain"
 	"github.com/pingcap/tidb/pkg/kv"
+	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
 	filter "github.com/pingcap/tidb/pkg/util/table-filter"
 	"go.uber.org/zap"
@@ -43,6 +44,9 @@ const (
 	// Using ASCII Unit Separator (US) character which never appears in SQL identifiers or expressions.
 	FilterSeparator = "\x1F"
 
+	// StaleTaskThresholdMinutes is the threshold in minutes to consider a running task as potentially stale
+	StaleTaskThresholdMinutes = 5
+
 	// lookupRegistrationSQLTemplate is the SQL template for looking up a registration by its parameters
 	lookupRegistrationSQLTemplate = `
 		SELECT id, status FROM %s.%s
@@ -83,6 +87,40 @@ const (
 		(filter_strings, filter_hash, start_ts, restored_ts, upstream_cluster_id,
 		 with_sys_table, status, cmd, task_start_time, last_heartbeat_time)
 		VALUES (%%?, MD5(%%?), %%?, %%?, %%?, %%?, 'running', %%?, %%?, %%?)`
+
+	// selectStaleRunningTasksSQLTemplate is the SQL template for finding potentially stale running tasks
+	selectStaleRunningTasksSQLTemplate = `
+		SELECT id, last_heartbeat_time
+		FROM %s.%s
+		WHERE status = 'running'
+		AND last_heartbeat_time < DATE_SUB(NOW(), INTERVAL %%? MINUTE)
+		ORDER BY id ASC`
+
+	// selectTaskHeartbeatSQLTemplate is the SQL template for getting a specific task's heartbeat time
+	selectTaskHeartbeatSQLTemplate = `
+		SELECT last_heartbeat_time
+		FROM %s.%s
+		WHERE id = %%?`
+
+	// selectConflictingTaskSQLTemplate is the SQL template for finding tasks with same parameters except restoredTS
+	selectConflictingTaskSQLTemplate = `
+		SELECT id, restored_ts, status FROM %s.%s
+		WHERE filter_hash = MD5(%%?)
+		AND start_ts = %%?
+		AND upstream_cluster_id = %%?
+		AND with_sys_table = %%?
+		AND cmd = %%?
+		AND restored_ts != %%?
+		AND status IN ('running', 'paused')
+		ORDER BY id DESC
+		LIMIT 1`
+
+	// transitionStaleTaskToPausedSQLTemplate is the SQL template for atomically transitioning a
+	// stale running task to paused
+	transitionStaleTaskToPausedSQLTemplate = `
+		UPDATE %s.%s
+		SET status = 'paused'
+		WHERE id = %%? AND status = 'running' AND last_heartbeat_time = %%?`
 )
 
 // TaskStatus represents the current state of a restore task
@@ -194,7 +232,23 @@ func (r *Registry) executeInTransaction(ctx context.Context, fn func(context.Con
 
 // ResumeOrCreateRegistration first looks for an existing registration with the given parameters.
 // If found and paused, it tries to resume it. Otherwise, it creates a new registration.
-func (r *Registry) ResumeOrCreateRegistration(ctx context.Context, info RegistrationInfo) (uint64, error) {
+// Returns: (taskID, resolvedRestoreTS, error)
+func (r *Registry) ResumeOrCreateRegistration(ctx context.Context, info RegistrationInfo,
+	isRestoredTSUserSpecified bool) (uint64, uint64, error) {
+	// resolve which restoredTS to use, handling auto-detection conflicts
+	resolvedRestoreTS, err := r.resolveRestoreTS(ctx, info, isRestoredTSUserSpecified)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	// update info with resolved restoredTS if different
+	if resolvedRestoreTS != info.RestoredTS {
+		log.Info("using resolved restoredTS from existing task",
+			zap.Uint64("original_restored_ts", info.RestoredTS),
+			zap.Uint64("resolved_restored_ts", resolvedRestoreTS))
+		info.RestoredTS = resolvedRestoreTS
+	}
+
 	filterStrings := strings.Join(info.FilterStrings, FilterSeparator)
 
 	log.Info("attempting to resume or create registration",
@@ -203,11 +257,12 @@ func (r *Registry) ResumeOrCreateRegistration(ctx context.Context, info Registra
 		zap.Uint64("restored_ts", info.RestoredTS),
 		zap.Uint64("upstream_cluster_id", info.UpstreamClusterID),
 		zap.Bool("with_sys_table", info.WithSysTable),
-		zap.String("cmd", info.Cmd))
+		zap.String("cmd", info.Cmd),
+		zap.Bool("is_restored_ts_user_specified", isRestoredTSUserSpecified))
 
 	var taskID uint64
 
-	err := r.executeInTransaction(ctx, func(ctx context.Context, execCtx sqlexec.RestrictedSQLExecutor,
+	err = r.executeInTransaction(ctx, func(ctx context.Context, execCtx sqlexec.RestrictedSQLExecutor,
 		sessionOpts []sqlexec.OptionFuncAlias) error {
 		// first look for an existing task with the same parameters
 		lookupSQL := fmt.Sprintf(lookupRegistrationSQLTemplate, RestoreRegistryDBName, RestoreRegistryTableName)
@@ -292,10 +347,10 @@ func (r *Registry) ResumeOrCreateRegistration(ctx context.Context, info Registra
 	})
 
 	if err != nil {
-		return 0, err
+		return 0, 0, err
 	}
 
-	return taskID, nil
+	return taskID, resolvedRestoreTS, nil
 }
 
 // updateTaskStatusConditional updates a task's status only if its current status matches the expected status
@@ -498,3 +553,252 @@ func (r *Registry) StopHeartbeatManager() {
 		log.Info("stopped heartbeat manager for restore task")
 	}
 }
+
+// resolveRestoreTS determines which restoredTS to use, handling conflicts with existing tasks
+// when restoredTS is not user-specified. Returns: (resolvedRestoreTS, error)
+func (r *Registry) resolveRestoreTS(ctx context.Context,
+	info RegistrationInfo, isRestoredTSUserSpecified bool) (uint64, error) {
+	// if restoredTS is user-specified, use it directly without any conflict resolution
+	if isRestoredTSUserSpecified {
+		log.Info("restoredTS is user-specified, using it directly",
+			zap.Uint64("restored_ts", info.RestoredTS))
+		return info.RestoredTS, nil
+	}
+
+	filterStrings := strings.Join(info.FilterStrings, FilterSeparator)
+
+	// look for tasks with same filter, startTS, cluster, sysTable, cmd but different restoredTS
+	execCtx := r.se.GetSessionCtx().GetRestrictedSQLExecutor()
+	ctx = kv.WithInternalSourceType(ctx, kv.InternalTxnBR)
+
+	checkSQL := fmt.Sprintf(selectConflictingTaskSQLTemplate, RestoreRegistryDBName, RestoreRegistryTableName)
+	rows, _, err := execCtx.ExecRestrictedSQL(ctx, nil, checkSQL,
+		filterStrings, info.StartTS, info.UpstreamClusterID, info.WithSysTable, info.Cmd, info.RestoredTS)
+	if err != nil {
+		return 0, errors.Annotate(err, "failed to check for existing tasks with different restoredTS")
+	}
+
+	// no conflicting task found, use the current restoredTS
+	if len(rows) == 0 {
+		log.Info("no existing tasks found with different restoredTS",
+			zap.Uint64("restored_ts", info.RestoredTS))
+		return info.RestoredTS, nil
+	}
+
+	conflictingTaskID := rows[0].GetUint64(0)
+	existingRestoredTS := rows[0].GetUint64(1)
+	existingStatus := rows[0].GetString(2)
+
+	log.Info("found existing task with different restoredTS",
+		zap.Uint64("existing_task_id", conflictingTaskID),
+		zap.Uint64("existing_restored_ts", existingRestoredTS),
+		zap.String("existing_status", existingStatus),
+		zap.Uint64("current_restored_ts", info.RestoredTS),
+		zap.Strings("filters", info.FilterStrings),
+		zap.Uint64("start_ts", info.StartTS))
+
+	// if existing task is paused, reuse its restoredTS
+	if existingStatus == string(TaskStatusPaused) {
+		log.Info("existing task is paused, reusing its restoredTS",
+			zap.Uint64("existing_task_id", conflictingTaskID),
+			zap.Uint64("existing_restored_ts", existingRestoredTS))
+		return existingRestoredTS, nil
+	}
+
+	// if existing task is running, check if it's stale
+	if existingStatus == string(TaskStatusRunning) {
+		log.Info("existing task is running, checking if it's stale",
+			zap.Uint64("existing_task_id", conflictingTaskID))
+
+		// First, get the current heartbeat time for atomic transition
+		selectHeartbeatSQL := fmt.Sprintf(selectTaskHeartbeatSQLTemplate, RestoreRegistryDBName, RestoreRegistryTableName)
+		heartbeatRows, _, err := execCtx.ExecRestrictedSQL(ctx, nil, selectHeartbeatSQL, conflictingTaskID)
+		if err != nil {
+			log.Warn("failed to get task heartbeat, using current restoredTS",
+				zap.Uint64("task_id", conflictingTaskID),
+				zap.Error(err))
+			return info.RestoredTS, nil
+		}
+
+		if len(heartbeatRows) == 0 {
+			log.Info("task not found during heartbeat check, using current restoredTS",
+				zap.Uint64("task_id", conflictingTaskID))
+			return info.RestoredTS, nil
+		}
+
+		currentHeartbeatTime := heartbeatRows[0].GetTime(0).String()
+
+		isStale, err := r.isTaskStale(ctx, conflictingTaskID)
+		if err != nil {
+			log.Warn("failed to check if task is stale, using current restoredTS",
+				zap.Uint64("task_id", conflictingTaskID),
+				zap.Error(err))
+			return info.RestoredTS, nil
+		}
+
+		if isStale {
+			log.Info("existing running task is stale, attempting to transition to paused",
+				zap.Uint64("existing_task_id", conflictingTaskID),
+				zap.Uint64("existing_restored_ts", existingRestoredTS))
+
+			// atomically transition the stale task to paused state
+			transitioned, transitionErr := r.transitionStaleTaskToPaused(ctx, conflictingTaskID, currentHeartbeatTime)
+			if transitionErr != nil {
+				log.Warn("failed to transition stale task to paused, using current restoredTS",
+					zap.Uint64("task_id", conflictingTaskID),
+					zap.Error(transitionErr))
+				return info.RestoredTS, nil
+			}
+
+			if transitioned {
+				log.Info("successfully transitioned stale task to paused, will reuse its restoredTS",
+					zap.Uint64("existing_task_id", conflictingTaskID),
+					zap.Uint64("existing_restored_ts", existingRestoredTS))
+				return existingRestoredTS, nil
+			}
+			log.Info("task was not transitioned (concurrent update), using current restoredTS",
+				zap.Uint64("existing_task_id", conflictingTaskID))
+			return info.RestoredTS, nil
+		}
+
+		log.Info("existing running task is active, using current restoredTS",
+			zap.Uint64("existing_task_id", conflictingTaskID))
+		return info.RestoredTS, nil
+	}
+
+	// existing task is in unexpected state, use current restoredTS
+	log.Warn("existing task is in unexpected state, using current restoredTS",
+		zap.Uint64("existing_task_id", conflictingTaskID),
+		zap.String("status", existingStatus))
+	return info.RestoredTS, nil
+}
+
+// isTaskStale checks if a running task is stale by waiting up to 5 minutes and checking if heartbeat updates
+func (r *Registry) isTaskStale(ctx context.Context, taskID uint64) (bool, error) {
+	execCtx := r.se.GetSessionCtx().GetRestrictedSQLExecutor()
+	ctx = kv.WithInternalSourceType(ctx, kv.InternalTxnBR)
+
+	// get initial heartbeat time
+	selectHeartbeatSQL := fmt.Sprintf(selectTaskHeartbeatSQLTemplate, RestoreRegistryDBName, RestoreRegistryTableName)
+	initialRows, _, err := execCtx.ExecRestrictedSQL(ctx, nil, selectHeartbeatSQL, taskID)
+	if err != nil {
+		return false, errors.Annotate(err, "failed to get initial heartbeat time")
+	}
+
+	if len(initialRows) == 0 {
+		return false, nil // task not found (might have been deleted), proceed with user's restoredTS
+	}
+
+	initialHeartbeatTime := initialRows[0].GetTime(0).String()
+
+	log.Info("checking if task is stale, will check heartbeat every minute up to 5 minutes",
+		zap.Uint64("task_id", taskID),
+		zap.String("initial_heartbeat", initialHeartbeatTime))
+
+	// check heartbeat every minute for up to 5 minutes
+	ticker := time.NewTicker(time.Minute)
+	defer ticker.Stop()
+
+	remainingMinutes := StaleTaskThresholdMinutes
+	for remainingMinutes > 0 {
+		select {
+		case <-ctx.Done():
+			return false, ctx.Err()
+		case <-ticker.C:
+			remainingMinutes--
+
+			// check heartbeat time at each tick
+			currentRows, _, err := execCtx.ExecRestrictedSQL(ctx, nil, selectHeartbeatSQL, taskID)
+			if err != nil {
+				log.Warn("failed to check heartbeat during stale check, assuming task is active",
+					zap.Uint64("task_id", taskID),
+					zap.Error(err))
+				return false, nil
+			}
+
+			if len(currentRows) == 0 {
+				return false, nil // task not found (might have been deleted), proceed with user's restoredTS
+			}
+
+			currentHeartbeatTime := currentRows[0].GetTime(0).String()
+
+			// if heartbeat changed, task is active - exit early
+			if currentHeartbeatTime != initialHeartbeatTime {
+				log.Info("task heartbeat updated, task is active",
+					zap.Uint64("task_id", taskID),
+					zap.String("initial_heartbeat", initialHeartbeatTime),
+					zap.String("current_heartbeat", currentHeartbeatTime),
+					zap.Int("minutes_waited", StaleTaskThresholdMinutes-remainingMinutes))
+				return false, nil
+			}
+
+			if remainingMinutes > 0 {
+				log.Info("task heartbeat unchanged, continuing to wait",
+					zap.Int("remaining_minutes", remainingMinutes),
+					zap.Uint64("task_id", taskID))
+			}
+		}
+	}
+
+	// if we get here, heartbeat hasn't changed for 5 minutes - task is stale
+	log.Info("task heartbeat unchanged for 5 minutes, task is stale",
+		zap.Uint64("task_id", taskID),
+		zap.String("initial_heartbeat", initialHeartbeatTime))
+
+	return true, nil
+}
+
+// transitionStaleTaskToPaused atomically transitions a stale running task to paused state
+// if the heartbeat timestamp hasn't changed. Returns whether the transition was successful.
+func (r *Registry) transitionStaleTaskToPaused(ctx context.Context, taskID uint64,
+	expectedHeartbeatTime string) (bool, error) {
+	log.Info("attempting to transition stale task to paused state",
+		zap.Uint64("task_id", taskID),
+		zap.String("expected_heartbeat", expectedHeartbeatTime))
+
+	var transitioned bool
+	err := r.executeInTransaction(ctx, func(ctx context.Context, execCtx sqlexec.RestrictedSQLExecutor,
+		sessionOpts []sqlexec.OptionFuncAlias) error {
+		// atomically update task to paused only if it's still running with the same heartbeat time
+		updateSQL := fmt.Sprintf(transitionStaleTaskToPausedSQLTemplate, RestoreRegistryDBName, RestoreRegistryTableName)
+
+		// We need to parse the heartbeat time string back to time.Time for the SQL query
+		// The expectedHeartbeatTime comes from MySQL's time format
+		expectedTime, parseErr := time.Parse("2006-01-02 15:04:05", expectedHeartbeatTime)
+		if parseErr != nil {
+			return errors.Annotatef(parseErr, "failed to parse expected heartbeat time: %s", expectedHeartbeatTime)
+		}
+
+		_, _, updateErr := execCtx.ExecRestrictedSQL(ctx, sessionOpts, updateSQL, taskID, expectedTime)
+		if updateErr != nil {
+			return errors.Annotate(updateErr, "failed to transition stale task to paused")
+		}
+
+		// Check if the task was actually transitioned by querying its current status
+		checkTaskSQL := fmt.Sprintf(
+			"SELECT status FROM %s.%s WHERE id = %%?", RestoreRegistryDBName, RestoreRegistryTableName)
+		var statusRows []chunk.Row
+		var checkErr error
+		statusRows, _, checkErr = execCtx.ExecRestrictedSQL(ctx, sessionOpts, checkTaskSQL, taskID)
+		if checkErr != nil {
+			return errors.Annotate(checkErr, "failed to check task status after transition attempt")
+		}
+
+		if len(statusRows) > 0 && statusRows[0].GetString(0) == string(TaskStatusPaused) {
+			transitioned = true
+			log.Info("successfully transitioned stale task to paused state",
+				zap.Uint64("task_id", taskID))
+		} else {
+			log.Info("task was not transitioned (either already changed state or heartbeat was updated)",
+				zap.Uint64("task_id", taskID))
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		return false, err
+	}
+
+	return transitioned, nil
+}