@SuppressWarnings("unchecked")
 private IngestProxyActionFilter buildFilter(int ingestNodes, int totalNodes) {
   ClusterState.Builder clusterState = new ClusterState.Builder(new ClusterName("_name"));
   DiscoveryNodes.Builder builder = new DiscoveryNodes.Builder();
   DiscoveryNode localNode = null;
   for (int i = 0; i < totalNodes; i++) {
     String nodeId = "node" + i;
     Map<String, String> attributes = new HashMap<>();
     if (i >= ingestNodes) {
       attributes.put("ingest", "false");
     } else if (randomBoolean()) {
       attributes.put("ingest", "true");
     }
     DiscoveryNode node =
         new DiscoveryNode(
             nodeId,
             nodeId,
             DummyTransportAddress.INSTANCE,
             attributes,
             VersionUtils.randomVersion(random()));
     builder.put(node);
     if (i == totalNodes - 1) {
       localNode = node;
     }
   }
   clusterState.nodes(builder);
   ClusterService clusterService = mock(ClusterService.class);
   when(clusterService.localNode()).thenReturn(localNode);
   when(clusterService.state()).thenReturn(clusterState.build());
   transportService = mock(TransportService.class);
   return new IngestProxyActionFilter(clusterService, transportService);
 }
  @Test
  public void testClusterStateSerialization() throws Exception {
    MetaData metaData =
        MetaData.builder()
            .put(IndexMetaData.builder("test").numberOfShards(10).numberOfReplicas(1))
            .build();

    RoutingTable routingTable = RoutingTable.builder().addAsNew(metaData.index("test")).build();

    DiscoveryNodes nodes =
        DiscoveryNodes.builder()
            .put(newNode("node1"))
            .put(newNode("node2"))
            .put(newNode("node3"))
            .localNodeId("node1")
            .masterNodeId("node2")
            .build();

    ClusterState clusterState =
        ClusterState.builder().nodes(nodes).metaData(metaData).routingTable(routingTable).build();

    AllocationService strategy = createAllocationService();
    clusterState =
        ClusterState.builder(clusterState)
            .routingTable(strategy.reroute(clusterState).routingTable())
            .build();

    ClusterState serializedClusterState =
        ClusterState.Builder.fromBytes(
            ClusterState.Builder.toBytes(clusterState), newNode("node1"));

    assertThat(
        serializedClusterState.routingTable().prettyPrint(),
        equalTo(clusterState.routingTable().prettyPrint()));
  }
 /** Creates a cluster state with no index */
 public static ClusterState stateWithNoShard() {
   int numberOfNodes = 2;
   DiscoveryNodes.Builder discoBuilder = DiscoveryNodes.builder();
   discoBuilder.localNodeId(newNode(0).getId());
   discoBuilder.masterNodeId(newNode(1).getId());
   ClusterState.Builder state = ClusterState.builder(new ClusterName("test"));
   state.nodes(discoBuilder);
   state.metaData(MetaData.builder().generateClusterUuidIfNeeded());
   state.routingTable(RoutingTable.builder().build());
   return state.build();
 }
  /**
   * Creates a cluster state where local node and master node can be specified
   *
   * @param localNode node in allNodes that is the local node
   * @param masterNode node in allNodes that is the master node. Can be null if no master exists
   * @param allNodes all nodes in the cluster
   * @return cluster state
   */
  public static ClusterState state(
      DiscoveryNode localNode, DiscoveryNode masterNode, DiscoveryNode... allNodes) {
    DiscoveryNodes.Builder discoBuilder = DiscoveryNodes.builder();
    for (DiscoveryNode node : allNodes) {
      discoBuilder.put(node);
    }
    if (masterNode != null) {
      discoBuilder.masterNodeId(masterNode.getId());
    }
    discoBuilder.localNodeId(localNode.getId());

    ClusterState.Builder state = ClusterState.builder(new ClusterName("test"));
    state.nodes(discoBuilder);
    state.metaData(MetaData.builder().generateClusterUuidIfNeeded());
    return state.build();
  }
  @Override
  protected void masterOperation(
      final ClusterStateRequest request,
      final ClusterState state,
      ActionListener<ClusterStateResponse> listener)
      throws ElasticsearchException {
    ClusterState currentState = clusterService.state();
    logger.trace("Serving cluster state request using version {}", currentState.version());
    ClusterState.Builder builder = ClusterState.builder(currentState.getClusterName());
    builder.version(currentState.version());
    if (request.nodes()) {
      builder.nodes(currentState.nodes());
    }
    if (request.routingTable()) {
      if (request.indices().length > 0) {
        RoutingTable.Builder routingTableBuilder = RoutingTable.builder();
        for (String filteredIndex : request.indices()) {
          if (currentState.routingTable().getIndicesRouting().containsKey(filteredIndex)) {
            routingTableBuilder.add(
                currentState.routingTable().getIndicesRouting().get(filteredIndex));
          }
        }
        builder.routingTable(routingTableBuilder);
      } else {
        builder.routingTable(currentState.routingTable());
      }
    }
    if (request.blocks()) {
      builder.blocks(currentState.blocks());
    }
    if (request.metaData()) {
      MetaData.Builder mdBuilder;
      if (request.indices().length == 0) {
        mdBuilder = MetaData.builder(currentState.metaData());
      } else {
        mdBuilder = MetaData.builder();
      }

      if (request.indices().length > 0) {
        String[] indices =
            currentState
                .metaData()
                .concreteIndices(IndicesOptions.lenientExpandOpen(), request.indices());
        for (String filteredIndex : indices) {
          IndexMetaData indexMetaData = currentState.metaData().index(filteredIndex);
          if (indexMetaData != null) {
            mdBuilder.put(indexMetaData, false);
          }
        }
      }

      builder.metaData(mdBuilder);
    }
    listener.onResponse(new ClusterStateResponse(clusterName, builder.build()));
  }
  /** Creates cluster state with several shards and one replica and all shards STARTED. */
  public static ClusterState stateWithAssignedPrimariesAndOneReplica(
      String index, int numberOfShards) {

    int numberOfNodes = 2; // we need a non-local master to test shard failures
    DiscoveryNodes.Builder discoBuilder = DiscoveryNodes.builder();
    for (int i = 0; i < numberOfNodes + 1; i++) {
      final DiscoveryNode node = newNode(i);
      discoBuilder = discoBuilder.put(node);
    }
    discoBuilder.localNodeId(newNode(0).getId());
    discoBuilder.masterNodeId(
        newNode(1).getId()); // we need a non-local master to test shard failures
    IndexMetaData indexMetaData =
        IndexMetaData.builder(index)
            .settings(
                Settings.builder()
                    .put(SETTING_VERSION_CREATED, Version.CURRENT)
                    .put(SETTING_NUMBER_OF_SHARDS, numberOfShards)
                    .put(SETTING_NUMBER_OF_REPLICAS, 1)
                    .put(SETTING_CREATION_DATE, System.currentTimeMillis()))
            .build();
    ClusterState.Builder state = ClusterState.builder(new ClusterName("test"));
    state.nodes(discoBuilder);
    state.metaData(MetaData.builder().put(indexMetaData, false).generateClusterUuidIfNeeded());
    IndexRoutingTable.Builder indexRoutingTableBuilder =
        IndexRoutingTable.builder(indexMetaData.getIndex());
    for (int i = 0; i < numberOfShards; i++) {
      RoutingTable.Builder routing = new RoutingTable.Builder();
      routing.addAsNew(indexMetaData);
      final ShardId shardId = new ShardId(index, "_na_", i);
      IndexShardRoutingTable.Builder indexShardRoutingBuilder =
          new IndexShardRoutingTable.Builder(shardId);
      indexShardRoutingBuilder.addShard(
          TestShardRouting.newShardRouting(
              index, i, newNode(0).getId(), null, null, true, ShardRoutingState.STARTED, null));
      indexShardRoutingBuilder.addShard(
          TestShardRouting.newShardRouting(
              index, i, newNode(1).getId(), null, null, false, ShardRoutingState.STARTED, null));
      indexRoutingTableBuilder.addIndexShard(indexShardRoutingBuilder.build());
    }
    state.routingTable(RoutingTable.builder().add(indexRoutingTableBuilder.build()).build());
    return state.build();
  }
    @Override
    @SuppressWarnings("unchecked")
    public <T extends TransportResponse> void sendRequest(
        DiscoveryNode node,
        String action,
        TransportRequest request,
        TransportRequestOptions options,
        TransportResponseHandler<T> handler) {
      if (TransportLivenessAction.NAME.equals(action)) {
        assertHeaders(threadPool);
        ((TransportResponseHandler<LivenessResponse>) handler)
            .handleResponse(new LivenessResponse(clusterName, node));
        return;
      }
      if (ClusterStateAction.NAME.equals(action)) {
        assertHeaders(threadPool);
        ClusterName cluster1 = new ClusterName("cluster1");
        ClusterState.Builder builder = ClusterState.builder(cluster1);
        // the sniffer detects only data nodes
        builder.nodes(
            DiscoveryNodes.builder()
                .put(
                    new DiscoveryNode(
                        "node_id",
                        address,
                        Collections.emptyMap(),
                        Collections.singleton(DiscoveryNode.Role.DATA),
                        Version.CURRENT)));
        ((TransportResponseHandler<ClusterStateResponse>) handler)
            .handleResponse(new ClusterStateResponse(cluster1, builder.build()));
        clusterStateLatch.countDown();
        return;
      }

      handler.handleException(new TransportException("", new InternalException(action)));
    }
  /**
   * Creates cluster state with and index that has one shard and #(replicaStates) replicas
   *
   * @param index name of the index
   * @param activePrimaryLocal if active primary should coincide with the local node in the cluster
   *     state
   * @param primaryState state of primary
   * @param replicaStates states of the replicas. length of this array determines also the number of
   *     replicas
   */
  public static ClusterState state(
      String index,
      boolean activePrimaryLocal,
      ShardRoutingState primaryState,
      ShardRoutingState... replicaStates) {
    final int numberOfReplicas = replicaStates.length;

    int numberOfNodes = numberOfReplicas + 1;
    if (primaryState == ShardRoutingState.RELOCATING) {
      numberOfNodes++;
    }
    for (ShardRoutingState state : replicaStates) {
      if (state == ShardRoutingState.RELOCATING) {
        numberOfNodes++;
      }
    }
    numberOfNodes = Math.max(2, numberOfNodes); // we need a non-local master to test shard failures
    final ShardId shardId = new ShardId(index, "_na_", 0);
    DiscoveryNodes.Builder discoBuilder = DiscoveryNodes.builder();
    Set<String> unassignedNodes = new HashSet<>();
    for (int i = 0; i < numberOfNodes + 1; i++) {
      final DiscoveryNode node = newNode(i);
      discoBuilder = discoBuilder.put(node);
      unassignedNodes.add(node.getId());
    }
    discoBuilder.localNodeId(newNode(0).getId());
    discoBuilder.masterNodeId(
        newNode(1).getId()); // we need a non-local master to test shard failures
    final int primaryTerm = 1 + randomInt(200);
    IndexMetaData indexMetaData =
        IndexMetaData.builder(index)
            .settings(
                Settings.builder()
                    .put(SETTING_VERSION_CREATED, Version.CURRENT)
                    .put(SETTING_NUMBER_OF_SHARDS, 1)
                    .put(SETTING_NUMBER_OF_REPLICAS, numberOfReplicas)
                    .put(SETTING_CREATION_DATE, System.currentTimeMillis()))
            .primaryTerm(0, primaryTerm)
            .build();

    RoutingTable.Builder routing = new RoutingTable.Builder();
    routing.addAsNew(indexMetaData);
    IndexShardRoutingTable.Builder indexShardRoutingBuilder =
        new IndexShardRoutingTable.Builder(shardId);

    String primaryNode = null;
    String relocatingNode = null;
    UnassignedInfo unassignedInfo = null;
    if (primaryState != ShardRoutingState.UNASSIGNED) {
      if (activePrimaryLocal) {
        primaryNode = newNode(0).getId();
        unassignedNodes.remove(primaryNode);
      } else {
        Set<String> unassignedNodesExecludingPrimary = new HashSet<>(unassignedNodes);
        unassignedNodesExecludingPrimary.remove(newNode(0).getId());
        primaryNode = selectAndRemove(unassignedNodesExecludingPrimary);
      }
      if (primaryState == ShardRoutingState.RELOCATING) {
        relocatingNode = selectAndRemove(unassignedNodes);
      }
    } else {
      unassignedInfo = new UnassignedInfo(UnassignedInfo.Reason.INDEX_CREATED, null);
    }
    indexShardRoutingBuilder.addShard(
        TestShardRouting.newShardRouting(
            index, 0, primaryNode, relocatingNode, null, true, primaryState, unassignedInfo));

    for (ShardRoutingState replicaState : replicaStates) {
      String replicaNode = null;
      relocatingNode = null;
      unassignedInfo = null;
      if (replicaState != ShardRoutingState.UNASSIGNED) {
        assert primaryNode != null : "a replica is assigned but the primary isn't";
        replicaNode = selectAndRemove(unassignedNodes);
        if (replicaState == ShardRoutingState.RELOCATING) {
          relocatingNode = selectAndRemove(unassignedNodes);
        }
      } else {
        unassignedInfo = new UnassignedInfo(UnassignedInfo.Reason.INDEX_CREATED, null);
      }
      indexShardRoutingBuilder.addShard(
          TestShardRouting.newShardRouting(
              index,
              shardId.id(),
              replicaNode,
              relocatingNode,
              null,
              false,
              replicaState,
              unassignedInfo));
    }

    ClusterState.Builder state = ClusterState.builder(new ClusterName("test"));
    state.nodes(discoBuilder);
    state.metaData(MetaData.builder().put(indexMetaData, false).generateClusterUuidIfNeeded());
    state.routingTable(
        RoutingTable.builder()
            .add(
                IndexRoutingTable.builder(indexMetaData.getIndex())
                    .addIndexShard(indexShardRoutingBuilder.build()))
            .build());
    return state.build();
  }
 @Override
 public void writeTo(StreamOutput out) throws IOException {
   ClusterState.Builder.writeTo(clusterState, out);
 }
 @Override
 public void readFrom(StreamInput in) throws IOException {
   clusterState = ClusterState.Builder.readFrom(in, nodesProvider.nodes().localNode());
 }
 @Override
 public void readFrom(StreamInput in) throws IOException {
   super.readFrom(in);
   clusterName = new ClusterName(in);
   clusterState = ClusterState.Builder.readFrom(in, null);
 }
  @Test
  public void tesStartedShardsMatching() {
    AllocationService allocation = createAllocationService();

    logger.info("--> building initial cluster state");
    final IndexMetaData indexMetaData =
        IndexMetaData.builder("test").numberOfShards(3).numberOfReplicas(0).build();
    ClusterState.Builder stateBuilder =
        ClusterState.builder(ClusterName.DEFAULT)
            .nodes(DiscoveryNodes.builder().put(newNode("node1")).put(newNode("node2")))
            .metaData(MetaData.builder().put(indexMetaData, false));

    final ImmutableShardRouting initShard =
        new ImmutableShardRouting(
            "test", 0, "node1", randomBoolean(), ShardRoutingState.INITIALIZING, 1);
    final ImmutableShardRouting startedShard =
        new ImmutableShardRouting(
            "test", 1, "node2", randomBoolean(), ShardRoutingState.STARTED, 1);
    final ImmutableShardRouting relocatingShard =
        new ImmutableShardRouting(
            "test", 2, "node1", "node2", randomBoolean(), ShardRoutingState.RELOCATING, 1);
    stateBuilder.routingTable(
        RoutingTable.builder()
            .add(
                IndexRoutingTable.builder("test")
                    .addIndexShard(
                        new IndexShardRoutingTable.Builder(initShard.shardId(), true)
                            .addShard(initShard)
                            .build())
                    .addIndexShard(
                        new IndexShardRoutingTable.Builder(startedShard.shardId(), true)
                            .addShard(startedShard)
                            .build())
                    .addIndexShard(
                        new IndexShardRoutingTable.Builder(relocatingShard.shardId(), true)
                            .addShard(relocatingShard)
                            .build())));

    ClusterState state = stateBuilder.build();

    logger.info("--> test starting of shard");

    RoutingAllocation.Result result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    initShard.index(),
                    initShard.id(),
                    initShard.currentNodeId(),
                    initShard.relocatingNodeId(),
                    initShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    randomInt())),
            false);
    assertTrue(
        "failed to start "
            + initShard
            + "\ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());
    assertTrue(
        initShard + "isn't started \ncurrent routing table:" + result.routingTable().prettyPrint(),
        result
                .routingTable()
                .index("test")
                .shard(initShard.id())
                .countWithState(ShardRoutingState.STARTED)
            == 1);

    logger.info("--> testing shard variants that shouldn't match the started shard");

    result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    initShard.index(),
                    initShard.id(),
                    initShard.currentNodeId(),
                    initShard.relocatingNodeId(),
                    !initShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    1)),
            false);
    assertFalse(
        "wrong primary flag shouldn't start shard "
            + initShard
            + "\ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());

    result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    initShard.index(),
                    initShard.id(),
                    "some_node",
                    initShard.currentNodeId(),
                    initShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    1)),
            false);
    assertFalse(
        "relocating shard from node shouldn't start shard "
            + initShard
            + "\ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());

    result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    initShard.index(),
                    initShard.id(),
                    initShard.currentNodeId(),
                    "some_node",
                    initShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    1)),
            false);
    assertFalse(
        "relocating shard to node shouldn't start shard "
            + initShard
            + "\ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());

    logger.info("--> testing double starting");

    result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    startedShard.index(),
                    startedShard.id(),
                    startedShard.currentNodeId(),
                    startedShard.relocatingNodeId(),
                    startedShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    1)),
            false);
    assertFalse(
        "duplicate starting of the same shard should be ignored \ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());

    logger.info("--> testing starting of relocating shards");
    result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    relocatingShard.index(),
                    relocatingShard.id(),
                    relocatingShard.relocatingNodeId(),
                    relocatingShard.currentNodeId(),
                    relocatingShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    randomInt())),
            false);
    assertTrue(
        "failed to start "
            + relocatingShard
            + "\ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());
    ShardRouting shardRouting =
        result.routingTable().index("test").shard(relocatingShard.id()).getShards().get(0);
    assertThat(shardRouting.state(), equalTo(ShardRoutingState.STARTED));
    assertThat(shardRouting.currentNodeId(), equalTo("node2"));
    assertThat(shardRouting.relocatingNodeId(), nullValue());

    logger.info("--> testing shard variants that shouldn't match the relocating shard");

    result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    relocatingShard.index(),
                    relocatingShard.id(),
                    relocatingShard.relocatingNodeId(),
                    relocatingShard.currentNodeId(),
                    !relocatingShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    1)),
            false);
    assertFalse(
        "wrong primary flag shouldn't start shard "
            + relocatingShard
            + "\ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());

    result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    relocatingShard.index(),
                    relocatingShard.id(),
                    "some_node",
                    relocatingShard.currentNodeId(),
                    relocatingShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    1)),
            false);
    assertFalse(
        "relocating shard to a different node shouldn't start shard "
            + relocatingShard
            + "\ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());

    result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    relocatingShard.index(),
                    relocatingShard.id(),
                    relocatingShard.relocatingNodeId(),
                    "some_node",
                    relocatingShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    1)),
            false);
    assertFalse(
        "relocating shard from a different node shouldn't start shard "
            + relocatingShard
            + "\ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());

    result =
        allocation.applyStartedShards(
            state,
            Arrays.asList(
                new ImmutableShardRouting(
                    relocatingShard.index(),
                    relocatingShard.id(),
                    relocatingShard.relocatingNodeId(),
                    relocatingShard.primary(),
                    ShardRoutingState.INITIALIZING,
                    1)),
            false);
    assertFalse(
        "non-relocating shard shouldn't start shard"
            + relocatingShard
            + "\ncurrent routing table:"
            + result.routingTable().prettyPrint(),
        result.changed());
  }
 private static ClusterState state(ClusterName clusterName) {
   ClusterState.Builder builder = ClusterState.builder(clusterName);
   builder.nodes(
       DiscoveryNodes.builder().put(new DiscoveryNode("node_id", address, Version.CURRENT)));
   return builder.build();
 }