diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 177626aec0..e6ed31031d 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -551,10 +551,10 @@ func runIngestor(args *serverArgs) error { } // Create context with timeout for graceful shutdown - _, cancel := context.WithTimeout(context.Background(), 30*time.Second) + shutdownCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - ingestor.Stop() + ingestor.Stop(shutdownCtx) common.Info(fmt.Sprintf("Ingestor %s shutdown complete", *args.name)) diff --git a/internal/entity/ingestion_task.go b/internal/entity/ingestion_task.go index 604388d5e9..9746368d70 100644 --- a/internal/entity/ingestion_task.go +++ b/internal/entity/ingestion_task.go @@ -37,13 +37,12 @@ func (IngestionTask) TableName() string { } type IngestionTaskLog struct { - ID int `gorm:"column:id;primaryKey;autoIncrement" json:"id"` - TaskID string `gorm:"column:task_id;size:32;not null;index" json:"task_id"` - Checkpoint JSONMap `gorm:"column:checkpoint;type:longtext;not null" json:"checkpoint"` - ComponentIndex int `gorm:"column:component_index" json:"component_index"` - Phase int `gorm:"column:phase" json:"phase"` - Component string `gorm:"column:component;size:64;index" json:"component"` - Message string `gorm:"column:message;type:text" json:"message"` + ID int `gorm:"column:id;primaryKey;autoIncrement" json:"id"` + TaskID string `gorm:"column:task_id;size:32;not null;index" json:"task_id"` + Checkpoint JSONMap `gorm:"column:checkpoint;type:longtext;not null" json:"checkpoint"` + Phase int `gorm:"column:phase" json:"phase"` + Component string `gorm:"column:component;size:64;index" json:"component"` + Message string `gorm:"column:message;type:text" json:"message"` BaseModel } diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index 7244352401..f86d87910f 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -94,8 +94,10 @@ import ( "encoding/json" "fmt" "log" + "os" "regexp" "slices" + "strconv" "strings" "time" @@ -107,11 +109,20 @@ import ( const ComponentNameTokenizer = "Tokenizer" -// tokenizerTimeout bounds the batched embedding call. Mirrors the -// python `@timeout(60)` decorator on `Tokenizer._embedding.embed_limiter` -// + `batch_encode` in tokenizer.py:92-104. Declared as a var so tests -// can shrink it; production wiring uses 60s. -var tokenizerTimeout = 60 * time.Second +// tokenizerTimeout returns the per-batch timeout for embedding API calls. +// Reads COMPONENT_EXEC_TIMEOUT_TOKENIZER env var (seconds); defaults to 600s +// (10 min) to match the canvas-level component timeout default. +// Invalid / non-positive values fall back to the default. +func tokenizerTimeout() time.Duration { + if v := os.Getenv("COMPONENT_EXEC_TIMEOUT_TOKENIZER"); v != "" { + if secs, err := strconv.Atoi(v); err == nil && secs > 0 { + return time.Duration(secs) * time.Second + } + } + return defaultTokenizerTimeout +} + +var defaultTokenizerTimeout = 600 * time.Second // tokenizerEmbeddingBatchSize mirrors Python's // settings.EMBEDDING_BATCH_SIZE default. @@ -461,7 +472,7 @@ func encodeWithTimeout(ctx context.Context, embedder Embedder, texts []string) ( results []EmbeddingResult encErr error ) - timeoutErr := runtime.WithTimeout(ctx, tokenizerTimeout, func(timeoutCtx context.Context) error { + timeoutErr := runtime.WithTimeout(ctx, tokenizerTimeout(), func(timeoutCtx context.Context) error { results, encErr = embedder.Encode(texts) return encErr }) diff --git a/internal/ingestion/component/tokenizer_test.go b/internal/ingestion/component/tokenizer_test.go index b3cf2467d8..486cf2b4a8 100644 --- a/internal/ingestion/component/tokenizer_test.go +++ b/internal/ingestion/component/tokenizer_test.go @@ -547,12 +547,10 @@ func (c *countMismatchedEmbedder) Encode(texts []string) ([]EmbeddingResult, err // asserts the component returns context.DeadlineExceeded. func TestTokenizerComponent_Invoke_HonorsTimeout(t *testing.T) { requireTokenizerPool(t) - prevTimeout := tokenizerTimeout - tokenizerTimeout = 50 * time.Millisecond - t.Cleanup(func() { tokenizerTimeout = prevTimeout }) + t.Setenv("COMPONENT_EXEC_TIMEOUT_TOKENIZER", "1") c, stub := withStubEmbedder(t, 4) - stub.delay = 500 * time.Millisecond + stub.delay = 2 * time.Second ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() diff --git a/internal/ingestion/pipeline/pipeline.go b/internal/ingestion/pipeline/pipeline.go index 4a7695ba51..243890f816 100644 --- a/internal/ingestion/pipeline/pipeline.go +++ b/internal/ingestion/pipeline/pipeline.go @@ -21,7 +21,6 @@ import ( "encoding/json" "errors" "fmt" - "sort" "time" "ragflow/internal/agent/canvas" @@ -99,17 +98,15 @@ func WithDocumentID(docID string) PipelineOption { } // ProgressEvent is a structured component lifecycle event emitted by the -// pipeline to a ProgressSink. The pipeline fills every field - including the -// ingestion-proprietary Index (from its componentIndexMap) and the Total -// denominator - so the sink needs no canvas knowledge to persist it. +// pipeline to a ProgressSink. The pipeline fills the task/document/component +// identity and phase/status message; the sink caches the denominator (total) +// from OnComponentTotal and needs no canvas knowledge. type ProgressEvent struct { TaskID string DocumentID string Component string Message string - Index int Phase int - Total int } // ProgressSink receives pipeline progress for durable persistence. It is the @@ -294,7 +291,7 @@ func (p *Pipeline) Run(ctx context.Context, inputs map[string]any, setups ...map // is nil when the DB is not initialized (unit tests, headless // runs), in which case TrackProgress is a no-op — progress is an // observability concern, not a data dependency. - runCtx = runtime.WithProgressCallback(runCtx, p.taskLogProgressCallback()) + runCtx = runtime.WithProgressCallback(runCtx, p.componentProgressCallback()) current := cloneMapOrEmpty(inputs) @@ -459,37 +456,17 @@ func finalizeResult(current, out map[string]any, runState *canvas.CanvasState) m return merged } -// componentIndexMap builds a deterministic cpnID → 0-based-index map for -// the task's canvas. Map iteration order is non-deterministic in Go, so -// the cpnIDs are sorted to keep the index stable across runs. The index is -// ingestion-proprietary and computed here, then carried on the pipeline-local -// ProgressEvent so the sink needs no canvas knowledge. -func (p *Pipeline) componentIndexMap() map[string]int { - ids := make([]string, 0, len(p.canvas.Components)) - for id := range p.canvas.Components { - ids = append(ids, id) - } - sort.Strings(ids) - m := make(map[string]int, len(ids)) - for i, id := range ids { - m[id] = i - } - return m -} - -// taskLogProgressCallback returns a runtime.ProgressCallback that forwards +// componentProgressCallback returns a runtime.ProgressCallback that forwards // every component lifecycle event (start/done/fail) to the pipeline's // ProgressSink. The sink owns all persistence; this callback only shapes the -// event - deriving the message string the frontend expects and the -// ingestion-proprietary component index - so the pipeline never touches the -// DAO layer. Returns nil when no sink is attached, leaving TrackProgress a -// no-op and the pipeline DB-independent (unit tests, headless runs). -func (p *Pipeline) taskLogProgressCallback() runtime.ProgressCallback { +// event - deriving the message string the frontend expects - so the pipeline +// never touches the DAO layer. Returns nil when no sink is attached, leaving +// TrackProgress a no-op and the pipeline DB-independent (unit tests, headless +// runs). +func (p *Pipeline) componentProgressCallback() runtime.ProgressCallback { if p.sink == nil { return nil } - indexMap := p.componentIndexMap() - total := len(p.canvas.Components) return func(ev runtime.ProgressEvent) { var msg string switch ev.Phase { @@ -509,9 +486,7 @@ func (p *Pipeline) taskLogProgressCallback() runtime.ProgressCallback { DocumentID: p.documentID, Component: ev.Component, Message: msg, - Index: indexMap[ev.Component], Phase: int(ev.Phase), - Total: total, }) } } diff --git a/internal/ingestion/pipeline/pipeline_test.go b/internal/ingestion/pipeline/pipeline_test.go index fadd2c1f19..fd0a2f4fab 100644 --- a/internal/ingestion/pipeline/pipeline_test.go +++ b/internal/ingestion/pipeline/pipeline_test.go @@ -411,9 +411,8 @@ func (r *recordingSink) OnComponentProgress(ev ProgressEvent) { } // TestPipelineRunForwardsProgressToSink verifies the pipeline reports the -// component-total denominator and each component lifecycle event to the -// injected ProgressSink, and carries task/document/total context on every -// event so the sink needs no canvas knowledge. +// component-total denominator once via OnComponentTotal and each component +// lifecycle event to the injected ProgressSink. func TestPipelineRunForwardsProgressToSink(t *testing.T) { stageA := &mockCanvasStage{output: map[string]any{"a": 1}} stageB := &mockCanvasStage{output: map[string]any{"b": 2}} @@ -463,9 +462,6 @@ func TestPipelineRunForwardsProgressToSink(t *testing.T) { if ev.DocumentID != "doc-sink" { t.Fatalf("event DocumentID = %q, want doc-sink", ev.DocumentID) } - if ev.Total != 3 { - t.Fatalf("event Total = %d, want 3", ev.Total) - } seen[ev.Component] = true } for _, want := range []string{"a", "b"} { diff --git a/internal/ingestion/service/execute_task_test.go b/internal/ingestion/service/execute_task_test.go index c9d5056509..9937ec51c0 100644 --- a/internal/ingestion/service/execute_task_test.go +++ b/internal/ingestion/service/execute_task_test.go @@ -95,7 +95,7 @@ func TestDefaultRunDocumentTask_RequiresConfiguredPipelineID(t *testing.T) { } } -func TestExecuteTask_PipelineRoutesToTaskHandler(t *testing.T) { +func TestExecuteTask_RunsDocumentTask(t *testing.T) { db := testutil.SetupTestDB(t) cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() @@ -106,12 +106,12 @@ func TestExecuteTask_PipelineRoutesToTaskHandler(t *testing.T) { ) ingestor := NewIngestor("test", 1, []string{"pdf"}) - var routedToPipeline bool + var runDocumentTaskCalled bool var gotTaskID string var gotProgress []float64 var gotMsgs []string ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error { - routedToPipeline = true + runDocumentTaskCalled = true gotTaskID = ingestionTask.ID wrapped := func(prog float64, msg string) { gotProgress = append(gotProgress, prog*100) @@ -129,8 +129,8 @@ func TestExecuteTask_PipelineRoutesToTaskHandler(t *testing.T) { ingestor.executeTask(taskCtx) - if !routedToPipeline { - t.Fatal("expected executeTask to route pipeline task to runDocumentTask") + if !runDocumentTaskCalled { + t.Fatal("expected executeTask to run runDocumentTask") } if gotTaskID != taskID { t.Fatalf("runDocumentTask got task ID %q, want %q", gotTaskID, taskID) diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index bf12c629b0..e86b59757a 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -21,7 +21,9 @@ import ( "errors" "fmt" "ragflow/internal/utility" + "strings" "sync" + "sync/atomic" "time" "ragflow/internal/common" @@ -34,6 +36,8 @@ import ( "github.com/cenkalti/backoff/v5" ) +const defaultHeartbeatInterval = 10 * time.Second + type Ingestor struct { id string name string @@ -47,8 +51,9 @@ type Ingestor struct { heartbeatInterval time.Duration // Runtime state - currentTasks map[string]*taskpkg.TaskContext - tasksMu sync.RWMutex + currentTasks map[string]struct{} // set of task IDs currently claimed by a worker + tasksMu sync.RWMutex + activeWorkers atomic.Int32 // number of worker goroutines currently in workerLoop // Shutdown channel - receive on this to trigger graceful shutdown ShutdownCh chan struct{} @@ -85,12 +90,12 @@ func NewIngestor(name string, maxConcurrency int32, supportedTypes []string) *In maxConcurrency: maxConcurrency, supportedDocTypes: supportedTypes, version: "1.0.0", - currentTasks: make(map[string]*taskpkg.TaskContext), + currentTasks: make(map[string]struct{}), taskChan: make(chan *taskpkg.TaskContext, maxConcurrency*2), ShutdownCh: make(chan struct{}, 1), ingestionTaskSvc: servicepkg.NewIngestionTaskService(), docState: newDocStateUpdater(), - heartbeatInterval: 10 * time.Second, + heartbeatInterval: defaultHeartbeatInterval, } ingestor.runDocumentTask = ingestor.defaultRunDocumentTask ingestor.cancelCheck = ingestor.defaultCancelCheck @@ -101,6 +106,11 @@ func (e *Ingestor) ID() string { return e.id } +// consumeErrorBackoff paces the consume loop when GetMessages returns an +// error, so a persistent MQ failure does not pin a CPU. The backoff is +// cancellable so a shutdown during backoff returns promptly. +const consumeErrorBackoff = 1 * time.Second + func (e *Ingestor) Start() error { common.Info(fmt.Sprintf("Ingestor %s initialized", e.id)) msgQueueEngine := engine.GetMessageQueueEngine() @@ -113,34 +123,57 @@ func (e *Ingestor) Start() error { go e.startWorkerPool() for { - var taskHandles []common.TaskHandle - taskHandles, err = msgQueueEngine.GetMessages(4) + // Graceful shutdown is the only condition under which the consume + // loop exits. Per-message processing failures never terminate the + // consumer: processMessage settles (ack/nack) each message itself. + if err := e.ctx.Err(); err != nil { + return nil + } + taskHandles, err := msgQueueEngine.GetMessages(4) if err != nil { common.Error("error consuming message", err) + select { + case <-time.After(consumeErrorBackoff): + case <-e.ctx.Done(): + return nil + } continue } for _, taskHandle := range taskHandles { - if err := e.processMessage(taskHandle); err != nil { - return err - } + e.processMessage(taskHandle) } } } // processMessage handles a single incoming MQ message: filter by type, // activate the task (state transition), guard against duplicate execution -// (claim), and enqueue to the worker pool (or backpressure-reject). -func (e *Ingestor) processMessage(handle common.TaskHandle) error { +// (claim), and enqueue to the worker pool (or backpressure-reject). It +// settles (ack/nack) every message itself and never returns an error: a +// single message can never terminate the consume loop. Only ctx cancellation +// (graceful shutdown) stops the consumer - see Start. +func (e *Ingestor) processMessage(handle common.TaskHandle) { taskMessage := handle.GetMessage() common.Info(fmt.Sprintf("Received task id: %s, type: %s", taskMessage.TaskID, taskMessage.TaskType)) + // Deferred claim release: if this function claims a task but the task + // is not successfully enqueued to the worker pool (e.g. backpressure, + // or a future error path added between claim and enqueue), the defer + // cleans up so the task can be reclaimed on MQ redelivery. When the + // task IS enqueued, claimedTaskID is cleared and executeTask's own + // defer takes ownership of the release. + var claimedTaskID string + defer func() { + if claimedTaskID != "" { + e.releaseTask(claimedTaskID) + } + }() + if taskMessage.TaskType != common.TaskTypeIngestionTask { common.Info(fmt.Sprintf("task %s is not an ingestion task", taskMessage.TaskID)) if err := handle.Ack(); err != nil { common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) - return err } - return nil + return } task, err := e.ingestionTaskSvc.StartRunning(taskMessage.TaskID) @@ -149,19 +182,24 @@ func (e *Ingestor) processMessage(handle common.TaskHandle) error { common.Warn(fmt.Sprintf("task %s not found, skipping", taskMessage.TaskID)) if ackErr := handle.Ack(); ackErr != nil { common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), ackErr) - return ackErr } - return nil + return } + // Recoverable activation failure (e.g. a DB blip): nack for + // redelivery instead of dropping the message or killing the + // consumer. common.Error(fmt.Sprintf("error setting task %s to running", taskMessage.TaskID), err) - return err + if nackErr := handle.Nack(); nackErr != nil { + common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), nackErr) + } + return } if task == nil { common.Info(fmt.Sprintf("task %s is already removed", taskMessage.TaskID)) if ackErr := handle.Ack(); ackErr != nil { - return ackErr + common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), ackErr) } - return nil + return } switch task.Status { @@ -169,11 +207,8 @@ func (e *Ingestor) processMessage(handle common.TaskHandle) error { common.Info(fmt.Sprintf("task %s is already %s", taskMessage.TaskID, task.Status)) if ackErr := handle.Ack(); ackErr != nil { common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), ackErr) - return ackErr } - return nil - case common.STOPPING, common.CREATED: - return fmt.Errorf("task %s is in unexpected status %s", taskMessage.TaskID, task.Status) + return case common.RUNNING: // Guard against MQ redelivery: if another worker in this // process is already processing this task, ack the redelivered @@ -183,10 +218,19 @@ func (e *Ingestor) processMessage(handle common.TaskHandle) error { taskMessage.TaskID, task.ID, task.DocumentID, task.DatasetID)) if ackErr := handle.Ack(); ackErr != nil { common.Error(fmt.Sprintf("error ack redelivered task %s", taskMessage.TaskID), ackErr) - return ackErr } - return nil + return } + claimedTaskID = task.ID + default: + // Unreachable given StartRunning normalizes every status to + // RUNNING/COMPLETED/STOPPED/FAILED, but defensive: ack-skip an + // unknown status instead of enqueuing it for execution. + common.Warn(fmt.Sprintf("task %s in unexpected status %s, ack-skip", taskMessage.TaskID, task.Status)) + if ackErr := handle.Ack(); ackErr != nil { + common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), ackErr) + } + return } // Construct TaskContext and carry the MQ handle so the worker can @@ -197,16 +241,15 @@ func (e *Ingestor) processMessage(handle common.TaskHandle) error { // Push to task channel; if full, reject the task (backpressure). select { case e.taskChan <- taskCtx: + claimedTaskID = "" // executeTask owns the release now common.Info(fmt.Sprintf("Task %s queued (channel: %d/%d)", task.ID, len(e.taskChan), cap(e.taskChan))) default: common.Info(fmt.Sprintf("No available slot for task %s, failed", task.ID)) - e.releaseTask(task.ID) + // claimedTaskID is still set; defer will call releaseTask. if nackErr := handle.Nack(); nackErr != nil { common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), nackErr) - return nackErr } } - return nil } func (e *Ingestor) startWorkerPool() { @@ -221,6 +264,8 @@ func (e *Ingestor) startWorkerPool() { func (e *Ingestor) workerLoop(id int32) { defer e.workerWg.Done() + defer e.activeWorkers.Add(-1) + e.activeWorkers.Add(1) common.Info(fmt.Sprintf("Worker %d started", id)) for { select { @@ -331,10 +376,7 @@ func (e *Ingestor) runTask(ctx context.Context, task *entity.IngestionTask) bool return e.markFailed(task.ID) } - _, err := backoff.Retry(ctx, func() (struct{}, error) { - return struct{}{}, e.ingestionTaskSvc.MarkCompleted(task.ID) - }, backoff.WithMaxTries(3)) - if err != nil { + if err := e.completeTask(ctx, task.ID); err != nil { common.Error(fmt.Sprintf("Task %s update status failed", task.ID), err) return false } @@ -343,6 +385,67 @@ func (e *Ingestor) runTask(ctx context.Context, task *entity.IngestionTask) bool return true } +// completeTask persists the task's terminal status after a successful pipeline. +// MarkCompleted is retried with backoff for transient (DB) failures only. A +// terminal transition failure - the task is no longer RUNNING because a +// concurrent stop (or another worker) moved it - is NOT retried: the pipeline +// already did the work, so completeOrSettle settles the task to its actual +// terminal state and the caller Acks instead of redelivering. +func (e *Ingestor) completeTask(ctx context.Context, taskID string) error { + _, err := backoff.Retry(ctx, func() (struct{}, error) { + return struct{}{}, e.completeOrSettle(taskID) + }, backoff.WithMaxTries(3)) + return err +} + +// completeOrSettle marks the task COMPLETED, or - if the transition is +// terminally invalid because the task is no longer RUNNING - settles it to its +// actual terminal state. Returns nil once the task is in any terminal state; +// returns a non-terminal (transient) error only for retry-worthy DB failures. +func (e *Ingestor) completeOrSettle(taskID string) error { + if err := e.ingestionTaskSvc.MarkCompleted(taskID); err != nil { + if isTerminalTransitionError(err) { + return e.settleToTerminal(taskID) + } + return err + } + return nil +} + +// isTerminalTransitionError reports whether err is a state-machine transition +// failure - an invalid transition or a lost optimistic CAS - meaning the task's +// status moved on and MarkCompleted will never succeed as-is. Not retry-worthy; +// the caller settles by the task's current status. +func isTerminalTransitionError(err error) bool { + var ite *servicepkg.InvalidTaskTransitionError + var tce *servicepkg.TaskStatusConflictError + return errors.As(err, &ite) || errors.As(err, &tce) +} + +// settleToTerminal finalizes a task whose MarkCompleted failed because it was +// no longer RUNNING. STOPPING is moved to STOPPED via markStopped (which also +// clears the Redis cancel flag so a future retry does not immediately +// re-cancel); already-terminal states (COMPLETED/STOPPED/FAILED) need no +// action. An unexpected status returns an error so the caller nacks and +// redelivery settles it. +func (e *Ingestor) settleToTerminal(taskID string) error { + task, err := e.ingestionTaskSvc.GetTask(taskID) + if err != nil { + return err + } + switch task.Status { + case common.STOPPING: + if !e.markStopped(taskID) { + return fmt.Errorf("task %s: settle to STOPPED failed", taskID) + } + return nil + case common.COMPLETED, common.STOPPED, common.FAILED: + return nil + default: + return fmt.Errorf("task %s in unexpected status %s after transition failure", taskID, task.Status) + } +} + // settleMessage runs body under a heartbeat, then settles the MQ message. The // heartbeat is stopped (and waited on) before ack/nack — see startHeartbeat. // terminal is derived from body's return value; on panic terminal defaults to @@ -500,7 +603,7 @@ func (e *Ingestor) claimTask(taskID string) bool { if _, ok := e.currentTasks[taskID]; ok { return false } - e.currentTasks[taskID] = nil // placeholder; replaced after scheduling + e.currentTasks[taskID] = struct{}{} return true } @@ -524,16 +627,11 @@ func (e *Ingestor) defaultRunDocumentTask(ctx context.Context, ingestionTask *en // The sink owns all document/ingestion_task_log/ingestion_task.component_total // writes for this run; inject it into the executor so the pipeline reports // progress to the service layer instead of touching the DAO directly. - sink := newProgressSink(e.ingestionTaskSvc) - result, err := taskpkg.NewTaskHandler(docTaskCtx). - WithPipelineExecutorFactory(func(c *taskpkg.TaskContext, canvasID string) (*taskpkg.PipelineExecutor, error) { - ex, err := taskpkg.NewPipelineExecutor(c, canvasID, 0) - if err != nil { - return nil, err - } - return ex.WithProgressSink(sink), nil - }). - Handle() + executor, err := taskpkg.NewPipelineExecutor(docTaskCtx, strings.TrimSpace(docTaskCtx.PipelineID), 0) + if err != nil { + return err + } + result, err := executor.WithProgressSink(newProgressSink(e.ingestionTaskSvc)).Execute(docTaskCtx.Ctx) if err != nil { return err } @@ -541,12 +639,33 @@ func (e *Ingestor) defaultRunDocumentTask(ctx context.Context, ingestionTask *en return nil } -// Stop gracefully shuts down the ingestor -func (e *Ingestor) Stop() { +// Stop gracefully shuts down the ingestor. It cancels the root context so +// idle workers exit immediately and in-flight pipelines abort at their next +// ctx.Err() check, then waits for workers to return. The wait is bounded by +// ctx: a stage that does not honor cancellation (e.g. a native CGO parse) +// would otherwise block workerWg.Wait() indefinitely; when ctx expires Stop +// returns and leaves the broker to redeliver any in-flight messages +// (at-least-once). Callers must pass a deadline-bearing context. +func (e *Ingestor) Stop(ctx context.Context) { common.Info(fmt.Sprintf("Stopping ingestor %s", e.id)) e.cancel() - // Wait for all workers to finish (they exit on ctx.Done()) - e.workerWg.Wait() - common.Info("All tasks completed") + waitDone := make(chan struct{}) + go func() { + e.workerWg.Wait() + close(waitDone) + }() + + select { + case <-waitDone: + common.Info("All tasks completed") + case <-ctx.Done(): + e.tasksMu.RLock() + ids := make([]string, 0, len(e.currentTasks)) + for id := range e.currentTasks { + ids = append(ids, id) + } + e.tasksMu.RUnlock() + common.Warn(fmt.Sprintf("Stop timed out with %d task(s) still in-flight (will be redelivered by broker): %v", len(ids), ids)) + } } diff --git a/internal/ingestion/service/ingestor_lifecycle_test.go b/internal/ingestion/service/ingestor_lifecycle_test.go index bdd1baf2b1..83f047efa6 100644 --- a/internal/ingestion/service/ingestor_lifecycle_test.go +++ b/internal/ingestion/service/ingestor_lifecycle_test.go @@ -17,19 +17,47 @@ package service import ( + "context" "testing" "time" + + "ragflow/internal/common" + "ragflow/internal/entity" + taskpkg "ragflow/internal/ingestion/task" + "ragflow/internal/ingestion/testutil" ) // TestStartWorkerPool_StartOnceIdempotent verifies that calling startWorkerPool -// twice only starts maxConcurrency workers (sync.Once gate). +// twice only starts maxConcurrency workers (sync.Once gate). It observes the +// active worker count directly: a broken sync.Once would double the worker +// pool and activeWorkers would exceed concurrency after the second call. func TestStartWorkerPool_StartOnceIdempotent(t *testing.T) { const concurrency int32 = 3 ingestor := NewIngestor("test-idempotent", concurrency, nil) - // Stop background worker loops immediately so we can count workers. + + ingestor.startWorkerPool() + // Wait for all workers to enter their loop (they block on the select + // since ctx is not cancelled and no tasks are queued). + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if ingestor.activeWorkers.Load() == concurrency { + break + } + time.Sleep(time.Millisecond) + } + if got := ingestor.activeWorkers.Load(); got != concurrency { + t.Fatalf("activeWorkers after first startWorkerPool = %d, want %d", got, concurrency) + } + + // Calling again must not start additional workers (sync.Once gate). + ingestor.startWorkerPool() + // Allow any erroneously-started workers to register, then re-check. + time.Sleep(50 * time.Millisecond) + if got := ingestor.activeWorkers.Load(); got != concurrency { + t.Fatalf("activeWorkers after second startWorkerPool = %d, want %d (sync.Once not idempotent)", got, concurrency) + } + ingestor.cancel() - ingestor.startWorkerPool() - ingestor.startWorkerPool() ingestor.workerWg.Wait() } @@ -44,7 +72,7 @@ func TestStop_GracefulShutdown(t *testing.T) { done := make(chan struct{}) go func() { - ingestor.Stop() + ingestor.Stop(context.Background()) close(done) }() @@ -55,3 +83,65 @@ func TestStop_GracefulShutdown(t *testing.T) { t.Fatal("Stop() timed out waiting for workers to exit") } } + +// TestStop_TimesOutWhenWorkerStuck verifies the B1 fix: when a worker is +// blocked in a stage that does not honor ctx cancellation (e.g. a native +// CGO parse), Stop returns once its deadline expires instead of hanging on +// workerWg.Wait() forever. The in-flight task is left for broker redelivery. +func TestStop_TimesOutWhenWorkerStuck(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + const concurrency int32 = 1 + ingestor := NewIngestor("test-stuck", concurrency, []string{"pdf"}) + ingestor.startWorkerPool() + + // runDocumentTask blocks on release and ignores ctx, simulating a + // non-cancellable native parse. started signals the worker is inside it. + release := make(chan struct{}) + started := make(chan struct{}) + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + close(started) + <-release + return nil + } + + // Seed the task RUNNING so runTask's MarkCompleted path is valid. + if err := db.Model(&entity.IngestionTask{}).Where("id = ?", taskID). + Update("status", common.RUNNING).Error; err != nil { + t.Fatalf("set task RUNNING: %v", err) + } + + taskCtx := taskpkg.NewTaskContextForScheduling(ingestor.ctx, &entity.IngestionTask{ + ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING, + }) + ingestor.taskChan <- taskCtx + + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("worker did not enter runDocumentTask") + } + + // Stop with a short deadline must return instead of hanging. + stopDone := make(chan struct{}) + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + ingestor.Stop(ctx) + close(stopDone) + }() + + select { + case <-stopDone: + // Stop returned within the deadline - the fix works. + case <-time.After(3 * time.Second): + t.Fatal("Stop() hung instead of returning on deadline") + } + + // Release the stuck worker so it finishes and the test goroutine stays clean. + close(release) + ingestor.workerWg.Wait() +} diff --git a/internal/ingestion/service/process_message_test.go b/internal/ingestion/service/process_message_test.go index 8b302c3df4..44ec2d7e47 100644 --- a/internal/ingestion/service/process_message_test.go +++ b/internal/ingestion/service/process_message_test.go @@ -23,10 +23,7 @@ func TestProcessMessage_NonIngestionTaskAcks(t *testing.T) { ingestor := NewIngestor("test", 1, []string{"pdf"}) handle := newFakeHandle("task-1", "not-ingestion") - err := ingestor.processMessage(handle) - if err != nil { - t.Fatalf("expected nil (continue), got: %v", err) - } + ingestor.processMessage(handle) if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { t.Fatalf("non-ingestion: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) } @@ -46,10 +43,7 @@ func TestProcessMessage_TaskNotFoundAcks(t *testing.T) { // No task seeded in DB — StartRunning returns ErrTaskNotFound. handle := newFakeHandle("no-such-task", common.TaskTypeIngestionTask) - err := ingestor.processMessage(handle) - if err != nil { - t.Fatalf("expected nil (continue), got: %v", err) - } + ingestor.processMessage(handle) if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { t.Fatalf("not-found: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) } @@ -72,10 +66,7 @@ func TestProcessMessage_AlreadyCompletedAcks(t *testing.T) { ingestor := NewIngestor("test", 1, []string{"pdf"}) handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) - err := ingestor.processMessage(handle) - if err != nil { - t.Fatalf("expected nil (continue), got: %v", err) - } + ingestor.processMessage(handle) if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { t.Fatalf("completed: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) } @@ -99,10 +90,7 @@ func TestProcessMessage_ClaimFailsAcks(t *testing.T) { handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) - err := ingestor.processMessage(handle) - if err != nil { - t.Fatalf("expected nil (continue), got: %v", err) - } + ingestor.processMessage(handle) if handle.acks.Load() != 1 || handle.nacks.Load() != 0 { t.Fatalf("claim-fail: expected 1 Ack/0 Nack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) } @@ -123,10 +111,7 @@ func TestProcessMessage_ClaimSucceedsEnqueues(t *testing.T) { ingestor := NewIngestor("test", 1, []string{"pdf"}) handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) - err := ingestor.processMessage(handle) - if err != nil { - t.Fatalf("expected nil (continue), got: %v", err) - } + ingestor.processMessage(handle) // Ack/Nack must not be called — settlement is deferred to the worker. if handle.acks.Load() != 0 || handle.nacks.Load() != 0 { t.Fatalf("enqueued: expected 0 Ack/0 Nack (deferred), got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) @@ -159,10 +144,7 @@ func TestProcessMessage_ChannelFullNacks(t *testing.T) { handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) - err := ingestor.processMessage(handle) - if err != nil { - t.Fatalf("expected nil (nack ok → continue), got: %v", err) - } + ingestor.processMessage(handle) if handle.nacks.Load() != 1 || handle.acks.Load() != 0 { t.Fatalf("channel-full: expected 1 Nack/0 Ack, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) } @@ -178,3 +160,32 @@ func TestProcessMessage_ChannelFullNacks(t *testing.T) { <-ingestor.taskChan } } + +// TestProcessMessage_StartRunningErrorNacks: when StartRunning returns a +// non-ErrTaskNotFound error (e.g. a DB blip), processMessage nacks the +// message for redelivery instead of killing the consumer (B2 fix). The +// consume loop's resilience relies on this never being a fatal return. +func TestProcessMessage_StartRunningErrorNacks(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + // Drop the tasks table so GetTask fails with a generic SQL error, not + // ErrRecordNotFound (which would map to ErrTaskNotFound and ack-skip). + if err := db.Migrator().DropTable(&entity.IngestionTask{}); err != nil { + t.Fatalf("drop table: %v", err) + } + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) + + ingestor.processMessage(handle) + + if handle.nacks.Load() != 1 || handle.acks.Load() != 0 { + t.Fatalf("start-running error: expected 1 Nack/0 Ack (redeliver), got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + if len(ingestor.taskChan) != 0 { + t.Fatal("expected no task enqueued on activation failure") + } +} diff --git a/internal/ingestion/service/progress_sink.go b/internal/ingestion/service/progress_sink.go index 1e295c4649..0146d3a23d 100644 --- a/internal/ingestion/service/progress_sink.go +++ b/internal/ingestion/service/progress_sink.go @@ -18,6 +18,7 @@ package service import ( "fmt" + "sync/atomic" "ragflow/internal/common" "ragflow/internal/dao" @@ -36,6 +37,11 @@ import ( type progressSink struct { taskSvc *servicepkg.IngestionTaskService docSvc docProgressSvc + // total is the component-count denominator cached from OnComponentTotal. + // It is Store-d once in the Run goroutine and Load-ed by OnComponentProgress, + // which eino fires from concurrent parallel-branch goroutines. Atomic because + // the two access paths share no other synchronization. + total atomic.Int64 } // docProgressSvc is the subset of *service.DocumentService the sink needs to @@ -49,40 +55,40 @@ type docProgressSvc interface { func newProgressSink(taskSvc *servicepkg.IngestionTaskService) *progressSink { // Eagerly construct the DocumentService so docSvc is immutable after this // point. eino's compose graph runs parallel branches concurrently, so - // OnComponentProgress (and thus docService) can fire from multiple - // goroutines; a lazy check-then-act here would be a data race. The sink - // owns no server-config dependency, so this is safe in any environment. + // OnComponentProgress (and thus docSvc) can fire from multiple goroutines; + // a lazy check-then-act here would be a data race. The sink owns no + // server-config dependency, so this is safe in any environment. return &progressSink{ taskSvc: taskSvc, docSvc: servicepkg.NewDocumentService(), } } -// docService returns the DocumentService bound at sink construction. It is an -// accessor (not lazy) so concurrent progress callbacks read a stable value. -func (s *progressSink) docService() docProgressSvc { - return s.docSvc -} - func (s *progressSink) OnComponentTotal(taskID string, total int) { + s.total.Store(int64(total)) if err := s.taskSvc.UpdateComponentTotal(taskID, total); err != nil { common.Error(fmt.Sprintf("progressSink: update component_total for task %s failed: %v", taskID, err), err) } } func (s *progressSink) OnComponentProgress(ev pipeline.ProgressEvent) { - if err := s.taskSvc.RecordComponentProgress(ev.TaskID, ev.Component, ev.Index, ev.Phase, ev.Message); err != nil { + if err := s.taskSvc.RecordComponentProgress(ev.TaskID, ev.Component, ev.Phase, ev.Message); err != nil { common.Error(fmt.Sprintf("progressSink: record component progress for task %s failed: %v", ev.TaskID, err), err) } if ev.DocumentID == "" { return } - agg, err := s.taskSvc.AggregateTaskProgress(ev.TaskID, ev.Total) - if err != nil || agg == nil || ev.Total <= 0 { + total := s.total.Load() + agg, err := s.taskSvc.AggregateTaskProgress(ev.TaskID, int(total)) + if err != nil { + common.Error(fmt.Sprintf("progressSink: aggregate task progress for task %s failed: %v", ev.TaskID, err), err) return } - progress, run := deriveDocumentProgress(agg, ev.Total) - if err := s.docService().UpdateRunProgress(ev.DocumentID, progress, run, ev.Message); err != nil { + if agg == nil || total <= 0 { + return + } + progress, run := deriveDocumentProgress(agg, int(total)) + if err := s.docSvc.UpdateRunProgress(ev.DocumentID, progress, run, ev.Message); err != nil { common.Error(fmt.Sprintf("progressSink: mirror progress to document %s for task %s failed: %v", ev.DocumentID, ev.TaskID, err), err) } } diff --git a/internal/ingestion/service/progress_sink_test.go b/internal/ingestion/service/progress_sink_test.go index 4b2d8e21d3..09f8abfb75 100644 --- a/internal/ingestion/service/progress_sink_test.go +++ b/internal/ingestion/service/progress_sink_test.go @@ -17,6 +17,7 @@ package service import ( + "runtime" "sync" "testing" @@ -63,12 +64,12 @@ func TestProgressSink_EagerlyConstructsDocumentService(t *testing.T) { // TestProgressSink_DocService_NoDataRace guards against regressing to lazy // DocumentService construction. eino's compose graph runs parallel branches // concurrently (compose/chain_parallel.go, branch.go), so the progress callback -// can fire from multiple goroutines; docService() must return a pre-built, -// immutable DocumentService, not lazily check-then-act on s.docSvc. +// can fire from multiple goroutines; docSvc must be a pre-built, immutable +// DocumentService, not lazily check-then-act on s.docSvc. // -// The race is hit directly on docService() rather than through -// OnComponentProgress because the latter serializes on the single test-DB -// connection before reaching docService(), which masks the race. +// The race is hit directly on docSvc rather than through OnComponentProgress +// because the latter serializes on the single test-DB connection before +// reaching docSvc, which masks the race. func TestProgressSink_DocService_NoDataRace(t *testing.T) { db := testutil.SetupTestDB(t) cleanup := testutil.ReplaceDBForTest(t, db) @@ -76,7 +77,7 @@ func TestProgressSink_DocService_NoDataRace(t *testing.T) { // Deliberately do NOT inject a stub docSvc: the sink's own DocumentService // must already be constructed (not lazily built mid-call) when the - // goroutines below race into docService(). + // goroutines below race into docSvc. sink := newProgressSink(servicepkg.NewIngestionTaskService()) const n = 30 @@ -87,7 +88,45 @@ func TestProgressSink_DocService_NoDataRace(t *testing.T) { go func() { defer wg.Done() <-start - _ = sink.docService() + _ = sink.docSvc + }() + } + close(start) + wg.Wait() +} + +// TestProgressSink_Total_NoDataRace guards the total denominator against being +// a non-atomic shared field. OnComponentTotal (writer, Run goroutine) and +// OnComponentProgress (reader, concurrent eino branches) share total; a plain +// int is a data race per the Go memory model. The read is hit directly on the +// field rather than through OnComponentProgress because the latter serializes +// on the single test-DB connection before reaching the read, masking the race. +func TestProgressSink_Total_NoDataRace(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + sink := newProgressSink(servicepkg.NewIngestionTaskService()) + + const n = 30 + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + sink.OnComponentTotal(taskID, 5) // writes s.total + }() + } + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + v := sink.total.Load() // reads s.total atomically + runtime.KeepAlive(v) }() } close(start) @@ -137,10 +176,8 @@ func TestProgressSinkPersistsViaService(t *testing.T) { TaskID: taskID, DocumentID: docID, Component: "Parser", - Index: 0, Phase: 1, Message: "Parser Done", - Total: 2, }) logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID(taskID) @@ -187,10 +224,8 @@ func TestProgressSinkEmptyDocumentIDSkipsMirror(t *testing.T) { sink.OnComponentProgress(pipeline.ProgressEvent{ TaskID: taskID, Component: "Chunker", - Index: 1, Phase: 1, Message: "Chunker Done", - Total: 2, }) logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID(taskID) diff --git a/internal/ingestion/service/real_consumer_pipeline_test.go b/internal/ingestion/service/real_consumer_pipeline_test.go index 820467be3e..9f8e4c7716 100644 --- a/internal/ingestion/service/real_consumer_pipeline_test.go +++ b/internal/ingestion/service/real_consumer_pipeline_test.go @@ -81,14 +81,14 @@ func TestRealConsumer_PipelineMessageRoutesToExecuteTask(t *testing.T) { ingestionTaskDAO := dao.NewIngestionTaskDAO() _, err = ingestionTaskDAO.UpdateStatusIfCurrent(taskMsg.TaskID, common.CREATED, common.RUNNING) if err != nil { - t.Fatalf("SetRunningByIngestor: %v", err) + t.Fatalf("UpdateStatusIfCurrent: %v", err) } task, err := ingestionTaskDAO.GetByID(taskMsg.TaskID) if err != nil || task == nil { t.Fatalf("task not found after publish: %s", taskMsg.TaskID) } if task.Status != common.RUNNING { - t.Fatalf("task status after SetRunningByIngestor = %s, want %s", task.Status, common.RUNNING) + t.Fatalf("task status after UpdateStatusIfCurrent = %s, want %s", task.Status, common.RUNNING) } ingestor := NewIngestor("queue-test", 1, []string{"pdf"}) diff --git a/internal/ingestion/service/run_task_test.go b/internal/ingestion/service/run_task_test.go index 57609889cd..b726907f07 100644 --- a/internal/ingestion/service/run_task_test.go +++ b/internal/ingestion/service/run_task_test.go @@ -204,10 +204,12 @@ func TestRunTask_ComponentTimeoutMarksFailed(t *testing.T) { } } -// TestRunTask_MarkCompletedFailure: when runDocumentTask succeeds but -// MarkCompleted fails (status conflict), runTask returns false (non-terminal) -// so the message is Nacked for retry. -func TestRunTask_MarkCompletedFailure(t *testing.T) { +// TestRunTask_AlreadyCompletedAcksNotRedelivers: when the pipeline succeeds but +// the task is already COMPLETED (e.g. another worker won a redelivery race), +// MarkCompleted's transition fails terminally. runTask must treat this as +// terminal and Ack - the work is done, redelivering would just ack-skip - and +// must NOT retry the deterministically-invalid transition. +func TestRunTask_AlreadyCompletedAcksNotRedelivers(t *testing.T) { db := testutil.SetupTestDB(t) cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() @@ -229,8 +231,8 @@ func TestRunTask_MarkCompletedFailure(t *testing.T) { ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING, }) - if terminal { - t.Fatal("expected false (non-terminal) on MarkCompleted failure") + if !terminal { + t.Fatal("expected true (terminal: task already COMPLETED, Ack instead of redeliver)") } // Task must still be COMPLETED (MarkCompleted failed to transition it). @@ -243,6 +245,44 @@ func TestRunTask_MarkCompletedFailure(t *testing.T) { } } +// TestRunTask_PipelineSucceedsConcurrentStopSettlesStopped: the pipeline +// finishes successfully, but a concurrent user stop (RequestStop) moved the +// task RUNNING->STOPPING just before MarkCompleted. The RUNNING->COMPLETED +// transition is now terminally invalid; runTask must settle the task to +// STOPPED and Ack (the pipeline already indexed the chunks) instead of +// retrying the invalid transition and Nacking for redelivery. +func TestRunTask_PipelineSucceedsConcurrentStopSettlesStopped(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = func(ctx context.Context, task *entity.IngestionTask) error { + // Simulate the user pressing Stop mid-pipeline: RUNNING->STOPPING. + if _, err := ingestor.ingestionTaskSvc.RequestStop(task.ID); err != nil { + t.Fatalf("RequestStop: %v", err) + } + return nil // pipeline still finishes successfully + } + + terminal := ingestor.runTask(context.Background(), &entity.IngestionTask{ + ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING, + }) + + if !terminal { + t.Fatal("expected true (terminal: settled to STOPPED, Ack)") + } + + task, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load task: %v", err) + } + if task.Status != common.STOPPED { + t.Fatalf("task status = %s, want STOPPED (settled from concurrent STOPPING)", task.Status) + } +} + // TestRunTask_SuccessfulCompletion: when everything succeeds, runTask returns // true (terminal) and the task is COMPLETED. func TestRunTask_SuccessfulCompletion(t *testing.T) { diff --git a/internal/ingestion/task/embedder.go b/internal/ingestion/task/embedder.go index be8a552d4b..f02fcfcebd 100644 --- a/internal/ingestion/task/embedder.go +++ b/internal/ingestion/task/embedder.go @@ -79,6 +79,9 @@ func newEmbedderResolver( if err != nil { return nil, err } + if model == nil { + return nil, fmt.Errorf("embedder: resolved embedding model is nil for embd_id=%s", embdID) + } return &embedder{model: model}, nil } } diff --git a/internal/ingestion/task/pipeline_e2e_test.go b/internal/ingestion/task/pipeline_e2e_test.go index 3b20fe00d4..7a9d487af7 100644 --- a/internal/ingestion/task/pipeline_e2e_test.go +++ b/internal/ingestion/task/pipeline_e2e_test.go @@ -161,7 +161,7 @@ func isPortOpen(host string, port int) bool { // Main E2E Test with Subtests for Elasticsearch and Infinity // ============================================================================= -func TestPipelineE2E_TaskHandlerToPipelineExecutor(t *testing.T) { +func TestPipelineE2E_PipelineExecutor(t *testing.T) { testCases := []struct { name string engineType engine.EngineType @@ -219,59 +219,54 @@ func TestPipelineE2E_TaskHandlerToPipelineExecutor(t *testing.T) { ) var capturedChunks [][]map[string]any - // Create TaskHandler with mocked DataflowService factory - handler := NewTaskHandler(taskCtx) - handler.WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { - svc := mustNewPipelineExecutor(t, ctx, canvasID, 0) + // Build a PipelineExecutor with mocked dependencies. + svc := mustNewPipelineExecutor(t, taskCtx, taskCtx.PipelineID, 0) - // Mock loadDSLFunc - svc.WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { - loadDSLCalled = true - return `{"nodes":[{"id":"test","type":"parser"}],"edges":[]}`, canvasID, nil - }) - - // Mock runPipelineFunc - returns test chunks (with vectors to skip embedding) - svc.WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { - runPipelineCalled = true - return map[string]any{ - "chunks": []map[string]any{ - { - "text": fmt.Sprintf("Hello world from E2E test with %s", tc.name), - "id": fmt.Sprintf("chunk_e2e_%s_1", lowerName), - "q_2_vec": []float64{0.1, 0.2}, // Pre-vectorized to skip embedding - }, - { - "text": fmt.Sprintf("Second chunk from E2E test with %s", tc.name), - "id": fmt.Sprintf("chunk_e2e_%s_2", lowerName), - "q_2_vec": []float64{0.3, 0.4}, // Pre-vectorized to skip embedding - }, - }, - EmbeddingTokenConsumptionKey: 100, - }, dsl, nil - }) - - // Use the injected DocEngine for insertChunks! - svc.WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { - insertChunksCalled = true - t.Logf("DocEngine InsertChunks called! baseName=%s datasetID=%s len(chunks)=%d", baseName, datasetID, len(chunks)) - ids, err := docEngine.InsertChunks(ctx, chunks, baseName, datasetID) - if err != nil { - t.Logf("WARNING: InsertChunks err=%v", err) - } - capturedChunks = append(capturedChunks, chunks) - return ids, err - }) - - return svc, nil + // Mock loadDSLFunc + svc.WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { + loadDSLCalled = true + return `{"nodes":[{"id":"test","type":"parser"}],"edges":[]}`, canvasID, nil }) - // Execute the task handler! - t.Logf("Calling TaskHandler.Handle()...") - _, err = handler.Handle() + // Mock runPipelineFunc - returns test chunks (with vectors to skip embedding) + svc.WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { + runPipelineCalled = true + return map[string]any{ + "chunks": []map[string]any{ + { + "text": fmt.Sprintf("Hello world from E2E test with %s", tc.name), + "id": fmt.Sprintf("chunk_e2e_%s_1", lowerName), + "q_2_vec": []float64{0.1, 0.2}, // Pre-vectorized to skip embedding + }, + { + "text": fmt.Sprintf("Second chunk from E2E test with %s", tc.name), + "id": fmt.Sprintf("chunk_e2e_%s_2", lowerName), + "q_2_vec": []float64{0.3, 0.4}, // Pre-vectorized to skip embedding + }, + }, + EmbeddingTokenConsumptionKey: 100, + }, dsl, nil + }) + + // Use the injected DocEngine for insertChunks! + svc.WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { + insertChunksCalled = true + t.Logf("DocEngine InsertChunks called! baseName=%s datasetID=%s len(chunks)=%d", baseName, datasetID, len(chunks)) + ids, err := docEngine.InsertChunks(ctx, chunks, baseName, datasetID) + if err != nil { + t.Logf("WARNING: InsertChunks err=%v", err) + } + capturedChunks = append(capturedChunks, chunks) + return ids, err + }) + + // Execute the pipeline! + t.Logf("Calling PipelineExecutor.Execute()...") + _, err = svc.Execute(taskCtx.Ctx) if err != nil { - t.Fatalf("TaskHandler.Handle failed: %v", err) + t.Fatalf("PipelineExecutor.Execute failed: %v", err) } - t.Logf("TaskHandler.Handle() complete!") + t.Logf("PipelineExecutor.Execute() complete!") // Verify all the expected calls happened if !loadDSLCalled { diff --git a/internal/ingestion/task/pipeline_executor.go b/internal/ingestion/task/pipeline_executor.go index 9878982a9d..4e3ff8041c 100644 --- a/internal/ingestion/task/pipeline_executor.go +++ b/internal/ingestion/task/pipeline_executor.go @@ -173,7 +173,6 @@ func (s *PipelineExecutor) Execute(ctx context.Context) (*PipelineResult, error) } func (s *PipelineExecutor) processOutput(ctx context.Context, pipelineOutput map[string]any) (*PipelineResult, error) { - taskStart := time.Now() if pipelineOutput == nil { return nil, nil } @@ -182,7 +181,7 @@ func (s *PipelineExecutor) processOutput(ctx context.Context, pipelineOutput map } chunks := NormalizeChunks(pipelineOutput) - if chunks == nil { + if len(chunks) == 0 { return nil, nil } @@ -195,12 +194,9 @@ func (s *PipelineExecutor) processOutput(ctx context.Context, pipelineOutput map time.Now(), ) - indexStart := time.Now() if err := s.indexWriter.Write(ctx, chunks); err != nil { return nil, err } - _ = time.Since(indexStart) - _ = time.Since(taskStart) return &PipelineResult{ DocID: s.taskCtx.Doc.ID, diff --git a/internal/ingestion/task/pipeline_executor_test.go b/internal/ingestion/task/pipeline_executor_test.go index d430f80510..14ec3cc157 100644 --- a/internal/ingestion/task/pipeline_executor_test.go +++ b/internal/ingestion/task/pipeline_executor_test.go @@ -372,6 +372,35 @@ func TestPipelineExecutor_Run_MainFlowWithStubs(t *testing.T) { } } +// TestPipelineExecutor_Execute_PropagatesContext verifies the ctx passed to +// Execute is the ctx received by runPipelineFunc - the task context must flow +// through to the pipeline run. +func TestPipelineExecutor_Execute_PropagatesContext(t *testing.T) { + type ctxKey string + const key ctxKey = "trace" + taskCtx := makeTaskCtx() + taskCtx.Ctx = context.WithValue(context.Background(), key, "task-ctx") + + svc := mustNewPipelineExecutor(t, taskCtx, "flow-1", 0). + WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { + return `{"nodes":[{"id":"n1"}],"edges":[]}`, canvasID, nil + }). + WithRunPipelineFunc(func(runCtx context.Context, dsl string) (map[string]any, string, error) { + if got := runCtx.Value(key); got != "task-ctx" { + t.Fatalf("runCtx value = %v, want task-ctx", got) + } + return map[string]any{"chunks": []map[string]any{{"text": "hello world"}}}, dsl, nil + }). + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }) + + if _, err := svc.Execute(taskCtx.Ctx); err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + // ============================================================================= // Stub implementations for testing // ============================================================================= diff --git a/internal/ingestion/task/real_consumer_test.go b/internal/ingestion/task/real_consumer_test.go index 0a9003a79e..a70ffc0ed7 100644 --- a/internal/ingestion/task/real_consumer_test.go +++ b/internal/ingestion/task/real_consumer_test.go @@ -101,11 +101,11 @@ func TestRealProducerConsumer(t *testing.T) { t.Fatalf("unexpected task type: %s", taskMsg.TaskType) } - // Mirrors Start():142-143 — SetRunningByIngestor + // Mirrors Start():142-143 — UpdateStatusIfCurrent ingestionTaskDAO := dao.NewIngestionTaskDAO() _, err = ingestionTaskDAO.UpdateStatusIfCurrent(taskMsg.TaskID, common.CREATED, common.RUNNING) if err != nil { - t.Fatalf("SetRunningByIngestor: %v", err) + t.Fatalf("UpdateStatusIfCurrent: %v", err) } task, err := ingestionTaskDAO.GetByID(taskMsg.TaskID) if err != nil { @@ -116,7 +116,7 @@ func TestRealProducerConsumer(t *testing.T) { taskHandle.Ack() return } - t.Logf("Consumer: SetRunningByIngestor status=%s", task.Status) + t.Logf("Consumer: UpdateStatusIfCurrent status=%s", task.Status) // Mirrors Start():167-180 — status check switch task.Status { @@ -143,27 +143,23 @@ func TestRealProducerConsumer(t *testing.T) { t.Logf("Consumer: Loaded Doc=%s Parser=%s KB=%s Tenant=%s", tc.Doc.ID, tc.Doc.ParserID, tc.KB.ID, tc.Tenant.ID) - handler := NewTaskHandler(tc) - handler.WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { - svc, err := NewPipelineExecutor(ctx, canvasID, 0) - if err != nil { - return nil, err - } - svc.WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { - return `{"nodes":[{"id":"test","type":"parser"}],"edges":[]}`, canvasID, nil - }) - svc.WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { - return nil, "", nil - }) - svc.WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - return nil, nil - }) - return svc, nil - }) - if _, err := handler.Handle(); err != nil { - t.Fatalf("Handle: %v", err) + svc, err := NewPipelineExecutor(tc, tc.PipelineID, 0) + if err != nil { + t.Fatalf("NewPipelineExecutor: %v", err) } - t.Log("Consumer: TaskHandler.Handle() — OK") + svc.WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { + return `{"nodes":[{"id":"test","type":"parser"}],"edges":[]}`, canvasID, nil + }) + svc.WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { + return nil, "", nil + }) + svc.WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }) + if _, err := svc.Execute(tc.Ctx); err != nil { + t.Fatalf("Execute: %v", err) + } + t.Log("Consumer: PipelineExecutor.Execute() - OK") // Mirrors executeTask — mark as completed if _, err := ingestionTaskDAO.UpdateStatusIfCurrent(task.ID, common.RUNNING, common.COMPLETED); err != nil { diff --git a/internal/ingestion/task/task_context_test.go b/internal/ingestion/task/task_context_test.go index dc2226dc36..37380de0ea 100644 --- a/internal/ingestion/task/task_context_test.go +++ b/internal/ingestion/task/task_context_test.go @@ -25,7 +25,8 @@ import ( func TestLoadFromIngestionTask_FallsBackToKnowledgebasePipelineID(t *testing.T) { db := testutil.SetupTestDB(t) - testutil.ReplaceDBForTest(t, db) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() if err := db.Create(&entity.Document{ ID: "doc-1", @@ -70,7 +71,8 @@ func TestLoadFromIngestionTask_FallsBackToKnowledgebasePipelineID(t *testing.T) func TestLoadFromIngestionTask_PrefersDocumentPipelineID(t *testing.T) { db := testutil.SetupTestDB(t) - testutil.ReplaceDBForTest(t, db) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() docPipelineID := "doc-flow-1" if err := db.Create(&entity.Document{ diff --git a/internal/ingestion/task/task_handler.go b/internal/ingestion/task/task_handler.go deleted file mode 100644 index 5cd7833445..0000000000 --- a/internal/ingestion/task/task_handler.go +++ /dev/null @@ -1,61 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package task - -import ( - "fmt" - "strings" -) - -// TaskHandler dispatches document processing tasks by task_type. -// Mirrors Python task_handler.py:handle(). -type TaskHandler struct { - ctx *TaskContext - newPipelineExecutor func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) -} - -// NewTaskHandler creates a TaskHandler for the given task context. -func NewTaskHandler(ctx *TaskContext) *TaskHandler { - return &TaskHandler{ - ctx: ctx, - newPipelineExecutor: func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { - return NewPipelineExecutor(ctx, canvasID, 0) - }, - } -} - -func (h *TaskHandler) WithPipelineExecutorFactory(factory func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error)) *TaskHandler { - h.newPipelineExecutor = factory - return h -} - -// Handle routes the task by type and executes the appropriate handler. -func (h *TaskHandler) Handle() (*PipelineResult, error) { - if h.ctx == nil { - return nil, fmt.Errorf("task handler: nil context") - } - - return h.handlePipeline() -} - -func (h *TaskHandler) handlePipeline() (*PipelineResult, error) { - svc, err := h.newPipelineExecutor(h.ctx, strings.TrimSpace(h.ctx.PipelineID)) - if err != nil { - return nil, err - } - return svc.Execute(h.ctx.Ctx) -} diff --git a/internal/ingestion/task/task_handler_test.go b/internal/ingestion/task/task_handler_test.go deleted file mode 100644 index ed687ffbcd..0000000000 --- a/internal/ingestion/task/task_handler_test.go +++ /dev/null @@ -1,178 +0,0 @@ -package task - -import ( - "context" - "strings" - "testing" - - "ragflow/internal/entity" -) - -func testStrPtr(s string) *string { return &s } - -func makeTaskHandlerTestContext(pipelineID string) *TaskContext { - return &TaskContext{ - IngestionTask: &entity.IngestionTask{ - ID: "task-1", - DocumentID: "doc-1", - }, - PipelineID: pipelineID, - Doc: entity.Document{ - ID: "doc-1", - KbID: "kb-1", - Name: testStrPtr("test-doc.pdf"), - ParserID: "naive", - ParserConfig: entity.JSONMap{}, - }, - KB: entity.Knowledgebase{ - ID: "kb-1", - TenantID: "tenant-1", - EmbdID: "embd-1", - }, - Tenant: entity.Tenant{ - ID: "tenant-1", - LLMID: "gpt-4", - }, - } -} - -func newNoopPipelineExecutor(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { - if strings.TrimSpace(canvasID) == "" { - canvasID = "flow-1" - } - svc, err := NewPipelineExecutor(ctx, canvasID, 0) - if err != nil { - return nil, err - } - svc = svc. - WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { - return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, canvasID, nil - }). - WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { - return map[string]any{ - "chunks": []map[string]any{{ - "text": "stub pipeline chunk", - "q_2_vec": []float64{0.1, 0.2}, - }}, - }, dsl, nil - }). - WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - return nil, nil - }). - WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { - return nil - }) - return svc, nil -} - -func newNoopTaskHandler(ctx *TaskContext) *TaskHandler { - return NewTaskHandler(ctx).WithPipelineExecutorFactory(newNoopPipelineExecutor) -} - -func TestTaskHandler_HandleRejectsNilContext(t *testing.T) { - if _, err := NewTaskHandler(nil).Handle(); err == nil { - t.Fatal("expected error for nil context") - } -} - -func TestTaskHandler_HandleRequiresPipelineID(t *testing.T) { - ctx := makeTaskHandlerTestContext("") - handler := NewTaskHandler(ctx) - if _, err := handler.Handle(); err == nil { - t.Fatal("expected error for empty pipeline id") - } -} - -func TestTaskHandler_HandleRunWithFactory(t *testing.T) { - ctx := makeTaskHandlerTestContext("flow-1") - ctx.Ctx = context.Background() - handler := NewTaskHandler(ctx).WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { - return newNoopPipelineExecutor(ctx, canvasID) - }) - if _, err := handler.Handle(); err != nil { - t.Fatalf("unexpected error: %v", err) - } -} - -func TestTaskHandler_Pipeline_UsesTaskContext(t *testing.T) { - ctx := makeTaskHandlerTestContext("flow-1") - type ctxKey string - const key ctxKey = "trace" - ctx.Ctx = context.WithValue(context.Background(), key, "task-ctx") - - handler := NewTaskHandler(ctx).WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { - return mustNewPipelineExecutor(t, ctx, canvasID, 0). - WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { - return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, canvasID, nil - }). - WithRunPipelineFunc(func(runCtx context.Context, dsl string) (map[string]any, string, error) { - if got := runCtx.Value(key); got != "task-ctx" { - t.Fatalf("runCtx value = %v, want task-ctx", got) - } - return map[string]any{"chunks": []map[string]any{{"text": "stub", "q_2_vec": []float64{0.1, 0.2}}}}, dsl, nil - }). - WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - return nil, nil - }). - WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }), nil - }) - if _, err := handler.Handle(); err != nil { - t.Fatalf("unexpected error: %v", err) - } -} - -func TestTaskHandler_Pipeline_ShowsProgressAndPipelineLog(t *testing.T) { - ctx := makeTaskHandlerTestContext("flow-1") - ctx.Ctx = context.Background() - ctx.Doc.PipelineID = testStrPtr("flow-1") - ctx.Doc.Name = testStrPtr("verify-pipeline.pdf") - - var pipelineCalled bool - var insertCalled bool - var logCreateCalls int - var insertedChunkCount int - - handler := NewTaskHandler(ctx).WithPipelineExecutorFactory(func(ctx *TaskContext, canvasID string) (*PipelineExecutor, error) { - svc := mustNewPipelineExecutor(t, ctx, canvasID, 0). - WithLoadDSLFunc(func(ctx context.Context, canvasID string) (string, string, error) { - return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, canvasID, nil - }). - WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { - pipelineCalled = true - return map[string]any{ - "chunks": []map[string]any{{ - "text": "stub pipeline chunk", - "q_2_vec": []float64{0.1, 0.2}, - }}, - }, dsl, nil - }). - WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { - insertCalled = true - insertedChunkCount = len(chunks) - return nil, nil - }). - WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { - logCreateCalls++ - return nil - }) - return svc, nil - }) - - if _, err := handler.Handle(); err != nil { - t.Fatalf("handler.Handle() error: %v", err) - } - - if !pipelineCalled { - t.Fatal("expected mock pipeline.run to be called") - } - if !insertCalled { - t.Fatal("expected insertChunks to be called") - } - if insertedChunkCount != 1 { - t.Fatalf("insertedChunkCount = %d, want 1", insertedChunkCount) - } - - if logCreateCalls != 1 { - t.Fatalf("logCreateCalls = %d, want 1", logCreateCalls) - } -} diff --git a/internal/service/ingestion_task_service.go b/internal/service/ingestion_task_service.go index 9d1ea29102..95891241c7 100644 --- a/internal/service/ingestion_task_service.go +++ b/internal/service/ingestion_task_service.go @@ -409,14 +409,13 @@ func (s *IngestionTaskService) UpdateComponentTotal(taskID string, total int) er // ingestion_task_log (phase: 0 started / 1 done / 2 errored). The row's // Checkpoint is empty; component progress and step checkpoints are distinct // row models sharing the same table. -func (s *IngestionTaskService) RecordComponentProgress(taskID, component string, index, phase int, message string) error { +func (s *IngestionTaskService) RecordComponentProgress(taskID, component string, phase int, message string) error { entry := &entity.IngestionTaskLog{ - TaskID: taskID, - Checkpoint: entity.JSONMap{}, - ComponentIndex: index, - Phase: phase, - Component: component, - Message: message, + TaskID: taskID, + Checkpoint: entity.JSONMap{}, + Phase: phase, + Component: component, + Message: message, } return s.ingestionTaskLogDAO.Create(entry) } diff --git a/internal/service/ingestion_task_service_test.go b/internal/service/ingestion_task_service_test.go index 2001dcef4b..a7deb63b06 100644 --- a/internal/service/ingestion_task_service_test.go +++ b/internal/service/ingestion_task_service_test.go @@ -494,7 +494,7 @@ func TestIngestionTaskServiceRecordComponentProgressAppendsRow(t *testing.T) { insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := NewIngestionTaskService() - if err := svc.RecordComponentProgress("task-1", "Parser", 0, 1, "Parser Done"); err != nil { + if err := svc.RecordComponentProgress("task-1", "Parser", 1, "Parser Done"); err != nil { t.Fatalf("RecordComponentProgress failed: %v", err) } logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID("task-1") @@ -505,7 +505,7 @@ func TestIngestionTaskServiceRecordComponentProgressAppendsRow(t *testing.T) { t.Fatalf("expected 1 log row, got %d", len(logs)) } row := logs[0] - if row.Component != "Parser" || row.ComponentIndex != 0 || row.Phase != 1 || row.Message != "Parser Done" { + if row.Component != "Parser" || row.Phase != 1 || row.Message != "Parser Done" { t.Fatalf("unexpected log row: %+v", row) } if len(row.Checkpoint) != 0 { @@ -519,10 +519,10 @@ func TestIngestionTaskServiceAggregateTaskProgressClassifiesByPhase(t *testing.T insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := NewIngestionTaskService() - if err := svc.RecordComponentProgress("task-1", "Parser", 0, 1, "Parser Done"); err != nil { + if err := svc.RecordComponentProgress("task-1", "Parser", 1, "Parser Done"); err != nil { t.Fatalf("record Parser: %v", err) } - if err := svc.RecordComponentProgress("task-1", "Chunker", 1, 0, "Chunker Started"); err != nil { + if err := svc.RecordComponentProgress("task-1", "Chunker", 0, "Chunker Started"); err != nil { t.Fatalf("record Chunker: %v", err) } agg, err := svc.AggregateTaskProgress("task-1", 2) @@ -546,7 +546,7 @@ func TestIngestionTaskServiceIncrementRunCountInitializesAndBumps(t *testing.T) if err := svc.IncrementRunCount("task-1"); err != nil { t.Fatalf("IncrementRunCount (first call) failed: %v", err) } - run, ok := findLastRunCount(t, "task-1") + run, ok := svc.lastRunCount("task-1") if !ok || run != 1 { t.Fatalf("run_count = %v (ok=%v), want 1", run, ok) } @@ -555,7 +555,7 @@ func TestIngestionTaskServiceIncrementRunCountInitializesAndBumps(t *testing.T) if err := svc.IncrementRunCount("task-1"); err != nil { t.Fatalf("IncrementRunCount (second call) failed: %v", err) } - run, _ = findLastRunCount(t, "task-1") + run, _ = svc.lastRunCount("task-1") if run != 2 { t.Fatalf("run_count after second bump = %v, want 2", run) } @@ -577,7 +577,7 @@ func TestIngestionTaskServiceIncrementRunCountSkippedCorruptedRunCount(t *testin if err := svc.IncrementRunCount("task-1"); err != nil { t.Fatalf("IncrementRunCount should skip corrupted value, got: %v", err) } - run, ok := findLastRunCount(t, "task-1") + run, ok := svc.lastRunCount("task-1") if !ok || run != 1 { t.Fatalf("run_count = %v (ok=%v), want 1", run, ok) } @@ -591,10 +591,10 @@ func TestIngestionTaskServiceIncrementRunCountRecoversFromComponentProgressLog(t // Simulate a previous run that created some component-progress logs // but died before recording a run_count row. The latest log has no run_count. svc := NewIngestionTaskService() - if err := svc.RecordComponentProgress("task-1", "Parser", 0, 1, "Parser Done"); err != nil { + if err := svc.RecordComponentProgress("task-1", "Parser", 1, "Parser Done"); err != nil { t.Fatalf("record Parser: %v", err) } - if err := svc.RecordComponentProgress("task-1", "Chunker", 1, 1, "Chunker Done"); err != nil { + if err := svc.RecordComponentProgress("task-1", "Chunker", 1, "Chunker Done"); err != nil { t.Fatalf("record Chunker: %v", err) } // Verify latest log has empty checkpoint (no run_count). @@ -610,7 +610,7 @@ func TestIngestionTaskServiceIncrementRunCountRecoversFromComponentProgressLog(t if err := svc.IncrementRunCount("task-1"); err != nil { t.Fatalf("IncrementRunCount failed: %v", err) } - run, ok := findLastRunCount(t, "task-1") + run, ok := svc.lastRunCount("task-1") if !ok || run != 1 { t.Fatalf("run_count = %v (ok=%v), want 1", run, ok) } @@ -638,7 +638,7 @@ func TestIngestionTaskServiceIncrementRunCountAccumulatesAcrossRetries(t *testin t.Fatalf("first IncrementRunCount: %v", err) } // Simulate first run: some components progress, then failure. - if err := svc.RecordComponentProgress("task-1", "Parser", 0, 1, "Parser Done"); err != nil { + if err := svc.RecordComponentProgress("task-1", "Parser", 1, "Parser Done"); err != nil { t.Fatalf("record Parser: %v", err) } @@ -647,7 +647,7 @@ func TestIngestionTaskServiceIncrementRunCountAccumulatesAcrossRetries(t *testin t.Fatalf("second IncrementRunCount: %v", err) } // More progress, then failure. - if err := svc.RecordComponentProgress("task-1", "Chunker", 1, 1, "Chunker Done"); err != nil { + if err := svc.RecordComponentProgress("task-1", "Chunker", 1, "Chunker Done"); err != nil { t.Fatalf("record Chunker: %v", err) } @@ -656,17 +656,34 @@ func TestIngestionTaskServiceIncrementRunCountAccumulatesAcrossRetries(t *testin t.Fatalf("third IncrementRunCount: %v", err) } - run, ok := findLastRunCount(t, "task-1") + run, ok := svc.lastRunCount("task-1") if !ok || run != 3 { t.Fatalf("run_count = %v (ok=%v), want 3", run, ok) } // ListAllForAdmin should still pick up the correct run_count. - allLogs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID("task-1") - if err != nil { - t.Fatalf("list all logs: %v", err) + status := "1" + if err := dao.DB.Create(&entity.User{ + ID: "user-1", + Email: "user-1@test.com", + Nickname: "user-1", + IsAuthenticated: "1", + IsActive: "1", + IsAnonymous: "0", + Status: &status, + }).Error; err != nil { + t.Fatalf("insert user: %v", err) + } + adminTasks, err := svc.ListAllForAdmin() + if err != nil { + t.Fatalf("ListAllForAdmin: %v", err) + } + if len(adminTasks) != 1 || adminTasks[0]["id"] != "task-1" { + t.Fatalf("ListAllForAdmin = %+v, want single task task-1", adminTasks) + } + if adminTasks[0]["run_count"] != 3 { + t.Fatalf("ListAllForAdmin run_count = %v, want 3", adminTasks[0]["run_count"]) } - t.Logf("total log rows: %d", len(allLogs)) } func TestIngestionTaskServiceMarkStoppedTransitionsStoppingTask(t *testing.T) { @@ -731,20 +748,3 @@ func TestDocumentServiceUpdateRunProgressMirrorsFields(t *testing.T) { t.Fatalf("progress_msg = %v, want halfway", doc.ProgressMsg) } } - -// findLastRunCount scans all logs for a task and returns the latest run_count. -// It avoids non-deterministic LatestLogByTaskID ordering when multiple rows -// share the same create_time. -func findLastRunCount(t *testing.T, taskID string) (int, bool) { - t.Helper() - logs, err := dao.NewIngestionTaskLogDAO().ListLogsByTaskID(taskID) - if err != nil { - t.Fatalf("list logs for %s: %v", taskID, err) - } - for i := len(logs) - 1; i >= 0; i-- { - if count, ok := common.GetInt(logs[i].Checkpoint[stepKeyRunCount]); ok { - return count, true - } - } - return 0, false -}