Fix starting and stopping query insights service

Signed-off-by: Siddhant Deshmukh <[email protected]>
opensearch-project · Jul 17, 2024 · d2ba20b · d2ba20b
1 parent e6293d3
commit d2ba20b
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 23 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
@@ -91,28 +91,27 @@ public QueryInsightsListener(final ClusterService clusterService, final QueryIns
      * and query insights services.
      *
      * @param metricType {@link MetricType}
-     * @param enabled    boolean
+     * @param isCurrentMetricEnabled boolean
      */
-    public void setEnableTopQueries(final MetricType metricType, final boolean enabled) {
-        boolean isInsightsServiceDisabled = !queryInsightsService.isEnabled();
+    public void setEnableTopQueries(final MetricType metricType, final boolean isCurrentMetricEnabled) {
+        boolean isTopNFeatureDisabled = !queryInsightsService.isTopNFeatureEnabled();
+        this.queryInsightsService.enableCollection(metricType, isCurrentMetricEnabled);
 
-        if (!enabled) {
+        if (!isCurrentMetricEnabled) {
             // disable QueryInsightsListener only if all metrics collections are disabled now
             // and search query metrics is disabled.
-            if (isInsightsServiceDisabled) {
+            if (isTopNFeatureDisabled) {
                 super.setEnabled(false);
-                this.queryInsightsService.stop();
+                queryInsightsService.checkAndStopQueryInsights();
             }
         } else {
             super.setEnabled(true);
             // restart QueryInsightsListener only if none of metrics collections is enabled before and
             // search query metrics is disabled before.
-            if (isInsightsServiceDisabled) {
-                this.queryInsightsService.stop();
-                this.queryInsightsService.start();
+            if (isTopNFeatureDisabled) {
+                queryInsightsService.checkAndRestartQueryInsights();
             }
         }
-
     }
 
     @Override

diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java b/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java
@@ -200,15 +200,51 @@ public boolean isCollectionEnabled(final MetricType metricType) {
      *
      * @return if query insights service is enabled
      */
-    public boolean isEnabled() {
+    public boolean isAnyFeatureEnabled() {
+        return isTopNFeatureEnabled() || isSearchQueryMetricsFeatureEnabled();
+    }
+
+    /**
+     * Check if top N enabled for any metric type
+     *
+     * @return if top N feature is enabled
+     */
+    public boolean isTopNFeatureEnabled() {
         for (MetricType t : MetricType.allMetricTypes()) {
             if (isCollectionEnabled(t)) {
                 return true;
             }
         }
+        return false;
+    }
+
+    /**
+     * Is search query metrics feature enabled.
+     * @return boolean flag
+     */
+    public boolean isSearchQueryMetricsFeatureEnabled() {
         return this.searchQueryMetricsEnabled;
     }
 
+    /**
+     * Stops query insights service if no features enabled
+     */
+    public void checkAndStopQueryInsights() {
+        if (!isAnyFeatureEnabled()) {
+            this.stop();
+        }
+    }
+
+    /**
+     * Restarts query insights service if any feature enabled
+     */
+    public void checkAndRestartQueryInsights() {
+        if (isAnyFeatureEnabled()) {
+            this.stop();
+            this.start();
+        }
+    }
+
     /**
      * Validate the window size config for a metricType
      *
@@ -274,15 +310,17 @@ public void setExporter(final MetricType type, final Settings settings) {
      * @param searchQueryMetricsEnabled boolean flag
      */
     public void setSearchQueryMetricsEnabled(boolean searchQueryMetricsEnabled) {
+        boolean oldsearchQueryMetricsEnabled = isSearchQueryMetricsFeatureEnabled();
         this.searchQueryMetricsEnabled = searchQueryMetricsEnabled;
-    }
-
-    /**
-     * Is search query metrics feature enabled.
-     * @return boolean flag
-     */
-    public boolean isSearchQueryMetricsEnabled() {
-        return this.searchQueryMetricsEnabled;
+        if (searchQueryMetricsEnabled) {
+            if (!oldsearchQueryMetricsEnabled) {
+                checkAndRestartQueryInsights();
+            }
+        } else {
+            if (oldsearchQueryMetricsEnabled) {
+                checkAndStopQueryInsights();
+            }
+        }
     }
 
     /**
@@ -307,7 +345,7 @@ public void validateExporterConfig(final MetricType type, final Settings setting
 
     @Override
     protected void doStart() {
-        if (isEnabled()) {
+        if (isAnyFeatureEnabled()) {
             scheduledFuture = threadPool.scheduleWithFixedDelay(
                 this::drainRecords,
                 QueryInsightsSettings.QUERY_RECORD_QUEUE_DRAIN_INTERVAL,

diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java
@@ -66,21 +66,21 @@ public void testClose() {
 
     public void testSearchQueryMetricsEnabled() {
         // Initially, searchQueryMetricsEnabled should be false and searchQueryCategorizer should be null
-        assertFalse(queryInsightsService.isSearchQueryMetricsEnabled());
+        assertFalse(queryInsightsService.isSearchQueryMetricsFeatureEnabled());
         assertNotNull(queryInsightsService.getSearchQueryCategorizer());
 
         // Enable search query metrics
         queryInsightsService.setSearchQueryMetricsEnabled(true);
 
         // Assert that searchQueryMetricsEnabled is true and searchQueryCategorizer is initialized
-        assertTrue(queryInsightsService.isSearchQueryMetricsEnabled());
+        assertTrue(queryInsightsService.isSearchQueryMetricsFeatureEnabled());
         assertNotNull(queryInsightsService.getSearchQueryCategorizer());
 
         // Disable search query metrics
         queryInsightsService.setSearchQueryMetricsEnabled(false);
 
         // Assert that searchQueryMetricsEnabled is false and searchQueryCategorizer is not null
-        assertFalse(queryInsightsService.isSearchQueryMetricsEnabled());
+        assertFalse(queryInsightsService.isSearchQueryMetricsFeatureEnabled());
         assertNotNull(queryInsightsService.getSearchQueryCategorizer());
 
     }