Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings
This repository was archived by the owner on May 11, 2024. It is now read-only.

Commit 8db9865

Browse files
chore: 多线程优化加载AdminApiKey
1 parent c634d34 commit 8db9865

File tree

2 files changed

+47
-14
lines changed

2 files changed

+47
-14
lines changed

‎src/main/java/com/gzhu/funai/config/FunAiConfig.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ public PaginationInterceptor paginationInterceptor() {
3939
public TaskExecutor queueThreadPool() {
4040
ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
4141
// 设置核心线程数
42-
executor.setCorePoolSize(8);
42+
executor.setCorePoolSize(15);
4343
// 设置最大线程数
4444
executor.setMaxPoolSize(1024);
4545
// 设置队列容量

‎src/main/java/com/gzhu/funai/service/impl/AdminApiKeyServiceImpl.java

Lines changed: 46 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package com.gzhu.funai.service.impl;
22

3-
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
43
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
54
import com.google.common.collect.ImmutableMap;
65
import com.gzhu.funai.api.openai.ChatGPTApi;
@@ -20,8 +19,12 @@
2019
import org.springframework.util.CollectionUtils;
2120

2221
import javax.annotation.Resource;
22+
import java.util.Comparator;
2323
import java.util.List;
2424
import java.util.Map;
25+
import java.util.concurrent.ConcurrentHashMap;
26+
import java.util.concurrent.CopyOnWriteArrayList;
27+
import java.util.concurrent.CountDownLatch;
2528
import java.util.stream.Collectors;
2629

2730
/**
@@ -107,20 +110,50 @@ public String roundRobinGetByType(ApiType apiTypes) {
107110
}
108111

109112
/**
110-
* 1 初始化数据并按照apikey的类型分组
113+
* 1 多线程判断apiKey是否能够被使用
111114
* 2 重置轮询下标
112-
* 定时任务:每隔60分钟执行一次
115+
* 定时任务:每隔1小时执行一次
113116
*/
114117
@Scheduled(initialDelay = TimeInterval.ZERO, fixedRate = TimeInterval.ONE_HOUR)
115118
@Override
116119
public void load(){
117-
Map<Integer, List<AdminApiKeyEntity>> collect = baseMapper.selectList(
118-
new QueryWrapper<AdminApiKeyEntity>().orderByDesc("priority")).stream()
119-
.filter(item -> filterInValidOpenAiApiKey(item))
120-
.collect(Collectors.groupingBy(AdminApiKeyEntity::getType)
121-
);
122-
this.cache = ImmutableMap.copyOf(collect);
120+
Map<Integer, List<AdminApiKeyEntity>> collect = new ConcurrentHashMap<>();
121+
List<AdminApiKeyEntity> adminApiKeyEntityList = baseMapper.selectList(null);
122+
123+
// 使用减少计数辅助类让主线程等待多线程执行完毕
124+
CountDownLatch countDownLatch = new CountDownLatch(adminApiKeyEntityList.size());
125+
for(AdminApiKeyEntity adminApiKeyEntity: adminApiKeyEntityList){
126+
queueThreadPool.execute(()->{
127+
try{
128+
if(isValidOpenAiApiKey(adminApiKeyEntity)){
129+
if(!collect.containsKey(adminApiKeyEntity.getType())){
130+
collect.putIfAbsent(adminApiKeyEntity.getType(), new CopyOnWriteArrayList<>());
131+
}
132+
collect.get(adminApiKeyEntity.getType()).add(adminApiKeyEntity);
133+
}
134+
}
135+
finally {
136+
countDownLatch.countDown();
137+
}
138+
});
139+
}
140+
try {
141+
countDownLatch.await();
142+
} catch (InterruptedException e) {
143+
log.error("error{}", e.getMessage());
144+
Thread.currentThread().interrupt();
145+
}
146+
147+
// 排序
148+
Map<Integer, List<AdminApiKeyEntity>> sortedCollect = collect.entrySet().stream().collect(
149+
Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().stream()
150+
// 按照优先级字段进行降序排序 再按照id进行升序排序
151+
.sorted(Comparator.comparing(AdminApiKeyEntity::getPriority).reversed().thenComparing(AdminApiKeyEntity::getId))
152+
.collect(Collectors.toList())));
153+
123154

155+
// 复制到缓存中
156+
this.cache = ImmutableMap.copyOf(sortedCollect);
124157
if(CollectionUtils.isEmpty(this.cache)){
125158
this.roundRobinIndex = new int[2];
126159
return;
@@ -150,7 +183,7 @@ public String getBestByType(ApiType apiTypes) {
150183
* @param adminApiKeyEntity
151184
* @return
152185
*/
153-
private boolean filterInValidOpenAiApiKey(AdminApiKeyEntity adminApiKeyEntity){
186+
private boolean isValidOpenAiApiKey(AdminApiKeyEntity adminApiKeyEntity){
154187
// 非openai类型,放行
155188
if(!ApiType.OPENAI.typeNo.equals(adminApiKeyEntity.getType())){
156189
return true;
@@ -168,7 +201,7 @@ private boolean filterInValidOpenAiApiKey(AdminApiKeyEntity adminApiKeyEntity){
168201
// 余额不足
169202
if(billingUsage.getTotalAmount().compareTo(billingUsage.getTotalUsage()) <= 0){
170203
log.error("{}的额度使用完毕!", adminApiKeyEntity.getName());
171-
queueThreadPool.execute(() -> baseMapper.deleteById(adminApiKeyEntity.getId()));
204+
baseMapper.deleteById(adminApiKeyEntity.getId());
172205
return false;
173206
}
174207

@@ -177,14 +210,14 @@ private boolean filterInValidOpenAiApiKey(AdminApiKeyEntity adminApiKeyEntity){
177210
adminApiKeyEntity.setTotalUsage(billingUsage.getTotalUsage());
178211
adminApiKeyEntity.setExpiredTime(billingUsage.getExpiredTime());
179212

180-
queueThreadPool.execute(() -> baseMapper.updateById(adminApiKeyEntity));
213+
baseMapper.updateById(adminApiKeyEntity);
181214
return true;
182215
}
183216
// 捕获 请求openai错误的异常, 删掉这个apiKey,不加载到缓存
184217
catch (BaseException e){
185218
log.error("apiKey:{}, error:{}",adminApiKeyEntity.getName(), e.getMsg());
186219
if(e.getCode() != OpenAiRespError.OPENAI_LIMIT_ERROR.code){
187-
queueThreadPool.execute(() -> baseMapper.deleteById(adminApiKeyEntity.getId()));
220+
baseMapper.deleteById(adminApiKeyEntity.getId());
188221
}
189222

190223
return false;

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /