Commit 8af6f27f authored by nanahira's avatar nanahira

add lock to account add

parent 7ff2e7ab
Pipeline #20232 passed with stages
in 4 minutes and 38 seconds
......@@ -11,6 +11,7 @@ import { AccountState } from './account-state';
import { AccountProviderService } from '../account-provider/account-provider.service';
import { AccountPoolStatusDto } from './account-pool-status.dto';
import { Interval } from '@nestjs/schedule';
import BetterLock from 'better-lock';
@Injectable()
export class AccountPoolService
......@@ -31,6 +32,9 @@ export class AccountPoolService
account: OpenAIAccount,
retry = true,
): Promise<boolean> {
if (this.accounts.has(account.email)) {
return false;
}
const proxy = await this.proxyPoolService.getProxy();
const state = new AccountState(
account,
......@@ -43,18 +47,18 @@ export class AccountPoolService
proxyServer: proxy,
}),
);
this.log(`Initializing account ${account.email}`);
this.log(`Adding account ${account.email}`);
const success = await state.init();
if (success) {
this.accounts.set(account.email, state);
this.log(`Initialized account ${account.email}`);
this.log(`Added account ${account.email}`);
return true;
} else if (retry) {
await Promise.all([
new Promise((r) => setTimeout(r, 5000)),
state.close(),
]);
return this.initAccount(account, retry);
return await this.initAccount(account, retry);
} else {
await state.close();
return false;
......@@ -87,7 +91,7 @@ export class AccountPoolService
await eval("import('chatgpt3')")
).ChatGPTAPIBrowser;
this.accountInfos = await this.accountProvider.get();
this.accountInfos.forEach((a) => this.initAccount(a));
this.accountInfos.forEach((a) => this.addAccount(a, true));
}
async onModuleDestroy() {
......@@ -130,11 +134,12 @@ export class AccountPoolService
return this.accounts.get(email) || this.randomAccount(exclude);
}
async addAccount(account: OpenAIAccount) {
if (this.accounts.has(account.email)) {
return false;
}
return this.initAccount(account, false);
private addAccountLock = new BetterLock();
async addAccount(account: OpenAIAccount, retry = false) {
return this.addAccountLock.acquire(account.email, () =>
this.initAccount(account, retry),
);
}
async removeAccount(email: string) {
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment