refactor: use cancellable timeouts via @apify/timeout (#62)

* refactor: use cancellable timeouts via `@apify/timeout`

Related:
apify/crawlee#1102
apify/crawlee#1216

* chore: fix spacing in package.json

* chore: require newer @apify/timeout

* chore: adjust rootDir for tests

* chore: add test for abortable timeouts

* refactor: do not try to close inactive browser controllers

Author: B4nan

jest.config.ts

@@ -5,7 +5,7 @@ export default async (): Promise<Config.InitialOptions> => ({
     preset: 'ts-jest',
     testEnvironment: 'node',
     testRunner: 'jest-circus/runner',
-    testTimeout: 20_000,
+    testTimeout: 30_000,
     collectCoverage: true,
     collectCoverageFrom: [
         '**/src/**/*.ts',

package.json

@@ -44,6 +44,7 @@
 	},
 	"dependencies": {
 		"@apify/log": "^1.1.1",
+		"@apify/timeout": "^0.2.1",
 		"fingerprint-generator": "^1.0.0",
 		"fingerprint-injector": "^1.0.0",
 		"lodash.merge": "^4.6.2",

src/abstract-classes/browser-controller.ts

@@ -1,5 +1,6 @@
 import { nanoid } from 'nanoid';
 import { TypedEmitter } from 'tiny-typed-emitter';
+import { tryCancel } from '@apify/timeout';
 import { BROWSER_CONTROLLER_EVENTS } from '../events';
 import { LaunchContext } from '../launch-context';
 import { log } from '../logger';
@@ -207,7 +208,9 @@ export abstract class BrowserController<
         this.totalPages++;
         await this.isActivePromise;
         const page = await this._newPage(pageOptions);
+        tryCancel();
         this.lastPageOpenedAt = Date.now();
+
         return page;
     }
 

src/browser-pool.ts

@@ -2,6 +2,7 @@ import pLimit from 'p-limit';
 import { nanoid } from 'nanoid';
 import ow from 'ow';
 import { TypedEmitter } from 'tiny-typed-emitter';
+import { addTimeoutToPromise, tryCancel } from '@apify/timeout';
 import { Fingerprint, FingerprintInjector } from 'fingerprint-injector';
 // eslint-disable-next-line @typescript-eslint/ban-ts-comment
 // @ts-expect-error no types for this package yet
@@ -12,7 +13,7 @@ import { BrowserPlugin } from './abstract-classes/browser-plugin';
 import { BROWSER_POOL_EVENTS } from './events';
 import { LaunchContext } from './launch-context';
 import { log } from './logger';
-import { addTimeoutToPromise, InferBrowserPluginArray, UnwrapPromise } from './utils';
+import { InferBrowserPluginArray, UnwrapPromise } from './utils';
 import { createFingerprintPreLaunchHook, createPrePageCreateHook, createPostPageCreateHook } from './fingerprinting/hooks';
 import { FingerprintGeneratorOptions } from './fingerprinting/types';
 
@@ -394,6 +395,7 @@ export class BrowserPool<
         return this.limiter(async () => {
             let browserController = this._pickBrowserWithFreeCapacity(browserPlugin);
             if (!browserController) browserController = await this._launchBrowser(id, { browserPlugin });
+            tryCancel();
 
             return this._createPageForBrowser(id, browserController, pageOptions, proxyUrl);
         });
@@ -417,6 +419,7 @@ export class BrowserPool<
         }
 
         const browserController = await this._launchBrowser(id, { launchOptions, browserPlugin });
+        tryCancel();
         return this._createPageForBrowser(id, browserController, pageOptions);
     }
 
@@ -501,6 +504,7 @@ export class BrowserPool<
         // It's not ideal though, we need to come up with a better API.
         // eslint-disable-next-line dot-notation -- accessing private property
         await browserController['isActivePromise'];
+        tryCancel();
 
         const finalPageOptions = browserController.launchContext.useIncognitoPages ? pageOptions : undefined;
 
@@ -509,15 +513,17 @@ export class BrowserPool<
         }
 
         await this._executeHooks(this.prePageCreateHooks, pageId, browserController, finalPageOptions);
+        tryCancel();
 
         let page: PageReturn;
 
         try {
             page = await addTimeoutToPromise(
-                browserController.newPage(finalPageOptions),
+                () => browserController.newPage(finalPageOptions),
                 this.operationTimeoutMillis,
                 'browserController.newPage() timed out.',
             ) as PageReturn;
+            tryCancel();
 
             this.pages.set(pageId, page);
             this.pageIds.set(page, pageId);
@@ -536,6 +542,7 @@ export class BrowserPool<
         }
 
         await this._executeHooks(this.postPageCreateHooks, page, browserController);
+        tryCancel();
 
         this.emit(BROWSER_POOL_EVENTS.PAGE_CREATED, page); // @TODO: CONSIDER renaming this event.
 
@@ -581,12 +588,9 @@ export class BrowserPool<
      */
     async closeAllBrowsers(): Promise<void> {
         const controllers = this._getAllBrowserControllers();
-
-        const promises: Promise<void>[] = [];
-
-        controllers.forEach((controller) => {
-            promises.push(controller.close());
-        });
+        const promises = [...controllers]
+            .filter((controller) => controller.isActive)
+            .map((controller) => controller.close());
 
         await Promise.all(promises);
     }
@@ -632,7 +636,9 @@ export class BrowserPool<
             // If the hooks or the launch fails, we need to delete the controller,
             // because otherwise it would be stuck in limbo without a browser.
             await this._executeHooks(this.preLaunchHooks, pageId, launchContext);
+            tryCancel();
             const browser = await browserPlugin.launch(launchContext);
+            tryCancel();
             browserController.assignBrowser(browser, launchContext);
         } catch (err) {
             this.activeBrowserControllers.delete(browserController);
@@ -656,6 +662,7 @@ export class BrowserPool<
             throw err;
         }
 
+        tryCancel();
         browserController.activate();
         this.emit(BROWSER_POOL_EVENTS.BROWSER_LAUNCHED, browserController);
 

src/playwright/playwright-controller.ts

@@ -1,4 +1,5 @@
 import type { Browser, BrowserType, Page } from 'playwright';
+import { tryCancel } from '@apify/timeout';
 import { BrowserController, Cookie } from '../abstract-classes/browser-controller';
 
 export class PlaywrightController extends BrowserController<BrowserType, Parameters<BrowserType['launch']>[0], Browser> {
@@ -27,6 +28,7 @@ export class PlaywrightController extends BrowserController<BrowserType, Paramet
         }
 
         const page = await this.browser.newPage(contextOptions);
+        tryCancel();
 
         page.once('close', () => {
             this.activePages--;

src/puppeteer/puppeteer-controller.ts

@@ -1,3 +1,4 @@
+import { tryCancel } from '@apify/timeout';
 import type Puppeteer from './puppeteer-proxy-per-page';
 import { BrowserController, Cookie } from '../abstract-classes/browser-controller';
 import { log } from '../logger';
@@ -29,13 +30,16 @@ export class PuppeteerController extends BrowserController<typeof Puppeteer> {
             }
 
             const context = await this.browser.createIncognitoBrowserContext(contextOptions);
+            tryCancel();
             const page = await context.newPage();
+            tryCancel();
 
             if (contextOptions.proxyUsername || contextOptions.proxyPassword) {
                 await page.authenticate({
                     username: contextOptions.proxyUsername ?? '',
                     password: contextOptions.proxyPassword ?? '',
                 });
+                tryCancel();
             }
 
             page.once('close', async () => {
@@ -46,12 +50,14 @@ export class PuppeteerController extends BrowserController<typeof Puppeteer> {
                 } catch (error: any) {
                     log.exception(error, 'Failed to close context.');
                 }
+                tryCancel();
             });
 
             return page;
         }
 
         const page = await this.browser.newPage();
+        tryCancel();
 
         page.once('close', () => {
             this.activePages--;

src/utils.ts

@@ -1,20 +1,6 @@
 import type { PlaywrightPlugin, PuppeteerPlugin } from '.';
 import type { BrowserPlugin } from './abstract-classes/browser-plugin';
 
-export function addTimeoutToPromise<T>(promise: Promise<T>, timeoutMillis: number, errorMessage: string): Promise<T> {
-    return new Promise(async (resolve, reject) => { // eslint-disable-line
-        const timeout = setTimeout(() => reject(new Error(errorMessage)), timeoutMillis);
-        try {
-            const data = await promise;
-            resolve(data);
-        } catch (err) {
-            reject(err);
-        } finally {
-            clearTimeout(timeout);
-        }
-    });
-};
-
 export type UnwrapPromise<T> = T extends PromiseLike<infer R> ? UnwrapPromise<R> : T;
 
 // eslint-disable-next-line @typescript-eslint/no-empty-function, @typescript-eslint/no-unused-vars

test/browser-pool.test.ts

@@ -5,6 +5,7 @@ import { promisify } from 'util';
 import { Server as ProxyChainServer } from 'proxy-chain';
 import puppeteer, { Page } from 'puppeteer';
 import playwright from 'playwright';
+import { addTimeoutToPromise } from '@apify/timeout';
 import { BrowserPool, PrePageCreateHook } from '../src/browser-pool';
 import { PuppeteerPlugin } from '../src/puppeteer/puppeteer-plugin';
 import { PlaywrightPlugin } from '../src/playwright/playwright-plugin';
@@ -89,6 +90,33 @@ describe('BrowserPool', () => {
             expect(page.close).toBeDefined();
         });
 
+        test('should allow early aborting in case of outer timeout', async () => {
+            const timeout = browserPool.operationTimeoutMillis;
+            browserPool.operationTimeoutMillis = 500;
+            // @ts-expect-error mocking private method
+            const spy = jest.spyOn(BrowserPool.prototype, '_executeHooks');
+
+            await browserPool.newPage();
+            expect(spy).toBeCalledTimes(4);
+            spy.mockReset();
+
+            await expect(addTimeoutToPromise(
+                () => browserPool.newPage(),
+                10,
+                'opening new page timed out',
+            )).rejects.toThrowError('opening new page timed out');
+
+            // We terminated early enough so only preLaunchHooks were not executed,
+            // thanks to `tryCancel()` calls after each await. If we did not run
+            // inside `addTimeoutToPromise()`, this would not work and we would get
+            // 4 calls instead of just one.
+            expect(spy).toBeCalledTimes(1);
+
+            spy.mockRestore();
+            browserPool.operationTimeoutMillis = timeout;
+            browserPool.retireAllBrowsers();
+        });
+
         test('proxy sugar syntax', async () => {
             const pool = new BrowserPool({
                 browserPlugins: [

test/tsconfig.json

@@ -1,7 +1,7 @@
 {
 	"extends": "../tsconfig.json",
 	"compilerOptions": {
-		"rootDir": "./",
+		"rootDir": "../",
 		"noEmit": true,
 		"allowJs": true
 	},