Skip to content

KAFKA-12487: Add support for cooperative consumer protocol with sink connectors #10563

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
kkonstantine merged 6 commits into from
Nov 10, 2021
Merged

KAFKA-12487: Add support for cooperative consumer protocol with sink connectors #10563

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
7708335
KAFKA-12487: Add support for cooperative consumer protocol with sink …
C0urante Apr 6, 2021
05aa4f6
KAFKA-12487: Fix sink task consumer pause/resume logic while paused f…
C0urante Apr 29, 2021
ec1c6db
KAFKA-12487: Fix unit test
C0urante May 12, 2021
6289b22
KAFKA-12487: Fix errant record reporter logic for multiple records wi…
C0urante Jul 16, 2021
9be00f6
KAFKA-12487: Address review feedback
C0urante Jul 16, 2021
db1f8f5
KAFKA-12487: Test cleanup
C0urante Jul 16, 2021
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
132 changes: 92 additions & 40 deletions connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import static java.util.Collections.singleton;
import static org.apache.kafka.connect.runtime.WorkerConfig.TOPIC_TRACKING_ENABLE_CONFIG;
Expand Down Expand Up @@ -125,6 +126,7 @@ public WorkerSinkTask(ConnectorTaskId id,
this.headerConverter = headerConverter;
this.transformationChain = transformationChain;
this.messageBatch = new ArrayList<>();
this.lastCommittedOffsets = new HashMap<>();
this.currentOffsets = new HashMap<>();
this.origOffsets = new HashMap<>();
this.pausedForRedelivery = false;
Expand Down Expand Up @@ -196,7 +198,7 @@ public void execute() {
initializeAndStart();
// Make sure any uncommitted data has been committed and the task has
// a chance to clean up its state
try (UncheckedCloseable suppressible = this::closePartitions) {
try (UncheckedCloseable suppressible = this::closeAllPartitions) {
while (!isStopping())
iteration();
} catch (WakeupException e) {
Expand Down Expand Up @@ -367,42 +369,53 @@ private void doCommit(Map<TopicPartition, OffsetAndMetadata> offsets, boolean cl
}

private void commitOffsets(long now, boolean closing) {
commitOffsets(now, closing, consumer.assignment());
}

private void commitOffsets(long now, boolean closing, Collection<TopicPartition> topicPartitions) {
if (workerErrantRecordReporter != null) {
log.trace("Awaiting all reported errors to be completed");
workerErrantRecordReporter.awaitAllFutures();
log.trace("Completed all reported errors");
log.trace("Awaiting reported errors for {} to be completed", topicPartitions);

@kkonstantine kkonstantine Jul 14, 2021

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.

I wonder if we want to print the actual list of partitions here, which might be long. And do it twice.
I see the same pattern is applied elsewhere. I understand the value of explicit listing.

@C0urante C0urante Jul 16, 2021

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.

That's fair; it may be useful to list once but probably not more than that. I'll remove the set of partitions from these log lines and add a single TRACE-level line at the top of the method that logs the set of partitions; at that level, I think the verbosity is acceptable, especially after this change goes out in case we need to do some debugging of unexpected behavior. LMK what you think, though.

workerErrantRecordReporter.awaitFutures(topicPartitions);
log.trace("Completed all reported errors for {}", topicPartitions);
}

if (currentOffsets.isEmpty())
Map<TopicPartition, OffsetAndMetadata> offsetsToCommit = currentOffsets.entrySet().stream()
.filter(e -> topicPartitions.contains(e.getKey()))
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

if (offsetsToCommit.isEmpty())
return;

committing = true;
commitSeqno += 1;
commitStarted = now;
sinkTaskMetricsGroup.recordOffsetSequenceNumber(commitSeqno);

Map<TopicPartition, OffsetAndMetadata> lastCommittedOffsets = this.lastCommittedOffsets.entrySet().stream()

@kkonstantine kkonstantine Jul 14, 2021

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.

shadowing of a member field this way is not ideal.
For example there's an older usage of this collection below. I assume it's the new local variable we want to use, but I can't be 100% sure.

Should we call this actuallyCommittedOffsets or similar?

@C0urante C0urante Jul 16, 2021

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.

The "actually committed offsets" (as in, the ones that we pass to the consumer) are computed a bit further on in the method. Maybe lastCommittedOffsetsForPartitions? Bit of a mouthful, will try to think of something better.

.filter(e -> offsetsToCommit.containsKey(e.getKey()))
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

final Map<TopicPartition, OffsetAndMetadata> taskProvidedOffsets;
try {
log.trace("{} Calling task.preCommit with current offsets: {}", this, currentOffsets);
taskProvidedOffsets = task.preCommit(new HashMap<>(currentOffsets));
log.trace("{} Calling task.preCommit with current offsets: {}", this, offsetsToCommit);
taskProvidedOffsets = task.preCommit(new HashMap<>(offsetsToCommit));
} catch (Throwable t) {
if (closing) {
log.warn("{} Offset commit failed during close", this);
onCommitCompleted(t, commitSeqno, null);
} else {
log.error("{} Offset commit failed, rewinding to last committed offsets", this, t);
for (Map.Entry<TopicPartition, OffsetAndMetadata> entry : lastCommittedOffsets.entrySet()) {
log.debug("{} Rewinding topic partition {} to offset {}", this, entry.getKey(), entry.getValue().offset());
consumer.seek(entry.getKey(), entry.getValue().offset());
}
currentOffsets = new HashMap<>(lastCommittedOffsets);

@kkonstantine kkonstantine Jul 14, 2021

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.

this change means we never create a fresh copy of the offsets map.
Is this correct? Is there a risk for the offsets to keep being added and never removed?

@C0urante C0urante Jul 16, 2021

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.

Good catch. I believe this would lead to incorrect metrics given the use of the currentOffsets field when invoking sinkTaskMetricsGroup::assignedOffsets, and may lead to issues with the end-of-life call to closePartitions as well.

I'll add cleanup logic in the consumer rebalance listener.

onCommitCompleted(t, commitSeqno, null);
currentOffsets.putAll(lastCommittedOffsets);
}
onCommitCompleted(t, commitSeqno, null);
return;
} finally {
if (closing) {
log.trace("{} Closing the task before committing the offsets: {}", this, currentOffsets);
task.close(currentOffsets.keySet());
log.trace("{} Closing the task before committing the offsets: {}", this, offsetsToCommit);
task.close(topicPartitions);
}
}

Expand All @@ -412,32 +425,36 @@ private void commitOffsets(long now, boolean closing) {
return;
}

final Map<TopicPartition, OffsetAndMetadata> commitableOffsets = new HashMap<>(lastCommittedOffsets);
Collection<TopicPartition> allAssignedTopicPartitions = consumer.assignment();
final Map<TopicPartition, OffsetAndMetadata> committableOffsets = new HashMap<>(lastCommittedOffsets);
for (Map.Entry<TopicPartition, OffsetAndMetadata> taskProvidedOffsetEntry : taskProvidedOffsets.entrySet()) {
final TopicPartition partition = taskProvidedOffsetEntry.getKey();
final OffsetAndMetadata taskProvidedOffset = taskProvidedOffsetEntry.getValue();
if (commitableOffsets.containsKey(partition)) {
if (committableOffsets.containsKey(partition)) {
long taskOffset = taskProvidedOffset.offset();
long currentOffset = currentOffsets.get(partition).offset();
long currentOffset = offsetsToCommit.get(partition).offset();
if (taskOffset <= currentOffset) {
commitableOffsets.put(partition, taskProvidedOffset);
committableOffsets.put(partition, taskProvidedOffset);
} else {
log.warn("{} Ignoring invalid task provided offset {}/{} -- not yet consumed, taskOffset={} currentOffset={}",
this, partition, taskProvidedOffset, taskOffset, currentOffset);
this, partition, taskProvidedOffset, taskOffset, currentOffset);
}
} else {
} else if (!allAssignedTopicPartitions.contains(partition)) {
log.warn("{} Ignoring invalid task provided offset {}/{} -- partition not assigned, assignment={}",
this, partition, taskProvidedOffset, consumer.assignment());
this, partition, taskProvidedOffset, allAssignedTopicPartitions);
} else {
log.debug("{} Ignoring task provided offset {}/{} -- topic partition not requested, requested={}",

@kkonstantine kkonstantine Jul 14, 2021

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.

what's a requested topic partition?
Also, above we mention just partition

@C0urante C0urante Jul 16, 2021

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.

Good catch, s/topic partition/partition/

"Requested" here means that although the partition is assigned to the task, it is not one of the partitions that we are currently committing offsets for.

this, partition, taskProvidedOffset, committableOffsets.keySet());
}
}

if (commitableOffsets.equals(lastCommittedOffsets)) {
if (committableOffsets.equals(lastCommittedOffsets)) {
log.debug("{} Skipping offset commit, no change since last commit", this);
onCommitCompleted(null, commitSeqno, null);
return;
}

doCommit(commitableOffsets, closing, commitSeqno);
doCommit(committableOffsets, closing, commitSeqno);
}


Expand Down Expand Up @@ -631,13 +648,31 @@ private void rewind() {
}

private void openPartitions(Collection<TopicPartition> partitions) {
sinkTaskMetricsGroup.recordPartitionCount(partitions.size());
updatePartitionCount();
task.open(partitions);
}

private void closePartitions() {
commitOffsets(time.milliseconds(), true);
sinkTaskMetricsGroup.recordPartitionCount(0);
private void closeAllPartitions() {
closePartitions(currentOffsets.keySet(), false);
}

private void closePartitions(Collection<TopicPartition> topicPartitions, boolean lost) {
if (!lost) {
commitOffsets(time.milliseconds(), true, topicPartitions);
} else {
log.trace("{} Closing the task as partitions have been lost: {}", this, topicPartitions);
task.close(topicPartitions);
if (workerErrantRecordReporter != null) {
log.trace("Cancelling reported errors for {}", topicPartitions);
workerErrantRecordReporter.cancelFutures(topicPartitions);

@ncliang ncliang Apr 29, 2021

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.

I'm not sure if cancelling the outstanding futures for error reporting is the right thing to do here. Would it be reasonable to await their completion for a reasonable amount of time before giving up?

@C0urante C0urante Apr 29, 2021 •
edited
Loading

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.

In this case, we've lost the partition assignment and won't be able to commit offsets for them, and since offset commits don't take place before outstanding error reports have completed, we know that no offsets for the records that caused these error reports have been committed either. So, we're guaranteed that the records that the task reported errors for will be redelivered, to whichever task now owns these partitions, regardless of whether we wait here for the reporting to complete.

@ncliang ncliang May 1, 2021

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.

I see. So the distinction is because in the revoked case we get a chance to await on the errant record reporter futures before we commit offsets for the revoked partitions, but in the lost case we've already lost the partition and don't get a chance to commit offsets. Since we already do not own the partition, we should not be reporting errors for it and should let the current owner take that responsibility. It should be noted that the cancelation is best effort, so there is a chance we duplicate reporting for the errant record.

log.trace("Cancelled all reported errors for {}", topicPartitions);
}
}
updatePartitionCount();
}

private void updatePartitionCount() {
sinkTaskMetricsGroup.recordPartitionCount(consumer.assignment().size());
}

@Override
Expand Down Expand Up @@ -670,8 +705,10 @@ private class HandleRebalance implements ConsumerRebalanceListener {
@Override
public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
log.debug("{} Partitions assigned {}", WorkerSinkTask.this, partitions);
lastCommittedOffsets = new HashMap<>();

@kkonstantine kkonstantine Jul 14, 2021

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.

similar question as above. This is a map that only grows now. Is this correct?

@C0urante C0urante Jul 16, 2021

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.

(Addressed above; it is almost certainly not correct)

currentOffsets = new HashMap<>();
if (partitions.isEmpty()) {
return;
}

for (TopicPartition tp : partitions) {
long pos = consumer.position(tp);
lastCommittedOffsets.put(tp, new OffsetAndMetadata(pos));
Expand All @@ -680,13 +717,13 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
}
sinkTaskMetricsGroup.assignedOffsets(currentOffsets);

// If we paused everything for redelivery (which is no longer relevant since we discarded the data), make
// If we paused everything for redelivery and all partitions for the failed deliveries have been revoked, make
// sure anything we paused that the task didn't request to be paused *and* which we still own is resumed.
// Also make sure our tracking of paused partitions is updated to remove any partitions we no longer own.
pausedForRedelivery = false;
pausedForRedelivery = pausedForRedelivery && !messageBatch.isEmpty();

@ncliang ncliang Apr 29, 2021

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.

I don't know if this change is required. The way I read the current implementation, we make sure that the paused partitions contain only assigned partitions in the block below, setting the paused partitions on context. We then rely on the code block in iteration() to resume partitions that should not be paused.

            } else if (!pausedForRedelivery) {
                resumeAll();
                onResume();
            }

Setting this to anything other than false causes us not to resume partitions which we own that were not explicitly requested to be paused.

@C0urante C0urante Apr 29, 2021

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.

Setting this to anything other than false causes us not to resume partitions which we own that were not explicitly requested to be paused.

Yep, that's intentional. The retry loop for sink tasks works like this:

  • Grab new batch of records from consumer, store in messageBatch
  • Deliver batch to task via SinkTask::put
  • Catch a RetriableException from the task
  • Pause all partitions on the consumer and set pausedForRedelivery to true
  • Redeliver batch (stored in the messageBatch field) to task via SinkTask::put continuously
  • Once a delivery attempt is successful, clear the batch stored in messageBatch and, since pausedForRedelivery is true, resume all partitions on the consumer and reset pausedForDelivery to false

When consumer rebalances meant a mass-revocation and mass-reassignment of partitions, we could get away with unconditionally resetting pausedForRedelivery to true during them, since offsets for any message batches going through the retry loop would not be committed and those records would be guaranteed to be redelivered after the rebalance. However, with incremental consumer rebalances, a task might not be forced to commit offsets for some (or even all) records in its current batch (if, for example, the new assignment only included additional partitions and no revocations). In that case, we have to be careful to remain paused so that the current batch can continue to be retried.

@C0urante C0urante Apr 29, 2021

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.

I'm actually wondering if there's a bug in the existing code; not sure pausedForRedelivery should be reset to false after a consumer rebalance. If the worker is in the middle of retrying a failed batch delivery, all the partitions will be paused, and after the rebalance, nothing explicitly resumes them.

During normal execution (when the connector isn't in the middle of redelivering a failed batch) this doesn't surface because partition consumption hasn't been paused anyways. But that assumption breaks down if we've paused to retry a failed batch.

I think it might be necessary to throw in a call to resumeAll() if pausedForRedelivery transitions from true to false as a result of a consumer rebalance (or, if we're lazy, we can just throw it in unconditionally at the end of each consumer rebalance if pausedForRedelivery is false since Consumer::resume is a no-op for partitions that weren't already paused).

@C0urante C0urante Apr 29, 2021

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.

Oooooh, one more edge case to consider--if new partitions are assigned while in the middle of retrying a failed batch, right now the worker might accidentally receive new records from those partitions, which would violate the assertion that new records can be received from the consumer if and only if the messageBatch field is empty.

I think we can (and should) handle this by invoking pauseAll() if pausedForRedelivery is still true after removing records from revoked partitions from the current batch.

@C0urante C0urante Apr 29, 2021 •
edited
Loading

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.

I've pushed a change that should address this gap and includes a new unit test for the pause/resume logic while paused for redelivery across consumer rebalances.


// Ensure that the paused partitions contains only assigned partitions and repause as necessary
context.pausedPartitions().retainAll(partitions);
context.pausedPartitions().retainAll(consumer.assignment());
if (shouldPause())
pauseAll();
else if (!context.pausedPartitions().isEmpty())
Expand All @@ -710,22 +747,35 @@ else if (!context.pausedPartitions().isEmpty())

@Override
public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
onPartitionsRemoved(partitions, false);
}

@Override
public void onPartitionsLost(Collection<TopicPartition> partitions) {
onPartitionsRemoved(partitions, true);
}

private void onPartitionsRemoved(Collection<TopicPartition> partitions, boolean lost) {
if (taskStopped) {
log.trace("Skipping partition revocation callback as task has already been stopped");
return;
}
log.debug("{} Partitions revoked", WorkerSinkTask.this);
log.debug("{} Partitions {}: {}", WorkerSinkTask.this, lost ? "lost" : "revoked", partitions);

if (partitions.isEmpty())
return;

try {
closePartitions();
sinkTaskMetricsGroup.clearOffsets();
closePartitions(partitions, lost);
sinkTaskMetricsGroup.clearOffsets(partitions);
} catch (RuntimeException e) {
// The consumer swallows exceptions raised in the rebalance listener, so we need to store
// exceptions and rethrow when poll() returns.
rebalanceException = e;
}

// Make sure we don't have any leftover data since offsets will be reset to committed positions
messageBatch.clear();
// Make sure we don't have any leftover data since offsets for these partitions will be reset to committed positions
messageBatch.removeIf(record -> partitions.contains(new TopicPartition(record.topic(), record.kafkaPartition())));

@kkonstantine kkonstantine Jul 14, 2021

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.

do you worry that this became more expensive now, especially in cases where we want to remove everything as before?

@C0urante C0urante Jul 16, 2021

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.

It certainly will be more expensive, although this will only occur in an edge case where the task has thrown a RetriableException from SinkTask::put, then before a follow-up invocation of SinkTask::put is able to succeed, a consumer rebalance takes place and partitions are revoked.

A naive improvement might be to convert messageBatch to a Map<TopicPartition, List<SinkRecord>>, which would allow us to quickly filter out records belonging to a given topic partition. But that would also be less efficient since we'd have to re-join those lists together before delivering records to the task, and would have to populate that map in the first place after retrieving the original list of records from Consumer::poll.

Not sure if there's a good way around this; if you believe it's a blocker I'm happy to spend some more time on it, though.

}
}

Expand Down Expand Up @@ -844,13 +894,15 @@ void recordCommittedOffsets(Map<TopicPartition, OffsetAndMetadata> offsets) {
void assignedOffsets(Map<TopicPartition, OffsetAndMetadata> offsets) {
consumedOffsets = new HashMap<>(offsets);
committedOffsets = offsets;
sinkRecordActiveCount.record(0.0);
computeSinkRecordLag();
}

void clearOffsets() {
consumedOffsets.clear();
committedOffsets.clear();
sinkRecordActiveCount.record(0.0);
void clearOffsets(Collection<TopicPartition> topicPartitions) {
topicPartitions.forEach(tp -> {
consumedOffsets.remove(tp);
committedOffsets.remove(tp);
});
computeSinkRecordLag();
}

void recordOffsetCommitSuccess() {
Expand Down
50 changes: 38 additions & 12 deletions ...ime/src/main/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.connect.errors.ConnectException;
import org.apache.kafka.connect.header.Header;
Expand All @@ -31,13 +32,17 @@
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.LinkedList;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

public class WorkerErrantRecordReporter implements ErrantRecordReporter {

Expand All @@ -49,7 +54,7 @@ public class WorkerErrantRecordReporter implements ErrantRecordReporter {
private final HeaderConverter headerConverter;

// Visible for testing
protected final LinkedList<Future<Void>> futures;
protected final Map<TopicPartition, Future<Void>> futures;

@kkonstantine kkonstantine Jul 14, 2021

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.

If we want this to be a concurrent map probably good idea to depict this in the declaration.
(same as you do in the tests below)

Suggested change
protected final Map<TopicPartition, Future<Void>> futures;
protected final ConcurrentMap<TopicPartition, Future<Void>> futures;

@C0urante C0urante Jul 16, 2021

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.

I like this style and agree with the suggestion; however, there's actually a bug in the PR right now where we only store one Future at a time per topic-partition, so will address in a follow-up commit instead of just clicking the Commit suggestion button.


public WorkerErrantRecordReporter(
RetryWithToleranceOperator retryWithToleranceOperator,
Expand All @@ -61,7 +66,7 @@ public WorkerErrantRecordReporter(
this.keyConverter = keyConverter;
this.valueConverter = valueConverter;
this.headerConverter = headerConverter;
this.futures = new LinkedList<>();
this.futures = new ConcurrentHashMap<>();
}

@Override
Expand Down Expand Up @@ -103,26 +108,47 @@ public Future<Void> report(SinkRecord record, Throwable error) {
Future<Void> future = retryWithToleranceOperator.executeFailed(Stage.TASK_PUT, SinkTask.class, consumerRecord, error);

if (!future.isDone()) {
futures.add(future);
futures.put(new TopicPartition(consumerRecord.topic(), consumerRecord.partition()), future);
}
return future;
}

/**
* Gets all futures returned by the sink records sent to Kafka by the errant
* record reporter. This function is intended to be used to block on all the errant record
* futures.
* Awaits the completion of all error reports for a given set of topic partitions
* @param topicPartitions the topic partitions to await reporter completion for
*/
public void awaitAllFutures() {
Future<?> future;
while ((future = futures.poll()) != null) {
public void awaitFutures(Collection<TopicPartition> topicPartitions) {
futuresFor(topicPartitions).forEach(future -> {
try {
future.get();
} catch (InterruptedException | ExecutionException e) {
log.error("Encountered an error while awaiting an errant record future's completion.");
log.error("Encountered an error while awaiting an errant record future's completion.", e);
throw new ConnectException(e);
}
}
});
}

/**
* Cancels all active error reports for a given set of topic partitions
* @param topicPartitions the topic partitions to cancel reporting for
*/
public void cancelFutures(Collection<TopicPartition> topicPartitions) {
futuresFor(topicPartitions).forEach(future -> {
try {
future.cancel(true);
} catch (Exception e) {
log.error("Encountered an error while cancelling an errant record future", e);
// No need to throw the exception here; it's enough to log an error message
}
});
}

// Removes and returns all futures for the given topic partitions from the set of currently-active futures
private Collection<Future<Void>> futuresFor(Collection<TopicPartition> topicPartitions) {
return topicPartitions.stream()
.map(futures::remove)
.filter(Objects::nonNull)
.collect(Collectors.toList());
}

/**
Expand Down
2 changes: 1 addition & 1 deletion ...runtime/src/test/java/org/apache/kafka/connect/integration/ErrantRecordSinkConnector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ public void put(Collection<SinkRecord> records) {
TopicPartition tp = cachedTopicPartitions
.computeIfAbsent(rec.topic(), v -> new HashMap<>())
.computeIfAbsent(rec.kafkaPartition(), v -> new TopicPartition(rec.topic(), rec.kafkaPartition()));
committedOffsets.put(tp, committedOffsets.getOrDefault(tp, 0L) + 1);
committedOffsets.put(tp, committedOffsets.getOrDefault(tp, 0) + 1);
reporter.report(rec, new Throwable());
}
}
Expand Down
2 changes: 1 addition & 1 deletion ...time/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,7 +261,7 @@ private boolean checkForPartitionAssignment() {
try {
ConnectorStateInfo info = connect.connectorStatus(CONNECTOR_NAME);
return info != null && info.tasks().size() == NUM_TASKS
&& connectorHandle.taskHandle(TASK_ID).partitionsAssigned() == 1;
&& connectorHandle.taskHandle(TASK_ID).numPartitionsAssigned() == 1;
} catch (Exception e) {
// Log the exception and return that the partitions were not assigned
log.error("Could not check connector state info.", e);
Expand Down
2 changes: 1 addition & 1 deletion ...ime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ private boolean checkForPartitionAssignment() {
try {
ConnectorStateInfo info = connect.connectorStatus(CONNECTOR_NAME);
return info != null && info.tasks().size() == NUM_TASKS
&& connectorHandle.tasks().stream().allMatch(th -> th.partitionsAssigned() == 1);
&& connectorHandle.tasks().stream().allMatch(th -> th.numPartitionsAssigned() == 1);
} catch (Exception e) {
// Log the exception and return that the partitions were not assigned
log.error("Could not check connector state info.", e);
Expand Down
Loading