diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java b/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java index 20f47db2262a4..d1580e3a7c1d7 100644 --- a/connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java +++ b/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 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); } } diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreTest.java index 71fe342ee281b..4cad728ae65f7 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreTest.java +++ b/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 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 secondStatusRead = new HashMap<>(); + secondStatusRead.put("worker_id", WORKER_ID); + secondStatusRead.put("state", "UNASSIGNED"); + secondStatusRead.put("generation", 9L); + + Map 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];