✨ feat: 加入限流算法，限制搜书时间

apps/query.js

@@ -11,6 +11,7 @@ import { CAT_LIMIT } from "../utils/constant.js";
 import { getZHelper, getBookDetail, getYiBook, getZBook } from "../utils/books.js";
 // 工具类
 import _ from "lodash";
+import TokenBucket from '../utils/token-bucket.js'
 
 export class query extends plugin {
     constructor() {
@@ -56,6 +57,10 @@ export class query extends plugin {
                     reg: "^#竹白(.*)",
                     fnc: "zhubaiSearch",
                 },
+                {
+                    reg: "^#测试",
+                    fnc: "test1",
+                },
             ],
         });
     }
@@ -303,6 +308,7 @@ export class query extends plugin {
         }
 
         // 集成易书、zBook
+        const searchBookFunc = async () => {
             try {
                 const bookList = await Promise.allSettled([
                     getYiBook(e, keyword),
@@ -337,6 +343,8 @@ export class query extends plugin {
                 logger.error(err);
                 e.reply("部分搜书正在施工🚧");
             }
+        }
+        await this.limitUserUse(e, searchBookFunc);
         return true;
     }
 
@@ -355,6 +363,7 @@ export class query extends plugin {
             id = /\d+/.exec(keyword)[0];
             source = "";
         }
+        await this.limitUserUse(e, async () => {
             try {
                 const res = await getBookDetail(e, id, source);
                 e.reply(await Bot.makeForwardMsg(res));
@@ -362,6 +371,7 @@ export class query extends plugin {
                 logger.error(err);
                 e.reply("搜书正在施工🚧");
             }
+        })
         return true;
     }
 
@@ -404,9 +414,29 @@ export class query extends plugin {
         return true;
     }
 
+    /**
+     * 限制用户调用（默认1分钟1次）
+     * @param e
+     * @param func
+     * @return {Promise<void>}
+     */
+    async limitUserUse(e, func) {
+        if (query.#tokenBucket.consume(e.user_id, 1)) {
+            await func();
+        } else {
+            e.reply(`🙅‍${e.nickname}你已经被限流，请稍后再试！`, true);
+        }
+    }
+
     // 删除标签
     removeTag(title) {
         const titleRex = /<[^>]+>/g;
         return title.replace(titleRex, "");
     }
+
+    /**
+     * 令牌桶 拿来限流
+     * @type {TokenBucket}
+     */
+    static #tokenBucket = new TokenBucket(1, 1);
 }

utils/token-bucket.js

@@ -0,0 +1,57 @@
+export default class TokenBucket {
+    constructor(rate, capacity) {
+        this.rate = rate / 60; // 修改为每分钟生成的令牌数量
+        this.capacity = capacity;
+        this.tokens = capacity;
+        this.tokens = new Map(); // 使用 Map 存储每个用户的令牌桶
+        this.lastTime = new Date().getTime();
+    }
+
+    /**
+     * 消耗令牌
+     * @param id     用户id
+     * @param count  请求次数
+     * @return {boolean}
+     */
+    consume(id, count = 1) {
+        const now = new Date().getTime();
+        const elapsed = now - this.lastTime;
+        const addedTokens = elapsed * (this.rate / 1000);
+        this.lastTime = now;
+
+        // 获取用户的令牌桶，如果不存在则创建一个新的令牌桶
+        let userTokens = this.tokens.get(id);
+        if (!userTokens) {
+            userTokens = { tokens: this.capacity, lastTime: now };
+            this.tokens.set(id, userTokens);
+        }
+
+        // 更新用户的令牌桶中的令牌数量
+        userTokens.tokens = Math.min(
+            userTokens.tokens + addedTokens,
+            this.capacity
+        );
+        userTokens.lastTime = now;
+
+        // 判断是否有足够的令牌
+        if (count <= userTokens.tokens) {
+            userTokens.tokens -= count;
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    /**
+     * 重置令牌
+     * @param newCapacity
+     */
+    resetCapacity(newCapacity) {
+        if (newCapacity >= this.tokens) {
+            this.capacity = newCapacity;
+            this.tokens = newCapacity;
+        } else {
+            throw new Error('New capacity cannot be less than current tokens');
+        }
+    }
+}