完善任务存储、权限及货源查询流程
Build Backend JAR / build (push) Has been cancelled

This commit is contained in:
supernijia
2026-08-14 22:49:15 +08:00
parent 5b1ccad40e
commit 7a7f1dfa21
21 changed files with 2584 additions and 566 deletions
@@ -15,6 +15,8 @@ public class TransientStorageProperties {
private int connectTimeoutSeconds = 10;
private int readTimeoutSeconds = 60;
private int writeTimeoutSeconds = 60;
private int callTimeoutSeconds = 90;
private int operationTimeoutSeconds = 120;
private int uploadMaxRetries = 3;
private int readMaxRetries = 3;
private int deleteMaxRetries = 3;
@@ -28,8 +30,10 @@ public class TransientStorageProperties {
private long failureWindowSeconds = 60;
private int failureWindowThreshold = 20;
private long failureCooldownMillis = 10000;
private int connectionPoolMaxIdle = 0;
private long connectionPoolKeepAliveMillis = 1;
private int dispatcherMaxRequests = 56;
private int dispatcherMaxRequestsPerHost = 56;
private int connectionPoolMaxIdle = 5;
private long connectionPoolKeepAliveMillis = 300000;
private long warnPayloadBytes = 5L * 1024 * 1024;
private long maxPayloadBytes = 50L * 1024 * 1024;
private long maxStoredPayloadBytes = 50L * 1024 * 1024;
@@ -38,4 +42,5 @@ public class TransientStorageProperties {
private String deleteRetryCron = "0 */5 * * * *";
private int deleteRetryQueueCapacity = 10000;
private int deleteRetryBatchSize = 200;
private int deleteRetryBatchTimeoutSeconds = 60;
}
@@ -7,8 +7,13 @@ import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
@Service
@@ -18,18 +23,30 @@ public class RustfsDeleteRetryService {
private final TransientStorageProperties properties;
private final RustfsObjectStorageService rustfsObjectStorageService;
private final ConcurrentLinkedQueue<String> queue = new ConcurrentLinkedQueue<>();
private final ConcurrentLinkedQueue<RetryItem> queue = new ConcurrentLinkedQueue<>();
private final ConcurrentHashMap<String, RetryItem> pending = new ConcurrentHashMap<>();
private final AtomicInteger queuedCount = new AtomicInteger();
private final Object admissionLock = new Object();
public void enqueue(String objectKey, Throwable cause) {
if (!properties.isDeleteRetryEnabled() || objectKey == null || objectKey.isBlank()) {
return;
}
RetryItem item = pending.computeIfAbsent(objectKey, RetryItem::new);
item.lastError = cause == null ? null : cause.getMessage();
item.lastFailedAt = LocalDateTime.now();
enqueueItem(item, false);
while (true) {
RetryItem item = findOrAdmit(objectKey, cause);
if (item == null) {
return;
}
synchronized (item) {
if (pending.get(objectKey) != item) {
continue;
}
item.generation++;
item.lastError = cause == null ? null : cause.getMessage();
item.lastFailedAt = LocalDateTime.now();
enqueueItemLocked(item);
return;
}
}
}
@Scheduled(cron = "${aiimage.transient-storage.delete-retry-cron:0 */5 * * * *}")
@@ -38,36 +55,50 @@ public class RustfsDeleteRetryService {
return;
}
int batchSize = Math.max(1, properties.getDeleteRetryBatchSize());
long batchDurationNanos = TimeUnit.SECONDS.toNanos(
Math.max(1L, properties.getDeleteRetryBatchTimeoutSeconds()));
long deadlineNanos = System.nanoTime() + batchDurationNanos;
int processed = 0;
int success = 0;
int failed = 0;
Set<String> attemptedObjectKeys = new HashSet<>();
List<RetryItem> retryNextInvocation = new ArrayList<>();
while (processed < batchSize) {
String objectKey = nextQueuedObjectKey(batchSize - processed);
if (objectKey == null) {
long remainingNanos = deadlineNanos - System.nanoTime();
if (remainingNanos <= 0L) {
break;
}
RetryItem item = pending.get(objectKey);
if (item == null) {
RetryAttempt attempt = nextAttempt();
if (attempt == null) {
break;
}
RetryItem item = attempt.item;
String objectKey = item.objectKey;
if (!attemptedObjectKeys.add(objectKey)) {
retryNextInvocation.add(item);
continue;
}
item.queued = false;
remainingNanos = deadlineNanos - System.nanoTime();
if (remainingNanos <= 0L) {
retryNextInvocation.add(item);
break;
}
processed++;
try {
rustfsObjectStorageService.deleteObjectFromRetry(objectKey);
pending.remove(objectKey);
rustfsObjectStorageService.deleteObjectFromRetry(objectKey, remainingNanos);
removeIfUnchanged(attempt);
success++;
} catch (Exception ex) {
failed++;
int failures = item.failures.incrementAndGet();
item.lastError = ex.getMessage();
item.lastFailedAt = LocalDateTime.now();
int failures = recordFailureIfUnchanged(attempt, ex);
log.warn("[rustfs] delete retry failed objectKey={} failures={} err={}", objectKey, failures, ex.getMessage());
enqueueItem(item, true);
retryNextInvocation.add(item);
}
}
retryNextInvocation.forEach(this::requeue);
if (processed > 0) {
log.info("[rustfs] delete retry batch completed processed={} success={} failed={} pending={} queued={}",
processed, success, failed, pending.size(), queuedCount.get());
log.info("[rustfs] delete retry batch completed processed={} success={} failed={} pending={}",
processed, success, failed, pending.size());
}
}
@@ -75,72 +106,94 @@ public class RustfsDeleteRetryService {
return pending.size();
}
private String nextQueuedObjectKey(int refillLimit) {
String objectKey = queue.poll();
if (objectKey != null) {
decrementQueuedCount();
refillQueue(refillLimit);
return objectKey;
}
if (refillQueue(refillLimit) <= 0) {
return null;
}
objectKey = queue.poll();
if (objectKey != null) {
decrementQueuedCount();
refillQueue(refillLimit);
}
return objectKey;
}
private int refillQueue(int maxItems) {
if (maxItems <= 0) {
return 0;
private RetryItem findOrAdmit(String objectKey, Throwable cause) {
RetryItem item = pending.get(objectKey);
if (item != null) {
return item;
}
int capacity = Math.max(1, properties.getDeleteRetryQueueCapacity());
if (queuedCount.get() >= capacity) {
return 0;
}
int added = 0;
for (RetryItem item : pending.values()) {
if (added >= maxItems || queuedCount.get() >= capacity) {
break;
synchronized (admissionLock) {
item = pending.get(objectKey);
if (item != null) {
return item;
}
if (enqueueItem(item, true)) {
added++;
if (pending.size() >= capacity) {
log.error("[rustfs] delete retry capacity reached, rejected objectKey={} pending={} capacity={} lastError={}",
objectKey, pending.size(), capacity, cause == null ? null : cause.getMessage());
return null;
}
item = new RetryItem(objectKey);
pending.put(objectKey, item);
return item;
}
return added;
}
private void decrementQueuedCount() {
queuedCount.updateAndGet(value -> Math.max(0, value - 1));
private RetryAttempt nextAttempt() {
while (true) {
RetryItem item = queue.poll();
if (item == null) {
return null;
}
synchronized (item) {
if (pending.get(item.objectKey) != item || !item.queued) {
continue;
}
item.queued = false;
return new RetryAttempt(item, item.generation);
}
}
}
private boolean enqueueItem(RetryItem item, boolean requeue) {
private void removeIfUnchanged(RetryAttempt attempt) {
RetryItem item = attempt.item;
synchronized (item) {
if (item.queued) {
return false;
if (pending.get(item.objectKey) != item || item.generation != attempt.generation) {
return;
}
int capacity = Math.max(1, properties.getDeleteRetryQueueCapacity());
if (queuedCount.get() >= capacity) {
log.error("[rustfs] delete retry queue full, deferred objectKey={} requeue={} pending={} queued={} capacity={} lastError={}",
item.objectKey, requeue, pending.size(), queuedCount.get(), capacity, item.lastError);
return false;
synchronized (admissionLock) {
pending.remove(item.objectKey, item);
}
item.queued = true;
queue.offer(item.objectKey);
queuedCount.incrementAndGet();
return true;
}
}
private int recordFailureIfUnchanged(RetryAttempt attempt, Exception ex) {
RetryItem item = attempt.item;
synchronized (item) {
if (pending.get(item.objectKey) != item || item.generation != attempt.generation) {
return item.failures.get();
}
int failures = item.failures.incrementAndGet();
item.lastError = ex.getMessage();
item.lastFailedAt = LocalDateTime.now();
return failures;
}
}
private void requeue(RetryItem item) {
synchronized (item) {
if (pending.get(item.objectKey) == item) {
enqueueItemLocked(item);
}
}
}
private void enqueueItemLocked(RetryItem item) {
if (!item.queued) {
item.queued = true;
queue.offer(item);
}
}
private record RetryAttempt(RetryItem item, long generation) {
}
private static class RetryItem {
private final String objectKey;
private final AtomicInteger failures = new AtomicInteger();
private volatile boolean queued;
private volatile String lastError;
private volatile LocalDateTime lastFailedAt;
private boolean queued;
private long generation;
private String lastError;
private LocalDateTime lastFailedAt;
private RetryItem(String objectKey) {
this.objectKey = objectKey;
@@ -11,6 +11,7 @@ import io.minio.RemoveObjectArgs;
import io.minio.StatObjectArgs;
import lombok.extern.slf4j.Slf4j;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import okhttp3.OkHttpClient;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Autowired;
@@ -22,6 +23,7 @@ import java.util.Objects;
import java.util.concurrent.Semaphore;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
@@ -78,35 +80,51 @@ public class RustfsObjectStorageService {
}
public String uploadText(String objectKey, String content, boolean verifyAfterUpload) {
long deadlineNanos = operationDeadlineNanos();
if (!isConfigured()) {
throw new IllegalStateException("transient storage is not configured");
}
rejectIfCircuitOpen(OP_UPLOAD, objectKey);
byte[] bytes = Objects.requireNonNullElse(content, "").getBytes(StandardCharsets.UTF_8);
recordPayloadBytes(bytes.length);
return executeWithRetry(OP_UPLOAD, objectKey, bytes.length, Math.max(1, properties.getUploadMaxRetries()), uploadSemaphore, () -> {
try (ByteArrayInputStream stream = new ByteArrayInputStream(bytes)) {
buildClient().putObject(PutObjectArgs.builder()
.bucket(properties.getBucket())
.object(objectKey)
.stream(stream, bytes.length, -1)
.contentType("application/json")
.build());
if (verifyAfterUpload) {
verifyObjectVisible(objectKey);
AtomicBoolean putCompleted = new AtomicBoolean();
try {
String uploadedObjectKey = executeWithRetry(OP_UPLOAD, objectKey, bytes.length,
Math.max(1, properties.getUploadMaxRetries()), uploadSemaphore, deadlineNanos, () -> {
try (ByteArrayInputStream stream = new ByteArrayInputStream(bytes)) {
buildClient(deadlineNanos).putObject(PutObjectArgs.builder()
.bucket(properties.getBucket())
.object(objectKey)
.stream(stream, bytes.length, -1)
.contentType("application/json")
.build());
putCompleted.set(true);
return objectKey;
}
return objectKey;
});
if (verifyAfterUpload) {
verifyObjectVisible(objectKey, deadlineNanos);
} else {
resetFailureWindow(OP_UPLOAD);
}
});
return uploadedObjectKey;
} catch (RuntimeException ex) {
if (putCompleted.get()) {
enqueueDeleteRetry(objectKey, ex);
}
throw ex;
}
}
public String readObjectAsString(String objectKey) {
long deadlineNanos = operationDeadlineNanos();
if (!isConfigured()) {
throw new IllegalStateException("transient storage is not configured");
}
rejectIfCircuitOpen(OP_READ, objectKey);
return executeWithRetry(OP_READ, objectKey, 0, Math.max(1, properties.getReadMaxRetries()), readSemaphore, () -> {
try (var stream = buildClient().getObject(GetObjectArgs.builder()
return executeWithRetry(OP_READ, objectKey, 0, Math.max(1, properties.getReadMaxRetries()),
readSemaphore, deadlineNanos, () -> {
try (var stream = buildClient(deadlineNanos).getObject(GetObjectArgs.builder()
.bucket(properties.getBucket())
.object(objectKey)
.build())) {
@@ -116,20 +134,25 @@ public class RustfsObjectStorageService {
}
public void deleteObject(String objectKey) {
deleteObject(objectKey, true);
deleteObject(objectKey, true, operationDeadlineNanos());
}
void deleteObjectFromRetry(String objectKey) {
deleteObject(objectKey, false);
deleteObject(objectKey, false, operationDeadlineNanos());
}
private void deleteObject(String objectKey, boolean enqueueRetryOnFailure) {
void deleteObjectFromRetry(String objectKey, long timeoutNanos) {
deleteObject(objectKey, false, deadlineFromNow(timeoutNanos));
}
private void deleteObject(String objectKey, boolean enqueueRetryOnFailure, long deadlineNanos) {
if (!isConfigured() || !notBlank(objectKey)) {
return;
}
try {
executeWithRetry(OP_DELETE, objectKey, 0, Math.max(1, properties.getDeleteMaxRetries()), deleteSemaphore, () -> {
buildClient().removeObject(RemoveObjectArgs.builder()
executeWithRetry(OP_DELETE, objectKey, 0, Math.max(1, properties.getDeleteMaxRetries()),
deleteSemaphore, deadlineNanos, () -> {
buildClient(deadlineNanos).removeObject(RemoveObjectArgs.builder()
.bucket(properties.getBucket())
.object(objectKey)
.build());
@@ -137,10 +160,7 @@ public class RustfsObjectStorageService {
});
} catch (RuntimeException ex) {
if (enqueueRetryOnFailure && properties.isDeleteRetryEnabled()) {
RustfsDeleteRetryService retryService = deleteRetryServiceProvider == null ? null : deleteRetryServiceProvider.getIfAvailable();
if (retryService != null) {
retryService.enqueue(objectKey, ex);
}
enqueueDeleteRetry(objectKey, ex);
}
throw ex;
}
@@ -151,48 +171,65 @@ public class RustfsObjectStorageService {
long bytes,
int maxRetries,
Semaphore semaphore,
long deadlineNanos,
CheckedSupplier<T> supplier) {
acquirePermit(operation, objectKey, semaphore);
long startedAt = System.nanoTime();
Exception last = null;
try {
for (int attempt = 1; attempt <= maxRetries; attempt++) {
for (int attempt = 1; attempt <= maxRetries; attempt++) {
checkDeadline(operation, objectKey, deadlineNanos);
rejectIfCircuitOpen(operation, objectKey);
acquirePermit(operation, objectKey, semaphore, deadlineNanos);
long delayMillis = 0L;
try {
try {
checkDeadline(operation, objectKey, deadlineNanos);
T result = supplier.get();
resetFailureWindow();
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;
evictIdleConnections();
recordOperation(operation, attempt < maxRetries ? "retry" : "failure", elapsedNanos(startedAt));
recordFailure(operation, objectKey, ex);
if (attempt < maxRetries) {
long delayMillis = retryDelayMillis(attempt);
delayMillis = retryDelayMillis(attempt);
log.warn("[rustfs] operation failed, retrying operation={} objectKey={} attempt={}/{} delayMs={} err={}",
operation, objectKey, attempt, maxRetries, delayMillis, ex.getMessage());
sleepQuietly(delayMillis, "retrying rustfs " + operation);
}
}
} finally {
semaphore.release();
}
if (last instanceof OperationTimeoutException timeoutException) {
throw timeoutException;
}
if (attempt < maxRetries) {
sleepQuietly(delayMillis, "retrying rustfs " + operation,
operation, objectKey, deadlineNanos);
}
} finally {
semaphore.release();
}
throw new IllegalStateException("failed to " + operation + " payload in transient storage", last);
}
private void acquirePermit(String operation, String objectKey, Semaphore semaphore) {
private void acquirePermit(String operation, String objectKey, Semaphore semaphore, long deadlineNanos) {
try {
checkDeadline(operation, objectKey, deadlineNanos);
boolean acquired;
long timeoutMillis = Math.max(0L, properties.getAcquirePermitTimeoutMillis());
if (timeoutMillis == 0L) {
acquired = semaphore.tryAcquire();
} else {
acquired = semaphore.tryAcquire(timeoutMillis, TimeUnit.MILLISECONDS);
long waitNanos = Math.min(TimeUnit.MILLISECONDS.toNanos(timeoutMillis),
remainingNanos(deadlineNanos));
acquired = semaphore.tryAcquire(waitNanos, TimeUnit.NANOSECONDS);
}
if (!acquired) {
checkDeadline(operation, objectKey, deadlineNanos);
recordOperation(operation, "rejected", 0L);
throw new IllegalStateException("rustfs " + operation + " concurrency limit reached: " + objectKey);
}
@@ -203,7 +240,7 @@ public class RustfsObjectStorageService {
}
}
private MinioClient buildClient() {
private MinioClient buildClient(long deadlineNanos) {
if (minioClientSupplier != null) {
return minioClientSupplier.get();
}
@@ -211,7 +248,7 @@ public class RustfsObjectStorageService {
.endpoint(properties.getEndpoint())
.credentials(properties.getAccessKeyId(), properties.getAccessKeySecret())
.region(properties.getRegion())
.httpClient(getHttpClient())
.httpClient(getHttpClient(deadlineNanos))
.build();
}
@@ -222,10 +259,15 @@ public class RustfsObjectStorageService {
}
synchronized (this) {
if (httpClient == null) {
Dispatcher dispatcher = new Dispatcher();
dispatcher.setMaxRequests(Math.max(1, properties.getDispatcherMaxRequests()));
dispatcher.setMaxRequestsPerHost(Math.max(1, properties.getDispatcherMaxRequestsPerHost()));
httpClient = new OkHttpClient.Builder()
.dispatcher(dispatcher)
.connectTimeout(Math.max(1, properties.getConnectTimeoutSeconds()), TimeUnit.SECONDS)
.readTimeout(Math.max(1, properties.getReadTimeoutSeconds()), TimeUnit.SECONDS)
.writeTimeout(Math.max(1, properties.getWriteTimeoutSeconds()), TimeUnit.SECONDS)
.callTimeout(Math.max(1, properties.getCallTimeoutSeconds()), TimeUnit.SECONDS)
.connectionPool(new ConnectionPool(
Math.max(0, properties.getConnectionPoolMaxIdle()),
Math.max(1L, properties.getConnectionPoolKeepAliveMillis()),
@@ -237,30 +279,49 @@ public class RustfsObjectStorageService {
}
}
private void evictIdleConnections() {
OkHttpClient current = httpClient;
if (current != null) {
current.connectionPool().evictAll();
OkHttpClient getHttpClient(long deadlineNanos) {
OkHttpClient client = getHttpClient();
if (deadlineNanos == Long.MAX_VALUE) {
return client;
}
long configuredTimeoutNanos = TimeUnit.SECONDS.toNanos(Math.max(1, properties.getCallTimeoutSeconds()));
long callTimeoutMillis = TimeUnit.NANOSECONDS.toMillis(
Math.min(configuredTimeoutNanos, remainingNanos(deadlineNanos)));
if (callTimeoutMillis <= 0L) {
throw new OperationTimeoutException("rustfs operation timeout before HTTP call");
}
return client.newBuilder()
.callTimeout(callTimeoutMillis, TimeUnit.MILLISECONDS)
.build();
}
private void verifyObjectVisible(String objectKey) {
private void verifyObjectVisible(String objectKey, long deadlineNanos) {
long startedAt = System.nanoTime();
RuntimeException last = null;
for (int attempt = 1; attempt <= 3; attempt++) {
checkDeadline(OP_STAT, objectKey, deadlineNanos);
try {
buildClient().statObject(StatObjectArgs.builder()
buildClient(deadlineNanos).statObject(StatObjectArgs.builder()
.bucket(properties.getBucket())
.object(objectKey)
.build());
checkDeadline(OP_STAT, objectKey, deadlineNanos);
resetFailureWindow(OP_STAT);
recordOperation(OP_STAT, "success", elapsedNanos(startedAt));
return;
} catch (Exception ex) {
if (ex instanceof OperationTimeoutException timeoutException) {
throw timeoutException;
}
last = new IllegalStateException("transient payload is not visible after upload: " + objectKey, ex);
evictIdleConnections();
recordOperation(OP_STAT, attempt < 3 ? "retry" : "failure", elapsedNanos(startedAt));
recordFailure(OP_STAT, objectKey, ex);
if (isCircuitOpen()) {
break;
}
if (attempt < 3) {
sleepQuietly(100L * attempt, "verifying transient payload upload");
sleepQuietly(100L * attempt, "verifying transient payload upload",
OP_STAT, objectKey, deadlineNanos);
}
}
}
@@ -270,6 +331,9 @@ public class RustfsObjectStorageService {
}
private void rejectIfCircuitOpen(String operation, String objectKey) {
if (!OP_UPLOAD.equals(operation) && !OP_READ.equals(operation)) {
return;
}
long openUntil = circuitOpenUntilMillis;
if (openUntil <= 0L) {
return;
@@ -282,6 +346,9 @@ public class RustfsObjectStorageService {
}
private void recordFailure(String operation, String objectKey, Exception ex) {
if (!participatesInFailureWindow(operation)) {
return;
}
long now = System.currentTimeMillis();
long windowMillis = Math.max(1L, properties.getFailureWindowSeconds()) * 1000L;
if (failureWindowStartedAtMillis <= 0L || now - failureWindowStartedAtMillis > windowMillis) {
@@ -290,7 +357,7 @@ public class RustfsObjectStorageService {
}
int failures = windowFailureCount.incrementAndGet();
int threshold = Math.max(1, properties.getFailureWindowThreshold());
if ((OP_UPLOAD.equals(operation) || OP_READ.equals(operation)) && failures >= threshold) {
if (failures >= threshold) {
long cooldownMillis = Math.max(0L, properties.getFailureCooldownMillis());
circuitOpenUntilMillis = now + cooldownMillis;
windowFailureCount.set(0);
@@ -300,7 +367,10 @@ public class RustfsObjectStorageService {
}
}
private void resetFailureWindow() {
private void resetFailureWindow(String operation) {
if (!participatesInFailureWindow(operation)) {
return;
}
windowFailureCount.set(0);
failureWindowStartedAtMillis = 0L;
circuitOpenUntilMillis = 0L;
@@ -317,16 +387,70 @@ public class RustfsObjectStorageService {
return delay;
}
private void sleepQuietly(long millis, String reason) {
private void sleepQuietly(long millis,
String reason,
String operation,
String objectKey,
long deadlineNanos) {
if (millis <= 0L) {
return;
}
checkDeadline(operation, objectKey, deadlineNanos);
long sleepNanos = Math.min(TimeUnit.MILLISECONDS.toNanos(millis), remainingNanos(deadlineNanos));
try {
Thread.sleep(millis);
TimeUnit.NANOSECONDS.sleep(sleepNanos);
} catch (InterruptedException ex) {
Thread.currentThread().interrupt();
throw new IllegalStateException("interrupted while " + reason, ex);
}
checkDeadline(operation, objectKey, deadlineNanos);
}
private long operationDeadlineNanos() {
long timeoutSeconds = properties.getOperationTimeoutSeconds();
return timeoutSeconds <= 0L
? Long.MAX_VALUE
: deadlineFromNow(TimeUnit.SECONDS.toNanos(timeoutSeconds));
}
private long deadlineFromNow(long timeoutNanos) {
long boundedTimeoutNanos = Math.min(Math.max(0L, timeoutNanos), Long.MAX_VALUE / 2L);
return System.nanoTime() + boundedTimeoutNanos;
}
private long remainingNanos(long deadlineNanos) {
return deadlineNanos == Long.MAX_VALUE ? Long.MAX_VALUE : Math.max(0L, deadlineNanos - System.nanoTime());
}
private void checkDeadline(String operation, String objectKey, long deadlineNanos) {
if (remainingNanos(deadlineNanos) <= 0L) {
throw new OperationTimeoutException(
"rustfs operation timeout operation=" + operation + " objectKey=" + objectKey);
}
}
private boolean participatesInFailureWindow(String operation) {
return OP_UPLOAD.equals(operation) || OP_READ.equals(operation) || OP_STAT.equals(operation);
}
private boolean isCircuitOpen() {
return System.currentTimeMillis() < circuitOpenUntilMillis;
}
private void enqueueDeleteRetry(String objectKey, RuntimeException failure) {
if (!properties.isDeleteRetryEnabled() || deleteRetryServiceProvider == null) {
return;
}
RustfsDeleteRetryService retryService = deleteRetryServiceProvider.getIfAvailable();
if (retryService == null) {
return;
}
try {
retryService.enqueue(objectKey, failure);
} catch (RuntimeException enqueueFailure) {
log.error("[rustfs] failed to enqueue delete compensation objectKey={} err={}",
objectKey, enqueueFailure.getMessage(), enqueueFailure);
}
}
private void recordPayloadBytes(long bytes) {
@@ -382,4 +506,10 @@ public class RustfsObjectStorageService {
private interface CheckedSupplier<T> {
T get() throws Exception;
}
private static class OperationTimeoutException extends IllegalStateException {
private OperationTimeoutException(String message) {
super(message);
}
}
}
@@ -1,9 +1,45 @@
package com.nanri.aiimage.modules.permission.mapper;
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.nanri.aiimage.modules.permission.model.entity.UserColumnPermissionEntity;
import org.apache.ibatis.annotations.Delete;
import org.apache.ibatis.annotations.Insert;
import org.apache.ibatis.annotations.Mapper;
import org.apache.ibatis.annotations.Param;
import org.apache.ibatis.annotations.Select;
import java.util.List;
@Mapper
public interface UserColumnPermissionMapper extends BaseMapper<UserColumnPermissionEntity> {
public interface UserColumnPermissionMapper {
@Insert("INSERT INTO user_column_permission (user_id, column_id) VALUES (#{userId}, #{columnId})")
int insert(UserColumnPermissionEntity entity);
@Select("SELECT user_id, column_id FROM user_column_permission WHERE user_id = #{userId}")
List<UserColumnPermissionEntity> selectByUserId(@Param("userId") Long userId);
@Select("SELECT user_id, column_id FROM user_column_permission WHERE column_id = #{columnId}")
List<UserColumnPermissionEntity> selectByColumnId(@Param("columnId") Long columnId);
@Select("SELECT COUNT(*) FROM user_column_permission WHERE user_id = #{userId} AND column_id = #{columnId}")
Long countByUserIdAndColumnId(@Param("userId") Long userId, @Param("columnId") Long columnId);
@Delete("DELETE FROM user_column_permission WHERE user_id = #{userId}")
int deleteByUserId(@Param("userId") Long userId);
@Delete("DELETE FROM user_column_permission WHERE column_id = #{columnId}")
int deleteByColumnId(@Param("columnId") Long columnId);
@Delete("""
<script>
DELETE FROM user_column_permission
WHERE user_id = #{userId}
AND column_id IN
<foreach collection="columnIds" item="columnId" open="(" separator="," close=")">
#{columnId}
</foreach>
</script>
""")
int deleteByUserIdAndColumnIds(@Param("userId") Long userId,
@Param("columnIds") List<Long> columnIds);
}
@@ -1,7 +1,6 @@
package com.nanri.aiimage.modules.permission.service;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.nanri.aiimage.common.exception.BusinessException;
import com.nanri.aiimage.modules.permission.mapper.AdminUserMapper;
import com.nanri.aiimage.modules.permission.mapper.PermissionMenuMapper;
@@ -144,8 +143,7 @@ public class PermissionMenuService {
}
// Do this explicitly even when the database FK is configured, so old
// installations without the FK are cleaned up as well.
userColumnPermissionMapper.delete(new LambdaUpdateWrapper<UserColumnPermissionEntity>()
.eq(UserColumnPermissionEntity::getColumnId, entity.getId()));
userColumnPermissionMapper.deleteByColumnId(entity.getId());
permissionMenuMapper.deleteById(entity.getId());
}
@@ -218,9 +216,7 @@ public class PermissionMenuService {
public List<ImageVideoDataPermissionUserVo> listImageVideoDataPermissionUsers(AdminUserEntity operator) {
ensureSuperAdminOperator(operator);
PermissionMenuEntity dataPermission = requireImageVideoDataPermission();
Set<Long> grantedUserIds = userColumnPermissionMapper.selectList(
new LambdaQueryWrapper<UserColumnPermissionEntity>()
.eq(UserColumnPermissionEntity::getColumnId, dataPermission.getId()))
Set<Long> grantedUserIds = userColumnPermissionMapper.selectByColumnId(dataPermission.getId())
.stream()
.map(UserColumnPermissionEntity::getUserId)
.filter(id -> id != null && id > 0)
@@ -246,7 +242,7 @@ public class PermissionMenuService {
if (requestedIds.stream().anyMatch(id -> !grantableUsers.containsKey(id))) {
throw new BusinessException("contains unknown or non-grantable user");
}
userColumnPermissionMapper.deleteByMap(Map.of("column_id", dataPermission.getId()));
userColumnPermissionMapper.deleteByColumnId(dataPermission.getId());
for (Long userId : requestedIds) {
UserColumnPermissionEntity grant = new UserColumnPermissionEntity();
grant.setUserId(userId);
@@ -260,9 +256,7 @@ public class PermissionMenuService {
ensureSuperAdminOperator(operator, "店铺数据任务");
PermissionMenuEntity dataPermission = requireDataPermission(
SHOP_DATA_CRAWL_DATA_PERMISSION_KEY, "店铺数据任务数据");
Set<Long> grantedUserIds = userColumnPermissionMapper.selectList(
new LambdaQueryWrapper<UserColumnPermissionEntity>()
.eq(UserColumnPermissionEntity::getColumnId, dataPermission.getId()))
Set<Long> grantedUserIds = userColumnPermissionMapper.selectByColumnId(dataPermission.getId())
.stream()
.map(UserColumnPermissionEntity::getUserId)
.filter(id -> id != null && id > 0)
@@ -289,7 +283,7 @@ public class PermissionMenuService {
if (requestedIds.stream().anyMatch(id -> !grantableUsers.containsKey(id))) {
throw new BusinessException("contains unknown or non-grantable user");
}
userColumnPermissionMapper.deleteByMap(Map.of("column_id", dataPermission.getId()));
userColumnPermissionMapper.deleteByColumnId(dataPermission.getId());
for (Long userId : requestedIds) {
UserColumnPermissionEntity grant = new UserColumnPermissionEntity();
grant.setUserId(userId);
@@ -389,10 +383,7 @@ public class PermissionMenuService {
if (!protectedIds.isEmpty()) {
for (Long protectedId : protectedIds) {
Long existingCount = userColumnPermissionMapper.selectCount(
new LambdaQueryWrapper<UserColumnPermissionEntity>()
.eq(UserColumnPermissionEntity::getUserId, userId)
.eq(UserColumnPermissionEntity::getColumnId, protectedId));
Long existingCount = userColumnPermissionMapper.countByUserIdAndColumnId(userId, protectedId);
if (existingCount != null && existingCount > 0) {
finalGrantIds.add(protectedId);
}
@@ -405,14 +396,10 @@ public class PermissionMenuService {
.forEach(finalGrantIds::add);
}
LambdaUpdateWrapper<UserColumnPermissionEntity> delete =
new LambdaUpdateWrapper<UserColumnPermissionEntity>()
.eq(UserColumnPermissionEntity::getUserId, userId);
if (normalizedType == null) {
userColumnPermissionMapper.delete(delete);
userColumnPermissionMapper.deleteByUserId(userId);
} else if (!scopedMenuIds.isEmpty()) {
userColumnPermissionMapper.delete(delete.in(
UserColumnPermissionEntity::getColumnId, scopedMenuIds));
userColumnPermissionMapper.deleteByUserIdAndColumnIds(userId, new ArrayList<>(scopedMenuIds));
}
for (Long columnId : finalGrantIds) {
UserColumnPermissionEntity entity = new UserColumnPermissionEntity();
@@ -439,9 +426,7 @@ public class PermissionMenuService {
}
private List<Long> loadDirectColumnIds(Long userId) {
List<UserColumnPermissionEntity> rows = userColumnPermissionMapper.selectList(
new LambdaQueryWrapper<UserColumnPermissionEntity>()
.eq(UserColumnPermissionEntity::getUserId, userId));
List<UserColumnPermissionEntity> rows = userColumnPermissionMapper.selectByUserId(userId);
if (rows == null || rows.isEmpty()) {
return List.of();
}
@@ -250,16 +250,28 @@ public class PublishTaskService {
if (lock == null) {
throw new BusinessException("task lock is busy");
}
List<String> storedPayloads = new ArrayList<>();
boolean[] cleanupAfterCommit = {false};
List<String> uploadedPayloads = new ArrayList<>();
PreparedResultSubmission prepared;
try {
transactionTemplate.executeWithoutResult(
status -> cleanupAfterCommit[0] = submitResultLocked(taskId, request, storedPayloads));
prepared = prepareResultSubmission(taskId, request, uploadedPayloads);
} catch (RuntimeException ex) {
deleteRolledBackPayloads(storedPayloads);
deleteUncommittedPayloads(uploadedPayloads);
throw ex;
}
if (cleanupAfterCommit[0]) {
SubmitResultCommit[] committedHolder = {null};
try {
transactionTemplate.executeWithoutResult(
status -> committedHolder[0] = submitResultLocked(taskId, prepared));
} catch (RuntimeException ex) {
deleteUncommittedPayloads(uploadedPayloads);
throw ex;
}
SubmitResultCommit committed = committedHolder[0];
Set<String> committedPayloads = committed == null ? Set.of() : committed.committedPayloads();
deleteUncommittedPayloads(uploadedPayloads.stream()
.filter(payload -> !committedPayloads.contains(payload))
.toList());
if (committed != null && committed.cleanupAfterCommit()) {
try {
deleteTransientResultChunks(taskId);
} catch (Exception ex) {
@@ -550,20 +562,43 @@ public class PublishTaskService {
.eq(TaskScopeStateEntity::getModuleType, MODULE_TYPE));
}
private void deleteRolledBackPayloads(List<String> storedPayloads) {
private void deleteUncommittedPayloads(List<String> storedPayloads) {
if (storedPayloads == null || storedPayloads.isEmpty()) {
return;
}
for (String storedPayload : storedPayloads) {
try {
if (isResultPayloadReferenced(storedPayload)) {
log.info("[publish] skip uncommitted payload cleanup because it is referenced pointer={}",
transientPayloadStorageService.extractPointer(storedPayload));
continue;
}
transientPayloadStorageService.deletePayloadIfPresent(storedPayload);
} catch (Exception ex) {
log.warn("[publish] rolled-back RustFS payload cleanup failed pointer={} msg={}",
log.warn("[publish] uncommitted RustFS payload cleanup failed pointer={} msg={}",
transientPayloadStorageService.extractPointer(storedPayload), safeMessage(ex));
}
}
}
private boolean isResultPayloadReferenced(String storedPayload) {
String pointer = transientPayloadStorageService.extractPointer(storedPayload);
if (pointer == null) {
return false;
}
Set<String> candidates = new LinkedHashSet<>();
candidates.add(storedPayload);
candidates.add(pointer);
try {
candidates.add(objectMapper.writeValueAsString(pointer));
} catch (Exception ignored) {
// Raw and extracted pointer values still cover the current publish format.
}
Long count = taskChunkMapper.selectCount(new LambdaQueryWrapper<TaskChunkEntity>()
.in(TaskChunkEntity::getPayloadJson, candidates));
return count != null && count > 0L;
}
private PreparedFile prepareFile(PublishSourceFileDto source) {
PreparedFile prepared = new PreparedFile(source);
prepared.shopName = FileUtil.mainName(source.getOriginalFilename()).trim();
@@ -709,21 +744,19 @@ public class PublishTaskService {
return new PersistedTask(task, result, savedFiles);
}
private boolean submitResultLocked(Long taskId,
PublishSubmitResultRequest request,
List<String> storedPayloads) {
private PreparedResultSubmission prepareResultSubmission(Long taskId,
PublishSubmitResultRequest request,
List<String> uploadedPayloads) {
FileTaskEntity task = requireTask(taskId, request.getUserId());
if (STATUS_SUCCESS.equals(task.getStatus())) {
return false;
return new PreparedResultSubmission(request.getUserId(), List.of());
}
if (STATUS_FAILED.equals(task.getStatus())) {
throw new BusinessException("任务已失败,拒绝继续回传");
}
if (STATUS_PENDING.equals(task.getStatus())) {
task.setStatus(STATUS_RUNNING);
}
Set<Long> submittedFileIds = new LinkedHashSet<>();
List<PreparedResultFile> preparedFiles = new ArrayList<>();
for (PublishResultFileDto incoming : request.getFiles()) {
if (incoming == null) {
throw new BusinessException("files 不能包含空对象");
@@ -733,14 +766,52 @@ public class PublishTaskService {
throw new BusinessException("同一文件不能在一次请求中重复提交");
}
if (STATUS_SUCCESS.equals(file.getStatus()) || STATUS_FAILED.equals(file.getStatus())) {
preparedFiles.add(new PreparedResultFile(file.getId(), null, null));
continue;
}
if (incoming.getError() != null && !incoming.getError().isBlank()) {
preparedFiles.add(new PreparedResultFile(file.getId(), incoming.getError().trim(), null));
continue;
}
preparedFiles.add(new PreparedResultFile(
file.getId(), null, prepareResultChunk(taskId, file, incoming, uploadedPayloads)));
}
return new PreparedResultSubmission(request.getUserId(), List.copyOf(preparedFiles));
}
private SubmitResultCommit submitResultLocked(Long taskId, PreparedResultSubmission prepared) {
FileTaskEntity task = requireTask(taskId, prepared.userId());
if (STATUS_SUCCESS.equals(task.getStatus())) {
return new SubmitResultCommit(false, Set.of());
}
if (STATUS_FAILED.equals(task.getStatus())) {
throw new BusinessException("任务已失败,拒绝继续回传");
}
if (STATUS_PENDING.equals(task.getStatus())) {
task.setStatus(STATUS_RUNNING);
}
Set<Long> submittedFileIds = new LinkedHashSet<>();
Set<String> committedPayloads = new LinkedHashSet<>();
for (PreparedResultFile preparedFile : prepared.files()) {
PublishFileEntity file = requireFile(taskId, preparedFile.fileId());
if (!submittedFileIds.add(file.getId())) {
throw new BusinessException("同一文件不能在一次请求中重复提交");
}
if (STATUS_SUCCESS.equals(file.getStatus()) || STATUS_FAILED.equals(file.getStatus())) {
continue;
}
if (preparedFile.error() != null) {
file.setStatus(STATUS_FAILED);
file.setProcessedRows(0);
file.setErrorMessage(incoming.getError().trim());
file.setErrorMessage(preparedFile.error());
} else {
ResultChunkReceipt receipt = persistResultChunk(taskId, file, incoming, storedPayloads);
PreparedResultChunk preparedChunk = preparedFile.chunk();
if (preparedChunk == null) {
throw new BusinessException("上架结果分片准备状态无效,请重试");
}
ResultChunkReceipt receipt = persistPreparedResultChunk(
taskId, file, preparedChunk, committedPayloads);
if (!receipt.completed()) {
updateReceivedProgress(taskId, file, receipt.receivedRowCount());
file.setStatus(STATUS_RUNNING);
@@ -750,7 +821,10 @@ public class PublishTaskService {
publishFileMapper.updateById(file);
continue;
}
List<PublishRowDto> rows = loadCompleteResultRows(taskId, receipt);
List<PublishRowDto> rows = preparedChunk.completeRows();
if (rows == null) {
throw new BusinessException("上架结果分片状态已变化,请重试");
}
validateCompleteResultRows(taskId, file.getId(), rows);
replaceRows(taskId, file.getId(), rows);
file.setStatus(STATUS_SUCCESS);
@@ -775,11 +849,11 @@ public class PublishTaskService {
if (terminalCount < files.size()) {
task.setStatus(STATUS_RUNNING);
fileTaskMapper.updateById(task);
return false;
return new SubmitResultCommit(false, Set.copyOf(committedPayloads));
}
if (successCount <= 0) {
markTaskAndResultFailed(task, result, "全部文件处理失败");
return true;
return new SubmitResultCommit(true, Set.copyOf(committedPayloads));
}
task.setStatus(STATUS_RUNNING);
@@ -787,7 +861,7 @@ public class PublishTaskService {
task.setFinishedAt(null);
fileTaskMapper.updateById(task);
taskFileJobService.enqueueAssembleResult(taskId, MODULE_TYPE, result.getId(), ownerScopeKey(taskId));
return false;
return new SubmitResultCommit(false, Set.copyOf(committedPayloads));
}
private List<PublishTaskDetailVo> loadTaskDetails(List<FileTaskEntity> tasks) {
@@ -1011,10 +1085,10 @@ public class PublishTaskService {
return file;
}
private ResultChunkReceipt persistResultChunk(Long taskId,
PublishFileEntity file,
PublishResultFileDto incoming,
List<String> storedPayloads) {
private PreparedResultChunk prepareResultChunk(Long taskId,
PublishFileEntity file,
PublishResultFileDto incoming,
List<String> uploadedPayloads) {
int chunkIndex = incoming.getChunkIndex() == null ? 1 : incoming.getChunkIndex();
int chunkTotal = incoming.getChunkTotal() == null ? 1 : incoming.getChunkTotal();
validateChunkMetadata(chunkIndex, chunkTotal);
@@ -1028,21 +1102,83 @@ public class PublishTaskService {
String payloadJson = writeJson(rows, "序列化上架结果分片失败");
String payloadHash = DigestUtil.sha256Hex(payloadJson);
TaskChunkEntity existing = findResultChunk(taskId, scopeHash, chunkIndex);
String storedPayload = null;
if (existing != null) {
validateExistingChunk(existing, chunkTotal, payloadHash);
int receivedChunkCount = countResultChunks(taskId, scopeHash);
int receivedRowCount = resolveReceivedRowCount(taskId, scopeHash, scope, 0, false);
persistResultScope(taskId, scopeKey, scopeHash, chunkTotal,
receivedChunkCount, receivedRowCount);
return new ResultChunkReceipt(scopeHash, chunkTotal,
receivedChunkCount >= chunkTotal, receivedRowCount);
} else {
ensureRustfsPayloadStorageEnabled();
storedPayload = transientPayloadStorageService.storeChunkPayloadVersioned(
MODULE_TYPE, taskId, scopeHash, chunkIndex, payloadJson);
requireRustfsPayload(storedPayload, "上架结果分片必须写入 RustFS");
uploadedPayloads.add(storedPayload);
}
ensureRustfsPayloadStorageEnabled();
String storedPayload = transientPayloadStorageService.storeChunkPayloadVersioned(
MODULE_TYPE, taskId, scopeHash, chunkIndex, payloadJson);
requireRustfsPayload(storedPayload, "上架结果分片必须写入 RustFS");
List<TaskChunkEntity> persistedChunks = listResultChunks(taskId, scopeHash);
List<TaskChunkEntity> expectedChunks = persistedChunks;
if (existing == null) {
expectedChunks = new ArrayList<>(expectedChunks);
expectedChunks.add(newResultChunk(
taskId, scopeKey, scopeHash, chunkIndex, chunkTotal, storedPayload, payloadHash));
expectedChunks.sort((left, right) -> Integer.compare(
Objects.requireNonNullElse(left.getChunkIndex(), 0),
Objects.requireNonNullElse(right.getChunkIndex(), 0)));
}
List<ResultChunkManifestEntry> manifest = expectedChunks.stream()
.map(this::toResultChunkManifestEntry)
.toList();
List<PublishRowDto> completeRows = expectedChunks.size() >= chunkTotal
? readCompleteResultRows(expectedChunks, chunkTotal)
: null;
int receivedRowCount = prepareReceivedRowCount(
scope, persistedChunks, existing == null ? rows.size() : 0, completeRows);
return new PreparedResultChunk(scopeKey, scopeHash, chunkIndex, chunkTotal,
payloadHash, storedPayload, receivedRowCount, manifest, completeRows);
}
private ResultChunkReceipt persistPreparedResultChunk(Long taskId,
PublishFileEntity file,
PreparedResultChunk prepared,
Set<String> committedPayloads) {
TaskScopeStateEntity scope = findResultScope(taskId, prepared.scopeHash());
validateChunkTotal(scope == null ? null : scope.getChunkTotal(), prepared.chunkTotal());
TaskChunkEntity existing = findResultChunk(taskId, prepared.scopeHash(), prepared.chunkIndex());
if (existing != null) {
validateExistingChunk(existing, prepared.chunkTotal(), prepared.payloadHash());
} else if (prepared.storedPayload() != null) {
TaskChunkEntity chunk = newResultChunk(taskId, prepared.scopeKey(), prepared.scopeHash(),
prepared.chunkIndex(), prepared.chunkTotal(), prepared.storedPayload(), prepared.payloadHash());
try {
taskChunkMapper.insert(chunk);
committedPayloads.add(prepared.storedPayload());
} catch (DuplicateKeyException ex) {
TaskChunkEntity winner = findResultChunk(taskId, prepared.scopeHash(), prepared.chunkIndex());
if (winner == null) {
throw new BusinessException("上架结果分片并发写入失败,请重试");
}
validateExistingChunk(winner, prepared.chunkTotal(), prepared.payloadHash());
}
}
List<TaskChunkEntity> actualChunks = listResultChunks(taskId, prepared.scopeHash());
validatePreparedResultManifest(actualChunks, prepared.manifest());
int receivedChunkCount = actualChunks.size();
int receivedRowCount = prepared.receivedRowCount();
persistResultScope(taskId, prepared.scopeKey(), prepared.scopeHash(), prepared.chunkTotal(),
receivedChunkCount, receivedRowCount);
log.info("[publish] result chunk received taskId={} fileId={} chunk={}/{} receivedChunks={} receivedRows={}",
taskId, file.getId(), prepared.chunkIndex(), prepared.chunkTotal(),
receivedChunkCount, receivedRowCount);
return new ResultChunkReceipt(
receivedChunkCount >= prepared.chunkTotal(), receivedRowCount);
}
private TaskChunkEntity newResultChunk(Long taskId,
String scopeKey,
String scopeHash,
int chunkIndex,
int chunkTotal,
String storedPayload,
String payloadHash) {
TaskChunkEntity chunk = new TaskChunkEntity();
chunk.setTaskId(taskId);
chunk.setModuleType(MODULE_TYPE);
@@ -1054,31 +1190,37 @@ public class PublishTaskService {
chunk.setPayloadHash(payloadHash);
chunk.setCreatedAt(LocalDateTime.now());
chunk.setUpdatedAt(LocalDateTime.now());
boolean inserted = false;
try {
taskChunkMapper.insert(chunk);
storedPayloads.add(storedPayload);
inserted = true;
} catch (DuplicateKeyException ex) {
transientPayloadStorageService.deletePayloadIfPresent(storedPayload);
TaskChunkEntity winner = findResultChunk(taskId, scopeHash, chunkIndex);
if (winner == null) {
throw new BusinessException("上架结果分片并发写入失败,请重试");
}
validateExistingChunk(winner, chunkTotal, payloadHash);
} catch (RuntimeException ex) {
transientPayloadStorageService.deletePayloadIfPresent(storedPayload);
throw ex;
}
return chunk;
}
int receivedChunkCount = countResultChunks(taskId, scopeHash);
int receivedRowCount = resolveReceivedRowCount(taskId, scopeHash, scope, rows.size(), inserted);
persistResultScope(taskId, scopeKey, scopeHash, chunkTotal,
receivedChunkCount, receivedRowCount);
log.info("[publish] result chunk received taskId={} fileId={} chunk={}/{} receivedChunks={} receivedRows={}",
taskId, file.getId(), chunkIndex, chunkTotal, receivedChunkCount, receivedRowCount);
return new ResultChunkReceipt(scopeHash, chunkTotal,
receivedChunkCount >= chunkTotal, receivedRowCount);
private List<TaskChunkEntity> listResultChunks(Long taskId, String scopeHash) {
List<TaskChunkEntity> chunks = taskChunkMapper.selectList(new LambdaQueryWrapper<TaskChunkEntity>()
.eq(TaskChunkEntity::getTaskId, taskId)
.eq(TaskChunkEntity::getModuleType, MODULE_TYPE)
.eq(TaskChunkEntity::getScopeHash, scopeHash)
.orderByAsc(TaskChunkEntity::getChunkIndex));
return chunks == null ? List.of() : chunks;
}
private ResultChunkManifestEntry toResultChunkManifestEntry(TaskChunkEntity chunk) {
return new ResultChunkManifestEntry(
chunk.getChunkIndex(), chunk.getChunkTotal(), chunk.getPayloadHash());
}
private void validatePreparedResultManifest(List<TaskChunkEntity> actualChunks,
List<ResultChunkManifestEntry> expectedManifest) {
if (actualChunks.size() != expectedManifest.size()) {
throw new BusinessException("上架结果分片状态已变化,请重试");
}
for (int i = 0; i < actualChunks.size(); i++) {
ResultChunkManifestEntry actual = toResultChunkManifestEntry(actualChunks.get(i));
ResultChunkManifestEntry expected = expectedManifest.get(i);
if (!Objects.equals(actual.chunkIndex(), expected.chunkIndex())
|| !Objects.equals(actual.chunkTotal(), expected.chunkTotal())
|| !Objects.equals(actual.payloadHash(), expected.payloadHash())) {
throw new BusinessException("上架结果分片状态已变化,请重试");
}
}
}
private void validateChunkMetadata(int chunkIndex, int chunkTotal) {
@@ -1123,31 +1265,20 @@ public class PublishTaskService {
}
}
private int countResultChunks(Long taskId, String scopeHash) {
Long count = taskChunkMapper.selectCount(new LambdaQueryWrapper<TaskChunkEntity>()
.eq(TaskChunkEntity::getTaskId, taskId)
.eq(TaskChunkEntity::getModuleType, MODULE_TYPE)
.eq(TaskChunkEntity::getScopeHash, scopeHash));
return count == null ? 0 : count.intValue();
}
/**
* Resolve the number of result rows received for a file without re-reading
* every payload on every callback. New callbacks keep the count in the
* existing scope state JSON; scopes created by older versions are repaired
* once by counting their stored chunks.
*/
private int resolveReceivedRowCount(Long taskId,
String scopeHash,
TaskScopeStateEntity scope,
int currentChunkRows,
boolean inserted) {
private int prepareReceivedRowCount(TaskScopeStateEntity scope,
List<TaskChunkEntity> expectedChunks,
int newChunkRows,
List<PublishRowDto> completeRows) {
Integer persisted = readReceivedRowCount(scope);
if (persisted != null) {
long next = (long) persisted + (inserted ? Math.max(0, currentChunkRows) : 0);
long next = (long) persisted + Math.max(0, newChunkRows);
return next > Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) Math.max(0, next);
}
return countReceivedResultRows(taskId, scopeHash);
if (completeRows != null) {
return completeRows.size();
}
long preparedRows = (long) countPreparedResultRows(expectedChunks) + Math.max(0, newChunkRows);
return preparedRows > Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) preparedRows;
}
private Integer readReceivedRowCount(TaskScopeStateEntity scope) {
@@ -1168,13 +1299,8 @@ public class PublishTaskService {
}
}
private int countReceivedResultRows(Long taskId, String scopeHash) {
List<TaskChunkEntity> chunks = taskChunkMapper.selectList(new LambdaQueryWrapper<TaskChunkEntity>()
.eq(TaskChunkEntity::getTaskId, taskId)
.eq(TaskChunkEntity::getModuleType, MODULE_TYPE)
.eq(TaskChunkEntity::getScopeHash, scopeHash)
.orderByAsc(TaskChunkEntity::getChunkIndex));
if (chunks == null || chunks.isEmpty()) {
private int countPreparedResultRows(List<TaskChunkEntity> chunks) {
if (chunks.isEmpty()) {
return 0;
}
TypeReference<List<PublishRowDto>> listType = new TypeReference<>() {
@@ -1268,13 +1394,8 @@ public class PublishTaskService {
file.setProcessedRows(progress);
}
private List<PublishRowDto> loadCompleteResultRows(Long taskId, ResultChunkReceipt receipt) {
List<TaskChunkEntity> chunks = taskChunkMapper.selectList(new LambdaQueryWrapper<TaskChunkEntity>()
.eq(TaskChunkEntity::getTaskId, taskId)
.eq(TaskChunkEntity::getModuleType, MODULE_TYPE)
.eq(TaskChunkEntity::getScopeHash, receipt.scopeHash())
.orderByAsc(TaskChunkEntity::getChunkIndex));
if (chunks == null || chunks.size() != receipt.chunkTotal()) {
private List<PublishRowDto> readCompleteResultRows(List<TaskChunkEntity> chunks, int chunkTotal) {
if (chunks.size() != chunkTotal) {
throw new BusinessException("上架结果分片尚未完整,暂不能合并");
}
@@ -1287,7 +1408,7 @@ public class PublishTaskService {
if (!Objects.equals(chunk.getChunkIndex(), expectedIndex)) {
throw new BusinessException("上架结果缺少第 " + expectedIndex + " 个分片");
}
validateChunkTotal(chunk.getChunkTotal(), receipt.chunkTotal());
validateChunkTotal(chunk.getChunkTotal(), chunkTotal);
List<PublishRowDto> chunkRows = readResultChunkRows(chunk, listType);
for (PublishRowDto row : chunkRows) {
rows.add(copyRequiredRow(row));
@@ -1737,9 +1858,36 @@ public class PublishTaskService {
private record TaskOptions(String publishCountry, List<String> syncCountries) {
}
private record ResultChunkReceipt(String scopeHash,
int chunkTotal,
boolean completed,
private record PreparedResultSubmission(Long userId,
List<PreparedResultFile> files) {
}
private record PreparedResultFile(Long fileId,
String error,
PreparedResultChunk chunk) {
}
private record PreparedResultChunk(String scopeKey,
String scopeHash,
int chunkIndex,
int chunkTotal,
String payloadHash,
String storedPayload,
int receivedRowCount,
List<ResultChunkManifestEntry> manifest,
List<PublishRowDto> completeRows) {
}
private record ResultChunkManifestEntry(Integer chunkIndex,
Integer chunkTotal,
String payloadHash) {
}
private record SubmitResultCommit(boolean cleanupAfterCommit,
Set<String> committedPayloads) {
}
private record ResultChunkReceipt(boolean completed,
int receivedRowCount) {
}
}
@@ -655,138 +655,44 @@ public class SimilarAsinTaskService {
}
public void submitResult(Long taskId, SimilarAsinSubmitResultRequest request) {
submitResultLocked(taskId, request);
try (TaskDistributedLockService.LockHandle ignored = requireTaskLock(taskId, TASK_LOCK_WAIT_MILLIS)) {
submitResultLocked(taskId, request);
}
}
private void submitResultLocked(Long taskId, SimilarAsinSubmitResultRequest request) {
if (transactionManager != null) {
SubmitContext context = inNewTransaction(() -> persistSubmittedChunk(taskId, request));
inNewTransaction(() -> {
completeSubmittedChunk(context);
return null;
});
scheduleCozePipelineForSubmittedChunk(context);
return;
if (request == null) {
throw new BusinessException("结果请求不能为空");
}
FileTaskEntity task = fileTaskMapper.selectById(taskId);
if (task == null || !MODULE_TYPE.equals(task.getModuleType())) {
throw new BusinessException("任务不存在");
PreparedSubmittedChunk prepared = prepareSubmittedChunk(taskId, request);
PersistSubmittedChunkResult persisted = transactionManager == null
? persistSubmittedChunk(prepared)
: inNewTransaction(() -> persistSubmittedChunk(prepared));
SubmitContext context = persisted.context();
// A thrown transaction may already have committed but lost its ACK. Keep
// that candidate; cleanup is only safe after a successful transaction
// has proved that another chunk won the unique key.
if (!persisted.payloadPersisted()) {
cleanupPreparedSubmittedChunkIfUnreferenced(prepared);
}
if (!STATUS_RUNNING.equals(task.getStatus())) {
throw new BusinessException("任务不是运行中状态");
}
ensureTaskOwnedByCurrentInstance(task, "submit result");
int chunkIndex = request.getChunkIndex() == null ? 0 : request.getChunkIndex();
int chunkTotal = request.getChunkTotal() == null ? 1 : request.getChunkTotal();
String scopeKey = firstNonBlank(request.getSubmissionId(), "task:" + taskId);
String scopeHash = DigestUtil.sha256Hex(scopeKey);
taskCacheService.touchTaskHeartbeat(taskId);
TaskChunkEntity existing = taskChunkMapper.selectOne(new LambdaQueryWrapper<TaskChunkEntity>()
.eq(TaskChunkEntity::getTaskId, taskId)
.eq(TaskChunkEntity::getModuleType, MODULE_TYPE)
.eq(TaskChunkEntity::getScopeHash, scopeHash)
.eq(TaskChunkEntity::getChunkIndex, chunkIndex)
.last("limit 1"));
if (existing == null) {
List<SimilarAsinResultRowDto> rawRows = flattenSubmittedRows(request);
String payloadJson = writeJson(rawRows, "结果序列化失败");
TaskChunkEntity chunk = new TaskChunkEntity();
chunk.setTaskId(taskId);
chunk.setModuleType(MODULE_TYPE);
chunk.setScopeKey(scopeKey);
chunk.setScopeHash(scopeHash);
chunk.setChunkIndex(chunkIndex);
chunk.setChunkTotal(chunkTotal);
// P0-2改用 versioned keychunk-{index}-{uuid}避免并发回调下 deterministic key
// 覆盖造成的 read chunk payload failed / 数据丢失loser 拿到的是自己独有的 object
// DuplicateKeyException 后可以安全删除不会影响 winner chunk
String storedPayload = transientPayloadStorageService.storeChunkPayloadVersioned(MODULE_TYPE, taskId, scopeHash, chunkIndex, payloadJson);
// P0-3判断本次 store 是否走了 RustFS local 兜底
boolean localFallback = transientPayloadStorageService.wasLastStoreLocalFallback();
chunk.setPayloadJson(storedPayload);
chunk.setPayloadHash(DigestUtil.sha256Hex(payloadJson));
chunk.setCreatedAt(LocalDateTime.now());
chunk.setUpdatedAt(LocalDateTime.now());
try {
taskChunkMapper.insert(chunk);
} catch (DuplicateKeyException ex) {
// P0-2versioned key 每次都是独立 objectloser storedPayload 不会被 winner 引用
// 这里直接清理掉以免 RustFS / local 上残留孤儿对象
log.info("[similar-asin] duplicate chunk inserted concurrently taskId={} scope={} chunk={} cleanupLoser={}",
taskId, scopeKey, chunkIndex, storedPayload != null);
try {
transientPayloadStorageService.deletePayloadIfPresent(storedPayload);
} catch (Exception cleanupEx) {
log.warn("[similar-asin] cleanup loser chunk payload failed taskId={} chunk={} err={}",
taskId, chunkIndex, cleanupEx.getMessage());
}
}
// P0-3本次 chunk 落到本地时把 task 锁到当前实例避免其他实例 assemble 时读不到
if (localFallback) {
bindTaskToCurrentOwnerForLocalFallback(task, scopeHash, chunkIndex);
}
if (persisted.finalizeResult() == null) {
taskCacheService.touchTaskHeartbeat(taskId);
} else {
log.info("[similar-asin] duplicate chunk ignored taskId={} scope={} chunk={}", taskId, scopeKey, chunkIndex);
}
TaskScopeStateEntity scope = taskScopeStateMapper.selectOne(new LambdaQueryWrapper<TaskScopeStateEntity>()
.eq(TaskScopeStateEntity::getTaskId, taskId)
.eq(TaskScopeStateEntity::getModuleType, MODULE_TYPE)
.eq(TaskScopeStateEntity::getScopeHash, scopeHash)
.last("limit 1"));
if (scope == null) {
scope = new TaskScopeStateEntity();
scope.setTaskId(taskId);
scope.setModuleType(MODULE_TYPE);
scope.setScopeKey(scopeKey);
scope.setScopeHash(scopeHash);
scope.setCreatedAt(LocalDateTime.now());
}
scope.setChunkTotal(chunkTotal);
scope.setReceivedChunkCount(resolveReceivedChunkProgress(taskId, scopeHash, scope.getChunkTotal()));
scope.setLastChunkAt(LocalDateTime.now());
scope.setLastError(request.getError());
scope.setCompleted(Boolean.TRUE.equals(request.getDone()) ? 1 : 0);
scope.setUpdatedAt(LocalDateTime.now());
scope.setStateJson("{\"phase\":\"RECEIVED\",\"coze\":\"PENDING\"}");
if (scope.getId() == null) {
taskScopeStateMapper.insert(scope);
} else {
taskScopeStateMapper.updateById(scope);
}
SubmitContext context = new SubmitContext(task, scopeKey, scopeHash, chunkIndex,
Boolean.TRUE.equals(request.getDone()), request.getError());
if (Boolean.TRUE.equals(request.getDone()) || request.getError() != null && !request.getError().isBlank()) {
finalizeTask(task, request.getError(), allRowCount(task), true);
} else {
task.setUpdatedAt(LocalDateTime.now());
fileTaskMapper.updateById(task);
applyFinalizeSideEffects(persisted.finalizeResult());
}
scheduleCozePipelineForSubmittedChunk(context);
}
@Transactional
public void deleteTask(Long taskId, Long userId) {
List<String> payloads;
try (TaskDistributedLockService.LockHandle ignored = requireTaskLock(taskId, TASK_LOCK_WAIT_MILLIS)) {
FileTaskEntity task = fileTaskMapper.selectById(taskId);
if (task == null || !MODULE_TYPE.equals(task.getModuleType()) || !Objects.equals(userId, task.getUserId())) {
throw new BusinessException("任务不存在");
payloads = transactionManager == null
? deleteTaskRecords(taskId, userId)
: inNewTransaction(() -> deleteTaskRecords(taskId, userId));
}
fileResultMapper.delete(new LambdaQueryWrapper<FileResultEntity>().eq(FileResultEntity::getTaskId, taskId).eq(FileResultEntity::getModuleType, MODULE_TYPE));
deleteTransientTaskPayloads(taskId);
taskScopeStateMapper.delete(new LambdaQueryWrapper<TaskScopeStateEntity>().eq(TaskScopeStateEntity::getTaskId, taskId).eq(TaskScopeStateEntity::getModuleType, MODULE_TYPE));
taskChunkMapper.delete(new LambdaQueryWrapper<TaskChunkEntity>().eq(TaskChunkEntity::getTaskId, taskId).eq(TaskChunkEntity::getModuleType, MODULE_TYPE));
deleteTransientPayloads(payloads, taskId);
taskCacheService.deleteTaskCache(taskId);
// 同步清理 task_file_job避免被删除任务遗留的 PENDING/FAILED 行被 TaskResultFileJobWorker 反复扫描出任务不存在
taskFileJobService.deleteTaskJobs(taskId, MODULE_TYPE);
fileTaskMapper.deleteById(taskId);
// P1-1任务被删除时一并清理滑窗记录防止内存泄漏
clearPoisonWindow(taskId);
}
}
public void deleteHistory(Long resultId, Long userId) {
@@ -799,6 +705,69 @@ public class SimilarAsinTaskService {
fileResultMapper.deleteById(resultId);
}
private List<String> deleteTaskRecords(Long taskId, Long userId) {
FileTaskEntity task = fileTaskMapper.selectById(taskId);
if (task == null || !MODULE_TYPE.equals(task.getModuleType()) || !Objects.equals(userId, task.getUserId())) {
throw new BusinessException("任务不存在");
}
List<String> payloads = collectTransientTaskPayloads(taskId);
fileResultMapper.delete(new LambdaQueryWrapper<FileResultEntity>()
.eq(FileResultEntity::getTaskId, taskId)
.eq(FileResultEntity::getModuleType, MODULE_TYPE));
taskScopeStateMapper.delete(new LambdaQueryWrapper<TaskScopeStateEntity>()
.eq(TaskScopeStateEntity::getTaskId, taskId)
.eq(TaskScopeStateEntity::getModuleType, MODULE_TYPE));
taskChunkMapper.delete(new LambdaQueryWrapper<TaskChunkEntity>()
.eq(TaskChunkEntity::getTaskId, taskId)
.eq(TaskChunkEntity::getModuleType, MODULE_TYPE));
taskFileJobService.deleteTaskJobs(taskId, MODULE_TYPE);
fileTaskMapper.deleteById(taskId);
return payloads;
}
private void deleteTransientPayloads(List<String> payloads, Long taskId) {
if (payloads == null || payloads.isEmpty()) {
return;
}
for (String payload : new LinkedHashSet<>(payloads)) {
try {
if (isTransientPayloadStillReferenced(payload)) {
log.info("[similar-asin] skip task payload delete because it is still referenced taskId={} pointer={}",
taskId, transientPayloadStorageService.extractPointer(payload));
continue;
}
transientPayloadStorageService.deletePayloadIfPresent(payload);
} catch (Exception ex) {
log.warn("[similar-asin] delete task payload deferred taskId={} pointer={} err={}",
taskId, transientPayloadStorageService.extractPointer(payload), ex.getMessage());
}
}
}
private boolean isTransientPayloadStillReferenced(String payload) {
String pointer = transientPayloadStorageService.extractPointer(payload);
if (pointer == null) {
return false;
}
LinkedHashSet<String> values = new LinkedHashSet<>();
values.add(payload);
values.add(pointer);
try {
values.add(objectMapper.writeValueAsString(pointer));
} catch (Exception ignored) {
}
Long chunkCount = taskChunkMapper.selectCount(new LambdaQueryWrapper<TaskChunkEntity>()
.in(TaskChunkEntity::getPayloadJson, values));
if (chunkCount != null && chunkCount > 0L) {
return true;
}
Long scopeCount = taskScopeStateMapper.selectCount(new LambdaQueryWrapper<TaskScopeStateEntity>()
.and(wrapper -> wrapper.in(TaskScopeStateEntity::getParsedPayloadJson, values)
.or()
.in(TaskScopeStateEntity::getStateJson, values)));
return scopeCount != null && scopeCount > 0L;
}
public String resolveResultDownloadUrl(Long resultId, Long userId) {
FileResultEntity row = fileResultMapper.selectById(resultId);
if (row == null || !MODULE_TYPE.equals(row.getModuleType()) || !Objects.equals(userId, row.getUserId())) {
@@ -873,14 +842,7 @@ public class SimilarAsinTaskService {
continue;
}
try (taskLockHandle) {
if (transactionManager != null) {
inNewTransaction(() -> {
finalizeStaleTask(task.getId(), "Python interrupted before uploading final similar ASIN result");
return null;
});
} else {
finalizeStaleTask(task.getId(), "Python interrupted before uploading final similar ASIN result");
}
finalizeStaleTask(task.getId(), "Python interrupted before uploading final similar ASIN result");
}
}
}
@@ -889,19 +851,6 @@ public class SimilarAsinTaskService {
if (taskId == null || taskId <= 0) {
return;
}
if (transactionManager != null) {
TaskDistributedLockService.LockHandle taskLockHandle = acquireTaskLock(taskId, 0L);
if (taskLockHandle == null) {
return;
}
try (taskLockHandle) {
inNewTransaction(() -> {
finalizeStaleTask(taskId, "Python interrupted before uploading final similar ASIN result");
return null;
});
}
return;
}
TaskDistributedLockService.LockHandle taskLockHandle = acquireTaskLock(taskId, 0L);
if (taskLockHandle == null) {
return;
@@ -943,7 +892,7 @@ public class SimilarAsinTaskService {
return updatedMillis <= thresholdMillis;
}
private SubmitContext persistSubmittedChunk(Long taskId, SimilarAsinSubmitResultRequest request) {
private PreparedSubmittedChunk prepareSubmittedChunk(Long taskId, SimilarAsinSubmitResultRequest request) {
FileTaskEntity task = fileTaskMapper.selectById(taskId);
if (task == null || !MODULE_TYPE.equals(task.getModuleType())) {
throw new BusinessException("任务不存在");
@@ -958,62 +907,132 @@ public class SimilarAsinTaskService {
boolean done = Boolean.TRUE.equals(request.getDone());
String scopeKey = firstNonBlank(request.getSubmissionId(), "task:" + taskId);
String scopeHash = DigestUtil.sha256Hex(scopeKey);
taskCacheService.touchTaskHeartbeat(taskId);
boolean terminalCallback = done || request.getError() != null && !request.getError().isBlank();
SubmittedTaskMetadata taskMetadata = terminalCallback
? readSubmittedTaskMetadata(task)
: null;
TaskChunkEntity existing = findSubmittedChunk(taskId, scopeHash, chunkIndex);
if (existing != null) {
return new PreparedSubmittedChunk(
taskId, scopeKey, scopeHash, chunkIndex, chunkTotal, done, request.getError(),
null, null, false, taskMetadata);
}
String payloadJson = writeJson(flattenSubmittedRows(request), "结果序列化失败");
String storedPayload = transientPayloadStorageService.storeChunkPayloadVersioned(
MODULE_TYPE, taskId, scopeHash, chunkIndex, payloadJson);
boolean localFallback = transientPayloadStorageService.wasLastStoreLocalFallback();
return new PreparedSubmittedChunk(
taskId, scopeKey, scopeHash, chunkIndex, chunkTotal, done, request.getError(),
payloadJson, storedPayload, localFallback, taskMetadata);
}
TaskChunkEntity existing = taskChunkMapper.selectOne(new LambdaQueryWrapper<TaskChunkEntity>()
private SubmittedTaskMetadata readSubmittedTaskMetadata(FileTaskEntity task) {
SimilarAsinParsedPayloadDto payload = readParsedPayload(task);
List<SimilarAsinSourceFileDto> sourceFiles = payload.getSourceFiles() == null
? List.of()
: List.copyOf(payload.getSourceFiles());
return new SubmittedTaskMetadata(rowCount(payload), sourceFiles);
}
private int rowCount(SimilarAsinParsedPayloadDto payload) {
if (payload == null) {
return 0;
}
if (payload.getAllItems() != null && !payload.getAllItems().isEmpty()) {
return payload.getAllItems().size();
}
if (payload.getItems() != null && !payload.getItems().isEmpty()) {
return payload.getItems().size();
}
if (payload.getGroups() != null && !payload.getGroups().isEmpty()) {
return payload.getGroups().stream()
.map(SimilarAsinParsedGroupVo::getItems)
.filter(Objects::nonNull)
.mapToInt(List::size)
.sum();
}
return 0;
}
private PersistSubmittedChunkResult persistSubmittedChunk(PreparedSubmittedChunk prepared) {
Long taskId = prepared.taskId();
FileTaskEntity task = fileTaskMapper.selectById(taskId);
if (task == null || !MODULE_TYPE.equals(task.getModuleType())) {
throw new BusinessException("任务不存在");
}
if (!STATUS_RUNNING.equals(task.getStatus())) {
throw new BusinessException("任务不是运行中状态");
}
ensureTaskOwnedByCurrentInstance(task, "submit result");
boolean payloadPersisted = prepared.storedPayload() == null;
TaskChunkEntity existing = findSubmittedChunk(taskId, prepared.scopeHash(), prepared.chunkIndex());
if (existing == null && prepared.storedPayload() != null) {
TaskChunkEntity chunk = new TaskChunkEntity();
chunk.setTaskId(taskId);
chunk.setModuleType(MODULE_TYPE);
chunk.setScopeKey(prepared.scopeKey());
chunk.setScopeHash(prepared.scopeHash());
chunk.setChunkIndex(prepared.chunkIndex());
chunk.setChunkTotal(prepared.chunkTotal());
chunk.setPayloadJson(prepared.storedPayload());
chunk.setPayloadHash(DigestUtil.sha256Hex(prepared.payloadJson()));
chunk.setCreatedAt(LocalDateTime.now());
chunk.setUpdatedAt(LocalDateTime.now());
try {
if (taskChunkMapper.insert(chunk) != 1) {
throw new IllegalStateException("Chunk insert affected no row taskId=" + taskId
+ " chunk=" + prepared.chunkIndex());
}
payloadPersisted = true;
} catch (DuplicateKeyException ex) {
log.info("[similar-asin] duplicate chunk inserted concurrently taskId={} scope={} chunk={} cleanupLoser={}",
taskId, prepared.scopeKey(), prepared.chunkIndex(), prepared.storedPayload() != null);
}
} else {
log.info("[similar-asin] duplicate chunk ignored taskId={} scope={} chunk={}",
taskId, prepared.scopeKey(), prepared.chunkIndex());
}
upsertScopeState(taskId, prepared.scopeKey(), prepared.scopeHash(), prepared.chunkTotal(),
prepared.error(), prepared.done(), false);
if (prepared.localFallback() && payloadPersisted) {
bindTaskToCurrentOwnerForLocalFallback(
task, prepared.scopeHash(), prepared.chunkIndex());
}
SubmitContext context = new SubmitContext(task, prepared.scopeKey(), prepared.scopeHash(), prepared.chunkIndex(),
prepared.done(), prepared.error(), prepared.taskMetadata());
FinalizeTaskResult finalizeResult = completeSubmittedChunk(context);
return new PersistSubmittedChunkResult(context, payloadPersisted, finalizeResult);
}
private TaskChunkEntity findSubmittedChunk(Long taskId, String scopeHash, Integer chunkIndex) {
return taskChunkMapper.selectOne(new LambdaQueryWrapper<TaskChunkEntity>()
.eq(TaskChunkEntity::getTaskId, taskId)
.eq(TaskChunkEntity::getModuleType, MODULE_TYPE)
.eq(TaskChunkEntity::getScopeHash, scopeHash)
.eq(TaskChunkEntity::getChunkIndex, chunkIndex)
.last("limit 1"));
if (existing == null) {
List<SimilarAsinResultRowDto> rawRows = flattenSubmittedRows(request);
String payloadJson = writeJson(rawRows, "结果序列化失败");
TaskChunkEntity chunk = new TaskChunkEntity();
chunk.setTaskId(taskId);
chunk.setModuleType(MODULE_TYPE);
chunk.setScopeKey(scopeKey);
chunk.setScopeHash(scopeHash);
chunk.setChunkIndex(chunkIndex);
chunk.setChunkTotal(chunkTotal);
// P0-2 submitResultLocked 同步切到 versioned keychunk-{index}-{uuid}
// 避免并发回调互相覆盖造成 read chunk payload failed
String storedPayload = transientPayloadStorageService.storeChunkPayloadVersioned(MODULE_TYPE, taskId, scopeHash, chunkIndex, payloadJson);
// P0-3拿到 pointer 后立刻读 ThreadLocal 标记决定是否需要把 task 锁到当前实例
boolean localFallback = transientPayloadStorageService.wasLastStoreLocalFallback();
chunk.setPayloadJson(storedPayload);
chunk.setPayloadHash(DigestUtil.sha256Hex(payloadJson));
chunk.setCreatedAt(LocalDateTime.now());
chunk.setUpdatedAt(LocalDateTime.now());
try {
taskChunkMapper.insert(chunk);
} catch (DuplicateKeyException ex) {
// P0-2versioned key loser storedPayload 不会被 winner 引用可以安全清理
log.info("[similar-asin] duplicate chunk inserted concurrently taskId={} scope={} chunk={} cleanupLoser={}",
taskId, scopeKey, chunkIndex, storedPayload != null);
try {
transientPayloadStorageService.deletePayloadIfPresent(storedPayload);
} catch (Exception cleanupEx) {
log.warn("[similar-asin] cleanup loser chunk payload failed taskId={} chunk={} err={}",
taskId, chunkIndex, cleanupEx.getMessage());
}
}
// P0-3fallback local 时把 task 与当前实例绑定确保后续 assemble 走对实例
if (localFallback) {
bindTaskToCurrentOwnerForLocalFallback(task, scopeHash, chunkIndex);
}
} else {
log.info("[similar-asin] duplicate chunk ignored taskId={} scope={} chunk={}", taskId, scopeKey, chunkIndex);
}
upsertScopeState(taskId, scopeKey, scopeHash, chunkTotal, request.getError(), done, false);
task.setUpdatedAt(LocalDateTime.now());
fileTaskMapper.updateById(task);
return new SubmitContext(task, scopeKey, scopeHash, chunkIndex, done, request.getError());
}
private void completeSubmittedChunk(SubmitContext context) {
private void cleanupPreparedSubmittedChunkIfUnreferenced(PreparedSubmittedChunk prepared) {
if (prepared == null || prepared.storedPayload() == null) {
return;
}
try {
if (isTransientPayloadStillReferenced(prepared.storedPayload())) {
log.info("[similar-asin] keep prepared chunk payload because it is referenced taskId={} chunk={}",
prepared.taskId(), prepared.chunkIndex());
return;
}
transientPayloadStorageService.deletePayloadIfPresent(prepared.storedPayload());
} catch (Exception ex) {
log.warn("[similar-asin] cleanup uncommitted chunk payload failed taskId={} chunk={} err={}",
prepared.taskId(), prepared.chunkIndex(), ex.getMessage());
}
}
private FinalizeTaskResult completeSubmittedChunk(SubmitContext context) {
FileTaskEntity task = fileTaskMapper.selectById(context.task().getId());
if (task == null || !MODULE_TYPE.equals(task.getModuleType())) {
throw new BusinessException("任务不存在");
@@ -1021,15 +1040,13 @@ public class SimilarAsinTaskService {
if (!STATUS_RUNNING.equals(task.getStatus())) {
log.info("[similar-asin] skip completion because task already finalized taskId={} status={}",
task.getId(), task.getStatus());
return;
return null;
}
upsertScopeState(task.getId(), context.scopeKey(), context.scopeHash(), null, context.error(), context.forceFlush(), false);
if (context.forceFlush() || context.error() != null && !context.error().isBlank()) {
finalizeTask(task, context.error(), allRowCount(task), true);
return;
return finalizeTask(task, context.error(), context.taskMetadata(), true);
}
task.setUpdatedAt(LocalDateTime.now());
fileTaskMapper.updateById(task);
touchJavaSideTaskActivity(task.getId());
return null;
}
private void submitCozeForSubmittedChunk(SubmitContext context) {
@@ -1118,7 +1135,11 @@ public class SimilarAsinTaskService {
if (tryRecoverTimedOutPythonTask(task)) {
return;
}
finalizeTask(task, error, allRowCount(task), true);
SubmittedTaskMetadata metadata = readSubmittedTaskMetadata(task);
FinalizeTaskResult result = transactionManager == null
? finalizeTask(task, error, metadata, true)
: inNewTransaction(() -> finalizeTask(task, error, metadata, true));
applyFinalizeSideEffects(result);
}
private boolean tryRecoverTimedOutPythonTask(FileTaskEntity task) {
@@ -1143,7 +1164,7 @@ public class SimilarAsinTaskService {
log.warn("[similar-asin] Python 超时恢复继续推进 Coze/文件收尾 taskId={} pendingCozeStates={} activeAssembleJobs={}",
taskId, pendingCozeStates, activeAssembleJobs);
}
submitCozeForSubmittedChunk(new SubmitContext(task, null, null, null, true, null));
submitCozeForSubmittedChunk(new SubmitContext(task, null, null, null, true, null, null));
touchJavaSideTaskActivity(taskId);
return true;
}
@@ -1220,7 +1241,7 @@ public class SimilarAsinTaskService {
scope.setReceivedChunkCount(countChunks(taskId, scopeHash));
scope.setLastChunkAt(now);
scope.setLastError(error);
scope.setCompleted(completed ? 1 : 0);
scope.setCompleted(completed || Integer.valueOf(1).equals(scope.getCompleted()) ? 1 : 0);
scope.setUpdatedAt(now);
scope.setStateJson(cozeDone
? "{\"phase\":\"RECEIVED\",\"coze\":\"DONE\"}"
@@ -1247,7 +1268,7 @@ public class SimilarAsinTaskService {
scope.setReceivedChunkCount(resolveReceivedChunkProgress(taskId, scopeHash, scope.getChunkTotal()));
scope.setLastChunkAt(now);
scope.setLastError(error);
scope.setCompleted(completed ? 1 : 0);
scope.setCompleted(completed || Integer.valueOf(1).equals(scope.getCompleted()) ? 1 : 0);
scope.setUpdatedAt(now);
scope.setStateJson(cozeDone
? "{\"phase\":\"RECEIVED\",\"coze\":\"DONE\"}"
@@ -1573,21 +1594,7 @@ public class SimilarAsinTaskService {
private int allRowCount(FileTaskEntity task) {
try {
SimilarAsinParsedPayloadDto payload = readParsedPayload(task);
if (payload.getAllItems() != null && !payload.getAllItems().isEmpty()) {
return payload.getAllItems().size();
}
if (payload.getItems() != null && !payload.getItems().isEmpty()) {
return payload.getItems().size();
}
if (payload.getGroups() != null && !payload.getGroups().isEmpty()) {
return payload.getGroups().stream()
.map(SimilarAsinParsedGroupVo::getItems)
.filter(Objects::nonNull)
.mapToInt(List::size)
.sum();
}
return 0;
return rowCount(readParsedPayload(task));
} catch (Exception ex) {
log.warn("[similar-asin] read all row count failed taskId={} err={}", task.getId(), ex.getMessage());
return 0;
@@ -1623,8 +1630,12 @@ public class SimilarAsinTaskService {
return Math.max(1, configured);
}
private void finalizeTask(FileTaskEntity task, String error, int rowCount, boolean assembleWorkbook) {
private FinalizeTaskResult finalizeTask(FileTaskEntity task,
String error,
SubmittedTaskMetadata taskMetadata,
boolean assembleWorkbook) {
String finalError = error;
int rowCount = taskMetadata == null ? 0 : taskMetadata.rowCount();
FileResultEntity result = null;
boolean hasPersistedResultRows = hasPersistedResultRows(task.getId());
List<FileResultEntity> rows = fileResultMapper.selectList(new LambdaQueryWrapper<FileResultEntity>()
@@ -1634,20 +1645,13 @@ public class SimilarAsinTaskService {
if (!rows.isEmpty()) {
result = rows.getFirst();
} else if (assembleWorkbook && hasPersistedResultRows) {
result = createResultRecordForAssembly(task, rowCount);
result = createResultRecordForAssembly(task, rowCount,
taskMetadata == null ? null : taskMetadata.sourceFiles());
}
if (result != null) {
boolean shouldAssembleResult = assembleWorkbook && (finalError == null || finalError.isBlank() || hasPersistedResultRows);
if (shouldAssembleResult && shouldAssembleSynchronously()) {
try {
// P0-3stale-recovery 同步 finalize 路径走这里确保
// assemble 之前缓冲的 cozeRows 已合并到 chunk
flushBufferedCozeResults(task.getId());
assembleResultWorkbook(task, result);
} catch (Exception ex) {
log.warn("[相似ASIN] 组装结果工作簿失败 任务ID={} 错误={}", task.getId(), ex.getMessage());
finalError = firstNonBlank(finalError, "生成相似ASIN结果失败");
}
throw new IllegalStateException("Synchronous result assembly is not allowed in a database transaction");
}
}
boolean failed = finalError != null && !finalError.isBlank();
@@ -1677,11 +1681,16 @@ public class SimilarAsinTaskService {
log.warn("[相似ASIN] 跳过失败任务的工作簿组装,因为没有已持久化结果行 任务ID={}", task.getId());
}
}
taskCacheService.deleteTaskCache(task.getId());
// P1-1任务进入终态SUCCESS/FAILED后清理滑窗记录避免长任务残留内存
// 处于 waitingForAssemble仍为 RUNNING等待异步 assemble时不清理 assemble 完真正进入终态再由后续路径触发
if (!waitingForAssemble) {
clearPoisonWindow(task.getId());
return new FinalizeTaskResult(task.getId(), !waitingForAssemble);
}
private void applyFinalizeSideEffects(FinalizeTaskResult result) {
if (result == null || result.taskId() == null) {
return;
}
taskCacheService.deleteTaskCache(result.taskId());
if (result.terminal()) {
clearPoisonWindow(result.taskId());
}
}
@@ -1716,12 +1725,19 @@ public class SimilarAsinTaskService {
private FileResultEntity createResultRecordForAssembly(FileTaskEntity task, int rowCount) {
SimilarAsinParsedPayloadDto payload = readParsedPayload(task);
List<SimilarAsinSourceFileDto> sourceFiles = payload.getSourceFiles() == null ? List.of() : payload.getSourceFiles();
return createResultRecordForAssembly(task, rowCount,
payload.getSourceFiles() == null ? List.of() : payload.getSourceFiles());
}
private FileResultEntity createResultRecordForAssembly(FileTaskEntity task,
int rowCount,
List<SimilarAsinSourceFileDto> sourceFiles) {
List<SimilarAsinSourceFileDto> resolvedSourceFiles = sourceFiles == null ? List.of() : sourceFiles;
FileResultEntity result = new FileResultEntity();
result.setTaskId(task.getId());
result.setModuleType(MODULE_TYPE);
result.setSourceFilename(buildAggregateSourceFilenameLabel(sourceFiles));
result.setSourceFileUrl(buildAggregateScopeKey(sourceFiles));
result.setSourceFilename(buildAggregateSourceFilenameLabel(resolvedSourceFiles));
result.setSourceFileUrl(buildAggregateScopeKey(resolvedSourceFiles));
result.setRowCount(rowCount);
result.setUserId(task.getUserId());
result.setCreatedAt(LocalDateTime.now());
@@ -4032,7 +4048,7 @@ public class SimilarAsinTaskService {
*/
private void bindTaskToCurrentOwnerForLocalFallback(FileTaskEntity task, String scopeHash, Integer chunkIndex) {
if (task == null || task.getId() == null) {
return;
throw new IllegalStateException("Cannot bind local fallback without a task ID");
}
String currentInstance = currentInstanceId();
try {
@@ -4054,28 +4070,30 @@ public class SimilarAsinTaskService {
update.setId(task.getId());
update.setResultJson(updatedJson);
update.setUpdatedAt(LocalDateTime.now());
fileTaskMapper.updateById(update);
if (fileTaskMapper.updateById(update) != 1) {
throw new IllegalStateException("Task owner update affected no row");
}
log.warn("[similar-asin] task bound to current instance due to rustfs fallback taskId={} instanceId={} chunk={}",
task.getId(), currentInstance, chunkIndex);
} else if (!Objects.equals(existingOwner, currentInstance)) {
log.error("[similar-asin] rustfs fallback on non-owner instance, chunk will be unreachable taskId={} chunk={} owner={} current={}",
task.getId(), chunkIndex, existingOwner, currentInstance);
throw new IllegalStateException("Local fallback task is owned by instance=" + existingOwner
+ " current=" + currentInstance);
}
} catch (Exception ex) {
log.warn("[similar-asin] bind owner for local fallback failed taskId={} err={}", task.getId(), ex.getMessage());
throw new IllegalStateException("Failed to bind local fallback task owner taskId=" + task.getId(), ex);
}
// fallback 信号写到 task_scope_state.state_json方便排查 / 后续告警钩子
if (scopeHash == null || scopeHash.isBlank()) {
return;
}
try {
if (scopeHash == null || scopeHash.isBlank()) {
return;
}
TaskScopeStateEntity scope = taskScopeStateMapper.selectOne(new LambdaQueryWrapper<TaskScopeStateEntity>()
.eq(TaskScopeStateEntity::getTaskId, task.getId())
.eq(TaskScopeStateEntity::getModuleType, MODULE_TYPE)
.eq(TaskScopeStateEntity::getScopeHash, scopeHash)
.last("limit 1"));
if (scope == null) {
return;
throw new IllegalStateException("Local fallback scope state does not exist");
}
ObjectNode stateNode;
String stateJson = scope.getStateJson();
@@ -4106,11 +4124,13 @@ public class SimilarAsinTaskService {
fallbackArr.add(tag);
scope.setStateJson(objectMapper.writeValueAsString(stateNode));
scope.setUpdatedAt(LocalDateTime.now());
taskScopeStateMapper.updateById(scope);
if (taskScopeStateMapper.updateById(scope) != 1) {
throw new IllegalStateException("Local fallback scope update affected no row");
}
}
} catch (Exception ex) {
log.warn("[similar-asin] write local-fallback state failed taskId={} chunk={} err={}",
task.getId(), chunkIndex, ex.getMessage());
throw new IllegalStateException("Failed to record local fallback state taskId="
+ task.getId() + " chunk=" + chunkIndex, ex);
}
}
@@ -6385,17 +6405,22 @@ public class SimilarAsinTaskService {
}
private void deleteTransientTaskPayloads(Long taskId) {
deleteTransientPayloads(collectTransientTaskPayloads(taskId), taskId);
}
private List<String> collectTransientTaskPayloads(Long taskId) {
if (taskId == null || taskId <= 0) {
return;
return List.of();
}
List<String> payloads = new ArrayList<>();
List<TaskScopeStateEntity> scopes = taskScopeStateMapper.selectList(new LambdaQueryWrapper<TaskScopeStateEntity>()
.select(TaskScopeStateEntity::getParsedPayloadJson, TaskScopeStateEntity::getStateJson)
.eq(TaskScopeStateEntity::getTaskId, taskId)
.eq(TaskScopeStateEntity::getModuleType, MODULE_TYPE));
if (scopes != null) {
for (TaskScopeStateEntity scope : scopes) {
transientPayloadStorageService.deletePayloadIfPresent(scope.getParsedPayloadJson());
transientPayloadStorageService.deletePayloadIfPresent(scope.getStateJson());
payloads.add(scope.getParsedPayloadJson());
payloads.add(scope.getStateJson());
}
}
List<TaskChunkEntity> chunks = taskChunkMapper.selectList(new LambdaQueryWrapper<TaskChunkEntity>()
@@ -6404,9 +6429,29 @@ public class SimilarAsinTaskService {
.eq(TaskChunkEntity::getModuleType, MODULE_TYPE));
if (chunks != null) {
for (TaskChunkEntity chunk : chunks) {
transientPayloadStorageService.deletePayloadIfPresent(chunk.getPayloadJson());
payloads.add(chunk.getPayloadJson());
}
}
payloads.removeIf(value -> value == null || value.isBlank());
return payloads;
}
private record PreparedSubmittedChunk(Long taskId,
String scopeKey,
String scopeHash,
Integer chunkIndex,
Integer chunkTotal,
boolean done,
String error,
String payloadJson,
String storedPayload,
boolean localFallback,
SubmittedTaskMetadata taskMetadata) {
}
private record PersistSubmittedChunkResult(SubmitContext context,
boolean payloadPersisted,
FinalizeTaskResult finalizeResult) {
}
private record SubmitContext(FileTaskEntity task,
@@ -6414,7 +6459,16 @@ public class SimilarAsinTaskService {
String scopeHash,
Integer chunkIndex,
boolean forceFlush,
String error) {
String error,
SubmittedTaskMetadata taskMetadata) {
}
private record SubmittedTaskMetadata(int rowCount,
List<SimilarAsinSourceFileDto> sourceFiles) {
}
private record FinalizeTaskResult(Long taskId,
boolean terminal) {
}
private record CozeBatchContext(Long jobId,
@@ -7,9 +7,12 @@ import com.nanri.aiimage.common.exception.BusinessException;
import com.nanri.aiimage.modules.task.mapper.TaskScopeStateMapper;
import com.nanri.aiimage.modules.task.model.entity.TaskScopeStateEntity;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.dao.DuplicateKeyException;
import org.springframework.stereotype.Service;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.transaction.support.TransactionTemplate;
import java.nio.charset.StandardCharsets;
import java.nio.file.Path;
@@ -21,6 +24,7 @@ import java.util.Map;
@Service
@RequiredArgsConstructor
@Slf4j
public class TaskScopePayloadStorageService {
private static final String LOCAL_BUFFER_DIR = "task-scope-buffer";
@@ -28,8 +32,8 @@ public class TaskScopePayloadStorageService {
private final TaskScopeStateMapper taskScopeStateMapper;
private final ObjectMapper objectMapper;
private final TransientPayloadStorageService transientPayloadStorageService;
private final PlatformTransactionManager transactionManager;
@Transactional
public <T> void saveScopePayload(Long taskId, String moduleType, String scopeKey, T payload) {
if (taskId == null || taskId <= 0 || isBlank(moduleType) || isBlank(scopeKey) || payload == null) {
return;
@@ -37,11 +41,60 @@ public class TaskScopePayloadStorageService {
String normalizedScopeKey = normalize(scopeKey);
String scopeHash = hash(normalizedScopeKey);
String payloadJson = writeJson(payload, "写入任务范围载荷失败");
String stateJson = transientPayloadStorageService.storeScopePayload(
String stateJson = transientPayloadStorageService.storeScopePayloadVersioned(
moduleType, taskId, scopeHash, payloadJson, true);
LocalDateTime now = LocalDateTime.now();
String replacedPayload;
try {
replacedPayload = inTransaction(() -> persistScopePayload(
taskId, moduleType, normalizedScopeKey, scopeHash, stateJson));
} catch (RuntimeException ex) {
deleteUploadedPayloadIfUnreferenced(taskId, moduleType, scopeHash, stateJson);
throw ex;
}
transientPayloadStorageService.deleteReplacedPayloadIfNeeded(replacedPayload, stateJson);
}
TaskScopeStateEntity state = getScopeState(taskId, moduleType, scopeHash);
private void deleteUploadedPayloadIfUnreferenced(Long taskId,
String moduleType,
String scopeHash,
String uploadedPayload) {
try {
TaskScopeStateEntity currentState = getScopeState(taskId, moduleType, scopeHash);
if (!isSafeToDelete(currentState, uploadedPayload)) {
return;
}
} catch (RuntimeException ex) {
log.warn("[task-scope-payload] skip failed-save cleanup because reference check failed taskId={} moduleType={} scopeHash={} err={}",
taskId, moduleType, scopeHash, ex.getMessage());
return;
}
transientPayloadStorageService.deletePayloadIfPresent(uploadedPayload);
}
private boolean isSafeToDelete(TaskScopeStateEntity state, String payload) {
String pointer = transientPayloadStorageService.extractPointer(payload);
if (pointer == null) {
return false;
}
if (state == null || isBlank(state.getStateJson())) {
return true;
}
String currentPointer = transientPayloadStorageService.extractPointer(state.getStateJson());
if (currentPointer == null) {
log.warn("[task-scope-payload] skip failed-save cleanup because current state pointer is unreadable taskId={} moduleType={} scopeHash={}",
state.getTaskId(), state.getModuleType(), state.getScopeHash());
return false;
}
return !pointer.equals(currentPointer);
}
private String persistScopePayload(Long taskId,
String moduleType,
String normalizedScopeKey,
String scopeHash,
String stateJson) {
LocalDateTime now = LocalDateTime.now();
TaskScopeStateEntity state = getScopeStateForUpdate(taskId, moduleType, scopeHash);
if (state == null) {
TaskScopeStateEntity entity = new TaskScopeStateEntity();
entity.setTaskId(taskId);
@@ -58,21 +111,22 @@ public class TaskScopePayloadStorageService {
entity.setUpdatedAt(now);
try {
taskScopeStateMapper.insert(entity);
return;
return null;
} catch (DuplicateKeyException ignored) {
state = getScopeState(taskId, moduleType, scopeHash);
state = getScopeStateForUpdate(taskId, moduleType, scopeHash);
}
}
if (state == null) {
throw new BusinessException("写入任务范围载荷失败");
}
transientPayloadStorageService.deleteReplacedPayloadIfNeeded(state.getStateJson(), stateJson);
String replacedPayload = state.getStateJson();
taskScopeStateMapper.update(null, new LambdaUpdateWrapper<TaskScopeStateEntity>()
.eq(TaskScopeStateEntity::getId, state.getId())
.set(TaskScopeStateEntity::getScopeKey, normalizedScopeKey)
.set(TaskScopeStateEntity::getStateJson, stateJson)
.set(TaskScopeStateEntity::getLastChunkAt, now)
.set(TaskScopeStateEntity::getUpdatedAt, now));
return replacedPayload;
}
public <T> T getScopePayload(Long taskId, String moduleType, String scopeKey, Class<T> clazz) {
@@ -279,6 +333,10 @@ public class TaskScopePayloadStorageService {
return false;
}
private <T> T inTransaction(java.util.function.Supplier<T> action) {
return new TransactionTemplate(transactionManager).execute(status -> action.get());
}
public Path getBufferedPayloadRoot() {
return Path.of(System.getProperty("java.io.tmpdir"), LOCAL_BUFFER_DIR);
}
@@ -291,6 +349,14 @@ public class TaskScopePayloadStorageService {
.last("limit 1"));
}
private TaskScopeStateEntity getScopeStateForUpdate(Long taskId, String moduleType, String scopeHash) {
return taskScopeStateMapper.selectOne(new LambdaQueryWrapper<TaskScopeStateEntity>()
.eq(TaskScopeStateEntity::getTaskId, taskId)
.eq(TaskScopeStateEntity::getModuleType, moduleType)
.eq(TaskScopeStateEntity::getScopeHash, scopeHash)
.last("limit 1 FOR UPDATE"));
}
private String resolveStatePayload(TaskScopeStateEntity state) {
if (state == null || isBlank(state.getStateJson())) {
return null;
@@ -34,10 +34,12 @@ public class ZiniaoAuthService {
private static final String CACHE_TYPE_COMPANY_ID = "COMPANY_ID";
private static final String CACHE_TYPE_STAFF_LIST = "STAFF_LIST";
private static final String CACHE_TYPE_USER_STORES = "USER_STORES";
private static final String CACHE_TYPE_INVALID_USER_STORES = "INVALID_USER_STORES";
private static final String CACHE_TYPE_SHOP_MATCH = "SHOP_MATCH";
private static final Duration COMPANY_ID_CACHE_TTL = Duration.ofHours(12);
private static final Duration STAFF_LIST_CACHE_TTL = Duration.ofMinutes(30);
private static final Duration USER_STORES_CACHE_TTL = Duration.ofMinutes(30);
private static final Duration INVALID_USER_STORES_CACHE_TTL = Duration.ofMinutes(30);
private static final Duration SHOP_MATCH_CACHE_TTL = Duration.ofMinutes(30);
public StoreMatchResult matchStoreByNameAcrossStaff(String targetShopName, Long preferUserId) {
@@ -139,6 +141,9 @@ public class ZiniaoAuthService {
}
public List<ZiniaoShopCacheDto> getOrLoadUserStoresForIndex(String apiKey, Long companyId, Long userId) {
if (isInvalidUserCachedForIndex(apiKey, companyId, userId)) {
return List.of();
}
return getOrLoadUserStores(apiKey, companyId, userId);
}
@@ -182,6 +187,11 @@ public boolean isIpWhitelistError(BusinessException ex) {
if (apiKey == null || apiKey.isBlank() || companyId == null || companyId <= 0 || userId == null || userId <= 0) {
return;
}
ziniaoTransientCacheService.put(
CACHE_TYPE_INVALID_USER_STORES,
buildUserScopeCacheKey(apiKey, companyId, userId),
Boolean.TRUE,
INVALID_USER_STORES_CACHE_TTL);
String staffCacheKey = buildApiKeyHash(apiKey) + ":" + companyId;
ziniaoTransientCacheService.getList(CACHE_TYPE_STAFF_LIST, staffCacheKey, ZiniaoStaffItemVo.class)
.ifPresent(staff -> {
@@ -197,6 +207,21 @@ public boolean isIpWhitelistError(BusinessException ex) {
ziniaoTransientCacheService.delete(CACHE_TYPE_USER_STORES, buildApiKeyHash(apiKey) + ":" + companyId + ":" + userId);
}
private boolean isInvalidUserCachedForIndex(String apiKey, Long companyId, Long userId) {
if (apiKey == null || apiKey.isBlank() || companyId == null || companyId <= 0 || userId == null || userId <= 0) {
return false;
}
return ziniaoTransientCacheService.get(
CACHE_TYPE_INVALID_USER_STORES,
buildUserScopeCacheKey(apiKey, companyId, userId),
Boolean.class)
.orElse(false);
}
private String buildUserScopeCacheKey(String apiKey, Long companyId, Long userId) {
return buildApiKeyHash(apiKey) + ":" + companyId + ":" + userId;
}
private boolean isSkippableUserStoresError(BusinessException ex) {
String message = ex == null ? null : ex.getMessage();
if (message == null || message.isBlank()) {
@@ -103,6 +103,8 @@ aiimage:
connect-timeout-seconds: ${AIIMAGE_TRANSIENT_STORAGE_CONNECT_TIMEOUT_SECONDS:10}
read-timeout-seconds: ${AIIMAGE_TRANSIENT_STORAGE_READ_TIMEOUT_SECONDS:60}
write-timeout-seconds: ${AIIMAGE_TRANSIENT_STORAGE_WRITE_TIMEOUT_SECONDS:60}
call-timeout-seconds: ${AIIMAGE_TRANSIENT_STORAGE_CALL_TIMEOUT_SECONDS:90}
operation-timeout-seconds: ${AIIMAGE_TRANSIENT_STORAGE_OPERATION_TIMEOUT_SECONDS:120}
upload-max-retries: ${AIIMAGE_TRANSIENT_STORAGE_UPLOAD_MAX_RETRIES:3}
read-max-retries: ${AIIMAGE_TRANSIENT_STORAGE_READ_MAX_RETRIES:3}
delete-max-retries: ${AIIMAGE_TRANSIENT_STORAGE_DELETE_MAX_RETRIES:3}
@@ -116,8 +118,10 @@ aiimage:
failure-window-seconds: ${AIIMAGE_TRANSIENT_STORAGE_FAILURE_WINDOW_SECONDS:60}
failure-window-threshold: ${AIIMAGE_TRANSIENT_STORAGE_FAILURE_WINDOW_THRESHOLD:20}
failure-cooldown-millis: ${AIIMAGE_TRANSIENT_STORAGE_FAILURE_COOLDOWN_MILLIS:10000}
connection-pool-max-idle: ${AIIMAGE_TRANSIENT_STORAGE_CONNECTION_POOL_MAX_IDLE:0}
connection-pool-keep-alive-millis: ${AIIMAGE_TRANSIENT_STORAGE_CONNECTION_POOL_KEEP_ALIVE_MILLIS:1}
dispatcher-max-requests: ${AIIMAGE_TRANSIENT_STORAGE_DISPATCHER_MAX_REQUESTS:56}
dispatcher-max-requests-per-host: ${AIIMAGE_TRANSIENT_STORAGE_DISPATCHER_MAX_REQUESTS_PER_HOST:56}
connection-pool-max-idle: ${AIIMAGE_TRANSIENT_STORAGE_CONNECTION_POOL_MAX_IDLE:5}
connection-pool-keep-alive-millis: ${AIIMAGE_TRANSIENT_STORAGE_CONNECTION_POOL_KEEP_ALIVE_MILLIS:300000}
warn-payload-bytes: ${AIIMAGE_TRANSIENT_STORAGE_WARN_PAYLOAD_BYTES:5242880}
max-payload-bytes: ${AIIMAGE_TRANSIENT_STORAGE_MAX_PAYLOAD_BYTES:52428800}
max-stored-payload-bytes: ${AIIMAGE_TRANSIENT_STORAGE_MAX_STORED_PAYLOAD_BYTES:52428800}
@@ -126,6 +130,7 @@ aiimage:
delete-retry-cron: ${AIIMAGE_TRANSIENT_STORAGE_DELETE_RETRY_CRON:0 */5 * * * *}
delete-retry-queue-capacity: ${AIIMAGE_TRANSIENT_STORAGE_DELETE_RETRY_QUEUE_CAPACITY:10000}
delete-retry-batch-size: ${AIIMAGE_TRANSIENT_STORAGE_DELETE_RETRY_BATCH_SIZE:200}
delete-retry-batch-timeout-seconds: ${AIIMAGE_TRANSIENT_STORAGE_DELETE_RETRY_BATCH_TIMEOUT_SECONDS:60}
storage:
local-temp-dir: ${AIIMAGE_STORAGE_LOCAL_TEMP_DIR:./data/tmp}
cleanup-enabled: ${AIIMAGE_STORAGE_CLEANUP_ENABLED:true}
@@ -0,0 +1,36 @@
SET @invalid_asin_brand_index_exists := (
SELECT COUNT(*)
FROM information_schema.STATISTICS
WHERE TABLE_SCHEMA = DATABASE()
AND TABLE_NAME = 'biz_invalid_asin_data'
AND INDEX_NAME = 'idx_brand'
);
SET @sql_drop_invalid_asin_brand_index := IF(
@invalid_asin_brand_index_exists > 0,
'ALTER TABLE biz_invalid_asin_data DROP INDEX idx_brand',
'SELECT 1'
);
PREPARE stmt_drop_invalid_asin_brand_index FROM @sql_drop_invalid_asin_brand_index;
EXECUTE stmt_drop_invalid_asin_brand_index;
DEALLOCATE PREPARE stmt_drop_invalid_asin_brand_index;
ALTER TABLE biz_invalid_asin_data
MODIFY COLUMN brand TEXT NULL COMMENT 'brand name';
SET @invalid_asin_brand_prefix_index_exists := (
SELECT COUNT(*)
FROM information_schema.STATISTICS
WHERE TABLE_SCHEMA = DATABASE()
AND TABLE_NAME = 'biz_invalid_asin_data'
AND INDEX_NAME = 'idx_brand_prefix'
);
SET @sql_add_invalid_asin_brand_prefix_index := IF(
@invalid_asin_brand_prefix_index_exists = 0,
'ALTER TABLE biz_invalid_asin_data ADD INDEX idx_brand_prefix (brand(191))',
'SELECT 1'
);
PREPARE stmt_add_invalid_asin_brand_prefix_index FROM @sql_add_invalid_asin_brand_prefix_index;
EXECUTE stmt_add_invalid_asin_brand_prefix_index;
DEALLOCATE PREPARE stmt_add_invalid_asin_brand_prefix_index;
@@ -3,38 +3,148 @@ package com.nanri.aiimage.modules.file.service.object;
import com.nanri.aiimage.config.TransientStorageProperties;
import org.junit.jupiter.api.Test;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
class RustfsDeleteRetryServiceTest {
@Test
void retryPendingDeletesRefillsDeferredItemsBeforeQueueDrainsCompletely() {
void pendingUniqueKeysNeverExceedCapacityUnderConcurrentAdmission() throws Exception {
TransientStorageProperties properties = new TransientStorageProperties();
properties.setDeleteRetryEnabled(true);
properties.setDeleteRetryQueueCapacity(2);
properties.setDeleteRetryBatchSize(1);
properties.setDeleteRetryQueueCapacity(3);
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
RustfsDeleteRetryService service = new RustfsDeleteRetryService(properties, rustfs);
ExecutorService executor = Executors.newFixedThreadPool(8);
CountDownLatch start = new CountDownLatch(1);
List<Future<?>> futures = new ArrayList<>();
try {
for (int i = 0; i < 20; i++) {
String objectKey = "object-" + i;
futures.add(executor.submit(() -> {
start.await();
service.enqueue(objectKey, new IllegalStateException("failed"));
return null;
}));
}
start.countDown();
for (Future<?> future : futures) {
future.get(5, TimeUnit.SECONDS);
}
} finally {
executor.shutdownNow();
}
assertEquals(3, service.pendingCount());
}
@Test
void failedKeyIsAttemptedOnlyOncePerScheduledInvocation() {
TransientStorageProperties properties = new TransientStorageProperties();
properties.setDeleteRetryEnabled(true);
properties.setDeleteRetryQueueCapacity(10);
properties.setDeleteRetryBatchSize(10);
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
doAnswer(invocation -> {
String objectKey = invocation.getArgument(0);
if ("b".equals(objectKey)) {
throw new IllegalStateException("still failing");
}
return null;
}).when(rustfs).deleteObjectFromRetry(org.mockito.ArgumentMatchers.anyString());
throw new IllegalStateException("still failing");
}).when(rustfs).deleteObjectFromRetry(eq("a"), anyLong());
RustfsDeleteRetryService service = new RustfsDeleteRetryService(properties, rustfs);
service.enqueue("a", new IllegalStateException("failed-a"));
service.enqueue("b", new IllegalStateException("failed-b"));
service.enqueue("c", new IllegalStateException("failed-c"));
service.retryPendingDeletes();
service.retryPendingDeletes();
verify(rustfs).deleteObjectFromRetry(eq("a"), anyLong());
assertEquals(1, service.pendingCount());
service.retryPendingDeletes();
verify(rustfs).deleteObjectFromRetry("c");
verify(rustfs, times(2)).deleteObjectFromRetry(eq("a"), anyLong());
assertEquals(1, service.pendingCount());
}
@Test
void successfulRetryDoesNotRemoveConcurrentReenqueue() throws Exception {
TransientStorageProperties properties = new TransientStorageProperties();
properties.setDeleteRetryEnabled(true);
properties.setDeleteRetryQueueCapacity(10);
properties.setDeleteRetryBatchSize(10);
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
CountDownLatch deleteStarted = new CountDownLatch(1);
CountDownLatch allowDeleteSuccess = new CountDownLatch(1);
doAnswer(invocation -> {
deleteStarted.countDown();
assertTrue(allowDeleteSuccess.await(5, TimeUnit.SECONDS));
return null;
}).when(rustfs).deleteObjectFromRetry(eq("a"), anyLong());
RustfsDeleteRetryService service = new RustfsDeleteRetryService(properties, rustfs);
service.enqueue("a", new IllegalStateException("first failure"));
ExecutorService executor = Executors.newSingleThreadExecutor();
try {
Future<?> retry = executor.submit(service::retryPendingDeletes);
assertTrue(deleteStarted.await(5, TimeUnit.SECONDS));
service.enqueue("a", new IllegalStateException("new failure while delete completes"));
allowDeleteSuccess.countDown();
retry.get(5, TimeUnit.SECONDS);
} finally {
allowDeleteSuccess.countDown();
executor.shutdownNow();
}
assertEquals(1, service.pendingCount());
service.retryPendingDeletes();
verify(rustfs, times(2)).deleteObjectFromRetry(eq("a"), anyLong());
assertEquals(0, service.pendingCount());
}
@Test
void batchDeadlineStopsBeforeNextItemAndPassesRemainingBudget() {
TransientStorageProperties properties = new TransientStorageProperties();
properties.setDeleteRetryEnabled(true);
properties.setDeleteRetryQueueCapacity(10);
properties.setDeleteRetryBatchSize(10);
properties.setDeleteRetryBatchTimeoutSeconds(1);
RustfsObjectStorageService rustfs = mock(RustfsObjectStorageService.class);
AtomicLong receivedBudgetNanos = new AtomicLong();
doAnswer(invocation -> {
long budgetNanos = invocation.getArgument(1);
receivedBudgetNanos.set(budgetNanos);
long deadlineNanos = System.nanoTime() + budgetNanos;
while (System.nanoTime() < deadlineNanos) {
Thread.onSpinWait();
}
throw new IllegalStateException("timed out");
}).when(rustfs).deleteObjectFromRetry(eq("a"), anyLong());
RustfsDeleteRetryService service = new RustfsDeleteRetryService(properties, rustfs);
service.enqueue("a", new IllegalStateException("failed-a"));
service.enqueue("b", new IllegalStateException("failed-b"));
long startedAt = System.nanoTime();
service.retryPendingDeletes();
long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedAt);
assertTrue(receivedBudgetNanos.get() > 0L);
assertTrue(receivedBudgetNanos.get() <= TimeUnit.SECONDS.toNanos(1));
assertTrue(elapsedMillis < 2000L, "batch should stop close to its configured deadline");
verify(rustfs, never()).deleteObjectFromRetry(eq("b"), anyLong());
assertEquals(2, service.pendingCount());
}
}
@@ -3,19 +3,34 @@ package com.nanri.aiimage.modules.file.service.object;
import com.nanri.aiimage.config.TransientStorageProperties;
import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
import io.minio.MinioClient;
import io.minio.PutObjectArgs;
import io.minio.RemoveObjectArgs;
import io.minio.StatObjectArgs;
import okhttp3.OkHttpClient;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.ObjectProvider;
import java.lang.reflect.Field;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
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.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
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.same;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -24,6 +39,9 @@ class RustfsObjectStorageServiceTest {
@Test
void httpClientUsesConfiguredConnectionPoolAndKeepsRetryEnabled() {
TransientStorageProperties properties = configuredProperties();
properties.setCallTimeoutSeconds(45);
properties.setDispatcherMaxRequests(12);
properties.setDispatcherMaxRequestsPerHost(9);
properties.setConnectionPoolMaxIdle(3);
properties.setConnectionPoolKeepAliveMillis(1500);
RustfsObjectStorageService service = new RustfsObjectStorageService(
@@ -33,6 +51,26 @@ class RustfsObjectStorageServiceTest {
assertTrue(client.retryOnConnectionFailure(), "RustFS OkHttp 应开启连接失败重试");
assertNotNull(client.connectionPool(), "RustFS OkHttp 应使用显式 connectionPool 配置");
assertEquals(TimeUnit.SECONDS.toMillis(45), client.callTimeoutMillis());
assertEquals(12, client.dispatcher().getMaxRequests());
assertEquals(9, client.dispatcher().getMaxRequestsPerHost());
}
@Test
void deadlineClientViewShortensCallTimeoutAndSharesResources() {
TransientStorageProperties properties = configuredProperties();
properties.setCallTimeoutSeconds(45);
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), emptyProvider(), null);
OkHttpClient shared = service.getHttpClient();
OkHttpClient bounded = service.getHttpClient(
System.nanoTime() + TimeUnit.SECONDS.toNanos(2));
assertTrue(bounded.callTimeoutMillis() > 0);
assertTrue(bounded.callTimeoutMillis() <= 2000);
assertTrue(shared.connectionPool() == bounded.connectionPool());
assertTrue(shared.dispatcher() == bounded.dispatcher());
}
@Test
@@ -106,6 +144,241 @@ class RustfsObjectStorageServiceTest {
assertTrue(rejected.getMessage().contains("cooldown active"));
}
@Test
void retryBackoffDoesNotHoldUploadPermit() throws Exception {
TransientStorageProperties properties = configuredProperties();
properties.setMaxConcurrentUploads(1);
properties.setAcquirePermitTimeoutMillis(300);
properties.setUploadMaxRetries(2);
properties.setBaseRetryDelayMillis(1000);
properties.setMaxRetryDelayMillis(1000);
MinioClient client = mock(MinioClient.class);
CountDownLatch firstAttemptFailed = new CountDownLatch(1);
AtomicInteger calls = new AtomicInteger();
doAnswer(invocation -> {
if (calls.incrementAndGet() == 1) {
firstAttemptFailed.countDown();
throw new IllegalStateException("first attempt failed");
}
return null;
}).when(client).putObject(any(PutObjectArgs.class));
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), emptyProvider(), () -> client);
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
Future<String> first = executor.submit(() -> service.uploadText("task/first.json", "{}", false));
assertTrue(firstAttemptFailed.await(2, TimeUnit.SECONDS));
Future<String> second = executor.submit(() -> service.uploadText("task/second.json", "{}", false));
assertEquals("task/second.json", second.get(800, TimeUnit.MILLISECONDS));
assertEquals("task/first.json", first.get(2, TimeUnit.SECONDS));
verify(client, times(3)).putObject(any(PutObjectArgs.class));
assertEquals(1, privateSemaphore(service, "uploadSemaphore").availablePermits());
} finally {
executor.shutdownNow();
}
}
@Test
void visibilityCheckDoesNotHoldUploadPermit() throws Exception {
TransientStorageProperties properties = configuredProperties();
properties.setMaxConcurrentUploads(1);
properties.setAcquirePermitTimeoutMillis(300);
MinioClient client = mock(MinioClient.class);
CountDownLatch statStarted = new CountDownLatch(1);
CountDownLatch releaseStat = new CountDownLatch(1);
doAnswer(invocation -> {
statStarted.countDown();
assertTrue(releaseStat.await(2, TimeUnit.SECONDS));
return null;
}).when(client).statObject(any(StatObjectArgs.class));
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), emptyProvider(), () -> client);
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
Future<String> first = executor.submit(() -> service.uploadText("task/first.json", "{}", true));
assertTrue(statStarted.await(2, TimeUnit.SECONDS));
Future<String> second = executor.submit(() -> service.uploadText("task/second.json", "{}", false));
assertEquals("task/second.json", second.get(800, TimeUnit.MILLISECONDS));
releaseStat.countDown();
assertEquals("task/first.json", first.get(2, TimeUnit.SECONDS));
assertEquals(1, privateSemaphore(service, "uploadSemaphore").availablePermits());
} finally {
releaseStat.countDown();
executor.shutdownNow();
}
}
@Test
void openedCircuitStopsRemainingOuterAttempts() {
TransientStorageProperties properties = configuredProperties();
properties.setUploadMaxRetries(3);
properties.setFailureWindowThreshold(1);
properties.setFailureCooldownMillis(10000);
AtomicInteger clientBuilds = new AtomicInteger();
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), emptyProvider(), () -> {
clientBuilds.incrementAndGet();
throw new IllegalStateException("rustfs down");
});
IllegalStateException ex = assertThrows(IllegalStateException.class,
() -> service.uploadText("task/a.json", "{}", false));
assertTrue(ex.getMessage().contains("cooldown active"));
assertEquals(1, clientBuilds.get());
}
@Test
void successfulDeleteDoesNotResetReadWriteFailureWindow() throws Exception {
TransientStorageProperties properties = configuredProperties();
properties.setUploadMaxRetries(1);
properties.setFailureWindowThreshold(2);
properties.setFailureCooldownMillis(10000);
MinioClient client = mock(MinioClient.class);
doThrow(new IllegalStateException("upload failed"))
.when(client).putObject(any(PutObjectArgs.class));
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), emptyProvider(), () -> client);
assertThrows(IllegalStateException.class,
() -> service.uploadText("task/first.json", "{}", false));
service.deleteObject("task/old.json");
assertThrows(IllegalStateException.class,
() -> service.uploadText("task/second.json", "{}", false));
IllegalStateException rejected = assertThrows(IllegalStateException.class,
() -> service.uploadText("task/third.json", "{}", false));
assertTrue(rejected.getMessage().contains("cooldown active"));
verify(client, times(2)).putObject(any(PutObjectArgs.class));
verify(client).removeObject(any(RemoveObjectArgs.class));
}
@Test
void failedDeleteDoesNotIncreaseReadWriteFailureWindow() throws Exception {
TransientStorageProperties properties = configuredProperties();
properties.setUploadMaxRetries(1);
properties.setDeleteMaxRetries(1);
properties.setFailureWindowThreshold(2);
properties.setFailureCooldownMillis(10000);
MinioClient client = mock(MinioClient.class);
doThrow(new IllegalStateException("upload failed"))
.when(client).putObject(any(PutObjectArgs.class));
doThrow(new IllegalStateException("delete failed"))
.when(client).removeObject(any(RemoveObjectArgs.class));
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), emptyProvider(), () -> client);
assertThrows(IllegalStateException.class,
() -> service.uploadText("task/first.json", "{}", false));
assertThrows(IllegalStateException.class, () -> service.deleteObject("task/old.json"));
IllegalStateException opensCircuit = assertThrows(IllegalStateException.class,
() -> service.uploadText("task/second.json", "{}", false));
assertTrue(opensCircuit.getMessage().contains("failed to upload"));
IllegalStateException rejected = assertThrows(IllegalStateException.class,
() -> service.uploadText("task/third.json", "{}", false));
assertTrue(rejected.getMessage().contains("cooldown active"));
verify(client, times(2)).putObject(any(PutObjectArgs.class));
}
@Test
void repeatedStatFailuresOpenCircuitAndEnqueueOneCompensationPerObject() throws Exception {
TransientStorageProperties properties = configuredProperties();
properties.setFailureWindowThreshold(5);
properties.setFailureCooldownMillis(10000);
RustfsDeleteRetryService retryService = mock(RustfsDeleteRetryService.class);
MinioClient client = mock(MinioClient.class);
doThrow(new IllegalStateException("stat failed"))
.when(client).statObject(any(StatObjectArgs.class));
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), provider(retryService), () -> client);
IllegalStateException firstFailure = assertThrows(IllegalStateException.class,
() -> service.uploadText("task/a.json", "{}", true));
IllegalStateException secondFailure = assertThrows(IllegalStateException.class,
() -> service.uploadText("task/b.json", "{}", true));
assertTrue(firstFailure.getMessage().contains("not visible"));
assertTrue(secondFailure.getMessage().contains("not visible"));
verify(client, times(5)).statObject(any(StatObjectArgs.class));
verify(retryService).enqueue(eq("task/a.json"), same(firstFailure));
verify(retryService).enqueue(eq("task/b.json"), same(secondFailure));
IllegalStateException rejected = assertThrows(IllegalStateException.class,
() -> service.uploadText("task/c.json", "{}", false));
assertTrue(rejected.getMessage().contains("cooldown active"));
verify(client, times(2)).putObject(any(PutObjectArgs.class));
}
@Test
void retryDeleteBudgetCutsOffLongBackoff() throws Exception {
TransientStorageProperties properties = configuredProperties();
properties.setDeleteMaxRetries(3);
properties.setBaseRetryDelayMillis(5000);
properties.setMaxRetryDelayMillis(5000);
MinioClient client = mock(MinioClient.class);
doThrow(new IllegalStateException("delete failed"))
.when(client).removeObject(any(RemoveObjectArgs.class));
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), emptyProvider(), () -> client);
long startedAt = System.nanoTime();
IllegalStateException failure = assertThrows(IllegalStateException.class,
() -> service.deleteObjectFromRetry("task/a.json", TimeUnit.MILLISECONDS.toNanos(200)));
long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedAt);
assertTrue(failure.getMessage().contains("operation timeout"));
assertTrue(elapsedMillis < 1000, "operation budget should cut off the five-second backoff");
verify(client).removeObject(any(RemoveObjectArgs.class));
}
@Test
void configuredOperationBudgetCutsOffLongBackoff() throws Exception {
TransientStorageProperties properties = configuredProperties();
properties.setOperationTimeoutSeconds(1);
properties.setUploadMaxRetries(3);
properties.setBaseRetryDelayMillis(5000);
properties.setMaxRetryDelayMillis(5000);
MinioClient client = mock(MinioClient.class);
doThrow(new IllegalStateException("upload failed"))
.when(client).putObject(any(PutObjectArgs.class));
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), emptyProvider(), () -> client);
long startedAt = System.nanoTime();
IllegalStateException failure = assertThrows(IllegalStateException.class,
() -> service.uploadText("task/a.json", "{}", false));
long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedAt);
assertTrue(failure.getMessage().contains("operation timeout"));
assertTrue(elapsedMillis < 2000, "configured budget should cut off the five-second backoff");
verify(client).putObject(any(PutObjectArgs.class));
}
@Test
void completedUploadIsEnqueuedWhenDeadlineExpiresAfterPut() throws Exception {
TransientStorageProperties properties = configuredProperties();
properties.setOperationTimeoutSeconds(1);
RustfsDeleteRetryService retryService = mock(RustfsDeleteRetryService.class);
MinioClient client = mock(MinioClient.class);
doAnswer(invocation -> {
Thread.sleep(1100);
return null;
}).when(client).putObject(any(PutObjectArgs.class));
RustfsObjectStorageService service = new RustfsObjectStorageService(
properties, emptyProvider(), provider(retryService), () -> client);
IllegalStateException failure = assertThrows(IllegalStateException.class,
() -> service.uploadText("task/a.json", "{}", false));
assertTrue(failure.getMessage().contains("operation timeout"));
verify(client).putObject(any(PutObjectArgs.class));
verify(retryService).enqueue(eq("task/a.json"), same(failure));
}
private static TransientStorageProperties configuredProperties() {
TransientStorageProperties properties = new TransientStorageProperties();
properties.setEndpoint("http://127.0.0.1:9000");
@@ -1,8 +1,5 @@
package com.nanri.aiimage.modules.permission.service;
import com.baomidou.mybatisplus.core.MybatisConfiguration;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import com.nanri.aiimage.common.exception.BusinessException;
import com.nanri.aiimage.modules.permission.mapper.AdminUserMapper;
import com.nanri.aiimage.modules.permission.mapper.PermissionMenuMapper;
@@ -17,11 +14,9 @@ import com.nanri.aiimage.modules.permission.model.vo.ImageVideoDataPermissionUse
import com.nanri.aiimage.modules.permission.model.vo.PermissionMenuItemVo;
import com.nanri.aiimage.modules.permission.model.vo.UserColumnIdsVo;
import org.junit.jupiter.api.Test;
import org.apache.ibatis.builder.MapperBuilderAssistant;
import org.mockito.ArgumentCaptor;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -43,7 +38,7 @@ class PermissionMenuServiceTest {
when(userMapper.selectById(7L)).thenReturn(new AdminUserEntity());
when(menuMapper.selectCount(any())).thenReturn(1L);
when(menuMapper.selectOne(any())).thenReturn(imageVideoPermission());
when(permissionMapper.selectCount(any())).thenReturn(1L);
when(permissionMapper.countByUserIdAndColumnId(any(), any())).thenReturn(1L);
UserColumnPermissionUpdateRequest request = new UserColumnPermissionUpdateRequest();
request.setColumnIds(List.of(2L));
@@ -65,7 +60,7 @@ class PermissionMenuServiceTest {
when(userMapper.selectById(8L)).thenReturn(new AdminUserEntity());
when(menuMapper.selectCount(any())).thenReturn(1L);
when(menuMapper.selectOne(any())).thenReturn(imageVideoPermission());
when(permissionMapper.selectCount(any())).thenReturn(0L);
when(permissionMapper.countByUserIdAndColumnId(any(), any())).thenReturn(0L);
UserColumnPermissionUpdateRequest request = new UserColumnPermissionUpdateRequest();
request.setColumnIds(List.of(75L));
@@ -83,7 +78,7 @@ class PermissionMenuServiceTest {
when(userMapper.selectById(7L)).thenReturn(new AdminUserEntity());
when(menuMapper.selectCount(any())).thenReturn(1L);
when(menuMapper.selectOne(any())).thenReturn(imageVideoDataPermission(), shopDataCrawlDataPermission());
when(permissionMapper.selectCount(any())).thenReturn(1L);
when(permissionMapper.countByUserIdAndColumnId(any(), any())).thenReturn(1L);
UserColumnPermissionUpdateRequest request = new UserColumnPermissionUpdateRequest();
request.setColumnIds(List.of(2L));
@@ -104,7 +99,7 @@ class PermissionMenuServiceTest {
PermissionMenuService service = new PermissionMenuService(menuMapper, permissionMapper, userMapper);
when(userMapper.selectById(9L)).thenReturn(user(9L, "normal", 0));
when(permissionMapper.selectList(any())).thenReturn(List.of(grant(9L, 1L)));
when(permissionMapper.selectByUserId(any())).thenReturn(List.of(grant(9L, 1L)));
when(menuMapper.selectList(any())).thenReturn(List.of(
menu(1L, null, "app", 10),
menu(2L, 1L, "app", 11),
@@ -126,7 +121,7 @@ class PermissionMenuServiceTest {
PermissionMenuService service = new PermissionMenuService(menuMapper, permissionMapper, userMapper);
when(userMapper.selectById(9L)).thenReturn(user(9L, "normal", 0));
when(permissionMapper.selectList(any())).thenReturn(List.of(grant(9L, 1L)));
when(permissionMapper.selectByUserId(any())).thenReturn(List.of(grant(9L, 1L)));
when(menuMapper.selectList(any())).thenReturn(List.of(
menu(1L, null, "app", 10),
menu(2L, 1L, "app", 11)));
@@ -144,7 +139,7 @@ class PermissionMenuServiceTest {
PermissionMenuService service = new PermissionMenuService(menuMapper, permissionMapper, userMapper);
when(userMapper.selectById(9L)).thenReturn(user(9L, "normal", 0));
when(permissionMapper.selectList(any())).thenReturn(List.of(grant(9L, 2L)));
when(permissionMapper.selectByUserId(any())).thenReturn(List.of(grant(9L, 2L)));
PermissionMenuEntity root = menu(1L, null, "app", 1);
root.setColumnKey("brand_front_tools");
PermissionMenuEntity leaf = menu(2L, 1L, "app", 2);
@@ -184,7 +179,7 @@ class PermissionMenuServiceTest {
PermissionMenuService service = new PermissionMenuService(menuMapper, permissionMapper, userMapper);
when(userMapper.selectById(1L)).thenReturn(user(1L, "super_admin", 1));
when(permissionMapper.selectList(any())).thenReturn(List.of(grant(1L, 2L)));
when(permissionMapper.selectByUserId(any())).thenReturn(List.of(grant(1L, 2L)));
when(menuMapper.selectList(any())).thenReturn(List.of(
menu(1L, null, "app", 1),
menu(2L, 1L, "app", 2)));
@@ -209,7 +204,7 @@ class PermissionMenuServiceTest {
List<PermissionMenuItemVo> effective = service.getUserColumnPermissions(1L, "app");
assertThat(effective).extracting(PermissionMenuItemVo::getId).containsExactly(1L, 2L);
verify(permissionMapper, never()).selectList(any());
verify(permissionMapper, never()).selectByUserId(any());
}
@Test
@@ -221,7 +216,7 @@ class PermissionMenuServiceTest {
AdminUserEntity currentUser = user(9L, "normal", 0);
when(userMapper.selectById(9L)).thenReturn(currentUser);
when(permissionMapper.selectList(any())).thenReturn(List.of(grant(9L, 1L)));
when(permissionMapper.selectByUserId(any())).thenReturn(List.of(grant(9L, 1L)));
when(menuMapper.selectList(any())).thenReturn(List.of(
menu(1L, null, "app", 10),
menu(2L, 1L, "app", 11)));
@@ -245,7 +240,7 @@ class PermissionMenuServiceTest {
assertThatThrownBy(() -> service.getUserColumnPermissions(currentUser, 10L, "app"))
.isInstanceOf(BusinessException.class)
.hasMessageContaining("管理员权限");
verify(permissionMapper, never()).selectList(any());
verify(permissionMapper, never()).selectByUserId(any());
}
@Test
@@ -327,7 +322,7 @@ class PermissionMenuServiceTest {
.isInstanceOf(BusinessException.class)
.hasMessageContaining("子菜单");
verify(menuMapper, never()).deleteById(any(Long.class));
verify(permissionMapper, never()).delete(any());
verify(permissionMapper, never()).deleteByColumnId(any());
}
@Test
@@ -346,7 +341,7 @@ class PermissionMenuServiceTest {
when(menuMapper.selectCount(any())).thenReturn(1L);
when(menuMapper.selectOne(any())).thenReturn(null);
when(menuMapper.selectList(any())).thenReturn(List.of(menu(1L, null, "app", 1)));
when(permissionMapper.selectList(any())).thenReturn(List.of(grant(10L, 1L)));
when(permissionMapper.selectByUserId(any())).thenReturn(List.of(grant(10L, 1L)));
UserColumnPermissionUpdateRequest request = new UserColumnPermissionUpdateRequest();
request.setColumnIds(List.of(2L));
@@ -354,7 +349,7 @@ class PermissionMenuServiceTest {
assertThatThrownBy(() -> service.updateUserColumnPermissions(operator, 20L, request))
.isInstanceOf(BusinessException.class)
.hasMessageContaining("自己已有");
verify(permissionMapper, never()).delete(any());
verify(permissionMapper, never()).deleteByUserId(any());
}
@Test
@@ -374,7 +369,7 @@ class PermissionMenuServiceTest {
when(menuMapper.selectList(any())).thenReturn(List.of(
menu(1L, null, "app", 1),
menu(2L, 1L, "app", 2)));
when(permissionMapper.selectList(any())).thenReturn(List.of(grant(10L, 1L)));
when(permissionMapper.selectByUserId(any())).thenReturn(List.of(grant(10L, 1L)));
UserColumnPermissionUpdateRequest request = new UserColumnPermissionUpdateRequest();
request.setColumnIds(List.of(2L));
@@ -401,7 +396,7 @@ class PermissionMenuServiceTest {
when(menuMapper.selectList(any())).thenReturn(List.of(
menu(1L, null, "app", 1),
menu(2L, null, "app", 2)));
when(permissionMapper.selectList(any()))
when(permissionMapper.selectByUserId(any()))
.thenReturn(List.of(grant(10L, 1L)), List.of(grant(20L, 2L)));
UserColumnPermissionUpdateRequest request = new UserColumnPermissionUpdateRequest();
request.setColumnIds(List.of(1L));
@@ -454,15 +449,7 @@ class PermissionMenuServiceTest {
service.updateUserColumnPermissions(operator, 9L, request, PermissionMenuService.MENU_TYPE_APP);
TableInfoHelper.initTableInfo(
new MapperBuilderAssistant(new MybatisConfiguration(), ""),
UserColumnPermissionEntity.class);
ArgumentCaptor<LambdaUpdateWrapper<UserColumnPermissionEntity>> deleted =
ArgumentCaptor.forClass(LambdaUpdateWrapper.class);
verify(permissionMapper).delete(deleted.capture());
assertThat(deleted.getValue().getSqlSegment()).contains("column_id", "IN");
assertThat(deleted.getValue().getParamNameValuePairs().values())
.contains(9L, 11L, 12L);
verify(permissionMapper).deleteByUserIdAndColumnIds(9L, List.of(11L, 12L));
ArgumentCaptor<UserColumnPermissionEntity> inserted =
ArgumentCaptor.forClass(UserColumnPermissionEntity.class);
verify(permissionMapper).insert(inserted.capture());
@@ -484,7 +471,7 @@ class PermissionMenuServiceTest {
.isInstanceOf(BusinessException.class)
.hasMessageContaining("超级管理员");
verify(menuMapper, never()).selectOne(any());
verify(permissionMapper, never()).delete(any());
verify(permissionMapper, never()).deleteByColumnId(any());
}
@Test
@@ -513,7 +500,7 @@ class PermissionMenuServiceTest {
normal.setUsername("normal");
when(menuMapper.selectOne(any())).thenReturn(imageVideoDataPermission());
when(permissionMapper.selectList(any())).thenReturn(List.of(grant(1L, 75L)));
when(permissionMapper.selectByColumnId(any())).thenReturn(List.of(grant(1L, 75L)));
when(userMapper.selectList(any())).thenReturn(List.of(operator, explicitAdmin, normal));
List<ImageVideoDataPermissionUserVo> result = service.listImageVideoDataPermissionUsers(operator);
@@ -539,7 +526,7 @@ class PermissionMenuServiceTest {
int grantedCount = service.updateImageVideoDataPermissionUsers(operator, List.of(2L));
assertThat(grantedCount).isEqualTo(1);
verify(permissionMapper).deleteByMap(Map.of("column_id", 75L));
verify(permissionMapper).deleteByColumnId(75L);
ArgumentCaptor<UserColumnPermissionEntity> inserted =
ArgumentCaptor.forClass(UserColumnPermissionEntity.class);
@@ -563,7 +550,7 @@ class PermissionMenuServiceTest {
.isInstanceOf(BusinessException.class)
.hasMessageContaining("超级管理员");
verify(menuMapper, never()).selectOne(any());
verify(permissionMapper, never()).delete(any());
verify(permissionMapper, never()).deleteByColumnId(any());
}
@Test
@@ -579,7 +566,7 @@ class PermissionMenuServiceTest {
normal.setUsername("normal");
when(menuMapper.selectOne(any())).thenReturn(shopDataCrawlDataPermission());
when(permissionMapper.selectList(any())).thenReturn(List.of(grant(1L, 76L)));
when(permissionMapper.selectByColumnId(any())).thenReturn(List.of(grant(1L, 76L)));
when(userMapper.selectList(any())).thenReturn(List.of(operator, admin, normal));
List<ImageVideoDataPermissionUserVo> users = service.listShopDataCrawlDataPermissionUsers(operator);
@@ -589,7 +576,7 @@ class PermissionMenuServiceTest {
assertThat(users.get(0).isGranted()).isTrue();
assertThat(users.get(1).isGranted()).isFalse();
assertThat(grantedCount).isEqualTo(1);
verify(permissionMapper).deleteByMap(Map.of("column_id", 76L));
verify(permissionMapper).deleteByColumnId(76L);
ArgumentCaptor<UserColumnPermissionEntity> inserted =
ArgumentCaptor.forClass(UserColumnPermissionEntity.class);
verify(permissionMapper).insert(inserted.capture());
@@ -41,6 +41,7 @@ import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.Spy;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.dao.DuplicateKeyException;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.TransactionCallback;
import org.springframework.transaction.support.TransactionTemplate;
@@ -53,9 +54,11 @@ import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -107,6 +110,8 @@ class PublishTaskServiceTest {
private final List<TaskChunkEntity> storedChunks = new ArrayList<>();
private final List<TaskScopeStateEntity> storedScopes = new ArrayList<>();
private final Map<String, String> rustfsPayloads = new LinkedHashMap<>();
private final AtomicBoolean transactionActive = new AtomicBoolean();
private final List<Boolean> storageTransactionStates = new ArrayList<>();
private int nextPayloadId;
@BeforeEach
@@ -115,11 +120,21 @@ class PublishTaskServiceTest {
lenient().when(instanceMetadata.getInstanceId()).thenReturn("instance-a");
lenient().when(transactionTemplate.execute(any())).thenAnswer(invocation -> {
TransactionCallback<?> callback = invocation.getArgument(0);
return callback.doInTransaction(null);
transactionActive.set(true);
try {
return callback.doInTransaction(null);
} finally {
transactionActive.set(false);
}
});
lenient().doAnswer(invocation -> {
Consumer<TransactionStatus> callback = invocation.getArgument(0);
callback.accept(null);
transactionActive.set(true);
try {
callback.accept(null);
} finally {
transactionActive.set(false);
}
return null;
}).when(transactionTemplate).executeWithoutResult(any());
}
@@ -399,6 +414,90 @@ class PublishTaskServiceTest {
verify(taskFileJobService).enqueueAssembleResult(
taskId, PublishTaskService.MODULE_TYPE, resultId,
"task:" + taskId + ":owner:instance-a");
assertFalse(storageTransactionStates.isEmpty());
assertTrue(storageTransactionStates.stream().noneMatch(Boolean::booleanValue));
}
@Test
void resultCallbackDeletesUploadedPayloadWhenTransactionFails() {
long taskId = 123L;
long fileId = 223L;
FileTaskEntity task = task(taskId, 7L, "RUNNING");
PublishFileEntity file = file(taskId, fileId, "RUNNING", "transaction-failure.xlsx");
TaskDistributedLockService.LockHandle lock = mock(TaskDistributedLockService.LockHandle.class);
when(taskDistributedLockService.acquire(PublishTaskService.MODULE_TYPE, taskId)).thenReturn(lock);
when(fileTaskMapper.selectById(taskId)).thenReturn(task);
when(publishFileMapper.selectById(fileId)).thenReturn(file);
when(publishFileMapper.selectList(any())).thenReturn(List.of(file));
when(fileResultMapper.selectOne(any())).thenReturn(result(taskId, 323L));
org.mockito.Mockito.reset(transactionTemplate);
doAnswer(invocation -> {
Consumer<TransactionStatus> callback = invocation.getArgument(0);
transactionActive.set(true);
try {
callback.accept(null);
} finally {
transactionActive.set(false);
}
storedChunks.clear();
storedScopes.clear();
throw new IllegalStateException("commit failed");
}).when(transactionTemplate).executeWithoutResult(any());
assertThrows(IllegalStateException.class, () -> service.submitResult(
taskId, chunkResultRequest(7L, fileId, 1, 2, List.of(row("1")))));
assertTrue(rustfsPayloads.isEmpty());
assertTrue(storedChunks.isEmpty());
assertEquals(List.of(false, false), storageTransactionStates);
org.mockito.InOrder order = org.mockito.Mockito.inOrder(
taskDistributedLockService, transientPayloadStorageService, transactionTemplate, lock);
order.verify(taskDistributedLockService).acquire(PublishTaskService.MODULE_TYPE, taskId);
order.verify(transientPayloadStorageService).storeChunkPayloadVersioned(
any(), any(), any(), any(), any());
order.verify(transactionTemplate).executeWithoutResult(any());
order.verify(transientPayloadStorageService).deletePayloadIfPresent(any());
order.verify(lock).close();
}
@Test
void resultCallbackDeletesLoserPayloadAfterDuplicateChunkCommit() {
long taskId = 124L;
long fileId = 224L;
FileTaskEntity task = task(taskId, 7L, "RUNNING");
PublishFileEntity file = file(taskId, fileId, "RUNNING", "duplicate-chunk.xlsx");
TaskDistributedLockService.LockHandle lock = mock(TaskDistributedLockService.LockHandle.class);
when(taskDistributedLockService.acquire(PublishTaskService.MODULE_TYPE, taskId)).thenReturn(lock);
when(fileTaskMapper.selectById(taskId)).thenReturn(task);
when(publishFileMapper.selectById(fileId)).thenReturn(file);
when(publishFileMapper.selectList(any())).thenReturn(List.of(file));
when(fileResultMapper.selectOne(any())).thenReturn(result(taskId, 324L));
when(taskChunkMapper.insert(any(TaskChunkEntity.class))).thenAnswer(invocation -> {
TaskChunkEntity attempted = invocation.getArgument(0);
TaskChunkEntity winner = new TaskChunkEntity();
winner.setId(999L);
winner.setTaskId(attempted.getTaskId());
winner.setModuleType(attempted.getModuleType());
winner.setScopeKey(attempted.getScopeKey());
winner.setScopeHash(attempted.getScopeHash());
winner.setChunkIndex(attempted.getChunkIndex());
winner.setChunkTotal(attempted.getChunkTotal());
winner.setPayloadJson("rustfs:test/publish/winner");
winner.setPayloadHash(attempted.getPayloadHash());
storedChunks.add(winner);
throw new DuplicateKeyException("duplicate chunk");
});
service.submitResult(taskId, chunkResultRequest(7L, fileId, 1, 2, List.of(row("1"))));
assertEquals(1, storedChunks.size());
assertEquals("rustfs:test/publish/winner", storedChunks.getFirst().getPayloadJson());
assertTrue(rustfsPayloads.isEmpty());
assertEquals(List.of(false, false), storageTransactionStates);
verify(transientPayloadStorageService).deletePayloadIfPresent(any());
verify(lock).close();
}
@Test
@@ -804,11 +903,14 @@ class PublishTaskServiceTest {
storedChunks.clear();
storedScopes.clear();
rustfsPayloads.clear();
storageTransactionStates.clear();
transactionActive.set(false);
nextPayloadId = 0;
lenient().when(transientPayloadStorageService.isSharedWriteEnabled()).thenReturn(true);
lenient().when(transientPayloadStorageService.storeChunkPayloadVersioned(
any(), any(), any(), any(), any())).thenAnswer(invocation -> {
storageTransactionStates.add(transactionActive.get());
String pointer = "rustfs:test/publish/chunk-" + (++nextPayloadId);
rustfsPayloads.put(pointer, invocation.getArgument(4));
return pointer;
@@ -818,6 +920,7 @@ class PublishTaskServiceTest {
return value != null && value.startsWith("rustfs:") ? value : null;
});
lenient().when(transientPayloadStorageService.resolvePayload(any(), any())).thenAnswer(invocation -> {
storageTransactionStates.add(transactionActive.get());
String pointer = invocation.getArgument(0);
String payload = rustfsPayloads.get(pointer);
if (payload == null) {
@@ -826,6 +929,7 @@ class PublishTaskServiceTest {
return payload;
});
lenient().doAnswer(invocation -> {
storageTransactionStates.add(transactionActive.get());
rustfsPayloads.remove(invocation.getArgument(0));
return null;
}).when(transientPayloadStorageService).deletePayloadIfPresent(any());
@@ -848,6 +952,19 @@ class PublishTaskServiceTest {
});
lenient().when(taskChunkMapper.selectCount(any())).thenAnswer(invocation -> {
LambdaQueryWrapper<TaskChunkEntity> query = invocation.getArgument(0);
query.getSqlSegment();
List<String> payloadValues = query.getParamNameValuePairs().values().stream()
.filter(String.class::isInstance)
.map(String.class::cast)
.filter(value -> value.contains("rustfs:")
|| value.contains("local:")
|| value.contains("oss:"))
.toList();
if (!payloadValues.isEmpty()) {
return storedChunks.stream()
.filter(chunk -> payloadValues.contains(chunk.getPayloadJson()))
.count();
}
Long taskId = queryLong(query);
String scopeHash = queryScopeHash(query);
return storedChunks.stream()
@@ -0,0 +1,366 @@
package com.nanri.aiimage.modules.similarasin.service;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.nanri.aiimage.common.service.DistributedJobLockService;
import com.nanri.aiimage.config.InstanceMetadata;
import com.nanri.aiimage.config.SimilarAsinProperties;
import com.nanri.aiimage.config.StorageProperties;
import com.nanri.aiimage.modules.coze.service.CozeCredentialPoolService;
import com.nanri.aiimage.modules.file.service.LocalFileStorageService;
import com.nanri.aiimage.modules.file.service.oss.OssStorageService;
import com.nanri.aiimage.modules.similarasin.client.SimilarAsinCozeClient;
import com.nanri.aiimage.modules.similarasin.mapper.SimilarAsinFilterConditionMapper;
import com.nanri.aiimage.modules.similarasin.model.dto.SimilarAsinParsedPayloadDto;
import com.nanri.aiimage.modules.similarasin.model.dto.SimilarAsinSourceFileDto;
import com.nanri.aiimage.modules.similarasin.model.dto.SimilarAsinSubmitResultRequest;
import com.nanri.aiimage.modules.similarasin.model.vo.SimilarAsinParsedRowVo;
import com.nanri.aiimage.modules.similarasin.util.SimilarAsinImageEmbedder;
import com.nanri.aiimage.modules.task.mapper.FileResultMapper;
import com.nanri.aiimage.modules.task.mapper.FileTaskMapper;
import com.nanri.aiimage.modules.task.mapper.TaskChunkMapper;
import com.nanri.aiimage.modules.task.mapper.TaskScopeStateMapper;
import com.nanri.aiimage.modules.task.model.entity.FileResultEntity;
import com.nanri.aiimage.modules.task.model.entity.FileTaskEntity;
import com.nanri.aiimage.modules.task.model.entity.TaskChunkEntity;
import com.nanri.aiimage.modules.task.model.entity.TaskScopeStateEntity;
import com.nanri.aiimage.modules.task.service.TaskDistributedLockService;
import com.nanri.aiimage.modules.task.service.TaskFileJobService;
import com.nanri.aiimage.modules.task.service.TaskProgressSnapshotService;
import com.nanri.aiimage.modules.task.service.TransientPayloadStorageService;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.Spy;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import com.baomidou.mybatisplus.core.MybatisConfiguration;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import org.apache.ibatis.builder.MapperBuilderAssistant;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@ExtendWith(MockitoExtension.class)
class SimilarAsinTaskServiceSubmitTest {
private static final Long TASK_ID = 21879L;
private static final String PARSED_POINTER = "rustfs:task-parsed/similar-asin/21879/payload.json";
private static final String CHUNK_POINTER = "rustfs:task-chunk/similar-asin/21879/chunk.json";
private static final String STORED_CHUNK_POINTER = "\"" + CHUNK_POINTER + "\"";
@Mock private LocalFileStorageService localFileStorageService;
@Mock private OssStorageService ossStorageService;
@Mock private StorageProperties storageProperties;
@Mock private FileTaskMapper fileTaskMapper;
@Mock private FileResultMapper fileResultMapper;
@Mock private TaskScopeStateMapper taskScopeStateMapper;
@Mock private TaskChunkMapper taskChunkMapper;
@Mock private SimilarAsinFilterConditionMapper filterConditionMapper;
@Spy private ObjectMapper objectMapper = new ObjectMapper();
@Mock private SimilarAsinCozeClient cozeClient;
@Mock private SimilarAsinTaskCacheService taskCacheService;
@Mock private SimilarAsinProperties properties;
@Mock private TaskFileJobService taskFileJobService;
@Mock private TaskDistributedLockService taskDistributedLockService;
@Mock private TaskProgressSnapshotService taskProgressSnapshotService;
@Mock private TransientPayloadStorageService transientPayloadStorageService;
@Mock private PlatformTransactionManager transactionManager;
@Mock private DistributedJobLockService distributedJobLockService;
@Mock private InstanceMetadata instanceMetadata;
@Mock private CozeCredentialPoolService cozeCredentialPoolService;
@Mock private SimilarAsinImageEmbedder imageEmbedder;
@Mock private SimilarAsinImagePrefetchService imagePrefetchService;
@Mock private TransactionStatus transactionStatus;
@InjectMocks private SimilarAsinTaskService service;
private final AtomicBoolean transactionActive = new AtomicBoolean();
@BeforeAll
static void initializeMybatisMetadata() {
MapperBuilderAssistant assistant = new MapperBuilderAssistant(new MybatisConfiguration(), "");
TableInfoHelper.initTableInfo(assistant, FileTaskEntity.class);
TableInfoHelper.initTableInfo(assistant, FileResultEntity.class);
TableInfoHelper.initTableInfo(assistant, TaskChunkEntity.class);
TableInfoHelper.initTableInfo(assistant, TaskScopeStateEntity.class);
}
@BeforeEach
void setUpTransactionAndLock() {
lenient().when(instanceMetadata.getInstanceId()).thenReturn("instance-a");
lenient().when(taskDistributedLockService.acquire(
eq(SimilarAsinTaskService.MODULE_TYPE),
anyLong(),
any(Duration.class),
eq(10_000L)))
.thenReturn(mock(TaskDistributedLockService.LockHandle.class));
lenient().when(transactionManager.getTransaction(any(TransactionDefinition.class)))
.thenAnswer(invocation -> {
transactionActive.set(true);
return transactionStatus;
});
lenient().doAnswer(invocation -> {
transactionActive.set(false);
return null;
}).when(transactionManager).commit(transactionStatus);
lenient().doAnswer(invocation -> {
transactionActive.set(false);
return null;
}).when(transactionManager).rollback(transactionStatus);
lenient().when(taskChunkMapper.selectList(any())).thenReturn(List.of());
lenient().when(fileResultMapper.selectList(any())).thenReturn(List.of());
}
@AfterEach
void shutdownExecutors() {
service.shutdownAssembleExecutor();
}
@Test
void doneCallbackReadsPayloadOnlyBeforeShortTransaction() throws Exception {
FileTaskEntity task = runningTask("instance-a");
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
configureNewChunkAndScope();
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
List<Boolean> storageCallTransactionStates = new ArrayList<>();
when(transientPayloadStorageService.resolvePayload(eq(PARSED_POINTER), anyString()))
.thenAnswer(invocation -> {
storageCallTransactionStates.add(transactionActive.get());
return parsedPayloadJson(List.of(new SimilarAsinParsedRowVo(), new SimilarAsinParsedRowVo()));
});
when(transientPayloadStorageService.storeChunkPayloadVersioned(
eq(SimilarAsinTaskService.MODULE_TYPE), eq(TASK_ID), anyString(), eq(0), anyString()))
.thenAnswer(invocation -> {
storageCallTransactionStates.add(transactionActive.get());
return STORED_CHUNK_POINTER;
});
when(transientPayloadStorageService.wasLastStoreLocalFallback()).thenReturn(false);
doAnswer(invocation -> {
assertTrue(transactionActive.get());
FileResultEntity result = invocation.getArgument(0);
result.setId(501L);
return 1;
}).when(fileResultMapper).insert(any(FileResultEntity.class));
doAnswer(invocation -> {
storageCallTransactionStates.add(transactionActive.get());
return null;
}).when(taskCacheService).deleteTaskCache(TASK_ID);
service.submitResult(TASK_ID, request(true));
assertFalse(storageCallTransactionStates.isEmpty());
assertTrue(storageCallTransactionStates.stream().noneMatch(Boolean::booleanValue));
verify(transientPayloadStorageService).resolvePayload(eq(PARSED_POINTER), anyString());
ArgumentCaptor<FileResultEntity> resultCaptor = ArgumentCaptor.forClass(FileResultEntity.class);
verify(fileResultMapper).insert(resultCaptor.capture());
assertEquals(2, resultCaptor.getValue().getRowCount());
assertEquals("germany.xlsx", resultCaptor.getValue().getSourceFilename());
}
@Test
void unknownCommitOutcomeDoesNotDeletePossiblyCommittedChunk() throws Exception {
FileTaskEntity task = runningTask("instance-a");
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
configureNewChunkAndScope();
when(taskChunkMapper.selectCount(any())).thenReturn(0L);
configureChunkStore(false);
doAnswer(invocation -> {
transactionActive.set(false);
throw new IllegalStateException("commit ACK lost");
}).when(transactionManager).commit(transactionStatus);
IllegalStateException thrown = assertThrows(IllegalStateException.class,
() -> service.submitResult(TASK_ID, request(false)));
assertEquals("commit ACK lost", thrown.getMessage());
verify(transientPayloadStorageService, never()).extractPointer(anyString());
verify(transientPayloadStorageService, never()).deletePayloadIfPresent(anyString());
verify(taskCacheService, never()).touchTaskHeartbeat(TASK_ID);
}
@Test
void duplicateAlreadyPersistedChunkDoesNotRunPayloadCleanup() throws Exception {
FileTaskEntity task = runningTask("instance-a");
task.setResultJson("{\"ownerInstanceId\":\"instance-a\"}");
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
when(taskChunkMapper.selectOne(any())).thenReturn(chunk("\"rustfs:winner\""));
configureScopeStorage(null);
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
service.submitResult(TASK_ID, request(false));
verify(transientPayloadStorageService, never()).extractPointer(anyString());
verify(transientPayloadStorageService, never()).deletePayloadIfPresent(anyString());
}
@Test
void localFallbackOwnerFailureRollsBackChunkAndKeepsCandidate() throws Exception {
FileTaskEntity task = runningTask(null);
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
configureNewChunkAndScope();
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
configureChunkStore(true);
when(fileTaskMapper.updateById(any(FileTaskEntity.class))).thenReturn(0);
IllegalStateException thrown = assertThrows(IllegalStateException.class,
() -> service.submitResult(TASK_ID, request(false)));
assertTrue(thrown.getMessage().contains("Failed to bind local fallback task owner"));
verify(taskChunkMapper).insert(any(TaskChunkEntity.class));
verify(transactionManager).rollback(transactionStatus);
verify(transactionManager, never()).commit(transactionStatus);
verify(transientPayloadStorageService, never()).deletePayloadIfPresent(anyString());
}
@Test
void duplicateLocalCandidateIsNotBoundAndReferencedPayloadIsKept() throws Exception {
FileTaskEntity task = runningTask(null);
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
TaskChunkEntity winner = chunk(STORED_CHUNK_POINTER);
when(taskChunkMapper.selectOne(any())).thenReturn(null, winner);
configureScopeStorage(null);
configureChunkStore(true);
when(transientPayloadStorageService.extractPointer(STORED_CHUNK_POINTER)).thenReturn(CHUNK_POINTER);
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
service.submitResult(TASK_ID, request(false));
verify(taskChunkMapper, never()).insert(any(TaskChunkEntity.class));
verify(fileTaskMapper, never()).updateById(any(FileTaskEntity.class));
verify(transientPayloadStorageService, never()).deletePayloadIfPresent(STORED_CHUNK_POINTER);
}
@Test
void repeatedNonFinalCallbackCannotClearCompletedScope() throws Exception {
FileTaskEntity task = runningTask("instance-a");
task.setResultJson("{\"ownerInstanceId\":\"instance-a\"}");
when(fileTaskMapper.selectById(TASK_ID)).thenReturn(task);
TaskChunkEntity existing = chunk("\"rustfs:winner\"");
when(taskChunkMapper.selectOne(any())).thenReturn(existing);
TaskScopeStateEntity scope = scope(1);
when(taskScopeStateMapper.selectOne(any())).thenReturn(scope);
when(taskChunkMapper.selectCount(any())).thenReturn(1L);
service.submitResult(TASK_ID, request(false));
assertEquals(1, scope.getCompleted());
verify(taskScopeStateMapper).updateById(scope);
verify(transientPayloadStorageService, never()).storeChunkPayloadVersioned(
anyString(), anyLong(), anyString(), any(), anyString());
}
private String parsedPayloadJson(List<SimilarAsinParsedRowVo> rows) throws Exception {
SimilarAsinParsedPayloadDto payload = new SimilarAsinParsedPayloadDto();
payload.setAllItems(rows);
payload.setItems(rows);
SimilarAsinSourceFileDto sourceFile = new SimilarAsinSourceFileDto();
sourceFile.setFileKey("uploads/germany.xlsx");
sourceFile.setOriginalFilename("germany.xlsx");
payload.setSourceFiles(List.of(sourceFile));
return objectMapper.writeValueAsString(payload);
}
private void configureNewChunkAndScope() {
AtomicReference<TaskChunkEntity> chunkRef = new AtomicReference<>();
when(taskChunkMapper.selectOne(any())).thenAnswer(invocation -> chunkRef.get());
doAnswer(invocation -> {
TaskChunkEntity chunk = invocation.getArgument(0);
chunk.setId(301L);
chunkRef.set(chunk);
return 1;
}).when(taskChunkMapper).insert(any(TaskChunkEntity.class));
configureScopeStorage(null);
}
private void configureScopeStorage(TaskScopeStateEntity initial) {
AtomicReference<TaskScopeStateEntity> scopeRef = new AtomicReference<>(initial);
when(taskScopeStateMapper.selectOne(any())).thenAnswer(invocation -> scopeRef.get());
doAnswer(invocation -> {
TaskScopeStateEntity scope = invocation.getArgument(0);
scope.setId(401L);
scopeRef.set(scope);
return 1;
}).when(taskScopeStateMapper).insert(any(TaskScopeStateEntity.class));
}
private void configureChunkStore(boolean localFallback) {
when(transientPayloadStorageService.storeChunkPayloadVersioned(
eq(SimilarAsinTaskService.MODULE_TYPE), eq(TASK_ID), anyString(), eq(0), anyString()))
.thenReturn(STORED_CHUNK_POINTER);
when(transientPayloadStorageService.wasLastStoreLocalFallback()).thenReturn(localFallback);
}
private FileTaskEntity runningTask(String owner) throws Exception {
FileTaskEntity task = new FileTaskEntity();
task.setId(TASK_ID);
task.setModuleType(SimilarAsinTaskService.MODULE_TYPE);
task.setStatus("RUNNING");
task.setUserId(7L);
String resultJson = "{\"parsedPayloadRef\":\"" + PARSED_POINTER + "\"";
if (owner != null) {
resultJson += ",\"ownerInstanceId\":\"" + owner + "\"";
}
task.setResultJson(resultJson + "}");
return task;
}
private SimilarAsinSubmitResultRequest request(boolean done) {
SimilarAsinSubmitResultRequest request = new SimilarAsinSubmitResultRequest();
request.setSubmissionId("similar-asin-21879");
request.setChunkIndex(0);
request.setChunkTotal(1);
request.setDone(done);
return request;
}
private TaskChunkEntity chunk(String payload) {
TaskChunkEntity chunk = new TaskChunkEntity();
chunk.setId(301L);
chunk.setTaskId(TASK_ID);
chunk.setModuleType(SimilarAsinTaskService.MODULE_TYPE);
chunk.setScopeHash("existing-scope");
chunk.setChunkIndex(0);
chunk.setChunkTotal(1);
chunk.setPayloadJson(payload);
return chunk;
}
private TaskScopeStateEntity scope(int completed) {
TaskScopeStateEntity scope = new TaskScopeStateEntity();
scope.setId(401L);
scope.setTaskId(TASK_ID);
scope.setModuleType(SimilarAsinTaskService.MODULE_TYPE);
scope.setScopeKey("similar-asin-21879");
scope.setScopeHash("existing-scope");
scope.setChunkTotal(1);
scope.setCompleted(completed);
scope.setStateJson("{\"phase\":\"RECEIVED\",\"coze\":\"PENDING\"}");
return scope;
}
}
@@ -0,0 +1,208 @@
package com.nanri.aiimage.modules.task.service;
import com.baomidou.mybatisplus.core.MybatisConfiguration;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.nanri.aiimage.modules.task.mapper.TaskScopeStateMapper;
import com.nanri.aiimage.modules.task.model.entity.TaskScopeStateEntity;
import org.apache.ibatis.builder.MapperBuilderAssistant;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.SimpleTransactionStatus;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
class TaskScopePayloadStorageServiceTest {
@BeforeAll
static void initializeMybatisMetadata() {
TableInfoHelper.initTableInfo(new MapperBuilderAssistant(new MybatisConfiguration(), ""),
TaskScopeStateEntity.class);
}
@Test
void uploadsBeforeOpeningTransactionAndCleansReplacedPayloadAfterCommit() {
TaskScopeStateMapper mapper = mock(TaskScopeStateMapper.class);
TransientPayloadStorageService payloadStorage = mock(TransientPayloadStorageService.class);
PlatformTransactionManager transactionManager = mock(PlatformTransactionManager.class);
TaskScopePayloadStorageService service = new TaskScopePayloadStorageService(
mapper, new ObjectMapper(), payloadStorage, transactionManager);
List<String> events = new ArrayList<>();
AtomicBoolean transactionActive = new AtomicBoolean();
TaskScopeStateEntity existing = new TaskScopeStateEntity();
existing.setId(9L);
existing.setTaskId(7L);
existing.setModuleType("PRICE_TRACK");
existing.setScopeHash("scope-hash");
existing.setStateJson("rustfs:old.json");
when(payloadStorage.storeScopePayloadVersioned(
eq("PRICE_TRACK"), eq(7L), any(), any(), eq(true)))
.thenAnswer(invocation -> {
assertFalse(transactionActive.get());
events.add("upload");
return "rustfs:new.json";
});
when(payloadStorage.extractPointer("rustfs:new.json")).thenReturn("rustfs:new.json");
when(payloadStorage.extractPointer("rustfs:old.json")).thenReturn("rustfs:old.json");
when(transactionManager.getTransaction(any(TransactionDefinition.class)))
.thenAnswer(invocation -> {
events.add("begin");
transactionActive.set(true);
return new SimpleTransactionStatus();
});
when(mapper.selectOne(any())).thenAnswer(invocation -> {
LambdaQueryWrapper<?> query = invocation.getArgument(0);
assertTrue(query.getSqlSegment().endsWith("limit 1 FOR UPDATE"));
return existing;
});
when(mapper.update(any(), any())).thenReturn(1);
doAnswer(invocation -> {
events.add("commit");
transactionActive.set(false);
return null;
}).when(transactionManager).commit(any(TransactionStatus.class));
doAnswer(invocation -> {
assertFalse(transactionActive.get());
events.add("cleanup");
return null;
}).when(payloadStorage).deleteReplacedPayloadIfNeeded("rustfs:old.json", "rustfs:new.json");
service.saveScopePayload(7L, "PRICE_TRACK", "shop-a", new Payload("ok"));
assertEquals(List.of("upload", "begin", "commit", "cleanup"), events);
verify(mapper).update(any(), any());
}
@Test
void deletesUploadedPayloadAfterRollbackWhenDatabaseConfirmsItIsUnreferenced() {
TaskScopeStateMapper mapper = mock(TaskScopeStateMapper.class);
TransientPayloadStorageService payloadStorage = mock(TransientPayloadStorageService.class);
PlatformTransactionManager transactionManager = mock(PlatformTransactionManager.class);
TaskScopePayloadStorageService service = new TaskScopePayloadStorageService(
mapper, new ObjectMapper(), payloadStorage, transactionManager);
List<String> events = new ArrayList<>();
AtomicBoolean transactionActive = new AtomicBoolean();
TaskScopeStateEntity existing = new TaskScopeStateEntity();
existing.setId(9L);
existing.setTaskId(7L);
existing.setModuleType("PRICE_TRACK");
existing.setScopeHash("scope-hash");
existing.setStateJson("rustfs:old.json");
when(payloadStorage.storeScopePayloadVersioned(
eq("PRICE_TRACK"), eq(7L), any(), any(), eq(true)))
.thenAnswer(invocation -> {
assertFalse(transactionActive.get());
events.add("upload");
return "rustfs:new.json";
});
when(payloadStorage.extractPointer("rustfs:new.json")).thenReturn("rustfs:new.json");
when(payloadStorage.extractPointer("rustfs:old.json")).thenReturn("rustfs:old.json");
when(transactionManager.getTransaction(any(TransactionDefinition.class)))
.thenAnswer(invocation -> {
events.add("begin");
transactionActive.set(true);
return new SimpleTransactionStatus();
});
when(mapper.selectOne(any())).thenAnswer(invocation -> {
events.add(transactionActive.get() ? "transaction-read" : "cleanup-read");
return existing;
});
when(mapper.update(any(), any())).thenThrow(new IllegalStateException("database write failed"));
doAnswer(invocation -> {
events.add("rollback");
transactionActive.set(false);
return null;
}).when(transactionManager).rollback(any(TransactionStatus.class));
doAnswer(invocation -> {
assertFalse(transactionActive.get());
events.add("delete");
return null;
}).when(payloadStorage).deletePayloadIfPresent("rustfs:new.json");
assertThrows(IllegalStateException.class,
() -> service.saveScopePayload(7L, "PRICE_TRACK", "shop-a", new Payload("ok")));
assertEquals(List.of("upload", "begin", "transaction-read", "rollback", "cleanup-read", "delete"), events);
verify(payloadStorage).deletePayloadIfPresent("rustfs:new.json");
}
@Test
void keepsUploadedPayloadWhenCommitThrowsButDatabaseReferencesIt() {
TaskScopeStateMapper mapper = mock(TaskScopeStateMapper.class);
TransientPayloadStorageService payloadStorage = mock(TransientPayloadStorageService.class);
PlatformTransactionManager transactionManager = mock(PlatformTransactionManager.class);
TaskScopePayloadStorageService service = new TaskScopePayloadStorageService(
mapper, new ObjectMapper(), payloadStorage, transactionManager);
List<String> events = new ArrayList<>();
AtomicBoolean transactionActive = new AtomicBoolean();
TaskScopeStateEntity previous = new TaskScopeStateEntity();
previous.setId(9L);
previous.setTaskId(7L);
previous.setModuleType("PRICE_TRACK");
previous.setScopeHash("scope-hash");
previous.setStateJson("rustfs:old.json");
TaskScopeStateEntity committed = new TaskScopeStateEntity();
committed.setId(9L);
committed.setTaskId(7L);
committed.setModuleType("PRICE_TRACK");
committed.setScopeHash("scope-hash");
committed.setStateJson("rustfs:new.json");
when(payloadStorage.storeScopePayloadVersioned(
eq("PRICE_TRACK"), eq(7L), any(), any(), eq(true)))
.thenAnswer(invocation -> {
assertFalse(transactionActive.get());
events.add("upload");
return "rustfs:new.json";
});
when(payloadStorage.extractPointer("rustfs:new.json")).thenReturn("rustfs:new.json");
when(payloadStorage.extractPointer("rustfs:old.json")).thenReturn("rustfs:old.json");
when(transactionManager.getTransaction(any(TransactionDefinition.class)))
.thenAnswer(invocation -> {
events.add("begin");
transactionActive.set(true);
return new SimpleTransactionStatus();
});
when(mapper.selectOne(any())).thenAnswer(invocation -> {
events.add(transactionActive.get() ? "transaction-read" : "cleanup-read");
return transactionActive.get() ? previous : committed;
});
when(mapper.update(any(), any())).thenReturn(1);
doAnswer(invocation -> {
events.add("commit");
transactionActive.set(false);
throw new IllegalStateException("commit outcome unknown");
}).when(transactionManager).commit(any(TransactionStatus.class));
IllegalStateException error = assertThrows(IllegalStateException.class,
() -> service.saveScopePayload(7L, "PRICE_TRACK", "shop-a", new Payload("ok")));
assertEquals("commit outcome unknown", error.getMessage());
assertEquals(List.of("upload", "begin", "transaction-read", "commit", "cleanup-read"), events);
verify(payloadStorage, never()).deletePayloadIfPresent("rustfs:new.json");
verify(payloadStorage, never()).deleteReplacedPayloadIfNeeded(any(), any());
}
private record Payload(String value) {
}
}
@@ -0,0 +1,63 @@
package com.nanri.aiimage.modules.ziniao.service;
import com.nanri.aiimage.config.ZiniaoProperties;
import com.nanri.aiimage.modules.ziniao.client.ZiniaoClient;
import com.nanri.aiimage.modules.ziniao.memory.service.ZiniaoTransientCacheService;
import com.nanri.aiimage.modules.ziniao.model.vo.ZiniaoStaffItemVo;
import org.junit.jupiter.api.Test;
import java.time.Duration;
import java.util.List;
import java.util.Optional;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
class ZiniaoAuthServiceTest {
@Test
void invalidUserCacheSkipsUpstreamStoreLookup() {
ZiniaoClient client = mock(ZiniaoClient.class);
ZiniaoTransientCacheService cache = mock(ZiniaoTransientCacheService.class);
ZiniaoAuthService service = service(client, cache);
when(cache.get(eq("INVALID_USER_STORES"), anyString(), eq(Boolean.class)))
.thenReturn(Optional.of(Boolean.TRUE));
assertTrue(service.getOrLoadUserStoresForIndex("api-key", 12L, 34L).isEmpty());
verify(client, never()).listUserStores(anyString(), any(), any());
}
@Test
void markingInvalidUserWritesNegativeCacheAndRemovesStaleStaffEntry() {
ZiniaoTransientCacheService cache = mock(ZiniaoTransientCacheService.class);
ZiniaoAuthService service = service(mock(ZiniaoClient.class), cache);
ZiniaoStaffItemVo invalid = new ZiniaoStaffItemVo();
invalid.setUserId(34L);
ZiniaoStaffItemVo valid = new ZiniaoStaffItemVo();
valid.setUserId(35L);
when(cache.getList(eq("STAFF_LIST"), anyString(), eq(ZiniaoStaffItemVo.class)))
.thenReturn(Optional.of(List.of(invalid, valid)));
service.evictInvalidUserForIndex("api-key", 12L, 34L);
verify(cache).put(eq("INVALID_USER_STORES"), anyString(), eq(Boolean.TRUE), eq(Duration.ofMinutes(30)));
verify(cache).put(eq("STAFF_LIST"), anyString(), eq(List.of(valid)), eq(Duration.ofMinutes(30)));
verify(cache).delete(eq("USER_STORES"), anyString());
}
private ZiniaoAuthService service(ZiniaoClient client, ZiniaoTransientCacheService cache) {
return new ZiniaoAuthService(
new ZiniaoProperties(),
client,
mock(ZiniaoSessionCacheService.class),
cache,
mock(ZiniaoApiKeyProvider.class));
}
}