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

Example 16 with BufferBuilder

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

the class RecordWriterTest method testIsAvailableOrNot.

/**
 * Tests that the RecordWriter is available iif the respective LocalBufferPool has at-least one
 * available buffer.
 */
@Test
public void testIsAvailableOrNot() throws Exception {
    // setup
    final NetworkBufferPool globalPool = new NetworkBufferPool(10, 128);
    final BufferPool localPool = globalPool.createBufferPool(1, 1, 1, Integer.MAX_VALUE);
    final ResultPartitionWriter resultPartition = new ResultPartitionBuilder().setBufferPoolFactory(() -> localPool).build();
    resultPartition.setup();
    final RecordWriter<?> recordWriter = createRecordWriter(resultPartition);
    try {
        // record writer is available because of initial available global pool
        assertTrue(recordWriter.getAvailableFuture().isDone());
        // request one buffer from the local pool to make it unavailable afterwards
        try (BufferBuilder bufferBuilder = localPool.requestBufferBuilder(0)) {
            assertNotNull(bufferBuilder);
            assertFalse(recordWriter.getAvailableFuture().isDone());
            // recycle the buffer to make the local pool available again
            final Buffer buffer = BufferBuilderTestUtils.buildSingleBuffer(bufferBuilder);
            buffer.recycleBuffer();
        }
        assertTrue(recordWriter.getAvailableFuture().isDone());
        assertEquals(recordWriter.AVAILABLE, recordWriter.getAvailableFuture());
    } finally {
        localPool.lazyDestroy();
        globalPool.destroy();
    }
}

Example 17 with BufferBuilder

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

the class DemultiplexingRecordDeserializerTest method testWatermarks.

/**
 * Tests that Watermarks are only forwarded when all watermarks are received.
 */
@Test
public void testWatermarks() throws IOException {
    DemultiplexingRecordDeserializer<Long> deserializer = DemultiplexingRecordDeserializer.create(new InputChannelInfo(0, 0), rescalingDescriptor(to(0, 1), array(mappings(to(0, 1), to(4, 5))), emptySet()), unused -> new SpillingAdaptiveSpanningRecordDeserializer<>(ioManager.getSpillingDirectoriesPaths()), unused -> RecordFilter.all());
    assertEquals(4, deserializer.getVirtualChannelSelectors().size());
    for (Iterator<SubtaskConnectionDescriptor> iterator = deserializer.getVirtualChannelSelectors().iterator(); iterator.hasNext(); ) {
        SubtaskConnectionDescriptor selector = iterator.next();
        MemorySegment memorySegment = allocateUnpooledSegment(128);
        try (BufferBuilder bufferBuilder = createBufferBuilder(memorySegment)) {
            final long ts = 42L + selector.getInputSubtaskIndex() + selector.getOutputSubtaskIndex();
            Buffer buffer = write(bufferBuilder, new Watermark(ts));
            deserializer.select(selector);
            deserializer.setNextBuffer(buffer);
        }
        if (iterator.hasNext()) {
            assertEquals(Collections.emptyList(), read(deserializer));
        } else {
            // last channel, min should be 42 + 0 + 0
            assertEquals(Arrays.asList(new Watermark(42)), read(deserializer));
        }
        assertTrue(memorySegment.isFreed());
    }
}

Example 18 with BufferBuilder

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

the class DemultiplexingRecordDeserializerTest method testAmbiguousChannels.

/**
 * Tests that {@link RecordFilter} are used correctly.
 */
@Test
public void testAmbiguousChannels() throws IOException {
    DemultiplexingRecordDeserializer<Long> deserializer = DemultiplexingRecordDeserializer.create(new InputChannelInfo(1, 0), rescalingDescriptor(to(41, 42), array(mappings(), mappings(to(2, 3), to(4, 5))), set(42)), unused -> new SpillingAdaptiveSpanningRecordDeserializer<>(ioManager.getSpillingDirectoriesPaths()), unused -> new RecordFilter(new ModSelector(2), LongSerializer.INSTANCE, 1));
    assertEquals(Sets.newSet(new SubtaskConnectionDescriptor(41, 2), new SubtaskConnectionDescriptor(41, 3), new SubtaskConnectionDescriptor(42, 2), new SubtaskConnectionDescriptor(42, 3)), deserializer.getVirtualChannelSelectors());
    for (int i = 0; i < 100; i++) {
        MemorySegment memorySegment = allocateUnpooledSegment(128);
        try (BufferBuilder bufferBuilder = createBufferBuilder(memorySegment)) {
            // add one even and one odd number
            Buffer buffer = writeLongs(bufferBuilder, i, i + 1L);
            SubtaskConnectionDescriptor selector = Iterables.get(deserializer.getVirtualChannelSelectors(), i / 10 % 2);
            deserializer.select(selector);
            deserializer.setNextBuffer(buffer);
            if (selector.getInputSubtaskIndex() == 41) {
                assertEquals(Arrays.asList((long) i, i + 1L), readLongs(deserializer));
            } else {
                // only odd should occur in output
                assertEquals(Arrays.asList(i / 2 * 2 + 1L), readLongs(deserializer));
            }
        }
        assertTrue(memorySegment.isFreed());
    }
}

Example 19 with BufferBuilder

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

the class DemultiplexingRecordDeserializerTest method testUpscale.

/**
 * Tests {@link SubtaskConnectionDescriptor} by mixing buffers from 4 different virtual
 * channels.
 */
@Test
public void testUpscale() throws IOException {
    DemultiplexingRecordDeserializer<Long> deserializer = DemultiplexingRecordDeserializer.create(new InputChannelInfo(2, 0), rescalingDescriptor(to(0, 1), array(mappings(), mappings(), mappings(to(2, 3), to(4, 5))), emptySet()), unused -> new SpillingAdaptiveSpanningRecordDeserializer<>(ioManager.getSpillingDirectoriesPaths()), unused -> RecordFilter.all());
    assertEquals(Sets.newSet(new SubtaskConnectionDescriptor(0, 2), new SubtaskConnectionDescriptor(0, 3), new SubtaskConnectionDescriptor(1, 2), new SubtaskConnectionDescriptor(1, 3)), deserializer.getVirtualChannelSelectors());
    for (int i = 0; i < 100; i++) {
        SubtaskConnectionDescriptor selector = Iterables.get(deserializer.getVirtualChannelSelectors(), random.nextInt(4));
        long start = selector.getInputSubtaskIndex() << 4 | selector.getOutputSubtaskIndex();
        MemorySegment memorySegment = allocateUnpooledSegment(128);
        try (BufferBuilder bufferBuilder = createBufferBuilder(memorySegment)) {
            Buffer buffer = writeLongs(bufferBuilder, start + 1L, start + 2L, start + 3L);
            deserializer.select(selector);
            deserializer.setNextBuffer(buffer);
        }
        assertEquals(Arrays.asList(start + 1L, start + 2L, start + 3L), readLongs(deserializer));
        assertTrue(memorySegment.isFreed());
    }
}

Example 20 with BufferBuilder

use of org.apache.flink.runtime.io.network.buffer.BufferBuilder 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;
}