Revert "Avoid publishing string set metrics on the Dataflow legacy ru…

…nner." (#32002) * Revert "Avoid publishing string set metrics on the Dataflow legacy runner. (#…" This reverts commit a226094. * unskip stingset * Fix ClassCastException
apache · Jul 26, 2024 · 7930a1f · 7930a1f
1 parent 4479281
commit 7930a1f
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 26 deletions.
diff --git a/runners/google-cloud-dataflow-java/build.gradle b/runners/google-cloud-dataflow-java/build.gradle
@@ -183,7 +183,6 @@ def commonLegacyExcludeCategories = [
  'org.apache.beam.sdk.testing.UsesExternalService',
  'org.apache.beam.sdk.testing.UsesDistributionMetrics',
  'org.apache.beam.sdk.testing.UsesGaugeMetrics',
- 'org.apache.beam.sdk.testing.UsesStringSetMetrics',
  'org.apache.beam.sdk.testing.UsesMultimapState',
  'org.apache.beam.sdk.testing.UsesTestStream',
  'org.apache.beam.sdk.testing.UsesParDoLifecycle',

diff --git a/...e-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowMetrics.java b/...e-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowMetrics.java
@@ -24,11 +24,11 @@
 import com.google.api.services.dataflow.model.JobMetrics;
 import com.google.api.services.dataflow.model.MetricUpdate;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 import org.apache.beam.model.pipeline.v1.RunnerApi;
 import org.apache.beam.sdk.metrics.DistributionResult;
 import org.apache.beam.sdk.metrics.GaugeResult;
@@ -191,7 +191,7 @@ private StringSetResult getStringSetValue(MetricUpdate metricUpdate) {
  if (metricUpdate.getSet() == null) {
  return StringSetResult.empty();
  }
- return StringSetResult.create(ImmutableSet.copyOf(((Set) metricUpdate.getSet())));
+ return StringSetResult.create(ImmutableSet.copyOf(((Collection) metricUpdate.getSet())));
  }
 
  private DistributionResult getDistributionValue(MetricUpdate metricUpdate) {

diff --git a/...rker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchModeExecutionContext.java b/...rker/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-31814): Remove once Dataflow legacy runner supports this.
- private final boolean populateStringSetMetrics;
-
  private BatchModeExecutionContext(
  CounterFactory counterFactory,
  Cache<?, WeightedValue<?>> dataCache,
  Cache<?, ?> logicalReferenceCache,
  ReaderFactory readerFactory,
  PipelineOptions options,
  DataflowExecutionStateTracker executionStateTracker,
- DataflowExecutionStateRegistry executionStateRegistry,
- boolean populateStringSetMetrics) {
+ DataflowExecutionStateRegistry executionStateRegistry) {
  super(
  counterFactory,
  createMetricsContainerRegistry(),
@@ -102,7 +97,6 @@ private BatchModeExecutionContext(
  this.dataCache = dataCache;
  this.containerRegistry =
  (MetricsContainerRegistry<MetricsContainerImpl>) getMetricsContainerRegistry();
- this.populateStringSetMetrics = populateStringSetMetrics;
  }
 
  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}. */
@@ -523,10 +515,7 @@ public Iterable<CounterUpdate> extractMetricUpdates(boolean isFinalUpdate) {
  update ->
  MetricsToCounterUpdateConverter.fromDistribution(
  update.getKey(), true, update.getUpdate())),
- FluentIterable.from(
- populateStringSetMetrics
- ? updates.stringSetUpdates()
- : Collections.emptyList())
+ FluentIterable.from(updates.stringSetUpdates())
  .transform(
  update ->
  MetricsToCounterUpdateConverter.fromStringSet(

diff --git a/.../src/main/java/org/apache/beam/runners/dataflow/worker/StreamingStepMetricsContainer.java b/.../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;
@@ -89,9 +88,6 @@ public class StreamingStepMetricsContainer implements MetricsContainer {
 
  private final Clock clock;
 
- // TODO(BEAM-31814): Remove once Dataflow legacy runner supports this.
- @VisibleForTesting boolean populateStringSetUpdates = false;
-
  private StreamingStepMetricsContainer(String stepName) {
  this.stepName = stepName;
  this.perWorkerCountersByFirstStaleTime = new ConcurrentHashMap<>();
@@ -191,8 +187,7 @@ public Histogram getPerWorkerHistogram(
  public Iterable<CounterUpdate> extractUpdates() {
  return counterUpdates()
  .append(distributionUpdates())
- .append(gaugeUpdates())
- .append(populateStringSetUpdates ? stringSetUpdates() : Collections.emptyList());
+ .append(gaugeUpdates().append(stringSetUpdates()));
  }
 
  private FluentIterable<CounterUpdate> counterUpdates() {

diff --git a/.../test/java/org/apache/beam/runners/dataflow/worker/StreamingStepMetricsContainerTest.java b/.../test/java/org/apache/beam/runners/dataflow/worker/StreamingStepMetricsContainerTest.java
@@ -292,7 +292,6 @@ public void testStringSetUpdateExtraction() {
  .setCumulative(false)
  .setStringList(new StringList().setElements(Arrays.asList("ab", "cd", "ef", "gh")));
 
- ((StreamingStepMetricsContainer) c1).populateStringSetUpdates = true;
  Iterable<CounterUpdate> updates = StreamingStepMetricsContainer.extractMetricUpdates(registry);
  assertThat(updates, containsInAnyOrder(name1Update));
 
@@ -315,7 +314,6 @@ public void testStringSetUpdateExtraction() {
  .setCumulative(false)
  .setStringList(new StringList().setElements(Arrays.asList("ij", "kl", "mn")));
 
- ((StreamingStepMetricsContainer) c2).populateStringSetUpdates = true;
  updates = StreamingStepMetricsContainer.extractMetricUpdates(registry);
  assertThat(updates, containsInAnyOrder(name1Update, name2Update));