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

Example 1 with TasksService

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

the class TransportKillJobsNodeActionTest method testKillIsCalledOnJobContextService.

@Test
public void testKillIsCalledOnJobContextService() throws Exception {
    TasksService tasksService = mock(TasksService.class, Answers.RETURNS_MOCKS);
    TransportKillJobsNodeAction transportKillJobsNodeAction = new TransportKillJobsNodeAction(tasksService, clusterService, MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, THREAD_POOL, clusterService.getClusterSettings()));
    List<UUID> toKill = List.of(UUID.randomUUID(), UUID.randomUUID());
    transportKillJobsNodeAction.nodeOperation(new KillJobsRequest(toKill, "dummy-user", null)).get(5, TimeUnit.SECONDS);
    verify(tasksService, times(1)).killJobs(toKill, "dummy-user", null);
}

Example 2 with TasksService

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

the class NodeFetchOperationTest method testSysOperationsIsClearedIfNothingToFetch.

@Test
public void testSysOperationsIsClearedIfNothingToFetch() throws Exception {
    ThreadPoolExecutor threadPoolExecutor = (ThreadPoolExecutor) Executors.newFixedThreadPool(1);
    try {
        JobsLogs jobsLogs = new JobsLogs(() -> true);
        NodeFetchOperation fetchOperation = new NodeFetchOperation(threadPoolExecutor, 2, jobsLogs, new TasksService(clusterService, jobsLogs), new NoopCircuitBreaker("dummy"));
        fetchOperation.fetch(UUID.randomUUID(), 1, null, true).get(5, TimeUnit.SECONDS);
        assertThat(StreamSupport.stream(jobsLogs.activeOperations().spliterator(), false).count(), is(0L));
    } finally {
        threadPoolExecutor.shutdown();
        threadPoolExecutor.awaitTermination(2, TimeUnit.SECONDS);
    }
}

Example 3 with TasksService

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

the class TransportDistributedResultActionTest method testKillIsInvokedIfContextIsNotFound.

@Test
public void testKillIsInvokedIfContextIsNotFound() throws Exception {
    TasksService tasksService = new TasksService(clusterService, new JobsLogs(() -> false));
    AtomicInteger numBroadcasts = new AtomicInteger(0);
    TransportKillJobsNodeAction killJobsAction = new TransportKillJobsNodeAction(tasksService, clusterService, mock(TransportService.class)) {

        @Override
        public void broadcast(KillJobsRequest request, ActionListener<Long> listener, Collection<String> excludedNodeIds) {
            numBroadcasts.incrementAndGet();
        }
    };
    TransportDistributedResultAction transportDistributedResultAction = new TransportDistributedResultAction(mock(Transports.class), tasksService, THREAD_POOL, mock(TransportService.class), clusterService, killJobsAction, BackoffPolicy.exponentialBackoff(TimeValue.ZERO, 0));
    StreamBucket.Builder builder = new StreamBucket.Builder(new Streamer[0], RamAccounting.NO_ACCOUNTING);
    try {
        transportDistributedResultAction.nodeOperation(new DistributedResultRequest(UUID.randomUUID(), 0, (byte) 0, 0, builder.build(), true)).get(5, TimeUnit.SECONDS);
        fail("nodeOperation call should fail with TaskMissing");
    } catch (ExecutionException e) {
        assertThat(e.getCause(), Matchers.instanceOf(TaskMissing.class));
    }
    assertThat(numBroadcasts.get(), is(1));
}

Example 4 with TasksService

use of io.crate.execution.jobs.TasksService 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 5 with TasksService

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

the class RemoteCollectorTest method prepare.

@Before
public void prepare() {
    MockitoAnnotations.initMocks(this);
    UUID jobId = UUID.randomUUID();
    RoutedCollectPhase collectPhase = new RoutedCollectPhase(jobId, 0, "remoteCollect", new Routing(Map.of("remoteNode", Map.of("dummyTable", IntArrayList.from(1)))), RowGranularity.DOC, Collections.singletonList(createReference("name", DataTypes.STRING)), Collections.emptyList(), WhereClause.MATCH_ALL.queryOrFallback(), DistributionInfo.DEFAULT_BROADCAST);
    transportJobAction = mock(TransportJobAction.class);
    TasksService tasksService = new TasksService(clusterService, new JobsLogs(() -> true));
    numBroadcastCalls = new AtomicInteger(0);
    transportKillJobsNodeAction = new TransportKillJobsNodeAction(tasksService, clusterService, mock(TransportService.class)) {

        @Override
        public void broadcast(KillJobsRequest request, ActionListener<Long> listener) {
            numBroadcastCalls.incrementAndGet();
        }
    };
    consumer = new TestingRowConsumer();
    remoteCollector = new RemoteCollector(jobId, new SessionSettings("dummyUser", SearchPath.createSearchPathFrom("dummySchema")), "localNode", "remoteNode", transportJobAction, transportKillJobsNodeAction, Runnable::run, tasksService, RamAccounting.NO_ACCOUNTING, consumer, collectPhase);
}