private Bucket collect(CollectPhase collectNode) throws Throwable {
ContextPreparer contextPreparer = internalCluster().getDataNodeInstance(ContextPreparer.class);
JobContextService contextService =
internalCluster().getDataNodeInstance(JobContextService.class);
SharedShardContexts sharedShardContexts =
new SharedShardContexts(internalCluster().getDataNodeInstance(IndicesService.class));
JobExecutionContext.Builder builder = contextService.newBuilder(collectNode.jobId());
NodeOperation nodeOperation =
NodeOperation.withDownstream(
collectNode, mock(ExecutionPhase.class), (byte) 0, "remoteNode");
List<ListenableFuture<Bucket>> results =
contextPreparer.prepareOnRemote(
collectNode.jobId(), ImmutableList.of(nodeOperation), builder, sharedShardContexts);
JobExecutionContext context = contextService.createContext(builder);
context.start();
return results.get(0).get(2, TimeUnit.SECONDS);
}
private List<Bucket> getBuckets(CollectPhase collectNode)
throws InterruptedException, java.util.concurrent.ExecutionException {
List<Bucket> results = new ArrayList<>();
for (String nodeName : internalCluster().getNodeNames()) {
ContextPreparer contextPreparer =
internalCluster().getInstance(ContextPreparer.class, nodeName);
JobContextService contextService =
internalCluster().getInstance(JobContextService.class, nodeName);
JobExecutionContext.Builder builder = contextService.newBuilder(collectNode.jobId());
ListenableFuture<Bucket> future =
contextPreparer.prepare(
collectNode.jobId(),
NodeOperation.withDownstream(collectNode, mock(ExecutionPhase.class)),
builder);
assert future != null;
JobExecutionContext context = contextService.createContext(builder);
context.start();
results.add(future.get());
}
return results;
}