From 5d5cf26a0e0c41475053ac828bf45b953765d6d7 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Thu, 25 Jun 2020 13:13:10 -0700
Subject: [PATCH] fix(logs): streaming logs from InjectedScriptPoll without
 exception (#2712)

We used to get undefined messages, because we were mistakenly
fulfilling the logs multiple times.
---
 src/dom.ts                     | 19 ++++++++-----------
 src/injected/injectedScript.ts | 31 +++++++++++++++++--------------
 src/types.ts                   |  5 +++--
 3 files changed, 28 insertions(+), 27 deletions(-)

diff --git a/src/dom.ts b/src/dom.ts
index 12cd205264..a2ab34a551 100644
--- a/src/dom.ts
+++ b/src/dom.ts
@@ -697,20 +697,17 @@ export class InjectedScriptPollHandler<T> {
     //   - no unnecessary work in the page;
     //   - no possible side effects after progress promsie rejects.
     this._progress.cleanupWhenAborted(() => this.cancel());
-    this._streamLogs(poll.evaluateHandle(poll => poll.logs));
+    this._streamLogs();
   }
 
-  private _streamLogs(logsPromise: Promise<js.JSHandle<types.InjectedScriptLogs>>) {
-    // We continuously get a chunk of logs, stream them to the progress and wait for the next chunk.
-    logsPromise.catch(e => null).then(logs => {
-      if (!logs || !this._poll || !this._progress.isRunning())
+  private async _streamLogs() {
+    while (this._poll && this._progress.isRunning()) {
+      const messages = await this._poll.evaluate(poll => poll.takeNextLogs()).catch(e => [] as string[]);
+      if (!this._poll || !this._progress.isRunning())
         return;
-      logs.evaluate(logs => logs.current).catch(e => [] as string[]).then(messages => {
-        for (const message of messages)
-          this._progress.logger.info(message);
-      });
-      this._streamLogs(logs.evaluateHandle(logs => logs.next));
-    });
+      for (const message of messages)
+        this._progress.logger.info(message);
+    }
   }
 
   async finishHandle(): Promise<js.SmartHandle<T>> {
diff --git a/src/injected/injectedScript.ts b/src/injected/injectedScript.ts
index cfc32d76e3..591169514d 100644
--- a/src/injected/injectedScript.ts
+++ b/src/injected/injectedScript.ts
@@ -158,14 +158,20 @@ export default class InjectedScript {
   }
 
   private _runAbortableTask<T>(task: (progess: types.InjectedScriptProgress) => Promise<T>): types.InjectedScriptPoll<T> {
-    let currentLogs: string[] = [];
-    let logReady = () => {};
-    const createLogsPromise = () => new Promise<types.InjectedScriptLogs>(fulfill => {
-      logReady = () => {
-        const current = currentLogs;
-        currentLogs = [];
-        fulfill({ current, next: createLogsPromise() });
-      };
+    let unsentLogs: string[] = [];
+    let takeNextLogsCallback: ((logs: string[]) => void) | undefined;
+    const logReady = () => {
+      if (!takeNextLogsCallback)
+        return;
+      takeNextLogsCallback(unsentLogs);
+      unsentLogs = [];
+      takeNextLogsCallback = undefined;
+    };
+
+    const takeNextLogs = () => new Promise<string[]>(fulfill => {
+      takeNextLogsCallback = fulfill;
+      if (unsentLogs.length)
+        logReady();
     });
 
     let lastLog = '';
@@ -173,7 +179,7 @@ export default class InjectedScript {
       aborted: false,
       log: (message: string) => {
         lastLog = message;
-        currentLogs.push(message);
+        unsentLogs.push(message);
         logReady();
       },
       logRepeating: (message: string) => {
@@ -182,14 +188,11 @@ export default class InjectedScript {
       },
     };
 
-    // It is important to create logs promise before running the poll to capture logs from the first run.
-    const logs = createLogsPromise();
-
     return {
-      logs,
+      takeNextLogs,
       result: task(progress),
       cancel: () => { progress.aborted = true; },
-      takeLastLogs: () => currentLogs,
+      takeLastLogs: () => unsentLogs,
     };
   }
 
diff --git a/src/types.ts b/src/types.ts
index 2c4d4ea3d4..bed3544f91 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -143,10 +143,11 @@ export type InjectedScriptProgress = {
   logRepeating: (message: string) => void,
 };
 
-export type InjectedScriptLogs = { current: string[], next: Promise<InjectedScriptLogs> };
 export type InjectedScriptPoll<T> = {
   result: Promise<T>,
-  logs: Promise<InjectedScriptLogs>,
+  // Takes more logs, waiting until at least one message is available.
+  takeNextLogs: () => Promise<string[]>,
+  // Takes all current logs without waiting.
   takeLastLogs: () => string[],
   cancel: () => void,
 };