Examples with ResultPartitionWriter org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter used on opensource projects

Example 26 with ResultPartitionWriter

use of org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter in project flink by apache.

the class StreamTaskFinalCheckpointsTest method testReportOperatorsFinishedInCheckpoint.

@Test
public void testReportOperatorsFinishedInCheckpoint() throws Exception {
    ResultPartition[] partitionWriters = new ResultPartition[2];
    try {
        for (int i = 0; i < partitionWriters.length; ++i) {
            partitionWriters[i] = PartitionTestUtils.createPartition(ResultPartitionType.PIPELINED_BOUNDED);
            partitionWriters[i].setup();
        }
        final CompletingCheckpointResponder checkpointResponder = new CompletingCheckpointResponder();
        try (StreamTaskMailboxTestHarness<String> testHarness = new StreamTaskMailboxTestHarnessBuilder<>(OneInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO).addInput(BasicTypeInfo.STRING_TYPE_INFO, 1).addAdditionalOutput(partitionWriters).setCheckpointResponder(checkpointResponder).modifyStreamConfig(config -> config.setCheckpointingEnabled(true)).setupOperatorChain(new StatefulOperator()).finishForSingletonOperatorChain(StringSerializer.INSTANCE).build()) {
            checkpointResponder.setHandlers(testHarness.streamTask::notifyCheckpointCompleteAsync, testHarness.streamTask::notifyCheckpointAbortAsync);
            // Trigger the first checkpoint before we call operators' finish method.
            CompletableFuture<Boolean> checkpointFuture = triggerCheckpoint(testHarness, 2);
            processMailTillCheckpointSucceeds(testHarness, checkpointFuture);
            assertEquals(2, testHarness.getTaskStateManager().getReportedCheckpointId());
            assertFalse(testHarness.getTaskStateManager().getJobManagerTaskStateSnapshotsByCheckpointId().get(2L).isTaskFinished());
            // Trigger the first checkpoint after we call operators' finish method.
            // The checkpoint is added to the mailbox and will be processed in the
            // mailbox loop after call operators' finish method in the afterInvoke()
            // method.
            testHarness.processEvent(new EndOfData(StopMode.DRAIN), 0, 0);
            checkpointFuture = triggerCheckpoint(testHarness, 4);
            checkpointFuture.thenAccept((ignored) -> {
                for (ResultPartition resultPartition : partitionWriters) {
                    resultPartition.onSubpartitionAllDataProcessed(0);
                }
            });
            testHarness.processAll();
            testHarness.finishProcessing();
            assertTrue(checkpointFuture.isDone());
            testHarness.getTaskStateManager().getWaitForReportLatch().await();
            assertTrue(testHarness.getTaskStateManager().getJobManagerTaskStateSnapshotsByCheckpointId().get(4L).isTaskFinished());
        }
    } finally {
        for (ResultPartitionWriter writer : partitionWriters) {
            if (writer != null) {
                writer.close();
            }
        }
    }
}

Example 27 with ResultPartitionWriter

use of org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter in project flink by apache.

the class StreamTaskMailboxTestHarnessBuilder method buildUnrestored.

public StreamTaskMailboxTestHarness<OUT> buildUnrestored() throws Exception {
    TestTaskStateManagerBuilder taskStateManagerBuilder = TestTaskStateManager.builder().setLocalRecoveryConfig(localRecoveryConfig).setCheckpointResponder(checkpointResponder);
    if (taskStateSnapshots != null) {
        taskStateManagerBuilder.setReportedCheckpointId(taskStateSnapshots.keySet().iterator().next()).setJobManagerTaskStateSnapshotsByCheckpointId(taskStateSnapshots);
    }
    TestTaskStateManager taskStateManager = taskStateManagerBuilder.build();
    StreamMockEnvironment streamMockEnvironment = new StreamMockEnvironment(new JobID(), new ExecutionAttemptID(), jobConfig, streamConfig.getConfiguration(), executionConfig, memorySize, new MockInputSplitProvider(), bufferSize, taskStateManager, collectNetworkEvents);
    streamMockEnvironment.setCheckpointResponder(taskStateManager.getCheckpointResponder());
    streamMockEnvironment.setTaskManagerInfo(taskManagerRuntimeInfo);
    initializeInputs(streamMockEnvironment);
    checkState(inputGates != null, "InputGates hasn't been initialised");
    StreamElementSerializer<OUT> outputStreamRecordSerializer = new StreamElementSerializer<>(outputSerializer);
    Queue<Object> outputList = new ArrayDeque<>();
    streamMockEnvironment.addOutput(outputList, outputStreamRecordSerializer);
    streamMockEnvironment.setTaskMetricGroup(taskMetricGroup);
    for (ResultPartitionWriter writer : additionalOutputs) {
        streamMockEnvironment.addOutput(writer);
    }
    StreamTask<OUT, ?> task = taskFactory.apply(streamMockEnvironment);
    return new StreamTaskMailboxTestHarness<>(task, outputList, inputGates, streamMockEnvironment);
}

Example 28 with ResultPartitionWriter

use of org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter in project flink by apache.

the class ResultPartitionTest method testAddOnReleasedPartition.

/**
 * Tests {@link ResultPartition#emitRecord} on a partition which has already been released.
 *
 * @param partitionType the result partition type to set up
 */
private void testAddOnReleasedPartition(ResultPartitionType partitionType) throws Exception {
    TestResultPartitionConsumableNotifier notifier = new TestResultPartitionConsumableNotifier();
    BufferWritingResultPartition bufferWritingResultPartition = createResultPartition(partitionType);
    ResultPartitionWriter partitionWriter = ConsumableNotifyingResultPartitionWriterDecorator.decorate(Collections.singleton(PartitionTestUtils.createPartitionDeploymentDescriptor(partitionType)), new ResultPartitionWriter[] { bufferWritingResultPartition }, new NoOpTaskActions(), new JobID(), notifier)[0];
    try {
        partitionWriter.release(null);
        // partitionWriter.emitRecord() should silently drop the given record
        partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 0);
    } finally {
        assertEquals(1, bufferWritingResultPartition.numBuffersOut.getCount());
        assertEquals(bufferSize, bufferWritingResultPartition.numBytesOut.getCount());
        // the buffer should be recycled for the result partition has already been released
        assertEquals(0, bufferWritingResultPartition.getBufferPool().bestEffortGetNumOfUsedBuffers());
        // should not have notified either
        notifier.check(null, null, null, 0);
    }
}

Example 29 with ResultPartitionWriter

use of org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter in project flink by apache.

the class ResultPartitionTest method testNumBytesProducedCounter.

private void testNumBytesProducedCounter(boolean isBroadcast) throws IOException {
    TestResultPartitionConsumableNotifier notifier = new TestResultPartitionConsumableNotifier();
    JobID jobId = new JobID();
    TaskActions taskActions = new NoOpTaskActions();
    BufferWritingResultPartition bufferWritingResultPartition = createResultPartition(ResultPartitionType.BLOCKING);
    ResultPartitionWriter partitionWriter = ConsumableNotifyingResultPartitionWriterDecorator.decorate(Collections.singleton(PartitionTestUtils.createPartitionDeploymentDescriptor(ResultPartitionType.BLOCKING)), new ResultPartitionWriter[] { bufferWritingResultPartition }, taskActions, jobId, notifier)[0];
    if (isBroadcast) {
        partitionWriter.broadcastRecord(ByteBuffer.allocate(bufferSize));
        assertEquals(bufferSize, bufferWritingResultPartition.numBytesProduced.getCount());
        assertEquals(2 * bufferSize, bufferWritingResultPartition.numBytesOut.getCount());
    } else {
        partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 0);
        assertEquals(bufferSize, bufferWritingResultPartition.numBytesProduced.getCount());
        assertEquals(bufferSize, bufferWritingResultPartition.numBytesOut.getCount());
    }
}

Example 30 with ResultPartitionWriter

use of org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter in project flink by apache.

the class ResultPartitionTest method testAddOnPartition.

/**
 * Tests {@link ResultPartition#emitRecord} on a working partition.
 *
 * @param partitionType the result partition type to set up
 */
private void testAddOnPartition(final ResultPartitionType partitionType) throws Exception {
    TestResultPartitionConsumableNotifier notifier = new TestResultPartitionConsumableNotifier();
    JobID jobId = new JobID();
    TaskActions taskActions = new NoOpTaskActions();
    BufferWritingResultPartition bufferWritingResultPartition = createResultPartition(partitionType);
    ResultPartitionWriter partitionWriter = ConsumableNotifyingResultPartitionWriterDecorator.decorate(Collections.singleton(PartitionTestUtils.createPartitionDeploymentDescriptor(partitionType)), new ResultPartitionWriter[] { bufferWritingResultPartition }, taskActions, jobId, notifier)[0];
    try {
        // partitionWriter.emitRecord() will allocate a new buffer and copies the record to it
        partitionWriter.emitRecord(ByteBuffer.allocate(bufferSize), 0);
    } finally {
        assertEquals(1, bufferWritingResultPartition.numBuffersOut.getCount());
        assertEquals(bufferSize, bufferWritingResultPartition.numBytesOut.getCount());
        assertEquals(1, bufferWritingResultPartition.getBufferPool().bestEffortGetNumOfUsedBuffers());
        // should have been notified for pipelined partitions
        if (partitionType.isPipelined()) {
            notifier.check(jobId, partitionWriter.getPartitionId(), taskActions, 1);
        }
    }
}