feat(engine): 增加运行熔断、并发短路取消及生命周期钩子
This commit is contained in:
unclesam committed
1 parent
693c107473
commit
b9218f2a6a
6 files changed
+565
-36
No files matched your search
@@ -96,7 +96,7 @@ func (cg *CompiledGraph[S]) Start(ctx context.Context, startNode string, initial
|
||||
NextNode: startNode,
|
||||
}
|
||||
|
||||
return cg.run(ctx, thread)
|
||||
return cg.run(ctx, thread, runOptions[S]{})
|
||||
}
|
||||
|
||||
// Resume 恢复执行一个被暂停的线程，并允许注入外部修改后的状态数据（例如人工审批修改后的结果）
|
||||
@@ -112,15 +112,38 @@ func (cg *CompiledGraph[S]) Resume(ctx context.Context, thread *Thread[S], modif
|
||||
thread.State = modifiedState
|
||||
thread.IsPaused = false
|
||||
|
||||
return cg.run(ctx, thread)
|
||||
return cg.run(ctx, thread, runOptions[S]{})
|
||||
}
|
||||
|
||||
// runOptions 封装引擎运行时的可选参数，由 Engine 包装器传入
|
||||
type runOptions[S any] struct {
|
||||
maxSteps int
|
||||
stateCloner func(S) S
|
||||
preNodeHook HookFn[S]
|
||||
postNodeHook HookFn[S]
|
||||
}
|
||||
|
||||
// run 是引擎内部循环调度器，负责驱动节点向前流转
|
||||
func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread[S], error) {
|
||||
func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S], opts runOptions[S]) (*Thread[S], error) {
|
||||
stepCount := 0
|
||||
|
||||
for {
|
||||
// 检查 Context，支持外部超时控制或取消
|
||||
if err := ctx.Err(); err != nil {
|
||||
return thread, err
|
||||
// 【熔断 1】检查步数上限，防御死循环（infinite loop protection）
|
||||
if opts.maxSteps > 0 {
|
||||
if stepCount >= opts.maxSteps {
|
||||
return thread, fmt.Errorf(
|
||||
"engine halt: max steps (%d) exceeded, possible infinite loop in graph",
|
||||
opts.maxSteps,
|
||||
)
|
||||
}
|
||||
stepCount++
|
||||
}
|
||||
|
||||
// 【熔断 2】检查 Context，支持外部超时控制或取消
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return thread, ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
currentNodeName := thread.NextNode
|
||||
@@ -130,51 +153,38 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread
|
||||
return thread, nil
|
||||
}
|
||||
|
||||
// 查找并执行当前节点
|
||||
// 查找当前节点
|
||||
nodeFn, exists := cg.nodes[currentNodeName]
|
||||
if !exists {
|
||||
return thread, fmt.Errorf("runtime error: node %q not found", currentNodeName)
|
||||
}
|
||||
|
||||
// 触发 PreNodeHook
|
||||
if opts.preNodeHook != nil {
|
||||
opts.preNodeHook(ctx, currentNodeName, thread.State)
|
||||
}
|
||||
|
||||
// 执行当前节点
|
||||
newState, err := nodeFn(ctx, thread.State)
|
||||
if err != nil {
|
||||
return thread, fmt.Errorf("node %q execution error: %w", currentNodeName, err)
|
||||
}
|
||||
thread.State = newState
|
||||
|
||||
// 触发 PostNodeHook
|
||||
if opts.postNodeHook != nil {
|
||||
opts.postNodeHook(ctx, currentNodeName, thread.State)
|
||||
}
|
||||
|
||||
// 计算下一个该执行的节点名称
|
||||
var nextNode string
|
||||
if step, isParallel := cg.parallels[currentNodeName]; isParallel {
|
||||
// 【处理并发分流】
|
||||
var wg sync.WaitGroup
|
||||
branches := make([]S, len(step.targets))
|
||||
errs := make([]error, len(step.targets))
|
||||
for i, target := range step.targets {
|
||||
wg.Add(1)
|
||||
// 启动并发协程运行分支节点。传入 thread.State 的副本（Go 默认是值拷贝）
|
||||
go func(idx int, targetNode string, stateCopy S) {
|
||||
defer wg.Done()
|
||||
if ctx.Err() != nil {
|
||||
errs[idx] = ctx.Err()
|
||||
return
|
||||
}
|
||||
nodeFn := cg.nodes[targetNode]
|
||||
resState, err := nodeFn(ctx, stateCopy)
|
||||
if err != nil {
|
||||
errs[idx] = err
|
||||
return
|
||||
}
|
||||
branches[idx] = resState
|
||||
}(i, target, thread.State)
|
||||
}
|
||||
// 等待所有分支协程执行完毕
|
||||
wg.Wait()
|
||||
// 检查是否有任何分支报错
|
||||
for _, err := range errs {
|
||||
// 【处理并发分流】使用短路取消机制
|
||||
branches, err := cg.runParallelBranches(ctx, step.targets, thread.State, opts.stateCloner)
|
||||
if err != nil {
|
||||
return thread, fmt.Errorf("parallel branch execution error: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 【状态合并】调用用户自定义的合并函数
|
||||
mergedState, err := step.merger(ctx, thread.State, branches)
|
||||
if err != nil {
|
||||
@@ -194,7 +204,7 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread
|
||||
nextNode = cg.edges[currentNodeName]
|
||||
}
|
||||
|
||||
// 更新快照中的“下一个节点”
|
||||
// 更新快照中的"下一个节点"
|
||||
thread.NextNode = nextNode
|
||||
|
||||
// 如果下一站是终点，直接进入下一次循环触发结束逻辑
|
||||
@@ -209,3 +219,82 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runParallelBranches 并发执行多个目标节点，并实现首个错误的短路取消。
|
||||
//
|
||||
// - 使用 context.WithCancel 派生子 Context，任意分支报错立即调用 cancel()，
|
||||
// 通知其他仍在运行的分支通过 ctx.Done() 优雅退出，避免 AI 算力空转。
|
||||
// - 使用缓冲 channel（容量 = 分支数）收集结果，防止 goroutine 在写入时永久阻塞（goroutine 泄漏）。
|
||||
// - stateCloner 为可选的深拷贝函数；若为 nil，则传入值拷贝（对含引用类型的状态存在数据竞争风险，需用户知悉）。
|
||||
func (cg *CompiledGraph[S]) runParallelBranches(
|
||||
ctx context.Context,
|
||||
targets []string,
|
||||
state S,
|
||||
stateCloner func(S) S,
|
||||
) ([]S, error) {
|
||||
// 派生可取消的子 Context：任意分支出错时，立即广播取消信号
|
||||
cancelCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel() // 确保函数退出时释放资源，不论成功或失败
|
||||
|
||||
type result struct {
|
||||
idx int
|
||||
state S
|
||||
err error
|
||||
}
|
||||
|
||||
// 缓冲通道容量 = 分支数，确保所有 goroutine 均可不阻塞地写入结果
|
||||
resultCh := make(chan result, len(targets))
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i, target := range targets {
|
||||
wg.Add(1)
|
||||
|
||||
// 按需深拷贝状态：有 stateCloner 则深拷贝，否则值拷贝（Go 默认语义）
|
||||
var stateCopy S
|
||||
if stateCloner != nil {
|
||||
stateCopy = stateCloner(state)
|
||||
} else {
|
||||
stateCopy = state
|
||||
}
|
||||
|
||||
go func(idx int, nodeName string, s S) {
|
||||
defer wg.Done()
|
||||
|
||||
// 在执行真正的工作之前，先检查是否已被其他分支取消
|
||||
select {
|
||||
case <-cancelCtx.Done():
|
||||
resultCh <- result{idx: idx, err: cancelCtx.Err()}
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
nodeFn := cg.nodes[nodeName]
|
||||
resState, err := nodeFn(cancelCtx, s)
|
||||
if err != nil {
|
||||
cancel() // 关键：一处报错，立即取消其余所有兄弟 goroutine
|
||||
resultCh <- result{idx: idx, err: err}
|
||||
return
|
||||
}
|
||||
resultCh <- result{idx: idx, state: resState}
|
||||
}(i, target, stateCopy)
|
||||
}
|
||||
|
||||
// 等待所有 goroutine 完成后关闭结果通道
|
||||
wg.Wait()
|
||||
close(resultCh)
|
||||
|
||||
// 按索引顺序收集结果，保证合并函数收到的 branches 顺序与 targets 一致
|
||||
branches := make([]S, len(targets))
|
||||
for r := range resultCh {
|
||||
if r.err != nil {
|
||||
// 过滤掉因取消产生的次生错误，只返回第一个真实错误
|
||||
if !errors.Is(r.err, context.Canceled) {
|
||||
return nil, r.err
|
||||
}
|
||||
} else {
|
||||
branches[r.idx] = r.state
|
||||
}
|
||||
}
|
||||
|
||||
return branches, nil
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package GopherGraph
|
||||
|
||||
import "context"
|
||||
|
||||
// Engine 是 CompiledGraph 的增强包装器，在完全兼容原有 API 的基础上，
|
||||
// 提供以下四项生产级能力：
|
||||
//
|
||||
// 1. StateCloner — 自定义深拷贝逻辑，消除并发分支的数据竞争隐患
|
||||
// 2. MaxSteps — 步数硬性熔断，防止图中的环路导致无限循环
|
||||
// 3. PreNodeHook — 节点执行前的生命周期钩子（日志、链路追踪等）
|
||||
// 4. PostNodeHook — 节点执行后的生命周期钩子（指标上报、流式推送等）
|
||||
//
|
||||
// 所有选项均为可选，通过链式调用注册。未注册的选项将退回到 CompiledGraph 的默认行为。
|
||||
//
|
||||
// 用法示例：
|
||||
//
|
||||
// engine := GopherGraph.NewEngine(compiledGraph).
|
||||
// WithStateCloner(func(s MyState) MyState {
|
||||
// clone := s
|
||||
// clone.Messages = append([]string{}, s.Messages...)
|
||||
// return clone
|
||||
// }).
|
||||
// WithMaxSteps(100).
|
||||
// WithPreNodeHook(func(ctx context.Context, name string, s MyState) {
|
||||
// log.Printf("[PRE] node=%s", name)
|
||||
// }).
|
||||
// WithPostNodeHook(func(ctx context.Context, name string, s MyState) {
|
||||
// log.Printf("[POST] node=%s value=%+v", name, s)
|
||||
// })
|
||||
//
|
||||
// thread, err := engine.Start(ctx, "start", initialState)
|
||||
type Engine[S any] struct {
|
||||
graph *CompiledGraph[S]
|
||||
stateCloner func(S) S
|
||||
maxSteps int
|
||||
preNodeHook HookFn[S]
|
||||
postNodeHook HookFn[S]
|
||||
}
|
||||
|
||||
// NewEngine 用一个已编译的图创建 Engine 实例。
|
||||
func NewEngine[S any](graph *CompiledGraph[S]) *Engine[S] {
|
||||
return &Engine[S]{graph: graph}
|
||||
}
|
||||
|
||||
// WithStateCloner 注册自定义深拷贝函数。
|
||||
//
|
||||
// 当图中存在并发分支（AddParallelEdges）时，引擎会在启动每个并发 goroutine 之前
|
||||
// 调用此函数对共享状态进行深拷贝，从根本上消除引用类型（切片、map、指针）的数据竞争。
|
||||
//
|
||||
// 若不注册，引擎退化为 Go 默认的值拷贝语义（浅拷贝）。
|
||||
// 对于状态中含有引用类型的场景，强烈建议注册此函数。
|
||||
func (e *Engine[S]) WithStateCloner(fn func(S) S) *Engine[S] {
|
||||
e.stateCloner = fn
|
||||
return e
|
||||
}
|
||||
|
||||
// WithMaxSteps 设置引擎单次运行（Start/Resume）允许的最大节点执行步数。
|
||||
//
|
||||
// 这是防御图中存在意外环路（infinite loop）的最后一道硬性熔断。
|
||||
// 超过上限后，引擎将返回明确的错误，而不是永久阻塞。
|
||||
//
|
||||
// 建议根据工作流的预期最大深度合理设置（如 100~1000）。
|
||||
// 传入 0 或负数表示不限制（与不调用此方法等效）。
|
||||
func (e *Engine[S]) WithMaxSteps(n int) *Engine[S] {
|
||||
e.maxSteps = n
|
||||
return e
|
||||
}
|
||||
|
||||
// WithPreNodeHook 注册在每个节点执行【之前】触发的钩子函数。
|
||||
//
|
||||
// 典型用途：
|
||||
// - 打印结构化日志（node 名称、输入状态）
|
||||
// - 记录 OpenTelemetry Span 的开始时间
|
||||
// - 向前端推送"正在执行 xxx 节点"的流式事件
|
||||
//
|
||||
// 钩子函数接收到的 state 是节点执行【前】的状态快照，请勿修改（只读语义）。
|
||||
// 钩子中的 panic 会向上传播，建议在内部做好 recover。
|
||||
func (e *Engine[S]) WithPreNodeHook(fn HookFn[S]) *Engine[S] {
|
||||
e.preNodeHook = fn
|
||||
return e
|
||||
}
|
||||
|
||||
// WithPostNodeHook 注册在每个节点执行【之后】触发的钩子函数。
|
||||
//
|
||||
// 典型用途：
|
||||
// - 打印节点的输出状态（用于 Debug）
|
||||
// - 上报节点耗时指标（Prometheus/DataDog）
|
||||
// - 向前端推送节点执行完毕的流式事件
|
||||
//
|
||||
// 钩子函数接收到的 state 是节点执行【后】更新过的状态，同为只读语义。
|
||||
func (e *Engine[S]) WithPostNodeHook(fn HookFn[S]) *Engine[S] {
|
||||
e.postNodeHook = fn
|
||||
return e
|
||||
}
|
||||
|
||||
// opts 将 Engine 的配置打包成 runOptions，传入底层 run() 调度器。
|
||||
func (e *Engine[S]) opts() runOptions[S] {
|
||||
return runOptions[S]{
|
||||
maxSteps: e.maxSteps,
|
||||
stateCloner: e.stateCloner,
|
||||
preNodeHook: e.preNodeHook,
|
||||
postNodeHook: e.postNodeHook,
|
||||
}
|
||||
}
|
||||
|
||||
// Start 从指定节点启动工作流，行为与 CompiledGraph.Start 完全一致，
|
||||
// 但会附加所有已注册的 Engine 增强选项（Hooks、MaxSteps、StateCloner）。
|
||||
func (e *Engine[S]) Start(ctx context.Context, startNode string, initialState S) (*Thread[S], error) {
|
||||
thread := &Thread[S]{
|
||||
State: initialState,
|
||||
NextNode: startNode,
|
||||
}
|
||||
return e.graph.run(ctx, thread, e.opts())
|
||||
}
|
||||
|
||||
// Resume 恢复执行一个被暂停的线程，行为与 CompiledGraph.Resume 完全一致，
|
||||
// 但会附加所有已注册的 Engine 增强选项（Hooks、MaxSteps、StateCloner）。
|
||||
func (e *Engine[S]) Resume(ctx context.Context, thread *Thread[S], modifiedState S) (*Thread[S], error) {
|
||||
if !thread.IsPaused {
|
||||
return nil, errNotPaused
|
||||
}
|
||||
if thread.IsFinished {
|
||||
return nil, errAlreadyFinished
|
||||
}
|
||||
|
||||
thread.State = modifiedState
|
||||
thread.IsPaused = false
|
||||
|
||||
return e.graph.run(ctx, thread, e.opts())
|
||||
}
|
||||
|
||||
// 复用与 CompiledGraph.Resume 中相同的哨兵错误，避免魔法字符串重复
|
||||
var (
|
||||
errNotPaused = errStr("cannot resume: thread is not paused")
|
||||
errAlreadyFinished = errStr("cannot resume: thread is already finished")
|
||||
)
|
||||
|
||||
type errStr string
|
||||
|
||||
func (e errStr) Error() string { return string(e) }
|
||||
@@ -0,0 +1,159 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
GopherGraph "github.com/unclesam-LY/GopherGraph"
|
||||
)
|
||||
|
||||
// AgentState 是这个示例工作流的状态
|
||||
type AgentState struct {
|
||||
Input string
|
||||
Draft string
|
||||
Reviewed string
|
||||
Messages []string // 含引用类型，演示 StateCloner 的必要性
|
||||
}
|
||||
|
||||
// deepClone 演示如何为含切片的状态做安全的深拷贝
|
||||
func deepClone(s AgentState) AgentState {
|
||||
clone := s
|
||||
clone.Messages = append([]string{}, s.Messages...) // 切片深拷贝
|
||||
return clone
|
||||
}
|
||||
|
||||
func main() {
|
||||
g := GopherGraph.NewGraph[AgentState]()
|
||||
|
||||
// 节点 1：起草
|
||||
g.AddNode("drafter", func(ctx context.Context, s AgentState) (AgentState, error) {
|
||||
time.Sleep(80 * time.Millisecond) // 模拟耗时 AI 调用
|
||||
s.Draft = fmt.Sprintf("【草稿】%s 的初版内容", s.Input)
|
||||
s.Messages = append(s.Messages, "drafter: 完成起草")
|
||||
return s, nil
|
||||
})
|
||||
|
||||
// 节点 2：审阅
|
||||
g.AddNode("reviewer", func(ctx context.Context, s AgentState) (AgentState, error) {
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
s.Reviewed = s.Draft + "（已审阅 ✓）"
|
||||
s.Messages = append(s.Messages, "reviewer: 审阅完成")
|
||||
return s, nil
|
||||
})
|
||||
|
||||
// 并发节点：同时翻译成英文和日文（演示 StateCloner）
|
||||
g.AddNode("translate_en", func(ctx context.Context, s AgentState) (AgentState, error) {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
s.Messages = append(s.Messages, fmt.Sprintf("translate_en: [EN] %s", s.Reviewed))
|
||||
return s, nil
|
||||
})
|
||||
g.AddNode("translate_ja", func(ctx context.Context, s AgentState) (AgentState, error) {
|
||||
time.Sleep(120 * time.Millisecond)
|
||||
s.Messages = append(s.Messages, fmt.Sprintf("translate_ja: [JA] %s", s.Reviewed))
|
||||
return s, nil
|
||||
})
|
||||
|
||||
// 节点：发布
|
||||
g.AddNode("publisher", func(ctx context.Context, s AgentState) (AgentState, error) {
|
||||
s.Messages = append(s.Messages, "publisher: 🎉 发布成功")
|
||||
return s, nil
|
||||
})
|
||||
|
||||
// 连线
|
||||
g.AddEdge("drafter", "reviewer")
|
||||
g.AddParallelEdges("reviewer", []string{"translate_en", "translate_ja"}, "publisher",
|
||||
func(ctx context.Context, parent AgentState, branches []AgentState) (AgentState, error) {
|
||||
// 合并所有分支的 Messages
|
||||
for _, b := range branches {
|
||||
parent.Messages = append(parent.Messages, b.Messages...)
|
||||
}
|
||||
parent.Messages = append(parent.Messages, "merger: 翻译分支已合并")
|
||||
return parent, nil
|
||||
},
|
||||
)
|
||||
|
||||
cg, err := g.Compile()
|
||||
if err != nil {
|
||||
log.Fatalf("编译失败: %v", err)
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// 使用 Engine 包装器，挂载所有增强能力
|
||||
// ============================================================
|
||||
var nodeTimings = make(map[string]time.Time)
|
||||
|
||||
engine := GopherGraph.NewEngine(cg).
|
||||
// 1. 注册深拷贝函数，消除并发分支对 Messages 切片的数据竞争
|
||||
WithStateCloner(deepClone).
|
||||
// 2. 设置最大步数熔断（防御意外死循环）
|
||||
WithMaxSteps(50).
|
||||
// 3. Pre Hook：记录节点开始时间 + 向控制台推送"流式进度"
|
||||
WithPreNodeHook(func(ctx context.Context, name string, s AgentState) {
|
||||
nodeTimings[name] = time.Now()
|
||||
fmt.Printf(" ⏳ [PRE ] %-15s | 消息数: %d\n", name, len(s.Messages))
|
||||
}).
|
||||
// 4. Post Hook：计算并打印节点耗时
|
||||
WithPostNodeHook(func(ctx context.Context, name string, s AgentState) {
|
||||
elapsed := time.Since(nodeTimings[name])
|
||||
fmt.Printf(" ✅ [POST] %-15s | 耗时: %v | 消息数: %d\n", name, elapsed.Round(time.Millisecond), len(s.Messages))
|
||||
})
|
||||
|
||||
fmt.Println("🚀 启动工作流...")
|
||||
fmt.Println("─────────────────────────────────────────────")
|
||||
|
||||
totalStart := time.Now()
|
||||
thread, err := engine.Start(context.Background(), "drafter", AgentState{
|
||||
Input: "GopherGraph 技术文档",
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatalf("运行失败: %v", err)
|
||||
}
|
||||
|
||||
fmt.Println("─────────────────────────────────────────────")
|
||||
fmt.Printf("🏁 工作流结束，总耗时: %v\n\n", time.Since(totalStart).Round(time.Millisecond))
|
||||
fmt.Println("📋 最终消息日志：")
|
||||
for i, msg := range thread.State.Messages {
|
||||
fmt.Printf(" %d. %s\n", i+1, msg)
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// 演示流式输出：通过 Context 注入 channel（不需要修改任何 NodeFn）
|
||||
// ============================================================
|
||||
fmt.Println("\n─────────────────────────────────────────────")
|
||||
fmt.Println("📡 演示：通过 Context 注入 channel 实现流式输出")
|
||||
fmt.Println("─────────────────────────────────────────────")
|
||||
|
||||
type streamCtxKey struct{}
|
||||
streamCh := make(chan string, 10)
|
||||
|
||||
// 在 Context 里注入 channel
|
||||
streamCtx := context.WithValue(context.Background(), streamCtxKey{}, streamCh)
|
||||
|
||||
// 启动一个 goroutine 消费流式事件（模拟前端 SSE/WebSocket）
|
||||
go func() {
|
||||
for msg := range streamCh {
|
||||
fmt.Printf(" 📨 [STREAM] %s\n", msg)
|
||||
}
|
||||
}()
|
||||
|
||||
// Post Hook 里从 Context 取出 channel，写入流式事件
|
||||
streamEngine := GopherGraph.NewEngine(cg).
|
||||
WithStateCloner(deepClone).
|
||||
WithPostNodeHook(func(ctx context.Context, name string, s AgentState) {
|
||||
ch, ok := ctx.Value(streamCtxKey{}).(chan string)
|
||||
if ok {
|
||||
ch <- fmt.Sprintf("节点 [%s] 执行完毕，当前消息数: %d", name, len(s.Messages))
|
||||
}
|
||||
})
|
||||
|
||||
_, err = streamEngine.Start(streamCtx, "drafter", AgentState{Input: "流式演示"})
|
||||
if err != nil {
|
||||
log.Fatalf("流式演示运行失败: %v", err)
|
||||
}
|
||||
close(streamCh)
|
||||
time.Sleep(50 * time.Millisecond) // 等待消费者 goroutine 打印完毕
|
||||
|
||||
fmt.Println("\n✨ 示例结束")
|
||||
}
|
||||
@@ -1,12 +1,13 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"GopherGraph"
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
GopherGraph "github.com/unclesam-LY/GopherGraph"
|
||||
)
|
||||
|
||||
// TranslationState 代表我们工作流的共享状态（State）
|
||||
|
||||
@@ -73,3 +73,8 @@ func (g *Graph[S]) AddParallelEdges(from string, targets []string, next string,
|
||||
merger: merger,
|
||||
}
|
||||
}
|
||||
|
||||
// HookFn 是节点执行前后的生命周期钩子函数签名。
|
||||
// nodeName 为当前执行的节点名称，state 为该节点执行时的状态快照（只读语义，请勿修改）。
|
||||
// 可通过此钩子接入日志、链路追踪（OpenTelemetry）、监控指标等能力，而无需修改任何 NodeFn。
|
||||
type HookFn[S any] func(ctx context.Context, nodeName string, state S)
|
||||
+135
@@ -317,3 +317,138 @@ func TestFileCheckpointer(t *testing.T) {
|
||||
t.Errorf("执行结果错误，最终值: %d, 是否结束: %t", thread.State.Value, thread.IsFinished)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMaxStepsCircuitBreaker 验证 Engine.WithMaxSteps 能在步数超限时熔断，防止死循环
|
||||
func TestMaxStepsCircuitBreaker(t *testing.T) {
|
||||
g := NewGraph[TestState]()
|
||||
|
||||
// 构造一个永远在 A -> B -> A 之间循环的图
|
||||
g.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
s.Value++
|
||||
return s, nil
|
||||
})
|
||||
g.AddNode("B", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
s.Value++
|
||||
return s, nil
|
||||
})
|
||||
|
||||
// A -> B -> A：形成闭环
|
||||
g.AddEdge("A", "B")
|
||||
g.AddEdge("B", "A")
|
||||
|
||||
cg, err := g.Compile()
|
||||
if err != nil {
|
||||
t.Fatalf("编译图失败: %v", err)
|
||||
}
|
||||
|
||||
// 设置最大步数为 5，图会在第 5 步触发熔断
|
||||
engine := NewEngine(cg).WithMaxSteps(5)
|
||||
_, err = engine.Start(context.Background(), "A", TestState{})
|
||||
if err == nil {
|
||||
t.Fatal("期望返回步数超限错误，实际无错误返回")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "max steps") {
|
||||
t.Errorf("期望错误信息包含 'max steps'，实际: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentCancellation 验证并发分支中一个节点报错时，其他兄弟节点能被及时取消
|
||||
func TestConcurrentCancellation(t *testing.T) {
|
||||
g := NewGraph[TestState]()
|
||||
|
||||
g.AddNode("start", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
return s, nil
|
||||
})
|
||||
// task_fail：立即返回错误
|
||||
g.AddNode("task_fail", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
return s, errors.New("task_fail intentional error")
|
||||
})
|
||||
// task_long：模拟耗时操作，应感知取消并提前退出
|
||||
g.AddNode("task_long", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
select {
|
||||
case <-time.After(5 * time.Second): // 5 秒，正常不应该跑完
|
||||
s.Value = 999
|
||||
return s, nil
|
||||
case <-ctx.Done():
|
||||
return s, ctx.Err()
|
||||
}
|
||||
})
|
||||
|
||||
g.AddParallelEdges("start", []string{"task_fail", "task_long"}, "", func(_ context.Context, parent TestState, _ []TestState) (TestState, error) {
|
||||
return parent, nil
|
||||
})
|
||||
|
||||
cg, err := g.Compile()
|
||||
if err != nil {
|
||||
t.Fatalf("编译图失败: %v", err)
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
_, err = cg.Start(context.Background(), "start", TestState{})
|
||||
elapsed := time.Since(start)
|
||||
|
||||
// 验证确实返回了错误
|
||||
if err == nil {
|
||||
t.Fatal("期望返回并发分支错误，实际无错误返回")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "task_fail intentional error") {
|
||||
t.Errorf("期望错误来自 task_fail，实际: %v", err)
|
||||
}
|
||||
// 验证 task_long 被快速取消，整个并发组的耗时远小于 5 秒
|
||||
if elapsed > 2*time.Second {
|
||||
t.Errorf("期望并发组在取消后快速结束（< 2s），实际耗时: %v", elapsed)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLifecycleHooks 验证 Engine 的 Pre/Post 钩子在每个节点执行前后均被正确触发
|
||||
func TestLifecycleHooks(t *testing.T) {
|
||||
g := NewGraph[TestState]()
|
||||
|
||||
g.AddNode("node1", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
s.Value += 1
|
||||
return s, nil
|
||||
})
|
||||
g.AddNode("node2", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
s.Value += 10
|
||||
return s, nil
|
||||
})
|
||||
g.AddEdge("node1", "node2")
|
||||
|
||||
cg, err := g.Compile()
|
||||
if err != nil {
|
||||
t.Fatalf("编译图失败: %v", err)
|
||||
}
|
||||
|
||||
var hookLog []string
|
||||
|
||||
engine := NewEngine(cg).
|
||||
WithPreNodeHook(func(ctx context.Context, name string, s TestState) {
|
||||
hookLog = append(hookLog, "pre:"+name)
|
||||
}).
|
||||
WithPostNodeHook(func(ctx context.Context, name string, s TestState) {
|
||||
hookLog = append(hookLog, "post:"+name)
|
||||
})
|
||||
|
||||
thread, err := engine.Start(context.Background(), "node1", TestState{})
|
||||
if err != nil {
|
||||
t.Fatalf("启动图失败: %v", err)
|
||||
}
|
||||
if !thread.IsFinished {
|
||||
t.Error("期望工作流运行结束")
|
||||
}
|
||||
// 最终值：node1(+1) -> node2(+10) = 11
|
||||
if thread.State.Value != 11 {
|
||||
t.Errorf("期望 Value 为 11，实际为 %d", thread.State.Value)
|
||||
}
|
||||
|
||||
// 验证 Hook 触发顺序：pre:node1 -> post:node1 -> pre:node2 -> post:node2
|
||||
expected := []string{"pre:node1", "post:node1", "pre:node2", "post:node2"}
|
||||
if len(hookLog) != len(expected) {
|
||||
t.Fatalf("期望 %d 次 hook 触发，实际 %d 次: %v", len(expected), len(hookLog), hookLog)
|
||||
}
|
||||
for i, entry := range expected {
|
||||
if hookLog[i] != entry {
|
||||
t.Errorf("第 %d 次 hook 期望 %q，实际 %q", i, entry, hookLog[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user