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.

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)
+}

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
 	}
 }
 

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")
 	}
 }