Use currentLag and position to set end offset estimate

Author: sjvanrossum

buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -629,7 +629,7 @@ class BeamModulePlugin implements Plugin<Project> {
     def jaxb_api_version = "2.3.3"
     def jsr305_version = "3.0.2"
     def everit_json_version = "1.14.2"
-    def kafka_version = "2.8.2"
+    def kafka_version = "3.9.1"
     def log4j2_version = "2.20.0"
     def nemo_version = "0.1"
     // [bomupgrader] determined by: io.grpc:grpc-netty, consistent with: google_cloud_platform_libraries_bom
@@ -836,7 +836,7 @@ class BeamModulePlugin implements Plugin<Project> {
         jupiter_api                                 : "org.junit.jupiter:junit-jupiter-api:$jupiter_version",
         jupiter_engine                              : "org.junit.jupiter:junit-jupiter-engine:$jupiter_version",
         jupiter_params                              : "org.junit.jupiter:junit-jupiter-params:$jupiter_version",
-        kafka                                       : "org.apache.kafka:kafka_2.11:$kafka_version",
+        kafka                                       : "org.apache.kafka:kafka_2.12:$kafka_version",
         kafka_clients                               : "org.apache.kafka:kafka-clients:$kafka_version",
         log4j                                       : "log4j:log4j:1.2.17",
         log4j_over_slf4j                            : "org.slf4j:log4j-over-slf4j:$slf4j_version",

sdks/java/io/kafka/build.gradle

@@ -36,9 +36,7 @@ ext {
 }
 
 def kafkaVersions = [
-    '282': "2.8.2",
-    '312': "3.1.2",
-    '390': "3.9.0",
+    '391': "3.9.1",
 ]
 
 kafkaVersions.each{k,v -> configurations.create("kafkaVersion$k")}

sdks/java/io/kafka/kafka-312/build.gradle

@@ -16,8 +16,8 @@
  * limitations under the License.
  */
 project.ext {
-    delimited="2.8.2"
-    undelimited="282"
+    delimited="3.9.1"
+    undelimited="391"
     sdfCompatible=true
 }
 

sdks/java/io/kafka/kafka-390/build.gradle

@@ -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,7 @@
 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.common.KafkaException;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.config.ConfigDef;
@@ -235,30 +234,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 +300,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 +309,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 +328,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 +439,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 +449,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 +465,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 +519,14 @@ 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 (KafkaException e) {
+        }
+
         try {
           remainingTimeout = remainingTimeout.minus(elapsed);
         } catch (ArithmeticException e) {
@@ -687,7 +634,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 +699,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;
 

sdks/java/io/kafka/kafka-282/build.gradle sdks/java/io/kafka/kafka-391/build.gradlesdks/java/io/kafka/kafka-282/build.gradle renamed to sdks/java/io/kafka/kafka-391/build.gradle

@@ -17,11 +17,11 @@
  */
 package org.apache.beam.sdk.io.kafka;
 
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.TimeUnit;
 import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
@@ -236,7 +236,7 @@ public synchronized void commitSync(Map<TopicPartition, OffsetAndMetadata> offse
     }
 
     @Override
-    public synchronized void close(long timeout, TimeUnit unit) {
+    public synchronized void close(Duration timeout) {
       // Ignore closing since we're using a single consumer.
     }
   }

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

@@ -344,20 +344,8 @@ project(":beam-test-gha").projectDir = file(".github")
 include("beam-validate-runner")
 project(":beam-validate-runner").projectDir = file(".test-infra/validate-runner")
 include("com.google.api.gax.batching")
-include("sdks:java:io:kafka:kafka-390")
-findProject(":sdks:java:io:kafka:kafka-390")?.name = "kafka-390"
-include("sdks:java:io:kafka:kafka-312")
-findProject(":sdks:java:io:kafka:kafka-312")?.name = "kafka-312"
-include("sdks:java:io:kafka:kafka-282")
-findProject(":sdks:java:io:kafka:kafka-282")?.name = "kafka-282"
-include("sdks:java:io:kafka:kafka-251")
-findProject(":sdks:java:io:kafka:kafka-251")?.name = "kafka-251"
-include("sdks:java:io:kafka:kafka-241")
-findProject(":sdks:java:io:kafka:kafka-241")?.name = "kafka-241"
-include("sdks:java:io:kafka:kafka-231")
-findProject(":sdks:java:io:kafka:kafka-231")?.name = "kafka-231"
-include("sdks:java:io:kafka:kafka-201")
-findProject(":sdks:java:io:kafka:kafka-201")?.name = "kafka-201"
+include("sdks:java:io:kafka:kafka-391")
+findProject(":sdks:java:io:kafka:kafka-391")?.name = "kafka-391"
 include("sdks:java:managed")
 findProject(":sdks:java:managed")?.name = "managed"
 include("sdks:java:io:iceberg")