From e61d9cbb4bf858f7ed4273d34747693109bddefd Mon Sep 17 00:00:00 2001
From: Nikolay Vitkov <34244704+Lightning00Blade@users.noreply.github.com>
Date: Wed, 31 May 2023 09:32:16 +0200
Subject: [PATCH] chore: refactor waitForNetworkIdle + waitForEvent (#10277)

---
 packages/puppeteer-core/src/api/Page.ts       | 21 +++++-----
 .../src/common/NetworkEventManager.ts         |  8 ++--
 .../src/common/bidi/NetworkManager.ts         |  2 +-
 packages/puppeteer-core/src/common/util.ts    | 40 ++++++-------------
 test/TestExpectations.json                    |  6 +++
 5 files changed, 33 insertions(+), 44 deletions(-)

diff --git a/packages/puppeteer-core/src/api/Page.ts b/packages/puppeteer-core/src/api/Page.ts
index a26d6d19..0c3621ad 100644
--- a/packages/puppeteer-core/src/api/Page.ts
+++ b/packages/puppeteer-core/src/api/Page.ts
@@ -1641,31 +1641,28 @@ export class Page extends EventEmitter {
     const idleDeferred = createDeferred<void>();
     const abortDeferred = createDeferred<Error>();
 
-    let idleTimer: NodeJS.Timeout;
+    let idleTimer: NodeJS.Timeout | undefined;
     const cleanup = () => {
-      idleTimer && clearTimeout(idleTimer);
+      clearTimeout(idleTimer);
       abortDeferred.reject(new Error('abort'));
     };
 
     const evaluate = () => {
-      idleTimer && clearTimeout(idleTimer);
+      clearTimeout(idleTimer);
+
       if (networkManager.inFlightRequestsCount() === 0) {
         idleTimer = setTimeout(idleDeferred.resolve, idleTime);
       }
     };
 
-    evaluate();
-
-    const eventHandler = () => {
-      evaluate();
-      return false;
-    };
-
     const listenToEvent = (event: symbol) => {
       return waitForEvent(
         networkManager,
         event,
-        eventHandler,
+        () => {
+          evaluate();
+          return false;
+        },
         timeout,
         abortDeferred.valueOrThrow()
       );
@@ -1677,6 +1674,8 @@ export class Page extends EventEmitter {
       listenToEvent(NetworkManagerEmittedEvents.RequestFailed),
     ];
 
+    evaluate();
+
     await Promise.race([
       idleDeferred.valueOrThrow(),
       ...eventPromises,
diff --git a/packages/puppeteer-core/src/common/NetworkEventManager.ts b/packages/puppeteer-core/src/common/NetworkEventManager.ts
index 2f2d6d6b..1dae2c3e 100644
--- a/packages/puppeteer-core/src/common/NetworkEventManager.ts
+++ b/packages/puppeteer-core/src/common/NetworkEventManager.ts
@@ -150,13 +150,13 @@ export class NetworkEventManager {
   }
 
   inFlightRequestsCount(): number {
-    let inProgressRequestCounter = 0;
-    for (const [, request] of this.#httpRequestsMap) {
+    let inFlightRequestCounter = 0;
+    for (const request of this.#httpRequestsMap.values()) {
       if (!request.response()) {
-        inProgressRequestCounter++;
+        inFlightRequestCounter++;
       }
     }
-    return inProgressRequestCounter;
+    return inFlightRequestCounter;
   }
 
   storeRequestWillBeSent(
diff --git a/packages/puppeteer-core/src/common/bidi/NetworkManager.ts b/packages/puppeteer-core/src/common/bidi/NetworkManager.ts
index 9a7895ae..8bc7215e 100644
--- a/packages/puppeteer-core/src/common/bidi/NetworkManager.ts
+++ b/packages/puppeteer-core/src/common/bidi/NetworkManager.ts
@@ -103,7 +103,7 @@ export class NetworkManager extends EventEmitter {
 
   inFlightRequestsCount(): number {
     let inFlightRequestCounter = 0;
-    for (const [, request] of this.#requestMap) {
+    for (const request of this.#requestMap.values()) {
       if (!request.response() || request._failureText) {
         inFlightRequestCounter++;
       }
diff --git a/packages/puppeteer-core/src/common/util.ts b/packages/puppeteer-core/src/common/util.ts
index c6390026..406ab568 100644
--- a/packages/puppeteer-core/src/common/util.ts
+++ b/packages/puppeteer-core/src/common/util.ts
@@ -22,6 +22,7 @@ import type {ElementHandle} from '../api/ElementHandle.js';
 import type {JSHandle} from '../api/JSHandle.js';
 import {Page} from '../api/Page.js';
 import {isNode} from '../environment.js';
+import {createDeferred} from '../puppeteer-core.js';
 import {assert} from '../util/assert.js';
 import {isErrorLike} from '../util/ErrorLike.js';
 
@@ -382,45 +383,28 @@ export async function waitForEvent<T>(
   timeout: number,
   abortPromise: Promise<Error>
 ): Promise<T> {
-  let eventTimeout: NodeJS.Timeout;
-  let resolveCallback: (value: T | PromiseLike<T>) => void;
-  let rejectCallback: (value: Error) => void;
-  const promise = new Promise<T>((resolve, reject) => {
-    resolveCallback = resolve;
-    rejectCallback = reject;
+  const deferred = createDeferred<T>({
+    message: `Timeout exceeded while waiting for event ${String(eventName)}`,
+    timeout,
   });
   const listener = addEventListener(emitter, eventName, async event => {
-    if (!(await predicate(event))) {
-      return;
+    if (await predicate(event)) {
+      deferred.resolve(event);
     }
-    resolveCallback(event);
   });
-  if (timeout) {
-    eventTimeout = setTimeout(() => {
-      rejectCallback(
-        new TimeoutError('Timeout exceeded while waiting for event')
-      );
-    }, timeout);
-  }
-  function cleanup(): void {
-    removeEventListeners([listener]);
-    clearTimeout(eventTimeout);
-  }
-  const result = await Promise.race([promise, abortPromise]).then(
+  return Promise.race([deferred.valueOrThrow(), abortPromise]).then(
     r => {
-      cleanup();
+      removeEventListeners([listener]);
+      if (isErrorLike(r)) {
+        throw r;
+      }
       return r;
     },
     error => {
-      cleanup();
+      removeEventListeners([listener]);
       throw error;
     }
   );
-  if (isErrorLike(result)) {
-    throw result;
-  }
-
-  return result;
 }
 
 /**
diff --git a/test/TestExpectations.json b/test/TestExpectations.json
index a1ce497c..f70df395 100644
--- a/test/TestExpectations.json
+++ b/test/TestExpectations.json
@@ -2099,6 +2099,12 @@
     "parameters": ["chrome", "webDriverBiDi"],
     "expectations": ["PASS"]
   },
+  {
+    "testIdPattern": "[page.spec] Page Page.waitForNetworkIdle should work",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["firefox", "webDriverBiDi"],
+    "expectations": ["FAIL"]
+  },
   {
     "testIdPattern": "[page.spec] Page Page.waitForNetworkIdle should work with aborted requests",
     "platforms": ["linux"],