/**
* Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state
* handle
*
* @throws IOException
* @throws RocksDBException
*/
private void restoreKVStateData() throws IOException, RocksDBException {
// for all key-groups in the current state handle...
for (Tuple2<Integer, Long> keyGroupOffset :
currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
long offset = keyGroupOffset.f1;
// not empty key-group?
if (0L != offset) {
currentStateHandleInStream.seek(offset);
boolean keyGroupHasMoreKeys = true;
// TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
int kvStateId = currentStateHandleInView.readShort();
ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
// insert all k/v pairs into DB
while (keyGroupHasMoreKeys) {
byte[] key =
BytePrimitiveArraySerializer.INSTANCE.deserialize(currentStateHandleInView);
byte[] value =
BytePrimitiveArraySerializer.INSTANCE.deserialize(currentStateHandleInView);
if (RocksDBSnapshotOperation.hasMetaDataFollowsFlag(key)) {
// clear the signal bit in the key to make it ready for insertion again
RocksDBSnapshotOperation.clearMetaDataFollowsFlag(key);
rocksDBKeyedStateBackend.db.put(handle, key, value);
// TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
kvStateId =
RocksDBSnapshotOperation.END_OF_KEY_GROUP_MARK
& currentStateHandleInView.readShort();
if (RocksDBSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
keyGroupHasMoreKeys = false;
} else {
handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
}
} else {
rocksDBKeyedStateBackend.db.put(handle, key, value);
}
}
}
}
}
/**
* Triggers an asynchronous snapshot of the keyed state backend from RocksDB. This snapshot can be
* canceled and is also stopped when the backend is closed through {@link #dispose()}. For each
* backend, this method must always be called by the same thread.
*
* @param checkpointId The Id of the checkpoint.
* @param timestamp The timestamp of the checkpoint.
* @param streamFactory The factory that we can use for writing our state to streams.
* @return Future to the state handle of the snapshot data.
* @throws Exception
*/
@Override
public RunnableFuture<KeyGroupsStateHandle> snapshot(
final long checkpointId, final long timestamp, final CheckpointStreamFactory streamFactory)
throws Exception {
long startTime = System.currentTimeMillis();
final RocksDBSnapshotOperation snapshotOperation =
new RocksDBSnapshotOperation(this, streamFactory);
// hold the db lock while operation on the db to guard us against async db disposal
synchronized (asyncSnapshotLock) {
if (kvStateInformation.isEmpty()) {
LOG.info(
"Asynchronous RocksDB snapshot performed on empty keyed state at "
+ timestamp
+ " . Returning null.");
return new DoneFuture<>(null);
}
if (db != null) {
snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
} else {
throw new IOException("RocksDB closed.");
}
}
// implementation of the async IO operation, based on FutureTask
AbstractAsyncIOCallable<
KeyGroupsStateHandle, CheckpointStreamFactory.CheckpointStateOutputStream>
ioCallable =
new AbstractAsyncIOCallable<
KeyGroupsStateHandle, CheckpointStreamFactory.CheckpointStateOutputStream>() {
@Override
public CheckpointStreamFactory.CheckpointStateOutputStream openIOHandle()
throws Exception {
snapshotOperation.openCheckpointStream();
return snapshotOperation.getOutStream();
}
@Override
public KeyGroupsStateHandle performOperation() throws Exception {
long startTime = System.currentTimeMillis();
synchronized (asyncSnapshotLock) {
try {
// hold the db lock while operation on the db to guard us against async db
// disposal
if (db == null) {
throw new IOException("RocksDB closed.");
}
snapshotOperation.writeDBSnapshot();
} finally {
snapshotOperation.closeCheckpointStream();
}
}
LOG.info(
"Asynchronous RocksDB snapshot ("
+ streamFactory
+ ", asynchronous part) in thread "
+ Thread.currentThread()
+ " took "
+ (System.currentTimeMillis() - startTime)
+ " ms.");
return snapshotOperation.getSnapshotResultStateHandle();
}
private void releaseSnapshotOperationResources(boolean canceled) {
// hold the db lock while operation on the db to guard us against async db disposal
synchronized (asyncSnapshotLock) {
snapshotOperation.releaseSnapshotResources(canceled);
}
}
@Override
public void done(boolean canceled) {
releaseSnapshotOperationResources(canceled);
}
};
LOG.info(
"Asynchronous RocksDB snapshot ("
+ streamFactory
+ ", synchronous part) in thread "
+ Thread.currentThread()
+ " took "
+ (System.currentTimeMillis() - startTime)
+ " ms.");
return AsyncStoppableTaskWithCallback.from(ioCallable);
}
