@Test
  public void cancelCommand() {
    AllocationService allocation =
        new AllocationService(
            settingsBuilder()
                .put("cluster.routing.allocation.disable_new_allocation", true)
                .put("cluster.routing.allocation.disable_allocation", true)
                .build());

    logger.info("--> building initial routing table");
    MetaData metaData =
        newMetaDataBuilder()
            .put(newIndexMetaDataBuilder("test").numberOfShards(1).numberOfReplicas(1))
            .build();
    RoutingTable routingTable = routingTable().addAsNew(metaData.index("test")).build();
    ClusterState clusterState =
        newClusterStateBuilder().metaData(metaData).routingTable(routingTable).build();

    logger.info("--> adding 3 nodes");
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .nodes(
                newNodesBuilder().put(newNode("node1")).put(newNode("node2")).put(newNode("node3")))
            .build();
    RoutingAllocationResult rerouteResult = allocation.reroute(clusterState);
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().shardsWithState(INITIALIZING).size(), equalTo(0));

    logger.info("--> allocating with primary flag set to true");
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new AllocateAllocationCommand(new ShardId("test", 0), "node1", true)));
    assertThat(rerouteResult.changed(), equalTo(true));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(INITIALIZING).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));

    logger.info("--> cancel primary allocation, make sure it fails...");
    try {
      allocation.reroute(
          clusterState,
          new AllocationCommands(
              new CancelAllocationCommand(new ShardId("test", 0), "node1", false)));
      assert false;
    } catch (ElasticSearchIllegalArgumentException e) {
    }

    logger.info("--> start the primary shard");
    rerouteResult =
        allocation.applyStartedShards(
            clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));

    logger.info("--> cancel primary allocation, make sure it fails...");
    try {
      allocation.reroute(
          clusterState,
          new AllocationCommands(
              new CancelAllocationCommand(new ShardId("test", 0), "node1", false)));
      assert false;
    } catch (ElasticSearchIllegalArgumentException e) {
    }

    logger.info("--> allocate the replica shard on on the second node");
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new AllocateAllocationCommand(new ShardId("test", 0), "node2", false)));
    assertThat(rerouteResult.changed(), equalTo(true));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node2").shardsWithState(INITIALIZING).size(), equalTo(1));

    logger.info("--> cancel the relocation allocation");
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new CancelAllocationCommand(new ShardId("test", 0), "node2", false)));
    assertThat(rerouteResult.changed(), equalTo(true));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));
    assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));

    logger.info("--> allocate the replica shard on on the second node");
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new AllocateAllocationCommand(new ShardId("test", 0), "node2", false)));
    assertThat(rerouteResult.changed(), equalTo(true));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node2").shardsWithState(INITIALIZING).size(), equalTo(1));

    logger.info("--> cancel the primary being replicated, make sure it fails");
    try {
      allocation.reroute(
          clusterState,
          new AllocationCommands(
              new CancelAllocationCommand(new ShardId("test", 0), "node1", false)));
      assert false;
    } catch (ElasticSearchIllegalArgumentException e) {
    }

    logger.info("--> start the replica shard");
    rerouteResult =
        allocation.applyStartedShards(
            clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node2").shardsWithState(STARTED).size(), equalTo(1));

    logger.info("--> cancel allocation of the replica shard");
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new CancelAllocationCommand(new ShardId("test", 0), "node2", false)));
    assertThat(rerouteResult.changed(), equalTo(true));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));
    assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));

    logger.info("--> allocate the replica shard on on the second node");
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new AllocateAllocationCommand(new ShardId("test", 0), "node2", false)));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(rerouteResult.changed(), equalTo(true));
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node2").shardsWithState(INITIALIZING).size(), equalTo(1));
    logger.info("--> start the replica shard");
    rerouteResult =
        allocation.applyStartedShards(
            clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node2").shardsWithState(STARTED).size(), equalTo(1));

    logger.info("--> cancel the primary allocation (with allow_primary set to true)");
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new CancelAllocationCommand(new ShardId("test", 0), "node1", true)));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(rerouteResult.changed(), equalTo(true));
    assertThat(
        clusterState.routingNodes().node("node2").shardsWithState(STARTED).get(0).primary(),
        equalTo(true));
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(0));
    assertThat(clusterState.routingNodes().node("node3").shards().size(), equalTo(0));
  }
  @Test
  public void moveShardCommand() {
    AllocationService allocation =
        new AllocationService(
            settingsBuilder().put("cluster.routing.allocation.concurrent_recoveries", 10).build());

    logger.info("creating an index with 1 shard, no replica");
    MetaData metaData =
        newMetaDataBuilder()
            .put(newIndexMetaDataBuilder("test").numberOfShards(1).numberOfReplicas(0))
            .build();
    RoutingTable routingTable = routingTable().addAsNew(metaData.index("test")).build();
    ClusterState clusterState =
        newClusterStateBuilder().metaData(metaData).routingTable(routingTable).build();

    logger.info("adding two nodes and performing rerouting");
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .nodes(newNodesBuilder().put(newNode("node1")).put(newNode("node2")))
            .build();
    RoutingAllocationResult rerouteResult = allocation.reroute(clusterState);
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();

    logger.info("start primary shard");
    rerouteResult =
        allocation.applyStartedShards(
            clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();

    logger.info("move the shard");
    String existingNodeId =
        clusterState.routingTable().index("test").shard(0).primaryShard().currentNodeId();
    String toNodeId;
    if ("node1".equals(existingNodeId)) {
      toNodeId = "node2";
    } else {
      toNodeId = "node1";
    }
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new MoveAllocationCommand(new ShardId("test", 0), existingNodeId, toNodeId)));
    assertThat(rerouteResult.changed(), equalTo(true));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(
        clusterState.routingNodes().node(existingNodeId).shards().get(0).state(),
        equalTo(ShardRoutingState.RELOCATING));
    assertThat(
        clusterState.routingNodes().node(toNodeId).shards().get(0).state(),
        equalTo(ShardRoutingState.INITIALIZING));

    logger.info("finish moving the shard");
    rerouteResult =
        allocation.applyStartedShards(
            clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();

    assertThat(clusterState.routingNodes().node(existingNodeId).shards().isEmpty(), equalTo(true));
    assertThat(
        clusterState.routingNodes().node(toNodeId).shards().get(0).state(),
        equalTo(ShardRoutingState.STARTED));
  }
  @Test
  public void allocateCommand() {
    AllocationService allocation =
        new AllocationService(
            settingsBuilder()
                .put("cluster.routing.allocation.disable_new_allocation", true)
                .put("cluster.routing.allocation.disable_allocation", true)
                .build());

    logger.info("--> building initial routing table");
    MetaData metaData =
        newMetaDataBuilder()
            .put(newIndexMetaDataBuilder("test").numberOfShards(1).numberOfReplicas(1))
            .build();
    RoutingTable routingTable = routingTable().addAsNew(metaData.index("test")).build();
    ClusterState clusterState =
        newClusterStateBuilder().metaData(metaData).routingTable(routingTable).build();

    logger.info("--> adding 3 nodes on same rack and do rerouting");
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .nodes(
                newNodesBuilder().put(newNode("node1")).put(newNode("node2")).put(newNode("node3")))
            .build();
    RoutingAllocationResult rerouteResult = allocation.reroute(clusterState);
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().shardsWithState(INITIALIZING).size(), equalTo(0));

    logger.info("--> allocating with primary flag set to false, should fail");
    try {
      allocation.reroute(
          clusterState,
          new AllocationCommands(
              new AllocateAllocationCommand(new ShardId("test", 0), "node1", false)));
      assert false;
    } catch (ElasticSearchIllegalArgumentException e) {
    }

    logger.info("--> allocating with primary flag set to true");
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new AllocateAllocationCommand(new ShardId("test", 0), "node1", true)));
    assertThat(rerouteResult.changed(), equalTo(true));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(INITIALIZING).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));

    logger.info("--> start the primary shard");
    rerouteResult =
        allocation.applyStartedShards(
            clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(0));

    logger.info("--> allocate the replica shard on the primary shard node, should fail");
    try {
      allocation.reroute(
          clusterState,
          new AllocationCommands(
              new AllocateAllocationCommand(new ShardId("test", 0), "node1", false)));
      assert false;
    } catch (ElasticSearchIllegalArgumentException e) {
    }

    logger.info("--> allocate the replica shard on on the second node");
    rerouteResult =
        allocation.reroute(
            clusterState,
            new AllocationCommands(
                new AllocateAllocationCommand(new ShardId("test", 0), "node2", false)));
    assertThat(rerouteResult.changed(), equalTo(true));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node2").shardsWithState(INITIALIZING).size(), equalTo(1));

    logger.info("--> start the replica shard");
    rerouteResult =
        allocation.applyStartedShards(
            clusterState, clusterState.routingNodes().shardsWithState(INITIALIZING));
    clusterState =
        newClusterStateBuilder()
            .state(clusterState)
            .routingTable(rerouteResult.routingTable())
            .build();
    assertThat(clusterState.routingNodes().node("node1").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node1").shardsWithState(STARTED).size(), equalTo(1));
    assertThat(clusterState.routingNodes().node("node2").shards().size(), equalTo(1));
    assertThat(
        clusterState.routingNodes().node("node2").shardsWithState(STARTED).size(), equalTo(1));

    logger.info("--> verify that we fail when there are no unassigned shards");
    try {
      allocation.reroute(
          clusterState,
          new AllocationCommands(
              new AllocateAllocationCommand(new ShardId("test", 0), "node3", false)));
      assert false;
    } catch (ElasticSearchIllegalArgumentException e) {
    }
  }