From f3b44d18bee8b6e2bf84adaaa259baac3cdf95eb Mon Sep 17 00:00:00 2001
From: Yury Semikhatsky <yurys@chromium.org>
Date: Fri, 9 Apr 2021 14:09:45 -0700
Subject: [PATCH] fix(screencast): wait for ffmpeg to finish before reporting
 video (#6167)

---
 src/server/browser.ts           |  4 ++--
 src/server/browserContext.ts    | 15 ++++++++-------
 src/server/chromium/crPage.ts   |  3 ++-
 src/server/firefox/ffBrowser.ts |  2 +-
 src/server/webkit/wkBrowser.ts  |  2 +-
 5 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/server/browser.ts b/src/server/browser.ts
index 7f4e80aaf7..14d441d8d3 100644
--- a/src/server/browser.ts
+++ b/src/server/browser.ts
@@ -107,10 +107,10 @@ export abstract class Browser extends SdkObject {
     });
   }
 
-  _videoFinished(videoId: string) {
+  _takeVideo(videoId: string): Artifact | undefined {
     const video = this._idToVideo.get(videoId);
     this._idToVideo.delete(videoId);
-    video?.artifact.reportFinished();
+    return video?.artifact;
   }
 
   _didClose() {
diff --git a/src/server/browserContext.ts b/src/server/browserContext.ts
index 983d19cfde..b9d426bb04 100644
--- a/src/server/browserContext.ts
+++ b/src/server/browserContext.ts
@@ -233,6 +233,14 @@ export abstract class BrowserContext extends SdkObject {
 
       await this.instrumentation.onContextWillDestroy(this);
 
+      // Cleanup.
+      const promises: Promise<void>[] = [];
+      for (const { context, artifact } of this._browser._idToVideo.values()) {
+        // Wait for the videos to finish.
+        if (context === this)
+          promises.push(artifact.finishedPromise());
+      }
+
       if (this._isPersistentContext) {
         // Close all the pages instead of the context,
         // because we cannot close the default context.
@@ -242,13 +250,6 @@ export abstract class BrowserContext extends SdkObject {
         await this._doClose();
       }
 
-      // Cleanup.
-      const promises: Promise<void>[] = [];
-      for (const { context, artifact } of this._browser._idToVideo.values()) {
-        // Wait for the videos to finish.
-        if (context === this)
-          promises.push(artifact.finishedPromise());
-      }
       // We delete downloads after context closure
       // so that browser does not write to the download file anymore.
       promises.push(this._deleteAllDownloads());
diff --git a/src/server/chromium/crPage.ts b/src/server/chromium/crPage.ts
index 7826bcffe4..08f9b644c8 100644
--- a/src/server/chromium/crPage.ts
+++ b/src/server/chromium/crPage.ts
@@ -900,9 +900,10 @@ class FrameSession {
     this._screencastId = null;
     const recorder = this._videoRecorder!;
     this._videoRecorder = null;
+    const video = this._crPage._browserContext._browser._takeVideo(screencastId);
     await this._stopScreencast(recorder);
     await recorder.stop().catch(() => {});
-    this._crPage._browserContext._browser._videoFinished(screencastId);
+    video?.reportFinished();
   }
 
   async _startScreencast(client: any, options: Protocol.Page.startScreencastParameters = {}) {
diff --git a/src/server/firefox/ffBrowser.ts b/src/server/firefox/ffBrowser.ts
index 1419571a4a..7c1baaaa74 100644
--- a/src/server/firefox/ffBrowser.ts
+++ b/src/server/firefox/ffBrowser.ts
@@ -134,7 +134,7 @@ export class FFBrowser extends Browser {
   }
 
   _onScreencastFinished(payload: Protocol.Browser.screencastFinishedPayload) {
-    this._videoFinished(payload.screencastId);
+    this._takeVideo(payload.screencastId)?.reportFinished();
   }
 }
 
diff --git a/src/server/webkit/wkBrowser.ts b/src/server/webkit/wkBrowser.ts
index 9a4f52cdda..679650569e 100644
--- a/src/server/webkit/wkBrowser.ts
+++ b/src/server/webkit/wkBrowser.ts
@@ -132,7 +132,7 @@ export class WKBrowser extends Browser {
   }
 
   _onScreencastFinished(payload: Protocol.Playwright.screencastFinishedPayload) {
-    this._videoFinished(payload.screencastId);
+    this._takeVideo(payload.screencastId)?.reportFinished();
   }
 
   _onPageProxyCreated(event: Protocol.Playwright.pageProxyCreatedPayload) {