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 index 27af3b67930f..684534434751 100644 --- 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 @@ -19,7 +19,6 @@ 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; @@ -27,7 +26,7 @@ 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 - notification) -> { - final @Nullable KafkaLatestOffsetEstimator value; - if (notification.getCause() == RemovalCause.COLLECTED - && (value = notification.getValue()) != null) { - value.close(); - } - }) .build( - new CacheLoader() { + new CacheLoader() { @Override - public KafkaLatestOffsetEstimator load( - final KafkaSourceDescriptor sourceDescriptor) { - LOG.info( - "Creating Kafka consumer for offset estimation for {}", - sourceDescriptor); - final Map config = - KafkaIOUtils.overrideBootstrapServersConfig( - consumerConfig, sourceDescriptor); - final Consumer 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 load( private final SerializableSupplier> avgRecordSizeCacheSupplier; - private final SerializableSupplier< - LoadingCache> + private final SerializableSupplier> latestOffsetEstimatorCacheSupplier; private final SerializableSupplier>> @@ -329,8 +310,7 @@ public Consumer load( private transient @MonotonicNonNull LoadingCache avgRecordSizeCache; - private transient @MonotonicNonNull LoadingCache< - KafkaSourceDescriptor, KafkaLatestOffsetEstimator> + private transient @MonotonicNonNull LoadingCache latestOffsetEstimatorCache; private transient @MonotonicNonNull LoadingCache> @@ -349,46 +329,6 @@ public Consumer 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 offsetConsumer; - private final Supplier offsetSupplier; - - KafkaLatestOffsetEstimator( - final Consumer 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 restrictionTracker( @Element KafkaSourceDescriptor kafkaSourceDescriptor, @Restriction OffsetRange restriction) { - final LoadingCache - latestOffsetEstimatorCache = this.latestOffsetEstimatorCache; + final LoadingCache latestOffsetEstimatorCache = + this.latestOffsetEstimatorCache; if (restriction.getTo() < Long.MAX_VALUE) { return new UnsplittableRestrictionTracker<>(new OffsetRangeTracker(restriction)); @@ -510,9 +450,10 @@ public UnsplittableRestrictionTracker 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 avgRecordSizeCache = this.avgRecordSizeCache; - final LoadingCache - latestOffsetEstimatorCache = this.latestOffsetEstimatorCache; + final LoadingCache latestOffsetEstimatorCache = + this.latestOffsetEstimatorCache; final LoadingCache> pollConsumerCache = this.pollConsumerCache; final MovingAvg avgRecordSize = avgRecordSizeCache.get(kafkaSourceDescriptor); - final KafkaLatestOffsetEstimator latestOffsetEstimator = - latestOffsetEstimatorCache.get(kafkaSourceDescriptor); + final AtomicLong latestOffsetEstimator = latestOffsetEstimatorCache.get(kafkaSourceDescriptor); final Consumer consumer = pollConsumerCache.get(kafkaSourceDescriptor); final Deserializer keyDeserializerInstance = Preconditions.checkStateNotNull(this.keyDeserializerInstance); @@ -580,6 +520,20 @@ public ProcessContinuation processElement( // Fetch the next records. final ConsumerRecords 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(Math.max(expectedOffset, 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 avgRecordSizeCache = this.avgRecordSizeCache; - final LoadingCache - latestOffsetEstimatorCache = this.latestOffsetEstimatorCache; + final LoadingCache latestOffsetEstimatorCache = + this.latestOffsetEstimatorCache; final LoadingCache> pollConsumerCache = this.pollConsumerCache;