Add Backlog Metrics to Kafka Splittable DoFn Implementation (#31281)

Naireen · Naireen · web-flow · commit fd4368f1c4ab · 2024-05-28T13:58:16.000-04:00
* Set backlog in gauge metric

* add backlog metrics to splittable dofn Kafka read implementatino

---------

Co-authored-by: Naireen &lt;naireenhussain@google.com&gt;
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
@@ -32,6 +32,7 @@
 import org.apache.beam.sdk.io.kafka.KafkaUnboundedReader.TimestampPolicyContext;
 import org.apache.beam.sdk.io.range.OffsetRange;
 import org.apache.beam.sdk.metrics.Distribution;
+import org.apache.beam.sdk.metrics.Gauge;
 import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.SerializableFunction;
@@ -222,6 +223,9 @@ private ReadFromKafkaDoFn(
 
   private transient @Nullable LoadingCache<TopicPartition, AverageRecordSize> avgRecordSize;
   private static final long DEFAULT_KAFKA_POLL_TIMEOUT = 2L;
+
+  private HashMap<String, Long> perPartitionBacklogMetrics = new HashMap<String, Long>();;
+
   @VisibleForTesting final long consumerPollingTimeout;
   @VisibleForTesting final DeserializerProvider<K> keyDeserializerProvider;
   @VisibleForTesting final DeserializerProvider<V> valueDeserializerProvider;
@@ -342,6 +346,13 @@ public double getSize(
     if (!avgRecordSize.asMap().containsKey(kafkaSourceDescriptor.getTopicPartition())) {
       return numRecords;
     }
+    if (offsetEstimatorCache != null) {
+      for (Map.Entry<TopicPartition, KafkaLatestOffsetEstimator> tp :
+          offsetEstimatorCache.entrySet()) {
+        perPartitionBacklogMetrics.put(tp.getKey().toString(), tp.getValue().estimate());
+      }
+    }
+
     return avgRecordSize.get(kafkaSourceDescriptor.getTopicPartition()).getTotalSize(numRecords);
   }
 
@@ -394,6 +405,13 @@ public ProcessContinuation processElement(
         Metrics.distribution(
             METRIC_NAMESPACE,
             RAW_SIZE_METRIC_PREFIX + kafkaSourceDescriptor.getTopicPartition().toString());
+    for (Map.Entry<String, Long> backlogSplit : perPartitionBacklogMetrics.entrySet()) {
+      Gauge backlog =
+          Metrics.gauge(
+              METRIC_NAMESPACE, RAW_SIZE_METRIC_PREFIX + "backlogBytes_" + backlogSplit.getKey());
+      backlog.set(backlogSplit.getValue());
+    }
+
     // Stop processing current TopicPartition when it's time to stop.
     if (checkStopReadingFn != null
         && checkStopReadingFn.apply(kafkaSourceDescriptor.getTopicPartition())) {
