1/*
2 * Copyright 2019 gRPC 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
17package edsbalancer
18
19import (
20	"context"
21	"fmt"
22	"testing"
23
24	"google.golang.org/grpc"
25	"google.golang.org/grpc/balancer"
26	"google.golang.org/grpc/connectivity"
27	"google.golang.org/grpc/resolver"
28	"google.golang.org/grpc/xds/internal"
29
30	corepb "github.com/envoyproxy/go-control-plane/envoy/api/v2/core"
31)
32
33const testSubConnsCount = 16
34
35var testSubConns []*testSubConn
36
37func init() {
38	for i := 0; i < testSubConnsCount; i++ {
39		testSubConns = append(testSubConns, &testSubConn{
40			id: fmt.Sprintf("sc%d", i),
41		})
42	}
43}
44
45type testSubConn struct {
46	id string
47}
48
49func (tsc *testSubConn) UpdateAddresses([]resolver.Address) {
50	panic("not implemented")
51}
52
53func (tsc *testSubConn) Connect() {
54}
55
56// Implement stringer to get human friendly error message.
57func (tsc *testSubConn) String() string {
58	return tsc.id
59}
60
61type testClientConn struct {
62	t *testing.T // For logging only.
63
64	newSubConnAddrsCh chan []resolver.Address // The last 10 []Address to create subconn.
65	newSubConnCh      chan balancer.SubConn   // The last 10 subconn created.
66	removeSubConnCh   chan balancer.SubConn   // The last 10 subconn removed.
67
68	newPickerCh chan balancer.V2Picker  // The last picker updated.
69	newStateCh  chan connectivity.State // The last state.
70
71	subConnIdx int
72}
73
74func newTestClientConn(t *testing.T) *testClientConn {
75	return &testClientConn{
76		t: t,
77
78		newSubConnAddrsCh: make(chan []resolver.Address, 10),
79		newSubConnCh:      make(chan balancer.SubConn, 10),
80		removeSubConnCh:   make(chan balancer.SubConn, 10),
81
82		newPickerCh: make(chan balancer.V2Picker, 1),
83		newStateCh:  make(chan connectivity.State, 1),
84	}
85}
86
87func (tcc *testClientConn) NewSubConn(a []resolver.Address, o balancer.NewSubConnOptions) (balancer.SubConn, error) {
88	sc := testSubConns[tcc.subConnIdx]
89	tcc.subConnIdx++
90
91	tcc.t.Logf("testClientConn: NewSubConn(%v, %+v) => %s", a, o, sc)
92	select {
93	case tcc.newSubConnAddrsCh <- a:
94	default:
95	}
96
97	select {
98	case tcc.newSubConnCh <- sc:
99	default:
100	}
101
102	return sc, nil
103}
104
105func (tcc *testClientConn) RemoveSubConn(sc balancer.SubConn) {
106	tcc.t.Logf("testClientCOnn: RemoveSubConn(%p)", sc)
107	select {
108	case tcc.removeSubConnCh <- sc:
109	default:
110	}
111}
112
113func (tcc *testClientConn) UpdateBalancerState(s connectivity.State, p balancer.Picker) {
114	tcc.t.Fatal("not implemented")
115}
116
117func (tcc *testClientConn) UpdateState(bs balancer.State) {
118	tcc.t.Logf("testClientConn: UpdateState(%v)", bs)
119	select {
120	case <-tcc.newStateCh:
121	default:
122	}
123	tcc.newStateCh <- bs.ConnectivityState
124
125	select {
126	case <-tcc.newPickerCh:
127	default:
128	}
129	tcc.newPickerCh <- bs.Picker
130}
131
132func (tcc *testClientConn) ResolveNow(resolver.ResolveNowOptions) {
133	panic("not implemented")
134}
135
136func (tcc *testClientConn) Target() string {
137	panic("not implemented")
138}
139
140type testServerLoad struct {
141	name string
142	d    float64
143}
144
145type testLoadStore struct {
146	callsStarted []internal.Locality
147	callsEnded   []internal.Locality
148	callsCost    []testServerLoad
149}
150
151func newTestLoadStore() *testLoadStore {
152	return &testLoadStore{}
153}
154
155func (*testLoadStore) CallDropped(category string) {
156	panic("not implemented")
157}
158
159func (tls *testLoadStore) CallStarted(l internal.Locality) {
160	tls.callsStarted = append(tls.callsStarted, l)
161}
162
163func (tls *testLoadStore) CallFinished(l internal.Locality, err error) {
164	tls.callsEnded = append(tls.callsEnded, l)
165}
166
167func (tls *testLoadStore) CallServerLoad(l internal.Locality, name string, d float64) {
168	tls.callsCost = append(tls.callsCost, testServerLoad{name: name, d: d})
169}
170
171func (*testLoadStore) ReportTo(ctx context.Context, cc *grpc.ClientConn, clusterName string, node *corepb.Node) {
172	panic("not implemented")
173}
174
175// isRoundRobin checks whether f's return value is roundrobin of elements from
176// want. But it doesn't check for the order. Note that want can contain
177// duplicate items, which makes it weight-round-robin.
178//
179// Step 1. the return values of f should form a permutation of all elements in
180// want, but not necessary in the same order. E.g. if want is {a,a,b}, the check
181// fails if f returns:
182//  - {a,a,a}: third a is returned before b
183//  - {a,b,b}: second b is returned before the second a
184//
185// If error is found in this step, the returned error contains only the first
186// iteration until where it goes wrong.
187//
188// Step 2. the return values of f should be repetitions of the same permutation.
189// E.g. if want is {a,a,b}, the check failes if f returns:
190//  - {a,b,a,b,a,a}: though it satisfies step 1, the second iteration is not
191//  repeating the first iteration.
192//
193// If error is found in this step, the returned error contains the first
194// iteration + the second iteration until where it goes wrong.
195func isRoundRobin(want []balancer.SubConn, f func() balancer.SubConn) error {
196	wantSet := make(map[balancer.SubConn]int) // SubConn -> count, for weighted RR.
197	for _, sc := range want {
198		wantSet[sc]++
199	}
200
201	// The first iteration: makes sure f's return values form a permutation of
202	// elements in want.
203	//
204	// Also keep the returns values in a slice, so we can compare the order in
205	// the second iteration.
206	gotSliceFirstIteration := make([]balancer.SubConn, 0, len(want))
207	for range want {
208		got := f()
209		gotSliceFirstIteration = append(gotSliceFirstIteration, got)
210		wantSet[got]--
211		if wantSet[got] < 0 {
212			return fmt.Errorf("non-roundrobin want: %v, result: %v", want, gotSliceFirstIteration)
213		}
214	}
215
216	// The second iteration should repeat the first iteration.
217	var gotSliceSecondIteration []balancer.SubConn
218	for i := 0; i < 2; i++ {
219		for _, w := range gotSliceFirstIteration {
220			g := f()
221			gotSliceSecondIteration = append(gotSliceSecondIteration, g)
222			if w != g {
223				return fmt.Errorf("non-roundrobin, first iter: %v, second iter: %v", gotSliceFirstIteration, gotSliceSecondIteration)
224			}
225		}
226	}
227
228	return nil
229}
230
231// testClosure is a test util for TestIsRoundRobin.
232type testClosure struct {
233	r []balancer.SubConn
234	i int
235}
236
237func (tc *testClosure) next() balancer.SubConn {
238	ret := tc.r[tc.i]
239	tc.i = (tc.i + 1) % len(tc.r)
240	return ret
241}
242
243func TestIsRoundRobin(t *testing.T) {
244	var (
245		sc1 = testSubConns[0]
246		sc2 = testSubConns[1]
247		sc3 = testSubConns[2]
248	)
249
250	testCases := []struct {
251		desc string
252		want []balancer.SubConn
253		got  []balancer.SubConn
254		pass bool
255	}{
256		{
257			desc: "0 element",
258			want: []balancer.SubConn{},
259			got:  []balancer.SubConn{},
260			pass: true,
261		},
262		{
263			desc: "1 element RR",
264			want: []balancer.SubConn{sc1},
265			got:  []balancer.SubConn{sc1, sc1, sc1, sc1},
266			pass: true,
267		},
268		{
269			desc: "1 element not RR",
270			want: []balancer.SubConn{sc1},
271			got:  []balancer.SubConn{sc1, sc2, sc1},
272			pass: false,
273		},
274		{
275			desc: "2 elements RR",
276			want: []balancer.SubConn{sc1, sc2},
277			got:  []balancer.SubConn{sc1, sc2, sc1, sc2, sc1, sc2},
278			pass: true,
279		},
280		{
281			desc: "2 elements RR different order from want",
282			want: []balancer.SubConn{sc2, sc1},
283			got:  []balancer.SubConn{sc1, sc2, sc1, sc2, sc1, sc2},
284			pass: true,
285		},
286		{
287			desc: "2 elements RR not RR, mistake in first iter",
288			want: []balancer.SubConn{sc1, sc2},
289			got:  []balancer.SubConn{sc1, sc1, sc1, sc2, sc1, sc2},
290			pass: false,
291		},
292		{
293			desc: "2 elements RR not RR, mistake in second iter",
294			want: []balancer.SubConn{sc1, sc2},
295			got:  []balancer.SubConn{sc1, sc2, sc1, sc1, sc1, sc2},
296			pass: false,
297		},
298		{
299			desc: "2 elements weighted RR",
300			want: []balancer.SubConn{sc1, sc1, sc2},
301			got:  []balancer.SubConn{sc1, sc1, sc2, sc1, sc1, sc2},
302			pass: true,
303		},
304		{
305			desc: "2 elements weighted RR different order",
306			want: []balancer.SubConn{sc1, sc1, sc2},
307			got:  []balancer.SubConn{sc1, sc2, sc1, sc1, sc2, sc1},
308			pass: true,
309		},
310
311		{
312			desc: "3 elements RR",
313			want: []balancer.SubConn{sc1, sc2, sc3},
314			got:  []balancer.SubConn{sc1, sc2, sc3, sc1, sc2, sc3, sc1, sc2, sc3},
315			pass: true,
316		},
317		{
318			desc: "3 elements RR different order",
319			want: []balancer.SubConn{sc1, sc2, sc3},
320			got:  []balancer.SubConn{sc3, sc2, sc1, sc3, sc2, sc1},
321			pass: true,
322		},
323		{
324			desc: "3 elements weighted RR",
325			want: []balancer.SubConn{sc1, sc1, sc1, sc2, sc2, sc3},
326			got:  []balancer.SubConn{sc1, sc2, sc3, sc1, sc2, sc1, sc1, sc2, sc3, sc1, sc2, sc1},
327			pass: true,
328		},
329		{
330			desc: "3 elements weighted RR not RR, mistake in first iter",
331			want: []balancer.SubConn{sc1, sc1, sc1, sc2, sc2, sc3},
332			got:  []balancer.SubConn{sc1, sc2, sc1, sc1, sc2, sc1, sc1, sc2, sc3, sc1, sc2, sc1},
333			pass: false,
334		},
335		{
336			desc: "3 elements weighted RR not RR, mistake in second iter",
337			want: []balancer.SubConn{sc1, sc1, sc1, sc2, sc2, sc3},
338			got:  []balancer.SubConn{sc1, sc2, sc3, sc1, sc2, sc1, sc1, sc1, sc3, sc1, sc2, sc1},
339			pass: false,
340		},
341	}
342	for _, tC := range testCases {
343		t.Run(tC.desc, func(t *testing.T) {
344			err := isRoundRobin(tC.want, (&testClosure{r: tC.got}).next)
345			if err == nil != tC.pass {
346				t.Errorf("want pass %v, want %v, got err %v", tC.pass, tC.want, err)
347			}
348		})
349	}
350}
351