@After
 public void tearDown() throws Exception {
   super.tearDown();
   serviceA.close();
   serviceB.close();
   threadPool.shutdown();
 }
  @Test
  public void testErrorMessage() {
    serviceA.registerHandler(
        "sayHelloException",
        new BaseTransportRequestHandler<StringMessageRequest>() {
          @Override
          public StringMessageRequest newInstance() {
            return new StringMessageRequest();
          }

          @Override
          public String executor() {
            return ThreadPool.Names.GENERIC;
          }

          @Override
          public void messageReceived(StringMessageRequest request, TransportChannel channel)
              throws Exception {
            assertThat("moshe", equalTo(request.message));
            throw new RuntimeException("bad message !!!");
          }
        });

    TransportFuture<StringMessageResponse> res =
        serviceB.submitRequest(
            nodeA,
            "sayHelloException",
            new StringMessageRequest("moshe"),
            new BaseTransportResponseHandler<StringMessageResponse>() {
              @Override
              public StringMessageResponse newInstance() {
                return new StringMessageResponse();
              }

              @Override
              public String executor() {
                return ThreadPool.Names.GENERIC;
              }

              @Override
              public void handleResponse(StringMessageResponse response) {
                assertThat("got response instead of exception", false, equalTo(true));
              }

              @Override
              public void handleException(TransportException exp) {
                assertThat("bad message !!!", equalTo(exp.getCause().getMessage()));
              }
            });

    try {
      res.txGet();
      assertThat("exception should be thrown", false, equalTo(true));
    } catch (Exception e) {
      assertThat("bad message !!!", equalTo(e.getCause().getMessage()));
    }

    serviceA.removeHandler("sayHelloException");
  }
  @Test
  public void testVersion_from0to0() throws Exception {
    serviceA.registerHandler(
        "/version",
        new BaseTransportRequestHandler<Version0Request>() {
          @Override
          public Version0Request newInstance() {
            return new Version0Request();
          }

          @Override
          public void messageReceived(Version0Request request, TransportChannel channel)
              throws Exception {
            assertThat(request.value1, equalTo(1));
            Version0Response response = new Version0Response();
            response.value1 = 1;
            channel.sendResponse(response);
          }

          @Override
          public String executor() {
            return ThreadPool.Names.SAME;
          }
        });

    Version0Request version0Request = new Version0Request();
    version0Request.value1 = 1;
    Version0Response version0Response =
        serviceA
            .submitRequest(
                nodeA,
                "/version",
                version0Request,
                new BaseTransportResponseHandler<Version0Response>() {
                  @Override
                  public Version0Response newInstance() {
                    return new Version0Response();
                  }

                  @Override
                  public void handleResponse(Version0Response response) {
                    assertThat(response.value1, equalTo(1));
                  }

                  @Override
                  public void handleException(TransportException exp) {
                    exp.printStackTrace();
                    fail();
                  }

                  @Override
                  public String executor() {
                    return ThreadPool.Names.SAME;
                  }
                })
            .txGet();

    assertThat(version0Response.value1, equalTo(1));
  }
 private static MockTransportService buildTransportService(
     Settings settings, ThreadPool threadPool) {
   MockTransportService transportService =
       MockTransportService.createNewService(settings, Version.CURRENT, threadPool, null);
   transportService.start();
   transportService.acceptIncomingRequests();
   return transportService;
 }
 protected void removeDisruption(
     DiscoveryNode node1,
     MockTransportService transportService1,
     DiscoveryNode node2,
     MockTransportService transportService2) {
   transportService1.clearRule(node2);
   transportService2.clearRule(node1);
 }
  public static MockNode createMockNode(
      String name,
      final Settings basSettings,
      @Nullable ClusterStateListener listener,
      ThreadPool threadPool,
      Logger logger,
      Map<String, MockNode> nodes)
      throws Exception {
    final Settings settings =
        Settings.builder()
            .put("name", name)
            .put(
                TransportService.TRACE_LOG_INCLUDE_SETTING.getKey(), "",
                TransportService.TRACE_LOG_EXCLUDE_SETTING.getKey(), "NOTHING")
            .put(basSettings)
            .build();

    MockTransportService service = buildTransportService(settings, threadPool);
    DiscoveryNode discoveryNode =
        DiscoveryNode.createLocal(
            settings,
            service.boundAddress().publishAddress(),
            NodeEnvironment.generateNodeId(settings));
    MockNode node = new MockNode(discoveryNode, service, listener, logger);
    node.action = buildPublishClusterStateAction(settings, service, () -> node.clusterState, node);
    final CountDownLatch latch = new CountDownLatch(nodes.size() * 2 + 1);
    TransportConnectionListener waitForConnection =
        new TransportConnectionListener() {
          @Override
          public void onNodeConnected(DiscoveryNode node) {
            latch.countDown();
          }

          @Override
          public void onNodeDisconnected(DiscoveryNode node) {
            fail("disconnect should not be called " + node);
          }
        };
    node.service.addConnectionListener(waitForConnection);
    for (MockNode curNode : nodes.values()) {
      curNode.service.addConnectionListener(waitForConnection);
      curNode.connectTo(node.discoveryNode);
      node.connectTo(curNode.discoveryNode);
    }
    node.connectTo(node.discoveryNode);
    assertThat(
        "failed to wait for all nodes to connect", latch.await(5, TimeUnit.SECONDS), equalTo(true));
    for (MockNode curNode : nodes.values()) {
      curNode.service.removeConnectionListener(waitForConnection);
    }
    node.service.removeConnectionListener(waitForConnection);
    if (nodes.put(name, node) != null) {
      fail("Node with the name " + name + " already exist");
    }
    return node;
  }
  @Test
  public void testNotifyOnShutdown() throws Exception {
    final CountDownLatch latch2 = new CountDownLatch(1);

    serviceA.registerHandler(
        "foobar",
        new BaseTransportRequestHandler<StringMessageRequest>() {
          @Override
          public StringMessageRequest newInstance() {
            return new StringMessageRequest();
          }

          @Override
          public String executor() {
            return ThreadPool.Names.GENERIC;
          }

          @Override
          public void messageReceived(StringMessageRequest request, TransportChannel channel) {

            try {
              latch2.await();
              logger.info("Stop ServiceB now");
              serviceB.stop();
            } catch (Exception e) {
              fail(e.getMessage());
            }
          }
        });
    TransportFuture<TransportResponse.Empty> foobar =
        serviceB.submitRequest(
            nodeA,
            "foobar",
            new StringMessageRequest(""),
            options(),
            EmptyTransportResponseHandler.INSTANCE_SAME);
    latch2.countDown();
    try {
      foobar.txGet();
      fail("TransportException expected");
    } catch (TransportException ex) {

    }
    serviceA.removeHandler("sayHelloTimeoutDelayedResponse");
  }
  @Test
  public void testDisconnectListener() throws Exception {
    final CountDownLatch latch = new CountDownLatch(1);
    TransportConnectionListener disconnectListener =
        new TransportConnectionListener() {
          @Override
          public void onNodeConnected(DiscoveryNode node) {
            fail(
                "node connected should not be called, all connection have been done previously, node: "
                    + node);
          }

          @Override
          public void onNodeDisconnected(DiscoveryNode node) {
            latch.countDown();
          }
        };
    serviceA.addConnectionListener(disconnectListener);
    serviceB.close();
    assertThat(latch.await(5, TimeUnit.SECONDS), equalTo(true));
  }
  @Before
  public void setUp() throws Exception {
    super.setUp();
    threadPool = new ThreadPool(getClass().getName());
    serviceA = build(ImmutableSettings.builder().put("name", "TS_A").build(), version0);
    nodeA =
        new DiscoveryNode(
            "TS_A",
            "TS_A",
            serviceA.boundAddress().publishAddress(),
            ImmutableMap.<String, String>of(),
            version0);
    serviceB = build(ImmutableSettings.builder().put("name", "TS_B").build(), version1);
    nodeB =
        new DiscoveryNode(
            "TS_B",
            "TS_B",
            serviceB.boundAddress().publishAddress(),
            ImmutableMap.<String, String>of(),
            version1);

    // wait till all nodes are properly connected and the event has been sent, so tests in this
    // class
    // will not get this callback called on the connections done in this setup
    final CountDownLatch latch = new CountDownLatch(4);
    TransportConnectionListener waitForConnection =
        new TransportConnectionListener() {
          @Override
          public void onNodeConnected(DiscoveryNode node) {
            latch.countDown();
          }

          @Override
          public void onNodeDisconnected(DiscoveryNode node) {
            fail("disconnect should not be called " + node);
          }
        };
    serviceA.addConnectionListener(waitForConnection);
    serviceB.addConnectionListener(waitForConnection);

    serviceA.connectToNode(nodeB);
    serviceA.connectToNode(nodeA);
    serviceB.connectToNode(nodeA);
    serviceB.connectToNode(nodeB);

    assertThat(
        "failed to wait for all nodes to connect", latch.await(5, TimeUnit.SECONDS), equalTo(true));
    serviceA.removeConnectionListener(waitForConnection);
    serviceB.removeConnectionListener(waitForConnection);
  }
  /**
   * This test triggers a corrupt index exception during finalization size if an empty commit point
   * is transferred during recovery we don't know the version of the segments_N file because it has
   * no segments we can take it from. This simulates recoveries from old indices or even without
   * checksums and makes sure if we fail during finalization we also check if the primary is ok.
   * Without the relevant checks this test fails with a RED cluster
   */
  public void testCorruptionOnNetworkLayerFinalizingRecovery()
      throws ExecutionException, InterruptedException, IOException {
    internalCluster().ensureAtLeastNumDataNodes(2);
    NodesStatsResponse nodeStats = client().admin().cluster().prepareNodesStats().get();
    List<NodeStats> dataNodeStats = new ArrayList<>();
    for (NodeStats stat : nodeStats.getNodes()) {
      if (stat.getNode().isDataNode()) {
        dataNodeStats.add(stat);
      }
    }

    assertThat(dataNodeStats.size(), greaterThanOrEqualTo(2));
    Collections.shuffle(dataNodeStats, random());
    NodeStats primariesNode = dataNodeStats.get(0);
    NodeStats unluckyNode = dataNodeStats.get(1);
    assertAcked(
        prepareCreate("test")
            .setSettings(
                Settings.builder()
                    .put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, "0")
                    .put(IndexMetaData.SETTING_NUMBER_OF_SHARDS, 1)
                    .put(
                        "index.routing.allocation.include._name", primariesNode.getNode().getName())
                    .put(
                        EnableAllocationDecider.INDEX_ROUTING_REBALANCE_ENABLE_SETTING.getKey(),
                        EnableAllocationDecider.Rebalance.NONE)
                    .put("index.allocation.max_retries", Integer.MAX_VALUE) // keep on retrying
                ));
    ensureGreen(); // allocated with empty commit
    final AtomicBoolean corrupt = new AtomicBoolean(true);
    final CountDownLatch hasCorrupted = new CountDownLatch(1);
    for (NodeStats dataNode : dataNodeStats) {
      MockTransportService mockTransportService =
          ((MockTransportService)
              internalCluster().getInstance(TransportService.class, dataNode.getNode().getName()));
      mockTransportService.addDelegate(
          internalCluster().getInstance(TransportService.class, unluckyNode.getNode().getName()),
          new MockTransportService.DelegateTransport(mockTransportService.original()) {

            @Override
            public void sendRequest(
                DiscoveryNode node,
                long requestId,
                String action,
                TransportRequest request,
                TransportRequestOptions options)
                throws IOException, TransportException {
              if (corrupt.get() && action.equals(RecoveryTargetService.Actions.FILE_CHUNK)) {
                RecoveryFileChunkRequest req = (RecoveryFileChunkRequest) request;
                byte[] array = BytesRef.deepCopyOf(req.content().toBytesRef()).bytes;
                int i = randomIntBetween(0, req.content().length() - 1);
                array[i] = (byte) ~array[i]; // flip one byte in the content
                hasCorrupted.countDown();
              }
              super.sendRequest(node, requestId, action, request, options);
            }
          });
    }

    Settings build =
        Settings.builder()
            .put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, "1")
            .put(
                "index.routing.allocation.include._name",
                primariesNode.getNode().getName() + "," + unluckyNode.getNode().getName())
            .build();
    client().admin().indices().prepareUpdateSettings("test").setSettings(build).get();
    client().admin().cluster().prepareReroute().get();
    hasCorrupted.await();
    corrupt.set(false);
    ensureGreen();
  }
  /* Test that shard is deleted in case ShardActiveRequest after relocation and next incoming cluster state is an index delete. */
  public void testShardCleanupIfShardDeletionAfterRelocationFailedAndIndexDeleted()
      throws Exception {
    final String node_1 = internalCluster().startNode();
    logger.info("--> creating index [test] with one shard and on replica");
    assertAcked(
        prepareCreate("test")
            .setSettings(
                Settings.builder()
                    .put(indexSettings())
                    .put(IndexMetaData.SETTING_NUMBER_OF_SHARDS, 1)
                    .put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, 0)));
    ensureGreen("test");
    ClusterState state = client().admin().cluster().prepareState().get().getState();
    Index index = state.metaData().index("test").getIndex();
    assertThat(Files.exists(shardDirectory(node_1, index, 0)), equalTo(true));
    assertThat(Files.exists(indexDirectory(node_1, index)), equalTo(true));

    final String node_2 = internalCluster().startDataOnlyNode(Settings.builder().build());
    assertFalse(client().admin().cluster().prepareHealth().setWaitForNodes("2").get().isTimedOut());

    assertThat(Files.exists(shardDirectory(node_1, index, 0)), equalTo(true));
    assertThat(Files.exists(indexDirectory(node_1, index)), equalTo(true));
    assertThat(Files.exists(shardDirectory(node_2, index, 0)), equalTo(false));
    assertThat(Files.exists(indexDirectory(node_2, index)), equalTo(false));

    // add a transport delegate that will prevent the shard active request to succeed the first time
    // after relocation has finished.
    // node_1 will then wait for the next cluster state change before it tries a next attempt to
    // delete the shard.
    MockTransportService transportServiceNode_1 =
        (MockTransportService) internalCluster().getInstance(TransportService.class, node_1);
    TransportService transportServiceNode_2 =
        internalCluster().getInstance(TransportService.class, node_2);
    final CountDownLatch shardActiveRequestSent = new CountDownLatch(1);
    transportServiceNode_1.addDelegate(
        transportServiceNode_2,
        new MockTransportService.DelegateTransport(transportServiceNode_1.original()) {
          @Override
          public void sendRequest(
              DiscoveryNode node,
              long requestId,
              String action,
              TransportRequest request,
              TransportRequestOptions options)
              throws IOException, TransportException {
            if (action.equals("internal:index/shard/exists")
                && shardActiveRequestSent.getCount() > 0) {
              shardActiveRequestSent.countDown();
              logger.info("prevent shard active request from being sent");
              throw new ConnectTransportException(node, "DISCONNECT: simulated");
            }
            super.sendRequest(node, requestId, action, request, options);
          }
        });

    logger.info("--> move shard from {} to {}, and wait for relocation to finish", node_1, node_2);
    internalCluster()
        .client()
        .admin()
        .cluster()
        .prepareReroute()
        .add(new MoveAllocationCommand("test", 0, node_1, node_2))
        .get();
    shardActiveRequestSent.await();
    ClusterHealthResponse clusterHealth =
        client().admin().cluster().prepareHealth().setWaitForNoRelocatingShards(true).get();
    assertThat(clusterHealth.isTimedOut(), equalTo(false));
    logClusterState();
    // delete the index. node_1 that still waits for the next cluster state update will then get the
    // delete index next.
    // it must still delete the shard, even if it cannot find it anymore in indicesservice
    client().admin().indices().prepareDelete("test").get();

    assertThat(waitForShardDeletion(node_1, index, 0), equalTo(false));
    assertThat(waitForIndexDeletion(node_1, index), equalTo(false));
    assertThat(Files.exists(shardDirectory(node_1, index, 0)), equalTo(false));
    assertThat(Files.exists(indexDirectory(node_1, index)), equalTo(false));
    assertThat(waitForShardDeletion(node_2, index, 0), equalTo(false));
    assertThat(waitForIndexDeletion(node_2, index), equalTo(false));
    assertThat(Files.exists(shardDirectory(node_2, index, 0)), equalTo(false));
    assertThat(Files.exists(indexDirectory(node_2, index)), equalTo(false));
  }
  public void testIndexCleanup() throws Exception {
    final String masterNode =
        internalCluster().startNode(Settings.builder().put(Node.NODE_DATA_SETTING.getKey(), false));
    final String node_1 =
        internalCluster()
            .startNode(Settings.builder().put(Node.NODE_MASTER_SETTING.getKey(), false));
    final String node_2 =
        internalCluster()
            .startNode(Settings.builder().put(Node.NODE_MASTER_SETTING.getKey(), false));
    logger.info("--> creating index [test] with one shard and on replica");
    assertAcked(
        prepareCreate("test")
            .setSettings(
                Settings.builder()
                    .put(indexSettings())
                    .put(IndexMetaData.SETTING_NUMBER_OF_SHARDS, 1)
                    .put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, 1)));
    ensureGreen("test");
    ClusterState state = client().admin().cluster().prepareState().get().getState();
    Index index = state.metaData().index("test").getIndex();

    logger.info("--> making sure that shard and its replica are allocated on node_1 and node_2");
    assertThat(Files.exists(shardDirectory(node_1, index, 0)), equalTo(true));
    assertThat(Files.exists(indexDirectory(node_1, index)), equalTo(true));
    assertThat(Files.exists(shardDirectory(node_2, index, 0)), equalTo(true));
    assertThat(Files.exists(indexDirectory(node_2, index)), equalTo(true));

    logger.info("--> starting node server3");
    final String node_3 =
        internalCluster()
            .startNode(Settings.builder().put(Node.NODE_MASTER_SETTING.getKey(), false));
    logger.info("--> running cluster_health");
    ClusterHealthResponse clusterHealth =
        client()
            .admin()
            .cluster()
            .prepareHealth()
            .setWaitForNodes("4")
            .setWaitForNoRelocatingShards(true)
            .get();
    assertThat(clusterHealth.isTimedOut(), equalTo(false));

    assertThat(Files.exists(shardDirectory(node_1, index, 0)), equalTo(true));
    assertThat(Files.exists(indexDirectory(node_1, index)), equalTo(true));
    assertThat(Files.exists(shardDirectory(node_2, index, 0)), equalTo(true));
    assertThat(Files.exists(indexDirectory(node_2, index)), equalTo(true));
    assertThat(Files.exists(shardDirectory(node_3, index, 0)), equalTo(false));
    assertThat(Files.exists(indexDirectory(node_3, index)), equalTo(false));

    logger.info("--> move shard from node_1 to node_3, and wait for relocation to finish");

    if (randomBoolean()) { // sometimes add cluster-state delay to trigger observers in
      // IndicesStore.ShardActiveRequestHandler
      SingleNodeDisruption disruption = new BlockClusterStateProcessing(node_3, random());
      internalCluster().setDisruptionScheme(disruption);
      MockTransportService transportServiceNode3 =
          (MockTransportService) internalCluster().getInstance(TransportService.class, node_3);
      CountDownLatch beginRelocationLatch = new CountDownLatch(1);
      CountDownLatch endRelocationLatch = new CountDownLatch(1);
      transportServiceNode3.addTracer(
          new ReclocationStartEndTracer(logger, beginRelocationLatch, endRelocationLatch));
      internalCluster()
          .client()
          .admin()
          .cluster()
          .prepareReroute()
          .add(new MoveAllocationCommand("test", 0, node_1, node_3))
          .get();
      // wait for relocation to start
      beginRelocationLatch.await();
      disruption.startDisrupting();
      // wait for relocation to finish
      endRelocationLatch.await();
      // wait a little so that cluster state observer is registered
      sleep(50);
      disruption.stopDisrupting();
    } else {
      internalCluster()
          .client()
          .admin()
          .cluster()
          .prepareReroute()
          .add(new MoveAllocationCommand("test", 0, node_1, node_3))
          .get();
    }
    clusterHealth =
        client().admin().cluster().prepareHealth().setWaitForNoRelocatingShards(true).get();
    assertThat(clusterHealth.isTimedOut(), equalTo(false));

    assertThat(waitForShardDeletion(node_1, index, 0), equalTo(false));
    assertThat(waitForIndexDeletion(node_1, index), equalTo(false));
    assertThat(Files.exists(shardDirectory(node_2, index, 0)), equalTo(true));
    assertThat(Files.exists(indexDirectory(node_2, index)), equalTo(true));
    assertThat(Files.exists(shardDirectory(node_3, index, 0)), equalTo(true));
    assertThat(Files.exists(indexDirectory(node_3, index)), equalTo(true));
  }
  @Test
  public void testCancellationCleansTempFiles() throws Exception {
    final String indexName = "test";

    final String p_node = internalCluster().startNode();

    client()
        .admin()
        .indices()
        .prepareCreate(indexName)
        .setSettings(
            Settings.builder()
                .put(
                    IndexMetaData.SETTING_NUMBER_OF_SHARDS,
                    1,
                    IndexMetaData.SETTING_NUMBER_OF_REPLICAS,
                    0))
        .get();

    internalCluster().startNodesAsync(2).get();

    List<IndexRequestBuilder> requests = new ArrayList<>();
    int numDocs = scaledRandomIntBetween(25, 250);
    for (int i = 0; i < numDocs; i++) {
      requests.add(client().prepareIndex(indexName, "type").setCreate(true).setSource("{}"));
    }
    indexRandom(true, requests);
    assertFalse(
        client()
            .admin()
            .cluster()
            .prepareHealth()
            .setWaitForNodes("3")
            .setWaitForGreenStatus()
            .get()
            .isTimedOut());
    flush();

    int allowedFailures = randomIntBetween(3, 10);
    logger.info("--> blocking recoveries from primary (allowed failures: [{}])", allowedFailures);
    CountDownLatch corruptionCount = new CountDownLatch(allowedFailures);
    ClusterService clusterService = internalCluster().getInstance(ClusterService.class, p_node);
    MockTransportService mockTransportService =
        (MockTransportService) internalCluster().getInstance(TransportService.class, p_node);
    for (DiscoveryNode node : clusterService.state().nodes()) {
      if (!node.equals(clusterService.localNode())) {
        mockTransportService.addDelegate(
            node, new RecoveryCorruption(mockTransportService.original(), corruptionCount));
      }
    }

    client()
        .admin()
        .indices()
        .prepareUpdateSettings(indexName)
        .setSettings(Settings.builder().put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, 1))
        .get();

    corruptionCount.await();

    logger.info("--> stopping replica assignment");
    assertAcked(
        client()
            .admin()
            .cluster()
            .prepareUpdateSettings()
            .setTransientSettings(
                Settings.builder()
                    .put(EnableAllocationDecider.CLUSTER_ROUTING_ALLOCATION_ENABLE, "none")));

    logger.info("--> wait for all replica shards to be removed, on all nodes");
    assertBusy(
        new Runnable() {
          @Override
          public void run() {
            for (String node : internalCluster().getNodeNames()) {
              if (node.equals(p_node)) {
                continue;
              }
              ClusterState state =
                  client(node).admin().cluster().prepareState().setLocal(true).get().getState();
              assertThat(
                  node + " indicates assigned replicas",
                  state
                      .getRoutingTable()
                      .index(indexName)
                      .shardsWithState(ShardRoutingState.UNASSIGNED)
                      .size(),
                  equalTo(1));
            }
          }
        });

    logger.info("--> verifying no temporary recoveries are left");
    for (String node : internalCluster().getNodeNames()) {
      NodeEnvironment nodeEnvironment = internalCluster().getInstance(NodeEnvironment.class, node);
      for (final Path shardLoc : nodeEnvironment.availableShardPaths(new ShardId(indexName, 0))) {
        if (Files.exists(shardLoc)) {
          assertBusy(
              new Runnable() {
                @Override
                public void run() {
                  try {
                    Files.walkFileTree(
                        shardLoc,
                        new SimpleFileVisitor<Path>() {
                          @Override
                          public FileVisitResult visitFile(Path file, BasicFileAttributes attrs)
                              throws IOException {
                            assertThat(
                                "found a temporary recovery file: " + file,
                                file.getFileName().toString(),
                                not(startsWith("recovery.")));
                            return FileVisitResult.CONTINUE;
                          }
                        });
                  } catch (IOException e) {
                    throw new AssertionError(
                        "failed to walk file tree starting at [" + shardLoc + "]", e);
                  }
                }
              });
        }
      }
    }
  }
  @Test
  public void testHelloWorldCompressed() {
    serviceA.registerHandler(
        "sayHello",
        new BaseTransportRequestHandler<StringMessageRequest>() {
          @Override
          public StringMessageRequest newInstance() {
            return new StringMessageRequest();
          }

          @Override
          public String executor() {
            return ThreadPool.Names.GENERIC;
          }

          @Override
          public void messageReceived(StringMessageRequest request, TransportChannel channel) {
            assertThat("moshe", equalTo(request.message));
            try {
              channel.sendResponse(
                  new StringMessageResponse("hello " + request.message),
                  TransportResponseOptions.options().withCompress(true));
            } catch (IOException e) {
              e.printStackTrace();
              assertThat(e.getMessage(), false, equalTo(true));
            }
          }
        });

    TransportFuture<StringMessageResponse> res =
        serviceB.submitRequest(
            nodeA,
            "sayHello",
            new StringMessageRequest("moshe"),
            TransportRequestOptions.options().withCompress(true),
            new BaseTransportResponseHandler<StringMessageResponse>() {
              @Override
              public StringMessageResponse newInstance() {
                return new StringMessageResponse();
              }

              @Override
              public String executor() {
                return ThreadPool.Names.GENERIC;
              }

              @Override
              public void handleResponse(StringMessageResponse response) {
                assertThat("hello moshe", equalTo(response.message));
              }

              @Override
              public void handleException(TransportException exp) {
                exp.printStackTrace();
                assertThat(
                    "got exception instead of a response: " + exp.getMessage(),
                    false,
                    equalTo(true));
              }
            });

    try {
      StringMessageResponse message = res.get();
      assertThat("hello moshe", equalTo(message.message));
    } catch (Exception e) {
      assertThat(e.getMessage(), false, equalTo(true));
    }

    serviceA.removeHandler("sayHello");
  }
  @Test
  public void testTimeoutSendExceptionWithDelayedResponse() throws Exception {
    serviceA.registerHandler(
        "sayHelloTimeoutDelayedResponse",
        new BaseTransportRequestHandler<StringMessageRequest>() {
          @Override
          public StringMessageRequest newInstance() {
            return new StringMessageRequest();
          }

          @Override
          public String executor() {
            return ThreadPool.Names.GENERIC;
          }

          @Override
          public void messageReceived(StringMessageRequest request, TransportChannel channel) {
            TimeValue sleep = TimeValue.parseTimeValue(request.message, null);
            try {
              Thread.sleep(sleep.millis());
            } catch (InterruptedException e) {
              // ignore
            }
            try {
              channel.sendResponse(new StringMessageResponse("hello " + request.message));
            } catch (IOException e) {
              e.printStackTrace();
              assertThat(e.getMessage(), false, equalTo(true));
            }
          }
        });
    final CountDownLatch latch = new CountDownLatch(1);
    TransportFuture<StringMessageResponse> res =
        serviceB.submitRequest(
            nodeA,
            "sayHelloTimeoutDelayedResponse",
            new StringMessageRequest("300ms"),
            options().withTimeout(100),
            new BaseTransportResponseHandler<StringMessageResponse>() {
              @Override
              public StringMessageResponse newInstance() {
                return new StringMessageResponse();
              }

              @Override
              public String executor() {
                return ThreadPool.Names.GENERIC;
              }

              @Override
              public void handleResponse(StringMessageResponse response) {
                latch.countDown();
                assertThat("got response instead of exception", false, equalTo(true));
              }

              @Override
              public void handleException(TransportException exp) {
                latch.countDown();
                assertThat(exp, instanceOf(ReceiveTimeoutTransportException.class));
              }
            });

    try {
      StringMessageResponse message = res.txGet();
      assertThat("exception should be thrown", false, equalTo(true));
    } catch (Exception e) {
      assertThat(e, instanceOf(ReceiveTimeoutTransportException.class));
    }
    latch.await();

    for (int i = 0; i < 10; i++) {
      final int counter = i;
      // now, try and send another request, this times, with a short timeout
      res =
          serviceB.submitRequest(
              nodeA,
              "sayHelloTimeoutDelayedResponse",
              new StringMessageRequest(counter + "ms"),
              options().withTimeout(3000),
              new BaseTransportResponseHandler<StringMessageResponse>() {
                @Override
                public StringMessageResponse newInstance() {
                  return new StringMessageResponse();
                }

                @Override
                public String executor() {
                  return ThreadPool.Names.GENERIC;
                }

                @Override
                public void handleResponse(StringMessageResponse response) {
                  assertThat("hello " + counter + "ms", equalTo(response.message));
                }

                @Override
                public void handleException(TransportException exp) {
                  exp.printStackTrace();
                  assertThat(
                      "got exception instead of a response for "
                          + counter
                          + ": "
                          + exp.getDetailedMessage(),
                      false,
                      equalTo(true));
                }
              });

      StringMessageResponse message = res.txGet();
      assertThat(message.message, equalTo("hello " + counter + "ms"));
    }

    serviceA.removeHandler("sayHelloTimeoutDelayedResponse");
  }
  @Test
  public void testTimeoutSendExceptionWithNeverSendingBackResponse() throws Exception {
    serviceA.registerHandler(
        "sayHelloTimeoutNoResponse",
        new BaseTransportRequestHandler<StringMessageRequest>() {
          @Override
          public StringMessageRequest newInstance() {
            return new StringMessageRequest();
          }

          @Override
          public String executor() {
            return ThreadPool.Names.GENERIC;
          }

          @Override
          public void messageReceived(StringMessageRequest request, TransportChannel channel) {
            assertThat("moshe", equalTo(request.message));
            // don't send back a response
            //                try {
            //                    channel.sendResponse(new StringMessage("hello " +
            // request.message));
            //                } catch (IOException e) {
            //                    e.printStackTrace();
            //                    assertThat(e.getMessage(), false, equalTo(true));
            //                }
          }
        });

    TransportFuture<StringMessageResponse> res =
        serviceB.submitRequest(
            nodeA,
            "sayHelloTimeoutNoResponse",
            new StringMessageRequest("moshe"),
            options().withTimeout(100),
            new BaseTransportResponseHandler<StringMessageResponse>() {
              @Override
              public StringMessageResponse newInstance() {
                return new StringMessageResponse();
              }

              @Override
              public String executor() {
                return ThreadPool.Names.GENERIC;
              }

              @Override
              public void handleResponse(StringMessageResponse response) {
                assertThat("got response instead of exception", false, equalTo(true));
              }

              @Override
              public void handleException(TransportException exp) {
                assertThat(exp, instanceOf(ReceiveTimeoutTransportException.class));
              }
            });

    try {
      StringMessageResponse message = res.txGet();
      assertThat("exception should be thrown", false, equalTo(true));
    } catch (Exception e) {
      assertThat(e, instanceOf(ReceiveTimeoutTransportException.class));
    }

    serviceA.removeHandler("sayHelloTimeoutNoResponse");
  }
  @Test
  public void testHostOnMessages() throws InterruptedException {
    final CountDownLatch latch = new CountDownLatch(2);
    final AtomicReference<TransportAddress> addressA = new AtomicReference<>();
    final AtomicReference<TransportAddress> addressB = new AtomicReference<>();
    serviceB.registerHandler(
        "action1",
        new TransportRequestHandler<TestRequest>() {
          @Override
          public TestRequest newInstance() {
            return new TestRequest();
          }

          @Override
          public void messageReceived(TestRequest request, TransportChannel channel)
              throws Exception {
            addressA.set(request.remoteAddress());
            channel.sendResponse(new TestResponse());
            latch.countDown();
          }

          @Override
          public String executor() {
            return ThreadPool.Names.SAME;
          }

          @Override
          public boolean isForceExecution() {
            return false;
          }
        });
    serviceA.sendRequest(
        nodeB,
        "action1",
        new TestRequest(),
        new TransportResponseHandler<TestResponse>() {
          @Override
          public TestResponse newInstance() {
            return new TestResponse();
          }

          @Override
          public void handleResponse(TestResponse response) {
            addressB.set(response.remoteAddress());
            latch.countDown();
          }

          @Override
          public void handleException(TransportException exp) {
            latch.countDown();
          }

          @Override
          public String executor() {
            return ThreadPool.Names.SAME;
          }
        });

    if (!latch.await(10, TimeUnit.SECONDS)) {
      fail("message round trip did not complete within a sensible time frame");
    }

    assertTrue(nodeA.address().sameHost(addressA.get()));
    assertTrue(nodeB.address().sameHost(addressB.get()));
  }
  /**
   * Tests corruption that happens on the network layer and that the primary does not get affected
   * by corruption that happens on the way to the replica. The file on disk stays uncorrupted
   */
  public void testCorruptionOnNetworkLayer() throws ExecutionException, InterruptedException {
    int numDocs = scaledRandomIntBetween(100, 1000);
    internalCluster().ensureAtLeastNumDataNodes(2);
    if (cluster().numDataNodes() < 3) {
      internalCluster()
          .startNode(
              Settings.builder()
                  .put(Node.NODE_DATA_SETTING.getKey(), true)
                  .put(Node.NODE_MASTER_SETTING.getKey(), false));
    }
    NodesStatsResponse nodeStats = client().admin().cluster().prepareNodesStats().get();
    List<NodeStats> dataNodeStats = new ArrayList<>();
    for (NodeStats stat : nodeStats.getNodes()) {
      if (stat.getNode().isDataNode()) {
        dataNodeStats.add(stat);
      }
    }

    assertThat(dataNodeStats.size(), greaterThanOrEqualTo(2));
    Collections.shuffle(dataNodeStats, random());
    NodeStats primariesNode = dataNodeStats.get(0);
    NodeStats unluckyNode = dataNodeStats.get(1);

    assertAcked(
        prepareCreate("test")
            .setSettings(
                Settings.builder()
                    .put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, "0")
                    .put(
                        IndexMetaData.SETTING_NUMBER_OF_SHARDS,
                        between(1, 4)) // don't go crazy here it must recovery fast
                    // This does corrupt files on the replica, so we can't check:
                    .put(MockFSIndexStore.INDEX_CHECK_INDEX_ON_CLOSE_SETTING.getKey(), false)
                    .put(
                        "index.routing.allocation.include._name", primariesNode.getNode().getName())
                    .put(
                        EnableAllocationDecider.INDEX_ROUTING_REBALANCE_ENABLE_SETTING.getKey(),
                        EnableAllocationDecider.Rebalance.NONE)));
    ensureGreen();
    IndexRequestBuilder[] builders = new IndexRequestBuilder[numDocs];
    for (int i = 0; i < builders.length; i++) {
      builders[i] = client().prepareIndex("test", "type").setSource("field", "value");
    }
    indexRandom(true, builders);
    ensureGreen();
    assertAllSuccessful(
        client()
            .admin()
            .indices()
            .prepareFlush()
            .setForce(true)
            .setWaitIfOngoing(true)
            .execute()
            .actionGet());
    // we have to flush at least once here since we don't corrupt the translog
    SearchResponse countResponse = client().prepareSearch().setSize(0).get();
    assertHitCount(countResponse, numDocs);
    final boolean truncate = randomBoolean();
    for (NodeStats dataNode : dataNodeStats) {
      MockTransportService mockTransportService =
          ((MockTransportService)
              internalCluster().getInstance(TransportService.class, dataNode.getNode().getName()));
      mockTransportService.addDelegate(
          internalCluster().getInstance(TransportService.class, unluckyNode.getNode().getName()),
          new MockTransportService.DelegateTransport(mockTransportService.original()) {

            @Override
            public void sendRequest(
                DiscoveryNode node,
                long requestId,
                String action,
                TransportRequest request,
                TransportRequestOptions options)
                throws IOException, TransportException {
              if (action.equals(RecoveryTargetService.Actions.FILE_CHUNK)) {
                RecoveryFileChunkRequest req = (RecoveryFileChunkRequest) request;
                if (truncate && req.length() > 1) {
                  BytesRef bytesRef = req.content().toBytesRef();
                  BytesArray array =
                      new BytesArray(bytesRef.bytes, bytesRef.offset, (int) req.length() - 1);
                  request =
                      new RecoveryFileChunkRequest(
                          req.recoveryId(),
                          req.shardId(),
                          req.metadata(),
                          req.position(),
                          array,
                          req.lastChunk(),
                          req.totalTranslogOps(),
                          req.sourceThrottleTimeInNanos());
                } else {
                  assert req.content().toBytesRef().bytes == req.content().toBytesRef().bytes
                      : "no internal reference!!";
                  final byte[] array = req.content().toBytesRef().bytes;
                  int i = randomIntBetween(0, req.content().length() - 1);
                  array[i] = (byte) ~array[i]; // flip one byte in the content
                }
              }
              super.sendRequest(node, requestId, action, request, options);
            }
          });
    }

    Settings build =
        Settings.builder()
            .put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, "1")
            .put("index.routing.allocation.include._name", "*")
            .build();
    client().admin().indices().prepareUpdateSettings("test").setSettings(build).get();
    client().admin().cluster().prepareReroute().get();
    ClusterHealthResponse actionGet =
        client()
            .admin()
            .cluster()
            .health(Requests.clusterHealthRequest("test").waitForGreenStatus())
            .actionGet();
    if (actionGet.isTimedOut()) {
      logger.info(
          "ensureGreen timed out, cluster state:\n{}\n{}",
          client().admin().cluster().prepareState().get().getState().prettyPrint(),
          client().admin().cluster().preparePendingClusterTasks().get().prettyPrint());
      assertThat("timed out waiting for green state", actionGet.isTimedOut(), equalTo(false));
    }
    // we are green so primaries got not corrupted.
    // ensure that no shard is actually allocated on the unlucky node
    ClusterStateResponse clusterStateResponse = client().admin().cluster().prepareState().get();
    for (IndexShardRoutingTable table :
        clusterStateResponse.getState().getRoutingTable().index("test")) {
      for (ShardRouting routing : table) {
        if (unluckyNode.getNode().getId().equals(routing.currentNodeId())) {
          assertThat(routing.state(), not(equalTo(ShardRoutingState.STARTED)));
          assertThat(routing.state(), not(equalTo(ShardRoutingState.RELOCATING)));
        }
      }
    }
    final int numIterations = scaledRandomIntBetween(5, 20);
    for (int i = 0; i < numIterations; i++) {
      SearchResponse response = client().prepareSearch().setSize(numDocs).get();
      assertHitCount(response, numDocs);
    }
  }
  @Test
  public void testMockUnresponsiveRule() {
    serviceA.registerHandler(
        "sayHello",
        new BaseTransportRequestHandler<StringMessageRequest>() {
          @Override
          public StringMessageRequest newInstance() {
            return new StringMessageRequest();
          }

          @Override
          public String executor() {
            return ThreadPool.Names.GENERIC;
          }

          @Override
          public void messageReceived(StringMessageRequest request, TransportChannel channel)
              throws Exception {
            assertThat("moshe", equalTo(request.message));
            throw new RuntimeException("bad message !!!");
          }
        });

    serviceB.addUnresponsiveRule(nodeA);

    TransportFuture<StringMessageResponse> res =
        serviceB.submitRequest(
            nodeA,
            "sayHello",
            new StringMessageRequest("moshe"),
            TransportRequestOptions.options().withTimeout(100),
            new BaseTransportResponseHandler<StringMessageResponse>() {
              @Override
              public StringMessageResponse newInstance() {
                return new StringMessageResponse();
              }

              @Override
              public String executor() {
                return ThreadPool.Names.GENERIC;
              }

              @Override
              public void handleResponse(StringMessageResponse response) {
                assertThat("got response instead of exception", false, equalTo(true));
              }

              @Override
              public void handleException(TransportException exp) {
                assertThat(exp, instanceOf(ReceiveTimeoutTransportException.class));
              }
            });

    try {
      res.txGet();
      assertThat("exception should be thrown", false, equalTo(true));
    } catch (Exception e) {
      assertThat(e, instanceOf(ReceiveTimeoutTransportException.class));
    }

    try {
      serviceB.connectToNode(nodeA);
      assertThat("exception should be thrown", false, equalTo(true));
    } catch (ConnectTransportException e) {
      // all is well
    }

    try {
      serviceB.connectToNodeLight(nodeA);
      assertThat("exception should be thrown", false, equalTo(true));
    } catch (ConnectTransportException e) {
      // all is well
    }

    serviceA.removeHandler("sayHello");
  }
  public void testDisconnectsWhileRecovering() throws Exception {
    final String indexName = "test";
    final Settings nodeSettings =
        Settings.builder()
            .put(RecoverySettings.INDICES_RECOVERY_RETRY_DELAY_NETWORK_SETTING.getKey(), "100ms")
            .put(RecoverySettings.INDICES_RECOVERY_INTERNAL_ACTION_TIMEOUT_SETTING.getKey(), "1s")
            .put(
                MockFSDirectoryService.RANDOM_PREVENT_DOUBLE_WRITE_SETTING.getKey(),
                false) // restarted recoveries will delete temp files and write them again
            .build();
    // start a master node
    internalCluster().startNode(nodeSettings);

    final String blueNodeName =
        internalCluster()
            .startNode(Settings.builder().put("node.attr.color", "blue").put(nodeSettings).build());
    final String redNodeName =
        internalCluster()
            .startNode(Settings.builder().put("node.attr.color", "red").put(nodeSettings).build());

    ClusterHealthResponse response =
        client().admin().cluster().prepareHealth().setWaitForNodes(">=3").get();
    assertThat(response.isTimedOut(), is(false));

    client()
        .admin()
        .indices()
        .prepareCreate(indexName)
        .setSettings(
            Settings.builder()
                .put(IndexMetaData.INDEX_ROUTING_INCLUDE_GROUP_SETTING.getKey() + "color", "blue")
                .put(IndexMetaData.SETTING_NUMBER_OF_SHARDS, 1)
                .put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, 0))
        .get();

    List<IndexRequestBuilder> requests = new ArrayList<>();
    int numDocs = scaledRandomIntBetween(25, 250);
    for (int i = 0; i < numDocs; i++) {
      requests.add(client().prepareIndex(indexName, "type").setSource("{}"));
    }
    indexRandom(true, requests);
    ensureSearchable(indexName);

    ClusterStateResponse stateResponse = client().admin().cluster().prepareState().get();
    final String blueNodeId =
        internalCluster().getInstance(ClusterService.class, blueNodeName).localNode().getId();

    assertFalse(stateResponse.getState().getRoutingNodes().node(blueNodeId).isEmpty());

    SearchResponse searchResponse = client().prepareSearch(indexName).get();
    assertHitCount(searchResponse, numDocs);

    String[] recoveryActions =
        new String[] {
          PeerRecoverySourceService.Actions.START_RECOVERY,
          PeerRecoveryTargetService.Actions.FILES_INFO,
          PeerRecoveryTargetService.Actions.FILE_CHUNK,
          PeerRecoveryTargetService.Actions.CLEAN_FILES,
          // RecoveryTarget.Actions.TRANSLOG_OPS, <-- may not be sent if already flushed
          PeerRecoveryTargetService.Actions.PREPARE_TRANSLOG,
          PeerRecoveryTargetService.Actions.FINALIZE
        };
    final String recoveryActionToBlock = randomFrom(recoveryActions);
    final boolean dropRequests = randomBoolean();
    logger.info(
        "--> will {} between blue & red on [{}]",
        dropRequests ? "drop requests" : "break connection",
        recoveryActionToBlock);

    MockTransportService blueMockTransportService =
        (MockTransportService) internalCluster().getInstance(TransportService.class, blueNodeName);
    MockTransportService redMockTransportService =
        (MockTransportService) internalCluster().getInstance(TransportService.class, redNodeName);
    TransportService redTransportService =
        internalCluster().getInstance(TransportService.class, redNodeName);
    TransportService blueTransportService =
        internalCluster().getInstance(TransportService.class, blueNodeName);
    final CountDownLatch requestBlocked = new CountDownLatch(1);

    blueMockTransportService.addDelegate(
        redTransportService,
        new RecoveryActionBlocker(
            dropRequests,
            recoveryActionToBlock,
            blueMockTransportService.original(),
            requestBlocked));
    redMockTransportService.addDelegate(
        blueTransportService,
        new RecoveryActionBlocker(
            dropRequests,
            recoveryActionToBlock,
            redMockTransportService.original(),
            requestBlocked));

    logger.info("--> starting recovery from blue to red");
    client()
        .admin()
        .indices()
        .prepareUpdateSettings(indexName)
        .setSettings(
            Settings.builder()
                .put(
                    IndexMetaData.INDEX_ROUTING_INCLUDE_GROUP_SETTING.getKey() + "color",
                    "red,blue")
                .put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, 1))
        .get();

    requestBlocked.await();

    logger.info("--> stopping to block recovery");
    blueMockTransportService.clearAllRules();
    redMockTransportService.clearAllRules();

    ensureGreen();
    searchResponse = client(redNodeName).prepareSearch(indexName).setPreference("_local").get();
    assertHitCount(searchResponse, numDocs);
  }
  @Test
  public void testNetworkPartitionDuringReplicaIndexOp() throws Exception {
    final String INDEX = "testidx";

    List<String> nodes = internalCluster().startNodesAsync(2, nodeSettings).get();

    // Create index test with 1 shard, 1 replica and ensure it is green
    createIndex(INDEX);
    ensureGreen(INDEX);

    // Disable allocation so the replica cannot be reallocated when it fails
    Settings s =
        ImmutableSettings.builder().put("cluster.routing.allocation.enable", "none").build();
    client().admin().cluster().prepareUpdateSettings().setTransientSettings(s).get();

    // Determine which node holds the primary shard
    ClusterState state = getNodeClusterState(nodes.get(0));
    IndexShardRoutingTable shard = state.getRoutingTable().index(INDEX).shard(0);
    String primaryNode;
    String replicaNode;
    if (shard.getShards().get(0).primary()) {
      primaryNode = nodes.get(0);
      replicaNode = nodes.get(1);
    } else {
      primaryNode = nodes.get(1);
      replicaNode = nodes.get(0);
    }
    logger.info("--> primary shard is on {}", primaryNode);

    // Index a document to make sure everything works well
    IndexResponse resp =
        internalCluster()
            .client(primaryNode)
            .prepareIndex(INDEX, "doc")
            .setSource("foo", "bar")
            .get();
    assertThat(
        "document exists on primary node",
        internalCluster()
            .client(primaryNode)
            .prepareGet(INDEX, "doc", resp.getId())
            .setPreference("_only_local")
            .get()
            .isExists(),
        equalTo(true));
    assertThat(
        "document exists on replica node",
        internalCluster()
            .client(replicaNode)
            .prepareGet(INDEX, "doc", resp.getId())
            .setPreference("_only_local")
            .get()
            .isExists(),
        equalTo(true));

    // Disrupt the network so indexing requests fail to replicate
    logger.info("--> preventing index/replica operations");
    TransportService mockTransportService =
        internalCluster().getInstance(TransportService.class, primaryNode);
    ((MockTransportService) mockTransportService)
        .addFailToSendNoConnectRule(
            internalCluster().getInstance(Discovery.class, replicaNode).localNode(),
            ImmutableSet.of(IndexAction.NAME + "[r]"));
    mockTransportService = internalCluster().getInstance(TransportService.class, replicaNode);
    ((MockTransportService) mockTransportService)
        .addFailToSendNoConnectRule(
            internalCluster().getInstance(Discovery.class, primaryNode).localNode(),
            ImmutableSet.of(IndexAction.NAME + "[r]"));

    logger.info("--> indexing into primary");
    // the replica shard should now be marked as failed because the replication operation will fail
    resp =
        internalCluster()
            .client(primaryNode)
            .prepareIndex(INDEX, "doc")
            .setSource("foo", "baz")
            .get();
    // wait until the cluster reaches an exact yellow state, meaning replica has failed
    assertBusy(
        new Runnable() {
          @Override
          public void run() {
            assertThat(
                client().admin().cluster().prepareHealth().get().getStatus(),
                equalTo(ClusterHealthStatus.YELLOW));
          }
        });
    assertThat(
        "document should still be indexed and available",
        client().prepareGet(INDEX, "doc", resp.getId()).get().isExists(),
        equalTo(true));

    state = getNodeClusterState(randomFrom(nodes.toArray(Strings.EMPTY_ARRAY)));
    RoutingNodes rn = state.routingNodes();
    logger.info(
        "--> counts: total: {}, unassigned: {}, initializing: {}, relocating: {}, started: {}",
        rn.shards(
                new Predicate<MutableShardRouting>() {
                  @Override
                  public boolean apply(
                      org.elasticsearch.cluster.routing.MutableShardRouting input) {
                    return true;
                  }
                })
            .size(),
        rn.shardsWithState(UNASSIGNED).size(),
        rn.shardsWithState(INITIALIZING).size(),
        rn.shardsWithState(RELOCATING).size(),
        rn.shardsWithState(STARTED).size());
    logger.info(
        "--> unassigned: {}, initializing: {}, relocating: {}, started: {}",
        rn.shardsWithState(UNASSIGNED),
        rn.shardsWithState(INITIALIZING),
        rn.shardsWithState(RELOCATING),
        rn.shardsWithState(STARTED));

    assertThat(
        "only a single shard is now active (replica should be failed and not reallocated)",
        rn.shardsWithState(STARTED).size(),
        equalTo(1));
  }
 public void connectTo(DiscoveryNode node) {
   service.connectToNode(node);
 }
  public void testPrimaryRelocationWhereRecoveryFails() throws Exception {
    Path dataPath = createTempDir();
    Settings nodeSettings =
        Settings.builder()
            .put("node.add_lock_id_to_custom_path", false)
            .put(Environment.PATH_SHARED_DATA_SETTING.getKey(), dataPath)
            .build();

    String node1 = internalCluster().startNode(nodeSettings);
    final String IDX = "test";

    Settings idxSettings =
        Settings.builder()
            .put(IndexMetaData.SETTING_NUMBER_OF_SHARDS, 1)
            .put(IndexMetaData.SETTING_NUMBER_OF_REPLICAS, 1)
            .put(IndexMetaData.SETTING_DATA_PATH, dataPath.toAbsolutePath().toString())
            .put(IndexMetaData.SETTING_SHADOW_REPLICAS, true)
            .put(IndexMetaData.SETTING_SHARED_FILESYSTEM, true)
            .build();

    prepareCreate(IDX).setSettings(idxSettings).addMapping("doc", "foo", "type=text").get();
    // Node1 has the primary, now node2 has the replica
    String node2 = internalCluster().startNode(nodeSettings);
    ensureGreen(IDX);
    flushAndRefresh(IDX);
    String node3 = internalCluster().startNode(nodeSettings);
    final AtomicInteger counter = new AtomicInteger(0);
    final CountDownLatch started = new CountDownLatch(1);

    final int numPhase1Docs = scaledRandomIntBetween(25, 200);
    final int numPhase2Docs = scaledRandomIntBetween(25, 200);
    final int numPhase3Docs = scaledRandomIntBetween(25, 200);
    final CountDownLatch phase1finished = new CountDownLatch(1);
    final CountDownLatch phase2finished = new CountDownLatch(1);
    final CountDownLatch phase3finished = new CountDownLatch(1);

    final AtomicBoolean keepFailing = new AtomicBoolean(true);

    MockTransportService mockTransportService =
        ((MockTransportService) internalCluster().getInstance(TransportService.class, node1));
    mockTransportService.addDelegate(
        internalCluster().getInstance(TransportService.class, node3),
        new MockTransportService.DelegateTransport(mockTransportService.original()) {

          @Override
          public void sendRequest(
              DiscoveryNode node,
              long requestId,
              String action,
              TransportRequest request,
              TransportRequestOptions options)
              throws IOException, TransportException {
            if (keepFailing.get()
                && action.equals(PeerRecoveryTargetService.Actions.TRANSLOG_OPS)) {
              logger.info("--> failing translog ops");
              throw new ElasticsearchException("failing on purpose");
            }
            super.sendRequest(node, requestId, action, request, options);
          }
        });

    Thread thread =
        new Thread() {
          @Override
          public void run() {
            started.countDown();
            while (counter.get() < (numPhase1Docs + numPhase2Docs + numPhase3Docs)) {
              final IndexResponse indexResponse =
                  client()
                      .prepareIndex(IDX, "doc", Integer.toString(counter.incrementAndGet()))
                      .setSource("foo", "bar")
                      .get();
              assertEquals(DocWriteResponse.Result.CREATED, indexResponse.getResult());
              final int docCount = counter.get();
              if (docCount == numPhase1Docs) {
                phase1finished.countDown();
              } else if (docCount == (numPhase1Docs + numPhase2Docs)) {
                phase2finished.countDown();
              }
            }
            logger.info("--> stopping indexing thread");
            phase3finished.countDown();
          }
        };
    thread.start();
    started.await();
    phase1finished.await(); // wait for a certain number of documents to be indexed
    logger.info("--> excluding {} from allocation", node1);
    // now prevent primary from being allocated on node 1 move to node_3
    Settings build =
        Settings.builder().put("index.routing.allocation.exclude._name", node1).build();
    client().admin().indices().prepareUpdateSettings(IDX).setSettings(build).execute().actionGet();
    // wait for more documents to be indexed post-recovery, also waits for
    // indexing thread to stop
    phase2finished.await();
    // stop failing
    keepFailing.set(false);
    // wait for more docs to be indexed
    phase3finished.await();
    ensureGreen(IDX);
    thread.join();
    logger.info("--> performing query");
    flushAndRefresh();

    SearchResponse resp = client().prepareSearch(IDX).setQuery(matchAllQuery()).get();
    assertHitCount(resp, counter.get());
  }
  @Test
  public void testMockFailToSendNoConnectRule() {
    serviceA.registerHandler(
        "sayHello",
        new BaseTransportRequestHandler<StringMessageRequest>() {
          @Override
          public StringMessageRequest newInstance() {
            return new StringMessageRequest();
          }

          @Override
          public String executor() {
            return ThreadPool.Names.GENERIC;
          }

          @Override
          public void messageReceived(StringMessageRequest request, TransportChannel channel)
              throws Exception {
            assertThat("moshe", equalTo(request.message));
            throw new RuntimeException("bad message !!!");
          }
        });

    serviceB.addFailToSendNoConnectRule(nodeA);

    TransportFuture<StringMessageResponse> res =
        serviceB.submitRequest(
            nodeA,
            "sayHello",
            new StringMessageRequest("moshe"),
            new BaseTransportResponseHandler<StringMessageResponse>() {
              @Override
              public StringMessageResponse newInstance() {
                return new StringMessageResponse();
              }

              @Override
              public String executor() {
                return ThreadPool.Names.GENERIC;
              }

              @Override
              public void handleResponse(StringMessageResponse response) {
                assertThat("got response instead of exception", false, equalTo(true));
              }

              @Override
              public void handleException(TransportException exp) {
                assertThat(exp.getCause().getMessage(), endsWith("DISCONNECT: simulated"));
              }
            });

    try {
      res.txGet();
      assertThat("exception should be thrown", false, equalTo(true));
    } catch (Exception e) {
      assertThat(e.getCause().getMessage(), endsWith("DISCONNECT: simulated"));
    }

    try {
      serviceB.connectToNode(nodeA);
      assertThat("exception should be thrown", false, equalTo(true));
    } catch (ConnectTransportException e) {
      // all is well
    }

    try {
      serviceB.connectToNodeLight(nodeA);
      assertThat("exception should be thrown", false, equalTo(true));
    } catch (ConnectTransportException e) {
      // all is well
    }

    serviceA.removeHandler("sayHello");
  }
  @Test
  public void testVoidMessageCompressed() {
    serviceA.registerHandler(
        "sayHello",
        new BaseTransportRequestHandler<TransportRequest.Empty>() {
          @Override
          public TransportRequest.Empty newInstance() {
            return TransportRequest.Empty.INSTANCE;
          }

          @Override
          public String executor() {
            return ThreadPool.Names.GENERIC;
          }

          @Override
          public void messageReceived(TransportRequest.Empty request, TransportChannel channel) {
            try {
              channel.sendResponse(
                  TransportResponse.Empty.INSTANCE,
                  TransportResponseOptions.options().withCompress(true));
            } catch (IOException e) {
              e.printStackTrace();
              assertThat(e.getMessage(), false, equalTo(true));
            }
          }
        });

    TransportFuture<TransportResponse.Empty> res =
        serviceB.submitRequest(
            nodeA,
            "sayHello",
            TransportRequest.Empty.INSTANCE,
            TransportRequestOptions.options().withCompress(true),
            new BaseTransportResponseHandler<TransportResponse.Empty>() {
              @Override
              public TransportResponse.Empty newInstance() {
                return TransportResponse.Empty.INSTANCE;
              }

              @Override
              public String executor() {
                return ThreadPool.Names.GENERIC;
              }

              @Override
              public void handleResponse(TransportResponse.Empty response) {}

              @Override
              public void handleException(TransportException exp) {
                exp.printStackTrace();
                assertThat(
                    "got exception instead of a response: " + exp.getMessage(),
                    false,
                    equalTo(true));
              }
            });

    try {
      TransportResponse.Empty message = res.get();
      assertThat(message, notNullValue());
    } catch (Exception e) {
      assertThat(e.getMessage(), false, equalTo(true));
    }

    serviceA.removeHandler("sayHello");
  }