task-66: 限制 RustFS 并发读写与重试的总资源预算

- 新增 aiimage.transient-storage.max-total-concurrent-operations 配置（默认 0 不启用）：跨 upload/read/delete 的总并发预算 - executeWithRetry 入口获取一个总许可，整次操作（含全部重试）全程占用，重试不额外消耗，超预算立即拒绝并计入 operation=total rejected 指标 - 失败/完成后 finally 释放总许可，无残留锁；负值/0 视为未启用保持原行为
2026-08-30 18:57:26 +08:00
parent e377814c48
commit e2d6875ef9
3 changed files with 331 additions and 33 deletions
@@ -23,6 +23,11 @@ public class TransientStorageProperties {
    private int maxConcurrentUploads = 16;
    private int maxConcurrentReads = 32;
    private int maxConcurrentDeletes = 8;
+    /**
+     * 跨操作类型的总并发预算（含重试期间）：一次操作全程占用一个总许可，
+     * 防止多任务叠加时读写与重试之和突破对后端的总压力上限。0/负值表示不启用。
+     */
+    private long maxTotalConcurrentOperations = 0;
    private long acquirePermitTimeoutMillis = 2000;
    private long baseRetryDelayMillis = 500;
    private long maxRetryDelayMillis = 5000;
@@ -35,6 +35,7 @@ public class RustfsObjectStorageService {
    private static final String OP_READ = "read";
    private static final String OP_DELETE = "delete";
    private static final String OP_STAT = "stat";
+    private static final String OP_TOTAL = "total";

    private final TransientStorageProperties properties;
    private final ObjectProvider<MeterRegistry> meterRegistryProvider;
@@ -43,6 +44,7 @@ public class RustfsObjectStorageService {
    private final Semaphore uploadSemaphore;
    private final Semaphore readSemaphore;
    private final Semaphore deleteSemaphore;
+    private final Semaphore totalSemaphore;
    private final AtomicInteger windowFailureCount = new AtomicInteger();
    private volatile long failureWindowStartedAtMillis;
    private volatile long circuitOpenUntilMillis;
@@ -66,6 +68,13 @@ public class RustfsObjectStorageService {
        this.uploadSemaphore = new Semaphore(Math.max(1, properties.getMaxConcurrentUploads()));
        this.readSemaphore = new Semaphore(Math.max(1, properties.getMaxConcurrentReads()));
        this.deleteSemaphore = new Semaphore(Math.max(1, properties.getMaxConcurrentDeletes()));
+        // 总预算为 0/负值时不启用（Semaphore(0) 的 tryAcquire 永远失败，须用启用开关区分）
+        long totalBudget = Math.max(0L, properties.getMaxTotalConcurrentOperations());
+        this.totalSemaphore = new Semaphore((int) Math.min(totalBudget, Integer.MAX_VALUE));
+    }
+
+    private boolean isTotalBudgetEnabled() {
+        return properties.getMaxTotalConcurrentOperations() > 0L;
    }

    public boolean isConfigured() {
@@ -185,45 +194,57 @@ public class RustfsObjectStorageService {
                                   CheckedSupplier<T> supplier) {
        long startedAt = System.nanoTime();
        Exception last = null;
-        for (int attempt = 1; attempt <= maxRetries; attempt++) {
-            checkDeadline(operation, objectKey, deadlineNanos);
-            rejectIfCircuitOpen(operation, objectKey);
-            acquirePermit(operation, objectKey, semaphore, deadlineNanos);
-            long delayMillis = 0L;
-            try {
+        // 总资源预算：一次操作（含全部重试）全程占用一个总许可，重试不额外消耗。
+        boolean totalAcquired = false;
+        if (isTotalBudgetEnabled()) {
+            acquirePermit(OP_TOTAL, objectKey, totalSemaphore, deadlineNanos);
+            totalAcquired = true;
+        }
+        try {
+            for (int attempt = 1; attempt <= maxRetries; attempt++) {
+                checkDeadline(operation, objectKey, deadlineNanos);
+                rejectIfCircuitOpen(operation, objectKey);
+                acquirePermit(operation, objectKey, semaphore, deadlineNanos);
+                long delayMillis = 0L;
                try {
-                    checkDeadline(operation, objectKey, deadlineNanos);
-                    T result = supplier.get();
-                    checkDeadline(operation, objectKey, deadlineNanos);
-                    if (!OP_UPLOAD.equals(operation)) {
-                        resetFailureWindow(operation);
-                    }
-                    recordOperation(operation, "success", elapsedNanos(startedAt));
-                    log.debug("[rustfs] operation success operation={} objectKey={} attempt={}/{} durationMs={} bytes={}",
-                            operation, objectKey, attempt, maxRetries, elapsedMillis(startedAt), bytes);
-                    return result;
-                } catch (Exception ex) {
-                    last = ex;
-                    recordOperation(operation, attempt < maxRetries ? "retry" : "failure", elapsedNanos(startedAt));
-                    recordFailure(operation, objectKey, ex);
-                    if (attempt < maxRetries) {
-                        delayMillis = retryDelayMillis(attempt);
-                        log.warn("[rustfs] operation failed, retrying operation={} objectKey={} attempt={}/{} delayMs={} err={}",
-                                operation, objectKey, attempt, maxRetries, delayMillis, ex.getMessage());
+                    try {
+                        checkDeadline(operation, objectKey, deadlineNanos);
+                        T result = supplier.get();
+                        checkDeadline(operation, objectKey, deadlineNanos);
+                        if (!OP_UPLOAD.equals(operation)) {
+                            resetFailureWindow(operation);
+                        }
+                        recordOperation(operation, "success", elapsedNanos(startedAt));
+                        log.debug("[rustfs] operation success operation={} objectKey={} attempt={}/{} durationMs={} bytes={}",
+                                operation, objectKey, attempt, maxRetries, elapsedMillis(startedAt), bytes);
+                        return result;
+                    } catch (Exception ex) {
+                        last = ex;
+                        recordOperation(operation, attempt < maxRetries ? "retry" : "failure", elapsedNanos(startedAt));
+                        recordFailure(operation, objectKey, ex);
+                        if (attempt < maxRetries) {
+                            delayMillis = retryDelayMillis(attempt);
+                            log.warn("[rustfs] operation failed, retrying operation={} objectKey={} attempt={}/{} delayMs={} err={}",
+                                    operation, objectKey, attempt, maxRetries, delayMillis, ex.getMessage());
+                        }
                    }
+                } finally {
+                    semaphore.release();
+                }
+                if (last instanceof OperationTimeoutException timeoutException) {
+                    throw timeoutException;
+                }
+                if (attempt < maxRetries) {
+                    sleepQuietly(delayMillis, "retrying rustfs " + operation,
+                            operation, objectKey, deadlineNanos);
                }
-            } finally {
-                semaphore.release();
            }
-            if (last instanceof OperationTimeoutException timeoutException) {
-                throw timeoutException;
-            }
-            if (attempt < maxRetries) {
-                sleepQuietly(delayMillis, "retrying rustfs " + operation,
-                        operation, objectKey, deadlineNanos);
+            throw new IllegalStateException("failed to " + operation + " payload in transient storage", last);
+        } finally {
+            if (totalAcquired) {
+                totalSemaphore.release();
            }
        }
-        throw new IllegalStateException("failed to " + operation + " payload in transient storage", last);
    }

    private void acquirePermit(String operation, String objectKey, Semaphore semaphore, long deadlineNanos) {
@@ -0,0 +1,272 @@
+package com.nanri.aiimage.modules.file.service.object;
+
+import com.nanri.aiimage.config.TransientStorageProperties;
+import io.micrometer.core.instrument.Counter;
+import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
+import io.minio.GetObjectArgs;
+import io.minio.GetObjectResponse;
+import io.minio.MinioClient;
+import io.minio.PutObjectArgs;
+import io.minio.StatObjectArgs;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentMatchers;
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Task 66：限制 RustFS 并发读写与重试的总资源预算，防止多任务叠加爆发。
+ * RustfsObjectStorageService 在分类信号量（upload/read/delete）之上增加总并发预算
+ * maxTotalConcurrentOperations：单次操作（含重试）全程占用一个总许可，
+ * 超预算立即拒绝（rejected 指标），失败/完成后释放，避免多任务叠加时
+ * 读写与重试叠加突破对后端的总压力上限。
+ */
+class RustfsTotalBudgetTest {
+
+    private TransientStorageProperties properties;
+    private SimpleMeterRegistry simpleRegistry;
+    private MinioClient client;
+    private ObjectProvider<MeterRegistry> meterRegistryProvider;
+    private ObjectProvider<RustfsDeleteRetryService> deleteRetryProvider;
+    private RustfsObjectStorageService service;
+
+    @BeforeEach
+    @SuppressWarnings("unchecked")
+    void setUp() {
+        properties = new TransientStorageProperties();
+        simpleRegistry = new SimpleMeterRegistry();
+        meterRegistryProvider = mock(ObjectProvider.class);
+        when(meterRegistryProvider.getIfAvailable()).thenReturn(simpleRegistry);
+        deleteRetryProvider = mock(ObjectProvider.class);
+        client = mock(MinioClient.class);
+        service = new RustfsObjectStorageService(properties, meterRegistryProvider, deleteRetryProvider,
+                () -> client);
+    }
+
+    @AfterEach
+    void tearDown() {
+        simpleRegistry.clear();
+    }
+
+    private static void setConfigured(TransientStorageProperties properties) {
+        properties.setEndpoint("http://rustfs.local:9000");
+        properties.setBucket("bucket");
+        properties.setAccessKeyId("ak");
+        properties.setAccessKeySecret("sk");
+        properties.setRegion("us-east-1");
+    }
+
+    private void stubUploadSuccess() {
+        try {
+            when(client.putObject(ArgumentMatchers.any(PutObjectArgs.class)))
+                    .thenReturn(mock(io.minio.ObjectWriteResponse.class));
+            when(client.statObject(ArgumentMatchers.any(StatObjectArgs.class))).thenReturn(null);
+        } catch (Exception ex) {
+            throw new IllegalStateException(ex);
+        }
+    }
+
+    private void stubReadSuccess(String content) throws Exception {
+        GetObjectResponse response = mock(GetObjectResponse.class);
+        when(response.readAllBytes()).thenReturn(content.getBytes(StandardCharsets.UTF_8));
+        when(client.getObject(ArgumentMatchers.any(GetObjectArgs.class))).thenReturn(response);
+    }
+
+    @SuppressWarnings("unchecked")
+    private Semaphore totalSemaphore() {
+        return (Semaphore) ReflectionTestUtils.getField(service, "totalSemaphore");
+    }
+
+    private static long rejectedTotal(MeterRegistry registry) {
+        Counter counter = registry.find("aiimage.rustfs.operation.total")
+                .tags("operation", "total", "result", "rejected").counter();
+        return counter == null ? 0L : (long) counter.count();
+    }
+
+    @Test
+    void test_task_066_rustfs_normal_default_path() throws Exception {
+        // 默认路径：配置总预算后正常上传成功，操作结束后总许可释放回池。
+        setConfigured(properties);
+        properties.setMaxTotalConcurrentOperations(8);
+        service = new RustfsObjectStorageService(properties, meterRegistryProvider, deleteRetryProvider,
+                () -> client);
+        stubUploadSuccess();
+
+        String objectKey = service.uploadText("task-parsed/test/1/scope/latest.json", "{\"a\":1}");
+
+        assertEquals("task-parsed/test/1/scope/latest.json", objectKey);
+        assertEquals(8, totalSemaphore().availablePermits(), "上传完成后总许可全部释放");
+        assertEquals(0L, rejectedTotal(simpleRegistry));
+    }
+
+    @Test
+    void test_task_066_rustfs_normal_multiple_items() throws Exception {
+        // 批量场景：多任务并发在总预算内全部成功，无拒绝无丢失。
+        setConfigured(properties);
+        properties.setMaxTotalConcurrentOperations(16);
+        service = new RustfsObjectStorageService(properties, meterRegistryProvider, deleteRetryProvider,
+                () -> client);
+        stubUploadSuccess();
+        int concurrency = 4;
+        CountDownLatch start = new CountDownLatch(1);
+        AtomicInteger successCount = new AtomicInteger();
+        Thread[] threads = new Thread[concurrency];
+        for (int i = 0; i < concurrency; i++) {
+            final int index = i;
+            threads[i] = new Thread(() -> {
+                try {
+                    start.await(2, TimeUnit.SECONDS);
+                    service.uploadText("task-parsed/test/" + index + "/scope/latest.json", "{\"i\":" + index + "}");
+                    successCount.incrementAndGet();
+                } catch (Exception ignored) {
+                }
+            });
+            threads[i].setDaemon(true);
+            threads[i].start();
+        }
+        start.countDown();
+        for (Thread thread : threads) {
+            thread.join(5_000);
+        }
+
+        assertEquals(concurrency, successCount.get(), "全部并发任务成功");
+        verify(client, times(concurrency)).putObject(ArgumentMatchers.any(PutObjectArgs.class));
+        assertEquals(16, totalSemaphore().availablePermits());
+        assertEquals(0L, rejectedTotal(simpleRegistry));
+    }
+
+    @Test
+    void test_task_066_rustfs_normal_repeated_operation_is_idempotent() throws Exception {
+        // 幂等：总预算 1 时串行重复操作每次都能获取并释放许可，不产生重复拒绝。
+        setConfigured(properties);
+        properties.setMaxTotalConcurrentOperations(1);
+        service = new RustfsObjectStorageService(properties, meterRegistryProvider, deleteRetryProvider,
+                () -> client);
+        stubUploadSuccess();
+
+        for (int i = 0; i < 3; i++) {
+            assertEquals("task-parsed/test/1/scope/latest.json",
+                    service.uploadText("task-parsed/test/1/scope/latest.json", "{\"x\":1}"));
+        }
+
+        verify(client, times(3)).putObject(ArgumentMatchers.any(PutObjectArgs.class));
+        assertEquals(1, totalSemaphore().availablePermits(), "每次操作后许可释放");
+        assertEquals(0L, rejectedTotal(simpleRegistry));
+    }
+
+    @Test
+    void test_task_066_rustfs_boundary_empty_input() throws Exception {
+        // 空输入：总预算未配置（0）时保持原行为，不启用总量限制。
+        setConfigured(properties);
+        stubUploadSuccess();
+
+        assertEquals("k", service.uploadText("k", "v"));
+
+        verify(client, times(1)).putObject(ArgumentMatchers.any(PutObjectArgs.class));
+        assertEquals(0L, rejectedTotal(simpleRegistry));
+    }
+
+    @Test
+    void test_task_066_rustfs_boundary_single_item() throws Exception {
+        // 单元素：总预算 1 时单次读取成功，许可释放，跨操作类型互不残留。
+        setConfigured(properties);
+        properties.setMaxTotalConcurrentOperations(1);
+        service = new RustfsObjectStorageService(properties, meterRegistryProvider, deleteRetryProvider,
+                () -> client);
+        stubReadSuccess("{}");
+
+        assertEquals("{}", service.readObjectAsString("task-parsed/test/1/scope/latest.json"));
+        assertEquals(1, totalSemaphore().availablePermits());
+    }
+
+    @Test
+    void test_task_066_rustfs_boundary_limit_and_overflow() throws Exception {
+        // 上限/超限：总预算 1 时第二个并发操作（不同分类）立即被拒绝，
+        // rejected 指标 +1，不发生无界排队；操作结束后许可释放。
+        setConfigured(properties);
+        properties.setMaxTotalConcurrentOperations(1);
+        properties.setAcquirePermitTimeoutMillis(0);
+        service = new RustfsObjectStorageService(properties, meterRegistryProvider, deleteRetryProvider,
+                () -> client);
+        stubUploadSuccess();
+        CountDownLatch entered = new CountDownLatch(1);
+        org.mockito.Mockito.doAnswer(invocation -> {
+            entered.countDown();
+            Thread.sleep(500);
+            throw new java.io.IOException("hold");
+        }).when(client).putObject(ArgumentMatchers.any(PutObjectArgs.class));
+
+        Thread first = new Thread(() -> {
+            try {
+                service.uploadText("block", "v");
+            } catch (Exception ignored) {
+            }
+        });
+        first.setDaemon(true);
+        first.start();
+        try {
+            assertTrue(entered.await(2, TimeUnit.SECONDS), "首线程已持有总预算");
+        } catch (InterruptedException ex) {
+            Thread.currentThread().interrupt();
+        }
+
+        assertThrows(IllegalStateException.class, () -> service.readObjectAsString("blocked"),
+                "跨操作类型并发超预算被拒绝");
+        assertEquals(1L, rejectedTotal(simpleRegistry), "拒绝计入 total rejected 指标");
+        first.interrupt();
+    }
+
+    @Test
+    void test_task_066_rustfs_invalid_input_rejected() throws Exception {
+        // 非法参数：总预算为负或 0 时视为未启用（构造不崩溃，行为与默认一致）。
+        setConfigured(properties);
+        properties.setMaxTotalConcurrentOperations(-5);
+        service = new RustfsObjectStorageService(properties, meterRegistryProvider, deleteRetryProvider,
+                () -> client);
+        stubUploadSuccess();
+
+        assertEquals("k", service.uploadText("k", "v"));
+        assertEquals(0L, rejectedTotal(simpleRegistry));
+    }
+
+    @Test
+    void test_task_066_rustfs_dependency_failure_releases_resources() throws Exception {
+        // 依赖失败：上传失败后总许可释放（错误可恢复），恢复后操作成功且无残留锁。
+        setConfigured(properties);
+        properties.setMaxTotalConcurrentOperations(1);
+        service = new RustfsObjectStorageService(properties, meterRegistryProvider, deleteRetryProvider,
+                () -> client);
+        try {
+            when(client.putObject(ArgumentMatchers.any(PutObjectArgs.class)))
+                    .thenThrow(new java.io.IOException("rustfs down"));
+            assertThrows(Exception.class, () -> service.uploadText("task-parsed/test/1/scope/latest.json", "{\"a\":1}"));
+        } finally {
+            try {
+                stubUploadSuccess();
+            } catch (Exception ignored) {
+            }
+        }
+        assertEquals(1, totalSemaphore().availablePermits(), "失败后总许可释放");
+
+        assertEquals("task-parsed/test/1/scope/latest.json",
+                service.uploadText("task-parsed/test/1/scope/latest.json", "{\"b\":2}"));
+        assertEquals(1, totalSemaphore().availablePermits());
+    }
+}