Examples with StageId io.trino.execution.StageId used on opensource projects

Example 6 with StageId

use of io.trino.execution.StageId in project trino by trinodb.

the class TestDirectExchangeClient method testTaskFailure.

@Test
public void testTaskFailure() throws Exception {
    DataSize maxResponseSize = DataSize.of(10, Unit.MEGABYTE);
    MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);
    TaskId task1 = new TaskId(new StageId("query", 1), 0, 0);
    TaskId task2 = new TaskId(new StageId("query", 1), 1, 0);
    TaskId task3 = new TaskId(new StageId("query", 1), 2, 0);
    TaskId task4 = new TaskId(new StageId("query", 1), 3, 0);
    URI location1 = URI.create("http://localhost:8080/1");
    URI location2 = URI.create("http://localhost:8080/2");
    URI location3 = URI.create("http://localhost:8080/3");
    URI location4 = URI.create("http://localhost:8080/4");
    processor.addPage(location1, createSerializedPage(1));
    processor.addPage(location4, createSerializedPage(2));
    processor.addPage(location4, createSerializedPage(3));
    TestingDirectExchangeBuffer buffer = new TestingDirectExchangeBuffer(DataSize.of(1, Unit.MEGABYTE));
    Set<TaskId> failedTasks = newConcurrentHashSet();
    CountDownLatch latch = new CountDownLatch(2);
    @SuppressWarnings("resource") DirectExchangeClient exchangeClient = new DirectExchangeClient("localhost", DataIntegrityVerification.ABORT, buffer, maxResponseSize, 1, new Duration(1, SECONDS), true, new TestingHttpClient(processor, scheduler), scheduler, new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), pageBufferClientCallbackExecutor, (taskId, failure) -> {
        failedTasks.add(taskId);
        latch.countDown();
    });
    assertThat(buffer.getAllTasks()).isEmpty();
    assertThat(buffer.getPages().asMap()).isEmpty();
    assertThat(buffer.getFinishedTasks()).isEmpty();
    assertThat(buffer.getFailedTasks().asMap()).isEmpty();
    assertFalse(buffer.isNoMoreTasks());
    exchangeClient.addLocation(task1, location1);
    assertThat(buffer.getAllTasks()).containsExactly(task1);
    assertTaskIsNotFinished(buffer, task1);
    processor.setComplete(location1);
    buffer.whenTaskFinished(task1).get(10, SECONDS);
    assertThat(buffer.getPages().get(task1)).hasSize(1);
    assertThat(buffer.getFinishedTasks()).containsExactly(task1);
    exchangeClient.addLocation(task2, location2);
    assertThat(buffer.getAllTasks()).containsExactlyInAnyOrder(task1, task2);
    assertTaskIsNotFinished(buffer, task2);
    RuntimeException randomException = new RuntimeException("randomfailure");
    processor.setFailed(location2, randomException);
    buffer.whenTaskFailed(task2).get(10, SECONDS);
    assertThat(buffer.getFinishedTasks()).containsExactly(task1);
    assertThat(buffer.getFailedTasks().keySet()).containsExactly(task2);
    assertThat(buffer.getPages().get(task2)).hasSize(0);
    exchangeClient.addLocation(task3, location3);
    assertThat(buffer.getAllTasks()).containsExactlyInAnyOrder(task1, task2, task3);
    assertTaskIsNotFinished(buffer, task2);
    assertTaskIsNotFinished(buffer, task3);
    TrinoException trinoException = new TrinoException(GENERIC_INTERNAL_ERROR, "generic internal error");
    processor.setFailed(location3, trinoException);
    buffer.whenTaskFailed(task3).get(10, SECONDS);
    assertThat(buffer.getFinishedTasks()).containsExactly(task1);
    assertThat(buffer.getFailedTasks().keySet()).containsExactlyInAnyOrder(task2, task3);
    assertThat(buffer.getPages().get(task2)).hasSize(0);
    assertThat(buffer.getPages().get(task3)).hasSize(0);
    assertTrue(latch.await(10, SECONDS));
    assertEquals(failedTasks, ImmutableSet.of(task2, task3));
    exchangeClient.addLocation(task4, location4);
    assertThat(buffer.getAllTasks()).containsExactlyInAnyOrder(task1, task2, task3, task4);
    assertTaskIsNotFinished(buffer, task4);
    processor.setComplete(location4);
    buffer.whenTaskFinished(task4).get(10, SECONDS);
    assertThat(buffer.getPages().get(task4)).hasSize(2);
    assertThat(buffer.getFinishedTasks()).containsExactlyInAnyOrder(task1, task4);
    assertFalse(exchangeClient.isFinished());
    buffer.setFinished(true);
    assertTrue(exchangeClient.isFinished());
    exchangeClient.close();
    assertEventually(() -> assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(0).getHttpRequestState(), "not scheduled", "httpRequestState"));
    assertEventually(() -> assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(1).getHttpRequestState(), "not scheduled", "httpRequestState"));
    assertEventually(() -> assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(2).getHttpRequestState(), "not scheduled", "httpRequestState"));
    assertEventually(() -> assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(3).getHttpRequestState(), "not scheduled", "httpRequestState"));
    assertThat(buffer.getFinishedTasks()).containsExactlyInAnyOrder(task1, task4);
    assertThat(buffer.getFailedTasks().keySet()).containsExactlyInAnyOrder(task2, task3);
    assertThat(buffer.getFailedTasks().asMap().get(task2)).hasSize(1);
    assertThat(buffer.getFailedTasks().asMap().get(task2).iterator().next()).isInstanceOf(TrinoTransportException.class);
    assertThat(buffer.getFailedTasks().asMap().get(task3)).hasSize(1);
    assertThat(buffer.getFailedTasks().asMap().get(task3).iterator().next()).isEqualTo(trinoException);
    assertTrue(exchangeClient.isFinished());
}

Example 7 with StageId

use of io.trino.execution.StageId in project trino by trinodb.

the class TestDirectExchangeClient method testStreamingClose.

@Test
public void testStreamingClose() throws Exception {
    DataSize maxResponseSize = DataSize.ofBytes(1);
    MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);
    URI location = URI.create("http://localhost:8080");
    processor.addPage(location, createPage(1));
    processor.addPage(location, createPage(2));
    processor.addPage(location, createPage(3));
    @SuppressWarnings("resource") DirectExchangeClient exchangeClient = new DirectExchangeClient("localhost", DataIntegrityVerification.ABORT, new StreamingDirectExchangeBuffer(scheduler, DataSize.ofBytes(1)), maxResponseSize, 1, new Duration(1, TimeUnit.MINUTES), true, new TestingHttpClient(processor, newCachedThreadPool(daemonThreadsNamed(getClass().getSimpleName() + "-testClose-%s"))), scheduler, new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), pageBufferClientCallbackExecutor, (taskId, failure) -> {
    });
    exchangeClient.addLocation(new TaskId(new StageId("query", 1), 0, 0), location);
    exchangeClient.noMoreLocations();
    // fetch a page
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(1));
    // close client while pages are still available
    exchangeClient.close();
    while (!exchangeClient.isFinished()) {
        MILLISECONDS.sleep(10);
    }
    assertTrue(exchangeClient.isFinished());
    assertNull(exchangeClient.pollPage());
    assertEquals(exchangeClient.getStatus().getBufferedPages(), 0);
    PageBufferClientStatus clientStatus = exchangeClient.getStatus().getPageBufferClientStatuses().get(0);
    assertEquals(clientStatus.getUri(), location);
    assertEquals(clientStatus.getState(), "closed", "status");
    assertEquals(clientStatus.getHttpRequestState(), "not scheduled", "httpRequestState");
}

Example 8 with StageId

use of io.trino.execution.StageId in project trino by trinodb.

the class TestDirectExchangeClient method testHappyPath.

@Test
public void testHappyPath() throws Exception {
    DataSize maxResponseSize = DataSize.of(10, Unit.MEGABYTE);
    MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);
    List<Slice> pages = ImmutableList.of(createSerializedPage(1), createSerializedPage(2), createSerializedPage(3));
    URI location = URI.create("http://localhost:8080");
    pages.forEach(page -> processor.addPage(location, page));
    processor.setComplete(location);
    TestingDirectExchangeBuffer buffer = new TestingDirectExchangeBuffer(DataSize.of(1, Unit.MEGABYTE));
    @SuppressWarnings("resource") DirectExchangeClient exchangeClient = new DirectExchangeClient("localhost", DataIntegrityVerification.ABORT, buffer, maxResponseSize, 1, new Duration(1, TimeUnit.MINUTES), true, new TestingHttpClient(processor, scheduler), scheduler, new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), pageBufferClientCallbackExecutor, (taskId, failure) -> {
    });
    assertThat(buffer.getAllTasks()).isEmpty();
    assertThat(buffer.getPages().asMap()).isEmpty();
    assertThat(buffer.getFinishedTasks()).isEmpty();
    assertThat(buffer.getFailedTasks().asMap()).isEmpty();
    assertFalse(buffer.isNoMoreTasks());
    TaskId taskId = new TaskId(new StageId("query", 1), 0, 0);
    exchangeClient.addLocation(taskId, location);
    assertThat(buffer.getAllTasks()).containsExactly(taskId);
    exchangeClient.noMoreLocations();
    assertTrue(buffer.isNoMoreTasks());
    buffer.whenTaskFinished(taskId).get(10, SECONDS);
    assertThat(buffer.getFinishedTasks()).containsExactly(taskId);
    assertThat(buffer.getPages().get(taskId)).hasSize(3);
    assertThat(buffer.getFailedTasks().asMap()).isEmpty();
    assertFalse(exchangeClient.isFinished());
    buffer.setFinished(true);
    assertTrue(exchangeClient.isFinished());
    DirectExchangeClientStatus status = exchangeClient.getStatus();
    assertEquals(status.getBufferedPages(), 0);
    // client should have sent only 3 requests: one to get all pages, one to acknowledge and one to get the done signal
    assertStatus(status.getPageBufferClientStatuses().get(0), location, "closed", 3, 3, 3, "not scheduled");
    exchangeClient.close();
    assertEventually(() -> assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(0).getHttpRequestState(), "not scheduled", "httpRequestState"));
    assertThat(buffer.getFinishedTasks()).containsExactly(taskId);
    assertThat(buffer.getFailedTasks().asMap()).isEmpty();
    assertThat(buffer.getPages().size()).isEqualTo(3);
}

Example 9 with StageId

use of io.trino.execution.StageId in project trino by trinodb.

the class TestDynamicFilterService method testDynamicFilterConsumerCallbackCount.

@Test
public void testDynamicFilterConsumerCallbackCount() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    DynamicFilterId filterId2 = new DynamicFilterId("df2");
    Set<DynamicFilterId> dynamicFilters = ImmutableSet.of(filterId1, filterId2);
    QueryId queryId = new QueryId("query");
    StageId stageId = new StageId(queryId, 0);
    dynamicFilterService.registerQuery(queryId, session, dynamicFilters, dynamicFilters, ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId, 0, 2);
    Map<DynamicFilterId, Domain> consumerCollectedFilters = new HashMap<>();
    AtomicInteger callbackCount = new AtomicInteger();
    dynamicFilterService.registerDynamicFilterConsumer(queryId, 0, dynamicFilters, domains -> {
        callbackCount.getAndIncrement();
        domains.forEach((filter, domain) -> assertNull(consumerCollectedFilters.put(filter, domain)));
    });
    assertTrue(consumerCollectedFilters.isEmpty());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 0, 0), ImmutableMap.of(filterId1, singleValue(INTEGER, 1L), filterId2, singleValue(INTEGER, 2L)));
    assertTrue(consumerCollectedFilters.isEmpty());
    // complete both filterId1 and filterId2
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 1, 0), ImmutableMap.of(filterId1, singleValue(INTEGER, 3L), filterId2, singleValue(INTEGER, 4L)));
    assertEquals(consumerCollectedFilters, ImmutableMap.of(filterId1, multipleValues(INTEGER, ImmutableList.of(1L, 3L)), filterId2, multipleValues(INTEGER, ImmutableList.of(2L, 4L))));
    // both filters should be received in single callback
    assertEquals(callbackCount.get(), 1);
    // register another consumer after both filters have been collected
    Map<DynamicFilterId, Domain> secondConsumerCollectedFilters = new HashMap<>();
    AtomicInteger secondCallbackCount = new AtomicInteger();
    dynamicFilterService.registerDynamicFilterConsumer(queryId, 0, dynamicFilters, domains -> {
        secondCallbackCount.getAndIncrement();
        domains.forEach((filter, domain) -> assertNull(secondConsumerCollectedFilters.put(filter, domain)));
    });
    assertEquals(secondConsumerCollectedFilters, ImmutableMap.of(filterId1, multipleValues(INTEGER, ImmutableList.of(1L, 3L)), filterId2, multipleValues(INTEGER, ImmutableList.of(2L, 4L))));
    // both filters should be received by second consumer in single callback
    assertEquals(secondCallbackCount.get(), 1);
    // first consumer should not receive callback again since it already got the completed filter
    assertEquals(callbackCount.get(), 1);
}

Example 10 with StageId

use of io.trino.execution.StageId in project trino by trinodb.

the class TestDynamicFilterService method testMultipleAttempts.

@Test
public void testMultipleAttempts() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId = new DynamicFilterId("df");
    QueryId queryId = new QueryId("query");
    StageId stageId = new StageId(queryId, 0);
    dynamicFilterService.registerQuery(queryId, session, ImmutableSet.of(filterId), ImmutableSet.of(filterId), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId, 0, 3);
    assertFalse(dynamicFilterService.getSummary(queryId, filterId).isPresent());
    // Collect DF from 2 tasks
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 0, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 1L)));
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 1, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 2L)));
    assertFalse(dynamicFilterService.getSummary(queryId, filterId).isPresent());
    // Register query retry
    dynamicFilterService.registerQueryRetry(queryId, 1);
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId, 1, 3);
    // Ignore update from previous attempt
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 2, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 3L)));
    assertFalse(dynamicFilterService.getSummary(queryId, filterId).isPresent());
    // Collect DF from 3 tasks in new attempt
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 0, 1), ImmutableMap.of(filterId, singleValue(INTEGER, 4L)));
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 1, 1), ImmutableMap.of(filterId, singleValue(INTEGER, 5L)));
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 2, 1), ImmutableMap.of(filterId, singleValue(INTEGER, 6L)));
    assertEquals(dynamicFilterService.getSummary(queryId, filterId), Optional.of(multipleValues(INTEGER, ImmutableList.of(4L, 5L, 6L))));
    DynamicFiltersStats stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 1);
    assertEquals(stats.getLazyDynamicFilters(), 1);
    assertEquals(stats.getReplicatedDynamicFilters(), 0);
    assertEquals(stats.getDynamicFilterDomainStats(), ImmutableList.of(new DynamicFilterDomainStats(filterId, getSimplifiedDomainString(4L, 6L, 3, INTEGER))));
}