Optimizing Producer Performance with Message Batching

Introduction to the Lesson

While our broker is now more resilient, the communication between the client (producer) and the broker is still inefficient. Currently, a producer sends a single message and waits for the broker to process it before sending the next. This "one-at-a-time" approach introduces significant overhead:

• High Network Latency: Each message incurs a full network round-trip time.
• Reduced Throughput: The rate of message production is limited by latency, not by the broker's processing capacity or network bandwidth.
• Increased Resource Usage: Each request involves TCP packet overhead, which adds up when sending many small messages.

Today, we'll solve this problem. Your learning outcome is to implement message batching for producers to improve throughput and reduce network overhead.

This involves two key concepts:

1. Understanding the core logic of batching, which balances throughput and latency using batch size and time-based triggers.
2. Modifying our protocol and broker logic to handle requests containing multiple messages.

This is a standard and vital optimization in all serious messaging systems, from Kafka to NSQ.

1. The Art of Batching: Balancing Throughput and Latency

At its core, batching is about collecting multiple small items to process them as a single, larger unit. For our producer, this means a client gathers several messages before sending them to the broker in one network request.

This introduces a classic engineering trade-off:

• Larger batches lead to higher throughput because the fixed cost of the network request is amortized over more messages.
• Waiting to fill a large batch can increase the end-to-end latency for any individual message sitting in that batch.

To manage this trade-off, producer clients in systems like Kafka use two parameters:

1. Batch Size (batch.size): The maximum number of messages (or bytes) to include in a batch. Once the batch is full, it's sent immediately.
2. Linger Time (linger.ms): A timeout. If a batch isn't full but has been waiting for this duration, it's sent anyway. This ensures messages don't get stuck waiting indefinitely during periods of low traffic.

This dual-trigger mechanism provides high throughput under heavy load (batches fill up quickly) and low latency under light load (batches are sent by the linger timeout).

This next article provides an excellent real-world analogy and explanation for this exact problem.

For a more detailed view, which should appeal to your interest in system design, the following state machine diagram shows the complete lifecycle of a batch within a sophisticated producer client.

2. Evolving Our Protocol for Batched Messages

To enable batching, we first need to update our communication protocol. The ProduceRequest must be able to carry a collection of messages, not just one.

Step 1: Define a Record

In the previous lesson, we used protocol.Message ambiguously. Let's clarify our domain. We'll introduce a Record struct that represents what the user wants to send: a key-value pair. The protocol.Message will remain our top-level on-the-wire envelope.

Create a new file internal/protocol/record.go:

// internal/protocol/record.go
package protocol

// Record represents the data unit sent by a producer. It consists of a key and a value.
type Record struct {
    Key   []byte
    Value []byte
}

Step 2: Update ProduceRequest

Now, let's modify the ProduceRequest to hold a slice of these records.

// internal/protocol/produce.go

// ProduceRequest is the payload for a produce message.
// It contains the topic name and a batch of records to be produced.
type ProduceRequest struct {
    TopicName string
    Records   []*Record
}

This is a significant change. Our old ProduceRequest held Key and Value directly. Now it holds a slice of Record pointers.

Step 3: Update Encoding and Decoding

We must update our serialization logic to handle this new structure. Let's use encoding/gob for simplicity.

In a production system, you'd design a more efficient binary format (e.g., length-prefixing fields), but gob is sufficient for our learning purposes and demonstrates the principle.

// internal/protocol/produce.go
package protocol

import (
    "bytes"
    "encoding/gob"
)

// ... Record struct definition here or in its own file ...

// ProduceRequest definition...

// EncodeProduceRequest serializes a ProduceRequest into a byte slice.
func EncodeProduceRequest(req *ProduceRequest) ([]byte, error) {
    var buf bytes.Buffer
    encoder := gob.NewEncoder(&buf)
    if err := encoder.Encode(req); err != nil {
        return nil, err
    }
    return buf.Bytes(), nil
}

// DecodeProduceRequest deserializes bytes into a ProduceRequest.
func DecodeProduceRequest(data []byte) (*ProduceRequest, error) {
    var req ProduceRequest
    buf := bytes.NewReader(data)
    decoder := gob.NewDecoder(buf)
    if err := decoder.Decode(&req); err != nil {
        return nil, err
    }
    return &req, nil
}

With these changes, our protocol now natively supports sending a batch of records in a single request.

3. Updating the Broker to Handle Batches

Now that our protocol supports batches, we need to update the broker's request handler. The change is surprisingly simple, thanks to the asynchronous worker pool we built in the last lesson.

The handleProduceRequest function's role is to receive the batched request and enqueue each individual record as a separate work item for our worker pool.

Let's refactor handleProduceRequest and processProduceWork.

// internal/broker/broker.go

// handleProduceRequest now decodes a request containing a batch of records.
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)
    }

    // A WaitGroup to wait for all records in the batch to be processed.
    var wg sync.WaitGroup
    wg.Add(len(req.Records))

    // An error channel to collect errors from workers.
    // Buffered to prevent goroutine leaks if we stop listening early.
    errChan := make(chan error, len(req.Records))

    for _, record := range req.Records {
        work := ProduceWork{
            TopicName:  req.TopicName,
            Record:     record,
            ResultChan: make(chan error, 1), // Buffered channel for the result
        }

        // Send work to the channel. This will block if the queue is full (backpressure).
        b.produceChan <- work

        // Spawn a goroutine to wait for the result of this specific record.
        go func() {
            defer wg.Done()
            if err := <-work.ResultChan; err != nil {
                errChan <- err
            }
        }()
    }

    // Wait for all records to be processed.
    wg.Wait()
    close(errChan)

    // Check if any errors occurred during processing.
    if len(errChan) > 0 {
        // In a real system, you'd return detailed errors. For now, we return the first.
        firstErr := <-errChan
        return fmt.Errorf("failed to process one or more records in batch: %w", firstErr)
    }

    // Send success response back to client (logic to be added).
    log.Printf("Batch of %d records for topic '%s' queued successfully", len(req.Records), req.TopicName)
    return nil
}

// processProduceWork now operates on a single record.
func (b *Broker) processProduceWork(work ProduceWork) error {
    b.mu.RLock()
    topic, ok := b.topics[work.TopicName]
    b.mu.RUnlock()
    if !ok {
        return fmt.Errorf("topic '%s' not found", work.TopicName)
    }
    
    partitionCount := len(topic.Partitions)
    if partitionCount == 0 {
        return fmt.Errorf("no partitions for topic '%s'", work.TopicName)
    }
    
    partitionID := b.partitioner.Balance(work.Record, partitionCount) // Pass the Record to the partitioner
    partition := topic.Partitions[partitionID]
    
    // Append the single record's value to the commit log
    _, err := partition.Log.Append(work.Record.Value) // Assuming Log.Append takes []byte
    if err != nil {
        return fmt.Errorf("failed to append to log for partition %d: %w", partition.ID, err)
    }
    
    return nil
}

We also need to update the ProduceWork struct and the Partitioner interface.

// internal/broker/work.go
type ProduceWork struct {
    TopicName  string
    Record     *protocol.Record
    ResultChan chan error
}

// internal/broker/partitioner.go
type Partitioner interface {
    Balance(record *protocol.Record, partitionCount int) int
}

// Update DefaultPartitioner to use record.Key
func (p *DefaultPartitioner) Balance(record *protocol.Record, partitionCount int) int {
    if record.Key == nil {
        // ... round-robin logic
    } else {
        // ... key-hashing logic using record.Key
    }
    // ...
}

This design is powerful. The network handler quickly unpacks the batch and offloads the real work, freeing it up to handle the next request. The backpressure mechanism we already built now applies to individual records, so if a large batch arrives that would overflow the queue, the handleProduceRequest goroutine will simply block partway through the loop until space becomes available.

4. A Note on Further Optimization: I/O Batching

We've implemented batching at the network layer, which is a massive improvement. However, our workers still write one record at a time to the disk. For ultimate performance, we could also batch these disk writes.

This article gives a hint about this very optimization.

Implementing this would involve changing our produceWorker to:

1. Read one item from produceChan.
2. Try to greedily read more items from the channel without blocking (e.g., using a select with a default case).
3. Collect all retrieved items into a slice.
4. Perform a single, larger write to the commit log.

We won't implement this today, as it adds complexity to the worker logic, but it's a logical next step in performance tuning and something to keep in mind for our final project optimizations.

Conclusion

Excellent work! You have now implemented one of the most critical performance optimizations for a message producer. By allowing clients to send batches of messages, you've drastically reduced network overhead and paved the way for much higher throughput.

Key Takeaways:

• Batching is a trade-off between throughput and latency, managed with batch.size and linger.ms parameters on the client side.
• Protocols must be designed to support batching, which we accomplished by changing our ProduceRequest to hold a slice of records.
• An asynchronous worker pool is a perfect backend for a batching API. The request handler can quickly iterate through a batch and enqueue individual work items, relying on the queue for backpressure.
• Batching can be applied at multiple layers: we've focused on network batching, but I/O batching is another powerful technique for optimization.

Next Up

We have now built a robust and efficient producer API. The broker can accept messages, partition them, and persist them asynchronously in batches. It's time to turn our attention to the other side of the equation: getting messages out of the broker. In the next lesson, you will implement the core logic to read a sequence of messages from a partition's log given a starting offset, which is the first step in building our consumer.