fix(deployment): find orphaned deployments directly so the sweep can close them in minutes

.helm/console-api-prod-mainnet-values.yaml

@@ -19,8 +19,11 @@ jobs:
       - ./dist/instrumentation.js
       - ./dist/console.js
       - refill-wallets
+  # Start minutes are offset off every other sweep's, mint-act's */10 included, so no two jobs start on the same minute.
   - name: cleanup-stale-deployments
-    schedule: "30 10,22 * * *" # 10:30 AM and 10:30 PM every day
+    schedule: "4,14,24,34,44,54 * * * *" # every 10 minutes
+    concurrencyPolicy: Forbid
+    activeDeadlineSeconds: 540
     command:
       - node
       - --require

.helm/console-api-staging-sandbox-values.yaml

@@ -17,8 +17,11 @@ jobs:
       - ./dist/instrumentation.js
       - ./dist/console.js
       - refill-wallets
+  # Start minutes are offset off every other sweep's so no two jobs start on the same minute.
   - name: cleanup-stale-deployments
-    schedule: "30 10 * * 2" # 10:30 AM every Tuesday
+    schedule: "4,14,24,34,44,54 * * * *" # every 10 minutes
+    concurrencyPolicy: Forbid
+    activeDeadlineSeconds: 540
     command:
       - node
       - --require

apps/api/src/app/console.ts

@@ -56,10 +56,11 @@ program
 program
   .command("cleanup-stale-deployments")
   .description("Close deployments without leases created at least 10min ago")
-  .option("-c, --concurrency <number>", "How many wallets is processed concurrently", value => z.number({ coerce: true }).optional().default(10).parse(value))
+  .option("-c, --concurrency <number>", "How many owners' closes run concurrently", value => z.number({ coerce: true }).optional().default(10).parse(value))
+  .option("-d, --dry-run", "Log which deployments would be closed without broadcasting", false)
   .action(async (options, command) => {
     await executeCliHandler(command.name(), async () => {
-      await container.resolve(TopUpDeploymentsController).cleanUpStaleDeployment(options);
+      return container.resolve(TopUpDeploymentsController).cleanUpStaleDeployment(options);
     });
   });
 

apps/api/src/billing/repositories/user-wallet/user-wallet.repository.integration.ts

@@ -314,6 +314,70 @@ describe(UserWalletRepository.name, () => {
     });
   });
 
+  describe("findManagedIteratively", () => {
+    it("yields every wallet that holds an address, in id order across batches", async () => {
+      const { createWallet, collectYielded } = await setupManagedWallets();
+      const first = await createWallet();
+      const second = await createWallet();
+      const third = await createWallet();
+
+      const yielded = await collectYielded([first, second, third], 2);
+
+      expect(yielded).toEqual([
+        { id: first.id, address: first.address },
+        { id: second.id, address: second.address },
+        { id: third.id, address: third.address }
+      ]);
+    });
+
+    it("leaves out a wallet that has no address yet, since it owns nothing on chain", async () => {
+      const { createWallet, collectYielded } = await setupManagedWallets();
+      const addressless = await createWallet({ withAddress: false });
+      const managed = await createWallet();
+
+      const yielded = await collectYielded([addressless, managed], 10);
+
+      expect(yielded).toEqual([{ id: managed.id, address: managed.address }]);
+    });
+
+    it("caps every batch at the batch size", async () => {
+      const { userWalletRepository, createWallet } = await setupManagedWallets();
+      await createWallet();
+      await createWallet();
+      await createWallet();
+      const sizes: number[] = [];
+
+      for await (const batch of userWalletRepository.findManagedIteratively({ batchSize: 2 })) {
+        sizes.push(batch.length);
+      }
+
+      expect(Math.max(...sizes)).toBeLessThanOrEqual(2);
+    });
+  });
+
+  async function setupManagedWallets() {
+    const userRepository = container.resolve(UserRepository);
+    const userWalletRepository = container.resolve(UserWalletRepository);
+
+    async function createWallet(input: { withAddress?: boolean } = {}) {
+      const user = await userRepository.create({ userId: faker.string.uuid() });
+      return await userWalletRepository.create({ userId: user.id, address: input.withAddress === false ? undefined : createAkashAddress() });
+    }
+
+    async function collectYielded(wallets: { id: number }[], batchSize: number) {
+      const ids = new Set(wallets.map(wallet => wallet.id));
+      const yielded = [];
+
+      for await (const batch of userWalletRepository.findManagedIteratively({ batchSize })) {
+        yielded.push(...batch.filter(wallet => ids.has(wallet.id)));
+      }
+
+      return yielded;
+    }
+
+    return { userWalletRepository, createWallet, collectYielded };
+  }
+
   async function setupDomain() {
     const userRepository = container.resolve(UserRepository);
     const userWalletRepository = container.resolve(UserWalletRepository);

apps/api/src/billing/repositories/user-wallet/user-wallet.repository.ts

@@ -1,6 +1,6 @@
 import { Trace } from "@akashnetwork/instrumentation";
 import subDays from "date-fns/subDays";
-import { and, count, eq, gt, inArray, isNotNull, isNull, lte, ne, notExists, or, sql } from "drizzle-orm";
+import { and, asc, count, eq, gt, inArray, isNotNull, isNull, lte, ne, notExists, or, sql } from "drizzle-orm";
 import { singleton } from "tsyringe";
 
 import { StripeTransactions } from "@src/billing/model-schemas";
@@ -30,6 +30,12 @@ export function isWalletInitialized(wallet: UserWalletOutput): wallet is WalletI
   return !!wallet.address;
 }
 
+/** All a sweep needs to close a deployment on a wallet's behalf: the address it owns on chain and the index the signer derives it from. */
+export interface ManagedWalletRef {
+  id: UserWalletOutput["id"];
+  address: string;
+}
+
 export interface TrialWindow {
   trialEndsAt: Date | null;
   trialDurationDays: number | null;
@@ -243,6 +249,27 @@ export class UserWalletRepository extends BaseRepository<ApiPgTables["UserWallet
     return this.toOutputList(await this.cursor.query.UserWallets.findMany({ where: this.whereAccessibleBy(inArray(this.table.address, addresses)) }));
   }
 
+  /** Keyset-paged on the primary key and projected to what a close needs, because a sweep reads every managed wallet to find the few that own an orphan. */
+  async *findManagedIteratively({ batchSize }: { batchSize: number }): AsyncGenerator<ManagedWalletRef[]> {
+    let cursor: number | undefined;
+
+    while (true) {
+      const batch = await this.cursor
+        .select({ id: this.table.id, address: this.table.address })
+        .from(this.table)
+        .where(this.whereAccessibleBy(and(isNotNull(this.table.address), ...(cursor === undefined ? [] : [gt(this.table.id, cursor)]))))
+        .orderBy(asc(this.table.id))
+        .limit(batchSize);
+
+      if (!batch.length) return;
+
+      yield batch as ManagedWalletRef[];
+
+      if (batch.length < batchSize) return;
+      cursor = batch[batch.length - 1].id;
+    }
+  }
+
   async findFirst() {
     return this.findOneBy();
   }

apps/api/src/deployment/controllers/deployment/top-up-deployments.controller.spec.ts

@@ -69,7 +69,7 @@ describe(TopUpDeploymentsController.name, () => {
   describe("cleanUpStaleDeployment", () => {
     it("should call the service to clean up stale deployments", async () => {
       const { controller, staleDeploymentsCleanerService } = setup();
-      const options = { concurrency: 5 };
+      const options = { concurrency: 5, dryRun: false };
 
       await controller.cleanUpStaleDeployment(options);
 

apps/api/src/deployment/controllers/deployment/top-up-deployments.controller.ts

@@ -30,7 +30,7 @@ export class TopUpDeploymentsController {
   }
 
   async cleanUpStaleDeployment(options: CleanUpStaleDeploymentsParams) {
-    await this.staleDeploymentsCleanerService.cleanup(options);
+    return await this.staleDeploymentsCleanerService.cleanup(options);
   }
 
   async notifyExpiringDeployments(options: DryRunOptions) {

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

@@ -186,7 +186,7 @@ describe(DeploymentRepository.name, () => {
       const { repository, owner, base } = setup();
       const deployment = await seedOpenDeployment(owner, { createdHeight: base - 1 });
 
-      const found = await repository.findStaleDeployments({ owner, staleBeforeHeight: base });
+      const found = await repository.findStaleDeployments({ owners: [owner], staleBeforeHeight: base });
 
       expect(found.map(stale => String(stale.dseq))).toEqual([deployment.dseq]);
     });
@@ -196,7 +196,7 @@ describe(DeploymentRepository.name, () => {
       const deployment = await seedOpenDeployment(owner, { createdHeight: base - 1_000 });
       await seedLease(deployment, { closedHeight: base - 1 });
 
-      const found = await repository.findStaleDeployments({ owner, staleBeforeHeight: base });
+      const found = await repository.findStaleDeployments({ owners: [owner], staleBeforeHeight: base });
 
       expect(found.map(stale => String(stale.dseq))).toEqual([deployment.dseq]);
     });
@@ -207,7 +207,7 @@ describe(DeploymentRepository.name, () => {
       await seedLease(deployment, { gseq: 1, closedHeight: base - 500 });
       await seedLease(deployment, { gseq: 2, closedHeight: base });
 
-      const found = await repository.findStaleDeployments({ owner, staleBeforeHeight: base });
+      const found = await repository.findStaleDeployments({ owners: [owner], staleBeforeHeight: base });
 
       expect(found).toEqual([]);
     });
@@ -218,7 +218,7 @@ describe(DeploymentRepository.name, () => {
       await seedLease(deployment, { gseq: 1, closedHeight: base - 500 });
       await seedLease(deployment, { gseq: 2 });
 
-      const found = await repository.findStaleDeployments({ owner, staleBeforeHeight: base });
+      const found = await repository.findStaleDeployments({ owners: [owner], staleBeforeHeight: base });
 
       expect(found).toEqual([]);
     });
@@ -227,7 +227,7 @@ describe(DeploymentRepository.name, () => {
       const { repository, owner, base } = setup();
       await seedOpenDeployment(owner, { createdHeight: base });
 
-      const found = await repository.findStaleDeployments({ owner, staleBeforeHeight: base });
+      const found = await repository.findStaleDeployments({ owners: [owner], staleBeforeHeight: base });
 
       expect(found).toEqual([]);
     });
@@ -236,7 +236,7 @@ describe(DeploymentRepository.name, () => {
       const { repository, owner, base } = setup();
       await seedOpenDeployment(owner, { createdHeight: base - 1_000, closedHeight: base - 500 });
 
-      const found = await repository.findStaleDeployments({ owner, staleBeforeHeight: base });
+      const found = await repository.findStaleDeployments({ owners: [owner], staleBeforeHeight: base });
 
       expect(found).toEqual([]);
     });
@@ -245,10 +245,40 @@ describe(DeploymentRepository.name, () => {
       const { repository, owner, base } = setup();
       await seedOpenDeployment(createAkashAddress(), { createdHeight: base - 1 });
 
-      const found = await repository.findStaleDeployments({ owner, staleBeforeHeight: base });
+      const found = await repository.findStaleDeployments({ owners: [owner], staleBeforeHeight: base });
 
       expect(found).toEqual([]);
     });
+
+    it("returns the orphans of every owner it is given, each labelled with its owner", async () => {
+      const { repository, owner, base } = setup();
+      const other = createAkashAddress();
+      const mine = await seedOpenDeployment(owner, { createdHeight: base - 1 });
+      const theirs = await seedOpenDeployment(other, { createdHeight: base - 1 });
+
+      const found = await repository.findStaleDeployments({ owners: [owner, other], staleBeforeHeight: base });
+
+      expect(found).toHaveLength(2);
+      expect(found).toContainEqual({ owner, dseq: mine.dseq });
+      expect(found).toContainEqual({ owner: other, dseq: theirs.dseq });
+    });
+
+    it("returns a deployment once however many of its leases closed before the cutoff", async () => {
+      const { repository, owner, base } = setup();
+      const deployment = await seedOpenDeployment(owner, { createdHeight: base - 1_000 });
+      await seedLease(deployment, { gseq: 1, closedHeight: base - 500 });
+      await seedLease(deployment, { gseq: 2, closedHeight: base - 400 });
+
+      const found = await repository.findStaleDeployments({ owners: [owner], staleBeforeHeight: base });
+
+      expect(found).toEqual([{ owner, dseq: deployment.dseq }]);
+    });
+
+    it("queries nothing when given no owners", async () => {
+      const { repository, base } = setup();
+
+      await expect(repository.findStaleDeployments({ owners: [], staleBeforeHeight: base })).resolves.toEqual([]);
+    });
   });
 
   describe("countActiveByOwner", () => {

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

@@ -8,7 +8,7 @@ import { CHAIN_DB } from "@src/chain";
 
 export interface StaleDeploymentsOptions {
   staleBeforeHeight: number;
-  owner: string;
+  owners: string[];
 }
 
 export interface ProviderCleanupOptions {
@@ -27,7 +27,7 @@ export interface DeploymentActivityWindow {
   endDate: string;
 }
 
-/** Heights reach the query through a literal, so anything but a plain integer is refused before it can be interpolated. */
+/** A missing height would bind as NULL and match nothing, so anything but a plain integer is refused before the query runs. */
 function asHeight(value: number): number {
   if (!Number.isSafeInteger(value)) {
     throw new TypeError(`Expected a block height, received ${value}`);
@@ -52,7 +52,11 @@ export interface DatabaseDeploymentListParams {
 }
 
 export interface StaleDeploymentsOutput {
-  dseq: number;
+  dseq: string;
+}
+
+export interface StaleDeployment extends StaleDeploymentsOutput {
+  owner: string;
 }
 
 export interface DeploymentKey {
@@ -117,32 +121,29 @@ export class DeploymentRepository {
     });
   }
 
-  async findStaleDeployments(options: StaleDeploymentsOptions): Promise<StaleDeploymentsOutput[]> {
-    const deployments = await Deployment.findAll({
-      attributes: ["dseq"],
-      include: [
-        {
-          model: Lease,
-          attributes: [],
-          required: false
-        }
-      ],
-      where: {
-        owner: options.owner,
-        createdHeight: {
-          [Op.lt]: options.staleBeforeHeight
-        },
-        closedHeight: null
-      },
-      group: ["deployment.dseq"],
-      having: literal(
-        `COUNT("leases"."deploymentId") FILTER (WHERE "leases"."closedHeight" IS NULL) = 0 ` +
-          `AND COALESCE(MAX("leases"."closedHeight"), 0) < ${asHeight(options.staleBeforeHeight)}`
-      ),
-      raw: true
-    });
+  /** Owners reach the query as one array rather than one call each, so a sweep costs a query per batch instead of a query per wallet. */
+  async findStaleDeployments(options: StaleDeploymentsOptions): Promise<StaleDeployment[]> {
+    if (options.owners.length === 0) return [];
 
-    return deployments ? (deployments as unknown as StaleDeploymentsOutput[]) : [];
+    const staleBeforeHeight = asHeight(options.staleBeforeHeight);
+
+    return await this.#chainDb.query<StaleDeployment>(
+      `/* deployment:staleByOwners */
+      SELECT d."owner", d."dseq"
+      FROM deployment d
+      JOIN unnest($1::text[]) AS o(owner) ON d."owner" = o.owner
+      WHERE d."closedHeight" IS NULL
+        AND d."createdHeight" < $2
+        AND NOT EXISTS (
+          SELECT 1 FROM lease live
+          WHERE live."deploymentId" = d."id" AND live."closedHeight" IS NULL
+        )
+        AND COALESCE((SELECT MAX(last."closedHeight") FROM lease last WHERE last."deploymentId" = d."id"), 0) < $2`,
+      {
+        bind: [options.owners, staleBeforeHeight],
+        type: QueryTypes.SELECT
+      }
+    );
   }
 
   async findDeploymentsForProvider(options: ProviderCleanupOptions): Promise<StaleDeploymentsOutput[]> {