diff --git a/apps/api/src/controllers/v1/concurrency-check.ts b/apps/api/src/controllers/v1/concurrency-check.ts new file mode 100644 index 00000000..6ed4fa55 --- /dev/null +++ b/apps/api/src/controllers/v1/concurrency-check.ts @@ -0,0 +1,25 @@ +import { authenticateUser } from "../auth"; +import { + ConcurrencyCheckParams, + ConcurrencyCheckResponse, + RequestWithAuth, +} from "./types"; +import { RateLimiterMode } from "../../types"; +import { Response } from "express"; +import { redisConnection } from "../../services/queue-service"; +// Basically just middleware and error wrapping +export async function concurrencyCheckController( + req: RequestWithAuth, + res: Response +) { + const concurrencyLimiterKey = "concurrency-limiter:" + req.params.teamId; + const now = Date.now(); + const activeJobsOfTeam = await redisConnection.zrangebyscore( + concurrencyLimiterKey, + now, + Infinity + ); + return res + .status(200) + .json({ success: true, concurrency: activeJobsOfTeam.length }); +} diff --git a/apps/api/src/controllers/v1/types.ts b/apps/api/src/controllers/v1/types.ts index 45db51b5..3781eb78 100644 --- a/apps/api/src/controllers/v1/types.ts +++ b/apps/api/src/controllers/v1/types.ts @@ -294,6 +294,17 @@ export type CrawlStatusParams = { jobId: string; }; +export type ConcurrencyCheckParams = { + teamId: string; +}; + +export type ConcurrencyCheckResponse = + | ErrorResponse + | { + success: true; + concurrency: number; + }; + export type CrawlStatusResponse = | ErrorResponse | { diff --git a/apps/api/src/routes/v1.ts b/apps/api/src/routes/v1.ts index 49a41ce7..b0ceceb4 100644 --- a/apps/api/src/routes/v1.ts +++ b/apps/api/src/routes/v1.ts @@ -16,6 +16,7 @@ import { isUrlBlocked } from "../scraper/WebScraper/utils/blocklist"; import { crawlCancelController } from "../controllers/v1/crawl-cancel"; import { Logger } from "../lib/logger"; import { scrapeStatusController } from "../controllers/v1/scrape-status"; +import { concurrencyCheckController } from "../controllers/v1/concurrency-check"; // import { crawlPreviewController } from "../../src/controllers/v1/crawlPreview"; // import { crawlJobStatusPreviewController } from "../../src/controllers/v1/status"; // import { searchController } from "../../src/controllers/v1/search"; @@ -140,11 +141,19 @@ v1Router.get( wrap(scrapeStatusController) ); +v1Router.get( + "/concurrency-check", + authMiddleware(RateLimiterMode.CrawlStatus), + wrap(concurrencyCheckController) +); + v1Router.ws( "/crawl/:jobId", crawlStatusWSController ); + + // v1Router.post("/crawlWebsitePreview", crawlPreviewController); diff --git a/apps/api/src/services/rate-limiter.test.ts b/apps/api/src/services/rate-limiter.test.ts index 3e252301..ba4a0a73 100644 --- a/apps/api/src/services/rate-limiter.test.ts +++ b/apps/api/src/services/rate-limiter.test.ts @@ -49,7 +49,7 @@ describe("Rate Limiter Service", () => { "nonexistent" as RateLimiterMode, "test-prefix:someToken" ); - expect(limiter).toBe(serverRateLimiter); + expect(limiter.points).toBe(serverRateLimiter.points); }); it("should return the correct rate limiter based on mode and plan", () => { @@ -210,7 +210,7 @@ describe("Rate Limiter Service", () => { "test-prefix:someToken", "starter" ); - expect(limiter2.points).toBe(3); + expect(limiter2.points).toBe(10); const limiter3 = getRateLimiter( "crawl" as RateLimiterMode, @@ -233,7 +233,7 @@ describe("Rate Limiter Service", () => { "test-prefix:someToken", "starter" ); - expect(limiter2.points).toBe(20); + expect(limiter2.points).toBe(100); const limiter3 = getRateLimiter( "scrape" as RateLimiterMode, @@ -263,14 +263,14 @@ describe("Rate Limiter Service", () => { "test-prefix:someToken", "starter" ); - expect(limiter2.points).toBe(20); + expect(limiter2.points).toBe(50); const limiter3 = getRateLimiter( "search" as RateLimiterMode, "test-prefix:someToken", "standard" ); - expect(limiter3.points).toBe(40); + expect(limiter3.points).toBe(50); }); it("should return the correct rate limiter for 'preview' mode", () => { diff --git a/apps/api/src/services/rate-limiter.ts b/apps/api/src/services/rate-limiter.ts index e0fc5646..21e05948 100644 --- a/apps/api/src/services/rate-limiter.ts +++ b/apps/api/src/services/rate-limiter.ts @@ -132,8 +132,22 @@ export function getRateLimiterPoints( token?: string, plan?: string, teamId?: string -) { +) : number { + const rateLimitConfig = RATE_LIMITS[mode]; // {default : 5} + + if (!rateLimitConfig) return RATE_LIMITS.account.default; + const points : number = + rateLimitConfig[makePlanKey(plan)] || rateLimitConfig.default; // 5 + return points; +} + +export function getRateLimiter( + mode: RateLimiterMode, + token?: string, + plan?: string, + teamId?: string + ) : RateLimiterRedis { if (token && testSuiteTokens.some(testToken => token.includes(testToken))) { return testSuiteRateLimiter; } @@ -145,22 +159,6 @@ export function getRateLimiterPoints( if(teamId && manual.includes(teamId)) { return manualRateLimiter; } - - const rateLimitConfig = RATE_LIMITS[mode]; // {default : 5} - - if (!rateLimitConfig) return serverRateLimiter; - - const points = - rateLimitConfig[makePlanKey(plan)] || rateLimitConfig.default || rateLimitConfig; // 5 - - return points; -} - -export function getRateLimiter( - mode: RateLimiterMode, - token?: string, - plan?: string, - teamId?: string -) { + return createRateLimiter(`${mode}-${makePlanKey(plan)}`, getRateLimiterPoints(mode, token, plan, teamId)); }