use of io.trino.execution.TaskId in project trino by trinodb.
the class TestDirectExchangeClient method testStreamingBufferLimit.
@Test
public void testStreamingBufferLimit() {
DataSize maxResponseSize = DataSize.ofBytes(1);
MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);
URI location = URI.create("http://localhost:8080");
// add a pages
processor.addPage(location, createPage(1));
processor.addPage(location, createPage(2));
processor.addPage(location, createPage(3));
processor.setComplete(location);
@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() + "-testBufferLimit-%s"))), scheduler, new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), pageBufferClientCallbackExecutor, (taskId, failure) -> {
});
exchangeClient.addLocation(new TaskId(new StageId("query", 1), 0, 0), location);
exchangeClient.noMoreLocations();
assertFalse(exchangeClient.isFinished());
long start = System.nanoTime();
// wait for a page to be fetched
do {
// there is no thread coordination here, so sleep is the best we can do
assertLessThan(Duration.nanosSince(start), new Duration(5, TimeUnit.SECONDS));
sleepUninterruptibly(100, MILLISECONDS);
} while (exchangeClient.getStatus().getBufferedPages() == 0);
// client should have sent a single request for a single page
assertEquals(exchangeClient.getStatus().getBufferedPages(), 1);
assertTrue(exchangeClient.getStatus().getBufferedBytes() > 0);
assertStatus(exchangeClient.getStatus().getPageBufferClientStatuses().get(0), location, "queued", 1, 1, 1, "not scheduled");
// remove the page and wait for the client to fetch another page
assertPageEquals(exchangeClient.pollPage(), createPage(1));
do {
assertLessThan(Duration.nanosSince(start), new Duration(5, TimeUnit.SECONDS));
sleepUninterruptibly(100, MILLISECONDS);
} while (exchangeClient.getStatus().getBufferedPages() == 0);
// client should have sent a single request for a single page
assertStatus(exchangeClient.getStatus().getPageBufferClientStatuses().get(0), location, "queued", 2, 2, 2, "not scheduled");
assertEquals(exchangeClient.getStatus().getBufferedPages(), 1);
assertTrue(exchangeClient.getStatus().getBufferedBytes() > 0);
// remove the page and wait for the client to fetch another page
assertPageEquals(exchangeClient.pollPage(), createPage(2));
do {
assertLessThan(Duration.nanosSince(start), new Duration(5, TimeUnit.SECONDS));
sleepUninterruptibly(100, MILLISECONDS);
} while (exchangeClient.getStatus().getBufferedPages() == 0);
// client should have sent a single request for a single page
assertStatus(exchangeClient.getStatus().getPageBufferClientStatuses().get(0), location, "queued", 3, 3, 3, "not scheduled");
assertEquals(exchangeClient.getStatus().getBufferedPages(), 1);
assertTrue(exchangeClient.getStatus().getBufferedBytes() > 0);
// remove last page
assertPageEquals(getNextPage(exchangeClient), createPage(3));
// wait for client to decide there are no more pages
assertNull(getNextPage(exchangeClient));
assertEquals(exchangeClient.getStatus().getBufferedPages(), 0);
assertTrue(exchangeClient.isFinished());
exchangeClient.close();
assertStatus(exchangeClient.getStatus().getPageBufferClientStatuses().get(0), location, "closed", 3, 5, 5, "not scheduled");
}
use of io.trino.execution.TaskId in project trino by trinodb.
the class TestDirectExchangeClient method testAddLocation.
@Test
public void testAddLocation() 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);
URI location1 = URI.create("http://localhost:8080/1");
URI location2 = URI.create("http://localhost:8080/2");
URI location3 = URI.create("http://localhost:8080/3");
processor.addPage(location1, createSerializedPage(1));
processor.addPage(location1, createSerializedPage(2));
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());
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(2);
assertThat(buffer.getFinishedTasks()).containsExactly(task1);
exchangeClient.addLocation(task2, location2);
assertThat(buffer.getAllTasks()).containsExactlyInAnyOrder(task1, task2);
assertTaskIsNotFinished(buffer, task2);
processor.setComplete(location2);
buffer.whenTaskFinished(task2).get(10, SECONDS);
assertThat(buffer.getFinishedTasks()).containsExactlyInAnyOrder(task1, task2);
assertThat(buffer.getPages().get(task2)).hasSize(0);
exchangeClient.addLocation(task3, location3);
assertThat(buffer.getAllTasks()).containsExactlyInAnyOrder(task1, task2, task3);
assertTaskIsNotFinished(buffer, task3);
exchangeClient.noMoreLocations();
assertTrue(buffer.isNoMoreTasks());
assertThat(buffer.getAllTasks()).containsExactlyInAnyOrder(task1, task2, task3);
assertTaskIsNotFinished(buffer, task3);
exchangeClient.close();
buffer.whenTaskFinished(task3).get(10, SECONDS);
assertThat(buffer.getFinishedTasks()).containsExactlyInAnyOrder(task1, task2, task3);
assertThat(buffer.getFailedTasks().asMap()).isEmpty();
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"));
assertTrue(exchangeClient.isFinished());
}
use of io.trino.execution.TaskId in project trino by trinodb.
the class TestDeduplicatingDirectExchangeBuffer method testRemoteTaskFailedError.
private void testRemoteTaskFailedError(RetryPolicy retryPolicy) {
// fail before noMoreTasks
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, retryPolicy)) {
TaskId taskId = createTaskId(0, 0);
buffer.addTask(taskId);
buffer.taskFailed(taskId, new TrinoException(REMOTE_TASK_FAILED, "Remote task failed"));
buffer.noMoreTasks();
assertFalse(buffer.isFinished());
assertFalse(buffer.isFailed());
assertBlocked(buffer.isBlocked());
assertNull(buffer.pollPage());
}
// fail after noMoreTasks
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, retryPolicy)) {
TaskId taskId = createTaskId(0, 0);
buffer.addTask(taskId);
buffer.noMoreTasks();
buffer.taskFailed(taskId, new TrinoException(REMOTE_TASK_FAILED, "Remote task failed"));
assertFalse(buffer.isFinished());
assertFalse(buffer.isFailed());
assertBlocked(buffer.isBlocked());
assertNull(buffer.pollPage());
}
}
use of io.trino.execution.TaskId in project trino by trinodb.
the class TestDeduplicatingDirectExchangeBuffer method testIsBlocked.
@Test
public void testIsBlocked() {
// immediate close
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, RetryPolicy.QUERY)) {
ListenableFuture<Void> blocked = buffer.isBlocked();
assertBlocked(blocked);
buffer.close();
assertNotBlocked(blocked);
}
// empty set of tasks
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, RetryPolicy.QUERY)) {
ListenableFuture<Void> blocked = buffer.isBlocked();
assertBlocked(blocked);
buffer.noMoreTasks();
assertNotBlocked(blocked);
}
// single task finishes before noMoreTasks
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, RetryPolicy.QUERY)) {
ListenableFuture<Void> blocked = buffer.isBlocked();
assertBlocked(blocked);
TaskId taskId = createTaskId(0, 0);
buffer.addTask(taskId);
assertBlocked(blocked);
buffer.taskFinished(taskId);
assertBlocked(blocked);
buffer.noMoreTasks();
assertNotBlocked(blocked);
}
// single task finishes after noMoreTasks
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, RetryPolicy.QUERY)) {
ListenableFuture<Void> blocked = buffer.isBlocked();
assertBlocked(blocked);
TaskId taskId = createTaskId(0, 0);
buffer.addTask(taskId);
assertBlocked(blocked);
buffer.noMoreTasks();
assertBlocked(blocked);
buffer.taskFinished(taskId);
assertNotBlocked(blocked);
}
// single task fails before noMoreTasks
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, RetryPolicy.QUERY)) {
ListenableFuture<Void> blocked = buffer.isBlocked();
assertBlocked(blocked);
TaskId taskId = createTaskId(0, 0);
buffer.addTask(taskId);
assertBlocked(blocked);
buffer.taskFailed(taskId, new RuntimeException());
assertBlocked(blocked);
buffer.noMoreTasks();
assertNotBlocked(blocked);
}
// single task fails after noMoreTasks
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, RetryPolicy.QUERY)) {
ListenableFuture<Void> blocked = buffer.isBlocked();
assertBlocked(blocked);
TaskId taskId = createTaskId(0, 0);
buffer.addTask(taskId);
assertBlocked(blocked);
buffer.noMoreTasks();
assertBlocked(blocked);
buffer.taskFailed(taskId, new RuntimeException());
assertNotBlocked(blocked);
}
// cancelled blocked future doesn't affect other blocked futures
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, RetryPolicy.QUERY)) {
ListenableFuture<Void> blocked1 = buffer.isBlocked();
ListenableFuture<Void> blocked2 = buffer.isBlocked();
assertBlocked(blocked1);
assertBlocked(blocked2);
blocked2.cancel(true);
assertBlocked(blocked1);
assertNotBlocked(blocked2);
}
}
use of io.trino.execution.TaskId in project trino by trinodb.
the class TestDeduplicatingDirectExchangeBuffer method testRemainingBufferCapacity.
@Test
public void testRemainingBufferCapacity() {
try (DirectExchangeBuffer buffer = createDeduplicatingDirectExchangeBuffer(DEFAULT_BUFFER_CAPACITY, RetryPolicy.QUERY)) {
assertFalse(buffer.isFinished());
TaskId taskId = createTaskId(0, 0);
buffer.addTask(taskId);
Slice page = utf8Slice("page");
buffer.addPages(taskId, ImmutableList.of(page));
assertEquals(buffer.getRemainingCapacityInBytes(), Long.MAX_VALUE);
}
}
Aggregations