🎈 perf: 简化限流算法 & feat: 增强哔哩哔哩解析健壮性,防止高并发环境下解析致使服务器宕机

This commit is contained in:
zhiyu1998 2023-04-13 11:27:43 +08:00
parent 1799a82dbd
commit f1e21accbe
3 changed files with 65 additions and 36 deletions

View File

@ -57,10 +57,6 @@ export class query extends plugin {
reg: "^#竹白(.*)", reg: "^#竹白(.*)",
fnc: "zhubaiSearch", fnc: "zhubaiSearch",
}, },
{
reg: "^#测试",
fnc: "test1",
},
], ],
}); });
} }
@ -438,5 +434,5 @@ export class query extends plugin {
* 令牌桶 拿来限流 * 令牌桶 拿来限流
* @type {TokenBucket} * @type {TokenBucket}
*/ */
static #tokenBucket = new TokenBucket(1, 1); static #tokenBucket = new TokenBucket(1, 1, 60);
} }

View File

@ -20,6 +20,7 @@ import { getBodianAudio, getBodianMv, getBodianMusicInfo } from "../utils/bodian
import { ChatGPTBrowserClient } from "@waylaidwanderer/chatgpt-api"; import { ChatGPTBrowserClient } from "@waylaidwanderer/chatgpt-api";
import { av2BV } from "../utils/bilibili-bv-av-convert.js"; import { av2BV } from "../utils/bilibili-bv-av-convert.js";
import querystring from "querystring"; import querystring from "querystring";
import TokenBucket from "../utils/token-bucket.js";
export class tools extends plugin { export class tools extends plugin {
constructor() { constructor() {
@ -326,6 +327,11 @@ export class tools extends plugin {
// bilibi解析 // bilibi解析
async bili(e) { 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 urlRex = /(?:https?:\/\/)?www\.bilibili\.com\/[A-Za-z\d._?%&+\-=\/#]*/g;
const bShortRex = /(http:|https:)\/\/b23.tv\/[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(); 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<void>}
*/
async limitUserUse(e, func) {
if (tools.#tokenBucket.consume(e.user_id, 1)) {
await func();
} else {
logger.warn(`解析被限制使用`, true);
}
}
/** /**
* 构造安全的命令 * 构造安全的命令
* @type {{existsPromptKey: string, existsTransKey: string}} * @type {{existsPromptKey: string, existsTransKey: string}}
@ -1218,4 +1238,10 @@ export class tools extends plugin {
existsTransKey: Object.keys(transMap).join("|"), existsTransKey: Object.keys(transMap).join("|"),
existsPromptKey: Object.keys(PROMPT_MAP).join("|").slice(0, -1), existsPromptKey: Object.keys(PROMPT_MAP).join("|").slice(0, -1),
}; };
/**
* 构造令牌桶防止解析致使服务器宕机默认限制5s调用一次
* @type {TokenBucket}
*/
static #tokenBucket = new TokenBucket(1, 1, 5);
} }

View File

@ -1,10 +1,32 @@
export default class TokenBucket { export default class TokenBucket {
constructor(rate, capacity) { constructor(rate, capacity, interval = 1, isMinute = false) {
this.rate = rate / 60; // 修改为每分钟生成的令牌数量 this.interval = interval; // 生成令牌的时间间隔
this.capacity = capacity; this.rate = isMinute ? rate / 60 : rate; // 修改为每分钟生成的令牌数量
this.tokens = capacity; this.capacity = capacity; // 令牌容量
this.tokens = capacity; // 令牌容量
this.tokens = new Map(); // 使用 Map 存储每个用户的令牌桶 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} * @return {boolean}
*/ */
consumeSingle(count = 1) { consumeSingle(count = 1) {
const now = new Date().getTime(); const { tokens, lastTime } = this.updateTokens(this.tokens, this.capacity, this.rate, this.lastTime, this.interval);
const elapsed = now - this.lastTime; // 更新令牌桶中的令牌数量
const addedTokens = elapsed * (this.rate / 1000 / 60); // 修改为每分钟生成的令牌数量 this.tokens = tokens;
this.tokens = Math.min(this.tokens + addedTokens, this.capacity);
this.lastTime = now;
// 判断请求是否能够被处理(即令牌桶中是否有足够的令牌)
if (count <= this.tokens) { if (count <= this.tokens) {
this.tokens -= count; this.tokens -= count;
return true; // 返回 true 表示请求被处理 return true; // 返回 true 表示请求被处理
@ -34,28 +55,14 @@ export default class TokenBucket {
* @return {boolean} * @return {boolean}
*/ */
consume(id, count = 1) { consume(id, count = 1) {
const now = new Date().getTime(); const { tokens: userTokens, lastTime: userLastTime } = this.tokens.get(id) || { tokens: this.capacity, lastTime: new Date().getTime() };
const elapsed = now - this.lastTime; const { tokens, lastTime } = this.updateTokens(userTokens, this.capacity, this.rate, userLastTime, this.interval);
const addedTokens = elapsed * (this.rate / 1000); // 更新令牌桶中的令牌数量
this.lastTime = now; this.tokens.set(id, { tokens, lastTime });
// 获取用户的令牌桶,如果不存在则创建一个新的令牌桶 // 判断请求是否能够被处理(即令牌桶中是否有足够的令牌)
let userTokens = this.tokens.get(id); if (count <= tokens) {
if (!userTokens) { this.tokens.set(id, { tokens: tokens - count, lastTime });
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; return true;
} else { } else {
return false; return false;