github.com/lingyao2333/mo-zero@v1.4.1/zrpc/internal/balancer/p2c/p2c_test.go (about)

     1  package p2c
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"runtime"
     7  	"strconv"
     8  	"sync"
     9  	"testing"
    10  
    11  	"github.com/lingyao2333/mo-zero/core/logx"
    12  	"github.com/lingyao2333/mo-zero/core/mathx"
    13  	"github.com/lingyao2333/mo-zero/core/stringx"
    14  	"github.com/stretchr/testify/assert"
    15  	"google.golang.org/grpc/balancer"
    16  	"google.golang.org/grpc/balancer/base"
    17  	"google.golang.org/grpc/codes"
    18  	"google.golang.org/grpc/resolver"
    19  	"google.golang.org/grpc/status"
    20  )
    21  
    22  func init() {
    23  	logx.Disable()
    24  }
    25  
    26  func TestP2cPicker_PickNil(t *testing.T) {
    27  	builder := new(p2cPickerBuilder)
    28  	picker := builder.Build(base.PickerBuildInfo{})
    29  	_, err := picker.Pick(balancer.PickInfo{
    30  		FullMethodName: "/",
    31  		Ctx:            context.Background(),
    32  	})
    33  	assert.NotNil(t, err)
    34  }
    35  
    36  func TestP2cPicker_Pick(t *testing.T) {
    37  	tests := []struct {
    38  		name       string
    39  		candidates int
    40  		err        error
    41  		threshold  float64
    42  	}{
    43  		{
    44  			name:       "empty",
    45  			candidates: 0,
    46  			err:        balancer.ErrNoSubConnAvailable,
    47  		},
    48  		{
    49  			name:       "single",
    50  			candidates: 1,
    51  			threshold:  0.9,
    52  		},
    53  		{
    54  			name:       "two",
    55  			candidates: 2,
    56  			threshold:  0.5,
    57  		},
    58  		{
    59  			name:       "multiple",
    60  			candidates: 100,
    61  			threshold:  0.95,
    62  		},
    63  	}
    64  
    65  	for _, test := range tests {
    66  		test := test
    67  		t.Run(test.name, func(t *testing.T) {
    68  			t.Parallel()
    69  
    70  			const total = 10000
    71  			builder := new(p2cPickerBuilder)
    72  			ready := make(map[balancer.SubConn]base.SubConnInfo)
    73  			for i := 0; i < test.candidates; i++ {
    74  				ready[mockClientConn{
    75  					id: stringx.Rand(),
    76  				}] = base.SubConnInfo{
    77  					Address: resolver.Address{
    78  						Addr: strconv.Itoa(i),
    79  					},
    80  				}
    81  			}
    82  
    83  			picker := builder.Build(base.PickerBuildInfo{
    84  				ReadySCs: ready,
    85  			})
    86  			var wg sync.WaitGroup
    87  			wg.Add(total)
    88  			for i := 0; i < total; i++ {
    89  				result, err := picker.Pick(balancer.PickInfo{
    90  					FullMethodName: "/",
    91  					Ctx:            context.Background(),
    92  				})
    93  				assert.Equal(t, test.err, err)
    94  
    95  				if test.err != nil {
    96  					return
    97  				}
    98  
    99  				if i%100 == 0 {
   100  					err = status.Error(codes.DeadlineExceeded, "deadline")
   101  				}
   102  
   103  				go func() {
   104  					runtime.Gosched()
   105  					result.Done(balancer.DoneInfo{
   106  						Err: err,
   107  					})
   108  					wg.Done()
   109  				}()
   110  			}
   111  
   112  			wg.Wait()
   113  			dist := make(map[interface{}]int)
   114  			conns := picker.(*p2cPicker).conns
   115  			for _, conn := range conns {
   116  				dist[conn.addr.Addr] = int(conn.requests)
   117  			}
   118  
   119  			entropy := mathx.CalcEntropy(dist)
   120  			assert.True(t, entropy > test.threshold, fmt.Sprintf("entropy is %f, less than %f",
   121  				entropy, test.threshold))
   122  		})
   123  	}
   124  }
   125  
   126  type mockClientConn struct {
   127  	// add random string member to avoid map key equality.
   128  	id string
   129  }
   130  
   131  func (m mockClientConn) UpdateAddresses(addresses []resolver.Address) {
   132  }
   133  
   134  func (m mockClientConn) Connect() {
   135  }