From 19fcb4a98124981e3764b18d8e10f7a9fd148daf Mon Sep 17 00:00:00 2001 From: Yingfeng Date: Fri, 3 Jul 2026 12:53:25 +0800 Subject: [PATCH] Fix harness DAG slow-branch test cased by nil initialization of pregel engine (#16591) --- internal/harness/core/workflow_graph_test.go | 126 ------------------- internal/harness/harness_test.go | 117 +++++++++++++++++ 2 files changed, 117 insertions(+), 126 deletions(-) diff --git a/internal/harness/core/workflow_graph_test.go b/internal/harness/core/workflow_graph_test.go index a27d0d7b2c..05cf9dca3e 100644 --- a/internal/harness/core/workflow_graph_test.go +++ b/internal/harness/core/workflow_graph_test.go @@ -953,132 +953,6 @@ func TestGraph_ParallelGraph_ConcurrentTenants(t *testing.T) { t.Logf("Concurrent parallel graph: %d tenants, %d agents each", numTenants, agentsPerGraph) } -// ===================================================================== -// Test 14: AllPredecessor with slow branches — verify merge waits for slowest -// ===================================================================== - -func TestGraph_DAG_SlowBranchMerge(t *testing.T) { - var mu sync.Mutex - var order []string - slowStarted := make(chan struct{}) - fastDone := make(chan struct{}) - releaseSlow := make(chan struct{}) - mergeRan := make(chan struct{}, 1) - record := func(name string) { - mu.Lock() - order = append(order, name) - mu.Unlock() - } - - sg := graph.NewStateGraph(&dagState{}) - - // Start → dispatch → {fast, slow} → merge - sg.AddNode("dispatch", func(ctx context.Context, state interface{}) (interface{}, error) { - s := state.(*dagState) - record("dispatch") - s.Messages = append(s.Messages, "dispatch done") - return s, nil - }) - sg.AddNode("fast", func(ctx context.Context, state interface{}) (interface{}, error) { - s := state.(*dagState) - record("fast") - s.Messages = append(s.Messages, "fast done") - close(fastDone) - return s, nil - }) - sg.AddNode("slow", func(ctx context.Context, state interface{}) (interface{}, error) { - s := state.(*dagState) - close(slowStarted) - <-releaseSlow - record("slow") - s.Messages = append(s.Messages, "slow done") - return s, nil - }) - sg.AddNode("merge", func(ctx context.Context, state interface{}) (interface{}, error) { - s := state.(*dagState) - select { - case mergeRan <- struct{}{}: - default: - } - record("merge") - s.Messages = append(s.Messages, "merge done") - return s, nil - }) - - sg.AddEdge(constants.Start, "dispatch") - sg.AddEdge("dispatch", "fast") - sg.AddEdge("dispatch", "slow") - sg.AddEdge("fast", "merge") - sg.AddEdge("slow", "merge") - sg.AddEdge("merge", constants.End) - - compiled, err := sg.Compile( - graph.WithNodeTriggerMode(types.NodeTriggerAllPredecessor), - graph.WithRecursionLimit(10), - ) - if err != nil { - t.Fatal(err) - } - - done := make(chan error, 1) - go func() { - _, err := compiled.Invoke(context.Background(), &dagState{ - Messages: []string{"start"}, - }) - done <- err - }() - - select { - case <-slowStarted: - case <-time.After(time.Second): - t.Fatal("slow branch never started") - } - - select { - case <-fastDone: - case <-time.After(time.Second): - t.Fatal("fast branch never completed") - } - - select { - case <-mergeRan: - t.Fatal("BUG: merge ran before slow branch completed") - case <-time.After(50 * time.Millisecond): - } - - close(releaseSlow) - - select { - case err = <-done: - case <-time.After(time.Second): - t.Fatal("compiled graph did not finish after slow branch release") - } - - if err != nil { - t.Fatal(err) - } - - t.Logf("Slow branch merge order=%v", order) - - if len(order) != 4 { - t.Errorf("expected 4 nodes (dispatch+fast+slow+merge), got %d: %v", len(order), order) - } - slowIdx, mergeIdx := -1, -1 - for i, name := range order { - switch name { - case "slow": - slowIdx = i - case "merge": - mergeIdx = i - } - } - if mergeIdx < slowIdx { - t.Errorf("BUG: merge before slow branch. merge=%d, slow=%d", mergeIdx, slowIdx) - } else { - t.Log("DAG slow-branch merge: merge correctly waited for slow branch") - } -} - // ===================================================================== // StateGraph Invoke Error Recovery Tests // ===================================================================== diff --git a/internal/harness/harness_test.go b/internal/harness/harness_test.go index 20f4e3a378..9fb57c2bbb 100644 --- a/internal/harness/harness_test.go +++ b/internal/harness/harness_test.go @@ -2,7 +2,11 @@ package harness import ( "context" + "sync" "testing" + "time" + + "ragflow/internal/harness/graph/types" ) func TestNewStateGraph(t *testing.T) { @@ -96,3 +100,116 @@ func TestDefaultRetryPolicy(t *testing.T) { t.Error("expected non-zero MaxAttempts in default retry policy") } } + +// TestGraph_DAG_SlowBranchMerge verifies AllPredecessor mode: merge waits for +// slow branch. Uses pregel engine (via harness.init() which sets PregelRunFunc) +// and channel-based state (pregel-compatible pattern). +func TestGraph_DAG_SlowBranchMerge(t *testing.T) { + var orderMu sync.Mutex + var order []string + slowStarted := make(chan struct{}) + fastDone := make(chan struct{}) + releaseSlow := make(chan struct{}) + mergeRan := make(chan struct{}, 1) + record := func(name string) { + orderMu.Lock() + order = append(order, name) + orderMu.Unlock() + } + + sg := NewStateGraph(map[string]any{}) + + // Start → dispatch → {fast, slow} → merge + sg.AddNode("dispatch", func(ctx context.Context, state any) (any, error) { + record("dispatch") + return map[string]any{"messages": []string{"dispatch done"}}, nil + }) + sg.AddNode("fast", func(ctx context.Context, state any) (any, error) { + record("fast") + close(fastDone) + return map[string]any{"messages": []string{"fast done"}}, nil + }) + sg.AddNode("slow", func(ctx context.Context, state any) (any, error) { + close(slowStarted) + <-releaseSlow + record("slow") + return map[string]any{"messages": []string{"slow done"}}, nil + }) + sg.AddNode("merge", func(ctx context.Context, state any) (any, error) { + select { + case mergeRan <- struct{}{}: + default: + } + record("merge") + return map[string]any{"messages": []string{"merge done"}}, nil + }) + + sg.AddEdge(Start, "dispatch") + sg.AddEdge("dispatch", "fast") + sg.AddEdge("dispatch", "slow") + sg.AddEdge("fast", "merge") + sg.AddEdge("slow", "merge") + sg.AddEdge("merge", End) + + compiled, err := sg.Compile( + WithNodeTriggerMode(types.NodeTriggerAllPredecessor), + WithRecursionLimit(10), + ) + if err != nil { + t.Fatal(err) + } + + done := make(chan error, 1) + go func() { + _, err := compiled.Invoke(context.Background(), map[string]any{}) + done <- err + }() + + select { + case <-slowStarted: + case <-time.After(time.Second): + t.Fatal("slow branch never started") + } + + select { + case <-fastDone: + case <-time.After(time.Second): + t.Fatal("fast branch never completed") + } + + select { + case <-mergeRan: + t.Fatal("BUG: merge ran before slow branch completed") + case <-time.After(50 * time.Millisecond): + } + + close(releaseSlow) + + select { + case err = <-done: + case <-time.After(time.Second): + t.Fatal("compiled graph did not finish after slow branch release") + } + + if err != nil { + t.Fatal(err) + } + + if len(order) != 4 { + t.Errorf("expected 4 nodes (dispatch+fast+slow+merge), got %d: %v", len(order), order) + } + slowIdx, mergeIdx := -1, -1 + for i, name := range order { + switch name { + case "slow": + slowIdx = i + case "merge": + mergeIdx = i + } + } + if mergeIdx < slowIdx { + t.Errorf("BUG: merge before slow branch. merge=%d, slow=%d", mergeIdx, slowIdx) + } else { + t.Log("DAG slow-branch merge: merge correctly waited for slow branch") + } +}