Fix appName option and count billing events via dailyEvents

.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=

lib/db/controller.ts

@@ -87,7 +87,8 @@ export class DatabaseController {
           useNewUrlParser: true,
           useUnifiedTopology: true,
           serverSelectionTimeoutMS: SERVER_SELECTION_TIMEOUT_MS,
-          ...(this.appName ? { appName: this.appName } : {}),
+          /** driver 3.x only recognizes the lowercase `appname` option key */
+          ...(this.appName ? { appname: this.appName } : {}),
         });
         this.db = this.connection.db();
 

workers/limiter/src/dbHelper.ts

@@ -3,6 +3,10 @@ import { PlanDBScheme, ProjectDBScheme, WorkspaceDBScheme } from '@hawk.so/types
 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,
@@ -145,19 +149,13 @@ export class DbHelper {
     since: number
   ): Promise<number> {
     try {
-      const repetitionsCollection = this.eventsDbConnection.collection('repetitions:' + project._id.toString());
-      const eventsCollection = this.eventsDbConnection.collection('events:' + project._id.toString());
-
       const query = {
         timestamp: {
           $gt: since,
         },
       };
 
-      const repetitionsCount = await repetitionsCollection.countDocuments(query);
-      const originalEventCount = await eventsCollection.countDocuments(query);
-
-      return repetitionsCount + originalEventCount;
+      return await this.getRawEventsCountByProject(project, query);
     } catch (e) {
       HawkCatcher.send(e);
       throw new CriticalError(e);
@@ -179,6 +177,75 @@ export class DbHelper {
       .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,
+   * since dailyEvents buckets have day granularity and lastChargeDate does not.
+   *
+   * @param project - project to check
+   * @param since - timestamp of the time from which we count the events
+   */
+  public async getEventsCountByProjectUsingDailyEvents(
+    project: ProjectDBScheme,
+    since: number
+  ): Promise<number> {
+    try {
+      const projectId = project._id.toString();
+      const dailyEventsCollection = this.eventsDbConnection.collection('dailyEvents:' + projectId);
+      const firstFullDayTimestamp = this.getFirstFullDailyEventsTimestamp(since);
+
+      const boundaryDayQuery = {
+        timestamp: {
+          $gt: since,
+          $lt: firstFullDayTimestamp,
+        },
+      };
+
+      const [boundaryDayCount, dailyCounters] = await Promise.all([
+        since < firstFullDayTimestamp
+          ? this.getRawEventsCountByProject(project, boundaryDayQuery)
+          : 0,
+        dailyEventsCollection
+          .aggregate<{ count: number }>([
+            { $match: { groupingTimestamp: { $gte: firstFullDayTimestamp } } },
+            {
+              $group: {
+                _id: null,
+                count: { $sum: '$count' },
+              },
+            },
+          ])
+          .toArray(),
+      ]);
+
+      const fullDaysCount = dailyCounters.length > 0 ? dailyCounters[0].count : 0;
+
+      return boundaryDayCount + fullDaysCount;
+    } catch (e) {
+      HawkCatcher.send(e);
+      throw new CriticalError(e);
+    }
+  }
+
+  /**
+   * 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 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.getEventsCountByProjectUsingDailyEvents(project, since)
+    ))
+      .then(sum);
+  }
+
   /**
    * Returns all projects from Database or projects of the specified workspace
    *
@@ -197,6 +264,52 @@ export class DbHelper {
     return this.projectsCollection.find(query).toArray();
   }
 
+  /**
+   * UTC midnight right after the given timestamp. Mirrors grouper's
+   * getMidnightByEventTimestamp, which fills the dailyEvents buckets.
+   *
+   * @param timestamp - unix timestamp in seconds
+   */
+  private getNextUtcMidnight(timestamp: number): number {
+    const date = new Date(timestamp * MS_IN_SEC);
+
+    date.setUTCDate(date.getUTCDate() + 1);
+    date.setUTCHours(0, 0, 0, 0);
+
+    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
    *

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
    *