diff --git a/.env.test b/.env.test index a60516a7..fe5007d6 100644 --- a/.env.test +++ b/.env.test @@ -36,8 +36,8 @@ ACCOUNTING_API_ENDPOINT=http://localhost:3999/graphql # Notify URL for sending Telegram notifications REPORT_NOTIFY_URL=http://mock.com/ -# Url for connecting to Redis -REDIS_URL=redis://localhost:6379 +# REDIS_URL is set automatically by jest.setup.redis-mock.js (redis-memory-server). +# Set REDIS_URL in the environment to use an external Redis instance instead. # Disable memoization in tests MEMOIZATION_TTL=-1 diff --git a/workers/grouper/.env.sample b/workers/grouper/.env.sample index 85889136..4216297f 100644 --- a/workers/grouper/.env.sample +++ b/workers/grouper/.env.sample @@ -1,2 +1,8 @@ # Url for connecting to Redis REDIS_URL=redis://redis:6379 + +# How often to refresh project rate limits from MongoDB (seconds) +PROJECTS_LIMITS_UPDATE_PERIOD=3600 + +# Redis hash key for per-project rate limit counters +REDIS_RATE_LIMITS_KEY=rate_limits diff --git a/workers/grouper/README.md b/workers/grouper/README.md index 69bf7327..a56f9ec5 100644 --- a/workers/grouper/README.md +++ b/workers/grouper/README.md @@ -3,11 +3,31 @@ Language-workers adds tasks for Group Worker in this format. Group Worker gets these tasks (events from language-workers) and saves it to the DB +## Rate limiting + +Per-project rate limits are enforced before events are saved to MongoDB. Limits are loaded from the accounts database (`rateLimitSettings` on plans, workspaces, and projects) and tracked in Redis hash `rate_limits` (`timestamp:count` per project). + +Environment variables: + +| Variable | Description | Default | +|----------|-------------|---------| +| `PROJECTS_LIMITS_UPDATE_PERIOD` | Cache refresh interval (seconds) | `3600` | +| `REDIS_RATE_LIMITS_KEY` | Redis hash key for counters | `rate_limits` | + +When the limit is exceeded, the event is dropped (message acked, no DB write) and `events-rate-limited` TimeSeries metrics are recorded. + ## How to run 1. Make sure you are in Workers root directory -3. `yarn install` -4. `yarn run-grouper` - +2. `nvm use` (requires Node 24, see `.nvmrc`) +3. Start dependencies: Redis and MongoDB (e.g. `docker compose up -d redis` from repo root) +4. `yarn install` +5. `yarn run-grouper` +## Tests +```bash +nvm use +yarn build +yarn test:grouper +``` diff --git a/workers/grouper/src/index.ts b/workers/grouper/src/index.ts index 8203e8d3..5eb4a998 100644 --- a/workers/grouper/src/index.ts +++ b/workers/grouper/src/index.ts @@ -22,10 +22,12 @@ import { MS_IN_SEC } from '../../../lib/utils/consts'; import TimeMs from '../../../lib/utils/time'; import DataFilter from './data-filter'; import RedisHelper from './redisHelper'; +import ProjectLimitsCache from './projectLimitsCache'; import { computeDelta } from './utils/repetitionDiff'; import { bucketTimestampMs } from './utils/bucketTimestamp'; import { rightTrim } from '../../../lib/utils/string'; import { hasValue } from '../../../lib/utils/hasValue'; +import { positiveIntEnv } from '../../../lib/utils/positiveIntEnv'; import GrouperMetrics from './metrics/grouperMetrics'; import GrouperMemoryMonitor from './metrics/memoryMonitor'; import SlowHandleDiagnostics, { SlowHandleSession } from './metrics/slowHandleDiagnostics'; @@ -63,6 +65,11 @@ const DAILY_METRICS_RETENTION_DAYS = 90; */ const MAX_CODE_LINE_LENGTH = 140; +/** + * Default interval for refreshing project limits cache (in seconds) + */ +const DEFAULT_PROJECTS_LIMITS_UPDATE_PERIOD_SECONDS = 3600; + /** * Worker for handling Javascript events */ @@ -92,6 +99,16 @@ export default class GrouperWorker extends Worker { */ private redis = new RedisHelper(); + /** + * Cached project rate limits loaded from accounts MongoDB + */ + private projectLimitsCache: ProjectLimitsCache; + + /** + * Interval for periodic project limits cache refresh + */ + private projectLimitsRefreshInterval: NodeJS.Timeout | null = null; + /** * Prometheus metrics facade. */ @@ -117,6 +134,14 @@ export default class GrouperWorker extends Worker { */ private handledTasksCount = 0; + /** + * Create grouper worker instance + */ + constructor() { + super(); + this.projectLimitsCache = new ProjectLimitsCache(this.accountsDb); + } + /** * Start consuming messages */ @@ -131,6 +156,21 @@ export default class GrouperWorker extends Worker { await this.redis.initialize(); console.log('redis initialized'); + await this.projectLimitsCache.refresh(); + + const limitsUpdatePeriodSeconds = positiveIntEnv( + process.env.PROJECTS_LIMITS_UPDATE_PERIOD, + DEFAULT_PROJECTS_LIMITS_UPDATE_PERIOD_SECONDS, + ); + + this.projectLimitsRefreshInterval = setInterval(() => { + /* eslint-disable-next-line no-void */ + void this.projectLimitsCache.refresh() + .catch((error) => { + this.logger.error('Failed to refresh project limits cache', error); + }); + }, limitsUpdatePeriodSeconds * MS_IN_SEC); + /** * Start periodic cache cleanup to prevent memory leaks from unbounded cache growth * Runs every 30 seconds to clear old cache entries @@ -155,6 +195,11 @@ export default class GrouperWorker extends Worker { this.cacheCleanupInterval = null; } + if (this.projectLimitsRefreshInterval) { + clearInterval(this.projectLimitsRefreshInterval); + this.projectLimitsRefreshInterval = null; + } + this.memoryMonitor.logShutdown(this.handledTasksCount); await super.finish(); this.prepareCache(); @@ -170,6 +215,16 @@ export default class GrouperWorker extends Worker { */ public async handle(task: GroupWorkerTask): Promise { try { + const withinLimit = await this.checkRateLimit(task.projectId); + + if (!withinLimit) { + this.grouperMetrics.incrementRateLimitedTotal(); + await this.recordProjectMetrics(task.projectId, 'events-rate-limited'); + this.logger.info(`[rate-limit] project=${task.projectId} title="${task.payload?.title}" dropped`); + + return; + } + await this.grouperMetrics.observeHandleDuration(async () => { await this.handleInternal(task); }); @@ -180,6 +235,30 @@ export default class GrouperWorker extends Worker { } } + /** + * Check and update per-project rate limit in Redis. + * + * @param projectId - project id + */ + private async checkRateLimit(projectId: string): Promise { + const limits = this.projectLimitsCache.getProjectLimits(projectId); + + if (!limits) { + this.logger.warn(`Project ${projectId} is not in the projects limits cache`); + } + + const eventsLimit = limits?.eventsLimit ?? 0; + const eventsPeriod = limits?.eventsPeriod ?? 0; + + try { + return await this.redis.updateRateLimit(projectId, eventsLimit, eventsPeriod); + } catch (error) { + this.logger.error(`Failed to update rate limit for project ${projectId}`, error); + + return false; + } + } + /** * Internal task handling function * @@ -301,7 +380,7 @@ export default class GrouperWorker extends Worker { this.grouperMetrics.incrementDuplicateRetriesTotal(); this.logger.info(`[saveEvent] project=${task.projectId} title="${task.payload.title}" duplicate key, retrying as repetition`); - await this.handle(task); + await this.handleInternal(task); return; } diff --git a/workers/grouper/src/metrics/grouperMetrics.ts b/workers/grouper/src/metrics/grouperMetrics.ts index 02a36306..dfc55bfc 100644 --- a/workers/grouper/src/metrics/grouperMetrics.ts +++ b/workers/grouper/src/metrics/grouperMetrics.ts @@ -135,6 +135,15 @@ export default class GrouperMetrics { }) ); + private readonly rateLimitedTotal = getOrCreateMetric( + 'hawk_grouper_events_rate_limited_total', + () => new client.Counter({ + name: 'hawk_grouper_events_rate_limited_total', + help: 'Number of events dropped due to rate limiting', + registers: [ register ], + }) + ); + /** * Measure top-level handle() duration. * @@ -207,6 +216,13 @@ export default class GrouperMetrics { this.duplicateRetriesTotal.inc(); } + /** + * Increment events dropped by rate limiting. + */ + public incrementRateLimitedTotal(): void { + this.rateLimitedTotal.inc(); + } + /** * Measure Mongo operation duration. * diff --git a/workers/grouper/src/projectLimitsCache.ts b/workers/grouper/src/projectLimitsCache.ts new file mode 100644 index 00000000..e8574777 --- /dev/null +++ b/workers/grouper/src/projectLimitsCache.ts @@ -0,0 +1,174 @@ +import type { Db } from 'mongodb'; +import { DatabaseController } from '../../../lib/db/controller'; +import createLogger from '../../../lib/logger'; + +/** + * Rate limit settings stored in MongoDB (plans, workspaces, projects). + */ +interface RateLimitSettingsDocument { + N?: number; + T?: number; +} + +interface ProjectDocument { + _id: { toString(): string }; + workspaceId?: { toString(): string }; + rateLimitSettings?: RateLimitSettingsDocument; +} + +interface PlanDocument { + _id: unknown; + rateLimitSettings?: RateLimitSettingsDocument; +} + +interface WorkspaceDocument { + _id: { toString(): string }; + rateLimitSettings?: RateLimitSettingsDocument; + plan: PlanDocument; +} + +/** + * Resolved rate limit for a project. + */ +export interface ProjectRateLimits { + eventsLimit: number; + eventsPeriod: number; +} + +const CONTEXT_TIMEOUT_MS = 5000; + +/** + * In-memory cache of project rate limits loaded from MongoDB accounts database. + * Priority: project → workspace → plan. + */ +export default class ProjectLimitsCache { + private projectLimits = new Map(); + + private readonly logger = createLogger(); + private readonly accountsDb: DatabaseController; + + /** + * @param accountsDb - accounts database controller + */ + constructor(accountsDb: DatabaseController) { + this.accountsDb = accountsDb; + } + + /** + * Returns cached rate limits for a project. + * + * @param projectId - project id + */ + public getProjectLimits(projectId: string): ProjectRateLimits | undefined { + return this.projectLimits.get(projectId); + } + + /** + * Reload all project limits from MongoDB. + */ + public async refresh(): Promise { + const db = await this.accountsDb.connect(); + const projectLimitsTmp = new Map(); + + const workspaceMap = await this.loadWorkspacesWithPlans(db); + + const projects = await db.collection('projects') + .find({}, { maxTimeMS: CONTEXT_TIMEOUT_MS }) + .toArray(); + + for (const project of projects) { + const projectId = project._id.toString(); + let finalLimits: ProjectRateLimits = { + eventsLimit: 0, + eventsPeriod: 0, + }; + + const workspaceId = project.workspaceId?.toString(); + const workspace = workspaceId ? workspaceMap.get(workspaceId) : undefined; + + if (workspace) { + finalLimits = { + eventsLimit: toNumber(workspace.plan?.rateLimitSettings?.N), + eventsPeriod: toNumber(workspace.plan?.rateLimitSettings?.T), + }; + + const workspaceLimit = toNumber(workspace.rateLimitSettings?.N); + const workspacePeriod = toNumber(workspace.rateLimitSettings?.T); + + if (workspaceLimit > 0) { + finalLimits.eventsLimit = workspaceLimit; + } + if (workspacePeriod > 0) { + finalLimits.eventsPeriod = workspacePeriod; + } + } + + const projectLimit = toNumber(project.rateLimitSettings?.N); + const projectPeriod = toNumber(project.rateLimitSettings?.T); + + if (projectLimit > 0) { + finalLimits.eventsLimit = projectLimit; + } + if (projectPeriod > 0) { + finalLimits.eventsPeriod = projectPeriod; + } + + projectLimitsTmp.set(projectId, finalLimits); + } + + this.projectLimits = projectLimitsTmp; + this.logger.debug(`Project limits cache refreshed with ${projectLimitsTmp.size} projects`); + } + + /** + * Load workspaces joined with tariff plans. + * + * @param db - accounts database + */ + private async loadWorkspacesWithPlans(db: Db): Promise> { + const workspaces = await db.collection('workspaces') + .aggregate([ + { + $lookup: { + from: 'plans', + localField: 'tariffPlanId', + foreignField: '_id', + as: 'plan', + }, + }, + { $unwind: '$plan' }, + ], { maxTimeMS: CONTEXT_TIMEOUT_MS }) + .toArray(); + + const workspaceMap = new Map(); + + for (const workspace of workspaces) { + workspaceMap.set(workspace._id.toString(), workspace); + } + + return workspaceMap; + } +} + +/** + * Coerce MongoDB numeric values to number. + * + * @param value - raw value from document + */ +function toNumber(value: unknown): number { + if (typeof value === 'number') { + return value; + } + + if (typeof value === 'object' && value !== null && 'toNumber' in value && typeof (value as { toNumber: () => number }).toNumber === 'function') { + return (value as { toNumber: () => number }).toNumber(); + } + + if (typeof value === 'string') { + const parsed = Number(value); + + return Number.isFinite(parsed) ? parsed : 0; + } + + return 0; +} diff --git a/workers/grouper/src/redisHelper.ts b/workers/grouper/src/redisHelper.ts index 75c92032..314695b8 100644 --- a/workers/grouper/src/redisHelper.ts +++ b/workers/grouper/src/redisHelper.ts @@ -2,6 +2,43 @@ import HawkCatcher from '@hawk.so/nodejs'; import type { RedisClientType } from 'redis'; import { createClient } from 'redis'; import createLogger from '../../../lib/logger'; +import { MS_IN_SEC } from '../../../lib/utils/consts'; + +const RATE_LIMITS_KEY = process.env.REDIS_RATE_LIMITS_KEY || 'rate_limits'; + +const UPDATE_RATE_LIMIT_SCRIPT = ` + local key = KEYS[1] + local field = ARGV[1] + local now = tonumber(ARGV[2]) + local limit = tonumber(ARGV[3]) + local period = tonumber(ARGV[4]) + + local current = redis.call('HGET', key, field) + if not current then + redis.call('HSET', key, field, now .. ':1') + return 1 + end + + local timestamp, count = string.match(current, '^(%d+):(%d+)$') + if not timestamp then + redis.call('HSET', key, field, now .. ':1') + return 1 + end + timestamp = tonumber(timestamp) + count = tonumber(count) + + if now - timestamp >= period then + redis.call('HSET', key, field, now .. ':1') + return 1 + end + + if count + 1 > limit then + return 0 + end + + redis.call('HSET', key, field, timestamp .. ':' .. (count + 1)) + return 1 +`; /** * Class with helper functions for working with Redis @@ -63,6 +100,38 @@ export default class RedisHelper { } } + /** + * Atomically checks and updates the per-project rate limit counter. + * + * @param projectId - project id used as hash field + * @param eventsLimit - max events allowed in the period (0 = unlimited) + * @param eventsPeriod - window size in seconds + * @returns true when the event is within the limit + */ + public async updateRateLimit( + projectId: string, + eventsLimit: number, + eventsPeriod: number, + ): Promise { + if (eventsLimit === 0) { + return true; + } + + const now = Math.floor(Date.now() / MS_IN_SEC); + + const result = await this.redisClient.eval(UPDATE_RATE_LIMIT_SCRIPT, { + keys: [ RATE_LIMITS_KEY ], + arguments: [ + projectId, + now.toString(), + eventsLimit.toString(), + eventsPeriod.toString(), + ], + }); + + return Number(result) === 1; + } + /** * Checks if a lock exists on the given group hash and identifier pair. If it does not exist, creates a lock. * Returns true if lock exists diff --git a/workers/grouper/tests/index.test.ts b/workers/grouper/tests/index.test.ts index c75351a9..96d42735 100644 --- a/workers/grouper/tests/index.test.ts +++ b/workers/grouper/tests/index.test.ts @@ -743,6 +743,23 @@ describe('GrouperWorker', () => { }); }); + describe('Rate limiting', () => { + test('drops events that exceed the project rate limit', async () => { + await projectsCollection.updateOne( + { _id: new mongodb.ObjectId(projectIdMock) }, + { $set: { rateLimitSettings: { N: 2, T: 3600 } } }, + ); + await (worker as any).projectLimitsCache.refresh(); + await redisClient.del('rate_limits'); + + await worker.handle(generateTask({ title: 'Rate limit event 1' })); + await worker.handle(generateTask({ title: 'Rate limit event 2' })); + await worker.handle(generateTask({ title: 'Rate limit event 3' })); + + expect(await eventsCollection.countDocuments()).toBe(2); + }); + }); + describe('Events-accepted metrics', () => { test('writes minutely, hourly, and daily samples after handling an event', async () => { const safeTsAddSpy = jest.spyOn((worker as any).redis, 'safeTsAdd'); diff --git a/workers/grouper/tests/redisHelper.rateLimit.test.ts b/workers/grouper/tests/redisHelper.rateLimit.test.ts new file mode 100644 index 00000000..f454d81f --- /dev/null +++ b/workers/grouper/tests/redisHelper.rateLimit.test.ts @@ -0,0 +1,94 @@ +import '../../../env-test'; +import RedisHelper from '../src/redisHelper'; +import type { RedisClientType } from 'redis'; +import { createClient } from 'redis'; +import { MS_IN_SEC } from '../../../lib/utils/consts'; + +function parseRateLimitValue(value: string): { timestamp: number; count: number } | null { + const match = /^(\d+):(\d+)$/.exec(value); + + if (!match) { + return null; + } + + return { + timestamp: Number(match[1]), + count: Number(match[2]), + }; +} + +describe('RedisHelper.updateRateLimit', () => { + let redisClient: RedisClientType; + let redisHelper: RedisHelper; + + beforeAll(async () => { + redisClient = createClient({ url: process.env.REDIS_URL }); + await redisClient.connect(); + redisHelper = new RedisHelper(); + await redisHelper.initialize(); + }); + + beforeEach(async () => { + await redisClient.del('rate_limits'); + }); + + afterAll(async () => { + await redisHelper.close(); + await redisClient.quit(); + }); + + test('allows events when limit is zero', async () => { + expect(await redisHelper.updateRateLimit('project-zero', 0, 60)).toBe(true); + expect(await redisHelper.updateRateLimit('project-zero', 0, 60)).toBe(true); + }); + + test('allows events until the limit is reached', async () => { + const projectId = 'project-limit'; + + expect(await redisHelper.updateRateLimit(projectId, 2, 3600)).toBe(true); + expect(await redisHelper.updateRateLimit(projectId, 2, 3600)).toBe(true); + expect(await redisHelper.updateRateLimit(projectId, 2, 3600)).toBe(false); + }); + + test('resets the counter after the period expires', async () => { + const projectId = 'project-reset'; + const now = Math.floor(Date.now() / MS_IN_SEC); + const prevTimestamp = now - 3601; + + await redisClient.hSet('rate_limits', projectId, `${prevTimestamp}:5`); + + expect(await redisHelper.updateRateLimit(projectId, 5, 3600)).toBe(true); + + const stored = await redisClient.hGet('rate_limits', projectId); + const counter = parseRateLimitValue(stored!); + + expect(counter).toEqual({ + timestamp: expect.any(Number), + count: 1, + }); + expect(counter!.timestamp).toBeGreaterThan(prevTimestamp); + expect(counter!.timestamp).toBeGreaterThanOrEqual(now); + }); + + test('keeps the same timestamp when the window is active and limit is reached', async () => { + const projectId = 'project-deny'; + const now = Math.floor(Date.now() / MS_IN_SEC); + const prevTimestamp = now - 30; + + await redisClient.hSet('rate_limits', projectId, `${prevTimestamp}:2`); + + expect(await redisHelper.updateRateLimit(projectId, 2, 3600)).toBe(false); + + const stored = await redisClient.hGet('rate_limits', projectId); + + expect(stored).toBe(`${prevTimestamp}:2`); + }); + + test('resets malformed counter values', async () => { + const projectId = 'project-malformed'; + + await redisClient.hSet('rate_limits', projectId, 'invalid'); + + expect(await redisHelper.updateRateLimit(projectId, 1, 3600)).toBe(true); + }); +}); diff --git a/workers/limiter/src/index.ts b/workers/limiter/src/index.ts index 6ed21cfe..71aa2d81 100644 --- a/workers/limiter/src/index.ts +++ b/workers/limiter/src/index.ts @@ -271,7 +271,7 @@ export default class LimiterWorker extends Worker { this.logger.info(`workspace ${workspace._id} events count since last charge date: ${workspaceEventsCount}`); const usedQuota = workspaceEventsCount / workspace.tariffPlan.eventsLimit; - const quotaNotification = NOTIFY_ABOUT_LIMIT.reverse().find(quota => quota < usedQuota); + const quotaNotification = NOTIFY_ABOUT_LIMIT.find(quota => quota < usedQuota); const shouldBeBlockedByQuota = usedQuota >= 1; const isAlreadyBlocked = workspace.isBlocked;