Examples with BufferAndAvailabilityProvider org.apache.flink.runtime.io.network.partition.consumer.TestInputChannel.BufferAndAvailabilityProvider used on opensource projects

Example 1 with BufferAndAvailabilityProvider

use of org.apache.flink.runtime.io.network.partition.consumer.TestInputChannel.BufferAndAvailabilityProvider in project flink by apache.

the class StreamTestSingleInputGate method setupInputChannels.

private TestInputChannel[] setupInputChannels() {
    TestInputChannel[] inputChannels = new TestInputChannel[numInputChannels];
    for (int i = 0; i < numInputChannels; i++) {
        final int channelIndex = i;
        final DataOutputSerializer dataOutputSerializer = new DataOutputSerializer(128);
        final SerializationDelegate<StreamElement> delegate = new SerializationDelegate<>(new StreamElementSerializer<T>(serializer));
        inputQueues[channelIndex] = new ConcurrentLinkedQueue<>();
        inputChannels[channelIndex] = new TestInputChannel(inputGate, i);
        final BufferAndAvailabilityProvider answer = () -> {
            ConcurrentLinkedQueue<InputValue<Object>> inputQueue = inputQueues[channelIndex];
            InputValue<Object> input;
            Buffer.DataType nextType;
            synchronized (inputQueue) {
                input = inputQueue.poll();
                nextType = !inputQueue.isEmpty() ? Buffer.DataType.DATA_BUFFER : Buffer.DataType.NONE;
            }
            if (input != null && input.isStreamEnd()) {
                inputChannels[channelIndex].setReleased();
                return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE, false), nextType, 0, 0));
            } else if (input != null && input.isDataEnd()) {
                return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(new EndOfData(StopMode.DRAIN), false), nextType, 0, 0));
            } else if (input != null && input.isStreamRecord()) {
                StreamElement inputElement = input.getStreamRecord();
                delegate.setInstance(inputElement);
                ByteBuffer serializedRecord = RecordWriter.serializeRecord(dataOutputSerializer, delegate);
                BufferBuilder bufferBuilder = createBufferBuilder(bufferSize);
                BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
                bufferBuilder.appendAndCommit(serializedRecord);
                bufferBuilder.finish();
                bufferBuilder.close();
                // Call getCurrentBuffer to ensure size is set
                return Optional.of(new BufferAndAvailability(bufferConsumer.build(), nextType, 0, 0));
            } else if (input != null && input.isEvent()) {
                AbstractEvent event = input.getEvent();
                if (event instanceof EndOfPartitionEvent) {
                    inputChannels[channelIndex].setReleased();
                }
                return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(event, false), nextType, 0, 0));
            } else {
                return Optional.empty();
            }
        };
        inputChannels[channelIndex].addBufferAndAvailability(answer);
    }
    return inputChannels;
}

Example 2 with BufferAndAvailabilityProvider

use of org.apache.flink.runtime.io.network.partition.consumer.TestInputChannel.BufferAndAvailabilityProvider in project flink by apache.

the class IteratorWrappingTestSingleInputGate method wrapIterator.

private IteratorWrappingTestSingleInputGate<T> wrapIterator(MutableObjectIterator<T> iterator) throws IOException, InterruptedException {
    inputIterator = iterator;
    serializer = new DataOutputSerializer(128);
    // The input iterator can produce an infinite stream. That's why we have to serialize each
    // record on demand and cannot do it upfront.
    final BufferAndAvailabilityProvider answer = new BufferAndAvailabilityProvider() {

        private boolean hasData = inputIterator.next(reuse) != null;

        @Override
        public Optional<BufferAndAvailability> getBufferAvailability() throws IOException {
            if (hasData) {
                ByteBuffer serializedRecord = RecordWriter.serializeRecord(serializer, reuse);
                BufferBuilder bufferBuilder = createBufferBuilder(bufferSize);
                BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
                bufferBuilder.appendAndCommit(serializedRecord);
                hasData = inputIterator.next(reuse) != null;
                // Call getCurrentBuffer to ensure size is set
                final Buffer.DataType nextDataType = hasData ? Buffer.DataType.DATA_BUFFER : Buffer.DataType.EVENT_BUFFER;
                return Optional.of(new BufferAndAvailability(bufferConsumer.build(), nextDataType, 0, 0));
            } else {
                inputChannel.setReleased();
                return Optional.of(new BufferAndAvailability(EventSerializer.toBuffer(EndOfPartitionEvent.INSTANCE, false), Buffer.DataType.NONE, 0, 0));
            }
        }
    };
    inputChannel.addBufferAndAvailability(answer);
    inputGate.setInputChannels(inputChannel);
    return this;
}