org.apache.flink.runtime.state.internal.InternalKvState Java Examples

/**
 * Test StateIncrementalVisitor with closed state map.
 */
@Test
public void testStateIncrementalVisitorWithClosedStateMap() {
	CopyOnWriteSkipListStateMap<Integer, Long, String> stateMap = createEmptyStateMap();
	// put some states
	for (long namespace = 0;  namespace < 15; namespace++) {
		for (int key = 0; key < 20; key++) {
			String state = String.valueOf(namespace * key);
			stateMap.put(key, namespace, state);
		}
	}
	InternalKvState.StateIncrementalVisitor<Integer, Long, String> closedVisitor =
		stateMap.getStateIncrementalVisitor(5);
	assertTrue(closedVisitor.hasNext());

	stateMap.close();
	// the visitor will be invalid after state map is closed
	assertFalse(closedVisitor.hasNext());
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
private static <V, K, N> List<V> getSerializedList(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeList(serializedValue, valueSerializer);
	}
}
 

/**
 * @see KeyedStateBackend
 */
@Override
@SuppressWarnings("unchecked")
public <N, S extends State, V> S getOrCreateKeyedState(
		final TypeSerializer<N> namespaceSerializer,
		StateDescriptor<S, V> stateDescriptor) throws Exception {
	checkNotNull(namespaceSerializer, "Namespace serializer");
	checkNotNull(keySerializer, "State key serializer has not been configured in the config. " +
			"This operation cannot use partitioned state.");

	InternalKvState<K, ?, ?> kvState = keyValueStatesByName.get(stateDescriptor.getName());
	if (kvState == null) {
		if (!stateDescriptor.isSerializerInitialized()) {
			stateDescriptor.initializeSerializerUnlessSet(executionConfig);
		}
		kvState = TtlStateFactory.createStateAndWrapWithTtlIfEnabled(
			namespaceSerializer, stateDescriptor, this, ttlTimeProvider);
		keyValueStatesByName.put(stateDescriptor.getName(), kvState);
		publishQueryableStateIfEnabled(stateDescriptor, kvState);
	}
	return (S) kvState;
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
private static <V, K, N> List<V> getSerializedList(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeList(serializedValue, valueSerializer);
	}
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
protected static <V, K, N> V getSerializedValue(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeValue(serializedValue, valueSerializer);
	}
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
protected static <V, K, N> V getSerializedValue(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeValue(serializedValue, valueSerializer);
	}
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
private static <V, K, N> List<V> getSerializedList(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeList(serializedValue, valueSerializer);
	}
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
protected static <V, K, N> V getSerializedValue(
		InternalKvState<K, N, V> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<V> valueSerializer) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeValue(serializedValue, valueSerializer);
	}
}
 

/**
 * @see KeyedStateBackend
 */
@Override
@SuppressWarnings("unchecked")
public <N, S extends State, V> S getOrCreateKeyedState(
		final TypeSerializer<N> namespaceSerializer,
		StateDescriptor<S, V> stateDescriptor) throws Exception {
	checkNotNull(namespaceSerializer, "Namespace serializer");
	checkNotNull(keySerializer, "State key serializer has not been configured in the config. " +
			"This operation cannot use partitioned state.");

	InternalKvState<K, ?, ?> kvState = keyValueStatesByName.get(stateDescriptor.getName());
	if (kvState == null) {
		if (!stateDescriptor.isSerializerInitialized()) {
			stateDescriptor.initializeSerializerUnlessSet(executionConfig);
		}
		kvState = TtlStateFactory.createStateAndWrapWithTtlIfEnabled(
			namespaceSerializer, stateDescriptor, this, ttlTimeProvider);
		keyValueStatesByName.put(stateDescriptor.getName(), kvState);
		publishQueryableStateIfEnabled(stateDescriptor, kvState);
	}
	return (S) kvState;
}
 

/**
 * @see KeyedStateBackend
 */
@Override
@SuppressWarnings("unchecked")
public <N, S extends State, V> S getOrCreateKeyedState(
		final TypeSerializer<N> namespaceSerializer,
		StateDescriptor<S, V> stateDescriptor) throws Exception {
	checkNotNull(namespaceSerializer, "Namespace serializer");
	checkNotNull(keySerializerProvider, "State key serializer has not been configured in the config. " +
			"This operation cannot use partitioned state.");

	InternalKvState<K, ?, ?> kvState = keyValueStatesByName.get(stateDescriptor.getName());
	if (kvState == null) {
		if (!stateDescriptor.isSerializerInitialized()) {
			stateDescriptor.initializeSerializerUnlessSet(executionConfig);
		}
		kvState = TtlStateFactory.createStateAndWrapWithTtlIfEnabled(
			namespaceSerializer, stateDescriptor, this, ttlTimeProvider);
		keyValueStatesByName.put(stateDescriptor.getName(), kvState);
		publishQueryableStateIfEnabled(stateDescriptor, kvState);
	}
	return (S) kvState;
}
 

private boolean isStateIteratorSupported(InternalKvState<?, ?, ?> originalState, int size) {
	boolean stateIteratorSupported = false;
	try {
		stateIteratorSupported = originalState.getStateIncrementalVisitor(size) != null;
	} catch (Throwable t) {
		// ignore
	}
	return stateIteratorSupported;
}
 

private boolean isStateIteratorSupported(InternalKvState<?, ?, ?> originalState, int size) {
	boolean stateIteratorSupported = false;
	try {
		stateIteratorSupported = originalState.getStateIncrementalVisitor(size) != null;
	} catch (Throwable t) {
		// ignore
	}
	return stateIteratorSupported;
}
 

@SuppressWarnings("unchecked")
private <OIS extends State, TTLS extends State, V, TTLV> TtlStateContext<OIS, V>
	createTtlStateContext(StateDescriptor<TTLS, TTLV> ttlDescriptor) throws Exception {

	ttlDescriptor.enableTimeToLive(stateDesc.getTtlConfig()); // also used by RocksDB backend for TTL compaction filter config
	OIS originalState = (OIS) stateBackend.createInternalState(
		namespaceSerializer, ttlDescriptor, getSnapshotTransformFactory());
	return new TtlStateContext<>(
		originalState, ttlConfig, timeProvider, (TypeSerializer<V>) stateDesc.getSerializer(),
		registerTtlIncrementalCleanupCallback((InternalKvState<?, ?, ?>) originalState));
}
 

private Runnable registerTtlIncrementalCleanupCallback(InternalKvState<?, ?, ?> originalState) {
	StateTtlConfig.IncrementalCleanupStrategy config =
		ttlConfig.getCleanupStrategies().getIncrementalCleanupStrategy();
	boolean cleanupConfigured = config != null && incrementalCleanup != null;
	boolean isCleanupActive = cleanupConfigured &&
		isStateIteratorSupported(originalState, incrementalCleanup.getCleanupSize());
	Runnable callback = isCleanupActive ? incrementalCleanup::stateAccessed : () -> { };
	if (isCleanupActive && config.runCleanupForEveryRecord()) {
		stateBackend.registerKeySelectionListener(stub -> callback.run());
	}
	return callback;
}
 

/**
 * Registers the KvState instance and returns the assigned ID.
 *
 * @param jobId            JobId the KvState instance belongs to
 * @param jobVertexId      JobVertexID the KvState instance belongs to
 * @param keyGroupRange    Key group range the KvState instance belongs to
 * @param registrationName Name under which the KvState is registered
 * @param kvState          KvState instance to be registered
 * @return Assigned KvStateID
 */
public KvStateID registerKvState(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		InternalKvState<?, ?, ?> kvState) {

	KvStateID kvStateId = new KvStateID();

	if (registeredKvStates.putIfAbsent(kvStateId, new KvStateEntry<>(kvState)) == null) {
		final KvStateRegistryListener listener = getKvStateRegistryListener(jobId);

		if (listener != null) {
			listener.notifyKvStateRegistered(
				jobId,
				jobVertexId,
				keyGroupRange,
				registrationName,
				kvStateId);
		}

		return kvStateId;
	} else {
		throw new IllegalStateException(
				"State \"" + registrationName + " \"(id=" + kvStateId + ") appears registered although it should not.");
	}
}
 

/**
 * TODO: NOTE: This method does a lot of work caching / retrieving states just to update the namespace.
 *       This method should be removed for the sake of namespaces being lazily fetched from the keyed
 *       state backend, or being set on the state directly.
 *
 * @see KeyedStateBackend
 */
@SuppressWarnings("unchecked")
@Override
public <N, S extends State> S getPartitionedState(
		final N namespace,
		final TypeSerializer<N> namespaceSerializer,
		final StateDescriptor<S, ?> stateDescriptor) throws Exception {

	checkNotNull(namespace, "Namespace");

	if (lastName != null && lastName.equals(stateDescriptor.getName())) {
		lastState.setCurrentNamespace(namespace);
		return (S) lastState;
	}

	InternalKvState<K, ?, ?> previous = keyValueStatesByName.get(stateDescriptor.getName());
	if (previous != null) {
		lastState = previous;
		lastState.setCurrentNamespace(namespace);
		lastName = stateDescriptor.getName();
		return (S) previous;
	}

	final S state = getOrCreateKeyedState(namespaceSerializer, stateDescriptor);
	final InternalKvState<K, N, ?> kvState = (InternalKvState<K, N, ?>) state;

	lastName = stateDescriptor.getName();
	lastState = kvState;
	kvState.setCurrentNamespace(namespace);

	return state;
}
 

private void publishQueryableStateIfEnabled(
	StateDescriptor<?, ?> stateDescriptor,
	InternalKvState<?, ?, ?> kvState) {
	if (stateDescriptor.isQueryable()) {
		if (kvStateRegistry == null) {
			throw new IllegalStateException("State backend has not been initialized for job.");
		}
		String name = stateDescriptor.getQueryableStateName();
		kvStateRegistry.registerKvState(keyGroupRange, name, kvState);
	}
}
 

/**
 * TODO: NOTE: This method does a lot of work caching / retrieving states just to update the namespace.
 *       This method should be removed for the sake of namespaces being lazily fetched from the keyed
 *       state backend, or being set on the state directly.
 *
 * @see KeyedStateBackend
 */
@SuppressWarnings("unchecked")
@Override
public <N, S extends State> S getPartitionedState(
		final N namespace,
		final TypeSerializer<N> namespaceSerializer,
		final StateDescriptor<S, ?> stateDescriptor) throws Exception {

	checkNotNull(namespace, "Namespace");

	if (lastName != null && lastName.equals(stateDescriptor.getName())) {
		lastState.setCurrentNamespace(namespace);
		return (S) lastState;
	}

	InternalKvState<K, ?, ?> previous = keyValueStatesByName.get(stateDescriptor.getName());
	if (previous != null) {
		lastState = previous;
		lastState.setCurrentNamespace(namespace);
		lastName = stateDescriptor.getName();
		return (S) previous;
	}

	final S state = getOrCreateKeyedState(namespaceSerializer, stateDescriptor);
	final InternalKvState<K, N, ?> kvState = (InternalKvState<K, N, ?>) state;

	lastName = stateDescriptor.getName();
	lastState = kvState;
	kvState.setCurrentNamespace(namespace);

	return state;
}
 

@SuppressWarnings("unchecked")
private <OIS extends State, TTLS extends State, V, TTLV> TtlStateContext<OIS, V>
	createTtlStateContext(StateDescriptor<TTLS, TTLV> ttlDescriptor) throws Exception {

	ttlDescriptor.enableTimeToLive(stateDesc.getTtlConfig()); // also used by RocksDB backend for TTL compaction filter config
	OIS originalState = (OIS) stateBackend.createInternalState(
		namespaceSerializer, ttlDescriptor, getSnapshotTransformFactory());
	return new TtlStateContext<>(
		originalState, ttlConfig, timeProvider, (TypeSerializer<V>) stateDesc.getSerializer(),
		registerTtlIncrementalCleanupCallback((InternalKvState<?, ?, ?>) originalState));
}
 

private Runnable registerTtlIncrementalCleanupCallback(InternalKvState<?, ?, ?> originalState) {
	StateTtlConfig.IncrementalCleanupStrategy config =
		ttlConfig.getCleanupStrategies().getIncrementalCleanupStrategy();
	boolean cleanupConfigured = config != null && incrementalCleanup != null;
	boolean isCleanupActive = cleanupConfigured &&
		isStateIteratorSupported(originalState, incrementalCleanup.getCleanupSize());
	Runnable callback = isCleanupActive ? incrementalCleanup::stateAccessed : () -> { };
	if (isCleanupActive && config.runCleanupForEveryRecord()) {
		stateBackend.registerKeySelectionListener(stub -> callback.run());
	}
	return callback;
}
 

private boolean isStateIteratorSupported(InternalKvState<?, ?, ?> originalState, int size) {
	boolean stateIteratorSupported = false;
	try {
		stateIteratorSupported = originalState.getStateIncrementalVisitor(size) != null;
	} catch (Throwable t) {
		// ignore
	}
	return stateIteratorSupported;
}
 

private Runnable registerTtlIncrementalCleanupCallback(InternalKvState<?, ?, ?> originalState) {
	StateTtlConfig.IncrementalCleanupStrategy config =
		ttlConfig.getCleanupStrategies().getIncrementalCleanupStrategy();
	boolean cleanupConfigured = config != null && incrementalCleanup != null;
	boolean isCleanupActive = cleanupConfigured &&
		isStateIteratorSupported(originalState, incrementalCleanup.getCleanupSize());
	Runnable callback = isCleanupActive ? incrementalCleanup::stateAccessed : () -> { };
	if (isCleanupActive && config.runCleanupForEveryRecord()) {
		stateBackend.registerKeySelectionListener(stub -> callback.run());
	}
	return callback;
}
 

/**
 * Registers the KvState instance and returns the assigned ID.
 *
 * @param jobId            JobId the KvState instance belongs to
 * @param jobVertexId      JobVertexID the KvState instance belongs to
 * @param keyGroupRange    Key group range the KvState instance belongs to
 * @param registrationName Name under which the KvState is registered
 * @param kvState          KvState instance to be registered
 * @return Assigned KvStateID
 */
public KvStateID registerKvState(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		InternalKvState<?, ?, ?> kvState) {

	KvStateID kvStateId = new KvStateID();

	if (registeredKvStates.putIfAbsent(kvStateId, new KvStateEntry<>(kvState)) == null) {
		final KvStateRegistryListener listener = getKvStateRegistryListener(jobId);

		if (listener != null) {
			listener.notifyKvStateRegistered(
				jobId,
				jobVertexId,
				keyGroupRange,
				registrationName,
				kvStateId);
		}

		return kvStateId;
	} else {
		throw new IllegalStateException(
				"State \"" + registrationName + " \"(id=" + kvStateId + ") appears registered although it should not.");
	}
}
 

/**
 * Returns the value by getting the serialized value and deserializing it
 * if it is not null.
 */
private static <UK, UV, K, N> Map<UK, UV> getSerializedMap(
		InternalKvState<K, N, Map<UK, UV>> kvState,
		K key,
		TypeSerializer<K> keySerializer,
		N namespace,
		TypeSerializer<N> namespaceSerializer,
		TypeSerializer<UK> userKeySerializer,
		TypeSerializer<UV> userValueSerializer
) throws Exception {

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key, keySerializer, namespace, namespaceSerializer);

	byte[] serializedValue = kvState.getSerializedValue(
			serializedKeyAndNamespace,
			kvState.getKeySerializer(),
			kvState.getNamespaceSerializer(),
			kvState.getValueSerializer()
	);

	if (serializedValue == null) {
		return null;
	} else {
		return KvStateSerializer.deserializeMap(serializedValue, userKeySerializer, userValueSerializer);
	}
}
 

@SuppressWarnings("unchecked")
<N, S extends State, V> S createState(
	StateDescriptor<S, V> stateDescriptor,
	@SuppressWarnings("SameParameterValue") N defaultNamespace) throws Exception {
	S state = keyedStateBackend.getOrCreateKeyedState(StringSerializer.INSTANCE, stateDescriptor);
	((InternalKvState<?, N, ?>) state).setCurrentNamespace(defaultNamespace);
	return state;
}
 

@SuppressWarnings("unchecked")
private <OIS extends State, TTLS extends State, V, TTLV> TtlStateContext<OIS, V>
	createTtlStateContext(StateDescriptor<TTLS, TTLV> ttlDescriptor) throws Exception {

	ttlDescriptor.enableTimeToLive(stateDesc.getTtlConfig()); // also used by RocksDB backend for TTL compaction filter config
	OIS originalState = (OIS) stateBackend.createInternalState(
		namespaceSerializer, ttlDescriptor, getSnapshotTransformFactory());
	return new TtlStateContext<>(
		originalState, ttlConfig, timeProvider, (TypeSerializer<V>) stateDesc.getSerializer(),
		registerTtlIncrementalCleanupCallback((InternalKvState<?, ?, ?>) originalState));
}
 

@SuppressWarnings("unchecked")
private <S extends State> StateDescriptor<S, Object> createState() throws Exception {
	StateDescriptor<S, Object> stateDescriptor = ctx().createStateDescriptor();
	stateDescriptor.enableTimeToLive(ttlConfig);
	ctx().ttlState =
		(InternalKvState<?, String, ?>) sbetc.createState(stateDescriptor, "defaultNamespace");
	return stateDescriptor;
}
 

private static <K, N, V> byte[] getSerializedValue(
		final KvStateEntry<K, N, V> entry,
		final byte[] serializedKeyAndNamespace) throws Exception {

	final InternalKvState<K, N, V> state = entry.getState();
	final KvStateInfo<K, N, V> infoForCurrentThread = entry.getInfoForCurrentThread();

	return state.getSerializedValue(
			serializedKeyAndNamespace,
			infoForCurrentThread.getKeySerializer(),
			infoForCurrentThread.getNamespaceSerializer(),
			infoForCurrentThread.getStateValueSerializer()
	);
}
 

/**
 * TODO: NOTE: This method does a lot of work caching / retrieving states just to update the namespace.
 *       This method should be removed for the sake of namespaces being lazily fetched from the keyed
 *       state backend, or being set on the state directly.
 *
 * @see KeyedStateBackend
 */
@SuppressWarnings("unchecked")
@Override
public <N, S extends State> S getPartitionedState(
		final N namespace,
		final TypeSerializer<N> namespaceSerializer,
		final StateDescriptor<S, ?> stateDescriptor) throws Exception {

	checkNotNull(namespace, "Namespace");

	if (lastName != null && lastName.equals(stateDescriptor.getName())) {
		lastState.setCurrentNamespace(namespace);
		return (S) lastState;
	}

	InternalKvState<K, ?, ?> previous = keyValueStatesByName.get(stateDescriptor.getName());
	if (previous != null) {
		lastState = previous;
		lastState.setCurrentNamespace(namespace);
		lastName = stateDescriptor.getName();
		return (S) previous;
	}

	final S state = getOrCreateKeyedState(namespaceSerializer, stateDescriptor);
	final InternalKvState<K, N, ?> kvState = (InternalKvState<K, N, ?>) state;

	lastName = stateDescriptor.getName();
	lastState = kvState;
	kvState.setCurrentNamespace(namespace);

	return state;
}
 

private void publishQueryableStateIfEnabled(
	StateDescriptor<?, ?> stateDescriptor,
	InternalKvState<?, ?, ?> kvState) {
	if (stateDescriptor.isQueryable()) {
		if (kvStateRegistry == null) {
			throw new IllegalStateException("State backend has not been initialized for job.");
		}
		String name = stateDescriptor.getQueryableStateName();
		kvStateRegistry.registerKvState(keyGroupRange, name, kvState);
	}
}