kitcn/dist/ratelimit/react/index.js

kitcn - React Query integration and CLI tools for Convex

'use client';
import { useConvex, useQuery } from "convex/react";
import { useCallback, useEffect, useMemo, useState } from "react";
import { makeFunctionReference } from "convex/server";

//#region src/ratelimit/core/calculate-rate-limit.ts
function calculateRatelimit(state, algorithm, now, count) {
	if (algorithm.kind === "fixedWindow") return calculateFixedWindow(state, algorithm, now, count);
	if (algorithm.kind === "tokenBucket") return calculateTokenBucket(state, algorithm, now, count);
	return calculateSlidingWindow(state, algorithm, now, count);
}
function calculateTokenBucket(state, config, now, count) {
	const ratePerMs = config.refillRate / config.interval;
	const initial = state ?? {
		value: config.maxTokens,
		ts: now
	};
	const elapsed = Math.max(0, now - initial.ts);
	const nextValue = Math.min(initial.value + elapsed * ratePerMs, config.maxTokens) - count;
	const retryAfter = nextValue < 0 ? Math.ceil(-nextValue / ratePerMs) : void 0;
	return {
		state: {
			value: nextValue,
			ts: now
		},
		retryAfter,
		remaining: Math.max(0, Math.floor(nextValue)),
		reset: retryAfter ? now + retryAfter : now,
		limit: config.maxTokens
	};
}
function calculateFixedWindow(state, config, now, count) {
	const windowStart = alignWindowStart(now, config.window, config.start);
	const initial = state ?? {
		value: config.capacity,
		ts: windowStart
	};
	const elapsedWindows = Math.max(0, Math.floor((now - initial.ts) / config.window));
	const replenished = Math.min(initial.value + config.limit * elapsedWindows, config.capacity);
	const ts = initial.ts + elapsedWindows * config.window;
	const nextValue = replenished - count;
	const retryAfter = nextValue < 0 ? ts + config.window * Math.ceil(-nextValue / config.limit) - now : void 0;
	return {
		state: {
			value: nextValue,
			ts
		},
		retryAfter,
		remaining: Math.max(0, Math.floor(nextValue)),
		reset: ts + config.window,
		limit: config.limit
	};
}
function calculateSlidingWindow(state, config, now, count) {
	const windowStart = alignWindowStart(now, config.window);
	const previousWindowStart = windowStart - config.window;
	const elapsedInWindow = now - windowStart;
	const previousWeight = Math.max(0, (config.window - elapsedInWindow) / config.window);
	let currentCount = 0;
	let previousCount = 0;
	if (state) {
		if (state.ts === windowStart) {
			currentCount = Math.max(0, state.value);
			if (state.auxTs === previousWindowStart) previousCount = Math.max(0, state.auxValue ?? 0);
		} else if (state.ts === previousWindowStart) previousCount = Math.max(0, state.value);
	}
	const projectedCurrent = currentCount + count;
	const projectedUsed = projectedCurrent + previousCount * previousWeight;
	const remaining = config.limit - projectedUsed;
	const retryAfter = remaining < 0 ? Math.max(1, config.window - elapsedInWindow) : void 0;
	return {
		state: {
			value: projectedCurrent,
			ts: windowStart,
			auxValue: previousCount,
			auxTs: previousWindowStart
		},
		retryAfter,
		remaining: Math.max(0, Math.floor(remaining)),
		reset: windowStart + config.window,
		limit: config.limit
	};
}
function alignWindowStart(now, window, start = 0) {
	const offsetNow = now - start;
	return start + Math.floor(offsetNow / window) * window;
}

//#endregion
//#region src/ratelimit/react/use-rate-limit.ts
function useRatelimit(getRatelimitValueQuery, options) {
	const [timeOffset, setTimeOffset] = useState(0);
	const [now, setNow] = useState(() => Date.now());
	const convex = useConvex();
	const getRatelimitValueQueryRef = useMemo(() => resolveGetRatelimitValueQuery(getRatelimitValueQuery), [getRatelimitValueQuery]);
	const { getServerTimeMutation, count, identifier, sampleShards } = options ?? {};
	const getServerTimeMutationRef = useMemo(() => getServerTimeMutation ? resolveGetServerTimeMutation(getServerTimeMutation) : void 0, [getServerTimeMutation]);
	useEffect(() => {
		if (!getServerTimeMutationRef) return;
		const clientTime = Date.now();
		convex.mutation(getServerTimeMutationRef, {}).then((serverTime) => {
			const latency = Date.now() - clientTime;
			setTimeOffset(serverTime - clientTime - latency / 2);
		});
	}, [convex, getServerTimeMutationRef]);
	const ratelimitData = useQuery(getRatelimitValueQueryRef, {
		identifier,
		sampleShards
	});
	const check = useCallback((ts, requestedCount) => {
		if (!ratelimitData) return;
		const serverTs = (ts ?? Date.now()) + timeOffset;
		const evaluation = evaluateSnapshot(ratelimitData, serverTs, requestedCount ?? count ?? 1);
		return {
			value: evaluation.value,
			ts: evaluation.ts - timeOffset,
			config: ratelimitData.config,
			shard: ratelimitData.shard,
			ok: evaluation.value >= 0,
			retryAt: evaluation.retryAfter ? serverTs + evaluation.retryAfter - timeOffset : void 0
		};
	}, [
		count,
		ratelimitData,
		timeOffset
	]);
	const current = check(now, count ?? 1);
	const response = useMemo(() => {
		if (!current) return {
			status: void 0,
			check
		};
		if (current.value < 0) return {
			status: {
				ok: false,
				retryAt: current.retryAt
			},
			check
		};
		return {
			status: {
				ok: true,
				retryAt: void 0
			},
			check
		};
	}, [check, current]);
	useEffect(() => {
		if (response.status?.ok !== false || !response.status.retryAt) return;
		const timeout = setTimeout(() => setNow(Date.now()), Math.max(0, response.status.retryAt - now));
		return () => clearTimeout(timeout);
	}, [
		now,
		response.status?.ok,
		response.status?.retryAt
	]);
	return response;
}
function resolveGetRatelimitValueQuery(ref) {
	if (typeof ref === "string") return makeFunctionReference(ref);
	return ref;
}
function resolveGetServerTimeMutation(ref) {
	if (typeof ref === "string") return makeFunctionReference(ref);
	return ref;
}
function evaluateSnapshot(snapshot, now, count) {
	const evaluated = calculateRatelimit(snapshot.config.kind === "slidingWindow" ? {
		value: Math.max(0, snapshot.config.limit - snapshot.value),
		ts: snapshot.ts
	} : {
		value: snapshot.value,
		ts: snapshot.ts
	}, snapshot.config, now, count);
	return {
		value: snapshot.config.kind === "slidingWindow" ? evaluated.retryAfter !== void 0 ? -1 : evaluated.remaining : evaluated.state.value,
		ts: evaluated.state.ts,
		retryAfter: evaluated.retryAfter
	};
}

//#endregion
export { useRatelimit };