Prometheus: Improve autocomplete performance and remove disabling of dynamic label lookup (#30199)

* processLabels: Use Sets instead of Array

* Add and update comment

* Limit autocomplete items to 10000

* Remove lookup treshold, limit display of items

* Update tests

* Add test

* Update public/app/plugins/datasource/prometheus/language_provider.ts

public/app/plugins/datasource/prometheus/components/PromQueryField.test.tsx

@@ -1,7 +1,7 @@
 // @ts-ignore
 import RCCascader from 'rc-cascader';
 import React from 'react';
-import PromQlLanguageProvider, { DEFAULT_LOOKUP_METRICS_THRESHOLD } from '../language_provider';
+import PromQlLanguageProvider from '../language_provider';
 import PromQueryField, { groupMetricsByPrefix, RECORDING_RULES_GROUP } from './PromQueryField';
 import { DataSourceInstanceSettings, dateTime } from '@grafana/data';
 import { PromOptions } from '../types';
@@ -254,7 +254,6 @@ function makeLanguageProvider(options: { metrics: string[][] }) {
     metrics: [],
     metricsMetadata: {},
     lookupsDisabled: false,
-    lookupMetricsThreshold: DEFAULT_LOOKUP_METRICS_THRESHOLD,
     start() {
       this.metrics = metricsStack.shift();
       return Promise.resolve([]);

public/app/plugins/datasource/prometheus/components/PromQueryField.tsx

@@ -209,9 +209,9 @@ class PromQueryField extends React.PureComponent<PromQueryFieldProps, PromQueryF
     let hint = hints.length > 0 ? hints[0] : null;
 
     // Hint for big disabled lookups
-    if (!hint && !datasource.lookupsDisabled && datasource.languageProvider.lookupsDisabled) {
+    if (!hint && datasource.lookupsDisabled) {
       hint = {
-        label: `Dynamic label lookup is disabled for datasources with more than ${datasource.languageProvider.lookupMetricsThreshold} metrics.`,
+        label: `Labels and metrics lookup was disabled in data source settings.`,
         type: 'INFO',
       };
     }

public/app/plugins/datasource/prometheus/language_provider.test.ts

@@ -226,7 +226,6 @@ describe('Language completion provider', () => {
   describe('label suggestions', () => {
     it('returns default label suggestions on label context and no metric', async () => {
       const instance = new LanguageProvider(datasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('{}');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(1).value;
@@ -246,7 +245,6 @@ describe('Language completion provider', () => {
         getTimeRange: () => ({ start: 0, end: 1 }),
       } as any) as PrometheusDatasource;
       const instance = new LanguageProvider(datasources);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('metric{}');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(7).value;
@@ -278,7 +276,6 @@ describe('Language completion provider', () => {
         getTimeRange: () => ({ start: 0, end: 1 }),
       } as any) as PrometheusDatasource;
       const instance = new LanguageProvider(datasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('{job1="foo",job2!="foo",job3=~"foo",__name__="metric",}');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(54).value;
@@ -299,7 +296,6 @@ describe('Language completion provider', () => {
           return { data: { data: ['value1', 'value2'] } };
         },
       } as any) as PrometheusDatasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('{job!=}');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(6).value;
@@ -321,7 +317,6 @@ describe('Language completion provider', () => {
 
     it('returns a refresher on label context and unavailable metric', async () => {
       const instance = new LanguageProvider(datasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('metric{}');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(7).value;
@@ -340,7 +335,6 @@ describe('Language completion provider', () => {
         ...datasource,
         metadataRequest: () => simpleMetricLabelsResponse,
       } as any) as PrometheusDatasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('metric{bar=ba}');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(13).value;
@@ -360,7 +354,6 @@ describe('Language completion provider', () => {
         ...datasource,
         metadataRequest: () => simpleMetricLabelsResponse,
       } as any) as PrometheusDatasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('sum(metric{foo="xx"}) by ()');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(26).value;
@@ -379,7 +372,6 @@ describe('Language completion provider', () => {
         ...datasource,
         metadataRequest: () => simpleMetricLabelsResponse,
       } as any) as PrometheusDatasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('sum(metric) by ()');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(16).value;
@@ -398,7 +390,6 @@ describe('Language completion provider', () => {
         ...datasource,
         metadataRequest: () => simpleMetricLabelsResponse,
       } as any) as PrometheusDatasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('sum(\nmetric\n)\nby ()');
       const aggregationTextBlock = value.document.getBlocks().get(3);
       const ed = new SlateEditor({ value });
@@ -424,7 +415,6 @@ describe('Language completion provider', () => {
         ...datasource,
         metadataRequest: () => simpleMetricLabelsResponse,
       } as any) as PrometheusDatasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('sum(rate(metric[1h])) by ()');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(26).value;
@@ -448,7 +438,6 @@ describe('Language completion provider', () => {
         ...datasource,
         metadataRequest: () => simpleMetricLabelsResponse,
       } as any) as PrometheusDatasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('sum(rate(metric{label1="value"}[1h])) by ()');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(42).value;
@@ -469,7 +458,6 @@ describe('Language completion provider', () => {
 
     it('returns no suggestions inside an unclear aggregation context using alternate syntax', async () => {
       const instance = new LanguageProvider(datasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('sum by ()');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(8).value;
@@ -488,7 +476,6 @@ describe('Language completion provider', () => {
         ...datasource,
         metadataRequest: () => simpleMetricLabelsResponse,
       } as any) as PrometheusDatasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('sum by () (metric)');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(8).value;
@@ -514,7 +501,6 @@ describe('Language completion provider', () => {
       } as any) as PrometheusDatasource;
 
       const instance = new LanguageProvider(datasource);
-      instance.lookupsDisabled = false;
       const value = Plain.deserialize('{}');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(1).value;
@@ -533,39 +519,14 @@ describe('Language completion provider', () => {
       expect((datasource.metadataRequest as Mock).mock.calls.length).toBe(2);
     });
   });
-
-  describe('dynamic lookup protection for big installations', () => {
-    it('dynamic lookup is enabled if number of metrics is reasonably low', async () => {
-      const datasource: PrometheusDatasource = ({
-        metadataRequest: () => ({ data: { data: ['foo'] as string[] } }),
-        getTimeRange: () => ({ start: 0, end: 1 }),
-      } as any) as PrometheusDatasource;
-
-      const instance = new LanguageProvider(datasource, { lookupMetricsThreshold: 1 });
-      expect(instance.lookupsDisabled).toBeTruthy();
-      await instance.start();
-      expect(instance.lookupsDisabled).toBeFalsy();
-    });
-
-    it('dynamic lookup is disabled if number of metrics is higher than threshold', async () => {
-      const datasource: PrometheusDatasource = ({
-        metadataRequest: () => ({ data: { data: ['foo', 'bar'] as string[] } }),
-        getTimeRange: () => ({ start: 0, end: 1 }),
-      } as any) as PrometheusDatasource;
-
-      const instance = new LanguageProvider(datasource, { lookupMetricsThreshold: 1 });
-      expect(instance.lookupsDisabled).toBeTruthy();
-      await instance.start();
-      expect(instance.lookupsDisabled).toBeTruthy();
-    });
-
-    it('does not issue label-based metadata requests when lookup is disabled', async () => {
+  describe('disabled metrics lookup', () => {
+    it('does not issue any metadata requests when lookup is disabled', async () => {
       const datasource: PrometheusDatasource = ({
         metadataRequest: jest.fn(() => ({ data: { data: ['foo', 'bar'] as string[] } })),
         getTimeRange: jest.fn(() => ({ start: 0, end: 1 })),
+        lookupsDisabled: true,
       } as any) as PrometheusDatasource;
-
-      const instance = new LanguageProvider(datasource, { lookupMetricsThreshold: 1 });
+      const instance = new LanguageProvider(datasource);
       const value = Plain.deserialize('{}');
       const ed = new SlateEditor({ value });
       const valueWithSelection = ed.moveForward(1).value;
@@ -575,15 +536,24 @@ describe('Language completion provider', () => {
         wrapperClasses: ['context-labels'],
         value: valueWithSelection,
       };
-      expect(instance.lookupsDisabled).toBeTruthy();
+
       expect((datasource.metadataRequest as Mock).mock.calls.length).toBe(0);
       await instance.start();
-      expect(instance.lookupsDisabled).toBeTruthy();
-      // Capture request count to metadata
-      const callCount = (datasource.metadataRequest as Mock).mock.calls.length;
-      expect((datasource.metadataRequest as Mock).mock.calls.length).toBeGreaterThan(0);
+      expect((datasource.metadataRequest as Mock).mock.calls.length).toBe(0);
       await instance.provideCompletionItems(args);
-      expect((datasource.metadataRequest as Mock).mock.calls.length).toBe(callCount);
+      expect((datasource.metadataRequest as Mock).mock.calls.length).toBe(0);
+    });
+    it('issues metadata requests when lookup is not disabled', async () => {
+      const datasource: PrometheusDatasource = ({
+        metadataRequest: jest.fn(() => ({ data: { data: ['foo', 'bar'] as string[] } })),
+        getTimeRange: jest.fn(() => ({ start: 0, end: 1 })),
+        lookupsDisabled: false,
+      } as any) as PrometheusDatasource;
+      const instance = new LanguageProvider(datasource);
+
+      expect((datasource.metadataRequest as Mock).mock.calls.length).toBe(0);
+      await instance.start();
+      expect((datasource.metadataRequest as Mock).mock.calls.length).toBeGreaterThan(0);
     });
   });
 });

public/app/plugins/datasource/prometheus/language_provider.ts

@@ -11,6 +11,8 @@ import {
   processHistogramLabels,
   processLabels,
   roundSecToMin,
+  addLimitInfo,
+  limitSuggestions,
 } from './language_utils';
 import PromqlSyntax, { FUNCTIONS, RATE_RANGES } from './promql';
 
@@ -21,7 +23,8 @@ const DEFAULT_KEYS = ['job', 'instance'];
 const EMPTY_SELECTOR = '{}';
 const HISTORY_ITEM_COUNT = 5;
 const HISTORY_COUNT_CUTOFF = 1000 * 60 * 60 * 24; // 24h
-export const DEFAULT_LOOKUP_METRICS_THRESHOLD = 10000; // number of metrics defining an installation that's too big
+// Max number of items (metrics, labels, values) that we display as suggestions. Prevents from running out of memory.
+export const SUGGESTIONS_LIMIT = 10000;
 
 const wrapLabel = (label: string): CompletionItem => ({ label });
 
@@ -66,8 +69,6 @@ export default class PromQlLanguageProvider extends LanguageProvider {
   metricsMetadata?: PromMetricsMetadata;
   startTask: Promise<any>;
   datasource: PrometheusDatasource;
-  lookupMetricsThreshold: number;
-  lookupsDisabled: boolean; // Dynamically set to true for big/slow instances
 
   /**
    *  Cache for labels of series. This is bit simplistic in the sense that it just counts responses each as a 1 and does
@@ -83,9 +84,6 @@ export default class PromQlLanguageProvider extends LanguageProvider {
     this.histogramMetrics = [];
     this.timeRange = { start: 0, end: 0 };
     this.metrics = [];
-    // Disable lookups until we know the instance is small enough
-    this.lookupMetricsThreshold = DEFAULT_LOOKUP_METRICS_THRESHOLD;
-    this.lookupsDisabled = true;
 
     Object.assign(this, initialValues);
   }
@@ -128,7 +126,6 @@ export default class PromQlLanguageProvider extends LanguageProvider {
     const url = `/api/v1/label/__name__/values?${params.toString()}`;
 
     this.metrics = await this.request(url, []);
-    this.lookupsDisabled = this.metrics.length > this.lookupMetricsThreshold;
     this.metricsMetadata = fixSummariesMetadata(await this.request('/api/v1/metadata', {}));
     this.processHistogramMetrics(this.metrics);
 
@@ -241,9 +238,10 @@ export default class PromQlLanguageProvider extends LanguageProvider {
     });
 
     if (metrics && metrics.length) {
+      const limitInfo = addLimitInfo(metrics);
       suggestions.push({
-        label: 'Metrics',
-        items: metrics.map(m => addMetricsMetadata(m, metricsMetadata)),
+        label: `Metrics${limitInfo}`,
+        items: limitSuggestions(metrics).map(m => addMetricsMetadata(m, metricsMetadata)),
       });
     }
 
@@ -314,7 +312,11 @@ export default class PromQlLanguageProvider extends LanguageProvider {
 
     const labelValues = await this.getLabelValues(selector);
     if (labelValues) {
-      suggestions.push({ label: 'Labels', items: Object.keys(labelValues).map(wrapLabel) });
+      const limitInfo = addLimitInfo(labelValues[0]);
+      suggestions.push({
+        label: `Labels${limitInfo}`,
+        items: Object.keys(labelValues).map(wrapLabel),
+      });
     }
     return result;
   };
@@ -376,8 +378,9 @@ export default class PromQlLanguageProvider extends LanguageProvider {
       // Label values
       if (labelKey && labelValues[labelKey]) {
         context = 'context-label-values';
+        const limitInfo = addLimitInfo(labelValues[labelKey]);
         suggestions.push({
-          label: `Label values for "${labelKey}"`,
+          label: `Label values for "${labelKey}"${limitInfo}`,
           items: labelValues[labelKey].map(wrapLabel),
         });
       }
@@ -390,7 +393,8 @@ export default class PromQlLanguageProvider extends LanguageProvider {
         if (possibleKeys.length) {
           context = 'context-labels';
           const newItems = possibleKeys.map(key => ({ label: key }));
-          const newSuggestion: CompletionItemGroup = { label: `Labels`, items: newItems };
+          const limitInfo = addLimitInfo(newItems);
+          const newSuggestion: CompletionItemGroup = { label: `Labels${limitInfo}`, items: newItems };
           suggestions.push(newSuggestion);
         }
       }
@@ -400,7 +404,7 @@ export default class PromQlLanguageProvider extends LanguageProvider {
   };
 
   async getLabelValues(selector: string, withName?: boolean) {
-    if (this.lookupsDisabled) {
+    if (this.datasource.lookupsDisabled) {
       return undefined;
     }
     try {

public/app/plugins/datasource/prometheus/language_utils.ts

@@ -1,5 +1,6 @@
 import { PromMetricsMetadata } from './types';
 import { addLabelToQuery } from './add_label_to_query';
+import { SUGGESTIONS_LIMIT } from './language_provider';
 
 export const RATE_RANGES = ['1m', '5m', '10m', '30m', '1h'];
 
@@ -19,26 +20,35 @@ export const processHistogramLabels = (labels: string[]) => {
 };
 
 export function processLabels(labels: Array<{ [key: string]: string }>, withName = false) {
-  const values: { [key: string]: string[] } = {};
-  labels.forEach(l => {
-    const { __name__, ...rest } = l;
+  // For processing we are going to use sets as they have significantly better performance than arrays
+  // After we process labels, we will convert sets to arrays and return object with label values in arrays
+  const valueSet: { [key: string]: Set<string> } = {};
+  labels.forEach(label => {
+    const { __name__, ...rest } = label;
     if (withName) {
-      values['__name__'] = values['__name__'] || [];
-      if (!values['__name__'].includes(__name__)) {
-        values['__name__'].push(__name__);
+      valueSet['__name__'] = valueSet['__name__'] || new Set();
+      if (!valueSet['__name__'].has(__name__)) {
+        valueSet['__name__'].add(__name__);
       }
     }
 
     Object.keys(rest).forEach(key => {
-      if (!values[key]) {
-        values[key] = [];
+      if (!valueSet[key]) {
+        valueSet[key] = new Set();
       }
-      if (!values[key].includes(rest[key])) {
-        values[key].push(rest[key]);
+      if (!valueSet[key].has(rest[key])) {
+        valueSet[key].add(rest[key]);
       }
     });
   });
-  return { values, keys: Object.keys(values) };
+
+  // valueArray that we are going to return in the object
+  const valueArray: { [key: string]: string[] } = {};
+  limitSuggestions(Object.keys(valueSet)).forEach(key => {
+    valueArray[key] = limitSuggestions(Array.from(valueSet[key]));
+  });
+
+  return { values: valueArray, keys: Object.keys(valueArray) };
 }
 
 // const cleanSelectorRegexp = /\{(\w+="[^"\n]*?")(,\w+="[^"\n]*?")*\}/;
@@ -193,3 +203,11 @@ export function roundMsToMin(milliseconds: number): number {
 export function roundSecToMin(seconds: number): number {
   return Math.floor(seconds / 60);
 }
+
+export function limitSuggestions(items: string[]) {
+  return items.slice(0, SUGGESTIONS_LIMIT);
+}
+
+export function addLimitInfo(items: any[] | undefined): string {
+  return items && items.length >= SUGGESTIONS_LIMIT ? `, limited to the first ${SUGGESTIONS_LIMIT} received items` : '';
+}