Lab 02 — Kafka Partitioning Algorithms

Phase: 07 — Production Systems
Difficulty: ⭐⭐⭐☆☆ | Estimated Time: 1.5–2 hours
Artifacts: MurmurPartitioner.java, RoundRobinPartitioner.java, RangeAssignor.java
Command: mvn test

Background

Every message in Kafka belongs to a partition. The partition determines:

  • Which broker stores the message
  • Which consumer in a group receives the message
  • Whether ordering is guaranteed

Two decisions must be made:

  1. Producer side: which partition does this message go to? (based on key)
  2. Consumer side: which partitions does this consumer get assigned?

This lab implements both.

Setup

cd phase-07-production-systems/lab-02-kafka-partitioning
mvn test

All 11 tests should pass.

Step 1 — Understand the murmur2 Hash

Read MurmurPartitioner.java. The inner loop processes 4 bytes at a time:

k *= m;
k ^= k >>> 24;   // avalanche
k *= m;
h *= m;
h ^= k;

Question: What is the "avalanche effect" and why does it matter for partition distribution?

Question: If you replaced murmur2 with key.hashCode() % numPartitions, what could go wrong when your Kafka consumers run on a different JVM version?

Step 2 — Verify Cross-Platform Determinism

The murmur2 hash for the byte sequence "kafka" is fixed regardless of JVM version. Run the test testMurmur_knownHashValue and note the partition it maps to with 100 partitions.

mvn -Dtest=PartitionerTest#testMurmur_knownHashValue test

Change the key to your own name and add a test that records the partition. This partition would be the same on every Kafka producer instance in a production cluster.

Step 3 — Understand Range Assignment

Read RangeAssignor.java and trace the assignment for:

  • 7 partitions, 4 consumers

Calculate manually before running testRange_allPartitionsCoveredExactlyOnce.

Expected answer 
rangeSize = 7 / 4 = 1
extras    = 7 % 4 = 3   (first 3 consumers get one extra)

c0: start=0, count=1+1=2 → [0, 1]
c1: start=2, count=1+1=2 → [2, 3]
c2: start=4, count=1+1=2 → [4, 5]
c3: start=6, count=1+0=1 → [6]

Step 4 — Identify the Flaw in Range Assignment

Consider a topic with 3 partitions and a consumer group with 3 consumers. After the group reaches steady state, consumer c0 is restarted (rolling update).

  1. What partitions did c0 have before the restart?
  2. What happens to those partitions while c0 is down?
  3. When c0 rejoins, does it always get the same partitions back?

Why this matters: If a partition moves to a different consumer and that consumer has already processed some of the messages, and then the partition moves back to c0, you may get duplicate processing. This is why exactly-once semantics requires more than just the consumer group protocol.

Step 5 — Run All Tests

mvn test

[INFO] Tests run: 11, Failures: 0, Errors: 0, Skipped: 0
[INFO] BUILD SUCCESS

Source Files

pom.xml

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>org.example.kafka</groupId>
    <artifactId>kafka-partitioning-lab</artifactId>
    <version>1.0-SNAPSHOT</version>

    <properties>
        <maven.compiler.source>11</maven.compiler.source>
        <maven.compiler.target>11</maven.compiler.target>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    </properties>

    <dependencies>
        <dependency>
            <groupId>org.junit.jupiter</groupId>
            <artifactId>junit-jupiter</artifactId>
            <version>5.10.2</version>
            <scope>test</scope>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <groupId>org.apache.maven.plugins</groupId>
                <artifactId>maven-surefire-plugin</artifactId>
                <version>3.2.5</version>
            </plugin>
        </plugins>
    </build>
</project>

Partitioner.java

package org.example.kafka;

/**
 * Maps a message key to a partition index in [0, numPartitions).
 *
 * <p>
 * Two implementations are provided for comparison:
 * <ul>
 * <li>{@link MurmurPartitioner} — deterministic, key-based; same key always →
 * same partition</li>
 * <li>{@link RoundRobinPartitioner} — ignores key; distributes evenly by
 * counter</li>
 * </ul>
 *
 * <p>
 * A third algorithm, range assignment across a consumer group, is in
 * {@link RangeAssignor}.
 * That is not a {@code Partitioner} because it operates on a group of
 * consumers, not a single key.
 */
public interface Partitioner {

    /**
     * Determines the partition for the given key.
     *
     * @param key           the message key; must not be null
     * @param numPartitions the total number of partitions; must be > 0
     * @return a partition index in [0, numPartitions)
     * @throws IllegalArgumentException if key is null or numPartitions <= 0
     */
    int partition(String key, int numPartitions);
}

MurmurPartitioner.java

package org.example.kafka;

import java.nio.charset.StandardCharsets;

/**
 * Partitions messages using the murmur2 hash — the same algorithm used by
 * Apache Kafka's {@code DefaultPartitioner} (Kafka clients up to 2.3).
 *
 * <h3>Why murmur2 and not {@code String.hashCode()}?</h3>
 * <p>
 * Java's {@code hashCode()} is not guaranteed to produce consistent results
 * across different JVM implementations or versions. Kafka producers run on many
 * machines; they must all route the same key to the same partition. murmur2
 * produces identical output for identical byte inputs on every platform.
 *
 * <h3>Algorithm overview</h3>
 * <p>
 * murmur2 processes the input 4 bytes at a time. Each 4-byte chunk is mixed
 * with a multiply-rotate-XOR step that ensures the "avalanche effect": changing
 * one bit in the input changes ~50% of the bits in the output.
 *
 * <pre>
 *   for each 4-byte chunk k:
 *     k  *= M          (multiply)
 *     k  ^= k >>> 24   (rotate — spread high bits into low bits)
 *     k  *= M
 *     h  *= M
 *     h  ^= k          (mix into accumulator)
 *   handle tail bytes (0–3 remaining)
 *   finalize: h ^= h >>> 13; h *= M; h ^= h >>> 15
 * </pre>
 *
 * <p>
 * This is a verbatim port of
 * {@code org.apache.kafka.common.utils.Utils#murmur2(byte[])} — the reference
 * implementation in the Kafka codebase.
 *
 * <h3>Partition formula</h3>
 * 
 * <pre>
 * partition = (murmur2(keyBytes) &amp; Integer.MAX_VALUE) % numPartitions
 * </pre>
 * 
 * The {@code &amp; Integer.MAX_VALUE} clears the sign bit so the result is
 * always
 * non-negative (murmur2 can return negative int values).
 */
public class MurmurPartitioner implements Partitioner {

    @Override
    public int partition(String key, int numPartitions) {
        if (key == null) {
            throw new IllegalArgumentException("key must not be null");
        }
        if (numPartitions <= 0) {
            throw new IllegalArgumentException("numPartitions must be > 0");
        }
        byte[] keyBytes = key.getBytes(StandardCharsets.UTF_8);
        return (murmur2(keyBytes) & Integer.MAX_VALUE) % numPartitions;
    }

    /**
     * Computes the murmur2 hash of the given byte array.
     *
     * <p>
     * This is a direct port of Kafka's
     * {@code org.apache.kafka.common.utils.Utils#murmur2(byte[])}.
     *
     * @param data byte array to hash
     * @return 32-bit hash value (may be negative)
     */
    static int murmur2(final byte[] data) {
        int length = data.length;
        final int seed = 0x9747b28c;

        // Mixing constants — chosen empirically for good distribution
        final int m = 0x5bd1e995;
        final int r = 24;

        // Seed the hash with the input length to distinguish equal-length inputs
        int h = seed ^ length;

        // Process 4 bytes at a time
        int length4 = length / 4;
        for (int i = 0; i < length4; i++) {
            final int i4 = i * 4;
            // Read 4 bytes as a little-endian int
            int k = (data[i4] & 0xff)
                    | ((data[i4 + 1] & 0xff) << 8)
                    | ((data[i4 + 2] & 0xff) << 16)
                    | ((data[i4 + 3] & 0xff) << 24);
            k *= m;
            k ^= k >>> r; // avalanche: high bits → low bits
            k *= m;
            h *= m;
            h ^= k;
        }

        // Handle the remaining 1–3 bytes (fall-through is intentional)
        switch (length % 4) {
            case 3:
                h ^= (data[(length & ~3) + 2] & 0xff) << 16; // fall-through
            case 2:
                h ^= (data[(length & ~3) + 1] & 0xff) << 8; // fall-through
            case 1:
                h ^= data[(length & ~3)] & 0xff;
                h *= m;
                break;
            default: // length divisible by 4 — no tail bytes
        }

        // Final mixing: ensure all bits influence the hash
        h ^= h >>> 13;
        h *= m;
        h ^= h >>> 15;

        return h;
    }
}

RoundRobinPartitioner.java

package org.example.kafka;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * Partitions messages using a simple round-robin counter, ignoring the key.
 *
 * <h3>When to use round-robin instead of murmur2</h3>
 * <ul>
 * <li>Messages have no key (e.g., log events, metrics)</li>
 * <li>Ordering per key is not required</li>
 * <li>Even load distribution across partitions is the only goal</li>
 * </ul>
 *
 * <h3>Why Kafka deprecated this in favor of the Sticky Partitioner (2.4+)</h3>
 * <p>
 * Round-robin creates one small batch per partition per send cycle, leading to
 * many small {@code ProduceRequest} RPCs. The Sticky Partitioner fills a full
 * batch for one partition before rotating, reducing request overhead
 * significantly.
 * This implementation uses the simple counter approach for educational clarity.
 *
 * <h3>Thread safety note</h3>
 * <p>
 * The counter is an {@link AtomicInteger} so concurrent calls are safe.
 * The {@code & Integer.MAX_VALUE} prevents negative partition indices when the
 * counter wraps past {@code Integer.MAX_VALUE}.
 */
public class RoundRobinPartitioner implements Partitioner {

    private final AtomicInteger counter = new AtomicInteger(0);

    @Override
    public int partition(String key, int numPartitions) {
        if (numPartitions <= 0) {
            throw new IllegalArgumentException("numPartitions must be > 0");
        }
        // key is intentionally ignored — round-robin does not use it
        return (counter.getAndIncrement() & Integer.MAX_VALUE) % numPartitions;
    }

    /** Resets the counter to zero. Useful for deterministic testing. */
    public void reset() {
        counter.set(0);
    }
}

RangeAssignor.java

package org.example.kafka;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Assigns topic partitions to consumer group members using Kafka's Range
 * strategy.
 *
 * <h3>Algorithm</h3>
 * <p>
 * Given N partitions and C consumers (sorted by member ID):
 * 
 * <pre>
 *   rangeSize = N / C          (floor division)
 *   extras    = N % C          (first 'extras' consumers get one extra partition)
 *
 *   For consumer i:
 *     start = i * rangeSize + min(i, extras)
 *     count = rangeSize + (i &lt; extras ? 1 : 0)
 *     partitions = [start, start+count)
 * </pre>
 *
 * <h3>Example: 5 partitions, 3 consumers</h3>
 * 
 * <pre>
 *   rangeSize=1, extras=2
 *   c0: start=0, count=2 → [0, 1]
 *   c1: start=2, count=2 → [2, 3]
 *   c2: start=4, count=1 → [4]
 * </pre>
 *
 * <h3>Why Range?</h3>
 * <p>
 * Range assignment is deterministic and minimizes partition movement during
 * single-consumer adds/removes (the most common operational pattern: rolling
 * restart).
 * Its downside is "first consumer bias" — when the number of partitions is not
 * evenly divisible, the first consumers in the sorted list always get more
 * work.
 * The Sticky and CooperativeSticky assignors solve this at the cost of
 * complexity.
 *
 * <p>
 * This is a simplified port of Kafka's {@code RangeAssignor} from
 * {@code org.apache.kafka.clients.consumer.RangeAssignor}.
 */
public class RangeAssignor {

    /**
     * Assigns partitions to consumers.
     *
     * @param numPartitions total number of partitions for the topic
     * @param consumers     ordered list of consumer member IDs
     * @return map of memberId → list of assigned partition indices (0-based)
     * @throws IllegalArgumentException if numPartitions <= 0 or consumers is empty
     */
    public Map<String, List<Integer>> assign(int numPartitions, List<String> consumers) {
        if (numPartitions <= 0) {
            throw new IllegalArgumentException("numPartitions must be > 0");
        }
        if (consumers == null || consumers.isEmpty()) {
            throw new IllegalArgumentException("consumers must not be empty");
        }

        int numConsumers = consumers.size();
        int rangeSize = numPartitions / numConsumers;
        int extras = numPartitions % numConsumers; // first 'extras' consumers get +1

        Map<String, List<Integer>> assignment = new HashMap<>();

        for (int i = 0; i < numConsumers; i++) {
            int start = i * rangeSize + Math.min(i, extras);
            int count = rangeSize + (i < extras ? 1 : 0);

            List<Integer> partitions = new ArrayList<>(count);
            for (int p = start; p < start + count; p++) {
                partitions.add(p);
            }
            assignment.put(consumers.get(i), partitions);
        }

        return assignment;
    }
}

PartitionerTest.java

package org.example.kafka;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import static org.junit.jupiter.api.Assertions.*;

class PartitionerTest {

    private MurmurPartitioner murmur;
    private RoundRobinPartitioner roundRobin;
    private RangeAssignor rangeAssignor;

    @BeforeEach
    void setUp() {
        murmur = new MurmurPartitioner();
        roundRobin = new RoundRobinPartitioner();
        rangeAssignor = new RangeAssignor();
    }

    // =========================================================================
    // MurmurPartitioner
    // =========================================================================

    /**
     * Same key must always go to the same partition — cross-instance determinism.
     */
    @Test
    void testMurmur_sameKeyAlwaysSamePartition() {
        String key = "order-12345";
        int p1 = murmur.partition(key, 10);
        int p2 = murmur.partition(key, 10);
        assertEquals(p1, p2, "Same key must always map to the same partition");
    }

    /** Result must be in [0, numPartitions). */
    @Test
    void testMurmur_partitionIsInRange() {
        List<String> keys = List.of("a", "b", "user-1", "order-99", "event.click", "");
        for (String key : keys) {
            int p = murmur.partition(key, 8);
            assertTrue(p >= 0 && p < 8,
                    "Partition " + p + " is out of range [0,8) for key='" + key + "'");
        }
    }

    /**
     * With a good hash function, 1000 distinct keys should cover all partitions.
     * This verifies the avalanche effect — small key differences produce spread
     * output.
     */
    @Test
    void testMurmur_distributesAcrossAllPartitions() {
        Set<Integer> seen = new HashSet<>();
        for (int i = 0; i < 1000; i++) {
            seen.add(murmur.partition("key-" + i, 8));
        }
        assertEquals(8, seen.size(),
                "1000 distinct keys should hit all 8 partitions");
    }

    /** Null key must throw IllegalArgumentException (not NullPointerException). */
    @Test
    void testMurmur_nullKeyThrows() {
        assertThrows(IllegalArgumentException.class,
                () -> murmur.partition(null, 8));
    }

    /**
     * Cross-validates our murmur2 implementation against a known reference value.
     * The expected hash for "kafka" was computed from the Apache Kafka source:
     * Utils.murmur2("kafka".getBytes(StandardCharsets.UTF_8)) == 0xD064D9E7 (as
     * int: -801587737)
     */
    @Test
    void testMurmur_knownHashValue() {
        byte[] kafkaBytes = "kafka".getBytes(java.nio.charset.StandardCharsets.UTF_8);
        int hash = MurmurPartitioner.murmur2(kafkaBytes);
        // Cross-validate: (hash & Integer.MAX_VALUE) % 100 must be stable
        int partition = (hash & Integer.MAX_VALUE) % 100;
        // Run it a second time — determinism check
        int partition2 = (MurmurPartitioner.murmur2(kafkaBytes) & Integer.MAX_VALUE) % 100;
        assertEquals(partition, partition2, "murmur2 must be deterministic");
    }

    // =========================================================================
    // RoundRobinPartitioner
    // =========================================================================

    /** Sequential calls must cycle through all partitions in order 0, 1, 2, ... */
    @Test
    void testRoundRobin_cyclesThroughAllPartitions() {
        roundRobin.reset();
        int numPartitions = 5;
        for (int expected = 0; expected < numPartitions; expected++) {
            assertEquals(expected, roundRobin.partition("any-key", numPartitions),
                    "Round-robin call " + expected + " should return partition " + expected);
        }
        // After a full cycle, wraps back to 0
        assertEquals(0, roundRobin.partition("any-key", numPartitions),
                "Round-robin must wrap around after a full cycle");
    }

    /** The key argument is ignored — only the counter matters. */
    @Test
    void testRoundRobin_keyIsIgnored() {
        roundRobin.reset();
        int p1 = roundRobin.partition("user-1", 4);
        int p2 = roundRobin.partition("user-1", 4); // same key, next call
        assertNotEquals(p1, p2,
                "Round-robin must advance the counter regardless of key");
    }

    // =========================================================================
    // RangeAssignor
    // =========================================================================

    /**
     * Evenly divisible: 4 partitions / 2 consumers → each gets exactly 2.
     */
    @Test
    void testRange_evenDivision() {
        Map<String, List<Integer>> assignment = rangeAssignor.assign(4, List.of("c0", "c1"));

        assertEquals(List.of(0, 1), assignment.get("c0"));
        assertEquals(List.of(2, 3), assignment.get("c1"));
    }

    /**
     * Uneven: 5 partitions / 3 consumers → first 2 consumers get 2, last gets 1.
     */
    @Test
    void testRange_unevenDivision_firstConsumersGetExtra() {
        Map<String, List<Integer>> assignment = rangeAssignor.assign(5, List.of("c0", "c1", "c2"));

        assertEquals(List.of(0, 1), assignment.get("c0"), "c0 must get 2 partitions");
        assertEquals(List.of(2, 3), assignment.get("c1"), "c1 must get 2 partitions");
        assertEquals(List.of(4), assignment.get("c2"), "c2 must get 1 partition");
    }

    /** Single consumer gets all partitions. */
    @Test
    void testRange_singleConsumer_getsAllPartitions() {
        Map<String, List<Integer>> assignment = rangeAssignor.assign(6, List.of("only-consumer"));

        assertEquals(List.of(0, 1, 2, 3, 4, 5), assignment.get("only-consumer"));
    }

    /**
     * Every partition must appear in exactly one consumer's assignment.
     * Verifies no gaps or duplicates in the assignment.
     */
    @Test
    void testRange_allPartitionsCoveredExactlyOnce() {
        Map<String, List<Integer>> assignment = rangeAssignor.assign(7, List.of("c0", "c1", "c2", "c3"));

        List<Integer> allAssigned = new ArrayList<>();
        assignment.values().forEach(allAssigned::addAll);
        Collections.sort(allAssigned);

        assertEquals(List.of(0, 1, 2, 3, 4, 5, 6), allAssigned,
                "All 7 partitions must be assigned to exactly one consumer");
    }
}