[Lens] Fully unmount React when flyout closes (#95359)

* [Lens] Fully unmount React when flyout closes

* Fix bug with editor frame unmounting

* Fix type

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

x-pack/plugins/lens/public/app_plugin/app.test.tsx

@@ -72,7 +72,7 @@ const { TopNavMenu } = navigationStartMock.ui;
 
 function createMockFrame(): jest.Mocked<EditorFrameInstance> {
   return {
-    mount: jest.fn((el, props) => {}),
+    mount: jest.fn(async (el, props) => {}),
     unmount: jest.fn(() => {}),
   };
 }

x-pack/plugins/lens/public/native_renderer/native_renderer.test.tsx

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import React from 'react';
+import React, { useEffect } from 'react';
 import { render } from 'react-dom';
 import { NativeRenderer } from './native_renderer';
 import { act } from 'react-dom/test-utils';
@@ -151,4 +151,102 @@ describe('native_renderer', () => {
     const containerElement: Element = mountpoint.firstElementChild!;
     expect(containerElement.nodeName).toBe('SPAN');
   });
+
+  it('should properly unmount a react element that is mounted inside the renderer', () => {
+    let isUnmounted = false;
+
+    function TestComponent() {
+      useEffect(() => {
+        return () => {
+          isUnmounted = true;
+        };
+      }, []);
+      return <>Hello</>;
+    }
+
+    renderAndTriggerHooks(
+      <NativeRenderer
+        render={(element) => {
+          // This render function mimics the most common usage inside Lens
+          render(<TestComponent />, element);
+        }}
+        nativeProps={{}}
+      />,
+      mountpoint
+    );
+
+    // Replaces the component at the mountpoint with nothing
+    renderAndTriggerHooks(<>Empty</>, mountpoint);
+
+    expect(isUnmounted).toBe(true);
+  });
+
+  it('should call the unmount function provided for non-react elements', () => {
+    const unmountCallback = jest.fn();
+
+    renderAndTriggerHooks(
+      <NativeRenderer
+        render={(element, props) => {
+          return unmountCallback;
+        }}
+        nativeProps={{}}
+      />,
+      mountpoint
+    );
+
+    // Replaces the component at the mountpoint with nothing
+    renderAndTriggerHooks(<>Empty</>, mountpoint);
+
+    expect(unmountCallback).toHaveBeenCalled();
+  });
+
+  it('should handle when the mount function is asynchronous without a cleanup fn', () => {
+    let isUnmounted = false;
+
+    function TestComponent() {
+      useEffect(() => {
+        return () => {
+          isUnmounted = true;
+        };
+      }, []);
+      return <>Hello</>;
+    }
+
+    renderAndTriggerHooks(
+      <NativeRenderer
+        render={async (element, props) => {
+          render(<TestComponent />, element);
+        }}
+        nativeProps={{}}
+      />,
+      mountpoint
+    );
+
+    // Replaces the component at the mountpoint with nothing
+    renderAndTriggerHooks(<>Empty</>, mountpoint);
+
+    expect(isUnmounted).toBe(true);
+  });
+
+  it('should handle when the mount function is asynchronous with a cleanup fn', async () => {
+    const unmountCallback = jest.fn();
+
+    renderAndTriggerHooks(
+      <NativeRenderer
+        render={async (element, props) => {
+          return unmountCallback;
+        }}
+        nativeProps={{}}
+      />,
+      mountpoint
+    );
+
+    // Schedule a promise cycle to update the DOM
+    await new Promise((resolve) => setTimeout(resolve, 0));
+
+    // Replaces the component at the mountpoint with nothing
+    renderAndTriggerHooks(<>Empty</>, mountpoint);
+
+    expect(unmountCallback).toHaveBeenCalled();
+  });
 });

x-pack/plugins/lens/public/native_renderer/native_renderer.tsx

@@ -5,10 +5,16 @@
  * 2.0.
  */
 
-import React, { HTMLAttributes } from 'react';
+import React, { HTMLAttributes, useEffect, useRef } from 'react';
+import { unmountComponentAtNode } from 'react-dom';
+
+type CleanupCallback = (el: Element) => void;
 
 export interface NativeRendererProps<T> extends HTMLAttributes<HTMLDivElement> {
-  render: (domElement: Element, props: T) => void;
+  render: (
+    domElement: Element,
+    props: T
+  ) => Promise<CleanupCallback | void> | CleanupCallback | void;
   nativeProps: T;
   tag?: string;
 }
@@ -19,11 +25,42 @@ export interface NativeRendererProps<T> extends HTMLAttributes<HTMLDivElement> {
  * By default the mountpoint element will be a div, this can be changed with the
  * `tag` prop.
  *
+ * If the rendered component tree was using React, we need to clean it up manually,
+ * otherwise the unmount event never happens. A future addition is for non-React components
+ * to get cleaned up, which could be added in the future.
+ *
  * @param props
  */
 export function NativeRenderer<T>({ render, nativeProps, tag, ...rest }: NativeRendererProps<T>) {
+  const elementRef = useRef<Element>();
+  const cleanupRef = useRef<((cleanupElement: Element) => void) | void>();
+  useEffect(() => {
+    return () => {
+      if (elementRef.current) {
+        if (cleanupRef.current && typeof cleanupRef.current === 'function') {
+          cleanupRef.current(elementRef.current);
+        }
+        unmountComponentAtNode(elementRef.current);
+      }
+    };
+  }, []);
   return React.createElement(tag || 'div', {
     ...rest,
-    ref: (el) => el && render(el, nativeProps),
+    ref: (el) => {
+      if (el) {
+        elementRef.current = el;
+        // Handles the editor frame renderer, which is async
+        const result = render(el, nativeProps);
+        if (result instanceof Promise) {
+          result.then((cleanup) => {
+            if (typeof cleanup === 'function') {
+              cleanupRef.current = cleanup;
+            }
+          });
+        } else if (typeof result === 'function') {
+          cleanupRef.current = result;
+        }
+      }
+    },
   });
 }

x-pack/plugins/lens/public/types.ts

@@ -64,7 +64,7 @@ export interface EditorFrameProps {
   showNoDataPopover: () => void;
 }
 export interface EditorFrameInstance {
-  mount: (element: Element, props: EditorFrameProps) => void;
+  mount: (element: Element, props: EditorFrameProps) => Promise<void>;
   unmount: () => void;
 }
 
@@ -190,10 +190,22 @@ export interface Datasource<T = unknown, P = unknown> {
   getLayers: (state: T) => string[];
   removeColumn: (props: { prevState: T; layerId: string; columnId: string }) => T;
 
-  renderDataPanel: (domElement: Element, props: DatasourceDataPanelProps<T>) => void;
-  renderDimensionTrigger: (domElement: Element, props: DatasourceDimensionTriggerProps<T>) => void;
-  renderDimensionEditor: (domElement: Element, props: DatasourceDimensionEditorProps<T>) => void;
-  renderLayerPanel: (domElement: Element, props: DatasourceLayerPanelProps<T>) => void;
+  renderDataPanel: (
+    domElement: Element,
+    props: DatasourceDataPanelProps<T>
+  ) => ((cleanupElement: Element) => void) | void;
+  renderDimensionTrigger: (
+    domElement: Element,
+    props: DatasourceDimensionTriggerProps<T>
+  ) => ((cleanupElement: Element) => void) | void;
+  renderDimensionEditor: (
+    domElement: Element,
+    props: DatasourceDimensionEditorProps<T>
+  ) => ((cleanupElement: Element) => void) | void;
+  renderLayerPanel: (
+    domElement: Element,
+    props: DatasourceLayerPanelProps<T>
+  ) => ((cleanupElement: Element) => void) | void;
   getDropProps: (
     props: DatasourceDimensionDropProps<T> & {
       groupId: string;
@@ -591,12 +603,18 @@ export interface Visualization<T = unknown> {
    * Popover contents that open when the user clicks the contextMenuIcon. This can be used
    * for extra configurability, such as for styling the legend or axis
    */
-  renderLayerContextMenu?: (domElement: Element, props: VisualizationLayerWidgetProps<T>) => void;
+  renderLayerContextMenu?: (
+    domElement: Element,
+    props: VisualizationLayerWidgetProps<T>
+  ) => ((cleanupElement: Element) => void) | void;
   /**
    * Toolbar rendered above the visualization. This is meant to be used to provide chart-level
    * settings for the visualization.
    */
-  renderToolbar?: (domElement: Element, props: VisualizationToolbarProps<T>) => void;
+  renderToolbar?: (
+    domElement: Element,
+    props: VisualizationToolbarProps<T>
+  ) => ((cleanupElement: Element) => void) | void;
   /**
    * Visualizations can provide a custom icon which will open a layer-specific popover
    * If no icon is provided, gear icon is default
@@ -626,7 +644,7 @@ export interface Visualization<T = unknown> {
   renderDimensionEditor?: (
     domElement: Element,
     props: VisualizationDimensionEditorProps<T>
-  ) => void;
+  ) => ((cleanupElement: Element) => void) | void;
 
   /**
    * The frame will call this function on all visualizations at different times. The