use of io.trino.execution.TaskId in project trino by trinodb.
the class TestDeduplicatingDirectExchangeBuffer method testExchangeManagerNotConfigured.
@Test
public void testExchangeManagerNotConfigured() {
// no overflow
try (DirectExchangeBuffer buffer = new DeduplicatingDirectExchangeBuffer(directExecutor(), DataSize.of(100, BYTE), RetryPolicy.QUERY, new ExchangeManagerRegistry(new ExchangeHandleResolver()), new QueryId("query"), createRandomExchangeId())) {
TaskId task = createTaskId(0, 0);
Slice page = createPage("1234", DataSize.of(10, BYTE));
buffer.addTask(task);
buffer.addPages(task, ImmutableList.of(page));
buffer.taskFinished(task);
buffer.noMoreTasks();
assertFalse(buffer.isFinished());
assertNotBlocked(buffer.isBlocked());
assertEquals(buffer.pollPage(), page);
assertNull(buffer.pollPage());
assertTrue(buffer.isFinished());
}
// overflow
try (DirectExchangeBuffer buffer = new DeduplicatingDirectExchangeBuffer(directExecutor(), DataSize.of(100, BYTE), RetryPolicy.QUERY, new ExchangeManagerRegistry(new ExchangeHandleResolver()), new QueryId("query"), createRandomExchangeId())) {
TaskId task = createTaskId(0, 0);
Slice page1 = utf8Slice("1234");
Slice page2 = utf8Slice("123456789");
assertThat(page1.getRetainedSize()).isLessThanOrEqualTo(100);
assertThat(page1.getRetainedSize() + page2.getRetainedSize()).isGreaterThan(100);
buffer.addTask(task);
buffer.addPages(task, ImmutableList.of(page1));
assertFalse(buffer.isFinished());
assertBlocked(buffer.isBlocked());
assertEquals(buffer.getRetainedSizeInBytes(), page1.getRetainedSize());
buffer.addPages(task, ImmutableList.of(page2));
assertFalse(buffer.isFinished());
assertTrue(buffer.isFailed());
assertNotBlocked(buffer.isBlocked());
assertEquals(buffer.getRetainedSizeInBytes(), 0);
assertEquals(buffer.getBufferedPageCount(), 0);
assertThatThrownBy(buffer::pollPage).isInstanceOf(TrinoException.class);
}
}
use of io.trino.execution.TaskId 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());
}
use of io.trino.execution.TaskId 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");
}
use of io.trino.execution.TaskId 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);
}
use of io.trino.execution.TaskId 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);
}
Aggregations