[FLINK-3359] Make RocksDB File Copies Asynchronous
Project: http://git-wip-us.apache.org/repos/asf/flink/repo Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/31310cd6 Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/31310cd6 Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/31310cd6 Branch: refs/heads/master Commit: 31310cd657381d4035cda9659d737ed440f90d26 Parents: c47cb7a Author: Aljoscha Krettek <[email protected]> Authored: Thu Feb 11 17:29:37 2016 +0100 Committer: Aljoscha Krettek <[email protected]> Committed: Fri Feb 12 22:59:14 2016 +0100 ---------------------------------------------------------------------- .../flink-statebackend-rocksdb/pom.xml | 15 +- .../streaming/state/AbstractRocksDBState.java | 161 +++++++++----- .../streaming/state/RocksDBFoldingState.java | 2 +- .../streaming/state/RocksDBListState.java | 2 +- .../streaming/state/RocksDBReducingState.java | 2 +- .../streaming/state/RocksDBStateBackend.java | 6 +- .../streaming/state/RocksDBValueState.java | 2 +- .../state/RocksDBAsyncKVSnapshotTest.java | 208 +++++++++++++++++++ .../state/RocksDBStateBackendTest.java | 4 +- .../state/AsynchronousKvStateSnapshot.java | 62 ++++++ .../runtime/state/StateBackendTestBase.java | 180 ++++++++++------ .../streaming/runtime/tasks/StreamTask.java | 26 ++- .../runtime/tasks/StreamTaskStateList.java | 58 +++--- .../tasks/OneInputStreamTaskTestHarness.java | 7 + .../tasks/StreamTaskAsyncCheckpointTest.java | 3 +- .../EventTimeWindowCheckpointingITCase.java | 6 + 16 files changed, 584 insertions(+), 160 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-contrib/flink-statebackend-rocksdb/pom.xml ---------------------------------------------------------------------- diff --git a/flink-contrib/flink-statebackend-rocksdb/pom.xml b/flink-contrib/flink-statebackend-rocksdb/pom.xml index 999c496..9c1601e 100644 --- a/flink-contrib/flink-statebackend-rocksdb/pom.xml +++ b/flink-contrib/flink-statebackend-rocksdb/pom.xml @@ -57,7 +57,6 @@ under the License. <artifactId>rocksdbjni</artifactId> <version>4.1.0</version> </dependency> - <dependency> <groupId>org.apache.flink</groupId> <artifactId>flink-runtime_2.10</artifactId> @@ -65,7 +64,17 @@ under the License. <type>test-jar</type> <scope>test</scope> </dependency> - + <dependency> + <groupId>org.apache.flink</groupId> + <artifactId>flink-streaming-java_2.10</artifactId> + <version>${project.version}</version> + <type>test-jar</type> + <scope>test</scope> + </dependency> + <dependency> + <groupId>junit</groupId> + <artifactId>junit</artifactId> + <version>4.11</version> + </dependency> </dependencies> - </project> http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java ---------------------------------------------------------------------- diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java index 7d3172a..6e4adf5 100644 --- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java +++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java @@ -24,11 +24,12 @@ import org.apache.flink.api.common.state.StateDescriptor; import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.core.memory.DataOutputView; import org.apache.flink.core.memory.DataOutputViewStreamWrapper; +import org.apache.flink.runtime.state.AsynchronousKvStateSnapshot; import org.apache.flink.runtime.state.KvState; import org.apache.flink.runtime.state.KvStateSnapshot; + import org.apache.flink.util.HDFSCopyFromLocal; import org.apache.flink.util.HDFSCopyToLocal; - import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; @@ -48,6 +49,7 @@ import java.io.ByteArrayOutputStream; import java.io.File; import java.io.IOException; import java.net.URI; +import java.util.UUID; import static java.util.Objects.requireNonNull; @@ -58,6 +60,11 @@ import static java.util.Objects.requireNonNull; * checkpointing/restoring the database and for disposal in the {@link #dispose()} method. The * concrete subclasses just use the RocksDB handle to store/retrieve state. * + * <p>State is checkpointed asynchronously. The synchronous part is drawing the actual backup + * from RocksDB, this is done in {@link #snapshot(long, long)}. This will return a + * {@link AsyncRocksDBSnapshot} that will perform the copying of the backup to the remote + * file system. + * * @param <K> The type of the key. * @param <N> The type of the namespace. * @param <S> The type of {@link State}. @@ -81,51 +88,59 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta protected N currentNamespace; /** Store it so that we can clean up in dispose() */ - protected final File dbPath; + protected final File basePath; + /** FileSystem path where checkpoints are stored */ protected final String checkpointPath; + /** Directory in "basePath" where the actual RocksDB data base instance stores its files */ + protected final File rocksDbPath; + /** Our RocksDB instance */ protected final RocksDB db; + /** * Creates a new RocksDB backed state. * * @param keySerializer The serializer for the keys. * @param namespaceSerializer The serializer for the namespace. - * @param dbPath The path on the local system where RocksDB data should be stored. + * @param basePath The path on the local system where RocksDB data should be stored. */ protected AbstractRocksDBState( TypeSerializer<K> keySerializer, TypeSerializer<N> namespaceSerializer, - File dbPath, + File basePath, String checkpointPath, Options options) { - + + rocksDbPath = new File(basePath, "db" + UUID.randomUUID().toString()); + this.keySerializer = requireNonNull(keySerializer); this.namespaceSerializer = namespaceSerializer; - this.dbPath = dbPath; + this.basePath = basePath; this.checkpointPath = checkpointPath; RocksDB.loadLibrary(); - if (!dbPath.exists()) { - if (!dbPath.mkdirs()) { + if (!basePath.exists()) { + if (!basePath.mkdirs()) { throw new RuntimeException("Could not create RocksDB data directory."); } } // clean it, this will remove the last part of the path but RocksDB will recreate it try { - File db = new File(dbPath, "db"); - LOG.warn("Deleting already existing db directory {}.", db); - FileUtils.deleteDirectory(db); + if (rocksDbPath.exists()) { + LOG.warn("Deleting already existing db directory {}.", rocksDbPath); + FileUtils.deleteDirectory(rocksDbPath); + } } catch (IOException e) { throw new RuntimeException("Error cleaning RocksDB data directory.", e); } try { - db = RocksDB.open(options, new File(dbPath, "db").getAbsolutePath()); + db = RocksDB.open(options, rocksDbPath.getAbsolutePath()); } catch (RocksDBException e) { throw new RuntimeException("Error while opening RocksDB instance.", e); } @@ -137,39 +152,56 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta * * @param keySerializer The serializer for the keys. * @param namespaceSerializer The serializer for the namespace. - * @param dbPath The path on the local system where RocksDB data should be stored. + * @param basePath The path on the local system where RocksDB data should be stored. * @param restorePath The path to a backup directory from which to restore RocksDb database. */ protected AbstractRocksDBState( TypeSerializer<K> keySerializer, TypeSerializer<N> namespaceSerializer, - File dbPath, + File basePath, String checkpointPath, String restorePath, Options options) { + rocksDbPath = new File(basePath, "db" + UUID.randomUUID().toString()); + RocksDB.loadLibrary(); + // clean it, this will remove the last part of the path but RocksDB will recreate it + try { + if (rocksDbPath.exists()) { + LOG.warn("Deleting already existing db directory {}.", rocksDbPath); + FileUtils.deleteDirectory(rocksDbPath); + } + } catch (IOException e) { + throw new RuntimeException("Error cleaning RocksDB data directory.", e); + } + try (BackupEngine backupEngine = BackupEngine.open(Env.getDefault(), new BackupableDBOptions(restorePath + "/"))) { - backupEngine.restoreDbFromLatestBackup(new File(dbPath, "db").getAbsolutePath(), new File(dbPath, "db").getAbsolutePath(), new RestoreOptions(true)); - FileUtils.deleteDirectory(new File(restorePath)); - } catch (RocksDBException|IOException|IllegalArgumentException e) { + backupEngine.restoreDbFromLatestBackup(rocksDbPath.getAbsolutePath(), rocksDbPath.getAbsolutePath(), new RestoreOptions(true)); + } catch (RocksDBException|IllegalArgumentException e) { throw new RuntimeException("Error while restoring RocksDB state from " + restorePath, e); + } finally { + try { + FileUtils.deleteDirectory(new File(restorePath)); + } catch (IOException e) { + LOG.error("Error cleaning up local restore directory " + restorePath, e); + } } this.keySerializer = requireNonNull(keySerializer); this.namespaceSerializer = namespaceSerializer; - this.dbPath = dbPath; + this.basePath = basePath; this.checkpointPath = checkpointPath; - if (!dbPath.exists()) { - if (!dbPath.mkdirs()) { + if (!basePath.exists()) { + if (!basePath.mkdirs()) { throw new RuntimeException("Could not create RocksDB data directory."); } } try { - db = RocksDB.open(options, new File(dbPath, "db").getAbsolutePath()); + db = RocksDB.open(options, rocksDbPath.getAbsolutePath()); } catch (RocksDBException e) { throw new RuntimeException("Error while opening RocksDB instance.", e); } @@ -209,49 +241,41 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta protected abstract AbstractRocksDBSnapshot<K, N, S, SD> createRocksDBSnapshot(URI backupUri, long checkpointId); @Override - public final AbstractRocksDBSnapshot<K, N, S, SD> snapshot(long checkpointId, long timestamp) throws Exception { - boolean success = false; + public final KvStateSnapshot<K, N, S, SD, RocksDBStateBackend> snapshot(final long checkpointId, long timestamp) throws Exception { - final File localBackupPath = new File(dbPath, "backup-" + checkpointId); + final File localBackupPath = new File(basePath, "local-chk-" + checkpointId); final URI backupUri = new URI(checkpointPath + "/chk-" + checkpointId); - try { - if (!localBackupPath.exists()) { - if (!localBackupPath.mkdirs()) { - throw new RuntimeException("Could not create local backup path " + localBackupPath); - } - } - try (BackupEngine backupEngine = BackupEngine.open(Env.getDefault(), new BackupableDBOptions(localBackupPath.getAbsolutePath()))) { - backupEngine.createNewBackup(db); + if (!localBackupPath.exists()) { + if (!localBackupPath.mkdirs()) { + throw new RuntimeException("Could not create local backup path " + localBackupPath); } + } - HDFSCopyFromLocal.copyFromLocal(localBackupPath, backupUri); - AbstractRocksDBSnapshot<K, N, S, SD> result = createRocksDBSnapshot(backupUri, checkpointId); - success = true; - return result; - } finally { - FileUtils.deleteDirectory(localBackupPath); - if (!success) { - FileSystem fs = FileSystem.get(backupUri, new Configuration()); - fs.delete(new Path(backupUri), true); - } + try (BackupEngine backupEngine = BackupEngine.open(Env.getDefault(), new BackupableDBOptions(localBackupPath.getAbsolutePath()))) { + backupEngine.createNewBackup(db); } + + return new AsyncRocksDBSnapshot<>( + localBackupPath, + backupUri, + checkpointId, + this); } @Override final public void dispose() { db.dispose(); try { - FileUtils.deleteDirectory(dbPath); + FileUtils.deleteDirectory(basePath); } catch (IOException e) { throw new RuntimeException("Error disposing RocksDB data directory.", e); } } - public static abstract class AbstractRocksDBSnapshot<K, N, S extends State, SD extends StateDescriptor<S, ?>> - implements KvStateSnapshot<K, N, S, SD, RocksDBStateBackend> - { + protected static abstract class AbstractRocksDBSnapshot<K, N, S extends State, SD extends StateDescriptor<S, ?>> + implements KvStateSnapshot<K, N, S, SD, RocksDBStateBackend> { private static final long serialVersionUID = 1L; private static final Logger LOG = LoggerFactory.getLogger(AbstractRocksDBSnapshot.class); @@ -287,6 +311,9 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta /** Hash of the StateDescriptor, for sanity checks */ protected final SD stateDesc; + /** + * Creates a new snapshot from the given state parameters. + */ public AbstractRocksDBSnapshot(File dbPath, String checkpointPath, URI backupUri, @@ -305,6 +332,9 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta this.namespaceSerializer = namespaceSerializer; } + /** + * Subclasses must implement this for creating a concrete RocksDB state. + */ protected abstract KvState<K, N, S, SD, RocksDBStateBackend> createRocksDBState( TypeSerializer<K> keySerializer, TypeSerializer<N> namespaceSerializer, @@ -336,8 +366,6 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta } } - FileSystem fs = FileSystem.get(backupUri, new Configuration()); - final File localBackupPath = new File(dbPath, "chk-" + checkpointId); if (localBackupPath.exists()) { @@ -365,5 +393,42 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta return 0; } } + + /** + * Upon snapshotting the RocksDB backup is created synchronously. The asynchronous part is + * copying the backup to a (possibly) remote filesystem. This is done in {@link #materialize()} + * of this class. + */ + private static class AsyncRocksDBSnapshot<K, N, S extends State, SD extends StateDescriptor<S, ?>> extends AsynchronousKvStateSnapshot<K, N, S, SD, RocksDBStateBackend> { + private static final long serialVersionUID = 1L; + private final File localBackupPath; + private final URI backupUri; + private final long checkpointId; + private transient AbstractRocksDBState<K, N, S, SD> state; + + public AsyncRocksDBSnapshot(File localBackupPath, + URI backupUri, + long checkpointId, + AbstractRocksDBState<K, N, S, SD> state) { + this.localBackupPath = localBackupPath; + this.backupUri = backupUri; + this.checkpointId = checkpointId; + this.state = state; + } + + @Override + public KvStateSnapshot<K, N, S, SD, RocksDBStateBackend> materialize() throws Exception { + try { + HDFSCopyFromLocal.copyFromLocal(localBackupPath, backupUri); + return state.createRocksDBSnapshot(backupUri, checkpointId); + } catch (Exception e) { + FileSystem fs = FileSystem.get(backupUri, new Configuration()); + fs.delete(new Path(backupUri), true); + throw e; + } finally { + FileUtils.deleteQuietly(localBackupPath); + } + } + } } http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java ---------------------------------------------------------------------- diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java index d7b75bd..1a3e69e 100644 --- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java +++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java @@ -156,7 +156,7 @@ public class RocksDBFoldingState<K, N, T, ACC> protected AbstractRocksDBSnapshot<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>> createRocksDBSnapshot( URI backupUri, long checkpointId) { - return new Snapshot<>(dbPath, checkpointPath, backupUri, checkpointId, keySerializer, namespaceSerializer, stateDesc); + return new Snapshot<>(basePath, checkpointPath, backupUri, checkpointId, keySerializer, namespaceSerializer, stateDesc); } private static class Snapshot<K, N, T, ACC> extends AbstractRocksDBSnapshot<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>> { http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java ---------------------------------------------------------------------- diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java index 6c55566..aa029ac 100644 --- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java +++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java @@ -154,7 +154,7 @@ public class RocksDBListState<K, N, V> URI backupUri, long checkpointId) { - return new Snapshot<>(dbPath, checkpointPath, backupUri, checkpointId, keySerializer, namespaceSerializer, stateDesc); + return new Snapshot<>(basePath, checkpointPath, backupUri, checkpointId, keySerializer, namespaceSerializer, stateDesc); } private static class Snapshot<K, N, V> extends http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java ---------------------------------------------------------------------- diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java index b7ba3c7..3fdfafe 100644 --- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java +++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java @@ -152,7 +152,7 @@ public class RocksDBReducingState<K, N, V> URI backupUri, long checkpointId) { - return new Snapshot<>(dbPath, checkpointPath, backupUri, checkpointId, keySerializer, namespaceSerializer, stateDesc); + return new Snapshot<>(basePath, checkpointPath, backupUri, checkpointId, keySerializer, namespaceSerializer, stateDesc); } private static class Snapshot<K, N, V> extends http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java ---------------------------------------------------------------------- diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java index 04bb17c..72a4c58 100644 --- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java +++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java @@ -68,19 +68,17 @@ public class RocksDBStateBackend extends AbstractStateBackend { private static final Logger LOG = LoggerFactory.getLogger(RocksDBStateBackend.class); - /** The checkpoint directory that copy the RocksDB backups to. */ + /** The checkpoint directory that we copy the RocksDB backups to. */ private final Path checkpointDirectory; /** The state backend that stores the non-partitioned state */ private final AbstractStateBackend nonPartitionedStateBackend; - - + /** Operator identifier that is used to uniqueify the RocksDB storage path. */ private String operatorIdentifier; /** JobID for uniquifying backup paths. */ private JobID jobId; - // DB storage directories http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java ---------------------------------------------------------------------- diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java index 7a19153..5f6eccd 100644 --- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java +++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java @@ -140,7 +140,7 @@ public class RocksDBValueState<K, N, V> URI backupUri, long checkpointId) { - return new Snapshot<>(dbPath, checkpointPath, backupUri, checkpointId, keySerializer, namespaceSerializer, stateDesc); + return new Snapshot<>(basePath, checkpointPath, backupUri, checkpointId, keySerializer, namespaceSerializer, stateDesc); } private static class Snapshot<K, N, V> http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java ---------------------------------------------------------------------- diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java new file mode 100644 index 0000000..260819e --- /dev/null +++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java @@ -0,0 +1,208 @@ +/* + * 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.contrib.streaming.state; + +import org.apache.flink.api.common.functions.MapFunction; +import org.apache.flink.api.common.state.ReducingStateDescriptor; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.BasicTypeInfo; +import org.apache.flink.api.common.typeutils.base.StringSerializer; +import org.apache.flink.api.common.typeutils.base.VoidSerializer; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.configuration.ConfigConstants; +import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter; +import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider; +import org.apache.flink.runtime.state.AsynchronousStateHandle; +import org.apache.flink.runtime.state.StateHandle; +import org.apache.flink.runtime.state.memory.MemoryStateBackend; +import org.apache.flink.runtime.taskmanager.OneShotLatch; +import org.apache.flink.streaming.api.graph.StreamConfig; +import org.apache.flink.streaming.api.operators.AbstractStreamOperator; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; +import org.apache.flink.streaming.api.watermark.Watermark; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask; +import org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness; +import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment; +import org.apache.flink.streaming.runtime.tasks.StreamTask; +import org.apache.flink.streaming.runtime.tasks.StreamTaskState; +import org.apache.flink.streaming.runtime.tasks.StreamTaskStateList; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.powermock.core.classloader.annotations.PrepareForTest; +import org.powermock.modules.junit4.PowerMockRunner; + +import java.io.File; +import java.lang.reflect.Field; +import java.util.UUID; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +/** + * Tests for asynchronous RocksDB Key/Value state checkpoints. + */ +@RunWith(PowerMockRunner.class) +@PrepareForTest(ResultPartitionWriter.class) +@SuppressWarnings("serial") +public class RocksDBAsyncKVSnapshotTest { + + /** + * This ensures that asynchronous state handles are actually materialized asynchonously. + * + * <p>We use latches to block at various stages and see if the code still continues through + * the parts that are not asynchronous. If the checkpoint is not done asynchronously the + * test will simply lock forever. + */ + @Test + public void testAsyncCheckpoints() throws Exception { + final OneShotLatch delayCheckpointLatch = new OneShotLatch(); + final OneShotLatch ensureCheckpointLatch = new OneShotLatch(); + + final OneInputStreamTask<String, String> task = new OneInputStreamTask<>(); + + final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO); + + testHarness.configureForKeyedStream(new KeySelector<String, String>() { + @Override + public String getKey(String value) throws Exception { + return value; + } + }, BasicTypeInfo.STRING_TYPE_INFO); + + StreamConfig streamConfig = testHarness.getStreamConfig(); + + File dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state"); + File chkDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "snapshots"); + + RocksDBStateBackend backend = new RocksDBStateBackend(chkDir.getAbsoluteFile().toURI(), new MemoryStateBackend()); + backend.setDbStoragePath(dbDir.getAbsolutePath()); + + streamConfig.setStateBackend(backend); + + streamConfig.setStreamOperator(new AsyncCheckpointOperator()); + + StreamMockEnvironment mockEnv = new StreamMockEnvironment( + testHarness.jobConfig, + testHarness.taskConfig, + testHarness.memorySize, + new MockInputSplitProvider(), + testHarness.bufferSize) { + + @Override + public void acknowledgeCheckpoint(long checkpointId) { + super.acknowledgeCheckpoint(checkpointId); + } + + @Override + public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) { + super.acknowledgeCheckpoint(checkpointId, state); + + // block on the latch, to verify that triggerCheckpoint returns below, + // even though the async checkpoint would not finish + try { + delayCheckpointLatch.await(); + } catch (InterruptedException e) { + e.printStackTrace(); + } + + assertTrue(state instanceof StreamTaskStateList); + StreamTaskStateList stateList = (StreamTaskStateList) state; + + // should be only one k/v state + StreamTaskState taskState = stateList.getState(this.getUserClassLoader())[0]; + assertEquals(1, taskState.getKvStates().size()); + assertTrue(taskState.getKvStates().get("count") instanceof AbstractRocksDBState.AbstractRocksDBSnapshot); + + // we now know that the checkpoint went through + ensureCheckpointLatch.trigger(); + } + }; + + testHarness.invoke(mockEnv); + + // wait for the task to be running + for (Field field: StreamTask.class.getDeclaredFields()) { + if (field.getName().equals("isRunning")) { + field.setAccessible(true); + while (!field.getBoolean(task)) { + Thread.sleep(10); + } + + } + } + + testHarness.processElement(new StreamRecord<>("Wohoo", 0)); + + task.triggerCheckpoint(42, 17); + + // now we allow the checkpoint + delayCheckpointLatch.trigger(); + + // wait for the checkpoint to go through + ensureCheckpointLatch.await(); + + testHarness.endInput(); + testHarness.waitForTaskCompletion(); + } + + + // ------------------------------------------------------------------------ + + public static class AsyncCheckpointOperator + extends AbstractStreamOperator<String> + implements OneInputStreamOperator<String, String> { + + @Override + public void open() throws Exception { + super.open(); + + // also get the state in open, this way we are sure that it was created before + // we trigger the test checkpoint + ValueState<String> state = getPartitionedState(null, + VoidSerializer.INSTANCE, + new ValueStateDescriptor<>("count", + StringSerializer.INSTANCE, "hello")); + + } + + @Override + public void processElement(StreamRecord<String> element) throws Exception { + // we also don't care + + ValueState<String> state = getPartitionedState(null, + VoidSerializer.INSTANCE, + new ValueStateDescriptor<>("count", + StringSerializer.INSTANCE, "hello")); + + state.update(element.getValue()); + } + + @Override + public void processWatermark(Watermark mark) throws Exception { + // not interested + } + } + + public static class DummyMapFunction<T> implements MapFunction<T, T> { + @Override + public T map(T value) { return value; } + } +} http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java ---------------------------------------------------------------------- diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java index fe933e0..3c77eb9 100644 --- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java +++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java @@ -37,8 +37,8 @@ public class RocksDBStateBackendTest extends StateBackendTestBase<RocksDBStateBa @Override protected RocksDBStateBackend getStateBackend() throws IOException { - dbDir = new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()); - chkDir = new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()); + dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state"); + chkDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "snapshots"); RocksDBStateBackend backend = new RocksDBStateBackend(chkDir.getAbsoluteFile().toURI(), new MemoryStateBackend()); backend.setDbStoragePath(dbDir.getAbsolutePath()); http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousKvStateSnapshot.java ---------------------------------------------------------------------- diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousKvStateSnapshot.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousKvStateSnapshot.java new file mode 100644 index 0000000..30a9c5a --- /dev/null +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousKvStateSnapshot.java @@ -0,0 +1,62 @@ +/* + * 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.runtime.state; + +import org.apache.flink.api.common.state.State; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.typeutils.TypeSerializer; + +/** + * {@link KvStateSnapshot} that asynchronously materializes the state that it represents. Instead + * of representing a materialized handle to state this would normally hold the (immutable) state + * internally and materializes it when {@link #materialize()} is called. + * + * @param <K> The type of the key + * @param <N> The type of the namespace + * @param <S> The type of the {@link State} + * @param <SD> The type of the {@link StateDescriptor} + * @param <Backend> The type of the backend that can restore the state from this snapshot. + */ +public abstract class AsynchronousKvStateSnapshot<K, N, S extends State, SD extends StateDescriptor<S, ?>, Backend extends AbstractStateBackend> implements KvStateSnapshot<K, N, S, SD, Backend> { + private static final long serialVersionUID = 1L; + + /** + * Materializes the state held by this {@code AsynchronousKvStateSnapshot}. + */ + public abstract KvStateSnapshot<K, N, S, SD, Backend> materialize() throws Exception; + + @Override + public final KvState<K, N, S, SD, Backend> restoreState( + Backend stateBackend, + TypeSerializer<K> keySerializer, + ClassLoader classLoader, + long recoveryTimestamp) throws Exception { + throw new RuntimeException("This should never be called and probably points to a bug."); + } + + @Override + public void discardState() throws Exception { + throw new RuntimeException("This should never be called and probably points to a bug."); + } + + @Override + public long getStateSize() throws Exception { + throw new RuntimeException("This should never be called and probably points to a bug."); + } +} http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java ---------------------------------------------------------------------- diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java index 27dee6a..6c3930b 100644 --- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java @@ -72,88 +72,96 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { @Test public void testValueState() throws Exception { - backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE); + backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE); - ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class, null); - kvId.initializeSerializerUnlessSet(new ExecutionConfig()); - - ValueState<String> state = backend.getPartitionedState(null, VoidSerializer.INSTANCE, kvId); + ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class, null); + kvId.initializeSerializerUnlessSet(new ExecutionConfig()); - @SuppressWarnings("unchecked") - KvState<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> kv = - (KvState<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B>) state; + ValueState<String> state = backend.getPartitionedState(null, VoidSerializer.INSTANCE, kvId); - // some modifications to the state - kv.setCurrentKey(1); - assertNull(state.value()); - state.update("1"); - kv.setCurrentKey(2); - assertNull(state.value()); - state.update("2"); - kv.setCurrentKey(1); - assertEquals("1", state.value()); + @SuppressWarnings("unchecked") + KvState<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> kv = + (KvState<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B>) state; - // draw a snapshot - KvStateSnapshot<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> snapshot1 = - kv.snapshot(682375462378L, 2); + // some modifications to the state + kv.setCurrentKey(1); + assertNull(state.value()); + state.update("1"); + kv.setCurrentKey(2); + assertNull(state.value()); + state.update("2"); + kv.setCurrentKey(1); + assertEquals("1", state.value()); - // make some more modifications - kv.setCurrentKey(1); - state.update("u1"); - kv.setCurrentKey(2); - state.update("u2"); - kv.setCurrentKey(3); - state.update("u3"); + // draw a snapshot + KvStateSnapshot<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> snapshot1 = + kv.snapshot(682375462378L, 2); - // draw another snapshot - KvStateSnapshot<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> snapshot2 = - kv.snapshot(682375462379L, 4); + if (snapshot1 instanceof AsynchronousKvStateSnapshot) { + snapshot1 = ((AsynchronousKvStateSnapshot<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B>) snapshot1).materialize(); + } - // validate the original state - kv.setCurrentKey(1); - assertEquals("u1", state.value()); - kv.setCurrentKey(2); - assertEquals("u2", state.value()); - kv.setCurrentKey(3); - assertEquals("u3", state.value()); + // make some more modifications + kv.setCurrentKey(1); + state.update("u1"); + kv.setCurrentKey(2); + state.update("u2"); + kv.setCurrentKey(3); + state.update("u3"); - kv.dispose(); + // draw another snapshot + KvStateSnapshot<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> snapshot2 = + kv.snapshot(682375462379L, 4); - // restore the first snapshot and validate it - KvState<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> restored1 = snapshot1.restoreState( - backend, - IntSerializer.INSTANCE, - this.getClass().getClassLoader(), 10); + if (snapshot2 instanceof AsynchronousKvStateSnapshot) { + snapshot2 = ((AsynchronousKvStateSnapshot<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B>) snapshot2).materialize(); + } - snapshot1.discardState(); + // validate the original state + kv.setCurrentKey(1); + assertEquals("u1", state.value()); + kv.setCurrentKey(2); + assertEquals("u2", state.value()); + kv.setCurrentKey(3); + assertEquals("u3", state.value()); - @SuppressWarnings("unchecked") - ValueState<String> restored1State = (ValueState<String>) restored1; + kv.dispose(); - restored1.setCurrentKey(1); - assertEquals("1", restored1State.value()); - restored1.setCurrentKey(2); - assertEquals("2", restored1State.value()); + // restore the first snapshot and validate it + KvState<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> restored1 = snapshot1.restoreState( + backend, + IntSerializer.INSTANCE, + this.getClass().getClassLoader(), 10); - restored1.dispose(); + snapshot1.discardState(); - // restore the second snapshot and validate it - KvState<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> restored2 = snapshot2.restoreState( - backend, - IntSerializer.INSTANCE, - this.getClass().getClassLoader(), 10); + @SuppressWarnings("unchecked") + ValueState<String> restored1State = (ValueState<String>) restored1; - snapshot2.discardState(); + restored1.setCurrentKey(1); + assertEquals("1", restored1State.value()); + restored1.setCurrentKey(2); + assertEquals("2", restored1State.value()); - @SuppressWarnings("unchecked") - ValueState<String> restored2State = (ValueState<String>) restored2; + restored1.dispose(); - restored2.setCurrentKey(1); - assertEquals("u1", restored2State.value()); - restored2.setCurrentKey(2); - assertEquals("u2", restored2State.value()); - restored2.setCurrentKey(3); - assertEquals("u3", restored2State.value()); + // restore the second snapshot and validate it + KvState<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> restored2 = snapshot2.restoreState( + backend, + IntSerializer.INSTANCE, + this.getClass().getClassLoader(), 10); + + snapshot2.discardState(); + + @SuppressWarnings("unchecked") + ValueState<String> restored2State = (ValueState<String>) restored2; + + restored2.setCurrentKey(1); + assertEquals("u1", restored2State.value()); + restored2.setCurrentKey(2); + assertEquals("u2", restored2State.value()); + restored2.setCurrentKey(3); + assertEquals("u3", restored2State.value()); } /** @@ -211,6 +219,11 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, ValueState<Long>, ValueStateDescriptor<Long>, B> snapshot1 = kv.snapshot(682375462378L, 2); + if (snapshot1 instanceof AsynchronousKvStateSnapshot) { + snapshot1 = ((AsynchronousKvStateSnapshot<Integer, Void, ValueState<Long>, ValueStateDescriptor<Long>, B>) snapshot1).materialize(); + } + + kv.dispose(); // restore the snapshot @@ -248,6 +261,11 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, ListState<String>, ListStateDescriptor<String>, B> snapshot1 = kv.snapshot(682375462378L, 2); + if (snapshot1 instanceof AsynchronousKvStateSnapshot) { + snapshot1 = ((AsynchronousKvStateSnapshot<Integer, Void, ListState<String>, ListStateDescriptor<String>, B>) snapshot1).materialize(); + } + + // make some more modifications kv.setCurrentKey(1); state.add("u1"); @@ -260,6 +278,10 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, ListState<String>, ListStateDescriptor<String>, B> snapshot2 = kv.snapshot(682375462379L, 4); + if (snapshot2 instanceof AsynchronousKvStateSnapshot) { + snapshot2 = ((AsynchronousKvStateSnapshot<Integer, Void, ListState<String>, ListStateDescriptor<String>, B>) snapshot2).materialize(); + } + // validate the original state kv.setCurrentKey(1); assertEquals("1,u1", joiner.join(state.get())); @@ -349,6 +371,10 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, B> snapshot1 = kv.snapshot(682375462378L, 2); + if (snapshot1 instanceof AsynchronousKvStateSnapshot) { + snapshot1 = ((AsynchronousKvStateSnapshot<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, B>) snapshot1).materialize(); + } + // make some more modifications kv.setCurrentKey(1); state.add("u1"); @@ -361,6 +387,10 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, B> snapshot2 = kv.snapshot(682375462379L, 4); + if (snapshot2 instanceof AsynchronousKvStateSnapshot) { + snapshot2 = ((AsynchronousKvStateSnapshot<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, B>) snapshot2).materialize(); + } + // validate the original state kv.setCurrentKey(1); assertEquals("1,u1", state.get()); @@ -451,6 +481,10 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, FoldingState<Integer, String>, FoldingStateDescriptor<Integer, String>, B> snapshot1 = kv.snapshot(682375462378L, 2); + if (snapshot1 instanceof AsynchronousKvStateSnapshot) { + snapshot1 = ((AsynchronousKvStateSnapshot<Integer, Void, FoldingState<Integer, String>, FoldingStateDescriptor<Integer, String>, B>) snapshot1).materialize(); + } + // make some more modifications kv.setCurrentKey(1); state.clear(); @@ -464,6 +498,10 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, FoldingState<Integer, String>, FoldingStateDescriptor<Integer, String>, B> snapshot2 = kv.snapshot(682375462379L, 4); + if (snapshot2 instanceof AsynchronousKvStateSnapshot) { + snapshot2 = ((AsynchronousKvStateSnapshot<Integer, Void, FoldingState<Integer, String>, FoldingStateDescriptor<Integer, String>, B>) snapshot2).materialize(); + } + // validate the original state kv.setCurrentKey(1); assertEquals("Fold-Initial:,101", state.get()); @@ -540,6 +578,10 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B> snapshot = kv.snapshot(682375462378L, System.currentTimeMillis()); + if (snapshot instanceof AsynchronousKvStateSnapshot) { + snapshot = ((AsynchronousKvStateSnapshot<Integer, Void, ValueState<String>, ValueStateDescriptor<String>, B>) snapshot).materialize(); + } + @SuppressWarnings("unchecked") TypeSerializer<Integer> fakeIntSerializer = (TypeSerializer<Integer>) (TypeSerializer<?>) FloatSerializer.INSTANCE; @@ -579,6 +621,10 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, ListState<String>, ListStateDescriptor<String>, B> snapshot = kv.snapshot(682375462378L, System.currentTimeMillis()); + if (snapshot instanceof AsynchronousKvStateSnapshot) { + snapshot = ((AsynchronousKvStateSnapshot<Integer, Void, ListState<String>, ListStateDescriptor<String>, B>) snapshot).materialize(); + } + kv.dispose(); @SuppressWarnings("unchecked") @@ -627,6 +673,10 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> { KvStateSnapshot<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, B> snapshot = kv.snapshot(682375462378L, System.currentTimeMillis()); + if (snapshot instanceof AsynchronousKvStateSnapshot) { + snapshot = ((AsynchronousKvStateSnapshot<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, B>) snapshot).materialize(); + } + kv.dispose(); @SuppressWarnings("unchecked") http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java ---------------------------------------------------------------------- diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java index 4e75b1c..9b5c64a 100644 --- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java +++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java @@ -19,6 +19,7 @@ package org.apache.flink.streaming.runtime.tasks; import java.io.Serializable; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Set; @@ -36,7 +37,9 @@ import org.apache.flink.runtime.io.network.api.CheckpointBarrier; import org.apache.flink.runtime.jobgraph.tasks.AbstractInvokable; import org.apache.flink.runtime.jobgraph.tasks.StatefulTask; import org.apache.flink.runtime.state.AbstractStateBackend; +import org.apache.flink.runtime.state.AsynchronousKvStateSnapshot; import org.apache.flink.runtime.state.AsynchronousStateHandle; +import org.apache.flink.runtime.state.KvStateSnapshot; import org.apache.flink.runtime.taskmanager.DispatcherThreadFactory; import org.apache.flink.runtime.util.event.EventListener; import org.apache.flink.streaming.api.graph.StreamConfig; @@ -198,7 +201,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>> // -------- Invoke -------- LOG.debug("Invoking {}", getName()); - + // first order of business is to give operators back their state restoreState(); @@ -474,6 +477,14 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>> if (state.getFunctionState() instanceof AsynchronousStateHandle) { hasAsyncStates = true; } + if (state.getKvStates() != null) { + for (KvStateSnapshot<?, ?, ?, ?, ?> kvSnapshot: state.getKvStates().values()) { + if (kvSnapshot instanceof AsynchronousKvStateSnapshot) { + hasAsyncStates = true; + } + } + } + states[i] = state.isEmpty() ? null : state; } } @@ -509,10 +520,22 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>> AsynchronousStateHandle<?> asyncState = (AsynchronousStateHandle<?>) state.getOperatorState(); state.setOperatorState(asyncState.materialize()); } + if (state.getKvStates() != null) { + Set<String> keys = state.getKvStates().keySet(); + HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates = state.getKvStates(); + for (String key: keys) { + if (kvStates.get(key) instanceof AsynchronousKvStateSnapshot) { + AsynchronousKvStateSnapshot<?, ?, ?, ?, ?> asyncHandle = (AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) kvStates.get(key); + kvStates.put(key, asyncHandle.materialize()); + } + } + } + } } StreamTaskStateList allStates = new StreamTaskStateList(states); getEnvironment().acknowledgeCheckpoint(checkpointId, allStates); + LOG.debug("Finished asynchronous checkpoints for checkpoint {} on task {}", checkpointId, getName()); } catch (Exception e) { if (isRunning()) { @@ -523,7 +546,6 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>> } } asyncCheckpointThreads.remove(this); - LOG.debug("Finished asynchronous checkpoints for checkpoint {} on task {}", checkpointId, getName()); } }; http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskStateList.java ---------------------------------------------------------------------- diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskStateList.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskStateList.java index e698db6..c9946e1 100644 --- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskStateList.java +++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskStateList.java @@ -33,11 +33,35 @@ public class StreamTaskStateList implements StateHandle<StreamTaskState[]> { /** The states for all operator */ private final StreamTaskState[] states; - private final long stateSize; - public StreamTaskStateList(StreamTaskState[] states) throws Exception { this.states = states; + } + public boolean isEmpty() { + for (StreamTaskState state : states) { + if (state != null) { + return false; + } + } + return true; + } + + @Override + public StreamTaskState[] getState(ClassLoader userCodeClassLoader) { + return states; + } + + @Override + public void discardState() throws Exception { + for (StreamTaskState state : states) { + if (state != null) { + state.discardState(); + } + } + } + + @Override + public long getStateSize() throws Exception { long sumStateSize = 0; if (states != null) { @@ -67,34 +91,6 @@ public class StreamTaskStateList implements StateHandle<StreamTaskState[]> { } // State size as sum of all state sizes - stateSize = sumStateSize; - } - - public boolean isEmpty() { - for (StreamTaskState state : states) { - if (state != null) { - return false; - } - } - return true; - } - - @Override - public StreamTaskState[] getState(ClassLoader userCodeClassLoader) { - return states; - } - - @Override - public void discardState() throws Exception { - for (StreamTaskState state : states) { - if (state != null) { - state.discardState(); - } - } - } - - @Override - public long getStateSize() throws Exception { - return stateSize; + return sumStateSize; } } http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java ---------------------------------------------------------------------- diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java index 8f8dd9f..145edc2 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java @@ -19,6 +19,8 @@ package org.apache.flink.streaming.runtime.tasks; import org.apache.flink.api.common.typeinfo.TypeInformation; import org.apache.flink.api.common.typeutils.TypeSerializer; +import org.apache.flink.api.java.ClosureCleaner; +import org.apache.flink.api.java.functions.KeySelector; import org.apache.flink.runtime.event.AbstractEvent; import org.apache.flink.runtime.io.network.partition.consumer.StreamTestSingleInputGate; @@ -99,5 +101,10 @@ public class OneInputStreamTaskTestHarness<IN, OUT> extends StreamTaskTestHarnes streamConfig.setTypeSerializerIn1(inputSerializer); } + public <K> void configureForKeyedStream(KeySelector<IN, K> keySelector, TypeInformation<K> keyType) { + ClosureCleaner.clean(keySelector, false); + streamConfig.setStatePartitioner(0, keySelector); + streamConfig.setStateKeySerializer(keyType.createSerializer(executionConfig)); + } } http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java ---------------------------------------------------------------------- diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java index 313188a..1830054 100644 --- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java +++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java @@ -52,7 +52,8 @@ public class StreamTaskAsyncCheckpointTest { * This ensures that asynchronous state handles are actually materialized asynchonously. * * <p>We use latches to block at various stages and see if the code still continues through - * the parts that are not asynchronous. + * the parts that are not asynchronous. If the checkpoint is not done asynchronously the + * test will simply lock forever. * @throws Exception */ @Test http://git-wip-us.apache.org/repos/asf/flink/blob/31310cd6/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java ---------------------------------------------------------------------- diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java index 9400bd7..d9ced4b 100644 --- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java +++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java @@ -51,10 +51,16 @@ import org.junit.BeforeClass; import org.junit.Rule; import org.junit.Test; import org.junit.rules.TemporaryFolder; +import org.junit.rules.TestRule; +import org.junit.runner.Description; import org.junit.runner.RunWith; import org.junit.runners.Parameterized; +import org.junit.runners.model.Statement; import java.io.IOException; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; import java.util.Arrays; import java.util.Collection; import java.util.HashMap;
