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 resolver
20
21import (
22	"testing"
23	"time"
24
25	"github.com/google/go-cmp/cmp"
26	"google.golang.org/grpc/internal/grpctest"
27	"google.golang.org/grpc/internal/serviceconfig"
28)
29
30type s struct {
31	grpctest.Tester
32}
33
34func Test(t *testing.T) {
35	grpctest.RunSubTests(t, s{})
36}
37
38type fakeConfigSelector struct {
39	selectConfig func(RPCInfo) (*RPCConfig, error)
40}
41
42func (f *fakeConfigSelector) SelectConfig(r RPCInfo) (*RPCConfig, error) {
43	return f.selectConfig(r)
44}
45
46func (s) TestSafeConfigSelector(t *testing.T) {
47	testRPCInfo := RPCInfo{Method: "test method"}
48
49	retChan1 := make(chan *RPCConfig)
50	retChan2 := make(chan *RPCConfig)
51	defer close(retChan1)
52	defer close(retChan2)
53
54	one := 1
55	two := 2
56
57	resp1 := &RPCConfig{MethodConfig: serviceconfig.MethodConfig{MaxReqSize: &one}}
58	resp2 := &RPCConfig{MethodConfig: serviceconfig.MethodConfig{MaxReqSize: &two}}
59
60	cs1Called := make(chan struct{}, 1)
61	cs2Called := make(chan struct{}, 1)
62
63	cs1 := &fakeConfigSelector{
64		selectConfig: func(r RPCInfo) (*RPCConfig, error) {
65			cs1Called <- struct{}{}
66			if diff := cmp.Diff(r, testRPCInfo); diff != "" {
67				t.Errorf("SelectConfig(%v) called; want %v\n  Diffs:\n%s", r, testRPCInfo, diff)
68			}
69			return <-retChan1, nil
70		},
71	}
72	cs2 := &fakeConfigSelector{
73		selectConfig: func(r RPCInfo) (*RPCConfig, error) {
74			cs2Called <- struct{}{}
75			if diff := cmp.Diff(r, testRPCInfo); diff != "" {
76				t.Errorf("SelectConfig(%v) called; want %v\n  Diffs:\n%s", r, testRPCInfo, diff)
77			}
78			return <-retChan2, nil
79		},
80	}
81
82	scs := &SafeConfigSelector{}
83	scs.UpdateConfigSelector(cs1)
84
85	cs1Returned := make(chan struct{})
86	go func() {
87		got, err := scs.SelectConfig(testRPCInfo) // blocks until send to retChan1
88		if err != nil || got != resp1 {
89			t.Errorf("SelectConfig(%v) = %v, %v; want %v, nil", testRPCInfo, got, err, resp1)
90		}
91		close(cs1Returned)
92	}()
93
94	// cs1 is blocked but should be called
95	select {
96	case <-time.After(500 * time.Millisecond):
97		t.Fatalf("timed out waiting for cs1 to be called")
98	case <-cs1Called:
99	}
100
101	// swap in cs2 now that cs1 is called
102	csSwapped := make(chan struct{})
103	go func() {
104		// wait awhile first to ensure cs1 could be called below.
105		time.Sleep(50 * time.Millisecond)
106		scs.UpdateConfigSelector(cs2) // Blocks until cs1 done
107		close(csSwapped)
108	}()
109
110	// Allow cs1 to return and cs2 to eventually be swapped in.
111	retChan1 <- resp1
112
113	cs1Done := false // set when cs2 is first called
114	for dl := time.Now().Add(150 * time.Millisecond); !time.Now().After(dl); {
115		gotConfigChan := make(chan *RPCConfig)
116		go func() {
117			cfg, _ := scs.SelectConfig(testRPCInfo)
118			gotConfigChan <- cfg
119		}()
120		select {
121		case <-time.After(500 * time.Millisecond):
122			t.Fatalf("timed out waiting for cs1 or cs2 to be called")
123		case <-cs1Called:
124			// Initially, before swapping to cs2, cs1 should be called
125			retChan1 <- resp1
126			go func() { <-gotConfigChan }()
127			if cs1Done {
128				t.Fatalf("cs1 called after cs2")
129			}
130		case <-cs2Called:
131			// Success! the new config selector is being called
132			if !cs1Done {
133				select {
134				case <-csSwapped:
135				case <-time.After(50 * time.Millisecond):
136					t.Fatalf("timed out waiting for UpdateConfigSelector to return")
137				}
138				select {
139				case <-cs1Returned:
140				case <-time.After(50 * time.Millisecond):
141					t.Fatalf("timed out waiting for cs1 to return")
142				}
143				cs1Done = true
144			}
145			retChan2 <- resp2
146			got := <-gotConfigChan
147			if diff := cmp.Diff(got, resp2); diff != "" {
148				t.Fatalf("SelectConfig(%v) = %v; want %v\n  Diffs:\n%s", testRPCInfo, got, resp2, diff)
149			}
150		}
151		time.Sleep(10 * time.Millisecond)
152	}
153	if !cs1Done {
154		t.Fatalf("timed out waiting for cs2 to be called")
155	}
156}
157