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

Example 1 with BufferPool

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

the class SingleInputGateTest method testBackwardsEventWithUninitializedChannel.

@Test
public void testBackwardsEventWithUninitializedChannel() throws Exception {
    // Setup environment
    final TaskEventDispatcher taskEventDispatcher = mock(TaskEventDispatcher.class);
    when(taskEventDispatcher.publish(any(ResultPartitionID.class), any(TaskEvent.class))).thenReturn(true);
    final ResultSubpartitionView iterator = mock(ResultSubpartitionView.class);
    when(iterator.getNextBuffer()).thenReturn(new Buffer(MemorySegmentFactory.allocateUnpooledSegment(1024), mock(BufferRecycler.class)));
    final ResultPartitionManager partitionManager = mock(ResultPartitionManager.class);
    when(partitionManager.createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferProvider.class), any(BufferAvailabilityListener.class))).thenReturn(iterator);
    // Setup reader with one local and one unknown input channel
    final IntermediateDataSetID resultId = new IntermediateDataSetID();
    final SingleInputGate inputGate = new SingleInputGate("Test Task Name", new JobID(), resultId, ResultPartitionType.PIPELINED, 0, 2, mock(TaskActions.class), new UnregisteredTaskMetricsGroup.DummyTaskIOMetricGroup());
    final BufferPool bufferPool = mock(BufferPool.class);
    when(bufferPool.getNumberOfRequiredMemorySegments()).thenReturn(2);
    inputGate.setBufferPool(bufferPool);
    // Local
    ResultPartitionID localPartitionId = new ResultPartitionID(new IntermediateResultPartitionID(), new ExecutionAttemptID());
    InputChannel local = new LocalInputChannel(inputGate, 0, localPartitionId, partitionManager, taskEventDispatcher, new UnregisteredTaskMetricsGroup.DummyTaskIOMetricGroup());
    // Unknown
    ResultPartitionID unknownPartitionId = new ResultPartitionID(new IntermediateResultPartitionID(), new ExecutionAttemptID());
    InputChannel unknown = new UnknownInputChannel(inputGate, 1, unknownPartitionId, partitionManager, taskEventDispatcher, mock(ConnectionManager.class), 0, 0, new UnregisteredTaskMetricsGroup.DummyTaskIOMetricGroup());
    // Set channels
    inputGate.setInputChannel(localPartitionId.getPartitionId(), local);
    inputGate.setInputChannel(unknownPartitionId.getPartitionId(), unknown);
    // Request partitions
    inputGate.requestPartitions();
    // Only the local channel can request
    verify(partitionManager, times(1)).createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferProvider.class), any(BufferAvailabilityListener.class));
    // Send event backwards and initialize unknown channel afterwards
    final TaskEvent event = new TestTaskEvent();
    inputGate.sendTaskEvent(event);
    // Only the local channel can send out the event
    verify(taskEventDispatcher, times(1)).publish(any(ResultPartitionID.class), any(TaskEvent.class));
    // After the update, the pending event should be send to local channel
    inputGate.updateInputChannel(new InputChannelDeploymentDescriptor(new ResultPartitionID(unknownPartitionId.getPartitionId(), unknownPartitionId.getProducerId()), ResultPartitionLocation.createLocal()));
    verify(partitionManager, times(2)).createSubpartitionView(any(ResultPartitionID.class), anyInt(), any(BufferProvider.class), any(BufferAvailabilityListener.class));
    verify(taskEventDispatcher, times(2)).publish(any(ResultPartitionID.class), any(TaskEvent.class));
}

Example 2 with BufferPool

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

the class RecordWriterTest method testClearBuffersAfterExceptionInPartitionWriter.

@Test
public void testClearBuffersAfterExceptionInPartitionWriter() throws Exception {
    NetworkBufferPool buffers = null;
    BufferPool bufferPool = null;
    try {
        buffers = new NetworkBufferPool(1, 1024, MemoryType.HEAP);
        bufferPool = spy(buffers.createBufferPool(1, Integer.MAX_VALUE));
        ResultPartitionWriter partitionWriter = mock(ResultPartitionWriter.class);
        when(partitionWriter.getBufferProvider()).thenReturn(checkNotNull(bufferPool));
        when(partitionWriter.getNumberOfOutputChannels()).thenReturn(1);
        // Recycle buffer and throw Exception
        doAnswer(new Answer<Void>() {

            @Override
            public Void answer(InvocationOnMock invocation) throws Throwable {
                Buffer buffer = (Buffer) invocation.getArguments()[0];
                buffer.recycle();
                throw new RuntimeException("Expected test Exception");
            }
        }).when(partitionWriter).writeBuffer(any(Buffer.class), anyInt());
        RecordWriter<IntValue> recordWriter = new RecordWriter<>(partitionWriter);
        try {
            // manual flush here doesn't test this case (see next).
            for (; ; ) {
                recordWriter.emit(new IntValue(0));
            }
        } catch (Exception e) {
            // Verify that the buffer is not part of the record writer state after a failure
            // to flush it out. If the buffer is still part of the record writer state, this
            // will fail, because the buffer has already been recycled. NOTE: The mock
            // partition writer needs to recycle the buffer to correctly test this.
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(1)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(1)).requestBufferBlocking();
        try {
            // Verify that manual flushing correctly clears the buffer.
            recordWriter.emit(new IntValue(0));
            recordWriter.flush();
            Assert.fail("Did not throw expected test Exception");
        } catch (Exception e) {
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(2)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(2)).requestBufferBlocking();
        try {
            // Verify that broadcast emit correctly clears the buffer.
            for (; ; ) {
                recordWriter.broadcastEmit(new IntValue(0));
            }
        } catch (Exception e) {
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(3)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(3)).requestBufferBlocking();
        try {
            // Verify that end of super step correctly clears the buffer.
            recordWriter.emit(new IntValue(0));
            recordWriter.broadcastEvent(EndOfSuperstepEvent.INSTANCE);
            Assert.fail("Did not throw expected test Exception");
        } catch (Exception e) {
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(4)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(4)).requestBufferBlocking();
        try {
            // Verify that broadcasting and event correctly clears the buffer.
            recordWriter.emit(new IntValue(0));
            recordWriter.broadcastEvent(new TestTaskEvent());
            Assert.fail("Did not throw expected test Exception");
        } catch (Exception e) {
            recordWriter.clearBuffers();
        }
        // Verify expected methods have been called
        verify(partitionWriter, times(5)).writeBuffer(any(Buffer.class), anyInt());
        verify(bufferPool, times(5)).requestBufferBlocking();
    } finally {
        if (bufferPool != null) {
            assertEquals(1, bufferPool.getNumberOfAvailableMemorySegments());
            bufferPool.lazyDestroy();
        }
        if (buffers != null) {
            assertEquals(1, buffers.getNumberOfAvailableMemorySegments());
            buffers.destroy();
        }
    }
}

Example 3 with BufferPool

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

the class NetworkEnvironmentTest method createSingleInputGateMock.

/**
	 * Helper to create a mock of a {@link SingleInputGate} for use by a {@link Task} inside
	 * {@link NetworkEnvironment#registerTask(Task)}.
	 *
	 * @param partitionType
	 * 		the consumed partition type
	 * @param channels
	 * 		the nummer of input channels
	 *
	 * @return mock with minimal functionality necessary by {@link NetworkEnvironment#registerTask(Task)}
	 */
private static SingleInputGate createSingleInputGateMock(final ResultPartitionType partitionType, final int channels) {
    SingleInputGate ig = mock(SingleInputGate.class);
    when(ig.getConsumedPartitionType()).thenReturn(partitionType);
    when(ig.getNumberOfInputChannels()).thenReturn(channels);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(final InvocationOnMock invocation) throws Throwable {
            BufferPool bp = invocation.getArgumentAt(0, BufferPool.class);
            if (partitionType == ResultPartitionType.PIPELINED_BOUNDED) {
                assertEquals(channels * 2 + 8, bp.getMaxNumberOfMemorySegments());
            } else {
                assertEquals(Integer.MAX_VALUE, bp.getMaxNumberOfMemorySegments());
            }
            return null;
        }
    }).when(ig).setBufferPool(any(BufferPool.class));
    return ig;
}

Example 4 with BufferPool

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

the class BarrierBufferMassiveRandomTest method testWithTwoChannelsAndRandomBarriers.

@Test
public void testWithTwoChannelsAndRandomBarriers() {
    IOManager ioMan = null;
    try {
        ioMan = new IOManagerAsync();
        BufferPool pool1 = new NetworkBufferPool(100, PAGE_SIZE, MemoryType.HEAP).createBufferPool(100, 100);
        BufferPool pool2 = new NetworkBufferPool(100, PAGE_SIZE, MemoryType.HEAP).createBufferPool(100, 100);
        RandomGeneratingInputGate myIG = new RandomGeneratingInputGate(new BufferPool[] { pool1, pool2 }, new BarrierGenerator[] { new CountBarrier(100000), new RandomBarrier(100000) });
        BarrierBuffer barrierBuffer = new BarrierBuffer(myIG, ioMan);
        for (int i = 0; i < 2000000; i++) {
            BufferOrEvent boe = barrierBuffer.getNextNonBlocked();
            if (boe.isBuffer()) {
                boe.getBuffer().recycle();
            }
        }
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    } finally {
        if (ioMan != null) {
            ioMan.shutdown();
        }
    }
}

Example 5 with BufferPool

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

the class NetworkEnvironment method registerTask.

// --------------------------------------------------------------------------------------------
//  Task operations
// --------------------------------------------------------------------------------------------
public void registerTask(Task task) throws IOException {
    final ResultPartition[] producedPartitions = task.getProducedPartitions();
    final ResultPartitionWriter[] writers = task.getAllWriters();
    if (writers.length != producedPartitions.length) {
        throw new IllegalStateException("Unequal number of writers and partitions.");
    }
    synchronized (lock) {
        if (isShutdown) {
            throw new IllegalStateException("NetworkEnvironment is shut down");
        }
        for (int i = 0; i < producedPartitions.length; i++) {
            final ResultPartition partition = producedPartitions[i];
            final ResultPartitionWriter writer = writers[i];
            // Buffer pool for the partition
            BufferPool bufferPool = null;
            try {
                int maxNumberOfMemorySegments = partition.getPartitionType().isBounded() ? partition.getNumberOfSubpartitions() * networkBuffersPerChannel + extraNetworkBuffersPerGate : Integer.MAX_VALUE;
                bufferPool = networkBufferPool.createBufferPool(partition.getNumberOfSubpartitions(), maxNumberOfMemorySegments);
                partition.registerBufferPool(bufferPool);
                resultPartitionManager.registerResultPartition(partition);
            } catch (Throwable t) {
                if (bufferPool != null) {
                    bufferPool.lazyDestroy();
                }
                if (t instanceof IOException) {
                    throw (IOException) t;
                } else {
                    throw new IOException(t.getMessage(), t);
                }
            }
            // Register writer with task event dispatcher
            taskEventDispatcher.registerWriterForIncomingTaskEvents(writer.getPartitionId(), writer);
        }
        // Setup the buffer pool for each buffer reader
        final SingleInputGate[] inputGates = task.getAllInputGates();
        for (SingleInputGate gate : inputGates) {
            BufferPool bufferPool = null;
            try {
                int maxNumberOfMemorySegments = gate.getConsumedPartitionType().isBounded() ? gate.getNumberOfInputChannels() * networkBuffersPerChannel + extraNetworkBuffersPerGate : Integer.MAX_VALUE;
                bufferPool = networkBufferPool.createBufferPool(gate.getNumberOfInputChannels(), maxNumberOfMemorySegments);
                gate.setBufferPool(bufferPool);
            } catch (Throwable t) {
                if (bufferPool != null) {
                    bufferPool.lazyDestroy();
                }
                if (t instanceof IOException) {
                    throw (IOException) t;
                } else {
                    throw new IOException(t.getMessage(), t);
                }
            }
        }
    }
}