super
107 lines
4.5 KiB
Java
107 lines
4.5 KiB
Java
package com.nanri.aiimage.config;
|
|
|
|
import com.nanri.aiimage.common.exception.BusinessException;
|
|
import com.nanri.aiimage.modules.task.service.TaskDistributedLockService;
|
|
import jakarta.servlet.http.HttpServletRequest;
|
|
import jakarta.servlet.http.HttpServletResponse;
|
|
import lombok.RequiredArgsConstructor;
|
|
import lombok.extern.slf4j.Slf4j;
|
|
import org.springframework.context.annotation.Configuration;
|
|
import org.springframework.web.servlet.HandlerInterceptor;
|
|
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
|
|
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
|
|
|
|
import java.util.Locale;
|
|
import java.util.Set;
|
|
import java.util.regex.Matcher;
|
|
import java.util.regex.Pattern;
|
|
|
|
@Configuration
|
|
@RequiredArgsConstructor
|
|
public class TaskOperationLockConfig implements WebMvcConfigurer {
|
|
|
|
private static final String LOCK_ATTRIBUTE = TaskOperationLockConfig.class.getName() + ".LOCK";
|
|
|
|
private final TaskDistributedLockService taskDistributedLockService;
|
|
|
|
@Override
|
|
public void addInterceptors(InterceptorRegistry registry) {
|
|
registry.addInterceptor(new TaskOperationLockInterceptor(taskDistributedLockService))
|
|
.addPathPatterns("/api/**");
|
|
}
|
|
|
|
public static boolean releaseRequestLock(HttpServletRequest request) {
|
|
if (request == null) {
|
|
return false;
|
|
}
|
|
Object lockHandle = request.getAttribute(LOCK_ATTRIBUTE);
|
|
if (lockHandle instanceof TaskDistributedLockService.LockHandle handle) {
|
|
request.removeAttribute(LOCK_ATTRIBUTE);
|
|
handle.close();
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
@Slf4j
|
|
private static final class TaskOperationLockInterceptor implements HandlerInterceptor {
|
|
|
|
private static final Pattern TASK_PATH = Pattern.compile(".*/api/([^/]+)/tasks/(\\d+)(?:/.*)?$");
|
|
private static final Pattern TASK_RESULT_PATH = Pattern.compile(".*/api/([^/]+)/tasks/(\\d+)/result/?$");
|
|
private static final Set<String> MUTATING_METHODS = Set.of("POST", "PUT", "PATCH", "DELETE");
|
|
private static final long RESULT_SUBMIT_WAIT_MILLIS = 30 * 1000L;
|
|
private static final long DELETE_WAIT_MILLIS = 60 * 1000L;
|
|
|
|
private final TaskDistributedLockService taskDistributedLockService;
|
|
|
|
private TaskOperationLockInterceptor(TaskDistributedLockService taskDistributedLockService) {
|
|
this.taskDistributedLockService = taskDistributedLockService;
|
|
}
|
|
|
|
@Override
|
|
public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
|
|
if (!MUTATING_METHODS.contains(request.getMethod())) {
|
|
return true;
|
|
}
|
|
String uri = request.getRequestURI();
|
|
if ("POST".equals(request.getMethod()) && TASK_RESULT_PATH.matcher(uri == null ? "" : uri).matches()) {
|
|
return true;
|
|
}
|
|
Matcher matcher = TASK_PATH.matcher(uri == null ? "" : uri);
|
|
if (!matcher.matches()) {
|
|
return true;
|
|
}
|
|
String moduleType = matcher.group(1).trim().toUpperCase(Locale.ROOT).replace('-', '_');
|
|
Long taskId = Long.valueOf(matcher.group(2));
|
|
long waitMillis = resolveWaitMillis(request.getMethod(), uri);
|
|
TaskDistributedLockService.LockHandle lockHandle =
|
|
taskDistributedLockService.acquire(moduleType, taskId, waitMillis);
|
|
if (lockHandle == null) {
|
|
log.info("[task-operation-lock] rejected busy task operation method={} uri={} moduleType={} taskId={} waitMillis={}",
|
|
request.getMethod(), uri, moduleType, taskId, waitMillis);
|
|
throw new BusinessException(40902, "TASK_LOCK_BUSY");
|
|
}
|
|
request.setAttribute(LOCK_ATTRIBUTE, lockHandle);
|
|
return true;
|
|
}
|
|
|
|
private long resolveWaitMillis(String method, String uri) {
|
|
if ("POST".equals(method) && TASK_RESULT_PATH.matcher(uri == null ? "" : uri).matches()) {
|
|
return RESULT_SUBMIT_WAIT_MILLIS;
|
|
}
|
|
if ("DELETE".equals(method)) {
|
|
return DELETE_WAIT_MILLIS;
|
|
}
|
|
return TaskDistributedLockService.DEFAULT_WAIT_MILLIS;
|
|
}
|
|
|
|
@Override
|
|
public void afterCompletion(HttpServletRequest request,
|
|
HttpServletResponse response,
|
|
Object handler,
|
|
Exception ex) {
|
|
releaseRequestLock(request);
|
|
}
|
|
}
|
|
}
|