fix: 冷静一下key

smilexizheng · Mar 19, 2024 · 208860f · 208860f
1 parent c51a7cf
commit 208860f
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 18 deletions.
diff --git a/src/main/java/com/coco/boot/constant/SysConstant.java b/src/main/java/com/coco/boot/constant/SysConstant.java
@@ -12,7 +12,7 @@ public interface SysConstant {
      */
     String GHU_ALIVE_KEY ="ghu_key:alive:";
     String GHU_NO_ALIVE_KEY ="ghu_key:noAlive:";
-    String GHU_COOLING_KEY ="ghu_key:cooling:";
+    String GHU_COOLING_KEY ="ghu_key:cooling:map";
 
     String GHU_RATE_LIMITER = "ghu-rate-limiter:";
     String USER_RATE_LIMITER = "user-rate-limiter:";

diff --git a/src/main/java/com/coco/boot/service/impl/CoCoPilotServiceImpl.java b/src/main/java/com/coco/boot/service/impl/CoCoPilotServiceImpl.java
@@ -17,6 +17,7 @@
 import lombok.extern.slf4j.Slf4j;
 import org.jetbrains.annotations.NotNull;
 import org.redisson.api.*;
+import org.redisson.api.map.event.EntryExpiredListener;
 import org.redisson.client.codec.StringCodec;
 import org.springframework.http.*;
 import org.springframework.stereotype.Service;
@@ -47,19 +48,6 @@ public class CoCoPilotServiceImpl implements CoCoPilotService {
 
     private final CoCoConfig coCoConfig;
 
-//    private static final HttpHeaders apiHeaders;
-//
-//    static {
-//        apiHeaders = new HttpHeaders();
-//        apiHeaders.set("Access-Control-Allow-Origin", "*");
-//        apiHeaders.set("Host", "api.cocopilot.com");
-//        apiHeaders.set("Editor-Version", "vscode/1.85.2");
-//        apiHeaders.set("Editor-Plugin-Version", "copilot-chat/0.11.1");
-//        apiHeaders.set("User-Agent", "GitHubCopilotChat/0.11.1");
-//        apiHeaders.set("Accept", "*/*");
-//        apiHeaders.set("Accept-Encoding", "gzip, deflate, br");
-//    }
-
     @Override
     public R<String> uploadGhu(String data) {
         if (!StringUtils.hasLength(data)) {
@@ -259,9 +247,32 @@ private ResponseEntity<String> getBaseProxyResponse(Object requestBody, String p
                 atomicLong.incrementAndGet();
                 return response;
             } else {
-                // 异步移除，并添加到不可用key
                 ghuAliveKey.remove(ghu);
-                redissonClient.getSet(GHU_NO_ALIVE_KEY, StringCodec.INSTANCE).addAsync(ghu);
+                if (response.getStatusCode() == HttpStatus.TOO_MANY_REQUESTS) {
+                    String retryAfter = response.getHeaders().getFirst("x-ratelimit-user-retry-after");
+                    // 默认 600秒
+                    long time = 120;
+                    if (StringUtil.isNotBlank(retryAfter)) {
+                        try {
+                            time = Long.parseLong(retryAfter);
+                        } catch (NumberFormatException e) {}
+                    }
+
+                    if (time > 1000) {
+                        redissonClient.getSet(GHU_NO_ALIVE_KEY, StringCodec.INSTANCE).addAsync(ghu);
+                    } else {
+                        RMapCache<String, Integer> collingMap = redissonClient.getMapCache(GHU_COOLING_KEY);
+                        if (!collingMap.isExists()) {
+                            collingMap.addListener((EntryExpiredListener<String, Integer>) event -> {
+                                // expired key
+                                redissonClient.getSet(GHU_ALIVE_KEY, StringCodec.INSTANCE).add(event.getKey());
+                            });
+                        }
+                        collingMap.put(ghu, 1, time+5, TimeUnit.SECONDS);
+                    }
+                } else {
+                    redissonClient.getSet(GHU_NO_ALIVE_KEY, StringCodec.INSTANCE).addAsync(ghu);
+                }
                 i++;
             }
         }

diff --git a/src/main/resources/application-docker.yml b/src/main/resources/application-docker.yml
@@ -49,10 +49,10 @@ coco:
   frequencyTime: ${COCO_FREQUENCY_TIME:1}
   #ghu频率数
   frequencyDegree: ${COCO_FREQUENCY_DEGREE:8}
-  #用户基础频率  5秒1次
+  #用户基础频率  5分钟10次
   userRateTime: ${COCO_USER_RATE_TIME:5}
   #用户基础频率数
-  userFrequencyDegree: ${COCO_USER_FREQUENCY_DEGREE:1}
+  userFrequencyDegree: ${COCO_USER_FREQUENCY_DEGREE:10}
   #用户token 有效期无请求接口 小时
   userTokenExpire: ${COCO_USER_TOKEN_EXPIRE:1}
   #允许用户的最低等级等级