UNDF: UNDF-2026-000000438

kafka-0007 — Kafka Streams StreamsPartitionAssignor: PriorityQueue.contains() O(T²) in task assignment loop

Metadata

Project: Apache Kafka
Component: streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
CWE: CWE-407 (Inefficient Algorithmic Complexity)
Severity: HIGH
Complexity: O(C × T²) → O(C × T) where C = consumers/threads, T = tasks
Hot path: assignTasksToThreads() is called on every Streams rebalance

Location

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
method: assignTasksToThreads()

Defective code — PriorityQueue.contains() inside nested loop (lines ~1321–1323)

final PriorityQueue<TaskId> unassignedTasks = new PriorityQueue<>(tasksToAssign);
// ...
for (final String consumer : consumers) {                    // O(C)
    for (final TaskId task : state.prevTasksByLag(consumer)) { // O(T_prev)
        if (unassignedTasks.contains(task)) {                  // O(T) — PriorityQueue linear scan
            // ...
        }
    }
}

PriorityQueue.contains() is O(N) — it performs a linear scan of the heap array. With C consumers, each having up to T previous tasks, this loop body fires C × T_prev times, each paying O(T) for the contains check: O(C × T_prev × T).

Second defective pattern — LinkedList.contains() in follow-up loop (lines ~1367–1371)

final Queue<String> consumersToFill = new LinkedList<>();
// ...
for (final Map.Entry<TaskId, String> taskEntry : unassignedTaskToPreviousOwner.entrySet()) { // O(T)
    final TaskId task = taskEntry.getKey();
    final String consumer = taskEntry.getValue();
    if (consumersToFill.contains(consumer) && unassignedTasks.contains(task)) { // O(C) + O(T)
        // ...
        consumersToFill.remove(consumer);  // O(C) LinkedList.remove
    }
}

Both consumersToFill.contains(consumer) (LinkedList, O(C)) and unassignedTasks.contains(task) (PriorityQueue, O(T)) are linear inside the O(T) outer loop: O(T × (C + T)) = O(T²).

Fix

// Replace PriorityQueue with a HashSet for O(1) membership testing.
// Keep a separate PriorityQueue only for ordered polling.
final PriorityQueue<TaskId> unassignedTasksOrdered = new PriorityQueue<>(tasksToAssign);
final Set<TaskId> unassignedTasksSet = new HashSet<>(tasksToAssign);  // O(1) contains

// For the consumersToFill loop:
// Replace LinkedList with LinkedHashSet — preserves insertion order, O(1) contains/remove
final Set<String> consumersToFill = new LinkedHashSet<>();

All callers:

unassignedTasks.contains(task) — replace with unassignedTasksSet.contains(task) → O(1)
unassignedTasks.remove(task) — remove from both ordered queue and set
unassignedTasks.poll() — poll from ordered queue, remove from set
consumersToFill.contains(consumer) — replace with LinkedHashSet.contains() → O(1)
consumersToFill.remove(consumer) — O(1) with LinkedHashSet
consumersToFill.offer/add(consumer) — O(1) with LinkedHashSet

Complexity analysis

Scenario	Before	After
C consumers, T tasks	O(C × T²)	O(C × T)
T=500 tasks, C=10 consumers	2,500,000 ops	5,000 ops
T=1000 tasks, C=20 consumers	20,000,000 ops	20,000 ops
Speedup at T=500	—	~500×

Notes

During a Kafka Streams application rebalance with many tasks (common in large deployments), assignTasksToThreads() is called for each client node. The combination of PriorityQueue.contains() inside nested loops creates quadratic scaling with the number of tasks. This was not visible at small scale but becomes a significant bottleneck in deployments with hundreds of stateful tasks per consumer group.

3.8 KiB Raw Blame History Unescape Escape