[expressions] caching (#180440)

src/plugins/data/common/search/expressions/eql.ts

@@ -58,6 +58,7 @@ export const getEqlFn = ({
     name,
     type: 'eql_raw_response',
     inputTypes: ['kibana_context', 'null'],
+    allowCache: true,
     help: i18n.translate('data.search.eql.help', {
       defaultMessage: 'Run Elasticsearch request',
     }),

src/plugins/data/common/search/expressions/esaggs/esaggs_fn.ts

@@ -57,6 +57,7 @@ export const getEsaggsMeta: () => Omit<EsaggsExpressionFunctionDefinition, 'fn'>
   name,
   type: 'datatable',
   inputTypes: ['kibana_context', 'null'],
+  allowCache: true,
   help: i18n.translate('data.functions.esaggs.help', {
     defaultMessage: 'Run AggConfig aggregation',
   }),

src/plugins/data/common/search/expressions/esdsl.ts

@@ -51,6 +51,7 @@ export const getEsdslFn = ({
     name,
     type: 'es_raw_response',
     inputTypes: ['kibana_context', 'null'],
+    allowCache: true,
     help: i18n.translate('data.search.esdsl.help', {
       defaultMessage: 'Run Elasticsearch request',
     }),

src/plugins/data/common/search/expressions/esql.ts

@@ -75,6 +75,7 @@ export const getEsqlFn = ({ getStartDependencies }: EsqlFnArguments) => {
     name: 'esql',
     type: 'datatable',
     inputTypes: ['kibana_context', 'null'],
+    allowCache: true,
     help: i18n.translate('data.search.esql.help', {
       defaultMessage: 'Queries Elasticsearch using ES|QL.',
     }),

src/plugins/data/common/search/expressions/essql.ts

@@ -66,6 +66,7 @@ export const getEssqlFn = ({ getStartDependencies }: EssqlFnArguments) => {
     name: 'essql',
     type: 'datatable',
     inputTypes: ['kibana_context', 'null'],
+    allowCache: true,
     help: i18n.translate('data.search.essql.help', {
       defaultMessage: 'Queries Elasticsearch using Elasticsearch SQL.',
     }),

src/plugins/expressions/common/execution/execution.ts

@@ -27,7 +27,7 @@ import {
   Subscription,
 } from 'rxjs';
 import { catchError, finalize, map, pluck, shareReplay, switchMap, tap } from 'rxjs';
-import { now, AbortError } from '@kbn/kibana-utils-plugin/common';
+import { now, AbortError, calculateObjectHash } from '@kbn/kibana-utils-plugin/common';
 import { Adapters } from '@kbn/inspector-plugin/common';
 import { Executor } from '../executor';
 import { createExecutionContainer, ExecutionContainer } from './container';
@@ -55,6 +55,10 @@ type UnwrapReturnType<Function extends (...args: any[]) => unknown> =
     ? UnwrapObservable<ReturnType<Function>>
     : Awaited<ReturnType<Function>>;
 
+export interface FunctionCacheItem {
+  value: unknown;
+  time: number;
+}
 /**
  * The result returned after an expression function execution.
  */
@@ -70,6 +74,8 @@ export interface ExecutionResult<Output> {
   result: Output;
 }
 
+const maxCacheSize = 1000;
+
 const createAbortErrorValue = () =>
   createError({
     message: 'The expression was aborted.',
@@ -235,6 +241,7 @@ export class Execution<
    * @private
    */
   private readonly childExecutions: Execution[] = [];
+  private cacheTimeout: number = 30000;
 
   /**
    * Contract is a public representation of `Execution` instances. Contract we
@@ -248,7 +255,11 @@ export class Execution<
     return this.context.inspectorAdapters;
   }
 
-  constructor(public readonly execution: ExecutionParams, private readonly logger?: Logger) {
+  constructor(
+    public readonly execution: ExecutionParams,
+    private readonly logger?: Logger,
+    private readonly functionCache: Map<string, FunctionCacheItem> = new Map()
+  ) {
     const { executor } = execution;
 
     this.contract = new ExecutionContract<Input, Output, InspectorAdapters>(this);
@@ -278,6 +289,7 @@ export class Execution<
         ? () => execution.params.kibanaRequest!
         : undefined,
       variables: execution.params.variables || {},
+      allowCache: this.execution.params.allowCache,
       types: executor.getTypes(),
       abortSignal: this.abortController.signal,
       inspectorAdapters,
@@ -454,42 +466,75 @@ export class Execution<
     input: unknown,
     args: Record<string, unknown>
   ): Observable<UnwrapReturnType<Fn['fn']>> {
-    return of(input).pipe(
-      map((currentInput) => this.cast(currentInput, fn.inputTypes)),
-      switchMap((normalizedInput) => of(fn.fn(normalizedInput, args, this.context))),
-      switchMap(
-        (fnResult) =>
-          (isObservable(fnResult)
-            ? fnResult
-            : from(isPromise(fnResult) ? fnResult : [fnResult])) as Observable<
-            UnwrapReturnType<Fn['fn']>
-          >
-      ),
-      map((output) => {
-        // Validate that the function returned the type it said it would.
-        // This isn't required, but it keeps function developers honest.
-        const returnType = getType(output);
-        const expectedType = fn.type;
-        if (expectedType && returnType !== expectedType) {
-          throw new Error(
-            `Function '${fn.name}' should return '${expectedType}',` +
-              ` actually returned '${returnType}'`
-          );
-        }
+    let hash: string | undefined;
+    let lastValue: unknown;
+    let completionFlag = false;
+
+    return of(input)
+      .pipe(
+        map((currentInput) => this.cast(currentInput, fn.inputTypes)),
+        switchMap((normalizedInput) => {
+          if (fn.allowCache && this.context.allowCache) {
+            hash = calculateObjectHash([
+              fn.name,
+              normalizedInput,
+              args,
+              this.context.getSearchContext(),
+            ]);
+          }
+          if (hash && this.functionCache.has(hash)) {
+            const cached = this.functionCache.get(hash);
+            if (cached && Date.now() - cached.time < this.cacheTimeout) {
+              return of(cached.value);
+            }
+          }
+          return of(fn.fn(normalizedInput, args, this.context));
+        }),
+        switchMap((fnResult) => {
+          return (
+            isObservable(fnResult) ? fnResult : from(isPromise(fnResult) ? fnResult : [fnResult])
+          ) as Observable<UnwrapReturnType<Fn['fn']>>;
+        }),
+        map((output) => {
+          // Validate that the function returned the type it said it would.
+          // This isn't required, but it keeps function developers honest.
+          const returnType = getType(output);
+          const expectedType = fn.type;
+          if (expectedType && returnType !== expectedType) {
+            throw new Error(
+              `Function '${fn.name}' should return '${expectedType}',` +
+                ` actually returned '${returnType}'`
+            );
+          }
 
-        // Validate the function output against the type definition's validate function.
-        const type = this.context.types[fn.type];
-        if (type && type.validate) {
-          try {
-            type.validate(output);
-          } catch (e) {
-            throw new Error(`Output of '${fn.name}' is not a valid type '${fn.type}': ${e}`);
+          // Validate the function output against the type definition's validate function.
+          const type = this.context.types[fn.type];
+          if (type && type.validate) {
+            try {
+              type.validate(output);
+            } catch (e) {
+              throw new Error(`Output of '${fn.name}' is not a valid type '${fn.type}': ${e}`);
+            }
           }
-        }
 
-        return output;
-      })
-    );
+          lastValue = output;
+
+          return output;
+        }),
+        finalize(() => {
+          if (completionFlag && hash) {
+            while (this.functionCache.size >= maxCacheSize) {
+              this.functionCache.delete(this.functionCache.keys().next().value);
+            }
+            this.functionCache.set(hash, { value: lastValue, time: Date.now() });
+          }
+        })
+      )
+      .pipe(
+        tap({
+          complete: () => (completionFlag = true), // Set flag true only on successful completion
+        })
+      );
   }
 
   public cast<Type = unknown>(value: unknown, toTypeNames?: string[]): Type {

src/plugins/expressions/common/execution/types.ts

@@ -35,6 +35,11 @@ export interface ExecutionContext<InspectorAdapters extends Adapters = Adapters>
    */
   types: Record<string, ExpressionType>;
 
+  /**
+   * Allow caching in the current execution.
+   */
+  allowCache?: boolean;
+
   /**
    * Adds ability to abort current execution.
    */

src/plugins/expressions/common/executor/executor.execution.test.ts

@@ -27,7 +27,8 @@ describe('Executor mocked execution tests', () => {
 
         expect(Execution).toHaveBeenCalledWith(
           expect.objectContaining({ expression: 'foo bar="baz"' }),
-          undefined
+          undefined,
+          expect.anything()
         );
       });
     });
@@ -40,7 +41,8 @@ describe('Executor mocked execution tests', () => {
 
         expect(Execution).toHaveBeenCalledWith(
           expect.not.objectContaining({ expression: expect.anything() }),
-          undefined
+          undefined,
+          expect.anything()
         );
       });
     });

src/plugins/expressions/common/executor/executor.test.ts

@@ -9,7 +9,7 @@
 import { Executor } from './executor';
 import * as expressionTypes from '../expression_types';
 import * as expressionFunctions from '../expression_functions';
-import { Execution } from '../execution';
+import { Execution, FunctionCacheItem } from '../execution';
 import { ExpressionAstFunction, parseExpression, formatExpression } from '../ast';
 import { MigrateFunction } from '@kbn/kibana-utils-plugin/common/persistable_state';
 import { SavedObjectReference } from '@kbn/core/types';
@@ -312,4 +312,92 @@ describe('Executor', () => {
       });
     });
   });
+
+  describe('caching', () => {
+    const functionCache: Map<string, FunctionCacheItem> = new Map();
+    const fakeCacheEntry = { time: Date.now(), value: 'test' };
+    let executor: Executor;
+
+    beforeAll(() => {
+      executor = new Executor(undefined, undefined, functionCache);
+      executor.registerFunction(expressionFunctions.variable);
+      expressionFunctions.theme.allowCache = true;
+      executor.registerFunction(expressionFunctions.theme);
+    });
+
+    afterEach(() => {
+      functionCache.clear();
+    });
+
+    it('caches the result of function', async () => {
+      await executor.run('theme size default=12', null, { allowCache: true }).toPromise();
+      expect(functionCache.size).toEqual(1);
+      const entry = functionCache.keys().next().value;
+      functionCache.set(entry, fakeCacheEntry);
+      const result = await executor
+        .run('theme size default=12', null, { allowCache: true })
+        .toPromise();
+      expect(functionCache.size).toEqual(1);
+      expect(result?.result).toEqual(fakeCacheEntry.value);
+    });
+
+    it('doesnt cache if allowCache flag is false', async () => {
+      await executor.run('theme size default=12', null, { allowCache: true }).toPromise();
+      expect(functionCache.size).toEqual(1);
+      const entry = functionCache.keys().next().value;
+      functionCache.set(entry, fakeCacheEntry);
+      const result = await executor
+        .run('theme size default=12', null, { allowCache: false })
+        .toPromise();
+      expect(functionCache.size).toEqual(1);
+      expect(result?.result).not.toEqual(fakeCacheEntry.value);
+    });
+
+    it('doesnt cache results of functions that have allowCache property set to false', async () => {
+      await executor.run('var name="test"', null, { allowCache: true }).toPromise();
+      expect(functionCache.size).toEqual(0);
+    });
+
+    describe('doesnt use cached version', () => {
+      const cachedVersion = { time: Date.now(), value: 'value' };
+
+      beforeAll(async () => {
+        await executor.run('theme size default=12', null, { allowCache: true }).toPromise();
+        expect(functionCache.size).toEqual(1);
+        const entry: string = Object.keys(functionCache)[0];
+        functionCache.set(entry, cachedVersion);
+      });
+
+      it('input changed', async () => {
+        const result = await executor
+          .run(
+            'theme size default=12',
+            {
+              type: 'kibana_context',
+              value: 'test',
+            },
+            { allowCache: true }
+          )
+          .toPromise();
+        expect(result).not.toEqual(cachedVersion);
+      });
+
+      it('arguments changed', async () => {
+        const result = await executor
+          .run('theme size default=14', null, { allowCache: true })
+          .toPromise();
+        expect(result).not.toEqual(cachedVersion);
+      });
+
+      it('search context changed', async () => {
+        const result = await executor
+          .run('theme size default=12', null, {
+            searchContext: { filters: [] },
+            allowCache: true,
+          })
+          .toPromise();
+        expect(result).not.toEqual(cachedVersion);
+      });
+    });
+  });
 });

src/plugins/expressions/common/executor/executor.ts

@@ -22,7 +22,12 @@ import {
 import { ExecutorState, ExecutorContainer } from './container';
 import { createExecutorContainer } from './container';
 import { AnyExpressionFunctionDefinition, ExpressionFunction } from '../expression_functions';
-import { Execution, ExecutionParams, ExecutionResult } from '../execution/execution';
+import {
+  Execution,
+  ExecutionParams,
+  ExecutionResult,
+  FunctionCacheItem,
+} from '../execution/execution';
 import { IRegistry } from '../types';
 import { ExpressionType } from '../expression_types/expression_type';
 import { AnyExpressionTypeDefinition } from '../expression_types/types';
@@ -109,10 +114,17 @@ export class Executor<Context extends Record<string, unknown> = Record<string, u
    */
   public readonly types: TypesRegistry;
 
-  constructor(private readonly logger?: Logger, state?: ExecutorState<Context>) {
+  private functionCache: Map<string, FunctionCacheItem>;
+
+  constructor(
+    private readonly logger?: Logger,
+    state?: ExecutorState<Context>,
+    functionCache: Map<string, FunctionCacheItem> = new Map()
+  ) {
     this.functions = new FunctionsRegistry(this as Executor);
     this.types = new TypesRegistry(this as Executor);
     this.container = createExecutorContainer<Context>(state);
+    this.functionCache = functionCache;
   }
 
   public get state(): ExecutorState<Context> {
@@ -189,12 +201,17 @@ export class Executor<Context extends Record<string, unknown> = Record<string, u
     const executionParams = {
       params,
       executor: this,
+      functionCache: this.functionCache,
     } as ExecutionParams;
 
     if (typeof ast === 'string') executionParams.expression = ast;
     else executionParams.ast = ast;
 
-    const execution = new Execution<Input, Output>(executionParams, this.logger);
+    const execution = new Execution<Input, Output>(
+      executionParams,
+      this.logger,
+      this.functionCache
+    );
 
     return execution;
   }