writer_test.go

package kafka

import (
	"context"
	"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
		function func(*testing.T)
	}{
		{
			scenario: "closing a writer right after creating it returns promptly with no error",
			function: testWriterClose,
		},

		{
			scenario: "writing 1 message through a writer using round-robin balancing produces 1 message to the first partition",
			function: testWriterRoundRobin1,
		},

		{
			scenario: "running out of max attempts should return an error",
			function: testWriterMaxAttemptsErr,
		},

		{
			scenario: "writing a message larger then the max bytes should return an error",
			function: testWriterMaxBytes,
		},

		{
			scenario: "writing a batch of message based on batch byte size",
			function: testWriterBatchBytes,
		},

		{
			scenario: "writing a batch of messages",
			function: testWriterBatchSize,
		},

		{
			scenario: "writing messsages with a small batch byte size",
			function: testWriterSmallBatchBytes,
		},
		{
			scenario: "setting a non default balancer on the writer",
			function: testWriterSetsRightBalancer,
		},
		{
			scenario: "setting RequiredAcks to None in Writer does not cause a panic",
			function: testWriterRequiredAcksNone,
		},
		{
			scenario: "writing messages to multiple topics",
			function: testWriterMultipleTopics,
		},
		{
			scenario: "writing messages without specifying a topic",
			function: testWriterMissingTopic,
		},
		{
			scenario: "specifying topic for message when already set for writer",
			function: testWriterUnexpectedMessageTopic,
		},
		{
			scenario: "writing a message to an invalid partition",
			function: testWriterInvalidPartition,
		},
	}

	for _, test := range tests {
		testFunc := test.function
		t.Run(test.scenario, func(t *testing.T) {
			t.Parallel()
			testFunc(t)
		})
	}
}

func newTestWriter(config WriterConfig) *Writer {
	if len(config.Brokers) == 0 {
		config.Brokers = []string{"localhost:9092"}
	}
	return NewWriter(config)
}

func testWriterClose(t *testing.T) {
	const topic = "test-writer-0"
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	w := newTestWriter(WriterConfig{
		Topic: topic,
	})

	if err := w.Close(); err != nil {
		t.Error(err)
	}
}

func testWriterRequiredAcksNone(t *testing.T) {
	topic := makeTopic()
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	transport := &Transport{}
	defer transport.CloseIdleConnections()

	writer := &Writer{
		Addr:         TCP("localhost:9092"),
		Topic:        topic,
		Balancer:     &RoundRobin{},
		RequiredAcks: RequireNone,
		Transport:    transport,
	}
	defer writer.Close()

	msg := Message{
		Key:   []byte("ThisIsAKey"),
		Value: []byte("Test message for required acks test"),
	}

	err := writer.WriteMessages(context.Background(), msg)
	if err != nil {
		t.Fatal(err)
	}
}

func testWriterSetsRightBalancer(t *testing.T) {
	const topic = "test-writer-1"
	balancer := &CRC32Balancer{}
	w := newTestWriter(WriterConfig{
		Topic:    topic,
		Balancer: balancer,
	})
	defer w.Close()

	if w.Balancer != balancer {
		t.Errorf("Balancer not set correctly")
	}
}

func testWriterRoundRobin1(t *testing.T) {
	const topic = "test-writer-1"
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	offset, err := readOffset(topic, 0)
	if err != nil {
		t.Fatal(err)
	}

	w := newTestWriter(WriterConfig{
		Topic:    topic,
		Balancer: &RoundRobin{},
	})
	defer w.Close()

	if err := w.WriteMessages(context.Background(), Message{
		Value: []byte("Hello World!"),
	}); err != nil {
		t.Error(err)
		return
	}

	msgs, err := readPartition(topic, 0, offset)
	if err != nil {
		t.Error("error reading partition", err)
		return
	}

	if len(msgs) != 1 {
		t.Error("bad messages in partition", msgs)
		return
	}

	for _, m := range msgs {
		if string(m.Value) != "Hello World!" {
			t.Error("bad messages in partition", msgs)
			break
		}
	}
}

func TestValidateWriter(t *testing.T) {
	tests := []struct {
		config       WriterConfig
		errorOccured bool
	}{
		{config: WriterConfig{}, errorOccured: true},
		{config: WriterConfig{Brokers: []string{"broker1", "broker2"}}, errorOccured: false},
		{config: WriterConfig{Brokers: []string{"broker1"}, Topic: "topic1"}, errorOccured: false},
	}
	for _, test := range tests {
		err := test.config.Validate()
		if test.errorOccured && err == nil {
			t.Fail()
		}
		if !test.errorOccured && err != nil {
			t.Fail()
		}
	}
}

func testWriterMaxAttemptsErr(t *testing.T) {
	topic := makeTopic()
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	w := newTestWriter(WriterConfig{
		Brokers:     []string{"localhost:9999"}, // nothing is listening here
		Topic:       topic,
		MaxAttempts: 3,
		Balancer:    &RoundRobin{},
	})
	defer w.Close()

	if err := w.WriteMessages(context.Background(), Message{
		Value: []byte("Hello World!"),
	}); err == nil {
		t.Error("expected error")
		return
	}
}

func testWriterMaxBytes(t *testing.T) {
	topic := makeTopic()
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	w := newTestWriter(WriterConfig{
		Topic:      topic,
		BatchBytes: 25,
	})
	defer w.Close()

	if err := w.WriteMessages(context.Background(), Message{
		Value: []byte("Hi"),
	}); err != nil {
		t.Error(err)
		return
	}

	firstMsg := []byte("Hello World!")
	secondMsg := []byte("LeftOver!")
	msgs := []Message{
		{
			Value: firstMsg,
		},
		{
			Value: secondMsg,
		},
	}
	if err := w.WriteMessages(context.Background(), msgs...); err == nil {
		t.Error("expected error")
		return
	} else if err != nil {
		switch e := err.(type) {
		case MessageTooLargeError:
			if string(e.Message.Value) != string(firstMsg) {
				t.Errorf("unxpected returned message. Expected: %s, Got %s", firstMsg, e.Message.Value)
				return
			}
			if len(e.Remaining) != 1 {
				t.Error("expected remaining errors; found none")
				return
			}
			if string(e.Remaining[0].Value) != string(secondMsg) {
				t.Errorf("unxpected returned message. Expected: %s, Got %s", secondMsg, e.Message.Value)
				return
			}
		default:
			t.Errorf("unexpected error: %s", err)
			return
		}
	}
}

// readOffset gets the latest offset for the given topic/partition.
func readOffset(topic string, partition int) (offset int64, err error) {
	var conn *Conn

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if conn, err = DialLeader(ctx, "tcp", "localhost:9092", topic, partition); err != nil {
		err = fmt.Errorf("readOffset, DialLeader: %w", err)
		return
	}
	defer conn.Close()

	offset, err = conn.ReadLastOffset()
	if err != nil {
		err = fmt.Errorf("readOffset, conn.ReadLastOffset: %w", err)
	}
	return
}

func readPartition(topic string, partition int, offset int64) (msgs []Message, err error) {
	var conn *Conn

	if conn, err = DialLeader(context.Background(), "tcp", "localhost:9092", topic, partition); err != nil {
		return
	}
	defer conn.Close()

	conn.Seek(offset, SeekAbsolute)
	conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	batch := conn.ReadBatch(0, 1000000000)
	defer batch.Close()

	for {
		var msg Message

		if msg, err = batch.ReadMessage(); err != nil {
			if err == io.EOF {
				err = nil
			}
			return
		}

		msgs = append(msgs, msg)
	}
}

func testWriterBatchBytes(t *testing.T) {
	topic := makeTopic()
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	offset, err := readOffset(topic, 0)
	if err != nil {
		t.Fatal(err)
	}

	w := newTestWriter(WriterConfig{
		Topic:        topic,
		BatchBytes:   48,
		BatchTimeout: math.MaxInt32 * time.Second,
		Balancer:     &RoundRobin{},
	})
	defer w.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	if err := w.WriteMessages(ctx, []Message{
		{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 != 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) != 4 {
		t.Error("bad messages in partition", msgs)
		return
	}

	for i, m := range msgs {
		if string(m.Value) == "M"+strconv.Itoa(i) {
			continue
		}
		t.Error("bad messages in partition", string(m.Value))
	}
}

func testWriterBatchSize(t *testing.T) {
	topic := makeTopic()
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	offset, err := readOffset(topic, 0)
	if err != nil {
		t.Fatal(err)
	}

	w := newTestWriter(WriterConfig{
		Topic:        topic,
		BatchSize:    2,
		BatchTimeout: math.MaxInt32 * time.Second,
		Balancer:     &RoundRobin{},
	})
	defer w.Close()

	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
	}...); err != nil {
		t.Error(err)
		return
	}

	if w.Stats().Writes > 1 {
		t.Error("didn't batch messages")
		return
	}
	msgs, err := readPartition(topic, 0, offset)
	if err != nil {
		t.Error("error reading partition", err)
		return
	}

	if len(msgs) != 2 {
		t.Error("bad messages in partition", msgs)
		return
	}

	for _, m := range msgs {
		if string(m.Value) == "Hi" || string(m.Value) == "By" {
			continue
		}
		t.Error("bad messages in partition", msgs)
	}
}

func testWriterSmallBatchBytes(t *testing.T) {
	topic := makeTopic()
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	offset, err := readOffset(topic, 0)
	if err != nil {
		t.Fatal(err)
	}

	w := newTestWriter(WriterConfig{
		Topic:        topic,
		BatchBytes:   25,
		BatchTimeout: 50 * time.Millisecond,
		Balancer:     &RoundRobin{},
	})
	defer w.Close()

	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
	}...); err != nil {
		t.Error(err)
		return
	}
	ws := w.Stats()
	if ws.Writes != 2 {
		t.Error("didn't batch messages; Writes: ", ws.Writes)
		return
	}
	msgs, err := readPartition(topic, 0, offset)
	if err != nil {
		t.Error("error reading partition", err)
		return
	}

	if len(msgs) != 2 {
		t.Error("bad messages in partition", msgs)
		return
	}

	for _, m := range msgs {
		if string(m.Value) == "Hi" || string(m.Value) == "By" {
			continue
		}
		t.Error("bad messages in partition", msgs)
	}
}

func testWriterMultipleTopics(t *testing.T) {
	topic1 := makeTopic()
	createTopic(t, topic1, 1)
	defer deleteTopic(t, topic1)

	offset1, err := readOffset(topic1, 0)
	if err != nil {
		t.Fatal(err)
	}

	topic2 := makeTopic()
	createTopic(t, topic2, 1)
	defer deleteTopic(t, topic2)

	offset2, err := readOffset(topic2, 0)
	if err != nil {
		t.Fatal(err)
	}

	w := newTestWriter(WriterConfig{
		Balancer: &RoundRobin{},
	})
	defer w.Close()

	msg1 := Message{Topic: topic1, Value: []byte("Hello")}
	msg2 := Message{Topic: topic2, Value: []byte("World")}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := w.WriteMessages(ctx, msg1, msg2); err != nil {
		t.Error(err)
		return
	}
	ws := w.Stats()
	if ws.Writes != 2 {
		t.Error("didn't batch messages; Writes: ", ws.Writes)
		return
	}

	msgs1, err := readPartition(topic1, 0, offset1)
	if err != nil {
		t.Error("error reading partition", err)
		return
	}
	if len(msgs1) != 1 {
		t.Error("bad messages in partition", msgs1)
		return
	}
	if string(msgs1[0].Value) != "Hello" {
		t.Error("bad message in partition", msgs1)
	}

	msgs2, err := readPartition(topic2, 0, offset2)
	if err != nil {
		t.Error("error reading partition", err)
		return
	}
	if len(msgs2) != 1 {
		t.Error("bad messages in partition", msgs2)
		return
	}
	if string(msgs2[0].Value) != "World" {
		t.Error("bad message in partition", msgs2)
	}
}

func testWriterMissingTopic(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	w := newTestWriter(WriterConfig{
		// no topic
		Balancer: &RoundRobin{},
	})
	defer w.Close()

	msg := Message{Value: []byte("Hello World")} // no topic

	if err := w.WriteMessages(ctx, msg); err == nil {
		t.Error("expected error")
		return
	}
}

func testWriterInvalidPartition(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	topic := makeTopic()
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	w := newTestWriter(WriterConfig{
		Topic:       topic,
		MaxAttempts: 1,                              // only try once to get the error back immediately
		Balancer:    &staticBalancer{partition: -1}, // intentionally invalid partition
	})
	defer w.Close()

	msg := Message{
		Value: []byte("Hello World!"),
	}

	// this call should return an error and not panic (see issue #517)
	if err := w.WriteMessages(ctx, msg); err == nil {
		t.Fatal("expected error attempting to write message")
	}
}

func testWriterUnexpectedMessageTopic(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	topic := makeTopic()
	createTopic(t, topic, 1)
	defer deleteTopic(t, topic)

	w := newTestWriter(WriterConfig{
		Topic:    topic,
		Balancer: &RoundRobin{},
	})
	defer w.Close()

	msg := Message{Topic: "should-fail", Value: []byte("Hello World")}

	if err := w.WriteMessages(ctx, msg); err == nil {
		t.Error("expected error")
		return
	}
}

type staticBalancer struct {
	partition int
}

func (b *staticBalancer) Balance(_ Message, partitions ...int) int {
	return b.partition
}