Examples with BufferProvider org.apache.flink.runtime.io.network.buffer.BufferProvider used on opensource projects

Example 6 with BufferProvider

use of org.apache.flink.runtime.io.network.buffer.BufferProvider in project flink by apache.

the class MockEnvironment method addOutput.

public void addOutput(final List<Record> outputList) {
    try {
        // The record-oriented writers wrap the buffer writer. We mock it
        // to collect the returned buffers and deserialize the content to
        // the output list
        BufferProvider mockBufferProvider = mock(BufferProvider.class);
        when(mockBufferProvider.requestBufferBlocking()).thenAnswer(new Answer<Buffer>() {

            @Override
            public Buffer answer(InvocationOnMock invocationOnMock) throws Throwable {
                return new Buffer(MemorySegmentFactory.allocateUnpooledSegment(bufferSize), mock(BufferRecycler.class));
            }
        });
        ResultPartitionWriter mockWriter = mock(ResultPartitionWriter.class);
        when(mockWriter.getNumberOfOutputChannels()).thenReturn(1);
        when(mockWriter.getBufferProvider()).thenReturn(mockBufferProvider);
        final Record record = new Record();
        final RecordDeserializer<Record> deserializer = new AdaptiveSpanningRecordDeserializer<Record>();
        // Add records from the buffer to the output list
        doAnswer(new Answer<Void>() {

            @Override
            public Void answer(InvocationOnMock invocationOnMock) throws Throwable {
                Buffer buffer = (Buffer) invocationOnMock.getArguments()[0];
                deserializer.setNextBuffer(buffer);
                while (deserializer.hasUnfinishedData()) {
                    RecordDeserializer.DeserializationResult result = deserializer.getNextRecord(record);
                    if (result.isFullRecord()) {
                        outputList.add(record.createCopy());
                    }
                    if (result == RecordDeserializer.DeserializationResult.LAST_RECORD_FROM_BUFFER || result == RecordDeserializer.DeserializationResult.PARTIAL_RECORD) {
                        break;
                    }
                }
                return null;
            }
        }).when(mockWriter).writeBuffer(any(Buffer.class), anyInt());
        outputs.add(mockWriter);
    } catch (Throwable t) {
        t.printStackTrace();
        fail(t.getMessage());
    }
}

Example 7 with BufferProvider

use of org.apache.flink.runtime.io.network.buffer.BufferProvider in project flink by apache.

the class PartitionRequestClientHandler method decodeBufferOrEvent.

private boolean decodeBufferOrEvent(RemoteInputChannel inputChannel, NettyMessage.BufferResponse bufferOrEvent, boolean isStagedBuffer) throws Throwable {
    boolean releaseNettyBuffer = true;
    try {
        if (bufferOrEvent.isBuffer()) {
            // IndexOutOfBoundsException.
            if (bufferOrEvent.getSize() == 0) {
                inputChannel.onEmptyBuffer(bufferOrEvent.sequenceNumber);
                return true;
            }
            BufferProvider bufferProvider = inputChannel.getBufferProvider();
            if (bufferProvider == null) {
                // receiver has been cancelled/failed
                cancelRequestFor(bufferOrEvent.receiverId);
                return isStagedBuffer;
            }
            while (true) {
                Buffer buffer = bufferProvider.requestBuffer();
                if (buffer != null) {
                    buffer.setSize(bufferOrEvent.getSize());
                    bufferOrEvent.getNettyBuffer().readBytes(buffer.getNioBuffer());
                    inputChannel.onBuffer(buffer, bufferOrEvent.sequenceNumber);
                    return true;
                } else if (bufferListener.waitForBuffer(bufferProvider, bufferOrEvent)) {
                    releaseNettyBuffer = false;
                    return false;
                } else if (bufferProvider.isDestroyed()) {
                    return isStagedBuffer;
                }
            }
        } else {
            // ---- Event -------------------------------------------------
            // TODO We can just keep the serialized data in the Netty buffer and release it later at the reader
            byte[] byteArray = new byte[bufferOrEvent.getSize()];
            bufferOrEvent.getNettyBuffer().readBytes(byteArray);
            MemorySegment memSeg = MemorySegmentFactory.wrap(byteArray);
            Buffer buffer = new Buffer(memSeg, FreeingBufferRecycler.INSTANCE, false);
            inputChannel.onBuffer(buffer, bufferOrEvent.sequenceNumber);
            return true;
        }
    } finally {
        if (releaseNettyBuffer) {
            bufferOrEvent.releaseBuffer();
        }
    }
}

Example 8 with BufferProvider

use of org.apache.flink.runtime.io.network.buffer.BufferProvider in project flink by apache.

the class RecordWriterTest method testClearBuffersAfterInterruptDuringBlockingBufferRequest.

// ---------------------------------------------------------------------------------------------
// Resource release tests
// ---------------------------------------------------------------------------------------------
/**
	 * Tests a fix for FLINK-2089.
	 *
	 * @see <a href="https://issues.apache.org/jira/browse/FLINK-2089">FLINK-2089</a>
	 */
@Test
public void testClearBuffersAfterInterruptDuringBlockingBufferRequest() throws Exception {
    ExecutorService executor = null;
    try {
        executor = Executors.newSingleThreadExecutor();
        final CountDownLatch sync = new CountDownLatch(2);
        final Buffer buffer = spy(TestBufferFactory.createBuffer(4));
        // Return buffer for first request, but block for all following requests.
        Answer<Buffer> request = new Answer<Buffer>() {

            @Override
            public Buffer answer(InvocationOnMock invocation) throws Throwable {
                sync.countDown();
                if (sync.getCount() == 1) {
                    return buffer;
                }
                final Object o = new Object();
                synchronized (o) {
                    while (true) {
                        o.wait();
                    }
                }
            }
        };
        BufferProvider bufferProvider = mock(BufferProvider.class);
        when(bufferProvider.requestBufferBlocking()).thenAnswer(request);
        ResultPartitionWriter partitionWriter = createResultPartitionWriter(bufferProvider);
        final RecordWriter<IntValue> recordWriter = new RecordWriter<IntValue>(partitionWriter);
        Future<?> result = executor.submit(new Callable<Void>() {

            @Override
            public Void call() throws Exception {
                IntValue val = new IntValue(0);
                try {
                    recordWriter.emit(val);
                    recordWriter.flush();
                    recordWriter.emit(val);
                } catch (InterruptedException e) {
                    recordWriter.clearBuffers();
                }
                return null;
            }
        });
        sync.await();
        // Interrupt the Thread.
        //
        // The second emit call requests a new buffer and blocks the thread.
        // When interrupting the thread at this point, clearing the buffers
        // should not recycle any buffer.
        result.cancel(true);
        recordWriter.clearBuffers();
        // Verify that buffer have been requested, but only one has been written out.
        verify(bufferProvider, times(2)).requestBufferBlocking();
        verify(partitionWriter, times(1)).writeBuffer(any(Buffer.class), anyInt());
        // Verify that the written out buffer has only been recycled once
        // (by the partition writer).
        assertTrue("Buffer not recycled.", buffer.isRecycled());
        verify(buffer, times(1)).recycle();
    } finally {
        if (executor != null) {
            executor.shutdown();
        }
    }
}

Example 9 with BufferProvider

use of org.apache.flink.runtime.io.network.buffer.BufferProvider in project flink by apache.

the class PartitionRequestClientHandlerTest method testAutoReadAfterUnsuccessfulStagedMessage.

/**
	 * Tests that an unsuccessful message decode call for a staged message
	 * does not leave the channel with auto read set to false.
	 */
@Test
@SuppressWarnings("unchecked")
public void testAutoReadAfterUnsuccessfulStagedMessage() throws Exception {
    PartitionRequestClientHandler handler = new PartitionRequestClientHandler();
    EmbeddedChannel channel = new EmbeddedChannel(handler);
    final AtomicReference<EventListener<Buffer>> listener = new AtomicReference<>();
    BufferProvider bufferProvider = mock(BufferProvider.class);
    when(bufferProvider.addListener(any(EventListener.class))).thenAnswer(new Answer<Boolean>() {

        @Override
        @SuppressWarnings("unchecked")
        public Boolean answer(InvocationOnMock invocation) throws Throwable {
            listener.set((EventListener<Buffer>) invocation.getArguments()[0]);
            return true;
        }
    });
    when(bufferProvider.requestBuffer()).thenReturn(null);
    InputChannelID channelId = new InputChannelID(0, 0);
    RemoteInputChannel inputChannel = mock(RemoteInputChannel.class);
    when(inputChannel.getInputChannelId()).thenReturn(channelId);
    // The 3rd staged msg has a null buffer provider
    when(inputChannel.getBufferProvider()).thenReturn(bufferProvider, bufferProvider, null);
    handler.addInputChannel(inputChannel);
    BufferResponse msg = createBufferResponse(createBuffer(true), 0, channelId);
    // Write 1st buffer msg. No buffer is available, therefore the buffer
    // should be staged and auto read should be set to false.
    assertTrue(channel.config().isAutoRead());
    channel.writeInbound(msg);
    // No buffer available, auto read false
    assertFalse(channel.config().isAutoRead());
    // Write more buffers... all staged.
    msg = createBufferResponse(createBuffer(true), 1, channelId);
    channel.writeInbound(msg);
    msg = createBufferResponse(createBuffer(true), 2, channelId);
    channel.writeInbound(msg);
    // Notify about buffer => handle 1st msg
    Buffer availableBuffer = createBuffer(false);
    listener.get().onEvent(availableBuffer);
    // Start processing of staged buffers (in run pending tasks). Make
    // sure that the buffer provider acts like it's destroyed.
    when(bufferProvider.addListener(any(EventListener.class))).thenReturn(false);
    when(bufferProvider.isDestroyed()).thenReturn(true);
    // Execute all tasks that are scheduled in the event loop. Further
    // eventLoop().execute() calls are directly executed, if they are
    // called in the scope of this call.
    channel.runPendingTasks();
    assertTrue(channel.config().isAutoRead());
}

Example 10 with BufferProvider

use of org.apache.flink.runtime.io.network.buffer.BufferProvider in project flink by apache.

the class RecordWriterTest method testBroadcastEventNoRecords.

/**
	 * Tests broadcasting events when no records have been emitted yet.
	 */
@Test
public void testBroadcastEventNoRecords() throws Exception {
    int numChannels = 4;
    int bufferSize = 32;
    @SuppressWarnings("unchecked") Queue<BufferOrEvent>[] queues = new Queue[numChannels];
    for (int i = 0; i < numChannels; i++) {
        queues[i] = new ArrayDeque<>();
    }
    BufferProvider bufferProvider = createBufferProvider(bufferSize);
    ResultPartitionWriter partitionWriter = createCollectingPartitionWriter(queues, bufferProvider);
    RecordWriter<ByteArrayIO> writer = new RecordWriter<>(partitionWriter, new RoundRobin<ByteArrayIO>());
    CheckpointBarrier barrier = new CheckpointBarrier(Integer.MAX_VALUE + 919192L, Integer.MAX_VALUE + 18828228L, CheckpointOptions.forFullCheckpoint());
    // No records emitted yet, broadcast should not request a buffer
    writer.broadcastEvent(barrier);
    verify(bufferProvider, times(0)).requestBufferBlocking();
    for (Queue<BufferOrEvent> queue : queues) {
        assertEquals(1, queue.size());
        BufferOrEvent boe = queue.remove();
        assertTrue(boe.isEvent());
        assertEquals(barrier, boe.getEvent());
    }
}