From 7014ac9b95c0db46606b6b3ca68d37b55d957c5a Mon Sep 17 00:00:00 2001
From: Chris Cowan <chris@elastic.co>
Date: Wed, 22 Dec 2021 09:59:08 -0700
Subject: [PATCH] Adding performance optimizations

---
 .../metric_threshold/lib/evaluate_rule.ts     | 61 +++++++++----------
 .../metric_threshold_executor.ts              |  1 -
 2 files changed, 30 insertions(+), 32 deletions(-)

diff --git a/x-pack/plugins/infra/server/lib/alerting/metric_threshold/lib/evaluate_rule.ts b/x-pack/plugins/infra/server/lib/alerting/metric_threshold/lib/evaluate_rule.ts
index ec45af644347c..970c684b98d37 100644
--- a/x-pack/plugins/infra/server/lib/alerting/metric_threshold/lib/evaluate_rule.ts
+++ b/x-pack/plugins/infra/server/lib/alerting/metric_threshold/lib/evaluate_rule.ts
@@ -7,7 +7,7 @@
 
 import moment from 'moment';
 import { ElasticsearchClient } from 'kibana/server';
-import { mapValues, first, last, isNaN, isNumber, isObject, has } from 'lodash';
+import { difference, mapValues, first, last, isNaN, isNumber, isObject, has } from 'lodash';
 import {
   isTooManyBucketsPreviewException,
   TOO_MANY_BUCKETS_PREVIEW_EXCEPTION,
@@ -22,6 +22,7 @@ import { UNGROUPED_FACTORY_KEY } from '../../common/utils';
 import { MetricExpressionParams, Comparator, Aggregators } from '../types';
 import { getElasticsearchMetricQuery } from './metric_query';
 import { createTimerange } from './create_timerange';
+import { doWork } from './do_work';
 
 interface AggregationWithoutIntervals {
   aggregatedValue: { value: number; values?: Array<{ key: number; value: number }> };
@@ -98,27 +99,22 @@ export const evaluateRule = <Params extends EvaluatedRuleParams = EvaluatedRuleP
       // If any previous groups are no longer being reported, backfill them with null values
       const currentGroups = Object.keys(currentValues);
 
-      const missingGroups = prevGroups.filter((g) => !currentGroups.includes(g));
+      const missingGroups = difference(prevGroups, currentGroups);
+
       if (currentGroups.length === 0 && missingGroups.length === 0) {
         missingGroups.push(UNGROUPED_FACTORY_KEY);
       }
       const backfillTimestamp =
         last(last(Object.values(currentValues)))?.key ?? new Date().toISOString();
-      const backfilledPrevGroups: Record<
-        string,
-        Array<{ key: string; value: number }>
-      > = missingGroups.reduce(
-        (result, group) => ({
-          ...result,
-          [group]: [
-            {
-              key: backfillTimestamp,
-              value: criterion.aggType === Aggregators.COUNT ? 0 : null,
-            },
-          ],
-        }),
-        {}
-      );
+      const backfilledPrevGroups: Record<string, Array<{ key: string; value: number | null }>> = {};
+      for (const group of missingGroups) {
+        backfilledPrevGroups[group] = [
+          {
+            key: backfillTimestamp,
+            value: criterion.aggType === Aggregators.COUNT ? 0 : null,
+          },
+        ];
+      }
       const currentValuesWithBackfilledPrevGroups = {
         ...currentValues,
         ...backfilledPrevGroups,
@@ -207,20 +203,23 @@ const getMetric: (
         bucketSelector,
         afterKeyHandler
       )) as Array<Aggregation & { key: Record<string, string>; doc_count: number }>;
-      const groupedResults = compositeBuckets.reduce(
-        (result, bucket) => ({
-          ...result,
-          [Object.values(bucket.key)
-            .map((value) => value)
-            .join(', ')]: getValuesFromAggregations(
-            bucket,
-            aggType,
-            dropPartialBucketsOptions,
-            calculatedTimerange,
-            bucket.doc_count
-          ),
-        }),
-        {}
+      const groupedResults: Record<string, any> = {};
+      await Promise.all(
+        compositeBuckets.map((bucket) => {
+          return doWork(() => {
+            const key = Object.values(bucket.key)
+              .map((value) => value)
+              .join(', ');
+            const value = getValuesFromAggregations(
+              bucket,
+              aggType,
+              dropPartialBucketsOptions,
+              calculatedTimerange,
+              bucket.doc_count
+            );
+            groupedResults[key] = value;
+          });
+        })
       );
       return groupedResults;
     }
diff --git a/x-pack/plugins/infra/server/lib/alerting/metric_threshold/metric_threshold_executor.ts b/x-pack/plugins/infra/server/lib/alerting/metric_threshold/metric_threshold_executor.ts
index c2bd9c5ca5de7..fbd12fa39ce87 100644
--- a/x-pack/plugins/infra/server/lib/alerting/metric_threshold/metric_threshold_executor.ts
+++ b/x-pack/plugins/infra/server/lib/alerting/metric_threshold/metric_threshold_executor.ts
@@ -250,7 +250,6 @@ export const createMetricThresholdExecutor = (libs: InfraBackendLibs) =>
         });
       }
     }
-
     return { groups, groupBy: params.groupBy, filterQuery: params.filterQuery };
   });