Examples with HttpRemoteTaskFactory io.trino.server.HttpRemoteTaskFactory used on opensource projects

Example 1 with HttpRemoteTaskFactory

use of io.trino.server.HttpRemoteTaskFactory in project trino by trinodb.

the class TestHttpRemoteTask method testOutboundDynamicFilters.

@Test(timeOut = 30_000)
public void testOutboundDynamicFilters() throws Exception {
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    DynamicFilterId filterId2 = new DynamicFilterId("df2");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", BIGINT);
    Symbol symbol2 = symbolAllocator.newSymbol("DF_SYMBOL2", BIGINT);
    SymbolReference df1 = symbol1.toSymbolReference();
    SymbolReference df2 = symbol2.toSymbolReference();
    ColumnHandle handle1 = new TestingColumnHandle("column1");
    ColumnHandle handle2 = new TestingColumnHandle("column2");
    QueryId queryId = new QueryId("test");
    TestingTaskResource testingTaskResource = new TestingTaskResource(new AtomicLong(System.nanoTime()), FailureScenario.NO_FAILURE);
    DynamicFilterService dynamicFilterService = new DynamicFilterService(PLANNER_CONTEXT.getMetadata(), PLANNER_CONTEXT.getFunctionManager(), new TypeOperators(), newDirectExecutorService());
    dynamicFilterService.registerQuery(queryId, TEST_SESSION, ImmutableSet.of(filterId1, filterId2), ImmutableSet.of(filterId1, filterId2), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(new StageId(queryId, 1), 0, 1);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId2, df2)), ImmutableMap.of(symbol1, handle1, symbol2, handle2), symbolAllocator.getTypes());
    // make sure initial dynamic filter is collected
    CompletableFuture<?> future = dynamicFilter.isBlocked();
    dynamicFilterService.addTaskDynamicFilters(new TaskId(new StageId(queryId.getId(), 1), 1, 0), ImmutableMap.of(filterId1, Domain.singleValue(BIGINT, 1L)));
    future.get();
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(handle1, Domain.singleValue(BIGINT, 1L))));
    // Create remote task after dynamic filter is created to simulate new nodes joining
    HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource, dynamicFilterService);
    RemoteTask remoteTask = createRemoteTask(httpRemoteTaskFactory, ImmutableSet.of(filterId1, filterId2));
    testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo());
    remoteTask.start();
    assertEventually(new Duration(10, SECONDS), () -> assertEquals(testingTaskResource.getDynamicFiltersSentCounter(), 1L));
    assertEquals(testingTaskResource.getCreateOrUpdateCounter(), 1L);
    // schedule a couple of splits to trigger task updates
    addSplit(remoteTask, testingTaskResource, 1);
    addSplit(remoteTask, testingTaskResource, 2);
    // make sure dynamic filter was sent in task updates only once
    assertEquals(testingTaskResource.getDynamicFiltersSentCounter(), 1L);
    assertEquals(testingTaskResource.getCreateOrUpdateCounter(), 3L);
    assertEquals(testingTaskResource.getLatestDynamicFilterFromCoordinator(), ImmutableMap.of(filterId1, Domain.singleValue(BIGINT, 1L)));
    future = dynamicFilter.isBlocked();
    dynamicFilterService.addTaskDynamicFilters(new TaskId(new StageId(queryId.getId(), 1), 1, 0), ImmutableMap.of(filterId2, Domain.singleValue(BIGINT, 2L)));
    future.get();
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(handle1, Domain.singleValue(BIGINT, 1L), handle2, Domain.singleValue(BIGINT, 2L))));
    // dynamic filter should be sent even though there were no further splits scheduled
    assertEventually(new Duration(10, SECONDS), () -> assertEquals(testingTaskResource.getDynamicFiltersSentCounter(), 2L));
    assertEquals(testingTaskResource.getCreateOrUpdateCounter(), 4L);
    // previously sent dynamic filter should not be repeated
    assertEquals(testingTaskResource.getLatestDynamicFilterFromCoordinator(), ImmutableMap.of(filterId2, Domain.singleValue(BIGINT, 2L)));
    httpRemoteTaskFactory.stop();
    dynamicFilterService.stop();
}

Example 2 with HttpRemoteTaskFactory

use of io.trino.server.HttpRemoteTaskFactory in project trino by trinodb.

the class TestHttpRemoteTask method testRegular.

@Test(timeOut = 30000)
public void testRegular() throws Exception {
    AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime());
    TestingTaskResource testingTaskResource = new TestingTaskResource(lastActivityNanos, FailureScenario.NO_FAILURE);
    HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource);
    RemoteTask remoteTask = createRemoteTask(httpRemoteTaskFactory, ImmutableSet.of());
    testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo());
    remoteTask.start();
    Lifespan lifespan = Lifespan.driverGroup(3);
    remoteTask.addSplits(ImmutableMultimap.of(TABLE_SCAN_NODE_ID, new Split(new CatalogName("test"), TestingSplit.createLocalSplit(), lifespan)));
    poll(() -> testingTaskResource.getTaskSplitAssignment(TABLE_SCAN_NODE_ID) != null);
    poll(() -> testingTaskResource.getTaskSplitAssignment(TABLE_SCAN_NODE_ID).getSplits().size() == 1);
    remoteTask.noMoreSplits(TABLE_SCAN_NODE_ID, lifespan);
    poll(() -> testingTaskResource.getTaskSplitAssignment(TABLE_SCAN_NODE_ID).getNoMoreSplitsForLifespan().size() == 1);
    remoteTask.noMoreSplits(TABLE_SCAN_NODE_ID);
    poll(() -> testingTaskResource.getTaskSplitAssignment(TABLE_SCAN_NODE_ID).isNoMoreSplits());
    remoteTask.cancel();
    poll(() -> remoteTask.getTaskStatus().getState().isDone());
    poll(() -> remoteTask.getTaskInfo().getTaskStatus().getState().isDone());
    httpRemoteTaskFactory.stop();
}

Example 3 with HttpRemoteTaskFactory

use of io.trino.server.HttpRemoteTaskFactory in project trino by trinodb.

the class TestHttpRemoteTask method testDynamicFilters.

@Test(timeOut = 30000)
public void testDynamicFilters() throws Exception {
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    DynamicFilterId filterId2 = new DynamicFilterId("df2");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", BIGINT);
    Symbol symbol2 = symbolAllocator.newSymbol("DF_SYMBOL2", BIGINT);
    SymbolReference df1 = symbol1.toSymbolReference();
    SymbolReference df2 = symbol2.toSymbolReference();
    ColumnHandle handle1 = new TestingColumnHandle("column1");
    ColumnHandle handle2 = new TestingColumnHandle("column2");
    QueryId queryId = new QueryId("test");
    TestingTaskResource testingTaskResource = new TestingTaskResource(new AtomicLong(System.nanoTime()), FailureScenario.NO_FAILURE);
    DynamicFilterService dynamicFilterService = new DynamicFilterService(PLANNER_CONTEXT.getMetadata(), PLANNER_CONTEXT.getFunctionManager(), new TypeOperators(), newDirectExecutorService());
    HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource, dynamicFilterService);
    RemoteTask remoteTask = createRemoteTask(httpRemoteTaskFactory, ImmutableSet.of());
    Map<DynamicFilterId, Domain> initialDomain = ImmutableMap.of(filterId1, Domain.singleValue(BIGINT, 1L));
    testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo());
    testingTaskResource.setDynamicFilterDomains(new VersionedDynamicFilterDomains(1L, initialDomain));
    dynamicFilterService.registerQuery(queryId, TEST_SESSION, ImmutableSet.of(filterId1, filterId2), ImmutableSet.of(filterId1, filterId2), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(new StageId(queryId, 1), 0, 1);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId2, df2)), ImmutableMap.of(symbol1, handle1, symbol2, handle2), symbolAllocator.getTypes());
    // make sure initial dynamic filters are collected
    CompletableFuture<?> future = dynamicFilter.isBlocked();
    remoteTask.start();
    future.get();
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(handle1, Domain.singleValue(BIGINT, 1L))));
    assertEquals(testingTaskResource.getDynamicFiltersFetchCounter(), 1);
    // make sure dynamic filters are not collected for every status update
    assertEventually(new Duration(15, SECONDS), () -> assertGreaterThanOrEqual(testingTaskResource.getStatusFetchCounter(), 3L));
    assertEquals(testingTaskResource.getDynamicFiltersFetchCounter(), 1L, testingTaskResource.getDynamicFiltersFetchRequests().toString());
    future = dynamicFilter.isBlocked();
    testingTaskResource.setDynamicFilterDomains(new VersionedDynamicFilterDomains(2L, ImmutableMap.of(filterId2, Domain.singleValue(BIGINT, 2L))));
    future.get();
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(handle1, Domain.singleValue(BIGINT, 1L), handle2, Domain.singleValue(BIGINT, 2L))));
    assertEquals(testingTaskResource.getDynamicFiltersFetchCounter(), 2L, testingTaskResource.getDynamicFiltersFetchRequests().toString());
    assertGreaterThanOrEqual(testingTaskResource.getStatusFetchCounter(), 4L);
    httpRemoteTaskFactory.stop();
    dynamicFilterService.stop();
}

Example 4 with HttpRemoteTaskFactory

use of io.trino.server.HttpRemoteTaskFactory in project trino by trinodb.

the class TestHttpRemoteTask method createHttpRemoteTaskFactory.

private static HttpRemoteTaskFactory createHttpRemoteTaskFactory(TestingTaskResource testingTaskResource, DynamicFilterService dynamicFilterService) {
    Bootstrap app = new Bootstrap(new JsonModule(), new HandleJsonModule(), new Module() {

        @Override
        public void configure(Binder binder) {
            binder.bind(JsonMapper.class).in(SINGLETON);
            binder.bind(Metadata.class).toInstance(createTestMetadataManager());
            jsonBinder(binder).addDeserializerBinding(Type.class).to(TypeDeserializer.class);
            jsonCodecBinder(binder).bindJsonCodec(TaskStatus.class);
            jsonCodecBinder(binder).bindJsonCodec(VersionedDynamicFilterDomains.class);
            jsonBinder(binder).addSerializerBinding(Block.class).to(BlockJsonSerde.Serializer.class);
            jsonBinder(binder).addDeserializerBinding(Block.class).to(BlockJsonSerde.Deserializer.class);
            jsonCodecBinder(binder).bindJsonCodec(TaskInfo.class);
            jsonCodecBinder(binder).bindJsonCodec(TaskUpdateRequest.class);
            jsonCodecBinder(binder).bindJsonCodec(FailTaskRequest.class);
            binder.bind(TypeManager.class).toInstance(TESTING_TYPE_MANAGER);
            binder.bind(BlockEncodingManager.class).in(SINGLETON);
            binder.bind(BlockEncodingSerde.class).to(InternalBlockEncodingSerde.class).in(SINGLETON);
        }

        @Provides
        private HttpRemoteTaskFactory createHttpRemoteTaskFactory(JsonMapper jsonMapper, JsonCodec<TaskStatus> taskStatusCodec, JsonCodec<VersionedDynamicFilterDomains> dynamicFilterDomainsCodec, JsonCodec<TaskInfo> taskInfoCodec, JsonCodec<TaskUpdateRequest> taskUpdateRequestCodec, JsonCodec<FailTaskRequest> failTaskRequestCodec) {
            JaxrsTestingHttpProcessor jaxrsTestingHttpProcessor = new JaxrsTestingHttpProcessor(URI.create("http://fake.invalid/"), testingTaskResource, jsonMapper);
            TestingHttpClient testingHttpClient = new TestingHttpClient(jaxrsTestingHttpProcessor.setTrace(TRACE_HTTP));
            testingTaskResource.setHttpClient(testingHttpClient);
            return new HttpRemoteTaskFactory(new QueryManagerConfig(), TASK_MANAGER_CONFIG, testingHttpClient, new TestSqlTaskManager.MockLocationFactory(), taskStatusCodec, dynamicFilterDomainsCodec, taskInfoCodec, taskUpdateRequestCodec, failTaskRequestCodec, new RemoteTaskStats(), dynamicFilterService);
        }
    });
    Injector injector = app.doNotInitializeLogging().quiet().initialize();
    return injector.getInstance(HttpRemoteTaskFactory.class);
}

Example 5 with HttpRemoteTaskFactory

use of io.trino.server.HttpRemoteTaskFactory in project trino by trinodb.

the class TestHttpRemoteTask method runTest.

private void runTest(FailureScenario failureScenario) throws Exception {
    AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime());
    TestingTaskResource testingTaskResource = new TestingTaskResource(lastActivityNanos, failureScenario);
    HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource);
    RemoteTask remoteTask = createRemoteTask(httpRemoteTaskFactory, ImmutableSet.of());
    testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo());
    remoteTask.start();
    waitUntilIdle(lastActivityNanos);
    httpRemoteTaskFactory.stop();
    assertTrue(remoteTask.getTaskStatus().getState().isDone(), format("TaskStatus is not in a done state: %s", remoteTask.getTaskStatus()));
    ErrorCode actualErrorCode = getOnlyElement(remoteTask.getTaskStatus().getFailures()).getErrorCode();
    switch(failureScenario) {
        case TASK_MISMATCH:
        case TASK_MISMATCH_WHEN_VERSION_IS_HIGH:
            assertTrue(remoteTask.getTaskInfo().getTaskStatus().getState().isDone(), format("TaskInfo is not in a done state: %s", remoteTask.getTaskInfo()));
            assertEquals(actualErrorCode, REMOTE_TASK_MISMATCH.toErrorCode());
            break;
        case REJECTED_EXECUTION:
            // for a rejection to occur, the http client must be shutdown, which means we will not be able to ge the final task info
            assertEquals(actualErrorCode, REMOTE_TASK_ERROR.toErrorCode());
            break;
        default:
            throw new UnsupportedOperationException();
    }
}