github.com/Jeffail/benthos/v3@v3.65.0/lib/input/socket_test.go (about) 1 package input 2 3 import ( 4 "errors" 5 "fmt" 6 "net" 7 "path/filepath" 8 "reflect" 9 "sync" 10 "testing" 11 "time" 12 13 "github.com/Jeffail/benthos/v3/lib/input/reader" 14 "github.com/Jeffail/benthos/v3/lib/log" 15 "github.com/Jeffail/benthos/v3/lib/message" 16 "github.com/Jeffail/benthos/v3/lib/metrics" 17 "github.com/Jeffail/benthos/v3/lib/response" 18 "github.com/Jeffail/benthos/v3/lib/types" 19 "github.com/stretchr/testify/assert" 20 "github.com/stretchr/testify/require" 21 ) 22 23 func TestSocketBasic(t *testing.T) { 24 tmpDir := t.TempDir() 25 26 ln, err := net.Listen("unix", filepath.Join(tmpDir, "benthos.sock")) 27 if err != nil { 28 t.Fatalf("failed to listen on a address: %v", err) 29 } 30 defer ln.Close() 31 32 conf := NewConfig() 33 conf.Socket.Network = ln.Addr().Network() 34 conf.Socket.Address = ln.Addr().String() 35 36 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 37 if err != nil { 38 t.Fatal(err) 39 } 40 41 defer func() { 42 rdr.CloseAsync() 43 if err := rdr.WaitForClose(time.Second); err != nil { 44 t.Error(err) 45 } 46 }() 47 48 conn, err := ln.Accept() 49 if err != nil { 50 t.Fatal(err) 51 } 52 53 wg := sync.WaitGroup{} 54 wg.Add(1) 55 go func() { 56 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 57 if _, cerr := conn.Write([]byte("foo\n")); cerr != nil { 58 t.Error(cerr) 59 } 60 if _, cerr := conn.Write([]byte("bar\n")); cerr != nil { 61 t.Error(cerr) 62 } 63 if _, cerr := conn.Write([]byte("baz\n")); cerr != nil { 64 t.Error(cerr) 65 } 66 wg.Done() 67 }() 68 69 readNextMsg := func() (types.Message, error) { 70 var msg types.Message 71 select { 72 case tran := <-rdr.TransactionChan(): 73 msg = tran.Payload.DeepCopy() 74 select { 75 case tran.ResponseChan <- response.NewAck(): 76 case <-time.After(time.Second): 77 return nil, errors.New("timed out") 78 } 79 case <-time.After(time.Second): 80 return nil, errors.New("timed out") 81 } 82 return msg, nil 83 } 84 85 exp := [][]byte{[]byte("foo")} 86 msg, err := readNextMsg() 87 if err != nil { 88 t.Fatal(err) 89 } 90 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 91 t.Errorf("Wrong message contents: %s != %s", act, exp) 92 } 93 94 exp = [][]byte{[]byte("bar")} 95 if msg, err = readNextMsg(); err != nil { 96 t.Fatal(err) 97 } 98 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 99 t.Errorf("Wrong message contents: %s != %s", act, exp) 100 } 101 102 exp = [][]byte{[]byte("baz")} 103 if msg, err = readNextMsg(); err != nil { 104 t.Fatal(err) 105 } 106 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 107 t.Errorf("Wrong message contents: %s != %s", act, exp) 108 } 109 110 wg.Wait() 111 conn.Close() 112 } 113 114 func TestSocketReconnect(t *testing.T) { 115 tmpDir := t.TempDir() 116 117 ln, err := net.Listen("unix", filepath.Join(tmpDir, "benthos.sock")) 118 if err != nil { 119 t.Fatalf("failed to listen on address: %v", err) 120 } 121 defer ln.Close() 122 123 conf := NewConfig() 124 conf.Socket.Network = ln.Addr().Network() 125 conf.Socket.Address = ln.Addr().String() 126 127 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 128 if err != nil { 129 t.Fatal(err) 130 } 131 132 defer func() { 133 rdr.CloseAsync() 134 if err := rdr.WaitForClose(time.Second); err != nil { 135 t.Error(err) 136 } 137 }() 138 139 conn, err := ln.Accept() 140 if err != nil { 141 t.Fatal(err) 142 } 143 144 wg := sync.WaitGroup{} 145 wg.Add(1) 146 go func() { 147 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 148 _, cerr := conn.Write([]byte("foo\n")) 149 if cerr != nil { 150 t.Error(cerr) 151 } 152 conn.Close() 153 conn, cerr = ln.Accept() 154 if cerr != nil { 155 t.Error(cerr) 156 } 157 if _, cerr := conn.Write([]byte("bar\n")); cerr != nil { 158 t.Error(cerr) 159 } 160 if _, cerr := conn.Write([]byte("baz\n")); cerr != nil { 161 t.Error(cerr) 162 } 163 wg.Done() 164 }() 165 166 readNextMsg := func() (types.Message, error) { 167 var msg types.Message 168 select { 169 case tran := <-rdr.TransactionChan(): 170 msg = tran.Payload.DeepCopy() 171 select { 172 case tran.ResponseChan <- response.NewAck(): 173 case <-time.After(time.Second): 174 return nil, errors.New("timed out") 175 } 176 case <-time.After(time.Second): 177 return nil, errors.New("timed out") 178 } 179 return msg, nil 180 } 181 182 exp := [][]byte{[]byte("foo")} 183 msg, err := readNextMsg() 184 if err != nil { 185 t.Fatal(err) 186 } 187 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 188 t.Errorf("Wrong message contents: %s != %s", act, exp) 189 } 190 191 exp = [][]byte{[]byte("bar")} 192 if msg, err = readNextMsg(); err != nil { 193 t.Fatal(err) 194 } 195 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 196 t.Errorf("Wrong message contents: %s != %s", act, exp) 197 } 198 199 exp = [][]byte{[]byte("baz")} 200 if msg, err = readNextMsg(); err != nil { 201 t.Fatal(err) 202 } 203 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 204 t.Errorf("Wrong message contents: %s != %s", act, exp) 205 } 206 207 wg.Wait() 208 conn.Close() 209 } 210 211 func TestSocketMultipart(t *testing.T) { 212 tmpDir := t.TempDir() 213 214 ln, err := net.Listen("unix", filepath.Join(tmpDir, "benthos.sock")) 215 if err != nil { 216 t.Fatalf("failed to listen on a port: %v", err) 217 } 218 defer ln.Close() 219 220 conf := NewConfig() 221 conf.Socket.Multipart = true 222 conf.Socket.Network = ln.Addr().Network() 223 conf.Socket.Address = ln.Addr().String() 224 225 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 226 if err != nil { 227 t.Fatal(err) 228 } 229 230 defer func() { 231 rdr.CloseAsync() 232 if err := rdr.WaitForClose(time.Second); err != nil { 233 t.Error(err) 234 } 235 }() 236 237 conn, err := ln.Accept() 238 if err != nil { 239 t.Fatal(err) 240 } 241 242 wg := sync.WaitGroup{} 243 wg.Add(1) 244 go func() { 245 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 246 if _, cerr := conn.Write([]byte("foo\n")); cerr != nil { 247 t.Error(cerr) 248 } 249 if _, cerr := conn.Write([]byte("bar\n")); cerr != nil { 250 t.Error(cerr) 251 } 252 if _, cerr := conn.Write([]byte("\n")); cerr != nil { 253 t.Error(cerr) 254 } 255 if _, cerr := conn.Write([]byte("baz\n\n")); cerr != nil { 256 t.Error(cerr) 257 } 258 wg.Done() 259 }() 260 261 readNextMsg := func() (types.Message, error) { 262 var msg types.Message 263 select { 264 case tran := <-rdr.TransactionChan(): 265 msg = tran.Payload.DeepCopy() 266 select { 267 case tran.ResponseChan <- response.NewAck(): 268 case <-time.After(time.Second): 269 return nil, errors.New("timed out") 270 } 271 case <-time.After(time.Second): 272 return nil, errors.New("timed out") 273 } 274 return msg, nil 275 } 276 277 exp := [][]byte{[]byte("foo"), []byte("bar")} 278 msg, err := readNextMsg() 279 if err != nil { 280 t.Fatal(err) 281 } 282 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 283 t.Errorf("Wrong message contents: %s != %s", act, exp) 284 } 285 286 exp = [][]byte{[]byte("baz")} 287 if msg, err = readNextMsg(); err != nil { 288 t.Fatal(err) 289 } 290 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 291 t.Errorf("Wrong message contents: %s != %s", act, exp) 292 } 293 294 wg.Wait() 295 conn.Close() 296 } 297 298 func TestSocketMultipartCustomDelim(t *testing.T) { 299 tmpDir := t.TempDir() 300 301 ln, err := net.Listen("unix", filepath.Join(tmpDir, "b.sock")) 302 if err != nil { 303 t.Fatalf("failed to listen on address: %v", err) 304 } 305 defer ln.Close() 306 307 conf := NewConfig() 308 conf.Socket.Multipart = true 309 conf.Socket.Delim = "@" 310 conf.Socket.Network = ln.Addr().Network() 311 conf.Socket.Address = ln.Addr().String() 312 313 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 314 if err != nil { 315 t.Fatal(err) 316 } 317 318 defer func() { 319 rdr.CloseAsync() 320 if err := rdr.WaitForClose(time.Second); err != nil { 321 t.Error(err) 322 } 323 }() 324 325 conn, err := ln.Accept() 326 if err != nil { 327 t.Fatal(err) 328 } 329 330 wg := sync.WaitGroup{} 331 wg.Add(1) 332 go func() { 333 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 334 if _, cerr := conn.Write([]byte("foo@")); cerr != nil { 335 t.Error(cerr) 336 } 337 if _, cerr := conn.Write([]byte("bar@")); cerr != nil { 338 t.Error(cerr) 339 } 340 if _, cerr := conn.Write([]byte("@")); cerr != nil { 341 t.Error(cerr) 342 } 343 if _, cerr := conn.Write([]byte("baz\n@@")); cerr != nil { 344 t.Error(cerr) 345 } 346 wg.Done() 347 }() 348 349 readNextMsg := func() (types.Message, error) { 350 var msg types.Message 351 select { 352 case tran := <-rdr.TransactionChan(): 353 msg = tran.Payload.DeepCopy() 354 select { 355 case tran.ResponseChan <- response.NewAck(): 356 case <-time.After(time.Second): 357 return nil, errors.New("timed out") 358 } 359 case <-time.After(time.Second): 360 return nil, errors.New("timed out") 361 } 362 return msg, nil 363 } 364 365 exp := [][]byte{[]byte("foo"), []byte("bar")} 366 msg, err := readNextMsg() 367 if err != nil { 368 t.Fatal(err) 369 } 370 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 371 t.Errorf("Wrong message contents: %s != %s", act, exp) 372 } 373 374 exp = [][]byte{[]byte("baz\n")} 375 if msg, err = readNextMsg(); err != nil { 376 t.Fatal(err) 377 } 378 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 379 t.Errorf("Wrong message contents: %s != %s", act, exp) 380 } 381 382 wg.Wait() 383 conn.Close() 384 } 385 386 func TestSocketMultipartShutdown(t *testing.T) { 387 tmpDir := t.TempDir() 388 389 ln, err := net.Listen("unix", filepath.Join(tmpDir, "benthos.sock")) 390 if err != nil { 391 t.Fatalf("failed to listen on address: %v", err) 392 } 393 defer ln.Close() 394 395 conf := NewConfig() 396 conf.Socket.Multipart = true 397 conf.Socket.Network = ln.Addr().Network() 398 conf.Socket.Address = ln.Addr().String() 399 400 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 401 if err != nil { 402 t.Fatal(err) 403 } 404 405 defer func() { 406 rdr.CloseAsync() 407 if err := rdr.WaitForClose(time.Second); err != nil { 408 t.Error(err) 409 } 410 }() 411 412 conn, err := ln.Accept() 413 if err != nil { 414 t.Fatal(err) 415 } 416 417 wg := sync.WaitGroup{} 418 wg.Add(1) 419 go func() { 420 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 421 if _, cerr := conn.Write([]byte("foo\n")); cerr != nil { 422 t.Error(cerr) 423 } 424 if _, cerr := conn.Write([]byte("bar\n")); cerr != nil { 425 t.Error(cerr) 426 } 427 if _, cerr := conn.Write([]byte("\n")); cerr != nil { 428 t.Error(cerr) 429 } 430 if _, cerr := conn.Write([]byte("baz\n")); cerr != nil { 431 t.Error(cerr) 432 } 433 conn.Close() 434 wg.Done() 435 }() 436 437 readNextMsg := func() (types.Message, error) { 438 var msg types.Message 439 select { 440 case tran := <-rdr.TransactionChan(): 441 msg = tran.Payload.DeepCopy() 442 select { 443 case tran.ResponseChan <- response.NewAck(): 444 case <-time.After(time.Second): 445 return nil, errors.New("timed out on ack") 446 } 447 case <-time.After(time.Second): 448 return nil, errors.New("timed out on read") 449 } 450 return msg, nil 451 } 452 453 exp := [][]byte{[]byte("foo"), []byte("bar")} 454 msg, err := readNextMsg() 455 if err != nil { 456 t.Fatal(err) 457 } 458 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 459 t.Errorf("Wrong message contents: %s != %s", act, exp) 460 } 461 462 exp = [][]byte{[]byte("baz")} 463 if msg, err = readNextMsg(); err != nil { 464 t.Fatal(err) 465 } 466 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 467 t.Errorf("Wrong message contents: %s != %s", act, exp) 468 } 469 470 wg.Wait() 471 } 472 473 func TestTCPSocketBasic(t *testing.T) { 474 ln, err := net.Listen("tcp", "127.0.0.1:0") 475 if err != nil { 476 if ln, err = net.Listen("tcp6", "[::1]:0"); err != nil { 477 t.Fatalf("failed to listen on a port: %v", err) 478 } 479 } 480 defer ln.Close() 481 482 conf := NewConfig() 483 conf.Socket.Network = "tcp" 484 conf.Socket.Address = ln.Addr().String() 485 486 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 487 if err != nil { 488 t.Fatal(err) 489 } 490 491 defer func() { 492 rdr.CloseAsync() 493 if err := rdr.WaitForClose(time.Second); err != nil { 494 t.Error(err) 495 } 496 }() 497 498 conn, err := ln.Accept() 499 if err != nil { 500 t.Fatal(err) 501 } 502 503 wg := sync.WaitGroup{} 504 wg.Add(1) 505 go func() { 506 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 507 if _, cerr := conn.Write([]byte("foo\n")); cerr != nil { 508 t.Error(cerr) 509 } 510 if _, cerr := conn.Write([]byte("bar\n")); cerr != nil { 511 t.Error(cerr) 512 } 513 if _, cerr := conn.Write([]byte("baz\n")); cerr != nil { 514 t.Error(cerr) 515 } 516 wg.Done() 517 }() 518 519 readNextMsg := func() (types.Message, error) { 520 var msg types.Message 521 select { 522 case tran := <-rdr.TransactionChan(): 523 msg = tran.Payload.DeepCopy() 524 select { 525 case tran.ResponseChan <- response.NewAck(): 526 case <-time.After(time.Second): 527 return nil, errors.New("timed out") 528 } 529 case <-time.After(time.Second): 530 return nil, errors.New("timed out") 531 } 532 return msg, nil 533 } 534 535 exp := [][]byte{[]byte("foo")} 536 msg, err := readNextMsg() 537 if err != nil { 538 t.Fatal(err) 539 } 540 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 541 t.Errorf("Wrong message contents: %s != %s", act, exp) 542 } 543 544 exp = [][]byte{[]byte("bar")} 545 if msg, err = readNextMsg(); err != nil { 546 t.Fatal(err) 547 } 548 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 549 t.Errorf("Wrong message contents: %s != %s", act, exp) 550 } 551 552 exp = [][]byte{[]byte("baz")} 553 if msg, err = readNextMsg(); err != nil { 554 t.Fatal(err) 555 } 556 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 557 t.Errorf("Wrong message contents: %s != %s", act, exp) 558 } 559 560 wg.Wait() 561 conn.Close() 562 } 563 564 func TestTCPSocketReconnect(t *testing.T) { 565 ln, err := net.Listen("tcp", "127.0.0.1:0") 566 if err != nil { 567 if ln, err = net.Listen("tcp6", "[::1]:0"); err != nil { 568 t.Fatalf("failed to listen on a port: %v", err) 569 } 570 } 571 defer ln.Close() 572 573 conf := NewConfig() 574 conf.Socket.Network = "tcp" 575 conf.Socket.Address = ln.Addr().String() 576 577 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 578 if err != nil { 579 t.Fatal(err) 580 } 581 582 defer func() { 583 rdr.CloseAsync() 584 if err := rdr.WaitForClose(time.Second); err != nil { 585 t.Error(err) 586 } 587 }() 588 589 conn, err := ln.Accept() 590 if err != nil { 591 t.Fatal(err) 592 } 593 594 wg := sync.WaitGroup{} 595 wg.Add(1) 596 go func() { 597 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 598 _, cerr := conn.Write([]byte("foo\n")) 599 if cerr != nil { 600 t.Error(cerr) 601 } 602 conn.Close() 603 conn, cerr = ln.Accept() 604 if cerr != nil { 605 t.Error(cerr) 606 } 607 if _, cerr := conn.Write([]byte("bar\n")); cerr != nil { 608 t.Error(cerr) 609 } 610 if _, cerr := conn.Write([]byte("baz\n")); cerr != nil { 611 t.Error(cerr) 612 } 613 wg.Done() 614 }() 615 616 readNextMsg := func() (types.Message, error) { 617 var msg types.Message 618 select { 619 case tran := <-rdr.TransactionChan(): 620 msg = tran.Payload.DeepCopy() 621 select { 622 case tran.ResponseChan <- response.NewAck(): 623 case <-time.After(time.Second): 624 return nil, errors.New("timed out") 625 } 626 case <-time.After(time.Second): 627 return nil, errors.New("timed out") 628 } 629 return msg, nil 630 } 631 632 exp := [][]byte{[]byte("foo")} 633 msg, err := readNextMsg() 634 if err != nil { 635 t.Fatal(err) 636 } 637 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 638 t.Errorf("Wrong message contents: %s != %s", act, exp) 639 } 640 641 exp = [][]byte{[]byte("bar")} 642 if msg, err = readNextMsg(); err != nil { 643 t.Fatal(err) 644 } 645 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 646 t.Errorf("Wrong message contents: %s != %s", act, exp) 647 } 648 649 exp = [][]byte{[]byte("baz")} 650 if msg, err = readNextMsg(); err != nil { 651 t.Fatal(err) 652 } 653 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 654 t.Errorf("Wrong message contents: %s != %s", act, exp) 655 } 656 657 wg.Wait() 658 conn.Close() 659 } 660 661 func TestTCPSocketMultipart(t *testing.T) { 662 ln, err := net.Listen("tcp", "127.0.0.1:0") 663 if err != nil { 664 if ln, err = net.Listen("tcp6", "[::1]:0"); err != nil { 665 t.Fatalf("failed to listen on a port: %v", err) 666 } 667 } 668 defer ln.Close() 669 670 conf := NewConfig() 671 conf.Socket.Network = "tcp" 672 conf.Socket.Multipart = true 673 conf.Socket.Address = ln.Addr().String() 674 675 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 676 if err != nil { 677 t.Fatal(err) 678 } 679 680 defer func() { 681 rdr.CloseAsync() 682 if err := rdr.WaitForClose(time.Second); err != nil { 683 t.Error(err) 684 } 685 }() 686 687 conn, err := ln.Accept() 688 if err != nil { 689 t.Fatal(err) 690 } 691 692 wg := sync.WaitGroup{} 693 wg.Add(1) 694 go func() { 695 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 696 if _, cerr := conn.Write([]byte("foo\n")); cerr != nil { 697 t.Error(cerr) 698 } 699 if _, cerr := conn.Write([]byte("bar\n")); cerr != nil { 700 t.Error(cerr) 701 } 702 if _, cerr := conn.Write([]byte("\n")); cerr != nil { 703 t.Error(cerr) 704 } 705 if _, cerr := conn.Write([]byte("baz\n\n")); cerr != nil { 706 t.Error(cerr) 707 } 708 wg.Done() 709 }() 710 711 readNextMsg := func() (types.Message, error) { 712 var msg types.Message 713 select { 714 case tran := <-rdr.TransactionChan(): 715 msg = tran.Payload.DeepCopy() 716 select { 717 case tran.ResponseChan <- response.NewAck(): 718 case <-time.After(time.Second): 719 return nil, errors.New("timed out") 720 } 721 case <-time.After(time.Second): 722 return nil, errors.New("timed out") 723 } 724 return msg, nil 725 } 726 727 exp := [][]byte{[]byte("foo"), []byte("bar")} 728 msg, err := readNextMsg() 729 if err != nil { 730 t.Fatal(err) 731 } 732 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 733 t.Errorf("Wrong message contents: %s != %s", act, exp) 734 } 735 736 exp = [][]byte{[]byte("baz")} 737 if msg, err = readNextMsg(); err != nil { 738 t.Fatal(err) 739 } 740 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 741 t.Errorf("Wrong message contents: %s != %s", act, exp) 742 } 743 744 wg.Wait() 745 conn.Close() 746 } 747 748 func TestTCPSocketMultipartCustomDelim(t *testing.T) { 749 ln, err := net.Listen("tcp", "127.0.0.1:0") 750 if err != nil { 751 if ln, err = net.Listen("tcp6", "[::1]:0"); err != nil { 752 t.Fatalf("failed to listen on a port: %v", err) 753 } 754 } 755 defer ln.Close() 756 757 conf := NewConfig() 758 conf.Socket.Network = "tcp" 759 conf.Socket.Multipart = true 760 conf.Socket.Delim = "@" 761 conf.Socket.Address = ln.Addr().String() 762 763 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 764 if err != nil { 765 t.Fatal(err) 766 } 767 768 defer func() { 769 rdr.CloseAsync() 770 if err := rdr.WaitForClose(time.Second); err != nil { 771 t.Error(err) 772 } 773 }() 774 775 conn, err := ln.Accept() 776 if err != nil { 777 t.Fatal(err) 778 } 779 780 wg := sync.WaitGroup{} 781 wg.Add(1) 782 go func() { 783 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 784 if _, cerr := conn.Write([]byte("foo@")); cerr != nil { 785 t.Error(cerr) 786 } 787 if _, cerr := conn.Write([]byte("bar@")); cerr != nil { 788 t.Error(cerr) 789 } 790 if _, cerr := conn.Write([]byte("@")); cerr != nil { 791 t.Error(cerr) 792 } 793 if _, cerr := conn.Write([]byte("baz\n@@")); cerr != nil { 794 t.Error(cerr) 795 } 796 wg.Done() 797 }() 798 799 readNextMsg := func() (types.Message, error) { 800 var msg types.Message 801 select { 802 case tran := <-rdr.TransactionChan(): 803 msg = tran.Payload.DeepCopy() 804 select { 805 case tran.ResponseChan <- response.NewAck(): 806 case <-time.After(time.Second): 807 return nil, errors.New("timed out") 808 } 809 case <-time.After(time.Second): 810 return nil, errors.New("timed out") 811 } 812 return msg, nil 813 } 814 815 exp := [][]byte{[]byte("foo"), []byte("bar")} 816 msg, err := readNextMsg() 817 if err != nil { 818 t.Fatal(err) 819 } 820 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 821 t.Errorf("Wrong message contents: %s != %s", act, exp) 822 } 823 824 exp = [][]byte{[]byte("baz\n")} 825 if msg, err = readNextMsg(); err != nil { 826 t.Fatal(err) 827 } 828 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 829 t.Errorf("Wrong message contents: %s != %s", act, exp) 830 } 831 832 wg.Wait() 833 conn.Close() 834 } 835 836 func TestTCPSocketMultipartShutdown(t *testing.T) { 837 ln, err := net.Listen("tcp", "127.0.0.1:0") 838 if err != nil { 839 if ln, err = net.Listen("tcp6", "[::1]:0"); err != nil { 840 t.Fatalf("failed to listen on a port: %v", err) 841 } 842 } 843 defer ln.Close() 844 845 conf := NewConfig() 846 conf.Socket.Network = "tcp" 847 conf.Socket.Multipart = true 848 conf.Socket.Address = ln.Addr().String() 849 850 rdr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop()) 851 if err != nil { 852 t.Fatal(err) 853 } 854 855 defer func() { 856 rdr.CloseAsync() 857 if err := rdr.WaitForClose(time.Second); err != nil { 858 t.Error(err) 859 } 860 }() 861 862 conn, err := ln.Accept() 863 if err != nil { 864 t.Fatal(err) 865 } 866 867 wg := sync.WaitGroup{} 868 wg.Add(1) 869 go func() { 870 conn.SetWriteDeadline(time.Now().Add(time.Second * 5)) 871 if _, cerr := conn.Write([]byte("foo\n")); cerr != nil { 872 t.Error(cerr) 873 } 874 if _, cerr := conn.Write([]byte("bar\n")); cerr != nil { 875 t.Error(cerr) 876 } 877 if _, cerr := conn.Write([]byte("\n")); cerr != nil { 878 t.Error(cerr) 879 } 880 if _, cerr := conn.Write([]byte("baz\n")); cerr != nil { 881 t.Error(cerr) 882 } 883 conn.Close() 884 wg.Done() 885 }() 886 887 readNextMsg := func() (types.Message, error) { 888 var msg types.Message 889 select { 890 case tran := <-rdr.TransactionChan(): 891 msg = tran.Payload.DeepCopy() 892 select { 893 case tran.ResponseChan <- response.NewAck(): 894 case <-time.After(time.Second): 895 return nil, errors.New("timed out on ack") 896 } 897 case <-time.After(time.Second): 898 return nil, errors.New("timed out on read") 899 } 900 return msg, nil 901 } 902 903 exp := [][]byte{[]byte("foo"), []byte("bar")} 904 msg, err := readNextMsg() 905 if err != nil { 906 t.Fatal(err) 907 } 908 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 909 t.Errorf("Wrong message contents: %s != %s", act, exp) 910 } 911 912 exp = [][]byte{[]byte("baz")} 913 if msg, err = readNextMsg(); err != nil { 914 t.Fatal(err) 915 } 916 if act := message.GetAllBytes(msg); !reflect.DeepEqual(exp, act) { 917 t.Errorf("Wrong message contents: %s != %s", act, exp) 918 } 919 920 wg.Wait() 921 } 922 923 func BenchmarkTCPSocketWithCutOff(b *testing.B) { 924 ln, err := net.Listen("tcp", "127.0.0.1:0") 925 if err != nil { 926 ln, err = net.Listen("tcp6", "[::1]:0") 927 require.NoError(b, err) 928 } 929 b.Cleanup(func() { 930 ln.Close() 931 }) 932 933 conf := NewConfig() 934 conf.Socket.Network = "tcp" 935 conf.Socket.Address = ln.Addr().String() 936 937 sRdr, err := newSocketClient(conf.Socket, log.Noop()) 938 require.NoError(b, err) 939 940 rdr, err := NewAsyncReader(TypeSocket, true, reader.NewAsyncCutOff(reader.NewAsyncPreserver(sRdr)), log.Noop(), metrics.Noop()) 941 require.NoError(b, err) 942 943 defer func() { 944 rdr.CloseAsync() 945 assert.NoError(b, rdr.WaitForClose(time.Second)) 946 }() 947 948 conn, err := ln.Accept() 949 require.NoError(b, err) 950 951 wg := sync.WaitGroup{} 952 wg.Add(1) 953 go func() { 954 conn.SetWriteDeadline(time.Now().Add(time.Second * 60)) 955 for i := 0; i < b.N; i++ { 956 _, cerr := fmt.Fprintf(conn, "hello world this is message %v\n", i) 957 assert.NoError(b, cerr) 958 } 959 wg.Done() 960 }() 961 962 readNextMsg := func() (string, error) { 963 var payload string 964 select { 965 case tran := <-rdr.TransactionChan(): 966 payload = string(tran.Payload.Get(0).Get()) 967 go func() { 968 select { 969 case tran.ResponseChan <- response.NewAck(): 970 case <-time.After(time.Second): 971 } 972 }() 973 case <-time.After(time.Second): 974 return "", errors.New("timed out") 975 } 976 return payload, nil 977 } 978 979 b.ReportAllocs() 980 b.ResetTimer() 981 982 for i := 0; i < b.N; i++ { 983 exp := fmt.Sprintf("hello world this is message %v", i) 984 act, err := readNextMsg() 985 assert.NoError(b, err) 986 assert.Equal(b, exp, act) 987 } 988 989 wg.Wait() 990 conn.Close() 991 } 992 993 func BenchmarkTCPSocketNoCutOff(b *testing.B) { 994 ln, err := net.Listen("tcp", "127.0.0.1:0") 995 if err != nil { 996 ln, err = net.Listen("tcp6", "[::1]:0") 997 require.NoError(b, err) 998 } 999 b.Cleanup(func() { 1000 ln.Close() 1001 }) 1002 1003 conf := NewConfig() 1004 conf.Socket.Network = "tcp" 1005 conf.Socket.Address = ln.Addr().String() 1006 1007 sRdr, err := newSocketClient(conf.Socket, log.Noop()) 1008 require.NoError(b, err) 1009 1010 rdr, err := NewAsyncReader(TypeSocket, true, reader.NewAsyncPreserver(sRdr), log.Noop(), metrics.Noop()) 1011 require.NoError(b, err) 1012 1013 defer func() { 1014 rdr.CloseAsync() 1015 assert.NoError(b, rdr.WaitForClose(time.Second)) 1016 }() 1017 1018 conn, err := ln.Accept() 1019 require.NoError(b, err) 1020 1021 wg := sync.WaitGroup{} 1022 wg.Add(1) 1023 go func() { 1024 conn.SetWriteDeadline(time.Now().Add(time.Second * 60)) 1025 for i := 0; i < b.N; i++ { 1026 _, cerr := fmt.Fprintf(conn, "hello world this is message %v\n", i) 1027 assert.NoError(b, cerr) 1028 } 1029 wg.Done() 1030 }() 1031 1032 readNextMsg := func() (string, error) { 1033 var payload string 1034 select { 1035 case tran := <-rdr.TransactionChan(): 1036 payload = string(tran.Payload.Get(0).Get()) 1037 go func() { 1038 select { 1039 case tran.ResponseChan <- response.NewAck(): 1040 case <-time.After(time.Second): 1041 } 1042 }() 1043 case <-time.After(time.Second): 1044 return "", errors.New("timed out") 1045 } 1046 return payload, nil 1047 } 1048 1049 b.ReportAllocs() 1050 b.ResetTimer() 1051 1052 for i := 0; i < b.N; i++ { 1053 exp := fmt.Sprintf("hello world this is message %v", i) 1054 act, err := readNextMsg() 1055 assert.NoError(b, err) 1056 assert.Equal(b, exp, act) 1057 } 1058 1059 wg.Wait() 1060 conn.Close() 1061 }