Merge branch 'master' into feat/grouper-slow-diagnostics

Author: Kuchizu

workers/limiter/src/dbHelper.ts

@@ -2,7 +2,7 @@ import { Collection, Db, ObjectId } from 'mongodb';
 import { ProjectDBScheme, WorkspaceDBScheme } from '@hawk.so/types';
 import { WorkspaceWithTariffPlan } from '../types';
 import HawkCatcher from '@hawk.so/nodejs';
-import { CriticalError } from '../../../lib/workerErrors';
+import { CriticalError, NonCriticalError } from '../../../lib/workerErrors';
 
 /**
  * Class that implements methods used for interaction between limiter and db
@@ -35,49 +35,24 @@ export class DbHelper {
   }
 
   /**
-   * Method that returns all workspaces with their tariff plans
+   * Method that yields all workspaces with their tariff plans
    */
-  public async getWorkspacesWithTariffPlans():Promise<WorkspaceWithTariffPlan[]>;
+  public getWorkspacesWithTariffPlans(): AsyncGenerator<WorkspaceWithTariffPlan>;
   /**
    * Method that returns workspace with its tariff plan by its id
    *
    * @param id - id of the workspace to fetch
    */
-  public async getWorkspacesWithTariffPlans(id: string):Promise<WorkspaceWithTariffPlan>;
+  public getWorkspacesWithTariffPlans(id: string): Promise<WorkspaceWithTariffPlan>;
   /**
-   * Returns workspace with its tariff plan by its id
-   *
-   * @param id - workspace id
+   * @param id - id of the workspace to fetch
    */
-  public async getWorkspacesWithTariffPlans(id?: string):Promise<WorkspaceWithTariffPlan[] | WorkspaceWithTariffPlan> {
-    /* eslint-disable-next-line */
-    const queue: any[] = [
-      {
-        $lookup: {
-          from: 'plans',
-          localField: 'tariffPlanId',
-          foreignField: '_id',
-          as: 'tariffPlan',
-        },
-      },
-      {
-        $unwind: {
-          path: '$tariffPlan',
-        },
-      },
-    ];
-
+  public getWorkspacesWithTariffPlans(id?: string): AsyncGenerator<WorkspaceWithTariffPlan> | Promise<WorkspaceWithTariffPlan> {
     if (id !== undefined) {
-      queue.unshift({
-        $match: {
-          _id: new ObjectId(id),
-        },
-      });
+      return this.getOneWorkspaceWithTariffPlan(id);
     }
 
-    const workspacesArray = await this.workspacesCollection.aggregate<WorkspaceWithTariffPlan>(queue).toArray();
-
-    return (id !== undefined) ? workspacesArray[0] : workspacesArray;
+    return this.yieldWorkspacesWithTariffPlans();
   }
 
   /**
@@ -172,4 +147,72 @@ export class DbHelper {
 
     return this.projectsCollection.find(query).toArray();
   }
+
+  /**
+   * Returns a single workspace with its tariff plan by id
+   *
+   * @param id - workspace id
+   */
+  private async getOneWorkspaceWithTariffPlan(id: string): Promise<WorkspaceWithTariffPlan> {
+    const pipeline = [
+      {
+        $match: {
+          _id: new ObjectId(id),
+        },
+      },
+      ...this.tariffPlanLookupPipeline(),
+    ];
+
+    const workspace = await this.workspacesCollection.aggregate<WorkspaceWithTariffPlan>(pipeline).next();
+
+    if (workspace === null) {
+      throw new NonCriticalError(`Workspace ${id} not found`, {
+        workspaceId: id,
+      });
+    }
+
+    return workspace;
+  }
+
+  /**
+   * Yields all workspaces with their tariff plans one by one
+   */
+  private async * yieldWorkspacesWithTariffPlans(): AsyncGenerator<WorkspaceWithTariffPlan> {
+    const pipeline = this.tariffPlanLookupPipeline();
+    const cursor = this.workspacesCollection.aggregate<WorkspaceWithTariffPlan>(pipeline);
+
+    for await (const workspace of cursor) {
+      yield workspace;
+    }
+  }
+
+  /* eslint-disable-next-line */
+  private tariffPlanLookupPipeline(): any[] {
+    return [
+      {
+        $lookup: {
+          from: 'plans',
+          localField: 'tariffPlanId',
+          foreignField: '_id',
+          as: 'tariffPlan',
+        },
+      },
+      {
+        $unwind: {
+          path: '$tariffPlan',
+        },
+      },
+      {
+        $project: {
+          _id: 1,
+          name: 1,
+          isBlocked: 1,
+          blockedDate: 1,
+          lastChargeDate: 1,
+          billingPeriodEventsCount: 1,
+          tariffPlan: 1,
+        },
+      },
+    ];
+  }
 }

workers/limiter/src/index.ts

@@ -189,16 +189,16 @@ export default class LimiterWorker extends Worker {
   private async handleRegularWorkspacesCheck(): Promise<void> {
     let message = '';
 
-    const workspaces = await this.dbHelper.getWorkspacesWithTariffPlans();
+    const workspaces = this.dbHelper.getWorkspacesWithTariffPlans();
 
     const updatedWorkspaces: WorkspaceWithTariffPlan[] = [];
 
-    await Promise.all(workspaces.map(async (workspace) => {
+    for await (const workspace of workspaces) {
       /**
        * If workspace is already blocked - do nothing
        */
       if (workspace.isBlocked) {
-        return;
+        continue;
       }
 
       const workspaceProjects = await this.dbHelper.getProjects(workspace._id.toString());
@@ -211,7 +211,7 @@ export default class LimiterWorker extends Worker {
        * If there are no projects to update - move on to next workspace
        */
       if (projectsToUpdate.length === 0) {
-        return;
+        continue;
       }
 
       /**
@@ -223,12 +223,12 @@ export default class LimiterWorker extends Worker {
         updatedWorkspace.isBlocked = true;
         updatedWorkspace.blockedDate = new Date();
 
-        this.redis.appendBannedProjects(projectIds);
+        await this.redis.appendBannedProjects(projectIds);
         message += this.formSingleWorkspaceMessage(updatedWorkspace, projectsToUpdate, 'blocked');
       }
-    }));
+    }
 
-    this.dbHelper.updateWorkspacesEventsCountAndIsBlocked(updatedWorkspaces);
+    await this.dbHelper.updateWorkspacesEventsCountAndIsBlocked(updatedWorkspaces);
 
     this.sendRegularReport(message);
   }

workers/limiter/tests/dbHelper.test.ts

@@ -159,16 +159,22 @@ describe('DbHelper', () => {
       /**
        * Act
        */
-      const result = await dbHelper.getWorkspacesWithTariffPlans();
+      const cursor = dbHelper.getWorkspacesWithTariffPlans();
+
+      const workspaces = [];
+
+      for await (const workspace of cursor) {
+        workspaces.push(workspace);
+      }
 
       /**
        * Assert
        */
-      expect(result).toHaveLength(2);
-      expect(result[0].tariffPlan).toBeDefined();
-      expect(result[1].tariffPlan).toBeDefined();
-      expect(result[0].tariffPlan.eventsLimit).toBe(10);
-      expect(result[1].tariffPlan.eventsLimit).toBe(10000);
+      expect(workspaces).toHaveLength(2);
+      expect(workspaces[0].tariffPlan).toBeDefined();
+      expect(workspaces[1].tariffPlan).toBeDefined();
+      expect(workspaces[0].tariffPlan.eventsLimit).toBe(10);
+      expect(workspaces[1].tariffPlan.eventsLimit).toBe(10000);
     });
 
     test('Should return single workspace with its tariff plan by id', async () => {

workers/limiter/types/index.ts

@@ -3,4 +3,9 @@ import { PlanDBScheme, WorkspaceDBScheme } from '@hawk.so/types';
 /**
  * Workspace with its tariff plan
  */
-export type WorkspaceWithTariffPlan = WorkspaceDBScheme & {tariffPlan: PlanDBScheme};
+export type WorkspaceWithTariffPlan = Pick<
+  WorkspaceDBScheme,
+  '_id' | 'name' | 'isBlocked' | 'blockedDate' | 'lastChargeDate' | 'billingPeriodEventsCount'
+> & {
+  tariffPlan: PlanDBScheme;
+};

workers/paymaster/src/index.ts

@@ -34,6 +34,33 @@ const DAYS_LEFT_ALERT = [3, 2, 1];
 // eslint-disable-next-line @typescript-eslint/no-magic-numbers
 const DAYS_AFTER_BLOCK_TO_REMIND = [1, 2, 3, 5, 7, 30];
 
+/**
+ * Bounds concurrent updateOne / addTask calls per subscription check tick.
+ */
+const WORKSPACE_PROCESSING_CONCURRENCY = 25;
+
+const WORKSPACE_CURSOR_BATCH_SIZE = 200;
+
+/**
+ * Keep in sync with fields read by `processWorkspaceSubscriptionCheck` and its helpers.
+ */
+const WORKSPACE_SUBSCRIPTION_PROJECTION = {
+  _id: 1,
+  name: 1,
+  tariffPlanId: 1,
+  lastChargeDate: 1,
+  paidUntil: 1,
+  isDebug: 1,
+  isBlocked: 1,
+  blockedDate: 1,
+  subscriptionId: 1,
+} as const;
+
+const PLAN_PROJECTION = {
+  _id: 1,
+  monthlyCharge: 1,
+} as const;
+
 /**
  * Worker to check workspaces subscription status and ban workspaces without actual subscription
  */
@@ -151,7 +178,10 @@ export default class PaymasterWorker extends Worker {
       throw new Error('Plans collection is not initialized');
     }
 
-    this.plans = await this.plansCollection.find({}).toArray();
+    this.plans = await this.plansCollection
+      .find({})
+      .project<PlanDBScheme>(PLAN_PROJECTION)
+      .toArray();
 
     if (this.plans.length === 0) {
       throw new Error('Please add tariff plans to the database');
@@ -195,28 +225,45 @@ export default class PaymasterWorker extends Worker {
    * Called periodically, enumerate through workspaces and check if today is a payday for workspace subscription
    */
   private async handleWorkspaceSubscriptionCheckEvent(): Promise<void> {
-    const workspaces = await this.workspaces.find({}).toArray();
+    const cursor = this.workspaces
+      .find({})
+      .project<WorkspaceDBScheme>(WORKSPACE_SUBSCRIPTION_PROJECTION)
+      .batchSize(WORKSPACE_CURSOR_BATCH_SIZE);
+
+    let batch: WorkspaceDBScheme[] = [];
 
-    await Promise.all(workspaces
-      .filter(workspace => {
+    const flush = async (currentBatch: WorkspaceDBScheme[]): Promise<void> => {
+      if (currentBatch.length === 0) {
+        return;
+      }
+
+      await Promise.all(currentBatch.map((workspace) => this.processWorkspaceSubscriptionCheck(workspace)));
+    };
+
+    try {
+      for await (const workspace of cursor) {
         /**
          * Skip workspace without lastChargeDate
          */
         if (!workspace.lastChargeDate) {
-          const error = new Error('[Paymaster] Workspace without lastChargeDate detected');
-
-          HawkCatcher.send(error, {
+          HawkCatcher.send(new Error('[Paymaster] Workspace without lastChargeDate detected'), {
             workspaceId: workspace._id.toString(),
           });
+          continue;
+        }
+
+        batch.push(workspace);
 
-          return false;
+        if (batch.length >= WORKSPACE_PROCESSING_CONCURRENCY) {
+          await flush(batch);
+          batch = [];
         }
+      }
 
-        return true;
-      })
-      .map(
-        (workspace) => this.processWorkspaceSubscriptionCheck(workspace)
-      ));
+      await flush(batch);
+    } finally {
+      await cursor.close();
+    }
   }
 
   /**

workers/paymaster/tests/index.test.ts

@@ -794,6 +794,54 @@ describe('PaymasterWorker', () => {
     MockDate.reset();
   });
 
+  test('Should process every workspace when there are several batches', async () => {
+    /**
+     * 50 > WORKSPACE_PROCESSING_CONCURRENCY (25), so the subscription check
+     * has to flush more than one batch.
+     */
+    const WORKSPACES_COUNT = 50;
+    const currentDate = new Date('2005-12-22');
+    const plan = createPlanMock({
+      monthlyCharge: 0,
+      isDefault: true,
+    });
+
+    const workspaces = Array.from({ length: WORKSPACES_COUNT }, () =>
+      createWorkspaceMock({
+        plan,
+        subscriptionId: null,
+        lastChargeDate: new Date('2005-11-22'),
+        isBlocked: false,
+        billingPeriodEventsCount: 0,
+      })
+    );
+
+    await tariffCollection.insertOne(plan);
+    await workspacesCollection.insertMany(workspaces);
+
+    MockDate.set(currentDate);
+
+    const worker = new PaymasterWorker();
+    const processSpy = jest
+      .spyOn(worker as any, 'processWorkspaceSubscriptionCheck')
+      .mockResolvedValue([null, false]);
+
+    await worker.start();
+    await worker.handle(WORKSPACE_SUBSCRIPTION_CHECK);
+    await worker.finish();
+
+    expect(processSpy).toHaveBeenCalledTimes(WORKSPACES_COUNT);
+
+    const calledIds = processSpy.mock.calls
+      .map((call) => (call[0] as WorkspaceDBScheme)._id.toString())
+      .sort();
+    const expectedIds = workspaces.map((w) => w._id.toString()).sort();
+
+    expect(calledIds).toEqual(expectedIds);
+
+    MockDate.reset();
+  });
+
   afterAll(async () => {
     await connection.close();
     MockDate.reset();