From 110ed630bdfcbde422745f9040412a8652d3af33 Mon Sep 17 00:00:00 2001 From: Lem0nTea2002 Date: Mon, 17 Aug 2026 21:27:18 +0800 Subject: [PATCH] fix(harness): honor node and engine retry policies in Pregel (#18388) ### Summary Pregel nodes created with `AddNodeWithOptions` retain a node-level `RetryPolicy`, but task construction dropped that policy. When a task had no policy, execution also bypassed the engine-level policy configured with `WithRetryPolicy` and instantiated the default policy directly. --- internal/harness/graph/pregel/engine.go | 46 +++++++------- internal/harness/graph/pregel/engine_test.go | 21 +++++++ .../pregel/pregel_retry_comprehensive_test.go | 63 +++++++++++++++++++ 3 files changed, 108 insertions(+), 22 deletions(-) diff --git a/internal/harness/graph/pregel/engine.go b/internal/harness/graph/pregel/engine.go index 4a71c8d413..f970c6c1f9 100644 --- a/internal/harness/graph/pregel/engine.go +++ b/internal/harness/graph/pregel/engine.go @@ -1161,12 +1161,7 @@ func (e *Engine) executeTasksAsync( return t.Func(ctx, convertedInput) } - // Use task's retry policy or default - retryPolicy := t.RetryPolicy - if retryPolicy == nil { - defaultPolicy := types.DefaultRetryPolicy() - retryPolicy = &defaultPolicy - } + retryPolicy := e.resolveRetryPolicy(t) // Execute with async pipeline resultCh := asyncPipeline.ExecuteNode(ctx, t.Name, executeFn, &RetryConfig{Policy: retryPolicy}) @@ -1239,13 +1234,7 @@ func (e *Engine) executeTask( input = e.mapToStateSchema(input) // Use RetryExecutor for retry logic - retryPolicy := task.RetryPolicy - if retryPolicy == nil { - defaultPolicy := types.DefaultRetryPolicy() - retryPolicy = &defaultPolicy - } - - retryExecutor := NewRetryExecutor(retryPolicy) + retryExecutor := NewRetryExecutor(e.resolveRetryPolicy(task)) // Define the function to execute executeFn := func(ctx context.Context) (any, error) { @@ -1546,12 +1535,24 @@ func (e *Engine) getTriggers(node *types.Node) []string { return node.Triggers } +func (e *Engine) resolveRetryPolicy(task *Task) *types.RetryPolicy { + if task.RetryPolicy != nil { + return task.RetryPolicy + } + if e.retryPolicy != nil { + return e.retryPolicy + } + defaultPolicy := types.DefaultRetryPolicy() + return &defaultPolicy +} + func (e *Engine) createTask(node *types.Node, state any, channels []string, triggers []string) *Task { task := &Task{ - ID: uuid.New().String(), - Name: node.Name, - Channels: channels, - Triggers: make(map[string]struct{}), + ID: uuid.New().String(), + Name: node.Name, + Channels: channels, + Triggers: make(map[string]struct{}), + RetryPolicy: node.RetryPolicy, } if node.Function != nil { task.Func = node.Function @@ -1566,11 +1567,12 @@ func (e *Engine) createTask(node *types.Node, state any, channels []string, trig // This is similar to Python's prepare_next_tasks with for_execution=False. func (e *Engine) createTaskInfo(node *types.Node, state any, channels []string, triggers []string) *Task { task := &Task{ - ID: uuid.New().String(), - Name: node.Name, - Channels: channels, - Triggers: make(map[string]struct{}), - Func: nil, + ID: uuid.New().String(), + Name: node.Name, + Channels: channels, + Triggers: make(map[string]struct{}), + RetryPolicy: node.RetryPolicy, + Func: nil, } for _, trigger := range triggers { task.Triggers[trigger] = struct{}{} diff --git a/internal/harness/graph/pregel/engine_test.go b/internal/harness/graph/pregel/engine_test.go index 58ae1e7e25..ee62919049 100644 --- a/internal/harness/graph/pregel/engine_test.go +++ b/internal/harness/graph/pregel/engine_test.go @@ -64,6 +64,27 @@ func TestNewEngine(t *testing.T) { } } +func TestEngine_TaskConstructionPreservesNodeRetryPolicy(t *testing.T) { + policy := &types.RetryPolicy{MaxAttempts: 1} + sg := graph.NewStateGraph(map[string]any{"value": ""}) + node := sg.AddNodeWithOptions("work", func(_ context.Context, state any) (any, error) { + return state, nil + }, types.NodeOptions{RetryPolicy: policy}) + engine := NewEngine(sg) + + tests := map[string]*Task{ + "execution": engine.createTask(node, nil, nil, nil), + "inspection": engine.createTaskInfo(node, nil, nil, nil), + } + for name, task := range tests { + t.Run(name, func(t *testing.T) { + if task.RetryPolicy != policy { + t.Fatalf("RetryPolicy = %p, want %p", task.RetryPolicy, policy) + } + }) + } +} + func TestEngine_RunSync(t *testing.T) { sg := newSimpleGraph(t) engine := NewEngine(sg, WithRecursionLimit(10)) diff --git a/internal/harness/graph/pregel/pregel_retry_comprehensive_test.go b/internal/harness/graph/pregel/pregel_retry_comprehensive_test.go index c1aa9f46e1..f051bb0d7f 100644 --- a/internal/harness/graph/pregel/pregel_retry_comprehensive_test.go +++ b/internal/harness/graph/pregel/pregel_retry_comprehensive_test.go @@ -53,6 +53,69 @@ func TestRetry_BackoffTiming(t *testing.T) { } } +func TestRetry_NodePolicyOverridesEnginePolicy(t *testing.T) { + var attempts atomic.Int32 + sg := graphPkg.NewStateGraph(map[string]any{}) + sg.AddChannel("value", channels.NewLastValue("")) + + nodePolicy := types.RetryPolicy{ + InitialInterval: 0, + BackoffFactor: 1, + MaxInterval: 0, + MaxAttempts: 1, + Jitter: false, + RetryOn: func(error) bool { return true }, + } + sg.AddNodeWithOptions("work", func(context.Context, any) (any, error) { + attempts.Add(1) + return nil, fmt.Errorf("fail") + }, types.NodeOptions{RetryPolicy: &nodePolicy}) + _ = sg.AddEdge(constants.Start, "work") + _ = sg.AddEdge("work", constants.End) + + enginePolicy := nodePolicy + enginePolicy.MaxAttempts = 3 + engine := NewEngine(sg, WithRecursionLimit(10), WithRetryPolicy(&enginePolicy)) + + _, err := engine.RunSync(context.Background(), map[string]any{"value": "x"}) + if err == nil { + t.Fatal("expected error") + } + if got := attempts.Load(); got != 1 { + t.Fatalf("expected node retry policy to stop after 1 attempt, got %d", got) + } +} + +func TestRetry_EnginePolicyFallback(t *testing.T) { + var attempts atomic.Int32 + sg := graphPkg.NewStateGraph(map[string]any{}) + sg.AddChannel("value", channels.NewLastValue("")) + sg.AddNode("work", func(context.Context, any) (any, error) { + attempts.Add(1) + return nil, fmt.Errorf("fail") + }) + _ = sg.AddEdge(constants.Start, "work") + _ = sg.AddEdge("work", constants.End) + + enginePolicy := types.RetryPolicy{ + InitialInterval: 0, + BackoffFactor: 1, + MaxInterval: 0, + MaxAttempts: 2, + Jitter: false, + RetryOn: func(error) bool { return true }, + } + engine := NewEngine(sg, WithRecursionLimit(10), WithRetryPolicy(&enginePolicy)) + + _, err := engine.RunSync(context.Background(), map[string]any{"value": "x"}) + if err == nil { + t.Fatal("expected error") + } + if got := attempts.Load(); got != 2 { + t.Fatalf("expected engine retry policy to stop after 2 attempts, got %d", got) + } +} + // ============================================================ // P0: Retry with jitter produces varying times // ============================================================