diff --git a/internal/application/service/image_multimodal.go b/internal/application/service/image_multimodal.go index 0f00e0121..215d0de53 100644 --- a/internal/application/service/image_multimodal.go +++ b/internal/application/service/image_multimodal.go @@ -151,6 +151,9 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e logger.Infof(ctx, "[ImageMultimodal] Dropping task chunk=%s knowledge=%s kb=%s image=%s", payload.ChunkID, payload.KnowledgeID, payload.KnowledgeBaseID, payload.ImageURL) + // Still count this image toward the parent finalize gate so a batch + // of dropped orphans cannot strand multimodal:pending forever. + s.checkAndFinalizeAllImages(ctx, payload) return nil } diff --git a/internal/application/service/image_multimodal_orphan_test.go b/internal/application/service/image_multimodal_orphan_test.go index aa7eeb69e..3154dcff1 100644 --- a/internal/application/service/image_multimodal_orphan_test.go +++ b/internal/application/service/image_multimodal_orphan_test.go @@ -9,7 +9,9 @@ import ( "github.com/Tencent/WeKnora/internal/application/repository" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/alicebob/miniredis/v2" "github.com/hibiken/asynq" + "github.com/redis/go-redis/v9" ) type orphanKnowledgeRepo struct { @@ -96,6 +98,58 @@ func TestImageMultimodalHandleDropsMissingKnowledge(t *testing.T) { } } +type orphanTaskEnqueuer struct { + interfaces.TaskEnqueuer + enqueued []*asynq.Task +} + +func (e *orphanTaskEnqueuer) Enqueue(task *asynq.Task, _ ...asynq.Option) (*asynq.TaskInfo, error) { + e.enqueued = append(e.enqueued, task) + return &asynq.TaskInfo{ID: "post-process"}, nil +} + +func TestImageMultimodalHandleDropFinalizesPendingCounter(t *testing.T) { + t.Parallel() + mr, err := miniredis.Run() + if err != nil { + t.Fatalf("start miniredis: %v", err) + } + t.Cleanup(mr.Close) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + + enqueuer := &orphanTaskEnqueuer{} + svc := &ImageMultimodalService{ + knowledgeRepo: &orphanKnowledgeRepo{err: repository.ErrKnowledgeNotFound}, + kbService: &orphanKBService{kb: &types.KnowledgeBase{ID: "kb-1"}}, + redisClient: rdb, + taskEnqueuer: enqueuer, + } + const knowledgeID = "missing" + redisKey := "multimodal:pending:" + knowledgeID + if err := rdb.Set(context.Background(), redisKey, 1, 0).Err(); err != nil { + t.Fatal(err) + } + payload, err := json.Marshal(types.ImageMultimodalPayload{ + TenantID: 1, + KnowledgeID: knowledgeID, + KnowledgeBaseID: "kb-1", + ImageURL: "minio://bucket/img.png", + }) + if err != nil { + t.Fatal(err) + } + if err := svc.Handle(context.Background(), asynq.NewTask(types.TypeImageMultimodal, payload)); err != nil { + t.Fatalf("orphan drop should succeed: %v", err) + } + if mr.Exists(redisKey) { + t.Fatal("pending counter should be cleared after drop finalize") + } + if len(enqueuer.enqueued) != 1 || enqueuer.enqueued[0].Type() != types.TypeKnowledgePostProcess { + t.Fatalf("expected post-process enqueue, got %d tasks", len(enqueuer.enqueued)) + } +} + func TestImageMultimodalHandlePropagatesTransientKnowledgeError(t *testing.T) { t.Parallel() dbErr := errors.New("db unavailable") diff --git a/internal/handler/system.go b/internal/handler/system.go index 98a31fe62..01d135563 100644 --- a/internal/handler/system.go +++ b/internal/handler/system.go @@ -1827,29 +1827,31 @@ func isKnowledgeGone(err error) bool { } // purgeOrphanRuntimeTask removes queue records for a task whose knowledge row -// is already gone. Best-effort: CancelTasksForKnowledge sweeps every queue -// state; ForceDeleteRuntimeTask is the fallback for the specific task ID. +// is already gone. CancelTasksForKnowledge sweeps sibling tasks best-effort; +// success is reported only once the requested task ID is gone. func (h *SystemHandler) purgeOrphanRuntimeTask( ctx context.Context, inspector interfaces.RuntimeTaskInspector, queue, taskID, knowledgeID string, ) bool { if h.taskInspector != nil && knowledgeID != "" { - deleted, _, err := h.taskInspector.CancelTasksForKnowledge(ctx, knowledgeID) - if err != nil { + if _, _, err := h.taskInspector.CancelTasksForKnowledge(ctx, knowledgeID); err != nil { logger.Warnf(ctx, "purge orphan tasks for knowledge %s: %v", knowledgeID, err) - } else if deleted > 0 { - return true } } supported, err := inspector.ForceDeleteRuntimeTask(ctx, queue, taskID) - if !supported || err != nil { - if err != nil { - logger.Warnf(ctx, "force delete orphan task queue=%s task=%s: %v", queue, taskID, err) - } + if !supported { return false } - return true + if err == nil { + return true + } + // The knowledge-wide sweep may have already removed this task ID. + if _, available, getErr := inspector.GetRuntimeTask(ctx, queue, taskID); available && getErr != nil { + return true + } + logger.Warnf(ctx, "force delete orphan task queue=%s task=%s: %v", queue, taskID, err) + return false } // MutateRuntimeTask executes a backend-advertised action after re-reading the diff --git a/internal/handler/system_runtime_test.go b/internal/handler/system_runtime_test.go index ef1728e9c..664466d79 100644 --- a/internal/handler/system_runtime_test.go +++ b/internal/handler/system_runtime_test.go @@ -3,6 +3,7 @@ package handler import ( "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "testing" @@ -85,6 +86,10 @@ type runtimeTaskTestInspector struct { cancelKnowledge string cancelDeleted int mutatedQueue string + forceDeleteErr error + getRuntimeErr error + getRuntimeErrFrom int + getRuntimeTaskCalls int } func (r *runtimeTaskTestInspector) CancelTasksForKnowledge( @@ -134,6 +139,10 @@ func (r *runtimeTaskTestInspector) ListRuntimeTasks( func (r *runtimeTaskTestInspector) GetRuntimeTask( _ context.Context, queue, taskID string, ) (*types.RuntimeTaskInfo, bool, error) { + r.getRuntimeTaskCalls++ + if r.getRuntimeErr != nil && r.getRuntimeTaskCalls >= r.getRuntimeErrFrom { + return nil, true, r.getRuntimeErr + } for i := range r.tasks { if r.tasks[i].ID == taskID { return &r.tasks[i], true, nil @@ -169,7 +178,7 @@ func (r *runtimeTaskTestInspector) ForceDeleteRuntimeTask( ) (bool, error) { r.mutatedQueue = queue r.forceDeleted = taskID - return true, nil + return true, r.forceDeleteErr } func TestGetRuntimeQueuesReportsIsolatedPoolCapacity(t *testing.T) { @@ -401,3 +410,67 @@ func TestRuntimeTaskCancelPurgesOrphanWhenKnowledgeGone(t *testing.T) { inspector.deletedTask, inspector.forceDeleted, inspector.cancelKnowledge) } } + +func TestRuntimeTaskCancelPurgesOrphanAfterSiblingSweep(t *testing.T) { + gin.SetMode(gin.TestMode) + inspector := &runtimeTaskTestInspector{ + cancelDeleted: 2, + tasks: []types.RuntimeTaskInfo{{ + ID: "task-orphan", Queue: types.QueueMultimodal, Type: types.TypeImageMultimodal, + State: types.RuntimeTaskArchived, TenantID: 42, KnowledgeID: "knowledge-gone", + AllowedActions: []types.RuntimeTaskAction{types.RuntimeTaskActionCancel}, + }}, + } + canceller := &runtimeKnowledgeCancelTest{err: repository.ErrKnowledgeNotFound} + handler := &SystemHandler{taskInspector: inspector, knowledgeSvc: canceller} + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Params = gin.Params{ + {Key: "queue", Value: types.QueueMultimodal}, + {Key: "task_id", Value: "task-orphan"}, + {Key: "action", Value: string(types.RuntimeTaskActionCancel)}, + } + ctx.Request = httptest.NewRequest(http.MethodPost, "/cancel", nil) + + handler.MutateRuntimeTask(ctx) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + if inspector.cancelKnowledge != "knowledge-gone" || inspector.forceDeleted != "task-orphan" { + t.Fatalf("expected sibling sweep then force delete, got cancel=%q force=%q", + inspector.cancelKnowledge, inspector.forceDeleted) + } +} + +func TestRuntimeTaskCancelPurgesOrphanWhenSweepAlreadyRemovedTask(t *testing.T) { + gin.SetMode(gin.TestMode) + inspector := &runtimeTaskTestInspector{ + cancelDeleted: 1, + forceDeleteErr: errors.New("task not found"), + getRuntimeErr: errors.New("task not found"), + getRuntimeErrFrom: 2, + tasks: []types.RuntimeTaskInfo{{ + ID: "task-orphan", Queue: types.QueueMultimodal, Type: types.TypeImageMultimodal, + State: types.RuntimeTaskRetry, TenantID: 42, KnowledgeID: "knowledge-gone", + AllowedActions: []types.RuntimeTaskAction{types.RuntimeTaskActionCancel}, + }}, + } + canceller := &runtimeKnowledgeCancelTest{err: repository.ErrKnowledgeNotFound} + handler := &SystemHandler{taskInspector: inspector, knowledgeSvc: canceller} + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Params = gin.Params{ + {Key: "queue", Value: types.QueueMultimodal}, + {Key: "task_id", Value: "task-orphan"}, + {Key: "action", Value: string(types.RuntimeTaskActionCancel)}, + } + ctx.Request = httptest.NewRequest(http.MethodPost, "/cancel", nil) + + handler.MutateRuntimeTask(ctx) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body=%s", recorder.Code, recorder.Body.String()) + } + if inspector.forceDeleted != "task-orphan" { + t.Fatalf("expected force delete attempt, got %q", inspector.forceDeleted) + } +}