| /* |
| * 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.test.checkpointing.utils; |
| |
| import org.apache.flink.api.common.restartstrategy.RestartStrategies; |
| import org.apache.flink.api.common.state.MapStateDescriptor; |
| import org.apache.flink.api.common.typeinfo.BasicTypeInfo; |
| import org.apache.flink.api.java.functions.KeySelector; |
| import org.apache.flink.api.java.tuple.Tuple2; |
| import org.apache.flink.configuration.Configuration; |
| import org.apache.flink.contrib.streaming.state.RocksDBStateBackend; |
| import org.apache.flink.runtime.state.StateBackendLoader; |
| import org.apache.flink.runtime.state.memory.MemoryStateBackend; |
| import org.apache.flink.streaming.api.datastream.BroadcastStream; |
| import org.apache.flink.streaming.api.datastream.KeyedStream; |
| import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; |
| import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction; |
| import org.apache.flink.streaming.api.functions.source.SourceFunction; |
| import org.apache.flink.testutils.migration.MigrationVersion; |
| import org.apache.flink.util.Collector; |
| |
| import org.junit.Assert; |
| import org.junit.Test; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.Parameterized; |
| |
| import java.util.Arrays; |
| import java.util.Collection; |
| import java.util.HashMap; |
| import java.util.Map; |
| |
| /** |
| * Migration ITCases for a stateful job with broadcast state. The tests are parameterized to |
| * (potentially) cover migrating for multiple previous Flink versions, as well as for different |
| * state backends. |
| */ |
| @RunWith(Parameterized.class) |
| public class StatefulJobWBroadcastStateMigrationITCase extends SavepointMigrationTestBase { |
| |
| private static final int NUM_SOURCE_ELEMENTS = 4; |
| |
| // TODO change this to PERFORM_SAVEPOINT to regenerate binary savepoints |
| // TODO Note: You should generate the savepoint based on the release branch instead of the |
| // master. |
| private final StatefulJobSavepointMigrationITCase.ExecutionMode executionMode = |
| StatefulJobSavepointMigrationITCase.ExecutionMode.VERIFY_SAVEPOINT; |
| |
| @Parameterized.Parameters(name = "Migrate Savepoint / Backend: {0}") |
| public static Collection<Tuple2<MigrationVersion, String>> parameters() { |
| return Arrays.asList( |
| Tuple2.of(MigrationVersion.v1_5, StateBackendLoader.MEMORY_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_5, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_6, StateBackendLoader.MEMORY_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_6, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_7, StateBackendLoader.MEMORY_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_7, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_8, StateBackendLoader.MEMORY_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_8, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_9, StateBackendLoader.MEMORY_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_9, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_10, StateBackendLoader.MEMORY_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_10, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_11, StateBackendLoader.MEMORY_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_11, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_12, StateBackendLoader.MEMORY_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_12, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_13, StateBackendLoader.MEMORY_STATE_BACKEND_NAME), |
| Tuple2.of(MigrationVersion.v1_13, StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME)); |
| } |
| |
| private final MigrationVersion testMigrateVersion; |
| private final String testStateBackend; |
| |
| public StatefulJobWBroadcastStateMigrationITCase( |
| Tuple2<MigrationVersion, String> testMigrateVersionAndBackend) throws Exception { |
| this.testMigrateVersion = testMigrateVersionAndBackend.f0; |
| this.testStateBackend = testMigrateVersionAndBackend.f1; |
| } |
| |
| @Test |
| public void testSavepoint() throws Exception { |
| |
| final int parallelism = 4; |
| |
| final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); |
| env.setRestartStrategy(RestartStrategies.noRestart()); |
| |
| switch (testStateBackend) { |
| case StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME: |
| env.setStateBackend(new RocksDBStateBackend(new MemoryStateBackend())); |
| break; |
| case StateBackendLoader.MEMORY_STATE_BACKEND_NAME: |
| env.setStateBackend(new MemoryStateBackend()); |
| break; |
| default: |
| throw new UnsupportedOperationException(); |
| } |
| |
| env.enableCheckpointing(500); |
| env.setParallelism(parallelism); |
| env.setMaxParallelism(parallelism); |
| |
| SourceFunction<Tuple2<Long, Long>> nonParallelSource; |
| SourceFunction<Tuple2<Long, Long>> nonParallelSourceB; |
| SourceFunction<Tuple2<Long, Long>> parallelSource; |
| SourceFunction<Tuple2<Long, Long>> parallelSourceB; |
| KeyedBroadcastProcessFunction< |
| Long, Tuple2<Long, Long>, Tuple2<Long, Long>, Tuple2<Long, Long>> |
| firstBroadcastFunction; |
| KeyedBroadcastProcessFunction< |
| Long, Tuple2<Long, Long>, Tuple2<Long, Long>, Tuple2<Long, Long>> |
| secondBroadcastFunction; |
| |
| final Map<Long, Long> expectedFirstState = new HashMap<>(); |
| expectedFirstState.put(0L, 0L); |
| expectedFirstState.put(1L, 1L); |
| expectedFirstState.put(2L, 2L); |
| expectedFirstState.put(3L, 3L); |
| |
| final Map<String, Long> expectedSecondState = new HashMap<>(); |
| expectedSecondState.put("0", 0L); |
| expectedSecondState.put("1", 1L); |
| expectedSecondState.put("2", 2L); |
| expectedSecondState.put("3", 3L); |
| |
| final Map<Long, String> expectedThirdState = new HashMap<>(); |
| expectedThirdState.put(0L, "0"); |
| expectedThirdState.put(1L, "1"); |
| expectedThirdState.put(2L, "2"); |
| expectedThirdState.put(3L, "3"); |
| |
| if (executionMode == StatefulJobSavepointMigrationITCase.ExecutionMode.PERFORM_SAVEPOINT) { |
| nonParallelSource = |
| new MigrationTestUtils.CheckpointingNonParallelSourceWithListState( |
| NUM_SOURCE_ELEMENTS); |
| nonParallelSourceB = |
| new MigrationTestUtils.CheckpointingNonParallelSourceWithListState( |
| NUM_SOURCE_ELEMENTS); |
| parallelSource = |
| new MigrationTestUtils.CheckpointingParallelSourceWithUnionListState( |
| NUM_SOURCE_ELEMENTS); |
| parallelSourceB = |
| new MigrationTestUtils.CheckpointingParallelSourceWithUnionListState( |
| NUM_SOURCE_ELEMENTS); |
| firstBroadcastFunction = new CheckpointingKeyedBroadcastFunction(); |
| secondBroadcastFunction = new CheckpointingKeyedSingleBroadcastFunction(); |
| } else if (executionMode |
| == StatefulJobSavepointMigrationITCase.ExecutionMode.VERIFY_SAVEPOINT) { |
| nonParallelSource = |
| new MigrationTestUtils.CheckingNonParallelSourceWithListState( |
| NUM_SOURCE_ELEMENTS); |
| nonParallelSourceB = |
| new MigrationTestUtils.CheckingNonParallelSourceWithListState( |
| NUM_SOURCE_ELEMENTS); |
| parallelSource = |
| new MigrationTestUtils.CheckingParallelSourceWithUnionListState( |
| NUM_SOURCE_ELEMENTS); |
| parallelSourceB = |
| new MigrationTestUtils.CheckingParallelSourceWithUnionListState( |
| NUM_SOURCE_ELEMENTS); |
| firstBroadcastFunction = |
| new CheckingKeyedBroadcastFunction(expectedFirstState, expectedSecondState); |
| secondBroadcastFunction = new CheckingKeyedSingleBroadcastFunction(expectedThirdState); |
| } else { |
| throw new IllegalStateException("Unknown ExecutionMode " + executionMode); |
| } |
| |
| KeyedStream<Tuple2<Long, Long>, Long> npStream = |
| env.addSource(nonParallelSource) |
| .uid("CheckpointingSource1") |
| .keyBy( |
| new KeySelector<Tuple2<Long, Long>, Long>() { |
| |
| private static final long serialVersionUID = |
| -4514793867774977152L; |
| |
| @Override |
| public Long getKey(Tuple2<Long, Long> value) throws Exception { |
| return value.f0; |
| } |
| }); |
| |
| KeyedStream<Tuple2<Long, Long>, Long> pStream = |
| env.addSource(parallelSource) |
| .uid("CheckpointingSource2") |
| .keyBy( |
| new KeySelector<Tuple2<Long, Long>, Long>() { |
| |
| private static final long serialVersionUID = |
| 4940496713319948104L; |
| |
| @Override |
| public Long getKey(Tuple2<Long, Long> value) throws Exception { |
| return value.f0; |
| } |
| }); |
| |
| final MapStateDescriptor<Long, Long> firstBroadcastStateDesc = |
| new MapStateDescriptor<>( |
| "broadcast-state-1", |
| BasicTypeInfo.LONG_TYPE_INFO, |
| BasicTypeInfo.LONG_TYPE_INFO); |
| |
| final MapStateDescriptor<String, Long> secondBroadcastStateDesc = |
| new MapStateDescriptor<>( |
| "broadcast-state-2", |
| BasicTypeInfo.STRING_TYPE_INFO, |
| BasicTypeInfo.LONG_TYPE_INFO); |
| |
| final MapStateDescriptor<Long, String> thirdBroadcastStateDesc = |
| new MapStateDescriptor<>( |
| "broadcast-state-3", |
| BasicTypeInfo.LONG_TYPE_INFO, |
| BasicTypeInfo.STRING_TYPE_INFO); |
| |
| BroadcastStream<Tuple2<Long, Long>> npBroadcastStream = |
| env.addSource(nonParallelSourceB) |
| .uid("BrCheckpointingSource1") |
| .broadcast(firstBroadcastStateDesc, secondBroadcastStateDesc); |
| |
| BroadcastStream<Tuple2<Long, Long>> pBroadcastStream = |
| env.addSource(parallelSourceB) |
| .uid("BrCheckpointingSource2") |
| .broadcast(thirdBroadcastStateDesc); |
| |
| npStream.connect(npBroadcastStream) |
| .process(firstBroadcastFunction) |
| .uid("BrProcess1") |
| .addSink(new MigrationTestUtils.AccumulatorCountingSink<>()); |
| |
| pStream.connect(pBroadcastStream) |
| .process(secondBroadcastFunction) |
| .uid("BrProcess2") |
| .addSink(new MigrationTestUtils.AccumulatorCountingSink<>()); |
| |
| if (executionMode == StatefulJobSavepointMigrationITCase.ExecutionMode.PERFORM_SAVEPOINT) { |
| executeAndSavepoint( |
| env, |
| "src/test/resources/" |
| + getBroadcastSavepointPath(testMigrateVersion, testStateBackend), |
| new Tuple2<>( |
| MigrationTestUtils.AccumulatorCountingSink.NUM_ELEMENTS_ACCUMULATOR, |
| 2 * NUM_SOURCE_ELEMENTS)); |
| } else { |
| restoreAndExecute( |
| env, |
| getResourceFilename( |
| getBroadcastSavepointPath(testMigrateVersion, testStateBackend)), |
| new Tuple2<>( |
| MigrationTestUtils.CheckingNonParallelSourceWithListState |
| .SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR, |
| 2), // we have 2 sources |
| new Tuple2<>( |
| MigrationTestUtils.CheckingParallelSourceWithUnionListState |
| .SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR, |
| 2 * parallelism), // we have 2 sources |
| new Tuple2<>( |
| MigrationTestUtils.AccumulatorCountingSink.NUM_ELEMENTS_ACCUMULATOR, |
| NUM_SOURCE_ELEMENTS * 2)); |
| } |
| } |
| |
| private String getBroadcastSavepointPath( |
| MigrationVersion savepointVersion, String backendType) { |
| switch (backendType) { |
| case StateBackendLoader.ROCKSDB_STATE_BACKEND_NAME: |
| return "new-stateful-broadcast-udf-migration-itcase-flink" |
| + savepointVersion |
| + "-rocksdb-savepoint"; |
| case StateBackendLoader.MEMORY_STATE_BACKEND_NAME: |
| return "new-stateful-broadcast-udf-migration-itcase-flink" |
| + savepointVersion |
| + "-savepoint"; |
| default: |
| throw new UnsupportedOperationException(); |
| } |
| } |
| |
| /** |
| * A simple {@link KeyedBroadcastProcessFunction} that puts everything on the broadcast side in |
| * the state. |
| */ |
| private static class CheckpointingKeyedBroadcastFunction |
| extends KeyedBroadcastProcessFunction< |
| Long, Tuple2<Long, Long>, Tuple2<Long, Long>, Tuple2<Long, Long>> { |
| |
| private static final long serialVersionUID = 1333992081671604521L; |
| |
| private MapStateDescriptor<Long, Long> firstStateDesc; |
| |
| private MapStateDescriptor<String, Long> secondStateDesc; |
| |
| @Override |
| public void open(Configuration parameters) throws Exception { |
| super.open(parameters); |
| |
| firstStateDesc = |
| new MapStateDescriptor<>( |
| "broadcast-state-1", |
| BasicTypeInfo.LONG_TYPE_INFO, |
| BasicTypeInfo.LONG_TYPE_INFO); |
| |
| secondStateDesc = |
| new MapStateDescriptor<>( |
| "broadcast-state-2", |
| BasicTypeInfo.STRING_TYPE_INFO, |
| BasicTypeInfo.LONG_TYPE_INFO); |
| } |
| |
| @Override |
| public void processElement( |
| Tuple2<Long, Long> value, ReadOnlyContext ctx, Collector<Tuple2<Long, Long>> out) |
| throws Exception { |
| out.collect(value); |
| } |
| |
| @Override |
| public void processBroadcastElement( |
| Tuple2<Long, Long> value, Context ctx, Collector<Tuple2<Long, Long>> out) |
| throws Exception { |
| ctx.getBroadcastState(firstStateDesc).put(value.f0, value.f1); |
| ctx.getBroadcastState(secondStateDesc).put(Long.toString(value.f0), value.f1); |
| } |
| } |
| |
| /** |
| * A simple {@link KeyedBroadcastProcessFunction} that puts everything on the broadcast side in |
| * the state. |
| */ |
| private static class CheckpointingKeyedSingleBroadcastFunction |
| extends KeyedBroadcastProcessFunction< |
| Long, Tuple2<Long, Long>, Tuple2<Long, Long>, Tuple2<Long, Long>> { |
| |
| private static final long serialVersionUID = 1333992081671604521L; |
| |
| private MapStateDescriptor<Long, String> stateDesc; |
| |
| @Override |
| public void open(Configuration parameters) throws Exception { |
| super.open(parameters); |
| |
| stateDesc = |
| new MapStateDescriptor<>( |
| "broadcast-state-3", |
| BasicTypeInfo.LONG_TYPE_INFO, |
| BasicTypeInfo.STRING_TYPE_INFO); |
| } |
| |
| @Override |
| public void processElement( |
| Tuple2<Long, Long> value, ReadOnlyContext ctx, Collector<Tuple2<Long, Long>> out) |
| throws Exception { |
| out.collect(value); |
| } |
| |
| @Override |
| public void processBroadcastElement( |
| Tuple2<Long, Long> value, Context ctx, Collector<Tuple2<Long, Long>> out) |
| throws Exception { |
| ctx.getBroadcastState(stateDesc).put(value.f0, Long.toString(value.f1)); |
| } |
| } |
| |
| /** |
| * A simple {@link KeyedBroadcastProcessFunction} that verifies the contents of the broadcast |
| * state after recovery. |
| */ |
| private static class CheckingKeyedBroadcastFunction |
| extends KeyedBroadcastProcessFunction< |
| Long, Tuple2<Long, Long>, Tuple2<Long, Long>, Tuple2<Long, Long>> { |
| |
| private static final long serialVersionUID = 1333992081671604521L; |
| |
| private final Map<Long, Long> expectedFirstState; |
| |
| private final Map<String, Long> expectedSecondState; |
| |
| private MapStateDescriptor<Long, Long> firstStateDesc; |
| |
| private MapStateDescriptor<String, Long> secondStateDesc; |
| |
| CheckingKeyedBroadcastFunction(Map<Long, Long> firstState, Map<String, Long> secondState) { |
| this.expectedFirstState = firstState; |
| this.expectedSecondState = secondState; |
| } |
| |
| @Override |
| public void open(Configuration parameters) throws Exception { |
| super.open(parameters); |
| |
| firstStateDesc = |
| new MapStateDescriptor<>( |
| "broadcast-state-1", |
| BasicTypeInfo.LONG_TYPE_INFO, |
| BasicTypeInfo.LONG_TYPE_INFO); |
| |
| secondStateDesc = |
| new MapStateDescriptor<>( |
| "broadcast-state-2", |
| BasicTypeInfo.STRING_TYPE_INFO, |
| BasicTypeInfo.LONG_TYPE_INFO); |
| } |
| |
| @Override |
| public void processElement( |
| Tuple2<Long, Long> value, ReadOnlyContext ctx, Collector<Tuple2<Long, Long>> out) |
| throws Exception { |
| |
| final Map<Long, Long> actualFirstState = new HashMap<>(); |
| for (Map.Entry<Long, Long> entry : |
| ctx.getBroadcastState(firstStateDesc).immutableEntries()) { |
| actualFirstState.put(entry.getKey(), entry.getValue()); |
| } |
| Assert.assertEquals(expectedFirstState, actualFirstState); |
| |
| final Map<String, Long> actualSecondState = new HashMap<>(); |
| for (Map.Entry<String, Long> entry : |
| ctx.getBroadcastState(secondStateDesc).immutableEntries()) { |
| actualSecondState.put(entry.getKey(), entry.getValue()); |
| } |
| Assert.assertEquals(expectedSecondState, actualSecondState); |
| |
| out.collect(value); |
| } |
| |
| @Override |
| public void processBroadcastElement( |
| Tuple2<Long, Long> value, Context ctx, Collector<Tuple2<Long, Long>> out) |
| throws Exception { |
| // now we do nothing as we just want to verify the contents of the broadcast state. |
| } |
| } |
| |
| /** |
| * A simple {@link KeyedBroadcastProcessFunction} that verifies the contents of the broadcast |
| * state after recovery. |
| */ |
| private static class CheckingKeyedSingleBroadcastFunction |
| extends KeyedBroadcastProcessFunction< |
| Long, Tuple2<Long, Long>, Tuple2<Long, Long>, Tuple2<Long, Long>> { |
| |
| private static final long serialVersionUID = 1333992081671604521L; |
| |
| private final Map<Long, String> expectedState; |
| |
| private MapStateDescriptor<Long, String> stateDesc; |
| |
| CheckingKeyedSingleBroadcastFunction(Map<Long, String> state) { |
| this.expectedState = state; |
| } |
| |
| @Override |
| public void open(Configuration parameters) throws Exception { |
| super.open(parameters); |
| |
| stateDesc = |
| new MapStateDescriptor<>( |
| "broadcast-state-3", |
| BasicTypeInfo.LONG_TYPE_INFO, |
| BasicTypeInfo.STRING_TYPE_INFO); |
| } |
| |
| @Override |
| public void processElement( |
| Tuple2<Long, Long> value, ReadOnlyContext ctx, Collector<Tuple2<Long, Long>> out) |
| throws Exception { |
| final Map<Long, String> actualState = new HashMap<>(); |
| for (Map.Entry<Long, String> entry : |
| ctx.getBroadcastState(stateDesc).immutableEntries()) { |
| actualState.put(entry.getKey(), entry.getValue()); |
| } |
| Assert.assertEquals(expectedState, actualState); |
| |
| out.collect(value); |
| } |
| |
| @Override |
| public void processBroadcastElement( |
| Tuple2<Long, Long> value, Context ctx, Collector<Tuple2<Long, Long>> out) |
| throws Exception { |
| // now we do nothing as we just want to verify the contents of the broadcast state. |
| } |
| } |
| } |
