From f8e576adb72d59ffd9bbaa3e745e47c4bc9f5207 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=A2=AB=E9=81=97=E5=BF=98=E7=9A=84=E8=AE=B0=E5=BF=86?= <32097237+YIWANG-sketch@users.noreply.github.com> Date: Thu, 12 Dec 2024 00:56:46 +0800 Subject: [PATCH] fix lock logic --- sessionManager.mjs | 143 +++++++++++++++++++++++++-------------------- 1 file changed, 81 insertions(+), 62 deletions(-) diff --git a/sessionManager.mjs b/sessionManager.mjs index 9f96b43..e16ca0c 100644 --- a/sessionManager.mjs +++ b/sessionManager.mjs @@ -1,88 +1,107 @@ +import { Mutex } from 'async-mutex'; + class SessionManager { constructor(provider) { this.sessions = provider.sessions; this.provider = provider; - - // 初始化账号锁定状态 + this.isCustomModeEnabled = process.env.USE_CUSTOM_MODE === 'true'; + this.isRotationEnabled = process.env.ENABLE_MODE_ROTATION === 'true'; + this.currentIndex = 0; + // 缓存用户名列表 + this.usernameList = Object.keys(this.sessions); for (const username in this.sessions) { - this.sessions[username].locked = false; + const session = this.sessions[username]; + session.locked = false; + session.requestCount = 0; + session.valid = true; + session.mutex = new Mutex(); // 创建互斥锁 + if (session.currentMode === undefined) { + session.currentMode = this.isCustomModeEnabled ? 'custom' : 'default'; + } + if (!session.modeStatus) { + session.modeStatus = { + default: true, + custom: true, + }; + } + session.rotationEnabled = true; + session.switchCounter = 0; + session.requestsInCurrentMode = 0; + session.lastDefaultThreshold = 0; + session.switchThreshold = this.provider.getRandomSwitchThreshold(session); } } // 获取所有可用且未锁定会话 - getAvailableSessions() { - if (this.provider && this.provider.currentMode) { - const currentMode = this.provider.currentMode; + async getAvailableSessions() { + const allSessionsLocked = this.usernameList.every( + (username) => this.sessions[username].locked + ); + if (allSessionsLocked) { + console.warn('所有会话都已被锁定，等待释放...'); + throw new Error('所有会话都已被锁定'); + } - let availableSessions = Object.keys(this.sessions).filter(username => { - const session = this.sessions[username]; - return session.valid && !session.locked && session.modeStatus && session.modeStatus[currentMode]; - }); + // 轮询选择下一个可用 + const totalSessions = this.usernameList.length; + for (let i = 0; i < totalSessions; i++) { + const index = (this.currentIndex + i) % totalSessions; + const username = this.usernameList[index]; + const session = this.sessions[username]; - if (availableSessions.length === 0) { - console.warn(`在模式 [${currentMode}] 下没有可用的会话。`); - - if (this.provider && typeof this.provider.switchMode === 'function') { - console.warn(`尝试切换模式...`); - this.provider.switchMode(); - const newMode = this.provider.currentMode; - - availableSessions = Object.keys(this.sessions).filter(username => { - const session = this.sessions[username]; - return session.valid && !session.locked && session.modeStatus && session.modeStatus[newMode]; - }); - - if (availableSessions.length === 0) { - // 返回所有可用且未锁定会话 - availableSessions = Object.keys(this.sessions).filter(username => { - const session = this.sessions[username]; - return session.valid && !session.locked; - }); + // 检查是否有效且未锁定 + if (session.valid && !session.locked) { + const result = await session.mutex.runExclusive(async () => { + // 再次检查锁定状态 + if (session.locked) { + return null; } - } else { - console.warn('提供者没有 switchMode 方法'); - availableSessions = Object.keys(this.sessions).filter(username => { - const session = this.sessions[username]; - return session.valid && !session.locked; - }); + // 检查模式状态 + if (session.modeStatus && session.modeStatus[session.currentMode]) { + session.locked = true; + session.requestCount++; + this.currentIndex = (index + 1) % totalSessions; + return { selectedUsername: username, modeSwitched: false }; + } else if (this.isCustomModeEnabled && this.isRotationEnabled && this.provider && typeof this.provider.switchMode === 'function') { + console.warn(`尝试为账号 ${username} 切换模式...`); + this.provider.switchMode(session); + session.rotationEnabled = false; + + if (session.modeStatus && session.modeStatus[session.currentMode]) { + session.locked = true; + session.requestCount++; + this.currentIndex = (index + 1) % totalSessions; + return { selectedUsername: username, modeSwitched: true }; + } + } + return null; + }); + + if (result) { + // 成功锁定会话，返回 + return result; } } - - return availableSessions; - } else { - console.warn('提供者没有 currentMode 属性'); - - return Object.keys(this.sessions).filter(username => { - const session = this.sessions[username]; - return session.valid && !session.locked; - }); } - } - // 随机选择一个可用会话 - getRandomSession() { - const availableSessions = this.getAvailableSessions(); - if (availableSessions.length === 0) { - throw new Error('没有可用的会话'); - } - const randomIndex = Math.floor(Math.random() * availableSessions.length); - const selectedUsername = availableSessions[randomIndex]; - // 锁定账号 - this.sessions[selectedUsername].locked = true; - return selectedUsername; + console.warn('没有可用的会话'); + throw new Error('没有可用的会话'); } // 释放账号锁定 - releaseSession(username) { - if (this.sessions[username]) { - this.sessions[username].locked = false; + async releaseSession(username) { + const session = this.sessions[username]; + if (session) { + await session.mutex.runExclusive(() => { + session.locked = false; + }); } } // 策略 - getSessionByStrategy(strategy = 'random') { - if (strategy === 'random') { - return this.getRandomSession(); + async getSessionByStrategy(strategy = 'round_robin') { + if (strategy === 'round_robin') { + return await this.getAvailableSessions(); } throw new Error(`未实现的策略: ${strategy}`); }