Topic Manager Partition Assignment

Introduction to the Lesson

With topic creation and deletion now handled, our broker can manage its core inventory. The next critical question is: how does data get distributed within a topic? A topic can have many partitions, and deciding which partition a message goes to is fundamental to achieving parallelism, scalability, and message ordering. This decision-making process is called partition assignment (or partitioning).

Your Goal for Today: Your learning outcome is to implement partition assignment logic within the topic manager.

Specifically, you will equip our Broker struct (the topic manager) with the logic to select partitions based on two common strategies used in systems like Kafka:

1. Round-Robin: For even load distribution.
2. Key-Based: For guaranteeing that messages with the same key always land in the same partition, preserving order.

This logic will be a cornerstone for the producer and consumer APIs we'll build in the upcoming lessons.

1. Understanding Partition Assignment Strategies

Before we write any code, it's essential to understand why different assignment strategies exist and what trade-offs they entail. Systems like Kafka don't just have one method; they offer several that are optimized for different use cases.

To get a solid overview of these strategies, please review the following resources. They describe the standard assignors provided by Kafka, which serve as the blueprint for our own implementation.

From these readings, the key takeaways are:

• Round-Robin: Distributes partitions one-by-one across all available consumers (or in our case, for all incoming produce requests). This is excellent for ensuring an even workload when message order isn't a concern.
• Key-Based (Partitioning): Ensures all messages with the same "key" (e.g., a user_id, order_id) are sent to the same partition. This is critical for use cases that require processing events for a specific entity in the order they were produced.
• Other Strategies (Range, Sticky): These are variations that optimize for specific goals, like co-locating partitions from different topics (Range) or minimizing partition movement during a rebalance (Sticky).

For our broker, we will implement the two most fundamental strategies: Round-Robin and Key-Based.

2. Implementing Round-Robin Assignment

To implement a round-robin strategy, we need to keep track of the last partition that was assigned for a given topic. A simple counter will do. Each time a partition is requested, we increment the counter and use the modulo operator to pick the next partition from the list.

Since multiple producers might be requesting partitions concurrently, this counter must be safe for concurrent access. The sync/atomic package is perfect for this.

First, let's update our Topic struct to include this counter. We'll likely have this struct defined in internal/broker/broker.go or a similar file.

// internal/broker/broker.go (or wherever Topic is defined)

import (
	"sync"
	"sync/atomic" // Import the atomic package
	// ... other imports
)

// ... CommitLog interface ...

// ... Partition struct ...

type Topic struct {
	Name string
	// A slice of partitions belonging to this topic.
	Partitions []*Partition

	// A counter for round-robin assignment. Must be accessed atomically.
	roundRobinCounter atomic.Uint64
}

// Update the NewTopic function to initialize the partitions slice
func NewTopic(name string, partitions []*Partition) *Topic {
	return &Topic{
		Name:       name,
		Partitions: partitions,
	}
}

// ... Broker struct ...

Note: In your CreateTopic method, where you instantiate a new Topic, ensure you are calling this updated NewTopic function. The roundRobinCounter will be zero-valued by default, which is a perfect starting point.

Now, let's add the AssignPartitionRoundRobin method to the Broker. This method will encapsulate the round-robin logic. This pattern is similar to how the distributed scheduler in the video you'll see below uses a roundRobinIndex to distribute tasks.

// internal/broker/broker.go

import (
	"fmt"
	// ... other imports
)

// ...

func (b *Broker) AssignPartitionRoundRobin(topicName string) (*Partition, error) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	topic, exists := b.topics[topicName]
	if !exists {
		return nil, fmt.Errorf("topic '%s' not found", topicName)
	}

	if len(topic.Partitions) == 0 {
		return nil, fmt.Errorf("topic '%s' has no partitions", topicName)
	}

	// Atomically increment the counter to get the next value.
	// Add() returns the new value. We subtract 1 to get a 0-based index.
	nextIdx := topic.roundRobinCounter.Add(1) - 1

	// Modulo the number of partitions to wrap around the slice.
	partitionIdx := nextIdx % uint64(len(topic.Partitions))

	return topic.Partitions[partitionIdx], nil
}

This implementation is clean and, critically, thread-safe. By using atomic.Uint64, we avoid the need for another mutex just to protect this counter. The broker's read lock (RLock) is sufficient to safely access the topics map.

3. Implementing Key-Based Assignment

For key-based assignment, the logic is deterministic: the same key must always map to the same partition. The standard way to achieve this is to:

1. Calculate a hash of the key.
2. Take the result of the hash modulo the number of partitions.
3. The result is the index of the partition to use.

The image below illustrates this process perfectly.

We'll use a non-cryptographic hash function from Go's standard library, hash/fnv, which is fast and provides good distribution.

Let's add the AssignPartitionByKey method to our Broker.

// internal/broker/broker.go

import (
	"fmt"
	"hash/fnv" // Import the FNV hash package
	// ... other imports
)

// ...

func (b *Broker) AssignPartitionByKey(topicName string, key []byte) (*Partition, error) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	topic, exists := b.topics[topicName]
	if !exists {
		return nil, fmt.Errorf("topic '%s' not found", topicName)
	}

	numPartitions := len(topic.Partitions)
	if numPartitions == 0 {
		return nil, fmt.Errorf("topic '%s' has no partitions", topicName)
	}

	// Create a new FNV-1a 32-bit hash.
	hasher := fnv.New32a()

	// Write the key to the hasher. FNV's Write method never returns an error.
	_, _ = hasher.Write(key)

	// Calculate the hash and get the partition index.
	hash := hasher.Sum32()
	partitionIdx := int(hash % uint32(numPartitions))

	return topic.Partitions[partitionIdx], nil
}

With these two methods, our Broker is now equipped to handle the two most important partition assignment strategies. This functionality is completely decoupled from the network layer or the storage layer, adhering to the clean separation of concerns we've been building.

If you're interested in seeing a similar round-robin distribution pattern applied in another context, this video segment on building a distributed task scheduler demonstrates the same core idea of using an index to select the next available "worker" from a pool.

Conclusion

In this lesson, you've added a crucial piece of intelligence to the broker's topic manager. By implementing partition assignment logic, you've defined how data will be distributed across the partitions of a topic, a decision that directly impacts the system's scalability and its ability to provide ordering guarantees.

Key Takeaways:

• Partition assignment strategies like Round-Robin and Key-Based serve different purposes: the former for load balancing, the latter for message ordering per key.
• Implementing a thread-safe round-robin counter can be done efficiently using the sync/atomic package.
• Key-based assignment is a deterministic process implemented by hashing the key and using the modulo operator against the number of partitions.
• Placing this logic within the Broker struct centralizes the functionality, making it available to any part of the system (like producer handlers or consumer group coordinators) that needs to select a partition.

Next Up

You've now reached the end of Module 6. You have a broker that can manage the complete lifecycle of topics (create, delete) and intelligently assign partitions.

In the next module, "Implementing the Producer API and Write Path," you will finally start accepting data from the outside world. We will build a TCP server that listens for incoming producer connections. When a message arrives, it will use the very AssignPartition... methods you just built to select a partition and then write the message to that partition's commit log. The theoretical pieces are falling into place, and it's time to build the data ingress path.