Examples with OutputBufferId io.trino.execution.buffer.OutputBuffers.OutputBufferId used on opensource projects

Example 1 with OutputBufferId

use of io.trino.execution.buffer.OutputBuffers.OutputBufferId in project trino by trinodb.

the class ArbitraryOutputBuffer method setOutputBuffers.

@Override
public void setOutputBuffers(OutputBuffers newOutputBuffers) {
    checkState(!Thread.holdsLock(this), "Cannot set output buffers while holding a lock on this");
    requireNonNull(newOutputBuffers, "newOutputBuffers is null");
    synchronized (this) {
        // ignore buffers added after query finishes, which can happen when a query is canceled
        // also ignore old versions, which is normal
        BufferState state = stateMachine.getState();
        if (state.isTerminal() || outputBuffers.getVersion() >= newOutputBuffers.getVersion()) {
            return;
        }
        // verify this is valid state change
        outputBuffers.checkValidTransition(newOutputBuffers);
        outputBuffers = newOutputBuffers;
        // add the new buffers
        for (OutputBufferId outputBufferId : outputBuffers.getBuffers().keySet()) {
            getBuffer(outputBufferId);
        }
        // Reset resume from position
        nextClientBufferIndex.set(0);
        // update state if no more buffers is set
        if (outputBuffers.isNoMoreBufferIds()) {
            stateMachine.noMoreBuffers();
        }
    }
    if (!stateMachine.getState().canAddBuffers()) {
        noMoreBuffers();
    }
    checkFlushComplete();
}

Example 2 with OutputBufferId

use of io.trino.execution.buffer.OutputBuffers.OutputBufferId in project trino by trinodb.

the class LazyOutputBuffer method setOutputBuffers.

@Override
public void setOutputBuffers(OutputBuffers newOutputBuffers) {
    Set<OutputBufferId> destroyedBuffers = ImmutableSet.of();
    List<PendingRead> pendingReads = ImmutableList.of();
    OutputBuffer outputBuffer = delegate;
    if (outputBuffer == null) {
        synchronized (this) {
            outputBuffer = delegate;
            if (outputBuffer == null) {
                // ignore set output if buffer was already destroyed or failed
                if (stateMachine.getState().isTerminal()) {
                    return;
                }
                switch(newOutputBuffers.getType()) {
                    case PARTITIONED:
                        outputBuffer = new PartitionedOutputBuffer(taskInstanceId, stateMachine, newOutputBuffers, maxBufferSize, memoryContextSupplier, executor);
                        break;
                    case BROADCAST:
                        outputBuffer = new BroadcastOutputBuffer(taskInstanceId, stateMachine, maxBroadcastBufferSize, memoryContextSupplier, executor, notifyStatusChanged);
                        break;
                    case ARBITRARY:
                        outputBuffer = new ArbitraryOutputBuffer(taskInstanceId, stateMachine, maxBufferSize, memoryContextSupplier, executor);
                        break;
                    case SPOOL:
                        ExchangeSinkInstanceHandle exchangeSinkInstanceHandle = newOutputBuffers.getExchangeSinkInstanceHandle().orElseThrow(() -> new IllegalArgumentException("exchange sink handle is expected to be present for buffer type EXTERNAL"));
                        ExchangeManager exchangeManager = exchangeManagerRegistry.getExchangeManager();
                        ExchangeSink exchangeSink = exchangeManager.createSink(exchangeSinkInstanceHandle, false);
                        outputBuffer = new SpoolingExchangeOutputBuffer(stateMachine, newOutputBuffers, exchangeSink, memoryContextSupplier);
                        break;
                    default:
                        throw new IllegalArgumentException("Unexpected output buffer type: " + newOutputBuffers.getType());
                }
                // process pending aborts and reads outside of synchronized lock
                destroyedBuffers = ImmutableSet.copyOf(this.destroyedBuffers);
                this.destroyedBuffers.clear();
                pendingReads = ImmutableList.copyOf(this.pendingReads);
                this.pendingReads.clear();
                // Must be assigned last to avoid a race condition with unsynchronized readers
                delegate = outputBuffer;
            }
        }
    }
    outputBuffer.setOutputBuffers(newOutputBuffers);
    // process pending aborts and reads outside of synchronized lock
    destroyedBuffers.forEach(outputBuffer::destroy);
    for (PendingRead pendingRead : pendingReads) {
        pendingRead.process(outputBuffer);
    }
}

Example 3 with OutputBufferId

use of io.trino.execution.buffer.OutputBuffers.OutputBufferId in project trino by trinodb.

the class TestArbitraryOutputBuffer method testResumeFromPreviousPosition.

@Test
public void testResumeFromPreviousPosition() {
    OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(ARBITRARY);
    OutputBufferId[] ids = new OutputBufferId[5];
    for (int i = 0; i < ids.length; i++) {
        ids[i] = new OutputBufferId(i);
        outputBuffers = outputBuffers.withBuffer(ids[i], i);
    }
    ArbitraryOutputBuffer buffer = createArbitraryBuffer(outputBuffers, sizeOfPages(5));
    assertEquals(buffer.getState(), OPEN);
    Map<OutputBufferId, ListenableFuture<BufferResult>> firstReads = new HashMap<>();
    for (OutputBufferId id : ids) {
        firstReads.put(id, buffer.get(id, 0L, sizeOfPages(1)));
    }
    // All must be blocked initially
    assertThat(firstReads.values()).noneMatch(Future::isDone);
    List<ListenableFuture<BufferResult>> secondReads = new ArrayList<>();
    for (int i = 0; i < ids.length; i++) {
        // add one page
        addPage(buffer, createPage(33));
        assertThat(secondReads).allMatch(future -> !future.isDone(), "No secondary reads should complete until after all first reads");
        List<OutputBufferId> completedIds = firstReads.entrySet().stream().filter(entry -> entry.getValue().isDone()).map(Map.Entry::getKey).collect(toList());
        assertEquals(completedIds.size(), 1, "One completed buffer read per page addition");
        OutputBufferId completed = completedIds.get(0);
        BufferResult result = getFuture(firstReads.remove(completed), NO_WAIT);
        // Store completion order of first for follow up sequence
        secondReads.add(buffer.get(completed, result.getNextToken(), sizeOfPages(1)));
    }
    // Test sanity
    assertEquals(secondReads.size(), ids.length);
    // Completion order should be identical to the first iteration at this point
    for (int i = 0; i < ids.length; i++) {
        // add one page
        addPage(buffer, createPage(33));
        assertTrue(secondReads.get(i).isDone(), "Invalid second read completion order at index: " + i);
    }
}

Example 4 with OutputBufferId

use of io.trino.execution.buffer.OutputBuffers.OutputBufferId in project trino by trinodb.

the class TestPartitionedOutputBufferManager method test.

@Test
public void test() {
    PartitionedOutputBufferManager hashOutputBufferManager = new PartitionedOutputBufferManager(FIXED_HASH_DISTRIBUTION, 4);
    // output buffers are set immediately when the manager is created
    assertOutputBuffers(hashOutputBufferManager.getOutputBuffers());
    // add buffers, which does not cause an error
    hashOutputBufferManager.addOutputBuffer(new OutputBufferId(0));
    assertOutputBuffers(hashOutputBufferManager.getOutputBuffers());
    hashOutputBufferManager.addOutputBuffer(new OutputBufferId(3));
    assertOutputBuffers(hashOutputBufferManager.getOutputBuffers());
    // try to a buffer out side of the partition range, which should result in an error
    assertThatThrownBy(() -> hashOutputBufferManager.addOutputBuffer(new OutputBufferId(5))).isInstanceOf(IllegalStateException.class).hasMessage("Unexpected new output buffer 5");
    assertOutputBuffers(hashOutputBufferManager.getOutputBuffers());
    // try to a buffer out side of the partition range, which should result in an error
    assertThatThrownBy(() -> hashOutputBufferManager.addOutputBuffer(new OutputBufferId(6))).isInstanceOf(IllegalStateException.class).hasMessage("Unexpected new output buffer 6");
    assertOutputBuffers(hashOutputBufferManager.getOutputBuffers());
}

Example 5 with OutputBufferId

use of io.trino.execution.buffer.OutputBuffers.OutputBufferId in project trino by trinodb.

the class TestBroadcastOutputBufferManager method test.

@Test
public void test() {
    BroadcastOutputBufferManager hashOutputBufferManager = new BroadcastOutputBufferManager();
    assertEquals(hashOutputBufferManager.getOutputBuffers(), createInitialEmptyOutputBuffers(BROADCAST));
    hashOutputBufferManager.addOutputBuffer(new OutputBufferId(0));
    OutputBuffers expectedOutputBuffers = createInitialEmptyOutputBuffers(BROADCAST).withBuffer(new OutputBufferId(0), BROADCAST_PARTITION_ID);
    assertEquals(hashOutputBufferManager.getOutputBuffers(), expectedOutputBuffers);
    hashOutputBufferManager.addOutputBuffer(new OutputBufferId(1));
    hashOutputBufferManager.addOutputBuffer(new OutputBufferId(2));
    expectedOutputBuffers = expectedOutputBuffers.withBuffer(new OutputBufferId(1), BROADCAST_PARTITION_ID);
    expectedOutputBuffers = expectedOutputBuffers.withBuffer(new OutputBufferId(2), BROADCAST_PARTITION_ID);
    assertEquals(hashOutputBufferManager.getOutputBuffers(), expectedOutputBuffers);
    // set no more buffers
    hashOutputBufferManager.addOutputBuffer(new OutputBufferId(3));
    hashOutputBufferManager.noMoreBuffers();
    expectedOutputBuffers = expectedOutputBuffers.withBuffer(new OutputBufferId(3), BROADCAST_PARTITION_ID);
    expectedOutputBuffers = expectedOutputBuffers.withNoMoreBufferIds();
    assertEquals(hashOutputBufferManager.getOutputBuffers(), expectedOutputBuffers);
    // try to add another buffer, which should not result in an error
    // and output buffers should not change
    hashOutputBufferManager.addOutputBuffer(new OutputBufferId(5));
    assertEquals(hashOutputBufferManager.getOutputBuffers(), expectedOutputBuffers);
    // try to set no more buffers again, which should not result in an error
    // and output buffers should not change
    hashOutputBufferManager.addOutputBuffer(new OutputBufferId(6));
    assertEquals(hashOutputBufferManager.getOutputBuffers(), expectedOutputBuffers);
}