test: 添加并发分流与合并逻辑测试
This commit is contained in:
unclesam committed
1 parent
1c0296fa80
commit
2bbd2b7b04
3 files changed
+156
-6
No files matched your search
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Thread 代表图执行的具体实例快照（线程上下文）
|
||||
@@ -21,6 +22,7 @@ type CompiledGraph[S any] struct {
|
||||
edges map[string]string
|
||||
conditional map[string]RouterFn[S]
|
||||
interrupts map[string]bool
|
||||
parallels map[string]parallelStep[S]
|
||||
}
|
||||
|
||||
// Compile 校验图结构的合法性，并生成可运行的 CompiledGraph。
|
||||
@@ -53,12 +55,37 @@ func (g *Graph[S]) Compile() (*CompiledGraph[S], error) {
|
||||
}
|
||||
}
|
||||
|
||||
// 校验并发连线边的有效性
|
||||
for from, step := range g.parallels {
|
||||
if _, exists := g.nodes[from]; !exists {
|
||||
return nil, fmt.Errorf("compile error: parallel edge origin %q does not exist", from)
|
||||
}
|
||||
if len(step.targets) == 0 {
|
||||
return nil, fmt.Errorf("compile error: parallel edge from %q has no targets", from)
|
||||
}
|
||||
|
||||
for _, target := range step.targets {
|
||||
if _, exists := g.nodes[target]; !exists {
|
||||
return nil, fmt.Errorf("compile error: parallel target %q from %q does not exist", target, from)
|
||||
}
|
||||
}
|
||||
if step.next != "" {
|
||||
if _, exists := g.nodes[step.next]; !exists {
|
||||
return nil, fmt.Errorf("compile error: parallel next node %q from %q does not exist", step.next, from)
|
||||
}
|
||||
}
|
||||
if step.merger == nil {
|
||||
return nil, fmt.Errorf("compile error: parallel step from %q has nil merger", from)
|
||||
}
|
||||
}
|
||||
|
||||
// 返回编译好的只读图
|
||||
return &CompiledGraph[S]{
|
||||
nodes: g.nodes,
|
||||
edges: g.edges,
|
||||
conditional: g.conditional,
|
||||
interrupts: g.interrupts,
|
||||
parallels: g.parallels,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -117,14 +144,51 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread
|
||||
|
||||
// 计算下一个该执行的节点名称
|
||||
var nextNode string
|
||||
if routerFn, ok := cg.conditional[currentNodeName]; ok {
|
||||
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 {
|
||||
if err != nil {
|
||||
return thread, fmt.Errorf("parallel branch execution error: %w", err)
|
||||
}
|
||||
}
|
||||
// 【状态合并】调用用户自定义的合并函数
|
||||
mergedState, err := step.merger(ctx, thread.State, branches)
|
||||
if err != nil {
|
||||
return thread, fmt.Errorf("parallel merger execution error: %w", err)
|
||||
}
|
||||
thread.State = mergedState
|
||||
nextNode = step.next
|
||||
} else if routerFn, ok := cg.conditional[currentNodeName]; ok {
|
||||
// 如果有条件路由函数，则通过路由函数动态计算去向
|
||||
next, err := routerFn(ctx, thread.State)
|
||||
if err != nil {
|
||||
return thread, fmt.Errorf("router for node %q execution error: %w", currentNodeName, err)
|
||||
}
|
||||
nextNode = next
|
||||
|
||||
} else {
|
||||
// 否则使用静态连线边
|
||||
nextNode = cg.edges[currentNodeName]
|
||||
|
||||
@@ -10,12 +10,23 @@ type NodeFn[S any] func(ctx context.Context, state S) (S, error)
|
||||
// 它根据当前状态 S，动态决定下一个要执行的节点名称（例如返回 "translator" 或 "reviewer"
|
||||
type RouterFn[S any] func(ctx context.Context, state S) (string, error)
|
||||
|
||||
// ParallelMergeFn 接收分流前的原始状态 parent，以及各个并发分支运行结束后的状态列表 branches，返回合并后的状态。
|
||||
type ParallelMergeFn[S any] func(ctx context.Context, parent S, branches []S) (S, error)
|
||||
|
||||
// parallelStep 内部结构，用于记录并发分支的配置
|
||||
type parallelStep[S any] struct {
|
||||
targets []string // 需要并发执行的目标节点列表
|
||||
next string // 合并后下一步要去的节点名称
|
||||
merger ParallelMergeFn[S] // 状态合并函数
|
||||
}
|
||||
|
||||
// Graph 是图的构建器，S 代表用户自定义的状态结构体。
|
||||
type Graph[S any] struct {
|
||||
nodes map[string]NodeFn[S] // 节点名称 -> 节点执行函数
|
||||
edges map[string]string // 普通单向边：源节点 -> 目标节点
|
||||
conditional map[string]RouterFn[S] // 条件路由边：源节点 -> 路由函数
|
||||
interrupts map[string]bool // 标记需要中断的节点：执行完该节点后挂起
|
||||
nodes map[string]NodeFn[S] // 节点名称 -> 节点执行函数
|
||||
edges map[string]string // 普通单向边：源节点 -> 目标节点
|
||||
conditional map[string]RouterFn[S] // 条件路由边：源节点 -> 路由函数
|
||||
interrupts map[string]bool // 标记需要中断的节点：执行完该节点后挂起
|
||||
parallels map[string]parallelStep[S] // 并发分支：源节点 -> 并发步骤配置
|
||||
}
|
||||
|
||||
// NewGraph 创建并初始化一个全新强类型的图构建器
|
||||
@@ -25,6 +36,7 @@ func NewGraph[S any]() *Graph[S] {
|
||||
edges: make(map[string]string),
|
||||
conditional: make(map[string]RouterFn[S]),
|
||||
interrupts: make(map[string]bool),
|
||||
parallels: make(map[string]parallelStep[S]), // 初始化 parallels
|
||||
}
|
||||
}
|
||||
|
||||
@@ -48,3 +60,16 @@ func (g *Graph[S]) AddConditionalEdges(from string, router RouterFn[S]) {
|
||||
func (g *Graph[S]) AddInterrupt(nodeName string) {
|
||||
g.interrupts[nodeName] = true
|
||||
}
|
||||
|
||||
// AddParallelEdges 建立并发分支。
|
||||
// from: 起始节点
|
||||
// targets: 需要并发执行的多个目标节点名称（如 "translate_en", "translate_fr"）
|
||||
// next: 所有并发节点运行完并合并后，下一步要去的节点（如 "publisher"）
|
||||
// merger: 用于合并各个并发分支状态的函数
|
||||
func (g *Graph[S]) AddParallelEdges(from string, targets []string, next string, merger ParallelMergeFn[S]) {
|
||||
g.parallels[from] = parallelStep[S]{
|
||||
targets: targets,
|
||||
next: next,
|
||||
merger: merger,
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,7 @@ package GopherGraph
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
@@ -190,3 +191,63 @@ func TestTimeoutCancellation(t *testing.T) {
|
||||
t.Errorf("期望返回 context.DeadlineExceeded 错误，实际返回: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestParallelExecution 测试并发分流与合并逻辑 (start -> [task1, task2 并发] -> merger -> end)
|
||||
func TestParallelExecution(t *testing.T) {
|
||||
g := NewGraph[TestState]()
|
||||
// 1. 注册节点
|
||||
g.AddNode("start", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
s.Log = append(s.Log, "start")
|
||||
return s, nil
|
||||
})
|
||||
g.AddNode("task1", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
s.Log = append(s.Log, "task1")
|
||||
s.Value += 10
|
||||
return s, nil
|
||||
})
|
||||
g.AddNode("task2", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
s.Log = append(s.Log, "task2")
|
||||
s.Value += 20
|
||||
return s, nil
|
||||
})
|
||||
g.AddNode("end", func(ctx context.Context, s TestState) (TestState, error) {
|
||||
s.Log = append(s.Log, "end")
|
||||
return s, nil
|
||||
})
|
||||
// 定义合并函数：将分支产生的 Value 相加，合并日志
|
||||
merger := func(ctx context.Context, parent TestState, branches []TestState) (TestState, error) {
|
||||
parent.Log = append(parent.Log, "merged")
|
||||
for _, b := range branches {
|
||||
parent.Value += b.Value
|
||||
parent.Log = append(parent.Log, b.Log...)
|
||||
}
|
||||
return parent, nil
|
||||
}
|
||||
// 2. 建立并发连线：从 start 分流并发执行 task1 和 task2，执行完后通过 merger 合并状态并去往 end 节点
|
||||
g.AddParallelEdges("start", []string{"task1", "task2"}, "end", merger)
|
||||
cg, err := g.Compile()
|
||||
if err != nil {
|
||||
t.Fatalf("编译图失败: %v", err)
|
||||
}
|
||||
// 3. 运行工作流
|
||||
thread, err := cg.Start(context.Background(), "start", TestState{Value: 0})
|
||||
if err != nil {
|
||||
t.Fatalf("启动图失败: %v", err)
|
||||
}
|
||||
// 4. 验证结果
|
||||
if !thread.IsFinished {
|
||||
t.Errorf("期望工作流正常结束")
|
||||
}
|
||||
// 验证 Value 是否被正确累加: task1 (10) + task2 (20) = 30
|
||||
if thread.State.Value != 30 {
|
||||
t.Errorf("期望 Value 为 30，实际为 %d", thread.State.Value)
|
||||
}
|
||||
// 验证日志中是否同时包含两个并发节点的执行结果（且顺序不固定，但都必须存在）
|
||||
logStr := strings.Join(thread.State.Log, ",")
|
||||
expectedLogs := []string{"start", "merged", "task1", "task2", "end"}
|
||||
for _, exp := range expectedLogs {
|
||||
if !strings.Contains(logStr, exp) {
|
||||
t.Errorf("期望日志中包含 %q，实际日志为: %v", exp, thread.State.Log)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user