diff --git a/internal/harness/graph/pregel/async.go b/internal/harness/graph/pregel/async.go index 2754169476..0d8eba4ce7 100644 --- a/internal/harness/graph/pregel/async.go +++ b/internal/harness/graph/pregel/async.go @@ -93,11 +93,24 @@ func (e *AsyncExecutor) Execute(ctx context.Context, name string, fn func(contex select { case <-e.workerPool: defer func() { e.workerPool <- struct{}{} }() - case <-ctx.Done(): + case <-task.Context.Done(): resultCh <- &asyncTaskResult{ TaskID: task.ID, Name: task.Name, - Err: ctx.Err(), + Err: task.Context.Err(), + } + return + } + // Synchronize with Cancel: another cancelled task may release a worker slot + // before this queued task is cancelled. + e.mu.Lock() + taskErr := task.Context.Err() + e.mu.Unlock() + if taskErr != nil { + resultCh <- &asyncTaskResult{ + TaskID: task.ID, + Name: task.Name, + Err: taskErr, } return } @@ -180,44 +193,10 @@ func (e *AsyncExecutor) ExecuteBatch(ctx context.Context, tasks []asyncTask) <-c // ExecuteWithRetry executes a task with retry logic. func (e *AsyncExecutor) ExecuteWithRetry(ctx context.Context, name string, fn func(context.Context) (any, error), retryConfig *RetryConfig) <-chan *asyncTaskResult { - resultCh := make(chan *asyncTaskResult, 1) - - taskCtx, cancel := context.WithCancel(ctx) - task := &asyncTask{ - ID: uuid.New().String(), - Name: name, - Context: taskCtx, - Cancel: cancel, - } - - e.mu.Lock() - e.activeTasks[task.ID] = task - e.mu.Unlock() - - go func() { - defer close(resultCh) - - executor := NewRetryExecutor(retryConfig.Policy) - - startTime := time.Now() - output, err := executor.Execute(task.Context, name, fn) - - result := &asyncTaskResult{ - TaskID: task.ID, - Name: name, - Output: output, - Err: err, - Duration: time.Since(startTime), - } - - e.mu.Lock() - delete(e.activeTasks, task.ID) - e.mu.Unlock() - - resultCh <- result - }() - - return resultCh + executor := NewRetryExecutor(retryConfig.Policy) + return e.Execute(ctx, name, func(taskCtx context.Context) (any, error) { + return executor.Execute(taskCtx, name, fn) + }) } // Cancel cancels all active tasks by invoking their cancel functions. diff --git a/internal/harness/graph/pregel/async_test.go b/internal/harness/graph/pregel/async_test.go new file mode 100644 index 0000000000..54a4d1bba5 --- /dev/null +++ b/internal/harness/graph/pregel/async_test.go @@ -0,0 +1,71 @@ +package pregel + +import ( + "context" + "errors" + "sync/atomic" + "testing" + + "ragflow/internal/harness/graph/types" +) + +func TestAsyncExecutor_ExecuteWithRetryHonorsWorkerPool(t *testing.T) { + executor := NewAsyncExecutor(1) + resultCh := executor.ExecuteWithRetry( + context.Background(), + "retry-task", + func(context.Context) (any, error) { + return len(executor.workerPool), nil + }, + &RetryConfig{Policy: &types.RetryPolicy{MaxAttempts: 1}}, + ) + + result, ok := <-resultCh + if !ok || result == nil { + t.Fatal("ExecuteWithRetry() returned no result") + } + if result.Err != nil { + t.Fatalf("ExecuteWithRetry() error = %v", result.Err) + } + if availableSlots, ok := result.Output.(int); !ok || availableSlots != 0 { + t.Fatalf("available worker slots during execution = %v, want 0", result.Output) + } +} + +func TestAsyncExecutor_CancelQueuedRetryTask(t *testing.T) { + executor := NewAsyncExecutor(1) + firstStarted := make(chan struct{}) + firstResultCh := executor.Execute(context.Background(), "blocking-task", func(ctx context.Context) (any, error) { + close(firstStarted) + <-ctx.Done() + return nil, ctx.Err() + }) + <-firstStarted + + var retryCalls atomic.Int32 + retryResultCh := executor.ExecuteWithRetry( + context.Background(), + "queued-retry-task", + func(context.Context) (any, error) { + retryCalls.Add(1) + return nil, nil + }, + &RetryConfig{Policy: &types.RetryPolicy{MaxAttempts: 1}}, + ) + if active := executor.GetActiveTaskCount(); active != 2 { + t.Fatalf("active tasks before cancellation = %d, want 2", active) + } + + executor.Cancel() + <-firstResultCh + retryResult, ok := <-retryResultCh + if !ok || retryResult == nil { + t.Fatal("queued retry task returned no result") + } + if !errors.Is(retryResult.Err, context.Canceled) { + t.Fatalf("queued retry task error = %v, want context.Canceled", retryResult.Err) + } + if calls := retryCalls.Load(); calls != 0 { + t.Fatalf("queued retry task calls after cancellation = %d, want 0", calls) + } +}