Nick: redlock cache in auth

This commit is contained in:
Nicolas 2024-08-12 13:37:47 -04:00
parent 5a778f2c22
commit 25a899eae3
6 changed files with 119 additions and 64 deletions

View File

@ -7,7 +7,15 @@ import { RateLimiterRedis } from "rate-limiter-flexible";
import { setTraceAttributes } from '@hyperdx/node-opentelemetry';
import { sendNotification } from "../services/notification/email_notification";
import { Logger } from "../lib/logger";
import { redlock } from "../../src/services/redlock";
import { getValue } from "../../src/services/redis";
import { setValue } from "../../src/services/redis";
import { validate } from 'uuid';
function normalizedApiIsUuid(potentialUuid: string): boolean {
// Check if the string is a valid UUID
return validate(potentialUuid);
}
export async function authenticateUser(req, res, mode?: RateLimiterMode): Promise<AuthResponse> {
return withAuth(supaAuthenticateUser)(req, res, mode);
}
@ -54,17 +62,72 @@ export async function supaAuthenticateUser(
let subscriptionData: { team_id: string, plan: string } | null = null;
let normalizedApi: string;
let team_id: string;
let cacheKey= "";
let redLockKey = "";
const lockTTL = 5000; // 5 seconds
let teamId: string | null = null;
let priceId: string | null = null;
if (token == "this_is_just_a_preview_token") {
rateLimiter = getRateLimiter(RateLimiterMode.Preview, token);
team_id = "preview";
teamId = "preview";
} else {
normalizedApi = parseApi(token);
if(!normalizedApiIsUuid(normalizedApi)){
return {
success: false,
error: "Unauthorized: Invalid token",
status: 401,
};
}
cacheKey = `api_key:${normalizedApi}`;
redLockKey = `redlock:${cacheKey}`;
try{
const lock = await redlock.acquire([redLockKey], lockTTL);
const { data, error } = await supabase_service.rpc(
'get_key_and_price_id_2', { api_key: normalizedApi }
);
try{
const teamIdPriceId = await getValue(cacheKey);
if(teamIdPriceId){
const { team_id, price_id } = JSON.parse(teamIdPriceId);
teamId = team_id;
priceId = price_id;
}
else{
const { data, error } = await supabase_service.rpc(
'get_key_and_price_id_2', { api_key: normalizedApi }
);
if(error){
Logger.error(`RPC ERROR (get_key_and_price_id_2): ${error.message}`);
return {
success: false,
error: "The server seems overloaded. Please contact hello@firecrawl.com if you aren't sending too many requests at once.",
status: 500,
};
}
if (!data || data.length === 0) {
Logger.warn(`Error fetching api key: ${error.message} or data is empty`);
// TODO: change this error code ?
return {
success: false,
error: "Unauthorized: Invalid token",
status: 401,
};
}
else {
teamId = data[0].team_id;
priceId = data[0].price_id;
}
}
}finally{
await lock.release();
}
}catch(error){
Logger.error(`Error acquiring the rate limiter lock: ${error}`);
}
// get_key_and_price_id_2 rpc definition:
// create or replace function get_key_and_price_id_2(api_key uuid)
// returns table(key uuid, team_id uuid, price_id text) as $$
@ -82,30 +145,12 @@ export async function supaAuthenticateUser(
// end;
// $$ language plpgsql;
if (error) {
Logger.warn(`Error fetching key and price_id: ${error.message}`);
} else {
// console.log('Key and Price ID:', data);
}
if (error || !data || data.length === 0) {
Logger.warn(`Error fetching api key: ${error.message} or data is empty`);
return {
success: false,
error: "Unauthorized: Invalid token",
status: 401,
};
}
const internal_team_id = data[0].team_id;
team_id = internal_team_id;
const plan = getPlanByPriceId(data[0].price_id);
const plan = getPlanByPriceId(priceId);
// HyperDX Logging
setTrace(team_id, normalizedApi);
setTrace(teamId, normalizedApi);
subscriptionData = {
team_id: team_id,
team_id: teamId,
plan: plan
}
switch (mode) {
@ -134,7 +179,7 @@ export async function supaAuthenticateUser(
}
}
const team_endpoint_token = token === "this_is_just_a_preview_token" ? iptoken : team_id;
const team_endpoint_token = token === "this_is_just_a_preview_token" ? iptoken : teamId;
try {
await rateLimiter.consume(team_endpoint_token);
@ -147,7 +192,13 @@ export async function supaAuthenticateUser(
const startDate = new Date();
const endDate = new Date();
endDate.setDate(endDate.getDate() + 7);
// await sendNotification(team_id, NotificationType.RATE_LIMIT_REACHED, startDate.toISOString(), endDate.toISOString());
// TODO: cache 429 for a few minuts
if(teamId && priceId && mode !== RateLimiterMode.Preview){
await setValue(cacheKey, JSON.stringify({team_id: teamId, price_id: priceId}), 60 * 5);
}
return {
success: false,
error: `Rate limit exceeded. Consumed points: ${rateLimiterRes.consumedPoints}, Remaining points: ${rateLimiterRes.remainingPoints}. Upgrade your plan at https://firecrawl.dev/pricing for increased rate limits or please retry after ${secs}s, resets at ${retryDate}`,

View File

@ -4,37 +4,12 @@ import { sendNotification } from "../notification/email_notification";
import { supabase_service } from "../supabase";
import { Logger } from "../../lib/logger";
import { getValue, setValue } from "../redis";
import Redlock from "redlock";
import Client from "ioredis";
import { redlock } from "../redlock";
const FREE_CREDITS = 500;
const redlock = new Redlock(
// You should have one client for each independent redis node
// or cluster.
[new Client(process.env.REDIS_RATE_LIMIT_URL)],
{
// The expected clock drift; for more details see:
// http://redis.io/topics/distlock
driftFactor: 0.01, // multiplied by lock ttl to determine drift time
// The max number of times Redlock will attempt to lock a resource
// before erroring.
retryCount: 5,
// the time in ms between attempts
retryDelay: 100, // time in ms
// the max time in ms randomly added to retries
// to improve performance under high contention
// see https://www.awsarchitectureblog.com/2015/03/backoff.html
retryJitter: 200, // time in ms
// The minimum remaining time on a lock before an extension is automatically
// attempted with the `using` API.
automaticExtensionThreshold: 500, // time in ms
}
);
export async function billTeam(team_id: string, credits: number) {
return withAuth(supaBillTeam)(team_id, credits);
}

View File

@ -0,0 +1,29 @@
import Redlock from "redlock";
import Client from "ioredis";
export const redlock = new Redlock(
// You should have one client for each independent redis node
// or cluster.
[new Client(process.env.REDIS_RATE_LIMIT_URL)],
{
// The expected clock drift; for more details see:
// http://redis.io/topics/distlock
driftFactor: 0.01, // multiplied by lock ttl to determine drift time
// The max number of times Redlock will attempt to lock a resource
// before erroring.
retryCount: 5,
// the time in ms between attempts
retryDelay: 100, // time in ms
// the max time in ms randomly added to retries
// to improve performance under high contention
// see https://www.awsarchitectureblog.com/2015/03/backoff.html
retryJitter: 200, // time in ms
// The minimum remaining time on a lock before an extension is automatically
// attempted with the `using` API.
automaticExtensionThreshold: 500, // time in ms
}
);

View File

@ -36,9 +36,9 @@ class FirecrawlApp {
* @param {Params | null} params - Additional parameters for the scrape request.
* @returns {Promise<ScrapeResponse>} The response from the scrape operation.
*/
scrapeUrl(url, params = null) {
var _a;
return __awaiter(this, void 0, void 0, function* () {
scrapeUrl(url_1) {
return __awaiter(this, arguments, void 0, function* (url, params = null) {
var _a;
const headers = {
"Content-Type": "application/json",
Authorization: `Bearer ${this.apiKey}`,
@ -79,8 +79,8 @@ class FirecrawlApp {
* @param {Params | null} params - Additional parameters for the search request.
* @returns {Promise<SearchResponse>} The response from the search operation.
*/
search(query, params = null) {
return __awaiter(this, void 0, void 0, function* () {
search(query_1) {
return __awaiter(this, arguments, void 0, function* (query, params = null) {
const headers = {
"Content-Type": "application/json",
Authorization: `Bearer ${this.apiKey}`,
@ -119,8 +119,8 @@ class FirecrawlApp {
* @param {string} idempotencyKey - Optional idempotency key for the request.
* @returns {Promise<CrawlResponse | any>} The response from the crawl operation.
*/
crawlUrl(url, params = null, waitUntilDone = true, pollInterval = 2, idempotencyKey) {
return __awaiter(this, void 0, void 0, function* () {
crawlUrl(url_1) {
return __awaiter(this, arguments, void 0, function* (url, params = null, waitUntilDone = true, pollInterval = 2, idempotencyKey) {
const headers = this.prepareHeaders(idempotencyKey);
let jsonData = { url };
if (params) {

View File

@ -1,12 +1,12 @@
{
"name": "@mendable/firecrawl-js",
"version": "0.0.34",
"version": "0.0.36",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "@mendable/firecrawl-js",
"version": "0.0.34",
"version": "0.0.36",
"license": "MIT",
"dependencies": {
"axios": "^1.6.8",

View File

@ -1,6 +1,6 @@
{
"name": "@mendable/firecrawl-js",
"version": "0.0.35",
"version": "0.0.36",
"description": "JavaScript SDK for Firecrawl API",
"main": "build/cjs/index.js",
"types": "types/index.d.ts",