From cbb24944e88c921b861b85d7c1b715c3ae0aabcd Mon Sep 17 00:00:00 2001 From: Hz_ Date: Thu, 2 Jul 2026 15:46:11 +0800 Subject: [PATCH] fix(go): clear task cancel signals and chunk counters on rerunWithDelete (#16544) --- internal/service/document.go | 101 +++++++++++++++++++++++++++++- internal/service/document_test.go | 92 +++++++++++++++++++++++++++ 2 files changed, 190 insertions(+), 3 deletions(-) diff --git a/internal/service/document.go b/internal/service/document.go index a1328a684a..76281fdced 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -54,6 +54,7 @@ import ( "github.com/google/uuid" "go.uber.org/zap" "gorm.io/gorm" + "gorm.io/gorm/clause" ) // DocumentService document service @@ -1235,8 +1236,8 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com } } - if rerunWithDelete && doc.Run != nil && *doc.Run == string(entity.TaskStatusDone) { - if err := s.clearKBChunkNumWhenRerun(doc); err != nil { + if rerunWithDelete { + if err := s.prepareDocumentRerunWithDelete(doc, kb.TenantID); err != nil { return common.CodeExceptionError, err } } @@ -1245,7 +1246,7 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com return common.CodeExceptionError, err } - if req.Delete { + if req.Delete && !rerunWithDelete { _, _ = s.taskDAO.DeleteByDocIDs([]string{doc.ID}) indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) if s.docEngine != nil { @@ -1325,6 +1326,100 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com return common.CodeSuccess, nil } +func (s *DocumentService) prepareDocumentRerunWithDelete(doc *entity.Document, tenantID string) error { + if doc == nil { + return fmt.Errorf("document is nil") + } + + s.cancelExistingParseTasksBestEffort(doc.ID) + + if _, err := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil { + return err + } + + if err := s.clearDocumentAndKBCountersForRerun(doc.ID, doc.KbID); err != nil { + return err + } + + if s.docEngine == nil { + return nil + } + + indexName := fmt.Sprintf("ragflow_%s", tenantID) + exists, err := s.docEngine.ChunkStoreExists(context.Background(), indexName, doc.KbID) + if err != nil { + return err + } + if !exists { + return nil + } + if _, err := s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": doc.ID}, indexName, doc.KbID); err != nil { + return err + } + return nil +} + +func (s *DocumentService) cancelExistingParseTasksBestEffort(docID string) { + tasks, err := s.taskDAO.GetByDocID(docID) + if err != nil { + common.Logger.Warn(fmt.Sprintf("cancelExistingParseTasksBestEffort: failed to get tasks for %s: %v", docID, err)) + return + } + redisClient := redis.Get() + if redisClient == nil { + return + } + for _, task := range tasks { + if task == nil { + continue + } + redisClient.Set(fmt.Sprintf("%s-cancel", task.ID), "x", 24*time.Hour) + } +} + +func (s *DocumentService) clearDocumentAndKBCountersForRerun(docID, kbID string) error { + return dao.DB.Transaction(func(tx *gorm.DB) error { + var current entity.Document + if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("id = ? AND kb_id = ?", docID, kbID). + First(¤t).Error; err != nil { + return err + } + + if current.TokenNum == 0 && current.ChunkNum == 0 && current.ProcessDuration == 0 { + return nil + } + + result := tx.Model(&entity.Document{}). + Where("id = ? AND kb_id = ?", docID, kbID). + Updates(map[string]interface{}{ + "token_num": 0, + "chunk_num": 0, + "process_duration": 0, + }) + if result.Error != nil { + return result.Error + } + if current.TokenNum == 0 && current.ChunkNum == 0 { + return nil + } + + result = tx.Model(&entity.Knowledgebase{}). + Where("id = ?", kbID). + Updates(map[string]interface{}{ + "token_num": gorm.Expr("token_num - ?", current.TokenNum), + "chunk_num": gorm.Expr("chunk_num - ?", current.ChunkNum), + }) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return fmt.Errorf("knowledgebase not found") + } + return nil + }) +} + func (s *DocumentService) countDoneDocuments(datasetID string) (int64, error) { var count int64 err := dao.GetDB().Model(&entity.Document{}). diff --git a/internal/service/document_test.go b/internal/service/document_test.go index 3c5889d652..7b11282834 100644 --- a/internal/service/document_test.go +++ b/internal/service/document_test.go @@ -190,6 +190,26 @@ type failingDeleteMetadataEngine struct { updateCalled bool } +type rerunDeleteDocEngine struct { + fakeChatDocEngine + deleteCalls int + condition map[string]interface{} + indexName string + datasetID string +} + +func (e *rerunDeleteDocEngine) ChunkStoreExists(context.Context, string, string) (bool, error) { + return true, nil +} + +func (e *rerunDeleteDocEngine) DeleteChunks(_ context.Context, condition map[string]interface{}, indexName string, datasetID string) (int64, error) { + e.deleteCalls++ + e.condition = condition + e.indexName = indexName + e.datasetID = datasetID + return 3, nil +} + type metadataDocEngine struct { fakeChatDocEngine records map[string]map[string]interface{} @@ -1686,6 +1706,78 @@ func TestResetDocumentForReparseSkipsSecondCounterDecrement(t *testing.T) { } } +func TestPrepareDocumentRerunWithDeleteClearsCountersTasksAndChunks(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) + insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5) + insertTestTask(t, "task-1", "doc-1") + + doc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + + engine := &rerunDeleteDocEngine{} + svc := testDocumentService(t) + svc.docEngine = engine + + if err := svc.prepareDocumentRerunWithDelete(doc, "tenant-1"); err != nil { + t.Fatalf("prepareDocumentRerunWithDelete failed: %v", err) + } + + updatedDoc, _ := dao.NewDocumentDAO().GetByID("doc-1") + if updatedDoc.TokenNum != 0 || updatedDoc.ChunkNum != 0 { + t.Fatalf("doc counters = token:%d chunk:%d, want zero", updatedDoc.TokenNum, updatedDoc.ChunkNum) + } + kb, _ := dao.NewKnowledgebaseDAO().GetByID("kb-1") + if kb.TokenNum != 0 || kb.ChunkNum != 0 { + t.Fatalf("kb counters = token:%d chunk:%d, want zero", kb.TokenNum, kb.ChunkNum) + } + var taskCount int64 + if err := dao.DB.Model(&entity.Task{}).Where("doc_id = ?", "doc-1").Count(&taskCount).Error; err != nil { + t.Fatalf("count tasks: %v", err) + } + if taskCount != 0 { + t.Fatalf("task count = %d, want zero", taskCount) + } + if engine.deleteCalls != 1 { + t.Fatalf("deleteCalls = %d, want 1", engine.deleteCalls) + } + if engine.indexName != "ragflow_tenant-1" || engine.datasetID != "kb-1" || engine.condition["doc_id"] != "doc-1" { + t.Fatalf("unexpected delete call: index=%s dataset=%s condition=%v", engine.indexName, engine.datasetID, engine.condition) + } +} + +func TestPrepareDocumentRerunWithDeleteIsIdempotentForStaleDocSnapshot(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) + insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5) + + staleDoc, err := dao.NewDocumentDAO().GetByID("doc-1") + if err != nil { + t.Fatalf("get doc: %v", err) + } + + svc := testDocumentService(t) + if err := svc.prepareDocumentRerunWithDelete(staleDoc, "tenant-1"); err != nil { + t.Fatalf("first prepareDocumentRerunWithDelete failed: %v", err) + } + if err := svc.prepareDocumentRerunWithDelete(staleDoc, "tenant-1"); err != nil { + t.Fatalf("second prepareDocumentRerunWithDelete failed: %v", err) + } + + doc, _ := dao.NewDocumentDAO().GetByID("doc-1") + if doc.TokenNum != 0 || doc.ChunkNum != 0 { + t.Fatalf("doc counters = token:%d chunk:%d, want zero", doc.TokenNum, doc.ChunkNum) + } + kb, _ := dao.NewKnowledgebaseDAO().GetByID("kb-1") + if kb.TokenNum != 0 || kb.ChunkNum != 0 { + t.Fatalf("kb counters = token:%d chunk:%d, want zero after duplicate prepare", kb.TokenNum, kb.ChunkNum) + } +} + func TestUpdateDatasetDocumentPropagatesMetadataDeleteFailure(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db)