Always collect available metrics in top queries service (#205)

Signed-off-by: Chenyang Ji <[email protected]>
opensearch-project · Jan 24, 2025 · a52586a · a52586a
1 parent bd6debd
commit a52586a
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 29 deletions.
diff --git a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java
@@ -220,10 +220,6 @@ public void onRequestFailure(final SearchPhaseContext context, final SearchReque
         constructSearchQueryRecord(context, searchRequestContext);
     }
 
-    private boolean shouldCollect(MetricType metricType) {
-        return queryInsightsService.isSearchQueryMetricsFeatureEnabled() || queryInsightsService.isCollectionEnabled(metricType);
-    }
-
     private void constructSearchQueryRecord(final SearchPhaseContext context, final SearchRequestContext searchRequestContext) {
         SearchTask searchTask = context.getTask();
         List<TaskResourceInfo> tasksResourceUsages = searchRequestContext.getPhaseResourceUsage();
@@ -240,28 +236,22 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final
         final SearchRequest request = context.getRequest();
         try {
             Map<MetricType, Measurement> measurements = new HashMap<>();
-            if (shouldCollect(MetricType.LATENCY)) {
-                measurements.put(
-                    MetricType.LATENCY,
-                    new Measurement(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - searchRequestContext.getAbsoluteStartNanos()))
-                );
-            }
-            if (shouldCollect(MetricType.CPU)) {
-                measurements.put(
-                    MetricType.CPU,
-                    new Measurement(
-                        tasksResourceUsages.stream().map(a -> a.getTaskResourceUsage().getCpuTimeInNanos()).mapToLong(Long::longValue).sum()
-                    )
-                );
-            }
-            if (shouldCollect(MetricType.MEMORY)) {
-                measurements.put(
-                    MetricType.MEMORY,
-                    new Measurement(
-                        tasksResourceUsages.stream().map(a -> a.getTaskResourceUsage().getMemoryInBytes()).mapToLong(Long::longValue).sum()
-                    )
-                );
-            }
+            measurements.put(
+                MetricType.LATENCY,
+                new Measurement(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - searchRequestContext.getAbsoluteStartNanos()))
+            );
+            measurements.put(
+                MetricType.CPU,
+                new Measurement(
+                    tasksResourceUsages.stream().map(a -> a.getTaskResourceUsage().getCpuTimeInNanos()).mapToLong(Long::longValue).sum()
+                )
+            );
+            measurements.put(
+                MetricType.MEMORY,
+                new Measurement(
+                    tasksResourceUsages.stream().map(a -> a.getTaskResourceUsage().getMemoryInBytes()).mapToLong(Long::longValue).sum()
+                )
+            );
 
             Map<Attribute, Object> attributes = new HashMap<>();
             attributes.put(Attribute.SEARCH_TYPE, request.searchType().toString().toLowerCase(Locale.ROOT));

diff --git a/...ain/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java b/...ain/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java
@@ -140,9 +140,15 @@ public void incrementSortCounter(double value, Tags tags, Map<MetricType, Measur
     }
 
     private void incrementAllHistograms(Tags tags, Map<MetricType, Measurement> measurements) {
-        queryTypeLatencyHistogram.record(measurements.get(MetricType.LATENCY).getMeasurement().doubleValue(), tags);
-        queryTypeCpuHistogram.record(measurements.get(MetricType.CPU).getMeasurement().doubleValue(), tags);
-        queryTypeMemoryHistogram.record(measurements.get(MetricType.MEMORY).getMeasurement().doubleValue(), tags);
+        if (measurements.containsKey(MetricType.LATENCY)) {
+            queryTypeLatencyHistogram.record(measurements.get(MetricType.LATENCY).getMeasurement().doubleValue(), tags);
+        }
+        if (measurements.containsKey(MetricType.CPU)) {
+            queryTypeCpuHistogram.record(measurements.get(MetricType.CPU).getMeasurement().doubleValue(), tags);
+        }
+        if (measurements.containsKey(MetricType.MEMORY)) {
+            queryTypeMemoryHistogram.record(measurements.get(MetricType.MEMORY).getMeasurement().doubleValue(), tags);
+        }
     }
 
     /**

diff --git a/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java b/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java
@@ -322,6 +322,9 @@ public String getId() {
      * @return the measurement object, or null if not found
      */
     public Number getMeasurement(final MetricType name) {
+        if (!measurements.containsKey(name)) {
+            return null;
+        }
         return measurements.get(name).getMeasurement();
     }
 
@@ -337,6 +340,10 @@ public Number getMeasurement(final MetricType name) {
      * @param numberToAdd The measurement number we want to add to the current measurement.
      */
     public void addMeasurement(final MetricType metricType, Number numberToAdd) {
+        if (!measurements.containsKey(metricType)) {
+            measurements.put(metricType, new Measurement(numberToAdd));
+            return;
+        }
         measurements.get(metricType).addMeasurement(numberToAdd);
     }
 
@@ -346,6 +353,9 @@ public void addMeasurement(final MetricType metricType, Number numberToAdd) {
      * @param aggregationType Aggregation type to set
      */
     public void setMeasurementAggregation(final MetricType name, AggregationType aggregationType) {
+        if (!measurements.containsKey(name)) {
+            return;
+        }
         measurements.get(name).setAggregationType(aggregationType);
     }