p2c_test.go 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  1. package p2c
  2. import (
  3. "context"
  4. "fmt"
  5. "runtime"
  6. "strconv"
  7. "sync"
  8. "testing"
  9. "github.com/stretchr/testify/assert"
  10. "github.com/tal-tech/go-zero/core/logx"
  11. "github.com/tal-tech/go-zero/core/mathx"
  12. "google.golang.org/grpc/balancer"
  13. "google.golang.org/grpc/codes"
  14. "google.golang.org/grpc/resolver"
  15. "google.golang.org/grpc/status"
  16. )
  17. func init() {
  18. logx.Disable()
  19. }
  20. func TestP2cPicker_PickNil(t *testing.T) {
  21. builder := new(p2cPickerBuilder)
  22. picker := builder.Build(nil)
  23. _, _, err := picker.Pick(context.Background(), balancer.PickInfo{
  24. FullMethodName: "/",
  25. Ctx: context.Background(),
  26. })
  27. assert.NotNil(t, err)
  28. }
  29. func TestP2cPicker_Pick(t *testing.T) {
  30. tests := []struct {
  31. name string
  32. candidates int
  33. threshold float64
  34. }{
  35. {
  36. name: "single",
  37. candidates: 1,
  38. threshold: 0.9,
  39. },
  40. {
  41. name: "two",
  42. candidates: 2,
  43. threshold: 0.5,
  44. },
  45. {
  46. name: "multiple",
  47. candidates: 100,
  48. threshold: 0.95,
  49. },
  50. }
  51. for _, test := range tests {
  52. test := test
  53. t.Run(test.name, func(t *testing.T) {
  54. t.Parallel()
  55. const total = 10000
  56. builder := new(p2cPickerBuilder)
  57. ready := make(map[resolver.Address]balancer.SubConn)
  58. for i := 0; i < test.candidates; i++ {
  59. ready[resolver.Address{
  60. Addr: strconv.Itoa(i),
  61. }] = new(mockClientConn)
  62. }
  63. picker := builder.Build(ready)
  64. var wg sync.WaitGroup
  65. wg.Add(total)
  66. for i := 0; i < total; i++ {
  67. _, done, err := picker.Pick(context.Background(), balancer.PickInfo{
  68. FullMethodName: "/",
  69. Ctx: context.Background(),
  70. })
  71. assert.Nil(t, err)
  72. if i%100 == 0 {
  73. err = status.Error(codes.DeadlineExceeded, "deadline")
  74. }
  75. go func() {
  76. runtime.Gosched()
  77. done(balancer.DoneInfo{
  78. Err: err,
  79. })
  80. wg.Done()
  81. }()
  82. }
  83. wg.Wait()
  84. dist := make(map[interface{}]int)
  85. conns := picker.(*p2cPicker).conns
  86. for _, conn := range conns {
  87. dist[conn.addr.Addr] = int(conn.requests)
  88. }
  89. entropy := mathx.CalcEntropy(dist)
  90. assert.True(t, entropy > test.threshold, fmt.Sprintf("entropy is %f, less than %f",
  91. entropy, test.threshold))
  92. })
  93. }
  94. }
  95. type mockClientConn struct {
  96. }
  97. func (m mockClientConn) UpdateAddresses(addresses []resolver.Address) {
  98. }
  99. func (m mockClientConn) Connect() {
  100. }