All files rateLimiter.ts

95.45% Statements 21/22
70% Branches 7/10
100% Functions 3/3
94.74% Lines 18/19

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55  8x                                             8x 8x 8x       800x 800x 286308x   800x 800x 16x     800x 800x     800x 800x   800x 800x 800x         8x  
import { ArgsType } from "./cache"
import { delay } from "./delay"
 
export interface RateLimiterOptions<F extends (...args: unknown[]) => Promise<unknown>> {
    rateLimitId: (...args: ArgsType<F>) => string
    concurrentRequests?: ((rateLimitId: string) => number) | number
    maxTotalRequests?: number
}
 
/**
 * Controls how many concurrent execution can be invoked. At any given time
 * @param errorHandler logic when to retry
 * @param request request function
 * @signature
 *    P.concurrent(fn, options)
 * @example
 *    const requestToEndpoint = async (endpoint: "A" | "B" | "C" | "D", data: any) => { ...  } 
 *    const rate = P.concurrent(requestToEndpoint, ({ rateLimitId: (endpoint) => endpoint, concurrentRequests: 2 })
 *    // Only two request are fired to endpoint A
 *    const endpoints = await Promise.all([{ endpoint: A, data: any }, ...].map(async (obj)=>{ 
 *      return rate(obj.endpoint, obj.data)
 *    }))
 * @category Utility, Promise
 */
export function concurrent<F extends (...args: any[]) => Promise<unknown>>(request: F, options: RateLimiterOptions<F>): F {
    const requestCount: Record<string, number> = {}
    const stats = {
        totalRequests: 0
    }
 
    const requestMiddleware = (async (...args) => {
        while (stats.totalRequests >= (options.maxTotalRequests || 50)) {
            await delay(0)
        }
        const namespace = options.rateLimitId(...args as ArgsType<F>)
        if (requestCount[namespace] == null) {
            requestCount[namespace] = 0
        }
 
        const maxRequests = typeof options.concurrentRequests === 'function' ? options.concurrentRequests(namespace) : options.concurrentRequests ?? 3
        while (requestCount[namespace] >= maxRequests) {
            await delay(0)
        }
        requestCount[namespace] += 1
        stats.totalRequests += 1
 
        return request(...args).finally(() => {
            requestCount[namespace] -= 1
            stats.totalRequests -= 1
        })
    }) as F
 
 
    return requestMiddleware
}