Publish process-wide metrics in Flink runner

runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/FlinkMetricContainerBase.java

@@ -32,6 +32,8 @@
 import org.apache.beam.sdk.metrics.MetricQueryResults;
 import org.apache.beam.sdk.metrics.MetricResult;
 import org.apache.beam.sdk.metrics.MetricResults;
+import org.apache.beam.sdk.metrics.MetricsContainer;
+import org.apache.beam.sdk.metrics.MetricsEnvironment;
 import org.apache.beam.sdk.metrics.MetricsFilter;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.flink.configuration.GlobalConfiguration;
@@ -49,12 +51,15 @@
  */
 abstract class FlinkMetricContainerBase {
 
+  private static final String PROCESS_WIDE_METRICS_STEP_NAME = "__process_wide_metrics";
+
   private static final String METRIC_KEY_SEPARATOR =
       GlobalConfiguration.loadConfiguration()
           .getOptional(MetricOptions.SCOPE_DELIMITER)
           .orElseGet(MetricOptions.SCOPE_DELIMITER::defaultValue);
 
   protected final MetricsContainerStepMap metricsContainers;
+  private final MetricsContainerImpl processWideMetricsContainer;
   private final Map<String, Counter> flinkCounterCache;
   private final Map<String, FlinkDistributionGauge> flinkDistributionGaugeCache;
   private final Map<String, FlinkGauge> flinkGaugeCache;
@@ -65,6 +70,8 @@ public FlinkMetricContainerBase(MetricGroup metricGroup) {
     this.flinkDistributionGaugeCache = new HashMap<>();
     this.flinkGaugeCache = new HashMap<>();
     this.metricsContainers = new MetricsContainerStepMap();
+    this.processWideMetricsContainer = MetricsContainerImpl.createProcessWideContainer();
+    GlobalMetricsUtils.setProcessWideMetricsContainer(processWideMetricsContainer);
     this.metricGroup = metricGroup;
   }
 
@@ -96,6 +103,23 @@ void updateMetrics(String stepName) {
     updateCounters(metricQueryResults.getCounters());
     updateDistributions(metricQueryResults.getDistributions());
     updateGauge(metricQueryResults.getGauges());
+    updateProcessWideMetrics();
+  }
+
+  private void updateProcessWideMetrics() {
+    MetricsContainer processWideContainer = MetricsEnvironment.getProcessWideContainer();
+    if (processWideContainer != processWideMetricsContainer) {
+      return;
+    }
+
+    MetricsContainerStepMap processWideMetrics = new MetricsContainerStepMap();
+    processWideMetrics.update(PROCESS_WIDE_METRICS_STEP_NAME, processWideMetricsContainer);
+    MetricQueryResults metricQueryResults =
+        asAttemptedOnlyMetricResults(processWideMetrics)
+            .queryMetrics(MetricsFilter.builder().addStep(PROCESS_WIDE_METRICS_STEP_NAME).build());
+    updateCounters(metricQueryResults.getCounters());
+    updateDistributions(metricQueryResults.getDistributions());
+    updateGauge(metricQueryResults.getGauges());
   }
 
   private void updateCounters(Iterable<MetricResult<Long>> counters) {

runners/flink/src/main/java/org/apache/beam/runners/flink/metrics/GlobalMetricsUtils.java

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink.metrics;
+
+import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
+import org.apache.beam.sdk.metrics.MetricsEnvironment;
+
+/** Utility class for managing process-wide Beam metrics in a Flink environment. */
+final class GlobalMetricsUtils {
+  private GlobalMetricsUtils() {}
+
+  static synchronized void setProcessWideMetricsContainer(MetricsContainerImpl container) {
+    if (MetricsEnvironment.getProcessWideContainer() == null) {
+      MetricsEnvironment.setProcessWideContainer(container);
+    }
+  }
+}

runners/flink/src/test/java/org/apache/beam/runners/flink/metrics/FlinkMetricContainerTest.java

@@ -25,6 +25,7 @@
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -39,19 +40,26 @@
 import org.apache.beam.runners.core.metrics.SimpleMonitoringInfoBuilder;
 import org.apache.beam.runners.flink.metrics.FlinkMetricContainerBase.FlinkDistributionGauge;
 import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.DelegatingCounter;
+import org.apache.beam.sdk.metrics.DelegatingDistribution;
+import org.apache.beam.sdk.metrics.DelegatingGauge;
 import org.apache.beam.sdk.metrics.Distribution;
 import org.apache.beam.sdk.metrics.DistributionResult;
 import org.apache.beam.sdk.metrics.Gauge;
 import org.apache.beam.sdk.metrics.GaugeResult;
 import org.apache.beam.sdk.metrics.MetricKey;
 import org.apache.beam.sdk.metrics.MetricName;
+import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.metrics.MetricsContainer;
+import org.apache.beam.sdk.metrics.MetricsEnvironment;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.flink.api.common.functions.RuntimeContext;
 import org.apache.flink.metrics.SimpleCounter;
 import org.apache.flink.metrics.groups.OperatorMetricGroup;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatcher;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
@@ -74,6 +82,11 @@ public void beforeTest() {
     container = new FlinkMetricContainer(runtimeContext);
   }
 
+  @After
+  public void afterTest() {
+    MetricsEnvironment.setProcessWideContainer(null);
+  }
+
   @Test
   public void testMetricNameGeneration() {
     MetricKey key = MetricKey.create("step", MetricName.named("namespace", "name"));
@@ -97,6 +110,91 @@ public void testCounter() {
     assertThat(flinkCounter.getCount(), is(2L));
   }
 
+  @Test
+  public void testProcessWideMetrics() {
+    SimpleCounter flinkCounter = new SimpleCounter();
+    when(metricGroup.counter("global.counter")).thenReturn(flinkCounter);
+    when(metricGroup.gauge(eq("global.gauge"), any()))
+        .thenAnswer(invocation -> invocation.getArgument(1));
+    when(metricGroup.gauge(eq("global.distribution"), any()))
+        .thenAnswer(invocation -> invocation.getArgument(1));
+
+    Counter counter = new DelegatingCounter(MetricName.named("global", "counter"), true);
+    Gauge gauge = new DelegatingGauge(MetricName.named("global", "gauge"), true);
+    Distribution distribution =
+        new DelegatingDistribution(MetricName.named("global", "distribution"), true);
+    counter.inc(5);
+    gauge.set(3);
+    distribution.update(11);
+    distribution.update(13);
+    MetricsContainerImpl processWideContainer =
+        (MetricsContainerImpl) MetricsEnvironment.getProcessWideContainer();
+    assertThat(
+        processWideContainer
+            .getGauge(MetricName.named("global", "gauge"))
+            .getCumulative()
+            .extractResult()
+            .getValue(),
+        is(3L));
+
+    container.updateMetrics("step");
+    assertThat(flinkCounter.getCount(), is(5L));
+    ArgumentCaptor<FlinkMetricContainerBase.FlinkGauge> gaugeCaptor =
+        ArgumentCaptor.forClass(FlinkMetricContainerBase.FlinkGauge.class);
+    ArgumentCaptor<FlinkDistributionGauge> distributionCaptor =
+        ArgumentCaptor.forClass(FlinkDistributionGauge.class);
+    verify(metricGroup).gauge(eq("global.gauge"), gaugeCaptor.capture());
+    verify(metricGroup).gauge(eq("global.distribution"), distributionCaptor.capture());
+    FlinkMetricContainerBase.FlinkGauge flinkGauge = gaugeCaptor.getValue();
+    FlinkDistributionGauge flinkDistributionGauge = distributionCaptor.getValue();
+    assertThat(flinkGauge.getValue(), is(3L));
+    assertThat(flinkDistributionGauge.getValue(), is(DistributionResult.create(24, 2, 11, 13)));
+
+    container.updateMetrics("step");
+    assertThat(flinkCounter.getCount(), is(5L));
+    assertThat(flinkGauge.getValue(), is(3L));
+    assertThat(flinkDistributionGauge.getValue(), is(DistributionResult.create(24, 2, 11, 13)));
+
+    counter.inc(2);
+    gauge.set(4);
+    distribution.update(17);
+    container.updateMetrics("step");
+    assertThat(flinkCounter.getCount(), is(7L));
+    assertThat(flinkGauge.getValue(), is(4L));
+    assertThat(flinkDistributionGauge.getValue(), is(DistributionResult.create(41, 3, 11, 17)));
+  }
+
+  @Test
+  public void testProcessWideMetricsArePublishedOnlyByOwningContainer() {
+    RuntimeContext otherRuntimeContext = org.mockito.Mockito.mock(RuntimeContext.class);
+    OperatorMetricGroup otherMetricGroup = org.mockito.Mockito.mock(OperatorMetricGroup.class);
+    when(otherRuntimeContext.<MetricsContainerStepMap, MetricsContainerStepMap>getAccumulator(
+            anyString()))
+        .thenReturn(new MetricsAccumulator());
+    when(otherRuntimeContext.getMetricGroup()).thenReturn(otherMetricGroup);
+    FlinkMetricContainer otherContainer = new FlinkMetricContainer(otherRuntimeContext);
+
+    Counter counter = new DelegatingCounter(MetricName.named("global", "counter"), true);
+    counter.inc(5);
+
+    otherContainer.updateMetrics("otherStep");
+    verify(otherMetricGroup, never()).counter("global.counter");
+
+    SimpleCounter flinkCounter = new SimpleCounter();
+    when(metricGroup.counter("global.counter")).thenReturn(flinkCounter);
+    container.updateMetrics("step");
+    assertThat(flinkCounter.getCount(), is(5L));
+  }
+
+  @Test
+  public void testUnscopedThreadMetricsAreNotPublishedAsProcessWideMetrics() {
+    Metrics.counter("global", "unscoped").inc(5);
+
+    container.updateMetrics("step");
+
+    verify(metricGroup, never()).counter("global.unscoped");
+  }
+
   @Test
   public void testGauge() {
     FlinkMetricContainerBase.FlinkGauge flinkGauge =