v2 IngestSendEvent rate limit (#1134)

matt-aitken · web-flow · commit bc7bbd45766f · 2024-05-31T10:00:03.000+01:00
* Easier to create a rate limiter, use it in the ApiRateLimiter. Upgraded the Upstash package

* Always prefix any rate limiter in Redis with “ratelimit:”

* By default log when the rate limit is hit

* Added rate limiting to IngestSendEvent

* Log out the EventRecord id

* Increase events.deliverScheduled attempts

* INGEST_EVENT_RATE_LIMIT_MAX is optional

* Removed old API rate limit code

* IngestSendEvent rate limiter is optional. Moved outside of the DB transaction

* Log a message out when the rate limiter is created

* Return undefined if the rate limit has been crossed
diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts
@@ -100,6 +100,10 @@ const EnvironmentSchema = z.object({
   API_RATE_LIMIT_REQUEST_LOGS_ENABLED: z.string().default("0"),
   API_RATE_LIMIT_REJECTION_LOGS_ENABLED: z.string().default("1"),
 
+  //Ingesting event rate limit
+  INGEST_EVENT_RATE_LIMIT_WINDOW: z.string().default("60s"),
+  INGEST_EVENT_RATE_LIMIT_MAX: z.coerce.number().int().optional(),
+
   //v3
   V3_ENABLED: z.string().default("false"),
   PROVIDER_SECRET: z.string().default("provider-secret"),
diff --git a/apps/webapp/app/services/apiRateLimit.server.ts b/apps/webapp/app/services/apiRateLimit.server.ts
@@ -1,45 +1,21 @@
 import { Ratelimit } from "@upstash/ratelimit";
 import { Request as ExpressRequest, Response as ExpressResponse, NextFunction } from "express";
-import Redis, { RedisOptions } from "ioredis";
+import { RedisOptions } from "ioredis";
 import { createHash } from "node:crypto";
 import { env } from "~/env.server";
 import { logger } from "./logger.server";
-
-function createRedisRateLimitClient(
-  redisOptions: RedisOptions
-): ConstructorParameters<typeof Ratelimit>[0]["redis"] {
-  const redis = new Redis(redisOptions);
-
-  return {
-    sadd: async <TData>(key: string, ...members: TData[]): Promise<number> => {
-      return redis.sadd(key, members as (string | number | Buffer)[]);
-    },
-    eval: <TArgs extends unknown[], TData = unknown>(
-      ...args: [script: string, keys: string[], args: TArgs]
-    ): Promise<TData> => {
-      const script = args[0];
-      const keys = args[1];
-      const argsArray = args[2];
-      return redis.eval(
-        script,
-        keys.length,
-        ...keys,
-        ...(argsArray as (string | Buffer | number)[])
-      ) as Promise<TData>;
-    },
-  };
-}
+import { Duration, Limiter, RateLimiter, createRedisRateLimitClient } from "./rateLimiter.server";
 
 type Options = {
+  redis?: RedisOptions;
+  keyPrefix: string;
+  pathMatchers: (RegExp | string)[];
+  pathWhiteList?: (RegExp | string)[];
+  limiter: Limiter;
   log?: {
     requests?: boolean;
     rejections?: boolean;
   };
-  redis: RedisOptions;
-  keyPrefix: string;
-  pathMatchers: (RegExp | string)[];
-  pathWhiteList?: (RegExp | string)[];
-  limiter: ConstructorParameters<typeof Ratelimit>[0]["limiter"];
 };
 
 //returns an Express middleware that rate limits using the Bearer token in the Authorization header
@@ -54,12 +30,12 @@ export function authorizationRateLimitMiddleware({
     requests: true,
   },
 }: Options) {
-  const rateLimiter = new Ratelimit({
-    redis: createRedisRateLimitClient(redis),
-    limiter: limiter,
-    ephemeralCache: new Map(),
-    analytics: false,
-    prefix: keyPrefix,
+  const rateLimiter = new RateLimiter({
+    redis,
+    keyPrefix,
+    limiter,
+    logSuccess: log.requests,
+    logFailure: log.rejections,
   });
 
   return async (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => {
@@ -135,27 +111,9 @@ export function authorizationRateLimitMiddleware({
     res.set("x-ratelimit-reset", reset.toString());
 
     if (success) {
-      if (log.requests) {
-        logger.info(`RateLimiter (${keyPrefix}): under rate limit`, {
-          limit,
-          reset,
-          remaining,
-          hashedAuthorizationValue,
-        });
-      }
       return next();
     }
 
-    if (log.rejections) {
-      logger.warn(`RateLimiter (${keyPrefix}): rate limit exceeded`, {
-        limit,
-        reset,
-        remaining,
-        pending,
-        hashedAuthorizationValue,
-      });
-    }
-
     res.setHeader("Content-Type", "application/problem+json");
     const secondsUntilReset = Math.max(0, (reset - new Date().getTime()) / 1000);
     return res.status(429).send(
@@ -167,6 +125,7 @@ export function authorizationRateLimitMiddleware({
           detail: `Rate limit exceeded ${remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
           reset,
           limit,
+          remaining,
           secondsUntilReset,
           error: `Rate limit exceeded ${remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
         },
@@ -177,18 +136,8 @@ export function authorizationRateLimitMiddleware({
   };
 }
 
-type Duration = Parameters<typeof Ratelimit.slidingWindow>[1];
-
 export const apiRateLimiter = authorizationRateLimitMiddleware({
-  keyPrefix: "ratelimit:api",
-  redis: {
-    port: env.REDIS_PORT,
-    host: env.REDIS_HOST,
-    username: env.REDIS_USERNAME,
-    password: env.REDIS_PASSWORD,
-    enableAutoPipelining: true,
-    ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
-  },
+  keyPrefix: "api",
   limiter: Ratelimit.slidingWindow(env.API_RATE_LIMIT_MAX, env.API_RATE_LIMIT_WINDOW as Duration),
   pathMatchers: [/^\/api/],
   // Allow /api/v1/tasks/:id/callback/:secret
diff --git a/apps/webapp/app/services/events/ingestSendEvent.server.ts b/apps/webapp/app/services/events/ingestSendEvent.server.ts
@@ -4,6 +4,9 @@ import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { workerQueue } from "~/services/worker.server";
 import { logger } from "../logger.server";
 import { EventRecord, ExternalAccount } from "@trigger.dev/database";
+import { Duration, RateLimiter } from "../rateLimiter.server";
+import { Ratelimit } from "@upstash/ratelimit";
+import { env } from "~/env.server";
 
 type UpdateEventInput = {
   tx: PrismaClientOrTransaction;
@@ -31,9 +34,19 @@ const EVENT_UPDATE_THRESHOLD_WINDOW_IN_MSECS = 5 * 1000; // 5 seconds
 
 export class IngestSendEvent {
   #prismaClient: PrismaClientOrTransaction;
+  #rateLimiter: RateLimiter | undefined;
 
   constructor(prismaClient: PrismaClientOrTransaction = prisma, private deliverEvents = true) {
     this.#prismaClient = prismaClient;
+    this.#rateLimiter = env.INGEST_EVENT_RATE_LIMIT_MAX
+      ? new RateLimiter({
+          keyPrefix: "ingestsendevent",
+          limiter: Ratelimit.slidingWindow(
+            env.INGEST_EVENT_RATE_LIMIT_MAX,
+            env.INGEST_EVENT_RATE_LIMIT_WINDOW as Duration
+          ),
+        })
+      : undefined;
   }
 
   #calculateDeliverAt(options?: SendEventOptions) {
@@ -65,7 +78,7 @@ export class IngestSendEvent {
         return;
       }
 
-      return await $transaction(this.#prismaClient, async (tx) => {
+      const createdEvent = await $transaction(this.#prismaClient, async (tx) => {
         const externalAccount = options?.accountId
           ? await tx.externalAccount.upsert({
               where: {
@@ -106,6 +119,23 @@ export class IngestSendEvent {
 
         return eventLog;
       });
+
+      if (!createdEvent) return;
+
+      //rate limit
+      const result = await this.#rateLimiter?.limit(environment.organizationId);
+      if (result && !result.success) {
+        logger.info("IngestSendEvent: Rate limit exceeded", {
+          eventRecordId: createdEvent.id,
+          organizationId: environment.organizationId,
+          reset: result.reset,
+          limit: result.limit,
+        });
+        return;
+      }
+
+      await this.enqueueWorkerEvent(this.#prismaClient, createdEvent);
+      return createdEvent;
     } catch (error) {
       const prismaError = PrismaErrorSchema.safeParse(error);
 
@@ -151,8 +181,6 @@ export class IngestSendEvent {
       },
     });
 
-    await this.enqueueWorkerEvent(tx, eventLog);
-
     return eventLog;
   }
 
@@ -177,8 +205,6 @@ export class IngestSendEvent {
       },
     });
 
-    await this.enqueueWorkerEvent(tx, updatedEventLog);
-
     return updatedEventLog;
   }
 
diff --git a/apps/webapp/app/services/rateLimiter.server.ts b/apps/webapp/app/services/rateLimiter.server.ts
@@ -0,0 +1,105 @@
+import { Ratelimit } from "@upstash/ratelimit";
+import Redis, { RedisOptions } from "ioredis";
+import { env } from "~/env.server";
+import { logger } from "./logger.server";
+
+type Options = {
+  redis?: RedisOptions;
+  keyPrefix: string;
+  limiter: Limiter;
+  logSuccess?: boolean;
+  logFailure?: boolean;
+};
+
+export type Limiter = ConstructorParameters<typeof Ratelimit>[0]["limiter"];
+export type Duration = Parameters<typeof Ratelimit.slidingWindow>[1];
+export type RateLimitResponse = Awaited<ReturnType<Ratelimit["limit"]>>;
+
+export class RateLimiter {
+  #ratelimit: Ratelimit;
+
+  constructor(private readonly options: Options) {
+    const { redis, keyPrefix, limiter } = options;
+    const prefix = `ratelimit:${keyPrefix}`;
+    this.#ratelimit = new Ratelimit({
+      redis: createRedisRateLimitClient(
+        redis ?? {
+          port: env.REDIS_PORT,
+          host: env.REDIS_HOST,
+          username: env.REDIS_USERNAME,
+          password: env.REDIS_PASSWORD,
+          enableAutoPipelining: true,
+          ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
+        }
+      ),
+      limiter,
+      ephemeralCache: new Map(),
+      analytics: false,
+      prefix,
+    });
+
+    logger.info(`RateLimiter (${keyPrefix}): initialized`, {
+      keyPrefix,
+      redisKeyspace: prefix,
+    });
+  }
+
+  async limit(identifier: string, rate = 1): Promise<RateLimitResponse> {
+    const result = this.#ratelimit.limit(identifier, { rate });
+    const { success, limit, reset, remaining } = await result;
+
+    if (success && this.options.logSuccess) {
+      logger.info(`RateLimiter (${this.options.keyPrefix}): under rate limit`, {
+        limit,
+        reset,
+        remaining,
+        identifier,
+      });
+    }
+
+    //log these by default
+    if (!success && this.options.logFailure !== false) {
+      logger.info(`RateLimiter (${this.options.keyPrefix}): rate limit exceeded`, {
+        limit,
+        reset,
+        remaining,
+        identifier,
+      });
+    }
+
+    return result;
+  }
+}
+
+export function createRedisRateLimitClient(
+  redisOptions: RedisOptions
+): ConstructorParameters<typeof Ratelimit>[0]["redis"] {
+  const redis = new Redis(redisOptions);
+
+  return {
+    sadd: async <TData>(key: string, ...members: TData[]): Promise<number> => {
+      return redis.sadd(key, members as (string | number | Buffer)[]);
+    },
+    hset: <TValue>(
+      key: string,
+      obj: {
+        [key: string]: TValue;
+      }
+    ): Promise<number> => {
+      return redis.hset(key, obj);
+    },
+    eval: <TArgs extends unknown[], TData = unknown>(
+      ...args: [script: string, keys: string[], args: TArgs]
+    ): Promise<TData> => {
+      const script = args[0];
+      const keys = args[1];
+      const argsArray = args[2];
+      return redis.eval(
+        script,
+        keys.length,
+        ...keys,
+        ...(argsArray as (string | Buffer | number)[])
+      ) as Promise<TData>;
+    },
+  };
+}
diff --git a/apps/webapp/app/services/worker.server.ts b/apps/webapp/app/services/worker.server.ts
@@ -309,7 +309,7 @@ function getWorkerQueue() {
       },
       "events.deliverScheduled": {
         priority: 0, // smaller number = higher priority
-        maxAttempts: 5,
+        maxAttempts: 8,
         handler: async ({ id, payload }, job) => {
           const service = new DeliverScheduledEventService();
 
diff --git a/apps/webapp/package.json b/apps/webapp/package.json
@@ -101,7 +101,7 @@
     "@trigger.dev/yalt": "workspace:*",
     "@types/pg": "8.6.6",
     "@uiw/react-codemirror": "^4.19.5",
-    "@upstash/ratelimit": "^1.0.1",
+    "@upstash/ratelimit": "^1.1.3",
     "@whatwg-node/fetch": "^0.9.14",
     "assert-never": "^1.2.1",
     "aws4fetch": "^1.0.18",
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
