fix(rrweb): port upstream shadow-root release on host removal (fixes #5419)

.changeset/release-detached-shadow-roots.md

@@ -0,0 +1,5 @@
+---
+'posthog-js': patch
+---
+
+Session replay releases a removed shadow DOM host as soon as it leaves the page, instead of keeping it and its detached subtree in memory until the next full snapshot.

packages/rrweb/rrweb/src/record/mutation.ts

@@ -868,6 +868,15 @@ export default class MutationBuffer {
         break;
       }
       case 'childList': {
+        // Runs before the blocked-parent return so hosts removed from a
+        // blocked container are released too. Still attached by the time the
+        // batch is processed means moved, not removed.
+        m.removedNodes.forEach((n) => {
+          if (!inDom(n)) {
+            this.shadowDomManager.removeShadowRootsIn(n);
+          }
+        });
+
         /**
          * Parent is blocked, ignore all child mutations
          */

packages/rrweb/rrweb/src/record/shadow-dom-manager.ts

@@ -1,5 +1,6 @@
 import type { MutationBufferParam } from '../types';
 import type {
+  listenerHandler,
   mutationCallBack,
   scrollCallback,
   SamplingStrategy,
@@ -22,14 +23,17 @@ type BypassOptions = Omit<
   sampling: SamplingStrategy;
 };
 
+type ObservedShadowRoot = { doc: Document; handlers: listenerHandler[] };
+
 export class ShadowDomManager {
-  private shadowDoms = new WeakSet<ShadowRoot>();
+  // Strong refs, so each root's entry must be dropped when its host is removed.
+  private shadowDoms = new Map<ShadowRoot, ObservedShadowRoot>();
   private mutationCb: mutationCallBack;
   private scrollCb: scrollCallback;
   private bypassOptions: BypassOptions;
   private mirror: Mirror;
-  // Handlers are tagged with the document that owns their shadow root so a
-  // single iframe can be torn down without disconnecting the rest of the page.
+  // attachShadow patches, tagged with the document they patch so a single
+  // iframe can be torn down without disconnecting the rest of the page.
   private restoreHandlers: { doc: Document; handler: () => void }[] = [];
 
   constructor(options: {
@@ -55,11 +59,12 @@ export class ShadowDomManager {
   public addShadowRoot(shadowRoot: ShadowRoot, doc: Document): void {
     if (!isNativeShadowDom(shadowRoot)) return;
     if (this.shadowDoms.has(shadowRoot)) return;
-    this.shadowDoms.add(shadowRoot);
     // Derive the owning document from the host so a shadow root nested in an
     // iframe is keyed to that iframe's document, not whatever the caller passed
     // (takeFullSnapshot's onSerialize hands us the top-level document).
     const ownerDoc = dom.host(shadowRoot)?.ownerDocument ?? doc;
+    const observed: ObservedShadowRoot = { doc: ownerDoc, handlers: [] };
+    this.shadowDoms.set(shadowRoot, observed);
     const { observer, buffer } = initMutationObserver(
       {
         ...this.bypassOptions,
@@ -70,32 +75,30 @@ export class ShadowDomManager {
       },
       shadowRoot,
     );
-    this.restoreHandlers.push({
-      doc: ownerDoc,
-      handler: () => {
-        observer.disconnect();
-        buffer.destroy();
-        // Release the canvas directly, not via buffer.reset(), per the recursion-guard unit test.
-        buffer.releaseCanvasManager();
-        const index = mutationBuffers.indexOf(buffer);
-        if (index !== -1) {
-          mutationBuffers.splice(index, 1);
-        }
-      },
+    observed.handlers.push(() => {
+      observer.disconnect();
+      buffer.destroy();
+      // Release the canvas directly, not via buffer.reset(), per the recursion-guard unit test.
+      buffer.releaseCanvasManager();
+      const index = mutationBuffers.indexOf(buffer);
+      if (index !== -1) {
+        mutationBuffers.splice(index, 1);
+      }
     });
-    this.restoreHandlers.push({
-      doc: ownerDoc,
-      handler: initScrollObserver({
+    observed.handlers.push(
+      initScrollObserver({
         ...this.bypassOptions,
         scrollCb: this.scrollCb,
         // https://gist.github.com/praveenpuglia/0832da687ed5a5d7a0907046c9ef1813
         // scroll is not allowed to pass the boundary, so we need to listen the shadow document
         doc: shadowRoot as unknown as Document,
         mirror: this.mirror,
       }),
-    });
+    );
     // Defer this to avoid adoptedStyleSheet events being created before the full snapshot is created or attachShadow action is recorded.
     setTimeout(() => {
+      // The root was released, or the manager reset, in the meantime.
+      if (this.shadowDoms.get(shadowRoot) !== observed) return;
       if (
         shadowRoot.adoptedStyleSheets &&
         shadowRoot.adoptedStyleSheets.length > 0
@@ -104,19 +107,37 @@ export class ShadowDomManager {
           shadowRoot.adoptedStyleSheets,
           this.mirror.getId(dom.host(shadowRoot)),
         );
-      this.restoreHandlers.push({
-        doc: ownerDoc,
-        handler: initAdoptedStyleSheetObserver(
+      observed.handlers.push(
+        initAdoptedStyleSheetObserver(
           {
             mirror: this.mirror,
             stylesheetManager: this.bypassOptions.stylesheetManager,
           },
           shadowRoot,
         ),
-      });
+      );
     }, 0);
   }
 
+  // Release every observed shadow root in a subtree that has left the DOM.
+  public removeShadowRootsIn(n: Node): void {
+    if (this.shadowDoms.size === 0) return;
+    const shadowRoot = dom.shadowRoot(n);
+    if (shadowRoot) {
+      if (this.shadowDoms.has(shadowRoot)) {
+        this.removeShadowRoot(shadowRoot);
+      }
+      this.removeShadowRootsIn(shadowRoot);
+    }
+    dom.childNodes(n).forEach((child) => this.removeShadowRootsIn(child));
+  }
+
+  private removeShadowRoot(shadowRoot: ShadowRoot): void {
+    const observed = this.shadowDoms.get(shadowRoot);
+    this.shadowDoms.delete(shadowRoot);
+    observed?.handlers.forEach(runHandler);
+  }
+
   /**
    * Monkey patch 'attachShadow' of an IFrameElement to observe newly added shadow doms.
    */
@@ -163,31 +184,36 @@ export class ShadowDomManager {
   }
 
   public reset(): void {
-    this.restoreHandlers.forEach(({ handler }) => {
-      try {
-        handler();
-      } catch (e) {
-        //
-      }
-    });
+    this.restoreHandlers.forEach(({ handler }) => runHandler(handler));
     this.restoreHandlers = [];
-    this.shadowDoms = new WeakSet();
+    const shadowDoms = this.shadowDoms;
+    this.shadowDoms = new Map();
+    shadowDoms.forEach(({ handlers }) => handlers.forEach(runHandler));
   }
 
   // Tear down only the shadow observers owned by `doc` (e.g. one iframe being removed), leaving the rest of the page's shadow observation intact.
   public resetForDoc(doc: Document): void {
     const remaining: { doc: Document; handler: () => void }[] = [];
     for (const entry of this.restoreHandlers) {
       if (entry.doc === doc) {
-        try {
-          entry.handler();
-        } catch (e) {
-          //
-        }
+        runHandler(entry.handler);
       } else {
         remaining.push(entry);
       }
     }
     this.restoreHandlers = remaining;
+    this.shadowDoms.forEach((observed, shadowRoot) => {
+      if (observed.doc === doc) {
+        this.removeShadowRoot(shadowRoot);
+      }
+    });
+  }
+}
+
+function runHandler(handler: listenerHandler) {
+  try {
+    handler();
+  } catch (e) {
+    //
   }
 }

packages/rrweb/rrweb/test/record.test.ts

@@ -1317,6 +1317,74 @@ describe('record', function (this: ISuite) {
     expect(serializedEvents).not.toContain('updated-private-value');
   });
 
+  it('keeps recording shadow content of a host that is moved or re-inserted', async () => {
+    await ctx.page.evaluate(async () => {
+      const { record } = (window as unknown as IWindow).rrweb;
+      record({ emit: (window as unknown as IWindow).emit });
+      const settle = () => new Promise((resolve) => setTimeout(resolve, 20));
+
+      const host = document.createElement('div');
+      host.attachShadow({ mode: 'open' }).innerHTML = '<span>initial</span>';
+      const other = document.createElement('section');
+      document.body.append(host, other);
+      await settle();
+      const span = host.shadowRoot!.querySelector('span')!;
+
+      // Removed and re-attached within one mutation batch.
+      host.remove();
+      other.appendChild(host);
+      await settle();
+      span.textContent = 'after-move';
+      await settle();
+
+      // Removed, then re-inserted in a later batch.
+      host.remove();
+      await settle();
+      document.body.appendChild(host);
+      await settle();
+      span.textContent = 'after-reinsert';
+      await settle();
+    });
+    await waitForRAF(ctx.page);
+
+    const serializedEvents = JSON.stringify(ctx.events);
+    expect(serializedEvents).toContain('after-move');
+    expect(serializedEvents).toContain('after-reinsert');
+  });
+
+  it('keeps shadow changes made while frozen on a host removed and re-inserted', async () => {
+    await ctx.page.evaluate(async () => {
+      const { record, freezePage } = (window as unknown as IWindow).rrweb;
+      record({ emit: (window as unknown as IWindow).emit });
+      const settle = () => new Promise((resolve) => setTimeout(resolve, 20));
+
+      const host = document.createElement('div');
+      host.attachShadow({ mode: 'open' }).innerHTML = '<span>initial</span>';
+      document.body.appendChild(host);
+      await settle();
+      const span = host.shadowRoot!.querySelector('span')!;
+
+      freezePage();
+      span.textContent = 'while-frozen';
+      await settle();
+      host.remove();
+      await settle();
+      document.body.appendChild(host);
+      await settle();
+      // 'unfreeze' happens upon a user event
+      document.body.click();
+      await settle();
+
+      span.textContent = 'after-unfreeze';
+      await settle();
+    });
+    await waitForRAF(ctx.page);
+
+    const serializedEvents = JSON.stringify(ctx.events);
+    expect(serializedEvents).toContain('while-frozen');
+    expect(serializedEvents).toContain('after-unfreeze');
+  });
+
   it('records mouse movement inside a same-origin iframe as MouseMove', async () => {
     await ctx.page.evaluate(async () => {
       const iframe = document.createElement('iframe');

packages/rrweb/rrweb/test/record/move-observer-source.test.ts

@@ -53,6 +53,7 @@ const createOptions = (doc: Document = document): observerParam =>
       addShadowRoot: vi.fn(),
       observeAttachShadow: vi.fn(),
       resetForDoc: vi.fn(),
+      removeShadowRootsIn: vi.fn(),
     },
     canvasManager: {
       acquire: vi.fn(),

packages/rrweb/rrweb/test/record/observer-teardown.test.ts

@@ -52,6 +52,7 @@ const createOptions = (doc: Document = document): observerParam =>
       addShadowRoot: vi.fn(),
       observeAttachShadow: vi.fn(),
       resetForDoc: vi.fn(),
+      removeShadowRootsIn: vi.fn(),
     },
     canvasManager: {
       acquire: vi.fn(),

packages/rrweb/rrweb/test/record/shadow-dom-manager.test.ts

@@ -216,4 +216,74 @@ describe('ShadowDomManager', () => {
     manager.resetForDoc(iframeDoc);
     expect(mutationBuffers.length).toBe(buffersBeforeAdd);
   });
+
+  it('removeShadowRootsIn() releases only the roots inside the removed subtree', () => {
+    const manager = createManager();
+
+    const removed = document.createElement('div');
+    const kept = document.createElement('div');
+    elements.push(removed, kept);
+    document.body.append(removed, kept);
+
+    const buffersBeforeAdd = mutationBuffers.length;
+    removed.attachShadow({ mode: 'open' });
+    kept.attachShadow({ mode: 'open' });
+    expect(mutationBuffers.length).toBe(buffersBeforeAdd + 2);
+
+    removed.remove();
+    manager.removeShadowRootsIn(removed);
+    expect(mutationBuffers.length).toBe(buffersBeforeAdd + 1);
+
+    // A second pass over the same subtree is a no-op.
+    manager.removeShadowRootsIn(removed);
+    expect(mutationBuffers.length).toBe(buffersBeforeAdd + 1);
+  });
+
+  it('resetForDoc() still releases iframe roots after a top-level host was released', () => {
+    const manager = createManager();
+
+    const topHost = document.createElement('div');
+    const iframe = document.createElement('iframe');
+    elements.push(topHost, iframe);
+    document.body.append(topHost, iframe);
+    const iframeDoc = iframe.contentDocument as Document;
+    const iframeHost = iframeDoc.createElement('div');
+    iframeDoc.body.appendChild(iframeHost);
+
+    const buffersBeforeAdd = mutationBuffers.length;
+    manager.addShadowRoot(topHost.attachShadow({ mode: 'open' }), document);
+    manager.addShadowRoot(iframeHost.attachShadow({ mode: 'open' }), document);
+    expect(mutationBuffers.length).toBe(buffersBeforeAdd + 2);
+
+    topHost.remove();
+    manager.removeShadowRootsIn(topHost);
+    expect(mutationBuffers.length).toBe(buffersBeforeAdd + 1);
+
+    manager.resetForDoc(iframeDoc);
+    expect(mutationBuffers.length).toBe(buffersBeforeAdd);
+  });
+
+  it('resetForDoc() leaves top-level roots observed when an iframe is torn down', () => {
+    const manager = createManager();
+
+    const topHost = document.createElement('div');
+    const iframe = document.createElement('iframe');
+    elements.push(topHost, iframe);
+    document.body.append(topHost, iframe);
+    const iframeDoc = iframe.contentDocument as Document;
+    const iframeHost = iframeDoc.createElement('div');
+    iframeDoc.body.appendChild(iframeHost);
+
+    const buffersBeforeAdd = mutationBuffers.length;
+    manager.addShadowRoot(topHost.attachShadow({ mode: 'open' }), document);
+    manager.addShadowRoot(iframeHost.attachShadow({ mode: 'open' }), document);
+
+    manager.resetForDoc(iframeDoc);
+    expect(mutationBuffers.length).toBe(buffersBeforeAdd + 1);
+
+    // The top-level root is still tracked, so it can still be released on removal.
+    topHost.remove();
+    manager.removeShadowRootsIn(topHost);
+    expect(mutationBuffers.length).toBe(buffersBeforeAdd);
+  });
 });