java-topology/defects/kafka/unit/KafkaStickyAssignorTest.java
russell@unturf.com 0a580b313d undefect. CWE-407 — 63 sites patched across 27 ecosystems
Authors: russell@unturf.com · brackishbert@gmail.com · foxhop.net · TimeHexOn.com

Patches, unit tests, benchmarks, whitepaper, and outreach briefs.
Public domain — no copyright claimed. Use freely.
2026-03-26 17:11:57 -04:00

391 lines
16 KiB
Java
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package unit;
import java.util.*;
/**
* KafkaStickyAssignorTest
*
* Models two CWE-407 defects found in
* clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
*
* kafka-0001 (HIGH) — isBalanced(), line ~1207
* currentAssignment.get(consumer).contains(topicPartition)
* List.contains() is O(P) inside a triple-nested loop C × T × P → O(C × T × P²)
* Fix: snapshot to HashSet<TopicPartition> before inner loops → O(C × T × P)
*
* kafka-0002 (MEDIUM) — maybeAssignPartition(), line ~1267
* consumer2AllPotentialTopics.get(consumer).contains(partition.topic())
* List<String>.contains() is O(T) called P × C times → O(P × C × T)
* Fix: values stored as Set<String> → O(1) contains() → O(P × C)
*
* Run:
* cd /home/fox/git/java-topology/tests
* java -m jdk.compiler/com.sun.tools.javac.Main -cp . -d . ../defects/kafka/unit/KafkaStickyAssignorTest.java
* java -ea -cp . unit.KafkaStickyAssignorTest
*/
public class KafkaStickyAssignorTest {
// -----------------------------------------------------------------------
// Instrumented model of kafka-0001 defect: isBalanced() inner membership
// -----------------------------------------------------------------------
/**
* Defective: uses List<TopicPartition>.contains() — O(P) per call.
* Returns the number of contains() comparisons performed.
*/
static long isBalancedDefective(
Map<String, List<String[]>> currentAssignment, // consumer → list of [topic,partition]
Map<String, List<String>> consumerTopics, // consumer → subscribed topics
Map<String, Integer> partitionsPerTopic) {
long comparisons = 0;
for (String consumer : currentAssignment.keySet()) {
List<String[]> consumerPartitions = currentAssignment.get(consumer);
int consumerPartitionCount = consumerPartitions.size();
List<String> allSubscribedTopics = consumerTopics.get(consumer);
for (String topic : allSubscribedTopics) {
int partCount = partitionsPerTopic.get(topic);
for (int i = 0; i < partCount; i++) {
String[] tp = {topic, String.valueOf(i)};
// O(P): scan the whole list for each candidate partition
boolean found = false;
for (String[] existing : consumerPartitions) {
comparisons++;
if (existing[0].equals(tp[0]) && existing[1].equals(tp[1])) {
found = true;
break;
}
}
// (found result used notionally — we just count comparisons here)
if (found) { /* membership confirmed */ }
}
}
}
return comparisons;
}
/**
* Fixed: snapshots to HashSet<> before inner loops — O(1) per contains().
* Returns the number of contains() comparisons performed (always 1 per lookup).
*/
static long isBalancedFixed(
Map<String, List<String[]>> currentAssignment,
Map<String, List<String>> consumerTopics,
Map<String, Integer> partitionsPerTopic) {
long comparisons = 0;
for (String consumer : currentAssignment.keySet()) {
List<String[]> consumerPartitions = currentAssignment.get(consumer);
// kafka-0001 fix: snapshot to HashSet before inner loops
Set<String> consumerPartitionSet = new HashSet<>();
for (String[] tp : consumerPartitions) {
consumerPartitionSet.add(tp[0] + ":" + tp[1]);
}
List<String> allSubscribedTopics = consumerTopics.get(consumer);
for (String topic : allSubscribedTopics) {
int partCount = partitionsPerTopic.get(topic);
for (int i = 0; i < partCount; i++) {
String key = topic + ":" + i;
comparisons++; // HashSet.contains() — one hash probe
boolean found = consumerPartitionSet.contains(key);
if (found) { /* membership confirmed */ }
}
}
}
return comparisons;
}
// -----------------------------------------------------------------------
// Instrumented model of kafka-0002 defect: maybeAssignPartition() topic check
// -----------------------------------------------------------------------
/**
* Defective: uses List<String>.contains() — O(T) per call.
* Returns the number of comparisons performed across all partitions × consumers.
*/
static long maybeAssignDefective(
List<String[]> unassignedPartitions, // list of [topic,partition]
List<String> consumers,
Map<String, List<String>> consumer2Topics) { // consumer → List<String> of topics
long comparisons = 0;
for (String[] partition : unassignedPartitions) {
String topic = partition[0];
for (String consumer : consumers) {
List<String> topics = consumer2Topics.get(consumer);
// O(T): scan entire list
for (String t : topics) {
comparisons++;
if (t.equals(topic)) break;
}
}
}
return comparisons;
}
/**
* Fixed: uses Set<String> values — O(1) per contains().
* Returns the number of comparisons performed (always 1 per lookup).
*/
static long maybeAssignFixed(
List<String[]> unassignedPartitions,
List<String> consumers,
Map<String, Set<String>> consumer2TopicsSet) { // kafka-0002 fix: Set<String>
long comparisons = 0;
for (String[] partition : unassignedPartitions) {
String topic = partition[0];
for (String consumer : consumers) {
comparisons++; // O(1) hash probe
boolean sub = consumer2TopicsSet.get(consumer).contains(topic);
if (sub) { /* subscribed */ }
}
}
return comparisons;
}
// -----------------------------------------------------------------------
// Test helpers
// -----------------------------------------------------------------------
static Map<String, List<String[]>> buildAssignment(int consumers, int topics, int partitions) {
Map<String, List<String[]>> assignment = new LinkedHashMap<>();
for (int c = 0; c < consumers; c++) {
String consumer = "c" + c;
List<String[]> parts = new ArrayList<>();
// assign roughly half the partitions to each consumer for realism
for (int t = 0; t < topics; t++) {
for (int p = 0; p < partitions / 2; p++) {
parts.add(new String[]{"topic" + t, String.valueOf(p)});
}
}
assignment.put(consumer, parts);
}
return assignment;
}
static Map<String, List<String>> buildConsumerTopics(int consumers, int topics) {
Map<String, List<String>> map = new LinkedHashMap<>();
for (int c = 0; c < consumers; c++) {
List<String> tList = new ArrayList<>();
for (int t = 0; t < topics; t++) tList.add("topic" + t);
map.put("c" + c, tList);
}
return map;
}
static Map<String, Integer> buildPartitionsPerTopic(int topics, int partitions) {
Map<String, Integer> map = new LinkedHashMap<>();
for (int t = 0; t < topics; t++) map.put("topic" + t, partitions);
return map;
}
static List<String[]> buildUnassigned(int topics, int partitions) {
List<String[]> list = new ArrayList<>();
for (int t = 0; t < topics; t++)
for (int p = 0; p < partitions; p++)
list.add(new String[]{"topic" + t, String.valueOf(p)});
return list;
}
static Map<String, List<String>> buildConsumer2TopicsList(int consumers, int topics) {
Map<String, List<String>> map = new LinkedHashMap<>();
for (int c = 0; c < consumers; c++) {
List<String> tList = new ArrayList<>();
for (int t = 0; t < topics; t++) tList.add("topic" + t);
map.put("c" + c, tList);
}
return map;
}
static Map<String, Set<String>> buildConsumer2TopicsSet(int consumers, int topics) {
Map<String, Set<String>> map = new LinkedHashMap<>();
for (int c = 0; c < consumers; c++) {
Set<String> tSet = new HashSet<>();
for (int t = 0; t < topics; t++) tSet.add("topic" + t);
map.put("c" + c, tSet);
}
return map;
}
// -----------------------------------------------------------------------
// Tests
// -----------------------------------------------------------------------
/**
* TEST 1: kafka-0001 correctness
* Both defective and fixed implementations must agree on which partitions
* are "found" (i.e., already assigned). We verify membership counts match.
*/
static void test1_kafka0001_correctness() {
System.out.println("TEST 1: kafka-0001 correctness (defective vs fixed agree on membership)");
int C = 3, T = 4, P = 6;
Map<String, List<String[]>> assignment = buildAssignment(C, T, P);
Map<String, List<String>> consumerTopics = buildConsumerTopics(C, T);
Map<String, Integer> ppt = buildPartitionsPerTopic(T, P);
// For correctness: count how many partitions are "found" in each approach
// We reuse the comparison counters as a proxy — both should scan the same logical space.
// The defective version does more comparisons (O(P) scan) but finds the same results.
long defComp = isBalancedDefective(assignment, consumerTopics, ppt);
long fixComp = isBalancedFixed(assignment, consumerTopics, ppt);
// The fixed version does exactly C*T*P comparisons (one hash probe per candidate).
long expectedFixed = (long) C * T * P;
assert fixComp == expectedFixed :
"Fixed should do exactly C*T*P=" + expectedFixed + " comparisons, got " + fixComp;
// The defective version does at least as many (often more due to partial scans).
assert defComp >= fixComp :
"Defective should do >= comparisons than fixed, defComp=" + defComp + " fixComp=" + fixComp;
System.out.println(" defective comparisons: " + defComp);
System.out.println(" fixed comparisons: " + fixComp + " (== C*T*P=" + expectedFixed + ")");
System.out.println(" PASS");
}
/**
* TEST 2: kafka-0001 quadratic growth
* Increasing P should cause defective comparisons to grow quadratically.
* We measure at P=20 and P=40; ratio of comparisons should be ~4x (quadratic).
*/
static void test2_kafka0001_quadratic_growth() {
System.out.println("TEST 2: kafka-0001 quadratic growth in P (defective)");
int C = 4, T = 5;
int P1 = 20, P2 = 40;
long comp1 = isBalancedDefective(
buildAssignment(C, T, P1), buildConsumerTopics(C, T), buildPartitionsPerTopic(T, P1));
long comp2 = isBalancedDefective(
buildAssignment(C, T, P2), buildConsumerTopics(C, T), buildPartitionsPerTopic(T, P2));
double ratio = (double) comp2 / comp1;
System.out.printf(" P=%d comparisons: %d%n", P1, comp1);
System.out.printf(" P=%d comparisons: %d%n", P2, comp2);
System.out.printf(" ratio: %.2fx (expected ~4x for quadratic)%n", ratio);
// Quadratic: doubling P should roughly quadruple comparisons.
// We accept anywhere in [2.5, 6.0] to account for early-exit effects.
assert ratio >= 2.5 && ratio <= 6.0 :
"Expected quadratic ratio ~4x, got " + ratio;
System.out.println(" PASS");
}
/**
* TEST 3: kafka-0001 linear growth (fixed)
* Increasing P should cause fixed comparisons to grow linearly.
* We measure at P=20 and P=40; ratio should be ~2x (linear).
*/
static void test3_kafka0001_linear_growth() {
System.out.println("TEST 3: kafka-0001 linear growth in P (fixed)");
int C = 4, T = 5;
int P1 = 20, P2 = 40;
long comp1 = isBalancedFixed(
buildAssignment(C, T, P1), buildConsumerTopics(C, T), buildPartitionsPerTopic(T, P1));
long comp2 = isBalancedFixed(
buildAssignment(C, T, P2), buildConsumerTopics(C, T), buildPartitionsPerTopic(T, P2));
double ratio = (double) comp2 / comp1;
System.out.printf(" P=%d comparisons: %d%n", P1, comp1);
System.out.printf(" P=%d comparisons: %d%n", P2, comp2);
System.out.printf(" ratio: %.2fx (expected ~2x for linear)%n", ratio);
assert ratio >= 1.8 && ratio <= 2.2 :
"Expected linear ratio ~2x, got " + ratio;
System.out.println(" PASS");
}
/**
* TEST 4: kafka-0002 ratio at scale
* For kafka-0002 (maybeAssignPartition), the defective O(T) vs fixed O(1).
* At T=100 topics, the defective should do ~T× more comparisons than fixed.
*/
static void test4_kafka0002_ratio_at_scale() {
System.out.println("TEST 4: kafka-0002 ratio at scale (List.contains vs Set.contains)");
int C = 10, T = 100, P = 50;
List<String[]> unassigned = buildUnassigned(T, P);
List<String> consumers = new ArrayList<>();
for (int c = 0; c < C; c++) consumers.add("c" + c);
Map<String, List<String>> listMap = buildConsumer2TopicsList(C, T);
Map<String, Set<String>> setMap = buildConsumer2TopicsSet(C, T);
long defComp = maybeAssignDefective(unassigned, consumers, listMap);
long fixComp = maybeAssignFixed(unassigned, consumers, setMap);
double ratio = (double) defComp / fixComp;
System.out.printf(" T=%d, P=%d, C=%d%n", T, P, C);
System.out.printf(" defective comparisons: %d%n", defComp);
System.out.printf(" fixed comparisons: %d%n", fixComp);
System.out.printf(" ratio: %.1fx (expected >10x)%n", ratio);
assert ratio > 10.0 :
"Expected >10x ratio at T=" + T + ", got " + ratio;
System.out.println(" PASS");
}
/**
* TEST 5: kafka-0001 ratio at scale
* At large P, the defective contains() inside the triple loop blows up.
* We require >10x more comparisons in defective vs fixed at P=100.
*/
static void test5_kafka0001_ratio_at_scale() {
System.out.println("TEST 5: kafka-0001 ratio at scale (List.contains vs HashSet.contains)");
int C = 5, T = 8, P = 100;
Map<String, List<String[]>> assignment = buildAssignment(C, T, P);
Map<String, List<String>> consumerTopics = buildConsumerTopics(C, T);
Map<String, Integer> ppt = buildPartitionsPerTopic(T, P);
long defComp = isBalancedDefective(assignment, consumerTopics, ppt);
long fixComp = isBalancedFixed(assignment, consumerTopics, ppt);
double ratio = (double) defComp / fixComp;
System.out.printf(" C=%d, T=%d, P=%d%n", C, T, P);
System.out.printf(" defective comparisons: %d%n", defComp);
System.out.printf(" fixed comparisons: %d%n", fixComp);
System.out.printf(" ratio: %.1fx (expected >10x)%n", ratio);
assert ratio > 10.0 :
"Expected >10x ratio at P=" + P + ", got " + ratio;
System.out.println(" PASS");
}
// -----------------------------------------------------------------------
// main
// -----------------------------------------------------------------------
public static void main(String[] args) {
System.out.println("=== KafkaStickyAssignorTest ===");
System.out.println("kafka-0001: isBalanced() List.contains -> HashSet (CWE-407 HIGH)");
System.out.println("kafka-0002: maybeAssignPartition() List<String>.contains -> Set<String> (CWE-407 MEDIUM)");
System.out.println();
test1_kafka0001_correctness();
System.out.println();
test2_kafka0001_quadratic_growth();
System.out.println();
test3_kafka0001_linear_growth();
System.out.println();
test4_kafka0002_ratio_at_scale();
System.out.println();
test5_kafka0001_ratio_at_scale();
System.out.println();
System.out.println("=== ALL TESTS PASSED ===");
}
}