feat(workflow): 恢复例程 + 心跳 + 启动/周期扫描

2026-08-25 16:44:29 +08:00
parent 788c835580
commit 95dc0d6052
4 changed files with 256 additions and 3 deletions
@@ -156,9 +156,26 @@ func (d *execWorkflowDao) UpdateRetry(ctx context.Context, id int64, retryable i
 	return err
 }

-// ListRecoverable 返回可恢复执行：僵尸运行中（status=1 且心跳陈旧）或可重试失败（status=3 且 retryable=1 且未耗尽）
+// GetByIdNoTenant 跨租户按 id 读取执行记录（不追加 tenant_id 过滤）。
+// 供恢复例程抢锁后、尚未知晓租户时重读 exec 使用；用户路径请继续用租户隔离的 GetById。
+func (d *execWorkflowDao) GetByIdNoTenant(ctx context.Context, id int64) (res *entity.ExecWorkflow, err error) {
+	r, err := gfdb.DB(ctx, public.DbNameBlackDeacon).Model(ctx, public.TableNameExecWorkflow).
+		NoTenantId(ctx).
+		Where(entity.ExecWorkflowCol.Id, id).
+		One()
+	if err != nil {
+		return
+	}
+	err = r.Struct(&res)
+	return
+}
+
+// ListRecoverable 返回可恢复执行：僵尸运行中（status=1 且心跳陈旧）或可重试失败（status=3 且 retryable=1 且未耗尽）。
+// 调用方是跨租户的恢复扫描（无 HTTP 用户），显式 NoTenantId 走系统级扫描，避免僵尸执行因租户过滤漏检。
+// ctx 必须携带 OTel span（NoTenantId 依赖 traceID 作为 gcache 标记键），恢复扫描由 scanAndRecover 保证。
 func (d *execWorkflowDao) ListRecoverable(ctx context.Context, now int64, staleBefore int64, maxRetry int) (res []*entity.ExecWorkflow, err error) {
 	r, err := gfdb.DB(ctx, public.DbNameBlackDeacon).Model(ctx, public.TableNameExecWorkflow).
+		NoTenantId(ctx).
 		Where(fmt.Sprintf("(%s = ? AND %s < ?) OR (%s = ? AND %s = 1 AND %s < ?)",
 			entity.ExecWorkflowCol.Status, entity.ExecWorkflowCol.LastHeartbeat,
 			entity.ExecWorkflowCol.Status, entity.ExecWorkflowCol.Retryable, entity.ExecWorkflowCol.RetryCount),
@@ -12,6 +12,7 @@ import (

 	"gitea.redpowerfuture.com/red-future/common/beans"
 	_ "github.com/gogf/gf/contrib/drivers/pgsql/v2"
+	"github.com/gogf/gf/v2/net/gtrace"
 )

 // TestExecWorkflowRetryColsRoundTrip 验证三新列读写 + ListRecoverable 三类判定（需 AI_AGENT_TEST_DB=1 且本地 PG 可用）
@@ -19,8 +20,11 @@ func TestExecWorkflowRetryColsRoundTrip(t *testing.T) {
 	if !isTestDBEnabled(t) {
 		t.Skip("skip: AI_AGENT_TEST_DB not set")
 	}
-	// gfdb 的 Hook 依赖 ctx 中注入用户信息（租户/创建人），否则 Insert 报 "token 数据为空"
-	ctx := context.WithValue(context.Background(), "user", &beans.User{
+	// gfdb 的 Hook 依赖 ctx 中注入用户信息（租户/创建人），否则 Insert 报 "token 数据为空"。
+	// 先造 span：ListRecoverable 链 NoTenantId 依赖 traceID 作 gcache 标记键，无 span 时 getTraceID 返回空 → NoTenantId 返回 nil 会 panic
+	ctx, span := gtrace.NewSpan(context.Background(), "test.recover")
+	defer span.End()
+	ctx = context.WithValue(ctx, "user", &beans.User{
 		UserName: "test-retry", TenantId: 1,
 	})
 	id, err := ExecWorkflowDao.Insert(ctx, &sessionDto.CreateWorkflowReq{
@@ -0,0 +1,172 @@
+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 }
+}
@@ -0,0 +1,60 @@
+package flow
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"ai-agent/workflow/consts/flow"
+	"ai-agent/workflow/model/entity"
+)
+
+func TestIsRecoverable(t *testing.T) {
+	now := time.Now().UnixMilli()
+	staleBefore := now - int64(heartbeatStaleAfter/time.Millisecond) - 1
+	fresh := now // 心跳新鲜
+
+	running := func(hb int64) *entity.ExecWorkflow {
+		return &entity.ExecWorkflow{Status: flow.FlowExecutionStatusRunning.Code(), LastHeartbeat: hb}
+	}
+	failed := func(retryable, retryCount int) *entity.ExecWorkflow {
+		return &entity.ExecWorkflow{Status: flow.FlowExecutionStatusFailed.Code(), Retryable: retryable, RetryCount: retryCount}
+	}
+
+	cases := []struct {
+		name string
+		exec *entity.ExecWorkflow
+		want bool
+	}{
+		{"nil", nil, false},
+		{"status=1 心跳陈旧 → 恢复", running(staleBefore), true},
+		{"status=1 心跳新鲜 → 不恢复", running(fresh), false},
+		{"status=1 心跳为0（老数据）→ 恢复", running(0), true},
+		{"status=3 retryable=1 未耗尽 → 恢复", failed(1, 0), true},
+		{"status=3 retryable=1 已耗尽 → 不恢复", failed(1, execMaxRetryCount), false},
+		{"status=3 retryable=0（用户取消）→ 不恢复", failed(0, 0), false},
+		{"status=2 成功 → 不恢复", &entity.ExecWorkflow{Status: flow.FlowExecutionStatusSuccess.Code()}, false},
+	}
+	for _, c := range cases {
+		if got := isRecoverable(c.exec, now); got != c.want {
+			t.Fatalf("%s: got %v want %v", c.name, got, c.want)
+		}
+	}
+}
+
+func TestShouldRetry(t *testing.T) {
+	canceled := context.Canceled
+	if shouldRetry(canceled) {
+		t.Fatalf("用户取消不应重试")
+	}
+	if shouldRetry(nil) {
+		t.Fatalf("nil 不应重试")
+	}
+	if shouldRetry(errExecAlreadyRunning) {
+		t.Fatalf("运行中标记不应重试")
+	}
+	if !shouldRetry(errors.New("模型调用失败")) {
+		t.Fatalf("程序报错应重试")
+	}
+}