@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();
}
}