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