Refactor FeatureFlags

CHANGELOG.md

@@ -31,6 +31,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ### Fixed
 - Fix bytes parameter on `_cat/recovery` ([#17598](https://github.com/opensearch-project/OpenSearch/pull/17598))
+- Fix slow performance of FeatureFlag checks ([#17611](https://github.com/opensearch-project/OpenSearch/pull/17611))
 
 ### Security
 

modules/mapper-extras/src/test/java/org/opensearch/index/mapper/ScaledFloatFieldMapperTests.java

@@ -36,15 +36,12 @@
 import org.apache.lucene.index.IndexableField;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.settings.Settings;
-import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.compositeindex.datacube.startree.StarTreeIndexSettings;
 import org.opensearch.plugins.Plugin;
-import org.junit.AfterClass;
-import org.junit.BeforeClass;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -98,16 +95,7 @@ public void testExistsQueryDocValuesDisabled() throws IOException {
         assertParseMinimalWarnings();
     }
 
-    @BeforeClass
-    public static void createMapper() {
-        FeatureFlags.initializeFeatureFlags(Settings.builder().put(STAR_TREE_INDEX, "true").build());
-    }
-
-    @AfterClass
-    public static void clearMapper() {
-        FeatureFlags.initializeFeatureFlags(Settings.EMPTY);
-    }
-
+    @LockFeatureFlag(STAR_TREE_INDEX)
     public void testScaledFloatWithStarTree() throws Exception {
 
         double scalingFactorField1 = randomDouble() * 100;

plugins/arrow-flight-rpc/src/internalClusterTest/java/org/opensearch/arrow/flight/ArrowFlightServerIT.java

@@ -14,26 +14,20 @@
 import org.opensearch.arrow.flight.bootstrap.FlightService;
 import org.opensearch.arrow.flight.bootstrap.FlightStreamPlugin;
 import org.opensearch.cluster.node.DiscoveryNode;
-import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.plugins.Plugin;
-import org.opensearch.test.FeatureFlagSetter;
 import org.opensearch.test.OpenSearchIntegTestCase;
-import org.junit.BeforeClass;
 
 import java.util.Collection;
 import java.util.Collections;
 import java.util.concurrent.TimeUnit;
 
+import static org.opensearch.common.util.FeatureFlags.ARROW_STREAMS;
+
 @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, numDataNodes = 5)
 public class ArrowFlightServerIT extends OpenSearchIntegTestCase {
 
     private FlightClientManager flightClientManager;
 
-    @BeforeClass
-    public static void setupFeatureFlags() {
-        FeatureFlagSetter.set(FeatureFlags.ARROW_STREAMS_SETTING.getKey());
-    }
-
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
         return Collections.singleton(FlightStreamPlugin.class);
@@ -48,6 +42,7 @@ public void setUp() throws Exception {
         flightClientManager = flightService.getFlightClientManager();
     }
 
+    @LockFeatureFlag(ARROW_STREAMS)
     public void testArrowFlightEndpoint() throws Exception {
         for (DiscoveryNode node : getClusterState().nodes()) {
             try (FlightClient flightClient = flightClientManager.getFlightClient(node.getId()).get()) {

plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/FlightStreamPluginTests.java

@@ -19,9 +19,7 @@
 import org.opensearch.common.network.NetworkService;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
-import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.plugins.SecureTransportSettingsProvider;
-import org.opensearch.test.FeatureFlagSetter;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ExecutorBuilder;
 import org.opensearch.threadpool.ThreadPool;
@@ -31,19 +29,18 @@
 import java.util.List;
 import java.util.function.Supplier;
 
-import static org.opensearch.common.util.FeatureFlags.ARROW_STREAMS_SETTING;
+import static org.opensearch.common.util.FeatureFlags.ARROW_STREAMS;
 import static org.opensearch.plugins.NetworkPlugin.AuxTransport.AUX_TRANSPORT_TYPES_KEY;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 public class FlightStreamPluginTests extends OpenSearchTestCase {
-    private Settings settings;
+    private final Settings settings = Settings.EMPTY;
     private ClusterService clusterService;
 
     @Override
     public void setUp() throws Exception {
         super.setUp();
-        settings = Settings.builder().put(ARROW_STREAMS_SETTING.getKey(), true).build();
         clusterService = mock(ClusterService.class);
         ClusterState clusterState = mock(ClusterState.class);
         DiscoveryNodes nodes = mock(DiscoveryNodes.class);
@@ -52,9 +49,8 @@ public void setUp() throws Exception {
         when(nodes.getLocalNodeId()).thenReturn("test-node");
     }
 
+    @LockFeatureFlag(ARROW_STREAMS)
     public void testPluginEnabled() throws IOException {
-        FeatureFlags.initializeFeatureFlags(settings);
-        FeatureFlagSetter.set(ARROW_STREAMS_SETTING.getKey());
         FlightStreamPlugin plugin = new FlightStreamPlugin(settings);
         Collection<Object> components = plugin.createComponents(
             null,

plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/bootstrap/FlightClientManagerTests.java

@@ -28,7 +28,6 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.transport.BoundTransportAddress;
 import org.opensearch.core.common.transport.TransportAddress;
-import org.opensearch.test.FeatureFlagSetter;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.client.Client;
@@ -55,6 +54,7 @@
 import io.netty.util.NettyRuntime;
 
 import static org.opensearch.arrow.flight.bootstrap.FlightClientManager.LOCATION_TIMEOUT_MS;
+import static org.opensearch.common.util.FeatureFlags.ARROW_STREAMS;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
@@ -64,6 +64,7 @@
 
 @SuppressWarnings("unchecked")
 public class FlightClientManagerTests extends OpenSearchTestCase {
+    private static FeatureFlags.TestUtils.FlagWriteLock ffLock = null;
 
     private static BufferAllocator allocator;
     private static EventLoopGroup elg;
@@ -78,6 +79,7 @@ public class FlightClientManagerTests extends OpenSearchTestCase {
 
     @BeforeClass
     public static void setupClass() throws Exception {
+        ffLock = new FeatureFlags.TestUtils.FlagWriteLock(ARROW_STREAMS);
         ServerConfig.init(Settings.EMPTY);
         allocator = new RootAllocator();
         elg = ServerConfig.createELG("test-grpc-worker-elg", NettyRuntime.availableProcessors() * 2);
@@ -89,7 +91,6 @@ public void setUp() throws Exception {
         super.setUp();
         locationUpdaterExecutor = Executors.newScheduledThreadPool(1);
 
-        FeatureFlagSetter.set(FeatureFlags.ARROW_STREAMS_SETTING.getKey());
         clusterService = mock(ClusterService.class);
         client = mock(Client.class);
         state = getDefaultState();
@@ -176,6 +177,7 @@ private DiscoveryNode createNode(String nodeId, String host, int port) throws Ex
     @AfterClass
     public static void tearClass() {
         allocator.close();
+        ffLock.close();
     }
 
     public void testGetFlightClientForExistingNode() {

plugins/arrow-flight-rpc/src/test/java/org/opensearch/arrow/flight/bootstrap/FlightServiceTests.java

@@ -19,7 +19,6 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.common.transport.TransportAddress;
-import org.opensearch.test.FeatureFlagSetter;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.client.Client;
@@ -32,10 +31,12 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import static org.opensearch.common.util.FeatureFlags.ARROW_STREAMS;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 public class FlightServiceTests extends OpenSearchTestCase {
+    FeatureFlags.TestUtils.FlagWriteLock ffLock = null;
 
     private Settings settings;
     private ClusterService clusterService;
@@ -47,7 +48,7 @@ public class FlightServiceTests extends OpenSearchTestCase {
     @Override
     public void setUp() throws Exception {
         super.setUp();
-        FeatureFlagSetter.set(FeatureFlags.ARROW_STREAMS_SETTING.getKey());
+        ffLock = new FeatureFlags.TestUtils.FlagWriteLock(ARROW_STREAMS);
         int availablePort = getBasePort(9500) + port.addAndGet(1);
         settings = Settings.EMPTY;
         localNode = createNode(availablePort);
@@ -147,6 +148,7 @@ public void testLifecycleStateTransitions() throws Exception {
     @Override
     public void tearDown() throws Exception {
         super.tearDown();
+        ffLock.close();
     }
 
     private DiscoveryNode createNode(int port) throws Exception {