org.elasticsearch.tasks.Task Java Examples

@SuppressWarnings("unchecked")
@Test
public void testApplyActionRequestMakesPassThroughCallToOnResponse() {
    //given the chain will be called and wrappers the listener
    doAnswer(new Answer<Object>() {
        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            ActionListener listener = (ActionListener) args[3];
            listener.onResponse(response);
            return null;
        }
    }).when(chain).proceed(any(Task.class), anyString(), any(ActionRequest.class), any(ActionListener.class));
    
    //when
    filter.apply(task, action, request, listener, chain );
    
    //then the original listener should be notified
    verify(listener).onResponse(response);
}
 

public final Task execute(Request request, final TaskListener<Response> listener) {
    final Task task = taskManager.register("transport", actionName, request);
    execute(task, request, new ActionListener<Response>() {
        @Override
        public void onResponse(Response response) {
            if (task != null) {
                taskManager.unregister(task);
            }
            listener.onResponse(task, response);
        }

        @Override
        public void onFailure(Throwable e) {
            if (task != null) {
                taskManager.unregister(task);
            }
            listener.onFailure(task, e);
        }
    });
    return task;
}
 

/**
 * Perform a test search request to validate the element prior to storing it.
 *
 * @param validation validation info
 * @param element the element stored
 * @param task the parent task
 * @param listener the action listener to write to
 * @param onSuccess action ro run when the validation is successfull
 */
private void validate(FeatureValidation validation,
                      StorableElement element,
                      Task task,
                      ActionListener<FeatureStoreResponse> listener,
                      Runnable onSuccess) {
    ValidatingLtrQueryBuilder ltrBuilder = new ValidatingLtrQueryBuilder(element,
            validation, factory);
    SearchRequestBuilder builder = new SearchRequestBuilder(client, SearchAction.INSTANCE);
    builder.setIndices(validation.getIndex());
    builder.setQuery(ltrBuilder);
    builder.setFrom(0);
    builder.setSize(20);
    // Bail out early and don't score the whole index.
    builder.setTerminateAfter(1000);
    builder.request().setParentTask(clusterService.localNode().getId(), task.getId());
    builder.execute(wrap((r) -> {
            if (r.getFailedShards() > 0) {
                ShardSearchFailure failure = r.getShardFailures()[0];
                throw new IllegalArgumentException("Validating the element caused " + r.getFailedShards() +
                        " shard failures, see root cause: " + failure.reason(), failure.getCause());
            }
            onSuccess.run();
        },
        (e) -> listener.onFailure(new IllegalArgumentException("Cannot store element, validation failed.", e))));
}
 

@Override
protected void doExecute(Task task, UpgradeRequest request, final ActionListener<UpgradeResponse> listener) {
    ActionListener<UpgradeResponse> settingsUpdateListener = new ActionListener<UpgradeResponse>() {
        @Override
        public void onResponse(UpgradeResponse upgradeResponse) {
            try {
                if (upgradeResponse.versions().isEmpty()) {
                    listener.onResponse(upgradeResponse);
                } else {
                    updateSettings(upgradeResponse, listener);
                }
            } catch (Throwable t) {
                listener.onFailure(t);
            }
        }

        @Override
        public void onFailure(Throwable e) {
            listener.onFailure(e);
        }
    };
    super.doExecute(task, request, settingsUpdateListener);
}
 

@Override
public Task register(String type, String action, TaskAwareRequest request) {
    Task task = super.register(type, action, request);
    for (MockTaskManagerListener listener : listeners) {
        try {
            listener.onTaskRegistered(task);
        } catch (Exception e) {
            logger.warn(
                (Supplier<?>) () -> new ParameterizedMessage(
                    "failed to notify task manager listener about registering the task with id {}",
                    task.getId()),
                e);
        }
    }
    return task;
}
 

@SuppressWarnings("unchecked")
public void testExecutionException() {
    TransportService transportService = new TransportService(
        Settings.EMPTY,
        mock(Transport.class),
        null,
        TransportService.NOOP_TRANSPORT_INTERCEPTOR,
        x -> null,
        null,
        Collections.emptySet()
    );

    ModelManager manager = mock(ModelManager.class);
    ThresholdResultTransportAction action = new ThresholdResultTransportAction(mock(ActionFilters.class), transportService, manager);
    doThrow(NullPointerException.class)
        .when(manager)
        .getThresholdingResult(any(String.class), any(String.class), anyDouble(), any(ActionListener.class));

    final PlainActionFuture<ThresholdResultResponse> future = new PlainActionFuture<>();
    ThresholdResultRequest request = new ThresholdResultRequest("123", "123-threshold", 2);
    action.doExecute(mock(Task.class), request, future);

    expectThrows(NullPointerException.class, () -> future.actionGet());
}
 

/**
 * this is method is called on the recovery source node
 * the target is requesting the head of a file it got a PutReplicaChunkRequest for.
 */
@Override
public void messageReceived(final GetBlobHeadRequest request, TransportChannel channel, Task task) throws Exception {

    final BlobTransferStatus transferStatus = blobTransferTarget.getActiveTransfer(request.transferId);
    assert transferStatus != null :
        "Received GetBlobHeadRequest for transfer" + request.transferId.toString() +
        "but don't have an activeTransfer with that id";

    final DiscoveryNode recipientNode = clusterService.state().getNodes().get(request.senderNodeId);
    final long bytesToSend = request.endPos;

    blobTransferTarget.gotAGetBlobHeadRequest(request.transferId);

    channel.sendResponse(TransportResponse.Empty.INSTANCE);

    threadPool.generic().execute(
        new PutHeadChunkRunnable(
            transferStatus.digestBlob(), bytesToSend, transportService, blobTransferTarget,
            recipientNode, request.transferId)
    );
}
 

@Override @SuppressWarnings("unchecked")
public void proceed(Task task, String actionName, ActionRequest request, ActionListener listener) {
    int i = index.getAndIncrement();
    try {
        if (i < this.action.filters.length) {
            this.action.filters[i].apply(task, actionName, request, listener, this);
        } else if (i == this.action.filters.length) {
            this.action.doExecute(task, (Request) request, new FilteredActionListener<Response>(actionName, listener, new ResponseFilterChain(this.action.filters, logger)));
        } else {
            listener.onFailure(new IllegalStateException("proceed was called too many times"));
        }
    } catch(Throwable t) {
        logger.trace("Error during transport action execution.", t);
        listener.onFailure(t);
    }
}
 

@Override
public void messageReceived(RecoveryFinalizeRecoveryRequest request, TransportChannel channel, Task task) throws Exception {
    try (RecoveryRef recoveryRef = onGoingRecoveries.getRecoverySafe(request.recoveryId(), request.shardId())) {
        final ActionListener<TransportResponse> listener =
            new HandledTransportAction.ChannelActionListener<>(channel, Actions.FINALIZE, request);
        recoveryRef.target().finalizeRecovery(
            request.globalCheckpoint(),
            ActionListener.wrap(
                nullVal -> listener.onResponse(TransportResponse.Empty.INSTANCE),
                listener::onFailure
            )
        );
    }
}
 

@Override
protected void doExecute(Task task, AddFeaturesToSetRequest request, ActionListener<AddFeaturesToSetResponse> listener) {
    if (!clusterService.state().routingTable().hasIndex(request.getStore())) {
        throw new IllegalArgumentException("Store [" + request.getStore() + "] does not exist, please create it first.");
    }
    new AsyncAction(task, request, listener, clusterService, searchAction, getAction, featureStoreAction).start();
}
 

AsyncSingleAction(Task task, Request request, ActionListener<Response> listener) {
    this.task = task;
    this.request = request;
    if (task != null) {
        request.setParentTask(clusterService.localNode().getId(), task.getId());
    }
    this.listener = listener;
}
 

@SuppressWarnings("unchecked")
private void givenAnExeptionOccurs() {
    doAnswer(new Answer<Object>() {
        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            Object[] args = invocation.getArguments();
            ActionListener listener = (ActionListener) args[3];
            listener.onFailure(exception);
            return null;
        }
    }).when(chain).proceed(any(Task.class), anyString(), any(ActionRequest.class), any(ActionListener.class));
}
 

public void processMessageReceived(Request request, TransportChannel channel) throws Exception {
    final Task task = taskManager.register(channel.getChannelType(), action, request);
    boolean success = false;
    try {
        handler.messageReceived(request, new TaskTransportChannel(taskManager, task, channel), task);
        success = true;
    } finally {
        if (success == false) {
            taskManager.unregister(task);
        }
    }
}
 

@Override
public void messageReceived(RecoveryPrepareForTranslogOperationsRequest request, TransportChannel channel, Task task) throws Exception {
    try (RecoveryRef recoveryRef = onGoingRecoveries.getRecoverySafe(request.recoveryId(), request.shardId())) {
        final ActionListener<TransportResponse> listener =
            new HandledTransportAction.ChannelActionListener<>(channel, Actions.PREPARE_TRANSLOG, request);
        recoveryRef.target().prepareForTranslogOperations(
            request.isFileBasedRecovery(),
            request.totalTranslogOps(),
            ActionListener.wrap(
                nullVal -> listener.onResponse(TransportResponse.Empty.INSTANCE),
                listener::onFailure)
        );
    }
}
 

@Override
public void logSucceededLogin(String effectiveUser, boolean securityadmin, String initiatingUser, TransportRequest request, String action, Task task) {

    if(!checkTransportFilter(Category.AUTHENTICATED, action, effectiveUser, request)) {
        return;
    }

    final TransportAddress remoteAddress = getRemoteAddress();
    final List<AuditMessage> msgs = RequestResolver.resolve(Category.AUTHENTICATED, getOrigin(), action, null, effectiveUser, securityadmin, initiatingUser,remoteAddress, request, getThreadContextHeaders(), task, resolver, clusterService, settings, logRequestBody, resolveIndices, resolveBulkRequests, opendistrosecurityIndex, excludeSensitiveHeaders, null);

    for(AuditMessage msg: msgs) {
        save(msg);
    }
}
 

@Override
public void logMissingPrivileges(String privilege, TransportRequest request, Task task) {
    final String action = null;

    if(!checkTransportFilter(Category.MISSING_PRIVILEGES, privilege, getUser(), request)) {
        return;
    }

    final TransportAddress remoteAddress = getRemoteAddress();
    final List<AuditMessage> msgs = RequestResolver.resolve(Category.MISSING_PRIVILEGES, getOrigin(), action, privilege, getUser(), null, null, remoteAddress, request, getThreadContextHeaders(), task, resolver, clusterService, settings, logRequestBody, resolveIndices, resolveBulkRequests, opendistrosecurityIndex, excludeSensitiveHeaders, null);

    for(AuditMessage msg: msgs) {
        save(msg);
    }
}
 

private IndexRequest buildIndexRequest(Task parentTask, FeatureStoreRequest request) throws IOException {
    StorableElement elt = request.getStorableElement();

    IndexRequest indexRequest = client.prepareIndex(request.getStore(), IndexFeatureStore.ES_TYPE, elt.id())
            .setCreate(request.getAction() == FeatureStoreRequest.Action.CREATE)
            .setRouting(request.getRouting())
            .setSource(IndexFeatureStore.toSource(elt))
            .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
            .request();
    indexRequest.setParentTask(clusterService.localNode().getId(), parentTask.getId());
    return indexRequest;
}
 

@Override
public void messageReceived(VerifyNodeRepositoryRequest request, TransportChannel channel, Task task) throws Exception {
    DiscoveryNode localNode = clusterService.state().nodes().getLocalNode();
    try {
        doVerify(request.repository, request.verificationToken, localNode);
    } catch (Exception ex) {
        LOGGER.warn(() -> new ParameterizedMessage("[{}] failed to verify repository", request.repository), ex);
        throw ex;
    }
    channel.sendResponse(TransportResponse.Empty.INSTANCE);
}
 

@SuppressWarnings("unchecked")
public void testNormal() {
    TransportService transportService = new TransportService(
        Settings.EMPTY,
        mock(Transport.class),
        null,
        TransportService.NOOP_TRANSPORT_INTERCEPTOR,
        x -> null,
        null,
        Collections.emptySet()
    );

    ModelManager manager = mock(ModelManager.class);
    ADCircuitBreakerService adCircuitBreakerService = mock(ADCircuitBreakerService.class);
    RCFResultTransportAction action = new RCFResultTransportAction(
        mock(ActionFilters.class),
        transportService,
        manager,
        adCircuitBreakerService
    );
    doAnswer(invocation -> {
        ActionListener<RcfResult> listener = invocation.getArgument(3);
        listener.onResponse(new RcfResult(0, 0, 25));
        return null;
    }).when(manager).getRcfResult(any(String.class), any(String.class), any(double[].class), any(ActionListener.class));

    when(adCircuitBreakerService.isOpen()).thenReturn(false);

    final PlainActionFuture<RCFResultResponse> future = new PlainActionFuture<>();
    RCFResultRequest request = new RCFResultRequest("123", "123-rcf-1", new double[] { 0 });
    action.doExecute(mock(Task.class), request, future);

    RCFResultResponse response = future.actionGet();
    assertEquals(0, response.getRCFScore(), 0.001);
    assertEquals(25, response.getForestSize(), 0.001);
}
 

@SuppressWarnings("unchecked")
public void testExecutionException() {
    TransportService transportService = new TransportService(
        Settings.EMPTY,
        mock(Transport.class),
        null,
        TransportService.NOOP_TRANSPORT_INTERCEPTOR,
        x -> null,
        null,
        Collections.emptySet()
    );

    ModelManager manager = mock(ModelManager.class);
    ADCircuitBreakerService adCircuitBreakerService = mock(ADCircuitBreakerService.class);
    RCFResultTransportAction action = new RCFResultTransportAction(
        mock(ActionFilters.class),
        transportService,
        manager,
        adCircuitBreakerService
    );
    doThrow(NullPointerException.class)
        .when(manager)
        .getRcfResult(any(String.class), any(String.class), any(double[].class), any(ActionListener.class));
    when(adCircuitBreakerService.isOpen()).thenReturn(false);

    final PlainActionFuture<RCFResultResponse> future = new PlainActionFuture<>();
    RCFResultRequest request = new RCFResultRequest("123", "123-rcf-1", new double[] { 0 });
    action.doExecute(mock(Task.class), request, future);

    expectThrows(NullPointerException.class, () -> future.actionGet());
}
 

private RestChannelConsumer sendTask(String localNodeId, Task task) {
    return channel -> {
        try (XContentBuilder builder = channel.newBuilder()) {
            builder.startObject();
            builder.field("task", localNodeId + ":" + task.getId());
            builder.endObject();
            channel.sendResponse(new BytesRestResponse(RestStatus.OK, builder));
        }
    };
}
 

@Override
public void messageReceived(StartedShardEntry request, TransportChannel channel, Task task) throws Exception {
    logger.debug("{} received shard started for [{}]", request.shardId, request);
    clusterService.submitStateUpdateTask(
        "shard-started " + request,
        request,
        ClusterStateTaskConfig.build(Priority.URGENT),
        shardStartedClusterStateTaskExecutor,
        shardStartedClusterStateTaskExecutor);
    channel.sendResponse(TransportResponse.Empty.INSTANCE);
}
 

@Override
public void messageReceived(BlobRecoveryChunkRequest request, TransportChannel channel, Task task) throws Exception {

    BlobRecoveryStatus onGoingRecovery = onGoingBlobRecoveries.get(request.recoveryId());
    if (onGoingRecovery == null) {
        // shard is getting closed on us
        throw new IllegalBlobRecoveryStateException("Could not retrieve onGoingRecoveryStatus");
    }


    BlobRecoveryTransferStatus transferStatus = onGoingRecovery.onGoingTransfers().get(request.transferId());
    BlobShard shard = onGoingRecovery.blobShard;
    if (onGoingRecovery.canceled()) {
        onGoingRecovery.sentCanceledToSource();
        throw new IndexShardClosedException(onGoingRecovery.shardId());
    }

    if (transferStatus == null) {
        throw new IndexShardClosedException(onGoingRecovery.shardId());
    }

    request.content().writeTo(transferStatus.outputStream());

    if (request.isLast()) {
        transferStatus.outputStream().close();
        Path baseDirectory = shard.blobContainer().getBaseDirectory();
        Path source = baseDirectory.resolve(transferStatus.sourcePath());
        Path target = baseDirectory.resolve(transferStatus.targetPath());

        Files.move(source, target, StandardCopyOption.ATOMIC_MOVE, StandardCopyOption.REPLACE_EXISTING);
        onGoingRecovery.onGoingTransfers().remove(request.transferId());
    }

    channel.sendResponse(TransportResponse.Empty.INSTANCE);
}
 

@Override
public <Request extends ActionRequest, Response extends ActionResponse> void apply(final Task task, final String action,
        final Request request, final ActionListener<Response> listener, final ActionFilterChain<Request, Response> chain) {
    if (!SearchAction.INSTANCE.name().equals(action)) {
        chain.proceed(task, action, request, listener);
        return;
    }

    final SearchRequest searchRequest = (SearchRequest) request;
    final ActionListener<Response> wrappedListener = DynamicRanker.getInstance().wrapActionListener(action, searchRequest, listener);
    chain.proceed(task, action, request, wrappedListener == null ? listener : wrappedListener);
}
 

@Override
protected void doExecute(Task task, CreateModelFromSetRequest request, ActionListener<CreateModelFromSetResponse> listener) {
    if (!clusterService.state().routingTable().hasIndex(request.getStore())) {
        throw new IllegalArgumentException("Store [" + request.getStore() + "] does not exist, please create it first.");
    }
    GetRequest getRequest = new GetRequest(request.getStore())
            .type(IndexFeatureStore.ES_TYPE)
            .id(StorableElement.generateId(StoredFeatureSet.TYPE, request.getFeatureSetName()));
    getRequest.setParentTask(clusterService.localNode().getId(), task.getId());
    getAction.execute(getRequest, ActionListener.wrap((r) -> this.doStore(task, r, request, listener), listener::onFailure));
}
 

@Override
public void messageReceived(final NodeRequest request, TransportChannel channel, Task task) throws Exception {
    List<ShardRouting> shards = request.getShards();
    final int totalShards = shards.size();
    if (logger.isTraceEnabled()) {
        logger.trace("[{}] executing operation on [{}] shards", actionName, totalShards);
    }
    final Object[] shardResultOrExceptions = new Object[totalShards];

    int shardIndex = -1;
    for (final ShardRouting shardRouting : shards) {
        shardIndex++;
        onShardOperation(request, shardResultOrExceptions, shardIndex, shardRouting);
    }

    List<BroadcastShardOperationFailedException> accumulatedExceptions = new ArrayList<>();
    List<ShardOperationResult> results = new ArrayList<>();
    for (int i = 0; i < totalShards; i++) {
        if (shardResultOrExceptions[i] instanceof BroadcastShardOperationFailedException) {
            accumulatedExceptions.add((BroadcastShardOperationFailedException) shardResultOrExceptions[i]);
        } else {
            results.add((ShardOperationResult) shardResultOrExceptions[i]);
        }
    }

    channel.sendResponse(new NodeResponse(request.getNodeId(), totalShards, results, accumulatedExceptions));
}
 

public TaskInfo(DiscoveryNode node, long id, String type, String action, String description, Task.Status status, long startTime,
                long runningTimeNanos, TaskId parentTaskId) {
    this.node = node;
    this.taskId = new TaskId(node.getId(), id);
    this.type = type;
    this.action = action;
    this.description = description;
    this.status = status;
    this.startTime = startTime;
    this.runningTimeNanos = runningTimeNanos;
    this.parentTaskId = parentTaskId;
}
 

@Override
public void messageReceived(NodeRequest request, TransportChannel channel, Task task) throws Exception {
    channel.sendResponse(nodeOperation(request, task));
}
 

@Override
public void messageReceived(ShardRequest request, TransportChannel channel, Task task) throws Exception {
    channel.sendResponse(shardOperation(request));
}
 

@Override
public Task createTask(long id, String type, String action, TaskId parentTaskId) {
    return new ReplicationTask(id, type, action, getDescription(), parentTaskId);
}