Examples with BufferResult io.trino.execution.buffer.BufferResult used on opensource projects

Example 1 with BufferResult

use of io.trino.execution.buffer.BufferResult in project trino by trinodb.

the class TestSqlTask method testBufferCloseOnFinish.

@Test
public void testBufferCloseOnFinish() throws Exception {
    SqlTask sqlTask = createInitialTask();
    OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds();
    updateTask(sqlTask, EMPTY_SPLIT_ASSIGNMENTS, outputBuffers);
    ListenableFuture<BufferResult> bufferResult = sqlTask.getTaskResults(OUT, 0, DataSize.of(1, MEGABYTE));
    assertFalse(bufferResult.isDone());
    // close the sources (no splits will ever be added)
    updateTask(sqlTask, ImmutableList.of(new SplitAssignment(TABLE_SCAN_NODE_ID, ImmutableSet.of(), true)), outputBuffers);
    // finish the task by calling abort on it
    sqlTask.destroyTaskResults(OUT);
    // buffer will be closed by cancel event (wait for event to fire)
    bufferResult.get(1, SECONDS);
    // verify the buffer is closed
    bufferResult = sqlTask.getTaskResults(OUT, 0, DataSize.of(1, MEGABYTE));
    assertTrue(bufferResult.isDone());
    assertTrue(bufferResult.get().isBufferComplete());
}

Example 2 with BufferResult

use of io.trino.execution.buffer.BufferResult in project trino by trinodb.

the class TestSqlTaskManager method testSimpleQuery.

@Test(timeOut = 30_000)
public void testSimpleQuery() throws Exception {
    try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) {
        TaskId taskId = TASK_ID;
        createTask(sqlTaskManager, taskId, ImmutableSet.of(SPLIT), createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds());
        TaskInfo taskInfo = sqlTaskManager.getTaskInfo(taskId, TaskStatus.STARTING_VERSION).get();
        assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FLUSHING);
        BufferResult results = sqlTaskManager.getTaskResults(taskId, OUT, 0, DataSize.of(1, Unit.MEGABYTE)).get();
        assertFalse(results.isBufferComplete());
        assertEquals(results.getSerializedPages().size(), 1);
        assertEquals(getSerializedPagePositionCount(results.getSerializedPages().get(0)), 1);
        for (boolean moreResults = true; moreResults; moreResults = !results.isBufferComplete()) {
            results = sqlTaskManager.getTaskResults(taskId, OUT, results.getToken() + results.getSerializedPages().size(), DataSize.of(1, Unit.MEGABYTE)).get();
        }
        assertTrue(results.isBufferComplete());
        assertEquals(results.getSerializedPages().size(), 0);
        // complete the task by calling destroy on it
        TaskInfo info = sqlTaskManager.destroyTaskResults(taskId, OUT);
        assertEquals(info.getOutputBuffers().getState(), BufferState.FINISHED);
        taskInfo = sqlTaskManager.getTaskInfo(taskId, taskInfo.getTaskStatus().getVersion()).get();
        assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED);
        taskInfo = sqlTaskManager.getTaskInfo(taskId);
        assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED);
    }
}

Example 3 with BufferResult

use of io.trino.execution.buffer.BufferResult in project trino by trinodb.

the class MockExchangeRequestProcessor method handle.

@Override
public Response handle(Request request) {
    if (request.getMethod().equalsIgnoreCase("DELETE")) {
        return new TestingResponse(HttpStatus.NO_CONTENT, ImmutableListMultimap.of(), new byte[0]);
    }
    // verify we got a data size and it parses correctly
    assertTrue(!request.getHeaders().get(InternalHeaders.TRINO_MAX_SIZE).isEmpty());
    DataSize maxSize = DataSize.valueOf(request.getHeader(InternalHeaders.TRINO_MAX_SIZE));
    assertEquals(maxSize, expectedMaxSize);
    RequestLocation requestLocation = new RequestLocation(request.getUri());
    URI location = requestLocation.getLocation();
    BufferResult result = buffers.getUnchecked(location).getPages(requestLocation.getSequenceId(), maxSize);
    byte[] bytes = new byte[0];
    HttpStatus status;
    if (!result.getSerializedPages().isEmpty()) {
        DynamicSliceOutput sliceOutput = new DynamicSliceOutput(64);
        sliceOutput.writeInt(SERIALIZED_PAGES_MAGIC);
        sliceOutput.writeLong(calculateChecksum(result.getSerializedPages()));
        sliceOutput.writeInt(result.getSerializedPages().size());
        for (Slice page : result.getSerializedPages()) {
            sliceOutput.writeBytes(page);
        }
        bytes = sliceOutput.slice().getBytes();
        status = HttpStatus.OK;
    } else {
        status = HttpStatus.NO_CONTENT;
    }
    return new TestingResponse(status, ImmutableListMultimap.<String, String>builder().put(CONTENT_TYPE, TRINO_PAGES).put(TRINO_TASK_INSTANCE_ID, String.valueOf(result.getTaskInstanceId())).put(TRINO_PAGE_TOKEN, String.valueOf(result.getToken())).put(TRINO_PAGE_NEXT_TOKEN, String.valueOf(result.getNextToken())).put(TRINO_BUFFER_COMPLETE, String.valueOf(result.isBufferComplete())).put(TRINO_TASK_FAILED, "false").build(), bytes);
}

Example 4 with BufferResult

use of io.trino.execution.buffer.BufferResult in project trino by trinodb.

the class TestSqlTask method testBufferNotCloseOnFail.

@Test(timeOut = 30_000)
public void testBufferNotCloseOnFail() throws Exception {
    SqlTask sqlTask = createInitialTask();
    updateTask(sqlTask, EMPTY_SPLIT_ASSIGNMENTS, createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds());
    ListenableFuture<BufferResult> bufferResult = sqlTask.getTaskResults(OUT, 0, DataSize.of(1, MEGABYTE));
    assertFalse(bufferResult.isDone());
    long taskStatusVersion = sqlTask.getTaskInfo().getTaskStatus().getVersion();
    sqlTask.failed(new Exception("test"));
    assertEquals(sqlTask.getTaskInfo(taskStatusVersion).get().getTaskStatus().getState(), TaskState.FAILED);
    // buffer will not be closed by fail event.  event is async so wait a bit for event to fire
    assertThatThrownBy(() -> bufferResult.get(1, SECONDS)).isInstanceOf(TimeoutException.class).hasMessageContaining("Waited 1 seconds");
    assertFalse(sqlTask.getTaskResults(OUT, 0, DataSize.of(1, MEGABYTE)).isDone());
}

Example 5 with BufferResult

use of io.trino.execution.buffer.BufferResult in project trino by trinodb.

the class TestSqlTask method testBufferCloseOnCancel.

@Test
public void testBufferCloseOnCancel() throws Exception {
    SqlTask sqlTask = createInitialTask();
    updateTask(sqlTask, EMPTY_SPLIT_ASSIGNMENTS, createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds());
    ListenableFuture<BufferResult> bufferResult = sqlTask.getTaskResults(OUT, 0, DataSize.of(1, MEGABYTE));
    assertFalse(bufferResult.isDone());
    sqlTask.cancel();
    assertEquals(sqlTask.getTaskInfo().getTaskStatus().getState(), TaskState.CANCELED);
    // buffer future will complete.. the event is async so wait a bit for event to propagate
    bufferResult.get(1, SECONDS);
    bufferResult = sqlTask.getTaskResults(OUT, 0, DataSize.of(1, MEGABYTE));
    assertTrue(bufferResult.isDone());
    assertTrue(bufferResult.get().isBufferComplete());
}