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

package flow

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

	"ai-agent/workflow/consts/flow"
	flowDao "ai-agent/workflow/dao/flow"
	sessionDao "ai-agent/workflow/dao/session"
	"ai-agent/workflow/model/entity"

	"gitea.redpowerfuture.com/red-future/common/beans"
	"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（覆盖长时间续跑）
	execMaxRetryCount   = 2                     // 整次执行最多自动重试 2 次（共 3 次尝试）
)

// errExecAlreadyRunning 用户触发时该执行已在运行（本节点/其它节点后台恢复），不新建执行
var errExecAlreadyRunning = errors.New("工作流正在执行中，不重复执行")

// shouldRetry 错误分类（spec §3）：用户取消不重试；其余（模型/DB/网络/超时/panic）程序报错重试
func shouldRetry(err error) bool {
	return err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, errExecAlreadyRunning)
}

// StartRecoveryLoop 启动恢复扫描：先立即扫一次，再周期扫描。
// 多节点各自扫描，靠 Redis 锁对同一 exec 抢占去重（spec §8.2/8.3）
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) {
	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 {
		go recoverExecution(context.WithoutCancel(scanCtx), r.Id)
	}
}

// recoverExecution 统一恢复例程（spec §7）：抢锁 → 判定 → 置运行中续跑。
// 两个触发源共用：启动/周期扫描、executeOrResume 对僵尸行的用户触发。
func recoverExecution(parentCtx context.Context, execId int64) {
	lock := NewRedisLock(fmt.Sprintf("workflow:exec:recover:%d", execId), int64(recoverLockTTL/time.Second))
	ctx, cancel := context.WithTimeout(context.WithoutCancel(parentCtx), 2*time.Hour)
	defer cancel()

	ok, err := lock.Acquire(ctx)
	if err != nil {
		g.Log().Errorf(ctx, "恢复抢锁失败 execId=%d: %v", execId, err)
		return
	}
	if !ok {
		g.Log().Infof(ctx, "恢复锁被其它节点持有，跳过 execId=%d", execId)
		return
	}
	defer lock.Release(context.Background())

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

	stop := startHeartbeat(ctx, execId)
	defer stop()

	// 置运行中并续跑（BuildExecution false → 断点续跑：checkpoint 有则从中继续，无则从头 + 段/memo 复用）
	saveCtx := context.WithoutCancel(parentCtx)
	// 恢复无 HTTP 用户，但图节点 lambda 的 INSERT（node_execution/flow_async_task/segment_result）走
	// insertHook 硬性要求 user；用 exec 所属租户合成系统用户 ctx（保留 span），供后续落库使用
	userCtx := context.WithValue(saveCtx, "user", &beans.User{UserName: "workflow-recover", TenantId: exec.TenantId})
	nodeGroupId := uuid.NewString() // 置运行中与图执行的节点组标识须一致
	if err := sessionDao.ExecWorkflowDao.ResetRunning(userCtx, execId, nodeGroupId); err != nil {
		g.Log().Errorf(ctx, "恢复置运行中失败 execId=%d: %v", execId, err)
		return
	}
	err = BuildExecution(userCtx, false, exec.FlowId, execId, nodeGroupId, exec.SessionId, exec.RequestParams)
	if err != nil {
		// 续跑失败：错误分类落库（retryable/retry_count），留给下一轮扫描决定是否再恢复
		if errors.Is(err, context.Canceled) {
			_ = sessionDao.ExecWorkflowDao.UpdateRetry(userCtx, execId, 0, exec.RetryCount)
		} else {
			_ = sessionDao.ExecWorkflowDao.UpdateRetry(userCtx, execId, 1, exec.RetryCount+1)
			// 终局清理（Task 5 用户裁定）：非用户取消且重试耗尽 → 执行永久失败，
			// 清理该 exec 残留的 flow_async_task 孤儿缓存，避免未来复用同一 execId 的运行误取到过期 done 结果
			if exec.RetryCount+1 >= execMaxRetryCount {
				_ = flowDao.FlowAsyncTaskDao.DeleteByExecution(userCtx, execId)
			}
		}
		recordWorkflow(userCtx, execId, 0, err)
		return
	}
	// 续跑成功：recordWorkflow 落 status=2；BuildExecution 已清理 checkpoint/segment_result/flow_async_task
	recordWorkflow(userCtx, execId, 0, nil)
}

// isRecoverable 判定可恢复（spec §3）：僵尸运行中（status=1 且心跳陈旧）或可重试失败（status=3 且 retryable=1 且未耗尽）。
// 与 ListRecoverable SQL 判定一致；last_heartbeat=0（老数据/默认）视为陈旧。
func isRecoverable(exec *entity.ExecWorkflow, nowMs int64) bool {
	if exec == nil || exec.Status == nil {
		return false
	}
	status := *exec.Status
	switch status {
	case *flow.FlowExecutionStatusRunning.Code():
		return exec.LastHeartbeat < nowMs-int64(heartbeatStaleAfter/time.Millisecond)
	case *flow.FlowExecutionStatusFailed.Code():
		return exec.Retryable == 1 && exec.RetryCount < execMaxRetryCount
	default:
		return false
	}
}

// startHeartbeat 后台心跳 goroutine：每 30s touch last_heartbeat，返回 stop 函数。
// 覆盖正常执行与恢复执行，崩溃前最后一次心跳即崩溃近似时间戳（spec §4）
func startHeartbeat(ctx context.Context, execId int64) func() {
	stopCh := make(chan struct{})
	done := make(chan struct{})
	go func() {
		defer close(done)
		ticker := time.NewTicker(heartbeatInterval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				if err := sessionDao.ExecWorkflowDao.TouchHeartbeat(ctx, execId); err != nil {
					g.Log().Warningf(ctx, "心跳落库失败 execId=%d: %v", execId, err)
				}
			case <-stopCh:
				return
			case <-ctx.Done():
				return
			}
		}
	}()
	return func() { close(stopCh); <-done }
}