github.com/cloudwego/kitex@v0.9.0/pkg/loadbalance/consist_test.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package loadbalance
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"strconv"
    23  	"testing"
    24  
    25  	"github.com/bytedance/gopkg/lang/fastrand"
    26  
    27  	"github.com/cloudwego/kitex/internal"
    28  	"github.com/cloudwego/kitex/internal/test"
    29  	"github.com/cloudwego/kitex/pkg/discovery"
    30  )
    31  
    32  type testCtxKey int
    33  
    34  const (
    35  	keyCtxKey testCtxKey = iota
    36  )
    37  
    38  func getKey(ctx context.Context, request interface{}) string {
    39  	if val, ok := ctx.Value(keyCtxKey).(string); ok {
    40  		return val
    41  	}
    42  	return "1234"
    43  }
    44  
    45  func newTestConsistentHashOption() ConsistentHashOption {
    46  	opt := NewConsistentHashOption(getKey)
    47  	opt.ExpireDuration = 0
    48  	return opt
    49  }
    50  
    51  func TestNewConsistBalancer(t *testing.T) {
    52  	opt := ConsistentHashOption{}
    53  	test.Panic(t, func() { NewConsistBalancer(opt) })
    54  }
    55  
    56  func TestConsistPicker_Next_Nil(t *testing.T) {
    57  	opt := newTestConsistentHashOption()
    58  	opt.GetKey = func(ctx context.Context, request interface{}) string {
    59  		return ""
    60  	}
    61  
    62  	insList := []discovery.Instance{
    63  		discovery.NewInstance("tcp", "addr1", 10, nil),
    64  	}
    65  	e := discovery.Result{
    66  		Cacheable: false,
    67  		CacheKey:  "",
    68  		Instances: insList,
    69  	}
    70  
    71  	cb := NewConsistBalancer(opt)
    72  	picker := cb.GetPicker(e)
    73  	test.Assert(t, picker.Next(context.TODO(), nil) == nil)
    74  
    75  	e = discovery.Result{
    76  		Cacheable: false,
    77  		CacheKey:  "",
    78  		Instances: nil,
    79  	}
    80  
    81  	cb = NewConsistBalancer(newTestConsistentHashOption())
    82  	picker = cb.GetPicker(e)
    83  	test.Assert(t, picker.Next(context.TODO(), nil) == nil)
    84  	test.Assert(t, cb.Name() == "consist")
    85  }
    86  
    87  func TestConsistPicker_Replica(t *testing.T) {
    88  	opt := NewConsistentHashOption(getKey)
    89  	opt.Replica = 1
    90  	opt.ExpireDuration = 0
    91  	opt.GetKey = func(ctx context.Context, request interface{}) string {
    92  		return "1234"
    93  	}
    94  	insList := makeNInstances(2, 10)
    95  	e := discovery.Result{
    96  		Cacheable: false,
    97  		CacheKey:  "",
    98  		Instances: insList,
    99  	}
   100  
   101  	cb := NewConsistBalancer(opt)
   102  	picker := cb.GetPicker(e)
   103  	first := picker.Next(context.TODO(), nil)
   104  	second := picker.Next(context.TODO(), nil)
   105  	test.Assert(t, first != second)
   106  }
   107  
   108  func TestConsistPicker_Next_NoCache(t *testing.T) {
   109  	opt := newTestConsistentHashOption()
   110  	ins := discovery.NewInstance("tcp", "addr1", 10, nil)
   111  	insList := []discovery.Instance{
   112  		ins,
   113  	}
   114  	e := discovery.Result{
   115  		Cacheable: false,
   116  		CacheKey:  "",
   117  		Instances: insList,
   118  	}
   119  
   120  	cb := NewConsistBalancer(opt)
   121  	picker := cb.GetPicker(e)
   122  	test.Assert(t, picker.Next(context.TODO(), nil) == ins)
   123  	test.Assert(t, picker.Next(context.TODO(), nil) == nil)
   124  }
   125  
   126  func TestConsistPicker_Next_NoCache_Consist(t *testing.T) {
   127  	opt := newTestConsistentHashOption()
   128  	insList := []discovery.Instance{
   129  		discovery.NewInstance("tcp", "addr1", 10, nil),
   130  		discovery.NewInstance("tcp", "addr2", 10, nil),
   131  		discovery.NewInstance("tcp", "addr3", 10, nil),
   132  		discovery.NewInstance("tcp", "addr4", 10, nil),
   133  		discovery.NewInstance("tcp", "addr5", 10, nil),
   134  	}
   135  	e := discovery.Result{
   136  		Cacheable: false,
   137  		CacheKey:  "",
   138  		Instances: insList,
   139  	}
   140  
   141  	cb := NewConsistBalancer(opt)
   142  	picker := cb.GetPicker(e)
   143  	ins := picker.Next(context.TODO(), nil)
   144  	for i := 0; i < 100; i++ {
   145  		picker := cb.GetPicker(e)
   146  		test.Assert(t, picker.Next(context.TODO(), nil) == ins)
   147  	}
   148  
   149  	cb = NewConsistBalancer(opt)
   150  	for i := 0; i < 100; i++ {
   151  		picker := cb.GetPicker(e)
   152  		test.Assert(t, picker.Next(context.TODO(), nil) == ins)
   153  	}
   154  }
   155  
   156  func TestConsistPicker_Next_Cache(t *testing.T) {
   157  	opt := newTestConsistentHashOption()
   158  	ins := discovery.NewInstance("tcp", "addr1", 10, nil)
   159  	insList := []discovery.Instance{
   160  		ins,
   161  	}
   162  	e := discovery.Result{
   163  		Cacheable: true,
   164  		CacheKey:  "4321",
   165  		Instances: insList,
   166  	}
   167  
   168  	cb := NewConsistBalancer(opt)
   169  	picker := cb.GetPicker(e)
   170  	test.Assert(t, picker.Next(context.TODO(), nil) == ins)
   171  }
   172  
   173  func TestConsistPicker_Next_Cache_Consist(t *testing.T) {
   174  	opt := newTestConsistentHashOption()
   175  	insList := []discovery.Instance{
   176  		discovery.NewInstance("tcp", "addr1", 10, nil),
   177  		discovery.NewInstance("tcp", "addr2", 10, nil),
   178  		discovery.NewInstance("tcp", "addr3", 10, nil),
   179  		discovery.NewInstance("tcp", "addr4", 10, nil),
   180  		discovery.NewInstance("tcp", "addr5", 10, nil),
   181  		discovery.NewInstance("tcp", "addr6", 10, nil),
   182  		discovery.NewInstance("tcp", "addr7", 10, nil),
   183  		discovery.NewInstance("tcp", "addr8", 10, nil),
   184  		discovery.NewInstance("tcp", "addr9", 10, nil),
   185  	}
   186  	e := discovery.Result{
   187  		Cacheable: true,
   188  		CacheKey:  "4321",
   189  		Instances: insList,
   190  	}
   191  
   192  	cb := NewConsistBalancer(opt)
   193  	picker := cb.GetPicker(e)
   194  	ins := picker.Next(context.TODO(), nil)
   195  	for i := 0; i < 100; i++ {
   196  		picker := cb.GetPicker(e)
   197  		test.Assert(t, picker.Next(context.TODO(), nil) == ins)
   198  	}
   199  
   200  	cb = NewConsistBalancer(opt)
   201  	for i := 0; i < 100; i++ {
   202  		picker := cb.GetPicker(e)
   203  		test.Assert(t, picker.Next(context.TODO(), nil) == ins)
   204  	}
   205  }
   206  
   207  func TestConsistBalance(t *testing.T) {
   208  	opt := ConsistentHashOption{
   209  		GetKey: func(ctx context.Context, request interface{}) string {
   210  			return strconv.Itoa(fastrand.Intn(100000))
   211  		},
   212  		Replica:        0,
   213  		VirtualFactor:  1000,
   214  		Weighted:       false,
   215  		ExpireDuration: 0,
   216  	}
   217  	inss := makeNInstances(10, 10)
   218  
   219  	m := make(map[discovery.Instance]int)
   220  	e := discovery.Result{
   221  		Cacheable: true,
   222  		CacheKey:  "4321",
   223  		Instances: inss,
   224  	}
   225  
   226  	cb := NewConsistBalancer(opt)
   227  	for i := 0; i < 100000; i++ {
   228  		picker := cb.GetPicker(e)
   229  		ins := picker.Next(context.TODO(), nil)
   230  		m[ins]++
   231  		if p, ok := picker.(internal.Reusable); ok {
   232  			p.Recycle()
   233  		}
   234  	}
   235  }
   236  
   237  func TestWeightedConsistBalance(t *testing.T) {
   238  	opt := ConsistentHashOption{
   239  		GetKey: func(ctx context.Context, request interface{}) string {
   240  			return strconv.Itoa(fastrand.Intn(100000))
   241  		},
   242  		Replica:        0,
   243  		VirtualFactor:  1000,
   244  		Weighted:       true,
   245  		ExpireDuration: 0,
   246  	}
   247  	inss := makeNInstances(10, 10)
   248  
   249  	m := make(map[discovery.Instance]int)
   250  	e := discovery.Result{
   251  		Cacheable: true,
   252  		CacheKey:  "4321",
   253  		Instances: inss,
   254  	}
   255  
   256  	cb := NewConsistBalancer(opt)
   257  	for i := 0; i < 100000; i++ {
   258  		picker := cb.GetPicker(e)
   259  		ins := picker.Next(context.TODO(), nil)
   260  		m[ins]++
   261  	}
   262  }
   263  
   264  func TestConsistPicker_Reblance(t *testing.T) {
   265  	opt := NewConsistentHashOption(getKey)
   266  	insList := makeNInstances(10, 10)
   267  	e := discovery.Result{
   268  		Cacheable: true,
   269  		CacheKey:  "4321",
   270  		Instances: insList[:5],
   271  	}
   272  
   273  	ctx := context.Background()
   274  	cb := NewConsistBalancer(opt)
   275  	record := make(map[string]discovery.Instance)
   276  	for i := 0; i < 10; i++ {
   277  		picker := cb.GetPicker(e)
   278  		key := strconv.Itoa(i)
   279  		ctx = context.WithValue(ctx, keyCtxKey, key)
   280  		record[key] = picker.Next(ctx, nil)
   281  	}
   282  	c := discovery.Change{
   283  		Result: e,
   284  		Added:  insList[5:],
   285  	}
   286  	cb.(Rebalancer).Rebalance(c)
   287  	for i := 0; i < 10; i++ {
   288  		picker := cb.GetPicker(e)
   289  		key := strconv.Itoa(i)
   290  		ctx = context.WithValue(ctx, keyCtxKey, key)
   291  		test.DeepEqual(t, record[key], picker.Next(ctx, nil))
   292  	}
   293  }
   294  
   295  func BenchmarkNewConsistPicker_NoCache(bb *testing.B) {
   296  	n := 10
   297  	balancer := NewConsistBalancer(newTestConsistentHashOption())
   298  	ctx := context.Background()
   299  
   300  	for i := 0; i < 4; i++ {
   301  		bb.Run(fmt.Sprintf("%dins", n), func(b *testing.B) {
   302  			inss := makeNInstances(n, 10)
   303  			e := discovery.Result{
   304  				Cacheable: false,
   305  				CacheKey:  "",
   306  				Instances: inss,
   307  			}
   308  			picker := balancer.GetPicker(e)
   309  			picker.Next(ctx, nil)
   310  			picker.(internal.Reusable).Recycle()
   311  			b.ReportAllocs()
   312  			b.ResetTimer()
   313  			for i := 0; i < b.N; i++ {
   314  				picker := balancer.GetPicker(e)
   315  				// picker.Next(ctx, nil)
   316  				if r, ok := picker.(internal.Reusable); ok {
   317  					r.Recycle()
   318  				}
   319  			}
   320  		})
   321  		n *= 10
   322  	}
   323  }
   324  
   325  func BenchmarkNewConsistPicker(bb *testing.B) {
   326  	n := 10
   327  	balancer := NewConsistBalancer(newTestConsistentHashOption())
   328  	ctx := context.Background()
   329  
   330  	for i := 0; i < 4; i++ {
   331  		bb.Run(fmt.Sprintf("%dins", n), func(b *testing.B) {
   332  			inss := makeNInstances(n, 10)
   333  			e := discovery.Result{
   334  				Cacheable: true,
   335  				CacheKey:  "test",
   336  				Instances: inss,
   337  			}
   338  			picker := balancer.GetPicker(e)
   339  			picker.Next(ctx, nil)
   340  			picker.(internal.Reusable).Recycle()
   341  			b.ReportAllocs()
   342  			b.ResetTimer()
   343  			for i := 0; i < b.N; i++ {
   344  				picker := balancer.GetPicker(e)
   345  				picker.Next(ctx, nil)
   346  				if r, ok := picker.(internal.Reusable); ok {
   347  					r.Recycle()
   348  				}
   349  			}
   350  		})
   351  		n *= 10
   352  	}
   353  }