From 2edc45b8083e1de0ac7c5c7e4520127d17046e1d Mon Sep 17 00:00:00 2001
From: Catherine Lee <55311782+c298lee@users.noreply.github.com>
Date: Wed, 10 Jan 2024 10:56:41 -0500
Subject: [PATCH] feat(replays): Add manual canvas snapshot function (#149)

Adds a snapshot canvas function that allows you to manually snapshot
canvas elements, which enables recording of 3d and webgl canvas

Requires https://github.com/getsentry/sentry-javascript/pull/10066
---
 .../record/observers/canvas/canvas-manager.ts | 158 ++++++++++++------
 1 file changed, 109 insertions(+), 49 deletions(-)

diff --git a/packages/rrweb/src/record/observers/canvas/canvas-manager.ts b/packages/rrweb/src/record/observers/canvas/canvas-manager.ts
index 6b09f56308..9af8c3b2f9 100644
--- a/packages/rrweb/src/record/observers/canvas/canvas-manager.ts
+++ b/packages/rrweb/src/record/observers/canvas/canvas-manager.ts
@@ -35,10 +35,12 @@ export interface CanvasManagerInterface {
   unfreeze(): void;
   lock(): void;
   unlock(): void;
+  snapshot(canvasElement?: HTMLCanvasElement): void;
 }
 
 export interface CanvasManagerConstructorOptions {
   recordCanvas: boolean;
+  isManualSnapshot?: boolean;
   mutationCb: canvasMutationCallback;
   win: IWindow;
   blockClass: blockClass;
@@ -65,11 +67,15 @@ export class CanvasManagerNoop implements CanvasManagerInterface {
   public unlock() {
     // noop
   }
+  public snapshot() {
+    // noop
+  }
 }
 
 export class CanvasManager implements CanvasManagerInterface {
   private pendingCanvasMutations: pendingCanvasMutationsMap = new Map();
   private rafStamps: RafStamps = { latestId: 0, invokeId: null };
+  private options: CanvasManagerConstructorOptions;
   private mirror: Mirror;
 
   private mutationCb: canvasMutationCallback;
@@ -110,6 +116,11 @@ export class CanvasManager implements CanvasManagerInterface {
     } = options;
     this.mutationCb = options.mutationCb;
     this.mirror = options.mirror;
+    this.options = options;
+
+    if (options.isManualSnapshot) {
+      return;
+    }
 
     callbackWrapper(() => {
       if (recordCanvas && sampling === 'all')
@@ -167,6 +178,90 @@ export class CanvasManager implements CanvasManagerInterface {
       unblockSelector,
       true,
     );
+    const rafId = this.takeSnapshot(
+      false,
+      fps,
+      win,
+      blockClass,
+      blockSelector,
+      unblockSelector,
+      options.dataURLOptions,
+    );
+
+    this.resetObservers = () => {
+      canvasContextReset();
+      cancelAnimationFrame(rafId);
+    };
+  }
+
+  private initCanvasMutationObserver(
+    win: IWindow,
+    blockClass: blockClass,
+    blockSelector: string | null,
+    unblockSelector: string | null,
+  ): void {
+    this.startRAFTimestamping();
+    this.startPendingCanvasMutationFlusher();
+
+    const canvasContextReset = initCanvasContextObserver(
+      win,
+      blockClass,
+      blockSelector,
+      unblockSelector,
+      false,
+    );
+    const canvas2DReset = initCanvas2DMutationObserver(
+      this.processMutation.bind(this),
+      win,
+      blockClass,
+      blockSelector,
+      unblockSelector,
+    );
+
+    const canvasWebGL1and2Reset = initCanvasWebGLMutationObserver(
+      this.processMutation.bind(this),
+      win,
+      blockClass,
+      blockSelector,
+      unblockSelector,
+      this.mirror,
+    );
+
+    this.resetObservers = () => {
+      canvasContextReset();
+      canvas2DReset();
+      canvasWebGL1and2Reset();
+    };
+  }
+
+  public snapshot(canvasElement?: HTMLCanvasElement) {
+    const { options } = this;
+    const rafId = this.takeSnapshot(
+      true,
+      options.sampling === 'all' ? 2 : options.sampling || 2,
+      options.win,
+      options.blockClass,
+      options.blockSelector,
+      options.unblockSelector,
+      options.dataURLOptions,
+      canvasElement,
+    );
+
+    this.resetObservers = () => {
+      cancelAnimationFrame(rafId);
+    };
+  }
+
+  private takeSnapshot(
+    isManualSnapshot: boolean,
+    fps: number,
+    win: IWindow,
+    blockClass: blockClass,
+    blockSelector: string | null,
+    unblockSelector: string | null,
+    dataURLOptions: DataURLOptions,
+    canvasElement?: HTMLCanvasElement,
+  ) {
     const snapshotInProgressMap: Map<number, boolean> = new Map();
     const worker = new Worker(getImageBitmapDataUrlWorkerURL());
     worker.onmessage = (e) => {
@@ -210,7 +305,13 @@ export class CanvasManager implements CanvasManagerInterface {
     let lastSnapshotTime = 0;
     let rafId: number;
 
-    const getCanvas = (): HTMLCanvasElement[] => {
+    const getCanvas = (
+      canvasElement?: HTMLCanvasElement,
+    ): HTMLCanvasElement[] => {
+      if (canvasElement) {
+        return [canvasElement];
+      }
+
       const matchedCanvas: HTMLCanvasElement[] = [];
       win.document.querySelectorAll('canvas').forEach((canvas) => {
         if (
@@ -232,11 +333,14 @@ export class CanvasManager implements CanvasManagerInterface {
       }
       lastSnapshotTime = timestamp;
 
-      getCanvas().forEach((canvas: HTMLCanvasElement) => {
+      getCanvas(canvasElement).forEach((canvas: HTMLCanvasElement) => {
         const id = this.mirror.getId(canvas);
         if (snapshotInProgressMap.get(id)) return;
         snapshotInProgressMap.set(id, true);
-        if (['webgl', 'webgl2'].includes((canvas as ICanvas).__context)) {
+        if (
+          !isManualSnapshot &&
+          ['webgl', 'webgl2'].includes((canvas as ICanvas).__context)
+        ) {
           // if the canvas hasn't been modified recently,
           // its contents won't be in memory and `createImageBitmap`
           // will return a transparent imageBitmap
@@ -267,7 +371,7 @@ export class CanvasManager implements CanvasManagerInterface {
                 bitmap,
                 width: canvas.width,
                 height: canvas.height,
-                dataURLOptions: options.dataURLOptions,
+                dataURLOptions,
               },
               [bitmap],
             );
@@ -282,51 +386,7 @@ export class CanvasManager implements CanvasManagerInterface {
     };
 
     rafId = onRequestAnimationFrame(takeCanvasSnapshots);
-
-    this.resetObservers = () => {
-      canvasContextReset();
-      cancelAnimationFrame(rafId);
-    };
-  }
-
-  private initCanvasMutationObserver(
-    win: IWindow,
-    blockClass: blockClass,
-    blockSelector: string | null,
-    unblockSelector: string | null,
-  ): void {
-    this.startRAFTimestamping();
-    this.startPendingCanvasMutationFlusher();
-
-    const canvasContextReset = initCanvasContextObserver(
-      win,
-      blockClass,
-      blockSelector,
-      unblockSelector,
-      false,
-    );
-    const canvas2DReset = initCanvas2DMutationObserver(
-      this.processMutation.bind(this),
-      win,
-      blockClass,
-      blockSelector,
-      unblockSelector,
-    );
-
-    const canvasWebGL1and2Reset = initCanvasWebGLMutationObserver(
-      this.processMutation.bind(this),
-      win,
-      blockClass,
-      blockSelector,
-      unblockSelector,
-      this.mirror,
-    );
-
-    this.resetObservers = () => {
-      canvasContextReset();
-      canvas2DReset();
-      canvasWebGL1and2Reset();
-    };
+    return rafId;
   }
 
   private startPendingCanvasMutationFlusher() {