feat(deployment): clean up trial deployments for a provider

(#502)

Author: baktun14

apps/api/package.json

@@ -54,9 +54,9 @@
     "@hono/swagger-ui": "0.2.1",
     "@hono/zod-openapi": "0.9.5",
     "@octokit/rest": "^18.12.0",
-    "@opentelemetry/instrumentation": "^0.54.0",
-    "@opentelemetry/instrumentation-http": "^0.54.0",
-    "@opentelemetry/sdk-node": "^0.54.0",
+    "@opentelemetry/instrumentation": "^0.54.2",
+    "@opentelemetry/instrumentation-http": "^0.54.2",
+    "@opentelemetry/sdk-node": "^0.54.2",
     "@sentry/node": "^7.55.2",
     "@supercharge/promise-pool": "^3.2.0",
     "@ucast/core": "^1.10.2",

apps/api/src/billing/services/provider-cleanup/provider-cleanup.service.ts

@@ -0,0 +1,90 @@
+import { LoggerService } from "@akashnetwork/logging";
+import { singleton } from "tsyringe";
+
+import { BillingConfig, InjectBillingConfig } from "@src/billing/providers";
+import { UserWalletOutput, UserWalletRepository } from "@src/billing/repositories";
+import { ManagedUserWalletService, RpcMessageService } from "@src/billing/services";
+import { ErrorService } from "@src/core/services/error/error.service";
+import { ProviderCleanupSummarizer } from "@src/deployment/lib/provider-cleanup-summarizer/provider-cleanup-summarizer";
+import { DeploymentRepository } from "@src/deployment/repositories/deployment/deployment.repository";
+import { TxSignerService } from "../tx-signer/tx-signer.service";
+
+export interface ProviderCleanupParams {
+  concurrency: number;
+  providerAddress: string;
+  dryRun: boolean;
+}
+
+@singleton()
+export class ProviderCleanupService {
+  private readonly logger = LoggerService.forContext(ProviderCleanupService.name);
+
+  constructor(
+    @InjectBillingConfig() private readonly config: BillingConfig,
+    private readonly userWalletRepository: UserWalletRepository,
+    private readonly managedUserWalletService: ManagedUserWalletService,
+    private readonly txSignerService: TxSignerService,
+    private readonly deploymentRepository: DeploymentRepository,
+    private readonly rpcMessageService: RpcMessageService,
+    private readonly errorService: ErrorService
+  ) {}
+
+  async cleanup(options: ProviderCleanupParams) {
+    const summary = new ProviderCleanupSummarizer();
+    await this.userWalletRepository.paginate({ query: { isTrialing: true }, limit: options.concurrency || 10 }, async wallets => {
+      const cleanUpAllWallets = wallets.map(async wallet => {
+        await this.errorService.execWithErrorHandler(
+          {
+            wallet,
+            event: "PROVIDER_CLEAN_UP_ERROR",
+            context: ProviderCleanupService.name
+          },
+          () => this.cleanUpForWallet(wallet, options, summary)
+        );
+      });
+
+      await Promise.all(cleanUpAllWallets);
+    });
+
+    this.logger.info({ event: "PROVIDER_CLEAN_UP_SUMMARY", summary: summary.summarize(), dryRun: options.dryRun });
+  }
+
+  private async cleanUpForWallet(wallet: UserWalletOutput, options: ProviderCleanupParams, summary: ProviderCleanupSummarizer) {
+    const client = await this.txSignerService.getClientForAddressIndex(wallet.id);
+    const deployments = await this.deploymentRepository.findDeploymentsForProvider({
+      owner: wallet.address,
+      provider: options.providerAddress
+    });
+
+    const closeAllWalletStaleDeployments = deployments.map(async deployment => {
+      const message = this.rpcMessageService.getCloseDeploymentMsg(wallet.address, deployment.dseq);
+      this.logger.info({ event: "PROVIDER_CLEAN_UP", params: { owner: wallet.address, dseq: deployment.dseq } });
+
+      try {
+        if (!options.dryRun) {
+          await client.signAndBroadcast([message]);
+          this.logger.info({ event: "PROVIDER_CLEAN_UP_SUCCESS" });
+        }
+      } catch (error) {
+        if (error.message.includes("not allowed to pay fees")) {
+          if (!options.dryRun) {
+            await this.managedUserWalletService.authorizeSpending({
+              address: wallet.address,
+              limits: {
+                fees: this.config.FEE_ALLOWANCE_REFILL_AMOUNT
+              }
+            });
+            await client.signAndBroadcast([message]);
+            this.logger.info({ event: "PROVIDER_CLEAN_UP_SUCCESS" });
+          }
+        } else {
+          throw error;
+        }
+      } finally {
+        summary.inc("deploymentCount");
+      }
+    });
+
+    await Promise.all(closeAllWalletStaleDeployments);
+  }
+}

apps/api/src/console.ts

@@ -14,6 +14,7 @@ import { chainDb } from "@src/db/dbConnection";
 import { TopUpDeploymentsController } from "@src/deployment/controllers/deployment/deployment.controller";
 import { UserController } from "@src/user/controllers/user/user.controller";
 import { UserConfigService } from "@src/user/services/user-config/user-config.service";
+import { ProviderController } from "./deployment/controllers/provider/provider.controller";
 
 const program = new Command();
 
@@ -42,13 +43,25 @@ program
 program
   .command("cleanup-stale-deployments")
   .description("Close deployments without leases created at least 10min ago")
-  .option("-c, --concurrency <number>", "How much wallets is processed concurrently", value => z.number({ coerce: true }).optional().default(10).parse(value))
+  .option("-c, --concurrency <number>", "How many wallets is processed concurrently", value => z.number({ coerce: true }).optional().default(10).parse(value))
   .action(async (options, command) => {
     await executeCliHandler(command.name(), async () => {
       await container.resolve(TopUpDeploymentsController).cleanUpStaleDeployment(options);
     });
   });
 
+program
+  .command("cleanup-provider-deployments")
+  .description("Close trial deployments for a provider")
+  .option("-c, --concurrency <number>", "How many wallets is processed concurrently", value => z.number({ coerce: true }).optional().default(10).parse(value))
+  .option("-d, --dry-run", "Dry run the trial provider cleanup", false)
+  .option("-p, --provider <string>", "Provider address", value => z.string().parse(value))
+  .action(async (options, command) => {
+    await executeCliHandler(command.name(), async () => {
+      await container.resolve(ProviderController).cleanupProviderDeployments(options);
+    });
+  });
+
 const userConfig = container.resolve(UserConfigService);
 program
   .command("cleanup-stale-anonymous-users")

apps/api/src/deployment/controllers/provider/provider.controller.ts

@@ -1,12 +1,20 @@
 import { singleton } from "tsyringe";
 
+import { ProviderCleanupParams, ProviderCleanupService } from "@src/billing/services/provider-cleanup/provider-cleanup.service";
 import { TrialProvidersService } from "@src/deployment/services/trial-providers/trial-providers.service";
 
 @singleton()
 export class ProviderController {
-  constructor(private readonly trialProvidersService: TrialProvidersService) {}
+  constructor(
+    private readonly trialProvidersService: TrialProvidersService,
+    private readonly providerCleanupService: ProviderCleanupService
+  ) {}
 
   async getTrialProviders(): Promise<string[]> {
     return await this.trialProvidersService.getTrialProviders();
   }
+
+  async cleanupProviderDeployments(options: ProviderCleanupParams) {
+    return await this.providerCleanupService.cleanup(options);
+  }
 }

apps/api/src/deployment/lib/provider-cleanup-summarizer/provider-cleanup-summarizer.ts

@@ -0,0 +1,25 @@
+interface ProviderCleanupSummary {
+  deploymentCount: number;
+}
+
+export class ProviderCleanupSummarizer {
+  private deploymentCount = 0;
+
+  inc(param: keyof ProviderCleanupSummary, value = 1) {
+    this[param] += value;
+  }
+
+  set(param: keyof ProviderCleanupSummary, value: number) {
+    this[param] = value;
+  }
+
+  get(param: keyof ProviderCleanupSummary) {
+    return this[param];
+  }
+
+  summarize(): ProviderCleanupSummary {
+    return {
+      deploymentCount: this.deploymentCount
+    };
+  }
+}

apps/api/src/deployment/repositories/deployment/deployment.repository.ts

@@ -7,6 +7,11 @@ export interface StaleDeploymentsOptions {
   owner: string;
 }
 
+export interface ProviderCleanupOptions {
+  owner: string;
+  provider: string;
+}
+
 export interface StaleDeploymentsOutput {
   dseq: number;
 }
@@ -37,4 +42,27 @@ export class DeploymentRepository {
 
     return deployments ? (deployments as unknown as StaleDeploymentsOutput[]) : [];
   }
+
+  async findDeploymentsForProvider(options: ProviderCleanupOptions): Promise<StaleDeploymentsOutput[]> {
+    const deployments = await Deployment.findAll({
+      attributes: ["dseq"],
+      where: {
+        owner: options.owner,
+        closedHeight: null
+      },
+      include: [
+        {
+          model: Lease,
+          attributes: [],
+          required: true,
+          where: {
+            provider: options.provider
+          }
+        }
+      ],
+      raw: true
+    });
+
+    return deployments ? (deployments as unknown as StaleDeploymentsOutput[]) : [];
+  }
 }

apps/provider-console/package.json

@@ -22,7 +22,6 @@
     "@mui/icons-material": "^5.11.11",
     "@mui/material": "^5.4.4",
     "@mui/material-nextjs": "^5.15.11",
-    "@opentelemetry/instrumentation-lru-memoizer": "^0.41.0",
     "@radix-ui/react-icons": "^1.3.0",
     "@sentry/nextjs": "^8.34.0",
     "@sentry/tracing": "^7.114.0",