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

Example 1 with TestProducerSource

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

the class PipelinedSubpartitionTest method testProduceConsume.

private void testProduceConsume(boolean isSlowProducer, boolean isSlowConsumer) throws Exception {
    // Config
    final int producerBufferPoolSize = 8;
    final int producerNumberOfBuffersToProduce = 128;
    // Producer behaviour
    final TestProducerSource producerSource = new TestProducerSource() {

        private BufferProvider bufferProvider = new TestPooledBufferProvider(producerBufferPoolSize);

        private int numberOfBuffers;

        @Override
        public BufferOrEvent getNextBufferOrEvent() throws Exception {
            if (numberOfBuffers == producerNumberOfBuffersToProduce) {
                return null;
            }
            final Buffer buffer = bufferProvider.requestBufferBlocking();
            final MemorySegment segment = buffer.getMemorySegment();
            int next = numberOfBuffers * (segment.size() / 4);
            for (int i = 0; i < segment.size(); i += 4) {
                segment.putInt(i, next);
                next++;
            }
            numberOfBuffers++;
            return new BufferOrEvent(buffer, 0);
        }
    };
    // Consumer behaviour
    final TestConsumerCallback consumerCallback = new TestConsumerCallback() {

        private int numberOfBuffers;

        @Override
        public void onBuffer(Buffer buffer) {
            final MemorySegment segment = buffer.getMemorySegment();
            int expected = numberOfBuffers * (segment.size() / 4);
            for (int i = 0; i < segment.size(); i += 4) {
                assertEquals(expected, segment.getInt(i));
                expected++;
            }
            numberOfBuffers++;
            buffer.recycle();
        }

        @Override
        public void onEvent(AbstractEvent event) {
        // Nothing to do in this test
        }
    };
    final PipelinedSubpartition subpartition = createSubpartition();
    TestSubpartitionConsumer consumer = new TestSubpartitionConsumer(isSlowConsumer, consumerCallback);
    final PipelinedSubpartitionView view = subpartition.createReadView(null, consumer);
    consumer.setSubpartitionView(view);
    Future<Boolean> producerResult = executorService.submit(new TestSubpartitionProducer(subpartition, isSlowProducer, producerSource));
    Future<Boolean> consumerResult = executorService.submit(consumer);
    // Wait for producer and consumer to finish
    producerResult.get();
    consumerResult.get();
}