org.apache.flink.runtime.state.KeyGroupRange Java Examples

static TaskStateSnapshot mockSubtaskState(
	JobVertexID jobVertexID,
	int index,
	KeyGroupRange keyGroupRange) throws IOException {

	OperatorStateHandle partitionableState = generatePartitionableStateHandle(jobVertexID, index, 2, 8, false);
	KeyGroupsStateHandle partitionedKeyGroupState = generateKeyGroupState(jobVertexID, keyGroupRange, false);

	TaskStateSnapshot subtaskStates = spy(new TaskStateSnapshot());
	OperatorSubtaskState subtaskState = spy(new OperatorSubtaskState(
		partitionableState, null, partitionedKeyGroupState, null, null, null)
	);

	subtaskStates.putSubtaskStateByOperatorID(OperatorID.fromJobVertexID(jobVertexID), subtaskState);

	return subtaskStates;
}
 

public <K> HeapKeyedStateBackend<K> createKeyedBackend(
	TypeSerializer<K> keySerializer,
	Collection<KeyedStateHandle> stateHandles) throws Exception {
	final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 15);
	final int numKeyGroups = keyGroupRange.getNumberOfKeyGroups();
	ExecutionConfig executionConfig = new ExecutionConfig();

	return new HeapKeyedStateBackendBuilder<>(
		mock(TaskKvStateRegistry.class),
		keySerializer,
		HeapStateBackendTestBase.class.getClassLoader(),
		numKeyGroups,
		keyGroupRange,
		executionConfig,
		TtlTimeProvider.DEFAULT,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(executionConfig),
		TestLocalRecoveryConfig.disabled(),
		new HeapPriorityQueueSetFactory(keyGroupRange, numKeyGroups, 128),
		async,
		new CloseableRegistry()).build();
}
 

/**
 * Determine the subset of {@link KeyGroupsStateHandle KeyGroupsStateHandles} with correct
 * key group index for the given subtask {@link KeyGroupRange}.
 *
 * <p>This is publicly visible to be used in tests.
 */
public static List<KeyedStateHandle> getKeyedStateHandles(
	Collection<? extends KeyedStateHandle> keyedStateHandles,
	KeyGroupRange subtaskKeyGroupRange) {

	List<KeyedStateHandle> subtaskKeyedStateHandles = new ArrayList<>(keyedStateHandles.size());

	for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
		KeyedStateHandle intersectedKeyedStateHandle = keyedStateHandle.getIntersection(subtaskKeyGroupRange);

		if (intersectedKeyedStateHandle != null) {
			subtaskKeyedStateHandles.add(intersectedKeyedStateHandle);
		}
	}

	return subtaskKeyedStateHandles;
}
 

HeapRestoreOperation(
	@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
	StateSerializerProvider<K> keySerializerProvider,
	ClassLoader userCodeClassLoader,
	Map<String, StateTable<K, ?, ?>> registeredKVStates,
	Map<String, HeapPriorityQueueSnapshotRestoreWrapper> registeredPQStates,
	CloseableRegistry cancelStreamRegistry,
	HeapPriorityQueueSetFactory priorityQueueSetFactory,
	@Nonnull KeyGroupRange keyGroupRange,
	int numberOfKeyGroups,
	HeapSnapshotStrategy<K> snapshotStrategy,
	InternalKeyContext<K> keyContext) {
	this.restoreStateHandles = restoreStateHandles;
	this.keySerializerProvider = keySerializerProvider;
	this.userCodeClassLoader = userCodeClassLoader;
	this.registeredKVStates = registeredKVStates;
	this.registeredPQStates = registeredPQStates;
	this.cancelStreamRegistry = cancelStreamRegistry;
	this.priorityQueueSetFactory = priorityQueueSetFactory;
	this.keyGroupRange = keyGroupRange;
	this.numberOfKeyGroups = numberOfKeyGroups;
	this.snapshotStrategy = snapshotStrategy;
	this.keyContext = keyContext;
}
 

static TaskStateSnapshot mockSubtaskState(
	JobVertexID jobVertexID,
	int index,
	KeyGroupRange keyGroupRange) throws IOException {

	OperatorStateHandle partitionableState = generatePartitionableStateHandle(jobVertexID, index, 2, 8, false);
	KeyGroupsStateHandle partitionedKeyGroupState = generateKeyGroupState(jobVertexID, keyGroupRange, false);

	TaskStateSnapshot subtaskStates = spy(new TaskStateSnapshot());
	OperatorSubtaskState subtaskState = spy(new OperatorSubtaskState(
		partitionableState, null, partitionedKeyGroupState, null)
	);

	subtaskStates.putSubtaskStateByOperatorID(OperatorID.fromJobVertexID(jobVertexID), subtaskState);

	return subtaskStates;
}
 

public MockKeyedStateBackendBuilder(
	TaskKvStateRegistry kvStateRegistry,
	TypeSerializer<K> keySerializer,
	ClassLoader userCodeClassLoader,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	ExecutionConfig executionConfig,
	TtlTimeProvider ttlTimeProvider,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	StreamCompressionDecorator keyGroupCompressionDecorator,
	CloseableRegistry cancelStreamRegistry) {
	super(
		kvStateRegistry,
		keySerializer,
		userCodeClassLoader,
		numberOfKeyGroups,
		keyGroupRange,
		executionConfig,
		ttlTimeProvider,
		stateHandles,
		keyGroupCompressionDecorator,
		cancelStreamRegistry);
}
 

/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange  Key group range to register
 * @param kvStateId      ID of the KvState instance at the key group index.
 * @param kvStateAddress Server address of the KvState instance at the key group index.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 */
public void registerKvState(KeyGroupRange keyGroupRange, KvStateID kvStateId, InetSocketAddress kvStateAddress) {

	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {

		if (kvStateIds[kgIdx] == null && kvStateAddresses[kgIdx] == null) {
			numRegisteredKeyGroups++;
		}

		kvStateIds[kgIdx] = kvStateId;
		kvStateAddresses[kgIdx] = kvStateAddress;
	}
}
 

public RocksFullSnapshotStrategy(
	@Nonnull RocksDB db,
	@Nonnull ResourceGuard rocksDBResourceGuard,
	@Nonnull TypeSerializer<K> keySerializer,
	@Nonnull LinkedHashMap<String, RocksDbKvStateInfo> kvStateInformation,
	@Nonnull KeyGroupRange keyGroupRange,
	@Nonnegative int keyGroupPrefixBytes,
	@Nonnull LocalRecoveryConfig localRecoveryConfig,
	@Nonnull CloseableRegistry cancelStreamRegistry,
	@Nonnull StreamCompressionDecorator keyGroupCompressionDecorator) {
	super(
		DESCRIPTION,
		db,
		rocksDBResourceGuard,
		keySerializer,
		kvStateInformation,
		keyGroupRange,
		keyGroupPrefixBytes,
		localRecoveryConfig,
		cancelStreamRegistry);

	this.keyGroupCompressionDecorator = keyGroupCompressionDecorator;
}
 

/**
 * Creates an empty {@link HeapPriorityQueueSet} with the requested initial capacity.
 *
 * @param elementPriorityComparator comparator for the priority of contained elements.
 * @param keyExtractor function to extract a key from the contained elements.
 * @param minimumCapacity the minimum and initial capacity of this priority queue.
 * @param keyGroupRange the key-group range of the elements in this set.
 * @param totalNumberOfKeyGroups the total number of key-groups of the job.
 */
@SuppressWarnings("unchecked")
public HeapPriorityQueueSet(
	@Nonnull PriorityComparator<T> elementPriorityComparator,
	@Nonnull KeyExtractorFunction<T> keyExtractor,
	@Nonnegative int minimumCapacity,
	@Nonnull KeyGroupRange keyGroupRange,
	@Nonnegative int totalNumberOfKeyGroups) {

	super(elementPriorityComparator, minimumCapacity);

	this.keyExtractor = keyExtractor;

	this.totalNumberOfKeyGroups = totalNumberOfKeyGroups;
	this.keyGroupRange = keyGroupRange;

	final int keyGroupsInLocalRange = keyGroupRange.getNumberOfKeyGroups();
	final int deduplicationSetSize = 1 + minimumCapacity / keyGroupsInLocalRange;
	this.deduplicationMapsByKeyGroup = new HashMap[keyGroupsInLocalRange];
	for (int i = 0; i < keyGroupsInLocalRange; ++i) {
		deduplicationMapsByKeyGroup[i] = new HashMap<>(deduplicationSetSize);
	}
}
 

@VisibleForTesting
public static KeyedStateHandle deserializeKeyedStateHandle(DataInputStream dis) throws IOException {
	final int type = dis.readByte();
	if (NULL_HANDLE == type) {
		return null;
	} else if (KEY_GROUPS_HANDLE == type) {
		int startKeyGroup = dis.readInt();
		int numKeyGroups = dis.readInt();
		KeyGroupRange keyGroupRange = KeyGroupRange.of(startKeyGroup, startKeyGroup + numKeyGroups - 1);
		long[] offsets = new long[numKeyGroups];
		for (int i = 0; i < numKeyGroups; ++i) {
			offsets[i] = dis.readLong();
		}
		KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(
			keyGroupRange, offsets);
		StreamStateHandle stateHandle = deserializeStreamStateHandle(dis);
		return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
	} else {
		throw new IllegalStateException("Reading invalid KeyedStateHandle, type: " + type);
	}
}
 

public RocksDBSnapshotStrategyBase(
	@Nonnull String description,
	@Nonnull RocksDB db,
	@Nonnull ResourceGuard rocksDBResourceGuard,
	@Nonnull TypeSerializer<K> keySerializer,
	@Nonnull LinkedHashMap<String, RocksDbKvStateInfo> kvStateInformation,
	@Nonnull KeyGroupRange keyGroupRange,
	@Nonnegative int keyGroupPrefixBytes,
	@Nonnull LocalRecoveryConfig localRecoveryConfig,
	@Nonnull CloseableRegistry cancelStreamRegistry) {

	super(description);
	this.db = db;
	this.rocksDBResourceGuard = rocksDBResourceGuard;
	this.keySerializer = keySerializer;
	this.kvStateInformation = kvStateInformation;
	this.keyGroupRange = keyGroupRange;
	this.keyGroupPrefixBytes = keyGroupPrefixBytes;
	this.localRecoveryConfig = localRecoveryConfig;
	this.cancelStreamRegistry = cancelStreamRegistry;
}
 

public RocksDBSnapshotStrategyBase(
	@Nonnull String description,
	@Nonnull RocksDB db,
	@Nonnull ResourceGuard rocksDBResourceGuard,
	@Nonnull TypeSerializer<K> keySerializer,
	@Nonnull LinkedHashMap<String, RocksDbKvStateInfo> kvStateInformation,
	@Nonnull KeyGroupRange keyGroupRange,
	@Nonnegative int keyGroupPrefixBytes,
	@Nonnull LocalRecoveryConfig localRecoveryConfig,
	@Nonnull CloseableRegistry cancelStreamRegistry) {

	super(description);
	this.db = db;
	this.rocksDBResourceGuard = rocksDBResourceGuard;
	this.keySerializer = keySerializer;
	this.kvStateInformation = kvStateInformation;
	this.keyGroupRange = keyGroupRange;
	this.keyGroupPrefixBytes = keyGroupPrefixBytes;
	this.localRecoveryConfig = localRecoveryConfig;
	this.cancelStreamRegistry = cancelStreamRegistry;
}
 

/**
 * Notifies the registry about an unregistered KvState instance.
 *
 * @param jobVertexId JobVertexID the KvState instance belongs to
 * @param keyGroupRange Key group index the KvState instance belongs to
 * @param registrationName Name under which the KvState has been registered
 * @throws IllegalArgumentException If another operator registered the state instance
 * @throws IllegalArgumentException If the registration name is not known
 */
public void notifyKvStateUnregistered(
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName) {

	KvStateLocation location = lookupTable.get(registrationName);

	if (location != null) {
		// Duplicate name if vertex IDs don't match
		if (!location.getJobVertexId().equals(jobVertexId)) {
			throw new IllegalArgumentException("Another operator (" +
					location.getJobVertexId() + ") registered the KvState " +
					"under '" + registrationName + "'.");
		}

		location.unregisterKvState(keyGroupRange);

		if (location.getNumRegisteredKeyGroups() == 0) {
			lookupTable.remove(registrationName);
		}
	} else {
		throw new IllegalArgumentException("Unknown registration name '" +
				registrationName + "'. " + "Probably registration/unregistration race.");
	}
}
 

/**
 * Determine the subset of {@link KeyGroupsStateHandle KeyGroupsStateHandles} with correct
 * key group index for the given subtask {@link KeyGroupRange}.
 *
 * <p>This is publicly visible to be used in tests.
 */
public static List<KeyedStateHandle> getKeyedStateHandles(
	Collection<? extends KeyedStateHandle> keyedStateHandles,
	KeyGroupRange subtaskKeyGroupRange) {

	List<KeyedStateHandle> subtaskKeyedStateHandles = new ArrayList<>(keyedStateHandles.size());

	for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
		KeyedStateHandle intersectedKeyedStateHandle = keyedStateHandle.getIntersection(subtaskKeyGroupRange);

		if (intersectedKeyedStateHandle != null) {
			subtaskKeyedStateHandles.add(intersectedKeyedStateHandle);
		}
	}

	return subtaskKeyedStateHandles;
}
 

private static <K, N> InternalTimerServiceImpl<K, N> createInternalTimerService(
	KeyGroupRange keyGroupsList,
	KeyContext keyContext,
	ProcessingTimeService processingTimeService,
	TypeSerializer<K> keySerializer,
	TypeSerializer<N> namespaceSerializer,
	PriorityQueueSetFactory priorityQueueSetFactory) {

	TimerSerializer<K, N> timerSerializer = new TimerSerializer<>(keySerializer, namespaceSerializer);

	return new InternalTimerServiceImpl<>(
		keyGroupsList,
		keyContext,
		processingTimeService,
		createTimerQueue("__test_processing_timers", timerSerializer, priorityQueueSetFactory),
		createTimerQueue("__test_event_timers", timerSerializer, priorityQueueSetFactory));
}
 

@Override
public void notifyKvStateRegistered(
	JobID jobId,
	JobVertexID jobVertexId,
	KeyGroupRange keyGroupRange,
	String registrationName,
	KvStateID kvStateId) {

	Object msg = new KvStateMessage.NotifyKvStateRegistered(
		jobId,
		jobVertexId,
		keyGroupRange,
		registrationName,
		kvStateId,
		kvStateServerAddress);

	jobManager.tell(msg);
}
 

/**
 * Choose the best state handle according to the {@link #STATE_HANDLE_EVALUATOR}
 * to init the initial db.
 *
 * @param restoreStateHandles The candidate state handles.
 * @param targetKeyGroupRange The target key group range.
 * @return The best candidate or null if no candidate was a good fit.
 */
@Nullable
public static KeyedStateHandle chooseTheBestStateHandleForInitial(
	@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
	@Nonnull KeyGroupRange targetKeyGroupRange) {

	KeyedStateHandle bestStateHandle = null;
	double bestScore = 0;
	for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
		double handleScore = STATE_HANDLE_EVALUATOR.apply(rawStateHandle, targetKeyGroupRange);
		if (handleScore > bestScore) {
			bestStateHandle = rawStateHandle;
			bestScore = handleScore;
		}
	}

	return bestStateHandle;
}
 

public static <K> RocksDBKeyedStateBackend<K> createKeyedStateBackend(
		RocksDBStateBackend rocksDbBackend,
		Environment env,
		TypeSerializer<K> keySerializer) throws IOException {

	return (RocksDBKeyedStateBackend<K>) rocksDbBackend.createKeyedStateBackend(
		env,
		env.getJobID(),
		"test_op",
		keySerializer,
		1,
		new KeyGroupRange(0, 0),
		env.getTaskKvStateRegistry(),
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		new CloseableRegistry());
}
 

@Test
public void testKeyGroupStartIndexSetting() {

	int startKeyGroupIdx = 7;
	int endKeyGroupIdx = 21;
	KeyGroupRange testKeyGroupList = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);

	TestKeyContext keyContext = new TestKeyContext();

	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();

	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		testKeyGroupList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		createQueueFactory());

	Assert.assertEquals(startKeyGroupIdx, service.getLocalKeyGroupRangeStartIdx());
}
 

MockKeyedStateBackend(
	TaskKvStateRegistry kvStateRegistry,
	TypeSerializer<K> keySerializer,
	ClassLoader userCodeClassLoader,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	ExecutionConfig executionConfig,
	TtlTimeProvider ttlTimeProvider,
	Map<String, Map<K, Map<Object, Object>>> stateValues,
	Map<String, StateSnapshotTransformer<Object>> stateSnapshotFilters,
	CloseableRegistry cancelStreamRegistry) {
	super(kvStateRegistry, keySerializer, userCodeClassLoader,
		numberOfKeyGroups, keyGroupRange, executionConfig, ttlTimeProvider, cancelStreamRegistry);
	this.stateValues = stateValues;
	this.stateSnapshotFilters = stateSnapshotFilters;
}
 

private KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<
		TestElement, KeyGroupHeapPQSet<TestElement>> newFactory(int initialCapacity) {

	return (keyGroupId, numKeyGroups, keyExtractorFunction, elementComparator) ->
		new KeyGroupHeapPQSet<>(
			elementComparator,
			keyExtractorFunction,
			initialCapacity,
			KeyGroupRange.of(keyGroupId, keyGroupId),
			numKeyGroups);
}
 

private static int getKeyInKeyGroupRange(KeyGroupRange range, int maxParallelism) {
	Random rand = new Random(System.currentTimeMillis());
	int result = rand.nextInt();
	while (!range.contains(KeyGroupRangeAssignment.assignToKeyGroup(result, maxParallelism))) {
		result = rand.nextInt();
	}
	return result;
}
 

private static InternalTimerServiceImpl<Integer, String> restoreTimerService(
		Map<Integer, byte[]> state,
		int snapshotVersion,
		Triggerable<Integer, String> triggerable,
		KeyContext keyContext,
		ProcessingTimeService processingTimeService,
		KeyGroupRange keyGroupsList,
		PriorityQueueSetFactory priorityQueueSetFactory) throws Exception {

	// create an empty service
	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		keyGroupsList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	// restore the timers
	for (Integer keyGroupIndex : keyGroupsList) {
		if (state.containsKey(keyGroupIndex)) {
			try (ByteArrayInputStream inputStream = new ByteArrayInputStream(state.get(keyGroupIndex))) {
				InternalTimersSnapshot<?, ?> restoredTimersSnapshot =
					InternalTimersSnapshotReaderWriters
						.getReaderForVersion(snapshotVersion, InternalTimerServiceImplTest.class.getClassLoader())
						.readTimersSnapshot(new DataInputViewStreamWrapper(inputStream));

				service.restoreTimersForKeyGroup(restoredTimersSnapshot, keyGroupIndex);
			}
		}
	}

	// initialize the service
	service.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, triggerable);
	return service;
}
 

public KeyGroupHeapPQSet(
	@Nonnull PriorityComparator<T> elementPriorityComparator,
	@Nonnull KeyExtractorFunction<T> keyExtractor,
	int minimumCapacity,
	@Nonnull KeyGroupRange keyGroupRange,
	int totalNumberOfKeyGroups) {
	super(elementPriorityComparator, keyExtractor, minimumCapacity, keyGroupRange, totalNumberOfKeyGroups);
	this.internalIndex = HeapPriorityQueueElement.NOT_CONTAINED;
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {

	TaskStateManager taskStateManager = env.getTaskStateManager();
	LocalRecoveryConfig localRecoveryConfig = taskStateManager.createLocalRecoveryConfig();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);

	return new HeapKeyedStateBackendBuilder<>(
		kvStateRegistry,
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		localRecoveryConfig,
		priorityQueueSetFactory,
		isUsingAsynchronousSnapshots(),
		cancelStreamRegistry).build();
}
 

@Override
public void notifyKvStateRegistered(JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId) {
	this.jobVertexID = jobVertexId;
	this.keyGroupIndex = keyGroupRange;
	this.registrationName = registrationName;
	this.kvStateId = kvStateId;
}
 

@Override
public void notifyKvStateUnregistered(JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName) {

}
 

InternalTimeServiceManager(
	KeyGroupRange localKeyGroupRange,
	KeyContext keyContext,
	PriorityQueueSetFactory priorityQueueSetFactory,
	ProcessingTimeService processingTimeService, boolean useLegacySynchronousSnapshots) {

	this.localKeyGroupRange = Preconditions.checkNotNull(localKeyGroupRange);
	this.priorityQueueSetFactory = Preconditions.checkNotNull(priorityQueueSetFactory);
	this.keyContext = Preconditions.checkNotNull(keyContext);
	this.processingTimeService = Preconditions.checkNotNull(processingTimeService);
	this.useLegacySynchronousSnapshots = useLegacySynchronousSnapshots;

	this.timerServices = new HashMap<>();
}
 

private KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<
		TestElement, KeyGroupHeapPQSet<TestElement>> newFactory(int initialCapacity) {

	return (keyGroupId, numKeyGroups, keyExtractorFunction, elementComparator) ->
		new KeyGroupHeapPQSet<>(
			elementComparator,
			keyExtractorFunction,
			initialCapacity,
			KeyGroupRange.of(keyGroupId, keyGroupId),
			numKeyGroups);
}
 

protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
	AbstractKeyedStateBackend<K> keyedStatedBackend,
	KeyContext keyContext, //the operator
	Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {

	if (keyedStatedBackend == null) {
		return null;
	}

	final KeyGroupRange keyGroupRange = keyedStatedBackend.getKeyGroupRange();

	final InternalTimeServiceManager<K> timeServiceManager = new InternalTimeServiceManager<>(
		keyGroupRange,
		keyContext,
		keyedStatedBackend,
		processingTimeService,
		keyedStatedBackend.requiresLegacySynchronousTimerSnapshots());

	// and then initialize the timer services
	for (KeyGroupStatePartitionStreamProvider streamProvider : rawKeyedStates) {
		int keyGroupIdx = streamProvider.getKeyGroupId();

		Preconditions.checkArgument(keyGroupRange.contains(keyGroupIdx),
			"Key Group " + keyGroupIdx + " does not belong to the local range.");

		timeServiceManager.restoreStateForKeyGroup(
			streamProvider.getStream(),
			keyGroupIdx, environment.getUserClassLoader());
	}

	return timeServiceManager;
}