This is an automated email from the ASF dual-hosted git repository.
stankiewicz 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 0b6d68278b8 [KafkaIO] Use consumer position and lag to estimate end
offsets in KafkaUnboundedReader (#39830)
0b6d68278b8 is described below
commit 0b6d68278b878f674f9a8aafe769cdc408a99391
Author: Steven van Rossum <[email protected]>
AuthorDate: Tue Oct 6 11:41:12 2026 +0200
[KafkaIO] Use consumer position and lag to estimate end offsets in
KafkaUnboundedReader (#39830)
* Use currentLag and position to set end offset estimate
* Replace broken synchronized reads/writes with atomic reads/writes
---
.../beam/sdk/io/kafka/KafkaUnboundedReader.java | 181 +++++++++++----------
.../org/apache/beam/sdk/io/kafka/KafkaIOTest.java | 27 ---
.../org/apache/beam/sdk/io/kafka/KafkaMocks.java | 37 -----
3 files changed, 95 insertions(+), 150 deletions(-)
diff --git
a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
index c3aa6f5418b..5e18ad83e2a 100644
---
a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
+++
b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
@@ -34,12 +34,12 @@ import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
-import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
import java.util.stream.Collectors;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark;
@@ -64,7 +64,9 @@ 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.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.WakeupException;
import org.apache.kafka.common.serialization.Deserializer;
@@ -145,18 +147,6 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
// Note that consumer is not thread safe, should not be accessed out side
consumerPollLoop().
consumerPollThread.submit(this::consumerPollLoop);
- // offsetConsumer setup :
- Map<String, Object> offsetConsumerConfig =
- KafkaIOUtils.getOffsetConsumerConfig(
- name, spec.getOffsetConsumerConfig(), spec.getConsumerConfig());
-
- offsetConsumer = spec.getConsumerFactoryFn().apply(offsetConsumerConfig);
-
- // Fetch offsets once before running periodically.
- updateLatestOffsets();
-
- offsetFetcherThread.scheduleAtFixedRate(
- this::updateLatestOffsets, 0, OFFSET_UPDATE_INTERVAL_SECONDS,
TimeUnit.SECONDS);
return advance();
}
@@ -210,7 +200,7 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
int recordSize =
(rawRecord.key() == null ? 0 : rawRecord.key().length)
+ (rawRecord.value() == null ? 0 : rawRecord.value().length);
- pState.recordConsumed(rawRecord.offset(), recordSize);
+ pState.recordConsumed(rawRecord.offset(), recordSize, Instant.now());
bytesRead.inc(recordSize);
bytesReadBySplit.inc(recordSize);
@@ -409,16 +399,6 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
private AtomicBoolean closed = new AtomicBoolean(false);
private Instant nextAllowedCommitFailLogTime = Instant.ofEpochMilli(0);
- // Backlog support :
- // Kafka consumer does not have an API to fetch latest offset for topic. We
need to seekToEnd()
- // then look at position(). Use another consumer to do this so that the
primary consumer does
- // not need to be interrupted. The latest offsets are fetched periodically
on a thread. This is
- // still a bit of a hack, but so far there haven't been any issues reported
by the users.
- private @Nullable Consumer<byte[], byte[]> offsetConsumer = null;
- private final ScheduledExecutorService offsetFetcherThread =
- Executors.newSingleThreadScheduledExecutor();
- private static final int OFFSET_UPDATE_INTERVAL_SECONDS = 1;
-
private static final long UNINITIALIZED_OFFSET = -1;
/** watermark before any records have been read. */
@@ -457,25 +437,35 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
// maintains state of each assigned partition (buffered records, consumed
offset, etc)
private static class PartitionState<K, V> {
+ @SuppressWarnings("rawtypes")
+ private static final AtomicReferenceFieldUpdater<PartitionState, Instant>
LAST_UPDATED =
+ AtomicReferenceFieldUpdater.newUpdater(PartitionState.class,
Instant.class, "lastUpdated");
+
private final TopicPartition topicPartition;
+ // Note that nextOffset and latestOffset are updated on different threads.
+ // Writes are ordered by a subsequent write of lastUpdated.
+ // Reads are ordered by a preceding read of lastUpdated.
private long nextOffset;
private long latestOffset;
- private Instant latestOffsetFetchTime;
+ private volatile Instant lastUpdated;
private Instant lastWatermark; // As returned by timestampPolicy
private final TimestampPolicy<K, V> timestampPolicy;
private Iterator<ConsumerRecord<byte[], byte[]>> recordIter =
Collections.emptyIterator();
- private KafkaIOUtils.MovingAvg avgRecordSize = new
KafkaIOUtils.MovingAvg();
+ // Note that avgRecordSize is indepedently synchronized, but also needs to
be ordered in
+ // relation to lastUpdated to avoid observing independent states.
+ private final KafkaIOUtils.MovingAvg avgRecordSize = new
KafkaIOUtils.MovingAvg();
PartitionState(
TopicPartition partition, long nextOffset, TimestampPolicy<K, V>
timestampPolicy) {
this.topicPartition = partition;
this.nextOffset = nextOffset;
this.latestOffset = UNINITIALIZED_OFFSET;
- this.latestOffsetFetchTime = BoundedWindow.TIMESTAMP_MIN_VALUE;
this.lastWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
this.timestampPolicy = timestampPolicy;
+ this.lastUpdated =
+ BoundedWindow.TIMESTAMP_MIN_VALUE; // volatile store (sequentially
consistent release)
}
public TopicPartition topicPartition() {
@@ -483,41 +473,51 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
}
// Update consumedOffset and avgRecordSize
- void recordConsumed(long offset, int size) {
- nextOffset = offset + 1;
- avgRecordSize.update(size);
+ void recordConsumed(long offset, int size, Instant consumeTime) {
+ nextOffset = offset + 1; // normal store
+ avgRecordSize.update(size); // independent ordered store (release)
+ LAST_UPDATED.lazySet(this, consumeTime); // ordered store (release)
}
- synchronized void setLatestOffset(long latestOffset, Instant fetchTime) {
- this.latestOffset = latestOffset;
- this.latestOffsetFetchTime = fetchTime;
- LOG.debug(
- "{}: latest offset update for {} : {} (consumer offset {}, avg
record size {})",
- this,
- topicPartition,
- latestOffset,
- nextOffset,
- avgRecordSize);
+ void setLatestOffset(long latestOffset, Instant fetchTime) {
+ this.latestOffset = latestOffset; // normal store
+ LAST_UPDATED.lazySet(this, fetchTime); // ordered store (release)
}
- synchronized long approxBacklogInBytes() {
+ long approxBacklogInBytes() {
// Note that is an estimate of uncompressed backlog.
long backlogMessageCount = backlogMessageCount();
if (backlogMessageCount == UnboundedReader.BACKLOG_UNKNOWN) {
return UnboundedReader.BACKLOG_UNKNOWN;
}
- return (long) (backlogMessageCount * avgRecordSize.get());
+ return (long)
+ (backlogMessageCount * avgRecordSize.get()); // independent volatile
load (acquire)
+ }
+
+ long backlogMessageCount() {
+ return backlogMessageCountInternal(this.lastUpdated); // volatile load
(acquire)
}
- synchronized long backlogMessageCount() {
- if (latestOffset < 0 || nextOffset < 0 || latestOffset < nextOffset) {
+ // This is ugly, but needed because mkTimestampPolicyContext would
otherwise issue two volatile
+ // reads that could be observed as holding different values.
+ // Note that the argument must always be the result of acquiring the
member lastUpdated.
+ @SuppressWarnings("ReferenceEquality")
+ private long backlogMessageCountInternal(Instant lastUpdated) {
+ final long nextOffset = this.nextOffset; // normal load
+ final long latestOffset = this.latestOffset; // normal load
+
+ if (lastUpdated == BoundedWindow.TIMESTAMP_MIN_VALUE
+ || latestOffset < 0
+ || nextOffset < 0
+ || latestOffset < nextOffset) {
return UnboundedReader.BACKLOG_UNKNOWN;
}
return latestOffset - nextOffset;
}
- synchronized TimestampPolicyContext mkTimestampPolicyContext() {
- return new TimestampPolicyContext(backlogMessageCount(),
latestOffsetFetchTime);
+ TimestampPolicyContext mkTimestampPolicyContext() {
+ final Instant lastUpdated = this.lastUpdated; // volatile load (acquire)
+ return new
TimestampPolicyContext(backlogMessageCountInternal(lastUpdated), lastUpdated);
}
Instant updateAndGetWatermark() {
@@ -586,9 +586,30 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
}
private void consumerPollLoop() {
- // Read in a loop and enqueue the batch of records, if any, to
availableRecordsQueue.
Consumer<byte[], byte[]> consumer =
Preconditions.checkStateNotNull(this.consumer);
+ List<TopicPartition> topicPartitions =
+ Preconditions.checkStateNotNull(source.getSpec().getTopicPartitions());
+
+ // Set initial latest offset for each partition.
+ try {
+ Instant fetchTime = Instant.now();
+ Map<TopicPartition, Long> endOffsets =
consumer.endOffsets(topicPartitions);
+ for (PartitionState<K, V> p : partitionStates) {
+ p.setLatestOffset(
+ endOffsets.getOrDefault(p.topicPartition, UNINITIALIZED_OFFSET),
fetchTime);
+ }
+ } catch (KafkaException e) {
+ if (!closed.get()) { // Ignore the exception if the reader is closed.
+ LOG.warn(
+ "{}: exception while fetching latest offset for partitions {}.
will be retried.",
+ this,
+ topicPartitions,
+ e);
+ }
+ // Don't update the latest offset.
+ }
+ // Read in a loop and enqueue the batch of records, if any, to
availableRecordsQueue.
try {
ConsumerRecords<byte[], byte[]> records = ConsumerRecords.empty();
while (!closed.get()) {
@@ -602,6 +623,32 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
kafkaTopic,
java.time.Duration.ofMillis(stopwatch.elapsed(TimeUnit.MILLISECONDS)));
}
+
+ // Fetch the current offsets and update the latest offset.
+ Instant fetchTime = Instant.now();
+ for (PartitionState<K, V> pState : partitionStates) {
+ try {
+ final long position = consumer.position(pState.topicPartition);
+ consumer
+ .currentLag(pState.topicPartition)
+ .ifPresent(lag -> pState.setLatestOffset(position + lag,
fetchTime));
+ } catch (InvalidOffsetException e) {
+ // The position is undefined or out of range.
+ pState.setLatestOffset(UNINITIALIZED_OFFSET, fetchTime);
+ } catch (KafkaException e) {
+ if (!closed.get()) { // Ignore the exception if the reader is
closed.
+ LOG.warn(
+ "{}: exception while fetching latest offset for
partitions {}. will be retried.",
+ this,
+ topicPartitions,
+ 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.
+ }
+ }
} else if (availableRecordsQueue.offer(
records, RECORDS_ENQUEUE_POLL_TIMEOUT.getMillis(),
TimeUnit.MILLISECONDS)) {
records = ConsumerRecords.empty();
@@ -789,37 +836,6 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
}
}
- // Update latest offset for each partition.
- // Called from setupInitialOffset() at the start and then periodically from
offsetFetcher thread.
- private void updateLatestOffsets() {
- Consumer<byte[], byte[]> offsetConsumer =
Preconditions.checkStateNotNull(this.offsetConsumer);
- List<TopicPartition> topicPartitions =
- Preconditions.checkStateNotNull(source.getSpec().getTopicPartitions());
- Instant fetchTime = Instant.now();
- try {
- Map<TopicPartition, Long> endOffsets =
offsetConsumer.endOffsets(topicPartitions);
- for (PartitionState<K, V> p : partitionStates) {
- p.setLatestOffset(
- Preconditions.checkStateNotNull(
- endOffsets.get(p.topicPartition),
- "No end offset found for partition %s.",
- p.topicPartition),
- fetchTime);
- }
- } catch (Exception e) {
- if (!closed.get()) { // Ignore the exception if the reader is closed.
- LOG.warn(
- "{}: exception while fetching latest offset for partitions {}.
will be retried.",
- this,
- topicPartitions,
- e);
- }
- // Don't update the latest offset.
- }
-
- LOG.debug("{}: backlog {}", this, getSplitBacklogBytes());
- }
-
private void reportBacklog() {
long splitBacklogBytes = getSplitBacklogBytes();
if (splitBacklogBytes < 0) {
@@ -861,7 +877,6 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
public void close() throws IOException {
closed.set(true);
consumerPollThread.shutdown();
- offsetFetcherThread.shutdown();
boolean isShutdown = false;
@@ -872,14 +887,9 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
if (consumer != null) {
consumer.wakeup();
}
- if (offsetConsumer != null) {
- offsetConsumer.wakeup();
- }
availableRecordsQueue.poll(); // drain unread batch, this unblocks
consumer thread.
try {
- isShutdown =
- consumerPollThread.awaitTermination(10, TimeUnit.SECONDS)
- && offsetFetcherThread.awaitTermination(10, TimeUnit.SECONDS);
+ isShutdown = consumerPollThread.awaitTermination(10, TimeUnit.SECONDS);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException(e); // not expected
@@ -896,7 +906,6 @@ class KafkaUnboundedReader<K, V> extends
UnboundedReader<KafkaRecord<K, V>> {
Closeables.close(keyDeserializerInstance, true);
Closeables.close(valueDeserializerInstance, true);
- Closeables.close(offsetConsumer, true);
Closeables.close(consumer, true);
}
diff --git
a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index 5e7a5b32643..1e320681414 100644
---
a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++
b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -78,7 +78,6 @@ import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.io.UnboundedSource;
import org.apache.beam.sdk.io.UnboundedSource.UnboundedReader;
import org.apache.beam.sdk.io.kafka.KafkaIO.Read.FakeFlinkPipelineOptions;
-import org.apache.beam.sdk.io.kafka.KafkaMocks.EndOffsetErrorConsumerFactory;
import org.apache.beam.sdk.io.kafka.KafkaMocks.SendErrorProducerFactory;
import org.apache.beam.sdk.metrics.DistributionResult;
import org.apache.beam.sdk.metrics.Lineage;
@@ -1659,32 +1658,6 @@ public class KafkaIOTest {
commitsEnqueuedMetrics.getCounters().iterator().next().getAttempted(),
greaterThan(0L));
}
- @Test
- public void testUnboundedReaderLogsCommitFailure() throws Exception {
-
- List<String> topics = ImmutableList.of("topic_a");
-
- EndOffsetErrorConsumerFactory endOffsetErrorConsumerFactory =
- new EndOffsetErrorConsumerFactory();
-
- UnboundedSource<KafkaRecord<Integer, Long>, KafkaCheckpointMark> source =
- KafkaIO.<Integer, Long>read()
- .withBootstrapServers("myServer1:9092,myServer2:9092")
- .withTopics(topics)
- .withConsumerFactoryFn(endOffsetErrorConsumerFactory)
- .withKeyDeserializer(IntegerDeserializer.class)
- .withValueDeserializer(LongDeserializer.class)
- .makeSource();
-
- UnboundedReader<KafkaRecord<Integer, Long>> reader =
source.createReader(null, null);
-
- reader.start();
-
- unboundedReaderExpectedLogs.verifyWarn("exception while fetching latest
offset for partitions");
-
- reader.close();
- }
-
@Test
public void testSink() {
// Simply read from kafka source and write to kafka sink. Then verify the
records
diff --git
a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaMocks.java
b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaMocks.java
index 1303f1da3bc..8f52e4848c9 100644
---
a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaMocks.java
+++
b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaMocks.java
@@ -18,27 +18,20 @@
package org.apache.beam.sdk.io.kafka;
import java.io.Serializable;
-import java.util.Collection;
-import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Future;
import java.util.stream.Collectors;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.values.KV;
-import org.apache.kafka.clients.consumer.Consumer;
-import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.MockConsumer;
-import org.apache.kafka.clients.consumer.OffsetResetStrategy;
import org.apache.kafka.clients.producer.Callback;
import org.apache.kafka.clients.producer.MockProducer;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.Node;
import org.apache.kafka.common.PartitionInfo;
-import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.IntegerSerializer;
import org.apache.kafka.common.serialization.LongSerializer;
@@ -67,36 +60,6 @@ public class KafkaMocks {
}
}
- public static final class EndOffsetErrorConsumerFactory
- implements SerializableFunction<Map<String, Object>, Consumer<byte[],
byte[]>> {
- public EndOffsetErrorConsumerFactory() {}
-
- @Override
- public MockConsumer<byte[], byte[]> apply(Map<String, Object> input) {
- final MockConsumer<byte[], byte[]> consumer;
- if (input.containsKey(ConsumerConfig.GROUP_ID_CONFIG)) {
- consumer =
- new MockConsumer<byte[], byte[]>(OffsetResetStrategy.EARLIEST) {
- @Override
- public synchronized Map<TopicPartition, Long> endOffsets(
- Collection<TopicPartition> partitions) {
- throw new KafkaException("fakeException");
- }
- };
- } else {
- consumer = new MockConsumer<byte[],
byte[]>(OffsetResetStrategy.EARLIEST);
- }
- consumer.updatePartitions(
- "topic_a",
- Collections.singletonList(
- new PartitionInfo("topic_a", 1, new Node(1, "myServer1", 9092),
null, null)));
- consumer.updateBeginningOffsets(
- Collections.singletonMap(new TopicPartition("topic_a", 1), 0L));
- consumer.updateEndOffsets(Collections.singletonMap(new
TopicPartition("topic_a", 1), 0L));
- return consumer;
- }
- }
-
public static final class PartitionGrowthMockConsumer extends
MockConsumer<byte[], byte[]>
implements Serializable {
