Java Code Examples for org.rocksdb.Snapshot#close()

@Override
public long getApproximateKeysInRange(final byte[] startKey, final byte[] endKey) {
    // TODO This is a sad code, the performance is too damn bad
    final Timer.Context timeCtx = getTimeContext("APPROXIMATE_KEYS");
    final Lock readLock = this.readWriteLock.readLock();
    readLock.lock();
    final Snapshot snapshot = this.db.getSnapshot();
    try (final ReadOptions readOptions = new ReadOptions()) {
        readOptions.setSnapshot(snapshot);
        try (final RocksIterator it = this.db.newIterator(readOptions)) {
            if (startKey == null) {
                it.seekToFirst();
            } else {
                it.seek(startKey);
            }
            long approximateKeys = 0;
            for (;;) {
                // The accuracy is 100, don't ask more
                for (int i = 0; i < 100; i++) {
                    if (!it.isValid()) {
                        return approximateKeys;
                    }
                    it.next();
                    ++approximateKeys;
                }
                if (endKey != null && BytesUtil.compare(it.key(), endKey) >= 0) {
                    return approximateKeys;
                }
            }
        }
    } finally {
        // Nothing to release, rocksDB never own the pointer for a snapshot.
        snapshot.close();
        // The pointer to the snapshot is released by the database instance.
        this.db.releaseSnapshot(snapshot);
        readLock.unlock();
        timeCtx.stop();
    }
}
 

/**
 * Migrate only the state value, that is the "value" that is stored in RocksDB. We don't migrate
 * the key here, which is made up of key group, key, namespace and map key
 * (in case of MapState).
 */
private <N, S extends State, SV> void migrateStateValues(
	StateDescriptor<S, SV> stateDesc,
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> stateMetaInfo) throws Exception {

	if (stateDesc.getType() == StateDescriptor.Type.MAP) {
		throw new StateMigrationException("The new serializer for a MapState requires state migration in order for the job to proceed." +
			" However, migration for MapState currently isn't supported.");
	}

	LOG.info(
		"Performing state migration for state {} because the state serializer's schema, i.e. serialization format, has changed.",
		stateDesc);

	// we need to get an actual state instance because migration is different
	// for different state types. For example, ListState needs to deal with
	// individual elements
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), this.getClass());
		throw new FlinkRuntimeException(message);
	}
	State state = stateFactory.createState(
		stateDesc,
		stateMetaInfo,
		RocksDBKeyedStateBackend.this);
	if (!(state instanceof AbstractRocksDBState)) {
		throw new FlinkRuntimeException(
			"State should be an AbstractRocksDBState but is " + state);
	}

	@SuppressWarnings("unchecked")
	AbstractRocksDBState<?, ?, SV> rocksDBState = (AbstractRocksDBState<?, ?, SV>) state;

	Snapshot rocksDBSnapshot = db.getSnapshot();
	try (
		RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(db, stateMetaInfo.f0);
		RocksDBWriteBatchWrapper batchWriter = new RocksDBWriteBatchWrapper(db, getWriteOptions())
	) {
		iterator.seekToFirst();

		DataInputDeserializer serializedValueInput = new DataInputDeserializer();
		DataOutputSerializer migratedSerializedValueOutput = new DataOutputSerializer(512);
		while (iterator.isValid()) {
			serializedValueInput.setBuffer(iterator.value());

			rocksDBState.migrateSerializedValue(
				serializedValueInput,
				migratedSerializedValueOutput,
				stateMetaInfo.f1.getPreviousStateSerializer(),
				stateMetaInfo.f1.getStateSerializer());

			batchWriter.put(stateMetaInfo.f0, iterator.key(), migratedSerializedValueOutput.getCopyOfBuffer());

			migratedSerializedValueOutput.clear();
			iterator.next();
		}
	} finally {
		db.releaseSnapshot(rocksDBSnapshot);
		rocksDBSnapshot.close();
	}
}
 

@Override
public byte[] jumpOver(final byte[] startKey, final long distance) {
    final Timer.Context timeCtx = getTimeContext("JUMP_OVER");
    final Lock readLock = this.readWriteLock.readLock();
    readLock.lock();
    final Snapshot snapshot = this.db.getSnapshot();
    try (final ReadOptions readOptions = new ReadOptions()) {
        readOptions.setSnapshot(snapshot);
        try (final RocksIterator it = this.db.newIterator(readOptions)) {
            if (startKey == null) {
                it.seekToFirst();
            } else {
                it.seek(startKey);
            }
            long approximateKeys = 0;
            for (;;) {
                byte[] lastKey = null;
                if (it.isValid()) {
                    lastKey = it.key();
                }
                // The accuracy is 100, don't ask more
                for (int i = 0; i < 100; i++) {
                    if (!it.isValid()) {
                        return lastKey;
                    }
                    it.next();
                    if (++approximateKeys >= distance) {
                        return it.key();
                    }
                }
            }
        }
    } finally {
        // Nothing to release, rocksDB never own the pointer for a snapshot.
        snapshot.close();
        // The pointer to the snapshot is released by the database instance.
        this.db.releaseSnapshot(snapshot);
        readLock.unlock();
        timeCtx.stop();
    }
}
 

void doCreateSstFiles(final Snapshot snapshot, final EnumMap<SstColumnFamily, File> sstFileTable,
                      final byte[] startKey, final byte[] endKey, final CompletableFuture<Void> future) {
    final Timer.Context timeCtx = getTimeContext("CREATE_SST_FILE");
    final Lock readLock = this.readWriteLock.readLock();
    readLock.lock();
    try {
        if (!this.shutdownLock.isAvailable()) {
            // KV store has shutdown, we do not release rocksdb's snapshot
            future.completeExceptionally(new StorageException("KV store has shutdown."));
            return;
        }
        try (final ReadOptions readOptions = new ReadOptions();
                final EnvOptions envOptions = new EnvOptions();
                final Options options = new Options().setMergeOperator(new StringAppendOperator())) {
            readOptions.setSnapshot(snapshot);
            for (final Map.Entry<SstColumnFamily, File> entry : sstFileTable.entrySet()) {
                final SstColumnFamily sstColumnFamily = entry.getKey();
                final File sstFile = entry.getValue();
                final ColumnFamilyHandle columnFamilyHandle = findColumnFamilyHandle(sstColumnFamily);
                try (final RocksIterator it = this.db.newIterator(columnFamilyHandle, readOptions);
                        final SstFileWriter sstFileWriter = new SstFileWriter(envOptions, options)) {
                    if (startKey == null) {
                        it.seekToFirst();
                    } else {
                        it.seek(startKey);
                    }
                    sstFileWriter.open(sstFile.getAbsolutePath());
                    long count = 0;
                    for (;;) {
                        if (!it.isValid()) {
                            break;
                        }
                        final byte[] key = it.key();
                        if (endKey != null && BytesUtil.compare(key, endKey) >= 0) {
                            break;
                        }
                        sstFileWriter.put(key, it.value());
                        ++count;
                        it.next();
                    }
                    if (count == 0) {
                        sstFileWriter.close();
                    } else {
                        sstFileWriter.finish();
                    }
                    LOG.info("Finish sst file {} with {} keys.", sstFile, count);
                } catch (final RocksDBException e) {
                    throw new StorageException("Fail to create sst file at path: " + sstFile, e);
                }
            }
            future.complete(null);
        } catch (final Throwable t) {
            future.completeExceptionally(t);
        } finally {
            // Nothing to release, rocksDB never own the pointer for a snapshot.
            snapshot.close();
            // The pointer to the snapshot is released by the database instance.
            this.db.releaseSnapshot(snapshot);
        }
    } finally {
        readLock.unlock();
        timeCtx.stop();
    }
}
 

/**
 * Migrate only the state value, that is the "value" that is stored in RocksDB. We don't migrate
 * the key here, which is made up of key group, key, namespace and map key
 * (in case of MapState).
 */
@SuppressWarnings("unchecked")
private <N, S extends State, SV> void migrateStateValues(
	StateDescriptor<S, SV> stateDesc,
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> stateMetaInfo) throws Exception {

	if (stateDesc.getType() == StateDescriptor.Type.MAP) {
		TypeSerializerSnapshot<SV> previousSerializerSnapshot = stateMetaInfo.f1.getPreviousStateSerializerSnapshot();
		checkState(previousSerializerSnapshot != null, "the previous serializer snapshot should exist.");
		checkState(previousSerializerSnapshot instanceof MapSerializerSnapshot, "previous serializer snapshot should be a MapSerializerSnapshot.");

		TypeSerializer<SV> newSerializer = stateMetaInfo.f1.getStateSerializer();
		checkState(newSerializer instanceof MapSerializer, "new serializer should be a MapSerializer.");

		MapSerializer<?, ?> mapSerializer = (MapSerializer<?, ?>) newSerializer;
		MapSerializerSnapshot<?, ?> mapSerializerSnapshot = (MapSerializerSnapshot<?, ?>) previousSerializerSnapshot;
		if (!checkMapStateKeySchemaCompatibility(mapSerializerSnapshot, mapSerializer)) {
			throw new StateMigrationException(
				"The new serializer for a MapState requires state migration in order for the job to proceed, since the key schema has changed. However, migration for MapState currently only allows value schema evolutions.");
		}
	}

	LOG.info(
		"Performing state migration for state {} because the state serializer's schema, i.e. serialization format, has changed.",
		stateDesc);

	// we need to get an actual state instance because migration is different
	// for different state types. For example, ListState needs to deal with
	// individual elements
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), this.getClass());
		throw new FlinkRuntimeException(message);
	}
	State state = stateFactory.createState(
		stateDesc,
		stateMetaInfo,
		RocksDBKeyedStateBackend.this);
	if (!(state instanceof AbstractRocksDBState)) {
		throw new FlinkRuntimeException(
			"State should be an AbstractRocksDBState but is " + state);
	}

	@SuppressWarnings("unchecked")
	AbstractRocksDBState<?, ?, SV> rocksDBState = (AbstractRocksDBState<?, ?, SV>) state;

	Snapshot rocksDBSnapshot = db.getSnapshot();
	try (
		RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(db, stateMetaInfo.f0);
		RocksDBWriteBatchWrapper batchWriter = new RocksDBWriteBatchWrapper(db, getWriteOptions())
	) {
		iterator.seekToFirst();

		DataInputDeserializer serializedValueInput = new DataInputDeserializer();
		DataOutputSerializer migratedSerializedValueOutput = new DataOutputSerializer(512);
		while (iterator.isValid()) {
			serializedValueInput.setBuffer(iterator.value());

			rocksDBState.migrateSerializedValue(
				serializedValueInput,
				migratedSerializedValueOutput,
				stateMetaInfo.f1.getPreviousStateSerializer(),
				stateMetaInfo.f1.getStateSerializer());

			batchWriter.put(stateMetaInfo.f0, iterator.key(), migratedSerializedValueOutput.getCopyOfBuffer());

			migratedSerializedValueOutput.clear();
			iterator.next();
		}
	} finally {
		db.releaseSnapshot(rocksDBSnapshot);
		rocksDBSnapshot.close();
	}
}
 

/**
 * Migrate only the state value, that is the "value" that is stored in RocksDB. We don't migrate
 * the key here, which is made up of key group, key, namespace and map key
 * (in case of MapState).
 */
@SuppressWarnings("unchecked")
private <N, S extends State, SV> void migrateStateValues(
	StateDescriptor<S, SV> stateDesc,
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> stateMetaInfo) throws Exception {

	if (stateDesc.getType() == StateDescriptor.Type.MAP) {
		TypeSerializerSnapshot<SV> previousSerializerSnapshot = stateMetaInfo.f1.getPreviousStateSerializerSnapshot();
		checkState(previousSerializerSnapshot != null, "the previous serializer snapshot should exist.");
		checkState(previousSerializerSnapshot instanceof MapSerializerSnapshot, "previous serializer snapshot should be a MapSerializerSnapshot.");

		TypeSerializer<SV> newSerializer = stateMetaInfo.f1.getStateSerializer();
		checkState(newSerializer instanceof MapSerializer, "new serializer should be a MapSerializer.");

		MapSerializer<?, ?> mapSerializer = (MapSerializer<?, ?>) newSerializer;
		MapSerializerSnapshot<?, ?> mapSerializerSnapshot = (MapSerializerSnapshot<?, ?>) previousSerializerSnapshot;
		if (!checkMapStateKeySchemaCompatibility(mapSerializerSnapshot, mapSerializer)) {
			throw new StateMigrationException(
				"The new serializer for a MapState requires state migration in order for the job to proceed, since the key schema has changed. However, migration for MapState currently only allows value schema evolutions.");
		}
	}

	LOG.info(
		"Performing state migration for state {} because the state serializer's schema, i.e. serialization format, has changed.",
		stateDesc);

	// we need to get an actual state instance because migration is different
	// for different state types. For example, ListState needs to deal with
	// individual elements
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), this.getClass());
		throw new FlinkRuntimeException(message);
	}
	State state = stateFactory.createState(
		stateDesc,
		stateMetaInfo,
		RocksDBKeyedStateBackend.this);
	if (!(state instanceof AbstractRocksDBState)) {
		throw new FlinkRuntimeException(
			"State should be an AbstractRocksDBState but is " + state);
	}

	@SuppressWarnings("unchecked")
	AbstractRocksDBState<?, ?, SV> rocksDBState = (AbstractRocksDBState<?, ?, SV>) state;

	Snapshot rocksDBSnapshot = db.getSnapshot();
	try (
		RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(db, stateMetaInfo.f0, readOptions);
		RocksDBWriteBatchWrapper batchWriter = new RocksDBWriteBatchWrapper(db, getWriteOptions(), getWriteBatchSize())
	) {
		iterator.seekToFirst();

		DataInputDeserializer serializedValueInput = new DataInputDeserializer();
		DataOutputSerializer migratedSerializedValueOutput = new DataOutputSerializer(512);
		while (iterator.isValid()) {
			serializedValueInput.setBuffer(iterator.value());

			rocksDBState.migrateSerializedValue(
				serializedValueInput,
				migratedSerializedValueOutput,
				stateMetaInfo.f1.getPreviousStateSerializer(),
				stateMetaInfo.f1.getStateSerializer());

			batchWriter.put(stateMetaInfo.f0, iterator.key(), migratedSerializedValueOutput.getCopyOfBuffer());

			migratedSerializedValueOutput.clear();
			iterator.next();
		}
	} finally {
		db.releaseSnapshot(rocksDBSnapshot);
		rocksDBSnapshot.close();
	}
}