/*
* Licensed to Elasticsearch under one or more contributor
* license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright
* ownership. Elasticsearch licenses this file to you under
* the Apache License, Version 2.0 (the "License"); you may
* not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.elasticsearch.routing;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.test.ESIntegTestCase;
import org.mockito.internal.util.collections.Sets;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
public class PartitionedRoutingIT extends ESIntegTestCase {
public void testVariousPartitionSizes() throws Exception {
for (int shards = 1; shards <= 4; shards++) {
for (int partitionSize = 1; partitionSize < shards; partitionSize++) {
String index = "index_" + shards + "_" + partitionSize;
client().admin().indices().prepareCreate(index)
.setSettings(Settings.builder()
.put("index.number_of_shards", shards)
.put("index.routing_partition_size", partitionSize))
.addMapping("type", "{\"type\":{\"_routing\":{\"required\":true}}}", XContentType.JSON)
.execute().actionGet();
ensureGreen();
Map<String, Set<String>> routingToDocumentIds = generateRoutedDocumentIds(index);
verifyGets(index, routingToDocumentIds);
verifyBroadSearches(index, routingToDocumentIds, shards);
verifyRoutedSearches(index, routingToDocumentIds, Sets.newSet(partitionSize));
}
}
}
public void testShrinking() throws Exception {
// creates random routing groups and repeatedly halves the index until it is down to 1 shard
// verifying that the count is correct for each shrunken index
final int partitionSize = 3;
final int originalShards = 8;
int currentShards = originalShards;
String index = "index_" + currentShards;
client().admin().indices().prepareCreate(index)
.setSettings(Settings.builder()
.put("index.number_of_shards", currentShards)
.put("index.routing_partition_size", partitionSize))
.addMapping("type", "{\"type\":{\"_routing\":{\"required\":true}}}", XContentType.JSON)
.execute().actionGet();
ensureGreen();
Map<String, Set<String>> routingToDocumentIds = generateRoutedDocumentIds(index);
while (true) {
int factor = originalShards / currentShards;
verifyGets(index, routingToDocumentIds);
verifyBroadSearches(index, routingToDocumentIds, currentShards);
// we need the floor and ceiling of the routing_partition_size / factor since the partition size of the shrunken
// index will be one of those, depending on the routing value
verifyRoutedSearches(index, routingToDocumentIds,
Math.floorDiv(partitionSize, factor) == 0 ?
Sets.newSet(1, 2) :
Sets.newSet(Math.floorDiv(partitionSize, factor), -Math.floorDiv(-partitionSize, factor)));
client().admin().indices().prepareUpdateSettings(index)
.setSettings(Settings.builder()
.put("index.routing.allocation.require._name", client().admin().cluster().prepareState().get().getState().nodes()
.getDataNodes().values().toArray(DiscoveryNode.class)[0].getName())
.put("index.blocks.write", true)).get();
ensureGreen();
currentShards = Math.floorDiv(currentShards, 2);
if (currentShards == 0) {
break;
}
String previousIndex = index;
index = "index_" + currentShards;
logger.info("--> shrinking index [" + previousIndex + "] to [" + index + "]");
client().admin().indices().prepareShrinkIndex(previousIndex, index)
.setSettings(Settings.builder()
.put("index.number_of_shards", currentShards)
.build()).get();
ensureGreen();
}
}
private void verifyRoutedSearches(String index, Map<String, Set<String>> routingToDocumentIds, Set<Integer> expectedShards) {
for (Map.Entry<String, Set<String>> routingEntry : routingToDocumentIds.entrySet()) {
String routing = routingEntry.getKey();
int expectedDocuments = routingEntry.getValue().size();
SearchResponse response = client().prepareSearch()
.setQuery(QueryBuilders.termQuery("_routing", routing))
.setRouting(routing)
.setIndices(index)
.setSize(100)
.execute().actionGet();
logger.info("--> routed search on index [" + index + "] visited [" + response.getTotalShards()
+ "] shards for routing [" + routing + "] and got hits [" + response.getHits().getTotalHits() + "]");
assertTrue(response.getTotalShards() + " was not in " + expectedShards + " for " + index,
expectedShards.contains(response.getTotalShards()));
assertEquals(expectedDocuments, response.getHits().getTotalHits());
Set<String> found = new HashSet<>();
response.getHits().forEach(h -> found.add(h.getId()));
assertEquals(routingEntry.getValue(), found);
}
}
private void verifyBroadSearches(String index, Map<String, Set<String>> routingToDocumentIds, int expectedShards) {
for (Map.Entry<String, Set<String>> routingEntry : routingToDocumentIds.entrySet()) {
String routing = routingEntry.getKey();
int expectedDocuments = routingEntry.getValue().size();
SearchResponse response = client().prepareSearch()
.setQuery(QueryBuilders.termQuery("_routing", routing))
.setIndices(index)
.setSize(100)
.execute().actionGet();
assertEquals(expectedShards, response.getTotalShards());
assertEquals(expectedDocuments, response.getHits().getTotalHits());
Set<String> found = new HashSet<>();
response.getHits().forEach(h -> found.add(h.getId()));
assertEquals(routingEntry.getValue(), found);
}
}
private void verifyGets(String index, Map<String, Set<String>> routingToDocumentIds) {
for (Map.Entry<String, Set<String>> routingEntry : routingToDocumentIds.entrySet()) {
String routing = routingEntry.getKey();
for (String id : routingEntry.getValue()) {
assertTrue(client().prepareGet(index, "type", id).setRouting(routing).execute().actionGet().isExists());
}
}
}
private Map<String, Set<String>> generateRoutedDocumentIds(String index) {
Map<String, Set<String>> routingToDocumentIds = new HashMap<>();
int numRoutingValues = randomIntBetween(5, 15);
for (int i = 0; i < numRoutingValues; i++) {
String routingValue = String.valueOf(i);
int numDocuments = randomIntBetween(10, 100);
routingToDocumentIds.put(String.valueOf(routingValue), new HashSet<>());
for (int k = 0; k < numDocuments; k++) {
String id = routingValue + "_" + String.valueOf(k);
routingToDocumentIds.get(routingValue).add(id);
client().prepareIndex(index, "type", id)
.setRouting(routingValue)
.setSource("foo", "bar")
.get();
}
}
client().admin().indices().prepareRefresh(index).get();
return routingToDocumentIds;
}
}