fix: 上传文档增加限流处理

This commit is contained in:
liu.zixi
2025-11-25 10:34:51 +08:00
committed by joshen
parent 8112aea110
commit 3ddc9d58f0
2 changed files with 67 additions and 2 deletions

View File

@@ -17,6 +17,11 @@ public class CaseAiDocumentAsyncHandler {
private final AtomicInteger currentTaskCount = new AtomicInteger(0);
/**
* 限流默认QPS 40
*/
private final TokenBucketRateLimiter rateLimiter = new TokenBucketRateLimiter(40);
@Autowired
@Qualifier("aiDocExecutor")
private ThreadPoolTaskExecutor aiDocExecutor;
@@ -39,8 +44,13 @@ public class CaseAiDocumentAsyncHandler {
currentTaskCount.incrementAndGet();
aiDocExecutor.submit(() -> {
processCases(cases, optTypeEnum);
currentTaskCount.decrementAndGet();
try {
// 限流
rateLimiter.acquire();
processCases(cases, optTypeEnum);
} finally {
currentTaskCount.decrementAndGet();
}
});
}
}

View File

@@ -0,0 +1,55 @@
package com.xboe.module.boecase.async;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
/**
* 令牌桶限流算法实现
*/
public class TokenBucketRateLimiter {
private final double permitsPerSecond; // 每秒生成的令牌数(即 TPS
private final AtomicLong nextFreeTicketMicros = new AtomicLong(0); // 下一个令牌可用的时间(微秒)
private final AtomicLong storedPermits = new AtomicLong(0); // 当前桶中存储的令牌数(本简化版不支持突发,可省略)
private static final long MICROSECONDS_PER_SECOND = 1_000_000L;
public TokenBucketRateLimiter(double permitsPerSecond) {
this.permitsPerSecond = permitsPerSecond;
this.nextFreeTicketMicros.set(System.nanoTime() / 1000); // 初始化为当前时间(微秒)
}
/**
* 获取一个令牌,阻塞直到可用
*/
public void acquire() {
long waitMicros = reserve(1);
if (waitMicros > 0) {
try {
long waitNanos = waitMicros * 1000; // 转为纳秒
TimeUnit.NANOSECONDS.sleep(waitNanos);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
}
}
/**
* 预留 1 个令牌,返回需要等待的微秒数
*/
private long reserve(int permits) {
long nowMicros = System.nanoTime() / 1000;
long nextFreeTicket = nextFreeTicketMicros.get();
long waitMicros = Math.max(0, nextFreeTicket - nowMicros);
long newNextFreeTicket = nowMicros + waitMicros + (long) (permits * MICROSECONDS_PER_SECOND / permitsPerSecond);
while (!nextFreeTicketMicros.compareAndSet(nextFreeTicket, newNextFreeTicket)) {
// CAS 失败,说明其他线程修改了时间,重试
nowMicros = System.nanoTime() / 1000;
nextFreeTicket = nextFreeTicketMicros.get();
waitMicros = Math.max(0, nextFreeTicket - nowMicros);
newNextFreeTicket = nowMicros + waitMicros + (long) (permits * MICROSECONDS_PER_SECOND / permitsPerSecond);
}
return waitMicros;
}
}