zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-12 06:35:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(go): clear task cancel signals and chunk counters on rerunWithDelete (#16544)

Browse Source
This commit is contained in:
Hz_
2026-07-02 15:46:11 +08:00
committed by GitHub
parent fa1b52ca74
commit cbb24944e8
2 changed files with 190 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
internal/service/document.go
View File

@@ -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(&current).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{}).

92
internal/service/document_test.go
View File

@@ -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)