Address Sam's comments

[Dataflow Streaming] Use isolated windmill streams based on job settings (apache#32503)

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

@@ -19,7 +19,6 @@
 
 import com.google.auto.value.AutoValue;
 import java.time.Duration;
-import java.time.Instant;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ConcurrentLinkedQueue;
@@ -30,7 +29,7 @@
 /** Stores and exports metrics for a batch of Kafka Client RPCs. */
 public interface KafkaMetrics {
 
-  void updateSuccessfulRpcMetrics(String topic, Instant start, Instant end);
+  void updateSuccessfulRpcMetrics(String topic, Duration elapsedTime);
 
   void updateKafkaMetrics();
 
@@ -39,7 +38,7 @@ class NoOpKafkaMetrics implements KafkaMetrics {
     private NoOpKafkaMetrics() {}
 
     @Override
-    public void updateSuccessfulRpcMetrics(String topic, Instant start, Instant end) {}
+    public void updateSuccessfulRpcMetrics(String topic, Duration elapsedTime) {}
 
     @Override
     public void updateKafkaMetrics() {}
@@ -77,15 +76,15 @@ public static KafkaMetricsImpl create() {
 
     /** Record the rpc status and latency of a successful Kafka poll RPC call. */
     @Override
-    public void updateSuccessfulRpcMetrics(String topic, Instant start, Instant end) {
+    public void updateSuccessfulRpcMetrics(String topic, Duration elapsedTime) {
       if (isWritable().get()) {
         ConcurrentLinkedQueue<Duration> latencies = perTopicRpcLatencies().get(topic);
         if (latencies == null) {
           latencies = new ConcurrentLinkedQueue<Duration>();
-          latencies.add(Duration.between(start, end));
+          latencies.add(elapsedTime);
           perTopicRpcLatencies().put(topic, latencies);
         } else {
-          latencies.add(Duration.between(start, end));
+          latencies.add(elapsedTime);
         }
       }
     }

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

@@ -53,6 +53,7 @@
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.util.Preconditions;
 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.Stopwatch;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.Closeables;
@@ -407,6 +408,8 @@ public long getSplitBacklogBytes() {
 
   // Created in each next batch, and updated at the end.
   public KafkaMetrics kafkaResults = KafkaSinkMetrics.kafkaMetrics();
+  private Stopwatch stopwatch = Stopwatch.createUnstarted();
+  private String kafkaTopic = "";
 
   @Override
   public String toString() {
@@ -518,6 +521,13 @@ String name() {
 
     List<TopicPartition> partitions =
         Preconditions.checkArgumentNotNull(source.getSpec().getTopicPartitions());
+
+    // Each source has a single unique topic.
+    for (TopicPartition topicPartition : partitions) {
+      this.kafkaTopic = topicPartition.topic();
+      break;
+    }
+
     List<PartitionState<K, V>> states = new ArrayList<>(partitions.size());
 
     if (checkpointMark != null) {
@@ -577,21 +587,12 @@ private void consumerPollLoop() {
       while (!closed.get()) {
         try {
           if (records.isEmpty()) {
-            // Each source has a single unique topic.
-            List<TopicPartition> topicPartitions = source.getSpec().getTopicPartitions();
-            Preconditions.checkStateNotNull(topicPartitions);
-            String topicName = "null"; // value will be overridden
-            for (TopicPartition topicPartition : topicPartitions) {
-              topicName = topicPartition.topic();
-              break;
-            }
-
-            java.time.Instant operationStartTime = java.time.Instant.now();
-            records = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
-            java.time.Instant operationEndTime = java.time.Instant.now();
 
+            stopwatch.start();
+            records = consumer.poll(KAFKA_POLL_TIMEOUT.getMillis());
+            stopwatch.stop();
             kafkaResults.updateSuccessfulRpcMetrics(
-                topicName, operationStartTime, operationEndTime);
+                kafkaTopic, java.time.Duration.ofMillis(stopwatch.elapsed(TimeUnit.MILLISECONDS)));
 
           } else if (availableRecordsQueue.offer(
               records, RECORDS_ENQUEUE_POLL_TIMEOUT.getMillis(), TimeUnit.MILLISECONDS)) {
@@ -616,6 +617,7 @@ private void consumerPollLoop() {
 
   private void commitCheckpointMark() {
     KafkaCheckpointMark checkpointMark = finalizedCheckpointMark.getAndSet(null);
+
     if (checkpointMark != null) {
       LOG.debug("{}: Committing finalized checkpoint {}", this, checkpointMark);
       Consumer<byte[], byte[]> consumer = Preconditions.checkStateNotNull(this.consumer);

sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaMetricsTest.java

@@ -22,7 +22,6 @@
 import static org.hamcrest.Matchers.equalTo;
 
 import java.time.Duration;
-import java.time.Instant;
 import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 import org.apache.beam.runners.core.metrics.MetricsContainerImpl;
@@ -67,7 +66,6 @@ public TestMetricsContainer() {
     @Override
     public Histogram getPerWorkerHistogram(
         MetricName metricName, HistogramData.BucketType bucketType) {
-      System.out.println("xxx metricName " + metricName);
       perWorkerHistograms.computeIfAbsent(KV.of(metricName, bucketType), kv -> new TestHistogram());
       return perWorkerHistograms.get(KV.of(metricName, bucketType));
     }
@@ -83,9 +81,8 @@ public void testNoOpKafkaMetrics() throws Exception {
     TestMetricsContainer testContainer = new TestMetricsContainer();
     MetricsEnvironment.setCurrentContainer(testContainer);
 
-    Instant t1 = Instant.now();
     KafkaMetrics results = KafkaMetrics.NoOpKafkaMetrics.getInstance();
-    results.updateSuccessfulRpcMetrics("test-topic", t1, t1.plus(Duration.ofMillis(10)));
+    results.updateSuccessfulRpcMetrics("test-topic", Duration.ofMillis(10));
 
     results.updateKafkaMetrics();
 
@@ -97,10 +94,9 @@ public void testKafkaRPCLatencyMetrics() throws Exception {
     TestMetricsContainer testContainer = new TestMetricsContainer();
     MetricsEnvironment.setCurrentContainer(testContainer);
 
-    Instant t1 = Instant.now();
     KafkaMetrics results = KafkaSinkMetrics.kafkaMetrics();
 
-    results.updateSuccessfulRpcMetrics("test-topic", t1, t1.plus(Duration.ofMillis(10)));
+    results.updateSuccessfulRpcMetrics("test-topic", Duration.ofMillis(10));
 
     results.updateKafkaMetrics();
     // RpcLatency*rpc_method:POLL;topic_name:test-topic