Examples with TestPooledBufferProvider org.apache.flink.runtime.io.network.util.TestPooledBufferProvider used on opensource projects

Example 1 with TestPooledBufferProvider

use of org.apache.flink.runtime.io.network.util.TestPooledBufferProvider in project flink by apache.

the class SpilledSubpartitionViewTest method testWriteConsume.

@Test
public void testWriteConsume() throws Exception {
    // Config
    final int numberOfBuffersToWrite = 512;
    // Setup
    final BufferFileWriter writer = createWriterAndWriteBuffers(IO_MANAGER, writerBufferPool, numberOfBuffersToWrite);
    writer.close();
    TestPooledBufferProvider viewBufferPool = new TestPooledBufferProvider(1);
    TestSubpartitionConsumer consumer = new TestSubpartitionConsumer(false, new TestConsumerCallback.RecyclingCallback());
    SpilledSubpartitionView view = new SpilledSubpartitionView(mock(ResultSubpartition.class), viewBufferPool.getMemorySegmentSize(), writer, // +1 for end-of-partition
    numberOfBuffersToWrite + 1, consumer);
    consumer.setSubpartitionView(view);
    // Consume subpartition
    consumer.call();
}

Example 2 with TestPooledBufferProvider

use of org.apache.flink.runtime.io.network.util.TestPooledBufferProvider in project flink by apache.

the class CancelPartitionRequestTest method testDuplicateCancel.

@Test
public void testDuplicateCancel() throws Exception {
    NettyServerAndClient serverAndClient = null;
    try {
        final TestPooledBufferProvider outboundBuffers = new TestPooledBufferProvider(16);
        ResultPartitionManager partitions = mock(ResultPartitionManager.class);
        ResultPartitionID pid = new ResultPartitionID();
        final CountDownLatch sync = new CountDownLatch(1);
        final ResultSubpartitionView view = spy(new InfiniteSubpartitionView(outboundBuffers, sync));
        // Return infinite subpartition
        when(partitions.createSubpartitionView(eq(pid), eq(0), any(BufferAvailabilityListener.class))).thenAnswer(new Answer<ResultSubpartitionView>() {

            @Override
            public ResultSubpartitionView answer(InvocationOnMock invocationOnMock) throws Throwable {
                BufferAvailabilityListener listener = (BufferAvailabilityListener) invocationOnMock.getArguments()[2];
                listener.notifyDataAvailable();
                return view;
            }
        });
        NettyProtocol protocol = new NettyProtocol(partitions, mock(TaskEventDispatcher.class));
        serverAndClient = initServerAndClient(protocol);
        Channel ch = connect(serverAndClient);
        // Request for non-existing input channel => results in cancel request
        InputChannelID inputChannelId = new InputChannelID();
        ch.writeAndFlush(new PartitionRequest(pid, 0, inputChannelId, Integer.MAX_VALUE)).await();
        // Wait for the notification
        if (!sync.await(TestingUtils.TESTING_DURATION.toMillis(), TimeUnit.MILLISECONDS)) {
            fail("Timed out after waiting for " + TestingUtils.TESTING_DURATION.toMillis() + " ms to be notified about cancelled partition.");
        }
        ch.writeAndFlush(new CancelPartitionRequest(inputChannelId)).await();
        ch.close();
        NettyTestUtil.awaitClose(ch);
        verify(view, times(1)).releaseAllResources();
    } finally {
        shutdown(serverAndClient);
    }
}

Example 3 with TestPooledBufferProvider

use of org.apache.flink.runtime.io.network.util.TestPooledBufferProvider in project flink by apache.

the class SpilledSubpartitionViewTest method testReadMultipleFilesWithSingleBufferPool.

@Test
public void testReadMultipleFilesWithSingleBufferPool() throws Exception {
    ExecutorService executor = null;
    BufferFileWriter[] writers = null;
    ResultSubpartitionView[] readers = null;
    try {
        executor = Executors.newCachedThreadPool();
        // Setup
        writers = new BufferFileWriter[] { createWriterAndWriteBuffers(IO_MANAGER, writerBufferPool, 512), createWriterAndWriteBuffers(IO_MANAGER, writerBufferPool, 512) };
        readers = new ResultSubpartitionView[writers.length];
        TestSubpartitionConsumer[] consumers = new TestSubpartitionConsumer[writers.length];
        BufferProvider inputBuffers = new TestPooledBufferProvider(2);
        ResultSubpartition parent = mock(ResultSubpartition.class);
        // Wait for writers to finish
        for (BufferFileWriter writer : writers) {
            writer.close();
        }
        // Create the views depending on the test configuration
        for (int i = 0; i < readers.length; i++) {
            consumers[i] = new TestSubpartitionConsumer(false, new TestConsumerCallback.RecyclingCallback());
            readers[i] = new SpilledSubpartitionView(parent, inputBuffers.getMemorySegmentSize(), writers[i], // +1 for end of partition event
            512 + 1, consumers[i]);
            consumers[i].setSubpartitionView(readers[i]);
        }
        final List<Future<Boolean>> results = Lists.newArrayList();
        // Submit the consuming tasks
        for (TestSubpartitionConsumer consumer : consumers) {
            results.add(executor.submit(consumer));
        }
        // Wait for the results
        for (Future<Boolean> res : results) {
            try {
                res.get(2, TimeUnit.MINUTES);
            } catch (TimeoutException e) {
                throw new TimeoutException("There has been a timeout in the test. This " + "indicates that there is a bug/deadlock in the tested subpartition " + "view.");
            }
        }
    } finally {
        if (writers != null) {
            for (BufferFileWriter writer : writers) {
                if (writer != null) {
                    writer.deleteChannel();
                }
            }
        }
        if (readers != null) {
            for (ResultSubpartitionView reader : readers) {
                if (reader != null) {
                    reader.releaseAllResources();
                }
            }
        }
        if (executor != null) {
            executor.shutdown();
        }
    }
}

Example 4 with TestPooledBufferProvider

use of org.apache.flink.runtime.io.network.util.TestPooledBufferProvider in project flink by apache.

the class CancelPartitionRequestTest method testCancelPartitionRequest.

/**
 * Verifies that requests for non-existing (failed/cancelled) input channels are properly
 * cancelled. The receiver receives data, but there is no input channel to receive the data.
 * This should cancel the request.
 */
@Test
public void testCancelPartitionRequest() throws Exception {
    NettyServerAndClient serverAndClient = null;
    try {
        TestPooledBufferProvider outboundBuffers = new TestPooledBufferProvider(16);
        ResultPartitionManager partitions = mock(ResultPartitionManager.class);
        ResultPartitionID pid = new ResultPartitionID();
        CountDownLatch sync = new CountDownLatch(1);
        final ResultSubpartitionView view = spy(new InfiniteSubpartitionView(outboundBuffers, sync));
        // Return infinite subpartition
        when(partitions.createSubpartitionView(eq(pid), eq(0), any(BufferAvailabilityListener.class))).thenAnswer(new Answer<ResultSubpartitionView>() {

            @Override
            public ResultSubpartitionView answer(InvocationOnMock invocationOnMock) throws Throwable {
                BufferAvailabilityListener listener = (BufferAvailabilityListener) invocationOnMock.getArguments()[2];
                listener.notifyDataAvailable();
                return view;
            }
        });
        NettyProtocol protocol = new NettyProtocol(partitions, mock(TaskEventDispatcher.class));
        serverAndClient = initServerAndClient(protocol);
        Channel ch = connect(serverAndClient);
        // Request for non-existing input channel => results in cancel request
        ch.writeAndFlush(new PartitionRequest(pid, 0, new InputChannelID(), Integer.MAX_VALUE)).await();
        // Wait for the notification
        if (!sync.await(TestingUtils.TESTING_DURATION.toMillis(), TimeUnit.MILLISECONDS)) {
            fail("Timed out after waiting for " + TestingUtils.TESTING_DURATION.toMillis() + " ms to be notified about cancelled partition.");
        }
        verify(view, times(1)).releaseAllResources();
    } finally {
        shutdown(serverAndClient);
    }
}

Example 5 with TestPooledBufferProvider

use of org.apache.flink.runtime.io.network.util.TestPooledBufferProvider in project flink by apache.

the class ServerTransportErrorHandlingTest method testRemoteClose.

/**
 * Verifies remote closes trigger the release of all resources.
 */
@Test
public void testRemoteClose() throws Exception {
    final TestPooledBufferProvider outboundBuffers = new TestPooledBufferProvider(16);
    final CountDownLatch sync = new CountDownLatch(1);
    final ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
    when(partitionManager.createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferAvailabilityListener.class))).thenAnswer(new Answer<ResultSubpartitionView>() {

        @Override
        public ResultSubpartitionView answer(InvocationOnMock invocationOnMock) throws Throwable {
            BufferAvailabilityListener listener = (BufferAvailabilityListener) invocationOnMock.getArguments()[2];
            listener.notifyDataAvailable();
            return new CancelPartitionRequestTest.InfiniteSubpartitionView(outboundBuffers, sync);
        }
    });
    NettyProtocol protocol = new NettyProtocol(partitionManager, mock(TaskEventDispatcher.class)) {

        @Override
        public ChannelHandler[] getClientChannelHandlers() {
            return new ChannelHandler[] { new NettyMessage.NettyMessageEncoder(), // Close on read
            new ChannelInboundHandlerAdapter() {

                @Override
                public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
                    ctx.channel().close();
                }
            } };
        }
    };
    NettyTestUtil.NettyServerAndClient serverAndClient = null;
    try {
        for (int retry = 0; retry < NETTY_INIT_MAX_RETRY_TIMES; retry++) {
            try {
                serverAndClient = initServerAndClient(protocol, createConfig());
                break;
            } catch (Exception e) {
                if (retry >= NETTY_INIT_MAX_RETRY_TIMES - 1) {
                    throw new RuntimeException("Failed to initialize netty server and client, retried " + retry + " times.", e);
                }
                if (e instanceof BindException || ExceptionUtils.findThrowableWithMessage(e, "Address already in use").isPresent()) {
                    continue;
                }
                throw e;
            }
        }
        Channel ch = connect(serverAndClient);
        // Write something to trigger close by server
        ch.writeAndFlush(new NettyMessage.PartitionRequest(new ResultPartitionID(), 0, new InputChannelID(), Integer.MAX_VALUE));
        // Wait for the notification
        if (!sync.await(TestingUtils.TESTING_DURATION.toMillis(), TimeUnit.MILLISECONDS)) {
            fail("Timed out after waiting for " + TestingUtils.TESTING_DURATION.toMillis() + " ms to be notified about released partition.");
        }
    } finally {
        shutdown(serverAndClient);
    }
}