| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| package org.apache.flink.iteration.operator.perround; |
| |
| import org.apache.flink.api.common.state.ListState; |
| import org.apache.flink.api.common.state.ListStateDescriptor; |
| import org.apache.flink.api.common.typeinfo.BasicTypeInfo; |
| import org.apache.flink.api.common.typeutils.base.IntSerializer; |
| import org.apache.flink.iteration.IterationListener; |
| import org.apache.flink.iteration.IterationRecord; |
| import org.apache.flink.iteration.operator.OperatorUtils; |
| import org.apache.flink.iteration.operator.OperatorWrapper; |
| import org.apache.flink.iteration.operator.WrapperOperatorFactory; |
| import org.apache.flink.iteration.operator.allround.LifeCycle; |
| import org.apache.flink.iteration.proxy.ProxyKeySelector; |
| import org.apache.flink.iteration.typeinfo.IterationRecordTypeInfo; |
| import org.apache.flink.runtime.checkpoint.CheckpointMetaData; |
| import org.apache.flink.runtime.checkpoint.CheckpointMetricsBuilder; |
| import org.apache.flink.runtime.checkpoint.CheckpointOptions; |
| import org.apache.flink.runtime.checkpoint.CheckpointType; |
| import org.apache.flink.runtime.checkpoint.TaskStateSnapshot; |
| import org.apache.flink.runtime.io.network.api.EndOfData; |
| import org.apache.flink.runtime.io.network.api.StopMode; |
| import org.apache.flink.runtime.jobgraph.OperatorID; |
| import org.apache.flink.runtime.state.CheckpointStorageLocationReference; |
| import org.apache.flink.runtime.state.OperatorStateCheckpointOutputStream; |
| import org.apache.flink.runtime.state.StateInitializationContext; |
| import org.apache.flink.runtime.state.StatePartitionStreamProvider; |
| import org.apache.flink.runtime.state.StateSnapshotContext; |
| import org.apache.flink.streaming.api.graph.StreamConfig; |
| import org.apache.flink.streaming.api.operators.AbstractStreamOperator; |
| import org.apache.flink.streaming.api.operators.BoundedOneInput; |
| import org.apache.flink.streaming.api.operators.OneInputStreamOperator; |
| import org.apache.flink.streaming.api.operators.Output; |
| import org.apache.flink.streaming.api.operators.SimpleOperatorFactory; |
| import org.apache.flink.streaming.api.operators.StreamOperator; |
| import org.apache.flink.streaming.api.operators.StreamOperatorFactory; |
| import org.apache.flink.streaming.api.operators.StreamOperatorParameters; |
| import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; |
| import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask; |
| import org.apache.flink.streaming.runtime.tasks.StreamTask; |
| import org.apache.flink.streaming.runtime.tasks.StreamTaskMailboxTestHarness; |
| import org.apache.flink.streaming.runtime.tasks.StreamTaskMailboxTestHarnessBuilder; |
| import org.apache.flink.util.Collector; |
| import org.apache.flink.util.TestLogger; |
| |
| import org.apache.commons.collections.IteratorUtils; |
| import org.apache.commons.io.output.CloseShieldOutputStream; |
| import org.junit.Test; |
| |
| import java.io.DataInputStream; |
| import java.io.DataOutputStream; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.Collections; |
| import java.util.List; |
| import java.util.stream.Collectors; |
| import java.util.stream.IntStream; |
| |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertNotNull; |
| import static org.junit.Assert.assertTrue; |
| |
| /** Tests the {@link OneInputPerRoundWrapperOperator}. */ |
| public class OneInputPerRoundWrapperOperatorTest extends TestLogger { |
| |
| private static final List<LifeCycle> LIFE_CYCLES = new ArrayList<>(); |
| |
| @Test |
| public void testProcessElementsAndEpochWatermarks() throws Exception { |
| StreamOperatorFactory<IterationRecord<Integer>> wrapperFactory = |
| new WrapperOperatorFactory<>( |
| SimpleOperatorFactory.of(new LifeCycleTrackingOneInputStreamOperator()), |
| new PerRoundOperatorWrapper<>()); |
| OperatorID operatorId = new OperatorID(); |
| |
| try (StreamTaskMailboxTestHarness<IterationRecord<Integer>> harness = |
| new StreamTaskMailboxTestHarnessBuilder<>( |
| OneInputStreamTask::new, |
| new IterationRecordTypeInfo<>(BasicTypeInfo.INT_TYPE_INFO)) |
| .addInput( |
| new IterationRecordTypeInfo<>(BasicTypeInfo.INT_TYPE_INFO), |
| 1, |
| new ProxyKeySelector<Integer, Integer>(x -> x % 2)) |
| .setupOutputForSingletonOperatorChain(wrapperFactory, operatorId) |
| .build()) { |
| harness.processElement(new StreamRecord<>(IterationRecord.newRecord(5, 1), 2)); |
| harness.processElement(new StreamRecord<>(IterationRecord.newRecord(6, 2), 3)); |
| |
| // Checks the output |
| assertEquals( |
| Arrays.asList( |
| new StreamRecord<>(IterationRecord.newRecord(5, 1), 2), |
| new StreamRecord<>(IterationRecord.newRecord(6, 2), 3)), |
| new ArrayList<>(harness.getOutput())); |
| |
| // Checks the other lifecycles. |
| harness.getStreamTask() |
| .triggerCheckpointOnBarrier( |
| new CheckpointMetaData(5, 2), |
| CheckpointOptions.alignedNoTimeout( |
| CheckpointType.CHECKPOINT, |
| CheckpointStorageLocationReference.getDefault()), |
| new CheckpointMetricsBuilder() |
| .setAlignmentDurationNanos(0) |
| .setBytesProcessedDuringAlignment(0)); |
| harness.processAll(); |
| |
| harness.getStreamTask().notifyCheckpointCompleteAsync(5); |
| harness.processAll(); |
| |
| harness.getStreamTask().notifyCheckpointAbortAsync(6, 5); |
| harness.processAll(); |
| |
| harness.getOutput().clear(); |
| harness.processElement( |
| new StreamRecord<>(IterationRecord.newEpochWatermark(1, "only-one"))); |
| harness.processElement( |
| new StreamRecord<>( |
| IterationRecord.newEpochWatermark( |
| IterationRecord.END_EPOCH_WATERMARK, "only-one"))); |
| |
| // Checks the output |
| assertEquals( |
| Arrays.asList( |
| new StreamRecord<>( |
| IterationRecord.newEpochWatermark( |
| 1, OperatorUtils.getUniqueSenderId(operatorId, 0))), |
| new StreamRecord<>( |
| IterationRecord.newEpochWatermark( |
| IterationRecord.END_EPOCH_WATERMARK, |
| OperatorUtils.getUniqueSenderId(operatorId, 0)))), |
| new ArrayList<>(harness.getOutput())); |
| |
| harness.processEvent(new EndOfData(StopMode.DRAIN), 0); |
| harness.endInput(); |
| harness.finishProcessing(); |
| |
| assertEquals( |
| Arrays.asList( |
| /* First wrapped operator */ |
| LifeCycle.SETUP, |
| LifeCycle.INITIALIZE_STATE, |
| LifeCycle.OPEN, |
| LifeCycle.PROCESS_ELEMENT, |
| /* second wrapped operator */ |
| LifeCycle.SETUP, |
| LifeCycle.INITIALIZE_STATE, |
| LifeCycle.OPEN, |
| LifeCycle.PROCESS_ELEMENT, |
| /* states */ |
| LifeCycle.PREPARE_SNAPSHOT_PRE_BARRIER, |
| LifeCycle.PREPARE_SNAPSHOT_PRE_BARRIER, |
| LifeCycle.SNAPSHOT_STATE, |
| LifeCycle.SNAPSHOT_STATE, |
| LifeCycle.NOTIFY_CHECKPOINT_COMPLETE, |
| LifeCycle.NOTIFY_CHECKPOINT_COMPLETE, |
| LifeCycle.NOTIFY_CHECKPOINT_ABORT, |
| LifeCycle.NOTIFY_CHECKPOINT_ABORT, |
| LifeCycle.EPOCH_WATERMARK_INCREMENTED, |
| LifeCycle.END_INPUT, |
| LifeCycle.FINISH, |
| LifeCycle.CLOSE, |
| LifeCycle.ITERATION_TERMINATION, |
| LifeCycle.END_INPUT, |
| LifeCycle.FINISH, |
| LifeCycle.CLOSE), |
| LIFE_CYCLES); |
| } |
| } |
| |
| @Test |
| public void testSnapshotAndRestore() throws Exception { |
| StreamOperatorFactory<IterationRecord<Integer>> wrapperFactory = |
| new RecordingOperatorFactory<>( |
| SimpleOperatorFactory.of(new StatefulOperator()), |
| new PerRoundOperatorWrapper<>()); |
| OperatorID operatorId = new OperatorID(); |
| |
| TaskStateSnapshot taskStateSnapshot; |
| try (StreamTaskMailboxTestHarness<IterationRecord<Integer>> harness = |
| new StreamTaskMailboxTestHarnessBuilder<>( |
| OneInputStreamTask::new, |
| new IterationRecordTypeInfo<>(BasicTypeInfo.INT_TYPE_INFO)) |
| .addInput(new IterationRecordTypeInfo<>(BasicTypeInfo.INT_TYPE_INFO)) |
| .setupOutputForSingletonOperatorChain(wrapperFactory, operatorId) |
| .build()) { |
| harness.getTaskStateManager().getWaitForReportLatch().reset(); |
| |
| harness.processElement(new StreamRecord<>(IterationRecord.newRecord(100, 0))); |
| harness.processElement(new StreamRecord<>(IterationRecord.newRecord(101, 1))); |
| harness.processElement(new StreamRecord<>(IterationRecord.newRecord(102, 2))); |
| harness.processElement(new StreamRecord<>(IterationRecord.newRecord(103, 3))); |
| harness.processElement(new StreamRecord<>(IterationRecord.newRecord(103, 4))); |
| harness.processElement( |
| new StreamRecord<>(IterationRecord.newEpochWatermark(0, "fake"))); |
| harness.getStreamTask() |
| .triggerCheckpointAsync( |
| new CheckpointMetaData(2, 1000), |
| CheckpointOptions.alignedNoTimeout( |
| CheckpointType.CHECKPOINT, |
| CheckpointStorageLocationReference.getDefault())); |
| harness.processAll(); |
| |
| harness.getTaskStateManager().getWaitForReportLatch().await(); |
| taskStateSnapshot = harness.getTaskStateManager().getLastJobManagerTaskStateSnapshot(); |
| } |
| |
| assertNotNull(taskStateSnapshot); |
| try (StreamTaskMailboxTestHarness<IterationRecord<Integer>> harness = |
| new StreamTaskMailboxTestHarnessBuilder<>( |
| OneInputStreamTask::new, |
| new IterationRecordTypeInfo<>(BasicTypeInfo.INT_TYPE_INFO)) |
| .setTaskStateSnapshot(2, taskStateSnapshot) |
| .addInput(new IterationRecordTypeInfo<>(BasicTypeInfo.INT_TYPE_INFO)) |
| .setupOutputForSingletonOperatorChain(wrapperFactory, operatorId) |
| .build()) { |
| AbstractPerRoundWrapperOperator<Integer, OneInputStreamOperator<Integer, Integer>> |
| wrapperOperator = |
| ((AbstractPerRoundWrapperOperator) RecordingOperatorFactory.latest); |
| assertEquals(0, wrapperOperator.getLatestEpochWatermark()); |
| assertEquals( |
| Arrays.asList(1, 2, 3, 4), |
| new ArrayList<>(wrapperOperator.getWrappedOperators().keySet())); |
| |
| for (OneInputStreamOperator<Integer, Integer> o : |
| wrapperOperator.getWrappedOperators().values()) { |
| StatefulOperator statefulOperator = (StatefulOperator) o; |
| assertTrue(statefulOperator.hasState); |
| if (statefulOperator.index % 2 == 0) { |
| assertEquals(Collections.emptyList(), statefulOperator.values); |
| } else { |
| assertEquals( |
| IntStream.range(0, statefulOperator.index) |
| .boxed() |
| .collect(Collectors.toList()), |
| statefulOperator.values); |
| } |
| } |
| } |
| } |
| |
| private static class RecordingOperatorFactory<OUT> extends WrapperOperatorFactory<OUT> { |
| |
| static StreamOperator<?> latest = null; |
| |
| public RecordingOperatorFactory( |
| StreamOperatorFactory<OUT> operatorFactory, |
| OperatorWrapper<OUT, IterationRecord<OUT>> wrapper) { |
| super(operatorFactory, wrapper); |
| } |
| |
| @Override |
| public <T extends StreamOperator<IterationRecord<OUT>>> T createStreamOperator( |
| StreamOperatorParameters<IterationRecord<OUT>> parameters) { |
| latest = super.createStreamOperator(parameters); |
| return (T) latest; |
| } |
| } |
| |
| private static class LifeCycleTrackingOneInputStreamOperator |
| extends AbstractStreamOperator<Integer> |
| implements OneInputStreamOperator<Integer, Integer>, |
| BoundedOneInput, |
| IterationListener<Integer> { |
| |
| @Override |
| public void setup( |
| StreamTask<?, ?> containingTask, |
| StreamConfig config, |
| Output<StreamRecord<Integer>> output) { |
| super.setup(containingTask, config, output); |
| LIFE_CYCLES.add(LifeCycle.SETUP); |
| } |
| |
| @Override |
| public void open() throws Exception { |
| super.open(); |
| LIFE_CYCLES.add(LifeCycle.OPEN); |
| } |
| |
| @Override |
| public void initializeState(StateInitializationContext context) throws Exception { |
| super.initializeState(context); |
| LIFE_CYCLES.add(LifeCycle.INITIALIZE_STATE); |
| } |
| |
| @Override |
| public void finish() throws Exception { |
| super.finish(); |
| LIFE_CYCLES.add(LifeCycle.FINISH); |
| } |
| |
| @Override |
| public void close() throws Exception { |
| super.close(); |
| LIFE_CYCLES.add(LifeCycle.CLOSE); |
| } |
| |
| @Override |
| public void prepareSnapshotPreBarrier(long checkpointId) throws Exception { |
| super.prepareSnapshotPreBarrier(checkpointId); |
| LIFE_CYCLES.add(LifeCycle.PREPARE_SNAPSHOT_PRE_BARRIER); |
| } |
| |
| @Override |
| public void snapshotState(StateSnapshotContext context) throws Exception { |
| super.snapshotState(context); |
| LIFE_CYCLES.add(LifeCycle.SNAPSHOT_STATE); |
| } |
| |
| @Override |
| public void notifyCheckpointComplete(long checkpointId) throws Exception { |
| super.notifyCheckpointComplete(checkpointId); |
| LIFE_CYCLES.add(LifeCycle.NOTIFY_CHECKPOINT_COMPLETE); |
| } |
| |
| @Override |
| public void notifyCheckpointAborted(long checkpointId) throws Exception { |
| super.notifyCheckpointAborted(checkpointId); |
| LIFE_CYCLES.add(LifeCycle.NOTIFY_CHECKPOINT_ABORT); |
| } |
| |
| @Override |
| public void processElement(StreamRecord<Integer> element) throws Exception { |
| output.collect(element); |
| LIFE_CYCLES.add(LifeCycle.PROCESS_ELEMENT); |
| } |
| |
| @Override |
| public void endInput() throws Exception { |
| LIFE_CYCLES.add(LifeCycle.END_INPUT); |
| } |
| |
| @Override |
| public void onEpochWatermarkIncremented( |
| int epochWatermark, Context context, Collector<Integer> collector) { |
| LIFE_CYCLES.add(LifeCycle.EPOCH_WATERMARK_INCREMENTED); |
| } |
| |
| @Override |
| public void onIterationTerminated(Context context, Collector<Integer> collector) { |
| LIFE_CYCLES.add(LifeCycle.ITERATION_TERMINATION); |
| } |
| } |
| |
| private static class StatefulOperator extends AbstractStreamOperator<Integer> |
| implements OneInputStreamOperator<Integer, Integer> { |
| |
| private static int nextIndex = 0; |
| |
| private ListState<Integer> indexState; |
| |
| private boolean hasState; |
| |
| private int index; |
| |
| private List<Integer> values; |
| |
| @Override |
| public void initializeState(StateInitializationContext context) throws Exception { |
| super.initializeState(context); |
| |
| indexState = |
| context.getOperatorStateStore() |
| .getListState( |
| new ListStateDescriptor<>("index", IntSerializer.INSTANCE)); |
| List<Integer> indexStateValues = IteratorUtils.toList(indexState.get().iterator()); |
| if (indexStateValues.size() == 0) { |
| hasState = false; |
| index = nextIndex++; |
| |
| if (index % 2 == 0) { |
| values = new ArrayList<>(); |
| } else { |
| values = IntStream.range(0, index).boxed().collect(Collectors.toList()); |
| } |
| |
| } else { |
| hasState = true; |
| index = indexStateValues.get(0); |
| values = new ArrayList<>(); |
| for (StatePartitionStreamProvider provider : context.getRawOperatorStateInputs()) { |
| try (DataInputStream dis = new DataInputStream(provider.getStream())) { |
| values.add(dis.readInt()); |
| } |
| } |
| } |
| } |
| |
| @Override |
| public void snapshotState(StateSnapshotContext context) throws Exception { |
| super.snapshotState(context); |
| |
| indexState.update(Collections.singletonList(index)); |
| |
| OperatorStateCheckpointOutputStream rawOutput = context.getRawOperatorStateOutput(); |
| for (int value : values) { |
| rawOutput.startNewPartition(); |
| try (DataOutputStream dos = |
| new DataOutputStream(new CloseShieldOutputStream(rawOutput))) { |
| dos.writeInt(value); |
| dos.flush(); |
| } |
| } |
| } |
| |
| @Override |
| public void processElement(StreamRecord<Integer> element) throws Exception {} |
| } |
| } |
