task-73: 本地文件作业队列 in-flight 去重和队列背压

2026-08-30 20:22:58 +08:00
parent f190896ea7
commit ec9811c8a5
2 changed files with 233 additions and 2 deletions
@@ -11,6 +11,9 @@ import org.springframework.beans.factory.annotation.Value;
 import org.springframework.core.task.TaskExecutor;
 import org.springframework.stereotype.Service;

+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+
@Slf4j
@Service
@RequiredArgsConstructor
@@ -26,6 +29,15 @@ public class TaskFileJobLocalDispatcher {
    @Value("${aiimage.result-file-job.local-dispatch-enabled:true}")
    private boolean localDispatchEnabled;

+    /**
+     * 已受理未完成的 jobId 集合：同 job 重复 dispatch 幂等返回；
+     * 数量达到上限时拒绝新提交（背压），防止执行队列无界堆积。
+     */
+    @Value("${aiimage.result-file-job.max-inflight-dispatch:64}")
+    private int maxInflightDispatch = 64;
+
+    private final Set<Long> inflightJobIds = ConcurrentHashMap.newKeySet();
+
    public boolean dispatch(Long jobId, Long taskId, String moduleType) {
        return dispatch(jobId, taskId, moduleType, false);
    }
@@ -39,15 +51,31 @@ public class TaskFileJobLocalDispatcher {
                    jobId, taskId, moduleType);
            return false;
        }
+        int inflightLimit = Math.max(1, maxInflightDispatch);
+        if (!inflightJobIds.contains(jobId) && inflightJobIds.size() >= inflightLimit) {
+            log.warn("[task-file-job] local dispatch backpressure, inflight limit reached jobId={} taskId={} moduleType={} inflight={} limit={}",
+                    jobId, taskId, moduleType, inflightJobIds.size(), inflightLimit);
+            return false;
+        }
+        if (!inflightJobIds.add(jobId)) {
+            log.info("[task-file-job] local dispatch skipped, job already inflight jobId={} taskId={} moduleType={}",
+                    jobId, taskId, moduleType);
+            return true;
+        }
        try {
            taskFileJobDispatchExecutor.execute(() -> processLocally(jobId, taskId, moduleType));
            return true;
        } catch (RuntimeException ex) {
+            inflightJobIds.remove(jobId);
            log.warn("[task-file-job] local dispatch executor rejected jobId={} taskId={} moduleType={} msg={}",
                    jobId, taskId, moduleType, ex.getMessage(), ex);
            if (force) {
-                processLocally(jobId, taskId, moduleType);
-                return true;
+                try {
+                    processLocally(jobId, taskId, moduleType);
+                    return true;
+                } finally {
+                    inflightJobIds.remove(jobId);
+                }
            }
            return false;
        }
@@ -71,6 +99,8 @@ public class TaskFileJobLocalDispatcher {
        } catch (Exception ex) {
            log.warn("[task-file-job] local dispatch failed jobId={} taskId={} moduleType={} msg={}",
                    jobId, taskId, moduleType, ex.getMessage(), ex);
+        } finally {
+            inflightJobIds.remove(jobId);
        }
    }
 }
@@ -0,0 +1,201 @@
+package com.nanri.aiimage.modules.task.service;
+
+import com.nanri.aiimage.modules.task.mapper.TaskFileJobMapper;
+import com.nanri.aiimage.modules.task.model.entity.TaskFileJobEntity;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.core.task.TaskExecutor;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Task 73：为本地文件作业队列增加 in-flight 去重和队列背压。
+ * TaskFileJobLocalDispatcher 以 in-flight 集合跟踪已受理未完成的 job：
+ * 同一 jobId 重复 dispatch 幂等返回（不重复入队）；in-flight 数量达到
+ * max-inflight-dispatch 上限时拒绝新提交（背压，调用方走同步兜底）；
+ * 执行完成或失败后释放 in-flight 名额；executor 拒绝时名额一并释放。
+ */
+class TaskFileJobLocalDispatcherTest {
+
+    private TaskFileJobMapper taskFileJobMapper;
+    private ObjectProvider<TaskResultFileJobWorker> workerProvider;
+    private TaskResultFileJobWorker worker;
+    private TaskExecutor executor;
+    private TaskFileJobLocalDispatcher dispatcher;
+
+    @BeforeEach
+    @SuppressWarnings("unchecked")
+    void setUp() {
+        taskFileJobMapper = mock(TaskFileJobMapper.class);
+        worker = mock(TaskResultFileJobWorker.class);
+        workerProvider = mock(ObjectProvider.class);
+        when(workerProvider.getIfAvailable()).thenReturn(worker);
+        executor = mock(TaskExecutor.class);
+        dispatcher = new TaskFileJobLocalDispatcher(taskFileJobMapper, workerProvider);
+        ReflectionTestUtils.setField(dispatcher, "taskFileJobDispatchExecutor", executor);
+        ReflectionTestUtils.setField(dispatcher, "maxInflightDispatch", 8);
+        // @Value 注解在纯单测中不生效，boolean 字段默认 false，需显式开启本地派发。
+        ReflectionTestUtils.setField(dispatcher, "localDispatchEnabled", true);
+    }
+
+    private TaskFileJobEntity job(Long id) {
+        TaskFileJobEntity job = new TaskFileJobEntity();
+        job.setId(id);
+        job.setTaskId(100L + id);
+        job.setModuleType("SIMILAR_ASIN");
+        return job;
+    }
+
+    private void stubJobFound(Long id) {
+        when(taskFileJobMapper.selectById(id)).thenReturn(job(id));
+    }
+
+    /** executor 捕获 runnable 不执行，模拟任务仍在 in-flight。 */
+    private List<Runnable> captureRunnables() {
+        List<Runnable> captured = new ArrayList<>();
+        doAnswer(invocation -> {
+            captured.add(invocation.getArgument(0));
+            return null;
+        }).when(executor).execute(any(Runnable.class));
+        return captured;
+    }
+
+    @Test
+    void test_task_073_dispatch_normal_default_path() {
+        // 默认路径：正常 dispatch 提交到执行器并返回 true。
+        stubJobFound(1L);
+        captureRunnables();
+
+        assertTrue(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN"));
+        verify(executor, times(1)).execute(any(Runnable.class));
+    }
+
+    @Test
+    void test_task_073_dispatch_normal_multiple_items() throws Exception {
+        // 批量场景：多个不同 jobId 依次 dispatch，互不冲突，全部受理。
+        stubJobFound(1L);
+        stubJobFound(2L);
+        stubJobFound(3L);
+        List<Runnable> captured = captureRunnables();
+
+        assertTrue(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN"));
+        assertTrue(dispatcher.dispatch(2L, 102L, "SIMILAR_ASIN"));
+        assertTrue(dispatcher.dispatch(3L, 103L, "SIMILAR_ASIN"));
+        assertEquals(3, captured.size(), "三个不同 job 都入队");
+
+        for (Runnable runnable : captured) {
+            runnable.run();
+        }
+        verify(worker, times(3)).process(any(TaskFileJobEntity.class));
+    }
+
+    @Test
+    void test_task_073_dispatch_normal_repeated_operation_is_idempotent() throws Exception {
+        // 幂等：同一 jobId 在 in-flight 期间重复 dispatch 不重复入队，返回 true 表示已受理。
+        stubJobFound(1L);
+        List<Runnable> captured = captureRunnables();
+
+        assertTrue(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN"));
+        assertTrue(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN"), "重复 dispatch 幂等受理");
+        assertEquals(1, captured.size(), "同一 job 只入队一次");
+
+        captured.get(0).run();
+        assertTrue(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN"), "完成后可再次受理");
+        assertEquals(2, captured.size());
+        captured.get(1).run();
+        verify(worker, times(2)).process(any(TaskFileJobEntity.class));
+    }
+
+    @Test
+    void test_task_073_dispatch_boundary_empty_input() {
+        // 空输入：null jobId 直接拒绝，不提交、不计数。
+        captureRunnables();
+
+        assertFalse(dispatcher.dispatch(null, null, null));
+        assertFalse(dispatcher.dispatch(0L, null, null));
+        verify(executor, never()).execute(any(Runnable.class));
+    }
+
+    @Test
+    void test_task_073_dispatch_boundary_single_item() throws Exception {
+        // 单元素：单 jobId dispatch 一次执行一次，不依赖批量路径。
+        stubJobFound(9L);
+        List<Runnable> captured = captureRunnables();
+
+        assertTrue(dispatcher.dispatch(9L, 109L, "SIMILAR_ASIN"));
+        captured.get(0).run();
+        verify(worker).process(any(TaskFileJobEntity.class));
+    }
+
+    @Test
+    void test_task_073_dispatch_boundary_limit_and_overflow() throws Exception {
+        // 上限/超限：in-flight 达到上限后新提交被拒绝（背压）；
+        // 先完成的 job 释放名额后再次提交成功。
+        stubJobFound(1L);
+        stubJobFound(2L);
+        stubJobFound(3L);
+        ReflectionTestUtils.setField(dispatcher, "maxInflightDispatch", 2);
+        List<Runnable> captured = captureRunnables();
+
+        assertTrue(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN"));
+        assertTrue(dispatcher.dispatch(2L, 102L, "SIMILAR_ASIN"));
+        assertFalse(dispatcher.dispatch(3L, 103L, "SIMILAR_ASIN"), "达到上限拒绝，背压生效");
+        assertEquals(2, captured.size(), "被拒绝的 job 不入队");
+
+        captured.get(0).run();
+        assertTrue(dispatcher.dispatch(3L, 103L, "SIMILAR_ASIN"), "名额释放后可再次提交");
+        assertEquals(3, captured.size());
+    }
+
+    @Test
+    void test_task_073_dispatch_invalid_input_rejected() {
+        // 非法参数：jobId<=0 直接拒绝；非法上限配置回退到 1（不崩溃、行为确定）。
+        stubJobFound(1L);
+        ReflectionTestUtils.setField(dispatcher, "maxInflightDispatch", 0);
+        captureRunnables();
+
+        assertFalse(dispatcher.dispatch(-1L, null, null), "非法 jobId 拒绝");
+        assertTrue(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN"), "上限 0 回退到 1，首个受理");
+        assertFalse(dispatcher.dispatch(2L, 102L, "SIMILAR_ASIN"), "回退上限已满，背压拒绝");
+        verify(executor, times(1)).execute(any(Runnable.class));
+    }
+
+    @Test
+    void test_task_073_dispatch_dependency_failure_releases_resources() {
+        // 依赖失败：executor 拒绝时释放 in-flight 名额，恢复后再次 dispatch 成功；
+        // force 路径同步兜底执行且不残留名额。
+        stubJobFound(1L);
+        stubJobFound(2L);
+        doThrow(new RuntimeException("executor full"))
+                .doNothing()
+                .doThrow(new RuntimeException("executor full"))
+                .doNothing()
+                .when(executor).execute(any(Runnable.class));
+
+        assertFalse(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN"), "executor 拒绝返回 false");
+        assertTrue(dispatcher.dispatch(2L, 102L, "SIMILAR_ASIN"), "名额已释放，再次提交成功");
+        verify(worker, never()).process(any(TaskFileJobEntity.class));
+
+        assertTrue(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN", true), "force 路径同步兜底");
+        verify(worker, times(1)).process(any(TaskFileJobEntity.class));
+
+        assertTrue(dispatcher.dispatch(1L, 101L, "SIMILAR_ASIN"), "force 执行完不残留 in-flight");
+        verify(executor, times(4)).execute(any(Runnable.class));
+    }
+}