From f1e21accbeae01136b0b09e43287c9cc85ab839b Mon Sep 17 00:00:00 2001 From: zhiyu1998 Date: Thu, 13 Apr 2023 11:27:43 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=88=20perf:=20=E7=AE=80=E5=8C=96?= =?UTF-8?q?=E9=99=90=E6=B5=81=E7=AE=97=E6=B3=95=20&=20=E2=9C=A8=20feat:=20?= =?UTF-8?q?=E5=A2=9E=E5=BC=BA=E5=93=94=E5=93=A9=E5=93=94=E5=93=A9=E8=A7=A3?= =?UTF-8?q?=E6=9E=90=E5=81=A5=E5=A3=AE=E6=80=A7=EF=BC=8C=E9=98=B2=E6=AD=A2?= =?UTF-8?q?=E9=AB=98=E5=B9=B6=E5=8F=91=E7=8E=AF=E5=A2=83=E4=B8=8B=E8=A7=A3?= =?UTF-8?q?=E6=9E=90=E8=87=B4=E4=BD=BF=E6=9C=8D=E5=8A=A1=E5=99=A8=E5=AE=95?= =?UTF-8?q?=E6=9C=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/query.js | 6 +--- apps/tools.js | 26 ++++++++++++++++ utils/token-bucket.js | 69 ++++++++++++++++++++++++------------------- 3 files changed, 65 insertions(+), 36 deletions(-) diff --git a/apps/query.js b/apps/query.js index 7e4f2c8..7a208cd 100644 --- a/apps/query.js +++ b/apps/query.js @@ -57,10 +57,6 @@ export class query extends plugin { reg: "^#竹白(.*)", fnc: "zhubaiSearch", }, - { - reg: "^#测试", - fnc: "test1", - }, ], }); } @@ -438,5 +434,5 @@ export class query extends plugin { * 令牌桶 拿来限流 * @type {TokenBucket} */ - static #tokenBucket = new TokenBucket(1, 1); + static #tokenBucket = new TokenBucket(1, 1, 60); } diff --git a/apps/tools.js b/apps/tools.js index c82d5b1..d9cdf52 100644 --- a/apps/tools.js +++ b/apps/tools.js @@ -20,6 +20,7 @@ import { getBodianAudio, getBodianMv, getBodianMusicInfo } from "../utils/bodian import { ChatGPTBrowserClient } from "@waylaidwanderer/chatgpt-api"; import { av2BV } from "../utils/bilibili-bv-av-convert.js"; import querystring from "querystring"; +import TokenBucket from "../utils/token-bucket.js"; export class tools extends plugin { constructor() { @@ -326,6 +327,11 @@ export class tools extends plugin { // bilibi解析 async bili(e) { + await this.limitUserUse(e, async () => { + await this.biliCore(e); + }); + } + async biliCore(e) { const urlRex = /(?:https?:\/\/)?www\.bilibili\.com\/[A-Za-z\d._?%&+\-=\/#]*/g; const bShortRex = /(http:|https:)\/\/b23.tv\/[A-Za-z\d._?%&+\-=\/#]*/g; let url = e.msg === undefined ? e.message.shift().data.replaceAll("\\", "") : e.msg.trim(); @@ -1210,6 +1216,20 @@ export class tools extends plugin { } } + /** + * 限制用户调用（默认1分钟1次） + * @param e + * @param func + * @return {Promise} + */ + async limitUserUse(e, func) { + if (tools.#tokenBucket.consume(e.user_id, 1)) { + await func(); + } else { + logger.warn(`解析被限制使用`, true); + } + } + /** * 构造安全的命令 * @type {{existsPromptKey: string, existsTransKey: string}} @@ -1218,4 +1238,10 @@ export class tools extends plugin { existsTransKey: Object.keys(transMap).join("|"), existsPromptKey: Object.keys(PROMPT_MAP).join("|").slice(0, -1), }; + + /** + * 构造令牌桶，防止解析致使服务器宕机（默认限制5s调用一次） + * @type {TokenBucket} + */ + static #tokenBucket = new TokenBucket(1, 1, 5); } diff --git a/utils/token-bucket.js b/utils/token-bucket.js index 156370c..1877d6e 100644 --- a/utils/token-bucket.js +++ b/utils/token-bucket.js @@ -1,10 +1,32 @@ export default class TokenBucket { - constructor(rate, capacity) { - this.rate = rate / 60; // 修改为每分钟生成的令牌数量 - this.capacity = capacity; - this.tokens = capacity; + constructor(rate, capacity, interval = 1, isMinute = false) { + this.interval = interval; // 生成令牌的时间间隔 + this.rate = isMinute ? rate / 60 : rate; // 修改为每分钟生成的令牌数量 + this.capacity = capacity; // 令牌容量 + this.tokens = capacity; // 令牌容量 this.tokens = new Map(); // 使用 Map 存储每个用户的令牌桶 - this.lastTime = new Date().getTime(); + this.lastTime = new Date().getTime(); // 上次使用时间 + + /** + * 核心算法 + * @param tokens + * @param capacity + * @param rate + * @param lastTime + * @param interval + * @param isMinute + * @return {{lastTime: number, tokens: number}} + */ + this.updateTokens = (tokens, capacity, rate, lastTime, interval) => { + // 计算从上次请求到现在经过的时间 + const now = new Date().getTime(); + const elapsed = now - lastTime; + // 根据时间计算出新生成的令牌数量 + const addedTokens = elapsed * (rate / interval / 1000); // 修改为每分钟生成的令牌数量 + tokens = Math.min(tokens + addedTokens, capacity); + lastTime = now; + return { tokens, lastTime }; + } } /** @@ -13,12 +35,11 @@ export default class TokenBucket { * @return {boolean} */ consumeSingle(count = 1) { - const now = new Date().getTime(); - const elapsed = now - this.lastTime; - const addedTokens = elapsed * (this.rate / 1000 / 60); // 修改为每分钟生成的令牌数量 - this.tokens = Math.min(this.tokens + addedTokens, this.capacity); - this.lastTime = now; + const { tokens, lastTime } = this.updateTokens(this.tokens, this.capacity, this.rate, this.lastTime, this.interval); + // 更新令牌桶中的令牌数量 + this.tokens = tokens; + // 判断请求是否能够被处理（即令牌桶中是否有足够的令牌） if (count <= this.tokens) { this.tokens -= count; return true; // 返回 true 表示请求被处理 @@ -34,28 +55,14 @@ export default class TokenBucket { * @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; + const { tokens: userTokens, lastTime: userLastTime } = this.tokens.get(id) || { tokens: this.capacity, lastTime: new Date().getTime() }; + const { tokens, lastTime } = this.updateTokens(userTokens, this.capacity, this.rate, userLastTime, this.interval); + // 更新令牌桶中的令牌数量 + this.tokens.set(id, { tokens, lastTime }); - // 获取用户的令牌桶，如果不存在则创建一个新的令牌桶 - 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; + // 判断请求是否能够被处理（即令牌桶中是否有足够的令牌） + if (count <= tokens) { + this.tokens.set(id, { tokens: tokens - count, lastTime }); return true; } else { return false;