Repository: segmentio/kafka-go (language: go)

Issue:
v0.4 Writer does not preserve order of messages
**Describe the bug**
When writing multiple messages using kafka.Writer, multiple batches will be created, and those batches will be processed asynchronously, causing messages to arrive Kafka in random order. e.g. If I write message 1, 2, 3 in one WriteMessages call, then we may create three batches for message 1, 2, 3 and they will reach Kafka in the order of 2, 3, 1.

**Kafka Version**
2.6.0

**To Reproduce**
Start Kafka locally using default configs, create a topic "test-topic", then use this code snippet:
```
package bugreport_test

import (
	"context"
	"math/rand"
	"strconv"
	"testing"
	"time"

	"github.com/golang/glog"
	"github.com/segmentio/kafka-go"
)

type balancer int

func (b balancer) Balance(msg kafka.Message, partitions ...int) (partition int) {
	return 0
}

func NewMessage(size int, key, topic string) kafka.Message {
	return kafka.Message{
		Key:     []byte(key),
		Value:   messageValue(size),
		Headers: []kafka.Header{},
		Time:    time.Now(),
	}
}

func messageValue(size int) []byte {
	val := make([]byte, size)
	for i := 0; i < size; i++ {
		val[i] = byte(random(0, 9))
	}
	return val
}

func random(min int, max int) int {
	return rand.Intn(max-min) + min
}

func TestWriter(t *testing.T) {
	wr := kafka.NewWriter(kafka.WriterConfig{
		Brokers:      []string{"localhost:9092"},
		Topic:        "test-topic",
		Dialer:       kafka.DefaultDialer,
		BatchSize:    1000,
		BatchTimeout: 500 * time.Millisecond,
		Balancer:     balancer(0),
	})
	var msgs = []kafka.Message{
		NewMessage(950000, "1", "test-topic"),
		NewMessage(950000, "2", "test-topic"),
		NewMessage(950000, "3", "test-topic"),
		NewMessage(950000, "4", "test-topic"),
		NewMessage(950000, "5", "test-topic"),
		NewMessage(950000, "6", "test-topic"),
		NewMessage(950000, "7", "test-topic"),
		NewMessage(950000, "8", "test-topic"),
		NewMessage(950000, "9", "test-topic"),
		NewMessage(950000, "10", "test-topic"),
	}
	rdr := kafka.NewReader(kafka.ReaderConfig{
		Brokers:     []string{"localhost:9092"},
		Topic:       "test-topic",
		Dialer:      kafka.DefaultDialer,
		GroupID:     strconv.FormatInt(time.Now().Unix(), 10),
		StartOffset: kafka.LastOffset,
	})
	ctx, _ := context.WithTimeout(context.Background(), 100*time.Second)
	go func() {
		for {
			msg, err := rdr.ReadMessage(ctx)
			if err != nil {
				glog.Errorf("err read msg: %s", err.Error())
			} else {
				glog.Infof("receive msg: key: %s, offset: %d", string(msg.Key), msg.Offset)
			}
		}
	}()
	time.Sleep(3 * time.Second)
	err := wr.WriteMessages(ctx, msgs...)
	if err != nil {
		glog.Errorf("err write msg: %s", err.Error())
	} else {
		glog.Info("write complete")
	}
	time.Sleep(10 * time.Second)
}
```
Run the test using ```go test -v './...' --logtostderr --v=1```, you will see something like:

I0622 14:36:34.674736 kafka-go_test.go:69] receive msg: key: 2, offset: 36
I0622 14:36:34.780102 kafka-go_test.go:69] receive msg: key: 7, offset: 37
I0622 14:36:34.885081 kafka-go_test.go:69] receive msg: key: 3, offset: 38
I0622 14:36:34.990090 kafka-go_test.go:69] receive msg: key: 5, offset: 39
I0622 14:36:35.095683 kafka-go_test.go:69] receive msg: key: 8, offset: 40
I0622 14:36:35.200772 kafka-go_test.go:69] receive msg: key: 6, offset: 41
I0622 14:36:35.303842 kafka-go_test.go:69] receive msg: key: 9, offset: 42
I0622 14:36:35.409054 kafka-go_test.go:69] receive msg: key: 4, offset: 43
I0622 14:36:35.530030 kafka-go_test.go:69] receive msg: key: 1, offset: 44
I0622 14:36:35.634247 kafka-go_test.go:69] receive msg: key: 10, offset: 45

**Expected behavior**
Since we call WriteMessages() in the order of 1 to 10, we expect to see
I0622 14:36:34.674736 kafka-go_test.go:69] receive msg: key: 1, offset: 36
I0622 14:36:34.780102 kafka-go_test.go:69] receive msg: key: 2, offset: 37
I0622 14:36:34.885081 kafka-go_test.go:69] receive msg: key: 3, offset: 38
I0622 14:36:34.990090 kafka-go_test.go:69] receive msg: key: 4, offset: 39
I0622 14:36:35.095683 kafka-go_test.go:69] receive msg: key: 5, offset: 40
I0622 14:36:35.200772 kafka-go_test.go:69] receive msg: key: 6, offset: 41
I0622 14:36:35.303842 kafka-go_test.go:69] receive msg: key: 7, offset: 42
I0622 14:36:35.409054 kafka-go_test.go:69] receive msg: key: 8, offset: 43
I0622 14:36:35.530030 kafka-go_test.go:69] receive msg: key: 9, offset: 44
I0622 14:36:35.634247 kafka-go_test.go:69] receive msg: key: 10, offset: 45

**Additional context**
We are using kafka-go v0.4.16



Discussion / hints:
Wouldn't each of those messages go into different partitions?
> Wouldn't each of those messages go into different partitions?

I just updated the writer to use a balancer that always returns partition 0 and got the same issue.
any news?
Hello,
Thank you for your patience.

The kafka.Writer type does not make any guarantees of message ordering, especially across batches. If more control of message ordering is needed, the lower level Produce API of the kafka.Client type can be used.
@rhansen2 I hope you might reconsider. No longer respecting ordering of messages in a single batch is a big breaking change and it seems like an unnecessary contract to break with clients. We have a lot of usage of your library, and we will need to fork or move to a different library, which is pretty disappointing. 

Fix patch:
diff --git a/writer.go b/writer.go
index ff619efed..5d44f75c7 100644
--- a/writer.go
+++ b/writer.go
@@ -189,9 +189,9 @@ type Writer struct {
 	closed uint32
 	group  sync.WaitGroup
 
-	// Manages the current batch being aggregated on the writer.
+	// Manages the current set of partition-topic writers.
 	mutex   sync.Mutex
-	batches map[topicPartition]*writeBatch
+	writers map[topicPartition]*partitionWriter
 
 	// writer stats are all made of atomic values, no need for synchronization.
 	// Use a pointer to ensure 64-bit alignment of the values. The once value is
@@ -511,15 +511,15 @@ func NewWriter(config WriterConfig) *Writer {
 // io.ErrClosedPipe.
 func (w *Writer) Close() error {
 	w.markClosed()
-	// If batches are pending, trigger them so messages get sent.
 	w.mutex.Lock()
 
-	for _, batch := range w.batches {
-		batch.trigger()
+	// close all writers to trigger any pending batches
+	for _, writer := range w.writers {
+		writer.close()
 	}
 
-	for partition := range w.batches {
-		delete(w.batches, partition)
+	for partition := range w.writers {
+		delete(w.writers, partition)
 	}
 
 	w.mutex.Unlock()
@@ -653,156 +653,27 @@ func (w *Writer) batchMessages(messages []Message, assignments map[topicPartitio
 		batches = make(map[*writeBatch][]int32, len(assignments))
 	}
 
-	batchSize := w.batchSize()
-	batchBytes := w.batchBytes()
-
 	w.mutex.Lock()
 	defer w.mutex.Unlock()
 
-	if w.batches == nil {
-		w.batches = map[topicPartition]*writeBatch{}
+	if w.writers == nil {
+		w.writers = map[topicPartition]*partitionWriter{}
 	}
 
 	for key, indexes := range assignments {
-		for _, i := range indexes {
-		assignMessage:
-			batch := w.batches[key]
-			if batch == nil {
-				batch = w.newWriteBatch(key)
-				w.batches[key] = batch
-			}
-			if !batch.add(messages[i], batchSize, batchBytes) {
-				batch.trigger()
-				delete(w.batches, key)
-				goto assignMessage
-			}
-			if batch.full(batchSize, batchBytes) {
-				batch.trigger()
-				delete(w.batches, key)
-			}
-			if !w.Async {
-				batches[batch] = append(batches[batch], i)
-			}
+		writer := w.writers[key]
+		if writer == nil {
+			writer = newPartitionWriter(w, key)
+			w.writers[key] = writer
 		}
-	}
+		wbatches := writer.writeMessages(messages, indexes)
 
-	return batches
-}
-
-func (w *Writer) newWriteBatch(key topicPartition) *writeBatch {
-	batch := newWriteBatch(time.Now(), w.batchTimeout())
-	w.group.Add(1)
-	go func() {
-		defer w.group.Done()
-		w.writeBatch(key, batch)
-	}()
-	return batch
-}
-
-func (w *Writer) writeBatch(key topicPartition, batch *writeBatch) {
-	// This goroutine has taken ownership of the batch, it is responsible
-	// for waiting for the batch to be ready (because it became full), or
-	// to timeout.
-	select {
-	case <-batch.timer.C:
-		// The batch timed out, we want to detach it from the writer to
-		// prevent more messages from being added.
-		w.mutex.Lock()
-		if batch == w.batches[key] {
-			delete(w.batches, key)
+		for batch, idxs := range wbatches {
+			batches[batch] = idxs
 		}
-		w.mutex.Unlock()
-
-	case <-batch.ready:
-		// The batch became full, it was removed from the writer and its
-		// ready channel was closed. We need to close the timer to avoid
-		// having it leak until it expires.
-		batch.timer.Stop()
 	}
 
-	stats := w.stats()
-	stats.batchTime.observe(int64(time.Since(batch.time)))
-	stats.batchSize.observe(int64(len(batch.msgs)))
-	stats.batchSizeBytes.observe(batch.bytes)
-
-	var res *ProduceResponse
-	var err error
-
-	for attempt, maxAttempts := 0, w.maxAttempts(); attempt < maxAttempts; attempt++ {
-		if attempt != 0 {
-			stats.retries.observe(1)
-			// TODO: should there be a way to asynchronously cancel this
-			// operation?
-			//
-			// * If all goroutines that added message to this batch have stopped
-			//   waiting for it, should we abort?
-			//
-			// * If the writer has been closed? It reduces the durability
-			//   guarantees to abort, but may be better to avoid long wait times
-			//   on close.
-			//
-			delay := backoff(attempt, 100*time.Millisecond, 1*time.Second)
-			w.withLogger(func(log Logger) {
-				log.Printf("backing off %s writing %d messages to %s (partition: %d)", delay, len(batch.msgs), key.topic, key.partition)
-			})
-			time.Sleep(delay)
-		}
-
-		w.withLogger(func(log Logger) {
-			log.Printf("writing %d messages to %s (partition: %d)", len(batch.msgs), key.topic, key.partition)
-		})
-
-		start := time.Now()
-		res, err = w.produce(key, batch)
-
-		stats.writes.observe(1)
-		stats.messages.observe(int64(len(batch.msgs)))
-		stats.bytes.observe(batch.bytes)
-		// stats.writeTime used to report the duration of WriteMessages, but the
-		// implementation was broken and reporting values in the nanoseconds
-		// range. In kafka-go 0.4, we recylced this value to instead report the
-		// duration of produce requests, and changed the stats.waitTime value to
-		// report the time that kafka has throttled the requests for.
-		stats.writeTime.observe(int64(time.Since(start)))
-
-		if res != nil {
-			err = res.Error
-			stats.waitTime.observe(int64(res.Throttle))
-		}
-
-		if err == nil {
-			break
-		}
-
-		stats.errors.observe(1)
-
-		w.withErrorLogger(func(log Logger) {
-			log.Printf("error writing messages to %s (partition %d): %s", key.topic, key.partition, err)
-		})
-
-		if !isTemporary(err) && !isTransientNetworkError(err) {
-			break
-		}
-	}
-
-	if res != nil {
-		for i := range batch.msgs {
-			m := &batch.msgs[i]
-			m.Topic = key.topic
-			m.Partition = int(key.partition)
-			m.Offset = res.BaseOffset + int64(i)
-
-			if m.Time.IsZero() {
-				m.Time = res.LogAppendTime
-			}
-		}
-	}
-
-	if w.Completion != nil {
-		w.Completion(batch.msgs, err)
-	}
-
-	batch.complete(err)
+	return batches
 }
 
 func (w *Writer) produce(key topicPartition, batch *writeBatch) (*ProduceResponse, error) {
@@ -992,6 +863,292 @@ func (w *Writer) chooseTopic(msg Message) (string, error) {
 	return w.Topic, nil
 }
 
+type batchQueue struct {
+	queue []*writeBatch
+
+	mutex sync.Mutex
+	cond  sync.Cond
+
+	closed bool
+}
+
+func (b *batchQueue) Put(batch *writeBatch) bool {
+	b.cond.L.Lock()
+	defer b.cond.L.Unlock()
+	defer b.cond.Broadcast()
+
+	if b.closed {
+		return false
+	}
+	b.queue = append(b.queue, batch)
+	return true
+}
+
+func (b *batchQueue) Get() *writeBatch {
+	b.cond.L.Lock()
+	defer b.cond.L.Unlock()
+
+	for len(b.queue) == 0 && !b.closed {
+		b.cond.Wait()
+	}
+
+	if len(b.queue) == 0 {
+		return nil
+	}
+
+	batch := b.queue[0]
+	b.queue[0] = nil
+	b.queue = b.queue[1:]
+
+	return batch
+}
+
+func (b *batchQueue) Close() {
+	b.cond.L.Lock()
+	defer b.cond.L.Unlock()
+	defer b.cond.Broadcast()
+
+	b.closed = true
+}
+
+func newBatchQueue(initialSize int) batchQueue {
+	bq := batchQueue{
+		queue: make([]*writeBatch, 0, initialSize),
+	}
+
+	bq.cond.L = &bq.mutex
+
+	return bq
+}
+
+// partitionWriter is a writer for a topic-partion pair. It maintains messaging order
+// across batches of messages.
+type partitionWriter struct {
+	meta  topicPartition
+	queue batchQueue
+
+	mutex     sync.Mutex
+	currBatch *writeBatch
+
+	group sync.WaitGroup
+
+	// reference to the writer that owns this batch. Used for the produce logic
+	// as well as stat tracking
+	w *Writer
+}
+
+func newPartitionWriter(w *Writer, key topicPartition) *partitionWriter {
+	writer := &partitionWriter{
+		meta:  key,
+		queue: newBatchQueue(10),
+		w:     w,
+	}
+	go func() {
+		writer.group.Add(1)
+		defer writer.group.Done()
+		writer.writeBatches()
+	}()
+
+	return writer
+}
+
+func (ptw *partitionWriter) writeBatches() {
+	for {
+		batch := ptw.queue.Get()
+
+		// The only time we can return nil is when the queue is closed
+		// and empty. If the queue is closed that means
+		// the Writer is closed so once we're here it's time to exit.
+		if batch == nil {
+			return
+		}
+
+		ptw.writeBatch(batch)
+
+	}
+}
+
+func (ptw *partitionWriter) writeMessages(msgs []Message, indexes []int32) map[*writeBatch][]int32 {
+	ptw.group.Add(1)
+	defer ptw.group.Done()
+
+	ptw.mutex.Lock()
+	defer ptw.mutex.Unlock()
+
+	batchSize := ptw.w.batchSize()
+	batchBytes := ptw.w.batchBytes()
+
+	var batches map[*writeBatch][]int32
+	if !ptw.w.Async {
+		batches = make(map[*writeBatch][]int32, 1)
+	}
+
+	for _, i := range indexes {
+	assignMessage:
+		batch := ptw.currBatch
+		if batch == nil {
+			batch = ptw.newWriteBatch()
+			ptw.currBatch = batch
+		}
+		if !batch.add(msgs[i], batchSize, batchBytes) {
+			batch.trigger()
+			ptw.queue.Put(batch)
+			ptw.currBatch = nil
+			goto assignMessage
+		}
+
+		if batch.full(batchSize, batchBytes) {
+			batch.trigger()
+			ptw.queue.Put(batch)
+			ptw.currBatch = nil
+		}
+
+		if !ptw.w.Async {
+			batches[batch] = append(batches[batch], i)
+		}
+	}
+	return batches
+}
+
+// ptw.w can be accessed here because this is called with the lock ptw.mutex already held.
+func (ptw *partitionWriter) newWriteBatch() *writeBatch {
+	batch := newWriteBatch(time.Now(), ptw.w.batchTimeout())
+	ptw.group.Add(1)
+	go func() {
+		defer ptw.group.Done()
+		ptw.awaitBatch(batch)
+	}()
+	return batch
+}
+
+// awaitBatch waits for a batch to either fill up or time out.
+// If the batch is full it only stops the timer, if the timer
+// expires it will queue the batch for writing if needed.
+func (ptw *partitionWriter) awaitBatch(batch *writeBatch) {
+	select {
+	case <-batch.timer.C:
+		ptw.mutex.Lock()
+		// detach the batch from the writer if we're still attached
+		// and queue for writing.
+		// Only the current batch can expire, all previous batches were already written to the queue.
+		// If writeMesseages locks pw.mutex after the timer fires but before this goroutine
+		// can lock pw.mutex it will either have filled the batch and enqueued it which will mean
+		// pw.currBatch != batch so we just move on.
+		// Otherwise, we detach the batch from the ptWriter and enqueue it for writing.
+		if ptw.currBatch == batch {
+			ptw.queue.Put(batch)
+			ptw.currBatch = nil
+		}
+		ptw.mutex.Unlock()
+	case <-batch.ready:
+		// The batch became full, it was removed from the ptwriter and its
+		// ready channel was closed. We need to close the timer to avoid
+		// having it leak until it expires.
+		batch.timer.Stop()
+	}
+}
+
+func (ptw *partitionWriter) writeBatch(batch *writeBatch) {
+	stats := ptw.w.stats()
+	stats.batchTime.observe(int64(time.Since(batch.time)))
+	stats.batchSize.observe(int64(len(batch.msgs)))
+	stats.batchSizeBytes.observe(batch.bytes)
+
+	var res *ProduceResponse
+	var err error
+	key := ptw.meta
+	for attempt, maxAttempts := 0, ptw.w.maxAttempts(); attempt < maxAttempts; attempt++ {
+		if attempt != 0 {
+			stats.retries.observe(1)
+			// TODO: should there be a way to asynchronously cancel this
+			// operation?
+			//
+			// * If all goroutines that added message to this batch have stopped
+			//   waiting for it, should we abort?
+			//
+			// * If the writer has been closed? It reduces the durability
+			//   guarantees to abort, but may be better to avoid long wait times
+			//   on close.
+			//
+			delay := backoff(attempt, 100*time.Millisecond, 1*time.Second)
+			ptw.w.withLogger(func(log Logger) {
+				log.Printf("backing off %s writing %d messages to %s (partition: %d)", delay, len(batch.msgs), key.topic, key.partition)
+			})
+			time.Sleep(delay)
+		}
+
+		ptw.w.withLogger(func(log Logger) {
+			log.Printf("writing %d messages to %s (partition: %d)", len(batch.msgs), key.topic, key.partition)
+		})
+
+		start := time.Now()
+		res, err = ptw.w.produce(key, batch)
+
+		stats.writes.observe(1)
+		stats.messages.observe(int64(len(batch.msgs)))
+		stats.bytes.observe(batch.bytes)
+		// stats.writeTime used to report the duration of WriteMessages, but the
+		// implementation was broken and reporting values in the nanoseconds
+		// range. In kafka-go 0.4, we recylced this value to instead report the
+		// duration of produce requests, and changed the stats.waitTime value to
+		// report the time that kafka has throttled the requests for.
+		stats.writeTime.observe(int64(time.Since(start)))
+
+		if res != nil {
+			err = res.Error
+			stats.waitTime.observe(int64(res.Throttle))
+		}
+
+		if err == nil {
+			break
+		}
+
+		stats.errors.observe(1)
+
+		ptw.w.withErrorLogger(func(log Logger) {
+			log.Printf("error writing messages to %s (partition %d): %s", key.topic, key.partition, err)
+		})
+
+		if !isTemporary(err) && !isTransientNetworkError(err) {
+			break
+		}
+	}
+
+	if res != nil {
+		for i := range batch.msgs {
+			m := &batch.msgs[i]
+			m.Topic = key.topic
+			m.Partition = int(key.partition)
+			m.Offset = res.BaseOffset + int64(i)
+
+			if m.Time.IsZero() {
+				m.Time = res.LogAppendTime
+			}
+		}
+	}
+
+	if ptw.w.Completion != nil {
+		ptw.w.Completion(batch.msgs, err)
+	}
+
+	batch.complete(err)
+}
+
+func (ptw *partitionWriter) close() {
+	ptw.mutex.Lock()
+	defer ptw.mutex.Unlock()
+
+	if ptw.currBatch != nil {
+		batch := ptw.currBatch
+		ptw.queue.Put(batch)
+		ptw.currBatch = nil
+		batch.trigger()
+	}
+	ptw.queue.Close()
+
+	ptw.group.Wait()
+}
+
 type writeBatch struct {
 	time  time.Time
 	msgs  []Message


Test patch:
diff --git a/writer_test.go b/writer_test.go
index 15e7212eb..f5b92b1ef 100644
--- a/writer_test.go
+++ b/writer_test.go
@@ -5,10 +5,93 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"strconv"
+	"sync"
 	"testing"
 	"time"
 )
 
+func TestBatchQueue(t *testing.T) {
+	tests := []struct {
+		scenario string
+		function func(*testing.T)
+	}{
+		{
+			scenario: "the remaining items in a queue can be gotten after closing",
+			function: testBatchQueueGetWorksAfterClose,
+		},
+		{
+			scenario: "putting into a closed queue fails",
+			function: testBatchQueuePutAfterCloseFails,
+		},
+		{
+			scenario: "putting into a queue awakes a goroutine in a get call",
+			function: testBatchQueuePutWakesSleepingGetter,
+		},
+	}
+
+	for _, test := range tests {
+		testFunc := test.function
+		t.Run(test.scenario, func(t *testing.T) {
+			t.Parallel()
+			testFunc(t)
+		})
+	}
+}
+
+func testBatchQueuePutWakesSleepingGetter(t *testing.T) {
+	bq := newBatchQueue(10)
+	var wg sync.WaitGroup
+	ready := make(chan struct{})
+	var batch *writeBatch
+	go func() {
+		wg.Add(1)
+		defer wg.Done()
+		close(ready)
+		batch = bq.Get()
+	}()
+	<-ready
+	bq.Put(newWriteBatch(time.Now(), time.Hour*100))
+	wg.Wait()
+	if batch == nil {
+		t.Fatal("got nil batch")
+	}
+}
+
+func testBatchQueuePutAfterCloseFails(t *testing.T) {
+	bq := newBatchQueue(10)
+	bq.Close()
+	if put := bq.Put(newWriteBatch(time.Now(), time.Hour*100)); put {
+		t.Fatal("put batch into closed queue")
+	}
+}
+
+func testBatchQueueGetWorksAfterClose(t *testing.T) {
+	bq := newBatchQueue(10)
+	enqueueBatches := []*writeBatch{
+		newWriteBatch(time.Now(), time.Hour*100),
+		newWriteBatch(time.Now(), time.Hour*100),
+	}
+
+	for _, batch := range enqueueBatches {
+		put := bq.Put(batch)
+		if !put {
+			t.Fatal("failed to put batch into queue")
+		}
+	}
+
+	bq.Close()
+
+	batchesGotten := 0
+	for batchesGotten != 2 {
+		dequeueBatch := bq.Get()
+		if dequeueBatch == nil {
+			t.Fatalf("no batch returned from get")
+		}
+		batchesGotten++
+	}
+}
+
 func TestWriter(t *testing.T) {
 	tests := []struct {
 		scenario string
@@ -123,7 +206,8 @@ func testWriterRequiredAcksNone(t *testing.T) {
 
 	msg := Message{
 		Key:   []byte("ThisIsAKey"),
-		Value: []byte("Test message for required acks test")}
+		Value: []byte("Test message for required acks test"),
+	}
 
 	err := writer.WriteMessages(context.Background(), msg)
 	if err != nil {
@@ -169,7 +253,6 @@ func testWriterRoundRobin1(t *testing.T) {
 	}
 
 	msgs, err := readPartition(topic, 0, offset)
-
 	if err != nil {
 		t.Error("error reading partition", err)
 		return
@@ -282,7 +365,7 @@ func testWriterMaxBytes(t *testing.T) {
 	}
 }
 
-// readOffset gets the latest offset for the given topic/partition
+// readOffset gets the latest offset for the given topic/partition.
 func readOffset(topic string, partition int) (offset int64, err error) {
 	var conn *Conn
 
@@ -350,39 +433,39 @@ func testWriterBatchBytes(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	if err := w.WriteMessages(ctx, []Message{
-		Message{Value: []byte("Hi")}, // 24 Bytes
-		Message{Value: []byte("By")}, // 24 Bytes
+		{Value: []byte("M0")}, // 24 Bytes
+		{Value: []byte("M1")}, // 24 Bytes
+		{Value: []byte("M2")}, // 24 Bytes
+		{Value: []byte("M3")}, // 24 Bytes
 	}...); err != nil {
 		t.Error(err)
 		return
 	}
 
-	if w.Stats().Writes > 1 {
-		t.Error("didn't batch messages")
+	if w.Stats().Writes != 2 {
+		t.Error("didn't create expected batches")
 		return
 	}
 	msgs, err := readPartition(topic, 0, offset)
-
 	if err != nil {
 		t.Error("error reading partition", err)
 		return
 	}
 
-	if len(msgs) != 2 {
+	if len(msgs) != 4 {
 		t.Error("bad messages in partition", msgs)
 		return
 	}
 
-	for _, m := range msgs {
-		if string(m.Value) == "Hi" || string(m.Value) == "By" {
+	for i, m := range msgs {
+		if string(m.Value) == "M"+strconv.Itoa(i) {
 			continue
 		}
-		t.Error("bad messages in partition", msgs)
+		t.Error("bad messages in partition", string(m.Value))
 	}
 }
 
 func testWriterBatchSize(t *testing.T) {
-
 	topic := makeTopic()
 	createTopic(t, topic, 1)
 	defer deleteTopic(t, topic)
@@ -402,7 +485,6 @@ func testWriterBatchSize(t *testing.T) {
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
-
 	if err := w.WriteMessages(ctx, []Message{
 		Message{Value: []byte("Hi")}, // 24 Bytes
 		Message{Value: []byte("By")}, // 24 Bytes
@@ -416,7 +498,6 @@ func testWriterBatchSize(t *testing.T) {
 		return
 	}
 	msgs, err := readPartition(topic, 0, offset)
-
 	if err != nil {
 		t.Error("error reading partition", err)
 		return
@@ -436,7 +517,6 @@ func testWriterBatchSize(t *testing.T) {
 }
 
 func testWriterSmallBatchBytes(t *testing.T) {
-
 	topic := makeTopic()
 	createTopic(t, topic, 1)
 	defer deleteTopic(t, topic)
@@ -469,7 +549,6 @@ func testWriterSmallBatchBytes(t *testing.T) {
 		return
 	}
 	msgs, err := readPartition(topic, 0, offset)
-
 	if err != nil {
 		t.Error("error reading partition", err)
 		return
@@ -489,7 +568,6 @@ func testWriterSmallBatchBytes(t *testing.T) {
 }
 
 func testWriterMultipleTopics(t *testing.T) {
-
 	topic1 := makeTopic()
 	createTopic(t, topic1, 1)
 	defer deleteTopic(t, topic1)