This is an automated email from the ASF dual-hosted git repository.
Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new a72d6cba37a [KafkaIO] Use consumer position and lag to estimate end
offsets (#39285)
a72d6cba37a is described below
commit a72d6cba37ad8c741005edd477ade137f9454b19
Author: Steven van Rossum <[email protected]>
AuthorDate: Thu Jul 23 16:33:20 2026 +0200
[KafkaIO] Use consumer position and lag to estimate end offsets (#39285)
* Use currentLag and position to set end offset estimate
* Ensure that end offset estimate can't precede consumed offset
---
.../beam/sdk/io/kafka/ReadFromKafkaDoFn.java | 114 ++++++---------------
1 file changed, 34 insertions(+), 80 deletions(-)
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 27af3b67930..68453443475 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 @@ package org.apache.beam.sdk.io.kafka;
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.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.UnsplittableRestrictionTrac
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.Consumer;
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 @@ abstract class ReadFromKafkaDoFn<K, V>
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 @@ abstract class ReadFromKafkaDoFn<K, V>
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 @@ abstract class ReadFromKafkaDoFn<K, V>
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 @@ abstract class ReadFromKafkaDoFn<K, V>
@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 @@ abstract class ReadFromKafkaDoFn<K, V>
@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 @@ abstract class ReadFromKafkaDoFn<K, V>
// 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 @@ abstract class ReadFromKafkaDoFn<K, V>
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 @@ abstract class ReadFromKafkaDoFn<K, V>
// 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 @@ abstract class ReadFromKafkaDoFn<K, V>
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 @@ abstract class ReadFromKafkaDoFn<K, V>
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;
