200 lines
7.4 KiB
Java
200 lines
7.4 KiB
Java
package unit;
|
|
|
|
import java.util.ArrayList;
|
|
import java.util.HashMap;
|
|
import java.util.HashSet;
|
|
import java.util.Iterator;
|
|
import java.util.List;
|
|
import java.util.Map;
|
|
import java.util.Set;
|
|
|
|
/**
|
|
* CWE-407 unit test: kafka-0011
|
|
*
|
|
* Models AbstractStickyAssignor.GeneralAssignmentBuilder.assignOwnedPartitions()
|
|
* topic subscription check.
|
|
*
|
|
* DEFECT: For each consumer's owned partition, calls
|
|
* consumerSubscription.topics().contains(partition.topic())
|
|
* where topics() returns List<String>.
|
|
* Total: O(C x P x T) where C=consumers, P=partitions/consumer, T=topics.
|
|
*
|
|
* FIX: Build a Map<String, Set<String>> consumer->topicSet before the loop.
|
|
* O(1) set.contains() per partition check.
|
|
*
|
|
* Asserts: slowOps > fastOps * 10 at C=100, P=50, T=50.
|
|
*/
|
|
public class Kafka0011StickyAssignorTest {
|
|
|
|
/**
|
|
* Simulates the defective assignOwnedPartitions loop.
|
|
* For each consumer, for each of their owned partitions, checks
|
|
* if the topic is still in their subscription (List<String>).
|
|
*/
|
|
static long slow(int numConsumers, int partitionsPerConsumer, int topicsPerConsumer) {
|
|
// Setup: each consumer owns partitionsPerConsumer partitions,
|
|
// subscribed to topicsPerConsumer topics.
|
|
List<List<String>> consumerTopics = new ArrayList<>();
|
|
List<List<String>> ownedPartitionTopics = new ArrayList<>();
|
|
|
|
for (int c = 0; c < numConsumers; c++) {
|
|
List<String> topics = new ArrayList<>();
|
|
for (int t = 0; t < topicsPerConsumer; t++) {
|
|
topics.add("topic-" + t);
|
|
}
|
|
consumerTopics.add(topics);
|
|
|
|
List<String> partitions = new ArrayList<>();
|
|
for (int p = 0; p < partitionsPerConsumer; p++) {
|
|
// Partitions are spread across topics; topic-(p % topicsPerConsumer)
|
|
partitions.add("topic-" + (p % topicsPerConsumer));
|
|
}
|
|
ownedPartitionTopics.add(partitions);
|
|
}
|
|
|
|
long ops = 0;
|
|
for (int c = 0; c < numConsumers; c++) {
|
|
List<String> topicList = consumerTopics.get(c);
|
|
for (String partitionTopic : ownedPartitionTopics.get(c)) {
|
|
// O(T) scan — the defect
|
|
for (String t : topicList) {
|
|
ops++;
|
|
if (t.equals(partitionTopic)) break;
|
|
}
|
|
}
|
|
}
|
|
return ops;
|
|
}
|
|
|
|
/**
|
|
* Simulates the patched version.
|
|
* Pre-converts each consumer's topic list to HashSet; O(1) per check.
|
|
*/
|
|
static long fast(int numConsumers, int partitionsPerConsumer, int topicsPerConsumer) {
|
|
List<Set<String>> consumerTopicSets = new ArrayList<>();
|
|
List<List<String>> ownedPartitionTopics = new ArrayList<>();
|
|
|
|
for (int c = 0; c < numConsumers; c++) {
|
|
Set<String> topicSet = new HashSet<>();
|
|
for (int t = 0; t < topicsPerConsumer; t++) {
|
|
topicSet.add("topic-" + t);
|
|
}
|
|
consumerTopicSets.add(topicSet);
|
|
|
|
List<String> partitions = new ArrayList<>();
|
|
for (int p = 0; p < partitionsPerConsumer; p++) {
|
|
partitions.add("topic-" + (p % topicsPerConsumer));
|
|
}
|
|
ownedPartitionTopics.add(partitions);
|
|
}
|
|
|
|
long ops = 0;
|
|
for (int c = 0; c < numConsumers; c++) {
|
|
Set<String> topicSet = consumerTopicSets.get(c);
|
|
for (String partitionTopic : ownedPartitionTopics.get(c)) {
|
|
ops++; // O(1) hash probe
|
|
topicSet.contains(partitionTopic);
|
|
}
|
|
}
|
|
return ops;
|
|
}
|
|
|
|
public static void main(String[] args) {
|
|
int passed = 0;
|
|
int total = 0;
|
|
|
|
// Test 1: 100 consumers, 50 partitions each, 50 topics each
|
|
{
|
|
total++;
|
|
int C = 100, P = 50, T = 50;
|
|
long sOps = slow(C, P, T);
|
|
long fOps = fast(C, P, T);
|
|
// slow: C * P * (avg scan depth) ≈ 100 * 50 * 25 = 125000
|
|
// fast: C * P = 5000
|
|
boolean ok = sOps > fOps * 10L;
|
|
System.out.printf("Test 1 [C=%d P=%d T=%d slow=%d fast=%d ratio=%.1fx]: %s%n",
|
|
C, P, T, sOps, fOps, (double) sOps / fOps, ok ? "PASS" : "FAIL");
|
|
if (ok) passed++;
|
|
}
|
|
|
|
// Test 2: larger scale
|
|
{
|
|
total++;
|
|
int C = 200, P = 100, T = 100;
|
|
long sOps = slow(C, P, T);
|
|
long fOps = fast(C, P, T);
|
|
boolean ok = sOps > fOps * 20L;
|
|
System.out.printf("Test 2 [C=%d P=%d T=%d slow=%d fast=%d ratio=%.1fx]: %s%n",
|
|
C, P, T, sOps, fOps, (double) sOps / fOps, ok ? "PASS" : "FAIL");
|
|
if (ok) passed++;
|
|
}
|
|
|
|
// Test 3: worst case — topics in partition list always last in topic list
|
|
{
|
|
total++;
|
|
int C = 50, T = 100;
|
|
// Each partition's topic is last in its consumer's topic list
|
|
List<List<String>> consumerTopics = new ArrayList<>();
|
|
List<List<String>> owned = new ArrayList<>();
|
|
for (int c = 0; c < C; c++) {
|
|
List<String> tl = new ArrayList<>();
|
|
for (int t = 0; t < T; t++) tl.add("other-" + t);
|
|
tl.add("last-topic");
|
|
consumerTopics.add(tl);
|
|
List<String> pl = new ArrayList<>();
|
|
for (int p = 0; p < 50; p++) pl.add("last-topic");
|
|
owned.add(pl);
|
|
}
|
|
|
|
long sOps = 0;
|
|
for (int c = 0; c < C; c++) {
|
|
List<String> topicList = consumerTopics.get(c);
|
|
for (String partTopic : owned.get(c)) {
|
|
for (String t : topicList) { sOps++; if (t.equals(partTopic)) break; }
|
|
}
|
|
}
|
|
// fast: pre-hashed
|
|
long fOps = 0;
|
|
List<Set<String>> sets = new ArrayList<>();
|
|
for (int c = 0; c < C; c++) {
|
|
Set<String> ts = new HashSet<>(consumerTopics.get(c));
|
|
sets.add(ts);
|
|
}
|
|
for (int c = 0; c < C; c++) {
|
|
for (String pt : owned.get(c)) { fOps++; sets.get(c).contains(pt); }
|
|
}
|
|
|
|
boolean ok = sOps > fOps * 50L;
|
|
System.out.printf("Test 3 worst-case [C=%d T=%d slow=%d fast=%d ratio=%.1fx]: %s%n",
|
|
C, T, sOps, fOps, (double) sOps / fOps, ok ? "PASS" : "FAIL");
|
|
if (ok) passed++;
|
|
}
|
|
|
|
// Test 4: correctness — both give same membership results
|
|
{
|
|
total++;
|
|
int C = 20, T = 30;
|
|
boolean ok = true;
|
|
for (int c = 0; c < C && ok; c++) {
|
|
List<String> tl = new ArrayList<>();
|
|
Set<String> ts = new HashSet<>();
|
|
for (int t = 0; t < T; t++) {
|
|
tl.add("t" + t);
|
|
ts.add("t" + t);
|
|
}
|
|
for (int p = 0; p < 40; p++) {
|
|
String pt = "t" + (p % (T + 5));
|
|
if (tl.contains(pt) != ts.contains(pt)) {
|
|
ok = false;
|
|
}
|
|
}
|
|
}
|
|
System.out.printf("Test 4 [correctness C=%d T=%d ok=%b]: %s%n",
|
|
C, T, ok, ok ? "PASS" : "FAIL");
|
|
if (ok) passed++;
|
|
}
|
|
|
|
System.out.printf("%d/%d PASS%n", passed, total);
|
|
if (passed != total) System.exit(1);
|
|
}
|
|
}
|