Skip to content

[KafkaIO] Use consumer position and lag to estimate end offsets #39285

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Abacn merged 2 commits into from
Jul 23, 2026
+34 −80
Merged

[KafkaIO] Use consumer position and lag to estimate end offsets #39285

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
20df38f
Use currentLag and position to set end offset estimate
sjvanrossum Nov 27, 2025
5946687
Ensure that end offset estimate can't precede consumed offset
sjvanrossum Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 34 additions & 80 deletions sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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 =
Expand Down Expand Up @@ -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[]>>>
Expand All @@ -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[]>>
Expand All @@ -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) {
Expand Down Expand Up @@ -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));
Expand All @@ -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));
Comment on lines +453 to +456

@gemini-code-assist gemini-code-assist Bot Jul 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Since latestOffsetEstimator is initialized to Long.MIN_VALUE, the GrowableOffsetRangeTracker will receive Long.MIN_VALUE as the end offset estimate before the first poll completes. This causes getProgress() to cap the end offset to the current position and report 0 remaining work, which can incorrectly signal to the runner that the partition is fully processed. Returning Long.MAX_VALUE when the estimate is unknown (i.e., Long.MIN_VALUE) represents an infinite/unknown end offset, which is much safer and more appropriate for an unbounded source.

    final AtomicLong latestOffsetEstimator =
        latestOffsetEstimatorCache.getUnchecked(kafkaSourceDescriptor);
    return new UnsplittableRestrictionTracker<>(
        new GrowableOffsetRangeTracker(
            restriction.getFrom(),
            () -> {
              long estimate = latestOffsetEstimator.get();
              return estimate == Long.MIN_VALUE ? Long.MAX_VALUE : estimate;
            }));

@sjvanrossum sjvanrossum Jul 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This explanation and suggestion look wrong to me. Provide a source reference for the runner behavior.
Docs for GrowableOffsetRangeTracker.RangeEndEstimator state:

Return {@code Long.MIN_VALUE} if an estimate can not be provided.

}

@ProcessElement
Expand All @@ -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);
Expand Down Expand Up @@ -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.
}
Comment thread
sjvanrossum marked this conversation as resolved.

try {
remainingTimeout = remainingTimeout.minus(elapsed);
} catch (ArithmeticException e) {
Expand Down Expand Up @@ -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());
Comment thread
sjvanrossum marked this conversation as resolved.
Expand Down Expand Up @@ -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;

Expand Down
Loading