KAFKA-12525: Ignoring stale status statuses when reading from Connect status topic

connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java

@@ -588,6 +588,25 @@ private void readTaskStatus(String key, byte[] value) {
         synchronized (this) {
             log.trace("Received task {} status update {}", id, status);
             CacheEntry<TaskStatus> entry = getOrAdd(id);
+
+            // During frequent rebalances, there could be a race condition because of which
+            // an UNASSIGNED state of a prior or same generation can be sent by a worker despite a
+            // RUNNING status by another worker because the first worker
+            // couldn't read the latest RUNNING status. This can lead to an inaccurate status
+            // representation even though the task might be actually running.
+            // Note that this could also mean that when a generation reset happens, and an
+            // UNASSIGNED status is sent, then it would be ignored if the current status is RUNNING
+            // at a higher generation. But since it will be followed by a RUNNING or a different
+            // status message(at the lower generation) soon after, the misrepresentation of the UNASSIGNED
+            // status would be short-lived in most cases.
+            if (status.state() == TaskStatus.State.UNASSIGNED
+                    && entry.get() != null
+                    && entry.get().state() == TaskStatus.State.RUNNING
+                    && entry.get().generation() >= status.generation()) {
+                log.trace("Ignoring stale status {} in favour of more upto date status {}", status, entry.get());
+                return;
+            }
+
             entry.put(status);
         }
     }

connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreTest.java

@@ -309,6 +309,58 @@ public void readTaskState() {
         assertEquals(status, store.get(TASK));
     }
 
+    @Test
+    public void readTaskStateShouldIgnoreStaleStatusesFromOtherWorkers() {
+        byte[] value = new byte[0];
+        String otherWorkerId = "anotherhost:8083";
+        String yetAnotherWorkerId = "yetanotherhost:8083";
+
+        // This worker sends a RUNNING status in the most recent generation
+        Map<String, Object> firstStatusRead = new HashMap<>();
+        firstStatusRead.put("worker_id", otherWorkerId);
+        firstStatusRead.put("state", "RUNNING");
+        firstStatusRead.put("generation", 10L);
+
+        // Another worker still ends up producing an UNASSIGNED status before it could
+        // read the newer RUNNING status from above belonging to an older generation.
+        Map<String, Object> secondStatusRead = new HashMap<>();
+        secondStatusRead.put("worker_id", WORKER_ID);
+        secondStatusRead.put("state", "UNASSIGNED");
+        secondStatusRead.put("generation", 9L);
+
+        Map<String, Object> thirdStatusRead = new HashMap<>();
+        thirdStatusRead.put("worker_id", yetAnotherWorkerId);
+        thirdStatusRead.put("state", "RUNNING");
+        thirdStatusRead.put("generation", 1L);
+
+        when(converter.toConnectData(STATUS_TOPIC, value))
+                .thenReturn(new SchemaAndValue(null, firstStatusRead))
+                .thenReturn(new SchemaAndValue(null, secondStatusRead))
+                .thenReturn(new SchemaAndValue(null, thirdStatusRead));
+
+        when(converter.fromConnectData(eq(STATUS_TOPIC), any(Schema.class), any(Struct.class)))
+                .thenReturn(value);
+
+        doAnswer(invocation -> {
+            ((Callback) invocation.getArgument(2)).onCompletion(null, null);
+            store.read(consumerRecord(2, "status-task-conn-0", value));
+            return null;
+        }).when(kafkaBasedLog).send(eq("status-task-conn-0"), eq(value), any(Callback.class));
+
+        store.read(consumerRecord(0, "status-task-conn-0", value));
+        store.read(consumerRecord(1, "status-task-conn-0", value));
+
+        // The latest task status should reflect RUNNING status from the newer generation
+        TaskStatus status = new TaskStatus(TASK, TaskStatus.State.RUNNING, otherWorkerId, 10);
+        assertEquals(status, store.get(TASK));
+
+        // This  status is from the another worker not necessarily belonging to the above group.
+        // In this case, the status should get updated irrespective of whatever status info was present before.
+        TaskStatus latestStatus = new TaskStatus(TASK, TaskStatus.State.RUNNING, yetAnotherWorkerId, 1);
+        store.put(latestStatus);
+        assertEquals(latestStatus, store.get(TASK));
+    }
+
     @Test
     public void deleteConnectorState() {
         final byte[] value = new byte[0];