refactor to remove per worker gauge method

Author: Naireen

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

@@ -75,12 +75,6 @@ public Gauge getGauge(MetricName metricName) {
     return getCurrentContainer().getGauge(metricName);
   }
 
-  @Override
-  public Gauge getPerWorkerGauge(MetricName metricName) {
-    Gauge gauge = getCurrentContainer().getPerWorkerGauge(metricName);
-    return gauge;
-  }
-
   @Override
   public StringSet getStringSet(MetricName metricName) {
     return getCurrentContainer().getStringSet(metricName);

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

@@ -75,9 +75,6 @@ public class StreamingStepMetricsContainer implements MetricsContainer {
 
   private MetricsMap<MetricName, GaugeCell> gauges = new MetricsMap<>(GaugeCell::new);
 
-  private final ConcurrentHashMap<MetricName, GaugeCell> perWorkerGauges =
-      new ConcurrentHashMap<>();
-
   private MetricsMap<MetricName, StringSetCell> stringSets = new MetricsMap<>(StringSetCell::new);
 
   private MetricsMap<MetricName, DeltaDistributionCell> distributions =
@@ -178,19 +175,6 @@ public Gauge getGauge(MetricName metricName) {
     return gauges.get(metricName);
   }
 
-  @Override
-  public Gauge getPerWorkerGauge(MetricName metricName) {
-    if (!enablePerWorkerMetrics) {
-      return MetricsContainer.super.getPerWorkerGauge(metricName);
-    }
-    Gauge val = perWorkerGauges.get(metricName);
-    if (val != null) {
-      return val;
-    }
-
-    return perWorkerGauges.computeIfAbsent(metricName, name -> new GaugeCell(metricName));
-  }
-
   @Override
   public StringSet getStringSet(MetricName metricName) {
     return stringSets.get(metricName);
@@ -388,7 +372,8 @@ private void deleteStaleCounters(
   @VisibleForTesting
   Iterable<PerStepNamespaceMetrics> extractPerWorkerMetricUpdates() {
     ConcurrentHashMap<MetricName, Long> counters = new ConcurrentHashMap<MetricName, Long>();
-    ConcurrentHashMap<MetricName, Long> gauges = new ConcurrentHashMap<MetricName, Long>();
+    ConcurrentHashMap<MetricName, Long> per_worker_gauges =
+        new ConcurrentHashMap<MetricName, Long>();
     ConcurrentHashMap<MetricName, LockFreeHistogram.Snapshot> histograms =
         new ConcurrentHashMap<MetricName, LockFreeHistogram.Snapshot>();
     HashSet<MetricName> currentZeroValuedCounters = new HashSet<MetricName>();
@@ -402,12 +387,17 @@ Iterable<PerStepNamespaceMetrics> extractPerWorkerMetricUpdates() {
           counters.put(k, val);
         });
 
-    perWorkerGauges.forEach(
+    gauges.forEach(
         (k, v) -> {
-          Long val = v.getCumulative().value();
-          gauges.put(k, val);
-          v.reset();
+          // Check if metric name has the per worker label set
+          if (k.getLabels().containsKey("PER_WORKER_METRIC")
+              && k.getLabels().get("PER_WORKER_METRIC").equals("true")) {
+            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));
@@ -416,7 +406,7 @@ Iterable<PerStepNamespaceMetrics> extractPerWorkerMetricUpdates() {
     deleteStaleCounters(currentZeroValuedCounters, Instant.now(clock));
 
     return MetricsToPerStepNamespaceMetricsConverter.convert(
-        stepName, counters, gauges, histograms, parsedPerWorkerMetricsCache);
+        stepName, counters, per_worker_gauges, histograms, parsedPerWorkerMetricsCache);
   }
 
   /**

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

@@ -31,13 +31,15 @@
 import com.google.api.services.dataflow.model.CounterStructuredName;
 import com.google.api.services.dataflow.model.CounterStructuredNameAndMetadata;
 import com.google.api.services.dataflow.model.CounterUpdate;
+import com.google.api.services.dataflow.model.DataflowGaugeValue;
 import com.google.api.services.dataflow.model.DataflowHistogramValue;
 import com.google.api.services.dataflow.model.DistributionUpdate;
 import com.google.api.services.dataflow.model.IntegerGauge;
 import com.google.api.services.dataflow.model.Linear;
 import com.google.api.services.dataflow.model.MetricValue;
 import com.google.api.services.dataflow.model.PerStepNamespaceMetrics;
 import com.google.api.services.dataflow.model.StringList;
+import com.google.common.collect.ImmutableMap;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
@@ -471,7 +473,7 @@ public void testExtractPerWorkerMetricUpdates_populatedMetrics() {
   }
 
   @Test
-  public void testExtractPerWorkerMetricUpdatesKafka_populatedMetrics() {
+  public void testExtractPerWorkerMetricUpdatesKafka_populatedHistogramMetrics() {
     StreamingStepMetricsContainer.setEnablePerWorkerMetrics(true);
 
     MetricName histogramMetricName = MetricName.named("KafkaSink", "histogram");
@@ -508,6 +510,48 @@ public void testExtractPerWorkerMetricUpdatesKafka_populatedMetrics() {
     assertThat(updates, containsInAnyOrder(histograms));
   }
 
+  @Test
+  public void testExtractPerWorkerMetricUpdatesKafka_populateGaugeMetrics() {
+    StreamingStepMetricsContainer.setEnablePerWorkerMetrics(true);
+
+    MetricName gaugeMetricName =
+        MetricName.named("KafkaSink", "gauge", ImmutableMap.of("PER_WORKER_METRIC", "true"));
+    c2.getGauge(gaugeMetricName).set(5L);
+
+    Iterable<PerStepNamespaceMetrics> updates =
+        StreamingStepMetricsContainer.extractPerWorkerMetricUpdates(registry);
+
+    DataflowGaugeValue gaugeValue = new DataflowGaugeValue();
+    gaugeValue.setValue(5L);
+
+    MetricValue expectedGauge =
+        new MetricValue()
+            .setMetric("gauge")
+            .setMetricLabels(new HashMap<>())
+            .setValueGauge64(gaugeValue);
+
+    PerStepNamespaceMetrics gauge =
+        new PerStepNamespaceMetrics()
+            .setOriginalStep("s2")
+            .setMetricsNamespace("KafkaSink")
+            .setMetricValues(Collections.singletonList(expectedGauge));
+
+    assertThat(updates, containsInAnyOrder(gauge));
+  }
+
+  @Test
+  public void testExtractPerWorkerMetricUpdatesKafka_gaugeMetricsDropped() {
+    StreamingStepMetricsContainer.setEnablePerWorkerMetrics(true);
+
+    MetricName gaugeMetricName =
+        MetricName.named("KafkaSink", "gauge", ImmutableMap.of("PER_WORKER_METRIC", "false"));
+    c2.getGauge(gaugeMetricName).set(5L);
+
+    Iterable<PerStepNamespaceMetrics> updates =
+        StreamingStepMetricsContainer.extractPerWorkerMetricUpdates(registry);
+    assertThat(updates, IsEmptyIterable.emptyIterable());
+  }
+
   @Test
   public void testExtractPerWorkerMetricUpdates_emptyMetrics() {
     StreamingStepMetricsContainer.setEnablePerWorkerMetrics(true);

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

@@ -23,9 +23,9 @@
 /** Implementation of {@link Gauge} that delegates to the instance for the current context. */
 @Internal
 public class DelegatingGauge implements Metric, Gauge, Serializable {
+
   private final MetricName name;
   private final boolean processWideContainer;
-  private final boolean perWorkerGauge;
 
   /**
    * Create a {@code DelegatingGauge} with {@code perWorkerGauge} and {@code processWideContainer}
@@ -34,30 +34,17 @@ public class DelegatingGauge implements Metric, Gauge, Serializable {
    * @param name Metric name for this metric.
    */
   public DelegatingGauge(MetricName name) {
-    this(name, false, false);
-  }
-
-  /**
-   * Create a {@code DelegatingGauge} with {@code perWorkerGauge} set to false.
-   *
-   * @param name Metric name for this metric.
-   * @param processWideContainer Whether this Gauge is stored in the ProcessWide container or the
-   *     current thread's container.
-   */
-  public DelegatingGauge(MetricName name, boolean processWideContainer) {
-    this(name, processWideContainer, false);
+    this(name, false);
   }
 
   /**
    * @param name Metric name for this metric.
    * @param processWideContainer Whether this gauge is stored in the ProcessWide container or the
    *     current thread's container.
-   * @param perWorkerGauge Whether this gauge refers to a perWorker metric or not.
    */
-  public DelegatingGauge(MetricName name, boolean processWideContainer, boolean perWorkerGauge) {
+  public DelegatingGauge(MetricName name, boolean processWideContainer) {
     this.name = name;
     this.processWideContainer = processWideContainer;
-    this.perWorkerGauge = perWorkerGauge;
   }
 
   /** Set the gauge. */
@@ -70,11 +57,7 @@ public void set(long n) {
     if (container == null) {
       return;
     }
-    if (perWorkerGauge) {
-      container.getPerWorkerGauge(name).set(n);
-    } else {
-      container.getGauge(name).set(n);
-    }
+    container.getGauge(name).set(n);
   }
 
   @Override

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

@@ -81,14 +81,6 @@ default Histogram getPerWorkerHistogram(
     return NoOpHistogram.getInstance();
   }
 
-  /**
-   * Return the {@link Gauge} that should be used for implementing the given per-worker {@code
-   * metricName} in this container.
-   */
-  default Gauge getPerWorkerGauge(MetricName metricName) {
-    return NoOpGauge.getInstance();
-  }
-
   /** Return the cumulative values for any metrics in this container as MonitoringInfos. */
   default Iterable<MetricsApi.MonitoringInfo> getMonitoringInfos() {
     throw new RuntimeException("getMonitoringInfos is not implemented on this MetricsContainer.");

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaMetrics.java

@@ -33,18 +33,13 @@
 /** Stores and exports metrics for a batch of Kafka Client RPCs. */
 public interface KafkaMetrics {
 
-  /*Used to update latency metrics, which is later used to update metrics container in another thread*/
   void updateSuccessfulRpcMetrics(String topic, Duration elapsedTime);
 
-  /*Used to update backlog metrics in current thread*/
   void updateBacklogBytes(String topic, int partitionId, long backlog);
 
   /*Used to update all metrics in container*/
   void updateKafkaMetrics();
 
-  /*Used to update backlog metrics, which is later used to update metrics container in another thread*/
-  // void recordBacklogBytes(String topic, int partitionId, long backlog);
-
   /** No-op implementation of {@code KafkaResults}. */
   class NoOpKafkaMetrics implements KafkaMetrics {
     private NoOpKafkaMetrics() {}
@@ -58,9 +53,6 @@ public void updateBacklogBytes(String topic, int partitionId, long backlog) {}
     @Override
     public void updateKafkaMetrics() {}
 
-    // @Override
-    // public void recordBacklogBytes(String topic, int partitionId, long backlog) {};
-
     private static NoOpKafkaMetrics singleton = new NoOpKafkaMetrics();
 
     static NoOpKafkaMetrics getInstance() {
@@ -132,7 +124,6 @@ public void updateSuccessfulRpcMetrics(String topic, Duration elapsedTime) {
     @Override
     public void updateBacklogBytes(String topicName, int partitionId, long backlog) {
       if (isWritable().get()) {
-        // Key by MetricName so info isn't lost
         MetricName metricName = KafkaSinkMetrics.getMetricGaugeName(topicName, partitionId);
         perTopicPartitionBacklogs().put(metricName, backlog);
       }
@@ -162,33 +153,11 @@ private void recordRpcLatencyMetrics() {
     /** This is for creating gauges from backlog bytes recorded previously. */
     private void recordBacklogBytesInternal() {
       for (Map.Entry<MetricName, Long> backlog : perTopicPartitionBacklogs().entrySet()) {
-        // MetricName perPartionBacklogName = KafkaSinkMetrics.getMetricGaugeName(topicName,
-        // partitionId);
-        // Gauge perPartition =
-        //       Metrics.gauge(KafkaSinkMetrics.getMetricGaugeName(topicName, partitionId));
-        // Use lambda for more readability?
-        // map.forEach((key, value) -> System.out.println(key + ": " + value));
         Gauge gauge = KafkaSinkMetrics.createBacklogGauge(backlog.getKey());
         gauge.set(backlog.getValue());
       }
     }
 
-    /**
-     * This is for recording backlog bytes on the current thread.
-     *
-     * @param topicName topicName
-     * @param partitionId partitionId for the topic Only included in the metric key if
-     *     'supportsMetricsDeletion' is enabled.
-     * @param backlogBytes backlog for the topic Only included in the metric key if
-     *     'supportsMetricsDeletion' is enabled.
-     */
-    // @Override
-    // public void recordBacklogBytes(String topicName, int partitionId, long backlogBytes) {
-    //   Gauge perPartition =
-    //       Metrics.gauge(KafkaSinkMetrics.getMetricGaugeName(topicName, partitionId));
-    //   perPartition.set(backlogBytes);
-    // }
-
     /**
      * Export all metrics recorded in this instance to the underlying {@code perWorkerMetrics}
      * containers. This function will only report metrics once per instance. Subsequent calls to

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSinkMetrics.java

@@ -24,6 +24,7 @@
 import org.apache.beam.sdk.metrics.Histogram;
 import org.apache.beam.sdk.metrics.LabeledMetricNameUtils;
 import org.apache.beam.sdk.metrics.MetricName;
+import org.apache.beam.sdk.metrics.NoOpGauge;
 import org.apache.beam.sdk.util.HistogramData;
 
 /**
@@ -76,19 +77,6 @@ public static Histogram createRPCLatencyHistogram(RpcMethod method, String topic
     return new DelegatingHistogram(metricName, buckets, false, true);
   }
 
-  /**
-   * Creates a {@link Gauge} metric to record per partition backlog with the name
-   *
-   * <p>'EstimatedBacklogSize*topic_name:{topic};partitionId:{partitionId};'.
-   *
-   * @param topic Kafka topic associated with this metric.
-   * @param partitionId partition id associated with this metric.
-   * @return Counter.
-   */
-  public static Gauge createBacklogGauge(String topic, int partitionId) {
-    return new DelegatingGauge(getMetricGaugeName(topic, partitionId), false, true);
-  }
-
   /**
    * Creates a {@link Gauge} metric to record per partition backlog with the name
    *
@@ -98,24 +86,14 @@ public static Gauge createBacklogGauge(String topic, int partitionId) {
    * @return Counter.
    */
   public static Gauge createBacklogGauge(MetricName name) {
-    // use label to differenciate between the type of gauge metric is created
-    // TODO(bug to clean this to make more consistent between the two runners)
-    // test if set to false, this doesn't occur
-    // && name.getLabels().get(MonitoringInfoConstants.Labels.PER_WORKER_METRIC)
+    // Use label to differenciate between the type of gauge metric is created
+    // TODO(34195):
     if (name.getLabels().containsKey(MonitoringInfoConstants.Labels.PER_WORKER_METRIC)
         && name.getLabels().get(MonitoringInfoConstants.Labels.PER_WORKER_METRIC).equals("true")) {
-      // return Metrics.gauge(name); // can this be a delgating gauge of false type?
-      // metric name always exists, so need a way to handle them differently
-      // second bollean should be false for u2 path, how to not add label just yet?
-      // for legacy
-      // return new DelegatingGauge(name, false, true);
-
-      // for runner v2
-      // investigate why it gest ton UW container, but not to DFE
-      // return Metrics.gauge(name);
-      return new DelegatingGauge(name, false, false);
+      return new DelegatingGauge(name, false);
     } else {
-      return new DelegatingGauge(name, false, true);
+      // Currently KafkaSink metrics only supports aggregated per worker metrics.
+      return NoOpGauge.getInstance();
     }
   }
 

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java

@@ -559,8 +559,6 @@ public ProcessContinuation processElement(
                         .doubleValue()
                     * avgRecordSize.estimateRecordByteSizeToOffsetCountRatio()));
         KafkaMetrics kafkaResults = KafkaSinkMetrics.kafkaMetrics();
-        // public void updateBacklogBytes(String topicName, int partitionId, long backlog) {
-
         kafkaResults.updateBacklogBytes(
             kafkaSourceDescriptor.getTopic(),
             kafkaSourceDescriptor.getPartition(),
@@ -571,9 +569,6 @@ public ProcessContinuation processElement(
                         .subtract(BigDecimal.valueOf(expectedOffset), MathContext.DECIMAL128)
                         .doubleValue()
                     * avgRecordSize.estimateRecordByteSizeToOffsetCountRatio()));
-        // flush imediately afterwards, can it be moved to finish bundle?
-        // create the same kind of metric? This creates a per Worker metric, we dont want that to be
-        // reusable between the two ios
         kafkaResults.updateKafkaMetrics();
       }
     }