@Test
  public void testKill() throws Exception {
    ClusterService clusterService = mock(ClusterService.class);
    OperationRouting operationRouting = mock(OperationRouting.class);

    mockShard(operationRouting, 1);
    mockShard(operationRouting, 2);
    mockShard(operationRouting, 3);
    when(clusterService.operationRouting()).thenReturn(operationRouting);

    final AtomicReference<ActionListener<ShardUpsertResponse>> ref = new AtomicReference<>();
    SymbolBasedTransportShardUpsertActionDelegate transportShardUpsertActionDelegate =
        new SymbolBasedTransportShardUpsertActionDelegate() {
          @Override
          public void execute(
              SymbolBasedShardUpsertRequest request, ActionListener<ShardUpsertResponse> listener) {
            ref.set(listener);
          }
        };

    TransportActionProvider transportActionProvider =
        mock(TransportActionProvider.class, Answers.RETURNS_DEEP_STUBS.get());
    when(transportActionProvider.symbolBasedTransportShardUpsertActionDelegate())
        .thenReturn(transportShardUpsertActionDelegate);

    BulkRetryCoordinator bulkRetryCoordinator = new BulkRetryCoordinator(ImmutableSettings.EMPTY);
    BulkRetryCoordinatorPool coordinatorPool = mock(BulkRetryCoordinatorPool.class);
    when(coordinatorPool.coordinator(any(ShardId.class))).thenReturn(bulkRetryCoordinator);

    SymbolBasedShardUpsertRequest.Builder builder =
        new SymbolBasedShardUpsertRequest.Builder(
            TimeValue.timeValueMillis(10),
            false,
            false,
            null,
            new Reference[] {fooRef},
            UUID.randomUUID());

    final SymbolBasedBulkShardProcessor<SymbolBasedShardUpsertRequest, ShardUpsertResponse>
        bulkShardProcessor =
            new SymbolBasedBulkShardProcessor<>(
                clusterService,
                mock(TransportBulkCreateIndicesAction.class),
                ImmutableSettings.EMPTY,
                coordinatorPool,
                false,
                1,
                builder,
                transportShardUpsertActionDelegate,
                UUID.randomUUID());
    assertThat(bulkShardProcessor.add("foo", "1", new Object[] {"bar1"}, null, null), is(true));
    bulkShardProcessor.kill(new CancellationException());
    // A CancellationException is thrown
    expectedException.expect(CancellationException.class);
    bulkShardProcessor.result().get();
    // it's not possible to add more
    assertThat(bulkShardProcessor.add("foo", "1", new Object[] {"bar1"}, null, null), is(false));
  }
  @Test
  public void testThatAddAfterFailureBlocksDueToRetry() throws Exception {
    ClusterService clusterService = mock(ClusterService.class);
    OperationRouting operationRouting = mock(OperationRouting.class);

    mockShard(operationRouting, 1);
    mockShard(operationRouting, 2);
    mockShard(operationRouting, 3);
    when(clusterService.operationRouting()).thenReturn(operationRouting);

    // listener will be executed 2 times, once for the successfully added row and once for the
    // failure
    final CountDownLatch listenerLatch = new CountDownLatch(2);
    final AtomicReference<ActionListener<ShardUpsertResponse>> ref = new AtomicReference<>();
    SymbolBasedTransportShardUpsertActionDelegate transportShardUpsertActionDelegate =
        new SymbolBasedTransportShardUpsertActionDelegate() {
          @Override
          public void execute(
              SymbolBasedShardUpsertRequest request, ActionListener<ShardUpsertResponse> listener) {
            ref.set(listener);
            listenerLatch.countDown();
          }
        };

    TransportActionProvider transportActionProvider =
        mock(TransportActionProvider.class, Answers.RETURNS_DEEP_STUBS.get());
    when(transportActionProvider.symbolBasedTransportShardUpsertActionDelegate())
        .thenReturn(transportShardUpsertActionDelegate);

    BulkRetryCoordinator bulkRetryCoordinator = new BulkRetryCoordinator(ImmutableSettings.EMPTY);
    BulkRetryCoordinatorPool coordinatorPool = mock(BulkRetryCoordinatorPool.class);
    when(coordinatorPool.coordinator(any(ShardId.class))).thenReturn(bulkRetryCoordinator);

    SymbolBasedShardUpsertRequest.Builder builder =
        new SymbolBasedShardUpsertRequest.Builder(
            TimeValue.timeValueMillis(10),
            false,
            false,
            null,
            new Reference[] {fooRef},
            UUID.randomUUID());

    final SymbolBasedBulkShardProcessor<SymbolBasedShardUpsertRequest, ShardUpsertResponse>
        bulkShardProcessor =
            new SymbolBasedBulkShardProcessor<>(
                clusterService,
                mock(TransportBulkCreateIndicesAction.class),
                ImmutableSettings.EMPTY,
                coordinatorPool,
                false,
                1,
                builder,
                transportShardUpsertActionDelegate,
                UUID.randomUUID());
    bulkShardProcessor.add("foo", "1", new Object[] {"bar1"}, null, null);
    final ActionListener<ShardUpsertResponse> listener = ref.get();

    listener.onFailure(new EsRejectedExecutionException());
    // wait, failure retry lock is done in decoupled thread
    listenerLatch.await(10, TimeUnit.SECONDS);

    final ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(2);
    try {
      final AtomicBoolean hadBlocked = new AtomicBoolean(false);
      final AtomicBoolean hasBlocked = new AtomicBoolean(true);
      final CountDownLatch latch = new CountDownLatch(1);
      scheduledExecutorService.execute(
          new Runnable() {
            @Override
            public void run() {
              scheduledExecutorService.schedule(
                  new Runnable() {
                    @Override
                    public void run() {
                      hadBlocked.set(hasBlocked.get());
                      latch.countDown();
                    }
                  },
                  10,
                  TimeUnit.MILLISECONDS);
              bulkShardProcessor.add("foo", "2", new Object[] {"bar2"}, null, null);
              hasBlocked.set(false);
            }
          });
      latch.await();
      assertTrue(hadBlocked.get());
    } finally {
      scheduledExecutorService.shutdownNow();
      bulkRetryCoordinator.close();
    }
  }