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 }