Loki: Cache extracted labels (#75842)

* add simple cache to extracted label values in completion provider

public/app/plugins/datasource/loki/LanguageProvider.ts

@@ -18,7 +18,7 @@ import {
   extractUnwrapLabelKeysFromDataFrame,
 } from './responseUtils';
 import syntax, { FUNCTIONS, PIPE_PARSERS, PIPE_OPERATORS } from './syntax';
-import { LokiQuery, LokiQueryType } from './types';
+import { ExtractedLabelKeys, LokiQuery, LokiQueryType } from './types';
 
 const DEFAULT_KEYS = ['job', 'namespace'];
 const EMPTY_SELECTOR = '{}';
@@ -459,13 +459,7 @@ export default class LokiLanguageProvider extends LanguageProvider {
     return labelValues ?? [];
   }
 
-  async getParserAndLabelKeys(selector: string): Promise<{
-    extractedLabelKeys: string[];
-    hasJSON: boolean;
-    hasLogfmt: boolean;
-    hasPack: boolean;
-    unwrapLabelKeys: string[];
-  }> {
+  async getParserAndLabelKeys(selector: string): Promise<ExtractedLabelKeys> {
     const series = await this.datasource.getDataSamples({ expr: selector, refId: 'data-samples' });
 
     if (!series.length) {

public/app/plugins/datasource/loki/components/monaco-query-field/monaco-completion-provider/CompletionDataProvider.test.ts

@@ -110,6 +110,57 @@ describe('CompletionDataProvider', () => {
 
   test('Returns the expected parser and label keys', async () => {
     expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+    expect(languageProvider.getParserAndLabelKeys).toHaveBeenCalledTimes(1);
+  });
+
+  test('Returns the expected parser and label keys, cache duplicate query', async () => {
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+
+    expect(languageProvider.getParserAndLabelKeys).toHaveBeenCalledTimes(1);
+  });
+
+  test('Returns the expected parser and label keys, unique query is not cached', async () => {
+    //1
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+
+    //2
+    expect(await completionProvider.getParserAndLabelKeys('unique')).toEqual(parserAndLabelKeys);
+    expect(await completionProvider.getParserAndLabelKeys('unique')).toEqual(parserAndLabelKeys);
+
+    // 3
+    expect(await completionProvider.getParserAndLabelKeys('uffdah')).toEqual(parserAndLabelKeys);
+
+    // 4
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+
+    expect(languageProvider.getParserAndLabelKeys).toHaveBeenCalledTimes(4);
+  });
+
+  test('Returns the expected parser and label keys, cache size is 2', async () => {
+    //1
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+
+    //2
+    expect(await completionProvider.getParserAndLabelKeys('unique')).toEqual(parserAndLabelKeys);
+    expect(await completionProvider.getParserAndLabelKeys('unique')).toEqual(parserAndLabelKeys);
+
+    // 2
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+    expect(languageProvider.getParserAndLabelKeys).toHaveBeenCalledTimes(2);
+
+    // 3
+    expect(await completionProvider.getParserAndLabelKeys('new')).toEqual(parserAndLabelKeys);
+    expect(await completionProvider.getParserAndLabelKeys('unique')).toEqual(parserAndLabelKeys);
+    expect(languageProvider.getParserAndLabelKeys).toHaveBeenCalledTimes(3);
+
+    // 4
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+    expect(await completionProvider.getParserAndLabelKeys('')).toEqual(parserAndLabelKeys);
+    expect(languageProvider.getParserAndLabelKeys).toHaveBeenCalledTimes(4);
   });
 
   test('Returns the expected series labels', async () => {

public/app/plugins/datasource/loki/components/monaco-query-field/monaco-completion-provider/CompletionDataProvider.ts

@@ -4,7 +4,7 @@ import { HistoryItem } from '@grafana/data';
 import { escapeLabelValueInExactSelector } from 'app/plugins/datasource/prometheus/language_utils';
 
 import LanguageProvider from '../../../LanguageProvider';
-import { LokiQuery } from '../../../types';
+import { ExtractedLabelKeys, LokiQuery } from '../../../types';
 
 import { Label } from './situation';
 
@@ -16,7 +16,10 @@ export class CompletionDataProvider {
   constructor(
     private languageProvider: LanguageProvider,
     private historyRef: HistoryRef = { current: [] }
-  ) {}
+  ) {
+    this.queryToLabelKeysCache = new Map();
+  }
+  private queryToLabelKeysCache: Map<string, ExtractedLabelKeys>;
 
   private buildSelector(labels: Label[]): string {
     const allLabelTexts = labels.map(
@@ -55,8 +58,35 @@ export class CompletionDataProvider {
     return data[labelName] ?? [];
   }
 
-  async getParserAndLabelKeys(logQuery: string) {
-    return await this.languageProvider.getParserAndLabelKeys(logQuery);
+  /**
+   * Runs a Loki query to extract label keys from the result.
+   * The result is cached for the query string.
+   *
+   * Since various "situations" in the monaco code editor trigger this function, it is prone to being called multiple times for the same query
+   * Here is a lightweight and simple cache to avoid calling the backend multiple times for the same query.
+   *
+   * @param logQuery
+   */
+  async getParserAndLabelKeys(logQuery: string): Promise<ExtractedLabelKeys> {
+    const EXTRACTED_LABEL_KEYS_MAX_CACHE_SIZE = 2;
+    const cachedLabelKeys = this.queryToLabelKeysCache.has(logQuery) ? this.queryToLabelKeysCache.get(logQuery) : null;
+    if (cachedLabelKeys) {
+      // cache hit! Serve stale result from cache
+      return cachedLabelKeys;
+    } else {
+      // If cache is larger than max size, delete the first (oldest) index
+      if (this.queryToLabelKeysCache.size >= EXTRACTED_LABEL_KEYS_MAX_CACHE_SIZE) {
+        // Make room in the cache for the fresh result by deleting the "first" index
+        const keys = this.queryToLabelKeysCache.keys();
+        const firstKey = keys.next().value;
+        this.queryToLabelKeysCache.delete(firstKey);
+      }
+      // Fetch a fresh result from the backend
+      const labelKeys = await this.languageProvider.getParserAndLabelKeys(logQuery);
+      // Add the result to the cache
+      this.queryToLabelKeysCache.set(logQuery, labelKeys);
+      return labelKeys;
+    }
   }
 
   async getSeriesLabels(labels: Label[]) {

public/app/plugins/datasource/loki/types.ts

@@ -84,4 +84,12 @@ export interface ContextFilter {
   description?: string;
 }
 
+export interface ExtractedLabelKeys {
+  extractedLabelKeys: string[];
+  hasJSON: boolean;
+  hasLogfmt: boolean;
+  hasPack: boolean;
+  unwrapLabelKeys: string[];
+}
+
 export type LokiGroupedRequest = { request: DataQueryRequest<LokiQuery>; partition: TimeRange[] };