Revert Skip BoundedTrie on Dataflow till service is have BoundedTrie apache#33921

Author: rohitsinha54

runners/google-cloud-dataflow-java/build.gradle

@@ -193,7 +193,6 @@ def commonLegacyExcludeCategories = [
   'org.apache.beam.sdk.testing.UsesParDoLifecycle', // doesn't support remote runner
   'org.apache.beam.sdk.testing.UsesMetricsPusher',
   'org.apache.beam.sdk.testing.UsesBundleFinalizer',
-  'org.apache.beam.sdk.testing.UsesBoundedTrieMetrics',
 ]
 
 def commonRunnerV2ExcludeCategories = [
@@ -207,7 +206,6 @@ def commonRunnerV2ExcludeCategories = [
         'org.apache.beam.sdk.testing.UsesTestStream',
         'org.apache.beam.sdk.testing.UsesTestStreamWithProcessingTime',
         'org.apache.beam.sdk.testing.UsesRequiresTimeSortedInput',
-        'org.apache.beam.sdk.testing.UsesBoundedTrieMetrics',
 ]
 
 // For the following test tasks using legacy worker, set workerHarnessContainerImage to empty to

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeExecutionContext.java

@@ -19,7 +19,6 @@
 
 import com.google.api.services.dataflow.model.CounterUpdate;
 import com.google.api.services.dataflow.model.SideInputInfo;
-import java.util.Collections;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 import org.apache.beam.runners.core.InMemoryStateInternals;
@@ -78,18 +77,14 @@ public class BatchModeExecutionContext
   protected static final String BIGQUERY_READ_THROTTLE_TIME_NAMESPACE =
       "org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl$StorageClientImpl";
 
-  // TODO(BEAM-33720): Remove once Dataflow legacy runner supports BoundedTries.
-  private final boolean populateBoundedTrieMetrics;
-
   private BatchModeExecutionContext(
       CounterFactory counterFactory,
       Cache<?, WeightedValue<?>> dataCache,
       Cache<?, ?> logicalReferenceCache,
       ReaderFactory readerFactory,
       PipelineOptions options,
       DataflowExecutionStateTracker executionStateTracker,
-      DataflowExecutionStateRegistry executionStateRegistry,
-      boolean populateBoundedTrieMetrics) {
+      DataflowExecutionStateRegistry executionStateRegistry) {
     super(
         counterFactory,
         createMetricsContainerRegistry(),
@@ -102,7 +97,6 @@ private BatchModeExecutionContext(
     this.dataCache = dataCache;
     this.containerRegistry =
         (MetricsContainerRegistry<MetricsContainerImpl>) getMetricsContainerRegistry();
-    this.populateBoundedTrieMetrics = populateBoundedTrieMetrics;
   }
 
   private static MetricsContainerRegistry<MetricsContainerImpl> createMetricsContainerRegistry() {
@@ -138,8 +132,7 @@ public static BatchModeExecutionContext forTesting(
             counterFactory,
             options,
             "test-work-item-id"),
-        stateRegistry,
-        true);
+        stateRegistry);
   }
 
   public static BatchModeExecutionContext forTesting(PipelineOptions options, String stageName) {
@@ -252,8 +245,7 @@ public static BatchModeExecutionContext create(
             counterFactory,
             options,
             workItemId),
-        executionStateRegistry,
-        false);
+        executionStateRegistry);
   }
 
   /** Create a new {@link StepContext}. */
@@ -528,10 +520,7 @@ public Iterable<CounterUpdate> extractMetricUpdates(boolean isFinalUpdate) {
                           update ->
                               MetricsToCounterUpdateConverter.fromStringSet(
                                   update.getKey(), true, update.getUpdate())),
-                  FluentIterable.from(
-                          populateBoundedTrieMetrics
-                              ? updates.boundedTrieUpdates()
-                              : Collections.emptyList())
+                  FluentIterable.from(updates.boundedTrieUpdates())
                       .transform(
                           update ->
                               MetricsToCounterUpdateConverter.fromBoundedTrie(

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingStepMetricsContainer.java

@@ -23,7 +23,6 @@
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.Collections;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -104,9 +103,6 @@ public class StreamingStepMetricsContainer implements MetricsContainer {
 
   private final Clock clock;
 
-  // TODO(BEAM-33720): Remove once Dataflow legacy runner supports BoundedTries.
-  @VisibleForTesting boolean populateBoundedTrieMetrics;
-
   private StreamingStepMetricsContainer(String stepName) {
     this.stepName = stepName;
     this.perWorkerCountersByFirstStaleTime = new ConcurrentHashMap<>();
@@ -217,7 +213,7 @@ public Iterable<CounterUpdate> extractUpdates() {
         .append(distributionUpdates())
         .append(gaugeUpdates())
         .append(stringSetUpdates())
-        .append(populateBoundedTrieMetrics ? boundedTrieUpdates() : Collections.emptyList());
+        .append(boundedTrieUpdates());
   }
 
   private FluentIterable<CounterUpdate> counterUpdates() {

runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingStepMetricsContainerTest.java

@@ -371,7 +371,6 @@ public void testBoundedTrieUpdateExtraction() {
             .setBoundedTrie(
                 MetricsToCounterUpdateConverter.getBoundedTrie(expectedName1.toProto()));
 
-    ((StreamingStepMetricsContainer) c1).populateBoundedTrieMetrics = true;
     Iterable<CounterUpdate> updates = StreamingStepMetricsContainer.extractMetricUpdates(registry);
     assertThat(updates, containsInAnyOrder(name1Update));
 
@@ -400,7 +399,6 @@ public void testBoundedTrieUpdateExtraction() {
             .setBoundedTrie(
                 MetricsToCounterUpdateConverter.getBoundedTrie(expectedName2.toProto()));
 
-    ((StreamingStepMetricsContainer) c2).populateBoundedTrieMetrics = true;
     updates = StreamingStepMetricsContainer.extractMetricUpdates(registry);
     assertThat(updates, containsInAnyOrder(name2Update));
 
@@ -412,7 +410,6 @@ public void testBoundedTrieUpdateExtraction() {
     name1Update.setBoundedTrie(
         MetricsToCounterUpdateConverter.getBoundedTrie(expectedName1.toProto()));
 
-    ((StreamingStepMetricsContainer) c1).populateBoundedTrieMetrics = true;
     updates = StreamingStepMetricsContainer.extractMetricUpdates(registry);
     assertThat(updates, containsInAnyOrder(name1Update));
   }