[pull] master from apify:master

packages/core/src/storages/throttling_request_manager.ts

@@ -6,7 +6,8 @@ import { z } from 'zod';
 import type { Configuration } from '../configuration.js';
 import { asyncifyIterable } from '../iterables.js';
 import type { CrawleeLogger } from '../log.js';
-import type { Request, Source } from '../request.js';
+import type { Source } from '../request.js';
+import { Request } from '../request.js';
 import { serviceLocator } from '../service_locator.js';
 import { normalizeHostname } from '../url.js';
 import { parseArgument, schemas } from '../validators.js';
@@ -208,6 +209,12 @@ function newDomainState(domain: string): DomainState {
 
 const DEFAULT_PERSIST_STATE_KEY = 'CRAWLEE_THROTTLED_DOMAINS';
 
+/**
+ * How many requests one fetch may move out of the wrapped manager before giving up on finding a dispatchable
+ * one. The next fetch resumes where it left off, and a migrated request is never walked again.
+ */
+const MAX_INNER_MIGRATIONS = 1000;
+
 /**
  * A request manager that wraps another one and paces requests per domain.
  *
@@ -304,6 +311,9 @@ export class ThrottlingRequestManager<T extends IRequestManager = IRequestManage
     /** Batches still being added in the background; keeps {@apilink ThrottlingRequestManager.checkReadiness} honest. */
     #inProgressBatchCount = 0;
 
+    /** Requests moved into a sub-queue, which both managers count. Subtracted below; not persisted. */
+    #migratedFromInner = 0;
+
     /** The latest {@link setExpectedRequestProcessingTimeSecs} hint, kept for a wrapped manager resolved later. */
     #expectedRequestProcessingSecs?: number;
 
@@ -767,6 +777,7 @@ export class ThrottlingRequestManager<T extends IRequestManager = IRequestManage
         this.#warnIfNotRoutable(requestLike);
 
         const manager = await this.#selectManagerOrThrow(requestLike.url ?? '');
+
         return manager.addRequest(requestLike, options);
     }
 
@@ -856,6 +867,7 @@ export class ThrottlingRequestManager<T extends IRequestManager = IRequestManage
         options?: RequestQueueOperationOptions,
     ): Promise<RequestQueueOperationInfo | null> {
         const manager = await this.#managerHolding(request);
+
         return manager.reclaimRequest(request, options);
     }
 
@@ -882,15 +894,15 @@ export class ThrottlingRequestManager<T extends IRequestManager = IRequestManage
     }
 
     async getTotalCount(): Promise<number> {
-        return this.#sumOverManagers((manager) => manager.getTotalCount());
+        return (await this.#sumOverManagers((manager) => manager.getTotalCount())) - this.#migratedFromInner;
     }
 
     async getPendingCount(): Promise<number> {
         return this.#sumOverManagers((manager) => manager.getPendingCount());
     }
 
     async getHandledCount(): Promise<number> {
-        return this.#sumOverManagers((manager) => manager.getHandledCount());
+        return (await this.#sumOverManagers((manager) => manager.getHandledCount())) - this.#migratedFromInner;
     }
 
     /**
@@ -941,9 +953,11 @@ export class ThrottlingRequestManager<T extends IRequestManager = IRequestManage
             }
         }
 
+        const inner = await this.#getInner();
+
         const probed = (
             await Promise.all([
-                (await this.#getInner()).checkReadiness(),
+                inner.checkReadiness(),
                 ...dispatchable.map(async (subManager) => (await subManager).checkReadiness()),
             ])
         ).reduce(joinRequestSourceStatuses);
@@ -1009,6 +1023,8 @@ export class ThrottlingRequestManager<T extends IRequestManager = IRequestManage
         const subManagers = await this.#getSubManagers();
         await Promise.all(subManagers.map(async (manager) => manager.purge?.()));
 
+        this.#migratedFromInner = 0;
+
         for (const state of this.domainStates.values()) {
             state.consecutive429Count = 0;
             state.backoffUntil = 0;
@@ -1079,7 +1095,7 @@ export class ThrottlingRequestManager<T extends IRequestManager = IRequestManage
             state.crawlDelayUntil = crawlDelayUntilBefore;
         }
 
-        const request = await (await this.#getInner()).fetchNextRequest<R>();
+        const request = await this.#fetchFromInner<R>();
 
         if (request !== null) {
             this.#inFlightFromInner.add(request.id ?? request.uniqueKey);
@@ -1101,6 +1117,59 @@ export class ThrottlingRequestManager<T extends IRequestManager = IRequestManage
         return request;
     }
 
+    /**
+     * The next request the wrapped manager can offer whose domain is not being held back.
+     *
+     * Its requests are not stored per domain, so they cannot be skipped the way a sub-queue can - the only
+     * way to find out which domain one belongs to is to take it out. One that turns out to be backing off is
+     * moved into its domain's sub-queue, where that domain's clocks hold it back like any other: dispatching
+     * it would have the crawler hammer the domain that just told us to wait, and putting it back would have
+     * every later fetch walk past it again.
+     */
+    async #fetchFromInner<R extends Dictionary = Dictionary>(): Promise<Request<R> | null> {
+        const inner = await this.#getInner();
+
+        for (let migrated = 0; migrated < MAX_INNER_MIGRATIONS; migrated++) {
+            const request = await inner.fetchNextRequest<R>();
+
+            if (request === null) {
+                return null;
+            }
+
+            const state = this.#getDomainState(request.url);
+
+            if (state === null || Date.now() >= throttledUntil(state)) {
+                return request;
+            }
+
+            await this.#migrateToSubQueue(request, inner);
+        }
+
+        return null;
+    }
+
+    /**
+     * Moves a request from the wrapped manager into its domain's sub-queue.
+     *
+     * Added there before it is marked handled here, so that a crash in between duplicates the request rather
+     * than dropping it, and `uniqueKey` absorbs the duplicate. A copy is what lands in the sub-queue: adding
+     * the request itself would overwrite the `id` the wrapped manager knows it by.
+     */
+    async #migrateToSubQueue(request: Request, inner: T): Promise<void> {
+        try {
+            const subManager = await this.#selectManagerOrThrow(request.url);
+            await subManager.addRequest(new Request({ ...request, id: undefined }));
+        } catch (error) {
+            // Whatever the reason, the request is out of the wrapped manager and in nobody's hands, so it
+            // goes back before the failure is raised.
+            await inner.reclaimRequest(request, { forefront: true });
+            throw error;
+        }
+
+        await inner.markRequestAsHandled(request);
+        this.#migratedFromInner += 1;
+    }
+
     async *[Symbol.asyncIterator]() {
         while (true) {
             const req = await this.fetchNextRequest();

test/core/crawlers/http_crawler.test.ts

@@ -693,6 +693,34 @@ test('a domain that never stops rate-limiting shuts the crawl down instead of ha
     expect(await crawler.getRequestManager().then((manager) => manager.getPendingCount())).toBe(1);
 }, 30_000);
 
+test('a domain is paced even when its requests come from the wrapped manager', async () => {
+    let hits = 0;
+    router.set('/from-a-list', (req, res) => {
+        hits++;
+        res.statusCode = 429;
+        res.end();
+    });
+
+    const requestList = await RequestList.open(null, [`${url}/from-a-list`]);
+
+    const crawler = new HttpCrawler({
+        requestManager: new ThrottlingRequestManager({
+            inner: await requestList.toTandem(await RequestQueue.open()),
+            domains: ['127.0.0.1'],
+            baseDelaySecs: 0.5,
+            maxDelaySecs: 1,
+            maxDomainStallSecs: 1,
+        }),
+        maxRequestRetries: 0,
+        requestHandler: async () => {},
+    });
+
+    await expect(crawler.run()).rejects.toThrow(PersistentRateLimitError);
+
+    // A handful of paced attempts, rather than one per turn of the task loop for as long as the crawl lives.
+    expect(hits).toBeLessThan(10);
+}, 30_000);
+
 test('`keepAlive` outlives a domain that never stops rate-limiting', async () => {
     router.set('/always-429-keep-alive', (req, res) => {
         res.statusCode = 429;

test/core/storages/throttling_request_manager.test.ts

@@ -245,6 +245,107 @@ describe('ThrottlingRequestManager', () => {
         expect(await inner.getPendingCount()).toBe(0);
     });
 
+    test('a request the inner manager holds is paced by its domain like any other', async () => {
+        const inner = await createQueue();
+        const manager = new ThrottlingRequestManager({ inner, domains: ['example.com'], baseDelaySecs: 0.5 });
+
+        // The same bypass as above, so the request belongs to a throttled domain without being stored per domain.
+        await inner.addRequest({ url: 'https://example.com/1' });
+
+        const request = (await manager.fetchNextRequest())!;
+        expect(manager.recordPacingSignal({ url: request.url, reason: 'rateLimited' })).toBe(true);
+        // What the crawler does with the `RequestThrottledError` a 429 raises: reclaim, no retry spent.
+        await manager.reclaimRequest(request);
+
+        // Handing it straight back would have the crawler hammer the domain that just turned it away, and
+        // nothing else would slow it down - a deferred request costs neither a retry nor session reputation.
+        expect(await manager.fetchNextRequest()).toBeNull();
+        // Nothing to dispatch until the backoff runs out, while the request is still outstanding work rather
+        // than lost - a manager reporting `ready` here would have the crawler spin, `finished` drop it.
+        expect((await manager.checkReadiness()).status).toBe('waiting');
+
+        const start = Date.now();
+        const again = await pollForNextRequest(manager);
+
+        expect(Date.now() - start).toBeGreaterThanOrEqual(400);
+        expect(again.url).toBe('https://example.com/1');
+    });
+
+    test('a domain being held back does not stop the inner manager serving the others', async () => {
+        const inner = await createQueue();
+        const manager = new ThrottlingRequestManager({ inner, domains: ['example.com'], baseDelaySecs: 60 });
+
+        await inner.addRequest({ url: 'https://example.com/held-back' });
+        await inner.addRequest({ url: 'https://other.com/free' });
+
+        const held = (await manager.fetchNextRequest())!;
+        expect(held.url).toBe('https://example.com/held-back');
+        manager.recordPacingSignal({ url: held.url, reason: 'rateLimited' });
+        await manager.reclaimRequest(held);
+
+        // The reclaimed request is still ahead of the other one, and skipping it must not mean skipping the queue.
+        const request = await manager.fetchNextRequest();
+        expect(request!.url).toBe('https://other.com/free');
+    });
+
+    test('the backlog a throttled domain has in the inner manager is moved out of it, not walked again', async () => {
+        const inner = await createQueue();
+        const manager = new ThrottlingRequestManager({ inner, domains: ['example.com'], baseDelaySecs: 60 });
+
+        // The same bypass as above, at the size that makes the cost visible.
+        for (let i = 0; i < 20; i++) {
+            await inner.addRequest({ url: `https://example.com/${i}` });
+        }
+        await inner.addRequest({ url: 'https://other.com/free' });
+
+        manager.recordPacingSignal({ url: 'https://example.com/0', reason: 'rateLimited' });
+
+        const fetchedFromInner = vitest.spyOn(inner, 'fetchNextRequest');
+        const reclaimedToInner = vitest.spyOn(inner, 'reclaimRequest');
+
+        expect((await manager.fetchNextRequest())!.url).toBe('https://other.com/free');
+
+        expect(fetchedFromInner).toHaveBeenCalledTimes(21);
+        expect(reclaimedToInner).not.toHaveBeenCalled();
+        // Nothing lost on the way, and nothing counted twice: moving a request marks it handled where it came
+        // from, so both managers hold a record of it.
+        expect(await manager.getPendingCount()).toBe(21);
+        expect(await manager.getTotalCount()).toBe(21);
+        expect(await manager.getHandledCount()).toBe(0);
+
+        fetchedFromInner.mockClear();
+        await manager.addRequest({ url: 'https://other.com/behind-it' });
+
+        // The point of moving it: enqueuing no longer re-arms a walk over the backlog, which on a queue whose
+        // every read is a round trip is the whole cost.
+        expect((await manager.fetchNextRequest())!.url).toBe('https://other.com/behind-it');
+        expect(fetchedFromInner).toHaveBeenCalledTimes(1);
+    });
+
+    test('a request whose migration fails goes back to the inner manager instead of being stranded', async () => {
+        const inner = await createQueue();
+        const manager = new ThrottlingRequestManager({
+            inner,
+            domains: ['example.com'],
+            baseDelaySecs: 60,
+            requestManagerOpener: async (identifier, options) => {
+                const queue = await RequestQueue.open(identifier, options);
+                queue.addRequest = async () => {
+                    throw new Error('sub-queue unavailable');
+                };
+                return queue;
+            },
+        });
+
+        await inner.addRequest({ url: 'https://example.com/held-back' });
+        manager.recordPacingSignal({ url: 'https://example.com/held-back', reason: 'rateLimited' });
+
+        await expect(manager.fetchNextRequest()).rejects.toThrow('sub-queue unavailable');
+
+        // Left in progress in a manager nobody will hand it back to, it would be lost for the rest of the run.
+        expect((await inner.fetchNextRequest())!.url).toBe('https://example.com/held-back');
+    });
+
     describe('a lazily-opened inner manager', () => {
         const throttling = { domains: ['example.com'] } satisfies Omit<
             ThrottlingRequestManagerOptions<RequestQueue>,
@@ -593,6 +694,31 @@ describe('ThrottlingRequestManager', () => {
             expect((await manager.checkReadiness()).status).not.toBe('stalled');
         });
 
+        test('work the inner manager holds keeps a stalling domain in view', async () => {
+            // Its own backoff, long enough that the sweep below leaves the wrapped manager blocked - otherwise
+            // the request it holds reads as work anyone can dispatch, which outranks the stall.
+            const manager = new ThrottlingRequestManager({
+                inner: await createQueue(),
+                domains: ['example.com'],
+                baseDelaySecs: 30,
+                maxDomainStallSecs: 30,
+            });
+            const inner = manager.innerManager!;
+
+            // Its sub-queue stays empty, so this request is the only sign the domain still has work left.
+            await inner.addRequest({ url: 'https://example.com/1' });
+
+            const request = (await manager.fetchNextRequest())!;
+            manager.recordPacingSignal({ url: request.url, reason: 'rateLimited' });
+            await manager.reclaimRequest(request);
+            stallFor(manager, 'example.com');
+
+            // Sweeping it is what finds the request; without it the crawl would wait on the domain forever.
+            await manager.fetchNextRequest();
+
+            expect((await manager.checkReadiness()).status).toBe('stalled');
+        });
+
         test('a domain that has run out of work is finished, not stalled', async () => {
             const manager = await stallingManager();
             manager.recordPacingSignal({ url: 'https://example.com/1', reason: 'rateLimited' });