Fix bug with editor frame unmounting

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

@@ -199,4 +199,54 @@ describe('native_renderer', () => {
 
     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

@@ -8,8 +8,13 @@
 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;
 }
@@ -32,7 +37,7 @@ export function NativeRenderer<T>({ render, nativeProps, tag, ...rest }: NativeR
   useEffect(() => {
     return () => {
       if (elementRef.current) {
-        if (cleanupRef.current) {
+        if (cleanupRef.current && typeof cleanupRef.current === 'function') {
           cleanupRef.current(elementRef.current);
         }
         unmountComponentAtNode(elementRef.current);
@@ -44,7 +49,17 @@ export function NativeRenderer<T>({ render, nativeProps, tag, ...rest }: NativeR
     ref: (el) => {
       if (el) {
         elementRef.current = el;
-        cleanupRef.current = render(el, nativeProps);
+        // 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;
 }