Dynamic feature flag with callback

Signed-off-by: Siddhant Deshmukh <[email protected]>
deshsidd · Oct 12, 2023 · e4eb9b4 · e4eb9b4
1 parent 26a7498
commit e4eb9b4
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java
@@ -44,7 +44,7 @@ public void categorize(SearchSourceBuilder source) {
     }
 
     private void incrementQuerySortCounters(List<SortBuilder<?>> sorts) {
-        if (sorts.size() > 0) {
+        if (sorts != null && sorts.size() > 0) {
             searchQueryCounters.sortCounter.add(1);
         }
     }

diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -60,7 +60,6 @@
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Setting.Property;
 import org.opensearch.common.unit.TimeValue;
-import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.util.concurrent.AtomicArray;
 import org.opensearch.common.util.concurrent.CountDown;
 import org.opensearch.core.action.ActionListener;
@@ -138,6 +137,13 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
         Property.NodeScope
     );
 
+    public static final Setting<Boolean> SEARCH_QUERY_CATEGORIZATION_ENABLED_SETTING = Setting.boolSetting(
+        "search.query.categorization.enabled",
+        false,
+        Setting.Property.NodeScope,
+        Setting.Property.Dynamic
+    );
+
     // cluster level setting for timeout based search cancellation. If search request level parameter is present then that will take
     // precedence over the cluster setting value
     public static final String SEARCH_CANCEL_AFTER_TIME_INTERVAL_SETTING_KEY = "search.cancel_after_time_interval";
@@ -170,6 +176,8 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
 
     private volatile boolean isRequestStatsEnabled;
 
+    private volatile boolean isSearchQueryCategorizationEnabled;
+
     private final SearchRequestStats searchRequestStats;
 
     private final MetricsRegistry metricsRegistry;
@@ -210,13 +218,20 @@ public TransportSearchAction(
         clusterService.getClusterSettings().addSettingsUpdateConsumer(SEARCH_REQUEST_STATS_ENABLED, this::setIsRequestStatsEnabled);
         this.searchRequestStats = searchRequestStats;
         this.metricsRegistry = metricsRegistry;
-        if (FeatureFlags.isEnabled(FeatureFlags.QUERY_CATEOGORIZATION)) {
+        this.isSearchQueryCategorizationEnabled = clusterService.getClusterSettings().get(SEARCH_QUERY_CATEGORIZATION_ENABLED_SETTING);
+        clusterService.getClusterSettings()
+            .addSettingsUpdateConsumer(SEARCH_QUERY_CATEGORIZATION_ENABLED_SETTING, this::setIsSearchQueryCategorizationEnabled);
+        if (isSearchQueryCategorizationEnabled) {
             this.searchQueryCategorizor = new SearchQueryCategorizor(metricsRegistry);
         } else {
             this.searchQueryCategorizor = null;
         }
     }
 
+    private void setIsSearchQueryCategorizationEnabled(boolean isSearchQueryCategorizationEnabled) {
+        this.isSearchQueryCategorizationEnabled = isSearchQueryCategorizationEnabled;
+    }
+
     private void setIsRequestStatsEnabled(boolean isRequestStatsEnabled) {
         this.isRequestStatsEnabled = isRequestStatsEnabled;
     }
@@ -445,7 +460,7 @@ private void executeRequest(
             return;
         }
 
-        if (FeatureFlags.isEnabled(FeatureFlags.QUERY_CATEOGORIZATION)) {
+        if (isSearchQueryCategorizationEnabled) {
             searchQueryCategorizor.categorize(searchRequest.source());
         }