Repository: apache/kafka
File: connect/runtime/src/main/java/org/apache/kafka/connect/storage/KafkaStatusBackingStore.java (Java)

Code before:
        }

        if (value == null) {
            log.trace("Removing task status for {}", id);
            remove(id);
            return;
        }

        TaskStatus status = parseTaskStatus(id, value);
        if (status == null) {
            log.warn("Failed to parse task status with key {}", key);
            return;
        }

        synchronized (this) {
            log.trace("Received task {} status update {}", id, status);
            CacheEntry<TaskStatus> entry = getOrAdd(id);
            entry.put(status);
        }
    }

    private void readTopicStatus(String key, byte[] value) {
        int delimiterPos = key.indexOf(':');
        int beginPos = TOPIC_STATUS_PREFIX.length();
        if (beginPos > delimiterPos) {
            log.warn("Discarding record with invalid topic status key {}", key);
            return;
        }

        String topic = key.substring(beginPos, delimiterPos);
        if (topic.isEmpty()) {
            log.warn("Discarding record with invalid topic status key containing empty topic {}", key);
            return;
        }

        beginPos = delimiterPos + TOPIC_STATUS_SEPARATOR.length();
        int endPos = key.length();
        if (beginPos > endPos) {
            log.warn("Discarding record with invalid topic status key {}", key);
            return;
        }

        String connector = key.substring(beginPos);
        if (connector.isEmpty()) {
            log.warn("Discarding record with invalid topic status key containing empty connector {}", key);
            return;
        }

        if (value == null) {
            log.trace("Removing status for topic {} and connector {}", topic, connector);
            removeTopic(topic, connector);

Review comment (bug):
`AbstractHerder::onShutdown` which is responsible for updating the status to `Unassigned` for a task after it has been shutdown on a worker uses `StatusBackingStore::putSafe`. This "safe" send implementation also does the same check (whether the "newer" status update is from an older generation of tasks) on pre-existing cache entries before writing to the underlying Kafka based log. However, it looks like the cache is only updated on reading back from the log (on the `KafkaBasedLog`'s work thread) and it isn't updated when doing the write itself. So is the cause for this bug that you're trying to fix essentially a race condition where the `KafkaStatusBackingStore` isn't able to read its own writes (at least immediately anyway)? If so, I wonder whether we could consider altering the design of the status backing store to update its cache on writes as well? Although I guess we could run into issues here as well since we're doing async writes so maybe it's overall better to check generations on reads through the log and ensure that the actual reflected status always corresponds to that of the latest generation. What do you think?

Edit: Never mind, I just realized that the status updates will be from different workers.

Code after:
        }

        if (value == null) {
            log.trace("Removing task status for {}", id);
            remove(id);
            return;
        }

        TaskStatus status = parseTaskStatus(id, value);
        if (status == null) {
            log.warn("Failed to parse task status with key {}", key);
            return;
        }

        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);
        }
    }

    private void readTopicStatus(String key, byte[] value) {
        int delimiterPos = key.indexOf(':');
        int beginPos = TOPIC_STATUS_PREFIX.length();
        if (beginPos > delimiterPos) {
            log.warn("Discarding record with invalid topic status key {}", key);
            return;
        }

        String topic = key.substring(beginPos, delimiterPos);
        if (topic.isEmpty()) {
            log.warn("Discarding record with invalid topic status key containing empty topic {}", key);