Examples with CancellableTask org.opensearch.tasks.CancellableTask used on opensource projects

Example 6 with CancellableTask

use of org.opensearch.tasks.CancellableTask in project OpenSearch by opensearch-project.

the class RequestHandlerRegistry method processMessageReceived.

public void processMessageReceived(Request request, TransportChannel channel) throws Exception {
    final Task task = taskManager.register(channel.getChannelType(), action, request);
    Releasable unregisterTask = () -> taskManager.unregister(task);
    try {
        if (channel instanceof TcpTransportChannel && task instanceof CancellableTask) {
            if (request instanceof ShardSearchRequest) {
                // on receiving request, update the inbound network time to reflect time spent in transit over the network
                ((ShardSearchRequest) request).setInboundNetworkTime(Math.max(0, System.currentTimeMillis() - ((ShardSearchRequest) request).getInboundNetworkTime()));
            }
            final TcpChannel tcpChannel = ((TcpTransportChannel) channel).getChannel();
            final Releasable stopTracking = taskManager.startTrackingCancellableChannelTask(tcpChannel, (CancellableTask) task);
            unregisterTask = Releasables.wrap(unregisterTask, stopTracking);
        }
        final TaskTransportChannel taskTransportChannel = new TaskTransportChannel(channel, unregisterTask);
        handler.messageReceived(request, taskTransportChannel, task);
        unregisterTask = null;
    } finally {
        Releasables.close(unregisterTask);
    }
}