Use currentLag and position to set end offset estimate

Author: sjvanrossum

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

@@ -19,15 +19,14 @@
 
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
-import java.io.Closeable;
 import java.math.BigDecimal;
 import java.math.MathContext;
 import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.function.Supplier;
+import java.util.concurrent.atomic.AtomicLong;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.io.kafka.KafkaIO.ReadSourceDescriptors;
 import org.apache.beam.sdk.io.kafka.KafkaIOUtils.MovingAvg;
@@ -49,7 +48,6 @@
 import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator;
 import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
-import org.apache.beam.sdk.util.ExpiringMemoizingSerializableSupplier;
 import org.apache.beam.sdk.util.MemoizingPerInstantiationSerializableSupplier;
 import org.apache.beam.sdk.util.Preconditions;
 import org.apache.beam.sdk.util.SerializableSupplier;
@@ -71,6 +69,8 @@
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.InvalidOffsetException;
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.config.ConfigDef;
@@ -235,30 +235,12 @@ public MovingAvg load(KafkaSourceDescriptor kafkaSourceDescriptor)
                 CacheBuilder.newBuilder()
                     .concurrencyLevel(Runtime.getRuntime().availableProcessors())
                     .weakValues()
-                    .removalListener(
-                        (RemovalNotification<KafkaSourceDescriptor, KafkaLatestOffsetEstimator>
-                                notification) -> {
-                          final @Nullable KafkaLatestOffsetEstimator value;
-                          if (notification.getCause() == RemovalCause.COLLECTED
-                              && (value = notification.getValue()) != null) {
-                            value.close();
-                          }
-                        })
                     .build(
-                        new CacheLoader<KafkaSourceDescriptor, KafkaLatestOffsetEstimator>() {
+                        new CacheLoader<KafkaSourceDescriptor, AtomicLong>() {
                           @Override
-                          public KafkaLatestOffsetEstimator load(
-                              final KafkaSourceDescriptor sourceDescriptor) {
-                            LOG.info(
-                                "Creating Kafka consumer for offset estimation for {}",
-                                sourceDescriptor);
-                            final Map<String, Object> config =
-                                KafkaIOUtils.overrideBootstrapServersConfig(
-                                    consumerConfig, sourceDescriptor);
-                            final Consumer<byte[], byte[]> consumer =
-                                consumerFactoryFn.apply(config);
-                            return new KafkaLatestOffsetEstimator(
-                                consumer, sourceDescriptor.getTopicPartition());
+                          public AtomicLong load(final KafkaSourceDescriptor sourceDescriptor) {
+                            LOG.info("Creating end offset estimator for {}", sourceDescriptor);
+                            return new AtomicLong(Long.MIN_VALUE);
                           }
                         }));
     this.pollConsumerCacheSupplier =
@@ -319,8 +301,7 @@ public Consumer<byte[], byte[]> load(
   private final SerializableSupplier<LoadingCache<KafkaSourceDescriptor, MovingAvg>>
       avgRecordSizeCacheSupplier;
 
-  private final SerializableSupplier<
-          LoadingCache<KafkaSourceDescriptor, KafkaLatestOffsetEstimator>>
+  private final SerializableSupplier<LoadingCache<KafkaSourceDescriptor, AtomicLong>>
       latestOffsetEstimatorCacheSupplier;
 
   private final SerializableSupplier<LoadingCache<KafkaSourceDescriptor, Consumer<byte[], byte[]>>>
@@ -329,8 +310,7 @@ public Consumer<byte[], byte[]> load(
   private transient @MonotonicNonNull LoadingCache<KafkaSourceDescriptor, MovingAvg>
       avgRecordSizeCache;
 
-  private transient @MonotonicNonNull LoadingCache<
-          KafkaSourceDescriptor, KafkaLatestOffsetEstimator>
+  private transient @MonotonicNonNull LoadingCache<KafkaSourceDescriptor, AtomicLong>
       latestOffsetEstimatorCache;
 
   private transient @MonotonicNonNull LoadingCache<KafkaSourceDescriptor, Consumer<byte[], byte[]>>
@@ -349,46 +329,6 @@ public Consumer<byte[], byte[]> load(
   @VisibleForTesting
   static final String RAW_SIZE_METRIC_PREFIX = KafkaUnboundedReader.RAW_SIZE_METRIC_PREFIX;
 
-  /**
-   * A {@link GrowableOffsetRangeTracker.RangeEndEstimator} which uses a Kafka {@link Consumer} to
-   * fetch backlog.
-   */
-  private static class KafkaLatestOffsetEstimator
-      implements GrowableOffsetRangeTracker.RangeEndEstimator, Closeable {
-    private final Consumer<byte[], byte[]> offsetConsumer;
-    private final Supplier<Long> offsetSupplier;
-
-    KafkaLatestOffsetEstimator(
-        final Consumer<byte[], byte[]> offsetConsumer, final TopicPartition topicPartition) {
-      this.offsetConsumer = offsetConsumer;
-      this.offsetSupplier =
-          new ExpiringMemoizingSerializableSupplier<>(
-              () -> {
-                try {
-                  return offsetConsumer
-                      .endOffsets(Collections.singleton(topicPartition))
-                      .getOrDefault(topicPartition, Long.MIN_VALUE);
-                } catch (Throwable t) {
-                  LOG.error("Failed to get end offset for {}", topicPartition, t);
-                  return Long.MIN_VALUE;
-                }
-              },
-              Duration.ofSeconds(1),
-              Long.MIN_VALUE,
-              Duration.ZERO);
-    }
-
-    @Override
-    public long estimate() {
-      return offsetSupplier.get();
-    }
-
-    @Override
-    public void close() {
-      offsetConsumer.close();
-    }
-  }
-
   @GetInitialRestriction
   @RequiresNonNull({"pollConsumerCache"})
   public OffsetRange initialRestriction(@Element KafkaSourceDescriptor kafkaSourceDescriptor) {
@@ -500,8 +440,8 @@ public double getSize(
   @RequiresNonNull({"latestOffsetEstimatorCache"})
   public UnsplittableRestrictionTracker<OffsetRange, Long> restrictionTracker(
       @Element KafkaSourceDescriptor kafkaSourceDescriptor, @Restriction OffsetRange restriction) {
-    final LoadingCache<KafkaSourceDescriptor, KafkaLatestOffsetEstimator>
-        latestOffsetEstimatorCache = this.latestOffsetEstimatorCache;
+    final LoadingCache<KafkaSourceDescriptor, AtomicLong> latestOffsetEstimatorCache =
+        this.latestOffsetEstimatorCache;
 
     if (restriction.getTo() < Long.MAX_VALUE) {
       return new UnsplittableRestrictionTracker<>(new OffsetRangeTracker(restriction));
@@ -510,9 +450,10 @@ public UnsplittableRestrictionTracker<OffsetRange, Long> restrictionTracker(
     // OffsetEstimators are cached for each topic-partition because they hold a stateful connection,
     // so we want to minimize the amount of connections that we start and track with Kafka. Another
     // point is that it has a memoized backlog, and this should make that more reusable estimations.
+    final AtomicLong latestOffsetEstimator =
+        latestOffsetEstimatorCache.getUnchecked(kafkaSourceDescriptor);
     return new UnsplittableRestrictionTracker<>(
-        new GrowableOffsetRangeTracker(
-            restriction.getFrom(), latestOffsetEstimatorCache.getUnchecked(kafkaSourceDescriptor)));
+        new GrowableOffsetRangeTracker(restriction.getFrom(), latestOffsetEstimator::get));
   }
 
   @ProcessElement
@@ -525,14 +466,13 @@ public ProcessContinuation processElement(
       throws Exception {
     final LoadingCache<KafkaSourceDescriptor, MovingAvg> avgRecordSizeCache =
         this.avgRecordSizeCache;
-    final LoadingCache<KafkaSourceDescriptor, KafkaLatestOffsetEstimator>
-        latestOffsetEstimatorCache = this.latestOffsetEstimatorCache;
+    final LoadingCache<KafkaSourceDescriptor, AtomicLong> latestOffsetEstimatorCache =
+        this.latestOffsetEstimatorCache;
     final LoadingCache<KafkaSourceDescriptor, Consumer<byte[], byte[]>> pollConsumerCache =
         this.pollConsumerCache;
 
     final MovingAvg avgRecordSize = avgRecordSizeCache.get(kafkaSourceDescriptor);
-    final KafkaLatestOffsetEstimator latestOffsetEstimator =
-        latestOffsetEstimatorCache.get(kafkaSourceDescriptor);
+    final AtomicLong latestOffsetEstimator = latestOffsetEstimatorCache.get(kafkaSourceDescriptor);
     final Consumer<byte[], byte[]> consumer = pollConsumerCache.get(kafkaSourceDescriptor);
     final Deserializer<K> keyDeserializerInstance =
         Preconditions.checkStateNotNull(this.keyDeserializerInstance);
@@ -580,6 +520,20 @@ public ProcessContinuation processElement(
         // Fetch the next records.
         final ConsumerRecords<byte[], byte[]> rawRecords = consumer.poll(remainingTimeout);
         final Duration elapsed = pollTimer.elapsed();
+        try {
+          final long position = consumer.position(topicPartition);
+          consumer
+              .currentLag(topicPartition)
+              .ifPresent(lag -> latestOffsetEstimator.lazySet(position + lag));
+        } catch (InvalidOffsetException e) {
+          // The position is undefined or out of range.
+          latestOffsetEstimator.lazySet(Long.MIN_VALUE);
+        } catch (KafkaException e) {
+          // Wakeups, interrupts, authentication failures, authorization failures, timeouts and
+          // unrecoverable errors can be ignored. This routine is reattempted on every iteration
+          // and fallback methods would likely trigger the same unrecoverable errors.
+        }
+
         try {
           remainingTimeout = remainingTimeout.minus(elapsed);
         } catch (ArithmeticException e) {
@@ -687,7 +641,7 @@ public ProcessContinuation processElement(
 
         final long estimatedBacklogBytes =
             (long)
-                (BigDecimal.valueOf(latestOffsetEstimator.estimate())
+                (BigDecimal.valueOf(latestOffsetEstimator.get())
                         .subtract(BigDecimal.valueOf(expectedOffset), MathContext.DECIMAL128)
                         .doubleValue()
                     * avgRecordSize.get());
@@ -752,8 +706,8 @@ public void setup() throws Exception {
   public void teardown() throws Exception {
     final LoadingCache<KafkaSourceDescriptor, MovingAvg> avgRecordSizeCache =
         this.avgRecordSizeCache;
-    final LoadingCache<KafkaSourceDescriptor, KafkaLatestOffsetEstimator>
-        latestOffsetEstimatorCache = this.latestOffsetEstimatorCache;
+    final LoadingCache<KafkaSourceDescriptor, AtomicLong> latestOffsetEstimatorCache =
+        this.latestOffsetEstimatorCache;
     final LoadingCache<KafkaSourceDescriptor, Consumer<byte[], byte[]>> pollConsumerCache =
         this.pollConsumerCache;