add the ability to control the number of shards and replicas in the restart stress test, also add a step to verify version numbers

modules/test/integration/src/test/java/org/elasticsearch/test/stress/rollingrestart/RollingRestartStressTest.java

@@ -24,6 +24,9 @@ import org.elasticsearch.action.admin.indices.status.IndexShardStatus;
 import org.elasticsearch.action.admin.indices.status.IndicesStatusResponse;
 import org.elasticsearch.action.admin.indices.status.ShardStatus;
 import org.elasticsearch.action.count.CountResponse;
+import org.elasticsearch.action.get.GetResponse;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.common.UUID;
 import org.elasticsearch.common.io.FileSystemUtils;
 import org.elasticsearch.common.logging.ESLogger;
@@ -38,11 +41,15 @@ import org.elasticsearch.env.NodeEnvironment;
 import org.elasticsearch.node.Node;
 import org.elasticsearch.node.NodeBuilder;
 import org.elasticsearch.node.internal.InternalNode;
+import org.elasticsearch.search.SearchHit;
 
 import java.io.File;
 import java.util.concurrent.atomic.AtomicLong;
 
+import static org.elasticsearch.common.settings.ImmutableSettings.*;
 import static org.elasticsearch.index.query.xcontent.QueryBuilders.*;
+import static org.hamcrest.MatcherAssert.*;
+import static org.hamcrest.Matchers.*;
 
 /**
  * @author kimchy (shay.banon)
@@ -51,6 +58,8 @@ public class RollingRestartStressTest {
 
     private final ESLogger logger = Loggers.getLogger(getClass());
 
+    private int numberOfShards = 5;
+    private int numberOfReplicas = 1;
     private int numberOfNodes = 4;
 
     private int textTokens = 150;
@@ -78,6 +87,16 @@ public class RollingRestartStressTest {
         return this;
     }
 
+    public RollingRestartStressTest numberOfShards(int numberOfShards) {
+        this.numberOfShards = numberOfShards;
+        return this;
+    }
+
+    public RollingRestartStressTest numberOfReplicas(int numberOfReplicas) {
+        this.numberOfReplicas = numberOfReplicas;
+        return this;
+    }
+
     public RollingRestartStressTest initialNumberOfDocs(long initialNumberOfDocs) {
         this.initialNumberOfDocs = initialNumberOfDocs;
         return this;
@@ -125,6 +144,11 @@ public class RollingRestartStressTest {
         }
         client = NodeBuilder.nodeBuilder().settings(settings).client(true).node();
 
+        client.client().admin().indices().prepareCreate("test").setSettings(settingsBuilder()
+                .put("index.number_of_shards", numberOfShards)
+                .put("index.number_of_replicas", numberOfReplicas)
+        ).execute().actionGet();
+
         logger.info("********** [START] INDEXING INITIAL DOCS");
         for (long i = 0; i < initialNumberOfDocs; i++) {
             indexDoc();
@@ -223,6 +247,37 @@ public class RollingRestartStressTest {
             }
         }
 
+        // scan all the docs, verify all have the same version based on the number of replicas
+        SearchResponse searchResponse = client.client().prepareSearch()
+                .setSearchType(SearchType.SCAN)
+                .setQuery(matchAllQuery())
+                .setSize(50)
+                .setScroll(TimeValue.timeValueMinutes(2))
+                .execute().actionGet();
+        logger.info("Verifying versions for {} hits...", searchResponse.hits().totalHits());
+
+        while (true) {
+            searchResponse = client.client().prepareSearchScroll(searchResponse.scrollId()).setScroll(TimeValue.timeValueMinutes(2)).execute().actionGet();
+            assertThat(searchResponse.failedShards(), equalTo(0));
+            for (SearchHit hit : searchResponse.hits()) {
+                long version = -1;
+                for (int i = 0; i < (numberOfReplicas + 1); i++) {
+                    GetResponse getResponse = client.client().prepareGet(hit.index(), hit.type(), hit.id()).execute().actionGet();
+                    if (version == -1) {
+                        version = getResponse.version();
+                    } else {
+                        if (version != getResponse.version()) {
+                            logger.warn("Doc {} has different version numbers {} and {}", hit.id(), version, getResponse.version());
+                        }
+                    }
+                }
+            }
+            if (searchResponse.hits().hits().length == 0) {
+                break;
+            }
+        }
+        logger.info("Done verifying versions");
+
         client.close();
         for (Node node : nodes) {
             node.close();
@@ -280,7 +335,7 @@ public class RollingRestartStressTest {
     public static void main(String[] args) throws Exception {
         System.setProperty("es.logger.prefix", "");
 
-        Settings settings = ImmutableSettings.settingsBuilder()
+        Settings settings = settingsBuilder()
                 .put("index.shard.check_index", true)
                 .put("gateway.type", "none")
                 .build();
@@ -288,13 +343,15 @@ public class RollingRestartStressTest {
         RollingRestartStressTest test = new RollingRestartStressTest()
                 .settings(settings)
                 .numberOfNodes(4)
+                .numberOfShards(5)
+                .numberOfReplicas(1)
                 .initialNumberOfDocs(1000)
                 .textTokens(150)
                 .numberOfFields(10)
-                .cleanNodeData(true)
+                .cleanNodeData(false)
                 .indexers(5)
                 .indexerThrottle(TimeValue.timeValueMillis(50))
-                .period(TimeValue.timeValueMinutes(10));
+                .period(TimeValue.timeValueMinutes(1));
 
         test.run();
     }