Partition Selection Strategies

Introduction to the Lesson

Today, we will refine this design by implementing a more robust and extensible approach to partition selection. While our current implementation works, real-world systems like Kafka use a more modular pattern: a pluggable Partitioner (or Balancer). This component is solely responsible for deciding which partition a message should go to, encapsulating all the strategic logic.

Your Goal for Today: Your learning outcome is to implement partition selection strategies, including round-robin and key-based hashing, by refactoring our existing logic into a formal Partitioner interface.

This lesson will cover:

1. The design benefits of using a Partitioner interface.
2. Implementing the round-robin strategy for even load distribution when a message has no key.
3. Implementing key-based hashing to guarantee message ordering for a given key.
4. Combining these strategies into a single DefaultPartitioner that mimics Kafka's default behavior.
5. Refactoring our broker to use this new, swappable component.

This change will make our broker's design cleaner and more aligned with production-grade systems, a key step toward your goal of mastering system design.

1. The Power of Partitioning

Before we dive into the implementation, it's crucial to understand why these strategies are so fundamental to systems like Kafka. Partitioning is the mechanism that enables parallelism and scalability. By splitting a topic's log across multiple partitions (and potentially multiple machines), the broker can handle a much higher volume of reads and writes.

The choice of partitioning strategy directly impacts the trade-offs between load distribution and message ordering. This short video from Confluent provides an excellent overview of these concepts.

As the video explains:

• No Key (Round-Robin): Messages are distributed evenly across all available partitions. This is great for maximizing throughput when the order of messages doesn't matter.
• With Key (Hashing): All messages with the same key are guaranteed to land in the same partition. This preserves the order of messages for that specific key (e.g., all events for a single customer_id).

Let's now implement this dual-strategy logic within a single, well-designed component.

2. Designing a Partitioner Interface

To make our partitioning logic modular and swappable, we'll define a Partitioner interface. This is a common pattern in Go and aligns with the "accept interfaces, return structs" principle you'll encounter later. The popular kafka-go library uses an almost identical pattern with its Balancer interface.

Let's define our interface. This should live in a new file, for instance internal/broker/partitioner.go.

// internal/broker/partitioner.go
package broker

import (
    "hash/crc32"
    "sync/atomic"

    "our-kafka/internal/protocol"
)

// Partitioner determines the partition to which a message is sent.
type Partitioner interface {
    // Balance takes a message and the number of partitions and returns
    // the target partition ID.
    Balance(msg *protocol.Message, partitionCount int) int
}

This simple interface gives us a powerful abstraction. We can now create different structs that implement this Balance method with various strategies (round-robin, hashing, least-connections, etc.) and plug them into our broker without changing the core request handling logic.

3. Implementing the Default Partitioner

We will create a DefaultPartitioner that implements the logic described in the Confluent video: round-robin for null keys, and hashing for non-null keys.

The Round-Robin Strategy

The core of round-robin is a simple counter that increments for each message and uses the modulo operator to cycle through the partitions.

partitionID = counter % numPartitions

Since our broker handles connections concurrently, multiple goroutines will be calling Balance at the same time. Therefore, we must use the sync/atomic package to safely increment the counter.

The following resource on building a load balancer demonstrates this exact pattern for distributing web requests. The principle is identical to distributing messages across partitions.

The Key-Based Hashing Strategy

For messages with a key, we need a deterministic way to map the key to a partition. We do this by hashing the key and then using the modulo operator.

partitionID = hash(key) % numPartitions

The choice of hash function matters. We don't need a cryptographically secure hash like SHA-256; we need a fast, non-cryptographic hash with good distribution. Go's standard library hash/crc32 is an excellent choice and is commonly used for this purpose.

Putting It Together

Now, let's combine these two strategies into our DefaultPartitioner.

// internal/broker/partitioner.go (continued)

// DefaultPartitioner implements a strategy that uses round-robin for messages
// with no key and key-based hashing for messages with a key.
type DefaultPartitioner struct {
    // counter is used for the round-robin strategy.
    counter uint64
}

// NewDefaultPartitioner creates a new instance of DefaultPartitioner.
func NewDefaultPartitioner() *DefaultPartitioner {
    return &DefaultPartitioner{}
}

func (p *DefaultPartitioner) Balance(msg *protocol.Message, partitionCount int) int {
    // Note: This assumes protocol.Message has Key []byte
    // from the ProduceRequest payload. You may need to pass
    // the deserialized ProduceRequest instead. For simplicity,
    // we'll assume the key is accessible.

    if len(msg.Key) == 0 {
        // No key, use round-robin.
        // Atomically increment the counter and get the new value.
        next := atomic.AddUint64(&p.counter, 1)
        return int(next % uint64(partitionCount))
    }

    // Key is present, use hashing.
    hash := crc32.ChecksumIEEE(msg.Key)
    return int(hash % uint32(partitionCount))
}

This single Balance function now elegantly handles both cases, providing a robust default behavior for our broker.

4. Integrating the Partitioner into the Broker

With our Partitioner interface and DefaultPartitioner implementation ready, the final step is to integrate it into the broker.

First, update the Broker struct to hold a Partitioner instead of the old partition assignment logic.

// internal/broker/broker.go

type Broker struct {
    // ... other fields like topics, mu ...
    partitioner Partitioner // New field
}

func NewBroker() *Broker {
    return &Broker{
        topics:      make(map[string]*Topic),
        partitioner: NewDefaultPartitioner(), // Initialize with our default
    }
}

Next, refactor the handleProduceRequest function to use the new partitioner. This will significantly simplify the code, removing the if/else block and delegating the decision entirely to the partitioner.

// internal/broker/broker.go

func (b *Broker) handleProduceRequest(conn net.Conn, msg *protocol.Message) error {
    req, err := protocol.DecodeProduceRequest(msg.Payload)
    if err != nil {
        return fmt.Errorf("failed to decode produce request: %w", err)
    }

    b.mu.RLock()
    topic, ok := b.topics[req.TopicName]
    b.mu.RUnlock()
    if !ok {
        // In a future lesson, we might add auto topic creation.
        // For now, we require topics to be created beforehand.
        return fmt.Errorf("topic '%s' not found", req.TopicName)
    }

    // --- Start of Refactoring ---

    // The old logic with if/else for AssignPartitionByKey/RoundRobin is removed.
    
    // New, simplified logic:
    partitionCount := len(topic.Partitions)
    if partitionCount == 0 {
        return fmt.Errorf("no partitions for topic '%s'", req.TopicName)
    }
    
    // We need to construct a protocol.Message with the key to pass to the partitioner.
    // Let's assume the request itself can be adapted or wrapped.
    // For this example, let's create a temporary message struct.
    partitionerMsg := &protocol.Message{Key: req.Key}

    partitionID := b.partitioner.Balance(partitionerMsg, partitionCount)
    partition := topic.Partitions[partitionID]
    
    // --- End of Refactoring ---


    // The append logic remains the same.
    _, err = partition.Log.Append(req.Value)
    if err != nil {
        return fmt.Errorf("failed to append to commit log for partition %d: %w", partition.ID, err)
    }
    
    log.Printf("Appended message to topic '%s', partition %d", req.TopicName, partition.ID)

    return nil
}

Note: The protocol.Message passed to Balance needs to contain the key. You may need a small adjustment to ensure the key from the ProduceRequest is available to the Balance method, either by passing the ProduceRequest itself or a message structure that contains the key.

This refactoring makes our system more robust. To change the partitioning strategy for the entire broker, we would now only need to change the single line in NewBroker where the partitioner is initialized.

Conclusion

Excellent work. You have successfully refactored your broker's partitioning logic to use a formal, extensible Partitioner interface. This is a significant step forward in your system's design maturity.

Key Takeaways:

• A Partitioner interface abstracts away the logic of message distribution, making the system more modular and easier to maintain.
• Round-robin is the ideal strategy for distributing load evenly when message order is not a concern, and it is implemented using a simple, atomically-incremented counter.
• Key-based hashing provides a crucial guarantee: all messages with the same key will be processed in order, as they are always sent to the same partition.
• Combining these into a DefaultPartitioner that checks for the presence of a key allows our broker to intelligently switch between these two fundamental strategies.

Next Up

Our producer handler is now well-designed, but it's still synchronous—it waits for the disk write to complete before processing the next request from a client. This can become a bottleneck.

In the next lesson, you will learn how to use buffered channels to create a bounded work queue for incoming produce requests. This will decouple network I/O from disk I/O, introducing asynchronicity and backpressure to dramatically improve the throughput and resilience of your broker.