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

Example 1 with TestConsumerCallback

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

the class PipelinedSubpartitionTest method testProduceConsume.

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

        private int numberOfBuffers;

        @Override
        public BufferAndChannel getNextBuffer() throws Exception {
            if (numberOfBuffers == producerNumberOfBuffersToProduce) {
                return null;
            }
            MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(bufferSize);
            int next = numberOfBuffers * (bufferSize / Integer.BYTES);
            for (int i = 0; i < bufferSize; i += 4) {
                segment.putInt(i, next);
                next++;
            }
            numberOfBuffers++;
            return new BufferAndChannel(segment.getArray(), 0);
        }
    };
    // Consumer behaviour
    final TestConsumerCallback consumerCallback = new TestConsumerCallback() {

        private int numberOfBuffers;

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

        @Override
        public void onEvent(AbstractEvent event) {
        // Nothing to do in this test
        }
    };
    final PipelinedSubpartition subpartition = createSubpartition();
    TestSubpartitionProducer producer = new TestSubpartitionProducer(subpartition, isSlowProducer, producerSource);
    TestSubpartitionConsumer consumer = new TestSubpartitionConsumer(isSlowConsumer, consumerCallback);
    final PipelinedSubpartitionView view = subpartition.createReadView(consumer);
    consumer.setSubpartitionView(view);
    CompletableFuture<Boolean> producerResult = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(producer::call), executorService);
    CompletableFuture<Boolean> consumerResult = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(consumer::call), executorService);
    FutureUtils.waitForAll(Arrays.asList(producerResult, consumerResult)).get(60_000L, TimeUnit.MILLISECONDS);
}