88
99package org .elasticsearch .cluster .allocation ;
1010
11- import com .carrotsearch .hppc .ObjectIntHashMap ;
12-
1311import org .apache .logging .log4j .LogManager ;
1412import org .apache .logging .log4j .Logger ;
1513import org .elasticsearch .action .admin .cluster .health .ClusterHealthResponse ;
2725import org .elasticsearch .test .ESIntegTestCase .ClusterScope ;
2826
2927import java .util .Arrays ;
28+ import java .util .HashMap ;
3029import java .util .List ;
30+ import java .util .Map ;
3131import java .util .concurrent .TimeUnit ;
3232import java .util .stream .Collectors ;
3333
@@ -95,14 +95,7 @@ public void testSimpleAwareness() throws Exception {
9595 assertThat ("Some indices not closed" , notClosedIndices , empty ());
9696
9797 // verify that we have all the primaries on node3
98- ObjectIntHashMap <String > counts = new ObjectIntHashMap <>();
99- for (IndexRoutingTable indexRoutingTable : clusterState .routingTable ()) {
100- for (IndexShardRoutingTable indexShardRoutingTable : indexRoutingTable ) {
101- for (ShardRouting shardRouting : indexShardRoutingTable ) {
102- counts .addTo (clusterState .nodes ().get (shardRouting .currentNodeId ()).getName (), 1 );
103- }
104- }
105- }
98+ Map <String , Integer > counts = computeShardCounts (clusterState );
10699 assertThat (counts .get (node3 ), equalTo (totalPrimaries ));
107100 }, 10 , TimeUnit .SECONDS );
108101 }
@@ -151,15 +144,8 @@ public void testAwarenessZones() {
151144 assertThat (health .isTimedOut (), equalTo (false ));
152145
153146 ClusterState clusterState = client ().admin ().cluster ().prepareState ().execute ().actionGet ().getState ();
154- ObjectIntHashMap <String > counts = new ObjectIntHashMap <>( );
147+ Map <String , Integer > counts = computeShardCounts ( clusterState );
155148
156- for (IndexRoutingTable indexRoutingTable : clusterState .routingTable ()) {
157- for (IndexShardRoutingTable indexShardRoutingTable : indexRoutingTable ) {
158- for (ShardRouting shardRouting : indexShardRoutingTable ) {
159- counts .addTo (clusterState .nodes ().get (shardRouting .currentNodeId ()).getName (), 1 );
160- }
161- }
162- }
163149 assertThat (counts .get (A_1 ), anyOf (equalTo (2 ), equalTo (3 )));
164150 assertThat (counts .get (B_1 ), anyOf (equalTo (2 ), equalTo (3 )));
165151 assertThat (counts .get (A_0 ), anyOf (equalTo (2 ), equalTo (3 )));
@@ -201,15 +187,8 @@ public void testAwarenessZonesIncrementalNodes() {
201187 .actionGet ();
202188 assertThat (health .isTimedOut (), equalTo (false ));
203189 ClusterState clusterState = client ().admin ().cluster ().prepareState ().execute ().actionGet ().getState ();
204- ObjectIntHashMap <String > counts = new ObjectIntHashMap <>( );
190+ Map <String , Integer > counts = computeShardCounts ( clusterState );
205191
206- for (IndexRoutingTable indexRoutingTable : clusterState .routingTable ()) {
207- for (IndexShardRoutingTable indexShardRoutingTable : indexRoutingTable ) {
208- for (ShardRouting shardRouting : indexShardRoutingTable ) {
209- counts .addTo (clusterState .nodes ().get (shardRouting .currentNodeId ()).getName (), 1 );
210- }
211- }
212- }
213192 assertThat (counts .get (A_0 ), equalTo (5 ));
214193 assertThat (counts .get (B_0 ), equalTo (5 ));
215194 logger .info ("--> starting another node in zone 'b'" );
@@ -240,16 +219,8 @@ public void testAwarenessZonesIncrementalNodes() {
240219
241220 assertThat (health .isTimedOut (), equalTo (false ));
242221 clusterState = client ().admin ().cluster ().prepareState ().execute ().actionGet ().getState ();
222+ counts = computeShardCounts (clusterState );
243223
244- counts = new ObjectIntHashMap <>();
245-
246- for (IndexRoutingTable indexRoutingTable : clusterState .routingTable ()) {
247- for (IndexShardRoutingTable indexShardRoutingTable : indexRoutingTable ) {
248- for (ShardRouting shardRouting : indexShardRoutingTable ) {
249- counts .addTo (clusterState .nodes ().get (shardRouting .currentNodeId ()).getName (), 1 );
250- }
251- }
252- }
253224 assertThat (counts .get (A_0 ), equalTo (5 ));
254225 assertThat (counts .get (B_0 ), equalTo (3 ));
255226 assertThat (counts .get (B_1 ), equalTo (2 ));
@@ -280,16 +251,7 @@ public void testAwarenessZonesIncrementalNodes() {
280251
281252 assertThat (health .isTimedOut (), equalTo (false ));
282253 clusterState = client ().admin ().cluster ().prepareState ().execute ().actionGet ().getState ();
283-
284- counts = new ObjectIntHashMap <>();
285-
286- for (IndexRoutingTable indexRoutingTable : clusterState .routingTable ()) {
287- for (IndexShardRoutingTable indexShardRoutingTable : indexRoutingTable ) {
288- for (ShardRouting shardRouting : indexShardRoutingTable ) {
289- counts .addTo (clusterState .nodes ().get (shardRouting .currentNodeId ()).getName (), 1 );
290- }
291- }
292- }
254+ counts = computeShardCounts (clusterState );
293255
294256 assertThat (counts .get (A_0 ), equalTo (5 ));
295257 assertThat (counts .get (B_0 ), equalTo (3 ));
@@ -315,16 +277,7 @@ public void testAwarenessZonesIncrementalNodes() {
315277
316278 assertThat (health .isTimedOut (), equalTo (false ));
317279 clusterState = client ().admin ().cluster ().prepareState ().execute ().actionGet ().getState ();
318-
319- counts = new ObjectIntHashMap <>();
320-
321- for (IndexRoutingTable indexRoutingTable : clusterState .routingTable ()) {
322- for (IndexShardRoutingTable indexShardRoutingTable : indexRoutingTable ) {
323- for (ShardRouting shardRouting : indexShardRoutingTable ) {
324- counts .addTo (clusterState .nodes ().get (shardRouting .currentNodeId ()).getName (), 1 );
325- }
326- }
327- }
280+ counts = computeShardCounts (clusterState );
328281
329282 assertThat (counts .get (A_0 ), equalTo (3 ));
330283 assertThat (counts .get (B_0 ), equalTo (3 ));
@@ -371,4 +324,17 @@ public void testForceAwarenessSettingValidation() {
371324 containsString ("[cluster.routing.allocation.awareness.force.attr.values.junk]" )
372325 );
373326 }
327+
328+ Map <String , Integer > computeShardCounts (ClusterState clusterState ) {
329+ Map <String , Integer > counts = new HashMap <>();
330+
331+ for (IndexRoutingTable indexRoutingTable : clusterState .routingTable ()) {
332+ for (IndexShardRoutingTable indexShardRoutingTable : indexRoutingTable ) {
333+ for (ShardRouting shardRouting : indexShardRoutingTable ) {
334+ counts .merge (clusterState .nodes ().get (shardRouting .currentNodeId ()).getName (), 1 , Integer ::sum );
335+ }
336+ }
337+ }
338+ return counts ;
339+ }
374340}
0 commit comments