mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-18 22:38:25 +08:00
@@ -617,6 +617,64 @@ func TestNATSTaskWorkerAcksUnclaimableMessage(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestNATSTaskWorkerRetriesStillScheduledUnclaimedTask(t *testing.T) {
|
||||
db := setupSyncerDB(t)
|
||||
insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
|
||||
insertSyncLog(t, db, "conn-1", "kb-1", "task-2", dao.TaskTypePrune)
|
||||
if err := db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusRunning).Error; err != nil {
|
||||
t.Fatalf("mark sync running: %v", err)
|
||||
}
|
||||
if err := db.Model(&entity.Connector{}).Where("id = ?", "conn-1").Update("status", dao.SyncStatusRunning).Error; err != nil {
|
||||
t.Fatalf("mark connector running: %v", err)
|
||||
}
|
||||
|
||||
taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
|
||||
scheduler := NewNATSScheduler(make(chan TaskEnvelope, 1), taskService, &fakeSyncTaskBroker{})
|
||||
worker := NewTaskWorker(make(chan TaskEnvelope, 1), taskService, newCoordinator(taskService, newTestRegistry(nil), &fakeSink{}, nil, fakeStore{}), NewConnectorLock()).WithScheduler(scheduler)
|
||||
handle := &fakeTaskHandle{msg: common.TaskMessage{TaskID: "task-2", TaskType: common.TaskTypeSyncer}}
|
||||
|
||||
worker.handle(t.Context(), TaskEnvelope{TaskID: "task-2", Handle: handle})
|
||||
|
||||
if handle.acks != 1 || handle.nacks != 0 {
|
||||
t.Fatalf("settlement acks=%d nacks=%d", handle.acks, handle.nacks)
|
||||
}
|
||||
scheduler.timerMu.Lock()
|
||||
timer := scheduler.timers["task-2"]
|
||||
if timer != nil {
|
||||
timer.Stop()
|
||||
delete(scheduler.timers, "task-2")
|
||||
}
|
||||
scheduler.timerMu.Unlock()
|
||||
if timer == nil {
|
||||
t.Fatalf("scheduled prune retry timer was not registered")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNATSTaskWorkerDoesNotRetryCompletedUnclaimedTask(t *testing.T) {
|
||||
db := setupSyncerDB(t)
|
||||
insertTaskContext(t, db, "conn-1", "kb-1", "task-1", dao.TaskTypeSync)
|
||||
if err := db.Model(&entity.SyncLogs{}).Where("id = ?", "task-1").Update("status", dao.SyncStatusDone).Error; err != nil {
|
||||
t.Fatalf("mark done: %v", err)
|
||||
}
|
||||
|
||||
taskService := service.NewSyncTaskService(dao.NewSyncTaskDAO(db))
|
||||
scheduler := NewNATSScheduler(make(chan TaskEnvelope, 1), taskService, &fakeSyncTaskBroker{})
|
||||
worker := NewTaskWorker(make(chan TaskEnvelope, 1), taskService, newCoordinator(taskService, newTestRegistry(nil), &fakeSink{}, nil, fakeStore{}), NewConnectorLock()).WithScheduler(scheduler)
|
||||
handle := &fakeTaskHandle{msg: common.TaskMessage{TaskID: "task-1", TaskType: common.TaskTypeSyncer}}
|
||||
|
||||
worker.handle(t.Context(), TaskEnvelope{TaskID: "task-1", Handle: handle})
|
||||
|
||||
if handle.acks != 1 || handle.nacks != 0 {
|
||||
t.Fatalf("settlement acks=%d nacks=%d", handle.acks, handle.nacks)
|
||||
}
|
||||
scheduler.timerMu.Lock()
|
||||
_, scheduled := scheduler.timers["task-1"]
|
||||
scheduler.timerMu.Unlock()
|
||||
if scheduled {
|
||||
t.Fatalf("completed task should not be scheduled for retry")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSameConnectorDifferentKBsRunInParallel verifies one datasource can sync into different KBs concurrently.
|
||||
func TestSameConnectorDifferentKBsRunInParallel(t *testing.T) {
|
||||
db := setupSyncerDB(t)
|
||||
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/service"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -91,6 +92,7 @@ func (w *TaskWorker) handle(ctx context.Context, envelope TaskEnvelope) {
|
||||
}
|
||||
if !claimed {
|
||||
_ = envelope.Handle.Ack() // this task has been claimed by other worker
|
||||
w.scheduleRetryIfTaskScheduled(ctx, envelope.TaskID, 3*time.Second)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -187,6 +189,21 @@ func (w *TaskWorker) scheduleRetry(ctx context.Context, taskID string, delay tim
|
||||
}
|
||||
}
|
||||
|
||||
func (w *TaskWorker) scheduleRetryIfTaskScheduled(ctx context.Context, taskID string, delay time.Duration) {
|
||||
if w.scheduler == nil || taskID == "" {
|
||||
return
|
||||
}
|
||||
taskContext, err := w.taskService.GetContext(ctx, taskID)
|
||||
if err != nil {
|
||||
common.Warn("syncer retry task lookup failed", zap.String("task_id", taskID), zap.Error(err))
|
||||
return
|
||||
}
|
||||
if taskContext.Task.Status != dao.SyncStatusSchedule {
|
||||
return
|
||||
}
|
||||
w.scheduleRetry(ctx, taskID, delay)
|
||||
}
|
||||
|
||||
// transientRetryDelay return retry delay
|
||||
func transientRetryDelay(attempts int64) time.Duration {
|
||||
if attempts < 1 {
|
||||
|
||||
Reference in New Issue
Block a user