mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
feat: add per-user rate limiting and concurrent job limits
Redis sliding window (sorted set) enforces per-user request rate limits via the rateLimitPerUser DB setting. Concurrent job limits checked at the HTTP layer before enqueue via maxConcurrentJobsPerUser setting. Both default to 0 (unlimited).
This commit is contained in:
@@ -3,15 +3,18 @@ import { mkdir, rm } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { extname, join } from "node:path";
|
||||
import { ANALYTICS_EVENTS, getBundleForTool, TOOL_BUNDLE_MAP, TOOLS } from "@snapotter/shared";
|
||||
import { and, inArray, sql } from "drizzle-orm";
|
||||
import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
|
||||
import type { z } from "zod";
|
||||
import { env } from "../config.js";
|
||||
import { db, schema } from "../db/index.js";
|
||||
import { enqueueToolJob, waitForJob } from "../jobs/enqueue.js";
|
||||
import { trackEvent } from "../lib/analytics.js";
|
||||
import { formatZodErrors, stripInternalPaths } from "../lib/errors.js";
|
||||
import { isToolInstalled } from "../lib/feature-status.js";
|
||||
import { getObjectBuffer, putObject } from "../lib/object-storage.js";
|
||||
import { resolveToolPool, shouldSkipSyncWindow } from "../lib/pool.js";
|
||||
import { getSettingNumber } from "../lib/settings-helpers.js";
|
||||
import { type ReceivedUpload, receiveUpload } from "../lib/upload-stream.js";
|
||||
import { InputValidationError } from "../modality/contract.js";
|
||||
import { inputHandlerFor } from "../modality/input-handler.js";
|
||||
@@ -441,6 +444,29 @@ export function createToolRoute<T>(app: FastifyInstance, config: ToolRouteConfig
|
||||
});
|
||||
}
|
||||
|
||||
// Check per-user concurrent job limit before enqueuing
|
||||
const userId = getAuthUser(request)?.id ?? null;
|
||||
const maxConcurrent = await getSettingNumber("maxConcurrentJobsPerUser", 0);
|
||||
if (maxConcurrent > 0 && userId) {
|
||||
const activeJobs = await db
|
||||
.select({ count: sql<number>`count(*)::int` })
|
||||
.from(schema.jobs)
|
||||
.where(
|
||||
and(
|
||||
sql`${schema.jobs.userId} = ${userId}`,
|
||||
inArray(schema.jobs.status, ["queued", "processing"]),
|
||||
),
|
||||
);
|
||||
|
||||
if (activeJobs[0].count >= maxConcurrent) {
|
||||
return reply.status(429).send({
|
||||
error: "Too many concurrent jobs. Please wait for existing jobs to complete.",
|
||||
activeJobs: activeJobs[0].count,
|
||||
limit: maxConcurrent,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
const startTime = Date.now();
|
||||
const pool = resolveToolPool(config.toolId);
|
||||
|
||||
@@ -451,7 +477,7 @@ export function createToolRoute<T>(app: FastifyInstance, config: ToolRouteConfig
|
||||
await enqueueToolJob({
|
||||
jobId,
|
||||
toolId: config.toolId,
|
||||
userId: getAuthUser(request)?.id ?? null,
|
||||
userId,
|
||||
pool,
|
||||
inputRefs,
|
||||
filename,
|
||||
|
||||
Reference in New Issue
Block a user