diff --git a/package-lock.json b/package-lock.json index 199608e..5e60502 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10,6 +10,7 @@ "license": "MIT", "dependencies": { "@isaacs/ttlcache": "^1.4.1", + "async-mutex": "^0.5.0", "bytes": "^3.1.2", "cheerio": "^1.0.0", "country-emoji": "^1.5.6", @@ -2132,6 +2133,15 @@ "integrity": "sha512-htCUDlxyyCLMgaM3xXg0C0LW2xqfuQ6p05pCEIsXuyQ+a1koYKTuBMzRNwmybfLgvJDMd0r1LTn4+E0Ti6C2AA==", "license": "MIT" }, + "node_modules/async-mutex": { + "version": "0.5.0", + "resolved": "https://registry.npmjs.org/async-mutex/-/async-mutex-0.5.0.tgz", + "integrity": "sha512-1A94B18jkJ3DYq284ohPxoXbfTA5HsQ7/Mf4DEhcyLx3Bz27Rh59iScbB6EPiP+B+joue6YCxcMXSbFC1tZKwA==", + "license": "MIT", + "dependencies": { + "tslib": "^2.4.0" + } + }, "node_modules/asynckit": { "version": "0.4.0", "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz", @@ -6912,6 +6922,12 @@ } } }, + "node_modules/tslib": { + "version": "2.8.1", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", + "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", + "license": "0BSD" + }, "node_modules/type-check": { "version": "0.4.0", "resolved": "https://registry.npmjs.org/type-check/-/type-check-0.4.0.tgz", diff --git a/package.json b/package.json index ab6c448..8cfd572 100644 --- a/package.json +++ b/package.json @@ -32,6 +32,7 @@ }, "dependencies": { "@isaacs/ttlcache": "^1.4.1", + "async-mutex": "^0.5.0", "bytes": "^3.1.2", "cheerio": "^1.0.0", "country-emoji": "^1.5.6", diff --git a/src/utils/Fetcher.test.ts b/src/utils/Fetcher.test.ts index 2d52421..ae1f064 100644 --- a/src/utils/Fetcher.test.ts +++ b/src/utils/Fetcher.test.ts @@ -162,14 +162,15 @@ describe('fetch', () => { await expect(fetcher.text(ctx, new URL('https://some-timeout-url.test/'), { timeout: 10 })).rejects.toBeInstanceOf(DOMException); }); - test('full queue throws an error', async () => { - fetchMock.get('https://some-full-queue-url.test/', 'some text'); + test('queues requests and throws if queue is full', async () => { + fetchMock.get('https://some-full-queue-url.test/', 200, { delay: 20 }); const allPromises = Promise.all([ - fetcher.text(ctx, new URL('https://some-full-queue-url.test/'), { queueLimit: 3 }), - fetcher.text(ctx, new URL('https://some-full-queue-url.test/'), { queueLimit: 3 }), - fetcher.text(ctx, new URL('https://some-full-queue-url.test/'), { queueLimit: 3 }), - fetcher.text(ctx, new URL('https://some-full-queue-url.test/'), { queueLimit: 3 }), + fetcher.text(ctx, new URL('https://some-full-queue-url.test/'), { queueLimit: 1, queueErrorLimit: 2 }), + fetcher.text(ctx, new URL('https://some-full-queue-url.test/'), { queueLimit: 1, queueErrorLimit: 2 }), + fetcher.text(ctx, new URL('https://some-full-queue-url.test/'), { queueLimit: 1, queueErrorLimit: 2 }), + fetcher.text(ctx, new URL('https://some-full-queue-url.test/'), { queueLimit: 1, queueErrorLimit: 2 }), + fetcher.text(ctx, new URL('https://some-full-queue-url.test/'), { queueLimit: 1, queueErrorLimit: 2 }), ]); await expect(allPromises).rejects.toBeInstanceOf(QueueIsFullError); diff --git a/src/utils/Fetcher.ts b/src/utils/Fetcher.ts index 549d14d..1b61e58 100644 --- a/src/utils/Fetcher.ts +++ b/src/utils/Fetcher.ts @@ -1,6 +1,7 @@ import CachePolicy from 'http-cache-semantics'; import TTLCache from '@isaacs/ttlcache'; import winston from 'winston'; +import { Mutex } from 'async-mutex'; import { Cookie, CookieJar } from 'tough-cookie'; import { Context, TIMEOUT } from '../types'; import { BlockedError, NotFoundError, QueueIsFullError } from '../error'; @@ -14,11 +15,6 @@ interface HttpCacheItem { body: string; } -interface HostQueue { - promise: Promise; - count: number; -} - interface FlareSolverrResult { status: string; message: string; @@ -46,6 +42,7 @@ interface FlareSolverrResult { type CustomRequestInit = RequestInit & { queueLimit?: number; + queueErrorLimit?: number; timeout?: number; }; @@ -53,10 +50,12 @@ export class Fetcher { private readonly logger: winston.Logger; private readonly httpCache: TTLCache; - private readonly hostQueues = new Map(); private readonly hostUserAgentMap = new Map(); private readonly cookieJar = new CookieJar(); + private readonly hostQueueCount = new Map(); + private readonly countMutex = new Mutex(); + constructor(logger: winston.Logger) { this.logger = logger; @@ -119,7 +118,7 @@ export class Fetcher { this.logger.info(`Query FlareSolverr for ${url.href}`, ctx); const body = { cmd: 'request.get', url: url.href, session: 'default' }; - const challengeResult = await (await this.queuedFetch(ctx, new URL(flareSolverrEndpoint), { method: 'POST', body: JSON.stringify(body), headers: { 'Content-Type': 'application/json' } })).json() as FlareSolverrResult; + const challengeResult = await (await this.queuedFetch(ctx, new URL(flareSolverrEndpoint), { method: 'POST', body: JSON.stringify(body), headers: { 'Content-Type': 'application/json' }, queueLimit: 1 })).json() as FlareSolverrResult; if (challengeResult.status !== 'ok') { this.logger.warn(`FlareSolverr issue: ${JSON.stringify(challengeResult)}`, ctx); throw new BlockedError('flaresolverr_failed', {}); @@ -221,35 +220,46 @@ export class Fetcher { return response; }; - private readonly queuedFetch = async (ctx: Context, url: URL, init?: CustomRequestInit): Promise => { - let queue = this.hostQueues.get(url.host); - if (!queue) { - queue = { promise: Promise.resolve(), count: 0 }; - this.hostQueues.set(url.host, queue); - } + private readonly getHostQueueCount = (host: string): number => this.hostQueueCount.get(host) ?? 0; - queue.count++; - - if (queue.count > (init?.queueLimit ?? 10)) { - throw new QueueIsFullError(); - } - - let resolveCurrent: () => void; - const currentPromise = new Promise( - (resolve) => { resolveCurrent = resolve; }, - ); - - const fetchPromise = queue.promise.then(async () => { - try { - return await this.fetchWithTimeout(ctx, url, init); - } finally { - resolveCurrent(); - queue.count--; + private readonly lockFetchSlot = async (host: string, queueErrorLimit: number) => { + await this.countMutex.runExclusive(() => { + if (this.getHostQueueCount(host) > queueErrorLimit) { + throw new QueueIsFullError(); } + + this.hostQueueCount.set(host, this.getHostQueueCount(host) + 1); }); + }; - queue.promise = currentPromise.catch(/* istanbul ignore next */ () => undefined); + private readonly unlockFetchSlot = async (host: string) => { + await this.countMutex.runExclusive(() => { + this.hostQueueCount.set(host, Math.max(0, this.getHostQueueCount(host) - 1)); + }); + }; - return fetchPromise; + private readonly waitForHostQueueCount = async (host: string, queueLimit: number, queueErrorLimit: number): Promise => { + while (this.getHostQueueCount(host) > queueLimit) { + if (this.getHostQueueCount(host) > queueErrorLimit) { + // Very unlikely to happen.. + throw new QueueIsFullError(); + } + + await new Promise(resolve => setTimeout(resolve, 100)); + } + }; + + private readonly queuedFetch = async (ctx: Context, url: URL, init?: CustomRequestInit): Promise => { + const queueLimit = init?.queueLimit ?? 5; + const queueErrorLimit = init?.queueErrorLimit ?? 10; + + await this.lockFetchSlot(url.host, queueErrorLimit); + await this.waitForHostQueueCount(url.host, queueLimit, queueErrorLimit); + + try { + return await this.fetchWithTimeout(ctx, url, init); + } finally { + await this.unlockFetchSlot(url.host); + } }; }