[FLINK-3359] Make RocksDB File Copies Asynchronous

flink-contrib/flink-statebackend-rocksdb/pom.xml

@@ -57,15 +57,24 @@ under the License.
 			<artifactId>rocksdbjni</artifactId>
 			<version>4.1.0</version>
 		</dependency>
-
 		<dependency>
 			<groupId>org.apache.flink</groupId>
 			<artifactId>flink-runtime_2.10</artifactId>
 			<version>${project.version}</version>
 			<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>

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.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.File;
 import java.io.IOException;
 import java.net.URI;
+import java.util.UUID;
 
 import static java.util.Objects.requireNonNull;
 
@@ -58,6 +60,11 @@
  * 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 @@ protected AbstractRocksDBState(
 	 *
 	 * @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 @@ final public void setCurrentNamespace(N namespace) {
 	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 static abstract class AbstractRocksDBSnapshot<K, N, S extends State, SD e
 		/** 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 AbstractRocksDBSnapshot(File dbPath,
 			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 final KvState<K, N, S, SD, RocksDBStateBackend> restoreState(
 				}
 			}
 
-			FileSystem fs = FileSystem.get(backupUri, new Configuration());
-
 			final File localBackupPath = new File(dbPath, "chk-" + checkpointId);
 
 			if (localBackupPath.exists()) {
@@ -365,5 +393,42 @@ public final long getStateSize() throws Exception {
 			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);
+			}
+		}
+	}
 }
 

flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java

@@ -156,7 +156,7 @@ public void add(T value) throws IOException {
 	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>> {

flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java

@@ -154,7 +154,7 @@ protected AbstractRocksDBSnapshot<K, N, ListState<V>, ListStateDescriptor<V>> cr
 			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 

flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java

@@ -152,7 +152,7 @@ protected AbstractRocksDBSnapshot<K, N, ReducingState<V>, ReducingStateDescripto
 			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 

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
 

flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java

@@ -140,7 +140,7 @@ protected AbstractRocksDBSnapshot<K, N, ValueState<V>, ValueStateDescriptor<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>