Separate /engine/ rate limiter

Author: matt-aitken

apps/webapp/app/entry.server.tsx

@@ -208,6 +208,7 @@ const sqsEventConsumer = singleton("sqsEventConsumer", getSharedSqsEventConsumer
 singleton("RunEngineEventBusHandlers", registerRunEngineEventBusHandlers);
 
 export { apiRateLimiter } from "./services/apiRateLimit.server";
+export { engineRateLimiter } from "./services/engineRateLimit.server";
 export { socketIo } from "./v3/handleSocketIo.server";
 export { wss } from "./v3/handleWebsockets.server";
 export { registryProxy } from "./v3/registryProxy.server";

apps/webapp/app/env.server.ts

@@ -506,6 +506,22 @@ const EnvironmentSchema = z.object({
     .string()
     .default(process.env.REDIS_TLS_DISABLED ?? "false"),
 
+  //API Rate limiting
+  /**
+   * @example "60s"
+   * @example "1m"
+   * @example "1h"
+   * @example "1d"
+   * @example "1000ms"
+   * @example "1000s"
+   */
+  RUN_ENGINE_RATE_LIMIT_REFILL_INTERVAL: z.string().default("10s"), // refill 250 tokens every 10 seconds
+  RUN_ENGINE_RATE_LIMIT_MAX: z.coerce.number().int().default(1200), // allow bursts of 750 requests
+  RUN_ENGINE_RATE_LIMIT_REFILL_RATE: z.coerce.number().int().default(400), // refix 250 tokens every 10 seconds
+  RUN_ENGINE_RATE_LIMIT_REQUEST_LOGS_ENABLED: z.string().default("0"),
+  RUN_ENGINE_RATE_LIMIT_REJECTION_LOGS_ENABLED: z.string().default("1"),
+  RUN_ENGINE_RATE_LIMIT_LIMITER_LOGS_ENABLED: z.string().default("0"),
+
   /** How long should the presence ttl last */
   DEV_PRESENCE_TTL_MS: z.coerce.number().int().default(30_000),
   DEV_PRESENCE_POLL_INTERVAL_MS: z.coerce.number().int().default(5_000),

apps/webapp/app/services/apiRateLimit.server.ts

@@ -59,9 +59,6 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
     "/api/v1/usage/ingest",
     "/api/v1/auth/jwt/claims",
     /^\/api\/v1\/runs\/[^\/]+\/attempts$/, // /api/v1/runs/$runFriendlyId/attempts
-    // run engine DEV endpoints
-    "/api/v1/dev/dequeue",
-    "/api/v1/dev/presence",
   ],
   log: {
     rejections: env.API_RATE_LIMIT_REJECTION_LOGS_ENABLED === "1",

apps/webapp/app/services/engineRateLimit.server.ts

@@ -0,0 +1,34 @@
+import { env } from "~/env.server";
+import { authenticateAuthorizationHeader } from "./apiAuth.server";
+import { authorizationRateLimitMiddleware } from "./authorizationRateLimitMiddleware.server";
+import { Duration } from "./rateLimiter.server";
+
+export const engineRateLimiter = authorizationRateLimitMiddleware({
+  redis: {
+    port: env.RATE_LIMIT_REDIS_PORT,
+    host: env.RATE_LIMIT_REDIS_HOST,
+    username: env.RATE_LIMIT_REDIS_USERNAME,
+    password: env.RATE_LIMIT_REDIS_PASSWORD,
+    tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true",
+    clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1",
+  },
+  keyPrefix: "engine",
+  defaultLimiter: {
+    type: "tokenBucket",
+    refillRate: env.RUN_ENGINE_RATE_LIMIT_REFILL_RATE,
+    interval: env.RUN_ENGINE_RATE_LIMIT_REFILL_INTERVAL as Duration,
+    maxTokens: env.RUN_ENGINE_RATE_LIMIT_MAX,
+  },
+  limiterCache: {
+    fresh: 60_000 * 10, // Data is fresh for 10 minutes
+    stale: 60_000 * 20, // Date is stale after 20 minutes
+  },
+  pathMatchers: [/^\/engine/],
+  // Allow /api/v1/tasks/:id/callback/:secret
+  pathWhiteList: [],
+  log: {
+    rejections: env.RUN_ENGINE_RATE_LIMIT_REJECTION_LOGS_ENABLED === "1",
+    requests: env.RUN_ENGINE_RATE_LIMIT_REQUEST_LOGS_ENABLED === "1",
+    limiter: env.RUN_ENGINE_RATE_LIMIT_LIMITER_LOGS_ENABLED === "1",
+  },
+});

apps/webapp/server.ts

@@ -43,6 +43,7 @@ if (process.env.HTTP_SERVER_DISABLED !== "true") {
   const wss: WebSocketServer | undefined = build.entry.module.wss;
   const registryProxy: RegistryProxy | undefined = build.entry.module.registryProxy;
   const apiRateLimiter: RateLimitMiddleware = build.entry.module.apiRateLimiter;
+  const engineRateLimiter: RateLimitMiddleware = build.entry.module.engineRateLimiter;
   const runWithHttpContext: RunWithHttpContextFunction = build.entry.module.runWithHttpContext;
 
   if (registryProxy && process.env.ENABLE_REGISTRY_PROXY === "true") {
@@ -95,6 +96,7 @@ if (process.env.HTTP_SERVER_DISABLED !== "true") {
     }
 
     app.use(apiRateLimiter);
+    app.use(engineRateLimiter);
 
     app.all(
       "*",