From 37b945c94d3f3f580128b0326ce335559de8a31c Mon Sep 17 00:00:00 2001 From: Jack Date: Mon, 17 Aug 2026 19:09:58 +0800 Subject: [PATCH] fix(ingestor): stop dropping tasks under burst parse backpressure (#18369) Fixes the defect where selecting many files (e.g. 20+) in one dataset and starting parsing at once leaves most of them stuck in `RUNNING` forever: a few parse, the rest never do. --- .../service/burst_backpressure_test.go | 224 ++++++++++++++++++ .../ingestion/service/ingestion_service.go | 36 ++- .../ingestion/service/process_message_test.go | 56 +++-- 3 files changed, 290 insertions(+), 26 deletions(-) create mode 100644 internal/ingestion/service/burst_backpressure_test.go diff --git a/internal/ingestion/service/burst_backpressure_test.go b/internal/ingestion/service/burst_backpressure_test.go new file mode 100644 index 0000000000..ad77fd7ecf --- /dev/null +++ b/internal/ingestion/service/burst_backpressure_test.go @@ -0,0 +1,224 @@ +// +// 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. +// +// +// This test reproduces the "20+ files selected, only some parse, the rest +// get stuck forever" report against the Go ingestor (canvas / pipeline path). +// +// Root cause it guards against: processMessage applies backpressure by +// NACKing the message when taskChan is full. The NATS consumer is configured +// with MaxDeliver: 16 and NO dead-letter subject (internal/engine/nats/nats.go), +// so a message NACKed 16 times is permanently dropped. By that point +// StartRunning has already flipped the task (and its document) to RUNNING in +// the DB, and the ingestor has no scan-and-re-enqueue path on completion — so +// the dropped document is stuck in RUNNING forever. +// +// The test drives the REAL processMessage + REAL worker pool, but substitutes +// a deterministic "broker" that mirrors NATS' contract: bounded redelivery +// (MaxDeliver) with no dead-letter. A slow worker models a saturated pipeline +// (e.g. LLM calls in a canvas). The invariant under test is that NO task is +// ever lost: all N delivered tasks eventually reach COMPLETED. +// +// Under the buggy (NACK-on-backpressure) code, the burst overflow is dropped +// after MaxDeliver and the test FAILS. After the fix (blocking send so the +// consume loop applies backpressure instead of dropping), every task is +// processed and the test PASSES. + +package service + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "ragflow/internal/common" + "ragflow/internal/entity" + "ragflow/internal/ingestion/testutil" + + "gorm.io/gorm" +) + +const ( + burstTaskCount = 20 + burstMaxDeliver = 16 // mirrors NATS MaxDeliver in internal/engine/nats/nats.go + burstWorkerMs = 50 // slow worker models a saturated pipeline +) + +// burstMsg models one in-flight broker message: the handle processMessage +// settles, plus how many times it has been delivered (1 = initial delivery). +type burstMsg struct { + handle *fakeTaskHandle + deliveries int +} + +// TestProcessMessage_BurstNoTaskLossUnderBackpressure drives a 20-file burst +// through the real processMessage + worker pool with a simulated NATS broker +// (bounded redelivery, no dead-letter). It asserts that every task completes +// and none are dropped. +func TestProcessMessage_BurstNoTaskLossUnderBackpressure(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + taskIDs := seedBurstTasks(t, db, burstTaskCount) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + // Slow worker: model a saturated pipeline so the channel is full during + // the burst, forcing the backpressure path. + ingestor.runDocumentTask = func(ctx context.Context, _ *entity.IngestionTask) error { + time.Sleep(burstWorkerMs * time.Millisecond) + return nil + } + ingestor.startWorkerPool() + defer ingestor.Stop(context.Background()) + + // Simulated NATS broker: pending holds messages waiting to be delivered. + var mu sync.Mutex + var pending []*burstMsg + dropped := 0 + + seedBroker := func() { + mu.Lock() + defer mu.Unlock() + for _, id := range taskIDs { + pending = append(pending, &burstMsg{ + handle: newFakeHandle(id, common.TaskTypeIngestionTask), + deliveries: 1, + }) + } + } + seedBroker() + + // Simulated consume loop: mirrors consumeLoop -> GetMessages(4) -> + // processMessage, with NATS-style bounded redelivery on Nack. + stop := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + mu.Lock() + if len(pending) == 0 { + mu.Unlock() + time.Sleep(5 * time.Millisecond) // idle backoff + continue + } + batchSize := 4 + if len(pending) < batchSize { + batchSize = len(pending) + } + batch := pending[:batchSize] + pending = pending[batchSize:] + mu.Unlock() + + for _, bm := range batch { + ingestor.processMessage(bm.handle) + if bm.handle.nacks.Load() > 0 { + mu.Lock() + bm.deliveries++ + if bm.deliveries < burstMaxDeliver { + // Redeliver (no delay — worst case for loss). + pending = append(pending, bm) + } else { + // NATS drops after MaxDeliver with no dead-letter: + // the task is permanently lost. + dropped++ + } + mu.Unlock() + } + } + } + }() + + // Wait until all tasks reach COMPLETED, or time out. + deadline := time.Now().Add(30 * time.Second) + completed := 0 + for time.Now().Before(deadline) { + completed = countTasksWithStatus(t, db, taskIDs, common.COMPLETED) + if completed == burstTaskCount { + break + } + time.Sleep(50 * time.Millisecond) + } + close(stop) + wg.Wait() + + if dropped > 0 { + t.Fatalf("BUG: %d/%d task(s) were permanently dropped (lost) under burst backpressure "+ + "with MaxDeliver=%d and no dead-letter; only %d/%d completed. "+ + "This is the 'files get stuck after the first few parse' defect.", + dropped, burstTaskCount, burstMaxDeliver, completed, burstTaskCount) + } + if completed != burstTaskCount { + t.Fatalf("expected all %d tasks to complete, got %d (stuck in RUNNING)", + burstTaskCount, completed) + } +} + +// seedBurstTasks creates one tenant/kb plus N (document, ingestion_task) pairs +// in CREATED status, mirroring how CreateAndEnqueue publishes a message for a +// freshly-created task. Returns the ingestion task IDs. +func seedBurstTasks(t *testing.T, db *gorm.DB, n int) []string { + t.Helper() + const tenantID, kbID = "burst-tenant", "burst-kb" + if err := db.Create(&entity.Tenant{ID: tenantID, LLMID: "gpt-4", Status: testutil.StrPtr("1")}).Error; err != nil { + t.Fatalf("create tenant: %v", err) + } + if err := db.Create(&entity.Knowledgebase{ + ID: kbID, TenantID: tenantID, EmbdID: "embd-1", Status: testutil.StrPtr("1"), ParserConfig: entity.JSONMap{}, + }).Error; err != nil { + t.Fatalf("create kb: %v", err) + } + + ids := make([]string, 0, n) + for i := 0; i < n; i++ { + docID := fmt.Sprintf("burst-doc-%d", i) + taskID := fmt.Sprintf("burst-task-%d", i) + loc := "doc_store/" + docID + if err := db.Create(&entity.Document{ + ID: docID, KbID: kbID, Name: &docID, ParserID: "naive", + ParserConfig: entity.JSONMap{}, PipelineID: testutil.StrPtr("flow-1"), + Status: testutil.StrPtr("1"), Type: "pdf", Location: &loc, + }).Error; err != nil { + t.Fatalf("create doc %s: %v", docID, err) + } + if err := db.Create(&entity.IngestionTask{ + ID: taskID, UserID: "u1", DocumentID: docID, DatasetID: kbID, Status: common.CREATED, + }).Error; err != nil { + t.Fatalf("create ingestion task %s: %v", taskID, err) + } + ids = append(ids, taskID) + } + return ids +} + +// countTasksWithStatus returns how many of the given task IDs have the given status. +func countTasksWithStatus(t *testing.T, db *gorm.DB, ids []string, status string) int { + t.Helper() + var count int64 + if err := db.Model(&entity.IngestionTask{}). + Where("id IN ? AND status = ?", ids, status). + Count(&count).Error; err != nil { + t.Fatalf("count tasks: %v", err) + } + return int(count) +} diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index b221a48299..46d882b824 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -335,11 +335,11 @@ func (e *Ingestor) processMessage(handle common.TaskHandle) { select { case e.taskChan <- taskCtx: common.Info(fmt.Sprintf("Memory task %s queued (channel: %d/%d)", taskMessage.TaskID, len(e.taskChan), cap(e.taskChan))) - default: - common.Info(fmt.Sprintf("No available slot for memory task %s, nack", taskMessage.TaskID)) - if nackErr := handle.Nack(); nackErr != nil { - common.Error(fmt.Sprintf("error nack memory task %s", taskMessage.TaskID), nackErr) - } + case <-e.ctx.Done(): + // Shutdown won the race: return without settling so the broker + // redelivers the memory task after restart. + common.Info(fmt.Sprintf("Ingestor shutting down; memory task %s not enqueued", taskMessage.TaskID)) + return } return } @@ -414,17 +414,29 @@ func (e *Ingestor) processMessage(handle common.TaskHandle) { taskCtx := taskpkg.NewTaskContextForScheduling(e.ctx, task) taskCtx.Handle = handle - // Push to task channel; if full, reject the task (backpressure). + // Push to the task channel. Use a blocking send so backpressure is + // applied at the consumer: the consume loop waits for a free worker slot + // instead of dropping the message. Dropping on backpressure is unsafe + // because StartRunning (above) has already flipped the task — and its + // document — to RUNNING in the DB, and there is no scan-and-re-enqueue + // path on completion. A Nack that exceeds the broker's MaxDeliver (16, + // with no dead-letter in nats.go) would permanently lose the task and + // leave the document stuck in RUNNING — the "files get stuck after the + // first few parse" defect. + // + // The in-flight claim (set just above) guards against a redelivery racing + // the blocked send: a duplicate delivery sees claimTask fail and is + // ack-skipped, so a blocking send cannot double-execute a task. 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)) - // 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) - } + case <-e.ctx.Done(): + // Shutdown won the race: release the claim and return without + // settling so the broker redelivers the message after restart + // rather than blocking the consume loop forever. + common.Info(fmt.Sprintf("Ingestor shutting down; releasing slot for task %s without ack", task.ID)) + return } } diff --git a/internal/ingestion/service/process_message_test.go b/internal/ingestion/service/process_message_test.go index 2a565c76de..56cda0456c 100644 --- a/internal/ingestion/service/process_message_test.go +++ b/internal/ingestion/service/process_message_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "testing" + "time" "ragflow/internal/common" "ragflow/internal/dao" @@ -309,10 +310,12 @@ func TestProcessMessage_ClaimSucceedsEnqueues(t *testing.T) { } } -// TestProcessMessage_ChannelFullNacks: when the task channel is at capacity -// backpressure rejects the task with Nack, releases the claim, and returns nil -// so the message is redelivered and a future attempt can re-claim it. -func TestProcessMessage_ChannelFullNacks(t *testing.T) { +// TestProcessMessage_ChannelFullBlocksUntilSlot: backpressure must NOT drop +// the task. When the task channel is at capacity, processMessage blocks on the +// send (consuming no slot, no Nack) until a worker frees one; the message is +// then enqueued and settled by the worker. Dropping on backpressure would +// permanently lose the task once the broker's MaxDeliver is exceeded. +func TestProcessMessage_ChannelFullBlocksUntilSlot(t *testing.T) { db := testutil.SetupTestDB(t) cleanup := testutil.ReplaceDBForTest(t, db) defer cleanup() @@ -326,20 +329,45 @@ func TestProcessMessage_ChannelFullNacks(t *testing.T) { handle := newFakeHandle(taskID, common.TaskTypeIngestionTask) - 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()) + // processMessage must BLOCK on the full channel: no ack, no nack, no + // immediate enqueue. + done := make(chan struct{}) + go func() { + ingestor.processMessage(handle) + close(done) + }() + select { + case <-done: + t.Fatal("processMessage returned while channel was full; it must block under backpressure") + case <-time.After(100 * time.Millisecond): + } + if handle.nacks.Load() != 0 || handle.acks.Load() != 0 { + t.Fatalf("expected 0 Ack/0 Nack while blocked, got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + if len(ingestor.taskChan) != cap(ingestor.taskChan) { + t.Fatalf("expected channel still full, got %d/%d", len(ingestor.taskChan), cap(ingestor.taskChan)) } - // Claim must be released so a future redelivery can re-claim it. - if !ingestor.claimTask(taskID) { - t.Fatal("claim was not released on channel-full — task would be stuck forever") + // Free a slot: the blocked processMessage must now enqueue the task. + <-ingestor.taskChan + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("processMessage did not enqueue after a slot freed") } - ingestor.releaseTask(taskID) - - // Drain the fillers. + if handle.nacks.Load() != 0 || handle.acks.Load() != 0 { + t.Fatalf("expected 0 Ack/0 Nack (settlement deferred to worker), got acks=%d nacks=%d", handle.acks.Load(), handle.nacks.Load()) + } + // The real task must now be in the channel. + found := false for i := 0; i < cap(ingestor.taskChan); i++ { - <-ingestor.taskChan + tc := <-ingestor.taskChan + if tc.IngestionTask != nil && tc.IngestionTask.ID == taskID { + found = true + } + } + if !found { + t.Fatal("real task was not enqueued after a slot freed") } }