diff --git a/.env.sample b/.env.sample index f2df9be2..a5ecb4dc 100644 --- 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/lib/db/controller.ts b/lib/db/controller.ts index 07de49ac..9cd24e18 100644 --- a/lib/db/controller.ts +++ b/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(); diff --git a/workers/limiter/src/dbHelper.ts b/workers/limiter/src/dbHelper.ts index 3af60b34..984dd335 100644 --- a/workers/limiter/src/dbHelper.ts +++ b/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 { 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 { + 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 { + 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): Promise { + 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 index 6ed21cfe..5030e553 100644 --- 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 { + 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 index f57be08d..dd26baa7 100644 --- a/workers/limiter/tests/dbHelper.test.ts +++ b/workers/limiter/tests/dbHelper.test.ts @@ -9,9 +9,20 @@ import HawkCatcher from '@hawk.so/nodejs'; /** * Constant of last charge date in all workspaces for tests + * 2020-04-01T12:00:00Z — intentionally not midnight-aligned */ const LAST_CHARGE_DATE = new Date(1585742400 * 1000); +/** + * Timestamp inside the boundary day of LAST_CHARGE_DATE (2020-04-01T16:00:00Z) + */ +const BOUNDARY_DAY_TIMESTAMP = 1585756800; + +/** + * UTC midnight right after LAST_CHARGE_DATE (2020-04-02T00:00:00Z) + */ +const NEXT_MIDNIGHT_AFTER_LAST_CHARGE = 1585785600; + describe('DbHelper', () => { let connection: MongoClient; let db: Db; @@ -66,8 +77,10 @@ describe('DbHelper', () => { /** * Returns mocked event for tests + * + * @param timestamp - event timestamp, defaults to the boundary day of LAST_CHARGE_DATE */ - const createEventMock = (): GroupedEventDBScheme => { + const createEventMock = (timestamp: number = BOUNDARY_DAY_TIMESTAMP): GroupedEventDBScheme => { return { catcherType: '', totalCount: 0, @@ -77,7 +90,7 @@ describe('DbHelper', () => { payload: { title: 'Mocked event', }, - timestamp: 1586892935, + timestamp, }; }; @@ -89,11 +102,13 @@ describe('DbHelper', () => { const fillDatabaseWithMockedData = async (parameters: { workspace?: WorkspaceDBScheme, project: ProjectDBScheme, - eventsToMock: number + eventsToMock: number, repetitionsToMock?: number, + dailyEventsToMock?: Array<{ groupingTimestamp: number; count: number }>, }): Promise => { const eventsCollection = db.collection(`events:${parameters.project._id.toString()}`); const repetitionsCollection = db.collection(`repetitions:${parameters.project._id.toString()}`); + const dailyEventsCollection = db.collection(`dailyEvents:${parameters.project._id.toString()}`); if (parameters.workspace) { await workspaceCollection.insertOne(parameters.workspace); @@ -104,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; @@ -114,6 +131,14 @@ describe('DbHelper', () => { } await repetitionsCollection.insertMany(mockedEvents); } + + if (parameters.dailyEventsToMock?.length > 0) { + await dailyEventsCollection.insertMany(parameters.dailyEventsToMock.map(bucket => ({ + groupHash: 'ade987831d0d0d167aeea685b49db164eb4e113fd027858eef7f69d049357f62', + groupingTimestamp: bucket.groupingTimestamp, + count: bucket.count, + }))); + } }; beforeAll(async () => { @@ -554,7 +579,7 @@ describe('DbHelper', () => { }); describe('getEventsCountByProject', () => { - test('Should count events and repetitions for a project', async () => { + test('Should count boundary-day events and repetitions for a project', async () => { /** * Arrange */ @@ -584,10 +609,180 @@ 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 + */ + const workspace = createWorkspaceMock({ + plan: mockedPlans.eventsLimit10, + billingPeriodEventsCount: 0, + lastChargeDate: new Date(), + }); + const project = createProjectMock({ workspaceId: workspace._id }); + + await fillDatabaseWithMockedData({ + workspace, + project, + eventsToMock: 2, + repetitionsToMock: 3, + dailyEventsToMock: [ + { + groupingTimestamp: NEXT_MIDNIGHT_AFTER_LAST_CHARGE, + count: 7, + }, + { + groupingTimestamp: NEXT_MIDNIGHT_AFTER_LAST_CHARGE + 86400, + count: 3, + }, + ], + }); + + const since = Math.floor(LAST_CHARGE_DATE.getTime() / MS_IN_SEC); + + /** + * Act + */ + const count = await dbHelper.getEventsCountByProjectUsingDailyEvents(project, since); + + /** + * Assert + */ + expect(count).toBe(15); // 2 events + 3 repetitions on the boundary day + 7 + 3 from dailyEvents + }); + + test('Should ignore raw docs outside the boundary day and dailyEvents buckets before it', 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: [ + /** bucket of the boundary day itself must not be counted */ + { + groupingTimestamp: NEXT_MIDNIGHT_AFTER_LAST_CHARGE - 86400, + count: 100, + }, + ], + }); + + const eventsCollection = db.collection(`events:${project._id.toString()}`); + + await eventsCollection.insertMany([ + /** before lastChargeDate */ + createEventMock(since - 100), + /** after the boundary day — counted via dailyEvents, not the raw scan */ + createEventMock(NEXT_MIDNIGHT_AFTER_LAST_CHARGE + 100), + ]); + + /** + * Act + */ + 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', () => { - test('Should count events and repetitions for multiple projects', async () => { + describe('getEventsCountByProjectsUsingDailyEvents', () => { + test('Should count events, repetitions and dailyEvents for multiple projects', async () => { /** * Arrange */ @@ -604,6 +799,12 @@ describe('DbHelper', () => { project: project1, eventsToMock: 5, repetitionsToMock: 5, + dailyEventsToMock: [ + { + groupingTimestamp: NEXT_MIDNIGHT_AFTER_LAST_CHARGE, + count: 6, + }, + ], }); await fillDatabaseWithMockedData({ project: project2, @@ -616,12 +817,12 @@ describe('DbHelper', () => { /** * Act */ - const count = await dbHelper.getEventsCountByProjects([project1, project2], since); + const count = await dbHelper.getEventsCountByProjectsUsingDailyEvents([project1, project2], since); /** * Assert */ - expect(count).toBe(16); // (5 + 5) + (3 + 3) events and repetitions + expect(count).toBe(22); // (5 + 5 + 6) + (3 + 3) }); }); diff --git a/workers/limiter/tests/index.test.ts b/workers/limiter/tests/index.test.ts index 8a7ebf7b..442fa9bc 100644 --- a/workers/limiter/tests/index.test.ts +++ b/workers/limiter/tests/index.test.ts @@ -27,6 +27,12 @@ const REGULAR_WORKSPACES_CHECK_EVENT: RegularWorkspacesCheckEvent = { */ const LAST_CHARGE_DATE = new Date(1585742400 * 1000); +/** + * Timestamp inside the boundary day of LAST_CHARGE_DATE (2020-04-01T16:00:00Z): + * such events are counted from the raw collections, later ones — via dailyEvents + */ +const BOUNDARY_DAY_TIMESTAMP = 1585756800; + describe('Limiter worker', () => { let connection: MongoClient; let db: Db; @@ -89,7 +95,7 @@ describe('Limiter worker', () => { usersAffected: 0, visitedBy: [], groupHash: 'ade987831d0d0d167aeea685b49db164eb4e113fd027858eef7f69d049357f62', - timestamp: 1586892935, + timestamp: BOUNDARY_DAY_TIMESTAMP, payload: { title: 'Mocked event', }, @@ -104,7 +110,7 @@ describe('Limiter worker', () => { const fillDatabaseWithMockedData = async (parameters: { workspace: WorkspaceDBScheme, project: ProjectDBScheme, - eventsToMock: number + eventsToMock: number, repetitionsToMock?: number, }): Promise => { const eventsCollection = db.collection(`events:${parameters.project._id.toString()}`);