Skip to content
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

Jump to bottom

[KafkaIO] Share size and offset estimators among ReadFromKafkaDoFn instances per PTransform #32928

Open
wants to merge 8 commits into
base: master
Choose a base branch
from
Open

[KafkaIO] Share size and offset estimators among ReadFromKafkaDoFn instances per PTransform #32928

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
fa71cd6
Explicitly update partitions and start/end offsets for relevant mock …
sjvanrossum Oct 24, 2024
891f2d1
Revert "Set backlog in gauge metric (#31137)"
sjvanrossum Oct 21, 2024
25e58b8
Revert "Add Backlog Metrics to Kafka Splittable DoFn Implementation …
sjvanrossum Oct 22, 2024
cf4817e
Call reportBacklog in nextBatch to report split metrics more often
sjvanrossum Oct 21, 2024
6a07d66
Report SDF metrics for the active split/partition after processing a …
sjvanrossum Oct 24, 2024
45583dd
Remove unused property, assignment in finalize will not be observed
sjvanrossum Oct 24, 2024
496886c
Use KafkaSourceDescriptor as cache key and log entry
sjvanrossum Oct 24, 2024
7df359f
Share AvgRecordSize and KafkaLatestOffsetEstimator caches among DoFns
sjvanrossum Oct 24, 2024
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
12 changes: 6 additions & 6 deletions .../src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/kafka/KafkaTestTable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,10 +138,6 @@ public synchronized void assign(final Collection<TopicPartition> assigned) {
.collect(Collectors.toList());
super.assign(realPartitions);
assignedPartitions.set(ImmutableList.copyOf(realPartitions));
for (TopicPartition tp : realPartitions) {
updateBeginningOffsets(ImmutableMap.of(tp, 0L));
updateEndOffsets(ImmutableMap.of(tp, (long) kafkaRecords.get(tp).size()));
}
}
// Override offsetsForTimes() in order to look up the offsets by timestamp.
@Override
Expand All @@ -163,8 +159,12 @@ public synchronized Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(
}
};

for (String topic : getTopics()) {
consumer.updatePartitions(topic, partitionInfoMap.get(topic));
for (Map.Entry<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> entry :
kafkaRecords.entrySet()) {
consumer.updatePartitions(
entry.getKey().topic(), partitionInfoMap.get(entry.getKey().topic()));
consumer.updateBeginningOffsets(ImmutableMap.of(entry.getKey(), 0L));
consumer.updateEndOffsets(ImmutableMap.of(entry.getKey(), (long) entry.getValue().size()));
}

Runnable recordEnqueueTask =
Expand Down
21 changes: 5 additions & 16 deletions sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -228,10 +227,6 @@ public boolean advance() throws IOException {
METRIC_NAMESPACE, RAW_SIZE_METRIC_PREFIX + pState.topicPartition.toString());
rawSizes.update(recordSize);

for (Map.Entry<String, Long> backlogSplit : perPartitionBacklogMetrics.entrySet()) {
backlogBytesOfSplit.set(backlogSplit.getValue());
}

// Pass metrics to container.
kafkaResults.updateKafkaMetrics();
return true;
Expand Down Expand Up @@ -350,7 +345,6 @@ public long getSplitBacklogBytes() {
private final Counter bytesReadBySplit;
private final Gauge backlogBytesOfSplit;
private final Gauge backlogElementsOfSplit;
private HashMap<String, Long> perPartitionBacklogMetrics = new HashMap<String, Long>();;
private final Counter checkpointMarkCommitsEnqueued =
Metrics.counter(METRIC_NAMESPACE, CHECKPOINT_MARK_COMMITS_ENQUEUED_METRIC);
// Checkpoint marks skipped in favor of newer mark (only the latest needs to be committed).
Expand Down Expand Up @@ -507,10 +501,6 @@ Instant updateAndGetWatermark() {
lastWatermark = timestampPolicy.getWatermark(mkTimestampPolicyContext());
return lastWatermark;
}

String name() {
return this.topicPartition.toString();
}
}

KafkaUnboundedReader(
Expand Down Expand Up @@ -555,16 +545,14 @@ String name() {
prevWatermark = Optional.of(new Instant(ckptMark.getWatermarkMillis()));
}

PartitionState<K, V> state =
new PartitionState<K, V>(
states.add(
new PartitionState<>(
tp,
nextOffset,
source
.getSpec()
.getTimestampPolicyFactory()
.createTimestampPolicy(tp, prevWatermark));
states.add(state);
perPartitionBacklogMetrics.put(state.name(), 0L);
.createTimestampPolicy(tp, prevWatermark)));
}

partitionStates = ImmutableList.copyOf(states);
Expand Down Expand Up @@ -681,6 +669,8 @@ private void nextBatch() throws IOException {

partitionStates.forEach(p -> p.recordIter = records.records(p.topicPartition).iterator());

reportBacklog();

// cycle through the partitions in order to interleave records from each.
curBatch = Iterators.cycle(new ArrayList<>(partitionStates));
}
Expand Down Expand Up @@ -754,7 +744,6 @@ private long getSplitBacklogMessageCount() {
if (pBacklog == UnboundedReader.BACKLOG_UNKNOWN) {
return UnboundedReader.BACKLOG_UNKNOWN;
}
perPartitionBacklogMetrics.put(p.name(), pBacklog);
backlogCount += pBacklog;
}

Expand Down
Loading
Loading