Examples with RootTask io.crate.execution.jobs.RootTask used on opensource projects

Example 1 with RootTask

use of io.crate.execution.jobs.RootTask in project crate by crate.

the class JobLauncher method getHandlerBucketReceivers.

private List<PageBucketReceiver> getHandlerBucketReceivers(RootTask rootTask, List<Tuple<ExecutionPhase, RowConsumer>> handlerPhases) {
    final List<PageBucketReceiver> pageBucketReceivers = new ArrayList<>(handlerPhases.size());
    for (Tuple<ExecutionPhase, ?> handlerPhase : handlerPhases) {
        Task ctx = rootTask.getTaskOrNull(handlerPhase.v1().phaseId());
        if (ctx instanceof DownstreamRXTask) {
            PageBucketReceiver pageBucketReceiver = ((DownstreamRXTask) ctx).getBucketReceiver((byte) 0);
            pageBucketReceivers.add(pageBucketReceiver);
        }
    }
    return pageBucketReceivers;
}

Example 2 with RootTask

use of io.crate.execution.jobs.RootTask in project crate by crate.

the class SQLIntegrationTestCase method assertNoTasksAreLeftOpen.

@After
public void assertNoTasksAreLeftOpen() throws Exception {
    final Field activeTasks = TasksService.class.getDeclaredField("activeTasks");
    final Field activeOperationsSb = TransportShardAction.class.getDeclaredField("activeOperations");
    activeTasks.setAccessible(true);
    activeOperationsSb.setAccessible(true);
    try {
        assertBusy(() -> {
            for (TasksService tasksService : internalCluster().getInstances(TasksService.class)) {
                try {
                    // noinspection unchecked
                    Map<UUID, RootTask> contexts = (Map<UUID, RootTask>) activeTasks.get(tasksService);
                    assertThat(contexts.size(), is(0));
                } catch (IllegalAccessException e) {
                    throw new RuntimeException(e);
                }
            }
            for (TransportShardUpsertAction action : internalCluster().getInstances(TransportShardUpsertAction.class)) {
                try {
                    ConcurrentHashMap<TaskId, KillableCallable<?>> operations = (ConcurrentHashMap<TaskId, KillableCallable<?>>) activeOperationsSb.get(action);
                    assertThat(operations.size(), is(0));
                } catch (IllegalAccessException e) {
                    throw new RuntimeException(e);
                }
            }
            for (TransportShardDeleteAction action : internalCluster().getInstances(TransportShardDeleteAction.class)) {
                try {
                    ConcurrentHashMap<TaskId, KillableCallable<?>> operations = (ConcurrentHashMap<TaskId, KillableCallable<?>>) activeOperationsSb.get(action);
                    assertThat(operations.size(), is(0));
                } catch (IllegalAccessException e) {
                    throw new RuntimeException(e);
                }
            }
        }, 10L, TimeUnit.SECONDS);
    } catch (AssertionError e) {
        StringBuilder errorMessageBuilder = new StringBuilder();
        String[] nodeNames = internalCluster().getNodeNames();
        for (String nodeName : nodeNames) {
            TasksService tasksService = internalCluster().getInstance(TasksService.class, nodeName);
            try {
                // noinspection unchecked
                Map<UUID, RootTask> contexts = (Map<UUID, RootTask>) activeTasks.get(tasksService);
                String contextsString = contexts.toString();
                if (!"{}".equals(contextsString)) {
                    errorMessageBuilder.append("## node: ");
                    errorMessageBuilder.append(nodeName);
                    errorMessageBuilder.append("\n");
                    errorMessageBuilder.append(contextsString);
                    errorMessageBuilder.append("\n");
                }
                contexts.clear();
            } catch (IllegalAccessException ex) {
                throw new RuntimeException(ex);
            }
        }
        throw new AssertionError(errorMessageBuilder.toString(), e);
    }
}

Example 3 with RootTask

use of io.crate.execution.jobs.RootTask in project crate by crate.

the class TransportJobAction method nodeOperation.

@Override
public CompletableFuture<JobResponse> nodeOperation(final JobRequest request) {
    RootTask.Builder contextBuilder = tasksService.newBuilder(request.jobId(), request.sessionSettings().userName(), request.coordinatorNodeId(), Collections.emptySet());
    SharedShardContexts sharedShardContexts = maybeInstrumentProfiler(request.enableProfiling(), contextBuilder);
    List<CompletableFuture<StreamBucket>> directResponseFutures = jobSetup.prepareOnRemote(request.sessionSettings(), request.nodeOperations(), contextBuilder, sharedShardContexts);
    try {
        RootTask context = tasksService.createTask(contextBuilder);
        context.start();
    } catch (Throwable t) {
        return CompletableFuture.failedFuture(t);
    }
    if (directResponseFutures.size() == 0) {
        return CompletableFuture.completedFuture(new JobResponse(List.of()));
    } else {
        return CompletableFutures.allAsList(directResponseFutures).thenApply(JobResponse::new);
    }
}

Example 4 with RootTask

use of io.crate.execution.jobs.RootTask in project crate by crate.

the class JobLauncher method setupTasks.

private void setupTasks(TransactionContext txnCtx, Map<String, Collection<NodeOperation>> operationByServer, List<ExecutionPhase> handlerPhases, List<RowConsumer> handlerConsumers) throws Throwable {
    assert handlerPhases.size() == handlerConsumers.size() : "handlerPhases size must match handlerConsumers size";
    String localNodeId = clusterService.localNode().getId();
    Collection<NodeOperation> localNodeOperations = operationByServer.remove(localNodeId);
    if (localNodeOperations == null) {
        localNodeOperations = Collections.emptyList();
    }
    // + 1 for localTask which is always created
    InitializationTracker initializationTracker = new InitializationTracker(operationByServer.size() + 1);
    List<Tuple<ExecutionPhase, RowConsumer>> handlerPhaseAndReceiver = createHandlerPhaseAndReceivers(handlerPhases, handlerConsumers, initializationTracker);
    RootTask.Builder builder = tasksService.newBuilder(jobId, txnCtx.sessionSettings().userName(), localNodeId, operationByServer.keySet());
    SharedShardContexts sharedShardContexts = maybeInstrumentProfiler(builder);
    List<CompletableFuture<StreamBucket>> directResponseFutures = jobSetup.prepareOnHandler(txnCtx.sessionSettings(), localNodeOperations, builder, handlerPhaseAndReceiver, sharedShardContexts);
    RootTask localTask = tasksService.createTask(builder);
    List<PageBucketReceiver> pageBucketReceivers = getHandlerBucketReceivers(localTask, handlerPhaseAndReceiver);
    int bucketIdx = 0;
    /*
         * If you touch anything here make sure the following tests pass with > 1k iterations:
         *
         * Seed: 112E1807417E925A - testInvalidPatternSyntax
         * Seed: Any              - testRegularSelectWithFewAvailableThreadsShouldNeverGetStuck
         * Seed: CC456FF5004F35D3 - testFailureOfJoinDownstream
         */
    if (!localNodeOperations.isEmpty() && !directResponseFutures.isEmpty()) {
        assert directResponseFutures.size() == pageBucketReceivers.size() : "directResponses size must match pageBucketReceivers";
        CompletableFutures.allAsList(directResponseFutures).whenComplete(BucketForwarder.asConsumer(pageBucketReceivers, bucketIdx, initializationTracker));
        bucketIdx++;
        try {
            // initializationTracker for localNodeOperations is triggered via SetBucketCallback
            localTask.start();
        } catch (Throwable t) {
            accountFailureForRemoteOperations(operationByServer, initializationTracker, handlerPhaseAndReceiver, t);
            return;
        }
    } else {
        try {
            localTask.start();
            initializationTracker.jobInitialized();
        } catch (Throwable t) {
            initializationTracker.jobInitializationFailed(t);
            accountFailureForRemoteOperations(operationByServer, initializationTracker, handlerPhaseAndReceiver, t);
            return;
        }
    }
    sendJobRequests(txnCtx, localNodeId, operationByServer, pageBucketReceivers, handlerPhaseAndReceiver, bucketIdx, initializationTracker);
}

Example 5 with RootTask

use of io.crate.execution.jobs.RootTask in project crate by crate.

the class NodeFetchOperation method fetch.

public CompletableFuture<? extends IntObjectMap<StreamBucket>> fetch(UUID jobId, int phaseId, @Nullable IntObjectMap<IntArrayList> docIdsToFetch, boolean closeTaskOnFinish) {
    if (docIdsToFetch == null) {
        if (closeTaskOnFinish) {
            tryCloseTask(jobId, phaseId);
        }
        jobsLogs.operationStarted(phaseId, jobId, "fetch", () -> -1);
        jobsLogs.operationFinished(phaseId, jobId, null);
        return CompletableFuture.completedFuture(new IntObjectHashMap<>(0));
    }
    RootTask context = tasksService.getTask(jobId);
    FetchTask fetchTask = context.getTask(phaseId);
    jobsLogs.operationStarted(phaseId, jobId, "fetch", () -> -1);
    BiConsumer<? super IntObjectMap<StreamBucket>, ? super Throwable> whenComplete = (res, err) -> {
        if (closeTaskOnFinish) {
            if (err == null) {
                fetchTask.close();
            } else {
                fetchTask.kill(err);
            }
        }
        if (err == null) {
            jobsLogs.operationFinished(phaseId, jobId, null);
        } else {
            jobsLogs.operationFinished(phaseId, jobId, SQLExceptions.messageOf(err));
        }
    };
    try {
        return doFetch(fetchTask, docIdsToFetch).whenComplete(whenComplete);
    } catch (Throwable t) {
        whenComplete.accept(null, t);
        return CompletableFuture.failedFuture(t);
    }
}