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

Example 1 with TestTaskEvent

use of org.apache.flink.runtime.io.network.util.TestTaskEvent 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 TestTaskEvent

use of org.apache.flink.runtime.io.network.util.TestTaskEvent 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 TestTaskEvent

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

the class EventSerializerTest method testSerializeDeserializeEvent.

@Test
public void testSerializeDeserializeEvent() throws Exception {
    AbstractEvent[] events = { EndOfPartitionEvent.INSTANCE, EndOfSuperstepEvent.INSTANCE, new CheckpointBarrier(1678L, 4623784L, CheckpointOptions.forFullCheckpoint()), new TestTaskEvent(Math.random(), 12361231273L), new CancelCheckpointMarker(287087987329842L) };
    for (AbstractEvent evt : events) {
        ByteBuffer serializedEvent = EventSerializer.toSerializedEvent(evt);
        assertTrue(serializedEvent.hasRemaining());
        AbstractEvent deserialized = EventSerializer.fromSerializedEvent(serializedEvent, getClass().getClassLoader());
        assertNotNull(deserialized);
        assertEquals(evt, deserialized);
    }
}