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