Add a flag to disable BoundedTrie metrics in Beam

apache · Feb 3, 2025 · 5aae10d · 5aae10d
1 parent 0230b5d
commit 5aae10d
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Metrics.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/Metrics.java
@@ -65,10 +65,13 @@ static class MetricsFlag {
     private static final AtomicReference<@Nullable MetricsFlag> INSTANCE = new AtomicReference<>();
     final boolean counterDisabled;
     final boolean stringSetDisabled;
+    final boolean boundedTrieDisabled;
 
-    private MetricsFlag(boolean counterDisabled, boolean stringSetDisabled) {
+    private MetricsFlag(
+        boolean counterDisabled, boolean stringSetDisabled, boolean boundedTrieDisabled) {
       this.counterDisabled = counterDisabled;
       this.stringSetDisabled = stringSetDisabled;
+      this.boundedTrieDisabled = boundedTrieDisabled;
     }
 
     static boolean counterDisabled() {
@@ -80,6 +83,11 @@ static boolean stringSetDisabled() {
       MetricsFlag flag = INSTANCE.get();
       return flag != null && flag.stringSetDisabled;
     }
+
+    static boolean boundedTrieDisabled() {
+      MetricsFlag flag = INSTANCE.get();
+      return flag != null && flag.boundedTrieDisabled;
+    }
   }
 
   /**
@@ -101,7 +109,13 @@ public static void setDefaultPipelineOptions(PipelineOptions options) {
       if (stringSetDisabled) {
         LOG.info("StringSet metrics are disabled");
       }
-      MetricsFlag.INSTANCE.compareAndSet(null, new MetricsFlag(counterDisabled, stringSetDisabled));
+      boolean boundedTrieDisabled =
+          ExperimentalOptions.hasExperiment(exp, "disableBoundedTrieMetrics");
+      if (boundedTrieDisabled) {
+        LOG.info("BoundedTrie metrics are disabled");
+      }
+      MetricsFlag.INSTANCE.compareAndSet(
+          null, new MetricsFlag(counterDisabled, stringSetDisabled, boundedTrieDisabled));
     }
   }
 
@@ -287,6 +301,9 @@ public MetricName getName() {
 
     @Override
     public void add(Iterable<String> values) {
+      if (MetricsFlag.boundedTrieDisabled()) {
+        return;
+      }
       MetricsContainer container = MetricsEnvironment.getCurrentContainer();
       if (container != null) {
         container.getBoundedTrie(name).add(values);

diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/metrics/MetricsTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/metrics/MetricsTest.java
@@ -263,6 +263,7 @@ public void testMetricsFlag() {
       Metrics.resetDefaultPipelineOptions();
       assertFalse(Metrics.MetricsFlag.counterDisabled());
       assertFalse(Metrics.MetricsFlag.stringSetDisabled());
+      assertFalse(Metrics.MetricsFlag.boundedTrieDisabled());
       PipelineOptions options =
           PipelineOptionsFactory.fromArgs("--experiments=disableCounterMetrics").create();
       Metrics.setDefaultPipelineOptions(options);
@@ -274,6 +275,12 @@ public void testMetricsFlag() {
       assertFalse(Metrics.MetricsFlag.counterDisabled());
       assertTrue(Metrics.MetricsFlag.stringSetDisabled());
       Metrics.resetDefaultPipelineOptions();
+      options = PipelineOptionsFactory.fromArgs("--experiments=disableBoundedTrieMetrics").create();
+      Metrics.setDefaultPipelineOptions(options);
+      assertFalse(Metrics.MetricsFlag.counterDisabled());
+      assertFalse(Metrics.MetricsFlag.stringSetDisabled());
+      assertTrue(Metrics.MetricsFlag.boundedTrieDisabled());
+      Metrics.resetDefaultPipelineOptions();
     }
   }