refactor histograms as per issue #34195

Author: Naireen

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

@@ -75,6 +75,11 @@ public Gauge getGauge(MetricName metricName) {
     return getCurrentContainer().getGauge(metricName);
   }
 
+  @Override
+  public Histogram getHistogram(MetricName metricName, HistogramData.BucketType bucketType) {
+    return getCurrentContainer().getHistogram(metricName, bucketType);
+  }
+
   @Override
   public StringSet getStringSet(MetricName metricName) {
     return getCurrentContainer().getStringSet(metricName);
@@ -84,10 +89,4 @@ public StringSet getStringSet(MetricName metricName) {
   public BoundedTrie getBoundedTrie(MetricName metricName) {
     return getCurrentContainer().getBoundedTrie(metricName);
   }
-
-  @Override
-  public Histogram getPerWorkerHistogram(
-      MetricName metricName, HistogramData.BucketType bucketType) {
-    return getCurrentContainer().getPerWorkerHistogram(metricName, bucketType);
-  }
 }

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

@@ -48,12 +48,14 @@
 import org.apache.beam.sdk.metrics.MetricKey;
 import org.apache.beam.sdk.metrics.MetricName;
 import org.apache.beam.sdk.metrics.MetricsContainer;
+import org.apache.beam.sdk.metrics.NoOpHistogram;
 import org.apache.beam.sdk.metrics.StringSet;
 import org.apache.beam.sdk.util.HistogramData;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Function;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicates;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.FluentIterable;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
@@ -76,6 +78,9 @@ public class StreamingStepMetricsContainer implements MetricsContainer {
 
   private MetricsMap<MetricName, GaugeCell> gauges = new MetricsMap<>(GaugeCell::new);
 
+  // Stores metrics that will be later aggregted per worker.
+  private MetricsMap<MetricName, GaugeCell> perWorkerGauges = new MetricsMap<>(GaugeCell::new);
+
   private MetricsMap<MetricName, StringSetCell> stringSets = new MetricsMap<>(StringSetCell::new);
 
   private MetricsMap<MetricName, DeltaDistributionCell> distributions =
@@ -84,6 +89,7 @@ public class StreamingStepMetricsContainer implements MetricsContainer {
   private MetricsMap<MetricName, BoundedTrieCell> boundedTries =
       new MetricsMap<>(BoundedTrieCell::new);
 
+  // Stores metrics that will be later aggregted per worker.
   private final ConcurrentHashMap<MetricName, LockFreeHistogram> perWorkerHistograms =
       new ConcurrentHashMap<>();
 
@@ -173,6 +179,9 @@ public Distribution getDistribution(MetricName metricName) {
 
   @Override
   public Gauge getGauge(MetricName metricName) {
+    if (MonitoringInfoConstants.isPerWorkerMetric(metricName)) {
+      return perWorkerGauges.get(metricName);
+    }
     return gauges.get(metricName);
   }
 
@@ -187,10 +196,10 @@ public BoundedTrie getBoundedTrie(MetricName metricName) {
   }
 
   @Override
-  public Histogram getPerWorkerHistogram(
-      MetricName metricName, HistogramData.BucketType bucketType) {
-    if (!enablePerWorkerMetrics) {
-      return MetricsContainer.super.getPerWorkerHistogram(metricName, bucketType);
+  public Histogram getHistogram(MetricName metricName, HistogramData.BucketType bucketType) {
+    // Currently worker only supports per worker histograms
+    if (!enablePerWorkerMetrics || !MonitoringInfoConstants.isPerWorkerMetric(metricName)) {
+      return NoOpHistogram.getInstance();
     }
 
     LockFreeHistogram val = perWorkerHistograms.get(metricName);
@@ -247,7 +256,6 @@ private FluentIterable<CounterUpdate> gaugeUpdates() {
                   return MetricsToCounterUpdateConverter.fromGauge(
                       MetricKey.create(stepName, entry.getKey()), value, timestamp);
                 } else {
-                  // add a test for this.
                   return null;
                 }
               }
@@ -377,11 +385,11 @@ private void deleteStaleCounters(
    */
   @VisibleForTesting
   Iterable<PerStepNamespaceMetrics> extractPerWorkerMetricUpdates() {
-    ConcurrentHashMap<MetricName, Long> counters = new ConcurrentHashMap<MetricName, Long>();
-    ConcurrentHashMap<MetricName, Long> per_worker_gauges =
-        new ConcurrentHashMap<MetricName, Long>();
-    ConcurrentHashMap<MetricName, LockFreeHistogram.Snapshot> histograms =
-        new ConcurrentHashMap<MetricName, LockFreeHistogram.Snapshot>();
+    ImmutableMap.Builder<MetricName, Long> counters = new ImmutableMap.Builder<MetricName, Long>();
+    ImmutableMap.Builder<MetricName, Long> per_worker_gauges =
+        new ImmutableMap.Builder<MetricName, Long>();
+    ImmutableMap.Builder<MetricName, LockFreeHistogram.Snapshot> per_worker_histograms =
+        new ImmutableMap.Builder<MetricName, LockFreeHistogram.Snapshot>();
     HashSet<MetricName> currentZeroValuedCounters = new HashSet<MetricName>();
     perWorkerCounters.forEach(
         (k, v) -> {
@@ -393,27 +401,26 @@ Iterable<PerStepNamespaceMetrics> extractPerWorkerMetricUpdates() {
           counters.put(k, val);
         });
 
-    gauges.forEach(
+    perWorkerGauges.forEach(
         (k, v) -> {
-          // Check if metric name has the per worker label set.
-          // TODO(Naireen): Populate local map with perWorkerMetrics so we don't need to check each
-          // time we update the metrics.
-          if (MonitoringInfoConstants.isPerWorkerMetric(k)) {
-            Long val = v.getCumulative().value();
-            per_worker_gauges.put(k, val);
-            v.reset();
-          }
+          Long val = v.getCumulative().value();
+          per_worker_gauges.put(k, val);
+          v.reset();
         });
 
     perWorkerHistograms.forEach(
         (k, v) -> {
-          v.getSnapshotAndReset().ifPresent(snapshot -> histograms.put(k, snapshot));
+          v.getSnapshotAndReset().ifPresent(snapshot -> per_worker_histograms.put(k, snapshot));
         });
 
     deleteStaleCounters(currentZeroValuedCounters, Instant.now(clock));
 
     return MetricsToPerStepNamespaceMetricsConverter.convert(
-        stepName, counters, per_worker_gauges, histograms, parsedPerWorkerMetricsCache);
+        stepName,
+        counters.build(),
+        per_worker_gauges.build(),
+        per_worker_histograms.build(),
+        parsedPerWorkerMetricsCache);
   }
 
   /**

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

@@ -421,17 +421,23 @@ public void testBoundedTrieUpdateExtraction() {
   public void testPerWorkerMetrics() {
     StreamingStepMetricsContainer.setEnablePerWorkerMetrics(false);
     MetricsContainer metricsContainer = registry.getContainer("test_step");
+    MetricName histogramMetricName =
+        MetricName.named("BigQuerySink", "histogram", ImmutableMap.of("PER_WORKER_METRIC", "true"));
+
     assertThat(
-        metricsContainer.getPerWorkerCounter(name1), sameInstance(NoOpCounter.getInstance()));
+        metricsContainer.getPerWorkerCounter(histogramMetricName),
+        sameInstance(NoOpCounter.getInstance()));
     HistogramData.BucketType testBucket = HistogramData.LinearBuckets.of(1, 1, 1);
     assertThat(
-        metricsContainer.getPerWorkerHistogram(name1, testBucket),
+        metricsContainer.getHistogram(histogramMetricName, testBucket),
         sameInstance(NoOpHistogram.getInstance()));
 
     StreamingStepMetricsContainer.setEnablePerWorkerMetrics(true);
-    assertThat(metricsContainer.getPerWorkerCounter(name1), not(instanceOf(NoOpCounter.class)));
     assertThat(
-        metricsContainer.getPerWorkerHistogram(name1, testBucket),
+        metricsContainer.getPerWorkerCounter(histogramMetricName),
+        not(instanceOf(NoOpCounter.class)));
+    assertThat(
+        metricsContainer.getHistogram(histogramMetricName, testBucket),
         not(instanceOf(NoOpHistogram.class)));
   }
 
@@ -441,9 +447,10 @@ public void testExtractPerWorkerMetricUpdates_populatedMetrics() {
     MetricName counterMetricName = MetricName.named("BigQuerySink", "counter");
     c1.getPerWorkerCounter(counterMetricName).inc(3);
 
-    MetricName histogramMetricName = MetricName.named("BigQuerySink", "histogram");
+    MetricName histogramMetricName =
+        MetricName.named("BigQuerySink", "histogram", ImmutableMap.of("PER_WORKER_METRIC", "true"));
     HistogramData.LinearBuckets linearBuckets = HistogramData.LinearBuckets.of(0, 10, 10);
-    c2.getPerWorkerHistogram(histogramMetricName, linearBuckets).update(5.0);
+    c2.getHistogram(histogramMetricName, linearBuckets).update(5.0);
 
     Iterable<PerStepNamespaceMetrics> updates =
         StreamingStepMetricsContainer.extractPerWorkerMetricUpdates(registry);
@@ -488,9 +495,10 @@ public void testExtractPerWorkerMetricUpdates_populatedMetrics() {
   public void testExtractPerWorkerMetricUpdatesKafka_populatedHistogramMetrics() {
     StreamingStepMetricsContainer.setEnablePerWorkerMetrics(true);
 
-    MetricName histogramMetricName = MetricName.named("KafkaSink", "histogram");
+    MetricName histogramMetricName =
+        MetricName.named("KafkaSink", "histogram", ImmutableMap.of("PER_WORKER_METRIC", "true"));
     HistogramData.LinearBuckets linearBuckets = HistogramData.LinearBuckets.of(0, 10, 10);
-    c2.getPerWorkerHistogram(histogramMetricName, linearBuckets).update(5.0);
+    c2.getHistogram(histogramMetricName, linearBuckets).update(5.0);
 
     Iterable<PerStepNamespaceMetrics> updates =
         StreamingStepMetricsContainer.extractPerWorkerMetricUpdates(registry);
@@ -552,13 +560,17 @@ public void testExtractPerWorkerMetricUpdatesKafka_populateGaugeMetrics() {
   }
 
   @Test
-  public void testExtractPerWorkerMetricUpdatesKafka_gaugeMetricsDropped() {
+  public void testExtractPerWorkerMetricUpdatesKafka_perWorkerMetricsDropped() {
     StreamingStepMetricsContainer.setEnablePerWorkerMetrics(true);
 
     MetricName gaugeMetricName =
         MetricName.named("KafkaSink", "gauge", ImmutableMap.of("PER_WORKER_METRIC", "false"));
     c2.getGauge(gaugeMetricName).set(5L);
 
+    MetricName histogramMetricName = MetricName.named("BigQuerySink", "histogram");
+    HistogramData.LinearBuckets linearBuckets = HistogramData.LinearBuckets.of(0, 10, 10);
+    c2.getHistogram(histogramMetricName, linearBuckets);
+
     Iterable<PerStepNamespaceMetrics> updates =
         StreamingStepMetricsContainer.extractPerWorkerMetricUpdates(registry);
     assertThat(updates, IsEmptyIterable.emptyIterable());
@@ -573,7 +585,7 @@ public void testExtractPerWorkerMetricUpdates_emptyMetrics() {
 
     MetricName histogramMetricName = MetricName.named("BigQuerySink", "histogram");
     HistogramData.LinearBuckets linearBuckets = HistogramData.LinearBuckets.of(0, 10, 10);
-    c2.getPerWorkerHistogram(histogramMetricName, linearBuckets);
+    c2.getHistogram(histogramMetricName, linearBuckets);
 
     Iterable<PerStepNamespaceMetrics> updates =
         StreamingStepMetricsContainer.extractPerWorkerMetricUpdates(registry);

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

@@ -62,7 +62,7 @@ public void testTranslateKnownPerWorkerCounters_malformedCounters() throws Excep
             .getName();
 
     HistogramData.BucketType linearBuckets = HistogramData.LinearBuckets.of(0, 10.0, 10);
-    metricsContainer.getPerWorkerHistogram(name, linearBuckets).update(10.0);
+    metricsContainer.getHistogram(name, linearBuckets).update(10.0);
 
     stageInfo.extractPerWorkerMetricValues();
     assertThat(stageInfo.throttledMsecs().getAggregate(), equalTo(0L));

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/DelegatingHistogram.java

@@ -28,37 +28,18 @@ public class DelegatingHistogram implements Metric, Histogram, Serializable {
   private final MetricName name;
   private final HistogramData.BucketType bucketType;
   private final boolean processWideContainer;
-  private final boolean perWorkerHistogram;
 
   /**
-   * Create a {@code DelegatingHistogram} with {@code perWorkerHistogram} set to false.
-   *
    * @param name Metric name for this metric.
    * @param bucketType Histogram bucketing strategy.
    * @param processWideContainer Whether this Counter is stored in the ProcessWide container or the
    *     current thread's container.
    */
   public DelegatingHistogram(
       MetricName name, HistogramData.BucketType bucketType, boolean processWideContainer) {
-    this(name, bucketType, processWideContainer, false);
-  }
-
-  /**
-   * @param name Metric name for this metric.
-   * @param bucketType Histogram bucketing strategy.
-   * @param processWideContainer Whether this Counter is stored in the ProcessWide container or the
-   *     current thread's container.
-   * @param perWorkerHistogram Whether this Histogram refers to a perWorker metric or not.
-   */
-  public DelegatingHistogram(
-      MetricName name,
-      HistogramData.BucketType bucketType,
-      boolean processWideContainer,
-      boolean perWorkerHistogram) {
     this.name = name;
     this.bucketType = bucketType;
     this.processWideContainer = processWideContainer;
-    this.perWorkerHistogram = perWorkerHistogram;
   }
 
   private Optional<Histogram> getHistogram() {
@@ -69,11 +50,8 @@ private Optional<Histogram> getHistogram() {
     if (container == null) {
       return Optional.empty();
     }
-    if (perWorkerHistogram) {
-      return Optional.of(container.getPerWorkerHistogram(name, bucketType));
-    } else {
-      return Optional.of(container.getHistogram(name, bucketType));
-    }
+    // LOG.info("Xxx get histogram");
+    return Optional.of(container.getHistogram(name, bucketType));
   }
 
   @Override

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/MetricsContainer.java

@@ -26,7 +26,6 @@
  * appropriate metrics interface for the "current" step.
  */
 public interface MetricsContainer extends Serializable {
-
   /**
    * Return the {@link Counter} that should be used for implementing the given {@code metricName} in
    * this container.
@@ -70,14 +69,6 @@ default Counter getPerWorkerCounter(MetricName metricName) {
    * in this container.
    */
   default Histogram getHistogram(MetricName metricName, HistogramData.BucketType bucketType) {
-    throw new RuntimeException("Histogram metric is not supported yet.");
-  }
-  /**
-   * Return the {@link Histogram} that should be used for implementing the given per-worker {@code
-   * metricName} in this container.
-   */
-  default Histogram getPerWorkerHistogram(
-      MetricName metricName, HistogramData.BucketType bucketType) {
     return NoOpHistogram.getInstance();
   }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/metrics/NoOpGauge.java

@@ -21,6 +21,7 @@
 import java.time.Instant;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import org.apache.beam.runners.core.metrics.MonitoringInfoConstants;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils.NestedCounter;
 import org.apache.beam.sdk.io.gcp.bigquery.RetryManager.Operation.Context;
 import org.apache.beam.sdk.metrics.Counter;
@@ -114,13 +115,14 @@ static Histogram createRPCLatencyHistogram(RpcMethod method) {
         LabeledMetricNameUtils.MetricNameBuilder.baseNameBuilder(RPC_LATENCY);
     nameBuilder.addLabel(RPC_METHOD, method.toString());
     MetricName metricName = nameBuilder.build(METRICS_NAMESPACE);
+    nameBuilder.addMetricLabel(MonitoringInfoConstants.Labels.PER_WORKER_METRIC, "true");
 
     // Create Exponential histogram buckets with the following parameters:
     // 0 scale, resulting in bucket widths with a size 2 growth factor.
     // 17 buckets, so the max latency of that can be stored is (2^17 millis ~= 130 seconds).
     HistogramData.BucketType buckets = HistogramData.ExponentialBuckets.of(0, 17);
 
-    return new DelegatingHistogram(metricName, buckets, false, true);
+    return new DelegatingHistogram(metricName, buckets, false);
   }
 
   /**

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetrics.java

@@ -83,7 +83,7 @@ public TestMetricsContainer() {
     }
 
     @Override
-    public TestHistogramCell getPerWorkerHistogram(
+    public TestHistogramCell getHistogram(
         MetricName metricName, HistogramData.BucketType bucketType) {
       perWorkerHistograms.computeIfAbsent(
           KV.of(metricName, bucketType), kv -> new TestHistogramCell(kv));

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetricsTest.java

@@ -72,9 +72,10 @@ public static Histogram createRPCLatencyHistogram(RpcMethod method, String topic
     nameBuilder.addLabel(TOPIC_LABEL, topic);
 
     MetricName metricName = nameBuilder.build(METRICS_NAMESPACE);
-    HistogramData.BucketType buckets = HistogramData.ExponentialBuckets.of(1, 17);
+    nameBuilder.addMetricLabel(MonitoringInfoConstants.Labels.PER_WORKER_METRIC, "true");
 
-    return new DelegatingHistogram(metricName, buckets, false, true);
+    HistogramData.BucketType buckets = HistogramData.ExponentialBuckets.of(1, 17);
+    return new DelegatingHistogram(metricName, buckets, false);
   }
 
   /**