use of org.apache.flink.runtime.io.network.NettyShuffleEnvironment in project flink by apache.
the class SourceStreamTaskTest method testTriggeringCheckpointAfterSourceThreadFinished.
@Test
public void testTriggeringCheckpointAfterSourceThreadFinished() throws Exception {
ResultPartition[] partitionWriters = new ResultPartition[2];
try (NettyShuffleEnvironment env = new NettyShuffleEnvironmentBuilder().setNumNetworkBuffers(partitionWriters.length * 2).build()) {
for (int i = 0; i < partitionWriters.length; ++i) {
partitionWriters[i] = PartitionTestUtils.createPartition(env, ResultPartitionType.PIPELINED_BOUNDED, 1);
partitionWriters[i].setup();
}
final CompletableFuture<Long> checkpointCompleted = new CompletableFuture<>();
try (StreamTaskMailboxTestHarness<String> testHarness = new StreamTaskMailboxTestHarnessBuilder<>(SourceStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO).modifyStreamConfig(config -> config.setCheckpointingEnabled(true)).setCheckpointResponder(new TestCheckpointResponder() {
@Override
public void acknowledgeCheckpoint(JobID jobID, ExecutionAttemptID executionAttemptID, long checkpointId, CheckpointMetrics checkpointMetrics, TaskStateSnapshot subtaskState) {
super.acknowledgeCheckpoint(jobID, executionAttemptID, checkpointId, checkpointMetrics, subtaskState);
checkpointCompleted.complete(checkpointId);
}
}).addAdditionalOutput(partitionWriters).setupOperatorChain(new StreamSource<>(new MockSource(0, 0, 1))).finishForSingletonOperatorChain(StringSerializer.INSTANCE).build()) {
testHarness.processAll();
CompletableFuture<Void> taskFinished = testHarness.getStreamTask().getCompletionFuture();
do {
testHarness.processAll();
} while (!taskFinished.isDone());
Future<Boolean> checkpointFuture = triggerCheckpoint(testHarness, 2);
// Notifies the result partition that all records are processed after the
// last checkpoint is triggered.
checkState(checkpointFuture instanceof CompletableFuture, "The trigger future should " + " be also CompletableFuture.");
((CompletableFuture<?>) checkpointFuture).thenAccept((ignored) -> {
for (ResultPartition resultPartition : partitionWriters) {
resultPartition.onSubpartitionAllDataProcessed(0);
}
});
checkpointCompleted.whenComplete((id, error) -> testHarness.getStreamTask().notifyCheckpointCompleteAsync(2));
testHarness.finishProcessing();
assertTrue(checkpointFuture.isDone());
// EndOfUserRecordEvent.
for (ResultPartition resultPartition : partitionWriters) {
assertEquals(3, resultPartition.getNumberOfQueuedBuffers());
}
}
} finally {
for (ResultPartitionWriter writer : partitionWriters) {
if (writer != null) {
writer.close();
}
}
}
}
use of org.apache.flink.runtime.io.network.NettyShuffleEnvironment in project flink by apache.
the class RecordWriterTest method createResultPartition.
public static ResultPartition createResultPartition(int bufferSize, int numSubpartitions) throws IOException {
NettyShuffleEnvironment env = new NettyShuffleEnvironmentBuilder().setBufferSize(bufferSize).build();
ResultPartition partition = createPartition(env, ResultPartitionType.PIPELINED, numSubpartitions);
partition.setup();
return partition;
}
use of org.apache.flink.runtime.io.network.NettyShuffleEnvironment in project flink by apache.
the class PipelinedApproximateSubpartitionTest method createResultPartition.
private static BufferWritingResultPartition createResultPartition() throws IOException {
NettyShuffleEnvironment network = new NettyShuffleEnvironmentBuilder().setNumNetworkBuffers(10).setBufferSize(BUFFER_SIZE).build();
ResultPartition resultPartition = createPartition(network, NoOpFileChannelManager.INSTANCE, ResultPartitionType.PIPELINED_APPROXIMATE, 2);
resultPartition.setup();
return (BufferWritingResultPartition) resultPartition;
}
Aggregations