KAFKA-9184: Redundant task creation and periodic rebalances after zombie Connect worker rejoins the group

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java

@@ -1538,10 +1538,11 @@ public void onAssigned(ExtendedAssignment assignment, int generation) {
             short priorProtocolVersion = currentProtocolVersion;
             DistributedHerder.this.currentProtocolVersion = member.currentProtocolVersion();
             log.info(
-                "Joined group at generation {} with protocol version {} and got assignment: {}",
+                "Joined group at generation {} with protocol version {} and got assignment: {} with rebalance delay: {}",
                 generation,
                 DistributedHerder.this.currentProtocolVersion,
-                assignment
+                assignment,
+                assignment.delay()
             );
             synchronized (DistributedHerder.this) {
                 DistributedHerder.this.assignment = assignment;
@@ -1611,12 +1612,13 @@ public void onRevoked(String leader, Collection<String> connectors, Collection<C
 
                 // The actual timeout for graceful task stop is applied in worker's stopAndAwaitTask method.
                 startAndStop(callables);
+                log.info("Finished stopping tasks in preparation for rebalance");
 
                 // Ensure that all status updates have been pushed to the storage system before rebalancing.
                 // Otherwise, we may inadvertently overwrite the state with a stale value after the rebalance
                 // completes.
                 statusBackingStore.flush();
-                log.info("Finished stopping tasks in preparation for rebalance");
+                log.info("Finished flushing status backing store in preparation for rebalance");
             } else {
                 log.info("Wasn't unable to resume work after last rebalance, can skip stopping connectors and tasks");
             }

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignor.java

@@ -153,6 +153,9 @@ private Long ensureLeaderConfig(long maxOffset, WorkerCoordinator coordinator) {
     protected Map<String, ByteBuffer> performTaskAssignment(String leaderId, long maxOffset,
                                                             Map<String, ExtendedWorkerState> memberConfigs,
                                                             WorkerCoordinator coordinator, short protocolVersion) {
+        log.debug("Performing task assignment during generation: {} with memberId: {}",
+                coordinator.generationId(), coordinator.memberId());
+
         // Base set: The previous assignment of connectors-and-tasks is a standalone snapshot that
         // can be used to calculate derived sets
         log.debug("Previous assignments: {}", previousAssignment);
@@ -350,6 +353,7 @@ protected void handleLostAssignments(ConnectorsAndTasks lostAssignments,
                                          ConnectorsAndTasks newSubmissions,
                                          List<WorkerLoad> completeWorkerAssignment) {
         if (lostAssignments.isEmpty()) {
+            resetDelay();
             return;
         }
 
@@ -359,40 +363,53 @@ protected void handleLostAssignments(ConnectorsAndTasks lostAssignments,
 
         if (scheduledRebalance > 0 && now >= scheduledRebalance) {
             // delayed rebalance expired and it's time to assign resources
+            log.debug("Delayed rebalance expired. Reassigning lost tasks");
             Optional<WorkerLoad> candidateWorkerLoad = Optional.empty();
             if (!candidateWorkersForReassignment.isEmpty()) {
                 candidateWorkerLoad = pickCandidateWorkerForReassignment(completeWorkerAssignment);
             }
 
             if (candidateWorkerLoad.isPresent()) {
                 WorkerLoad workerLoad = candidateWorkerLoad.get();
+                log.debug("A candidate worker has been found to assign lost tasks: {}", workerLoad.worker());
                 lostAssignments.connectors().forEach(workerLoad::assign);
                 lostAssignments.tasks().forEach(workerLoad::assign);
             } else {
+                log.debug("No single candidate worker was found to assign lost tasks. Treating lost tasks as new tasks");
                 newSubmissions.connectors().addAll(lostAssignments.connectors());
                 newSubmissions.tasks().addAll(lostAssignments.tasks());
             }
-            candidateWorkersForReassignment.clear();
-            scheduledRebalance = 0;
-            delay = 0;
+            resetDelay();
         } else {
             candidateWorkersForReassignment
                     .addAll(candidateWorkersForReassignment(completeWorkerAssignment));
             if (now < scheduledRebalance) {
                 // a delayed rebalance is in progress, but it's not yet time to reassign
                 // unaccounted resources
                 delay = calculateDelay(now);
+                log.debug("Delayed rebalance in progress. Task reassignment is postponed. New computed rebalance delay: {}", delay);
             } else {
                 // This means scheduledRebalance == 0
                 // We could also also extract the current minimum delay from the group, to make
                 // independent of consecutive leader failures, but this optimization is skipped
                 // at the moment
                 delay = maxDelay;
+                log.debug("Resetting rebalance delay to the max: {}. scheduledRebalance: {} now: {} diff scheduledRebalance - now: {}",
+                        delay, scheduledRebalance, now, scheduledRebalance - now);
             }
             scheduledRebalance = now + delay;
         }
     }
 
+    private void resetDelay() {
+        candidateWorkersForReassignment.clear();
+        scheduledRebalance = 0;
+        if (delay != 0) {
+            log.debug("Resetting delay from previous value: {} to 0", delay);
+        }
+        delay = 0;
+    }
+
     private Set<String> candidateWorkersForReassignment(List<WorkerLoad> completeWorkerAssignment) {
         return completeWorkerAssignment.stream()
                 .filter(WorkerLoad::isEmpty)

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java

@@ -66,6 +66,7 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
     private volatile ConnectProtocolCompatibility currentConnectProtocol;
     private final ConnectAssignor eagerAssignor;
     private final ConnectAssignor incrementalAssignor;
+    private final int coordinatorDiscoveryTimeoutMs;
 
     /**
      * Initialize the coordination manager.
@@ -98,6 +99,7 @@ public WorkerCoordinator(GroupRebalanceConfig config,
         this.incrementalAssignor = new IncrementalCooperativeAssignor(logContext, time, maxDelay);
         this.eagerAssignor = new EagerAssignor(logContext);
         this.currentConnectProtocol = protocolCompatibility;
+        this.coordinatorDiscoveryTimeoutMs = config.heartbeatIntervalMs;
     }
 
     @Override
@@ -124,7 +126,19 @@ public void poll(long timeout) {
 
         do {
             if (coordinatorUnknown()) {
-                ensureCoordinatorReady(time.timer(Long.MAX_VALUE));
+                log.debug("Broker coordinator is marked unknown. Attempting discovery with a timeout of {}ms",
+                        coordinatorDiscoveryTimeoutMs);
+                if (ensureCoordinatorReady(time.timer(coordinatorDiscoveryTimeoutMs))) {
+                    log.debug("Broker coordinator is ready");
+                } else {
+                    log.debug("Can not connect to broker coordinator");
+                    if (assignmentSnapshot != null && !assignmentSnapshot.failed()) {
+                        log.info("Broker coordinator was unreachable for {}ms. Revoking previous assignment {} to " +
+                                "avoid running tasks while not being a member the group", coordinatorDiscoveryTimeoutMs, assignmentSnapshot);
+                        listener.onRevoked(assignmentSnapshot.leader(), assignmentSnapshot.connectors(), assignmentSnapshot.tasks());
+                        assignmentSnapshot = null;
+                    }
+                }
                 now = time.milliseconds();
             }
 
@@ -227,6 +241,16 @@ public String memberId() {
         return JoinGroupRequest.UNKNOWN_MEMBER_ID;
     }
 
+    /**
+     * Return the current generation. The generation refers to this worker's knowledge with
+     * respect to which  generation is the latest one and, therefore, this information is local.
+     *
+     * @return the generation ID or -1 if no generation is defined
+     */
+    public int generationId() {
+        return super.generation().generationId;
+    }
+
     private boolean isLeader() {
         return assignmentSnapshot != null && memberId().equals(assignmentSnapshot.leader());
     }

connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMember.java

@@ -151,6 +151,10 @@ public void stop() {
         stop(false);
     }
 
+    /**
+     * Ensure that the connection to the broker coordinator is up and that the worker is an
+     * active member of the group.
+     */
     public void ensureActive() {
         coordinator.poll(0);
     }

connect/runtime/src/test/java/org/apache/kafka/connect/integration/ConnectWorkerIntegrationTest.java

@@ -17,6 +17,7 @@
 package org.apache.kafka.connect.integration;
 
 import org.apache.kafka.connect.runtime.AbstractStatus;
+import org.apache.kafka.connect.runtime.distributed.DistributedConfig;
 import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
 import org.apache.kafka.connect.storage.StringConverter;
 import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
@@ -65,29 +66,27 @@ public class ConnectWorkerIntegrationTest {
     private static final int NUM_WORKERS = 3;
     private static final String CONNECTOR_NAME = "simple-source";
 
+    private EmbeddedConnectCluster.Builder connectBuilder;
     private EmbeddedConnectCluster connect;
+    Map<String, String> workerProps = new HashMap<>();
+    Properties brokerProps = new Properties();
 
     @Before
     public void setup() throws IOException {
         // setup Connect worker properties
-        Map<String, String> workerProps = new HashMap<>();
         workerProps.put(OFFSET_COMMIT_INTERVAL_MS_CONFIG, String.valueOf(OFFSET_COMMIT_INTERVAL_MS));
         workerProps.put(CONNECTOR_CLIENT_POLICY_CLASS_CONFIG, "All");
 
         // setup Kafka broker properties
-        Properties brokerProps = new Properties();
         brokerProps.put("auto.create.topics.enable", String.valueOf(false));
 
         // build a Connect cluster backed by Kafka and Zk
-        connect = new EmbeddedConnectCluster.Builder()
+        connectBuilder = new EmbeddedConnectCluster.Builder()
                 .name("connect-cluster")
                 .numWorkers(NUM_WORKERS)
                 .workerProps(workerProps)
                 .brokerProps(brokerProps)
-                .build();
-
-        // start the clusters
-        connect.start();
+                .maskExitProcedures(true); // true is the default, setting here as example
     }
 
     @After
@@ -102,6 +101,10 @@ public void close() {
      */
     @Test
     public void testAddAndRemoveWorker() throws Exception {
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
         int numTasks = 4;
         // create test topic
         connect.kafka().createTopic("test-topic", NUM_TOPIC_PARTITIONS);
@@ -149,6 +152,10 @@ public void testAddAndRemoveWorker() throws Exception {
      */
     @Test
     public void testRestartFailedTask() throws Exception {
+        connect = connectBuilder.build();
+        // start the clusters
+        connect.start();
+
         int numTasks = 1;
 
         // Properties for the source connector. The task should fail at startup due to the bad broker address.
@@ -180,6 +187,56 @@ public void testRestartFailedTask() throws Exception {
             CONNECTOR_SETUP_DURATION_MS, "Connector tasks are not all in running state.");
     }
 
+    /**
+     * Verify that a set of tasks restarts correctly after a broker goes offline and back online
+     */
+    @Test
+    public void testBrokerCoordinator() throws Exception {
+        workerProps.put(DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG, String.valueOf(5000));
+        connect = connectBuilder.workerProps(workerProps).build();
+        // start the clusters
+        connect.start();
+        int numTasks = 4;
+        // create test topic
+        connect.kafka().createTopic("test-topic", NUM_TOPIC_PARTITIONS);
+
+        // setup up props for the sink connector
+        Map<String, String> props = new HashMap<>();
+        props.put(CONNECTOR_CLASS_CONFIG, MonitorableSourceConnector.class.getSimpleName());
+        props.put(TASKS_MAX_CONFIG, String.valueOf(numTasks));
+        props.put("topic", "test-topic");
+        props.put("throughput", String.valueOf(1));
+        props.put("messages.per.poll", String.valueOf(10));
+        props.put(KEY_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+        props.put(VALUE_CONVERTER_CLASS_CONFIG, StringConverter.class.getName());
+
+        waitForCondition(() -> assertWorkersUp(NUM_WORKERS).orElse(false),
+                WORKER_SETUP_DURATION_MS, "Initial group of workers did not start in time.");
+
+        // start a source connector
+        connect.configureConnector(CONNECTOR_NAME, props);
+
+        waitForCondition(() -> assertConnectorAndTasksRunning(CONNECTOR_NAME, numTasks).orElse(false),
+                CONNECTOR_SETUP_DURATION_MS, "Connector tasks did not start in time.");
+
+        connect.kafka().stopOnlyKafka();
+
+        waitForCondition(() -> assertWorkersUp(NUM_WORKERS).orElse(false),
+                WORKER_SETUP_DURATION_MS, "Group of workers did not remain the same after broker shutdown");
+
+        connect.kafka().startOnlyKafkaOnSamePorts();
+
+        // Allow for the workers to discover that the coordinator is unavailable
+        Thread.sleep(TimeUnit.SECONDS.toMillis(10));
+
+        waitForCondition(() -> assertWorkersUp(NUM_WORKERS).orElse(false),
+                WORKER_SETUP_DURATION_MS, "Group of workers did not remain the same within the "
+                        + "designated time.");
+
+        waitForCondition(() -> assertConnectorAndTasksRunning(CONNECTOR_NAME, numTasks).orElse(false),
+                CONNECTOR_SETUP_DURATION_MS, "Connector tasks did not start in time.");
+    }
+
     /**
      * Confirm that the requested number of workers is up and running.
      *

connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignorTest.java

@@ -170,6 +170,8 @@ public void testTaskAssignmentWhenWorkerJoins() {
 
         verify(coordinator, times(rebalanceNum)).configSnapshot();
         verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
     }
 
     @Test
@@ -234,6 +236,8 @@ public void testTaskAssignmentWhenWorkerLeavesPermanently() {
 
         verify(coordinator, times(rebalanceNum)).configSnapshot();
         verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
     }
 
     @Test
@@ -314,6 +318,8 @@ public void testTaskAssignmentWhenWorkerBounces() {
 
         verify(coordinator, times(rebalanceNum)).configSnapshot();
         verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
     }
 
     @Test
@@ -369,6 +375,8 @@ public void testTaskAssignmentWhenLeaderLeavesPermanently() {
 
         verify(coordinator, times(rebalanceNum)).configSnapshot();
         verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
     }
 
     @Test
@@ -438,6 +446,8 @@ public void testTaskAssignmentWhenLeaderBounces() {
 
         verify(coordinator, times(rebalanceNum)).configSnapshot();
         verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
     }
 
     @Test
@@ -507,6 +517,8 @@ public void testTaskAssignmentWhenFirstAssignmentAttemptFails() {
 
         verify(coordinator, times(rebalanceNum)).configSnapshot();
         verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
     }
 
     @Test
@@ -563,6 +575,8 @@ public void testTaskAssignmentWhenSubsequentAssignmentAttemptFails() {
 
         verify(coordinator, times(rebalanceNum)).configSnapshot();
         verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
     }
 
     @Test
@@ -596,6 +610,8 @@ public void testTaskAssignmentWhenConnectorsAreDeleted() {
 
         verify(coordinator, times(rebalanceNum)).configSnapshot();
         verify(coordinator, times(rebalanceNum)).leaderState(any());
+        verify(coordinator, times(rebalanceNum)).generationId();
+        verify(coordinator, times(rebalanceNum)).memberId();
     }
 
     @Test