@@ -42,14 +42,23 @@ class TMockedPinger {
4242
4343 TDuration operator ()(const Ydb::Discovery::EndpointInfo& endpoint) const {
4444 auto it = EndpointByAdress_.find (endpoint.address ());
45- if (it == EndpointByAdress_.end ()) {
45+ if (it == EndpointByAdress_.end () || Blacklist_. contains (endpoint. address ()) ) {
4646 return TDuration::Max ();
4747 }
4848 return it->second .Ping ();
4949 }
5050
51+ void BanEndpoint (const std::string& adress) {
52+ Blacklist_.insert (adress);
53+ }
54+
55+ void UnbanEndpoint (const std::string& adress) {
56+ Blacklist_.erase (adress);
57+ }
58+
5159private:
5260 mutable std::unordered_map<std::string, TMockedEndpoint> EndpointByAdress_;
61+ std::unordered_set<std::string> Blacklist_;
5362};
5463
5564std::vector<TDuration> GenerateMeasures (size_t count, int minMs, int maxMs, std::mt19937& gen) {
@@ -62,7 +71,6 @@ std::vector<TDuration> GenerateMeasures(size_t count, int minMs, int maxMs, std:
6271 return measures;
6372}
6473
65-
6674Y_UNIT_TEST_SUITE (LocalDCDetectionTest) {
6775 Y_UNIT_TEST (Basic) {
6876 Ydb::Discovery::ListEndpointsResult endpoints;
@@ -73,20 +81,23 @@ Y_UNIT_TEST_SUITE(LocalDCDetectionTest) {
7381 const std::vector<std::string> endpointsB = {" B1" , " B2" , " B3" };
7482 const std::vector<std::string> endpointsC = {" C1" , " C2" , " C3" , " C4" , " C5" , " C6" , " C7" };
7583
84+ const std::size_t epoches = 3 ;
85+ const std::size_t measuresAmount = 10 * epoches;
86+
7687 for (const auto & ep : endpointsA) {
77- mockData[ep] = GenerateMeasures (10 , 20 , 30 , gen);
88+ mockData[ep] = GenerateMeasures (measuresAmount , 20 , 30 , gen);
7889 auto & endpoint = *endpoints.add_endpoints ();
7990 endpoint.set_location (" A" );
8091 endpoint.set_address (ep);
8192 }
8293 for (const auto & ep : endpointsB) {
83- mockData[ep] = GenerateMeasures (10 , 30 , 45 , gen);
94+ mockData[ep] = GenerateMeasures (measuresAmount , 30 , 45 , gen);
8495 auto & endpoint = *endpoints.add_endpoints ();
8596 endpoint.set_location (" B" );
8697 endpoint.set_address (ep);
8798 }
8899 for (const auto & ep : endpointsC) {
89- mockData[ep] = GenerateMeasures (8 , 50 , 70 , gen);
100+ mockData[ep] = GenerateMeasures (measuresAmount , 50 , 70 , gen);
90101 auto & endpoint = *endpoints.add_endpoints ();
91102 endpoint.set_location (" C" );
92103 endpoint.set_address (ep);
@@ -95,44 +106,141 @@ Y_UNIT_TEST_SUITE(LocalDCDetectionTest) {
95106 std::function<TDuration (const Ydb::Discovery::EndpointInfo& endpoint)> pinger = TMockedPinger (mockData);
96107 TLocalDCDetector detector (pinger);
97108
109+ for (std::size_t i = 0 ; i < epoches; ++i) {
110+ detector.DetectLocalDC (endpoints);
111+ UNIT_ASSERT_VALUES_EQUAL (true , detector.IsLocalDC (" A" ));
112+ UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" B" ));
113+ UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" C" ));
114+ }
115+ }
116+
117+ Y_UNIT_TEST (SingleLocation) {
118+ Ydb::Discovery::ListEndpointsResult endpoints;
119+ std::unordered_map<std::string, std::vector<TDuration>> mockData;
120+ std::mt19937 gen (std::random_device{}());
121+
122+ const std::vector<std::string> endpointsA = {" A1" , " A2" , " A3" };
123+
124+ for (const auto & ep : endpointsA) {
125+ mockData[ep] = GenerateMeasures (10 , 20 , 30 , gen);
126+ auto & endpoint = *endpoints.add_endpoints ();
127+ endpoint.set_location (" A" );
128+ endpoint.set_address (ep);
129+ }
130+
131+ std::function<TDuration (const Ydb::Discovery::EndpointInfo& endpoint)> pinger = TMockedPinger (mockData);
132+ TLocalDCDetector detector (pinger);
133+
98134 detector.DetectLocalDC (endpoints);
99135
100- UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" C" ));
101- UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" B" ));
102136 UNIT_ASSERT_VALUES_EQUAL (true , detector.IsLocalDC (" A" ));
103137 }
104138
105- Y_UNIT_TEST (Fallback ) {
139+ Y_UNIT_TEST (UnavailableLocalDC ) {
106140 Ydb::Discovery::ListEndpointsResult endpoints;
107141 std::unordered_map<std::string, std::vector<TDuration>> mockData;
142+ std::mt19937 gen (std::random_device{}());
108143
109- const std::vector<std::string> endpointsA = {" A1" , " A2" , " A3" };
110- const std::vector<std::string> endpointsB = {" B1" , " B2" , " B3" };
111- const std::vector<std::string> endpointsC = {" C1" , " C2" , " C3" };
144+ const std::vector<std::string> endpointsA = {" A1" , " A2" , " A3" , " A4" , " A5" , " A6" , " A7" };
145+ const std::vector<std::string> endpointsB = {" B1" , " B2" , " B3" , " B4" , " B5" , " B6" , " B7" };
146+ const std::vector<std::string> endpointsC = {" C1" , " C2" , " C3" , " C4" , " C5" , " C6" , " C7" };
147+
148+ const std::size_t epoches = 3 ;
149+ const std::size_t measuresAmount = 10 * epoches;
112150
113151 for (const auto & ep : endpointsA) {
152+ mockData[ep] = GenerateMeasures (measuresAmount, 20 , 30 , gen);
114153 auto & endpoint = *endpoints.add_endpoints ();
115154 endpoint.set_location (" A" );
116155 endpoint.set_address (ep);
117156 }
118157 for (const auto & ep : endpointsB) {
158+ mockData[ep] = GenerateMeasures (measuresAmount, 30 , 45 , gen);
119159 auto & endpoint = *endpoints.add_endpoints ();
120160 endpoint.set_location (" B" );
121161 endpoint.set_address (ep);
122162 }
123163 for (const auto & ep : endpointsC) {
164+ mockData[ep] = GenerateMeasures (measuresAmount, 50 , 70 , gen);
124165 auto & endpoint = *endpoints.add_endpoints ();
125166 endpoint.set_location (" C" );
126167 endpoint.set_address (ep);
127168 }
128169
129- std::function<TDuration (const Ydb::Discovery::EndpointInfo& endpoint)> pinger = TMockedPinger (mockData);
170+ TMockedPinger mockPinger (mockData);
171+ std::function<TDuration (const Ydb::Discovery::EndpointInfo& endpoint)> pinger =
172+ [&mockPinger](const Ydb::Discovery::EndpointInfo& endpoint) { return mockPinger (endpoint); };
173+
130174 TLocalDCDetector detector (pinger);
131175
132176 detector.DetectLocalDC (endpoints);
177+ UNIT_ASSERT_VALUES_EQUAL (true , detector.IsLocalDC (" A" ));
178+ UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" B" ));
179+ UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" C" ));
133180
134- UNIT_ASSERT_VALUES_EQUAL (true , detector.IsLocalDC (" C" ));
181+ for (const auto & ep : endpointsA) {
182+ mockPinger.BanEndpoint (ep);
183+ }
184+
185+ detector.DetectLocalDC (endpoints);
186+ UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" A" ));
135187 UNIT_ASSERT_VALUES_EQUAL (true , detector.IsLocalDC (" B" ));
188+ UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" C" ));
189+
190+ for (const auto & ep : endpointsA) {
191+ mockPinger.UnbanEndpoint (ep);
192+ }
193+
194+ detector.DetectLocalDC (endpoints);
136195 UNIT_ASSERT_VALUES_EQUAL (true , detector.IsLocalDC (" A" ));
196+ UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" B" ));
197+ UNIT_ASSERT_VALUES_EQUAL (false , detector.IsLocalDC (" C" ));
198+ }
199+
200+ Y_UNIT_TEST (OfflineDCs) {
201+ Ydb::Discovery::ListEndpointsResult endpoints;
202+ std::unordered_map<std::string, std::vector<TDuration>> mockData;
203+ std::mt19937 gen (std::random_device{}());
204+
205+ const std::vector<std::string> endpointsA = {" A1" , " A2" , " A3" , " A4" , " A5" , " A6" , " A7" };
206+ const std::vector<std::string> endpointsB = {" B1" , " B2" , " B3" , " B4" , " B5" , " B6" , " B7" };
207+ const std::vector<std::string> endpointsC = {" C1" , " C2" , " C3" , " C4" , " C5" , " C6" , " C7" };
208+
209+ for (const auto & ep : endpointsA) {
210+ auto & endpoint = *endpoints.add_endpoints ();
211+ endpoint.set_location (" A" );
212+ endpoint.set_address (ep);
213+ }
214+ for (const auto & ep : endpointsB) {
215+ auto & endpoint = *endpoints.add_endpoints ();
216+ endpoint.set_location (" B" );
217+ endpoint.set_address (ep);
218+ }
219+ for (const auto & ep : endpointsC) {
220+ auto & endpoint = *endpoints.add_endpoints ();
221+ endpoint.set_location (" C" );
222+ endpoint.set_address (ep);
223+ }
224+
225+ TMockedPinger mockPinger (mockData);
226+ std::function<TDuration (const Ydb::Discovery::EndpointInfo& endpoint)> pinger =
227+ [&mockPinger](const Ydb::Discovery::EndpointInfo& endpoint) { return mockPinger (endpoint); };
228+
229+ TLocalDCDetector detector (pinger);
230+
231+ for (const auto & ep : endpointsA) {
232+ mockPinger.BanEndpoint (ep);
233+ }
234+ for (const auto & ep : endpointsB) {
235+ mockPinger.BanEndpoint (ep);
236+ }
237+ for (const auto & ep : endpointsC) {
238+ mockPinger.BanEndpoint (ep);
239+ }
240+
241+ detector.DetectLocalDC (endpoints);
242+ UNIT_ASSERT_VALUES_EQUAL (true , detector.IsLocalDC (" A" ));
243+ UNIT_ASSERT_VALUES_EQUAL (true , detector.IsLocalDC (" B" ));
244+ UNIT_ASSERT_VALUES_EQUAL (true , detector.IsLocalDC (" C" ));
137245 }
138246}
0 commit comments