[Flink] Throttle Flink metric updates from source readers

CHANGES.md

@@ -68,6 +68,7 @@
 
 ## New Features / Improvements
 
+* (Java) Flink runner: sources now forward Beam metrics to Flink at most once per second instead of after every record, which was a significant share of the CPU of reading. Metrics in the pipeline result are not affected ([#40607](https://github.com/apache/beam/issues/40607)).
 * (Go) Added `wait.On`, which delays each input window until the corresponding windows in its signal PCollections have closed ([#39909](https://github.com/apache/beam/issues/39909)).
 * (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)).
 * The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)).

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

@@ -19,6 +19,8 @@
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.concurrent.TimeUnit;
+import java.util.function.LongSupplier;
 import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
 import org.apache.beam.runners.flink.FlinkPipelineOptions;
 import org.apache.beam.sdk.io.Source;
@@ -27,29 +29,51 @@
 
 /**
  * Util for invoking {@link Source.Reader} methods that might require a {@link MetricsContainerImpl}
- * to be active. Source.Reader decorator which registers {@link MetricsContainerImpl}. It update
- * metrics to Flink metric and accumulator in start and advance.
+ * to be active. Source.Reader decorator which registers {@link MetricsContainerImpl}. It updates
+ * Flink metrics after start and advance.
+ *
+ * <p>Updating Flink metrics walks all the metrics of the task, which is expensive compared to
+ * reading a single record. Flink metrics are therefore updated at most once every {@link
+ * #FLINK_METRICS_UPDATE_INTERVAL_MILLIS}, on the first call, and whenever the reader has no more
+ * data. The metrics reported in the pipeline result are collected separately and are not affected.
  */
 public class ReaderInvocationUtil<OutputT, ReaderT extends Source.Reader<OutputT>> {
 
+  static final long FLINK_METRICS_UPDATE_INTERVAL_MILLIS = 1000;
+
+  private static final long FLINK_METRICS_UPDATE_INTERVAL_NANOS =
+      TimeUnit.MILLISECONDS.toNanos(FLINK_METRICS_UPDATE_INTERVAL_MILLIS);
+
   private final String stepName;
   private final FlinkMetricContainerBase container;
   private final Boolean enableMetrics;
+  private final LongSupplier nanoTime;
+  private boolean metricsUpdated;
+  private long lastMetricsUpdateNanos;
 
   public ReaderInvocationUtil(
       String stepName, PipelineOptions options, FlinkMetricContainerBase container) {
+    this(stepName, options, container, System::nanoTime);
+  }
+
+  ReaderInvocationUtil(
+      String stepName,
+      PipelineOptions options,
+      FlinkMetricContainerBase container,
+      LongSupplier nanoTime) {
     FlinkPipelineOptions flinkPipelineOptions = options.as(FlinkPipelineOptions.class);
     this.stepName = stepName;
     this.enableMetrics = !flinkPipelineOptions.getDisableMetrics();
     this.container = container;
+    this.nanoTime = nanoTime;
   }
 
   public boolean invokeStart(ReaderT reader) throws IOException {
     if (enableMetrics) {
       try (Closeable ignored =
           MetricsEnvironment.scopedMetricsContainer(container.getMetricsContainer(stepName))) {
         boolean result = reader.start();
-        container.updateMetrics(stepName);
+        maybeUpdateMetrics(result);
         return result;
       }
     } else {
@@ -62,11 +86,24 @@ public boolean invokeAdvance(ReaderT reader) throws IOException {
       try (Closeable ignored =
           MetricsEnvironment.scopedMetricsContainer(container.getMetricsContainer(stepName))) {
         boolean result = reader.advance();
-        container.updateMetrics(stepName);
+        maybeUpdateMetrics(result);
         return result;
       }
     } else {
       return reader.advance();
     }
   }
+
+  private void maybeUpdateMetrics(boolean dataAvailable) {
+    long now = nanoTime.getAsLong();
+    // Always update on the first call, so that the metrics are registered in Flink right away, and
+    // when the reader has no data, so that the final values of a finished reader are reported.
+    if (!metricsUpdated
+        || !dataAvailable
+        || now - lastMetricsUpdateNanos >= FLINK_METRICS_UPDATE_INTERVAL_NANOS) {
+      container.updateMetrics(stepName);
+      metricsUpdated = true;
+      lastMetricsUpdateNanos = now;
+    }
+  }
 }

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

@@ -0,0 +1,115 @@
+/*
+ * 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 static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
+import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import org.apache.beam.sdk.io.Source;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.junit.Before;
+import org.junit.Test;
+
+/** Tests {@link ReaderInvocationUtil}. */
+public class ReaderInvocationUtilTest {
+
+  private static final String STEP = "step";
+  private static final long INTERVAL_NANOS =
+      TimeUnit.MILLISECONDS.toNanos(ReaderInvocationUtil.FLINK_METRICS_UPDATE_INTERVAL_MILLIS);
+
+  private final AtomicLong clock = new AtomicLong();
+  private FlinkMetricContainerBase container;
+  private Source.Reader<String> reader;
+
+  @Before
+  @SuppressWarnings("unchecked")
+  public void setUp() {
+    container = mock(FlinkMetricContainerBase.class);
+    when(container.getMetricsContainer(STEP)).thenReturn(new MetricsContainerImpl(STEP));
+    reader = mock(Source.Reader.class);
+  }
+
+  private ReaderInvocationUtil<String, Source.Reader<String>> invoker(boolean disableMetrics) {
+    FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);
+    options.setDisableMetrics(disableMetrics);
+    return new ReaderInvocationUtil<>(STEP, options, container, clock::get);
+  }
+
+  @Test
+  public void updatesFlinkMetricsAtMostOncePerInterval() throws IOException {
+    when(reader.advance()).thenReturn(true);
+    ReaderInvocationUtil<String, Source.Reader<String>> invoker = invoker(false);
+
+    // The first call registers the metrics in Flink.
+    invoker.invokeAdvance(reader);
+    verify(container, times(1)).updateMetrics(STEP);
+
+    for (int i = 0; i < 1000; i++) {
+      invoker.invokeAdvance(reader);
+    }
+    verify(container, times(1)).updateMetrics(STEP);
+
+    clock.addAndGet(INTERVAL_NANOS);
+    invoker.invokeAdvance(reader);
+    invoker.invokeAdvance(reader);
+    verify(container, times(2)).updateMetrics(STEP);
+
+    clock.addAndGet(INTERVAL_NANOS);
+    invoker.invokeAdvance(reader);
+    verify(container, times(3)).updateMetrics(STEP);
+  }
+
+  @Test
+  public void updatesFlinkMetricsWhenTheReaderHasNoData() throws IOException {
+    when(reader.start()).thenReturn(true);
+    when(reader.advance()).thenReturn(true, true, false);
+    ReaderInvocationUtil<String, Source.Reader<String>> invoker = invoker(false);
+
+    invoker.invokeStart(reader);
+    verify(container, times(1)).updateMetrics(STEP);
+
+    invoker.invokeAdvance(reader);
+    invoker.invokeAdvance(reader);
+    verify(container, times(1)).updateMetrics(STEP);
+
+    invoker.invokeAdvance(reader);
+    verify(container, times(2)).updateMetrics(STEP);
+  }
+
+  @Test
+  public void doesNotUpdateFlinkMetricsWhenDisabled() throws IOException {
+    when(reader.start()).thenReturn(false);
+    when(reader.advance()).thenReturn(false);
+    ReaderInvocationUtil<String, Source.Reader<String>> invoker = invoker(true);
+
+    invoker.invokeStart(reader);
+    clock.addAndGet(INTERVAL_NANOS);
+    invoker.invokeAdvance(reader);
+
+    verify(container, never()).updateMetrics(STEP);
+    verify(container, never()).getMetricsContainer(STEP);
+  }
+}