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.
391 lines
16 KiB
Java
391 lines
16 KiB
Java
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 ===");
|
||
}
|
||
}
|