1/* 2 * 3 * Copyright 2020 gRPC authors. 4 * 5 * Licensed under the Apache License, Version 2.0 (the "License"); 6 * you may not use this file except in compliance with the License. 7 * You may obtain a copy of the License at 8 * 9 * http://www.apache.org/licenses/LICENSE-2.0 10 * 11 * Unless required by applicable law or agreed to in writing, software 12 * distributed under the License is distributed on an "AS IS" BASIS, 13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 * See the License for the specific language governing permissions and 15 * limitations under the License. 16 * 17 */ 18 19package testutils 20 21import ( 22 "sync" 23 24 "google.golang.org/grpc/internal/wrr" 25) 26 27// testWRR is a deterministic WRR implementation. 28// 29// The real implementation does random WRR. testWRR makes the balancer behavior 30// deterministic and easier to test. 31// 32// With {a: 2, b: 3}, the Next() results will be {a, a, b, b, b}. 33type testWRR struct { 34 itemsWithWeight []struct { 35 item interface{} 36 weight int64 37 } 38 length int 39 40 mu sync.Mutex 41 idx int // The index of the item that will be picked 42 count int64 // The number of times the current item has been picked. 43} 44 45// NewTestWRR return a WRR for testing. It's deterministic instead of random. 46func NewTestWRR() wrr.WRR { 47 return &testWRR{} 48} 49 50func (twrr *testWRR) Add(item interface{}, weight int64) { 51 twrr.itemsWithWeight = append(twrr.itemsWithWeight, struct { 52 item interface{} 53 weight int64 54 }{item: item, weight: weight}) 55 twrr.length++ 56} 57 58func (twrr *testWRR) Next() interface{} { 59 twrr.mu.Lock() 60 iww := twrr.itemsWithWeight[twrr.idx] 61 twrr.count++ 62 if twrr.count >= iww.weight { 63 twrr.idx = (twrr.idx + 1) % twrr.length 64 twrr.count = 0 65 } 66 twrr.mu.Unlock() 67 return iww.item 68} 69