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

Example 11 with SimpleLocalMemoryContext

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

the class TestDirectExchangeClient method testStreamingAddLocation.

@Test(timeOut = 10000)
public void testStreamingAddLocation() throws Exception {
    DataSize maxResponseSize = DataSize.of(10, Unit.MEGABYTE);
    MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);
    @SuppressWarnings("resource") DirectExchangeClient exchangeClient = new DirectExchangeClient("localhost", DataIntegrityVerification.ABORT, new StreamingDirectExchangeBuffer(scheduler, DataSize.of(32, Unit.MEGABYTE)), maxResponseSize, 1, new Duration(1, TimeUnit.MINUTES), true, new TestingHttpClient(processor, newCachedThreadPool(daemonThreadsNamed(getClass().getSimpleName() + "-testAddLocation-%s"))), scheduler, new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), pageBufferClientCallbackExecutor, (taskId, failure) -> {
    });
    URI location1 = URI.create("http://localhost:8081/foo");
    processor.addPage(location1, createPage(1));
    processor.addPage(location1, createPage(2));
    processor.addPage(location1, createPage(3));
    processor.setComplete(location1);
    exchangeClient.addLocation(new TaskId(new StageId("query", 1), 0, 0), location1);
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(1));
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(2));
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(3));
    assertNull(exchangeClient.pollPage());
    ListenableFuture<Void> firstBlocked = exchangeClient.isBlocked();
    assertFalse(tryGetFutureValue(firstBlocked, 10, MILLISECONDS).isPresent());
    assertFalse(firstBlocked.isDone());
    assertNull(exchangeClient.pollPage());
    ListenableFuture<Void> secondBlocked = exchangeClient.isBlocked();
    assertFalse(tryGetFutureValue(secondBlocked, 10, MILLISECONDS).isPresent());
    assertFalse(secondBlocked.isDone());
    assertNull(exchangeClient.pollPage());
    ListenableFuture<Void> thirdBlocked = exchangeClient.isBlocked();
    assertFalse(tryGetFutureValue(thirdBlocked, 10, MILLISECONDS).isPresent());
    assertFalse(thirdBlocked.isDone());
    thirdBlocked.cancel(true);
    assertTrue(thirdBlocked.isDone());
    assertFalse(tryGetFutureValue(firstBlocked, 10, MILLISECONDS).isPresent());
    assertFalse(firstBlocked.isDone());
    assertFalse(tryGetFutureValue(secondBlocked, 10, MILLISECONDS).isPresent());
    assertFalse(secondBlocked.isDone());
    assertFalse(exchangeClient.isFinished());
    URI location2 = URI.create("http://localhost:8082/bar");
    processor.addPage(location2, createPage(4));
    processor.addPage(location2, createPage(5));
    processor.addPage(location2, createPage(6));
    processor.setComplete(location2);
    exchangeClient.addLocation(new TaskId(new StageId("query", 1), 1, 0), location2);
    tryGetFutureValue(firstBlocked, 5, SECONDS);
    assertTrue(firstBlocked.isDone());
    tryGetFutureValue(secondBlocked, 5, SECONDS);
    assertTrue(secondBlocked.isDone());
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(4));
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(5));
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(6));
    assertFalse(tryGetFutureValue(exchangeClient.isBlocked(), 10, MILLISECONDS).isPresent());
    assertFalse(exchangeClient.isFinished());
    exchangeClient.noMoreLocations();
    // receive a final GONE response, so just spin until it's closed or the test times out.
    while (!exchangeClient.isFinished()) {
        Thread.sleep(1);
    }
    exchangeClient.close();
    ImmutableMap<URI, PageBufferClientStatus> statuses = uniqueIndex(exchangeClient.getStatus().getPageBufferClientStatuses(), PageBufferClientStatus::getUri);
    assertStatus(statuses.get(location1), location1, "closed", 3, 3, 3, "not scheduled");
    assertStatus(statuses.get(location2), location2, "closed", 3, 3, 3, "not scheduled");
}

Example 12 with SimpleLocalMemoryContext

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

the class TestDirectExchangeClient method testDeduplication.

@Test
public void testDeduplication() throws Exception {
    DataSize maxResponseSize = DataSize.of(10, Unit.MEGABYTE);
    MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);
    TaskId taskP0A0 = new TaskId(new StageId("query", 1), 0, 0);
    TaskId taskP1A0 = new TaskId(new StageId("query", 1), 1, 0);
    TaskId taskP0A1 = new TaskId(new StageId("query", 1), 0, 1);
    URI locationP0A0 = URI.create("http://localhost:8080/1");
    URI locationP1A0 = URI.create("http://localhost:8080/2");
    URI locationP0A1 = URI.create("http://localhost:8080/3");
    processor.addPage(locationP1A0, createSerializedPage(1));
    processor.addPage(locationP0A1, createSerializedPage(2));
    processor.addPage(locationP0A1, createSerializedPage(3));
    @SuppressWarnings("resource") DirectExchangeClient exchangeClient = new DirectExchangeClient("localhost", DataIntegrityVerification.ABORT, new DeduplicatingDirectExchangeBuffer(scheduler, DataSize.of(1, Unit.KILOBYTE), RetryPolicy.QUERY, new ExchangeManagerRegistry(new ExchangeHandleResolver()), new QueryId("query"), createRandomExchangeId()), maxResponseSize, 1, new Duration(1, SECONDS), true, new TestingHttpClient(processor, scheduler), scheduler, new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), pageBufferClientCallbackExecutor, (taskId, failure) -> {
    });
    exchangeClient.addLocation(taskP0A0, locationP0A0);
    exchangeClient.addLocation(taskP1A0, locationP1A0);
    exchangeClient.addLocation(taskP0A1, locationP0A1);
    processor.setComplete(locationP0A0);
    // Failing attempt 0. Results from all tasks for attempt 0 must be discarded.
    processor.setFailed(locationP1A0, new RuntimeException("failure"));
    processor.setComplete(locationP0A1);
    assertFalse(exchangeClient.isFinished());
    assertThatThrownBy(() -> exchangeClient.isBlocked().get(50, MILLISECONDS)).isInstanceOf(TimeoutException.class);
    exchangeClient.noMoreLocations();
    exchangeClient.isBlocked().get(10, SECONDS);
    List<Page> pages = new ArrayList<>();
    while (!exchangeClient.isFinished()) {
        Slice page = exchangeClient.pollPage();
        if (page == null) {
            break;
        }
        pages.add(PAGES_SERDE.deserialize(page));
    }
    assertThat(pages).hasSize(2);
    assertThat(pages.stream().map(Page::getPositionCount).collect(toImmutableSet())).containsAll(ImmutableList.of(2, 3));
    assertEventually(() -> assertTrue(exchangeClient.isFinished()));
    assertEventually(() -> {
        assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(0).getHttpRequestState(), "not scheduled", "httpRequestState");
        assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(1).getHttpRequestState(), "not scheduled", "httpRequestState");
        assertEquals(exchangeClient.getStatus().getPageBufferClientStatuses().get(2).getHttpRequestState(), "not scheduled", "httpRequestState");
    });
    exchangeClient.close();
}

Example 13 with SimpleLocalMemoryContext

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

the class TestDirectExchangeClient method testStreamingHappyPath.

@Test
public void testStreamingHappyPath() {
    DataSize maxResponseSize = DataSize.of(10, Unit.MEGABYTE);
    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));
    processor.setComplete(location);
    @SuppressWarnings("resource") DirectExchangeClient exchangeClient = new DirectExchangeClient("localhost", DataIntegrityVerification.ABORT, 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();
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(1));
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(2));
    assertFalse(exchangeClient.isFinished());
    assertPageEquals(getNextPage(exchangeClient), createPage(3));
    assertNull(getNextPage(exchangeClient));
    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();
}

Example 14 with SimpleLocalMemoryContext

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

the class TestDirectExchangeClient method testDeduplicationTaskFailure.

@Test
public void testDeduplicationTaskFailure() throws Exception {
    DataSize maxResponseSize = DataSize.of(10, Unit.MEGABYTE);
    MockExchangeRequestProcessor processor = new MockExchangeRequestProcessor(maxResponseSize);
    TaskId attempt0Task1 = new TaskId(new StageId("query", 1), 0, 0);
    TaskId attempt1Task1 = new TaskId(new StageId("query", 1), 1, 0);
    TaskId attempt1Task2 = new TaskId(new StageId("query", 1), 2, 0);
    URI attempt0Task1Location = URI.create("http://localhost:8080/1/0");
    URI attempt1Task1Location = URI.create("http://localhost:8080/1/1");
    URI attempt1Task2Location = URI.create("http://localhost:8080/2/1");
    processor.setFailed(attempt0Task1Location, new RuntimeException("randomfailure"));
    processor.addPage(attempt1Task1Location, createPage(1));
    processor.setComplete(attempt1Task1Location);
    processor.setFailed(attempt1Task2Location, new RuntimeException("randomfailure"));
    DeduplicatingDirectExchangeBuffer buffer = new DeduplicatingDirectExchangeBuffer(scheduler, DataSize.of(1, Unit.MEGABYTE), RetryPolicy.QUERY, new ExchangeManagerRegistry(new ExchangeHandleResolver()), new QueryId("query"), createRandomExchangeId());
    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(attempt0Task1, attempt0Task1Location);
    assertFalse(tryGetFutureValue(exchangeClient.isBlocked(), 10, MILLISECONDS).isPresent());
    exchangeClient.addLocation(attempt1Task1, attempt1Task1Location);
    exchangeClient.addLocation(attempt1Task2, attempt1Task2Location);
    assertFalse(tryGetFutureValue(exchangeClient.isBlocked(), 10, MILLISECONDS).isPresent());
    exchangeClient.noMoreLocations();
    exchangeClient.isBlocked().get(10, SECONDS);
    assertThatThrownBy(() -> getNextPage(exchangeClient)).hasMessageContaining("Encountered too many errors talking to a worker node");
    assertFalse(exchangeClient.isFinished());
}