diff --git a/internal/ingestion/service/doc_state.go b/internal/ingestion/service/doc_state.go index 92b83fa3aa..e4a293028b 100644 --- a/internal/ingestion/service/doc_state.go +++ b/internal/ingestion/service/doc_state.go @@ -33,7 +33,7 @@ import ( type docStateSvc interface { GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]any, error) SetDocumentMetadata(ctx context.Context, docID string, meta map[string]any) error - IncrementChunkNum(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error + ApplyDocCounts(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error } // docStateUpdater applies a pipeline run's results to document state: it @@ -61,8 +61,8 @@ func (u *docStateUpdater) apply(ctx context.Context, r *taskpkg.PipelineResult) common.Warn(fmt.Sprintf("failed to update document metadata: %v", err)) } } - if err := u.docSvc.IncrementChunkNum(ctx, r.DocID, r.KbID, r.ChunkCount, r.TokenConsumption, r.Duration); err != nil { - common.Warn(fmt.Sprintf("failed to increment chunk num: %v", err)) + if err := u.docSvc.ApplyDocCounts(ctx, r.DocID, r.KbID, r.ChunkCount, r.TokenConsumption, r.Duration); err != nil { + common.Warn(fmt.Sprintf("failed to apply doc counts: %v", err)) } } diff --git a/internal/ingestion/service/doc_state_test.go b/internal/ingestion/service/doc_state_test.go index 5ea98f4cdc..786b64fe24 100644 --- a/internal/ingestion/service/doc_state_test.go +++ b/internal/ingestion/service/doc_state_test.go @@ -52,7 +52,7 @@ func (s *stubDocStateSvc) SetDocumentMetadata(ctx context.Context, docID string, return nil } -func (s *stubDocStateSvc) IncrementChunkNum(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error { +func (s *stubDocStateSvc) ApplyDocCounts(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error { s.incrementCalled = true s.gotDocID = docID s.gotKbID = kbID diff --git a/internal/ingestion/service/redelivery_counter_test.go b/internal/ingestion/service/redelivery_counter_test.go new file mode 100644 index 0000000000..87335263a9 --- /dev/null +++ b/internal/ingestion/service/redelivery_counter_test.go @@ -0,0 +1,131 @@ +// +// 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 service + +import ( + "context" + "testing" + + "gorm.io/gorm" + + "ragflow/internal/common" + "ragflow/internal/entity" + taskpkg "ragflow/internal/ingestion/task" + "ragflow/internal/ingestion/testutil" +) + +// The two tests below guard counter idempotency across at-least-once redelivery. +// runTask applies the pipeline's chunk/token counts to the document and its +// knowledgebase (runDocumentTask -> docState.apply -> ApplyDocCounts) before it +// marks the task complete, so a redelivered task re-runs the pipeline. Because +// ApplyDocCounts sets the document's own counts and rolls only the delta into the +// knowledge base aggregate, re-applying the same run contributes zero: the same +// task, processed twice, applies its counters once. The chunk store is likewise +// idempotent (upsert by deterministic id). + +const ( + rcChunks int64 = 5 + rcTokens int64 = 100 +) + +// applyResult stands in for defaultRunDocumentTask: a successful pipeline run +// that applies its result to the doc + KB counters, as docState.apply does after +// Execute returns. +func applyResult(ingestor *Ingestor, docID, kbID string) func(context.Context, *entity.IngestionTask) error { + return func(ctx context.Context, _ *entity.IngestionTask) error { + ingestor.docState.apply(ctx, &taskpkg.PipelineResult{ + DocID: docID, + KbID: kbID, + ChunkCount: int(rcChunks), + TokenConsumption: int(rcTokens), + Duration: 1, + }) + return nil + } +} + +// assertCountersAppliedOnce fails when the doc/KB counters reflect more than a +// single application of the pipeline result - i.e. a redelivery re-counted. +func assertCountersAppliedOnce(t *testing.T, db *gorm.DB, kbID, docID string) { + t.Helper() + var kb entity.Knowledgebase + if err := db.First(&kb, "id = ?", kbID).Error; err != nil { + t.Fatalf("load kb: %v", err) + } + var doc entity.Document + if err := db.First(&doc, "id = ?", docID).Error; err != nil { + t.Fatalf("load doc: %v", err) + } + if kb.ChunkNum != rcChunks || kb.TokenNum != rcTokens { + t.Errorf("kb counters = (chunk %d, token %d), want (%d, %d) - redelivery double-counted", kb.ChunkNum, kb.TokenNum, rcChunks, rcTokens) + } + if doc.ChunkNum != rcChunks || doc.TokenNum != rcTokens { + t.Errorf("doc counters = (chunk %d, token %d), want (%d, %d) - redelivery double-counted", doc.ChunkNum, doc.TokenNum, rcChunks, rcTokens) + } +} + +func rcMsg(taskID, docID, kbID string) *entity.IngestionTask { + return &entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: kbID, Status: common.RUNNING} +} + +// TestRunTask_RedeliveryOfCompletedTaskCountsOnce: the first delivery fully +// completes (task -> COMPLETED), then the broker redelivers the same message +// because its Ack was lost. runTask has no already-completed guard, so it +// re-runs the pipeline; the counter application must stay idempotent. +func TestRunTask_RedeliveryOfCompletedTaskCountsOnce(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, kbID, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = applyResult(ingestor, docID, kbID) + + // First delivery: parse succeeds, counters applied once, task -> COMPLETED. + if terminal := ingestor.runTask(context.Background(), rcMsg(taskID, docID, kbID)); !terminal { + t.Fatalf("expected the first delivery to complete (terminal=true)") + } + // Redelivery (Ack lost): the same message is processed again. Must NOT re-count. + ingestor.runTask(context.Background(), rcMsg(taskID, docID, kbID)) + + assertCountersAppliedOnce(t, db, kbID, docID) +} + +// TestRunTask_RedeliveryAfterIncompleteRunCountsOnce: the crash/nack window. A +// prior run applied the counters but died before MarkCompleted, so the task is +// left RUNNING and the broker redelivers it; the redelivery re-runs and +// completes, and the counters must not be applied twice. +func TestRunTask_RedeliveryAfterIncompleteRunCountsOnce(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + _, kbID, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1")) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + ingestor.runDocumentTask = applyResult(ingestor, docID, kbID) + + // Prior run: counters applied, but the task never completed (crash before + // MarkCompleted) - the task row is left RUNNING, so the broker redelivers. + applyResult(ingestor, docID, kbID)(context.Background(), nil) + + // Redelivery of the still-RUNNING task: re-runs and completes. + if terminal := ingestor.runTask(context.Background(), rcMsg(taskID, docID, kbID)); !terminal { + t.Fatalf("expected the redelivery run to complete (terminal=true)") + } + + assertCountersAppliedOnce(t, db, kbID, docID) +} diff --git a/internal/service/document/apply_doc_counts_test.go b/internal/service/document/apply_doc_counts_test.go new file mode 100644 index 0000000000..66e4dc1a18 --- /dev/null +++ b/internal/service/document/apply_doc_counts_test.go @@ -0,0 +1,126 @@ +// +// 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 document + +import ( + "context" + "testing" + + "ragflow/internal/entity" +) + +// TestApplyDocCounts_ReparseCarriesDelta checks that re-applying a document with +// a different count - a re-parse - moves the knowledge base aggregate by the +// delta to the new absolute value, rather than summing the two runs. +func TestApplyDocCounts_ReparseCarriesDelta(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + if err := db.Create(&entity.Knowledgebase{ID: "kb-1"}).Error; err != nil { + t.Fatalf("create kb: %v", err) + } + if err := db.Create(&entity.Document{ID: "doc-1", KbID: "kb-1", ParserConfig: entity.JSONMap{}}).Error; err != nil { + t.Fatalf("create doc: %v", err) + } + svc := testDocumentService(t) + ctx := context.Background() + + // First parse produces 5 chunks / 100 tokens; a re-parse produces 7 / 140. + if err := svc.ApplyDocCounts(ctx, "doc-1", "kb-1", 5, 100, 1); err != nil { + t.Fatalf("first apply: %v", err) + } + if err := svc.ApplyDocCounts(ctx, "doc-1", "kb-1", 7, 140, 1); err != nil { + t.Fatalf("reparse apply: %v", err) + } + + // The document holds the new absolute counts; the KB aggregate follows the + // delta to 7/140, not 5+7 / 100+140. + var kb entity.Knowledgebase + if err := db.First(&kb, "id = ?", "kb-1").Error; err != nil { + t.Fatalf("load kb: %v", err) + } + var doc entity.Document + if err := db.First(&doc, "id = ?", "doc-1").Error; err != nil { + t.Fatalf("load doc: %v", err) + } + if kb.ChunkNum != 7 || kb.TokenNum != 140 { + t.Errorf("kb = (chunk %d, token %d), want (7, 140)", kb.ChunkNum, kb.TokenNum) + } + if doc.ChunkNum != 7 || doc.TokenNum != 140 { + t.Errorf("doc = (chunk %d, token %d), want (7, 140)", doc.ChunkNum, doc.TokenNum) + } +} + +// TestApplyDocCounts_KBAggregateClampsAtZero checks that an aggregate driven +// below zero from an inconsistent starting state clamps to 0 instead of +// underflowing. +func TestApplyDocCounts_KBAggregateClampsAtZero(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + // The KB aggregate is inconsistent: below this document's own contribution. + if err := db.Create(&entity.Knowledgebase{ID: "kb-1", ChunkNum: 2, TokenNum: 40}).Error; err != nil { + t.Fatalf("create kb: %v", err) + } + if err := db.Create(&entity.Document{ID: "doc-1", KbID: "kb-1", ChunkNum: 5, TokenNum: 100, ParserConfig: entity.JSONMap{}}).Error; err != nil { + t.Fatalf("create doc: %v", err) + } + svc := testDocumentService(t) + + // A run that clears the document (0 chunks) drives the aggregate delta to -5; + // 2 - 5 is negative and must clamp to 0 rather than underflow. + if err := svc.ApplyDocCounts(context.Background(), "doc-1", "kb-1", 0, 0, 1); err != nil { + t.Fatalf("apply: %v", err) + } + var kb entity.Knowledgebase + if err := db.First(&kb, "id = ?", "kb-1").Error; err != nil { + t.Fatalf("load kb: %v", err) + } + if kb.ChunkNum != 0 || kb.TokenNum != 0 { + t.Errorf("kb = (chunk %d, token %d), want (0, 0) - aggregate must clamp, not underflow", kb.ChunkNum, kb.TokenNum) + } +} + +// TestApplyDocCounts_ProcessDurationIsAbsolute checks that process_duration is +// the last run's value, not a cumulative sum: a later run replaces the prior +// value rather than adding to it. +func TestApplyDocCounts_ProcessDurationIsAbsolute(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + if err := db.Create(&entity.Knowledgebase{ID: "kb-1"}).Error; err != nil { + t.Fatalf("create kb: %v", err) + } + if err := db.Create(&entity.Document{ID: "doc-1", KbID: "kb-1", ParserConfig: entity.JSONMap{}}).Error; err != nil { + t.Fatalf("create doc: %v", err) + } + svc := testDocumentService(t) + ctx := context.Background() + + // A run sets process_duration to its own value; a later run replaces it rather + // than accumulating, so the stored value is the last run's, not the 3.5+1.25 sum. + if err := svc.ApplyDocCounts(ctx, "doc-1", "kb-1", 5, 100, 3.5); err != nil { + t.Fatalf("first apply: %v", err) + } + if err := svc.ApplyDocCounts(ctx, "doc-1", "kb-1", 5, 100, 1.25); err != nil { + t.Fatalf("second apply: %v", err) + } + var doc entity.Document + if err := db.First(&doc, "id = ?", "doc-1").Error; err != nil { + t.Fatalf("load doc: %v", err) + } + if doc.ProcessDuration != 1.25 { + t.Errorf("process_duration = %v, want 1.25 (last run's value, not the 4.75 sum)", doc.ProcessDuration) + } +} diff --git a/internal/service/document/document_crud.go b/internal/service/document/document_crud.go index 3121d3df16..4be9956b68 100644 --- a/internal/service/document/document_crud.go +++ b/internal/service/document/document_crud.go @@ -13,6 +13,7 @@ import ( "ragflow/internal/storage" "gorm.io/gorm" + "gorm.io/gorm/clause" ) // Accessible reports whether docID belongs to a knowledge base @@ -142,31 +143,49 @@ func (s *DocumentService) UpdateDocument(ctx context.Context, id string, req *Up return s.documentDAO.Update(ctx, dao.DB, document) } -// IncrementChunkNum atomically increments chunk/token counters on the document and its knowledge base in a transaction -func (s *DocumentService) IncrementChunkNum(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error { +// ApplyDocCounts records a pipeline run's chunk/token/duration counts on the +// document and rolls the change into its knowledge base aggregate. It is +// idempotent per document: the document row holds this document's own counts, so +// re-applying the same run - e.g. an at-least-once task redelivery that re-runs +// the pipeline - sets the same document values and adjusts the knowledge base by +// a zero delta. It also carries a changed count (a re-parse producing a +// different result) correctly, since only the delta reaches the aggregate. The +// document row is locked so the read-modify-write of its current counts is +// atomic against a concurrent apply, and the aggregate is clamped at zero. +func (s *DocumentService) ApplyDocCounts(ctx context.Context, docID, kbID string, chunkNum, tokenNum int, duration float64) error { return dao.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - // Update document + var doc entity.Document + if err := tx.WithContext(ctx).Clauses(clause.Locking{Strength: "UPDATE"}). + Where("id = ? AND kb_id = ?", docID, kbID). + First(&doc).Error; err != nil { + return err + } + + chunkDelta := int64(chunkNum) - doc.ChunkNum + tokenDelta := int64(tokenNum) - doc.TokenNum + + // Set the document to this run's absolute counts. if err := tx.WithContext(ctx).Model(&entity.Document{}). Where("id = ? AND kb_id = ?", docID, kbID). Updates(map[string]interface{}{ - "chunk_num": gorm.Expr("chunk_num + ?", int64(chunkNum)), - "token_num": gorm.Expr("token_num + ?", int64(tokenNum)), - "process_duration": gorm.Expr("process_duration + ?", duration), + "chunk_num": int64(chunkNum), + "token_num": int64(tokenNum), + "process_duration": duration, }).Error; err != nil { return err } - // Update knowledgebase - if err := tx.WithContext(ctx).Model(&entity.Knowledgebase{}). + // Roll only the delta into the knowledge base aggregate; a re-applied run + // contributes zero. Clamp at zero so a stale aggregate cannot go negative. + if chunkDelta == 0 && tokenDelta == 0 { + return nil + } + return tx.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("id = ?", kbID). Updates(map[string]interface{}{ - "chunk_num": gorm.Expr("chunk_num + ?", int64(chunkNum)), - "token_num": gorm.Expr("token_num + ?", int64(tokenNum)), - }).Error; err != nil { - return err - } - - return nil + "chunk_num": gorm.Expr("CASE WHEN chunk_num + ? >= 0 THEN chunk_num + ? ELSE 0 END", chunkDelta, chunkDelta), + "token_num": gorm.Expr("CASE WHEN token_num + ? >= 0 THEN token_num + ? ELSE 0 END", tokenDelta, tokenDelta), + }).Error }) }