diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java index 1229fec7204b1..49ccb0dcfccfc 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java @@ -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; @@ -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; @@ -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) { @@ -367,13 +369,22 @@ private void doCommit(Map offsets, boolean cl } private void commitOffsets(long now, boolean closing) { + commitOffsets(now, closing, consumer.assignment()); + } + + private void commitOffsets(long now, boolean closing, Collection topicPartitions) { + log.trace("Committing offsets for partitions {}", 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 to be completed"); + workerErrantRecordReporter.awaitFutures(topicPartitions); + log.trace("Completed reported errors"); } - if (currentOffsets.isEmpty()) + Map 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; @@ -381,28 +392,31 @@ private void commitOffsets(long now, boolean closing) { commitStarted = now; sinkTaskMetricsGroup.recordOffsetSequenceNumber(commitSeqno); + Map lastCommittedOffsetsForPartitions = this.lastCommittedOffsets.entrySet().stream() + .filter(e -> offsetsToCommit.containsKey(e.getKey())) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + final Map 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 entry : lastCommittedOffsets.entrySet()) { + for (Map.Entry entry : lastCommittedOffsetsForPartitions.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); - onCommitCompleted(t, commitSeqno, null); + currentOffsets.putAll(lastCommittedOffsetsForPartitions); } + 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); } } @@ -412,32 +426,36 @@ private void commitOffsets(long now, boolean closing) { return; } - final Map commitableOffsets = new HashMap<>(lastCommittedOffsets); + Collection allAssignedTopicPartitions = consumer.assignment(); + final Map committableOffsets = new HashMap<>(lastCommittedOffsetsForPartitions); for (Map.Entry 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 {}/{} -- partition not requested, requested={}", + this, partition, taskProvidedOffset, committableOffsets.keySet()); } } - if (commitableOffsets.equals(lastCommittedOffsets)) { + if (committableOffsets.equals(lastCommittedOffsetsForPartitions)) { log.debug("{} Skipping offset commit, no change since last commit", this); onCommitCompleted(null, commitSeqno, null); return; } - doCommit(commitableOffsets, closing, commitSeqno); + doCommit(committableOffsets, closing, commitSeqno); } @@ -598,10 +616,12 @@ private void deliverMessages() { } } catch (RetriableException e) { log.error("{} RetriableException from SinkTask:", this, e); - // If we're retrying a previous batch, make sure we've paused all topic partitions so we don't get new data, - // but will still be able to poll in order to handle user-requested timeouts, keep group membership, etc. - pausedForRedelivery = true; - pauseAll(); + if (!pausedForRedelivery) { + // If we're retrying a previous batch, make sure we've paused all topic partitions so we don't get new data, + // but will still be able to poll in order to handle user-requested timeouts, keep group membership, etc. + pausedForRedelivery = true; + pauseAll(); + } // Let this exit normally, the batch will be reprocessed on the next loop. } catch (Throwable t) { log.error("{} Task threw an uncaught and unrecoverable exception. Task is being killed and will not " @@ -631,13 +651,32 @@ private void rewind() { } private void openPartitions(Collection 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 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); + log.trace("Cancelled all reported errors for {}", topicPartitions); + } + topicPartitions.forEach(currentOffsets::remove); + } + updatePartitionCount(); + } + + private void updatePartitionCount() { + sinkTaskMetricsGroup.recordPartitionCount(consumer.assignment().size()); } @Override @@ -670,8 +709,7 @@ private class HandleRebalance implements ConsumerRebalanceListener { @Override public void onPartitionsAssigned(Collection partitions) { log.debug("{} Partitions assigned {}", WorkerSinkTask.this, partitions); - lastCommittedOffsets = new HashMap<>(); - currentOffsets = new HashMap<>(); + for (TopicPartition tp : partitions) { long pos = consumer.position(tp); lastCommittedOffsets.put(tp, new OffsetAndMetadata(pos)); @@ -680,17 +718,29 @@ public void onPartitionsAssigned(Collection partitions) { } sinkTaskMetricsGroup.assignedOffsets(currentOffsets); - // If we paused everything for redelivery (which is no longer relevant since we discarded the data), 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; - - // Ensure that the paused partitions contains only assigned partitions and repause as necessary - context.pausedPartitions().retainAll(partitions); - if (shouldPause()) + boolean wasPausedForRedelivery = pausedForRedelivery; + pausedForRedelivery = wasPausedForRedelivery && !messageBatch.isEmpty(); + if (pausedForRedelivery) { + // Re-pause here in case we picked up new partitions in the rebalance pauseAll(); - else if (!context.pausedPartitions().isEmpty()) - consumer.pause(context.pausedPartitions()); + } else { + // 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. + if (wasPausedForRedelivery) { + resumeAll(); + } + // Ensure that the paused partitions contains only assigned partitions and repause as necessary + context.pausedPartitions().retainAll(consumer.assignment()); + if (shouldPause()) + pauseAll(); + else if (!context.pausedPartitions().isEmpty()) + consumer.pause(context.pausedPartitions()); + } + + if (partitions.isEmpty()) { + return; + } // Instead of invoking the assignment callback on initialization, we guarantee the consumer is ready upon // task start. Since this callback gets invoked during that initial setup before we've started the task, we @@ -710,22 +760,35 @@ else if (!context.pausedPartitions().isEmpty()) @Override public void onPartitionsRevoked(Collection partitions) { + onPartitionsRemoved(partitions, false); + } + + @Override + public void onPartitionsLost(Collection partitions) { + onPartitionsRemoved(partitions, true); + } + + private void onPartitionsRemoved(Collection 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()))); } } @@ -844,13 +907,15 @@ void recordCommittedOffsets(Map offsets) { void assignedOffsets(Map 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 topicPartitions) { + topicPartitions.forEach(tp -> { + consumedOffsets.remove(tp); + committedOffsets.remove(tp); + }); + computeSinkRecordLag(); } void recordOffsetCommitSuccess() { diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporter.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporter.java index aa1d0be36f493..ed48f79158d1f 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporter.java +++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporter.java @@ -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; @@ -31,13 +32,18 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.LinkedList; +import java.util.ArrayList; +import java.util.Collection; import java.util.List; import java.util.Optional; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; 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 { @@ -49,7 +55,7 @@ public class WorkerErrantRecordReporter implements ErrantRecordReporter { private final HeaderConverter headerConverter; // Visible for testing - protected final LinkedList> futures; + protected final ConcurrentMap>> futures; public WorkerErrantRecordReporter( RetryWithToleranceOperator retryWithToleranceOperator, @@ -61,7 +67,7 @@ public WorkerErrantRecordReporter( this.keyConverter = keyConverter; this.valueConverter = valueConverter; this.headerConverter = headerConverter; - this.futures = new LinkedList<>(); + this.futures = new ConcurrentHashMap<>(); } @Override @@ -103,26 +109,49 @@ public Future report(SinkRecord record, Throwable error) { Future future = retryWithToleranceOperator.executeFailed(Stage.TASK_PUT, SinkTask.class, consumerRecord, error); if (!future.isDone()) { - futures.add(future); + TopicPartition partition = new TopicPartition(consumerRecord.topic(), consumerRecord.partition()); + futures.computeIfAbsent(partition, p -> new ArrayList<>()).add(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 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 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> futuresFor(Collection topicPartitions) { + return topicPartitions.stream() + .map(futures::remove) + .filter(Objects::nonNull) + .flatMap(List::stream) + .collect(Collectors.toList()); } /** diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrantRecordSinkConnector.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrantRecordSinkConnector.java index 0fe2f88083883..251c67cef691c 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrantRecordSinkConnector.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrantRecordSinkConnector.java @@ -53,7 +53,7 @@ public void put(Collection 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()); } } diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java index b6211edc6b39c..b3dd9a097eda1 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ErrorHandlingIntegrationTest.java @@ -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); diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.java index 6f8d8a1596783..23a87c2c0fe49 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/ExampleConnectIntegrationTest.java @@ -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); diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java index 7b9afa4419008..5733199a0434e 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/MonitorableSinkConnector.java @@ -29,10 +29,8 @@ import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Set; /** * A sink connector that is used in Apache Kafka integration tests to verify the behavior of the @@ -91,12 +89,10 @@ public static class MonitorableSinkTask extends SinkTask { private String connectorName; private String taskId; TaskHandle taskHandle; - Set assignments; - Map committedOffsets; + Map committedOffsets; Map> cachedTopicPartitions; public MonitorableSinkTask() { - this.assignments = new HashSet<>(); this.committedOffsets = new HashMap<>(); this.cachedTopicPartitions = new HashMap<>(); } @@ -117,9 +113,15 @@ public void start(Map props) { @Override public void open(Collection partitions) { - log.debug("Opening {} partitions", partitions.size()); - assignments.addAll(partitions); - taskHandle.partitionsAssigned(partitions.size()); + log.debug("Opening partitions {}", partitions); + taskHandle.partitionsAssigned(partitions); + } + + @Override + public void close(Collection partitions) { + log.debug("Closing partitions {}", partitions); + taskHandle.partitionsRevoked(partitions); + partitions.forEach(committedOffsets::remove); } @Override @@ -129,26 +131,22 @@ public void put(Collection 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); log.trace("Task {} obtained record (key='{}' value='{}')", taskId, rec.key(), rec.value()); } } @Override public Map preCommit(Map offsets) { - for (TopicPartition tp : assignments) { - Long recordsSinceLastCommit = committedOffsets.get(tp); - if (recordsSinceLastCommit == null) { - log.warn("preCommit was called with topic-partition {} that is not included " - + "in the assignments of this task {}", tp, assignments); - } else { - taskHandle.commit(recordsSinceLastCommit.intValue()); - log.error("Forwarding to framework request to commit additional {} for {}", - recordsSinceLastCommit, tp); - taskHandle.commit((int) (long) recordsSinceLastCommit); - committedOffsets.put(tp, 0L); + taskHandle.partitionsCommitted(offsets.keySet()); + offsets.forEach((tp, offset) -> { + int recordsSinceLastCommit = committedOffsets.getOrDefault(tp, 0); + if (recordsSinceLastCommit != 0) { + taskHandle.commit(recordsSinceLastCommit); + log.debug("Forwarding to framework request to commit {} records for {}", recordsSinceLastCommit, tp); + committedOffsets.put(tp, 0); } - } + }); return offsets; } diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/SinkConnectorsIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/SinkConnectorsIntegrationTest.java new file mode 100644 index 0000000000000..a8bfbb291ffa1 --- /dev/null +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/SinkConnectorsIntegrationTest.java @@ -0,0 +1,321 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.connect.integration; + +import org.apache.kafka.clients.consumer.CooperativeStickyAssignor; +import org.apache.kafka.clients.consumer.RoundRobinAssignor; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.connect.sink.SinkRecord; +import org.apache.kafka.connect.storage.StringConverter; +import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster; +import org.apache.kafka.test.IntegrationTest; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Objects; +import java.util.Properties; +import java.util.Set; +import java.util.function.Consumer; + +import static org.apache.kafka.clients.consumer.ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG; +import static org.apache.kafka.clients.consumer.ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG; +import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG; +import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX; +import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG; +import static org.apache.kafka.connect.runtime.SinkConnectorConfig.TOPICS_CONFIG; +import static org.apache.kafka.connect.runtime.WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG; +import static org.apache.kafka.connect.runtime.WorkerConfig.KEY_CONVERTER_CLASS_CONFIG; +import static org.apache.kafka.connect.runtime.WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG; +import static org.apache.kafka.test.TestUtils.waitForCondition; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +/** + * Integration test for sink connectors + */ +@Category(IntegrationTest.class) +public class SinkConnectorsIntegrationTest { + + private static final int NUM_TASKS = 1; + private static final int NUM_WORKERS = 1; + private static final String CONNECTOR_NAME = "connect-integration-test-sink"; + private static final long TASK_CONSUME_TIMEOUT_MS = 10_000L; + + private EmbeddedConnectCluster connect; + + @Before + public void setup() throws Exception { + Map workerProps = new HashMap<>(); + // permit all Kafka client overrides; required for testing different consumer partition assignment strategies + workerProps.put(CONNECTOR_CLIENT_POLICY_CLASS_CONFIG, "All"); + + // setup Kafka broker properties + Properties brokerProps = new Properties(); + brokerProps.put("auto.create.topics.enable", "false"); + brokerProps.put("delete.topic.enable", "true"); + + // build a Connect cluster backed by Kafka and Zk + connect = new EmbeddedConnectCluster.Builder() + .name("connect-cluster") + .numWorkers(NUM_WORKERS) + .workerProps(workerProps) + .brokerProps(brokerProps) + .build(); + connect.start(); + connect.assertions().assertAtLeastNumWorkersAreUp(NUM_WORKERS, "Initial group of workers did not start in time."); + } + + @After + public void close() { + // delete connector handle + RuntimeHandles.get().deleteConnector(CONNECTOR_NAME); + + // stop all Connect, Kafka and Zk threads. + connect.stop(); + } + + @Test + public void testEagerConsumerPartitionAssignment() throws Exception { + final String topic1 = "topic1", topic2 = "topic2", topic3 = "topic3"; + final TopicPartition tp1 = new TopicPartition(topic1, 0), tp2 = new TopicPartition(topic2, 0), tp3 = new TopicPartition(topic3, 0); + final Collection topics = Arrays.asList(topic1, topic2, topic3); + + Map connectorProps = baseSinkConnectorProps(String.join(",", topics)); + // Need an eager assignor here; round robin is as good as any + connectorProps.put( + CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX + PARTITION_ASSIGNMENT_STRATEGY_CONFIG, + RoundRobinAssignor.class.getName()); + // After deleting a topic, offset commits will fail for it; reduce the timeout here so that the test doesn't take forever to proceed past that point + connectorProps.put( + CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX + DEFAULT_API_TIMEOUT_MS_CONFIG, + "5000"); + + final Set consumedRecordValues = new HashSet<>(); + Consumer onPut = record -> assertTrue("Task received duplicate record from Connect", consumedRecordValues.add(Objects.toString(record.value()))); + ConnectorHandle connector = RuntimeHandles.get().connectorHandle(CONNECTOR_NAME); + TaskHandle task = connector.taskHandle(CONNECTOR_NAME + "-0", onPut); + + connect.configureConnector(CONNECTOR_NAME, connectorProps); + connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS, "Connector tasks did not start in time."); + + // None of the topics has been created yet; the task shouldn't be assigned any partitions + assertEquals(0, task.numPartitionsAssigned()); + + Set expectedRecordValues = new HashSet<>(); + Set expectedAssignment = new HashSet<>(); + + connect.kafka().createTopic(topic1, 1); + expectedAssignment.add(tp1); + connect.kafka().produce(topic1, "t1v1"); + expectedRecordValues.add("t1v1"); + + waitForCondition( + () -> expectedRecordValues.equals(consumedRecordValues), + TASK_CONSUME_TIMEOUT_MS, + "Task did not receive records in time"); + assertEquals(1, task.timesAssigned(tp1)); + assertEquals(0, task.timesRevoked(tp1)); + assertEquals(expectedAssignment, task.assignment()); + + connect.kafka().createTopic(topic2, 1); + expectedAssignment.add(tp2); + connect.kafka().produce(topic2, "t2v1"); + expectedRecordValues.add("t2v1"); + connect.kafka().produce(topic2, "t1v2"); + expectedRecordValues.add("t1v2"); + + waitForCondition( + () -> expectedRecordValues.equals(consumedRecordValues), + TASK_CONSUME_TIMEOUT_MS, + "Task did not receive records in time"); + assertEquals(2, task.timesAssigned(tp1)); + assertEquals(1, task.timesRevoked(tp1)); + assertEquals(1, task.timesCommitted(tp1)); + assertEquals(1, task.timesAssigned(tp2)); + assertEquals(0, task.timesRevoked(tp2)); + assertEquals(expectedAssignment, task.assignment()); + + connect.kafka().createTopic(topic3, 1); + expectedAssignment.add(tp3); + connect.kafka().produce(topic3, "t3v1"); + expectedRecordValues.add("t3v1"); + connect.kafka().produce(topic2, "t2v2"); + expectedRecordValues.add("t2v2"); + connect.kafka().produce(topic2, "t1v3"); + expectedRecordValues.add("t1v3"); + + expectedAssignment.add(tp3); + waitForCondition( + () -> expectedRecordValues.equals(consumedRecordValues), + TASK_CONSUME_TIMEOUT_MS, + "Task did not receive records in time"); + assertEquals(3, task.timesAssigned(tp1)); + assertEquals(2, task.timesRevoked(tp1)); + assertEquals(2, task.timesCommitted(tp1)); + assertEquals(2, task.timesAssigned(tp2)); + assertEquals(1, task.timesRevoked(tp2)); + assertEquals(1, task.timesCommitted(tp2)); + assertEquals(1, task.timesAssigned(tp3)); + assertEquals(0, task.timesRevoked(tp3)); + assertEquals(expectedAssignment, task.assignment()); + + connect.kafka().deleteTopic(topic1); + expectedAssignment.remove(tp1); + connect.kafka().produce(topic3, "t3v2"); + expectedRecordValues.add("t3v2"); + connect.kafka().produce(topic2, "t2v3"); + expectedRecordValues.add("t2v3"); + + waitForCondition( + () -> expectedRecordValues.equals(consumedRecordValues) && expectedAssignment.equals(task.assignment()), + TASK_CONSUME_TIMEOUT_MS, + "Timed out while waiting for task to receive records and updated topic partition assignment"); + assertEquals(3, task.timesAssigned(tp1)); + assertEquals(3, task.timesRevoked(tp1)); + assertEquals(3, task.timesCommitted(tp1)); + assertEquals(3, task.timesAssigned(tp2)); + assertEquals(2, task.timesRevoked(tp2)); + assertEquals(2, task.timesCommitted(tp2)); + assertEquals(2, task.timesAssigned(tp3)); + assertEquals(1, task.timesRevoked(tp3)); + assertEquals(1, task.timesCommitted(tp3)); + } + + @Test + public void testCooperativeConsumerPartitionAssignment() throws Exception { + final String topic1 = "topic1", topic2 = "topic2", topic3 = "topic3"; + final TopicPartition tp1 = new TopicPartition(topic1, 0), tp2 = new TopicPartition(topic2, 0), tp3 = new TopicPartition(topic3, 0); + final Collection topics = Arrays.asList(topic1, topic2, topic3); + + Map connectorProps = baseSinkConnectorProps(String.join(",", topics)); + // Need an eager assignor here; round robin is as good as any + connectorProps.put( + CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX + PARTITION_ASSIGNMENT_STRATEGY_CONFIG, + CooperativeStickyAssignor.class.getName()); + // After deleting a topic, offset commits will fail for it; reduce the timeout here so that the test doesn't take forever to proceed past that point + connectorProps.put( + CONNECTOR_CLIENT_CONSUMER_OVERRIDES_PREFIX + DEFAULT_API_TIMEOUT_MS_CONFIG, + "5000"); + + final Set consumedRecordValues = new HashSet<>(); + Consumer onPut = record -> assertTrue("Task received duplicate record from Connect", consumedRecordValues.add(Objects.toString(record.value()))); + ConnectorHandle connector = RuntimeHandles.get().connectorHandle(CONNECTOR_NAME); + TaskHandle task = connector.taskHandle(CONNECTOR_NAME + "-0", onPut); + + connect.configureConnector(CONNECTOR_NAME, connectorProps); + connect.assertions().assertConnectorAndAtLeastNumTasksAreRunning(CONNECTOR_NAME, NUM_TASKS, "Connector tasks did not start in time."); + + // None of the topics has been created yet; the task shouldn't be assigned any partitions + assertEquals(0, task.numPartitionsAssigned()); + + Set expectedRecordValues = new HashSet<>(); + Set expectedAssignment = new HashSet<>(); + + connect.kafka().createTopic(topic1, 1); + expectedAssignment.add(tp1); + connect.kafka().produce(topic1, "t1v1"); + expectedRecordValues.add("t1v1"); + + waitForCondition( + () -> expectedRecordValues.equals(consumedRecordValues), + TASK_CONSUME_TIMEOUT_MS, + "Task did not receive records in time"); + assertEquals(1, task.timesAssigned(tp1)); + assertEquals(0, task.timesRevoked(tp1)); + assertEquals(expectedAssignment, task.assignment()); + + connect.kafka().createTopic(topic2, 1); + expectedAssignment.add(tp2); + connect.kafka().produce(topic2, "t2v1"); + expectedRecordValues.add("t2v1"); + connect.kafka().produce(topic2, "t1v2"); + expectedRecordValues.add("t1v2"); + + waitForCondition( + () -> expectedRecordValues.equals(consumedRecordValues), + TASK_CONSUME_TIMEOUT_MS, + "Task did not receive records in time"); + assertEquals(1, task.timesAssigned(tp1)); + assertEquals(0, task.timesRevoked(tp1)); + assertEquals(0, task.timesCommitted(tp1)); + assertEquals(1, task.timesAssigned(tp2)); + assertEquals(0, task.timesRevoked(tp2)); + assertEquals(expectedAssignment, task.assignment()); + + connect.kafka().createTopic(topic3, 1); + expectedAssignment.add(tp3); + connect.kafka().produce(topic3, "t3v1"); + expectedRecordValues.add("t3v1"); + connect.kafka().produce(topic2, "t2v2"); + expectedRecordValues.add("t2v2"); + connect.kafka().produce(topic2, "t1v3"); + expectedRecordValues.add("t1v3"); + + expectedAssignment.add(tp3); + waitForCondition( + () -> expectedRecordValues.equals(consumedRecordValues), + TASK_CONSUME_TIMEOUT_MS, + "Task did not receive records in time"); + assertEquals(1, task.timesAssigned(tp1)); + assertEquals(0, task.timesRevoked(tp1)); + assertEquals(0, task.timesCommitted(tp1)); + assertEquals(1, task.timesAssigned(tp2)); + assertEquals(0, task.timesRevoked(tp2)); + assertEquals(0, task.timesCommitted(tp2)); + assertEquals(1, task.timesAssigned(tp3)); + assertEquals(0, task.timesRevoked(tp3)); + assertEquals(expectedAssignment, task.assignment()); + + connect.kafka().deleteTopic(topic1); + expectedAssignment.remove(tp1); + connect.kafka().produce(topic3, "t3v2"); + expectedRecordValues.add("t3v2"); + connect.kafka().produce(topic2, "t2v3"); + expectedRecordValues.add("t2v3"); + + waitForCondition( + () -> expectedRecordValues.equals(consumedRecordValues) && expectedAssignment.equals(task.assignment()), + TASK_CONSUME_TIMEOUT_MS, + "Timed out while waiting for task to receive records and updated topic partition assignment"); + assertEquals(1, task.timesAssigned(tp1)); + assertEquals(1, task.timesRevoked(tp1)); + assertEquals(1, task.timesCommitted(tp1)); + assertEquals(1, task.timesAssigned(tp2)); + assertEquals(0, task.timesRevoked(tp2)); + assertEquals(0, task.timesCommitted(tp2)); + assertEquals(1, task.timesAssigned(tp3)); + assertEquals(0, task.timesRevoked(tp3)); + assertEquals(0, task.timesCommitted(tp3)); + } + + private Map baseSinkConnectorProps(String topics) { + Map props = new HashMap<>(); + props.put(CONNECTOR_CLASS_CONFIG, MonitorableSinkConnector.class.getSimpleName()); + props.put(TASKS_MAX_CONFIG, String.valueOf(NUM_TASKS)); + props.put(TOPICS_CONFIG, topics); + props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName()); + props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName()); + return props; + } +} diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TaskHandle.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TaskHandle.java index b63c08bf36811..35836faf8d5a8 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TaskHandle.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TaskHandle.java @@ -16,15 +16,19 @@ */ package org.apache.kafka.connect.integration; +import org.apache.kafka.common.TopicPartition; import org.apache.kafka.connect.errors.DataException; import org.apache.kafka.connect.sink.SinkRecord; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Collection; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; +import java.util.stream.Collectors; import java.util.stream.IntStream; /** @@ -37,7 +41,7 @@ public class TaskHandle { private final String taskId; private final ConnectorHandle connectorHandle; - private final AtomicInteger partitionsAssigned = new AtomicInteger(0); + private final ConcurrentMap partitions = new ConcurrentHashMap<>(); private final StartAndStopCounter startAndStopCounter = new StartAndStopCounter(); private final Consumer consumer; @@ -124,19 +128,74 @@ public void expectedCommits(int expected) { } /** - * Set the number of partitions assigned to this task. + * Adds a set of partitions to the (sink) task's assignment * - * @param numPartitions number of partitions + * @param partitions the newly-assigned partitions */ - public void partitionsAssigned(int numPartitions) { - partitionsAssigned.set(numPartitions); + public void partitionsAssigned(Collection partitions) { + partitions.forEach(partition -> this.partitions.computeIfAbsent(partition, PartitionHistory::new).assigned()); } /** - * @return the number of topic partitions assigned to this task. + * Removes a set of partitions to the (sink) task's assignment + * + * @param partitions the newly-revoked partitions + */ + public void partitionsRevoked(Collection partitions) { + partitions.forEach(partition -> this.partitions.computeIfAbsent(partition, PartitionHistory::new).revoked()); + } + + /** + * Records offset commits for a (sink) task's partitions + * + * @param partitions the committed partitions + */ + public void partitionsCommitted(Collection partitions) { + partitions.forEach(partition -> this.partitions.computeIfAbsent(partition, PartitionHistory::new).committed()); + } + + /** + * @return the complete set of partitions currently assigned to this (sink) task + */ + public Collection assignment() { + return partitions.values().stream() + .filter(PartitionHistory::isAssigned) + .map(PartitionHistory::topicPartition) + .collect(Collectors.toSet()); + } + + /** + * @return the number of topic partitions assigned to this (sink) task. + */ + public int numPartitionsAssigned() { + return assignment().size(); + } + + /** + * Returns the number of times the partition has been assigned to this (sink) task. + * @param partition the partition + * @return the number of times it has been assigned; may be 0 if never assigned + */ + public int timesAssigned(TopicPartition partition) { + return partitions.computeIfAbsent(partition, PartitionHistory::new).timesAssigned(); + } + + /** + * Returns the number of times the partition has been revoked from this (sink) task. + * @param partition the partition + * @return the number of times it has been revoked; may be 0 if never revoked + */ + public int timesRevoked(TopicPartition partition) { + return partitions.computeIfAbsent(partition, PartitionHistory::new).timesRevoked(); + } + + /** + * Returns the number of times the framework has committed offsets for this partition + * @param partition the partition + * @return the number of times it has been committed; may be 0 if never committed */ - public int partitionsAssigned() { - return partitionsAssigned.get(); + public int timesCommitted(TopicPartition partition) { + return partitions.computeIfAbsent(partition, PartitionHistory::new).timesCommitted(); } /** @@ -253,4 +312,50 @@ public String toString() { "taskId='" + taskId + '\'' + '}'; } + + private static class PartitionHistory { + private final TopicPartition topicPartition; + private boolean assigned = false; + private int timesAssigned = 0; + private int timesRevoked = 0; + private int timesCommitted = 0; + + public PartitionHistory(TopicPartition topicPartition) { + this.topicPartition = topicPartition; + } + + public void assigned() { + timesAssigned++; + assigned = true; + } + + public void revoked() { + timesRevoked++; + assigned = false; + } + + public void committed() { + timesCommitted++; + } + + public TopicPartition topicPartition() { + return topicPartition; + } + + public boolean isAssigned() { + return assigned; + } + + public int timesAssigned() { + return timesAssigned; + } + + public int timesRevoked() { + return timesRevoked; + } + + public int timesCommitted() { + return timesCommitted; + } + } } diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TransformationIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TransformationIntegrationTest.java index 7b71f2f5f70a5..02d8c7f71b1ac 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TransformationIntegrationTest.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TransformationIntegrationTest.java @@ -236,7 +236,7 @@ public void testFilterOnTombstonesWithSinkConnector() throws Exception { props.put(PREDICATES_CONFIG + ".barPredicate.type", RecordIsTombstone.class.getSimpleName()); // expect only half the records to be consumed by the connector - connectorHandle.expectedCommits(numRecords); + connectorHandle.expectedCommits(numRecords / 2); connectorHandle.expectedRecords(numRecords / 2); // start a sink connector diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java index b7492d31f502f..bcf25beee5b85 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java @@ -83,7 +83,9 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; import java.util.regex.Pattern; +import java.util.stream.Collectors; import static java.util.Arrays.asList; import static java.util.Collections.singleton; @@ -117,6 +119,9 @@ public class WorkerSinkTaskTest { private static final TopicPartition TOPIC_PARTITION2 = new TopicPartition(TOPIC, PARTITION2); private static final TopicPartition TOPIC_PARTITION3 = new TopicPartition(TOPIC, PARTITION3); + private static final Set INITIAL_ASSIGNMENT = + new HashSet<>(Arrays.asList(TOPIC_PARTITION, TOPIC_PARTITION2)); + private static final Map TASK_PROPS = new HashMap<>(); static { TASK_PROPS.put(SinkConnector.TOPICS_CONFIG, TOPIC); @@ -199,9 +204,8 @@ public void testStartPaused() throws Exception { expectTaskGetTopic(true); expectPollInitialAssignment(); - Set partitions = new HashSet<>(asList(TOPIC_PARTITION, TOPIC_PARTITION2)); - EasyMock.expect(consumer.assignment()).andReturn(partitions); - consumer.pause(partitions); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT); + consumer.pause(INITIAL_ASSIGNMENT); PowerMock.expectLastCall(); PowerMock.replayAll(); @@ -233,14 +237,12 @@ public void testPause() throws Exception { sinkTask.put(EasyMock.anyObject()); EasyMock.expectLastCall(); - Set partitions = new HashSet<>(asList(TOPIC_PARTITION, TOPIC_PARTITION2)); - // Pause statusListener.onPause(taskId); EasyMock.expectLastCall(); expectConsumerWakeup(); - EasyMock.expect(consumer.assignment()).andReturn(partitions); - consumer.pause(partitions); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT); + consumer.pause(INITIAL_ASSIGNMENT); PowerMock.expectLastCall(); // Offset commit as requested when pausing; No records returned by consumer.poll() @@ -254,11 +256,11 @@ public void testPause() throws Exception { statusListener.onResume(taskId); EasyMock.expectLastCall(); expectConsumerWakeup(); - EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(asList(TOPIC_PARTITION, TOPIC_PARTITION2))); - consumer.resume(singleton(TOPIC_PARTITION)); - PowerMock.expectLastCall(); - consumer.resume(singleton(TOPIC_PARTITION2)); - PowerMock.expectLastCall(); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2); + INITIAL_ASSIGNMENT.forEach(tp -> { + consumer.resume(Collections.singleton(tp)); + PowerMock.expectLastCall(); + }); expectConsumerPoll(1); expectConversionAndTransformation(1); @@ -338,11 +340,12 @@ public void testShutdown() throws Exception { sinkTask.stop(); PowerMock.expectLastCall(); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT); // WorkerSinkTask::close consumer.close(); PowerMock.expectLastCall().andAnswer(() -> { rebalanceListener.getValue().onPartitionsRevoked( - asList(TOPIC_PARTITION, TOPIC_PARTITION2) + INITIAL_ASSIGNMENT ); return null; }); @@ -377,9 +380,8 @@ public void testPollRedelivery() throws Exception { sinkTask.put(EasyMock.capture(records)); EasyMock.expectLastCall().andThrow(new RetriableException("retry")); // Pause - HashSet partitions = new HashSet<>(asList(TOPIC_PARTITION, TOPIC_PARTITION2)); - EasyMock.expect(consumer.assignment()).andReturn(partitions); - consumer.pause(partitions); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT); + consumer.pause(INITIAL_ASSIGNMENT); PowerMock.expectLastCall(); // Retry delivery should succeed @@ -387,11 +389,11 @@ public void testPollRedelivery() throws Exception { sinkTask.put(EasyMock.capture(records)); EasyMock.expectLastCall(); // And unpause - EasyMock.expect(consumer.assignment()).andReturn(partitions); - consumer.resume(singleton(TOPIC_PARTITION)); - PowerMock.expectLastCall(); - consumer.resume(singleton(TOPIC_PARTITION2)); - PowerMock.expectLastCall(); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT); + INITIAL_ASSIGNMENT.forEach(tp -> { + consumer.resume(singleton(tp)); + PowerMock.expectLastCall(); + }); PowerMock.replayAll(); @@ -437,6 +439,111 @@ public void testPollRedelivery() throws Exception { PowerMock.verifyAll(); } + @Test + public void testPollRedeliveryWithConsumerRebalance() throws Exception { + createTask(initialState); + + expectInitializeTask(); + expectTaskGetTopic(true); + expectPollInitialAssignment(); + + // If a retriable exception is thrown, we should redeliver the same batch, pausing the consumer in the meantime + expectConsumerPoll(1); + expectConversionAndTransformation(1); + sinkTask.put(EasyMock.anyObject()); + EasyMock.expectLastCall().andThrow(new RetriableException("retry")); + // Pause + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT); + consumer.pause(INITIAL_ASSIGNMENT); + PowerMock.expectLastCall(); + + // Empty consumer poll (all partitions are paused) with rebalance; one new partition is assigned + EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer( + () -> { + rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet()); + rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3)); + return ConsumerRecords.empty(); + }); + Set newAssignment = new HashSet<>(Arrays.asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3)); + EasyMock.expect(consumer.assignment()).andReturn(newAssignment).times(3); + EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(FIRST_OFFSET); + sinkTask.open(Collections.singleton(TOPIC_PARTITION3)); + EasyMock.expectLastCall(); + // All partitions are re-paused in order to pause any newly-assigned partitions so that redelivery efforts can continue + consumer.pause(newAssignment); + EasyMock.expectLastCall(); + sinkTask.put(EasyMock.anyObject()); + EasyMock.expectLastCall().andThrow(new RetriableException("retry")); + + // Next delivery attempt fails again + expectConsumerPoll(0); + sinkTask.put(EasyMock.anyObject()); + EasyMock.expectLastCall().andThrow(new RetriableException("retry")); + + // Non-empty consumer poll; all initially-assigned partitions are revoked in rebalance, and new partitions are allowed to resume + ConsumerRecord newRecord = new ConsumerRecord<>(TOPIC, PARTITION3, FIRST_OFFSET, RAW_KEY, RAW_VALUE); + EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer( + () -> { + rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT); + rebalanceListener.getValue().onPartitionsAssigned(Collections.emptyList()); + return new ConsumerRecords<>(Collections.singletonMap(TOPIC_PARTITION3, Collections.singletonList(newRecord))); + }); + newAssignment = Collections.singleton(TOPIC_PARTITION3); + EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(newAssignment)).times(3); + final Map offsets = INITIAL_ASSIGNMENT.stream() + .collect(Collectors.toMap(Function.identity(), tp -> new OffsetAndMetadata(FIRST_OFFSET))); + sinkTask.preCommit(offsets); + EasyMock.expectLastCall().andReturn(offsets); + sinkTask.close(INITIAL_ASSIGNMENT); + EasyMock.expectLastCall(); + // All partitions are resumed, as all previously paused-for-redelivery partitions were revoked + newAssignment.forEach(tp -> { + consumer.resume(Collections.singleton(tp)); + EasyMock.expectLastCall(); + }); + expectConversionAndTransformation(1); + sinkTask.put(EasyMock.anyObject()); + EasyMock.expectLastCall(); + + PowerMock.replayAll(); + + workerTask.initialize(TASK_CONFIG); + workerTask.initializeAndStart(); + workerTask.iteration(); + workerTask.iteration(); + workerTask.iteration(); + workerTask.iteration(); + workerTask.iteration(); + + PowerMock.verifyAll(); + } + + @Test + public void testErrorInRebalancePartitionLoss() throws Exception { + RuntimeException exception = new RuntimeException("Revocation error"); + + createTask(initialState); + + expectInitializeTask(); + expectTaskGetTopic(true); + expectPollInitialAssignment(); + expectRebalanceLossError(exception); + + PowerMock.replayAll(); + + workerTask.initialize(TASK_CONFIG); + workerTask.initializeAndStart(); + workerTask.iteration(); + try { + workerTask.iteration(); + fail("Poll should have raised the rebalance exception"); + } catch (RuntimeException e) { + assertEquals(exception, e); + } + + PowerMock.verifyAll(); + } + @Test public void testErrorInRebalancePartitionRevocation() throws Exception { RuntimeException exception = new RuntimeException("Revocation error"); @@ -489,6 +596,74 @@ public void testErrorInRebalancePartitionAssignment() throws Exception { PowerMock.verifyAll(); } + @Test + public void testPartialRevocationAndAssignment() throws Exception { + createTask(initialState); + + expectInitializeTask(); + expectTaskGetTopic(true); + expectPollInitialAssignment(); + + EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer( + () -> { + rebalanceListener.getValue().onPartitionsRevoked(Collections.singleton(TOPIC_PARTITION)); + rebalanceListener.getValue().onPartitionsAssigned(Collections.emptySet()); + return ConsumerRecords.empty(); + }); + EasyMock.expect(consumer.assignment()).andReturn(Collections.singleton(TOPIC_PARTITION)).times(2); + final Map offsets = new HashMap<>(); + offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET)); + sinkTask.preCommit(offsets); + EasyMock.expectLastCall().andReturn(offsets); + sinkTask.close(Collections.singleton(TOPIC_PARTITION)); + EasyMock.expectLastCall(); + sinkTask.put(Collections.emptyList()); + EasyMock.expectLastCall(); + + EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer( + () -> { + rebalanceListener.getValue().onPartitionsRevoked(Collections.emptySet()); + rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION3)); + return ConsumerRecords.empty(); + }); + EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(Arrays.asList(TOPIC_PARTITION2, TOPIC_PARTITION3))).times(2); + EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(FIRST_OFFSET); + sinkTask.open(Collections.singleton(TOPIC_PARTITION3)); + EasyMock.expectLastCall(); + sinkTask.put(Collections.emptyList()); + EasyMock.expectLastCall(); + + EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer( + () -> { + rebalanceListener.getValue().onPartitionsLost(Collections.singleton(TOPIC_PARTITION3)); + rebalanceListener.getValue().onPartitionsAssigned(Collections.singleton(TOPIC_PARTITION)); + return ConsumerRecords.empty(); + }); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(4); + sinkTask.close(Collections.singleton(TOPIC_PARTITION3)); + EasyMock.expectLastCall(); + EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(FIRST_OFFSET); + sinkTask.open(Collections.singleton(TOPIC_PARTITION)); + EasyMock.expectLastCall(); + sinkTask.put(Collections.emptyList()); + EasyMock.expectLastCall(); + + PowerMock.replayAll(); + + workerTask.initialize(TASK_CONFIG); + workerTask.initializeAndStart(); + // First iteration--first call to poll, first consumer assignment + workerTask.iteration(); + // Second iteration--second call to poll, partial consumer revocation + workerTask.iteration(); + // Third iteration--third call to poll, partial consumer assignment + workerTask.iteration(); + // Fourth iteration--fourth call to poll, one partition lost; can't commit offsets for it, one new partition assigned + workerTask.iteration(); + + PowerMock.verifyAll(); + } + @Test public void testWakeupInCommitSyncCausesRetry() throws Exception { createTask(initialState); @@ -502,8 +677,6 @@ public void testWakeupInCommitSyncCausesRetry() throws Exception { sinkTask.put(EasyMock.anyObject()); EasyMock.expectLastCall(); - final List partitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2); - final Map offsets = new HashMap<>(); offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1)); offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET)); @@ -518,29 +691,26 @@ public void testWakeupInCommitSyncCausesRetry() throws Exception { consumer.commitSync(EasyMock.>anyObject()); EasyMock.expectLastCall(); - sinkTask.close(new HashSet<>(partitions)); + sinkTask.close(INITIAL_ASSIGNMENT); EasyMock.expectLastCall(); - EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(FIRST_OFFSET); - EasyMock.expect(consumer.position(TOPIC_PARTITION2)).andReturn(FIRST_OFFSET); + INITIAL_ASSIGNMENT.forEach(tp -> EasyMock.expect(consumer.position(tp)).andReturn(FIRST_OFFSET)); - sinkTask.open(partitions); + sinkTask.open(INITIAL_ASSIGNMENT); EasyMock.expectLastCall(); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(5); EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer( () -> { - rebalanceListener.getValue().onPartitionsRevoked(partitions); - rebalanceListener.getValue().onPartitionsAssigned(partitions); + rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT); + rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT); return ConsumerRecords.empty(); }); - EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(partitions)); - - consumer.resume(Collections.singleton(TOPIC_PARTITION)); - EasyMock.expectLastCall(); - - consumer.resume(Collections.singleton(TOPIC_PARTITION2)); - EasyMock.expectLastCall(); + INITIAL_ASSIGNMENT.forEach(tp -> { + consumer.resume(Collections.singleton(tp)); + EasyMock.expectLastCall(); + }); statusListener.onResume(taskId); EasyMock.expectLastCall(); @@ -604,6 +774,8 @@ public void testWakeupNotThrownDuringShutdown() throws Exception { sinkTask.put(EasyMock.eq(Collections.emptyList())); EasyMock.expectLastCall(); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(1); + final Map offsets = new HashMap<>(); offsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1)); offsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET)); @@ -650,6 +822,8 @@ public void testRequestCommit() throws Exception { sinkTask.preCommit(offsets); EasyMock.expectLastCall().andReturn(offsets); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2); + final Capture callback = EasyMock.newCapture(); consumer.commitAsync(EasyMock.eq(offsets), EasyMock.capture(callback)); EasyMock.expectLastCall().andAnswer(() -> { @@ -770,7 +944,7 @@ public void testPreCommit() throws Exception { sinkTask.preCommit(workerCurrentOffsets); EasyMock.expectLastCall().andReturn(taskOffsets); // Expect extra invalid topic partition to be filtered, which causes the consumer assignment to be logged - EasyMock.expect(consumer.assignment()).andReturn(workerCurrentOffsets.keySet()); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2); final Capture callback = EasyMock.newCapture(); consumer.commitAsync(EasyMock.eq(committableOffsets), EasyMock.capture(callback)); EasyMock.expectLastCall().andAnswer(() -> { @@ -823,6 +997,8 @@ public void testIgnoredCommit() throws Exception { workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1)); workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET)); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2); + // iter 3 sinkTask.preCommit(workerCurrentOffsets); EasyMock.expectLastCall().andReturn(workerStartingOffsets); @@ -874,6 +1050,8 @@ public void testLongRunningCommitWithoutTimeout() throws Exception { workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1)); workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET)); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2); + // iter 3 - note that we return the current offset to indicate they should be committed sinkTask.preCommit(workerCurrentOffsets); EasyMock.expectLastCall().andReturn(workerCurrentOffsets); @@ -947,6 +1125,8 @@ public void testSinkTasksHandleCloseErrors() throws Exception { expectInitializeTask(); expectTaskGetTopic(true); + expectPollInitialAssignment(); + // Put one message through the task to get some offsets to commit expectConsumerPoll(1); expectConversionAndTransformation(1); @@ -990,6 +1170,8 @@ public void testSuppressCloseErrors() throws Exception { expectInitializeTask(); expectTaskGetTopic(true); + expectPollInitialAssignment(); + // Put one message through the task to get some offsets to commit expectConsumerPoll(1); expectConversionAndTransformation(1); @@ -1052,7 +1234,7 @@ public void testCommitWithOutOfOrderCallback() throws Exception { workerCurrentOffsets.put(TOPIC_PARTITION, new OffsetAndMetadata(FIRST_OFFSET + 1)); workerCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET)); - final List originalPartitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2); + final List originalPartitions = new ArrayList<>(INITIAL_ASSIGNMENT); final List rebalancedPartitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3); final Map rebalanceOffsets = new HashMap<>(); rebalanceOffsets.put(TOPIC_PARTITION, workerCurrentOffsets.get(TOPIC_PARTITION)); @@ -1064,6 +1246,8 @@ public void testCommitWithOutOfOrderCallback() throws Exception { postRebalanceCurrentOffsets.put(TOPIC_PARTITION2, new OffsetAndMetadata(FIRST_OFFSET)); postRebalanceCurrentOffsets.put(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 2)); + EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(originalPartitions)).times(2); + // iter 3 - note that we return the current offset to indicate they should be committed sinkTask.preCommit(workerCurrentOffsets); EasyMock.expectLastCall().andReturn(workerCurrentOffsets); @@ -1126,7 +1310,7 @@ public void testCommitWithOutOfOrderCallback() throws Exception { EasyMock.expectLastCall().andReturn(workerCurrentOffsets); sinkTask.put(EasyMock.anyObject()); EasyMock.expectLastCall(); - sinkTask.close(workerCurrentOffsets.keySet()); + sinkTask.close(new ArrayList<>(workerCurrentOffsets.keySet())); EasyMock.expectLastCall(); consumer.commitSync(workerCurrentOffsets); EasyMock.expectLastCall(); @@ -1138,9 +1322,10 @@ public void testCommitWithOutOfOrderCallback() throws Exception { EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(offsetTp1); EasyMock.expect(consumer.position(TOPIC_PARTITION2)).andReturn(offsetTp2); EasyMock.expect(consumer.position(TOPIC_PARTITION3)).andReturn(offsetTp3); + EasyMock.expect(consumer.assignment()).andReturn(new HashSet<>(rebalancedPartitions)).times(6); // onPartitionsAssigned - step 2 - sinkTask.open(rebalancedPartitions); + sinkTask.open(EasyMock.eq(rebalancedPartitions)); EasyMock.expectLastCall(); // onPartitionsAssigned - step 3 rewind @@ -1259,6 +1444,8 @@ public void testDeliveryWithMutatingTransform() throws Exception { sinkTask.preCommit(offsets); EasyMock.expectLastCall().andReturn(offsets); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2); + final Capture callback = EasyMock.newCapture(); consumer.commitAsync(EasyMock.eq(offsets), EasyMock.capture(callback)); EasyMock.expectLastCall().andAnswer(() -> { @@ -1371,9 +1558,8 @@ public void testTopicsRegex() throws Exception { expectPollInitialAssignment(); - Set partitions = new HashSet<>(asList(TOPIC_PARTITION, TOPIC_PARTITION2)); - EasyMock.expect(consumer.assignment()).andReturn(partitions); - consumer.pause(partitions); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT); + consumer.pause(INITIAL_ASSIGNMENT); PowerMock.expectLastCall(); PowerMock.replayAll(); @@ -1538,10 +1724,19 @@ private void expectInitializeTask() throws Exception { PowerMock.expectLastCall(); } - private void expectRebalanceRevocationError(RuntimeException e) { - final List partitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2); + private void expectRebalanceLossError(RuntimeException e) { + sinkTask.close(new HashSet<>(INITIAL_ASSIGNMENT)); + EasyMock.expectLastCall().andThrow(e); + + EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer( + () -> { + rebalanceListener.getValue().onPartitionsLost(INITIAL_ASSIGNMENT); + return ConsumerRecords.empty(); + }); + } - sinkTask.close(new HashSet<>(partitions)); + private void expectRebalanceRevocationError(RuntimeException e) { + sinkTask.close(new HashSet<>(INITIAL_ASSIGNMENT)); EasyMock.expectLastCall().andThrow(e); sinkTask.preCommit(EasyMock.anyObject()); @@ -1549,15 +1744,13 @@ private void expectRebalanceRevocationError(RuntimeException e) { EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer( () -> { - rebalanceListener.getValue().onPartitionsRevoked(partitions); + rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT); return ConsumerRecords.empty(); }); } private void expectRebalanceAssignmentError(RuntimeException e) { - final List partitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2); - - sinkTask.close(new HashSet<>(partitions)); + sinkTask.close(INITIAL_ASSIGNMENT); EasyMock.expectLastCall(); sinkTask.preCommit(EasyMock.anyObject()); @@ -1566,29 +1759,29 @@ private void expectRebalanceAssignmentError(RuntimeException e) { EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(FIRST_OFFSET); EasyMock.expect(consumer.position(TOPIC_PARTITION2)).andReturn(FIRST_OFFSET); - sinkTask.open(partitions); + sinkTask.open(INITIAL_ASSIGNMENT); EasyMock.expectLastCall().andThrow(e); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(3); EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer( () -> { - rebalanceListener.getValue().onPartitionsRevoked(partitions); - rebalanceListener.getValue().onPartitionsAssigned(partitions); + rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT); + rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT); return ConsumerRecords.empty(); }); } private void expectPollInitialAssignment() { - final List partitions = asList(TOPIC_PARTITION, TOPIC_PARTITION2); - - sinkTask.open(partitions); + sinkTask.open(INITIAL_ASSIGNMENT); EasyMock.expectLastCall(); + EasyMock.expect(consumer.assignment()).andReturn(INITIAL_ASSIGNMENT).times(2); + EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(() -> { - rebalanceListener.getValue().onPartitionsAssigned(partitions); + rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT); return ConsumerRecords.empty(); }); - EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(FIRST_OFFSET); - EasyMock.expect(consumer.position(TOPIC_PARTITION2)).andReturn(FIRST_OFFSET); + INITIAL_ASSIGNMENT.forEach(tp -> EasyMock.expect(consumer.position(tp)).andReturn(FIRST_OFFSET)); sinkTask.put(Collections.emptyList()); EasyMock.expectLastCall(); diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java index 7b4474bc65644..9dec1a1ddb323 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java @@ -65,6 +65,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.Set; import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; @@ -92,6 +93,8 @@ public class WorkerSinkTaskThreadedTest extends ThreadedTest { private static final TopicPartition TOPIC_PARTITION2 = new TopicPartition(TOPIC, PARTITION2); private static final TopicPartition TOPIC_PARTITION3 = new TopicPartition(TOPIC, PARTITION3); private static final TopicPartition UNASSIGNED_TOPIC_PARTITION = new TopicPartition(TOPIC, 200); + private static final Set INITIAL_ASSIGNMENT = new HashSet<>(Arrays.asList( + TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3)); private static final Map TASK_PROPS = new HashMap<>(); private static final long TIMESTAMP = 42L; @@ -202,6 +205,7 @@ public void testCommit() throws Exception { expectInitializeTask(); expectTaskGetTopic(true); expectPollInitialAssignment(); + expectConsumerAssignment(INITIAL_ASSIGNMENT).times(2); // Make each poll() take the offset commit interval Capture> capturedRecords @@ -236,6 +240,7 @@ public void testCommitFailure() throws Exception { expectInitializeTask(); expectTaskGetTopic(true); expectPollInitialAssignment(); + expectConsumerAssignment(INITIAL_ASSIGNMENT); Capture> capturedRecords = expectPolls(WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_DEFAULT); expectOffsetCommit(1L, new RuntimeException(), null, 0, true); @@ -276,6 +281,7 @@ public void testCommitSuccessFollowedByFailure() throws Exception { expectInitializeTask(); expectTaskGetTopic(true); expectPollInitialAssignment(); + expectConsumerAssignment(INITIAL_ASSIGNMENT).times(3); Capture> capturedRecords = expectPolls(WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_DEFAULT); expectOffsetCommit(1L, null, null, 0, true); expectOffsetCommit(2L, new RuntimeException(), null, 0, true); @@ -315,6 +321,7 @@ public void testCommitConsumerFailure() throws Exception { expectInitializeTask(); expectTaskGetTopic(true); expectPollInitialAssignment(); + expectConsumerAssignment(INITIAL_ASSIGNMENT).times(2); Capture> capturedRecords = expectPolls(WorkerConfig.OFFSET_COMMIT_INTERVAL_MS_DEFAULT); @@ -347,6 +354,7 @@ public void testCommitTimeout() throws Exception { expectInitializeTask(); expectTaskGetTopic(true); expectPollInitialAssignment(); + expectConsumerAssignment(INITIAL_ASSIGNMENT).times(2); // Cut down amount of time to pass in each poll so we trigger exactly 1 offset commit Capture> capturedRecords @@ -483,6 +491,7 @@ public void testRewindOnRebalanceDuringPoll() throws Exception { expectInitializeTask(); expectTaskGetTopic(true); expectPollInitialAssignment(); + expectConsumerAssignment(INITIAL_ASSIGNMENT).times(2); expectRebalanceDuringPoll().andAnswer(() -> { Map offsets = sinkTaskContext.getValue().offsets(); @@ -515,13 +524,13 @@ private void expectInitializeTask() throws Exception { } private void expectPollInitialAssignment() throws Exception { - final List partitions = Arrays.asList(TOPIC_PARTITION, TOPIC_PARTITION2, TOPIC_PARTITION3); + expectConsumerAssignment(INITIAL_ASSIGNMENT).times(2); - sinkTask.open(partitions); + sinkTask.open(INITIAL_ASSIGNMENT); EasyMock.expectLastCall(); EasyMock.expect(consumer.poll(Duration.ofMillis(EasyMock.anyLong()))).andAnswer(() -> { - rebalanceListener.getValue().onPartitionsAssigned(partitions); + rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT); return ConsumerRecords.empty(); }); EasyMock.expect(consumer.position(TOPIC_PARTITION)).andReturn(FIRST_OFFSET); @@ -532,6 +541,10 @@ private void expectPollInitialAssignment() throws Exception { EasyMock.expectLastCall(); } + private IExpectationSetters> expectConsumerAssignment(Set assignment) { + return EasyMock.expect(consumer.assignment()).andReturn(assignment); + } + private void expectStopTask() throws Exception { sinkTask.stop(); PowerMock.expectLastCall(); diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporterTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporterTest.java index 07a4f9e8f1197..2d78297db5def 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporterTest.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/errors/WorkerErrantRecordReporterTest.java @@ -17,6 +17,7 @@ package org.apache.kafka.connect.runtime.errors; +import org.apache.kafka.common.TopicPartition; import org.apache.kafka.connect.sink.SinkRecord; import org.apache.kafka.connect.storage.Converter; import org.apache.kafka.connect.storage.HeaderConverter; @@ -27,6 +28,9 @@ import org.powermock.core.classloader.annotations.PowerMockIgnore; import org.powermock.modules.junit4.PowerMockRunner; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; import java.util.concurrent.CompletableFuture; import static org.junit.Assert.assertFalse; @@ -62,13 +66,16 @@ public void setup() { } @Test - public void testGetAllFutures() { + public void testGetFutures() { + Collection topicPartitions = new ArrayList<>(); assertTrue(reporter.futures.isEmpty()); for (int i = 0; i < 4; i++) { - reporter.futures.add(CompletableFuture.completedFuture(null)); + TopicPartition topicPartition = new TopicPartition("topic", i); + topicPartitions.add(topicPartition); + reporter.futures.put(topicPartition, Collections.singletonList(CompletableFuture.completedFuture(null))); } assertFalse(reporter.futures.isEmpty()); - reporter.awaitAllFutures(); + reporter.awaitFutures(topicPartitions); assertTrue(reporter.futures.isEmpty()); } } diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java index a8328c5069460..8bf8775e25bac 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java @@ -366,6 +366,19 @@ public void createTopic(String topic, int partitions, int replication, Map