org.apache.flink.contrib.streaming.state.RocksDBOperationUtils Java Examples

/**
 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
 */
private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
	List<StateMetaInfoSnapshot> stateMetaInfoSnapshots,
	boolean registerTtlCompactFilter) {

	List<ColumnFamilyDescriptor> columnFamilyDescriptors =
		new ArrayList<>(stateMetaInfoSnapshots.size());

	for (StateMetaInfoSnapshot stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
		RegisteredStateMetaInfoBase metaInfoBase =
			RegisteredStateMetaInfoBase.fromMetaInfoSnapshot(stateMetaInfoSnapshot);
		ColumnFamilyDescriptor columnFamilyDescriptor = RocksDBOperationUtils.createColumnFamilyDescriptor(
			metaInfoBase, columnFamilyOptionsFactory, registerTtlCompactFilter ? ttlCompactFiltersManager : null);
		columnFamilyDescriptors.add(columnFamilyDescriptor);
	}
	return columnFamilyDescriptors;
}
 

/**
 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
 */
private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
	List<StateMetaInfoSnapshot> stateMetaInfoSnapshots,
	boolean registerTtlCompactFilter) {

	List<ColumnFamilyDescriptor> columnFamilyDescriptors =
		new ArrayList<>(stateMetaInfoSnapshots.size());

	for (StateMetaInfoSnapshot stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
		RegisteredStateMetaInfoBase metaInfoBase =
			RegisteredStateMetaInfoBase.fromMetaInfoSnapshot(stateMetaInfoSnapshot);
		ColumnFamilyDescriptor columnFamilyDescriptor = RocksDBOperationUtils.createColumnFamilyDescriptor(
			metaInfoBase, columnFamilyOptionsFactory, registerTtlCompactFilter ? ttlCompactFiltersManager : null);
		columnFamilyDescriptors.add(columnFamilyDescriptor);
	}
	return columnFamilyDescriptors;
}
 

/**
 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
 */
private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
	List<StateMetaInfoSnapshot> stateMetaInfoSnapshots,
	boolean registerTtlCompactFilter) {

	List<ColumnFamilyDescriptor> columnFamilyDescriptors =
		new ArrayList<>(stateMetaInfoSnapshots.size());

	for (StateMetaInfoSnapshot stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
		RegisteredStateMetaInfoBase metaInfoBase =
			RegisteredStateMetaInfoBase.fromMetaInfoSnapshot(stateMetaInfoSnapshot);
		ColumnFamilyDescriptor columnFamilyDescriptor = RocksDBOperationUtils.createColumnFamilyDescriptor(
			metaInfoBase, columnFamilyOptionsFactory, registerTtlCompactFilter ? ttlCompactFiltersManager : null);
		columnFamilyDescriptors.add(columnFamilyDescriptor);
	}
	return columnFamilyDescriptors;
}
 

void openDB() throws IOException {
	db = RocksDBOperationUtils.openDB(
		dbPath,
		columnFamilyDescriptors,
		columnFamilyHandles,
		RocksDBOperationUtils.createColumnFamilyOptions(columnFamilyOptionsFactory, "default"),
		dbOptions);
	// remove the default column family which is located at the first index
	defaultColumnFamilyHandle = columnFamilyHandles.remove(0);
	// init native metrics monitor if configured
	nativeMetricMonitor = nativeMetricOptions.isEnabled() ?
		new RocksDBNativeMetricMonitor(nativeMetricOptions, metricGroup, db) : null;
}
 

@Override
public void close() {
	List<ColumnFamilyOptions> columnFamilyOptions = new ArrayList<>(columnFamilyDescriptors.size() + 1);
	columnFamilyDescriptors.forEach((cfd) -> columnFamilyOptions.add(cfd.getOptions()));
	RocksDBOperationUtils.addColumnFamilyOptionsToCloseLater(columnFamilyOptions, defaultColumnFamilyHandle);
	IOUtils.closeQuietly(defaultColumnFamilyHandle);
	IOUtils.closeAllQuietly(columnFamilyHandles);
	IOUtils.closeQuietly(db);
	IOUtils.closeAllQuietly(columnFamilyOptions);
}
 

private RestoredDBInstance restoreDBInstanceFromStateHandle(
	IncrementalRemoteKeyedStateHandle restoreStateHandle,
	Path temporaryRestoreInstancePath) throws Exception {

	try (RocksDBStateDownloader rocksDBStateDownloader =
			new RocksDBStateDownloader(numberOfTransferringThreads)) {
		rocksDBStateDownloader.transferAllStateDataToDirectory(
			restoreStateHandle,
			temporaryRestoreInstancePath,
			cancelStreamRegistry);
	}

	KeyedBackendSerializationProxy<K> serializationProxy = readMetaData(restoreStateHandle.getMetaStateHandle());
	// read meta data
	List<StateMetaInfoSnapshot> stateMetaInfoSnapshots = serializationProxy.getStateMetaInfoSnapshots();

	List<ColumnFamilyDescriptor> columnFamilyDescriptors =
		createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots, false);

	List<ColumnFamilyHandle> columnFamilyHandles =
		new ArrayList<>(stateMetaInfoSnapshots.size() + 1);

	RocksDB restoreDb = RocksDBOperationUtils.openDB(
		temporaryRestoreInstancePath.getPath(),
		columnFamilyDescriptors,
		columnFamilyHandles,
		RocksDBOperationUtils.createColumnFamilyOptions(columnFamilyOptionsFactory, "default"),
		dbOptions);

	return new RestoredDBInstance(restoreDb, columnFamilyHandles, columnFamilyDescriptors, stateMetaInfoSnapshots);
}
 

void openDB() throws IOException {
	db = RocksDBOperationUtils.openDB(
		dbPath,
		columnFamilyDescriptors,
		columnFamilyHandles,
		RocksDBOperationUtils.createColumnFamilyOptions(columnFamilyOptionsFactory, "default"),
		dbOptions);
	// remove the default column family which is located at the first index
	defaultColumnFamilyHandle = columnFamilyHandles.remove(0);
	// init native metrics monitor if configured
	nativeMetricMonitor = nativeMetricOptions.isEnabled() ?
		new RocksDBNativeMetricMonitor(nativeMetricOptions, metricGroup, db) : null;
}
 

@Override
public void close() {
	List<ColumnFamilyOptions> columnFamilyOptions = new ArrayList<>(columnFamilyDescriptors.size() + 1);
	columnFamilyDescriptors.forEach((cfd) -> columnFamilyOptions.add(cfd.getOptions()));
	RocksDBOperationUtils.addColumnFamilyOptionsToCloseLater(columnFamilyOptions, defaultColumnFamilyHandle);
	IOUtils.closeQuietly(defaultColumnFamilyHandle);
	IOUtils.closeAllQuietly(columnFamilyHandles);
	IOUtils.closeQuietly(db);
	IOUtils.closeAllQuietly(columnFamilyOptions);
}
 

private RestoredDBInstance restoreDBInstanceFromStateHandle(
	IncrementalRemoteKeyedStateHandle restoreStateHandle,
	Path temporaryRestoreInstancePath) throws Exception {

	try (RocksDBStateDownloader rocksDBStateDownloader =
			new RocksDBStateDownloader(numberOfTransferringThreads)) {
		rocksDBStateDownloader.transferAllStateDataToDirectory(
			restoreStateHandle,
			temporaryRestoreInstancePath,
			cancelStreamRegistry);
	}

	KeyedBackendSerializationProxy<K> serializationProxy = readMetaData(restoreStateHandle.getMetaStateHandle());
	// read meta data
	List<StateMetaInfoSnapshot> stateMetaInfoSnapshots = serializationProxy.getStateMetaInfoSnapshots();

	List<ColumnFamilyDescriptor> columnFamilyDescriptors =
		createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots, false);

	List<ColumnFamilyHandle> columnFamilyHandles =
		new ArrayList<>(stateMetaInfoSnapshots.size() + 1);

	RocksDB restoreDb = RocksDBOperationUtils.openDB(
		temporaryRestoreInstancePath.getPath(),
		columnFamilyDescriptors,
		columnFamilyHandles,
		RocksDBOperationUtils.createColumnFamilyOptions(columnFamilyOptionsFactory, "default"),
		dbOptions);

	return new RestoredDBInstance(restoreDb, columnFamilyHandles, columnFamilyDescriptors, stateMetaInfoSnapshots);
}
 

void openDB() throws IOException {
	db = RocksDBOperationUtils.openDB(
		dbPath,
		columnFamilyDescriptors,
		columnFamilyHandles,
		RocksDBOperationUtils.createColumnFamilyOptions(columnFamilyOptionsFactory, "default"),
		dbOptions);
	// remove the default column family which is located at the first index
	defaultColumnFamilyHandle = columnFamilyHandles.remove(0);
	// init native metrics monitor if configured
	nativeMetricMonitor = nativeMetricOptions.isEnabled() ?
		new RocksDBNativeMetricMonitor(nativeMetricOptions, metricGroup, db) : null;
}
 

private RestoredDBInstance(
	@Nonnull RocksDB db,
	@Nonnull List<ColumnFamilyHandle> columnFamilyHandles,
	@Nonnull List<ColumnFamilyDescriptor> columnFamilyDescriptors,
	@Nonnull List<StateMetaInfoSnapshot> stateMetaInfoSnapshots) {
	this.db = db;
	this.defaultColumnFamilyHandle = columnFamilyHandles.remove(0);
	this.columnFamilyHandles = columnFamilyHandles;
	this.columnFamilyDescriptors = columnFamilyDescriptors;
	this.stateMetaInfoSnapshots = stateMetaInfoSnapshots;
	this.readOptions = RocksDBOperationUtils.createTotalOrderSeekReadOptions();
}
 

@Override
public void close() {
	List<ColumnFamilyOptions> columnFamilyOptions = new ArrayList<>(columnFamilyDescriptors.size() + 1);
	columnFamilyDescriptors.forEach((cfd) -> columnFamilyOptions.add(cfd.getOptions()));
	RocksDBOperationUtils.addColumnFamilyOptionsToCloseLater(columnFamilyOptions, defaultColumnFamilyHandle);
	IOUtils.closeQuietly(defaultColumnFamilyHandle);
	IOUtils.closeAllQuietly(columnFamilyHandles);
	IOUtils.closeQuietly(db);
	IOUtils.closeAllQuietly(columnFamilyOptions);
	IOUtils.closeQuietly(readOptions);
}
 

private RestoredDBInstance restoreDBInstanceFromStateHandle(
	IncrementalRemoteKeyedStateHandle restoreStateHandle,
	Path temporaryRestoreInstancePath) throws Exception {

	try (RocksDBStateDownloader rocksDBStateDownloader =
			new RocksDBStateDownloader(numberOfTransferringThreads)) {
		rocksDBStateDownloader.transferAllStateDataToDirectory(
			restoreStateHandle,
			temporaryRestoreInstancePath,
			cancelStreamRegistry);
	}

	KeyedBackendSerializationProxy<K> serializationProxy = readMetaData(restoreStateHandle.getMetaStateHandle());
	// read meta data
	List<StateMetaInfoSnapshot> stateMetaInfoSnapshots = serializationProxy.getStateMetaInfoSnapshots();

	List<ColumnFamilyDescriptor> columnFamilyDescriptors =
		createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots, false);

	List<ColumnFamilyHandle> columnFamilyHandles =
		new ArrayList<>(stateMetaInfoSnapshots.size() + 1);

	RocksDB restoreDb = RocksDBOperationUtils.openDB(
		temporaryRestoreInstancePath.toString(),
		columnFamilyDescriptors,
		columnFamilyHandles,
		RocksDBOperationUtils.createColumnFamilyOptions(columnFamilyOptionsFactory, "default"),
		dbOptions);

	return new RestoredDBInstance(restoreDb, columnFamilyHandles, columnFamilyDescriptors, stateMetaInfoSnapshots);
}
 

/**
 * Recovery from multi incremental states with rescaling. For rescaling, this method creates a temporary
 * RocksDB instance for a key-groups shard. All contents from the temporary instance are copied into the
 * real restore instance and then the temporary instance is discarded.
 */
private void restoreWithRescaling(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {

	// Prepare for restore with rescaling
	KeyedStateHandle initialHandle = RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(
		restoreStateHandles, keyGroupRange);

	// Init base DB instance
	if (initialHandle != null) {
		restoreStateHandles.remove(initialHandle);
		initDBWithRescaling(initialHandle);
	} else {
		openDB();
	}

	// Transfer remaining key-groups from temporary instance into base DB
	byte[] startKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
	RocksDBKeySerializationUtils.serializeKeyGroup(keyGroupRange.getStartKeyGroup(), startKeyGroupPrefixBytes);

	byte[] stopKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
	RocksDBKeySerializationUtils.serializeKeyGroup(keyGroupRange.getEndKeyGroup() + 1, stopKeyGroupPrefixBytes);

	for (KeyedStateHandle rawStateHandle : restoreStateHandles) {

		if (!(rawStateHandle instanceof IncrementalRemoteKeyedStateHandle)) {
			throw new IllegalStateException("Unexpected state handle type, " +
				"expected " + IncrementalRemoteKeyedStateHandle.class +
				", but found " + rawStateHandle.getClass());
		}

		Path temporaryRestoreInstancePath = new Path(instanceBasePath.getAbsolutePath() + UUID.randomUUID().toString());
		try (RestoredDBInstance tmpRestoreDBInfo = restoreDBInstanceFromStateHandle(
			(IncrementalRemoteKeyedStateHandle) rawStateHandle,
			temporaryRestoreInstancePath);
			RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper(this.db)) {

			List<ColumnFamilyDescriptor> tmpColumnFamilyDescriptors = tmpRestoreDBInfo.columnFamilyDescriptors;
			List<ColumnFamilyHandle> tmpColumnFamilyHandles = tmpRestoreDBInfo.columnFamilyHandles;

			// iterating only the requested descriptors automatically skips the default column family handle
			for (int i = 0; i < tmpColumnFamilyDescriptors.size(); ++i) {
				ColumnFamilyHandle tmpColumnFamilyHandle = tmpColumnFamilyHandles.get(i);

				ColumnFamilyHandle targetColumnFamilyHandle = getOrRegisterStateColumnFamilyHandle(
					null, tmpRestoreDBInfo.stateMetaInfoSnapshots.get(i))
					.columnFamilyHandle;

				try (RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(tmpRestoreDBInfo.db, tmpColumnFamilyHandle)) {

					iterator.seek(startKeyGroupPrefixBytes);

					while (iterator.isValid()) {

						if (RocksDBIncrementalCheckpointUtils.beforeThePrefixBytes(iterator.key(), stopKeyGroupPrefixBytes)) {
							writeBatchWrapper.put(targetColumnFamilyHandle, iterator.key(), iterator.value());
						} else {
							// Since the iterator will visit the record according to the sorted order,
							// we can just break here.
							break;
						}

						iterator.next();
					}
				} // releases native iterator resources
			}
		} finally {
			cleanUpPathQuietly(temporaryRestoreInstancePath);
		}
	}
}
 

/**
 * Recovery from multi incremental states with rescaling. For rescaling, this method creates a temporary
 * RocksDB instance for a key-groups shard. All contents from the temporary instance are copied into the
 * real restore instance and then the temporary instance is discarded.
 */
private void restoreWithRescaling(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {

	// Prepare for restore with rescaling
	KeyedStateHandle initialHandle = RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(
		restoreStateHandles, keyGroupRange);

	// Init base DB instance
	if (initialHandle != null) {
		restoreStateHandles.remove(initialHandle);
		initDBWithRescaling(initialHandle);
	} else {
		openDB();
	}

	// Transfer remaining key-groups from temporary instance into base DB
	byte[] startKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
	RocksDBKeySerializationUtils.serializeKeyGroup(keyGroupRange.getStartKeyGroup(), startKeyGroupPrefixBytes);

	byte[] stopKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
	RocksDBKeySerializationUtils.serializeKeyGroup(keyGroupRange.getEndKeyGroup() + 1, stopKeyGroupPrefixBytes);

	for (KeyedStateHandle rawStateHandle : restoreStateHandles) {

		if (!(rawStateHandle instanceof IncrementalRemoteKeyedStateHandle)) {
			throw new IllegalStateException("Unexpected state handle type, " +
				"expected " + IncrementalRemoteKeyedStateHandle.class +
				", but found " + rawStateHandle.getClass());
		}

		Path temporaryRestoreInstancePath = new Path(instanceBasePath.getAbsolutePath() + UUID.randomUUID().toString());
		try (RestoredDBInstance tmpRestoreDBInfo = restoreDBInstanceFromStateHandle(
			(IncrementalRemoteKeyedStateHandle) rawStateHandle,
			temporaryRestoreInstancePath);
			RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper(this.db)) {

			List<ColumnFamilyDescriptor> tmpColumnFamilyDescriptors = tmpRestoreDBInfo.columnFamilyDescriptors;
			List<ColumnFamilyHandle> tmpColumnFamilyHandles = tmpRestoreDBInfo.columnFamilyHandles;

			// iterating only the requested descriptors automatically skips the default column family handle
			for (int i = 0; i < tmpColumnFamilyDescriptors.size(); ++i) {
				ColumnFamilyHandle tmpColumnFamilyHandle = tmpColumnFamilyHandles.get(i);

				ColumnFamilyHandle targetColumnFamilyHandle = getOrRegisterStateColumnFamilyHandle(
					null, tmpRestoreDBInfo.stateMetaInfoSnapshots.get(i))
					.columnFamilyHandle;

				try (RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(tmpRestoreDBInfo.db, tmpColumnFamilyHandle)) {

					iterator.seek(startKeyGroupPrefixBytes);

					while (iterator.isValid()) {

						if (RocksDBIncrementalCheckpointUtils.beforeThePrefixBytes(iterator.key(), stopKeyGroupPrefixBytes)) {
							writeBatchWrapper.put(targetColumnFamilyHandle, iterator.key(), iterator.value());
						} else {
							// Since the iterator will visit the record according to the sorted order,
							// we can just break here.
							break;
						}

						iterator.next();
					}
				} // releases native iterator resources
			}
		} finally {
			cleanUpPathQuietly(temporaryRestoreInstancePath);
		}
	}
}
 

/**
 * Recovery from multi incremental states with rescaling. For rescaling, this method creates a temporary
 * RocksDB instance for a key-groups shard. All contents from the temporary instance are copied into the
 * real restore instance and then the temporary instance is discarded.
 */
private void restoreWithRescaling(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {

	// Prepare for restore with rescaling
	KeyedStateHandle initialHandle = RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(
		restoreStateHandles, keyGroupRange);

	// Init base DB instance
	if (initialHandle != null) {
		restoreStateHandles.remove(initialHandle);
		initDBWithRescaling(initialHandle);
	} else {
		openDB();
	}

	// Transfer remaining key-groups from temporary instance into base DB
	byte[] startKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
	RocksDBKeySerializationUtils.serializeKeyGroup(keyGroupRange.getStartKeyGroup(), startKeyGroupPrefixBytes);

	byte[] stopKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
	RocksDBKeySerializationUtils.serializeKeyGroup(keyGroupRange.getEndKeyGroup() + 1, stopKeyGroupPrefixBytes);

	for (KeyedStateHandle rawStateHandle : restoreStateHandles) {

		if (!(rawStateHandle instanceof IncrementalRemoteKeyedStateHandle)) {
			throw new IllegalStateException("Unexpected state handle type, " +
				"expected " + IncrementalRemoteKeyedStateHandle.class +
				", but found " + rawStateHandle.getClass());
		}

		Path temporaryRestoreInstancePath = instanceBasePath.getAbsoluteFile().toPath().resolve(UUID.randomUUID().toString());
		try (RestoredDBInstance tmpRestoreDBInfo = restoreDBInstanceFromStateHandle(
			(IncrementalRemoteKeyedStateHandle) rawStateHandle,
			temporaryRestoreInstancePath);
			RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper(this.db, writeBatchSize)) {

			List<ColumnFamilyDescriptor> tmpColumnFamilyDescriptors = tmpRestoreDBInfo.columnFamilyDescriptors;
			List<ColumnFamilyHandle> tmpColumnFamilyHandles = tmpRestoreDBInfo.columnFamilyHandles;

			// iterating only the requested descriptors automatically skips the default column family handle
			for (int i = 0; i < tmpColumnFamilyDescriptors.size(); ++i) {
				ColumnFamilyHandle tmpColumnFamilyHandle = tmpColumnFamilyHandles.get(i);

				ColumnFamilyHandle targetColumnFamilyHandle = getOrRegisterStateColumnFamilyHandle(
					null, tmpRestoreDBInfo.stateMetaInfoSnapshots.get(i))
					.columnFamilyHandle;

				try (RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(tmpRestoreDBInfo.db, tmpColumnFamilyHandle, tmpRestoreDBInfo.readOptions)) {

					iterator.seek(startKeyGroupPrefixBytes);

					while (iterator.isValid()) {

						if (RocksDBIncrementalCheckpointUtils.beforeThePrefixBytes(iterator.key(), stopKeyGroupPrefixBytes)) {
							writeBatchWrapper.put(targetColumnFamilyHandle, iterator.key(), iterator.value());
						} else {
							// Since the iterator will visit the record according to the sorted order,
							// we can just break here.
							break;
						}

						iterator.next();
					}
				} // releases native iterator resources
			}
		} finally {
			cleanUpPathQuietly(temporaryRestoreInstancePath);
		}
	}
}