Fix asset loading in Replayer

packages/rrweb-snapshot/src/rebuild.ts

@@ -266,8 +266,22 @@ function buildNode(
             name === 'src' &&
             options.assetManager
           ) {
-            // TODO: do something with the asset manager
-            console.log('WIP! Please implement me!');
+            const originalValue = value.toString();
+            node.setAttribute(name, originalValue);
+            void options.assetManager
+              .whenReady(value.toString())
+              .then((status) => {
+                if (
+                  status.status === 'loaded' &&
+                  node.getAttribute('src') === originalValue
+                ) {
+                  node.setAttribute(name, status.url);
+                } else {
+                  console.log(
+                    `failed to load asset: ${originalValue}, ${status.status}`,
+                  );
+                }
+              });
           } else {
             node.setAttribute(name, value.toString());
           }
@@ -390,6 +404,7 @@ export function buildNodeWithSN(
     hackCss = true,
     afterAppend,
     cache,
+    assetManager,
   } = options;
   /**
    * Add a check to see if the node is already in the mirror. If it is, we can skip the whole process.
@@ -404,7 +419,7 @@ export function buildNodeWithSN(
     // For safety concern, check if the node in mirror is the same as the node we are trying to build
     if (isNodeMetaEqual(meta, n)) return mirror.getNode(n.id);
   }
-  let node = buildNode(n, { doc, hackCss, cache });
+  let node = buildNode(n, { doc, hackCss, cache, assetManager });
   if (!node) {
     return null;
   }
@@ -456,6 +471,7 @@ export function buildNodeWithSN(
         hackCss,
         afterAppend,
         cache,
+        assetManager,
       });
       if (!childNode) {
         console.warn('Failed to rebuild', childN);
@@ -544,6 +560,7 @@ function rebuild(
     afterAppend?: (n: Node, id: number) => unknown;
     cache: BuildCache;
     mirror: Mirror;
+    assetManager?: RebuildAssetManagerInterface;
   },
 ): Node | null {
   const {
@@ -553,6 +570,7 @@ function rebuild(
     afterAppend,
     cache,
     mirror = new Mirror(),
+    assetManager,
   } = options;
   const node = buildNodeWithSN(n, {
     doc,
@@ -561,6 +579,7 @@ function rebuild(
     hackCss,
     afterAppend,
     cache,
+    assetManager,
   });
   visit(mirror, (visitedNode) => {
     if (onVisit) {

packages/rrweb-snapshot/src/snapshot.ts

@@ -1183,7 +1183,10 @@ export function serializeNodeWithId(
           });
 
           if (serializedIframeNode) {
-            onIframeLoad(n as HTMLIFrameElement, serializedIframeNode);
+            onIframeLoad(
+              n as HTMLIFrameElement,
+              serializedIframeNode as serializedElementNodeWithId,
+            );
           }
         }
       },
@@ -1227,7 +1230,10 @@ export function serializeNodeWithId(
           });
 
           if (serializedLinkNode) {
-            onStylesheetLoad(n as HTMLLinkElement, serializedLinkNode);
+            onStylesheetLoad(
+              n as HTMLLinkElement,
+              serializedLinkNode as serializedElementNodeWithId,
+            );
           }
         }
       },

packages/rrweb/src/replay/assets/index.ts

@@ -1,16 +1,26 @@
-import type { RebuildAssetManagerInterface, assetEvent } from '@rrweb/types';
+import type {
+  RebuildAssetManagerFinalStatus,
+  RebuildAssetManagerInterface,
+  RebuildAssetManagerStatus,
+  assetEvent,
+} from '@rrweb/types';
 import { deserializeArg } from '../canvas/deserialize-args';
 
 export default class AssetManager implements RebuildAssetManagerInterface {
   private originalToObjectURLMap: Map<string, string> = new Map();
   private loadingURLs: Set<string> = new Set();
   private failedURLs: Set<string> = new Set();
+  private callbackMap: Map<
+    string,
+    Array<(status: RebuildAssetManagerFinalStatus) => void>
+  > = new Map();
 
   public async add(event: assetEvent) {
     const { data } = event;
     const { url, payload, failed } = { payload: false, failed: false, ...data };
     if (failed) {
       this.failedURLs.add(url);
+      this.executeCallbacks(url, { status: 'failed' });
       return;
     }
     this.loadingURLs.add(url);
@@ -28,15 +38,45 @@ export default class AssetManager implements RebuildAssetManagerInterface {
     const objectURL = URL.createObjectURL(result);
     this.originalToObjectURLMap.set(url, objectURL);
     this.loadingURLs.delete(url);
+    this.executeCallbacks(url, { status: 'loaded', url: objectURL });
   }
 
-  public get(
+  private executeCallbacks(
     url: string,
-  ):
-    | { status: 'loading' }
-    | { status: 'loaded'; url: string }
-    | { status: 'failed' }
-    | { status: 'unknown' } {
+    status: RebuildAssetManagerFinalStatus,
+  ) {
+    const callbacks = this.callbackMap.get(url);
+    while (callbacks && callbacks.length > 0) {
+      const callback = callbacks.pop();
+      if (!callback) {
+        break;
+      }
+      callback(status);
+    }
+  }
+
+  public async whenReady(url: string): Promise<RebuildAssetManagerFinalStatus> {
+    const currentStatus = this.get(url);
+    if (
+      currentStatus.status === 'loaded' ||
+      currentStatus.status === 'failed'
+    ) {
+      return currentStatus;
+    }
+    let resolve: (status: RebuildAssetManagerFinalStatus) => void;
+    const promise = new Promise<RebuildAssetManagerFinalStatus>((r) => {
+      resolve = r;
+    });
+    if (!this.callbackMap.has(url)) {
+      this.callbackMap.set(url, []);
+    }
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    this.callbackMap.get(url)!.push(resolve!);
+
+    return promise;
+  }
+
+  public get(url: string): RebuildAssetManagerStatus {
     const result = this.originalToObjectURLMap.get(url);
 
     if (result) {
@@ -62,4 +102,20 @@ export default class AssetManager implements RebuildAssetManagerInterface {
       status: 'unknown',
     };
   }
+
+  public reset(): void {
+    this.originalToObjectURLMap.forEach((objectURL) => {
+      URL.revokeObjectURL(objectURL);
+    });
+    this.originalToObjectURLMap.clear();
+    this.loadingURLs.clear();
+    this.failedURLs.clear();
+    this.callbackMap.forEach((callbacks) => {
+      while (callbacks.length > 0) {
+        const cb = callbacks.pop();
+        if (cb) cb({ status: 'reset' });
+      }
+    });
+    this.callbackMap.clear();
+  }
 }

packages/rrweb/src/replay/index.ts

@@ -137,6 +137,9 @@ export class Replayer {
   private cache: BuildCache = createCache();
 
   private imageMap: Map<eventWithTime | string, HTMLImageElement> = new Map();
+
+  private assetManager = new AssetManager();
+
   private canvasEventMap: Map<eventWithTime, canvasMutationParam> = new Map();
 
   private mirror: Mirror = createMirror();
@@ -717,6 +720,11 @@ export class Replayer {
           }
         };
         break;
+      case EventType.Asset:
+        castFn = () => {
+          void this.assetManager.add(event);
+        };
+        break;
       default:
     }
     const wrappedCastFn = () => {
@@ -788,6 +796,8 @@ export class Replayer {
       }
     };
 
+    void this.preloadAllAssets();
+
     /**
      * Normally rebuilding full snapshot should not be under virtual dom environment.
      * But if the order of data events has some issues, it might be possible.
@@ -804,6 +814,7 @@ export class Replayer {
       afterAppend,
       cache: this.cache,
       mirror: this.mirror,
+      assetManager: this.assetManager,
     });
     afterAppend(this.iframe.contentDocument, event.data.node.id);
 
@@ -827,7 +838,6 @@ export class Replayer {
     if (this.config.UNSAFE_replayCanvas) {
       void this.preloadAllImages();
     }
-    void this.preloadAllAssets();
   }
 
   private insertStyleRules(
@@ -999,11 +1009,10 @@ export class Replayer {
    * Process all asset events and preload them
    */
   private async preloadAllAssets(): Promise<void[]> {
-    const assetManager = new AssetManager();
     const promises: Promise<void>[] = [];
     for (const event of this.service.state.context.events) {
       if (event.type === EventType.Asset) {
-        promises.push(assetManager.add(event));
+        promises.push(this.assetManager.add(event));
       }
     }
     return Promise.all(promises);

packages/rrweb/test/replay/asset-unit.test.ts

@@ -33,6 +33,7 @@ describe('AssetManager', () => {
 
   afterEach(() => {
     jest.restoreAllMocks();
+    jest.useRealTimers();
   });
 
   afterAll(() => {
@@ -97,4 +98,46 @@ describe('AssetManager', () => {
 
     expect(assetManager.get(url)).toEqual({ status: 'unknown' });
   });
+
+  it('should execute hook when an asset is added', async () => {
+    jest.useFakeTimers();
+    const url = 'https://example.com/image.png';
+    const event: assetEvent = {
+      type: EventType.Asset,
+      data: {
+        url,
+        payload: examplePayload,
+      },
+    };
+    void assetManager.add(event);
+    const promise = assetManager.whenReady(url);
+
+    jest.spyOn(URL, 'createObjectURL').mockReturnValue('objectURL');
+
+    jest.runAllTimers();
+
+    await expect(promise).resolves.toEqual({
+      status: 'loaded',
+      url: 'objectURL',
+    });
+  });
+
+  it('should send status reset to callbacks when reset', async () => {
+    jest.useFakeTimers();
+    const url = 'https://example.com/image.png';
+    const event: assetEvent = {
+      type: EventType.Asset,
+      data: {
+        url,
+        payload: examplePayload,
+      },
+    };
+    void assetManager.add(event);
+    const promise = assetManager.whenReady(url);
+
+    assetManager.reset();
+    jest.runAllTimers();
+
+    await expect(promise).resolves.toEqual({ status: 'reset' });
+  });
 });

packages/types/src/index.ts

@@ -719,16 +719,25 @@ export type TakeTypedKeyValues<Obj extends object, Type> = Pick<
   TakeTypeHelper<Obj, Type>[keyof TakeTypeHelper<Obj, Type>]
 >;
 
-export abstract class RebuildAssetManagerInterface {
+export type RebuildAssetManagerResetStatus = { status: 'reset' };
+export type RebuildAssetManagerUnknownStatus = { status: 'unknown' };
+export type RebuildAssetManagerLoadingStatus = { status: 'loading' };
+export type RebuildAssetManagerLoadedStatus = { status: 'loaded'; url: string };
+export type RebuildAssetManagerFailedStatus = { status: 'failed' };
+export type RebuildAssetManagerFinalStatus =
+  | RebuildAssetManagerLoadedStatus
+  | RebuildAssetManagerFailedStatus
+  | RebuildAssetManagerResetStatus;
+export type RebuildAssetManagerStatus =
+  | RebuildAssetManagerUnknownStatus
+  | RebuildAssetManagerLoadingStatus
+  | RebuildAssetManagerFinalStatus;
+
+export declare abstract class RebuildAssetManagerInterface {
   abstract add(event: assetEvent): Promise<void>;
-
-  abstract get(
-    url: string,
-  ):
-    | { status: 'loading' }
-    | { status: 'loaded'; url: string }
-    | { status: 'failed' }
-    | { status: 'unknown' };
+  abstract get(url: string): RebuildAssetManagerStatus;
+  abstract whenReady(url: string): Promise<RebuildAssetManagerFinalStatus>;
+  abstract reset(): void;
 }
 
 export enum NodeType {