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

@Override
public KeyGroupStatePartitionStreamProvider next() {

	if (!hasNext()) {

		throw new NoSuchElementException("Iterator exhausted");
	}

	Tuple2<Integer, Long> keyGroupOffset = currentOffsetsIterator.next();
	try {
		if (null == currentStream) {
			openCurrentStream();
		}

		currentStream.seek(keyGroupOffset.f1);
		return new KeyGroupStatePartitionStreamProvider(currentStream, keyGroupOffset.f0);

	} catch (IOException ioex) {
		return new KeyGroupStatePartitionStreamProvider(ioex, keyGroupOffset.f0);
	}
}
 

@Override
public KeyGroupStatePartitionStreamProvider next() {

	if (!hasNext()) {

		throw new NoSuchElementException("Iterator exhausted");
	}

	Tuple2<Integer, Long> keyGroupOffset = currentOffsetsIterator.next();
	try {
		if (null == currentStream) {
			openCurrentStream();
		}

		currentStream.seek(keyGroupOffset.f1);
		return new KeyGroupStatePartitionStreamProvider(currentStream, keyGroupOffset.f0);

	} catch (IOException ioex) {
		return new KeyGroupStatePartitionStreamProvider(ioex, keyGroupOffset.f0);
	}
}
 

@Test
public void getKeyedStateStreams() throws Exception {

	int readKeyGroupCount = 0;

	for (KeyGroupStatePartitionStreamProvider stateStreamProvider
			: initializationContext.getRawKeyedStateInputs()) {

		Assert.assertNotNull(stateStreamProvider);

		try (InputStream is = stateStreamProvider.getStream()) {
			DataInputView div = new DataInputViewStreamWrapper(is);
			int val = div.readInt();
			++readKeyGroupCount;
			Assert.assertEquals(stateStreamProvider.getKeyGroupId(), val);
		}
	}

	Assert.assertEquals(writtenKeyGroups, readKeyGroupCount);
}
 

public void initializeOperatorState(CheckpointedStreamOperator streamOperator) throws Exception {
	CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = context.rawKeyedStateInputs();
	CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = context.rawOperatorStateInputs();

	try {
		StateInitializationContext initializationContext = new StateInitializationContextImpl(
			context.isRestored(), // information whether we restore or start for the first time
			operatorStateBackend, // access to operator state backend
			keyedStateStore, // access to keyed state backend
			keyedStateInputs, // access to keyed state stream
			operatorStateInputs); // access to operator state stream

		streamOperator.initializeState(initializationContext);
	} finally {
		closeFromRegistry(operatorStateInputs, closeableRegistry);
		closeFromRegistry(keyedStateInputs, closeableRegistry);
	}
}
 

@Test
public void getKeyedStateStreams() throws Exception {

	int readKeyGroupCount = 0;

	for (KeyGroupStatePartitionStreamProvider stateStreamProvider
			: initializationContext.getRawKeyedStateInputs()) {

		Assert.assertNotNull(stateStreamProvider);

		try (InputStream is = stateStreamProvider.getStream()) {
			DataInputView div = new DataInputViewStreamWrapper(is);
			int val = div.readInt();
			++readKeyGroupCount;
			Assert.assertEquals(stateStreamProvider.getKeyGroupId(), val);
		}
	}

	Assert.assertEquals(writtenKeyGroups, readKeyGroupCount);
}
 

@Override
public KeyGroupStatePartitionStreamProvider next() {

	if (!hasNext()) {

		throw new NoSuchElementException("Iterator exhausted");
	}

	Tuple2<Integer, Long> keyGroupOffset = currentOffsetsIterator.next();
	try {
		if (null == currentStream) {
			openCurrentStream();
		}

		currentStream.seek(keyGroupOffset.f1);
		return new KeyGroupStatePartitionStreamProvider(currentStream, keyGroupOffset.f0);

	} catch (IOException ioex) {
		return new KeyGroupStatePartitionStreamProvider(ioex, keyGroupOffset.f0);
	}
}
 

@Test
public void getKeyedStateStreams() throws Exception {

	int readKeyGroupCount = 0;

	for (KeyGroupStatePartitionStreamProvider stateStreamProvider
			: initializationContext.getRawKeyedStateInputs()) {

		Assert.assertNotNull(stateStreamProvider);

		try (InputStream is = stateStreamProvider.getStream()) {
			DataInputView div = new DataInputViewStreamWrapper(is);
			int val = div.readInt();
			++readKeyGroupCount;
			Assert.assertEquals(stateStreamProvider.getKeyGroupId(), val);
		}
	}

	Assert.assertEquals(writtenKeyGroups, readKeyGroupCount);
}
 

protected CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs(
	Iterator<StateObjectCollection<KeyedStateHandle>> restoreStateAlternatives) {

	if (restoreStateAlternatives.hasNext()) {
		Collection<KeyedStateHandle> rawKeyedState = restoreStateAlternatives.next();

		// TODO currently this does not support local state recovery, so we expect there is only one handle.
		Preconditions.checkState(
			!restoreStateAlternatives.hasNext(),
			"Local recovery is currently not implemented for raw keyed state, but found state alternative.");

		if (rawKeyedState != null) {
			Collection<KeyGroupsStateHandle> keyGroupsStateHandles = transform(rawKeyedState);
			final CloseableRegistry closeableRegistry = new CloseableRegistry();

			return new CloseableIterable<KeyGroupStatePartitionStreamProvider>() {
				@Override
				public void close() throws IOException {
					closeableRegistry.close();
				}

				@Override
				public Iterator<KeyGroupStatePartitionStreamProvider> iterator() {
					return new KeyGroupStreamIterator(keyGroupsStateHandles.iterator(), closeableRegistry);
				}
			};
		}
	}

	return CloseableIterable.empty();
}
 

StreamOperatorStateContextImpl(
	boolean restored,
	OperatorStateBackend operatorStateBackend,
	AbstractKeyedStateBackend<?> keyedStateBackend,
	InternalTimeServiceManager<?> internalTimeServiceManager,
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs,
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs) {

	this.restored = restored;
	this.operatorStateBackend = operatorStateBackend;
	this.keyedStateBackend = keyedStateBackend;
	this.internalTimeServiceManager = internalTimeServiceManager;
	this.rawOperatorStateInputs = rawOperatorStateInputs;
	this.rawKeyedStateInputs = rawKeyedStateInputs;
}
 

@Test
public void close() throws Exception {

	int count = 0;
	int stopCount = NUM_HANDLES / 2;
	boolean isClosed = false;

	try {
		for (KeyGroupStatePartitionStreamProvider stateStreamProvider
				: initializationContext.getRawKeyedStateInputs()) {
			Assert.assertNotNull(stateStreamProvider);

			if (count == stopCount) {
				closableRegistry.close();
				isClosed = true;
			}

			try (InputStream is = stateStreamProvider.getStream()) {
				DataInputView div = new DataInputViewStreamWrapper(is);
				try {
					int val = div.readInt();
					Assert.assertEquals(stateStreamProvider.getKeyGroupId(), val);
					if (isClosed) {
						Assert.fail("Close was ignored: stream");
					}
					++count;
				} catch (IOException ioex) {
					if (!isClosed) {
						throw ioex;
					}
				}
			}
		}
		Assert.fail("Close was ignored: registry");
	} catch (IOException iex) {
		Assert.assertTrue(isClosed);
		Assert.assertEquals(stopCount, count);
	}

}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	keyedStateBackend = (AbstractKeyedStateBackend<?>) getKeyedStateBackend();
	operatorStateBackend = getOperatorStateBackend();
	rawOperatorStateInputs =
		(CloseableIterable<StatePartitionStreamProvider>) context.getRawOperatorStateInputs();
	rawKeyedStateInputs =
		(CloseableIterable<KeyGroupStatePartitionStreamProvider>) context.getRawKeyedStateInputs();
	super.initializeState(context);
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
  super.initializeState(context);

  final IOManager ioManager = getContainingTask().getEnvironment().getIOManager();
  final int maxParallelism = getRuntimeContext().getMaxNumberOfParallelSubtasks();

  this.reusable = new StreamRecord<>(null);

  //
  // Initialize the unbounded feedback logger
  //
  @SuppressWarnings("unchecked")
  UnboundedFeedbackLogger<T> feedbackLogger =
      (UnboundedFeedbackLogger<T>)
          Loggers.unboundedSpillableLogger(
              ioManager,
              maxParallelism,
              totalMemoryUsedForFeedbackCheckpointing,
              elementSerializer,
              keySelector);

  this.feedbackLogger = feedbackLogger;
  //
  // we first must reply previously check-pointed envelopes before we start
  // processing any new envelopes.
  //
  for (KeyGroupStatePartitionStreamProvider keyedStateInput : context.getRawKeyedStateInputs()) {
    this.feedbackLogger.replyLoggedEnvelops(keyedStateInput.getStream(), this);
  }
  //
  // now we can start processing new messages. We do so by registering ourselves as a
  // FeedbackConsumer
  //
  registerFeedbackConsumer(new MailboxExecutorFacade(mailboxExecutor, "Feedback Consumer"));
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
  super.initializeState(context);

  final IOManager ioManager = getContainingTask().getEnvironment().getIOManager();
  final int maxParallelism = getRuntimeContext().getMaxNumberOfParallelSubtasks();

  this.reusable = new StreamRecord<>(null);

  //
  // Initialize the unbounded feedback logger
  //
  @SuppressWarnings("unchecked")
  UnboundedFeedbackLoggerFactory<T> feedbackLoggerFactory =
      (UnboundedFeedbackLoggerFactory<T>)
          Loggers.unboundedSpillableLoggerFactory(
              ioManager,
              maxParallelism,
              totalMemoryUsedForFeedbackCheckpointing,
              elementSerializer,
              keySelector);

  this.checkpoints = new Checkpoints<>(feedbackLoggerFactory::create);

  //
  // we first must reply previously check-pointed envelopes before we start
  // processing any new envelopes.
  //
  UnboundedFeedbackLogger<T> logger = feedbackLoggerFactory.create();
  for (KeyGroupStatePartitionStreamProvider keyedStateInput : context.getRawKeyedStateInputs()) {
    logger.replyLoggedEnvelops(keyedStateInput.getStream(), this);
  }
  //
  // now we can start processing new messages. We do so by registering ourselves as a
  // FeedbackConsumer
  //
  registerFeedbackConsumer(new MailboxExecutorFacade(mailboxExecutor, "Feedback Consumer"));
}
 

protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
	AbstractKeyedStateBackend<K> keyedStatedBackend,
	KeyContext keyContext, //the operator
	ProcessingTimeService processingTimeService,
	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;
}
 

@Override
public void initializeState(StateInitializationContext controller) throws Exception {
	keyedStateBackend = (AbstractKeyedStateBackend<?>) getKeyedStateBackend();
	operatorStateBackend = getOperatorStateBackend();
	rawOperatorStateInputs =
		(CloseableIterable<StatePartitionStreamProvider>) controller.getRawOperatorStateInputs();
	rawKeyedStateInputs =
		(CloseableIterable<KeyGroupStatePartitionStreamProvider>) controller.getRawKeyedStateInputs();
	super.initializeState(controller);
}
 

StreamOperatorStateContextImpl(
	boolean restored,
	OperatorStateBackend operatorStateBackend,
	AbstractKeyedStateBackend<?> keyedStateBackend,
	InternalTimeServiceManager<?> internalTimeServiceManager,
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs,
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs) {

	this.restored = restored;
	this.operatorStateBackend = operatorStateBackend;
	this.keyedStateBackend = keyedStateBackend;
	this.internalTimeServiceManager = internalTimeServiceManager;
	this.rawOperatorStateInputs = rawOperatorStateInputs;
	this.rawKeyedStateInputs = rawKeyedStateInputs;
}
 

@Test
public void close() throws Exception {

	int count = 0;
	int stopCount = NUM_HANDLES / 2;
	boolean isClosed = false;

	try {
		for (KeyGroupStatePartitionStreamProvider stateStreamProvider
				: initializationContext.getRawKeyedStateInputs()) {
			Assert.assertNotNull(stateStreamProvider);

			if (count == stopCount) {
				closableRegistry.close();
				isClosed = true;
			}

			try (InputStream is = stateStreamProvider.getStream()) {
				DataInputView div = new DataInputViewStreamWrapper(is);
				try {
					int val = div.readInt();
					Assert.assertEquals(stateStreamProvider.getKeyGroupId(), val);
					if (isClosed) {
						Assert.fail("Close was ignored: stream");
					}
					++count;
				} catch (IOException ioex) {
					if (!isClosed) {
						throw ioex;
					}
				}
			}
		}
		Assert.fail("Close was ignored: registry");
	} catch (IOException iex) {
		Assert.assertTrue(isClosed);
		Assert.assertEquals(stopCount, count);
	}

}
 

private StreamTaskStateInitializer streamTaskStateManager(
	StateBackend stateBackend,
	JobManagerTaskRestore jobManagerTaskRestore,
	boolean createTimerServiceManager) {

	JobID jobID = new JobID(42L, 43L);
	ExecutionAttemptID executionAttemptID = new ExecutionAttemptID(23L, 24L);
	TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();

	TaskLocalStateStore taskLocalStateStore = new TestTaskLocalStateStore();

	TaskStateManager taskStateManager = TaskStateManagerImplTest.taskStateManager(
		jobID,
		executionAttemptID,
		checkpointResponderMock,
		jobManagerTaskRestore,
		taskLocalStateStore);

	DummyEnvironment dummyEnvironment = new DummyEnvironment("test-task", 1, 0);
	dummyEnvironment.setTaskStateManager(taskStateManager);

	if (createTimerServiceManager) {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend);
	} else {
		return new StreamTaskStateInitializerImpl(
			dummyEnvironment,
			stateBackend) {
			@Override
			protected <K> InternalTimeServiceManager<K> internalTimeServiceManager(
				AbstractKeyedStateBackend<K> keyedStatedBackend,
				KeyContext keyContext,
				ProcessingTimeService processingTimeService,
				Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
				return null;
			}
		};
	}
}
 

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;
}
 

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;
}
 

protected CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs(
	Iterator<StateObjectCollection<KeyedStateHandle>> restoreStateAlternatives) {

	if (restoreStateAlternatives.hasNext()) {
		Collection<KeyedStateHandle> rawKeyedState = restoreStateAlternatives.next();

		// TODO currently this does not support local state recovery, so we expect there is only one handle.
		Preconditions.checkState(
			!restoreStateAlternatives.hasNext(),
			"Local recovery is currently not implemented for raw keyed state, but found state alternative.");

		if (rawKeyedState != null) {
			Collection<KeyGroupsStateHandle> keyGroupsStateHandles = transform(rawKeyedState);
			final CloseableRegistry closeableRegistry = new CloseableRegistry();

			return new CloseableIterable<KeyGroupStatePartitionStreamProvider>() {
				@Override
				public void close() throws IOException {
					closeableRegistry.close();
				}

				@Override
				public Iterator<KeyGroupStatePartitionStreamProvider> iterator() {
					return new KeyGroupStreamIterator(keyGroupsStateHandles.iterator(), closeableRegistry);
				}
			};
		}
	}

	return CloseableIterable.empty();
}
 

StreamOperatorStateContextImpl(
	boolean restored,
	OperatorStateBackend operatorStateBackend,
	AbstractKeyedStateBackend<?> keyedStateBackend,
	InternalTimeServiceManager<?> internalTimeServiceManager,
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs,
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs) {

	this.restored = restored;
	this.operatorStateBackend = operatorStateBackend;
	this.keyedStateBackend = keyedStateBackend;
	this.internalTimeServiceManager = internalTimeServiceManager;
	this.rawOperatorStateInputs = rawOperatorStateInputs;
	this.rawKeyedStateInputs = rawKeyedStateInputs;
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	keyedStateBackend = (AbstractKeyedStateBackend<?>) getKeyedStateBackend();
	operatorStateBackend = getOperatorStateBackend();
	rawOperatorStateInputs =
		(CloseableIterable<StatePartitionStreamProvider>) context.getRawOperatorStateInputs();
	rawKeyedStateInputs =
		(CloseableIterable<KeyGroupStatePartitionStreamProvider>) context.getRawKeyedStateInputs();
	super.initializeState(context);
}
 

@Test
public void close() throws Exception {

	int count = 0;
	int stopCount = NUM_HANDLES / 2;
	boolean isClosed = false;

	try {
		for (KeyGroupStatePartitionStreamProvider stateStreamProvider
				: initializationContext.getRawKeyedStateInputs()) {
			Assert.assertNotNull(stateStreamProvider);

			if (count == stopCount) {
				closableRegistry.close();
				isClosed = true;
			}

			try (InputStream is = stateStreamProvider.getStream()) {
				DataInputView div = new DataInputViewStreamWrapper(is);
				try {
					int val = div.readInt();
					Assert.assertEquals(stateStreamProvider.getKeyGroupId(), val);
					if (isClosed) {
						Assert.fail("Close was ignored: stream");
					}
					++count;
				} catch (IOException ioex) {
					if (!isClosed) {
						throw ioex;
					}
				}
			}
		}
		Assert.fail("Close was ignored: registry");
	} catch (IOException iex) {
		Assert.assertTrue(isClosed);
		Assert.assertEquals(stopCount, count);
	}

}
 

protected CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs(
	Iterator<StateObjectCollection<KeyedStateHandle>> restoreStateAlternatives) {

	if (restoreStateAlternatives.hasNext()) {
		Collection<KeyedStateHandle> rawKeyedState = restoreStateAlternatives.next();

		// TODO currently this does not support local state recovery, so we expect there is only one handle.
		Preconditions.checkState(
			!restoreStateAlternatives.hasNext(),
			"Local recovery is currently not implemented for raw keyed state, but found state alternative.");

		if (rawKeyedState != null) {
			Collection<KeyGroupsStateHandle> keyGroupsStateHandles = transform(rawKeyedState);
			final CloseableRegistry closeableRegistry = new CloseableRegistry();

			return new CloseableIterable<KeyGroupStatePartitionStreamProvider>() {
				@Override
				public void close() throws IOException {
					closeableRegistry.close();
				}

				@Override
				public Iterator<KeyGroupStatePartitionStreamProvider> iterator() {
					return new KeyGroupStreamIterator(keyGroupsStateHandles.iterator(), closeableRegistry);
				}
			};
		}
	}

	return CloseableIterable.empty();
}
 

@Test
public void testNoRestore() throws Exception {

	MemoryStateBackend stateBackend = spy(new MemoryStateBackend(1024));

	// No job manager provided state to restore
	StreamTaskStateInitializer streamTaskStateManager = streamTaskStateManager(stateBackend, null, true);

	OperatorID operatorID = new OperatorID(47L, 11L);
	AbstractStreamOperator<?> streamOperator = mock(AbstractStreamOperator.class);
	when(streamOperator.getOperatorID()).thenReturn(operatorID);

	TypeSerializer<?> typeSerializer = new IntSerializer();
	CloseableRegistry closeableRegistry = new CloseableRegistry();

	StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(
		streamOperator.getOperatorID(),
		streamOperator.getClass().getSimpleName(),
		streamOperator,
		typeSerializer,
		closeableRegistry,
		new UnregisteredMetricsGroup());

	OperatorStateBackend operatorStateBackend = stateContext.operatorStateBackend();
	AbstractKeyedStateBackend<?> keyedStateBackend = stateContext.keyedStateBackend();
	InternalTimeServiceManager<?> timeServiceManager = stateContext.internalTimerServiceManager();
	CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = stateContext.rawKeyedStateInputs();
	CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = stateContext.rawOperatorStateInputs();

	Assert.assertEquals(false, stateContext.isRestored());
	Assert.assertNotNull(operatorStateBackend);
	Assert.assertNotNull(keyedStateBackend);
	Assert.assertNotNull(timeServiceManager);
	Assert.assertNotNull(keyedStateInputs);
	Assert.assertNotNull(operatorStateInputs);

	checkCloseablesRegistered(
		closeableRegistry,
		operatorStateBackend,
		keyedStateBackend,
		keyedStateInputs,
		operatorStateInputs);

	Assert.assertFalse(keyedStateInputs.iterator().hasNext());
	Assert.assertFalse(operatorStateInputs.iterator().hasNext());
}
 

@Override
public StreamOperatorStateContext streamOperatorStateContext(
	@Nonnull OperatorID operatorID,
	@Nonnull String operatorClassName,
	@Nonnull KeyContext keyContext,
	@Nullable TypeSerializer<?> keySerializer,
	@Nonnull CloseableRegistry streamTaskCloseableRegistry,
	@Nonnull MetricGroup metricGroup) throws Exception {

	TaskInfo taskInfo = environment.getTaskInfo();
	OperatorSubtaskDescriptionText operatorSubtaskDescription =
		new OperatorSubtaskDescriptionText(
			operatorID,
			operatorClassName,
			taskInfo.getIndexOfThisSubtask(),
			taskInfo.getNumberOfParallelSubtasks());

	final String operatorIdentifierText = operatorSubtaskDescription.toString();

	final PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates =
		taskStateManager.prioritizedOperatorState(operatorID);

	AbstractKeyedStateBackend<?> keyedStatedBackend = null;
	OperatorStateBackend operatorStateBackend = null;
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs = null;
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs = null;
	InternalTimeServiceManager<?> timeServiceManager;

	try {

		// -------------- Keyed State Backend --------------
		keyedStatedBackend = keyedStatedBackend(
			keySerializer,
			operatorIdentifierText,
			prioritizedOperatorSubtaskStates,
			streamTaskCloseableRegistry,
			metricGroup);

		// -------------- Operator State Backend --------------
		operatorStateBackend = operatorStateBackend(
			operatorIdentifierText,
			prioritizedOperatorSubtaskStates,
			streamTaskCloseableRegistry);

		// -------------- Raw State Streams --------------
		rawKeyedStateInputs = rawKeyedStateInputs(
			prioritizedOperatorSubtaskStates.getPrioritizedRawKeyedState().iterator());
		streamTaskCloseableRegistry.registerCloseable(rawKeyedStateInputs);

		rawOperatorStateInputs = rawOperatorStateInputs(
			prioritizedOperatorSubtaskStates.getPrioritizedRawOperatorState().iterator());
		streamTaskCloseableRegistry.registerCloseable(rawOperatorStateInputs);

		// -------------- Internal Timer Service Manager --------------
		timeServiceManager = internalTimeServiceManager(keyedStatedBackend, keyContext, rawKeyedStateInputs);

		// -------------- Preparing return value --------------

		return new StreamOperatorStateContextImpl(
			prioritizedOperatorSubtaskStates.isRestored(),
			operatorStateBackend,
			keyedStatedBackend,
			timeServiceManager,
			rawOperatorStateInputs,
			rawKeyedStateInputs);
	} catch (Exception ex) {

		// cleanup if something went wrong before results got published.
		if (keyedStatedBackend != null) {
			if (streamTaskCloseableRegistry.unregisterCloseable(keyedStatedBackend)) {
				IOUtils.closeQuietly(keyedStatedBackend);
			}
			// release resource (e.g native resource)
			keyedStatedBackend.dispose();
		}

		if (operatorStateBackend != null) {
			if (streamTaskCloseableRegistry.unregisterCloseable(operatorStateBackend)) {
				IOUtils.closeQuietly(operatorStateBackend);
			}
			operatorStateBackend.dispose();
		}

		if (streamTaskCloseableRegistry.unregisterCloseable(rawKeyedStateInputs)) {
			IOUtils.closeQuietly(rawKeyedStateInputs);
		}

		if (streamTaskCloseableRegistry.unregisterCloseable(rawOperatorStateInputs)) {
			IOUtils.closeQuietly(rawOperatorStateInputs);
		}

		throw new Exception("Exception while creating StreamOperatorStateContext.", ex);
	}
}
 

@Override
public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	final StreamTaskStateInitializer streamTaskStateManager = super.createStreamTaskStateInitializer();
	return (operatorID, operatorClassName, processingTimeService, keyContext, keySerializer, closeableRegistry, metricGroup) -> {

		final StreamOperatorStateContext controller = streamTaskStateManager.streamOperatorStateContext(
			operatorID,
			operatorClassName,
			processingTimeService,
			keyContext,
			keySerializer,
			closeableRegistry,
			metricGroup);

		return new StreamOperatorStateContext() {
			@Override
			public boolean isRestored() {
				return controller.isRestored();
			}

			@Override
			public OperatorStateBackend operatorStateBackend() {
				return controller.operatorStateBackend();
			}

			@Override
			public AbstractKeyedStateBackend<?> keyedStateBackend() {
				return controller.keyedStateBackend();
			}

			@Override
			public InternalTimeServiceManager<?> internalTimerServiceManager() {
				InternalTimeServiceManager<?> timeServiceManager = controller.internalTimerServiceManager();
				return timeServiceManager != null ? spy(timeServiceManager) : null;
			}

			@Override
			public CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs() {
				return replaceWithSpy(controller.rawOperatorStateInputs());
			}

			@Override
			public CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs() {
				return replaceWithSpy(controller.rawKeyedStateInputs());
			}

			public <T extends Closeable> T replaceWithSpy(T closeable) {
				T spyCloseable = spy(closeable);
				if (closeableRegistry.unregisterCloseable(closeable)) {
					try {
						closeableRegistry.registerCloseable(spyCloseable);
					} catch (IOException e) {
						throw new RuntimeException(e);
					}
				}
				return spyCloseable;
			}
		};
	};
}
 

@Test
public void testNoRestore() throws Exception {

	MemoryStateBackend stateBackend = spy(new MemoryStateBackend(1024));

	// No job manager provided state to restore
	StreamTaskStateInitializer streamTaskStateManager = streamTaskStateManager(stateBackend, null, true);

	OperatorID operatorID = new OperatorID(47L, 11L);
	AbstractStreamOperator<?> streamOperator = mock(AbstractStreamOperator.class);
	when(streamOperator.getOperatorID()).thenReturn(operatorID);

	TypeSerializer<?> typeSerializer = new IntSerializer();
	CloseableRegistry closeableRegistry = new CloseableRegistry();

	StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(
		streamOperator.getOperatorID(),
		streamOperator.getClass().getSimpleName(),
		new TestProcessingTimeService(),
		streamOperator,
		typeSerializer,
		closeableRegistry,
		new UnregisteredMetricsGroup());

	OperatorStateBackend operatorStateBackend = stateContext.operatorStateBackend();
	AbstractKeyedStateBackend<?> keyedStateBackend = stateContext.keyedStateBackend();
	InternalTimeServiceManager<?> timeServiceManager = stateContext.internalTimerServiceManager();
	CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = stateContext.rawKeyedStateInputs();
	CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = stateContext.rawOperatorStateInputs();

	Assert.assertEquals(false, stateContext.isRestored());
	Assert.assertNotNull(operatorStateBackend);
	Assert.assertNotNull(keyedStateBackend);
	Assert.assertNotNull(timeServiceManager);
	Assert.assertNotNull(keyedStateInputs);
	Assert.assertNotNull(operatorStateInputs);

	checkCloseablesRegistered(
		closeableRegistry,
		operatorStateBackend,
		keyedStateBackend,
		keyedStateInputs,
		operatorStateInputs);

	Assert.assertFalse(keyedStateInputs.iterator().hasNext());
	Assert.assertFalse(operatorStateInputs.iterator().hasNext());
}
 

@Override
public CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs() {
	return rawKeyedStateInputs;
}