Add opt-in dailyEvents limiter counter

Kuchizu · Kuchizu · commit 12ee807ce7d8 · 2026-07-08T11:47:07.000+03:00
diff --git a/.env.sample b/.env.sample
@@ -55,5 +55,8 @@ HAWK_CATCHER_TOKEN=
 ## If true, Grouper worker will send messages about new events to Notifier worker
 IS_NOTIFIER_WORKER_ENABLED=false
 
+## Comma-separated workspace ids that should use dailyEvents counters in Limiter quota checks. Use * for all workspaces.
+LIMITER_DAILY_EVENTS_COUNTER_WORKSPACE_IDS=
+
 ## Url for telegram notifications about workspace blocks and unblocks 
 TELEGRAM_LIMITER_CHAT_URL=
diff --git a/workers/limiter/src/dbHelper.ts b/workers/limiter/src/dbHelper.ts
@@ -4,6 +4,9 @@ import { WorkspaceWithTariffPlan } from '../types';
 import HawkCatcher from '@hawk.so/nodejs';
 import { CriticalError, NonCriticalError } from '../../../lib/workerErrors';
 import { MS_IN_SEC } from '../../../lib/utils/consts';
+import TimeMs from '../../../lib/utils/time';
+
+const SEC_IN_DAY = TimeMs.DAY / TimeMs.SECOND;
 
 const WORKSPACE_PROJECTION = {
   _id: 1,
@@ -138,6 +141,45 @@ export class DbHelper {
   /**
    * Returns total event counts for last billing period
    *
+   * @param project - project to check
+   * @param since - timestamp of the time from which we count the events
+   */
+  public async getEventsCountByProject(
+    project: ProjectDBScheme,
+    since: number
+  ): Promise<number> {
+    try {
+      const query = {
+        timestamp: {
+          $gt: since,
+        },
+      };
+
+      return await this.getRawEventsCountByProject(project, query);
+    } catch (e) {
+      HawkCatcher.send(e);
+      throw new CriticalError(e);
+    }
+  }
+
+  /**
+   * Calculates total events count for all provided projects since the specific date
+   *
+   * @param projects - projects to calculate for
+   * @param since - timestamp of the time from which we count the events
+   */
+  public async getEventsCountByProjects(projects: ProjectDBScheme[], since: number): Promise<number> {
+    const sum = (array: number[]): number => array.reduce((acc, val) => acc + val, 0);
+
+    return Promise.all(projects.map(
+      project => this.getEventsCountByProject(project, since)
+    ))
+      .then(sum);
+  }
+
+  /**
+   * Returns total event counts for last billing period using dailyEvents counters.
+   *
    * Full days are summed from dailyEvents per-day counters (grouper
    * increments `count` for originals and repetitions alike); only the
    * partial day containing `since` is counted from the raw collections,
@@ -146,39 +188,42 @@ export class DbHelper {
    * @param project - project to check
    * @param since - timestamp of the time from which we count the events
    */
-  public async getEventsCountByProject(
+  public async getEventsCountByProjectUsingDailyEvents(
     project: ProjectDBScheme,
     since: number
   ): Promise<number> {
     try {
       const projectId = project._id.toString();
-      const repetitionsCollection = this.eventsDbConnection.collection('repetitions:' + projectId);
-      const eventsCollection = this.eventsDbConnection.collection('events:' + projectId);
       const dailyEventsCollection = this.eventsDbConnection.collection('dailyEvents:' + projectId);
-
-      const sinceNextMidnight = this.getNextUtcMidnight(since);
+      const firstFullDayTimestamp = this.getFirstFullDailyEventsTimestamp(since);
 
       const boundaryDayQuery = {
         timestamp: {
           $gt: since,
-          $lt: sinceNextMidnight,
+          $lt: firstFullDayTimestamp,
         },
       };
 
-      const [repetitionsCount, originalEventCount, dailyCounters] = await Promise.all([
-        repetitionsCollection.countDocuments(boundaryDayQuery),
-        eventsCollection.countDocuments(boundaryDayQuery),
+      const [boundaryDayCount, dailyCounters] = await Promise.all([
+        since < firstFullDayTimestamp
+          ? this.getRawEventsCountByProject(project, boundaryDayQuery)
+          : 0,
         dailyEventsCollection
           .aggregate<{ count: number }>([
-            { $match: { groupingTimestamp: { $gte: sinceNextMidnight } } },
-            { $group: { _id: null, count: { $sum: '$count' } } },
+            { $match: { groupingTimestamp: { $gte: firstFullDayTimestamp } } },
+            {
+              $group: {
+                _id: null,
+                count: { $sum: '$count' },
+              },
+            },
           ])
           .toArray(),
       ]);
 
       const fullDaysCount = dailyCounters.length > 0 ? dailyCounters[0].count : 0;
 
-      return repetitionsCount + originalEventCount + fullDaysCount;
+      return boundaryDayCount + fullDaysCount;
     } catch (e) {
       HawkCatcher.send(e);
       throw new CriticalError(e);
@@ -187,15 +232,16 @@ export class DbHelper {
 
   /**
    * Calculates total events count for all provided projects since the specific date
+   * using dailyEvents counters for full days.
    *
    * @param projects - projects to calculate for
    * @param since - timestamp of the time from which we count the events
    */
-  public async getEventsCountByProjects(projects: ProjectDBScheme[], since: number): Promise<number> {
+  public async getEventsCountByProjectsUsingDailyEvents(projects: ProjectDBScheme[], since: number): Promise<number> {
     const sum = (array: number[]): number => array.reduce((acc, val) => acc + val, 0);
 
     return Promise.all(projects.map(
-      project => this.getEventsCountByProject(project, since)
+      project => this.getEventsCountByProjectUsingDailyEvents(project, since)
     ))
       .then(sum);
   }
@@ -233,6 +279,37 @@ export class DbHelper {
     return date.getTime() / MS_IN_SEC;
   }
 
+  /**
+   * Returns first dailyEvents bucket that can be safely used without counting
+   * events before the requested timestamp.
+   *
+   * @param timestamp - unix timestamp in seconds
+   */
+  private getFirstFullDailyEventsTimestamp(timestamp: number): number {
+    const midnight = timestamp - (timestamp % SEC_IN_DAY);
+
+    return timestamp === midnight ? timestamp : this.getNextUtcMidnight(timestamp);
+  }
+
+  /**
+   * Counts raw original events and repetitions for the passed query.
+   *
+   * @param project - project to check
+   * @param query - MongoDB timestamp query
+   */
+  private async getRawEventsCountByProject(project: ProjectDBScheme, query: Record<string, unknown>): Promise<number> {
+    const projectId = project._id.toString();
+    const repetitionsCollection = this.eventsDbConnection.collection('repetitions:' + projectId);
+    const eventsCollection = this.eventsDbConnection.collection('events:' + projectId);
+
+    const [repetitionsCount, originalEventCount] = await Promise.all([
+      repetitionsCollection.countDocuments(query),
+      eventsCollection.countDocuments(query),
+    ]);
+
+    return repetitionsCount + originalEventCount;
+  }
+
   /**
    * Returns plan from cache, refetches once on miss
    *
diff --git a/workers/limiter/src/index.ts b/workers/limiter/src/index.ts
@@ -30,6 +30,12 @@ const NOTIFY_ABOUT_LIMIT = [
   0.95,
 ];
 
+/**
+ * Enables dailyEvents-based quota counting for listed workspaces.
+ * Use comma-separated workspace ids or `*` to enable it for every workspace.
+ */
+const DAILY_EVENTS_COUNTER_WORKSPACE_IDS = process.env.LIMITER_DAILY_EVENTS_COUNTER_WORKSPACE_IDS || '';
+
 /**
  * Worker for checking current total events count in workspaces and limits events receiving if workspace exceed the limit
  */
@@ -266,7 +272,7 @@ export default class LimiterWorker extends Worker {
 
     const since = Math.floor(new Date(workspace.lastChargeDate).getTime() / MS_IN_SEC);
 
-    const workspaceEventsCount = await this.dbHelper.getEventsCountByProjects(projects, since);
+    const workspaceEventsCount = await this.getWorkspaceEventsCount(workspace, projects, since);
 
     this.logger.info(`workspace ${workspace._id} events count since last charge date: ${workspaceEventsCount}`);
 
@@ -328,6 +334,48 @@ export default class LimiterWorker extends Worker {
     };
   }
 
+  /**
+   * Returns workspace events count using the default raw counter or the
+   * dailyEvents-based counter when it is explicitly enabled for the workspace.
+   *
+   * @param workspace - workspace to count events for
+   * @param projects - workspace projects
+   * @param since - timestamp of the time from which we count the events
+   */
+  private async getWorkspaceEventsCount(
+    workspace: WorkspaceWithTariffPlan,
+    projects: ProjectDBScheme[],
+    since: number
+  ): Promise<number> {
+    if (!this.shouldUseDailyEventsCounter(workspace._id.toString())) {
+      return this.dbHelper.getEventsCountByProjects(projects, since);
+    }
+
+    try {
+      return await this.dbHelper.getEventsCountByProjectsUsingDailyEvents(projects, since);
+    } catch (error) {
+      HawkCatcher.send(error, {
+        workspaceId: workspace._id.toString(),
+      });
+
+      return this.dbHelper.getEventsCountByProjects(projects, since);
+    }
+  }
+
+  /**
+   * Checks whether dailyEvents-based quota counting is enabled for the workspace.
+   *
+   * @param workspaceId - workspace id
+   */
+  private shouldUseDailyEventsCounter(workspaceId: string): boolean {
+    const enabledWorkspaceIds = DAILY_EVENTS_COUNTER_WORKSPACE_IDS
+      .split(',')
+      .map(id => id.trim())
+      .filter(Boolean);
+
+    return enabledWorkspaceIds.includes('*') || enabledWorkspaceIds.includes(workspaceId);
+  }
+
   /**
    * Method that formats project list to html used in report messages
    *
diff --git a/workers/limiter/tests/dbHelper.test.ts b/workers/limiter/tests/dbHelper.test.ts
@@ -102,7 +102,7 @@ describe('DbHelper', () => {
   const fillDatabaseWithMockedData = async (parameters: {
     workspace?: WorkspaceDBScheme,
     project: ProjectDBScheme,
-    eventsToMock: number
+    eventsToMock: number,
     repetitionsToMock?: number,
     dailyEventsToMock?: Array<{ groupingTimestamp: number; count: number }>,
   }): Promise<void> => {
@@ -119,7 +119,9 @@ describe('DbHelper', () => {
     for (let i = 0; i < parameters.eventsToMock; i++) {
       mockedEvents.push(createEventMock());
     }
-    await eventsCollection.insertMany(mockedEvents);
+    if (mockedEvents.length > 0) {
+      await eventsCollection.insertMany(mockedEvents);
+    }
 
     mockedEvents.length = 0;
 
@@ -608,6 +610,47 @@ describe('DbHelper', () => {
       expect(count).toBe(10); // 5 events + 5 repetitions
     });
 
+    test('Should keep raw counting as default and include raw docs after the boundary day', async () => {
+      /**
+       * Arrange
+       */
+      const workspace = createWorkspaceMock({
+        plan: mockedPlans.eventsLimit10,
+        billingPeriodEventsCount: 0,
+        lastChargeDate: new Date(),
+      });
+      const project = createProjectMock({ workspaceId: workspace._id });
+      const since = Math.floor(LAST_CHARGE_DATE.getTime() / MS_IN_SEC);
+
+      await fillDatabaseWithMockedData({
+        workspace,
+        project,
+        eventsToMock: 1,
+        dailyEventsToMock: [
+          {
+            groupingTimestamp: NEXT_MIDNIGHT_AFTER_LAST_CHARGE,
+            count: 100,
+          },
+        ],
+      });
+
+      await db.collection(`events:${project._id.toString()}`).insertOne(
+        createEventMock(NEXT_MIDNIGHT_AFTER_LAST_CHARGE + 100)
+      );
+
+      /**
+       * Act
+       */
+      const count = await dbHelper.getEventsCountByProject(project, since);
+
+      /**
+       * Assert
+       */
+      expect(count).toBe(2);
+    });
+  });
+
+  describe('getEventsCountByProjectUsingDailyEvents', () => {
     test('Should add per-day counters from dailyEvents for days after the boundary day', async () => {
       /**
        * Arrange
@@ -641,7 +684,7 @@ describe('DbHelper', () => {
       /**
        * Act
        */
-      const count = await dbHelper.getEventsCountByProject(project, since);
+      const count = await dbHelper.getEventsCountByProjectUsingDailyEvents(project, since);
 
       /**
        * Assert
@@ -686,16 +729,59 @@ describe('DbHelper', () => {
       /**
        * Act
        */
-      const count = await dbHelper.getEventsCountByProject(project, since);
+      const count = await dbHelper.getEventsCountByProjectUsingDailyEvents(project, since);
 
       /**
        * Assert
        */
       expect(count).toBe(1); // only the single boundary-day event
     });
+
+    test('Should count dailyEvents bucket at since when since is already UTC midnight', async () => {
+      /**
+       * Arrange
+       */
+      const workspace = createWorkspaceMock({
+        plan: mockedPlans.eventsLimit10,
+        billingPeriodEventsCount: 0,
+        lastChargeDate: new Date(),
+      });
+      const project = createProjectMock({ workspaceId: workspace._id });
+      const since = NEXT_MIDNIGHT_AFTER_LAST_CHARGE;
+
+      await fillDatabaseWithMockedData({
+        workspace,
+        project,
+        eventsToMock: 0,
+        dailyEventsToMock: [
+          {
+            groupingTimestamp: since,
+            count: 4,
+          },
+          {
+            groupingTimestamp: since + 86400,
+            count: 3,
+          },
+        ],
+      });
+
+      await db.collection(`events:${project._id.toString()}`).insertOne(
+        createEventMock(since + 100)
+      );
+
+      /**
+       * Act
+       */
+      const count = await dbHelper.getEventsCountByProjectUsingDailyEvents(project, since);
+
+      /**
+       * Assert
+       */
+      expect(count).toBe(7);
+    });
   });
 
-  describe('getEventsCountByProjects', () => {
+  describe('getEventsCountByProjectsUsingDailyEvents', () => {
     test('Should count events, repetitions and dailyEvents for multiple projects', async () => {
       /**
        * Arrange
@@ -731,7 +817,7 @@ describe('DbHelper', () => {
       /**
        * Act
        */
-      const count = await dbHelper.getEventsCountByProjects([project1, project2], since);
+      const count = await dbHelper.getEventsCountByProjectsUsingDailyEvents([project1, project2], since);
 
       /**
        * Assert
diff --git a/workers/limiter/tests/index.test.ts b/workers/limiter/tests/index.test.ts
