[8.x] [Lens][Embeddable] Fix memory leak on ES|QL variables subscription (#210826) (#211171)

# Backport This will backport the following commits from `main` to `8.x`: - [[Lens][Embeddable] Fix memory leak on ES|QL variables subscription (#210826)](https://github.com/elastic/kibana/pull/210826)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Marco Liberati <dej611@users.noreply.github.com>
2025-04-24 09:48:58 -04:00 · 2025-02-14 22:10:18 +11:00 · 2025-02-14 22:10:18 +11:00 · 7c45530aab
commit 7c45530aab
parent 9f58dbbc77
3 changed files with 102 additions and 36 deletions
--- a/x-pack/platform/plugins/shared/lens/public/react_embeddable/data_loader.test.ts
+++ b/x-pack/platform/plugins/shared/lens/public/react_embeddable/data_loader.test.ts
@ -5,7 +5,7 @@
 * 2.0.
 */
 import { faker } from '@faker-js/faker';
-import { loadEmbeddableData } from './data_loader';
+import { ReloadReason, loadEmbeddableData } from './data_loader';
 import {
  createUnifiedSearchApi,
  getLensApiMock,
@ -35,11 +35,16 @@ import {
 import { PublishesSearchSession } from '@kbn/presentation-publishing/interfaces/fetch/publishes_search_session';
 import { isObject } from 'lodash';
 import { createMockDatasource, defaultDoc } from '../mocks';
+import { ESQLControlVariable, ESQLVariableType } from '@kbn/esql-validation-autocomplete';
+import * as Logger from './logger';
+import { buildObservableVariable } from './helper';

 jest.mock('@kbn/interpreter', () => ({
  toExpression: jest.fn().mockReturnValue('expression'),
 }));

+const loggerFn = jest.spyOn(Logger, 'addLog');
+
 // Mock it for now, later investigate why the real one is not triggering here on tests
 jest.mock('@kbn/presentation-publishing', () => {
  const original = jest.requireActual('@kbn/presentation-publishing');
@ -82,9 +87,9 @@ type ChangeFnType = ({
      searchSessionId$: BehaviorSubject<string>;
    };
  services: LensEmbeddableStartServices;
-}) => Promise<void | boolean>;
+}) => Promise<void | ReloadReason | false>;

-async function expectRerenderOnDataLoder(
+async function expectRerenderOnDataLoader(
  changeFn: ChangeFnType,
  runtimeState: LensRuntimeState = { attributes: getLensAttributesMock() },
  {
@ -97,6 +102,7 @@ async function expectRerenderOnDataLoder(
        filters$: BehaviorSubject<Filter[] | undefined>;
        query$: BehaviorSubject<Query | AggregateQuery | undefined>;
        timeRange$: BehaviorSubject<TimeRange | undefined>;
+        esqlVariables$: BehaviorSubject<ESQLControlVariable[] | undefined>;
      } & LensOverrides
    >;
    internalApiOverrides?: Partial<LensInternalApi>;
@ -124,7 +130,10 @@ async function expectRerenderOnDataLoder(
    parentApi,
  };
  const getState = jest.fn(() => runtimeState);
-  const internalApi = getLensInternalApiMock(internalApiOverrides);
+  const internalApi = getLensInternalApiMock({
+    ...internalApiOverrides,
+    attributes$: buildObservableVariable(runtimeState.attributes)[0],
+  });
  const services = {
    ...makeEmbeddableServices(new BehaviorSubject<string>(''), undefined, {
      visOverrides: { id: 'lnsXY' },
@ -153,11 +162,22 @@ async function expectRerenderOnDataLoder(
    services,
  });
  // fallback to true if undefined is returned
-  const expectRerender = result ?? true;
+  const expectRerender = result === false ? false : true;
+  // Add an advanced check if provided: the reload reason
+  const rerenderReason = typeof result === 'string' ? result : undefined;
  // there's a debounce, so skip to the next tick
  jest.advanceTimersByTime(200);
  // unsubscribe to all observables before checking
  cleanup();
+
+  if (expectRerender && rerenderReason) {
+    const reloadCalls = loggerFn.mock.calls.filter((call) =>
+      call[0].startsWith('Embeddable reload reason')
+    );
+    expect(reloadCalls[reloadCalls.length - 1][0]).toBe(
+      `Embeddable reload reason: ${rerenderReason}`
+    );
+  }
  // now check if the re-render has been dispatched
  expect(internalApi.dispatchRenderStart).toHaveBeenCalledTimes(expectRerender ? 2 : 1);
 }
@ -177,40 +197,45 @@ describe('Data Loader', () => {
  });
  afterAll(() => {
    jest.useRealTimers();
+    loggerFn.mockRestore();
  });

  beforeEach(() => jest.clearAllMocks());

  it('should re-render once on filter change', async () => {
-    await expectRerenderOnDataLoder(async ({ api }) => {
+    await expectRerenderOnDataLoader(async ({ api }) => {
      (api.filters$ as BehaviorSubject<Filter[]>).next([
        { meta: { alias: 'test', negate: false, disabled: false } },
      ]);
+      return 'searchContext';
    });
  });

  it('should re-render once on search session change', async () => {
-    await expectRerenderOnDataLoder(async ({ api }) => {
+    await expectRerenderOnDataLoader(async ({ api }) => {
      // dispatch a new searchSessionId

      (
        api.parentApi as unknown as { searchSessionId$: BehaviorSubject<string | undefined> }
      ).searchSessionId$.next('newSessionId');
+
+      return 'searchContext';
    });
  });

  it('should re-render once on attributes change', async () => {
-    await expectRerenderOnDataLoder(async ({ internalApi }) => {
+    await expectRerenderOnDataLoader(async ({ internalApi }) => {
      // trigger a change by changing the title in the attributes
      (internalApi.attributes$ as BehaviorSubject<LensDocument | undefined>).next({
        ...internalApi.attributes$.getValue(),
        title: faker.lorem.word(),
      });
+      return 'attributes';
    });
  });

  it('should re-render when dashboard view/edit mode changes if dynamic actions are set', async () => {
-    await expectRerenderOnDataLoder(async ({ api, getState }) => {
+    await expectRerenderOnDataLoader(async ({ api, getState }) => {
      getState.mockReturnValue({
        attributes: getLensAttributesMock(),
        enhancements: {
@ -221,11 +246,13 @@ describe('Data Loader', () => {
      });
      // trigger a change by changing the title in the attributes
      (api.viewMode$ as BehaviorSubject<ViewMode | undefined>).next('view');
+
+      return 'viewMode';
    });
  });

  it('should not re-render when dashboard view/edit mode changes if dynamic actions are not set', async () => {
-    await expectRerenderOnDataLoder(async ({ api }) => {
+    await expectRerenderOnDataLoader(async ({ api }) => {
      // the default get state does not have dynamic actions
      // trigger a change by changing the title in the attributes
      (api.viewMode$ as BehaviorSubject<ViewMode | undefined>).next('view');
@ -238,7 +265,7 @@ describe('Data Loader', () => {
    const query: Query = { language: 'kquery', query: '' };
    const filters: Filter[] = [{ meta: { alias: 'test', negate: false, disabled: false } }];

-    await expectRerenderOnDataLoder(
+    await expectRerenderOnDataLoader(
      async ({ internalApi }) => {
        await waitForValue(
          internalApi.expressionParams$,
@ -258,7 +285,7 @@ describe('Data Loader', () => {
  });

  it('should pass render mode to expression', async () => {
-    await expectRerenderOnDataLoder(async ({ internalApi }) => {
+    await expectRerenderOnDataLoader(async ({ internalApi }) => {
      await waitForValue(
        internalApi.expressionParams$,
        (v: unknown) => isObject(v) && 'renderMode' in v
@ -295,7 +322,7 @@ describe('Data Loader', () => {
      ],
    };

-    await expectRerenderOnDataLoder(
+    await expectRerenderOnDataLoader(
      async ({ internalApi }) => {
        await waitForValue(
          internalApi.expressionParams$,
@ -327,7 +354,7 @@ describe('Data Loader', () => {
  });

  it('should call onload after rerender and onData$ call', async () => {
-    await expectRerenderOnDataLoder(async ({ parentApi, internalApi, api }) => {
+    await expectRerenderOnDataLoader(async ({ parentApi, internalApi, api }) => {
      expect(parentApi.onLoad).toHaveBeenLastCalledWith(true);

      await waitForValue(
@ -347,7 +374,7 @@ describe('Data Loader', () => {
  });

  it('should initialize dateViews api with deduped list of index patterns', async () => {
-    await expectRerenderOnDataLoder(
+    await expectRerenderOnDataLoader(
      async ({ internalApi }) => {
        await waitForValue(
          internalApi.dataViews$,
@ -374,7 +401,7 @@ describe('Data Loader', () => {
  });

  it('should override noPadding in the display options if noPadding is set in the embeddable input', async () => {
-    await expectRerenderOnDataLoder(async ({ internalApi }) => {
+    await expectRerenderOnDataLoader(async ({ internalApi }) => {
      await waitForValue(
        internalApi.expressionParams$,
        (v: unknown) => isObject(v) && 'expression' in v && typeof v.expression != null
@ -387,18 +414,19 @@ describe('Data Loader', () => {
  });

  it('should reload only once when the attributes or savedObjectId and the search context change at the same time', async () => {
-    await expectRerenderOnDataLoder(async ({ internalApi, api }) => {
+    await expectRerenderOnDataLoader(async ({ internalApi, api }) => {
      // trigger a change by changing the title in the attributes
      (internalApi.attributes$ as BehaviorSubject<LensDocument | undefined>).next({
        ...internalApi.attributes$.getValue(),
        title: faker.lorem.word(),
      });
      (api.savedObjectId$ as BehaviorSubject<string | undefined>).next('newSavedObjectId');
+      return 'savedObjectId';
    });
  });

  it('should pass over the overrides as variables', async () => {
-    await expectRerenderOnDataLoder(
+    await expectRerenderOnDataLoader(
      async ({ internalApi }) => {
        await waitForValue(
          internalApi.expressionParams$,
@ -430,7 +458,7 @@ describe('Data Loader', () => {
  });

  it('should catch missing dataView errors correctly', async () => {
-    await expectRerenderOnDataLoder(
+    await expectRerenderOnDataLoader(
      async ({ internalApi }) => {
        // wait for the error to appear
        await waitForValue(internalApi.blockingError$);
@ -478,4 +506,45 @@ describe('Data Loader', () => {
      }
    );
  });
+
+  it('should re-render on ES|QL variable changes', async () => {
+    const baseAttributes = getLensAttributesMock();
+    await expectRerenderOnDataLoader(
+      async ({ internalApi }) => {
+        (internalApi.esqlVariables$ as BehaviorSubject<ESQLControlVariable[]>).next([
+          { key: 'foo', value: faker.database.column(), type: ESQLVariableType.FIELDS },
+        ]);
+        return 'ESQLvariables';
+      },
+      {
+        attributes: getLensAttributesMock({
+          state: { ...baseAttributes.state, query: { esql: 'from index | where $foo > 0' } },
+        }),
+      }
+    );
+  });
+
+  it('should not re-render on ES|QL variable identical changes', async () => {
+    const baseAttributes = getLensAttributesMock();
+    const variables: ESQLControlVariable[] = [
+      { key: 'foo', value: faker.database.column(), type: ESQLVariableType.FIELDS },
+    ];
+    await expectRerenderOnDataLoader(
+      async ({ internalApi }) => {
+        (internalApi.esqlVariables$ as BehaviorSubject<ESQLControlVariable[]>).next(variables);
+        // no rerender
+        return false;
+      },
+      {
+        attributes: getLensAttributesMock({
+          state: { ...baseAttributes.state, query: { esql: 'from index | where $foo > 0' } },
+        }),
+      },
+      {
+        internalApiOverrides: {
+          esqlVariables$: buildObservableVariable<ESQLControlVariable[]>(variables)[0],
+        },
+      }
+    );
+  });
 });
--- a/x-pack/platform/plugins/shared/lens/public/react_embeddable/data_loader.ts
+++ b/x-pack/platform/plugins/shared/lens/public/react_embeddable/data_loader.ts
@ -31,17 +31,18 @@ import { buildUserMessagesHelpers } from './user_messages/api';
 import { getLogError } from './expressions/telemetry';
 import type { SharingSavedObjectProps, UserMessagesDisplayLocationId } from '../types';
 import { apiHasLensComponentCallbacks } from './type_guards';
-import { getRenderMode, getParentContext } from './helper';
+import { getRenderMode, getParentContext, buildObservableVariable } from './helper';
 import { addLog } from './logger';
 import { getUsedDataViews } from './expressions/update_data_views';
 import { getMergedSearchContext } from './expressions/merged_search_context';
+import { getEmbeddableVariables } from './initializers/utils';

 const blockingMessageDisplayLocations: UserMessagesDisplayLocationId[] = [
  'visualization',
  'visualizationOnEmbeddable',
 ];

-type ReloadReason =
+export type ReloadReason =
  | 'ESQLvariables'
  | 'attributes'
  | 'savedObjectId'
@ -118,6 +119,8 @@ export function loadEmbeddableData(
    }
  };

+  const [controlESQLVariables$] = buildObservableVariable<ESQLControlVariable[]>([]);
+
  async function reload(
    // make reload easier to debug
    sourceId: ReloadReason
@ -189,11 +192,9 @@ export function loadEmbeddableData(
      callbacks
    );

-    const esqlVariables = internalApi?.esqlVariables$?.getValue();
-
    const searchContext = getMergedSearchContext(
      currentState,
-      getSearchContext(parentApi, esqlVariables),
+      getSearchContext(parentApi, controlESQLVariables$?.getValue()),
      api.timeRange$,
      parentApi,
      services
@ -260,7 +261,7 @@ export function loadEmbeddableData(
  const mergedSubscriptions = merge(
    // on search context change, reload
    fetch$(api).pipe(map(() => 'searchContext' as ReloadReason)),
-    internalApi?.esqlVariables$.pipe(
+    controlESQLVariables$.pipe(
      waitUntilChanged(),
      map(() => 'ESQLvariables' as ReloadReason)
    ),
@ -294,6 +295,12 @@ export function loadEmbeddableData(

  const subscriptions: Subscription[] = [
    mergedSubscriptions.pipe(debounceTime(0)).subscribe(reload),
+    // In case of changes to the dashboard ES|QL controls, re-map them
+    internalApi.esqlVariables$.subscribe((newVariables: ESQLControlVariable[]) => {
+      const query = internalApi.attributes$.getValue().state?.query;
+      const esqlVariables = getEmbeddableVariables(query, newVariables) ?? [];
+      controlESQLVariables$.next(esqlVariables);
+    }),
    // make sure to reload on viewMode change
    api.viewMode$.subscribe(() => {
      // only reload if drilldowns are set
--- a/x-pack/platform/plugins/shared/lens/public/react_embeddable/initializers/initialize_internal_api.ts
+++ b/x-pack/platform/plugins/shared/lens/public/react_embeddable/initializers/initialize_internal_api.ts
@ -8,7 +8,6 @@
 import { BehaviorSubject } from 'rxjs';
 import { initializeTitleManager } from '@kbn/presentation-publishing';
 import { apiPublishesESQLVariables } from '@kbn/esql-variables-types';
-import type { ESQLControlVariable } from '@kbn/esql-validation-autocomplete';
 import type { DataView } from '@kbn/data-views-plugin/common';
 import { buildObservableVariable, createEmptyLensState } from '../helper';
 import type {
@ -21,7 +20,6 @@ import type {
 } from '../types';
 import { apiHasAbortController, apiHasLensComponentProps } from '../type_guards';
 import type { UserMessage } from '../../types';
-import { getEmbeddableVariables } from './utils';

 export function initializeInternalApi(
  initialState: LensRuntimeState,
@ -73,21 +71,13 @@ export function initializeInternalApi(
    apiPublishesESQLVariables(parentApi) ? parentApi.esqlVariables$ : []
  );

-  const query = initialState.attributes.state.query;
-
-  const panelEsqlVariables$ = new BehaviorSubject<ESQLControlVariable[]>([]);
-  esqlVariables$.subscribe((newVariables) => {
-    const esqlVariables = getEmbeddableVariables(query, newVariables) ?? [];
-    panelEsqlVariables$.next(esqlVariables);
-  });
-
  // No need to expose anything at public API right now, that would happen later on
  // where each initializer will pick what it needs and publish it
  return {
    attributes$,
    overrides$,
    disableTriggers$,
-    esqlVariables$: panelEsqlVariables$,
+    esqlVariables$,
    dataLoading$,
    hasRenderCompleted$,
    expressionParams$,