[Lens] Fully unmount React when flyout closes

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,52 @@ 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();
+  });
 });

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

@@ -5,7 +5,8 @@
  * 2.0.
  */
 
-import React, { HTMLAttributes } from 'react';
+import React, { HTMLAttributes, useEffect, useRef } from 'react';
+import { unmountComponentAtNode } from 'react-dom';
 
 export interface NativeRendererProps<T> extends HTMLAttributes<HTMLDivElement> {
   render: (domElement: Element, props: T) => void;
@@ -19,11 +20,32 @@ 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) {
+          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;
+        cleanupRef.current = render(el, nativeProps);
+      }
+    },
   });
 }

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

@@ -185,10 +185,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;
@@ -585,12 +597,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
@@ -620,7 +638,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