mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-06 11:28:38 +08:00
Fix harness DAG slow-branch test cased by nil initialization of pregel engine (#16591)
This commit is contained in:
Yingfeng
@@ -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
|
||||
// =====================================================================
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user