Examples with SimpleLocalMemoryContext io.trino.memory.context.SimpleLocalMemoryContext used on opensource projects

Example 1 with SimpleLocalMemoryContext

use of io.trino.memory.context.SimpleLocalMemoryContext in project trino by trinodb.

the class TestDirectExchangeClient method setUpDataCorruption.

private DirectExchangeClient setUpDataCorruption(DataIntegrityVerification dataIntegrityVerification, URI location) {
    DataSize maxResponseSize = DataSize.of(10, Unit.MEGABYTE);
    MockExchangeRequestProcessor delegate = new MockExchangeRequestProcessor(maxResponseSize);
    delegate.addPage(location, createPage(1));
    delegate.addPage(location, createPage(2));
    delegate.setComplete(location);
    TestingHttpClient.Processor processor = new TestingHttpClient.Processor() {

        private int completedRequests;

        private TestingResponse savedResponse;

        @Override
        public synchronized Response handle(Request request) throws Exception {
            if (completedRequests == 0) {
                verify(savedResponse == null);
                TestingResponse response = (TestingResponse) delegate.handle(request);
                checkState(response.getStatusCode() == HttpStatus.OK.code(), "Unexpected status code: %s", response.getStatusCode());
                ListMultimap<String, String> headers = response.getHeaders().entries().stream().collect(toImmutableListMultimap(entry -> entry.getKey().toString(), Map.Entry::getValue));
                byte[] bytes = toByteArray(response.getInputStream());
                checkState(bytes.length > 42, "too short");
                savedResponse = new TestingResponse(HttpStatus.OK, headers, bytes.clone());
                // corrupt
                bytes[42]++;
                completedRequests++;
                return new TestingResponse(HttpStatus.OK, headers, bytes);
            }
            if (completedRequests == 1) {
                verify(savedResponse != null);
                Response response = savedResponse;
                savedResponse = null;
                completedRequests++;
                return response;
            }
            completedRequests++;
            return delegate.handle(request);
        }
    };
    DirectExchangeClient exchangeClient = new DirectExchangeClient("localhost", dataIntegrityVerification, new StreamingDirectExchangeBuffer(scheduler, DataSize.of(32, Unit.MEGABYTE)), maxResponseSize, 1, new Duration(1, TimeUnit.MINUTES), true, new TestingHttpClient(processor, scheduler), scheduler, new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), pageBufferClientCallbackExecutor, (taskId, failure) -> {
    });
    exchangeClient.addLocation(new TaskId(new StageId("query", 1), 0, 0), location);
    exchangeClient.noMoreLocations();
    return exchangeClient;
}

Example 2 with SimpleLocalMemoryContext

use of io.trino.memory.context.SimpleLocalMemoryContext in project trino by trinodb.

the class TestDirectExchangeClient method testStreamingTaskFailure.

@Test
public void testStreamingTaskFailure() {
    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);
    URI location1 = URI.create("http://localhost:8080/1");
    URI location2 = URI.create("http://localhost:8080/2");
    processor.addPage(location1, createPage(1));
    StreamingDirectExchangeBuffer buffer = new StreamingDirectExchangeBuffer(scheduler, DataSize.of(1, Unit.MEGABYTE));
    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) -> {
    });
    exchangeClient.addLocation(task1, location1);
    exchangeClient.addLocation(task2, location2);
    assertPageEquals(getNextPage(exchangeClient), createPage(1));
    processor.setComplete(location1);
    assertFalse(tryGetFutureValue(exchangeClient.isBlocked(), 10, MILLISECONDS).isPresent());
    RuntimeException randomException = new RuntimeException("randomfailure");
    processor.setFailed(location2, randomException);
    assertThatThrownBy(() -> getNextPage(exchangeClient)).hasMessageContaining("Encountered too many errors talking to a worker node");
    assertFalse(exchangeClient.isFinished());
}

Example 3 with SimpleLocalMemoryContext

use of io.trino.memory.context.SimpleLocalMemoryContext 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");
}

Example 4 with SimpleLocalMemoryContext

use of io.trino.memory.context.SimpleLocalMemoryContext 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());
}

Example 5 with SimpleLocalMemoryContext

use of io.trino.memory.context.SimpleLocalMemoryContext 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());
}