1/*
2 *
3 * Copyright 2018 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 grpclb
20
21import (
22	"fmt"
23	"sync"
24	"testing"
25	"time"
26
27	"google.golang.org/grpc/balancer"
28	"google.golang.org/grpc/resolver"
29)
30
31type mockSubConn struct {
32	balancer.SubConn
33}
34
35type mockClientConn struct {
36	balancer.ClientConn
37
38	mu       sync.Mutex
39	subConns map[balancer.SubConn]resolver.Address
40}
41
42func newMockClientConn() *mockClientConn {
43	return &mockClientConn{
44		subConns: make(map[balancer.SubConn]resolver.Address),
45	}
46}
47
48func (mcc *mockClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
49	sc := &mockSubConn{}
50	mcc.mu.Lock()
51	defer mcc.mu.Unlock()
52	mcc.subConns[sc] = addrs[0]
53	return sc, nil
54}
55
56func (mcc *mockClientConn) RemoveSubConn(sc balancer.SubConn) {
57	mcc.mu.Lock()
58	defer mcc.mu.Unlock()
59	delete(mcc.subConns, sc)
60}
61
62const testCacheTimeout = 100 * time.Millisecond
63
64func checkMockCC(mcc *mockClientConn, scLen int) error {
65	mcc.mu.Lock()
66	defer mcc.mu.Unlock()
67	if len(mcc.subConns) != scLen {
68		return fmt.Errorf("mcc = %+v, want len(mcc.subConns) = %v", mcc.subConns, scLen)
69	}
70	return nil
71}
72
73func checkCacheCC(ccc *lbCacheClientConn, sccLen, sctaLen int) error {
74	ccc.mu.Lock()
75	defer ccc.mu.Unlock()
76	if len(ccc.subConnCache) != sccLen {
77		return fmt.Errorf("ccc = %+v, want len(ccc.subConnCache) = %v", ccc.subConnCache, sccLen)
78	}
79	if len(ccc.subConnToAddr) != sctaLen {
80		return fmt.Errorf("ccc = %+v, want len(ccc.subConnToAddr) = %v", ccc.subConnToAddr, sctaLen)
81	}
82	return nil
83}
84
85// Test that SubConn won't be immediately removed.
86func TestLBCacheClientConnExpire(t *testing.T) {
87	mcc := newMockClientConn()
88	if err := checkMockCC(mcc, 0); err != nil {
89		t.Fatal(err)
90	}
91
92	ccc := newLBCacheClientConn(mcc)
93	ccc.timeout = testCacheTimeout
94	if err := checkCacheCC(ccc, 0, 0); err != nil {
95		t.Fatal(err)
96	}
97
98	sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
99	// One subconn in MockCC.
100	if err := checkMockCC(mcc, 1); err != nil {
101		t.Fatal(err)
102	}
103	// No subconn being deleted, and one in CacheCC.
104	if err := checkCacheCC(ccc, 0, 1); err != nil {
105		t.Fatal(err)
106	}
107
108	ccc.RemoveSubConn(sc)
109	// One subconn in MockCC before timeout.
110	if err := checkMockCC(mcc, 1); err != nil {
111		t.Fatal(err)
112	}
113	// One subconn being deleted, and one in CacheCC.
114	if err := checkCacheCC(ccc, 1, 1); err != nil {
115		t.Fatal(err)
116	}
117
118	// Should all become empty after timeout.
119	var err error
120	for i := 0; i < 2; i++ {
121		time.Sleep(testCacheTimeout)
122		err = checkMockCC(mcc, 0)
123		if err != nil {
124			continue
125		}
126		err = checkCacheCC(ccc, 0, 0)
127		if err != nil {
128			continue
129		}
130	}
131	if err != nil {
132		t.Fatal(err)
133	}
134}
135
136// Test that NewSubConn with the same address of a SubConn being removed will
137// reuse the SubConn and cancel the removing.
138func TestLBCacheClientConnReuse(t *testing.T) {
139	mcc := newMockClientConn()
140	if err := checkMockCC(mcc, 0); err != nil {
141		t.Fatal(err)
142	}
143
144	ccc := newLBCacheClientConn(mcc)
145	ccc.timeout = testCacheTimeout
146	if err := checkCacheCC(ccc, 0, 0); err != nil {
147		t.Fatal(err)
148	}
149
150	sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
151	// One subconn in MockCC.
152	if err := checkMockCC(mcc, 1); err != nil {
153		t.Fatal(err)
154	}
155	// No subconn being deleted, and one in CacheCC.
156	if err := checkCacheCC(ccc, 0, 1); err != nil {
157		t.Fatal(err)
158	}
159
160	ccc.RemoveSubConn(sc)
161	// One subconn in MockCC before timeout.
162	if err := checkMockCC(mcc, 1); err != nil {
163		t.Fatal(err)
164	}
165	// One subconn being deleted, and one in CacheCC.
166	if err := checkCacheCC(ccc, 1, 1); err != nil {
167		t.Fatal(err)
168	}
169
170	// Recreate the old subconn, this should cancel the deleting process.
171	sc, _ = ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
172	// One subconn in MockCC.
173	if err := checkMockCC(mcc, 1); err != nil {
174		t.Fatal(err)
175	}
176	// No subconn being deleted, and one in CacheCC.
177	if err := checkCacheCC(ccc, 0, 1); err != nil {
178		t.Fatal(err)
179	}
180
181	var err error
182	// Should not become empty after 2*timeout.
183	time.Sleep(2 * testCacheTimeout)
184	err = checkMockCC(mcc, 1)
185	if err != nil {
186		t.Fatal(err)
187	}
188	err = checkCacheCC(ccc, 0, 1)
189	if err != nil {
190		t.Fatal(err)
191	}
192
193	// Call remove again, will delete after timeout.
194	ccc.RemoveSubConn(sc)
195	// One subconn in MockCC before timeout.
196	if err := checkMockCC(mcc, 1); err != nil {
197		t.Fatal(err)
198	}
199	// One subconn being deleted, and one in CacheCC.
200	if err := checkCacheCC(ccc, 1, 1); err != nil {
201		t.Fatal(err)
202	}
203
204	// Should all become empty after timeout.
205	for i := 0; i < 2; i++ {
206		time.Sleep(testCacheTimeout)
207		err = checkMockCC(mcc, 0)
208		if err != nil {
209			continue
210		}
211		err = checkCacheCC(ccc, 0, 0)
212		if err != nil {
213			continue
214		}
215	}
216	if err != nil {
217		t.Fatal(err)
218	}
219}
220
221// Test that if the timer to remove a SubConn fires at the same time NewSubConn
222// cancels the timer, it doesn't cause deadlock.
223func TestLBCache_RemoveTimer_New_Race(t *testing.T) {
224	mcc := newMockClientConn()
225	if err := checkMockCC(mcc, 0); err != nil {
226		t.Fatal(err)
227	}
228
229	ccc := newLBCacheClientConn(mcc)
230	ccc.timeout = time.Nanosecond
231	if err := checkCacheCC(ccc, 0, 0); err != nil {
232		t.Fatal(err)
233	}
234
235	sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
236	// One subconn in MockCC.
237	if err := checkMockCC(mcc, 1); err != nil {
238		t.Fatal(err)
239	}
240	// No subconn being deleted, and one in CacheCC.
241	if err := checkCacheCC(ccc, 0, 1); err != nil {
242		t.Fatal(err)
243	}
244
245	done := make(chan struct{})
246
247	go func() {
248		for i := 0; i < 1000; i++ {
249			// Remove starts a timer with 1 ns timeout, the NewSubConn will race
250			// with with the timer.
251			ccc.RemoveSubConn(sc)
252			sc, _ = ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
253		}
254		close(done)
255	}()
256
257	select {
258	case <-time.After(time.Second):
259		t.Fatalf("Test didn't finish within 1 second. Deadlock")
260	case <-done:
261	}
262}
263