From 2bbd2b7b0442e241cfc1549662f0fcfcb2fbc56c Mon Sep 17 00:00:00 2001 From: unclesam <2036253946@qq.com> Date: Thu, 11 Jun 2026 02:00:34 +0800 Subject: [PATCH] =?UTF-8?q?test:=20=E6=B7=BB=E5=8A=A0=E5=B9=B6=E5=8F=91?= =?UTF-8?q?=E5=88=86=E6=B5=81=E4=B8=8E=E5=90=88=E5=B9=B6=E9=80=BB=E8=BE=91?= =?UTF-8?q?=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- engine.go | 68 +++++++++++++++++++++++++++++++++++++++++++++++++-- graph.go | 33 ++++++++++++++++++++++--- graph_test.go | 61 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 156 insertions(+), 6 deletions(-) diff --git a/engine.go b/engine.go index dd2c2a6..64e9f74 100644 --- a/engine.go +++ b/engine.go @@ -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] diff --git a/graph.go b/graph.go index abe9b36..24ce836 100644 --- a/graph.go +++ b/graph.go @@ -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, + } +} diff --git a/graph_test.go b/graph_test.go index 3d6b8da..454df8f 100644 --- a/graph_test.go +++ b/graph_test.go @@ -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) + } + } +}