Nick: fixed more stuff

This commit is contained in:
Nicolas 2024-10-01 16:04:39 -03:00
parent c0541cc990
commit 18f9cd09e1
5 changed files with 66 additions and 23 deletions

View File

@ -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<ConcurrencyCheckParams, undefined, undefined>,
res: Response<ConcurrencyCheckResponse>
) {
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 });
}

View File

@ -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
| {

View File

@ -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);

View File

@ -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", () => {

View File

@ -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));
}