feat(wallet): use smaphore to control flow for init trial

akash-network · Dec 14, 2024 · 63a04dc · 63a04dc
1 parent bf4b4da
commit 63a04dc
Showing 1 changed file with 47 additions and 15 deletions.
diff --git a/apps/api/src/billing/services/wallet-initializer/wallet-initializer.service.ts b/apps/api/src/billing/services/wallet-initializer/wallet-initializer.service.ts
@@ -3,34 +3,66 @@ import { singleton } from "tsyringe";
 import { AuthService } from "@src/auth/services/auth.service";
 import { UserWalletInput, UserWalletRepository } from "@src/billing/repositories";
 import { ManagedUserWalletService } from "@src/billing/services";
+import { Sema } from "async-sema";
 
 @singleton()
 export class WalletInitializerService {
+  private readonly semaphores = new Map<string, Sema>();
+
   constructor(
     private readonly walletManager: ManagedUserWalletService,
     private readonly userWalletRepository: UserWalletRepository,
     private readonly authService: AuthService
   ) {}
 
-  async initializeAndGrantTrialLimits(userId: UserWalletInput["userId"]) {
-    let userWallet = await this.userWalletRepository.findOneByUserId(userId);
+  private getSemaphore(userId: string): Sema {
+    let semaphore = this.semaphores.get(userId);
+    if (!semaphore) {
+      semaphore = new Sema(1);
+      this.semaphores.set(userId, semaphore);
+    }
+    return semaphore;
+  }
+
+  private async waitForCompletion(userId: string): Promise<void> {
+    const semaphore = this.getSemaphore(userId);
 
-    if (!userWallet) {
-      userWallet = await this.userWalletRepository.accessibleBy(this.authService.ability, "create").create({ userId });
+    while (!(await semaphore.tryAcquire())) {
+      await new Promise(resolve => setTimeout(resolve, 100));
     }
 
-    const wallet = await this.walletManager.createAndAuthorizeTrialSpending({ addressIndex: userWallet.id });
-    userWallet = await this.userWalletRepository.updateById(
-      userWallet.id,
-      {
-        address: wallet.address,
-        deploymentAllowance: wallet.limits.deployment,
-        feeAllowance: wallet.limits.fees
-      },
-      { returning: true }
-    );
+    semaphore.release();
+  }
+
+  async initializeAndGrantTrialLimits(userId: UserWalletInput["userId"]) {
+    // Wait for any existing operation to complete
+    await this.waitForCompletion(userId);
+
+    const semaphore = this.getSemaphore(userId);
+    await semaphore.acquire();
 
-    return this.userWalletRepository.toPublic(userWallet);
+    try {
+      let userWallet = await this.userWalletRepository.findOneByUserId(userId);
+
+      if (!userWallet) {
+        userWallet = await this.userWalletRepository.accessibleBy(this.authService.ability, "create").create({ userId });
+
+        const wallet = await this.walletManager.createAndAuthorizeTrialSpending({ addressIndex: userWallet.id });
+        userWallet = await this.userWalletRepository.updateById(
+          userWallet.id,
+          {
+            address: wallet.address,
+            deploymentAllowance: wallet.limits.deployment,
+            feeAllowance: wallet.limits.fees
+          },
+          { returning: true }
+        );
+      }
+
+      return this.userWalletRepository.toPublic(userWallet);
+    } finally {
+      semaphore.release();
+    }
   }
 
   async initialize(userId: UserWalletInput["userId"]) {