Merge pull request #36077: [KafkaIO] Only update size metrics once per batch

Author: kennknowles

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

@@ -168,18 +168,23 @@ private void setAvg(final double value) {
       AVG.lazySet(this, Double.doubleToRawLongBits(value));
     }
 
-    private long incrementAndGetNumUpdates() {
-      final long nextNumUpdates = Math.min(MOVING_AVG_WINDOW, numUpdates + 1);
-      numUpdates = nextNumUpdates;
-      return nextNumUpdates;
+    public void update(final double quantity) {
+      final double prevAvg = getAvg(); // volatile load (acquire)
+
+      final long nextNumUpdates = numUpdates + 1; // normal load
+      final double nextAvg = prevAvg + (quantity - prevAvg) / nextNumUpdates;
+
+      numUpdates = Math.min(MOVING_AVG_WINDOW, nextNumUpdates); // normal store
+      setAvg(nextAvg); // ordered store (release)
     }
 
-    public void update(final double quantity) {
+    public void update(final double sum, final long count) {
       final double prevAvg = getAvg(); // volatile load (acquire)
 
-      final long nextNumUpdates = incrementAndGetNumUpdates(); // normal load/store
-      final double nextAvg = prevAvg + (quantity - prevAvg) / nextNumUpdates; // normal load/store
+      final long nextNumUpdates = numUpdates + count; // normal load
+      final double nextAvg = prevAvg + (sum / count - prevAvg) * ((double) count / nextNumUpdates);
 
+      numUpdates = Math.min(MOVING_AVG_WINDOW, nextNumUpdates); // normal store
       setAvg(nextAvg); // ordered store (release)
     }
 

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

@@ -660,57 +660,70 @@ public ProcessContinuation processElement(
 
         // Visible progress within the consumer polling timeout.
         // Partially or fully claim and process records in this batch.
-        for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
-          if (!tracker.tryClaim(rawRecord.offset())) {
-            consumer.seek(topicPartition, rawRecord.offset());
-            consumer.pause(Collections.singleton(topicPartition));
+        long rawSizesSum = 0L;
+        long rawSizesCount = 0L;
+        long rawSizesMin = Long.MAX_VALUE;
+        long rawSizesMax = Long.MIN_VALUE;
+        try {
+          for (ConsumerRecord<byte[], byte[]> rawRecord : rawRecords) {
+            if (!tracker.tryClaim(rawRecord.offset())) {
+              consumer.seek(topicPartition, rawRecord.offset());
+              consumer.pause(Collections.singleton(topicPartition));
 
-            return ProcessContinuation.stop();
-          }
-          expectedOffset = rawRecord.offset() + 1;
-          try {
-            KafkaRecord<K, V> kafkaRecord =
-                new KafkaRecord<>(
-                    rawRecord.topic(),
-                    rawRecord.partition(),
-                    rawRecord.offset(),
-                    ConsumerSpEL.getRecordTimestamp(rawRecord),
-                    ConsumerSpEL.getRecordTimestampType(rawRecord),
-                    ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
-                    ConsumerSpEL.deserializeKey(keyDeserializerInstance, rawRecord),
-                    ConsumerSpEL.deserializeValue(valueDeserializerInstance, rawRecord));
-            int recordSize =
-                (rawRecord.key() == null ? 0 : rawRecord.key().length)
-                    + (rawRecord.value() == null ? 0 : rawRecord.value().length);
-            avgRecordSize.update(recordSize);
-            rawSizes.update(recordSize);
-            Instant outputTimestamp;
-            // The outputTimestamp and watermark will be computed by timestampPolicy, where the
-            // WatermarkEstimator should be a manual one.
-            if (timestampPolicy != null) {
-              TimestampPolicyContext context =
-                  updateWatermarkManually(timestampPolicy, watermarkEstimator, tracker);
-              outputTimestamp = timestampPolicy.getTimestampForRecord(context, kafkaRecord);
-            } else {
-              Preconditions.checkStateNotNull(this.extractOutputTimestampFn);
-              outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+              return ProcessContinuation.stop();
             }
-            receiver
-                .get(recordTag)
-                .outputWithTimestamp(KV.of(kafkaSourceDescriptor, kafkaRecord), outputTimestamp);
-          } catch (SerializationException e) {
-            // This exception should only occur during the key and value deserialization when
-            // creating the Kafka Record
-            badRecordRouter.route(
-                receiver,
-                rawRecord,
-                null,
-                e,
-                "Failure deserializing Key or Value of Kakfa record reading from Kafka");
-            if (timestampPolicy != null) {
-              updateWatermarkManually(timestampPolicy, watermarkEstimator, tracker);
+            expectedOffset = rawRecord.offset() + 1;
+            try {
+              KafkaRecord<K, V> kafkaRecord =
+                  new KafkaRecord<>(
+                      rawRecord.topic(),
+                      rawRecord.partition(),
+                      rawRecord.offset(),
+                      ConsumerSpEL.getRecordTimestamp(rawRecord),
+                      ConsumerSpEL.getRecordTimestampType(rawRecord),
+                      ConsumerSpEL.hasHeaders() ? rawRecord.headers() : null,
+                      ConsumerSpEL.deserializeKey(keyDeserializerInstance, rawRecord),
+                      ConsumerSpEL.deserializeValue(valueDeserializerInstance, rawRecord));
+              int recordSize =
+                  (rawRecord.key() == null ? 0 : rawRecord.key().length)
+                      + (rawRecord.value() == null ? 0 : rawRecord.value().length);
+              rawSizesSum = rawSizesSum + recordSize;
+              rawSizesCount = rawSizesCount + 1L;
+              rawSizesMin = Math.min(rawSizesMin, recordSize);
+              rawSizesMax = Math.max(rawSizesMax, recordSize);
+              Instant outputTimestamp;
+              // The outputTimestamp and watermark will be computed by timestampPolicy, where the
+              // WatermarkEstimator should be a manual one.
+              if (timestampPolicy != null) {
+                TimestampPolicyContext context =
+                    updateWatermarkManually(timestampPolicy, watermarkEstimator, tracker);
+                outputTimestamp = timestampPolicy.getTimestampForRecord(context, kafkaRecord);
+              } else {
+                Preconditions.checkStateNotNull(this.extractOutputTimestampFn);
+                outputTimestamp = extractOutputTimestampFn.apply(kafkaRecord);
+              }
+              receiver
+                  .get(recordTag)
+                  .outputWithTimestamp(KV.of(kafkaSourceDescriptor, kafkaRecord), outputTimestamp);
+            } catch (SerializationException e) {
+              // This exception should only occur during the key and value deserialization when
+              // creating the Kafka Record
+              badRecordRouter.route(
+                  receiver,
+                  rawRecord,
+                  null,
+                  e,
+                  "Failure deserializing Key or Value of Kakfa record reading from Kafka");
+              if (timestampPolicy != null) {
+                updateWatermarkManually(timestampPolicy, watermarkEstimator, tracker);
+              }
             }
           }
+        } finally {
+          if (rawSizesCount > 0L) {
+            avgRecordSize.update(rawSizesSum, rawSizesCount);
+            rawSizes.update(rawSizesSum, rawSizesCount, rawSizesMin, rawSizesMax);
+          }
         }
 
         // Non-visible progress within the consumer polling timeout.