github.com/deanMdreon/kafka-go@v0.4.32/writer_test.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"io"
     8  	"math"
     9  	"strconv"
    10  	"sync"
    11  	"testing"
    12  	"time"
    13  )
    14  
    15  func TestBatchQueue(t *testing.T) {
    16  	tests := []struct {
    17  		scenario string
    18  		function func(*testing.T)
    19  	}{
    20  		{
    21  			scenario: "the remaining items in a queue can be gotten after closing",
    22  			function: testBatchQueueGetWorksAfterClose,
    23  		},
    24  		{
    25  			scenario: "putting into a closed queue fails",
    26  			function: testBatchQueuePutAfterCloseFails,
    27  		},
    28  		{
    29  			scenario: "putting into a queue awakes a goroutine in a get call",
    30  			function: testBatchQueuePutWakesSleepingGetter,
    31  		},
    32  	}
    33  
    34  	for _, test := range tests {
    35  		testFunc := test.function
    36  		t.Run(test.scenario, func(t *testing.T) {
    37  			t.Parallel()
    38  			testFunc(t)
    39  		})
    40  	}
    41  }
    42  
    43  func testBatchQueuePutWakesSleepingGetter(t *testing.T) {
    44  	bq := newBatchQueue(10)
    45  	var wg sync.WaitGroup
    46  	ready := make(chan struct{})
    47  	var batch *writeBatch
    48  	go func() {
    49  		wg.Add(1)
    50  		defer wg.Done()
    51  		close(ready)
    52  		batch = bq.Get()
    53  	}()
    54  	<-ready
    55  	bq.Put(newWriteBatch(time.Now(), time.Hour*100))
    56  	wg.Wait()
    57  	if batch == nil {
    58  		t.Fatal("got nil batch")
    59  	}
    60  }
    61  
    62  func testBatchQueuePutAfterCloseFails(t *testing.T) {
    63  	bq := newBatchQueue(10)
    64  	bq.Close()
    65  	if put := bq.Put(newWriteBatch(time.Now(), time.Hour*100)); put {
    66  		t.Fatal("put batch into closed queue")
    67  	}
    68  }
    69  
    70  func testBatchQueueGetWorksAfterClose(t *testing.T) {
    71  	bq := newBatchQueue(10)
    72  	enqueueBatches := []*writeBatch{
    73  		newWriteBatch(time.Now(), time.Hour*100),
    74  		newWriteBatch(time.Now(), time.Hour*100),
    75  	}
    76  
    77  	for _, batch := range enqueueBatches {
    78  		put := bq.Put(batch)
    79  		if !put {
    80  			t.Fatal("failed to put batch into queue")
    81  		}
    82  	}
    83  
    84  	bq.Close()
    85  
    86  	batchesGotten := 0
    87  	for batchesGotten != 2 {
    88  		dequeueBatch := bq.Get()
    89  		if dequeueBatch == nil {
    90  			t.Fatalf("no batch returned from get")
    91  		}
    92  		batchesGotten++
    93  	}
    94  }
    95  
    96  func TestWriter(t *testing.T) {
    97  	tests := []struct {
    98  		scenario string
    99  		function func(*testing.T)
   100  	}{
   101  		{
   102  			scenario: "closing a writer right after creating it returns promptly with no error",
   103  			function: testWriterClose,
   104  		},
   105  
   106  		{
   107  			scenario: "writing 1 message through a writer using round-robin balancing produces 1 message to the first partition",
   108  			function: testWriterRoundRobin1,
   109  		},
   110  
   111  		{
   112  			scenario: "running out of max attempts should return an error",
   113  			function: testWriterMaxAttemptsErr,
   114  		},
   115  
   116  		{
   117  			scenario: "writing a message larger then the max bytes should return an error",
   118  			function: testWriterMaxBytes,
   119  		},
   120  
   121  		{
   122  			scenario: "writing a batch of message based on batch byte size",
   123  			function: testWriterBatchBytes,
   124  		},
   125  
   126  		{
   127  			scenario: "writing a batch of messages",
   128  			function: testWriterBatchSize,
   129  		},
   130  
   131  		{
   132  			scenario: "writing messsages with a small batch byte size",
   133  			function: testWriterSmallBatchBytes,
   134  		},
   135  		{
   136  			scenario: "setting a non default balancer on the writer",
   137  			function: testWriterSetsRightBalancer,
   138  		},
   139  		{
   140  			scenario: "setting RequiredAcks to None in Writer does not cause a panic",
   141  			function: testWriterRequiredAcksNone,
   142  		},
   143  		{
   144  			scenario: "writing messages to multiple topics",
   145  			function: testWriterMultipleTopics,
   146  		},
   147  		{
   148  			scenario: "writing messages without specifying a topic",
   149  			function: testWriterMissingTopic,
   150  		},
   151  		{
   152  			scenario: "specifying topic for message when already set for writer",
   153  			function: testWriterUnexpectedMessageTopic,
   154  		},
   155  		{
   156  			scenario: "writing a message to an invalid partition",
   157  			function: testWriterInvalidPartition,
   158  		},
   159  		{
   160  			scenario: "writing a message to a non-existant topic creates the topic",
   161  			function: testWriterAutoCreateTopic,
   162  		},
   163  	}
   164  
   165  	for _, test := range tests {
   166  		testFunc := test.function
   167  		t.Run(test.scenario, func(t *testing.T) {
   168  			t.Parallel()
   169  			testFunc(t)
   170  		})
   171  	}
   172  }
   173  
   174  func newTestWriter(config WriterConfig) *Writer {
   175  	if len(config.Brokers) == 0 {
   176  		config.Brokers = []string{"localhost:9092"}
   177  	}
   178  	return NewWriter(config)
   179  }
   180  
   181  func testWriterClose(t *testing.T) {
   182  	const topic = "test-writer-0"
   183  	createTopic(t, topic, 1)
   184  	defer deleteTopic(t, topic)
   185  
   186  	w := newTestWriter(WriterConfig{
   187  		Topic: topic,
   188  	})
   189  
   190  	if err := w.Close(); err != nil {
   191  		t.Error(err)
   192  	}
   193  }
   194  
   195  func testWriterRequiredAcksNone(t *testing.T) {
   196  	topic := makeTopic()
   197  	createTopic(t, topic, 1)
   198  	defer deleteTopic(t, topic)
   199  
   200  	transport := &Transport{}
   201  	defer transport.CloseIdleConnections()
   202  
   203  	writer := &Writer{
   204  		Addr:         TCP("localhost:9092"),
   205  		Topic:        topic,
   206  		Balancer:     &RoundRobin{},
   207  		RequiredAcks: RequireNone,
   208  		Transport:    transport,
   209  	}
   210  	defer writer.Close()
   211  
   212  	msg := Message{
   213  		Key:   []byte("ThisIsAKey"),
   214  		Value: []byte("Test message for required acks test"),
   215  	}
   216  
   217  	err := writer.WriteMessages(context.Background(), msg)
   218  	if err != nil {
   219  		t.Fatal(err)
   220  	}
   221  }
   222  
   223  func testWriterSetsRightBalancer(t *testing.T) {
   224  	const topic = "test-writer-1"
   225  	balancer := &CRC32Balancer{}
   226  	w := newTestWriter(WriterConfig{
   227  		Topic:    topic,
   228  		Balancer: balancer,
   229  	})
   230  	defer w.Close()
   231  
   232  	if w.Balancer != balancer {
   233  		t.Errorf("Balancer not set correctly")
   234  	}
   235  }
   236  
   237  func testWriterRoundRobin1(t *testing.T) {
   238  	const topic = "test-writer-1"
   239  	createTopic(t, topic, 1)
   240  	defer deleteTopic(t, topic)
   241  
   242  	offset, err := readOffset(topic, 0)
   243  	if err != nil {
   244  		t.Fatal(err)
   245  	}
   246  
   247  	w := newTestWriter(WriterConfig{
   248  		Topic:    topic,
   249  		Balancer: &RoundRobin{},
   250  	})
   251  	defer w.Close()
   252  
   253  	if err := w.WriteMessages(context.Background(), Message{
   254  		Value: []byte("Hello World!"),
   255  	}); err != nil {
   256  		t.Error(err)
   257  		return
   258  	}
   259  
   260  	msgs, err := readPartition(topic, 0, offset)
   261  	if err != nil {
   262  		t.Error("error reading partition", err)
   263  		return
   264  	}
   265  
   266  	if len(msgs) != 1 {
   267  		t.Error("bad messages in partition", msgs)
   268  		return
   269  	}
   270  
   271  	for _, m := range msgs {
   272  		if string(m.Value) != "Hello World!" {
   273  			t.Error("bad messages in partition", msgs)
   274  			break
   275  		}
   276  	}
   277  }
   278  
   279  func TestValidateWriter(t *testing.T) {
   280  	tests := []struct {
   281  		config       WriterConfig
   282  		errorOccured bool
   283  	}{
   284  		{config: WriterConfig{}, errorOccured: true},
   285  		{config: WriterConfig{Brokers: []string{"broker1", "broker2"}}, errorOccured: false},
   286  		{config: WriterConfig{Brokers: []string{"broker1"}, Topic: "topic1"}, errorOccured: false},
   287  	}
   288  	for _, test := range tests {
   289  		err := test.config.Validate()
   290  		if test.errorOccured && err == nil {
   291  			t.Fail()
   292  		}
   293  		if !test.errorOccured && err != nil {
   294  			t.Fail()
   295  		}
   296  	}
   297  }
   298  
   299  func testWriterMaxAttemptsErr(t *testing.T) {
   300  	topic := makeTopic()
   301  	createTopic(t, topic, 1)
   302  	defer deleteTopic(t, topic)
   303  
   304  	w := newTestWriter(WriterConfig{
   305  		Brokers:     []string{"localhost:9999"}, // nothing is listening here
   306  		Topic:       topic,
   307  		MaxAttempts: 3,
   308  		Balancer:    &RoundRobin{},
   309  	})
   310  	defer w.Close()
   311  
   312  	if err := w.WriteMessages(context.Background(), Message{
   313  		Value: []byte("Hello World!"),
   314  	}); err == nil {
   315  		t.Error("expected error")
   316  		return
   317  	}
   318  }
   319  
   320  func testWriterMaxBytes(t *testing.T) {
   321  	topic := makeTopic()
   322  	createTopic(t, topic, 1)
   323  	defer deleteTopic(t, topic)
   324  
   325  	w := newTestWriter(WriterConfig{
   326  		Topic:      topic,
   327  		BatchBytes: 25,
   328  	})
   329  	defer w.Close()
   330  
   331  	if err := w.WriteMessages(context.Background(), Message{
   332  		Value: []byte("Hi"),
   333  	}); err != nil {
   334  		t.Error(err)
   335  		return
   336  	}
   337  
   338  	firstMsg := []byte("Hello World!")
   339  	secondMsg := []byte("LeftOver!")
   340  	msgs := []Message{
   341  		{
   342  			Value: firstMsg,
   343  		},
   344  		{
   345  			Value: secondMsg,
   346  		},
   347  	}
   348  	if err := w.WriteMessages(context.Background(), msgs...); err == nil {
   349  		t.Error("expected error")
   350  		return
   351  	} else if err != nil {
   352  		switch e := err.(type) {
   353  		case MessageTooLargeError:
   354  			if string(e.Message.Value) != string(firstMsg) {
   355  				t.Errorf("unxpected returned message. Expected: %s, Got %s", firstMsg, e.Message.Value)
   356  				return
   357  			}
   358  			if len(e.Remaining) != 1 {
   359  				t.Error("expected remaining errors; found none")
   360  				return
   361  			}
   362  			if string(e.Remaining[0].Value) != string(secondMsg) {
   363  				t.Errorf("unxpected returned message. Expected: %s, Got %s", secondMsg, e.Message.Value)
   364  				return
   365  			}
   366  		default:
   367  			t.Errorf("unexpected error: %s", err)
   368  			return
   369  		}
   370  	}
   371  }
   372  
   373  // readOffset gets the latest offset for the given topic/partition.
   374  func readOffset(topic string, partition int) (offset int64, err error) {
   375  	var conn *Conn
   376  
   377  	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
   378  	defer cancel()
   379  
   380  	if conn, err = DialLeader(ctx, "tcp", "localhost:9092", topic, partition); err != nil {
   381  		err = fmt.Errorf("readOffset, DialLeader: %w", err)
   382  		return
   383  	}
   384  	defer conn.Close()
   385  
   386  	offset, err = conn.ReadLastOffset()
   387  	if err != nil {
   388  		err = fmt.Errorf("readOffset, conn.ReadLastOffset: %w", err)
   389  	}
   390  	return
   391  }
   392  
   393  func readPartition(topic string, partition int, offset int64) (msgs []Message, err error) {
   394  	var conn *Conn
   395  
   396  	if conn, err = DialLeader(context.Background(), "tcp", "localhost:9092", topic, partition); err != nil {
   397  		return
   398  	}
   399  	defer conn.Close()
   400  
   401  	conn.Seek(offset, SeekAbsolute)
   402  	conn.SetReadDeadline(time.Now().Add(10 * time.Second))
   403  	batch := conn.ReadBatch(0, 1000000000)
   404  	defer batch.Close()
   405  
   406  	for {
   407  		var msg Message
   408  
   409  		if msg, err = batch.ReadMessage(); err != nil {
   410  			if err == io.EOF {
   411  				err = nil
   412  			}
   413  			return
   414  		}
   415  
   416  		msgs = append(msgs, msg)
   417  	}
   418  }
   419  
   420  func testWriterBatchBytes(t *testing.T) {
   421  	topic := makeTopic()
   422  	createTopic(t, topic, 1)
   423  	defer deleteTopic(t, topic)
   424  
   425  	offset, err := readOffset(topic, 0)
   426  	if err != nil {
   427  		t.Fatal(err)
   428  	}
   429  
   430  	w := newTestWriter(WriterConfig{
   431  		Topic:        topic,
   432  		BatchBytes:   48,
   433  		BatchTimeout: math.MaxInt32 * time.Second,
   434  		Balancer:     &RoundRobin{},
   435  	})
   436  	defer w.Close()
   437  
   438  	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
   439  	defer cancel()
   440  	if err := w.WriteMessages(ctx, []Message{
   441  		{Value: []byte("M0")}, // 24 Bytes
   442  		{Value: []byte("M1")}, // 24 Bytes
   443  		{Value: []byte("M2")}, // 24 Bytes
   444  		{Value: []byte("M3")}, // 24 Bytes
   445  	}...); err != nil {
   446  		t.Error(err)
   447  		return
   448  	}
   449  
   450  	if w.Stats().Writes != 2 {
   451  		t.Error("didn't create expected batches")
   452  		return
   453  	}
   454  	msgs, err := readPartition(topic, 0, offset)
   455  	if err != nil {
   456  		t.Error("error reading partition", err)
   457  		return
   458  	}
   459  
   460  	if len(msgs) != 4 {
   461  		t.Error("bad messages in partition", msgs)
   462  		return
   463  	}
   464  
   465  	for i, m := range msgs {
   466  		if string(m.Value) == "M"+strconv.Itoa(i) {
   467  			continue
   468  		}
   469  		t.Error("bad messages in partition", string(m.Value))
   470  	}
   471  }
   472  
   473  func testWriterBatchSize(t *testing.T) {
   474  	topic := makeTopic()
   475  	createTopic(t, topic, 1)
   476  	defer deleteTopic(t, topic)
   477  
   478  	offset, err := readOffset(topic, 0)
   479  	if err != nil {
   480  		t.Fatal(err)
   481  	}
   482  
   483  	w := newTestWriter(WriterConfig{
   484  		Topic:        topic,
   485  		BatchSize:    2,
   486  		BatchTimeout: math.MaxInt32 * time.Second,
   487  		Balancer:     &RoundRobin{},
   488  	})
   489  	defer w.Close()
   490  
   491  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
   492  	defer cancel()
   493  	if err := w.WriteMessages(ctx, []Message{
   494  		{Value: []byte("Hi")}, // 24 Bytes
   495  		{Value: []byte("By")}, // 24 Bytes
   496  	}...); err != nil {
   497  		t.Error(err)
   498  		return
   499  	}
   500  
   501  	if w.Stats().Writes > 1 {
   502  		t.Error("didn't batch messages")
   503  		return
   504  	}
   505  	msgs, err := readPartition(topic, 0, offset)
   506  	if err != nil {
   507  		t.Error("error reading partition", err)
   508  		return
   509  	}
   510  
   511  	if len(msgs) != 2 {
   512  		t.Error("bad messages in partition", msgs)
   513  		return
   514  	}
   515  
   516  	for _, m := range msgs {
   517  		if string(m.Value) == "Hi" || string(m.Value) == "By" {
   518  			continue
   519  		}
   520  		t.Error("bad messages in partition", msgs)
   521  	}
   522  }
   523  
   524  func testWriterSmallBatchBytes(t *testing.T) {
   525  	topic := makeTopic()
   526  	createTopic(t, topic, 1)
   527  	defer deleteTopic(t, topic)
   528  
   529  	offset, err := readOffset(topic, 0)
   530  	if err != nil {
   531  		t.Fatal(err)
   532  	}
   533  
   534  	w := newTestWriter(WriterConfig{
   535  		Topic:        topic,
   536  		BatchBytes:   25,
   537  		BatchTimeout: 50 * time.Millisecond,
   538  		Balancer:     &RoundRobin{},
   539  	})
   540  	defer w.Close()
   541  
   542  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
   543  	defer cancel()
   544  	if err := w.WriteMessages(ctx, []Message{
   545  		{Value: []byte("Hi")}, // 24 Bytes
   546  		{Value: []byte("By")}, // 24 Bytes
   547  	}...); err != nil {
   548  		t.Error(err)
   549  		return
   550  	}
   551  	ws := w.Stats()
   552  	if ws.Writes != 2 {
   553  		t.Error("didn't batch messages; Writes: ", ws.Writes)
   554  		return
   555  	}
   556  	msgs, err := readPartition(topic, 0, offset)
   557  	if err != nil {
   558  		t.Error("error reading partition", err)
   559  		return
   560  	}
   561  
   562  	if len(msgs) != 2 {
   563  		t.Error("bad messages in partition", msgs)
   564  		return
   565  	}
   566  
   567  	for _, m := range msgs {
   568  		if string(m.Value) == "Hi" || string(m.Value) == "By" {
   569  			continue
   570  		}
   571  		t.Error("bad messages in partition", msgs)
   572  	}
   573  }
   574  
   575  func testWriterMultipleTopics(t *testing.T) {
   576  	topic1 := makeTopic()
   577  	createTopic(t, topic1, 1)
   578  	defer deleteTopic(t, topic1)
   579  
   580  	offset1, err := readOffset(topic1, 0)
   581  	if err != nil {
   582  		t.Fatal(err)
   583  	}
   584  
   585  	topic2 := makeTopic()
   586  	createTopic(t, topic2, 1)
   587  	defer deleteTopic(t, topic2)
   588  
   589  	offset2, err := readOffset(topic2, 0)
   590  	if err != nil {
   591  		t.Fatal(err)
   592  	}
   593  
   594  	w := newTestWriter(WriterConfig{
   595  		Balancer: &RoundRobin{},
   596  	})
   597  	defer w.Close()
   598  
   599  	msg1 := Message{Topic: topic1, Value: []byte("Hello")}
   600  	msg2 := Message{Topic: topic2, Value: []byte("World")}
   601  
   602  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
   603  	defer cancel()
   604  	if err := w.WriteMessages(ctx, msg1, msg2); err != nil {
   605  		t.Error(err)
   606  		return
   607  	}
   608  	ws := w.Stats()
   609  	if ws.Writes != 2 {
   610  		t.Error("didn't batch messages; Writes: ", ws.Writes)
   611  		return
   612  	}
   613  
   614  	msgs1, err := readPartition(topic1, 0, offset1)
   615  	if err != nil {
   616  		t.Error("error reading partition", err)
   617  		return
   618  	}
   619  	if len(msgs1) != 1 {
   620  		t.Error("bad messages in partition", msgs1)
   621  		return
   622  	}
   623  	if string(msgs1[0].Value) != "Hello" {
   624  		t.Error("bad message in partition", msgs1)
   625  	}
   626  
   627  	msgs2, err := readPartition(topic2, 0, offset2)
   628  	if err != nil {
   629  		t.Error("error reading partition", err)
   630  		return
   631  	}
   632  	if len(msgs2) != 1 {
   633  		t.Error("bad messages in partition", msgs2)
   634  		return
   635  	}
   636  	if string(msgs2[0].Value) != "World" {
   637  		t.Error("bad message in partition", msgs2)
   638  	}
   639  }
   640  
   641  func testWriterMissingTopic(t *testing.T) {
   642  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
   643  	defer cancel()
   644  
   645  	w := newTestWriter(WriterConfig{
   646  		// no topic
   647  		Balancer: &RoundRobin{},
   648  	})
   649  	defer w.Close()
   650  
   651  	msg := Message{Value: []byte("Hello World")} // no topic
   652  
   653  	if err := w.WriteMessages(ctx, msg); err == nil {
   654  		t.Error("expected error")
   655  		return
   656  	}
   657  }
   658  
   659  func testWriterInvalidPartition(t *testing.T) {
   660  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
   661  	defer cancel()
   662  
   663  	topic := makeTopic()
   664  	createTopic(t, topic, 1)
   665  	defer deleteTopic(t, topic)
   666  
   667  	w := newTestWriter(WriterConfig{
   668  		Topic:       topic,
   669  		MaxAttempts: 1,                              // only try once to get the error back immediately
   670  		Balancer:    &staticBalancer{partition: -1}, // intentionally invalid partition
   671  	})
   672  	defer w.Close()
   673  
   674  	msg := Message{
   675  		Value: []byte("Hello World!"),
   676  	}
   677  
   678  	// this call should return an error and not panic (see issue #517)
   679  	if err := w.WriteMessages(ctx, msg); err == nil {
   680  		t.Fatal("expected error attempting to write message")
   681  	}
   682  }
   683  
   684  func testWriterUnexpectedMessageTopic(t *testing.T) {
   685  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
   686  	defer cancel()
   687  
   688  	topic := makeTopic()
   689  	createTopic(t, topic, 1)
   690  	defer deleteTopic(t, topic)
   691  
   692  	w := newTestWriter(WriterConfig{
   693  		Topic:    topic,
   694  		Balancer: &RoundRobin{},
   695  	})
   696  	defer w.Close()
   697  
   698  	msg := Message{Topic: "should-fail", Value: []byte("Hello World")}
   699  
   700  	if err := w.WriteMessages(ctx, msg); err == nil {
   701  		t.Error("expected error")
   702  		return
   703  	}
   704  }
   705  
   706  func testWriterAutoCreateTopic(t *testing.T) {
   707  	topic := makeTopic()
   708  	// Assume it's going to get created.
   709  	defer deleteTopic(t, topic)
   710  
   711  	w := newTestWriter(WriterConfig{
   712  		Topic:    topic,
   713  		Balancer: &RoundRobin{},
   714  	})
   715  	defer w.Close()
   716  
   717  	msg := Message{Key: []byte("key"), Value: []byte("Hello World")}
   718  
   719  	var err error
   720  	const retries = 5
   721  	for i := 0; i < retries; i++ {
   722  		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
   723  		defer cancel()
   724  		err = w.WriteMessages(ctx, msg)
   725  		if errors.Is(err, LeaderNotAvailable) || errors.Is(err, context.DeadlineExceeded) {
   726  			time.Sleep(time.Millisecond * 250)
   727  			continue
   728  		}
   729  
   730  		if err != nil {
   731  			t.Errorf("unexpected error %v", err)
   732  			return
   733  		}
   734  	}
   735  	if err != nil {
   736  		t.Errorf("unable to create topic %v", err)
   737  	}
   738  }
   739  
   740  type staticBalancer struct {
   741  	partition int
   742  }
   743  
   744  func (b *staticBalancer) Balance(_ Message, partitions ...int) int {
   745  	return b.partition
   746  }