github.com/rbisecke/kafka-go@v0.4.27/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 }