ai-agent/workflow/service/flow/exec_recover.go

package flow

import (
	"context"
	"errors"
	"fmt"
	"time"

	sessionDao "ai-agent/workflow/dao/session"

	"gitea.redpowerfuture.com/red-future/common/beans"
	"gitea.redpowerfuture.com/red-future/common/utils"
	wsCommon "gitea.redpowerfuture.com/red-future/common/websocket"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/gtrace"
	"github.com/google/uuid"
)

// 恢复相关常量（可参数化调整）
const (
	heartbeatInterval   = 30 * time.Second      // 心跳 touch 间隔
	heartbeatStaleAfter = 2 * heartbeatInterval // 心跳陈旧阈值（>2×间隔），小于此视为僵尸
	recoverScanInterval = 30 * time.Second      // 周期扫描间隔
	recoverLockTTL      = 15 * time.Minute      // 恢复锁 TTL（覆盖长时间续跑）
	recoverLockPhase    = 2 * time.Minute       // 抢锁/前置判定阶段总时限（Redis SET + 一次 DB 读，毫秒级）
	recoverExecTimeout  = 12 * time.Hour        // 单次恢复执行最长时长；超时按可重试失败落库交下一轮扫描
	heartbeatMaxFail    = 2                     // 心跳连续失败达到陈旧阈值（2×30s=60s）即租约丢失，中止执行
	execMaxRetryCount   = 2                     // 整次执行最多自动重试 2 次（共 3 次尝试）
)

// StartRecoveryLoop 启动恢复扫描：先立即扫一次，再周期扫描。
// 多节点各自扫描，靠 Redis 锁对同一 exec 抢占去重（Redis 锁只管瞬时互斥、运行期靠 DB 心跳，见《工作流执行并发仲裁设计.md》§1/§2）
func StartRecoveryLoop(ctx context.Context) {
	go func() {
		scanAndRecover(ctx)
		ticker := time.NewTicker(recoverScanInterval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				scanAndRecover(ctx)
			case <-ctx.Done():
				return
			}
		}
	}()
}

// scanAndRecover 扫描可恢复执行并逐个异步恢复（不阻塞扫描循环）。
// 恢复运行在无 HTTP 用户的后台：ListRecoverable 走跨租户 NoTenantId，需要 ctx 携带 OTel span
// （NoTenantId 以 traceID 为 gcache 标记键，无 span 时 getTraceID 返回空 → NoTenantId 返回 nil 会 panic）
func scanAndRecover(ctx context.Context) {
	// 优雅关停期间不再捞起：连接 ctx 刚被取消、exec 正在落终态，避免本进程内部用合成用户重跑
	if IsShuttingDown() {
		return
	}
	scanCtx, span := gtrace.NewSpan(ctx, "workflow.recover.scan")
	defer span.End()
	now := time.Now().UnixMilli()
	rows, err := sessionDao.ExecWorkflowDao.ListRecoverable(scanCtx, now, now-int64(heartbeatStaleAfter/time.Millisecond), execMaxRetryCount)
	if err != nil {
		g.Log().Errorf(scanCtx, "扫描可恢复执行失败: %v", err)
		return
	}
	for _, r := range rows {
		// 扫描触发无附着连接（attach=nil）：hub 在确认可恢复后按 exec 的 session+flow 建立
		go recoverExecution(context.WithoutCancel(scanCtx), r.Id, nil)
	}
}

// recoverExecution 统一恢复例程：抢锁 → 判定（isRecoverable）→ 条件重置抢权 → 置运行中续跑
// （触发源/抢权闸/输家纪律见《工作流执行并发仲裁设计.md》§1/§2/§4）。
// 两个触发源共用：启动/周期扫描（attach=nil）、executeOrResume 对僵尸行的用户触发（attach 携带连接）。
func recoverExecution(parentCtx context.Context, execId int64, attach *execAttach) {
	// 优雅关停期间不再启动新的恢复执行（周期扫描已有守卫；用户 executeOrResume 触发路径这里兜底）
	if IsShuttingDown() {
		return
	}
	// 在函数最前面创建并登记本执行的 cancel：trackExecRun 绑定最终控制执行的 cancel，
	// 提前登记保证优雅关停（SetShuttingDown→cancelAllExecRuns）快照必然覆盖到本执行（已登记），
	// 或本执行在开始前置判定前自行发现关停标记放弃——不留"已置 status=1 却无执行"的无主记录。
	topCtx, topCancel := context.WithCancel(context.WithoutCancel(parentCtx))
	defer topCancel()
	finish := trackExecRun(topCancel)
	defer finish()
	// 登记后复查关停标记：已置位说明 cancelAllExecRuns 快照早于本登记（未覆盖本执行），
	// 此时本执行尚未走 ResetRunningIfRecoverable（不会留下 status=1 无主执行），直接放弃；
	// 若置位发生在复查之后，本执行已被快照覆盖，随关停取消并在错误分类路径落终态（status=3/retryable=1）
	if IsShuttingDown() {
		return
	}

	// 事件中枢（attach 路径）：用户点击执行触发恢复时，复用连接已订阅的候选 hub（executeOrResume 传入），
	// 让该连接订阅到后续节点进度与终态、并能通过 workflow_cancel 永久取消（hub.CancelByUser→topCancel）。
	// 此处只订阅不 TryOwn：所有权由"抢到锁+重置权"的执行者确定；提前退出路径（锁输/不可恢复/重置失败）
	// 不关 hub，保证最终执行者与连接指向同一 hub 实例，取消/进度不因竞态丢失。
	var hub *execHub
	if attach != nil {
		if attach.hub != nil {
			hub = attach.hub
		} else {
			hub, _ = registerHubIfAbsent(attach.sessionId, attach.flowId, newExecHub(attach.sessionId, attach.flowId))
		}
		subscribeConnToHub(attach.conn, hub)
	}

	lockKey := fmt.Sprintf("workflow:exec:recover:%d", execId)
	// 抢锁/前置判定用独立短超时 ctx：该阶段只做 Redis SET + 一次 DB 读 + 一次条件重置，应毫秒级完成
	lockCtx, lockCancel := context.WithTimeout(context.WithoutCancel(parentCtx), recoverLockPhase)
	defer lockCancel()
	// 恢复体整体包进 utils.WithLock（自动续期 + 单次尝试，替代本地对象形态锁 redis_lock.go）：
	//   - 自动续期：15min TTL 覆盖整段执行。长执行原本就靠心跳陈旧 + 条件重置防双跑，
	//     锁持满只是让其它节点提前「锁忙跳过」；节点崩溃续期停 → TTL 过期 → 其它节点照常捞起；
	//   - 单次尝试（retryTimes=1）：原「被其它节点持有就跳过」语义，抢不到不等待。
	ok, err := utils.WithLock(lockCtx, lockKey, int64(recoverLockTTL/time.Second), func(_ context.Context) error {

		// 抢锁后重读：此刻租户未知（该 exec 可能属于任意租户），必须跨租户读
		exec, err := sessionDao.ExecWorkflowDao.GetByIdNoTenant(lockCtx, execId)
		if err != nil || exec == nil {
			return nil
		}
		if !isRecoverable(exec, time.Now().UnixMilli()) {
			return nil // 锁等待期间状态已变化（其它节点已恢复/用户取消/已耗尽）
		}

		// 事件中枢（scan 路径，无附着连接）：此刻才知道 exec 的 session+flow，建 hub 供后续用户连接附着
		//（所有权在重置成功后统一接管，见下）
		if attach == nil {
			hub, _ = registerHubIfAbsent(exec.SessionId, exec.FlowId, newExecHub(exec.SessionId, exec.FlowId))
		}

		// 置运行中并续跑（BuildExecution false → 断点续跑：checkpoint 有则从中继续，无则从头 + 段/memo 复用）。
		// 心跳在条件重置成功后才启动：避免对未抢到重置权的 exec 空跑心跳
		saveCtx := context.WithoutCancel(parentCtx)
		// 恢复无 HTTP 用户，但图节点 lambda 的 INSERT（node_execution/flow_async_task/segment_result）走
		// insertHook 硬性要求 user，且节点模型调用外发 model-gateway 带 X-User-Info 要过单次调用最低
		// 余额门禁（须 user.Id>0）——用 exec 所属租户合成系统用户 ctx（Id 取创建时落库的 user_id，
		// creator 仅 userName 推不回数字 id；旧记录 user_id=0 时其恢复续跑会被该门禁拦截），保留 span。
		userCtx := context.WithValue(saveCtx, "user", &beans.User{Id: uint64(exec.UserId), UserName: exec.Creator, TenantId: exec.TenantId})
		// 恢复 = 同一逻辑运行继续：复用 exec 记录的组读其 checkpoint 续跑（不换组，否则读不到断点从头跑）。
		// 仅旧记录无组时新造并随重置持久化；置运行中与图执行的组标识须一致。
		nodeGroupId := exec.NodeGroupId
		if nodeGroupId == "" {
			nodeGroupId = uuid.NewString()
		}
		// 条件重置（原子防与用户断点续跑双跑）：仅当仍可恢复（status=3 或 status=1 心跳陈旧）时才抢到重置权
		staleBeforeMs := time.Now().UnixMilli() - int64(heartbeatStaleAfter/time.Millisecond)
		reset, err := sessionDao.ExecWorkflowDao.ResetRunningIfRecoverable(userCtx, execId, nodeGroupId, staleBeforeMs)
		if err != nil {
			g.Log().Errorf(lockCtx, "恢复置运行中失败 execId=%d: %v", execId, err)
			return nil
		}
		if !reset {
			// 状态已被其它路径（用户 reExecute/execute 或其它节点恢复）抢先重置：放弃续跑，状态由持有方收敛，不落终态
			g.Log().Infof(lockCtx, "execId=%d 已被其它路径抢先重置为运行中，跳过恢复", execId)
			return nil
		}

		// 事件中枢所有权：此时已抢到锁+重置权，本 goroutine 是实际执行者，接管 hub（TryOwn 原子置 owned+cancel）。
		// 已有持有者（同 session+flow 另有执行在跑）则本恢复不持有 hub：锁与重置权已到手，放弃会留 status=1
		// 无主，继续执行但进度不广播（恢复仍正常完成落终态）。
		if hub != nil {
			if !hub.TryOwn(topCancel) {
				hub = nil
			} else {
				defer hub.Close()
				// 恢复建立前用户已取消（占位期 workflow_cancel 只置标志、cancel 尚未设置）：立即中止，
				// 走下方 UserCancelled 分类写永久取消（retryable=0）
				if hub.UserCancelled() {
					topCancel()
				}
			}
		}

		// 执行与心跳绑定同一 execCtx（带 recoverExecTimeout 上限）：
		// 心跳停止（超时/进程死/租约丢失）与 BuildExecution 中止必须同步，否则出现
		// "心跳已过期但执行还活着" 的窗口，被其它节点扫描恢复导致双跑。
		// 心跳连续失败达陈旧阈值时回调 execCancel 中止执行（心跳与执行同生共死）。
		// 从函数顶部登记的 topCtx 派生执行 ctx：注入用户信息（供 insertHook 落库）+ 12h 执行超时上限。
		// 关停时 cancelAllExecRuns 取消 topCancel → 本 ctx 随之取消，BuildExecution 中止后走下方错误分类落终态。
		execCtx, execCancel := context.WithTimeout(context.WithValue(topCtx, "user", &beans.User{Id: uint64(exec.UserId), UserName: exec.Creator, TenantId: exec.TenantId}), recoverExecTimeout)
		defer execCancel()
		stop := startHeartbeat(execCtx, execId, execCancel)
		defer stop()
		// 图节点进度经 hub 广播（无 hub 时 getProgressHub 返回 nil，reporter nil 安全）
		progressCtx := execCtx
		if hub != nil {
			progressCtx = context.WithValue(execCtx, wsProgressCtxKey{}, hub)
		}
		err = BuildExecution(progressCtx, false, exec.FlowId, execId, nodeGroupId, exec.SessionId, exec.RequestParams)
		if err != nil {
			// 用户显式取消（附着连接 workflow_cancel）：永久取消，retryable=0，恢复扫描不再捞起，
			// 杜绝"取消→恢复→再取消"循环；产物（checkpoint/段/异步缓存）保留，落"用户已终止执行"。
			// 与 WS 路径一致：失败/取消不清，统一由"成功尾部 / 下一次 forceNewRun 起跑前"清理
			if hub != nil && hub.UserCancelled() {
				retryable, retryCnt := 0, 0
				recordWorkflow(userCtx, execId, 0, context.Canceled, &retryable, &retryCnt)
				hub.Publish(&wsCommon.WsPushMsg{Type: "error", Message: errWorkflowTerminated})
				return nil
			}
			// 续跑失败：恢复例程无用户，任意错误（含执行超时/租约丢失取消/模型/DB/网络/panic）
			// 一律 retryable=1 交下一轮扫描决定是否再恢复；重试耗尽才终局失败。
			// 产物（checkpoint/段/异步缓存）保留不在此清：重试耗尽后用户仍可手动同参数续跑
			// 复用已成功段/异步结果，换参数则由 forceNewRun 起跑前统一清理（见 BuildExecution）
			retryable, retryCnt := 1, exec.RetryCount+1
			// 优雅关停导致的取消：换错误标记让 recordWorkflow 写"程序关停中断"（与 WS 路径一致），
			// 仍 retryable=1 下次启动扫描捞起续跑；非关停的取消（租约丢失/执行超时）保留原错误
			if errors.Is(err, context.Canceled) && IsShuttingDown() {
				err = errInterruptedByShutdown
			}
			recordWorkflow(userCtx, execId, 0, err, &retryable, &retryCnt)
			// 终态广播（附着连接；scan 路径无连接则无人接收）
			if hub != nil {
				hub.Publish(&wsCommon.WsPushMsg{Type: "error", Message: "工作流执行失败", Error: err.Error()})
			}
			return nil
		}
		// 续跑成功：recordWorkflow 落 status=2；BuildExecution 已清理 checkpoint/segment_result/flow_async_task
		recordWorkflow(userCtx, execId, 0, nil, nil, nil)
		// 终态广播：把本次执行保存的结果文件路径一并推给前端
		if hub != nil {
			hub.Publish(&wsCommon.WsPushMsg{Type: "flow_complete", Message: "工作流执行完成", Data: map[string]interface{}{
				"resultFileUrls": workflowResultFileUrls(userCtx, execId),
			}})
		}
		return nil
	})
	if err != nil {
		g.Log().Errorf(lockCtx, "恢复抢锁失败 execId=%d: %v", execId, err)
		return
	}
	if !ok {
		g.Log().Infof(lockCtx, "恢复锁被其它节点持有，跳过 execId=%d", execId)
		return
	}
}