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