You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@flink.apache.org by StefanRRichter <gi...@git.apache.org> on 2018/01/04 13:54:17 UTC

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

GitHub user StefanRRichter opened a pull request:

    https://github.com/apache/flink/pull/5239

    [FLINK-8360] Implement task-local state recovery

    ## What is the purpose of the change
    
    *This changes introduces the task-local recovery feature. The main idea is to have a secondary, local copy of the checkpointed state, while there is still a primary copy in DFS that we report to the checkpoint coordinator.
    
    Recovery can attempt to restore from the secondary local copy, if available, to save network bandwidth. This requires that the assignment from tasks to slots is as sticky is possible.
    
    For starters, we will implement this feature for all managed keyed states and can easily enhance it to all other state types (e.g. operator state) later, because the basic infrastructure is already in place. This PR is on top of #4745.*
    
    
    ## Brief change log
    
      - *Introduced `TaskExecutorLocalStateStoresManager` per task manager. This class manages one `TaskLocalStateStore` for each task running on the task manager.*
      - *`TaskLocalStateStore` stores and provides the local state for one task. Reporting of checkpointed states goes through this class. The primary state handles are forwarded to the checkpoint coordinator, the optional secondary (local) state is stored in the local store.*
      - *`LocalRecoveryDirectoryProvider` is used by `TaskLocalStateStore` to manage the local state directory/ies.*
    - *`StreamTaskStateManager` uses the `TaskLocalStateStore` to restore state for its operators.*
    - *File-based local state is created through `DuplicatingCheckpointOutputStream`, a stream that duplicates writes into two internal streams - typically one primary against a DFS and one secondary against local FS.*
    -*RocksDB's incremental checkpoints are not just based on one file, but on a directory. As we do not require reference counting for local files (we can use hardlinks), we can deal with checkpoint directories as a whole. We introduced `IncrementalLocalKeyedStateHandle extends DirectoryKeyedStateHandle` for this purpose.*
    
    ## Verifying this change
    
    This change added tests and can be verified as follows:
    
      - *If we activate the local recovery feature on the state backend (via `#setLocalRecoveryMode(...)`, and introduce task failure by user-code exception, we should observe that managed keyed state is recovered from the local FS through the logs.*
    
    ## Does this pull request potentially affect one of the following parts:
    
      - Dependencies (does it add or upgrade a dependency): (no)
      - The public API, i.e., is any changed class annotated with `@Public(Evolving)`: (yes)
      - The serializers: (yes, slightly and in a way that should not matter.)
      - The runtime per-record code paths (performance sensitive): (no)
      - Anything that affects deployment or recovery: JobManager (and its components), Checkpointing, Yarn/Mesos, ZooKeeper: (yes)
      - The S3 file system connector: (no)
    
    ## Documentation
    
      - Does this pull request introduce a new feature? (yes)
      - If yes, how is the feature documented? (Documentation pending)


You can merge this pull request into a Git repository by running:

    $ git pull https://github.com/StefanRRichter/flink task-local-recovery

Alternatively you can review and apply these changes as the patch at:

    https://github.com/apache/flink/pull/5239.patch

To close this pull request, make a commit to your master/trunk branch
with (at least) the following in the commit message:

    This closes #5239
    
----
commit 769aa8619f7da4512e51bfc13110459b50d67c0c
Author: Stefan Richter <s....@...>
Date:   2017-08-21T12:31:38Z

    [FLINK-7719] [checkpoints] Send checkpoint id to task as part of deployment descriptor when resuming

commit 194d289eb5d67494fb7242dda3589171e0bcd4f6
Author: Stefan Richter <s....@...>
Date:   2017-09-27T13:10:52Z

    [FLINK-7720] [checkpoints] Centralize creation of backends and state related resources

commit 4beb98837fbb1c99d608b6b5ceea1fba1989fbd8
Author: Stefan Richter <s....@...>
Date:   2017-11-28T17:05:41Z

    smaller refactorings of method and class names

commit 2dc9f30bcad35ef2a63a4014e2fa7f9b6411d811
Author: Stefan Richter <s....@...>
Date:   2017-11-28T17:06:05Z

    introduce state object collection for better OO.

commit 3c9d39d8ca4c1370467d5596f2c3ef51b9aff043
Author: Stefan Richter <s....@...>
Date:   2017-11-29T09:59:05Z

    Split generic type in Snapshotable interface

commit 6cf67e5fc04997873d2d9f4abfaa6be962bb464f
Author: Stefan Richter <s....@...>
Date:   2017-11-29T14:33:32Z

    Introduce explicit snapshot result type with 2 locations

commit 5d829f235fa433af4311ca0c943843b879ed4627
Author: Stefan Richter <s....@...>
Date:   2017-11-30T11:00:48Z

    Deep compare handle broken. TODO could use matcher instead

commit 238ea9881900f26c6e78af093e96fad348cab002
Author: Stefan Richter <s....@...>
Date:   2017-12-01T11:13:43Z

    javadoc on new classes.

commit 5d9c414cec2bc9410ed0972d37390200e915e325
Author: Stefan Richter <s....@...>
Date:   2017-12-04T14:12:00Z

    Preparing, test related refactorings

commit adbe507c24eda44088745159444d1ae5011e2bda
Author: Stefan Richter <s....@...>
Date:   2017-12-04T14:12:12Z

    Additional unit test for forwarding to/from TaskStateManager and StateObjectCollection.

commit 5a8b43e33918b84aec98904a77061d69ec1ff51c
Author: Stefan Richter <s....@...>
Date:   2017-12-05T17:46:07Z

    Introduce configuration for local recovery.

commit 0c77cf022fd10ee60c2a5ffe649f4a3cec2fb5fa
Author: Stefan Richter <s....@...>
Date:   2017-12-06T11:00:12Z

    Introduce snapshot strategy and some connected refactorings.

commit 48ac41b78cfede2b3bc34f0f9d61f9e2374aff3b
Author: Stefan Richter <s....@...>
Date:   2017-12-06T16:52:59Z

    file based local recovery for heap keyed backend

commit f4946175060b0b05b463b13b2e94f8cbd5457392
Author: Stefan Richter <s....@...>
Date:   2017-12-11T11:01:54Z

    [BUGFIX] Lifecycle of operator close/dispose is of w.r.t finished flag.

commit e4bcc7f4aa67a9f91874d5ae765517ab8aa9166f
Author: Stefan Richter <s....@...>
Date:   2017-12-11T14:08:29Z

    Integration fixup.

commit 95f60dff6407f9f5e59ed975b61ece3071c3daae
Author: Stefan Richter <s....@...>
Date:   2017-12-12T16:07:08Z

    introduce disposable interface

commit 37eb9b7ce99171a20c07ece1dc11b2fc3122292d
Author: Stefan Richter <s....@...>
Date:   2017-12-14T10:07:51Z

    Implement restore path for file-based local snapshots

commit a12c1c207e1a93db43bb4077646fbd8eb6c6b13b
Author: Stefan Richter <s....@...>
Date:   2017-12-14T18:04:21Z

    Cleanup organization improvements

commit 9724491fe4c8ec0b1e318e24a76db2e96541a94e
Author: Stefan Richter <s....@...>
Date:   2017-12-14T18:05:00Z

    RocksDB local incremental snapshots

commit e9811c0735809021af7c8175bacab3ce6c046a22
Author: Stefan Richter <s....@...>
Date:   2017-12-15T15:01:55Z

    Recovery path for RocksDB local incremental snapshots.

commit 588f32fd5d91db22cbf7b14c4726a64408defd2a
Author: Stefan Richter <s....@...>
Date:   2017-12-18T16:15:55Z

    Additional unit tests (WIP).

commit c1b5a8337944bcae3ac1fc012e8c16bc0e6265a7
Author: Stefan Richter <s....@...>
Date:   2017-12-30T19:23:48Z

    Abandon old OperatorStateHandles (replace by newer OperatorSubtaskState).

commit e7e602768497b3cad0567ad1edf58c4bf67d1018
Author: Stefan Richter <s....@...>
Date:   2018-01-02T13:21:04Z

    Additional unit tests (WIP) - 4.

commit 3fff0889dafefdb7cb0920b692ed9fdb3645d3b4
Author: Stefan Richter <s....@...>
Date:   2018-01-02T16:38:33Z

    Change directory structure so that we can delete local state directories based on allocation id in the (flip-6) future.

commit 958baaa5c118b431c18df298f3f138133ac42466
Author: Stefan Richter <s....@...>
Date:   2018-01-02T17:30:48Z

    Make checkstyle happy :-)

----


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159834127
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java ---
    @@ -145,8 +145,8 @@
     
     	private volatile Throwable failureCause;          // once assigned, never changes
     
    -	/** The handle to the state that the task gets on restore */
    -	private volatile TaskStateSnapshot taskState;
    +	/** Information to restore the task onr recovery, such as checkpoint id and task state snapshot */
    --- End diff --
    
    `onr` typo


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168527958
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -59,58 +71,55 @@ public TaskExecutorLocalStateStoresManager(File[] localStateRootDirectories) {
     				}
     			}
     		}
    -	}
     
    -	public TaskLocalStateStore localStateStoreForTask(
    -		JobID jobId,
    -		JobVertexID jobVertexID,
    -		int subtaskIndex) {
    +	}
     
    -		Preconditions.checkNotNull(jobId);
    -		final JobVertexSubtaskKey taskKey = new JobVertexSubtaskKey(jobVertexID, subtaskIndex);
    +	public TaskLocalStateStore localStateStoreForSubtask(
    +		@Nonnull JobID jobId,
    +		@Nonnull AllocationID allocationID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex) {
     
     		final Map<JobVertexSubtaskKey, TaskLocalStateStore> taskStateManagers =
    -			this.taskStateStoresByJobID.computeIfAbsent(jobId, k -> new HashMap<>());
    +			this.taskStateStoresByAllocationID.computeIfAbsent(allocationID, k -> new ConcurrentHashMap<>());
    +
    +		final JobVertexSubtaskKey taskKey = new JobVertexSubtaskKey(jobVertexID, subtaskIndex);
     
     		return taskStateManagers.computeIfAbsent(
    -			taskKey, k -> new TaskLocalStateStore(jobId, jobVertexID, subtaskIndex, localStateRootDirectories));
    +			taskKey,
    +			k -> new TaskLocalStateStore(jobId, jobVertexID, subtaskIndex, localStateRootDirectories, discardExecutor));
     	}
     
    -	public void releaseJob(JobID jobID) {
    +	public void releaseLocalStateForAllocationId(@Nonnull AllocationID allocationID) {
     
    -		Map<JobVertexSubtaskKey, TaskLocalStateStore> cleanupLocalStores = taskStateStoresByJobID.remove(jobID);
    +		Map<JobVertexSubtaskKey, TaskLocalStateStore> cleanupLocalStores =
    +			taskStateStoresByAllocationID.remove(allocationID);
     
     		if (cleanupLocalStores != null) {
    -//			doRelease(cleanupLocalStores.values());
    +			doRelease(cleanupLocalStores.values());
     		}
     	}
     
     	public void releaseAll() {
     
    -		for (Map<JobVertexSubtaskKey, TaskLocalStateStore> stateStoreMap : taskStateStoresByJobID.values()) {
    -//			doRelease(stateStoreMap.values());
    +		for (Map<JobVertexSubtaskKey, TaskLocalStateStore> stateStoreMap : taskStateStoresByAllocationID.values()) {
    +			doRelease(stateStoreMap.values());
     		}
     
    -		taskStateStoresByJobID.clear();
    +		taskStateStoresByAllocationID.clear();
     	}
     
    -	private void doRelease(Iterable<TaskLocalStateStore> toRelease) throws Exception {
    +	private void doRelease(Iterable<TaskLocalStateStore> toRelease) {
     
     		if (toRelease != null) {
     
    -			Exception collectedExceptions = null;
    -
     			for (TaskLocalStateStore stateStore : toRelease) {
     				try {
     					stateStore.dispose();
     				} catch (Exception disposeEx) {
    -					collectedExceptions = ExceptionUtils.firstOrSuppressed(disposeEx, collectedExceptions);
    +					LOG.warn("Exception while disposing local state store", disposeEx);
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168777497
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java ---
    @@ -0,0 +1,265 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.checkpoint;
    +
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.Preconditions;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.util.ArrayList;
    +import java.util.Collections;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.function.BiFunction;
    +
    +/**
    + * This class is a wrapper over multiple alternative {@link OperatorSubtaskState} that are (partial) substitutes for
    + * each other and imposes a priority ordering over all alternatives for the different states which define an order in
    + * which the operator should attempt to restore the state from them. One OperatorSubtaskState is considered as the
    + * "ground truth" about which state should be represented. Alternatives may be complete or partial substitutes for
    + * the "ground truth" with a higher priority (if they had a lower alternative, they would not really be alternatives).
    + * Substitution is determined on a per-sub-state basis.
    + */
    +public class PrioritizedOperatorSubtaskState {
    +
    +	/** Singleton instance for an empty, non-restored operator state. */
    +	private static final PrioritizedOperatorSubtaskState EMPTY_NON_RESTORED_INSTANCE =
    +		new PrioritizedOperatorSubtaskState(new OperatorSubtaskState(), Collections.emptyList(), false);
    +
    +	/** List of prioritized snapshot alternatives for managed operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedManagedOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for raw operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedRawOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for managed keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedManagedKeyedState;
    +
    +	/** List of prioritized snapshot alternatives for raw keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedRawKeyedState;
    +
    +	/** Signal flag if this represents state for a restored operator. */
    +	private final boolean restored;
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority) {
    +		this(jobManagerState, alternativesByPriority, true);
    +	}
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority,
    +		boolean restored) {
    +
    +		Preconditions.checkNotNull(jobManagerState, "Job manager state is null.");
    +		int size = Preconditions.checkNotNull(alternativesByPriority, "Alternative states are null.").size();
    +
    +		this.restored = restored;
    +
    +		List<StateObjectCollection<OperatorStateHandle>> managedOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> managedKeyedAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<OperatorStateHandle>> rawOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> rawKeyedAlternatives = new ArrayList<>(size);
    +
    +		for (OperatorSubtaskState subtaskState : alternativesByPriority) {
    +
    +			if (subtaskState != null) {
    +				managedKeyedAlternatives.add(subtaskState.getManagedKeyedState());
    +				rawKeyedAlternatives.add(subtaskState.getRawKeyedState());
    +				managedOperatorAlternatives.add(subtaskState.getManagedOperatorState());
    +				rawOperatorAlternatives.add(subtaskState.getRawOperatorState());
    +			}
    +		}
    +
    +		// Key-groups should match.
    +		BiFunction<KeyedStateHandle, KeyedStateHandle, Boolean> keyedStateApprover =
    +			(ref, alt) -> ref.getKeyGroupRange().equals(alt.getKeyGroupRange());
    +
    +		// State meta data should match.
    +		BiFunction<OperatorStateHandle, OperatorStateHandle, Boolean> operatorStateApprover =
    +			(ref, alt) -> ref.getStateNameToPartitionOffsets().equals(alt.getStateNameToPartitionOffsets());
    +
    +		this.prioritizedManagedKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedKeyedState(),
    +			managedKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedRawKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawKeyedState(),
    +			rawKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedManagedOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedOperatorState(),
    +			managedOperatorAlternatives,
    +			operatorStateApprover);
    +
    +		this.prioritizedRawOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawOperatorState(),
    +			rawOperatorAlternatives,
    +			operatorStateApprover);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed operator state, in the order in which
    +	 * we should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedManagedOperatorState() {
    +		return prioritizedManagedOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw operator state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedRawOperatorState() {
    +		return prioritizedRawOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedManagedKeyedState() {
    +		return prioritizedManagedKeyedState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedRawKeyedState() {
    +		return prioritizedRawKeyedState.iterator();
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns the managed operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerManagedOperatorState() {
    +		return lastElement(prioritizedManagedOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the raw operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerRawOperatorState() {
    +		return lastElement(prioritizedRawOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the managed keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerManagedKeyedState() {
    +		return lastElement(prioritizedManagedKeyedState);
    +	}
    +
    +	/**
    +	 * Returns the raw keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerRawKeyedState() {
    +		return lastElement(prioritizedRawKeyedState);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns true if this was created for a restored operator, false otherwise. Restored operators are operators that
    +	 * participated in a previous checkpoint, even if they did not emit any state snapshots.
    +	 */
    +	public boolean isRestored() {
    +		return restored;
    +	}
    +
    +
    +	/**
    +	 * This helper method resolves the dependencies between the ground truth of the operator state obtained from the
    +	 * job manager and potential alternatives for recovery, e.g. from a task-local source.
    +	 */
    +	protected <T extends StateObject> List<StateObjectCollection<T>> resolvePrioritizedAlternatives(
    +		StateObjectCollection<T> jobManagerState,
    +		List<StateObjectCollection<T>> alternativesByPriority,
    +		BiFunction<T, T, Boolean> approveFun) {
    +
    +		// Nothing to resolve if there are no alternatives, or the ground truth has already no state, or if we can
    +		// assume that a rescaling happened because we find more than one handle in the JM state (this is more a sanity
    +		// check).
    +		if (alternativesByPriority == null
    +			|| alternativesByPriority.isEmpty()
    +			|| !jobManagerState.hasState()
    +			|| jobManagerState.size() != 1) {
    +
    --- End diff --
    
    Exactly, local recovery is for checkpoints not savepoints. Rescaling from checkpoint/savepoint involves a restart of the job, which means that all local state is cleaned up anyways because the job was canceled.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168169824
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -1016,7 +1017,7 @@ void releaseResources(boolean canceled) {
     	}
     
     	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    --- End diff --
    
    Will change it back to `Collection`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168753362
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    +
    +- RocksDBStateBackend: task-local recovery is supported for keyed state. For *full checkpoints*, state is duplicated to a local file. This can introduce additional write costs
    +and occupies local disk space. For *incremental snapshots*, the local state is based on RocksDB's native checkpointing mechanism. This mechanism is also used as the first step
    +to create the primary copy, which means that in this case no additional cost is introduced for create the secondary copy. We simply keep the native checkpoint directory around
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165323982
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/LocalStateForwardingTest.java ---
    @@ -0,0 +1,213 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.runtime.tasks;
    +
    +import org.apache.flink.api.common.ExecutionConfig;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.configuration.Configuration;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetrics;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.StateObjectCollection;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
    +import org.apache.flink.runtime.state.DoneFuture;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.SnapshotResult;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.runtime.state.TaskLocalStateStore;
    +import org.apache.flink.runtime.state.TaskStateManagerImpl;
    +import org.apache.flink.runtime.state.TestTaskStateManager;
    +import org.apache.flink.runtime.taskmanager.TestCheckpointResponder;
    +import org.apache.flink.streaming.api.operators.OperatorSnapshotFutures;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +
    +import java.util.HashMap;
    +import java.util.Map;
    +import java.util.concurrent.Future;
    +import java.util.concurrent.RunnableFuture;
    +import java.util.concurrent.atomic.AtomicBoolean;
    +
    +import static org.mockito.Mockito.mock;
    +
    +/**
    + * Test for forwarding of state reporting to and from {@link org.apache.flink.runtime.state.TaskStateManager}.
    + */
    +public class LocalStateForwardingTest {
    --- End diff --
    
    Same here with `TestLogger`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169895777
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -2123,9 +2162,21 @@ void takeSnapshot() throws Exception {
     				checkpointId,
     				sstFiles,
     				miscFiles,
    -				metaStateHandle);
    +				metaStateHandle.getJobManagerOwnedSnapshot());
    +
    +			DirectoryStateHandle directoryStateHandle = localBackupDirectory.completeSnapshotAndGetHandle();
    +			StreamStateHandle taskLocalSnapshotMetaDataStateHandle = metaStateHandle.getTaskLocalSnapshot();
    +			IncrementalLocalKeyedStateHandle directoryKeyedStateHandle =
    +				directoryStateHandle != null && taskLocalSnapshotMetaDataStateHandle != null ?
    +					new IncrementalLocalKeyedStateHandle(
    +						stateBackend.backendUID,
    +						checkpointId,
    +						directoryStateHandle,
    +						stateBackend.keyGroupRange,
    +						taskLocalSnapshotMetaDataStateHandle) :
    +					null;
     
    -			return new SnapshotResult<>(incrementalKeyedStateHandle, null);
    +			return new SnapshotResult<>(incrementalKeyedStateHandle, directoryKeyedStateHandle);
    --- End diff --
    
    Exactly 👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169047804
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    👍 will change this even a bit more.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160962776
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamTaskStateManagerImplTest.java ---
    @@ -0,0 +1,323 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.api.common.typeutils.TypeSerializer;
    +import org.apache.flink.api.common.typeutils.base.IntSerializer;
    +import org.apache.flink.core.fs.CloseableRegistry;
    +import org.apache.flink.runtime.checkpoint.JobManagerTaskRestore;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.checkpoint.savepoint.CheckpointTestUtils;
    +import org.apache.flink.runtime.execution.Environment;
    +import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
    +import org.apache.flink.runtime.query.TaskKvStateRegistry;
    +import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
    +import org.apache.flink.runtime.state.CheckpointStreamFactory;
    +import org.apache.flink.runtime.state.KeyGroupRange;
    +import org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.OperatorStateBackend;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.StateBackend;
    +import org.apache.flink.runtime.state.StatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.TaskStateManager;
    +import org.apache.flink.runtime.state.TaskStateManagerImplTest;
    +import org.apache.flink.runtime.state.memory.MemoryStateBackend;
    +import org.apache.flink.runtime.taskmanager.TestCheckpointResponder;
    +import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService;
    +import org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService;
    +import org.apache.flink.util.CloseableIterable;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import javax.annotation.Nullable;
    +
    +import java.io.Closeable;
    +import java.io.IOException;
    +import java.util.Collections;
    +import java.util.Random;
    +
    +import static org.mockito.Matchers.any;
    +import static org.mockito.Matchers.anyInt;
    +import static org.mockito.Matchers.eq;
    +import static org.mockito.Mockito.mock;
    +import static org.mockito.Mockito.spy;
    +import static org.mockito.Mockito.verify;
    +import static org.mockito.Mockito.when;
    +
    +/**
    + * Test for {@link StreamTaskStateManagerImpl}.
    + */
    +public class StreamTaskStateManagerImplTest {
    +
    +	@Test
    +	public void testNoRestore() throws Exception {
    +
    +		MemoryStateBackend stateBackend = spy(new MemoryStateBackend(1024));
    +
    +		// No job manager provided state to restore
    +		StreamTaskStateManager 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,
    +			typeSerializer,
    +			closeableRegistry);
    +
    +		verify(stateBackend).createKeyedStateBackend(
    +			any(Environment.class),
    +			any(JobID.class),
    +			any(String.class),
    +			eq(typeSerializer),
    +			anyInt(),
    +			any(KeyGroupRange.class),
    +			any(TaskKvStateRegistry.class));
    +
    +		verify(stateBackend).createOperatorStateBackend(
    +			any(Environment.class),
    +			any(String.class));
    +
    +		verify(stateBackend).createStreamFactory(
    +			any(JobID.class),
    +			any(String.class));
    +
    +		OperatorStateBackend operatorStateBackend = stateContext.operatorStateBackend();
    +		AbstractKeyedStateBackend<?> keyedStateBackend = stateContext.keyedStateBackend();
    +		InternalTimeServiceManager<?, ?> timeServiceManager = stateContext.internalTimerServiceManager();
    +		CheckpointStreamFactory streamFactory = stateContext.checkpointStreamFactory();
    +		CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = stateContext.rawKeyedStateInputs();
    +		CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = stateContext.rawOperatorStateInputs();
    +
    +		Assert.assertEquals(false, stateContext.isRestored());
    +		Assert.assertNotNull(operatorStateBackend);
    --- End diff --
    
    Shouldn't you make some additional assertions on those fields besides checking not null? Or is it enough on this level to have just those not nulls?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168191066
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStoreImpl.java ---
    @@ -0,0 +1,298 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
    +
    +import java.io.File;
    +import java.io.IOException;
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.HashMap;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
    +
    +/**
    + * Main implementation of a {@link TaskLocalStateStore}.
    + */
    +public class TaskLocalStateStoreImpl implements TaskLocalStateStore {
    +
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStoreImpl.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	@VisibleForTesting
    +	static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    +
    +	/** JobID from the owning subtask. */
    +	@Nonnull
    +	private final JobID jobID;
    +
    +	/** AllocationID of the owning slot. */
    +	@Nonnull
    +	private final AllocationID allocationID;
    +
    +	/** JobVertexID of the owning subtask. */
    +	@Nonnull
    +	private final JobVertexID jobVertexID;
    +
    +	/** Subtask index of the owning subtask. */
    +	@Nonnegative
    +	private final int subtaskIndex;
    +
    +	/** The configured mode for local recovery. */
    +	@Nonnull
    +	private final LocalRecoveryConfig localRecoveryConfig;
    +
    +	/** Executor that runs the discarding of released state objects. */
    +	@Nonnull
    +	private final Executor discardExecutor;
    +
    +	/** Lock for synchronisation on the storage map and the discarded status. */
    +	@Nonnull
    +	private final Object lock;
    +
    +	/** Status flag if this store was already discarded. */
    +	@GuardedBy("lock")
    +	private boolean discarded;
    +
    +	/** Maps checkpoint ids to local TaskStateSnapshots. */
    +	@Nonnull
    +	@GuardedBy("lock")
    +	private final SortedMap<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
    +	public TaskLocalStateStoreImpl(
    +		@Nonnull JobID jobID,
    +		@Nonnull AllocationID allocationID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex,
    +		@Nonnull LocalRecoveryConfig localRecoveryConfig,
    +		@Nonnull Executor discardExecutor) {
    +
    +		this.jobID = jobID;
    +		this.allocationID = allocationID;
    +		this.jobVertexID = jobVertexID;
    +		this.subtaskIndex = subtaskIndex;
    +		this.discardExecutor = discardExecutor;
    +		this.lock = new Object();
    +		this.storedTaskStateByCheckpointID = new TreeMap<>();
    +		this.discarded = false;
    +		this.localRecoveryConfig = localRecoveryConfig;
    +	}
    +
    +	@Override
    +	public void storeLocalState(
    +		@Nonnegative long checkpointId,
    +		@Nullable TaskStateSnapshot localState) {
    +
    +		if (localState == null) {
    +			localState = NULL_DUMMY;
    +		}
    +
    +		LOG.info("Storing local state for checkpoint {}.", checkpointId);
    +		LOG.debug("Local state for checkpoint {} is {}.", checkpointId, localState);
    +
    +		Map<Long, TaskStateSnapshot> toDiscard = new HashMap<>(MAX_RETAINED_SNAPSHOTS);
    +
    +		synchronized (lock) {
    +			if (discarded) {
    +				// we ignore late stores and simply discard the state.
    +				toDiscard.put(checkpointId, localState);
    +			} else {
    +				TaskStateSnapshot previous =
    +					storedTaskStateByCheckpointID.put(checkpointId, localState);
    +
    +				if (previous != null) {
    +					toDiscard.put(checkpointId, previous);
    +				}
    +
    +				// remove from history.
    +				while (storedTaskStateByCheckpointID.size() > MAX_RETAINED_SNAPSHOTS) {
    +					Long removeCheckpointID = storedTaskStateByCheckpointID.firstKey();
    --- End diff --
    
    If we have to remove local completed checkpoint here, I think `MAX_RETAINED_SNAPSHOTS` should bigger than `max concurrent checkpoint number.` 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168169266
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java ---
    @@ -3404,12 +3405,16 @@ public String fold(String acc, Integer value) throws Exception {
     		}
     	}
     
    -	protected KeyedStateHandle runSnapshot(RunnableFuture<KeyedStateHandle> snapshotRunnableFuture) throws Exception {
    +	protected KeyedStateHandle runSnapshot(
    +		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunnableFuture) throws Exception {
    +
     		if(!snapshotRunnableFuture.isDone()) {
     			Thread runner = new Thread(snapshotRunnableFuture);
     			runner.start();
     		}
    -		return snapshotRunnableFuture.get();
    +
    +		SnapshotResult<KeyedStateHandle> snapshotResult = snapshotRunnableFuture.get();
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168734460
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    This should also be remove when canceled is true.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168779365
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java ---
    @@ -0,0 +1,265 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.checkpoint;
    +
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.Preconditions;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.util.ArrayList;
    +import java.util.Collections;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.function.BiFunction;
    +
    +/**
    + * This class is a wrapper over multiple alternative {@link OperatorSubtaskState} that are (partial) substitutes for
    + * each other and imposes a priority ordering over all alternatives for the different states which define an order in
    + * which the operator should attempt to restore the state from them. One OperatorSubtaskState is considered as the
    + * "ground truth" about which state should be represented. Alternatives may be complete or partial substitutes for
    + * the "ground truth" with a higher priority (if they had a lower alternative, they would not really be alternatives).
    + * Substitution is determined on a per-sub-state basis.
    + */
    +public class PrioritizedOperatorSubtaskState {
    +
    +	/** Singleton instance for an empty, non-restored operator state. */
    +	private static final PrioritizedOperatorSubtaskState EMPTY_NON_RESTORED_INSTANCE =
    +		new PrioritizedOperatorSubtaskState(new OperatorSubtaskState(), Collections.emptyList(), false);
    +
    +	/** List of prioritized snapshot alternatives for managed operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedManagedOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for raw operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedRawOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for managed keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedManagedKeyedState;
    +
    +	/** List of prioritized snapshot alternatives for raw keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedRawKeyedState;
    +
    +	/** Signal flag if this represents state for a restored operator. */
    +	private final boolean restored;
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority) {
    +		this(jobManagerState, alternativesByPriority, true);
    +	}
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority,
    +		boolean restored) {
    +
    +		Preconditions.checkNotNull(jobManagerState, "Job manager state is null.");
    +		int size = Preconditions.checkNotNull(alternativesByPriority, "Alternative states are null.").size();
    +
    +		this.restored = restored;
    +
    +		List<StateObjectCollection<OperatorStateHandle>> managedOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> managedKeyedAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<OperatorStateHandle>> rawOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> rawKeyedAlternatives = new ArrayList<>(size);
    +
    +		for (OperatorSubtaskState subtaskState : alternativesByPriority) {
    +
    +			if (subtaskState != null) {
    +				managedKeyedAlternatives.add(subtaskState.getManagedKeyedState());
    +				rawKeyedAlternatives.add(subtaskState.getRawKeyedState());
    +				managedOperatorAlternatives.add(subtaskState.getManagedOperatorState());
    +				rawOperatorAlternatives.add(subtaskState.getRawOperatorState());
    +			}
    +		}
    +
    +		// Key-groups should match.
    +		BiFunction<KeyedStateHandle, KeyedStateHandle, Boolean> keyedStateApprover =
    +			(ref, alt) -> ref.getKeyGroupRange().equals(alt.getKeyGroupRange());
    +
    +		// State meta data should match.
    +		BiFunction<OperatorStateHandle, OperatorStateHandle, Boolean> operatorStateApprover =
    +			(ref, alt) -> ref.getStateNameToPartitionOffsets().equals(alt.getStateNameToPartitionOffsets());
    +
    +		this.prioritizedManagedKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedKeyedState(),
    +			managedKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedRawKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawKeyedState(),
    +			rawKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedManagedOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedOperatorState(),
    +			managedOperatorAlternatives,
    +			operatorStateApprover);
    +
    +		this.prioritizedRawOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawOperatorState(),
    +			rawOperatorAlternatives,
    +			operatorStateApprover);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed operator state, in the order in which
    +	 * we should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedManagedOperatorState() {
    +		return prioritizedManagedOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw operator state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedRawOperatorState() {
    +		return prioritizedRawOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedManagedKeyedState() {
    +		return prioritizedManagedKeyedState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedRawKeyedState() {
    +		return prioritizedRawKeyedState.iterator();
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns the managed operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerManagedOperatorState() {
    +		return lastElement(prioritizedManagedOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the raw operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerRawOperatorState() {
    +		return lastElement(prioritizedRawOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the managed keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerManagedKeyedState() {
    +		return lastElement(prioritizedManagedKeyedState);
    +	}
    +
    +	/**
    +	 * Returns the raw keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerRawKeyedState() {
    +		return lastElement(prioritizedRawKeyedState);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns true if this was created for a restored operator, false otherwise. Restored operators are operators that
    +	 * participated in a previous checkpoint, even if they did not emit any state snapshots.
    +	 */
    +	public boolean isRestored() {
    +		return restored;
    +	}
    +
    +
    +	/**
    +	 * This helper method resolves the dependencies between the ground truth of the operator state obtained from the
    +	 * job manager and potential alternatives for recovery, e.g. from a task-local source.
    +	 */
    +	protected <T extends StateObject> List<StateObjectCollection<T>> resolvePrioritizedAlternatives(
    +		StateObjectCollection<T> jobManagerState,
    +		List<StateObjectCollection<T>> alternativesByPriority,
    +		BiFunction<T, T, Boolean> approveFun) {
    +
    +		// Nothing to resolve if there are no alternatives, or the ground truth has already no state, or if we can
    +		// assume that a rescaling happened because we find more than one handle in the JM state (this is more a sanity
    +		// check).
    +		if (alternativesByPriority == null
    +			|| alternativesByPriority.isEmpty()
    +			|| !jobManagerState.hasState()
    +			|| jobManagerState.size() != 1) {
    +
    --- End diff --
    
    That might be, but there is still not a fundamental problem with the approach if we change the matching/creation of the effective state handles in the future, once it is required. For the time being, keeping it simple keeps maintenance simpler.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168169295
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskStateManager.java ---
    @@ -47,10 +47,11 @@
     	 * @param checkpointMetrics task level metrics for the checkpoint.
     	 * @param acknowledgedState the reported states from the owning task.
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165309670
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizer.java ---
    @@ -0,0 +1,83 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.SnapshotResult;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.FutureUtil;
    +
    +import java.util.concurrent.ExecutionException;
    +
    +/**
    + * TODO write comment.
    + */
    +public class OperatorSnapshotFinalizer {
    +
    +	private final OperatorSubtaskState jobManagerOwnedState;
    +	private final OperatorSubtaskState taskLocalState;
    +
    +	public OperatorSnapshotFinalizer(
    +		OperatorSnapshotFutures snapshotFutures) throws ExecutionException, InterruptedException {
    +
    +		SnapshotResult<KeyedStateHandle> keyedManaged =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getKeyedStateManagedFuture());
    +
    +		SnapshotResult<KeyedStateHandle> keyedRaw =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getKeyedStateRawFuture());
    +
    +		SnapshotResult<OperatorStateHandle> operatorManaged =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getOperatorStateManagedFuture());
    +
    +		SnapshotResult<OperatorStateHandle> operatorRaw =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getOperatorStateRawFuture());
    +
    +		jobManagerOwnedState = new OperatorSubtaskState(
    +			extractJobManagerOwned(operatorManaged),
    +			extractJobManagerOwned(operatorRaw),
    +			extractJobManagerOwned(keyedManaged),
    +			extractJobManagerOwned(keyedRaw)
    +		);
    +
    +		taskLocalState = new OperatorSubtaskState(
    +			extractTaskLocal(operatorManaged),
    +			extractTaskLocal(operatorRaw),
    +			extractTaskLocal(keyedManaged),
    +			extractTaskLocal(keyedRaw)
    +		);
    +	}
    +
    +	public OperatorSubtaskState getTaskLocalState() {
    +		return taskLocalState;
    +	}
    +
    +	public OperatorSubtaskState getJobManagerOwnedState() {
    +		return jobManagerOwnedState;
    +	}
    +
    +	private <T extends StateObject> T extractJobManagerOwned(SnapshotResult<T> snapshotResult) {
    +		return snapshotResult != null ? snapshotResult.getJobManagerOwnedSnapshot() : null;
    --- End diff --
    
    Same here with the `null` values.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168179034
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskStateManagerImplTest.java ---
    @@ -97,6 +104,39 @@ public void testStateReportingAndRetrieving() {
     		Assert.assertNull(taskStateManager.operatorStates(operatorID_3));
     	}
     
    +	/**
    +	 * This tests if the {@link TaskStateManager} properly returns the the subtask local state dir from the
    +	 * corresponding {@link TaskLocalStateStore}.
    +	 */
    +	@Test
    +	public void testForwardingSubtaskLocalStateBaseDirFromLocalStateStore() throws IOException {
    +		JobID jobID = new JobID(42L, 43L);
    +		JobVertexID jobVertexID = new JobVertexID(12L, 34L);
    +		ExecutionAttemptID executionAttemptID = new ExecutionAttemptID(23L, 24L);
    +		TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();
    +
    +		TemporaryFolder tmpFolder = new TemporaryFolder();
    +
    +		try {
    +			tmpFolder.create();
    +			TaskLocalStateStore taskLocalStateStore =
    +				new TaskLocalStateStore(jobID, jobVertexID, 13, tmpFolder.newFolder());
    +
    +			TaskStateManager taskStateManager = taskStateManager(
    +				jobID,
    +				executionAttemptID,
    +				checkpointResponderMock,
    +				null,
    +				taskLocalStateStore);
    +
    +			Assert.assertEquals(
    +				taskLocalStateStore.getSubtaskLocalStateBaseDirectory(),
    +				taskStateManager.getSubtaskLocalStateBaseDirectory());
    --- End diff --
    
    Already changed by later commits.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168752064
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165323575
  
    --- Diff: flink-core/src/test/java/org/apache/flink/util/MethodForwardingTestUtil.java ---
    @@ -0,0 +1,146 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.util;
    +
    +import org.mockito.Mockito;
    +
    +import java.lang.reflect.Array;
    +import java.lang.reflect.Method;
    +import java.util.Collections;
    +import java.util.Set;
    +import java.util.function.Function;
    +
    +import static org.junit.Assert.fail;
    +import static org.mockito.Mockito.reset;
    +import static org.mockito.Mockito.spy;
    +
    +/**
    + * Helper class with a method that attempt to automatically test method forwarding between a delegate and a wrapper.
    + */
    +public class MethodForwardingTestUtil {
    +
    +	/**
    +	 * This is a best effort automatic test for method forwarding between a delegate and its wrapper, where the wrapper
    +	 * class is a subtype of the delegate. This ignores methods that are inherited from Object.
    +	 *
    +	 * @param delegateClass the class for the delegate.
    +	 * @param wrapperFactory factory that produces a wrapper from a delegate.
    +	 * @param <D> type of the delegate
    +	 * @param <W> type of the wrapper
    +	 */
    +	public static <D, W> void testMethodForwarding(
    +		Class<D> delegateClass,
    +		Function<D, W> wrapperFactory) {
    +		testMethodForwarding(delegateClass, wrapperFactory, Collections.emptySet());
    +	}
    +
    +	/**
    +	 * This is a best effort automatic test for method forwarding between a delegate and its wrapper, where the wrapper
    +	 * class is a subtype of the delegate. Methods can be remapped in case that the implementation does not call the
    +	 * original method. Remapping to null skips the method. This ignores methods that are inherited from Object.
    +	 *
    +	 * @param delegateClass the class for the delegate.
    +	 * @param wrapperFactory factory that produces a wrapper from a delegate.
    +	 * @param skipMethodSet set of methods to ignore.
    +	 * @param <D> type of the delegate
    +	 * @param <W> type of the wrapper
    +	 */
    +	public static <D, W> void testMethodForwarding(
    +		Class<D> delegateClass,
    +		Function<D, W> wrapperFactory,
    +		Set<Method> skipMethodSet) {
    +
    +		Preconditions.checkNotNull(delegateClass);
    +		Preconditions.checkNotNull(wrapperFactory);
    +		Preconditions.checkNotNull(skipMethodSet);
    +
    +		D delegate = spy(delegateClass);
    +		W wrapper = wrapperFactory.apply(delegate);
    +
    +		// ensure that wrapper is a subtype of delegate
    +		Preconditions.checkArgument(delegateClass.isAssignableFrom(wrapper.getClass()));
    +
    +		for (Method delegateMethod : delegateClass.getMethods()) {
    +
    +			if (checkSkipMethodForwardCheck(delegateMethod, skipMethodSet)) {
    +				continue;
    +			}
    +
    +			try {
    +				// find the correct method to substitute the bridge for erased generic types.
    +				// if this doesn't work, the user need to exclude the method and write an additional test.
    +				Method wrapperMethod = wrapper.getClass().getDeclaredMethod(
    +					delegateMethod.getName(),
    +					delegateMethod.getParameterTypes());
    +
    +				// things get a bit fuzzy here, best effort to find a match but this might end up with a wrong method.
    +				if (wrapperMethod.isBridge()) {
    +					for (Method method : wrapper.getClass().getDeclaredMethods()) {
    +						if (!method.isBridge()
    +							&& method.getName().equals(wrapperMethod.getName())
    +							&& method.getParameterCount() == wrapperMethod.getParameterCount()) {
    +							wrapperMethod = method;
    +							break;
    +						}
    +					}
    +				}
    +
    +				Class<?>[] parameterTypes = wrapperMethod.getParameterTypes();
    +				Object[] arguments = new Object[parameterTypes.length];
    +				for (int j = 0; j < arguments.length; j++) {
    +					Class<?> parameterType = parameterTypes[j];
    +					if (parameterType.isArray()) {
    +						arguments[j] = Array.newInstance(parameterType.getComponentType(), 0);
    +					} else if (parameterType.isPrimitive()) {
    +						arguments[j] = 0;
    +					} else {
    +						arguments[j] = Mockito.mock(parameterType);
    +					}
    +				}
    +
    +				wrapperMethod.invoke(wrapper, arguments);
    +				delegateMethod.invoke(Mockito.verify(delegate, Mockito.times(1)), arguments);
    +				reset(delegate);
    +			} catch (Exception ex) {
    +				ex.printStackTrace();
    +				fail("Forwarding test failed: " + ex.getMessage());
    --- End diff --
    
    I think we should let the exception propagate.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165308709
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java ---
    @@ -3404,12 +3405,16 @@ public String fold(String acc, Integer value) throws Exception {
     		}
     	}
     
    -	protected KeyedStateHandle runSnapshot(RunnableFuture<KeyedStateHandle> snapshotRunnableFuture) throws Exception {
    +	protected KeyedStateHandle runSnapshot(
    +		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunnableFuture) throws Exception {
    +
     		if(!snapshotRunnableFuture.isDone()) {
     			Thread runner = new Thread(snapshotRunnableFuture);
     			runner.start();
     		}
    -		return snapshotRunnableFuture.get();
    +
    +		SnapshotResult<KeyedStateHandle> snapshotResult = snapshotRunnableFuture.get();
    --- End diff --
    
    Why are we spawning a new `Thread` if we block the current thread here? Wouldn't it be better to simply call `run` on the `RunnableFuture` if it is not yet done?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169062887
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    Aha, closing it.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165102131
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/DefaultOperatorStateBackend.java ---
    @@ -325,8 +325,7 @@ public OperatorStateHandle performOperation() throws Exception {
     		return task;
     	}
     
    -	@Override
    -	public void restore(Collection<OperatorStateHandle> restoreSnapshots) throws Exception {
    +	public void restore(StateObjectCollection<OperatorStateHandle> restoreSnapshots) throws Exception {
    --- End diff --
    
    Why can't it remain a `Collection`?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168778376
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -266,11 +305,16 @@ public RocksDBKeyedStateBackend(
     		RocksIterator iterator = db.newIterator(columnInfo.f0);
     		iterator.seekToFirst();
     
    -		Iterable<K> iterable = () -> new RocksIteratorWrapper<>(iterator, state, keySerializer, keyGroupPrefixBytes);
    +		Iterable<K> iterable = () -> new RocksIteratorToJavaIteratorAdapter<>(iterator, state, keySerializer, keyGroupPrefixBytes);
     		Stream<K> targetStream = StreamSupport.stream(iterable.spliterator(), false);
    --- End diff --
    
    Well, maybe, except if this is also only used with void-namespace.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169726050
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/DirectoryKeyedStateHandle.java ---
    @@ -0,0 +1,105 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import javax.annotation.Nonnull;
    +
    +/**
    + * This class is a keyed state handle based on a directory. It combines a {@link DirectoryStateHandle} and a
    + * {@link KeyGroupRange}.
    + */
    +public class DirectoryKeyedStateHandle implements KeyedStateHandle {
    +
    +	@Nonnull
    +	private final DirectoryStateHandle directoryStateHandle;
    +
    +	@Nonnull
    +	private final KeyGroupRange keyGroupRange;
    +
    +	public DirectoryKeyedStateHandle(
    +		@Nonnull DirectoryStateHandle directoryStateHandle,
    +		@Nonnull KeyGroupRange keyGroupRange) {
    +
    +		this.directoryStateHandle = directoryStateHandle;
    +		this.keyGroupRange = keyGroupRange;
    +	}
    +
    +	@Nonnull
    +	public DirectoryStateHandle getDirectoryStateHandle() {
    +		return directoryStateHandle;
    +	}
    +
    +	@Nonnull
    +	@Override
    +	public KeyGroupRange getKeyGroupRange() {
    +		return keyGroupRange;
    +	}
    +
    +	@Override
    +	public void discardState() throws Exception {
    +		directoryStateHandle.discardState();
    +	}
    +
    +	@Override
    +	public long getStateSize() {
    +		return directoryStateHandle.getStateSize();
    +	}
    +
    +	@Override
    +	public KeyedStateHandle getIntersection(KeyGroupRange keyGroupRange) {
    +		return keyGroupRange.getIntersection(keyGroupRange).getNumberOfKeyGroups() > 0 ? this : null;
    --- End diff --
    
    This looks like a bug. `keyGroupRange.getIntersection(keyGroupRange)` should be `keyGroupRange`. I think in general it is a good idea to not shadow local fields by function parameters.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168485458
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStoreImpl.java ---
    @@ -0,0 +1,298 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
    +
    +import java.io.File;
    +import java.io.IOException;
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.HashMap;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
    +
    +/**
    + * Main implementation of a {@link TaskLocalStateStore}.
    + */
    +public class TaskLocalStateStoreImpl implements TaskLocalStateStore {
    +
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStoreImpl.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	@VisibleForTesting
    +	static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    +
    +	/** JobID from the owning subtask. */
    +	@Nonnull
    +	private final JobID jobID;
    +
    +	/** AllocationID of the owning slot. */
    +	@Nonnull
    +	private final AllocationID allocationID;
    +
    +	/** JobVertexID of the owning subtask. */
    +	@Nonnull
    +	private final JobVertexID jobVertexID;
    +
    +	/** Subtask index of the owning subtask. */
    +	@Nonnegative
    +	private final int subtaskIndex;
    +
    +	/** The configured mode for local recovery. */
    +	@Nonnull
    +	private final LocalRecoveryConfig localRecoveryConfig;
    +
    +	/** Executor that runs the discarding of released state objects. */
    +	@Nonnull
    +	private final Executor discardExecutor;
    +
    +	/** Lock for synchronisation on the storage map and the discarded status. */
    +	@Nonnull
    +	private final Object lock;
    +
    +	/** Status flag if this store was already discarded. */
    +	@GuardedBy("lock")
    +	private boolean discarded;
    +
    +	/** Maps checkpoint ids to local TaskStateSnapshots. */
    +	@Nonnull
    +	@GuardedBy("lock")
    +	private final SortedMap<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
    +	public TaskLocalStateStoreImpl(
    +		@Nonnull JobID jobID,
    +		@Nonnull AllocationID allocationID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex,
    +		@Nonnull LocalRecoveryConfig localRecoveryConfig,
    +		@Nonnull Executor discardExecutor) {
    +
    +		this.jobID = jobID;
    +		this.allocationID = allocationID;
    +		this.jobVertexID = jobVertexID;
    +		this.subtaskIndex = subtaskIndex;
    +		this.discardExecutor = discardExecutor;
    +		this.lock = new Object();
    +		this.storedTaskStateByCheckpointID = new TreeMap<>();
    +		this.discarded = false;
    +		this.localRecoveryConfig = localRecoveryConfig;
    +	}
    +
    +	@Override
    +	public void storeLocalState(
    +		@Nonnegative long checkpointId,
    +		@Nullable TaskStateSnapshot localState) {
    +
    +		if (localState == null) {
    +			localState = NULL_DUMMY;
    +		}
    +
    +		LOG.info("Storing local state for checkpoint {}.", checkpointId);
    +		LOG.debug("Local state for checkpoint {} is {}.", checkpointId, localState);
    +
    +		Map<Long, TaskStateSnapshot> toDiscard = new HashMap<>(MAX_RETAINED_SNAPSHOTS);
    +
    +		synchronized (lock) {
    +			if (discarded) {
    +				// we ignore late stores and simply discard the state.
    +				toDiscard.put(checkpointId, localState);
    +			} else {
    +				TaskStateSnapshot previous =
    +					storedTaskStateByCheckpointID.put(checkpointId, localState);
    +
    +				if (previous != null) {
    +					toDiscard.put(checkpointId, previous);
    +				}
    +
    +				// remove from history.
    +				while (storedTaskStateByCheckpointID.size() > MAX_RETAINED_SNAPSHOTS) {
    +					Long removeCheckpointID = storedTaskStateByCheckpointID.firstKey();
    --- End diff --
    
    Yes, this pruning can be removed without replacement.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162582649
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -1041,7 +1045,13 @@ public void restore(Collection<KeyedStateHandle> restoreState) throws Exception
     
     	@Override
     	public void notifyCheckpointComplete(long completedCheckpointId) {
    +
    +		if (!enableIncrementalCheckpointing) {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170253088
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -1629,9 +1630,9 @@ private void closeLocalRegistry() {
     
     		private boolean isWithLocalRecovery(
     			CheckpointOptions.CheckpointType checkpointType,
    -			RocksDBStateBackend.LocalRecoveryMode recoveryMode) {
    +			LocalRecoveryConfig.LocalRecoveryMode recoveryMode) {
     			// we use local recovery when it is activated and we are not taking a savepoint.
    -			return RocksDBStateBackend.LocalRecoveryMode.ENABLE_FILE_BASED.equals(recoveryMode)
    +			return LocalRecoveryConfig.LocalRecoveryMode.ENABLE_FILE_BASED.equals(recoveryMode)
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168509527
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamWithResultProvider.java ---
    @@ -0,0 +1,175 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.state.filesystem.FixFileFsStateOutputStream;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import java.io.Closeable;
    +import java.io.File;
    +import java.io.IOException;
    +
    +/**
    + * Interface that provides access to a CheckpointStateOutputStream and a method to provide the {@link SnapshotResult}.
    + * This abstracts from different ways that a result is obtained from checkpoint output streams.
    + */
    +public interface CheckpointStreamWithResultProvider extends Closeable {
    +
    +	SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +		KeyGroupRangeOffsets kgOffs) throws IOException;
    +
    +	CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream();
    +
    +	@Override
    +	default void close() throws IOException {
    +		CheckpointStreamFactory.CheckpointStateOutputStream outputStream = getCheckpointOutputStream();
    +		if (outputStream != null) {
    +			outputStream.close();
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that only creates the
    +	 * primary/remote/jm-owned state.
    +	 */
    +	class PrimaryStreamOnly implements CheckpointStreamWithResultProvider {
    +
    +		private final CheckpointStreamFactory.CheckpointStateOutputStream outputStream;
    +
    +		public PrimaryStreamOnly(CheckpointStreamFactory.CheckpointStateOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets kgOffs) throws IOException {
    +
    +			StreamStateHandle streamStateHandle = outputStream.closeAndGetHandle();
    +			KeyGroupsStateHandle primaryStateHandle = new KeyGroupsStateHandle(kgOffs, streamStateHandle);
    +			return new SnapshotResult<>(primaryStateHandle, null);
    +		}
    +
    +		@Override
    +		public CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that creates both, the
    +	 * primary/remote/jm-owned state and the secondary/local/tm-owned state.
    +	 */
    +	class PrimaryAndSecondaryStream implements CheckpointStreamWithResultProvider {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		private final DuplicatingCheckpointOutputStream outputStream;
    +
    +		public PrimaryAndSecondaryStream(DuplicatingCheckpointOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets offsets) throws IOException {
    +
    +			final StreamStateHandle primaryStreamStateHandle;
    +
    +			try {
    +				primaryStreamStateHandle = outputStream.closeAndGetPrimaryHandle();
    +			} catch (IOException primaryEx) {
    +				try {
    +					outputStream.close();
    +				} catch (IOException closeEx) {
    +					primaryEx = ExceptionUtils.firstOrSuppressed(closeEx, primaryEx);
    +				}
    +				throw primaryEx;
    +			}
    +
    +			StreamStateHandle secondaryStreamStateHandle = null;
    +
    +			try {
    +				secondaryStreamStateHandle = outputStream.closeAndGetSecondaryHandle();
    +			} catch (IOException secondaryEx) {
    +				LOG.warn("Exception from secondary/local checkpoint stream.", secondaryEx);
    +			}
    +
    +			if (primaryStreamStateHandle != null) {
    +
    +				final KeyGroupsStateHandle primaryKeyGroupsStateHandle =
    +					new KeyGroupsStateHandle(offsets, primaryStreamStateHandle);
    +
    +				final KeyGroupsStateHandle secondaryKeyGroupsStateHandle = secondaryStreamStateHandle != null ?
    +					new KeyGroupsStateHandle(offsets, secondaryStreamStateHandle) :
    +					null;
    +
    +				return new SnapshotResult<>(primaryKeyGroupsStateHandle, secondaryKeyGroupsStateHandle);
    +			}
    +
    +			return null;
    +		}
    +
    +		@Override
    +		public DuplicatingCheckpointOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	class Factory {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170271924
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileBasedStateOutputStream.java ---
    @@ -0,0 +1,155 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state.filesystem;
    +
    +import org.apache.flink.core.fs.FSDataOutputStream;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.FileSystem.WriteMode;
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.state.CheckpointStreamFactory.CheckpointStateOutputStream;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.IOException;
    +
    +import static org.apache.flink.util.Preconditions.checkNotNull;
    +
    +/**
    + * A {@link CheckpointStateOutputStream} that writes into a specified file and
    + * returns a {@link FileStateHandle} upon closing.
    + *
    + * <p>Unlike the {@link org.apache.flink.runtime.state.filesystem.FsCheckpointStreamFactory.FsCheckpointStateOutputStream},
    + * this stream does not have a threshold below which it returns a memory byte stream handle,
    + * and does not create random files, but writes to a specified file.
    + */
    +public final class FileBasedStateOutputStream extends CheckpointStateOutputStream {
    +
    +	private static final Logger LOG = LoggerFactory.getLogger(FileBasedStateOutputStream.class);
    +
    +	// ------------------------------------------------------------------------
    +
    +	private final FSDataOutputStream out;
    +
    +	private final Path path;
    +
    +	private final FileSystem fileSystem;
    +
    +	private volatile boolean closed;
    +
    +
    +	public FileBasedStateOutputStream(FileSystem fileSystem, Path path) throws IOException {
    +		this.fileSystem = checkNotNull(fileSystem);
    +		this.path = checkNotNull(path);
    +
    +		this.out = fileSystem.create(path, WriteMode.NO_OVERWRITE);
    +	}
    +
    +	// ------------------------------------------------------------------------
    +	//  I/O
    +	// ------------------------------------------------------------------------
    +
    +	@Override
    +	public final void write(int b) throws IOException {
    +		out.write(b);
    --- End diff --
    
    The second one does not 4 bytes as the argument would make you think, but extracts the lower byte from the int. From the doc of `OutputStream`:
    ```
        * Writes the specified byte to this output stream. The general
         * contract for <code>write</code> is that one byte is written
         * to the output stream.
    (...)
    ```


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168479650
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -55,25 +54,24 @@
     	/** */
     	private final Map<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
     
    -	/** This is the base directory for all local state of the subtask that owns this {@link TaskLocalStateStore}. */
    -	private final File subtaskLocalStateBaseDirectory;
    +	/** The root directories for all local state of this {@link TaskLocalStateStore}. */
    +	private final File[] localStateRootDirectories;
     
     	public TaskLocalStateStore(
     		JobID jobID,
     		JobVertexID jobVertexID,
     		int subtaskIndex,
    -		File localStateRootDir) {
    +		File[] localStateRootDirectories) {
    --- End diff --
    
    Indentation wrong.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162585199
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/util/AbstractStreamOperatorTestHarness.java ---
    @@ -132,32 +141,49 @@ public AbstractStreamOperatorTestHarness(
     				1024,
     				new Configuration(),
     				new ExecutionConfig(),
    +				new TestTaskStateManager(),
     				maxParallelism,
     				parallelism,
     				subtaskIndex));
     	}
     
     	public AbstractStreamOperatorTestHarness(
     			StreamOperator<OUT> operator,
    -			final Environment environment) throws Exception {
    +			Environment env) throws Exception {
     		this.operator = operator;
     		this.outputList = new ConcurrentLinkedQueue<>();
     		this.sideOutputLists = new HashMap<>();
     
    -		Configuration underlyingConfig = environment.getTaskConfiguration();
    +		Configuration underlyingConfig = env.getTaskConfiguration();
     		this.config = new StreamConfig(underlyingConfig);
     		this.config.setCheckpointingEnabled(true);
     		this.config.setOperatorID(new OperatorID());
    -		this.executionConfig = environment.getExecutionConfig();
    +		this.executionConfig = env.getExecutionConfig();
     		this.closableRegistry = new CloseableRegistry();
     		this.checkpointLock = new Object();
     
    -		this.environment = Preconditions.checkNotNull(environment);
    +		Preconditions.checkNotNull(env);
    +
    +		MockUtil mockUtil = new MockUtil();
    +
    +		if (!mockUtil.isMock(env) && !mockUtil.isSpy(env)) {
    +			env = spy(env);
    +		}
    +
    +		this.environment = env;
    +
    +		this.taskStateManager = new TestTaskStateManager(
    +			env.getJobID(),
    +			env.getExecutionId());
    +
    +		when(this.environment.getTaskStateManager()).thenReturn(this.taskStateManager);
    --- End diff --
    
    👍 I made all tests use the `MockEnvironment`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160724253
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamTaskStateManagerImpl.java ---
    @@ -0,0 +1,640 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.api.common.TaskInfo;
    +import org.apache.flink.api.common.typeutils.TypeSerializer;
    +import org.apache.flink.api.java.tuple.Tuple2;
    +import org.apache.flink.core.fs.CloseableRegistry;
    +import org.apache.flink.core.fs.FSDataInputStream;
    +import org.apache.flink.core.memory.DataInputViewStreamWrapper;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.execution.Environment;
    +import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
    +import org.apache.flink.runtime.state.CheckpointStreamFactory;
    +import org.apache.flink.runtime.state.DefaultOperatorStateBackend;
    +import org.apache.flink.runtime.state.KeyGroupRange;
    +import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
    +import org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.KeyGroupsStateHandle;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.OperatorStateBackend;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.StateBackend;
    +import org.apache.flink.runtime.state.StatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.StreamStateHandle;
    +import org.apache.flink.runtime.state.TaskStateManager;
    +import org.apache.flink.runtime.util.OperatorSubtaskDescriptionText;
    +import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService;
    +import org.apache.flink.util.CloseableIterable;
    +import org.apache.flink.util.Preconditions;
    +
    +import org.apache.commons.io.IOUtils;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.IOException;
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.NoSuchElementException;
    +
    +/**
    + * This class is the main implementation of a {@link StreamTaskStateManager}. This class obtains the state to create
    + * {@link StreamOperatorStateContext} objects for stream operators from the {@link TaskStateManager} of the task that
    + * runs the stream task and hence the operator.
    + *
    + * <p>This implementation operates on top a {@link TaskStateManager}, from which it receives everything required to
    + * restore state in the backends from checkpoints or savepoints.
    + */
    +public class StreamTaskStateManagerImpl implements StreamTaskStateManager {
    +
    +	/**
    +	 * The environment of the task. This is required as parameter to construct state backends via their factory.
    +	 */
    +	private final Environment environment;
    +
    +	/** This processing time service is required to construct an internal timer service manager. */
    +	private final ProcessingTimeService processingTimeService;
    +
    +	/** The state manager of the tasks provides the information used to restore potential previous state. */
    +	private final TaskStateManager taskStateManager;
    +
    +	/** This object is the factory for everything related to state backends and checkpointing. */
    +	private final StateBackend stateBackend;
    +
    +	public StreamTaskStateManagerImpl(
    +		Environment environment,
    +		StateBackend stateBackend,
    +		ProcessingTimeService processingTimeService) {
    +
    +		this.environment = environment;
    +		this.taskStateManager = Preconditions.checkNotNull(environment.getTaskStateManager());
    +		this.stateBackend = Preconditions.checkNotNull(stateBackend);
    +		this.processingTimeService = processingTimeService;
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	@Override
    +	public StreamOperatorStateContext streamOperatorStateContext(
    +		AbstractStreamOperator<?> operator,
    +		TypeSerializer<?> keySerializer,
    +		CloseableRegistry streamTaskCloseableRegistry) throws Exception {
    +
    +		TaskInfo taskInfo = environment.getTaskInfo();
    +		OperatorSubtaskDescriptionText operatorSubtaskDescription =
    +			new OperatorSubtaskDescriptionText(
    +				operator.getOperatorID(),
    +				operator.getClass(),
    +				taskInfo.getIndexOfThisSubtask(),
    +				taskInfo.getNumberOfParallelSubtasks());
    +
    +		final String operatorIdentifierText = operatorSubtaskDescription.toString();
    +
    +		final OperatorSubtaskState operatorSubtaskStateFromJobManager =
    +			taskStateManager.operatorStates(operator.getOperatorID());
    +
    +		final boolean restoring = (operatorSubtaskStateFromJobManager != null);
    +
    +		AbstractKeyedStateBackend<?> keyedStatedBackend = null;
    +		OperatorStateBackend operatorStateBackend = null;
    +		CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs = null;
    +		CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs = null;
    +		CheckpointStreamFactory checkpointStreamFactory = null;
    +		InternalTimeServiceManager<?, ?> timeServiceManager = null;
    +
    +		try {
    +
    +			// -------------- Keyed State Backend --------------
    +			keyedStatedBackend = keyedStatedBackend(
    +				keySerializer,
    +				operatorIdentifierText,
    +				operatorSubtaskStateFromJobManager,
    +				streamTaskCloseableRegistry);
    +
    +			// -------------- Operator State Backend --------------
    +			operatorStateBackend = operatorStateBackend(
    +				operatorIdentifierText,
    +				operatorSubtaskStateFromJobManager,
    +				streamTaskCloseableRegistry);
    +
    +			// -------------- Raw State Streams --------------
    +			rawKeyedStateInputs = rawKeyedStateInputs(operatorSubtaskStateFromJobManager);
    +			streamTaskCloseableRegistry.registerCloseable(rawKeyedStateInputs);
    +
    +			rawOperatorStateInputs = rawOperatorStateInputs(operatorSubtaskStateFromJobManager);
    +			streamTaskCloseableRegistry.registerCloseable(rawOperatorStateInputs);
    +
    +			// -------------- Checkpoint Stream Factory --------------
    +			checkpointStreamFactory = streamFactory(operatorIdentifierText);
    +			streamTaskCloseableRegistry.registerCloseable(checkpointStreamFactory);
    +
    +			// -------------- Internal Timer Service Manager --------------
    +			timeServiceManager = internalTimeServiceManager(keyedStatedBackend, operator, rawKeyedStateInputs);
    +
    +			// -------------- Preparing return value --------------
    +
    +			return new StreamOperatorStateContextImpl(
    +				restoring,
    +				operatorStateBackend,
    +				keyedStatedBackend,
    +				timeServiceManager,
    +				rawOperatorStateInputs,
    +				rawKeyedStateInputs,
    +				checkpointStreamFactory);
    +		} catch (Exception ex) {
    +
    +			// cleanup if something went wrong before results got published.
    +			if (streamTaskCloseableRegistry.unregisterCloseable(keyedStatedBackend)) {
    +				IOUtils.closeQuietly(keyedStatedBackend);
    +			}
    +
    +			if (streamTaskCloseableRegistry.unregisterCloseable(operatorStateBackend)) {
    +				IOUtils.closeQuietly(keyedStatedBackend);
    +			}
    +
    +			if (streamTaskCloseableRegistry.unregisterCloseable(rawKeyedStateInputs)) {
    +				IOUtils.closeQuietly(rawKeyedStateInputs);
    +			}
    +
    +			if (streamTaskCloseableRegistry.unregisterCloseable(rawOperatorStateInputs)) {
    +				IOUtils.closeQuietly(rawOperatorStateInputs);
    +			}
    +
    +			if (streamTaskCloseableRegistry.unregisterCloseable(rawOperatorStateInputs)) {
    +				IOUtils.closeQuietly(rawOperatorStateInputs);
    +			}
    +
    +			throw new Exception("Exception while creating StreamOperatorStateContext.", ex);
    +		}
    +	}
    +
    +	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<>(
    +			keyedStatedBackend.getNumberOfKeyGroups(),
    +			keyGroupRange,
    +			keyContext,
    +			processingTimeService);
    +
    +		// 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(
    +				new DataInputViewStreamWrapper(streamProvider.getStream()),
    +				keyGroupIdx, environment.getUserClassLoader());
    +		}
    +
    +		return timeServiceManager;
    +	}
    +
    +	protected OperatorStateBackend operatorStateBackend(
    +		String operatorIdentifierText,
    +		OperatorSubtaskState operatorSubtaskStateFromJobManager,
    +		CloseableRegistry backendCloseableRegistry) throws Exception {
    +
    +		//TODO search in local state for a local recovery opportunity.
    +
    +		return createOperatorStateBackendFromJobManagerState(
    +			operatorIdentifierText,
    +			operatorSubtaskStateFromJobManager,
    +			backendCloseableRegistry);
    +	}
    +
    +	protected <K> AbstractKeyedStateBackend<K> keyedStatedBackend(
    --- End diff --
    
    this is more of a general remark, can you name the methods as `getSomething()` or `createSomething()` to communicate what they are actually doing? I had to couple of times go to the implementation because it wasn't (at least for me :( ) clear from the context which one of it is it.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169727315
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotDirectory.java ---
    @@ -0,0 +1,172 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.FileStatus;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.IOException;
    +import java.util.concurrent.atomic.AtomicReference;
    +
    +/**
    + * This class represents a directory that is the target for a state snapshot. This class provides some method that
    + * simplify resource management when dealing with such directories, e.g. it can produce a {@link DirectoryStateHandle}
    + * when the snapshot is completed and disposal considers whether or not a snapshot was already completed. For a
    + * completed snapshot, the ownership for cleanup is transferred to the created directory state handle. For incomplete
    + * snapshots, calling {@link #deleteIfIncompleteSnapshot()} will delete the underlying directory resource.
    + */
    +public class SnapshotDirectory {
    +
    +	/**
    +	 * Lifecycle stages of a snapshot directory.
    +	 */
    +	enum State {
    +		ONGOING, COMPLETED, DELETED
    +	}
    +
    +	/** This path describes the underlying directory for the snapshot. */
    +	@Nonnull
    +	private final Path directory;
    +
    +	/** The filesystem that contains the snapshot directory. */
    +	@Nonnull
    +	private final FileSystem fileSystem;
    --- End diff --
    
    Do we need `FileSystem` if we have `directory` which is of type `Path`?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160725404
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java ---
    @@ -214,49 +217,55 @@ public MetricGroup getMetricGroup() {
     	}
     
     	@Override
    -	public final void initializeState(OperatorSubtaskState stateHandles) throws Exception {
    +	public final void initializeState() throws Exception {
     
    -		Collection<KeyedStateHandle> keyedStateHandlesRaw = null;
    -		Collection<OperatorStateHandle> operatorStateHandlesRaw = null;
    -		Collection<OperatorStateHandle> operatorStateHandlesBackend = null;
    +		final TypeSerializer<?> keySerializer = config.getStateKeySerializer(getUserCodeClassloader());
     
    -		boolean restoring = (null != stateHandles);
    +		final StreamTask<?, ?> containingTask =
    +			Preconditions.checkNotNull(getContainingTask());
    +		final CloseableRegistry streamTaskCloseableRegistry =
    +			Preconditions.checkNotNull(containingTask.getCancelables());
    +		final StreamTaskStateManager streamTaskStateManager =
    +			Preconditions.checkNotNull(containingTask.getStreamTaskStateManager());
     
    -		initKeyedState(); //TODO we should move the actual initialization of this from StreamTask to this class
    -
    -		if (getKeyedStateBackend() != null && timeServiceManager == null) {
    -			timeServiceManager = new InternalTimeServiceManager<>(
    -				getKeyedStateBackend().getNumberOfKeyGroups(),
    -				getKeyedStateBackend().getKeyGroupRange(),
    +		final StreamOperatorStateContext context =
    --- End diff --
    
    Do we need this one additional level of indirection in form of `StreamOperatorStateContext`? Couldn't `streamTaskStateManager` return the `StateInitializationContext` directly? Or do you think otherwise? (I'm not sure on this comment)


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168736826
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    Hmm...strange Sorry, but I still cannot see the proper line. If this is not working, maybe just write the correct line number to a comment?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162583540
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/StateUtil.java ---
    @@ -71,4 +72,22 @@ public static void discardStateFuture(RunnableFuture<? extends StateObject> stat
     			}
     		}
     	}
    +
    +	/**
    +	 * Discards the given state collection future by first trying to cancel it. If this is not possible, then
    +	 * the state object contained in the future is calculated and afterwards discarded.
    +	 *
    +	 * @param stateCollectionFuture to be discarded
    +	 * @throws Exception if the discard operation failed
    +	 */
    +	public static void discardStateCollectionFuture(
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168168418
  
    --- Diff: flink-core/src/main/java/org/apache/flink/configuration/ConfigConstants.java ---
    @@ -316,6 +316,11 @@
     	 */
     	public static final String TASK_MANAGER_REFUSED_REGISTRATION_PAUSE = "taskmanager.refused-registration-pause";
     
    +	/**
    +	 * The config parameter defining the root directories for storing file-based state for local recovery.
    +	 */
    +	public static final String TASK_MANAGER_LOCAL_STATE_ROOT_DIR_KEY = "taskmanager.local-state-root.dir";
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168923652
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    Strange sorry, my comment does not match the highlighted lines when line number > 1400, the problem I described above is on line: 1878 
    ```java
    private void writeKVStateMetaData() throws IOException {
      // ...
            for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    				stateBackend.kvStateInformation.entrySet()) {
            }
      //...
    }
    ```


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168501599
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -19,92 +19,224 @@
     package org.apache.flink.runtime.state;
     
     import org.apache.flink.api.common.JobID;
    -import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
     import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
     import org.apache.flink.runtime.jobgraph.JobVertexID;
    -import org.apache.flink.util.ExceptionUtils;
     
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
     import javax.annotation.Nonnull;
     import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
     
     import java.io.File;
    -import java.util.HashMap;
    +import java.util.Arrays;
    +import java.util.Iterator;
     import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
     
     /**
      * This class will service as a task-manager-level local storage for local checkpointed state. The purpose is to provide
      * access to a state that is stored locally for a faster recovery compared to the state that is stored remotely in a
      * stable store DFS. For now, this storage is only complementary to the stable storage and local state is typically
      * lost in case of machine failures. In such cases (and others), client code of this class must fall back to using the
      * slower but highly available store.
    - *
    - * TODO this is currently a placeholder / mock that still must be implemented!
      */
     public class TaskLocalStateStore {
     
    -	/** */
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStore.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	private static final int MAX_RETAINED_SNAPSHOTS = 5;
    --- End diff --
    
    How did you come up with this number? Should this be configurable?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165342796
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/util/JvmExitOnFatalErrorTest.java ---
    @@ -137,9 +138,12 @@ public static void main(String[] args) throws Exception {
     
     			System.err.println("creating task");
     
    +			TemporaryFolder temporaryFolder = new TemporaryFolder();
    --- End diff --
    
    Making it a `ClassRule` or a `Rule` will take care of the clean up.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168169074
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTerminationTest.java ---
    @@ -227,21 +229,13 @@ protected void cleanup() throws Exception {
     			// has been stopped
     			CLEANUP_LATCH.trigger();
     
    -			// wait until handle async exception has been called to proceed with the termination of the
    -			// StreamTask
    -			HANDLE_ASYNC_EXCEPTION_LATCH.await();
    +			// wait until all async checkpoint threads are terminated, so that no more exceptions can be reported
    +			Assert.assertTrue(getAsyncOperationsThreadPool().awaitTermination(30L, TimeUnit.SECONDS));
    --- End diff --
    
    I happens through the StreamTask, before cleanups.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168528716
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskStateManagerImplTest.java ---
    @@ -109,10 +111,10 @@ public void testStateReportingAndRetrieving() {
     			executionAttemptID,
     			checkpointResponderMock,
     			taskRestore);
    -
    -		Assert.assertTrue(jmOperatorSubtaskState_1 == taskStateManager.operatorStates(operatorID_1));
    -		Assert.assertTrue(jmOperatorSubtaskState_2 == taskStateManager.operatorStates(operatorID_2));
    -		Assert.assertNull(taskStateManager.operatorStates(operatorID_3));
    +//TODO
    +//		Assert.assertTrue(jmOperatorSubtaskState_1 == taskStateManager.prioritizedOperatorState(operatorID_1));
    +//		Assert.assertTrue(jmOperatorSubtaskState_2 == taskStateManager.prioritizedOperatorState(operatorID_2));
    --- End diff --
    
    Fixed in later commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168514963
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskStateManagerImplTest.java ---
    @@ -148,22 +150,22 @@ public void testForwardingSubtaskLocalStateBaseDirFromLocalStateStore() throws I
     				taskLocalStateStore);
     
     			LocalRecoveryDirectoryProvider directoryProviderFromTaskLocalStateStore =
    -				taskLocalStateStore.createLocalRecoveryRootDirectoryProvider();
    +				taskLocalStateStore.getLocalRecoveryRootDirectoryProvider();
     
     			LocalRecoveryDirectoryProvider directoryProviderFromTaskStateManager =
     				taskStateManager.createLocalRecoveryRootDirectoryProvider();
     
     
     			for (int i = 0; i < 10; ++i) {
     				Assert.assertEquals(rootDirs[i % rootDirs.length],
    -					directoryProviderFromTaskLocalStateStore.nextRootDirectory());
    +					directoryProviderFromTaskLocalStateStore.rootDirectory(i));
     				Assert.assertEquals(rootDirs[i % rootDirs.length],
    -					directoryProviderFromTaskStateManager.nextRootDirectory());
    +					directoryProviderFromTaskStateManager.rootDirectory(i));
     			}
     
    -			Assert.assertEquals(
    -				directoryProviderFromTaskLocalStateStore.getSubtaskSpecificPath(),
    -				directoryProviderFromTaskStateManager.getSubtaskSpecificPath());
    +//			Assert.assertEquals(
    +//				directoryProviderFromTaskLocalStateStore.getSubtaskSpecificPath(),
    +//				directoryProviderFromTaskStateManager.getSubtaskSpecificPath());
    --- End diff --
    
    Why did you comment this part out? Can it be removed?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168923464
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    I know this part is an old code block and wasn't be change in this PR. But it is a bit odds to me, kvStateInformation is not a copied object and it can be changed when writeKVStateMetaData()is invoking ... If I am not wrong, this is a serious bug. Am I misunderstand something?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169047628
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159865703
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java ---
    @@ -669,31 +676,14 @@ else if (current == ExecutionState.CANCELING) {
     			Environment env = new RuntimeEnvironment(
     				jobId, vertexId, executionId, executionConfig, taskInfo,
     				jobConfiguration, taskConfiguration, userCodeClassLoader,
    -				memoryManager, ioManager, broadcastVariableManager,
    +				memoryManager, ioManager, broadcastVariableManager, taskStateManager,
    --- End diff --
    
    please reformat this call


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168168381
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -46,26 +52,63 @@
     	/** */
     	private final int subtaskIndex;
     
    +	/** */
    +	private final Map<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
    +	/** This is the base directory for all local state of the subtask that owns this {@link TaskLocalStateStore}. */
    +	private final File subtaskLocalStateBaseDirectory;
    +
     	public TaskLocalStateStore(
     		JobID jobID,
     		JobVertexID jobVertexID,
    -		int subtaskIndex) {
    +		int subtaskIndex,
    +		File localStateRootDir) {
     
     		this.jobID = jobID;
     		this.jobVertexID = jobVertexID;
     		this.subtaskIndex = subtaskIndex;
    +		this.storedTaskStateByCheckpointID = new HashMap<>();
    +		this.subtaskLocalStateBaseDirectory =
    +			new File(localStateRootDir, createSubtaskPath(jobID, jobVertexID, subtaskIndex));
    +	}
    +
    +	static String createSubtaskPath(JobID jobID, JobVertexID jobVertexID, int subtaskIndex) {
    +		return "jid-" + jobID + "_vtx-" + jobVertexID + "_sti-" + subtaskIndex;
     	}
     
     	public void storeLocalState(
     		@Nonnull CheckpointMetaData checkpointMetaData,
     		@Nullable TaskStateSnapshot localState) {
     
    -		if (localState != null) {
    -			throw new UnsupportedOperationException("Implement this before actually providing local state!");
    +		TaskStateSnapshot previous =
    +			storedTaskStateByCheckpointID.put(checkpointMetaData.getCheckpointId(), localState);
    +
    +		if (previous != null) {
    +			throw new IllegalStateException("Found previously registered local state for checkpoint with id " +
    +				checkpointMetaData.getCheckpointId() + "! This indicated a problem.");
     		}
     	}
     
    -	public void dispose() {
    -		//TODO
    +	public void dispose() throws Exception {
    +
    +		Exception collectedException = null;
    +
    +		for (TaskStateSnapshot snapshot : storedTaskStateByCheckpointID.values()) {
    +			try {
    +				snapshot.discardState();
    +			} catch (Exception discardEx) {
    +				collectedException = ExceptionUtils.firstOrSuppressed(discardEx, collectedException);
    +			}
    +		}
    +
    +		if (collectedException != null) {
    +			throw collectedException;
    +		}
    +
    +		FileUtils.deleteDirectoryQuietly(subtaskLocalStateBaseDirectory);
    --- End diff --
    
    This already works different in later commits.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160943448
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java ---
    @@ -316,7 +292,38 @@ private void initOperatorState(Collection<OperatorStateHandle> operatorStateHand
     	 * @throws Exception An exception in this method causes the operator to fail.
     	 */
     	@Override
    -	public void close() throws Exception {}
    +	public void close() throws Exception {
    --- End diff --
    
    `org.apache.flink.streaming.runtime.operators.GenericWriteAheadSink#close` doesn't invoke `super.close()`. Is it also a bug?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170038223
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java ---
    @@ -46,9 +48,15 @@ public void testCreationFromConfig() throws Exception {
     
     		final Configuration config = new Configuration();
     
    -		final String rootDirString = "localStateRoot1,localStateRoot2,localStateRoot3";
    +		String tmpDir = System.getProperty("java.io.tmpdir") + File.separator;
    --- End diff --
    
    I think it's better to use a `TemporaryFolder` if there are any files created to ensure a proper clean up at the end of the test.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162583508
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/BarrierBufferTest.java ---
    @@ -1483,11 +1482,6 @@ long getLastReportedBytesBufferedInAlignment() {
     			return lastReportedBytesBufferedInAlignment;
     		}
     
    -		@Override
    --- End diff --
    
    This PR changed nothing about that, so it should be a separate change.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170038982
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskLocalStateStoreImplTest.java ---
    @@ -59,12 +59,18 @@ public void before() throws Exception {
     		this.subtaskIdx = 0;
     		this.rootDirs = new File[]{temporaryFolder.newFolder(), temporaryFolder.newFolder()};
     
    +		LocalRecoveryDirectoryProviderImpl directoryProvider =
    +			new LocalRecoveryDirectoryProviderImpl(rootDirs, jobID, allocationID, jobVertexID, subtaskIdx);
    +
    +		LocalRecoveryConfig localRecoveryConfig =
    +			new LocalRecoveryConfig(LocalRecoveryConfig.LocalRecoveryMode.DISABLED, directoryProvider);
    --- End diff --
    
    It's a bit strange that we have to pass in a `directoryProvider`, even though the `LocalRecoveryMode` is disabled. I think it would be better to have something like `LocalRecoveryConfig.disabled()` and `LocalRecoveryConfig.fileBased(directoryProvider)`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168483656
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java ---
    @@ -100,21 +109,27 @@ public void testSubtaskStateStoreDirectoryCreation() throws Exception {
     		int subtaskIdx = 42;
     		TemporaryFolder tmp = new TemporaryFolder();
     		try {
    +
     			tmp.create();
    -			File localStoreDir = tmp.newFolder();
    +			File[] rootDirs = {tmp.newFolder(), tmp.newFolder(), tmp.newFolder()};
     			TaskExecutorLocalStateStoresManager storesManager =
    -				new TaskExecutorLocalStateStoresManager(localStoreDir);
    +				new TaskExecutorLocalStateStoresManager(rootDirs);
     
     			TaskLocalStateStore taskLocalStateStore =
     				storesManager.localStateStoreForTask(jobID, jobVertexID, subtaskIdx);
     
    -			Assert.assertEquals(
    -				new File(localStoreDir, TaskLocalStateStore.createSubtaskPath(jobID, jobVertexID, subtaskIdx)),
    -				taskLocalStateStore.getSubtaskLocalStateBaseDirectory());
    +			LocalRecoveryDirectoryProvider directoryProvider =
    +				taskLocalStateStore.createLocalRecoveryRootDirectoryProvider();
    +
    +			for (int i = 0; i < 10; ++i) {
    +				Assert.assertEquals(
    +					rootDirs[(i & Integer.MAX_VALUE) % rootDirs.length],
    +					directoryProvider.nextRootDirectory());
    +			}
    --- End diff --
    
    Aren't we mixing here things which are tested. I think the `LocalRecoveryDirectoryProvider` deserves its dedicated test.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168779711
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java ---
    @@ -0,0 +1,265 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.checkpoint;
    +
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.Preconditions;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.util.ArrayList;
    +import java.util.Collections;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.function.BiFunction;
    +
    +/**
    + * This class is a wrapper over multiple alternative {@link OperatorSubtaskState} that are (partial) substitutes for
    + * each other and imposes a priority ordering over all alternatives for the different states which define an order in
    + * which the operator should attempt to restore the state from them. One OperatorSubtaskState is considered as the
    + * "ground truth" about which state should be represented. Alternatives may be complete or partial substitutes for
    + * the "ground truth" with a higher priority (if they had a lower alternative, they would not really be alternatives).
    + * Substitution is determined on a per-sub-state basis.
    + */
    +public class PrioritizedOperatorSubtaskState {
    +
    +	/** Singleton instance for an empty, non-restored operator state. */
    +	private static final PrioritizedOperatorSubtaskState EMPTY_NON_RESTORED_INSTANCE =
    +		new PrioritizedOperatorSubtaskState(new OperatorSubtaskState(), Collections.emptyList(), false);
    +
    +	/** List of prioritized snapshot alternatives for managed operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedManagedOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for raw operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedRawOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for managed keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedManagedKeyedState;
    +
    +	/** List of prioritized snapshot alternatives for raw keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedRawKeyedState;
    +
    +	/** Signal flag if this represents state for a restored operator. */
    +	private final boolean restored;
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority) {
    +		this(jobManagerState, alternativesByPriority, true);
    +	}
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority,
    +		boolean restored) {
    +
    +		Preconditions.checkNotNull(jobManagerState, "Job manager state is null.");
    +		int size = Preconditions.checkNotNull(alternativesByPriority, "Alternative states are null.").size();
    +
    +		this.restored = restored;
    +
    +		List<StateObjectCollection<OperatorStateHandle>> managedOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> managedKeyedAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<OperatorStateHandle>> rawOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> rawKeyedAlternatives = new ArrayList<>(size);
    +
    +		for (OperatorSubtaskState subtaskState : alternativesByPriority) {
    +
    +			if (subtaskState != null) {
    +				managedKeyedAlternatives.add(subtaskState.getManagedKeyedState());
    +				rawKeyedAlternatives.add(subtaskState.getRawKeyedState());
    +				managedOperatorAlternatives.add(subtaskState.getManagedOperatorState());
    +				rawOperatorAlternatives.add(subtaskState.getRawOperatorState());
    +			}
    +		}
    +
    +		// Key-groups should match.
    +		BiFunction<KeyedStateHandle, KeyedStateHandle, Boolean> keyedStateApprover =
    +			(ref, alt) -> ref.getKeyGroupRange().equals(alt.getKeyGroupRange());
    +
    +		// State meta data should match.
    +		BiFunction<OperatorStateHandle, OperatorStateHandle, Boolean> operatorStateApprover =
    +			(ref, alt) -> ref.getStateNameToPartitionOffsets().equals(alt.getStateNameToPartitionOffsets());
    +
    +		this.prioritizedManagedKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedKeyedState(),
    +			managedKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedRawKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawKeyedState(),
    +			rawKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedManagedOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedOperatorState(),
    +			managedOperatorAlternatives,
    +			operatorStateApprover);
    +
    +		this.prioritizedRawOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawOperatorState(),
    +			rawOperatorAlternatives,
    +			operatorStateApprover);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed operator state, in the order in which
    +	 * we should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedManagedOperatorState() {
    +		return prioritizedManagedOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw operator state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedRawOperatorState() {
    +		return prioritizedRawOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedManagedKeyedState() {
    +		return prioritizedManagedKeyedState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedRawKeyedState() {
    +		return prioritizedRawKeyedState.iterator();
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns the managed operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerManagedOperatorState() {
    +		return lastElement(prioritizedManagedOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the raw operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerRawOperatorState() {
    +		return lastElement(prioritizedRawOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the managed keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerManagedKeyedState() {
    +		return lastElement(prioritizedManagedKeyedState);
    +	}
    +
    +	/**
    +	 * Returns the raw keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerRawKeyedState() {
    +		return lastElement(prioritizedRawKeyedState);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns true if this was created for a restored operator, false otherwise. Restored operators are operators that
    +	 * participated in a previous checkpoint, even if they did not emit any state snapshots.
    +	 */
    +	public boolean isRestored() {
    +		return restored;
    +	}
    +
    +
    +	/**
    +	 * This helper method resolves the dependencies between the ground truth of the operator state obtained from the
    +	 * job manager and potential alternatives for recovery, e.g. from a task-local source.
    +	 */
    +	protected <T extends StateObject> List<StateObjectCollection<T>> resolvePrioritizedAlternatives(
    +		StateObjectCollection<T> jobManagerState,
    +		List<StateObjectCollection<T>> alternativesByPriority,
    +		BiFunction<T, T, Boolean> approveFun) {
    +
    +		// Nothing to resolve if there are no alternatives, or the ground truth has already no state, or if we can
    +		// assume that a rescaling happened because we find more than one handle in the JM state (this is more a sanity
    +		// check).
    +		if (alternativesByPriority == null
    +			|| alternativesByPriority.isEmpty()
    +			|| !jobManagerState.hasState()
    +			|| jobManagerState.size() != 1) {
    +
    --- End diff --
    
    Agreed.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169649287
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
    --- End diff --
    
    That might be true, but again this comment has nothing to do with this PR - only indentation was changed. I would strongly suggest that this should not become a collection for all possible suggestions to other changes. Opening a jira for this idea would be much better, because this thread is already very long and posting unrelated suggestions like this here make things more confusing and increases the likelihood that something will be overlooked. Would that be ok?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165349688
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -363,1691 +372,1780 @@ public int getKeyGroupPrefixBytes() {
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointOptions.CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at " +
    -					checkpointTimestamp + " . Returning null.");
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    -
    -		snapshotOperation.takeSnapshot();
    -
    -		return new FutureTask<SnapshotResult<KeyedStateHandle>>(
    -			() -> snapshotOperation.materializeSnapshot()
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    +		try {
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    +				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +				restoreOperation.restore(restoreState);
    +			} else {
    +				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +				restoreOperation.doRestore(restoreState);
     			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at " + timestamp +
    -					" . Returning null.");
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<SnapshotResult<KeyedStateHandle>> ioCallable =
    -			new AbstractAsyncCallableWithResources<SnapshotResult<KeyedStateHandle>>() {
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -				@Override
    -				public SnapshotResult<KeyedStateHandle> performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in last position.
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					KeyGroupsStateHandle stateHandle = snapshotOperation.getSnapshotResultStateHandle();
    -					return new SnapshotResult<>(stateHandle, null);
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot (" + streamFactory + ", synchronous part) in thread " +
    -			Thread.currentThread() + " took " + (System.currentTimeMillis() - startTime) + " ms.");
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    -
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(checkpointId, checkpointTimeStamp);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			FSDataInputStream inputStream = null;
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			try {
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		private void readStateData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    -			}
    -		}
    -	}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
    +			}
    +		}
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// read state data
    +			Path restoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			try {
    +				final Map<StateHandleID, StreamStateHandle> sstFiles =
    +					restoreStateHandle.getSharedState();
    +				final Map<StateHandleID, StreamStateHandle> miscFiles =
    +					restoreStateHandle.getPrivateState();
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				readAllStateData(sstFiles, restoreInstancePath);
    +				readAllStateData(miscFiles, restoreInstancePath);
     
    -		/** Local filesystem for the RocksDB backup. */
    -		private FileSystem backupFileSystem;
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
     
    -		/** Local path for the RocksDB backup. */
    -		private Path backupPath;
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					new ArrayList<>(1 + stateMetaInfoSnapshots.size());
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +				for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +						stateBackend.columnOptions);
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +					columnFamilyDescriptors.add(columnFamilyDescriptor);
    +					stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
    +				}
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +				if (hasExtraKeys) {
     
    -		private StreamStateHandle metaStateHandle = null;
    +					List<ColumnFamilyHandle> columnFamilyHandles =
    +						new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +					try (RocksDB restoreDb = stateBackend.openDB(
    +						restoreInstancePath.getPath(),
    +						columnFamilyDescriptors,
    +						columnFamilyHandles)) {
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +						try {
    +							// iterating only the requested descriptors automatically skips the default column family handle
    +							for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +								ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +								ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +								RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    -			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    +								Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +									stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			try {
    -				final byte[] buffer = new byte[8 * 1024];
    +								if (null == registeredStateMetaInfoEntry) {
     
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    +									RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +										new RegisteredKeyedBackendStateMetaInfo<>(
    +											stateMetaInfoSnapshot.getStateType(),
    +											stateMetaInfoSnapshot.getName(),
    +											stateMetaInfoSnapshot.getNamespaceSerializer(),
    +											stateMetaInfoSnapshot.getStateSerializer());
     
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(checkpointId, checkpointTimestamp);
    -				closeableRegistry.registerCloseable(outputStream);
    +									registeredStateMetaInfoEntry =
    +										new Tuple2<>(
    +											stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +											stateMetaInfo);
     
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +									stateBackend.kvStateInformation.put(
    +										stateMetaInfoSnapshot.getName(),
    +										registeredStateMetaInfoEntry);
    +								}
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +								ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +								try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    +									int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +									byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +									for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +										startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +									}
     
    -			} finally {
    +									iterator.seek(startKeyGroupPrefixBytes);
     
    -				if (closeableRegistry.unregisterCloseable(inputStream)) {
    -					inputStream.close();
    -				}
    +									while (iterator.isValid()) {
     
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
    -			}
    -		}
    +										int keyGroup = 0;
    +										for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +											keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +										}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    +										if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +											stateBackend.db.put(targetColumnFamilyHandle,
    +												iterator.key(), iterator.value());
    +										}
     
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(checkpointId, checkpointTimestamp);
    -				closeableRegistry.registerCloseable(outputStream);
    +										iterator.next();
    +									}
    +								} // releases native iterator resources
    +							}
    +						} finally {
    +							//release native tmp db column family resources
    +							for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles) {
    +								IOUtils.closeQuietly(columnFamilyHandle);
    +							}
    +						}
    +					} // releases native tmp db resources
    +				} else {
    +					// pick up again the old backend id, so the we can reference existing state
    +					stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +					LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +						stateBackend.operatorIdentifier, stateBackend.backendUID);
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +					// create hard links in the instance directory
    +					if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +						throw new IOException("Could not create RocksDB data directory.");
    +					}
     
    -				serializationProxy.write(out);
    +					createFileHardLinksInRestorePath(sstFiles, restoreInstancePath);
    +					createFileHardLinksInRestorePath(miscFiles, restoreInstancePath);
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					List<ColumnFamilyHandle> columnFamilyHandles =
    +						new ArrayList<>(1 + columnFamilyDescriptors.size());
    +
    +					stateBackend.db = stateBackend.openDB(
    +						stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +						columnFamilyDescriptors, columnFamilyHandles);
    +
    +					// extract and store the default column family which is located at the last index
    +					stateBackend.defaultColumnFamily = columnFamilyHandles.remove(columnFamilyHandles.size() - 1);
    +
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
    +
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +							new RegisteredKeyedBackendStateMetaInfo<>(
    +								stateMetaInfoSnapshot.getStateType(),
    +								stateMetaInfoSnapshot.getName(),
    +								stateMetaInfoSnapshot.getNamespaceSerializer(),
    +								stateMetaInfoSnapshot.getStateSerializer());
    +
    +						stateBackend.kvStateInformation.put(
    +							stateMetaInfoSnapshot.getName(),
    +							new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +					}
    +
    +					// use the restore sst files as the base for succeeding checkpoints
    +					synchronized (stateBackend.materializedSstFiles) {
    +						stateBackend.materializedSstFiles.put(restoreStateHandle.getCheckpointId(), sstFiles.keySet());
    +					}
    +
    +					stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
     				}
    -				return result;
     			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    +				FileSystem restoreFileSystem = restoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(restoreInstancePath)) {
    +					restoreFileSystem.delete(restoreInstancePath, true);
     				}
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    -
    -			final long lastCompletedCheckpoint;
    +		private void readAllStateData(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				readStateData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
    +		}
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +		private void createFileHardLinksInRestorePath(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    +			for (StateHandleID stateHandleID : stateHandleMap.keySet()) {
    +				String newSstFileName = stateHandleID.toString();
    +				File restoreFile = new File(restoreInstancePath.getPath(), newSstFileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath, newSstFileName);
    +				Files.createLink(targetFile.toPath(), restoreFile.toPath());
     			}
    +		}
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
     			}
     
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
    -		}
    -
    -		SnapshotResult<KeyedStateHandle> materializeSnapshot() throws Exception {
    -
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    -
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    -
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    -
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    -
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    +				if (!(rawStateHandle instanceof IncrementalKeyedStateHandle)) {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    -			}
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    +				IncrementalKeyedStateHandle keyedStateHandle = (IncrementalKeyedStateHandle) rawStateHandle;
    +
    +				restoreInstance(keyedStateHandle, hasExtraKeys);
     			}
    +		}
    +	}
     
    -			IncrementalKeyedStateHandle incrementalKeyedStateHandle = new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -			return new SnapshotResult<>(incrementalKeyedStateHandle, null);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB data base and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		void stop() {
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
     
    -		void releaseResources(boolean canceled) {
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
     
    -			dbLease.close();
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    -				}
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
     			}
     		}
    -	}
     
    -	@Override
    -	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(DEFAULT_COLUMN_FAMILY_NAME_BYTES, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		final ColumnFamilyHandle columnFamily;
     
     		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    -			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -
    -		if (!enableIncrementalCheckpointing) {
    -			return;
    -		}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -		synchronized (materializedSstFiles) {
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    -	}
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    -	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    -
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	/**
    +	 * Only visible for testing, DO NOT USE.
    +	 */
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
    +
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
    +
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
    +
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
     		}
     
    -		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    -		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		return count;
    +	}
     
    -			rocksDBKeyedStateBackend.createDB();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
     					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    -				}
    +				});
     			}
     		}
     
    -		/**
    -		 * Restore one key groups state handle.
    -		 */
    -		private void restoreKeyGroupsInStateHandle()
    -			throws IOException, StateMigrationException, RocksDBException {
    -			try {
    -				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    -				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    -				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    -				restoreKVStateMetaData();
    -				restoreKVStateData();
    -			} finally {
    -				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    -					IOUtils.closeQuietly(currentStateHandleInStream);
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
     				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
     			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    -		 *
    -		 * @throws IOException
    -		 * @throws ClassNotFoundException
    -		 * @throws RocksDBException
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			serializationProxy.read(currentStateHandleInView);
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -			// check for key serializer compatibility; this also reconfigures the
    -			// key serializer to be compatible, if it is required and is possible
    -			if (CompatibilityUtil.resolveCompatibilityResult(
    -				serializationProxy.getKeySerializer(),
    -				UnloadableDummyTypeSerializer.class,
    -				serializationProxy.getKeySerializerConfigSnapshot(),
    -				rocksDBKeyedStateBackend.keySerializer)
    -				.isRequiresMigration()) {
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator = heap.poll();
    +					newKVState = currentSubIterator.getIterator() != rocksIterator;
    +					detectNewKeyGroup(oldKey);
    +				}
    +			} else {
    +				IOUtils.closeQuietly(rocksIterator);
     
    -				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    -				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    -					"Aborting now since state migration is currently not available");
    +				if (heap.isEmpty()) {
    +					currentSubIterator = null;
    +					valid = false;
    +				} else {
    +					currentSubIterator = heap.poll();
    +					newKVState = true;
    +					detectNewKeyGroup(oldKey);
    +				}
     			}
    +		}
     
    -			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    -				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
    -
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    -				serializationProxy.getStateMetaInfoSnapshots();
    -			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    -			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
    +		private boolean isDifferentKeyGroup(byte[] a, byte[] b) {
    +			return 0 != compareKeyGroupsForByteArrays(a, b, keyGroupPrefixByteCount);
    +		}
     
    -			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
    +		private void detectNewKeyGroup(byte[] oldKey) {
    +			if (isDifferentKeyGroup(oldKey, currentSubIterator.currentKey)) {
    +				newKeyGroup = true;
    +			}
    +		}
     
    -				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    -					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
    +		/**
    +		 * Returns the key-group for the current key.
    +		 * @return key-group for the current key
    +		 */
    +		public int keyGroup() {
    +			int result = 0;
    +			//big endian decode
    +			for (int i = 0; i < keyGroupPrefixByteCount; ++i) {
    +				result <<= 8;
    +				result |= (currentSubIterator.currentKey[i] & 0xFF);
    +			}
    +			return result;
    +		}
     
    -				if (registeredColumn == null) {
    -					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		public byte[] key() {
    +			return currentSubIterator.getCurrentKey();
    +		}
     
    -					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    -						nameBytes,
    -						rocksDBKeyedStateBackend.columnOptions);
    +		public byte[] value() {
    +			return currentSubIterator.getIterator().value();
    +		}
     
    -					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    -						new RegisteredKeyedBackendStateMetaInfo<>(
    -							restoredMetaInfo.getStateType(),
    -							restoredMetaInfo.getName(),
    -							restoredMetaInfo.getNamespaceSerializer(),
    -							restoredMetaInfo.getStateSerializer());
    +		/**
    +		 * Returns the Id of the k/v state to which the current key belongs.
    +		 * @return Id of K/V state to which the current key belongs.
    +		 */
    +		public int kvStateId() {
    +			return currentSubIterator.getKvStateId();
    +		}
     
    -					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restored
    --- End diff --
    
    I skimmed over the changes and I assume that you've mostly copied the code from one place to another without bigger changes. But it's really hard to see any real changes due to the combined movement of the code.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159836477
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java ---
    @@ -145,8 +145,8 @@
     
     	private volatile Throwable failureCause;          // once assigned, never changes
     
    -	/** The handle to the state that the task gets on restore */
    -	private volatile TaskStateSnapshot taskState;
    +	/** Information to restore the task onr recovery, such as checkpoint id and task state snapshot */
    +	private volatile TaskRestore taskRestore;
    --- End diff --
    
    I would add `SerializableOptional` class, make this field `SerializableOptional<TaskRestore>` and pass it down to `Task` to avoid nulls. 
    
    If you do not like this, please at least add `@Nullable` annotations to this field and subsequent usages.
    
    Or maybe embed `isEmpty` logic into `TaskRestore`.
      


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159884153
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/MultiStreamStateHandle.java ---
    @@ -1,104 +0,0 @@
    -/*
    --- End diff --
    
    Yes, they are deleted in later commits anyways.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168752042
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168750923
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    `localBackupDirectory` also can be removed if `canceled` is `true`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165323823
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateObjectCollectionTest.java ---
    @@ -0,0 +1,64 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.checkpoint;
    +
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.MethodForwardingTestUtil;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.Collections;
    +import java.util.function.Function;
    +
    +import static org.mockito.Mockito.mock;
    +
    +/**
    + * Tests for {@link StateObjectCollection}.
    + */
    +public class StateObjectCollectionTest {
    --- End diff --
    
    should extend `TestLogger`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170263293
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -186,6 +210,34 @@ private void doRelease(Iterable<TaskLocalStateStore> toRelease) {
     		return localRecoveryMode;
     	}
     
    +	private File[] allocationBaseDirectories(AllocationID allocationID) {
    +		File[] allocationDirectories = new File[localStateRootDirectories.length];
    +		for (int i = 0; i < localStateRootDirectories.length; ++i) {
    +			allocationDirectories[i] = new File(localStateRootDirectories[i], allocationSubDirString(allocationID));
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170266913
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/DuplicatingCheckpointOutputStream.java ---
    @@ -0,0 +1,287 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import java.io.IOException;
    +
    +/**
    + * A CheckpointStateOutputStream that wraps a primary and a secondary CheckpointStateOutputStream and duplicates
    + * all writes into both streams. This stream applies buffering to reduce the amount of dual-method calling. Furthermore,
    + * exceptions that happen in interactions with the secondary stream are not exposed, until the user calls
    + * {@link #closeAndGetSecondaryHandle()}. In contrast to that, exceptions from interactions with the primary stream
    + * are immediately returned to the user. This class is used to write state for local recovery as a local (secondary)
    + * copy of the (primary) state snapshot that is written to a (slower but highly-available) remote filesystem.
    + */
    +public class DuplicatingCheckpointOutputStream extends CheckpointStreamFactory.CheckpointStateOutputStream {
    +
    +	/** Flag if the positional alignment of both streams is checked after each operation. */
    +	private static final boolean STRICT_ALIGNMENT_CHECKS = false;
    +
    +	/** Default buffer size of 8KB. */
    +	private static final int DEFAULT_BUFFER_SIZER = 8 * 1024;
    +
    +	/** Write buffer. */
    +	private final byte[] buffer;
    +
    +	/** Position in the write buffer. */
    +	private int bufferIdx;
    +
    +	/**
    +	 * Primary stream for writing the checkpoint data. Failures from this stream are forwarded.
    +	 */
    +	private final CheckpointStreamFactory.CheckpointStateOutputStream primaryOutputStream;
    +
    +	/**
    +	 * Primary stream for writing the checkpoint data. Failures from this stream are not forwarded until
    +	 * {@link #closeAndGetSecondaryHandle()}.
    +	 */
    +	private final CheckpointStreamFactory.CheckpointStateOutputStream secondaryOutputStream;
    +
    +	/**
    +	 * Stores a potential exception that occurred while interacting with {@link #secondaryOutputStream}
    +	 */
    +	private IOException secondaryStreamException;
    +
    +	public DuplicatingCheckpointOutputStream(
    +		CheckpointStreamFactory.CheckpointStateOutputStream primaryOutputStream,
    +		CheckpointStreamFactory.CheckpointStateOutputStream secondaryOutputStream) throws IOException {
    +		this(primaryOutputStream, secondaryOutputStream, DEFAULT_BUFFER_SIZER);
    +	}
    +
    +	public DuplicatingCheckpointOutputStream(
    +		CheckpointStreamFactory.CheckpointStateOutputStream primaryOutputStream,
    +		CheckpointStreamFactory.CheckpointStateOutputStream secondaryOutputStream,
    +		int bufferSize) throws IOException {
    +
    +		this.primaryOutputStream = Preconditions.checkNotNull(primaryOutputStream);
    +		this.secondaryOutputStream = Preconditions.checkNotNull(secondaryOutputStream);
    +
    +		this.buffer = new byte[bufferSize];
    +		this.bufferIdx = 0;
    +
    +		this.secondaryStreamException = null;
    +
    +		checkForAlignedStreamPositions();
    +	}
    +
    +	@Override
    +	public void write(int b) throws IOException {
    +
    +		if (buffer.length <= bufferIdx) {
    +			flushInternalBuffer();
    +		}
    +
    +		buffer[bufferIdx] = (byte) b;
    +		++bufferIdx;
    +	}
    +
    +	@Override
    +	public void write(byte[] b) throws IOException {
    +
    +		write(b, 0, b.length);
    +	}
    +
    +	@Override
    +	public void write(byte[] b, int off, int len) throws IOException {
    +
    +		if (buffer.length <= len) {
    +
    +			flushInternalBuffer();
    +			writeThroughInternal(b, off, len);
    +		} else {
    +
    +			if (buffer.length < len + bufferIdx) {
    +				flushInternalBuffer();
    +			}
    +
    +			System.arraycopy(b, off, buffer, bufferIdx, len);
    +			bufferIdx += len;
    +		}
    +	}
    +
    +	@Override
    +	public long getPos() throws IOException {
    +		final long referencePos = primaryOutputStream.getPos();
    +		return referencePos + bufferIdx;
    +	}
    +
    +	@Override
    +	public void flush() throws IOException {
    +
    +		flushInternalBuffer();
    +		primaryOutputStream.flush();
    +
    +		if (secondaryStreamException == null) {
    +			try {
    +				secondaryOutputStream.flush();
    +			} catch (IOException flushEx) {
    +				handleSecondaryStreamOnException(flushEx);
    +			}
    +		}
    +	}
    +
    +	@Override
    +	public void sync() throws IOException {
    +
    +		flushInternalBuffer();
    +		primaryOutputStream.sync();
    +
    +		if (secondaryStreamException == null) {
    +			try {
    +				secondaryOutputStream.sync();
    +			} catch (IOException syncEx) {
    +				handleSecondaryStreamOnException(syncEx);
    +			}
    +		}
    +	}
    +
    +	@Override
    +	public void close() throws IOException {
    +
    +		IOException exCollector = null;
    +
    +		try {
    +			flushInternalBuffer();
    +		} catch (IOException flushEx) {
    +			exCollector = flushEx;
    +		}
    +
    +		try {
    +			primaryOutputStream.close();
    +		} catch (IOException closeEx) {
    +			exCollector = ExceptionUtils.firstOrSuppressed(closeEx, exCollector);
    +		}
    +
    +		if (secondaryStreamException == null) {
    +			try {
    +				secondaryOutputStream.close();
    +			} catch (IOException closeEx) {
    +				handleSecondaryStreamOnException(closeEx);
    +			}
    +		}
    +
    +		if (exCollector != null) {
    +			throw exCollector;
    +		}
    +	}
    +
    +	private void checkForAlignedStreamPositions() throws IOException {
    +
    +		final long primaryPos = primaryOutputStream.getPos();
    +
    +		if (secondaryStreamException != null) {
    +			return;
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160102220
  
    --- Diff: flink-contrib/flink-storm/src/test/java/org/apache/flink/storm/wrappers/BoltWrapperTest.java ---
    @@ -375,11 +376,13 @@ public void declareOutputFields(OutputFieldsDeclarer declarer) {
     		when(env.getMetricGroup()).thenReturn(new UnregisteredTaskMetricsGroup());
     		when(env.getTaskManagerInfo()).thenReturn(new TestingTaskManagerRuntimeInfo());
     
    +		final CloseableRegistry closeableRegistry = new CloseableRegistry();
     		StreamTask<?, ?> mockTask = mock(StreamTask.class);
     		when(mockTask.getCheckpointLock()).thenReturn(new Object());
     		when(mockTask.getConfiguration()).thenReturn(new StreamConfig(new Configuration()));
     		when(mockTask.getEnvironment()).thenReturn(env);
     		when(mockTask.getExecutionConfig()).thenReturn(execConfig);
    +		when(mockTask.getCancelables()).thenReturn(closeableRegistry);
    --- End diff --
    
    I completely share your dislike of Mockito. Overall this PR has already a lot of changes and I suggest to address those orthogonal test changes that remove Mockito in a separate task and PR.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170250183
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStoreImpl.java ---
    @@ -0,0 +1,298 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
    +
    +import java.io.File;
    +import java.io.IOException;
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.HashMap;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
    +
    +/**
    + * Main implementation of a {@link TaskLocalStateStore}.
    + */
    +public class TaskLocalStateStoreImpl implements TaskLocalStateStore {
    +
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStoreImpl.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	@VisibleForTesting
    +	static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    +
    +	/** JobID from the owning subtask. */
    +	@Nonnull
    +	private final JobID jobID;
    +
    +	/** AllocationID of the owning slot. */
    +	@Nonnull
    +	private final AllocationID allocationID;
    +
    +	/** JobVertexID of the owning subtask. */
    +	@Nonnull
    +	private final JobVertexID jobVertexID;
    +
    +	/** Subtask index of the owning subtask. */
    +	@Nonnegative
    +	private final int subtaskIndex;
    +
    +	/** The configured mode for local recovery. */
    +	@Nonnull
    +	private final LocalRecoveryConfig localRecoveryConfig;
    +
    +	/** Executor that runs the discarding of released state objects. */
    +	@Nonnull
    +	private final Executor discardExecutor;
    +
    +	/** Lock for synchronisation on the storage map and the discarded status. */
    +	@Nonnull
    +	private final Object lock;
    +
    +	/** Status flag if this store was already discarded. */
    +	@GuardedBy("lock")
    +	private boolean discarded;
    +
    +	/** Maps checkpoint ids to local TaskStateSnapshots. */
    +	@Nonnull
    +	@GuardedBy("lock")
    +	private final SortedMap<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
    +	public TaskLocalStateStoreImpl(
    +		@Nonnull JobID jobID,
    +		@Nonnull AllocationID allocationID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex,
    +		@Nonnull LocalRecoveryConfig localRecoveryConfig,
    +		@Nonnull Executor discardExecutor) {
    +
    +		this.jobID = jobID;
    +		this.allocationID = allocationID;
    +		this.jobVertexID = jobVertexID;
    +		this.subtaskIndex = subtaskIndex;
    +		this.discardExecutor = discardExecutor;
    +		this.lock = new Object();
    +		this.storedTaskStateByCheckpointID = new TreeMap<>();
    +		this.discarded = false;
    +		this.localRecoveryConfig = localRecoveryConfig;
    +	}
    +
    +	@Override
    +	public void storeLocalState(
    +		@Nonnegative long checkpointId,
    +		@Nullable TaskStateSnapshot localState) {
    +
    +		if (localState == null) {
    +			localState = NULL_DUMMY;
    +		}
    +
    +		LOG.info("Storing local state for checkpoint {}.", checkpointId);
    +		LOG.debug("Local state for checkpoint {} is {}.", checkpointId, localState);
    +
    +		Map<Long, TaskStateSnapshot> toDiscard = new HashMap<>(MAX_RETAINED_SNAPSHOTS);
    +
    +		synchronized (lock) {
    +			if (discarded) {
    +				// we ignore late stores and simply discard the state.
    +				toDiscard.put(checkpointId, localState);
    +			} else {
    +				TaskStateSnapshot previous =
    +					storedTaskStateByCheckpointID.put(checkpointId, localState);
    +
    +				if (previous != null) {
    +					toDiscard.put(checkpointId, previous);
    +				}
    +
    +				// remove from history.
    +				while (storedTaskStateByCheckpointID.size() > MAX_RETAINED_SNAPSHOTS) {
    +					Long removeCheckpointID = storedTaskStateByCheckpointID.firstKey();
    +					TaskStateSnapshot snapshot =
    +						storedTaskStateByCheckpointID.remove(removeCheckpointID);
    +					toDiscard.put(removeCheckpointID, snapshot);
    +				}
    +			}
    +		}
    +
    +		asyncDiscardLocalStateForCollection(toDiscard.entrySet());
    +	}
    +
    +	@Override
    +	@Nullable
    +	public TaskStateSnapshot retrieveLocalState(long checkpointID) {
    +		synchronized (lock) {
    +			TaskStateSnapshot snapshot = storedTaskStateByCheckpointID.get(checkpointID);
    +			return snapshot != NULL_DUMMY ? snapshot : null;
    +		}
    +	}
    +
    +	@Override
    +	@Nonnull
    +	public LocalRecoveryConfig getLocalRecoveryConfig() {
    +		return localRecoveryConfig;
    +	}
    +
    +	@Override
    +	public void confirmCheckpoint(long confirmedCheckpointId) {
    +
    +		LOG.debug("Received confirmation for checkpoint {}. Starting to prune history.", confirmedCheckpointId);
    +
    +		final List<Map.Entry<Long, TaskStateSnapshot>> toRemove = new ArrayList<>(MAX_RETAINED_SNAPSHOTS);
    +
    +		synchronized (lock) {
    +
    +			Iterator<Map.Entry<Long, TaskStateSnapshot>> entryIterator =
    +				storedTaskStateByCheckpointID.entrySet().iterator();
    +
    +			// remove entries for outdated checkpoints and discard their state.
    +			while (entryIterator.hasNext()) {
    +
    +				Map.Entry<Long, TaskStateSnapshot> snapshotEntry = entryIterator.next();
    +				long entryCheckpointId = snapshotEntry.getKey();
    +
    +				if (entryCheckpointId < confirmedCheckpointId) {
    +					toRemove.add(snapshotEntry);
    +					entryIterator.remove();
    +				} else {
    +					// we can stop because the map is sorted.
    +					break;
    +				}
    +			}
    +		}
    +
    +		asyncDiscardLocalStateForCollection(toRemove);
    +	}
    +
    +	/**
    +	 * Disposes the state of all local snapshots managed by this object.
    +	 */
    +	public void dispose() {
    +
    --- End diff --
    
    If the user cancels the job, then the allocated slot should be released, thus, also releasing the local state data.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168498124
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateSnapshot.java ---
    @@ -75,7 +79,10 @@ public OperatorSubtaskState getSubtaskStateByOperatorID(OperatorID operatorID) {
     	 * Maps the given operator id to the given subtask state. Returns the subtask state of a previous mapping, if such
     	 * a mapping existed or null otherwise.
     	 */
    -	public OperatorSubtaskState putSubtaskStateByOperatorID(OperatorID operatorID, OperatorSubtaskState state) {
    +	public OperatorSubtaskState putSubtaskStateByOperatorID(
    --- End diff --
    
    `@Nullable` missing.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168752606
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    +
    +- RocksDBStateBackend: task-local recovery is supported for keyed state. For *full checkpoints*, state is duplicated to a local file. This can introduce additional write costs
    +and occupies local disk space. For *incremental snapshots*, the local state is based on RocksDB's native checkpointing mechanism. This mechanism is also used as the first step
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165331645
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java ---
    @@ -686,4 +714,53 @@ static void resetRocksDBLoadedFlag() throws Exception {
     		initField.setAccessible(true);
     		initField.setBoolean(null, false);
     	}
    +
    +	/**
    +	 * This enum represents the different modes for local recovery.
    +	 */
    +	public enum LocalRecoveryMode {
    +		DISABLED, ENABLE_FILE_BASED
    +	}
    +
    +	/**
    +	 * This class encapsulates the configuration for local recovery of this backend.
    +	 */
    +	public static final class LocalRecoveryConfig implements Serializable {
    +
    +		private static final long serialVersionUID = 1L;
    +		private static final LocalRecoveryConfig DISABLED_SINGLETON =
    +			new LocalRecoveryConfig(LocalRecoveryMode.DISABLED, null);
    +
    +		private final LocalRecoveryMode localRecoveryMode;
    +		private final File localStateDirectory;
    +
    +		LocalRecoveryConfig(LocalRecoveryMode localRecoveryMode, File localStateDirectory) {
    +			this.localRecoveryMode = Preconditions.checkNotNull(localRecoveryMode);
    +			this.localStateDirectory = localStateDirectory;
    +			if (LocalRecoveryMode.ENABLE_FILE_BASED.equals(localRecoveryMode) && localStateDirectory == null) {
    +				throw new IllegalStateException("Local state directory must be specified if local recovery mode is " +
    +					LocalRecoveryMode.ENABLE_FILE_BASED);
    +			}
    +		}
    +
    +		public LocalRecoveryMode getLocalRecoveryMode() {
    +			return localRecoveryMode;
    +		}
    +
    +		public File getLocalStateDirectory() {
    --- End diff --
    
    @Nullable


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169896622
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/DirectoryKeyedStateHandle.java ---
    @@ -0,0 +1,105 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import javax.annotation.Nonnull;
    +
    +/**
    + * This class is a keyed state handle based on a directory. It combines a {@link DirectoryStateHandle} and a
    + * {@link KeyGroupRange}.
    + */
    +public class DirectoryKeyedStateHandle implements KeyedStateHandle {
    +
    +	@Nonnull
    +	private final DirectoryStateHandle directoryStateHandle;
    +
    +	@Nonnull
    +	private final KeyGroupRange keyGroupRange;
    +
    +	public DirectoryKeyedStateHandle(
    +		@Nonnull DirectoryStateHandle directoryStateHandle,
    +		@Nonnull KeyGroupRange keyGroupRange) {
    +
    +		this.directoryStateHandle = directoryStateHandle;
    +		this.keyGroupRange = keyGroupRange;
    +	}
    +
    +	@Nonnull
    +	public DirectoryStateHandle getDirectoryStateHandle() {
    +		return directoryStateHandle;
    +	}
    +
    +	@Nonnull
    +	@Override
    +	public KeyGroupRange getKeyGroupRange() {
    +		return keyGroupRange;
    +	}
    +
    +	@Override
    +	public void discardState() throws Exception {
    +		directoryStateHandle.discardState();
    +	}
    +
    +	@Override
    +	public long getStateSize() {
    +		return directoryStateHandle.getStateSize();
    +	}
    +
    +	@Override
    +	public KeyedStateHandle getIntersection(KeyGroupRange keyGroupRange) {
    +		return keyGroupRange.getIntersection(keyGroupRange).getNumberOfKeyGroups() > 0 ? this : null;
    --- End diff --
    
     👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168494204
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java ---
    @@ -553,27 +616,36 @@ public boolean supportsAsynchronousSnapshots() {
     					metaInfoSnapshots,
     					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, keyGroupCompressionDecorator));
     
    +			LocalRecoveryDirectoryProvider directoryProvider =
    +				FsStateBackend.LocalRecoveryMode.ENABLE_FILE_BASED.equals(localRecoveryConfig.getLocalRecoveryMode()) ?
    +					localRecoveryConfig.getLocalStateDirectories() :
    +					null;
    +
    +			final ThrowingSupplier<CheckpointStreamWithResultProvider, Exception> checkpointStreamSupplier =
    +				() -> new CheckpointStreamWithResultProvider.Factory().create(
    +					checkpointId,
    +					CheckpointedStateScope.EXCLUSIVE,
    --- End diff --
    
    Already changed in later commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165341901
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java ---
    @@ -0,0 +1,122 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.configuration.ConfigConstants;
    +import org.apache.flink.configuration.Configuration;
    +import org.apache.flink.runtime.clusterframework.types.ResourceID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +import org.apache.flink.runtime.taskexecutor.TaskManagerServices;
    +import org.apache.flink.runtime.taskexecutor.TaskManagerServicesConfiguration;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +import org.junit.rules.TemporaryFolder;
    +
    +import java.io.File;
    +import java.net.InetAddress;
    +
    +public class TaskExecutorLocalStateStoresManagerTest {
    +
    +	/**
    +	 * This tests that the creation of {@link TaskManagerServices} correctly creates the local state root directory
    +	 * for the {@link TaskExecutorLocalStateStoresManager} with the configured root directory.
    +	 */
    +	@Test
    +	public void testCreationFromConfig() throws Exception {
    +
    +		final Configuration config = new Configuration();
    +
    +		final String rootDirString = "localStateRoot";
    +		config.setString(ConfigConstants.TASK_MANAGER_LOCAL_STATE_ROOT_DIR_KEY, rootDirString);
    +
    +		final ResourceID tmResourceID = ResourceID.generate();
    +
    +		TaskManagerServicesConfiguration taskManagerServicesConfiguration =
    +			TaskManagerServicesConfiguration.fromConfiguration(config, InetAddress.getLocalHost(), true);
    +
    +		TaskManagerServices taskManagerServices =
    +			TaskManagerServices.fromConfiguration(taskManagerServicesConfiguration, tmResourceID);
    +
    +		TaskExecutorLocalStateStoresManager taskStateManager = taskManagerServices.getTaskStateManager();
    +
    +		Assert.assertEquals(
    +			new File(rootDirString, TaskManagerServices.LOCAL_STATE_SUB_DIRECTORY_ROOT),
    +			taskStateManager.getLocalStateRootDirectory());
    +
    +		Assert.assertEquals("localState", TaskManagerServices.LOCAL_STATE_SUB_DIRECTORY_ROOT);
    +	}
    +
    +	/**
    +	 * This tests that the creation of {@link TaskManagerServices} correctly falls back to the first tmp directory of
    +	 * the IOManager as default for the local state root directory.
    +	 */
    +	@Test
    +	public void testCreationFromConfigDefault() throws Exception {
    +
    +		final Configuration config = new Configuration();
    +
    +		final ResourceID tmResourceID = ResourceID.generate();
    +
    +		TaskManagerServicesConfiguration taskManagerServicesConfiguration =
    +			TaskManagerServicesConfiguration.fromConfiguration(config, InetAddress.getLocalHost(), true);
    +
    +		TaskManagerServices taskManagerServices =
    +			TaskManagerServices.fromConfiguration(taskManagerServicesConfiguration, tmResourceID);
    +
    +		TaskExecutorLocalStateStoresManager taskStateManager = taskManagerServices.getTaskStateManager();
    +
    +		Assert.assertEquals(
    +			new File(taskManagerServicesConfiguration.getTmpDirPaths()[0], TaskManagerServices.LOCAL_STATE_SUB_DIRECTORY_ROOT),
    +			taskStateManager.getLocalStateRootDirectory());
    +	}
    +
    +	/**
    +	 * This tests that the {@link TaskExecutorLocalStateStoresManager} creates {@link TaskLocalStateStore} that have
    +	 * a properly initialized local state base directory.
    +	 */
    +	@Test
    +	public void testSubtaskStateStoreDirectoryCreation() throws Exception {
    +
    +		JobID jobID = new JobID();
    +		JobVertexID jobVertexID = new JobVertexID();
    +		int subtaskIdx = 42;
    +		TemporaryFolder tmp = new TemporaryFolder();
    --- End diff --
    
    Why not making it a `ClassRule`?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170238910
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/DuplicatingCheckpointOutputStream.java ---
    @@ -0,0 +1,287 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import java.io.IOException;
    +
    +/**
    + * A CheckpointStateOutputStream that wraps a primary and a secondary CheckpointStateOutputStream and duplicates
    + * all writes into both streams. This stream applies buffering to reduce the amount of dual-method calling. Furthermore,
    + * exceptions that happen in interactions with the secondary stream are not exposed, until the user calls
    + * {@link #closeAndGetSecondaryHandle()}. In contrast to that, exceptions from interactions with the primary stream
    + * are immediately returned to the user. This class is used to write state for local recovery as a local (secondary)
    + * copy of the (primary) state snapshot that is written to a (slower but highly-available) remote filesystem.
    + */
    +public class DuplicatingCheckpointOutputStream extends CheckpointStreamFactory.CheckpointStateOutputStream {
    +
    +	/** Flag if the positional alignment of both streams is checked after each operation. */
    +	private static final boolean STRICT_ALIGNMENT_CHECKS = false;
    +
    +	/** Default buffer size of 8KB. */
    +	private static final int DEFAULT_BUFFER_SIZER = 8 * 1024;
    +
    +	/** Write buffer. */
    +	private final byte[] buffer;
    +
    +	/** Position in the write buffer. */
    +	private int bufferIdx;
    +
    +	/**
    +	 * Primary stream for writing the checkpoint data. Failures from this stream are forwarded.
    +	 */
    +	private final CheckpointStreamFactory.CheckpointStateOutputStream primaryOutputStream;
    +
    +	/**
    +	 * Primary stream for writing the checkpoint data. Failures from this stream are not forwarded until
    +	 * {@link #closeAndGetSecondaryHandle()}.
    +	 */
    +	private final CheckpointStreamFactory.CheckpointStateOutputStream secondaryOutputStream;
    +
    +	/**
    +	 * Stores a potential exception that occurred while interacting with {@link #secondaryOutputStream}
    +	 */
    +	private IOException secondaryStreamException;
    +
    +	public DuplicatingCheckpointOutputStream(
    +		CheckpointStreamFactory.CheckpointStateOutputStream primaryOutputStream,
    +		CheckpointStreamFactory.CheckpointStateOutputStream secondaryOutputStream) throws IOException {
    +		this(primaryOutputStream, secondaryOutputStream, DEFAULT_BUFFER_SIZER);
    +	}
    +
    +	public DuplicatingCheckpointOutputStream(
    +		CheckpointStreamFactory.CheckpointStateOutputStream primaryOutputStream,
    +		CheckpointStreamFactory.CheckpointStateOutputStream secondaryOutputStream,
    +		int bufferSize) throws IOException {
    +
    +		this.primaryOutputStream = Preconditions.checkNotNull(primaryOutputStream);
    +		this.secondaryOutputStream = Preconditions.checkNotNull(secondaryOutputStream);
    +
    +		this.buffer = new byte[bufferSize];
    +		this.bufferIdx = 0;
    +
    +		this.secondaryStreamException = null;
    +
    +		checkForAlignedStreamPositions();
    +	}
    +
    +	@Override
    +	public void write(int b) throws IOException {
    +
    +		if (buffer.length <= bufferIdx) {
    +			flushInternalBuffer();
    +		}
    +
    +		buffer[bufferIdx] = (byte) b;
    +		++bufferIdx;
    +	}
    +
    +	@Override
    +	public void write(byte[] b) throws IOException {
    +
    +		write(b, 0, b.length);
    +	}
    +
    +	@Override
    +	public void write(byte[] b, int off, int len) throws IOException {
    +
    +		if (buffer.length <= len) {
    +
    +			flushInternalBuffer();
    +			writeThroughInternal(b, off, len);
    +		} else {
    +
    +			if (buffer.length < len + bufferIdx) {
    +				flushInternalBuffer();
    +			}
    +
    +			System.arraycopy(b, off, buffer, bufferIdx, len);
    +			bufferIdx += len;
    +		}
    +	}
    +
    +	@Override
    +	public long getPos() throws IOException {
    +		final long referencePos = primaryOutputStream.getPos();
    +		return referencePos + bufferIdx;
    +	}
    +
    +	@Override
    +	public void flush() throws IOException {
    +
    +		flushInternalBuffer();
    +		primaryOutputStream.flush();
    +
    +		if (secondaryStreamException == null) {
    +			try {
    +				secondaryOutputStream.flush();
    +			} catch (IOException flushEx) {
    +				handleSecondaryStreamOnException(flushEx);
    --- End diff --
    
    Why not just catch `throwable` for `secondaryOutputStream` ? Maybe some `RuntimeException` could throw out from it.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r167356561
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    --- End diff --
    
    secondary cop**ies**


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160956609
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StateInitializationContextImplTest.java ---
    @@ -124,14 +140,69 @@ public void setUp() throws Exception {
     			operatorStateHandles.add(operatorStateHandle);
     		}
     
    +		OperatorSubtaskState operatorSubtaskState = new OperatorSubtaskState(
    +			Collections.emptyList(),
    +			operatorStateHandles,
    +			Collections.emptyList(),
    +			keyedStateHandles);
    +
    +		OperatorID operatorID = new OperatorID();
    +		TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    +		taskStateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
    +
    +		JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(0L, taskStateSnapshot);
    +
    +		TaskStateManager manager = new TaskStateManagerImpl(
    +			new JobID(),
    +			new ExecutionAttemptID(),
    +			mock(TaskLocalStateStore.class),
    +			jobManagerTaskRestore,
    +			mock(CheckpointResponder.class));
    +
    +		DummyEnvironment environment = new DummyEnvironment(
    +			"test",
    +			1,
    +			0,
    +			prev);
    +
    +		environment.setTaskStateManager(manager);
    +
    +		StateBackend stateBackend = new MemoryStateBackend(1024);
    +		StreamTaskStateManager streamTaskStateManager = new StreamTaskStateManagerImpl(
    +			environment,
    +			stateBackend,
    +			mock(ProcessingTimeService.class)) {
    +
    +			@Override
    +			protected <K> InternalTimeServiceManager<?, K> internalTimeServiceManager(
    +				AbstractKeyedStateBackend<K> keyedStatedBackend,
    +				KeyContext keyContext,
    +				Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
    +
    +				// We do not initialize a timer service manager here, because it would already consume the raw keyed
    +				// state as part of initialization. For the purpose of this test, we want an unconsumed raw keyed
    +				// stream.
    +				return null;
    +			}
    +		};
    +
    +		AbstractStreamOperator<?> mockOperator = mock(AbstractStreamOperator.class);
    +		when(mockOperator.getOperatorID()).thenReturn(operatorID);
    +
    +		StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(
    +			mockOperator,
    --- End diff --
    
    Does the `StreamOperatorStateContext` needs to have `AbstractStreamOperator` as a parameter? Couldn't we cut this dependency? 
    
    This mockito usage suggests that this `streamOperatorStateContext` method is receiving more in the parameter compared to what it needs. 
    
    Rephrasing this concern, is it legal and does it make sense for the `streamOperatorStateContext(operator, ...)` method to invoke one of the ~30 `operator` methods? Like `operator.dispose()`, `operator.snapshoState()`, `operator.getCurrentKey()`, ... ?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165325793
  
    --- Diff: flink-core/src/test/java/org/apache/flink/util/MethodForwardingTestUtil.java ---
    @@ -0,0 +1,146 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.util;
    +
    +import org.mockito.Mockito;
    +
    +import java.lang.reflect.Array;
    +import java.lang.reflect.Method;
    +import java.util.Collections;
    +import java.util.Set;
    +import java.util.function.Function;
    +
    +import static org.junit.Assert.fail;
    +import static org.mockito.Mockito.reset;
    +import static org.mockito.Mockito.spy;
    +
    +/**
    + * Helper class with a method that attempt to automatically test method forwarding between a delegate and a wrapper.
    + */
    +public class MethodForwardingTestUtil {
    +
    +	/**
    +	 * This is a best effort automatic test for method forwarding between a delegate and its wrapper, where the wrapper
    +	 * class is a subtype of the delegate. This ignores methods that are inherited from Object.
    +	 *
    +	 * @param delegateClass the class for the delegate.
    +	 * @param wrapperFactory factory that produces a wrapper from a delegate.
    +	 * @param <D> type of the delegate
    +	 * @param <W> type of the wrapper
    +	 */
    +	public static <D, W> void testMethodForwarding(
    +		Class<D> delegateClass,
    +		Function<D, W> wrapperFactory) {
    +		testMethodForwarding(delegateClass, wrapperFactory, Collections.emptySet());
    +	}
    +
    +	/**
    +	 * This is a best effort automatic test for method forwarding between a delegate and its wrapper, where the wrapper
    +	 * class is a subtype of the delegate. Methods can be remapped in case that the implementation does not call the
    +	 * original method. Remapping to null skips the method. This ignores methods that are inherited from Object.
    +	 *
    +	 * @param delegateClass the class for the delegate.
    +	 * @param wrapperFactory factory that produces a wrapper from a delegate.
    +	 * @param skipMethodSet set of methods to ignore.
    +	 * @param <D> type of the delegate
    +	 * @param <W> type of the wrapper
    +	 */
    +	public static <D, W> void testMethodForwarding(
    +		Class<D> delegateClass,
    +		Function<D, W> wrapperFactory,
    +		Set<Method> skipMethodSet) {
    +
    +		Preconditions.checkNotNull(delegateClass);
    +		Preconditions.checkNotNull(wrapperFactory);
    +		Preconditions.checkNotNull(skipMethodSet);
    +
    +		D delegate = spy(delegateClass);
    +		W wrapper = wrapperFactory.apply(delegate);
    +
    +		// ensure that wrapper is a subtype of delegate
    +		Preconditions.checkArgument(delegateClass.isAssignableFrom(wrapper.getClass()));
    +
    +		for (Method delegateMethod : delegateClass.getMethods()) {
    +
    +			if (checkSkipMethodForwardCheck(delegateMethod, skipMethodSet)) {
    +				continue;
    +			}
    +
    +			try {
    +				// find the correct method to substitute the bridge for erased generic types.
    +				// if this doesn't work, the user need to exclude the method and write an additional test.
    +				Method wrapperMethod = wrapper.getClass().getDeclaredMethod(
    +					delegateMethod.getName(),
    +					delegateMethod.getParameterTypes());
    +
    +				// things get a bit fuzzy here, best effort to find a match but this might end up with a wrong method.
    +				if (wrapperMethod.isBridge()) {
    +					for (Method method : wrapper.getClass().getDeclaredMethods()) {
    +						if (!method.isBridge()
    +							&& method.getName().equals(wrapperMethod.getName())
    +							&& method.getParameterCount() == wrapperMethod.getParameterCount()) {
    +							wrapperMethod = method;
    +							break;
    +						}
    +					}
    +				}
    +
    +				Class<?>[] parameterTypes = wrapperMethod.getParameterTypes();
    +				Object[] arguments = new Object[parameterTypes.length];
    +				for (int j = 0; j < arguments.length; j++) {
    +					Class<?> parameterType = parameterTypes[j];
    +					if (parameterType.isArray()) {
    +						arguments[j] = Array.newInstance(parameterType.getComponentType(), 0);
    +					} else if (parameterType.isPrimitive()) {
    +						arguments[j] = 0;
    --- End diff --
    
    Why is a primitive always of type `int`?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165105360
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotResult.java ---
    @@ -0,0 +1,83 @@
    +package org.apache.flink.runtime.state;/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    --- End diff --
    
    License and package seem in the wrong order.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168778178
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -266,11 +305,16 @@ public RocksDBKeyedStateBackend(
     		RocksIterator iterator = db.newIterator(columnInfo.f0);
     		iterator.seekToFirst();
     
    -		Iterable<K> iterable = () -> new RocksIteratorWrapper<>(iterator, state, keySerializer, keyGroupPrefixBytes);
    +		Iterable<K> iterable = () -> new RocksIteratorToJavaIteratorAdapter<>(iterator, state, keySerializer, keyGroupPrefixBytes);
     		Stream<K> targetStream = StreamSupport.stream(iterable.spliterator(), false);
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168506302
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskStateManagerImplTest.java ---
    @@ -109,10 +111,10 @@ public void testStateReportingAndRetrieving() {
     			executionAttemptID,
     			checkpointResponderMock,
     			taskRestore);
    -
    -		Assert.assertTrue(jmOperatorSubtaskState_1 == taskStateManager.operatorStates(operatorID_1));
    -		Assert.assertTrue(jmOperatorSubtaskState_2 == taskStateManager.operatorStates(operatorID_2));
    -		Assert.assertNull(taskStateManager.operatorStates(operatorID_3));
    +//TODO
    +//		Assert.assertTrue(jmOperatorSubtaskState_1 == taskStateManager.prioritizedOperatorState(operatorID_1));
    +//		Assert.assertTrue(jmOperatorSubtaskState_2 == taskStateManager.prioritizedOperatorState(operatorID_2));
    --- End diff --
    
    Unfinished TODO?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168473866
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamWithResultProvider.java ---
    @@ -0,0 +1,175 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.state.filesystem.FixFileFsStateOutputStream;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import java.io.Closeable;
    +import java.io.File;
    +import java.io.IOException;
    +
    +/**
    + * Interface that provides access to a CheckpointStateOutputStream and a method to provide the {@link SnapshotResult}.
    + * This abstracts from different ways that a result is obtained from checkpoint output streams.
    + */
    +public interface CheckpointStreamWithResultProvider extends Closeable {
    +
    +	SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +		KeyGroupRangeOffsets kgOffs) throws IOException;
    +
    +	CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream();
    +
    +	@Override
    +	default void close() throws IOException {
    +		CheckpointStreamFactory.CheckpointStateOutputStream outputStream = getCheckpointOutputStream();
    +		if (outputStream != null) {
    --- End diff --
    
    It looks as if the implementations in this class all return non null values.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168497612
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/DuplicatingCheckpointOutputStreamTest.java ---
    @@ -0,0 +1,242 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.FSDataInputStream;
    +import org.apache.flink.runtime.testutils.CommonTestUtils;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import java.io.IOException;
    +import java.util.Random;
    +
    +public class DuplicatingCheckpointOutputStreamTest {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168168810
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/StateObjectCollectionTest.java ---
    @@ -0,0 +1,64 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.checkpoint;
    +
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.MethodForwardingTestUtil;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.Collections;
    +import java.util.function.Function;
    +
    +import static org.mockito.Mockito.mock;
    +
    +/**
    + * Tests for {@link StateObjectCollection}.
    + */
    +public class StateObjectCollectionTest {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168514624
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -20,35 +20,47 @@
     
     import org.apache.flink.annotation.VisibleForTesting;
     import org.apache.flink.api.common.JobID;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
     import org.apache.flink.runtime.jobgraph.JobVertexID;
    -import org.apache.flink.util.ExceptionUtils;
    -import org.apache.flink.util.Preconditions;
     
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
     import javax.annotation.Nonnull;
     
     import java.io.File;
    -import java.util.HashMap;
     import java.util.Map;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
     
     /**
      * This class holds the all {@link TaskLocalStateStore} objects for a task executor (manager).
    - *
    - * TODO: this still still work in progress and partially still acts as a placeholder.
      */
     public class TaskExecutorLocalStateStoresManager {
     
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskExecutorLocalStateStoresManager.class);
    +
     	/**
     	 * This map holds all local state stores for tasks running on the task manager / executor that own the instance of
    -	 * this.
    +	 * this. Maps from allocation id to all the subtask's local state stores.
     	 */
    -	private final Map<JobID, Map<JobVertexSubtaskKey, TaskLocalStateStore>> taskStateStoresByJobID;
    +	private final Map<AllocationID, Map<JobVertexSubtaskKey, TaskLocalStateStore>> taskStateStoresByAllocationID;
     
     	/** This is the root directory for all local state of this task manager / executor. */
     	private final File[] localStateRootDirectories;
     
    -	public TaskExecutorLocalStateStoresManager(File[] localStateRootDirectories) {
    -		this.taskStateStoresByJobID = new HashMap<>();
    -		this.localStateRootDirectories = Preconditions.checkNotNull(localStateRootDirectories);
    +	/** Executor that runs the discarding of released state objects. */
    +	private final Executor discardExecutor;
    +
    +	public TaskExecutorLocalStateStoresManager(
    +		@Nonnull File[] localStateRootDirectories,
    +		@Nonnull Executor discardExecutor) {
    --- End diff --
    
    At least it is highlighting potential problems in the IDE and you can instruct the compiler to introduce runtime checks if you wish so (if you compile with Intellij this is the default).


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168502352
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -19,92 +19,224 @@
     package org.apache.flink.runtime.state;
     
     import org.apache.flink.api.common.JobID;
    -import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
     import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
     import org.apache.flink.runtime.jobgraph.JobVertexID;
    -import org.apache.flink.util.ExceptionUtils;
     
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
     import javax.annotation.Nonnull;
     import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
     
     import java.io.File;
    -import java.util.HashMap;
    +import java.util.Arrays;
    +import java.util.Iterator;
     import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
     
     /**
      * This class will service as a task-manager-level local storage for local checkpointed state. The purpose is to provide
      * access to a state that is stored locally for a faster recovery compared to the state that is stored remotely in a
      * stable store DFS. For now, this storage is only complementary to the stable storage and local state is typically
      * lost in case of machine failures. In such cases (and others), client code of this class must fall back to using the
      * slower but highly available store.
    - *
    - * TODO this is currently a placeholder / mock that still must be implemented!
      */
     public class TaskLocalStateStore {
     
    -	/** */
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStore.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	private static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    +
    +	/** JobID from the owning subtask. */
     	private final JobID jobID;
     
    -	/** */
    +	/** JobVertexID of the owning subtask. */
     	private final JobVertexID jobVertexID;
     
    -	/** */
    +	/** Subtask index of the owning subtask. */
     	private final int subtaskIndex;
     
    -	/** */
    -	private final Map<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    -
     	/** The root directories for all local state of this {@link TaskLocalStateStore}. */
     	private final File[] localStateRootDirectories;
     
    +	/** Executor that runs the discarding of released state objects. */
    +	private final Executor discardExecutor;
    +
    +	/** Lock for synchronisation on the storage map and the discarded status. */
    +	private final Object lock;
    +
    +	/** Status flag if this store was already discarded. */
    +	@GuardedBy("lock")
    +	private boolean discarded;
    +
    +	/** Maps checkpoint ids to local TaskStateSnapshots. */
    +	@GuardedBy("lock")
    +	private final SortedMap<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
     	public TaskLocalStateStore(
    -		JobID jobID,
    -		JobVertexID jobVertexID,
    -		int subtaskIndex,
    -		File[] localStateRootDirectories) {
    +		@Nonnull JobID jobID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex,
    +		@Nonnull File[] localStateRootDirectories,
    +		@Nonnull Executor discardExecutor) {
     
     		this.jobID = jobID;
     		this.jobVertexID = jobVertexID;
     		this.subtaskIndex = subtaskIndex;
    -		this.storedTaskStateByCheckpointID = new HashMap<>();
     		this.localStateRootDirectories = localStateRootDirectories;
    +		this.discardExecutor = discardExecutor;
    +		this.lock = new Object();
    +		this.storedTaskStateByCheckpointID = new TreeMap<>();
    +		this.discarded = false;
     	}
     
    +	@Nonnull
     	protected String createSubtaskPath() {
     		return jobID + File.separator + jobVertexID + File.separator + subtaskIndex;
     	}
     
    +	/**
    +	 * Stores the local state for the given checkpoint id.
    +	 *
    +	 * @param checkpointId id for the checkpoint that created the local state that will be stored.
    +	 * @param localState the local state to store.
    +	 */
     	public void storeLocalState(
    -		@Nonnull CheckpointMetaData checkpointMetaData,
    +		@Nonnegative long checkpointId,
     		@Nullable TaskStateSnapshot localState) {
     
    -		TaskStateSnapshot previous =
    -			storedTaskStateByCheckpointID.put(checkpointMetaData.getCheckpointId(), localState);
    +		if (localState == null) {
    +			localState = NULL_DUMMY;
    +		}
     
    -		if (previous != null) {
    -			throw new IllegalStateException("Found previously registered local state for checkpoint with id " +
    -				checkpointMetaData.getCheckpointId() + "! This indicated a problem.");
    +		LOG.info("Storing local state for checkpoint {}.", checkpointId);
    +		LOG.debug("Local state for checkpoint {} is {}.", checkpointId, localState);
    +
    +		synchronized (lock) {
    +			if (discarded) {
    +				// we ignore late stores and simply discard the state.
    +				discardStateObject(localState, checkpointId);
    +			} else {
    +				TaskStateSnapshot previous =
    +					storedTaskStateByCheckpointID.put(checkpointId, localState);
    +
    +				if (previous != null) {
    +					// this should never happen.
    +					discardStateObject(previous, checkpointId);
    +					throw new IllegalStateException("Found previously registered local state for checkpoint with id " +
    +						checkpointId + "! This indicated a problem.");
    +				}
    +
    +				// prune history.
    +				while (storedTaskStateByCheckpointID.size() > MAX_RETAINED_SNAPSHOTS) {
    +					Long removeCheckpointID = storedTaskStateByCheckpointID.firstKey();
    +					TaskStateSnapshot snapshot =
    +						storedTaskStateByCheckpointID.remove(storedTaskStateByCheckpointID.firstKey());
    --- End diff --
    
    I think we can directly pass in `removeCheckpointID`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169727620
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotDirectory.java ---
    @@ -0,0 +1,172 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.FileStatus;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.IOException;
    +import java.util.concurrent.atomic.AtomicReference;
    +
    +/**
    + * This class represents a directory that is the target for a state snapshot. This class provides some method that
    + * simplify resource management when dealing with such directories, e.g. it can produce a {@link DirectoryStateHandle}
    + * when the snapshot is completed and disposal considers whether or not a snapshot was already completed. For a
    + * completed snapshot, the ownership for cleanup is transferred to the created directory state handle. For incomplete
    + * snapshots, calling {@link #deleteIfIncompleteSnapshot()} will delete the underlying directory resource.
    + */
    +public class SnapshotDirectory {
    +
    +	/**
    +	 * Lifecycle stages of a snapshot directory.
    +	 */
    +	enum State {
    +		ONGOING, COMPLETED, DELETED
    +	}
    +
    +	/** This path describes the underlying directory for the snapshot. */
    +	@Nonnull
    +	private final Path directory;
    +
    +	/** The filesystem that contains the snapshot directory. */
    +	@Nonnull
    +	private final FileSystem fileSystem;
    +
    +	/** This reference tracks the lifecycle state of the snapshot directory. */
    +	@Nonnull
    +	private AtomicReference<State> state;
    +
    +	public SnapshotDirectory(@Nonnull Path directory, @Nonnull FileSystem fileSystem) {
    +		this.directory = directory;
    +		this.fileSystem = fileSystem;
    +		this.state = new AtomicReference<>(State.ONGOING);
    +	}
    +
    +	public SnapshotDirectory(@Nonnull Path directory) throws IOException {
    +		this(directory, directory.getFileSystem());
    +	}
    +
    +	@Nonnull
    +	public Path getDirectory() {
    +		return directory;
    +	}
    +
    +	public boolean mkdirs() throws IOException {
    +		return fileSystem.mkdirs(directory);
    +	}
    +
    +	@Nonnull
    +	public FileSystem getFileSystem() {
    +		return fileSystem;
    +	}
    +
    +	public boolean exists() throws IOException {
    +		return fileSystem.exists(directory);
    +	}
    +
    +	/**
    +	 * List the statuses of the files/directories in the snapshot directory.
    +	 *
    +	 * @return the statuses of the files/directories in the given path.
    +	 * @throws IOException if there is a problem creating the file statuses.
    +	 */
    +	public FileStatus[] listStatus() throws IOException {
    +		return fileSystem.listStatus(directory);
    +	}
    +
    +	/**
    +	 * Calling this method completes the snapshot into the snapshot directory and creates a corresponding
    +	 * {@link DirectoryStateHandle} that points to the snapshot directory. Calling this method will also change the
    +	 * lifecycle state from "ongoing" to "completed". If the state was already deleted, an {@link IOException} is
    +	 * thrown.
    +	 *
    +	 * @return a directory state handle that points to the snapshot directory.
    +	 * @throws IOException if the state of this snapshot directory object is different from "ongoing".
    +	 */
    +	public DirectoryStateHandle completeSnapshotAndGetHandle() throws IOException {
    +		if (state.compareAndSet(State.ONGOING, State.COMPLETED)) {
    +			return new DirectoryStateHandle(directory, fileSystem);
    +		} else {
    +			throw new IOException("Expected state " + State.ONGOING + " but found state " + state.get());
    +		}
    +	}
    +
    +	/**
    +	 * Calling this method will attempt delete the underlying snapshot directory recursively, if the state was
    +	 * "ongoing". In this case, the state will be set to "deleted" as a result of this call.
    +	 *
    +	 * @return <code>true</code> if delete is successful, <code>false</code> otherwise.
    +	 * @throws IOException if an exception happens during the delete.
    +	 */
    +	public boolean deleteIfIncompleteSnapshot() throws IOException {
    +		return state.compareAndSet(State.ONGOING, State.DELETED) && fileSystem.delete(directory, true);
    +	}
    +
    +	public boolean isSnapshotOngoing() {
    +		return State.ONGOING.equals(state.get());
    --- End diff --
    
    enum values should be compared using `==`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168735436
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    Aha, something blind my eyes... I moved these comment to where they are.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168168770
  
    --- Diff: flink-core/src/test/java/org/apache/flink/util/MethodForwardingTestUtil.java ---
    @@ -0,0 +1,146 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.util;
    +
    +import org.mockito.Mockito;
    +
    +import java.lang.reflect.Array;
    +import java.lang.reflect.Method;
    +import java.util.Collections;
    +import java.util.Set;
    +import java.util.function.Function;
    +
    +import static org.junit.Assert.fail;
    +import static org.mockito.Mockito.reset;
    +import static org.mockito.Mockito.spy;
    +
    +/**
    + * Helper class with a method that attempt to automatically test method forwarding between a delegate and a wrapper.
    + */
    +public class MethodForwardingTestUtil {
    +
    +	/**
    +	 * This is a best effort automatic test for method forwarding between a delegate and its wrapper, where the wrapper
    +	 * class is a subtype of the delegate. This ignores methods that are inherited from Object.
    +	 *
    +	 * @param delegateClass the class for the delegate.
    +	 * @param wrapperFactory factory that produces a wrapper from a delegate.
    +	 * @param <D> type of the delegate
    +	 * @param <W> type of the wrapper
    +	 */
    +	public static <D, W> void testMethodForwarding(
    +		Class<D> delegateClass,
    +		Function<D, W> wrapperFactory) {
    +		testMethodForwarding(delegateClass, wrapperFactory, Collections.emptySet());
    +	}
    +
    +	/**
    +	 * This is a best effort automatic test for method forwarding between a delegate and its wrapper, where the wrapper
    +	 * class is a subtype of the delegate. Methods can be remapped in case that the implementation does not call the
    +	 * original method. Remapping to null skips the method. This ignores methods that are inherited from Object.
    +	 *
    +	 * @param delegateClass the class for the delegate.
    +	 * @param wrapperFactory factory that produces a wrapper from a delegate.
    +	 * @param skipMethodSet set of methods to ignore.
    +	 * @param <D> type of the delegate
    +	 * @param <W> type of the wrapper
    +	 */
    +	public static <D, W> void testMethodForwarding(
    +		Class<D> delegateClass,
    +		Function<D, W> wrapperFactory,
    +		Set<Method> skipMethodSet) {
    +
    +		Preconditions.checkNotNull(delegateClass);
    +		Preconditions.checkNotNull(wrapperFactory);
    +		Preconditions.checkNotNull(skipMethodSet);
    +
    +		D delegate = spy(delegateClass);
    +		W wrapper = wrapperFactory.apply(delegate);
    +
    +		// ensure that wrapper is a subtype of delegate
    +		Preconditions.checkArgument(delegateClass.isAssignableFrom(wrapper.getClass()));
    +
    +		for (Method delegateMethod : delegateClass.getMethods()) {
    +
    +			if (checkSkipMethodForwardCheck(delegateMethod, skipMethodSet)) {
    +				continue;
    +			}
    +
    +			try {
    +				// find the correct method to substitute the bridge for erased generic types.
    +				// if this doesn't work, the user need to exclude the method and write an additional test.
    +				Method wrapperMethod = wrapper.getClass().getDeclaredMethod(
    +					delegateMethod.getName(),
    +					delegateMethod.getParameterTypes());
    +
    +				// things get a bit fuzzy here, best effort to find a match but this might end up with a wrong method.
    +				if (wrapperMethod.isBridge()) {
    +					for (Method method : wrapper.getClass().getDeclaredMethods()) {
    +						if (!method.isBridge()
    +							&& method.getName().equals(wrapperMethod.getName())
    +							&& method.getParameterCount() == wrapperMethod.getParameterCount()) {
    +							wrapperMethod = method;
    +							break;
    +						}
    +					}
    +				}
    +
    +				Class<?>[] parameterTypes = wrapperMethod.getParameterTypes();
    +				Object[] arguments = new Object[parameterTypes.length];
    +				for (int j = 0; j < arguments.length; j++) {
    +					Class<?> parameterType = parameterTypes[j];
    +					if (parameterType.isArray()) {
    +						arguments[j] = Array.newInstance(parameterType.getComponentType(), 0);
    +					} else if (parameterType.isPrimitive()) {
    +						arguments[j] = 0;
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165104909
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotResult.java ---
    @@ -0,0 +1,83 @@
    +package org.apache.flink.runtime.state;/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +import org.apache.flink.util.ExceptionUtils;
    +
    +/**
    + *
    --- End diff --
    
    Missing JavaDocs.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160717297
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperatorStateContext.java ---
    @@ -0,0 +1,73 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
    +import org.apache.flink.runtime.state.CheckpointStreamFactory;
    +import org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.OperatorStateBackend;
    +import org.apache.flink.runtime.state.StatePartitionStreamProvider;
    +import org.apache.flink.util.CloseableIterable;
    +
    +/**
    + * This interface represents a context from which a stream operator can initialize everything connected state to
    --- End diff --
    
    Possibly a typo: `s/"can initialize everything connected state to state"/"can initialize everything connected to state"/g`


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165340082
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -46,26 +52,63 @@
     	/** */
     	private final int subtaskIndex;
     
    +	/** */
    +	private final Map<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
    +	/** This is the base directory for all local state of the subtask that owns this {@link TaskLocalStateStore}. */
    +	private final File subtaskLocalStateBaseDirectory;
    +
     	public TaskLocalStateStore(
     		JobID jobID,
     		JobVertexID jobVertexID,
    -		int subtaskIndex) {
    +		int subtaskIndex,
    +		File localStateRootDir) {
     
     		this.jobID = jobID;
     		this.jobVertexID = jobVertexID;
     		this.subtaskIndex = subtaskIndex;
    +		this.storedTaskStateByCheckpointID = new HashMap<>();
    +		this.subtaskLocalStateBaseDirectory =
    +			new File(localStateRootDir, createSubtaskPath(jobID, jobVertexID, subtaskIndex));
    +	}
    +
    +	static String createSubtaskPath(JobID jobID, JobVertexID jobVertexID, int subtaskIndex) {
    +		return "jid-" + jobID + "_vtx-" + jobVertexID + "_sti-" + subtaskIndex;
     	}
     
     	public void storeLocalState(
     		@Nonnull CheckpointMetaData checkpointMetaData,
     		@Nullable TaskStateSnapshot localState) {
     
    -		if (localState != null) {
    -			throw new UnsupportedOperationException("Implement this before actually providing local state!");
    +		TaskStateSnapshot previous =
    +			storedTaskStateByCheckpointID.put(checkpointMetaData.getCheckpointId(), localState);
    +
    +		if (previous != null) {
    +			throw new IllegalStateException("Found previously registered local state for checkpoint with id " +
    +				checkpointMetaData.getCheckpointId() + "! This indicated a problem.");
     		}
     	}
     
    -	public void dispose() {
    -		//TODO
    +	public void dispose() throws Exception {
    +
    +		Exception collectedException = null;
    +
    +		for (TaskStateSnapshot snapshot : storedTaskStateByCheckpointID.values()) {
    +			try {
    +				snapshot.discardState();
    +			} catch (Exception discardEx) {
    +				collectedException = ExceptionUtils.firstOrSuppressed(discardEx, collectedException);
    +			}
    +		}
    +
    +		if (collectedException != null) {
    +			throw collectedException;
    +		}
    +
    +		FileUtils.deleteDirectoryQuietly(subtaskLocalStateBaseDirectory);
    --- End diff --
    
    Is there a way to retry the non discarded state handles based on this directory? If not, then we could delete it also in case of a failure.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168500911
  
    --- Diff: flink-runtime/src/main/scala/org/apache/flink/runtime/taskmanager/TaskManager.scala ---
    @@ -120,15 +120,16 @@ import scala.language.postfixOps
      *      requires a clean JVM.
      */
     class TaskManager(
    -    protected val config: TaskManagerConfiguration,
    -    protected val resourceID: ResourceID,
    -    protected val location: TaskManagerLocation,
    -    protected val memoryManager: MemoryManager,
    -    protected val ioManager: IOManager,
    -    protected val network: NetworkEnvironment,
    -    protected val numberOfSlots: Int,
    -    protected val highAvailabilityServices: HighAvailabilityServices,
    -    protected val taskManagerMetricGroup: TaskManagerMetricGroup)
    +                   protected val config: TaskManagerConfiguration,
    +                   protected val resourceID: ResourceID,
    +                   protected val location: TaskManagerLocation,
    +                   protected val memoryManager: MemoryManager,
    +                   protected val ioManager: IOManager,
    +                   protected val network: NetworkEnvironment,
    +                   protected val taskManagerLocalStateStoresManager: TaskExecutorLocalStateStoresManager,
    +                   protected val numberOfSlots: Int,
    +                   protected val highAvailabilityServices: HighAvailabilityServices,
    +                   protected val taskManagerMetricGroup: TaskManagerMetricGroup)
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168480873
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/DuplicatingCheckpointOutputStreamTest.java ---
    @@ -0,0 +1,242 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.FSDataInputStream;
    +import org.apache.flink.runtime.testutils.CommonTestUtils;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import java.io.IOException;
    +import java.util.Random;
    +
    +public class DuplicatingCheckpointOutputStreamTest {
    --- End diff --
    
    Missing `TestLogger`


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168495596
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalRecoveryConfigBase.java ---
    @@ -0,0 +1,37 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +/**
    + * Base class for backend-specific configurations for task-local state recovery. This gives access to the root
    + * directories into which all file-based snapshots can be written.
    + */
    +public class LocalRecoveryConfigBase {
    +
    +	/** Encapsulates the root directories and the subtask-specific path. */
    +	private final LocalRecoveryDirectoryProvider localStateDirectories;
    --- End diff --
    
    Changed in later commit, so that it even is `Nonnull`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159880011
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamFactory.java ---
    @@ -43,9 +44,9 @@ CheckpointStateOutputStream createCheckpointStateOutputStream(
     	 * Closes the stream factory, releasing all internal resources, but does not delete any
     	 * persistent checkpoint data.
     	 *
    -	 * @throws Exception Exceptions can be forwarded and will be logged by the system
    +	 * @throws IOException Exceptions can be forwarded and will be logged by the system
     	 */
    -	void close() throws Exception;
    +	void close() throws IOException;
    --- End diff --
    
    What's this change for? Is it used anywhere? It seems to me like `close()` is just a dead code and this method could be dropped, or did I missed something?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168484993
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    I think the comment does not match the highlighted lines, can you please check?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168229825
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizer.java ---
    @@ -0,0 +1,83 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.SnapshotResult;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.FutureUtil;
    +
    +import java.util.concurrent.ExecutionException;
    +
    +/**
    + * TODO write comment.
    + */
    +public class OperatorSnapshotFinalizer {
    +
    +	private final OperatorSubtaskState jobManagerOwnedState;
    +	private final OperatorSubtaskState taskLocalState;
    +
    +	public OperatorSnapshotFinalizer(
    +		OperatorSnapshotFutures snapshotFutures) throws ExecutionException, InterruptedException {
    +
    +		SnapshotResult<KeyedStateHandle> keyedManaged =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getKeyedStateManagedFuture());
    +
    +		SnapshotResult<KeyedStateHandle> keyedRaw =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getKeyedStateRawFuture());
    +
    +		SnapshotResult<OperatorStateHandle> operatorManaged =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getOperatorStateManagedFuture());
    +
    +		SnapshotResult<OperatorStateHandle> operatorRaw =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getOperatorStateRawFuture());
    +
    +		jobManagerOwnedState = new OperatorSubtaskState(
    +			extractJobManagerOwned(operatorManaged),
    +			extractJobManagerOwned(operatorRaw),
    +			extractJobManagerOwned(keyedManaged),
    +			extractJobManagerOwned(keyedRaw)
    +		);
    +
    +		taskLocalState = new OperatorSubtaskState(
    +			extractTaskLocal(operatorManaged),
    +			extractTaskLocal(operatorRaw),
    +			extractTaskLocal(keyedManaged),
    +			extractTaskLocal(keyedRaw)
    +		);
    +	}
    +
    +	public OperatorSubtaskState getTaskLocalState() {
    +		return taskLocalState;
    +	}
    +
    +	public OperatorSubtaskState getJobManagerOwnedState() {
    +		return jobManagerOwnedState;
    +	}
    +
    +	private <T extends StateObject> T extractJobManagerOwned(SnapshotResult<T> snapshotResult) {
    +		return snapshotResult != null ? snapshotResult.getJobManagerOwnedSnapshot() : null;
    --- End diff --
    
    This method is no longer required, with the changes from other comments that make sure `SnapshotResult` is never `null`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168501441
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -19,92 +19,224 @@
     package org.apache.flink.runtime.state;
     
     import org.apache.flink.api.common.JobID;
    -import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
     import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
     import org.apache.flink.runtime.jobgraph.JobVertexID;
    -import org.apache.flink.util.ExceptionUtils;
     
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
     import javax.annotation.Nonnull;
     import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
     
     import java.io.File;
    -import java.util.HashMap;
    +import java.util.Arrays;
    +import java.util.Iterator;
     import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
     
     /**
      * This class will service as a task-manager-level local storage for local checkpointed state. The purpose is to provide
      * access to a state that is stored locally for a faster recovery compared to the state that is stored remotely in a
      * stable store DFS. For now, this storage is only complementary to the stable storage and local state is typically
      * lost in case of machine failures. In such cases (and others), client code of this class must fall back to using the
      * slower but highly available store.
    - *
    - * TODO this is currently a placeholder / mock that still must be implemented!
      */
     public class TaskLocalStateStore {
     
    -	/** */
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStore.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	private static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    --- End diff --
    
    Can this be static?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170039184
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TestLocalDirectoryProvider.java ---
    @@ -0,0 +1,40 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +
    +import java.io.File;
    +
    +/**
    + * Easy-to-construct implementation of {@link LocalRecoveryDirectoryProvider} to use in tests.
    + */
    +public class TestLocalDirectoryProvider extends LocalRecoveryDirectoryProviderImpl {
    +
    +	public TestLocalDirectoryProvider() {
    +		super(
    +			new File(System.getProperty("java.io.tmpdir")),
    --- End diff --
    
    I wouldn't use the tmp dir, but instead pass in a pass. Then the user can instantiate `TemporaryFolder` to make sure that things are cleaned up.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170228635
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -89,6 +90,22 @@ public TaskExecutorLocalStateStoresManager(
     				}
     			}
     		}
    +
    +		// install a shutdown hook
    +		this.shutdownHook = new Thread("TaskExecutorLocalStateStoresManager shutdown hook") {
    --- End diff --
    
    I made a separate commit that introduces `ShudownHookUtil`. Used it to deduplicated the same code in like 20 places...


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160913576
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskStateManagerImpl.java ---
    @@ -0,0 +1,115 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetrics;
    +import org.apache.flink.runtime.checkpoint.JobManagerTaskRestore;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +import org.apache.flink.runtime.taskmanager.CheckpointResponder;
    +
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +
    +/**
    + * This class is the default implementation of {@link TaskStateManager} and collaborates with the job manager
    + * through {@link CheckpointResponder}) as well as a task-manager-local state store. Like this, client code does
    + * not have to deal with the differences between remote or local state on recovery because this class handles both
    + * cases transparently.
    + *
    + * Reported state is tagged by clients so that this class can properly forward to the right receiver for the
    + * checkpointed state.
    + *
    + * TODO: all interaction with local state store must still be implemented! It is currently just a placeholder.
    + */
    +public class TaskStateManagerImpl implements TaskStateManager {
    +
    +	/** The id of the job for which this manager was created, can report, and recover. */
    +	private final JobID jobId;
    +
    +	/** The execution attempt id that this manager reports for. */
    +	private final ExecutionAttemptID executionAttemptID;
    +
    +	/** The data given by the job manager to restore the job. This is not set for a new job without previous state. */
    +	private final JobManagerTaskRestore jobManagerTaskRestore;
    +
    +	/** The local state store to which this manager reports local state snapshots. */
    +	private final TaskLocalStateStore localStateStore;
    +
    +	/** The checkpoint responder through which this manager can report to the job manager. */
    +	private final CheckpointResponder checkpointResponder;
    +
    +	public TaskStateManagerImpl(
    +		JobID jobId,
    +		ExecutionAttemptID executionAttemptID,
    +		TaskLocalStateStore localStateStore,
    +		JobManagerTaskRestore jobManagerTaskRestore,
    +		CheckpointResponder checkpointResponder) {
    +
    +		this.jobId = jobId;
    +		this.localStateStore = localStateStore;
    --- End diff --
    
    can you add `checkNotNulls` on the parameters? Or use `Optional<>` or mark them `@Nullable` accordingly?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168169149
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTerminationTest.java ---
    @@ -227,21 +229,13 @@ protected void cleanup() throws Exception {
     			// has been stopped
     			CLEANUP_LATCH.trigger();
     
    -			// wait until handle async exception has been called to proceed with the termination of the
    -			// StreamTask
    -			HANDLE_ASYNC_EXCEPTION_LATCH.await();
    --- End diff --
    
    Already done in a later commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165341653
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java ---
    @@ -0,0 +1,122 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.configuration.ConfigConstants;
    +import org.apache.flink.configuration.Configuration;
    +import org.apache.flink.runtime.clusterframework.types.ResourceID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +import org.apache.flink.runtime.taskexecutor.TaskManagerServices;
    +import org.apache.flink.runtime.taskexecutor.TaskManagerServicesConfiguration;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +import org.junit.rules.TemporaryFolder;
    +
    +import java.io.File;
    +import java.net.InetAddress;
    +
    +public class TaskExecutorLocalStateStoresManagerTest {
    +
    +	/**
    +	 * This tests that the creation of {@link TaskManagerServices} correctly creates the local state root directory
    +	 * for the {@link TaskExecutorLocalStateStoresManager} with the configured root directory.
    +	 */
    +	@Test
    +	public void testCreationFromConfig() throws Exception {
    +
    +		final Configuration config = new Configuration();
    +
    +		final String rootDirString = "localStateRoot";
    +		config.setString(ConfigConstants.TASK_MANAGER_LOCAL_STATE_ROOT_DIR_KEY, rootDirString);
    +
    +		final ResourceID tmResourceID = ResourceID.generate();
    +
    +		TaskManagerServicesConfiguration taskManagerServicesConfiguration =
    +			TaskManagerServicesConfiguration.fromConfiguration(config, InetAddress.getLocalHost(), true);
    +
    +		TaskManagerServices taskManagerServices =
    +			TaskManagerServices.fromConfiguration(taskManagerServicesConfiguration, tmResourceID);
    +
    +		TaskExecutorLocalStateStoresManager taskStateManager = taskManagerServices.getTaskStateManager();
    +
    +		Assert.assertEquals(
    +			new File(rootDirString, TaskManagerServices.LOCAL_STATE_SUB_DIRECTORY_ROOT),
    +			taskStateManager.getLocalStateRootDirectory());
    +
    +		Assert.assertEquals("localState", TaskManagerServices.LOCAL_STATE_SUB_DIRECTORY_ROOT);
    --- End diff --
    
    This assertion seems to be redundant.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169062508
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    Thanks, but like I said, I have this already done in a hotfix that I will commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165309220
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java ---
    @@ -3404,12 +3405,16 @@ public String fold(String acc, Integer value) throws Exception {
     		}
     	}
     
    -	protected KeyedStateHandle runSnapshot(RunnableFuture<KeyedStateHandle> snapshotRunnableFuture) throws Exception {
    +	protected KeyedStateHandle runSnapshot(
    +		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunnableFuture) throws Exception {
    +
     		if(!snapshotRunnableFuture.isDone()) {
     			Thread runner = new Thread(snapshotRunnableFuture);
     			runner.start();
     		}
    -		return snapshotRunnableFuture.get();
    +
    +		SnapshotResult<KeyedStateHandle> snapshotResult = snapshotRunnableFuture.get();
    +		return snapshotResult != null ? snapshotResult.getJobManagerOwnedSnapshot() : null;
    --- End diff --
    
    By enforcing that `snapshotResult` is always not null, we could get rid of all these ternary statements.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168752264
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168176995
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java ---
    @@ -0,0 +1,122 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.configuration.ConfigConstants;
    +import org.apache.flink.configuration.Configuration;
    +import org.apache.flink.runtime.clusterframework.types.ResourceID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +import org.apache.flink.runtime.taskexecutor.TaskManagerServices;
    +import org.apache.flink.runtime.taskexecutor.TaskManagerServicesConfiguration;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +import org.junit.rules.TemporaryFolder;
    +
    +import java.io.File;
    +import java.net.InetAddress;
    +
    +public class TaskExecutorLocalStateStoresManagerTest {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168500772
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -59,58 +71,55 @@ public TaskExecutorLocalStateStoresManager(File[] localStateRootDirectories) {
     				}
     			}
     		}
    -	}
     
    -	public TaskLocalStateStore localStateStoreForTask(
    -		JobID jobId,
    -		JobVertexID jobVertexID,
    -		int subtaskIndex) {
    +	}
     
    -		Preconditions.checkNotNull(jobId);
    -		final JobVertexSubtaskKey taskKey = new JobVertexSubtaskKey(jobVertexID, subtaskIndex);
    +	public TaskLocalStateStore localStateStoreForSubtask(
    +		@Nonnull JobID jobId,
    +		@Nonnull AllocationID allocationID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex) {
     
     		final Map<JobVertexSubtaskKey, TaskLocalStateStore> taskStateManagers =
    -			this.taskStateStoresByJobID.computeIfAbsent(jobId, k -> new HashMap<>());
    +			this.taskStateStoresByAllocationID.computeIfAbsent(allocationID, k -> new ConcurrentHashMap<>());
    +
    +		final JobVertexSubtaskKey taskKey = new JobVertexSubtaskKey(jobVertexID, subtaskIndex);
     
     		return taskStateManagers.computeIfAbsent(
    -			taskKey, k -> new TaskLocalStateStore(jobId, jobVertexID, subtaskIndex, localStateRootDirectories));
    +			taskKey,
    +			k -> new TaskLocalStateStore(jobId, jobVertexID, subtaskIndex, localStateRootDirectories, discardExecutor));
     	}
     
    -	public void releaseJob(JobID jobID) {
    +	public void releaseLocalStateForAllocationId(@Nonnull AllocationID allocationID) {
     
    -		Map<JobVertexSubtaskKey, TaskLocalStateStore> cleanupLocalStores = taskStateStoresByJobID.remove(jobID);
    +		Map<JobVertexSubtaskKey, TaskLocalStateStore> cleanupLocalStores =
    +			taskStateStoresByAllocationID.remove(allocationID);
     
     		if (cleanupLocalStores != null) {
    -//			doRelease(cleanupLocalStores.values());
    +			doRelease(cleanupLocalStores.values());
     		}
     	}
     
     	public void releaseAll() {
     
    -		for (Map<JobVertexSubtaskKey, TaskLocalStateStore> stateStoreMap : taskStateStoresByJobID.values()) {
    -//			doRelease(stateStoreMap.values());
    +		for (Map<JobVertexSubtaskKey, TaskLocalStateStore> stateStoreMap : taskStateStoresByAllocationID.values()) {
    +			doRelease(stateStoreMap.values());
     		}
     
    -		taskStateStoresByJobID.clear();
    +		taskStateStoresByAllocationID.clear();
     	}
     
    -	private void doRelease(Iterable<TaskLocalStateStore> toRelease) throws Exception {
    +	private void doRelease(Iterable<TaskLocalStateStore> toRelease) {
     
     		if (toRelease != null) {
     
    -			Exception collectedExceptions = null;
    -
     			for (TaskLocalStateStore stateStore : toRelease) {
     				try {
     					stateStore.dispose();
     				} catch (Exception disposeEx) {
    -					collectedExceptions = ExceptionUtils.firstOrSuppressed(disposeEx, collectedExceptions);
    +					LOG.warn("Exception while disposing local state store", disposeEx);
    --- End diff --
    
    I think it would be interesting to know which LocalStateStore could not be disposed.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168529359
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TestTaskStateManager.java ---
    @@ -118,12 +121,38 @@ public void reportTaskStateSnapshots(
     		}
     	}
     
    +	@Nonnull
     	@Override
    -	public OperatorSubtaskState operatorStates(OperatorID operatorID) {
    -		TaskStateSnapshot taskStateSnapshot = getLastJobManagerTaskStateSnapshot();
    -		return taskStateSnapshot != null ? taskStateSnapshot.getSubtaskStateByOperatorID(operatorID) : null;
    +	public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {
    +		TaskStateSnapshot jmTaskStateSnapshot = getLastJobManagerTaskStateSnapshot();
    +		TaskStateSnapshot tmTaskStateSnapshot = getLastTaskManagerTaskStateSnapshot();
    +
    +		OperatorSubtaskState jmOpState = null;
    +		List<OperatorSubtaskState> tmStateCollection = null;
    --- End diff --
    
    Changed in later commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169651017
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
    --- End diff --
    
    Got it, ok!


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170026181
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -89,6 +90,22 @@ public TaskExecutorLocalStateStoresManager(
     				}
     			}
     		}
    +
    +		// install a shutdown hook
    +		this.shutdownHook = new Thread("TaskExecutorLocalStateStoresManager shutdown hook") {
    --- End diff --
    
    It might make sense to reuse `BlobUtils#addShutdownHook` or to refactor this and then reuse.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168486610
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamWithResultProvider.java ---
    @@ -0,0 +1,175 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.state.filesystem.FixFileFsStateOutputStream;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import java.io.Closeable;
    +import java.io.File;
    +import java.io.IOException;
    +
    +/**
    + * Interface that provides access to a CheckpointStateOutputStream and a method to provide the {@link SnapshotResult}.
    + * This abstracts from different ways that a result is obtained from checkpoint output streams.
    + */
    +public interface CheckpointStreamWithResultProvider extends Closeable {
    +
    +	SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +		KeyGroupRangeOffsets kgOffs) throws IOException;
    +
    +	CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream();
    +
    +	@Override
    +	default void close() throws IOException {
    +		CheckpointStreamFactory.CheckpointStateOutputStream outputStream = getCheckpointOutputStream();
    +		if (outputStream != null) {
    +			outputStream.close();
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that only creates the
    +	 * primary/remote/jm-owned state.
    +	 */
    +	class PrimaryStreamOnly implements CheckpointStreamWithResultProvider {
    +
    +		private final CheckpointStreamFactory.CheckpointStateOutputStream outputStream;
    +
    +		public PrimaryStreamOnly(CheckpointStreamFactory.CheckpointStateOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets kgOffs) throws IOException {
    +
    +			StreamStateHandle streamStateHandle = outputStream.closeAndGetHandle();
    +			KeyGroupsStateHandle primaryStateHandle = new KeyGroupsStateHandle(kgOffs, streamStateHandle);
    +			return new SnapshotResult<>(primaryStateHandle, null);
    +		}
    +
    +		@Override
    +		public CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that creates both, the
    +	 * primary/remote/jm-owned state and the secondary/local/tm-owned state.
    +	 */
    +	class PrimaryAndSecondaryStream implements CheckpointStreamWithResultProvider {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		private final DuplicatingCheckpointOutputStream outputStream;
    +
    +		public PrimaryAndSecondaryStream(DuplicatingCheckpointOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets offsets) throws IOException {
    +
    +			final StreamStateHandle primaryStreamStateHandle;
    +
    +			try {
    +				primaryStreamStateHandle = outputStream.closeAndGetPrimaryHandle();
    +			} catch (IOException primaryEx) {
    +				try {
    +					outputStream.close();
    +				} catch (IOException closeEx) {
    +					primaryEx = ExceptionUtils.firstOrSuppressed(closeEx, primaryEx);
    +				}
    +				throw primaryEx;
    +			}
    +
    +			StreamStateHandle secondaryStreamStateHandle = null;
    +
    +			try {
    +				secondaryStreamStateHandle = outputStream.closeAndGetSecondaryHandle();
    +			} catch (IOException secondaryEx) {
    +				LOG.warn("Exception from secondary/local checkpoint stream.", secondaryEx);
    +			}
    +
    +			if (primaryStreamStateHandle != null) {
    +
    +				final KeyGroupsStateHandle primaryKeyGroupsStateHandle =
    +					new KeyGroupsStateHandle(offsets, primaryStreamStateHandle);
    +
    +				final KeyGroupsStateHandle secondaryKeyGroupsStateHandle = secondaryStreamStateHandle != null ?
    +					new KeyGroupsStateHandle(offsets, secondaryStreamStateHandle) :
    +					null;
    +
    +				return new SnapshotResult<>(primaryKeyGroupsStateHandle, secondaryKeyGroupsStateHandle);
    +			}
    +
    +			return null;
    +		}
    +
    +		@Override
    +		public DuplicatingCheckpointOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	class Factory {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		public CheckpointStreamWithResultProvider create(
    +			long checkpointId,
    +			long timestamp,
    +			CheckpointStreamFactory primaryStreamFactory,
    +			LocalRecoveryDirectoryProvider secondaryStreamDirProvider) throws Exception {
    +
    +			CheckpointStreamFactory.CheckpointStateOutputStream primaryOut =
    +				primaryStreamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
    +
    +			CheckpointStreamFactory.CheckpointStateOutputStream secondaryOut;
    +
    +			if (secondaryStreamDirProvider != null) {
    --- End diff --
    
    Why do we allow passing in `null` values? Wouldn't it be easier to have two methods for the simple and duplicated stream?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162583721
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskStateManager.java ---
    @@ -0,0 +1,62 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetrics;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +
    +/**
    + * This interface provides methods to report and retrieve state for a task.
    + *
    + * <p>When a checkpoint or savepoint is triggered on a task, it will create snapshots for all stream operator instances
    + * it owns. All operator snapshots from the task are then reported via this interface. A typical implementation will
    + * dispatch and forward the reported state information to interested parties such as the checkpoint coordinator or a
    + * local state store.
    + *
    + * <p>This interface also offers the complementary method that provides access to previously saved state of operator
    + * instances in the task for restore purposes.
    + */
    +public interface TaskStateManager extends CheckpointListener {
    --- End diff --
    
    I prefer to keep it as is for the reasons we discussed offline.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162584371
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java ---
    @@ -316,7 +292,38 @@ private void initOperatorState(Collection<OperatorStateHandle> operatorStateHand
     	 * @throws Exception An exception in this method causes the operator to fail.
     	 */
     	@Override
    -	public void close() throws Exception {}
    +	public void close() throws Exception {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168477254
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalRecoveryConfigBase.java ---
    @@ -0,0 +1,37 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +/**
    + * Base class for backend-specific configurations for task-local state recovery. This gives access to the root
    + * directories into which all file-based snapshots can be written.
    + */
    +public class LocalRecoveryConfigBase {
    +
    +	/** Encapsulates the root directories and the subtask-specific path. */
    +	private final LocalRecoveryDirectoryProvider localStateDirectories;
    --- End diff --
    
    Should be `Nullable` according to usage.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165329369
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java ---
    @@ -104,10 +107,8 @@
     	@Nullable
     	private OptionsFactory optionsFactory;
     
    -	/** True if incremental checkpointing is enabled.
    -	 * Null if not yet set, in which case the configuration values will be used. */
    -	@Nullable
    -	private Boolean enableIncrementalCheckpointing;
    +	/** True if incremental checkpointing is enabled. */
    --- End diff --
    
    Please update comment.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168762457
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -266,11 +305,16 @@ public RocksDBKeyedStateBackend(
     		RocksIterator iterator = db.newIterator(columnInfo.f0);
     		iterator.seekToFirst();
     
    -		Iterable<K> iterable = () -> new RocksIteratorWrapper<>(iterator, state, keySerializer, keyGroupPrefixBytes);
    +		Iterable<K> iterable = () -> new RocksIteratorToJavaIteratorAdapter<>(iterator, state, keySerializer, keyGroupPrefixBytes);
     		Stream<K> targetStream = StreamSupport.stream(iterable.spliterator(), false);
    --- End diff --
    
    Yes, that seems a little odd, but I suggest to make this a separate issue because this PR has not introduced or changed this, except for that the class was renamed.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165107880
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotResult.java ---
    @@ -0,0 +1,83 @@
    +package org.apache.flink.runtime.state;/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +import org.apache.flink.util.ExceptionUtils;
    +
    +/**
    + *
    + */
    +public class SnapshotResult<T extends StateObject> implements StateObject {
    --- End diff --
    
    That way we would also get rid of all the null checks.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162584402
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperatorStateContext.java ---
    @@ -0,0 +1,73 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
    +import org.apache.flink.runtime.state.CheckpointStreamFactory;
    +import org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.OperatorStateBackend;
    +import org.apache.flink.runtime.state.StatePartitionStreamProvider;
    +import org.apache.flink.util.CloseableIterable;
    +
    +/**
    + * This interface represents a context from which a stream operator can initialize everything connected state to
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162584893
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamOperatorSnapshotRestoreTest.java ---
    @@ -98,7 +101,26 @@ public Integer getKey(Integer value) throws Exception {
     				TypeInformation.of(Integer.class),
     				MAX_PARALLELISM,
     				1 /* num subtasks */,
    -				0 /* subtask index */);
    +				0 /* subtask index */) {
    +
    +			@Override
    +			protected StreamTaskStateManager createStreamTaskStateManager(
    +				Environment env,
    +				StateBackend stateBackend,
    +				ProcessingTimeService processingTimeService) {
    +
    +				return new StreamTaskStateManagerImpl(env, stateBackend, processingTimeService) {
    +					@Override
    +					protected <K> InternalTimeServiceManager<?, K> internalTimeServiceManager(
    +						AbstractKeyedStateBackend<K> keyedStatedBackend,
    +						KeyContext keyContext,
    +						Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
    +
    +						return null;
    --- End diff --
    
    Yes, this is a workaround, because this is a dedicated test for raw keyed states - which is otherwise consumed immediately to restore a timer service. This will fail because we wrote only test data to the stream, not the data for a timer service.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168168842
  
    --- Diff: flink-core/src/test/java/org/apache/flink/util/MethodForwardingTestUtil.java ---
    @@ -0,0 +1,146 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.util;
    +
    +import org.mockito.Mockito;
    +
    +import java.lang.reflect.Array;
    +import java.lang.reflect.Method;
    +import java.util.Collections;
    +import java.util.Set;
    +import java.util.function.Function;
    +
    +import static org.junit.Assert.fail;
    +import static org.mockito.Mockito.reset;
    +import static org.mockito.Mockito.spy;
    +
    +/**
    + * Helper class with a method that attempt to automatically test method forwarding between a delegate and a wrapper.
    + */
    +public class MethodForwardingTestUtil {
    +
    +	/**
    +	 * This is a best effort automatic test for method forwarding between a delegate and its wrapper, where the wrapper
    +	 * class is a subtype of the delegate. This ignores methods that are inherited from Object.
    +	 *
    +	 * @param delegateClass the class for the delegate.
    +	 * @param wrapperFactory factory that produces a wrapper from a delegate.
    +	 * @param <D> type of the delegate
    +	 * @param <W> type of the wrapper
    +	 */
    +	public static <D, W> void testMethodForwarding(
    +		Class<D> delegateClass,
    +		Function<D, W> wrapperFactory) {
    +		testMethodForwarding(delegateClass, wrapperFactory, Collections.emptySet());
    +	}
    +
    +	/**
    +	 * This is a best effort automatic test for method forwarding between a delegate and its wrapper, where the wrapper
    +	 * class is a subtype of the delegate. Methods can be remapped in case that the implementation does not call the
    +	 * original method. Remapping to null skips the method. This ignores methods that are inherited from Object.
    +	 *
    +	 * @param delegateClass the class for the delegate.
    +	 * @param wrapperFactory factory that produces a wrapper from a delegate.
    +	 * @param skipMethodSet set of methods to ignore.
    +	 * @param <D> type of the delegate
    +	 * @param <W> type of the wrapper
    +	 */
    +	public static <D, W> void testMethodForwarding(
    +		Class<D> delegateClass,
    +		Function<D, W> wrapperFactory,
    +		Set<Method> skipMethodSet) {
    +
    +		Preconditions.checkNotNull(delegateClass);
    +		Preconditions.checkNotNull(wrapperFactory);
    +		Preconditions.checkNotNull(skipMethodSet);
    +
    +		D delegate = spy(delegateClass);
    +		W wrapper = wrapperFactory.apply(delegate);
    +
    +		// ensure that wrapper is a subtype of delegate
    +		Preconditions.checkArgument(delegateClass.isAssignableFrom(wrapper.getClass()));
    +
    +		for (Method delegateMethod : delegateClass.getMethods()) {
    +
    +			if (checkSkipMethodForwardCheck(delegateMethod, skipMethodSet)) {
    +				continue;
    +			}
    +
    +			try {
    +				// find the correct method to substitute the bridge for erased generic types.
    +				// if this doesn't work, the user need to exclude the method and write an additional test.
    +				Method wrapperMethod = wrapper.getClass().getDeclaredMethod(
    +					delegateMethod.getName(),
    +					delegateMethod.getParameterTypes());
    +
    +				// things get a bit fuzzy here, best effort to find a match but this might end up with a wrong method.
    +				if (wrapperMethod.isBridge()) {
    +					for (Method method : wrapper.getClass().getDeclaredMethods()) {
    +						if (!method.isBridge()
    +							&& method.getName().equals(wrapperMethod.getName())
    +							&& method.getParameterCount() == wrapperMethod.getParameterCount()) {
    +							wrapperMethod = method;
    +							break;
    +						}
    +					}
    +				}
    +
    +				Class<?>[] parameterTypes = wrapperMethod.getParameterTypes();
    +				Object[] arguments = new Object[parameterTypes.length];
    +				for (int j = 0; j < arguments.length; j++) {
    +					Class<?> parameterType = parameterTypes[j];
    +					if (parameterType.isArray()) {
    +						arguments[j] = Array.newInstance(parameterType.getComponentType(), 0);
    +					} else if (parameterType.isPrimitive()) {
    +						arguments[j] = 0;
    +					} else {
    +						arguments[j] = Mockito.mock(parameterType);
    +					}
    +				}
    +
    +				wrapperMethod.invoke(wrapper, arguments);
    +				delegateMethod.invoke(Mockito.verify(delegate, Mockito.times(1)), arguments);
    +				reset(delegate);
    +			} catch (Exception ex) {
    +				ex.printStackTrace();
    +				fail("Forwarding test failed: " + ex.getMessage());
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165309957
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/OperatorSnapshotFinalizer.java ---
    @@ -0,0 +1,83 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.SnapshotResult;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.FutureUtil;
    +
    +import java.util.concurrent.ExecutionException;
    +
    +/**
    + * TODO write comment.
    + */
    +public class OperatorSnapshotFinalizer {
    +
    +	private final OperatorSubtaskState jobManagerOwnedState;
    +	private final OperatorSubtaskState taskLocalState;
    +
    +	public OperatorSnapshotFinalizer(
    +		OperatorSnapshotFutures snapshotFutures) throws ExecutionException, InterruptedException {
    +
    +		SnapshotResult<KeyedStateHandle> keyedManaged =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getKeyedStateManagedFuture());
    +
    +		SnapshotResult<KeyedStateHandle> keyedRaw =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getKeyedStateRawFuture());
    +
    +		SnapshotResult<OperatorStateHandle> operatorManaged =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getOperatorStateManagedFuture());
    +
    +		SnapshotResult<OperatorStateHandle> operatorRaw =
    +			FutureUtil.runIfNotDoneAndGet(snapshotFutures.getOperatorStateRawFuture());
    +
    +		jobManagerOwnedState = new OperatorSubtaskState(
    +			extractJobManagerOwned(operatorManaged),
    +			extractJobManagerOwned(operatorRaw),
    +			extractJobManagerOwned(keyedManaged),
    +			extractJobManagerOwned(keyedRaw)
    +		);
    +
    +		taskLocalState = new OperatorSubtaskState(
    +			extractTaskLocal(operatorManaged),
    +			extractTaskLocal(operatorRaw),
    +			extractTaskLocal(keyedManaged),
    +			extractTaskLocal(keyedRaw)
    +		);
    +	}
    +
    +	public OperatorSubtaskState getTaskLocalState() {
    +		return taskLocalState;
    +	}
    +
    +	public OperatorSubtaskState getJobManagerOwnedState() {
    +		return jobManagerOwnedState;
    +	}
    +
    +	private <T extends StateObject> T extractJobManagerOwned(SnapshotResult<T> snapshotResult) {
    +		return snapshotResult != null ? snapshotResult.getJobManagerOwnedSnapshot() : null;
    --- End diff --
    
    We should add `@Nullable` annotation to make sure that this method can return a `null` value.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168891314
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    +
    +- RocksDBStateBackend: task-local recovery is supported for keyed state. For *full checkpoints*, state is duplicated to a local file. This can introduce additional write costs
    +and occupies local disk space. For *incremental snapshots*, the local state is based on RocksDB's native checkpointing mechanism. This mechanism is also used as the first step
    --- End diff --
    
    got you. Thanks for the clarification


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168735119
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    An `exists()` check can be performed before `mkdirs()`, if it exists we can remove it.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160941763
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/taskexecutor/TaskExecutorTest.java ---
    @@ -210,10 +211,11 @@ public void testHeartbeatTimeoutWithJobManager() throws Exception {
     			tmConfig,
     			taskManagerLocation,
     			mock(MemoryManager.class),
    -			mock(IOManager.class),
    +			mock(IOManager.class),mock(TaskExecutorLocalStateStoresManager.class),
    --- End diff --
    
    Can all of the `mock(TaskExecutorLocalStateStoresManager.class)` be replaced with: `new TaskExecutorLocalStateStoresManager();` (with an empty array of for `localStateRootDirectories`) or maybe `new TaskExecutorLocalStateStoresManager(sameThreadExecutor());`?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165106582
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotResult.java ---
    @@ -0,0 +1,83 @@
    +package org.apache.flink.runtime.state;/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +import org.apache.flink.util.ExceptionUtils;
    +
    +/**
    + *
    + */
    +public class SnapshotResult<T extends StateObject> implements StateObject {
    +
    +	private static final long serialVersionUID = 1L;
    +
    +	private final T jobManagerOwnedSnapshot;
    +	private final T taskLocalSnapshot;
    --- End diff --
    
    Maybe annotate as `Nullable`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168510181
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/BackendRestorerProcedure.java ---
    @@ -0,0 +1,143 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.core.fs.CloseableRegistry;
    +import org.apache.flink.runtime.checkpoint.StateObjectCollection;
    +import org.apache.flink.runtime.state.Snapshotable;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.Disposable;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +import org.apache.flink.util.ThrowingSupplier;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.Closeable;
    +import java.io.IOException;
    +import java.util.Iterator;
    +
    +/**
    + * This class implements the logic that creates (and potentially restores) a state backend. The restore logic
    + * considers multiple, prioritized options of snapshots to restore from, where all of the options should recreate
    + * the same state for the backend. When we fail to restore from the snapshot with the highest priority (typically
    + * the "fastest" to restore), we fallback to the next snapshot with the next highest priority. We also take care
    + * of cleaning up from failed restore attempts. We only reattempt when the problem occurs during the restore call
    + * and will only stop after all snapshot alternatives are exhausted and all failed.
    + *
    + * @param <T> type of the restored backend.
    + * @param <S> type of the supplied snapshots from which the backend restores.
    + */
    +public class BackendRestorerProcedure<
    +	T extends Closeable & Disposable & Snapshotable<?, StateObjectCollection<S>>,
    +	S extends StateObject> {
    +
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(BackendRestorerProcedure.class);
    +
    +	/** Factory for new, fresh backends without state. */
    +	private final ThrowingSupplier<T, Exception> instanceSupplier;
    +
    +	/** This registry is used so that recovery can participate in the task lifecycle, i.e. can be canceled. */
    +	private final CloseableRegistry backendCloseableRegistry;
    +
    +	/**
    +	 * Creates a new backend restorer using the given backend supplier and the closeable registry.
    +	 *
    +	 * @param instanceSupplier factory function for new, empty backend instances.
    +	 * @param backendCloseableRegistry registry to allow participation in task lifecycle, e.g. react to cancel.
    +	 */
    +	public BackendRestorerProcedure(
    +		@Nonnull ThrowingSupplier<T, Exception> instanceSupplier,
    +		@Nonnull CloseableRegistry backendCloseableRegistry) {
    +
    +		this.instanceSupplier = Preconditions.checkNotNull(instanceSupplier);
    +		this.backendCloseableRegistry = Preconditions.checkNotNull(backendCloseableRegistry);
    +	}
    +
    +	/**
    +	 * Creates a new state backend and restores it from the provided set of state snapshot alternatives.
    +	 *
    +	 * @param restoreOptions iterator over a prioritized set of state snapshot alternatives for recovery.
    +	 * @return the created (and restored) state backend.
    +	 * @throws Exception if the backend could not be created or restored.
    +	 */
    +	public @Nonnull T createAndRestore(@Nonnull Iterator<StateObjectCollection<S>> restoreOptions) throws Exception {
    +
    +		StateObjectCollection<S> restoreState = null;
    +
    +		boolean retry;
    +
    +		while (true) {
    --- End diff --
    
    I think we could get rid of the endless loop here by simply using the iterator.
    ```
    while (restoreOptions.hasNext()) {
      final StateObjectCollection<S> restoreState = restoreOptions.next();
    
      try {
         ...
      } catch (Exception e) {
        ...
      }
    }
    
    throw new FlinkException("Could not restore from any of the provided restore options");
    ```


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169895663
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -1552,10 +1556,33 @@ public IncrementalSnapshotStrategy() {
     				return DoneFuture.nullValue();
     			}
     
    +			SnapshotDirectory snapshotDirectory;
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by asfgit <gi...@git.apache.org>.
Github user asfgit closed the pull request at:

    https://github.com/apache/flink/pull/5239


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170251605
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -270,6 +277,8 @@ public RocksDBKeyedStateBackend(
     			new IncrementalSnapshotStrategy() :
     			new FullSnapshotStrategy();
     
    +		this.writeOptions = new WriteOptions().setDisableWAL(false);
    --- End diff --
    
    We enable WAL in contrast to the previous state. This is intended, right?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160136572
  
    --- Diff: flink-contrib/flink-storm/src/test/java/org/apache/flink/storm/wrappers/BoltWrapperTest.java ---
    @@ -375,11 +376,13 @@ public void declareOutputFields(OutputFieldsDeclarer declarer) {
     		when(env.getMetricGroup()).thenReturn(new UnregisteredTaskMetricsGroup());
     		when(env.getTaskManagerInfo()).thenReturn(new TestingTaskManagerRuntimeInfo());
     
    +		final CloseableRegistry closeableRegistry = new CloseableRegistry();
     		StreamTask<?, ?> mockTask = mock(StreamTask.class);
     		when(mockTask.getCheckpointLock()).thenReturn(new Object());
     		when(mockTask.getConfiguration()).thenReturn(new StreamConfig(new Configuration()));
     		when(mockTask.getEnvironment()).thenReturn(env);
     		when(mockTask.getExecutionConfig()).thenReturn(execConfig);
    +		when(mockTask.getCancelables()).thenReturn(closeableRegistry);
    --- End diff --
    
    I don't like the idea of postponing such things. I really would like it to be done in this pr (as a separate commit, ideally at the bottom to avoid modifying same lines of code twice, but as a last resort it could also be the last commit). Otherwise we will forget about it and while this only one added line it increases our technological debt and makes our code base a tiny bit worse then it used to be before - both of those things are quite dangerous. 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159865180
  
    --- Diff: flink-core/src/main/java/org/apache/flink/util/CloseableIterator.java ---
    @@ -18,11 +18,38 @@
     
     package org.apache.flink.util;
     
    -import org.apache.flink.annotation.Internal;
    +import java.io.Closeable;
    +import java.io.IOException;
    +import java.util.Iterator;
     
     /**
    - * Tagging interface for migration related classes.
    + * This interface represents an iterator that is also closeable.
    + *
    + * @param <T> type of the iterated objects.
      */
    -@Internal
    -public interface Migration {
    +public interface CloseableIterator<T> extends Iterator<T>, Closeable {
    --- End diff --
    
    I can not find usages of this in this commit. Is it used here somehow? If not, please remove it/extract to separate commit/move to appropriate commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169728016
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotDirectory.java ---
    @@ -0,0 +1,172 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.FileStatus;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.IOException;
    +import java.util.concurrent.atomic.AtomicReference;
    +
    +/**
    + * This class represents a directory that is the target for a state snapshot. This class provides some method that
    + * simplify resource management when dealing with such directories, e.g. it can produce a {@link DirectoryStateHandle}
    + * when the snapshot is completed and disposal considers whether or not a snapshot was already completed. For a
    + * completed snapshot, the ownership for cleanup is transferred to the created directory state handle. For incomplete
    + * snapshots, calling {@link #deleteIfIncompleteSnapshot()} will delete the underlying directory resource.
    + */
    +public class SnapshotDirectory {
    +
    +	/**
    +	 * Lifecycle stages of a snapshot directory.
    +	 */
    +	enum State {
    +		ONGOING, COMPLETED, DELETED
    +	}
    +
    +	/** This path describes the underlying directory for the snapshot. */
    +	@Nonnull
    +	private final Path directory;
    +
    +	/** The filesystem that contains the snapshot directory. */
    +	@Nonnull
    +	private final FileSystem fileSystem;
    +
    +	/** This reference tracks the lifecycle state of the snapshot directory. */
    +	@Nonnull
    +	private AtomicReference<State> state;
    +
    +	public SnapshotDirectory(@Nonnull Path directory, @Nonnull FileSystem fileSystem) {
    +		this.directory = directory;
    +		this.fileSystem = fileSystem;
    +		this.state = new AtomicReference<>(State.ONGOING);
    +	}
    +
    +	public SnapshotDirectory(@Nonnull Path directory) throws IOException {
    +		this(directory, directory.getFileSystem());
    +	}
    +
    +	@Nonnull
    +	public Path getDirectory() {
    +		return directory;
    +	}
    +
    +	public boolean mkdirs() throws IOException {
    +		return fileSystem.mkdirs(directory);
    +	}
    +
    +	@Nonnull
    +	public FileSystem getFileSystem() {
    +		return fileSystem;
    +	}
    +
    +	public boolean exists() throws IOException {
    +		return fileSystem.exists(directory);
    +	}
    +
    +	/**
    +	 * List the statuses of the files/directories in the snapshot directory.
    +	 *
    +	 * @return the statuses of the files/directories in the given path.
    +	 * @throws IOException if there is a problem creating the file statuses.
    +	 */
    +	public FileStatus[] listStatus() throws IOException {
    +		return fileSystem.listStatus(directory);
    +	}
    +
    +	/**
    +	 * Calling this method completes the snapshot into the snapshot directory and creates a corresponding
    +	 * {@link DirectoryStateHandle} that points to the snapshot directory. Calling this method will also change the
    +	 * lifecycle state from "ongoing" to "completed". If the state was already deleted, an {@link IOException} is
    +	 * thrown.
    +	 *
    +	 * @return a directory state handle that points to the snapshot directory.
    +	 * @throws IOException if the state of this snapshot directory object is different from "ongoing".
    +	 */
    +	public DirectoryStateHandle completeSnapshotAndGetHandle() throws IOException {
    +		if (state.compareAndSet(State.ONGOING, State.COMPLETED)) {
    +			return new DirectoryStateHandle(directory, fileSystem);
    +		} else {
    +			throw new IOException("Expected state " + State.ONGOING + " but found state " + state.get());
    +		}
    +	}
    +
    +	/**
    +	 * Calling this method will attempt delete the underlying snapshot directory recursively, if the state was
    +	 * "ongoing". In this case, the state will be set to "deleted" as a result of this call.
    +	 *
    +	 * @return <code>true</code> if delete is successful, <code>false</code> otherwise.
    +	 * @throws IOException if an exception happens during the delete.
    +	 */
    +	public boolean deleteIfIncompleteSnapshot() throws IOException {
    +		return state.compareAndSet(State.ONGOING, State.DELETED) && fileSystem.delete(directory, true);
    +	}
    +
    +	public boolean isSnapshotOngoing() {
    +		return State.ONGOING.equals(state.get());
    +	}
    +
    +	@Override
    +	public boolean equals(Object o) {
    +		if (this == o) {
    +			return true;
    +		}
    +		if (o == null || getClass() != o.getClass()) {
    +			return false;
    +		}
    +
    +		SnapshotDirectory that = (SnapshotDirectory) o;
    +
    +		return directory.equals(that.directory);
    +	}
    +
    +	@Override
    +	public int hashCode() {
    +		return directory.hashCode();
    +	}
    +
    +	@Override
    +	public String toString() {
    +		return "SnapshotDirectory{" +
    +			"directory=" + directory +
    +			", state=" + state +
    +			'}';
    +	}
    +
    +	/**
    +	 * Helper method that creates a temporary snapshot directory for the given path. This subclass of snapshot directory
    +	 * will always return "null" as result of {@link #completeSnapshotAndGetHandle()} and can never reach the lifecycle
    +	 * state "completed".
    +	 */
    +	public static SnapshotDirectory temporary(Path path) throws IOException {
    +
    +		return new SnapshotDirectory(path) {
    --- End diff --
    
    Not sure whether an anonymous class is strictly necessary here.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159866220
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/BarrierBufferTest.java ---
    @@ -1483,11 +1482,6 @@ long getLastReportedBytesBufferedInAlignment() {
     			return lastReportedBytesBufferedInAlignment;
     		}
     
    -		@Override
    --- End diff --
    
    I think you could replace `ValidatingCheckpointHandler` with `CheckpointSequenceValidator` and remove this code duplication.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165344648
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java ---
    @@ -630,19 +506,210 @@ public int numStateEntries(Object namespace) {
     		return sum;
     	}
     
    -	public <N, V> StateTable<K, N, V> newStateTable(RegisteredKeyedBackendStateMetaInfo<N, V> newMetaInfo) {
    -		return asynchronousSnapshots ?
    -				new CopyOnWriteStateTable<>(this, newMetaInfo) :
    -				new NestedMapsStateTable<>(this, newMetaInfo);
    -	}
    -
     	@Override
     	public boolean supportsAsynchronousSnapshots() {
    -		return asynchronousSnapshots;
    +		return snapshotStrategy.isAsynchronous();
     	}
     
     	@VisibleForTesting
     	public FsStateBackend.LocalRecoveryConfig getLocalRecoveryConfig() {
     		return localRecoveryConfig;
     	}
    +
    +	/**
    +	 * Base class for the snapshots of the heap backend that outlines the algorithm and offers some hooks to realize
    +	 * the concrete strategies.
    +	 */
    +	private abstract class HeapSnapshotStrategy implements SnapshotStrategy<SnapshotResult<KeyedStateHandle>> {
    +
    +		@Override
    +		public RunnableFuture<SnapshotResult<KeyedStateHandle>> performSnapshot(
    +			long checkpointId,
    +			long timestamp,
    +			CheckpointStreamFactory streamFactory,
    +			CheckpointOptions checkpointOptions) throws Exception {
    +
    +			if (!hasRegisteredState()) {
    +				return DoneFuture.nullValue();
    +			}
    +
    +			long syncStartTime = System.currentTimeMillis();
    +
    +			Preconditions.checkState(stateTables.size() <= Short.MAX_VALUE,
    +				"Too many KV-States: " + stateTables.size() +
    +					". Currently at most " + Short.MAX_VALUE + " states are supported");
    +
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    +				new ArrayList<>(stateTables.size());
    +
    +			final Map<String, Integer> kVStateToId = new HashMap<>(stateTables.size());
    +
    +			final Map<StateTable<K, ?, ?>, StateTableSnapshot> cowStateStableSnapshots =
    +				new HashedMap(stateTables.size());
    +
    +			for (Map.Entry<String, StateTable<K, ?, ?>> kvState : stateTables.entrySet()) {
    +				kVStateToId.put(kvState.getKey(), kVStateToId.size());
    +				StateTable<K, ?, ?> stateTable = kvState.getValue();
    +				if (null != stateTable) {
    +					metaInfoSnapshots.add(stateTable.getMetaInfo().snapshot());
    +					cowStateStableSnapshots.put(stateTable, stateTable.createSnapshot());
    +				}
    +			}
    +
    +			final KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(
    +					keySerializer,
    +					metaInfoSnapshots,
    +					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, keyGroupCompressionDecorator));
    +
    +			//--------------------------------------------------- this becomes the end of sync part
    +
    +			// implementation of the async IO operation, based on FutureTask
    +			final AbstractAsyncCallableWithResources<SnapshotResult<KeyedStateHandle>> ioCallable =
    +				new AbstractAsyncCallableWithResources<SnapshotResult<KeyedStateHandle>>() {
    +
    +					CheckpointStreamFactory.CheckpointStateOutputStream stream = null;
    +
    +					@Override
    +					protected void acquireResources() throws Exception {
    +						stream = streamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
    +						cancelStreamRegistry.registerCloseable(stream);
    +					}
    +
    +					@Override
    +					protected void releaseResources() throws Exception {
    +
    +						if (cancelStreamRegistry.unregisterCloseable(stream)) {
    +							IOUtils.closeQuietly(stream);
    +							stream = null;
    +						}
    +
    +						for (StateTableSnapshot tableSnapshot : cowStateStableSnapshots.values()) {
    +							tableSnapshot.release();
    +						}
    +					}
    +
    +					@Override
    +					protected void stopOperation() throws Exception {
    +						if (cancelStreamRegistry.unregisterCloseable(stream)) {
    +							IOUtils.closeQuietly(stream);
    +							stream = null;
    +						}
    +					}
    +
    +					@Override
    +					protected SnapshotResult<KeyedStateHandle> performOperation() throws Exception {
    +
    +						long startTime = System.currentTimeMillis();
    +
    +						CheckpointStreamFactory.CheckpointStateOutputStream localStream = this.stream;
    +
    +						DataOutputViewStreamWrapper outView = new DataOutputViewStreamWrapper(localStream);
    +						serializationProxy.write(outView);
    +
    +						long[] keyGroupRangeOffsets = new long[keyGroupRange.getNumberOfKeyGroups()];
    +
    +						for (int keyGroupPos = 0; keyGroupPos < keyGroupRange.getNumberOfKeyGroups(); ++keyGroupPos) {
    +							int keyGroupId = keyGroupRange.getKeyGroupId(keyGroupPos);
    +							keyGroupRangeOffsets[keyGroupPos] = localStream.getPos();
    +							outView.writeInt(keyGroupId);
    +
    +							for (Map.Entry<String, StateTable<K, ?, ?>> kvState : stateTables.entrySet()) {
    +								OutputStream kgCompressionOut = keyGroupCompressionDecorator.decorateWithCompression(localStream);
    --- End diff --
    
    Can we put this in a try with resource statement?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168184825
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -363,1691 +372,1780 @@ public int getKeyGroupPrefixBytes() {
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointOptions.CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at " +
    -					checkpointTimestamp + " . Returning null.");
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    -
    -		snapshotOperation.takeSnapshot();
    -
    -		return new FutureTask<SnapshotResult<KeyedStateHandle>>(
    -			() -> snapshotOperation.materializeSnapshot()
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    +		try {
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    +				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +				restoreOperation.restore(restoreState);
    +			} else {
    +				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +				restoreOperation.doRestore(restoreState);
     			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at " + timestamp +
    -					" . Returning null.");
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<SnapshotResult<KeyedStateHandle>> ioCallable =
    -			new AbstractAsyncCallableWithResources<SnapshotResult<KeyedStateHandle>>() {
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -				@Override
    -				public SnapshotResult<KeyedStateHandle> performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in last position.
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					KeyGroupsStateHandle stateHandle = snapshotOperation.getSnapshotResultStateHandle();
    -					return new SnapshotResult<>(stateHandle, null);
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot (" + streamFactory + ", synchronous part) in thread " +
    -			Thread.currentThread() + " took " + (System.currentTimeMillis() - startTime) + " ms.");
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    -
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(checkpointId, checkpointTimeStamp);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
    --- End diff --
    
    In general, I think that makes sense by now, but I would do this in a different PR.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168745303
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStoreImpl.java ---
    @@ -0,0 +1,298 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
    +
    +import java.io.File;
    +import java.io.IOException;
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.HashMap;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
    +
    +/**
    + * Main implementation of a {@link TaskLocalStateStore}.
    + */
    +public class TaskLocalStateStoreImpl implements TaskLocalStateStore {
    +
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStoreImpl.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	@VisibleForTesting
    +	static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    +
    +	/** JobID from the owning subtask. */
    +	@Nonnull
    +	private final JobID jobID;
    +
    +	/** AllocationID of the owning slot. */
    +	@Nonnull
    +	private final AllocationID allocationID;
    +
    +	/** JobVertexID of the owning subtask. */
    +	@Nonnull
    +	private final JobVertexID jobVertexID;
    +
    +	/** Subtask index of the owning subtask. */
    +	@Nonnegative
    +	private final int subtaskIndex;
    +
    +	/** The configured mode for local recovery. */
    +	@Nonnull
    +	private final LocalRecoveryConfig localRecoveryConfig;
    +
    +	/** Executor that runs the discarding of released state objects. */
    +	@Nonnull
    +	private final Executor discardExecutor;
    +
    +	/** Lock for synchronisation on the storage map and the discarded status. */
    +	@Nonnull
    +	private final Object lock;
    +
    +	/** Status flag if this store was already discarded. */
    +	@GuardedBy("lock")
    +	private boolean discarded;
    +
    +	/** Maps checkpoint ids to local TaskStateSnapshots. */
    +	@Nonnull
    +	@GuardedBy("lock")
    +	private final SortedMap<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
    +	public TaskLocalStateStoreImpl(
    +		@Nonnull JobID jobID,
    +		@Nonnull AllocationID allocationID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex,
    +		@Nonnull LocalRecoveryConfig localRecoveryConfig,
    +		@Nonnull Executor discardExecutor) {
    +
    +		this.jobID = jobID;
    +		this.allocationID = allocationID;
    +		this.jobVertexID = jobVertexID;
    +		this.subtaskIndex = subtaskIndex;
    +		this.discardExecutor = discardExecutor;
    +		this.lock = new Object();
    +		this.storedTaskStateByCheckpointID = new TreeMap<>();
    +		this.discarded = false;
    +		this.localRecoveryConfig = localRecoveryConfig;
    +	}
    +
    +	@Override
    +	public void storeLocalState(
    +		@Nonnegative long checkpointId,
    +		@Nullable TaskStateSnapshot localState) {
    +
    +		if (localState == null) {
    +			localState = NULL_DUMMY;
    +		}
    +
    +		LOG.info("Storing local state for checkpoint {}.", checkpointId);
    +		LOG.debug("Local state for checkpoint {} is {}.", checkpointId, localState);
    +
    +		Map<Long, TaskStateSnapshot> toDiscard = new HashMap<>(MAX_RETAINED_SNAPSHOTS);
    +
    +		synchronized (lock) {
    +			if (discarded) {
    +				// we ignore late stores and simply discard the state.
    +				toDiscard.put(checkpointId, localState);
    +			} else {
    +				TaskStateSnapshot previous =
    +					storedTaskStateByCheckpointID.put(checkpointId, localState);
    +
    +				if (previous != null) {
    +					toDiscard.put(checkpointId, previous);
    +				}
    +
    +				// remove from history.
    +				while (storedTaskStateByCheckpointID.size() > MAX_RETAINED_SNAPSHOTS) {
    +					Long removeCheckpointID = storedTaskStateByCheckpointID.firstKey();
    --- End diff --
    
    I don't think that should matter too much, the data will already be deleted once a checkpoint completes.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168753503
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    +
    +- RocksDBStateBackend: task-local recovery is supported for keyed state. For *full checkpoints*, state is duplicated to a local file. This can introduce additional write costs
    +and occupies local disk space. For *incremental snapshots*, the local state is based on RocksDB's native checkpointing mechanism. This mechanism is also used as the first step
    +to create the primary copy, which means that in this case no additional cost is introduced for create the secondary copy. We simply keep the native checkpoint directory around
    +instead of deleting it after uploading to the distributed store. This local copy can share active files with the working directory of RocksDB (via hard links), so for active
    +files also no additional disk space is consumed for task-local recovery with incremental snapshots.
    +
    +### Allocation-preserving scheduling
    +
    +Task-local recovery assumes allocation-preserving task scheduling under failures, which was introduced as part of FLIP-6 and works as follows. Each task remembers its previous
    +allocation and *requests the exact same slot* to restart in recovery. If this slot is not available, the task will request a *new, fresh slot* from the resource manager. This way,
    +if a task manager is no longer available, a task that cannot return to its previous location *will not drive other recovering tasks out of their previous slots*. Our reasoning is
    +that the previous slot can only disappear when a task manager is no longer available, and in this case *some* task has to request a new slot anyways. With our scheduling strategy
    +we give the maximum number of tasks a chance recover from their local state and avoid the cascading effect of tasks stealing their previous slots from one another.
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169726507
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/DirectoryStateHandle.java ---
    @@ -0,0 +1,101 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.IOException;
    +
    +/**
    + * This state handle represents a directory. This class is, for example, used to represent the directory of RocksDB's
    + * native checkpoint directories for local recovery.
    + */
    +public class DirectoryStateHandle implements StateObject {
    +
    +	/** Serial version. */
    +	private static final long serialVersionUID = 1L;
    +
    +	/** The path that describes the directory. */
    +	@Nonnull
    +	private final Path directory;
    +
    +	/** The filesystem that contains the directory described by path. */
    +	@Nonnull
    +	private final FileSystem fileSystem;
    --- End diff --
    
    `FileSystem` is not serializable and should be removed from the `DirectoryStateHandle`. `directory` should give you all you need.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168168982
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TestTaskStateManager.java ---
    @@ -67,8 +69,11 @@ public TestTaskStateManager(
     		this.jobId = jobId;
     		this.executionAttemptID = executionAttemptID;
     		this.checkpointResponder = checkpointResponder;
    -		this.taskStateSnapshotsByCheckpointId = new HashMap<>();
    +		this.jobManagerTaskStateSnapshotsByCheckpointId = new HashMap<>();
    +		this.taskManagerTaskStateSnapshotsByCheckpointId = new HashMap<>();
     		this.reportedCheckpointId = -1L;
    +		this.subtaskLocalStateBaseDirectory =
    +			new File(System.getProperty("java.io.tmpdir"), "testLocalState_" + UUID.randomUUID());
    --- End diff --
    
    This is already changed in later commits.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168497476
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -55,25 +54,24 @@
     	/** */
     	private final Map<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
     
    -	/** This is the base directory for all local state of the subtask that owns this {@link TaskLocalStateStore}. */
    -	private final File subtaskLocalStateBaseDirectory;
    +	/** The root directories for all local state of this {@link TaskLocalStateStore}. */
    +	private final File[] localStateRootDirectories;
     
     	public TaskLocalStateStore(
     		JobID jobID,
     		JobVertexID jobVertexID,
     		int subtaskIndex,
    -		File localStateRootDir) {
    +		File[] localStateRootDirectories) {
     
     		this.jobID = jobID;
     		this.jobVertexID = jobVertexID;
     		this.subtaskIndex = subtaskIndex;
     		this.storedTaskStateByCheckpointID = new HashMap<>();
    -		this.subtaskLocalStateBaseDirectory =
    -			new File(localStateRootDir, createSubtaskPath(jobID, jobVertexID, subtaskIndex));
    +		this.localStateRootDirectories = localStateRootDirectories;
    --- End diff --
    
    Changed in later commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170034046
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -1629,9 +1630,9 @@ private void closeLocalRegistry() {
     
     		private boolean isWithLocalRecovery(
     			CheckpointOptions.CheckpointType checkpointType,
    -			RocksDBStateBackend.LocalRecoveryMode recoveryMode) {
    +			LocalRecoveryConfig.LocalRecoveryMode recoveryMode) {
     			// we use local recovery when it is activated and we are not taking a savepoint.
    -			return RocksDBStateBackend.LocalRecoveryMode.ENABLE_FILE_BASED.equals(recoveryMode)
    +			return LocalRecoveryConfig.LocalRecoveryMode.ENABLE_FILE_BASED.equals(recoveryMode)
    --- End diff --
    
    I prefer for enum comparisons `==`, but we had this discussion before ;-)


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168752518
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168778447
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java ---
    @@ -0,0 +1,265 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.checkpoint;
    +
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.Preconditions;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.util.ArrayList;
    +import java.util.Collections;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.function.BiFunction;
    +
    +/**
    + * This class is a wrapper over multiple alternative {@link OperatorSubtaskState} that are (partial) substitutes for
    + * each other and imposes a priority ordering over all alternatives for the different states which define an order in
    + * which the operator should attempt to restore the state from them. One OperatorSubtaskState is considered as the
    + * "ground truth" about which state should be represented. Alternatives may be complete or partial substitutes for
    + * the "ground truth" with a higher priority (if they had a lower alternative, they would not really be alternatives).
    + * Substitution is determined on a per-sub-state basis.
    + */
    +public class PrioritizedOperatorSubtaskState {
    +
    +	/** Singleton instance for an empty, non-restored operator state. */
    +	private static final PrioritizedOperatorSubtaskState EMPTY_NON_RESTORED_INSTANCE =
    +		new PrioritizedOperatorSubtaskState(new OperatorSubtaskState(), Collections.emptyList(), false);
    +
    +	/** List of prioritized snapshot alternatives for managed operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedManagedOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for raw operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedRawOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for managed keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedManagedKeyedState;
    +
    +	/** List of prioritized snapshot alternatives for raw keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedRawKeyedState;
    +
    +	/** Signal flag if this represents state for a restored operator. */
    +	private final boolean restored;
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority) {
    +		this(jobManagerState, alternativesByPriority, true);
    +	}
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority,
    +		boolean restored) {
    +
    +		Preconditions.checkNotNull(jobManagerState, "Job manager state is null.");
    +		int size = Preconditions.checkNotNull(alternativesByPriority, "Alternative states are null.").size();
    +
    +		this.restored = restored;
    +
    +		List<StateObjectCollection<OperatorStateHandle>> managedOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> managedKeyedAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<OperatorStateHandle>> rawOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> rawKeyedAlternatives = new ArrayList<>(size);
    +
    +		for (OperatorSubtaskState subtaskState : alternativesByPriority) {
    +
    +			if (subtaskState != null) {
    +				managedKeyedAlternatives.add(subtaskState.getManagedKeyedState());
    +				rawKeyedAlternatives.add(subtaskState.getRawKeyedState());
    +				managedOperatorAlternatives.add(subtaskState.getManagedOperatorState());
    +				rawOperatorAlternatives.add(subtaskState.getRawOperatorState());
    +			}
    +		}
    +
    +		// Key-groups should match.
    +		BiFunction<KeyedStateHandle, KeyedStateHandle, Boolean> keyedStateApprover =
    +			(ref, alt) -> ref.getKeyGroupRange().equals(alt.getKeyGroupRange());
    +
    +		// State meta data should match.
    +		BiFunction<OperatorStateHandle, OperatorStateHandle, Boolean> operatorStateApprover =
    +			(ref, alt) -> ref.getStateNameToPartitionOffsets().equals(alt.getStateNameToPartitionOffsets());
    +
    +		this.prioritizedManagedKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedKeyedState(),
    +			managedKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedRawKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawKeyedState(),
    +			rawKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedManagedOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedOperatorState(),
    +			managedOperatorAlternatives,
    +			operatorStateApprover);
    +
    +		this.prioritizedRawOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawOperatorState(),
    +			rawOperatorAlternatives,
    +			operatorStateApprover);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed operator state, in the order in which
    +	 * we should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedManagedOperatorState() {
    +		return prioritizedManagedOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw operator state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedRawOperatorState() {
    +		return prioritizedRawOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedManagedKeyedState() {
    +		return prioritizedManagedKeyedState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedRawKeyedState() {
    +		return prioritizedRawKeyedState.iterator();
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns the managed operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerManagedOperatorState() {
    +		return lastElement(prioritizedManagedOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the raw operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerRawOperatorState() {
    +		return lastElement(prioritizedRawOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the managed keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerManagedKeyedState() {
    +		return lastElement(prioritizedManagedKeyedState);
    +	}
    +
    +	/**
    +	 * Returns the raw keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerRawKeyedState() {
    +		return lastElement(prioritizedRawKeyedState);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns true if this was created for a restored operator, false otherwise. Restored operators are operators that
    +	 * participated in a previous checkpoint, even if they did not emit any state snapshots.
    +	 */
    +	public boolean isRestored() {
    +		return restored;
    +	}
    +
    +
    +	/**
    +	 * This helper method resolves the dependencies between the ground truth of the operator state obtained from the
    +	 * job manager and potential alternatives for recovery, e.g. from a task-local source.
    +	 */
    +	protected <T extends StateObject> List<StateObjectCollection<T>> resolvePrioritizedAlternatives(
    +		StateObjectCollection<T> jobManagerState,
    +		List<StateObjectCollection<T>> alternativesByPriority,
    +		BiFunction<T, T, Boolean> approveFun) {
    +
    +		// Nothing to resolve if there are no alternatives, or the ground truth has already no state, or if we can
    +		// assume that a rescaling happened because we find more than one handle in the JM state (this is more a sanity
    +		// check).
    +		if (alternativesByPriority == null
    +			|| alternativesByPriority.isEmpty()
    +			|| !jobManagerState.hasState()
    +			|| jobManagerState.size() != 1) {
    +
    --- End diff --
    
    Aha, got it. I wonder if someday flink will support online rescaling.. but it a another topic, ignore me...


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168712774
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/util/JvmExitOnFatalErrorTest.java ---
    @@ -137,9 +138,12 @@ public static void main(String[] args) throws Exception {
     
     			System.err.println("creating task");
     
    +			TemporaryFolder temporaryFolder = new TemporaryFolder();
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r167354956
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    --- End diff --
    
    recovery time**~s~**


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168753015
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    +
    +- RocksDBStateBackend: task-local recovery is supported for keyed state. For *full checkpoints*, state is duplicated to a local file. This can introduce additional write costs
    +and occupies local disk space. For *incremental snapshots*, the local state is based on RocksDB's native checkpointing mechanism. This mechanism is also used as the first step
    --- End diff --
    
    In the way we use the terms snapshot and checkpoint in Flink, it is not wrong: a snapshot is what one subtask produces from its state. A checkpoint is the collection of all task's snapshots.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168177813
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManagerTest.java ---
    @@ -0,0 +1,122 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.configuration.ConfigConstants;
    +import org.apache.flink.configuration.Configuration;
    +import org.apache.flink.runtime.clusterframework.types.ResourceID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +import org.apache.flink.runtime.taskexecutor.TaskManagerServices;
    +import org.apache.flink.runtime.taskexecutor.TaskManagerServicesConfiguration;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +import org.junit.rules.TemporaryFolder;
    +
    +import java.io.File;
    +import java.net.InetAddress;
    +
    +public class TaskExecutorLocalStateStoresManagerTest {
    +
    +	/**
    +	 * This tests that the creation of {@link TaskManagerServices} correctly creates the local state root directory
    +	 * for the {@link TaskExecutorLocalStateStoresManager} with the configured root directory.
    +	 */
    +	@Test
    +	public void testCreationFromConfig() throws Exception {
    +
    +		final Configuration config = new Configuration();
    +
    +		final String rootDirString = "localStateRoot";
    +		config.setString(ConfigConstants.TASK_MANAGER_LOCAL_STATE_ROOT_DIR_KEY, rootDirString);
    +
    +		final ResourceID tmResourceID = ResourceID.generate();
    +
    +		TaskManagerServicesConfiguration taskManagerServicesConfiguration =
    +			TaskManagerServicesConfiguration.fromConfiguration(config, InetAddress.getLocalHost(), true);
    +
    +		TaskManagerServices taskManagerServices =
    +			TaskManagerServices.fromConfiguration(taskManagerServicesConfiguration, tmResourceID);
    +
    +		TaskExecutorLocalStateStoresManager taskStateManager = taskManagerServices.getTaskStateManager();
    +
    +		Assert.assertEquals(
    +			new File(rootDirString, TaskManagerServices.LOCAL_STATE_SUB_DIRECTORY_ROOT),
    +			taskStateManager.getLocalStateRootDirectory());
    +
    +		Assert.assertEquals("localState", TaskManagerServices.LOCAL_STATE_SUB_DIRECTORY_ROOT);
    +	}
    +
    +	/**
    +	 * This tests that the creation of {@link TaskManagerServices} correctly falls back to the first tmp directory of
    +	 * the IOManager as default for the local state root directory.
    +	 */
    +	@Test
    +	public void testCreationFromConfigDefault() throws Exception {
    +
    +		final Configuration config = new Configuration();
    +
    +		final ResourceID tmResourceID = ResourceID.generate();
    +
    +		TaskManagerServicesConfiguration taskManagerServicesConfiguration =
    +			TaskManagerServicesConfiguration.fromConfiguration(config, InetAddress.getLocalHost(), true);
    +
    +		TaskManagerServices taskManagerServices =
    +			TaskManagerServices.fromConfiguration(taskManagerServicesConfiguration, tmResourceID);
    +
    +		TaskExecutorLocalStateStoresManager taskStateManager = taskManagerServices.getTaskStateManager();
    +
    +		Assert.assertEquals(
    +			new File(taskManagerServicesConfiguration.getTmpDirPaths()[0], TaskManagerServices.LOCAL_STATE_SUB_DIRECTORY_ROOT),
    +			taskStateManager.getLocalStateRootDirectory());
    +	}
    +
    +	/**
    +	 * This tests that the {@link TaskExecutorLocalStateStoresManager} creates {@link TaskLocalStateStore} that have
    +	 * a properly initialized local state base directory.
    +	 */
    +	@Test
    +	public void testSubtaskStateStoreDirectoryCreation() throws Exception {
    +
    +		JobID jobID = new JobID();
    +		JobVertexID jobVertexID = new JobVertexID();
    +		int subtaskIdx = 42;
    +		TemporaryFolder tmp = new TemporaryFolder();
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169725430
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamWithResultProvider.java ---
    @@ -172,4 +167,30 @@ public CheckpointStreamWithResultProvider create(
     			return new CheckpointStreamWithResultProvider.PrimaryStreamOnly(primaryOut);
     		}
     	}
    +
    +	/**
    +	 * Helper method that takes a {@link SnapshotResult<StreamStateHandle>} and a {@link KeyGroupRangeOffsets} and
    +	 * creates a {@link SnapshotResult<KeyGroupsStateHandle>} by combining the key groups offsets with all the
    +	 * present stream state handles.
    +	 */
    +	static SnapshotResult<KeyedStateHandle> toKeyedStateHandleSnapshotResult(
    +		@Nullable SnapshotResult<StreamStateHandle> snapshotResult,
    --- End diff --
    
    I think this parameter should not be `Nullable`. If it is null, then we don't have to call this method.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168752445
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168826365
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -266,11 +305,16 @@ public RocksDBKeyedStateBackend(
     		RocksIterator iterator = db.newIterator(columnInfo.f0);
     		iterator.seekToFirst();
     
    -		Iterable<K> iterable = () -> new RocksIteratorWrapper<>(iterator, state, keySerializer, keyGroupPrefixBytes);
    +		Iterable<K> iterable = () -> new RocksIteratorToJavaIteratorAdapter<>(iterator, state, keySerializer, keyGroupPrefixBytes);
     		Stream<K> targetStream = StreamSupport.stream(iterable.spliterator(), false);
    --- End diff --
    
    After tracing the invoked path, I found this seem to be a new function. And there is a  block bug I found that relate to it ... [5513](https://github.com/apache/flink/pull/5513)


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160948522
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TestTaskStateManager.java ---
    @@ -0,0 +1,185 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.core.testutils.OneShotLatch;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetrics;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +import org.apache.flink.runtime.taskmanager.CheckpointResponder;
    +
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +
    +import java.util.HashMap;
    +import java.util.Map;
    +
    +/**
    + * Implementation of {@link TaskStateManager} for tests.
    + */
    +public class TestTaskStateManager implements TaskStateManager {
    +
    +	private long reportedCheckpointId;
    +
    +	private JobID jobId;
    +	private ExecutionAttemptID executionAttemptID;
    +
    +	private final Map<Long, TaskStateSnapshot> taskStateSnapshotsByCheckpointId;
    +	private CheckpointResponder checkpointResponder;
    +	private OneShotLatch waitForReportLatch;
    +
    +	public TestTaskStateManager() {
    +		this(null, null, null);
    --- End diff --
    
    Does this constructor necessary have to have those nulls?  Can not those values be properly initialised?
    
    By the way. Isn't this the symptom of a my previous remark that `TaskStateManager` should be decomposed into two/three separate entities? 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162582618
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -831,11 +834,12 @@ private StreamStateHandle materializeStateData(Path filePath) throws Exception {
     				return result;
     
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +
    +				if (closeableRegistry.unregisterCloseable(inputStream)) {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168717057
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamTaskStateInitializerImpl.java ---
    @@ -223,155 +223,110 @@ public StreamOperatorStateContext streamOperatorStateContext(
     
     	protected OperatorStateBackend operatorStateBackend(
     		String operatorIdentifierText,
    -		OperatorSubtaskState operatorSubtaskStateFromJobManager,
    +		PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
     		CloseableRegistry backendCloseableRegistry) throws Exception {
     
    -		//TODO search in local state for a local recovery opportunity.
    +		BackendRestorerProcedure<OperatorStateBackend, OperatorStateHandle> backendRestorer =
    +			new BackendRestorerProcedure<>(
    +				() -> stateBackend.createOperatorStateBackend(environment, operatorIdentifierText),
    +				backendCloseableRegistry);
     
    -		return createOperatorStateBackendFromJobManagerState(
    -			operatorIdentifierText,
    -			operatorSubtaskStateFromJobManager,
    -			backendCloseableRegistry);
    +		return backendRestorer.createAndRestore(prioritizedOperatorSubtaskStates.getPrioritizedManagedOperatorState());
     	}
     
     	protected <K> AbstractKeyedStateBackend<K> keyedStatedBackend(
     		TypeSerializer<K> keySerializer,
     		String operatorIdentifierText,
    -		OperatorSubtaskState operatorSubtaskStateFromJobManager,
    +		PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
     		CloseableRegistry backendCloseableRegistry) throws Exception {
     
     		if (keySerializer == null) {
     			return null;
     		}
     
    -		//TODO search in local state for a local recovery opportunity.
    +		TaskInfo taskInfo = environment.getTaskInfo();
    +
    +		final KeyGroupRange keyGroupRange = KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(
    +			taskInfo.getMaxNumberOfParallelSubtasks(),
    +			taskInfo.getNumberOfParallelSubtasks(),
    +			taskInfo.getIndexOfThisSubtask());
     
    -		return createKeyedStatedBackendFromJobManagerState(
    -			keySerializer,
    -			operatorIdentifierText,
    -			operatorSubtaskStateFromJobManager,
    -			backendCloseableRegistry);
    +		BackendRestorerProcedure<AbstractKeyedStateBackend<K>, KeyedStateHandle> backendRestorer =
    +			new BackendRestorerProcedure<>(
    +				() -> stateBackend.createKeyedStateBackend(
    +					environment,
    +					environment.getJobID(),
    +					operatorIdentifierText,
    +					keySerializer,
    +					taskInfo.getMaxNumberOfParallelSubtasks(),
    +					keyGroupRange,
    +					environment.getTaskKvStateRegistry()),
    +				backendCloseableRegistry);
    +
    +		return backendRestorer.createAndRestore(prioritizedOperatorSubtaskStates.getPrioritizedManagedKeyedState());
     	}
     
     	protected CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs(
    -		OperatorSubtaskState operatorSubtaskStateFromJobManager) {
    +		Iterator<StateObjectCollection<OperatorStateHandle>> restoreStateAlternatives) {
     
    -		if (operatorSubtaskStateFromJobManager != null) {
    +		if (restoreStateAlternatives.hasNext()) {
     
     			final CloseableRegistry closeableRegistry = new CloseableRegistry();
     
    -			Collection<OperatorStateHandle> rawOperatorState =
    -				operatorSubtaskStateFromJobManager.getRawOperatorState();
    -
    -			return new CloseableIterable<StatePartitionStreamProvider>() {
    -				@Override
    -				public void close() throws IOException {
    -					closeableRegistry.close();
    -				}
    -
    -				@Nonnull
    -				@Override
    -				public Iterator<StatePartitionStreamProvider> iterator() {
    -					return new OperatorStateStreamIterator(
    -						DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME,
    -						rawOperatorState.iterator(), closeableRegistry);
    -				}
    -			};
    -		}
    -
    -		return CloseableIterable.empty();
    -	}
    -
    -	protected CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs(
    -		OperatorSubtaskState operatorSubtaskStateFromJobManager) {
    +			Collection<OperatorStateHandle> rawOperatorState = restoreStateAlternatives.next();
    +			// TODO currently this does not support local state recovery, so we expect there is only one handle.
    +			Preconditions.checkState(!restoreStateAlternatives.hasNext());
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165343607
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotStrategy.java ---
    @@ -0,0 +1,50 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.runtime.checkpoint.CheckpointOptions;
    +
    +import java.util.concurrent.RunnableFuture;
    +
    +/**
    + * Interface for different snapshot approaches in state backends. Implementing classes should ideally be stateless or at
    + * least threadsafe, i.e. this is a functional interface and is can be called in parallel by multiple checkpoints.
    + *
    + * @param <S> type of the returned state object that represents the result of the snapshot operation.
    + */
    +public interface SnapshotStrategy<S extends StateObject> {
    --- End diff --
    
    If this interface is functional, then we could also annotate it with `@FunctionalInterface`


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168486771
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java ---
    @@ -727,19 +729,17 @@ static void resetRocksDBLoadedFlag() throws Exception {
     	/**
     	 * This class encapsulates the configuration for local recovery of this backend.
     	 */
    -	public static final class LocalRecoveryConfig implements Serializable {
    +	public static final class LocalRecoveryConfig extends LocalRecoveryConfigBase {
     
    -		private static final long serialVersionUID = 1L;
    --- End diff --
    
    Not sure whether this is ok, tbh.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168752196
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168231423
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/OperatorStateBackendTest.java ---
    @@ -355,10 +355,11 @@ public void testSnapshotEmpty() throws Exception {
     		CheckpointStreamFactory streamFactory =
     				abstractStateBackend.createStreamFactory(new JobID(), "testOperator");
     
    -		RunnableFuture<OperatorStateHandle> snapshot =
    +		RunnableFuture<SnapshotResult<OperatorStateHandle>> snapshot =
     				operatorStateBackend.snapshot(0L, 0L, streamFactory, CheckpointOptions.forCheckpoint());
     
    -		OperatorStateHandle stateHandle = FutureUtil.runIfNotDoneAndGet(snapshot);
    +		SnapshotResult<OperatorStateHandle> snapshotResult = FutureUtil.runIfNotDoneAndGet(snapshot);
    +		OperatorStateHandle stateHandle = snapshotResult != null ? snapshotResult.getJobManagerOwnedSnapshot() : null;
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169899133
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotDirectory.java ---
    @@ -0,0 +1,172 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.FileStatus;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.IOException;
    +import java.util.concurrent.atomic.AtomicReference;
    +
    +/**
    + * This class represents a directory that is the target for a state snapshot. This class provides some method that
    + * simplify resource management when dealing with such directories, e.g. it can produce a {@link DirectoryStateHandle}
    + * when the snapshot is completed and disposal considers whether or not a snapshot was already completed. For a
    + * completed snapshot, the ownership for cleanup is transferred to the created directory state handle. For incomplete
    + * snapshots, calling {@link #deleteIfIncompleteSnapshot()} will delete the underlying directory resource.
    + */
    +public class SnapshotDirectory {
    +
    +	/**
    +	 * Lifecycle stages of a snapshot directory.
    +	 */
    +	enum State {
    +		ONGOING, COMPLETED, DELETED
    +	}
    +
    +	/** This path describes the underlying directory for the snapshot. */
    +	@Nonnull
    +	private final Path directory;
    +
    +	/** The filesystem that contains the snapshot directory. */
    +	@Nonnull
    +	private final FileSystem fileSystem;
    +
    +	/** This reference tracks the lifecycle state of the snapshot directory. */
    +	@Nonnull
    +	private AtomicReference<State> state;
    +
    +	public SnapshotDirectory(@Nonnull Path directory, @Nonnull FileSystem fileSystem) {
    +		this.directory = directory;
    +		this.fileSystem = fileSystem;
    +		this.state = new AtomicReference<>(State.ONGOING);
    +	}
    +
    +	public SnapshotDirectory(@Nonnull Path directory) throws IOException {
    +		this(directory, directory.getFileSystem());
    +	}
    +
    +	@Nonnull
    +	public Path getDirectory() {
    +		return directory;
    +	}
    +
    +	public boolean mkdirs() throws IOException {
    +		return fileSystem.mkdirs(directory);
    +	}
    +
    +	@Nonnull
    +	public FileSystem getFileSystem() {
    +		return fileSystem;
    +	}
    +
    +	public boolean exists() throws IOException {
    +		return fileSystem.exists(directory);
    +	}
    +
    +	/**
    +	 * List the statuses of the files/directories in the snapshot directory.
    +	 *
    +	 * @return the statuses of the files/directories in the given path.
    +	 * @throws IOException if there is a problem creating the file statuses.
    +	 */
    +	public FileStatus[] listStatus() throws IOException {
    +		return fileSystem.listStatus(directory);
    +	}
    +
    +	/**
    +	 * Calling this method completes the snapshot into the snapshot directory and creates a corresponding
    +	 * {@link DirectoryStateHandle} that points to the snapshot directory. Calling this method will also change the
    +	 * lifecycle state from "ongoing" to "completed". If the state was already deleted, an {@link IOException} is
    +	 * thrown.
    +	 *
    +	 * @return a directory state handle that points to the snapshot directory.
    +	 * @throws IOException if the state of this snapshot directory object is different from "ongoing".
    +	 */
    +	public DirectoryStateHandle completeSnapshotAndGetHandle() throws IOException {
    +		if (state.compareAndSet(State.ONGOING, State.COMPLETED)) {
    +			return new DirectoryStateHandle(directory, fileSystem);
    +		} else {
    +			throw new IOException("Expected state " + State.ONGOING + " but found state " + state.get());
    +		}
    +	}
    +
    +	/**
    +	 * Calling this method will attempt delete the underlying snapshot directory recursively, if the state was
    +	 * "ongoing". In this case, the state will be set to "deleted" as a result of this call.
    +	 *
    +	 * @return <code>true</code> if delete is successful, <code>false</code> otherwise.
    +	 * @throws IOException if an exception happens during the delete.
    +	 */
    +	public boolean deleteIfIncompleteSnapshot() throws IOException {
    +		return state.compareAndSet(State.ONGOING, State.DELETED) && fileSystem.delete(directory, true);
    +	}
    +
    +	public boolean isSnapshotOngoing() {
    +		return State.ONGOING.equals(state.get());
    --- End diff --
    
    I think this is only a question of personal style, there is no guideline that says this and it basically does the same if equals is called on the constant.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168495783
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalRecoveryDirectoryProvider.java ---
    @@ -0,0 +1,112 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.util.Preconditions;
    +
    +import java.io.File;
    +import java.io.Serializable;
    +import java.util.Arrays;
    +import java.util.UUID;
    +import java.util.concurrent.atomic.AtomicInteger;
    +
    +/**
    + * Provides root directories and the subtask-specific path to build directories for file-based local recovery. Calls
    + * to {@link #nextRootDirectory()} rotate over all available root directories. this class is thread-safe.
    + */
    +public class LocalRecoveryDirectoryProvider implements Serializable {
    +
    +	private static final long serialVersionUID = 1L;
    +
    +	/** All available root directories that this can potentially deliver. */
    +	private final File[] rootDirectories;
    +
    +	/** A subtask-specific string that describes a path to it's local state from a given checkpoint root. */
    +	private final String subtaskSpecificPath;
    +
    +	/** Current index for rotational selection of delivered root dirs. */
    +	private AtomicInteger idx;
    +
    +	public LocalRecoveryDirectoryProvider(File rootDir, String subtaskSpecificPath) {
    +		this(new File[]{rootDir}, subtaskSpecificPath);
    +	}
    +
    +	public LocalRecoveryDirectoryProvider(File[] rootDirectories, String subtaskSpecificPath) {
    +
    +		this.rootDirectories = Preconditions.checkNotNull(rootDirectories);
    +		this.subtaskSpecificPath = Preconditions.checkNotNull(subtaskSpecificPath);
    +		Preconditions.checkArgument(rootDirectories.length > 0);
    +
    +		for (File directory : rootDirectories) {
    +			Preconditions.checkNotNull(directory);
    +			if (!directory.isDirectory()) {
    +				throw new IllegalStateException("Local recovery root directory " + directory + " does not exist!");
    +			}
    +		}
    +
    +		this.idx = new AtomicInteger(-1);
    +	}
    +
    +	/**
    +	 * Returns the next root directory w.r.t. our rotation over all available root dirs.
    +	 */
    +	public File nextRootDirectory() {
    +		return rootDirectories[(idx.incrementAndGet() & Integer.MAX_VALUE) % rootDirectories.length];
    --- End diff --
    
    :D


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168773128
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -266,11 +305,16 @@ public RocksDBKeyedStateBackend(
     		RocksIterator iterator = db.newIterator(columnInfo.f0);
     		iterator.seekToFirst();
     
    -		Iterable<K> iterable = () -> new RocksIteratorWrapper<>(iterator, state, keySerializer, keyGroupPrefixBytes);
    +		Iterable<K> iterable = () -> new RocksIteratorToJavaIteratorAdapter<>(iterator, state, keySerializer, keyGroupPrefixBytes);
     		Stream<K> targetStream = StreamSupport.stream(iterable.spliterator(), false);
    --- End diff --
    
    Aha, I would like to open a issue with a PR for this if you don't mind?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168514168
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -166,20 +190,38 @@ public TaskStateSnapshot retrieveLocalState(long checkpointID) {
     	 * Disposes the state of all local snapshots managed by this object.
     	 */
     	public void dispose() {
    +
    +		Collection<Map.Entry<Long, TaskStateSnapshot>> statesCopy;
    +
     		synchronized (lock) {
    -			for (Map.Entry<Long, TaskStateSnapshot> entry : storedTaskStateByCheckpointID.entrySet()) {
    -				discardStateObject(entry.getValue(), entry.getKey());
    -			}
     			discarded = true;
    +			statesCopy = new ArrayList<>(storedTaskStateByCheckpointID.entrySet());
     		}
    +
    +		discardExecutor.execute(() -> {
    --- End diff --
    
    I think it would be good for these kind of asynchronous clean up operations to return a `CompletableFuture`. This future could then be returned by the `dispose` method. The benefit would be that the caller would know when the clean up has completed and, thus, would be safe to shut down the `discardExecutor`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165101966
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -1016,7 +1017,7 @@ void releaseResources(boolean canceled) {
     	}
     
     	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    --- End diff --
    
    Why is this necessary?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168753623
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    +
    +- RocksDBStateBackend: task-local recovery is supported for keyed state. For *full checkpoints*, state is duplicated to a local file. This can introduce additional write costs
    +and occupies local disk space. For *incremental snapshots*, the local state is based on RocksDB's native checkpointing mechanism. This mechanism is also used as the first step
    +to create the primary copy, which means that in this case no additional cost is introduced for create the secondary copy. We simply keep the native checkpoint directory around
    +instead of deleting it after uploading to the distributed store. This local copy can share active files with the working directory of RocksDB (via hard links), so for active
    +files also no additional disk space is consumed for task-local recovery with incremental snapshots.
    +
    +### Allocation-preserving scheduling
    +
    +Task-local recovery assumes allocation-preserving task scheduling under failures, which was introduced as part of FLIP-6 and works as follows. Each task remembers its previous
    +allocation and *requests the exact same slot* to restart in recovery. If this slot is not available, the task will request a *new, fresh slot* from the resource manager. This way,
    +if a task manager is no longer available, a task that cannot return to its previous location *will not drive other recovering tasks out of their previous slots*. Our reasoning is
    +that the previous slot can only disappear when a task manager is no longer available, and in this case *some* task has to request a new slot anyways. With our scheduling strategy
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168524601
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamWithResultProvider.java ---
    @@ -0,0 +1,175 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.state.filesystem.FixFileFsStateOutputStream;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import java.io.Closeable;
    +import java.io.File;
    +import java.io.IOException;
    +
    +/**
    + * Interface that provides access to a CheckpointStateOutputStream and a method to provide the {@link SnapshotResult}.
    + * This abstracts from different ways that a result is obtained from checkpoint output streams.
    + */
    +public interface CheckpointStreamWithResultProvider extends Closeable {
    +
    +	SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +		KeyGroupRangeOffsets kgOffs) throws IOException;
    +
    +	CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream();
    +
    +	@Override
    +	default void close() throws IOException {
    +		CheckpointStreamFactory.CheckpointStateOutputStream outputStream = getCheckpointOutputStream();
    +		if (outputStream != null) {
    +			outputStream.close();
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that only creates the
    +	 * primary/remote/jm-owned state.
    +	 */
    +	class PrimaryStreamOnly implements CheckpointStreamWithResultProvider {
    +
    +		private final CheckpointStreamFactory.CheckpointStateOutputStream outputStream;
    +
    +		public PrimaryStreamOnly(CheckpointStreamFactory.CheckpointStateOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets kgOffs) throws IOException {
    +
    +			StreamStateHandle streamStateHandle = outputStream.closeAndGetHandle();
    +			KeyGroupsStateHandle primaryStateHandle = new KeyGroupsStateHandle(kgOffs, streamStateHandle);
    +			return new SnapshotResult<>(primaryStateHandle, null);
    +		}
    +
    +		@Override
    +		public CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that creates both, the
    +	 * primary/remote/jm-owned state and the secondary/local/tm-owned state.
    +	 */
    +	class PrimaryAndSecondaryStream implements CheckpointStreamWithResultProvider {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		private final DuplicatingCheckpointOutputStream outputStream;
    +
    +		public PrimaryAndSecondaryStream(DuplicatingCheckpointOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets offsets) throws IOException {
    +
    +			final StreamStateHandle primaryStreamStateHandle;
    +
    +			try {
    +				primaryStreamStateHandle = outputStream.closeAndGetPrimaryHandle();
    +			} catch (IOException primaryEx) {
    +				try {
    +					outputStream.close();
    +				} catch (IOException closeEx) {
    +					primaryEx = ExceptionUtils.firstOrSuppressed(closeEx, primaryEx);
    +				}
    +				throw primaryEx;
    +			}
    +
    +			StreamStateHandle secondaryStreamStateHandle = null;
    +
    +			try {
    +				secondaryStreamStateHandle = outputStream.closeAndGetSecondaryHandle();
    +			} catch (IOException secondaryEx) {
    +				LOG.warn("Exception from secondary/local checkpoint stream.", secondaryEx);
    +			}
    +
    +			if (primaryStreamStateHandle != null) {
    +
    +				final KeyGroupsStateHandle primaryKeyGroupsStateHandle =
    +					new KeyGroupsStateHandle(offsets, primaryStreamStateHandle);
    +
    +				final KeyGroupsStateHandle secondaryKeyGroupsStateHandle = secondaryStreamStateHandle != null ?
    +					new KeyGroupsStateHandle(offsets, secondaryStreamStateHandle) :
    +					null;
    +
    +				return new SnapshotResult<>(primaryKeyGroupsStateHandle, secondaryKeyGroupsStateHandle);
    +			}
    +
    +			return null;
    +		}
    +
    +		@Override
    +		public DuplicatingCheckpointOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	class Factory {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		public CheckpointStreamWithResultProvider create(
    +			long checkpointId,
    +			long timestamp,
    +			CheckpointStreamFactory primaryStreamFactory,
    +			LocalRecoveryDirectoryProvider secondaryStreamDirProvider) throws Exception {
    +
    +			CheckpointStreamFactory.CheckpointStateOutputStream primaryOut =
    +				primaryStreamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
    +
    +			CheckpointStreamFactory.CheckpointStateOutputStream secondaryOut;
    +
    +			if (secondaryStreamDirProvider != null) {
    +				try {
    +					File currentLocalRecoveryRoot = secondaryStreamDirProvider.nextRootDirectory();
    +					File outFile = new File(
    +						currentLocalRecoveryRoot,
    +						secondaryStreamDirProvider.specificFileForCheckpointId(checkpointId));
    +					Path outPath = new Path(outFile.toURI());
    +					secondaryOut = new FixFileFsStateOutputStream(outPath.getFileSystem(), outPath);
    +					final DuplicatingCheckpointOutputStream effectiveOutStream =
    +						new DuplicatingCheckpointOutputStream(primaryOut, secondaryOut);
    +
    +					return new CheckpointStreamWithResultProvider.PrimaryAndSecondaryStream(effectiveOutStream);
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168514913
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -19,92 +19,224 @@
     package org.apache.flink.runtime.state;
     
     import org.apache.flink.api.common.JobID;
    -import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
     import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
     import org.apache.flink.runtime.jobgraph.JobVertexID;
    -import org.apache.flink.util.ExceptionUtils;
     
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
     import javax.annotation.Nonnull;
     import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
     
     import java.io.File;
    -import java.util.HashMap;
    +import java.util.Arrays;
    +import java.util.Iterator;
     import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
     
     /**
      * This class will service as a task-manager-level local storage for local checkpointed state. The purpose is to provide
      * access to a state that is stored locally for a faster recovery compared to the state that is stored remotely in a
      * stable store DFS. For now, this storage is only complementary to the stable storage and local state is typically
      * lost in case of machine failures. In such cases (and others), client code of this class must fall back to using the
      * slower but highly available store.
    - *
    - * TODO this is currently a placeholder / mock that still must be implemented!
      */
     public class TaskLocalStateStore {
     
    -	/** */
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStore.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	private static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162582925
  
    --- Diff: flink-core/src/main/java/org/apache/flink/util/CloseableIterator.java ---
    @@ -18,11 +18,38 @@
     
     package org.apache.flink.util;
     
    -import org.apache.flink.annotation.Internal;
    +import java.io.Closeable;
    +import java.io.IOException;
    +import java.util.Iterator;
     
     /**
    - * Tagging interface for migration related classes.
    + * This interface represents an iterator that is also closeable.
    + *
    + * @param <T> type of the iterated objects.
      */
    -@Internal
    -public interface Migration {
    +public interface CloseableIterator<T> extends Iterator<T>, Closeable {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168773814
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -266,11 +305,16 @@ public RocksDBKeyedStateBackend(
     		RocksIterator iterator = db.newIterator(columnInfo.f0);
     		iterator.seekToFirst();
     
    -		Iterable<K> iterable = () -> new RocksIteratorWrapper<>(iterator, state, keySerializer, keyGroupPrefixBytes);
    +		Iterable<K> iterable = () -> new RocksIteratorToJavaIteratorAdapter<>(iterator, state, keySerializer, keyGroupPrefixBytes);
     		Stream<K> targetStream = StreamSupport.stream(iterable.spliterator(), false);
    --- End diff --
    
    Fine with me, but since this seems to have not caused any trouble so far, you might first want to check if the namespace parameter is actually required. Maybe it can just be removed as parameter because the using code wants to iterate all keys in all namespace.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168168705
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/LocalStateForwardingTest.java ---
    @@ -0,0 +1,213 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.runtime.tasks;
    +
    +import org.apache.flink.api.common.ExecutionConfig;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.configuration.Configuration;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetrics;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.StateObjectCollection;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
    +import org.apache.flink.runtime.state.DoneFuture;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.SnapshotResult;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.runtime.state.TaskLocalStateStore;
    +import org.apache.flink.runtime.state.TaskStateManagerImpl;
    +import org.apache.flink.runtime.state.TestTaskStateManager;
    +import org.apache.flink.runtime.taskmanager.TestCheckpointResponder;
    +import org.apache.flink.streaming.api.operators.OperatorSnapshotFutures;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +
    +import java.util.HashMap;
    +import java.util.Map;
    +import java.util.concurrent.Future;
    +import java.util.concurrent.RunnableFuture;
    +import java.util.concurrent.atomic.AtomicBoolean;
    +
    +import static org.mockito.Mockito.mock;
    +
    +/**
    + * Test for forwarding of state reporting to and from {@link org.apache.flink.runtime.state.TaskStateManager}.
    + */
    +public class LocalStateForwardingTest {
    +
    +	/**
    +	 * This tests the forwarding of jm and tm-local state from the futures reported by the backends, through the
    +	 * async checkpointing thread to the {@link org.apache.flink.runtime.state.TaskStateManager}.
    +	 */
    +	@Test
    +	public void testForwardingFromSnapshotToTaskStateManager() throws Exception {
    +
    +		TestTaskStateManager taskStateManager = new TestTaskStateManager();
    +
    +		StreamMockEnvironment streamMockEnvironment = new StreamMockEnvironment(
    +			new Configuration(),
    +			new Configuration(),
    +			new ExecutionConfig(),
    +			1024*1024,
    +			new MockInputSplitProvider(),
    +			0,
    +			taskStateManager);
    +
    +		StreamTask testStreamTask = new StreamTaskTest.NoOpStreamTask(streamMockEnvironment);
    +		CheckpointMetaData checkpointMetaData = new CheckpointMetaData(0L, 0L);
    +		CheckpointMetrics checkpointMetrics = new CheckpointMetrics();
    +
    +		Map<OperatorID, OperatorSnapshotFutures> snapshots = new HashMap<>(1);
    +		OperatorSnapshotFutures osFuture = new OperatorSnapshotFutures();
    +
    +		osFuture.setKeyedStateManagedFuture(createSnapshotResult(KeyedStateHandle.class));
    +		osFuture.setKeyedStateRawFuture(createSnapshotResult(KeyedStateHandle.class));
    +		osFuture.setOperatorStateManagedFuture(createSnapshotResult(OperatorStateHandle.class));
    +		osFuture.setOperatorStateRawFuture(createSnapshotResult(OperatorStateHandle.class));
    +
    +		OperatorID operatorID = new OperatorID();
    +		snapshots.put(operatorID, osFuture);
    +
    +		StreamTask.AsyncCheckpointRunnable checkpointRunnable =
    +			new StreamTask.AsyncCheckpointRunnable(
    +				testStreamTask,
    +				snapshots,
    +				checkpointMetaData,
    +				checkpointMetrics,
    +				0L);
    +
    +		checkpointRunnable.run();
    +
    +		TaskStateSnapshot lastJobManagerTaskStateSnapshot = taskStateManager.getLastJobManagerTaskStateSnapshot();
    +		TaskStateSnapshot lastTaskManagerTaskStateSnapshot = taskStateManager.getLastTaskManagerTaskStateSnapshot();
    +
    +		OperatorSubtaskState jmState =
    +			lastJobManagerTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
    +
    +		OperatorSubtaskState tmState =
    +			lastTaskManagerTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
    +
    +		performCheck(osFuture.getKeyedStateManagedFuture(), jmState.getManagedKeyedState(), tmState.getManagedKeyedState());
    +		performCheck(osFuture.getKeyedStateRawFuture(), jmState.getRawKeyedState(), tmState.getRawKeyedState());
    +		performCheck(osFuture.getOperatorStateManagedFuture(), jmState.getManagedOperatorState(), tmState.getManagedOperatorState());
    +		performCheck(osFuture.getOperatorStateRawFuture(), jmState.getRawOperatorState(), tmState.getRawOperatorState());
    +	}
    +
    +	/**
    +	 * This tests that state that was reported to the {@link org.apache.flink.runtime.state.TaskStateManager} is also
    +	 * reported to {@link org.apache.flink.runtime.taskmanager.CheckpointResponder} and {@link TaskLocalStateStore}.
    +	 */
    +	@Test
    +	public void testForwardingFromTaskStateManagerToResponderAndTaskLocalStateStore() {
    +
    +		final JobID jobID = new JobID();
    +		final ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
    +		final CheckpointMetaData checkpointMetaData = new CheckpointMetaData(42L, 4711L);
    +		final CheckpointMetrics checkpointMetrics = new CheckpointMetrics();
    +		final int subtaskIdx = 42;
    +		JobVertexID jobVertexID = new JobVertexID();
    +
    +		TaskStateSnapshot jm = new TaskStateSnapshot();
    +		TaskStateSnapshot tm = new TaskStateSnapshot();
    +
    +		final AtomicBoolean jmReported = new AtomicBoolean(false);
    +		final AtomicBoolean tmReported = new AtomicBoolean(false);
    +
    +		TestCheckpointResponder checkpointResponder = new TestCheckpointResponder() {
    +
    +			@Override
    +			public void acknowledgeCheckpoint(
    +				JobID lJobID,
    +				ExecutionAttemptID lExecutionAttemptID,
    +				long lCheckpointId,
    +				CheckpointMetrics lCheckpointMetrics,
    +				TaskStateSnapshot lSubtaskState) {
    +
    +				Assert.assertEquals(jobID, lJobID);
    +				Assert.assertEquals(executionAttemptID, lExecutionAttemptID);
    +				Assert.assertEquals(checkpointMetaData.getCheckpointId(), lCheckpointId);
    +				Assert.assertEquals(checkpointMetrics, lCheckpointMetrics);
    +				jmReported.set(true);
    +			}
    +		};
    +
    +		TaskLocalStateStore taskLocalStateStore = new TaskLocalStateStore(jobID, jobVertexID, subtaskIdx) {
    +			@Override
    +			public void storeLocalState(
    +				@Nonnull CheckpointMetaData checkpointMetaData,
    +				@Nullable TaskStateSnapshot localState) {
    +
    +				Assert.assertEquals(tm, localState);
    +				tmReported.set(true);
    +			}
    +		};
    +
    +		TaskStateManagerImpl taskStateManager =
    +			new TaskStateManagerImpl(
    +				jobID,
    +				executionAttemptID,
    +				taskLocalStateStore,
    +				null,
    +				checkpointResponder);
    +
    +		taskStateManager.reportTaskStateSnapshots(
    +			checkpointMetaData,
    +			checkpointMetrics,
    +			jm,
    +			tm);
    +
    +		Assert.assertTrue("Reporting for JM state was not called.", jmReported.get());
    +		Assert.assertTrue("Reporting for TM state was not called.", tmReported.get());
    +	}
    +
    +	private static <T extends StateObject> void performCheck(
    +		Future<SnapshotResult<T>> cc,
    +		StateObjectCollection<T> jm,
    +		StateObjectCollection<T> tm) {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169638326
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
    --- End diff --
    
    This can be optimized. If a `restoreInstance` satisfy `startKeyGroup >= backend.keygroup.startKeyGroup() && endKeyGroup <= backend.keygroup.endKeyGroup()` we can open it directly as the target rocksdb, and use it for the rest of the recovery. (for now, we have to iterate all restored rocksdbs.)


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165112749
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/StateSnapshotContextSynchronousImpl.java ---
    @@ -130,7 +136,7 @@ public OperatorStateCheckpointOutputStream getRawOperatorStateOutput() throws Ex
     	}
     
     	private <T extends StreamStateHandle> void closeAndUnregisterStream(
    -		NonClosingCheckpointOutputStream<T> stream) throws IOException {
    +		NonClosingCheckpointOutputStream<? extends T> stream) throws IOException {
    --- End diff --
    
    Should be sufficient to have `NonClosingCheckpointOutputStream<? extends StreamStateHandle>`


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165332129
  
    --- Diff: flink-core/src/main/java/org/apache/flink/configuration/ConfigConstants.java ---
    @@ -316,6 +316,11 @@
     	 */
     	public static final String TASK_MANAGER_REFUSED_REGISTRATION_PAUSE = "taskmanager.refused-registration-pause";
     
    +	/**
    +	 * The config parameter defining the root directories for storing file-based state for local recovery.
    +	 */
    +	public static final String TASK_MANAGER_LOCAL_STATE_ROOT_DIR_KEY = "taskmanager.local-state-root.dir";
    --- End diff --
    
    Using `ConfigConstants` is discouraged. Please introduce a `ConfigOption` if not already done in a later commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r167355160
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    --- End diff --
    
    'we' refers to 'each task'? Better to be explicit about it


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159883287
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -831,11 +834,12 @@ private StreamStateHandle materializeStateData(Path filePath) throws Exception {
     				return result;
     
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +
    +				if (closeableRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
     
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    +				if (closeableRegistry.unregisterCloseable(outputStream)) {
    --- End diff --
    
    Yes, that is what it is used for. In case of cancelation, we want all blocking ops to terminate asap using this mechanism.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159877880
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointListener.java ---
    @@ -18,11 +18,14 @@
     package org.apache.flink.runtime.state;
     
     
    +import org.apache.flink.annotation.PublicEvolving;
    +
     /**
      * This interface must be implemented by functions/operations that want to receive
      * a commit notification once a checkpoint has been completely acknowledged by all
      * participants.
      */
    +@PublicEvolving
    --- End diff --
    
    Separate commit?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168482624
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/LocalRecoveryDirectoryProviderTest.java ---
    @@ -0,0 +1,95 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.junit.After;
    +import org.junit.Assert;
    +import org.junit.Before;
    +import org.junit.Test;
    +import org.junit.rules.TemporaryFolder;
    +
    +import java.io.File;
    +import java.io.IOException;
    +
    +public class LocalRecoveryDirectoryProviderTest {
    +
    +	private static final String SPECIFIC_PATH = "specificPath";
    +
    +	private TemporaryFolder tmpFolder;
    --- End diff --
    
    Adding `@ClassRule` would make deletion unnecessary.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168717401
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskStateManagerImplTest.java ---
    @@ -148,22 +150,22 @@ public void testForwardingSubtaskLocalStateBaseDirFromLocalStateStore() throws I
     				taskLocalStateStore);
     
     			LocalRecoveryDirectoryProvider directoryProviderFromTaskLocalStateStore =
    -				taskLocalStateStore.createLocalRecoveryRootDirectoryProvider();
    +				taskLocalStateStore.getLocalRecoveryRootDirectoryProvider();
     
     			LocalRecoveryDirectoryProvider directoryProviderFromTaskStateManager =
     				taskStateManager.createLocalRecoveryRootDirectoryProvider();
     
     
     			for (int i = 0; i < 10; ++i) {
     				Assert.assertEquals(rootDirs[i % rootDirs.length],
    -					directoryProviderFromTaskLocalStateStore.nextRootDirectory());
    +					directoryProviderFromTaskLocalStateStore.rootDirectory(i));
     				Assert.assertEquals(rootDirs[i % rootDirs.length],
    -					directoryProviderFromTaskStateManager.nextRootDirectory());
    +					directoryProviderFromTaskStateManager.rootDirectory(i));
     			}
     
    -			Assert.assertEquals(
    -				directoryProviderFromTaskLocalStateStore.getSubtaskSpecificPath(),
    -				directoryProviderFromTaskStateManager.getSubtaskSpecificPath());
    +//			Assert.assertEquals(
    +//				directoryProviderFromTaskLocalStateStore.getSubtaskSpecificPath(),
    +//				directoryProviderFromTaskStateManager.getSubtaskSpecificPath());
    --- End diff --
    
    removed in later commit


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168775173
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java ---
    @@ -0,0 +1,265 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.checkpoint;
    +
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.Preconditions;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.util.ArrayList;
    +import java.util.Collections;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.function.BiFunction;
    +
    +/**
    + * This class is a wrapper over multiple alternative {@link OperatorSubtaskState} that are (partial) substitutes for
    + * each other and imposes a priority ordering over all alternatives for the different states which define an order in
    + * which the operator should attempt to restore the state from them. One OperatorSubtaskState is considered as the
    + * "ground truth" about which state should be represented. Alternatives may be complete or partial substitutes for
    + * the "ground truth" with a higher priority (if they had a lower alternative, they would not really be alternatives).
    + * Substitution is determined on a per-sub-state basis.
    + */
    +public class PrioritizedOperatorSubtaskState {
    +
    +	/** Singleton instance for an empty, non-restored operator state. */
    +	private static final PrioritizedOperatorSubtaskState EMPTY_NON_RESTORED_INSTANCE =
    +		new PrioritizedOperatorSubtaskState(new OperatorSubtaskState(), Collections.emptyList(), false);
    +
    +	/** List of prioritized snapshot alternatives for managed operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedManagedOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for raw operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedRawOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for managed keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedManagedKeyedState;
    +
    +	/** List of prioritized snapshot alternatives for raw keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedRawKeyedState;
    +
    +	/** Signal flag if this represents state for a restored operator. */
    +	private final boolean restored;
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority) {
    +		this(jobManagerState, alternativesByPriority, true);
    +	}
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority,
    +		boolean restored) {
    +
    +		Preconditions.checkNotNull(jobManagerState, "Job manager state is null.");
    +		int size = Preconditions.checkNotNull(alternativesByPriority, "Alternative states are null.").size();
    +
    +		this.restored = restored;
    +
    +		List<StateObjectCollection<OperatorStateHandle>> managedOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> managedKeyedAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<OperatorStateHandle>> rawOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> rawKeyedAlternatives = new ArrayList<>(size);
    +
    +		for (OperatorSubtaskState subtaskState : alternativesByPriority) {
    +
    +			if (subtaskState != null) {
    +				managedKeyedAlternatives.add(subtaskState.getManagedKeyedState());
    +				rawKeyedAlternatives.add(subtaskState.getRawKeyedState());
    +				managedOperatorAlternatives.add(subtaskState.getManagedOperatorState());
    +				rawOperatorAlternatives.add(subtaskState.getRawOperatorState());
    +			}
    +		}
    +
    +		// Key-groups should match.
    +		BiFunction<KeyedStateHandle, KeyedStateHandle, Boolean> keyedStateApprover =
    +			(ref, alt) -> ref.getKeyGroupRange().equals(alt.getKeyGroupRange());
    +
    +		// State meta data should match.
    +		BiFunction<OperatorStateHandle, OperatorStateHandle, Boolean> operatorStateApprover =
    +			(ref, alt) -> ref.getStateNameToPartitionOffsets().equals(alt.getStateNameToPartitionOffsets());
    +
    +		this.prioritizedManagedKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedKeyedState(),
    +			managedKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedRawKeyedState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawKeyedState(),
    +			rawKeyedAlternatives,
    +			keyedStateApprover);
    +
    +		this.prioritizedManagedOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getManagedOperatorState(),
    +			managedOperatorAlternatives,
    +			operatorStateApprover);
    +
    +		this.prioritizedRawOperatorState = resolvePrioritizedAlternatives(
    +			jobManagerState.getRawOperatorState(),
    +			rawOperatorAlternatives,
    +			operatorStateApprover);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed operator state, in the order in which
    +	 * we should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedManagedOperatorState() {
    +		return prioritizedManagedOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw operator state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<OperatorStateHandle>> getPrioritizedRawOperatorState() {
    +		return prioritizedRawOperatorState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the managed keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedManagedKeyedState() {
    +		return prioritizedManagedKeyedState.iterator();
    +	}
    +
    +	/**
    +	 * Returns an iterator over all alternative snapshots to restore the raw keyed state, in the order in which we
    +	 * should attempt to restore.
    +	 */
    +	@Nonnull
    +	public Iterator<StateObjectCollection<KeyedStateHandle>> getPrioritizedRawKeyedState() {
    +		return prioritizedRawKeyedState.iterator();
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns the managed operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerManagedOperatorState() {
    +		return lastElement(prioritizedManagedOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the raw operator state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<OperatorStateHandle> getJobManagerRawOperatorState() {
    +		return lastElement(prioritizedRawOperatorState);
    +	}
    +
    +	/**
    +	 * Returns the managed keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerManagedKeyedState() {
    +		return lastElement(prioritizedManagedKeyedState);
    +	}
    +
    +	/**
    +	 * Returns the raw keyed state from the job manager, which represents the ground truth about what this state
    +	 * should represent. This is the alternative with lowest priority.
    +	 */
    +	@Nonnull
    +	public StateObjectCollection<KeyedStateHandle> getJobManagerRawKeyedState() {
    +		return lastElement(prioritizedRawKeyedState);
    +	}
    +
    +	// -----------------------------------------------------------------------------------------------------------------
    +
    +	/**
    +	 * Returns true if this was created for a restored operator, false otherwise. Restored operators are operators that
    +	 * participated in a previous checkpoint, even if they did not emit any state snapshots.
    +	 */
    +	public boolean isRestored() {
    +		return restored;
    +	}
    +
    +
    +	/**
    +	 * This helper method resolves the dependencies between the ground truth of the operator state obtained from the
    +	 * job manager and potential alternatives for recovery, e.g. from a task-local source.
    +	 */
    +	protected <T extends StateObject> List<StateObjectCollection<T>> resolvePrioritizedAlternatives(
    +		StateObjectCollection<T> jobManagerState,
    +		List<StateObjectCollection<T>> alternativesByPriority,
    +		BiFunction<T, T, Boolean> approveFun) {
    +
    +		// Nothing to resolve if there are no alternatives, or the ground truth has already no state, or if we can
    +		// assume that a rescaling happened because we find more than one handle in the JM state (this is more a sanity
    +		// check).
    +		if (alternativesByPriority == null
    +			|| alternativesByPriority.isEmpty()
    +			|| !jobManagerState.hasState()
    +			|| jobManagerState.size() != 1) {
    +
    --- End diff --
    
    This means that current local recovery doesn't support rescale?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162584071
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TestTaskStateManager.java ---
    @@ -0,0 +1,185 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.core.testutils.OneShotLatch;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetrics;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +import org.apache.flink.runtime.taskmanager.CheckpointResponder;
    +
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +
    +import java.util.HashMap;
    +import java.util.Map;
    +
    +/**
    + * Implementation of {@link TaskStateManager} for tests.
    + */
    +public class TestTaskStateManager implements TaskStateManager {
    +
    +	private long reportedCheckpointId;
    +
    +	private JobID jobId;
    +	private ExecutionAttemptID executionAttemptID;
    +
    +	private final Map<Long, TaskStateSnapshot> taskStateSnapshotsByCheckpointId;
    +	private CheckpointResponder checkpointResponder;
    +	private OneShotLatch waitForReportLatch;
    +
    +	public TestTaskStateManager() {
    +		this(null, null, null);
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165318377
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTerminationTest.java ---
    @@ -227,21 +229,13 @@ protected void cleanup() throws Exception {
     			// has been stopped
     			CLEANUP_LATCH.trigger();
     
    -			// wait until handle async exception has been called to proceed with the termination of the
    -			// StreamTask
    -			HANDLE_ASYNC_EXCEPTION_LATCH.await();
    +			// wait until all async checkpoint threads are terminated, so that no more exceptions can be reported
    +			Assert.assertTrue(getAsyncOperationsThreadPool().awaitTermination(30L, TimeUnit.SECONDS));
    --- End diff --
    
    Where do we trigger the shut down of the `getAsyncOperationsThreadpool`?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168485028
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    I think the comment does not match the highlighted lines, can you please check?


---

[GitHub] flink issue #5239: [FLINK-8360] Implement task-local state recovery

Posted by StephanEwen <gi...@git.apache.org>.
Github user StephanEwen commented on the issue:

    https://github.com/apache/flink/pull/5239
  
    I did a peer review and walk through the code.
    Overall, the design look good, +1 for that!
    
    Some comments:
      - I would argue to change the way that these recovery options are configured. Currently, this goes through methods on one state backend objects, i.e., *configuration in code*. Because that recovery aspect is a really an "ops-related" aspect of running a Flink job (or a broader streaming platform), it should not be configured in code, but in the config. I found it helpful to thing that settings in code are for what concerns the application developers, settings in the config for what concerns the people that run Flink. They may be the same person in the end, but even then it is helpful because they are frequently are in different stages of the application development and deployment. Configurations are more easy to "standardize on", like "we want all applications in that group to enable local recovery".
    
      - One thing I am not yet 100% sure of is how this will interact in the future with RocksDB's optimized local recovery. I assume that checkpoints will in the future always use incremental snapshots. For such, there is no stream of bytes to store locally in addition. The files are already local and immutable. Here, the RocksDB snapshot should probably directly go through the local recovery directory, and the diff files would be persisted from there (the complete snapshot, which consists only of hardlinks to the files that are also in the work directory, would be retained though). Is the assumption that this is a "retain data structure" style mechanism, bespoke for each state backend, similar as retaining the heap copy-on-write table for the Heap State Backend?
    
    Now, since this PR is already complicated and needs a heavy rebase, I would be okay with doing that in another PR, if there is commitment to do this soon (before the 1.5 release branch is cut).
    
    Slightly off topic: This code has a very distinct style of using many `@Nonnull` annotations. Other newer parts of the code (the once that use annotations) follow the contract "non-null unless annotated with `@Nullable`". I don't ask to change this. Would be good to actually have a discussion and come up with a recommended style to agree on for the future.



---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170253853
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskLocalStateStoreImplTest.java ---
    @@ -59,12 +59,18 @@ public void before() throws Exception {
     		this.subtaskIdx = 0;
     		this.rootDirs = new File[]{temporaryFolder.newFolder(), temporaryFolder.newFolder()};
     
    +		LocalRecoveryDirectoryProviderImpl directoryProvider =
    +			new LocalRecoveryDirectoryProviderImpl(rootDirs, jobID, allocationID, jobVertexID, subtaskIdx);
    +
    +		LocalRecoveryConfig localRecoveryConfig =
    +			new LocalRecoveryConfig(LocalRecoveryConfig.LocalRecoveryMode.DISABLED, directoryProvider);
    --- End diff --
    
    So wouldn't that mean that with `LocalRecoveryConfig.disabled()`, the directory provider is again nullable? I think that `LocalRecoveryDirectoryProvider` is a very lightweight object, that also always has a useful default configuration.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168742980
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    Sorry I found these seem to be the problem of browser..... I'm trying to clear the cookies of my browser。 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168751784
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165098878
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StateObjectCollection.java ---
    @@ -0,0 +1,179 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.checkpoint;
    +
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.runtime.state.StateUtil;
    +
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.Collections;
    +import java.util.Iterator;
    +import java.util.function.Predicate;
    +
    +/**
    + * @param <T> type of the contained state objects.
    --- End diff --
    
    JavaDoc not complete.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168495158
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/DuplicatingCheckpointOutputStream.java ---
    @@ -0,0 +1,281 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import java.io.IOException;
    +
    +/**
    + * A CheckpointStateOutputStream that wraps a primary and a secondary CheckpointStateOutputStream and duplicates
    + * all writes into both streams. This stream applies buffering to reduce the amount of dual-method calling. Furthermore,
    + * exceptions that happen in interactions with the secondary stream are not exposed, until the user calls
    + * {@link #closeAndGetSecondaryHandle()}. In contrast to that, exceptions from interactions with the primary stream
    + * are immediately returned to the user. This class is used to write state for local recovery as a local (secondary)
    + * copy of the (primary) state snapshot that is written to a (slower but highly-available) remote filesystem.
    + */
    +public class DuplicatingCheckpointOutputStream extends CheckpointStreamFactory.CheckpointStateOutputStream {
    +
    +	private static final boolean STRICT_ALIGNMENT_CHECKS = false;
    +	private static final int DEFAULT_BUFFER_SIZER = 8 * 1024;
    +
    +	private final byte[] buffer;
    +	private int count;
    +
    +	/**
    +	 * Primary stream for writing the checkpoint data. Failures from this stream are forwarded.
    +	 */
    +	private final CheckpointStreamFactory.CheckpointStateOutputStream primaryOutputStream;
    +
    +	/**
    +	 * Primary stream for writing the checkpoint data. Failures from this stream are not forwarded until
    +	 * {@link #closeAndGetSecondaryHandle()}.
    +	 */
    +	private final CheckpointStreamFactory.CheckpointStateOutputStream secondaryOutputStream;
    +
    +	/**
    +	 * Stores a potential exception that occurred while interacting with {@link #secondaryOutputStream}
    +	 */
    +	private IOException secondaryStreamException;
    +
    +	public DuplicatingCheckpointOutputStream(
    +		CheckpointStreamFactory.CheckpointStateOutputStream primaryOutputStream,
    +		CheckpointStreamFactory.CheckpointStateOutputStream secondaryOutputStream) throws IOException {
    +		this(primaryOutputStream, secondaryOutputStream, DEFAULT_BUFFER_SIZER);
    +	}
    +
    +	public DuplicatingCheckpointOutputStream(
    +		CheckpointStreamFactory.CheckpointStateOutputStream primaryOutputStream,
    +		CheckpointStreamFactory.CheckpointStateOutputStream secondaryOutputStream,
    +		int bufferSize) throws IOException {
    +
    +		this.primaryOutputStream = Preconditions.checkNotNull(primaryOutputStream);
    +		this.secondaryOutputStream = Preconditions.checkNotNull(secondaryOutputStream);
    +
    +		this.buffer = new byte[bufferSize];
    +		this.count = 0;
    +
    +		this.secondaryStreamException = null;
    +
    +		checkForAlignedStreamPositions();
    +	}
    +
    +	@Override
    +	public void write(int b) throws IOException {
    +
    +		if (buffer.length <= count) {
    +			flushInternalBuffer();
    +		}
    +
    +		buffer[count] = (byte) b;
    +		++count;
    +	}
    +
    +	@Override
    +	public void write(byte[] b) throws IOException {
    +
    +		write(b, 0, b.length);
    +	}
    +
    +	@Override
    +	public void write(byte[] b, int off, int len) throws IOException {
    +
    +		if (buffer.length <= len) {
    +
    +			flushInternalBuffer();
    +			writeThroughInternal(b, off, len);
    +		} else {
    +
    +			if (buffer.length < len + count) {
    +				flushInternalBuffer();
    +			}
    +
    +			System.arraycopy(b, off, buffer, count, len);
    +			count += len;
    +		}
    +	}
    +
    +	@Override
    +	public long getPos() throws IOException {
    +		final long referencePos = primaryOutputStream.getPos();
    +		return referencePos + count;
    +	}
    +
    +	@Override
    +	public void flush() throws IOException {
    +
    +		flushInternalBuffer();
    +		primaryOutputStream.flush();
    +
    +		if (secondaryStreamException == null) {
    +			try {
    +				secondaryOutputStream.flush();
    +			} catch (IOException flushEx) {
    +				handleSecondaryStreamOnException(flushEx);
    +			}
    +		}
    +	}
    +
    +	@Override
    +	public void sync() throws IOException {
    +
    +		flushInternalBuffer();
    +		primaryOutputStream.sync();
    +
    +		if (secondaryStreamException == null) {
    +			try {
    +				secondaryOutputStream.sync();
    +			} catch (IOException syncEx) {
    +				handleSecondaryStreamOnException(syncEx);
    +			}
    +		}
    +	}
    +
    +	@Override
    +	public void close() throws IOException {
    +
    +		IOException exCollector = null;
    +
    +		try {
    +			flushInternalBuffer();
    +		} catch (IOException flushEx) {
    +			exCollector = flushEx;
    +		}
    +
    +		try {
    +			primaryOutputStream.close();
    +		} catch (IOException closeEx) {
    +			exCollector = ExceptionUtils.firstOrSuppressed(closeEx, exCollector);
    +		}
    +
    +		if (secondaryStreamException == null) {
    +			try {
    +				secondaryOutputStream.close();
    +			} catch (IOException closeEx) {
    +				handleSecondaryStreamOnException(closeEx);
    +			}
    +		}
    +
    +		if (exCollector != null) {
    +			throw exCollector;
    +		}
    +	}
    +
    +	private void checkForAlignedStreamPositions() throws IOException {
    +
    +		final long primaryPos = primaryOutputStream.getPos();
    +
    +		if (secondaryStreamException != null) {
    +			return;
    +		}
    +
    +		try {
    +			final long secondaryPos = secondaryOutputStream.getPos();
    +
    +			if (primaryPos != secondaryPos) {
    +				handleSecondaryStreamOnException(
    +					new IOException("Stream positions are out of sync between primary stream and secondary stream. " +
    +						"Reported positions are " + primaryPos + " (primary) and " + secondaryPos + " (secondary)."));
    +			}
    +		} catch (IOException posEx) {
    +			handleSecondaryStreamOnException(posEx);
    +		}
    +	}
    +
    +	private void flushInternalBuffer() throws IOException {
    +
    +		if (count > 0) {
    +			writeThroughInternal(buffer, 0, count);
    +			count = 0;
    +		}
    +	}
    +
    +	private void writeThroughInternal(byte[] b, int off, int len) throws IOException {
    +
    +		if (STRICT_ALIGNMENT_CHECKS) {
    +			checkForAlignedStreamPositions();
    +		}
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168752498
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -18,81 +18,267 @@
     
     package org.apache.flink.runtime.state;
     
    +import org.apache.flink.annotation.VisibleForTesting;
     import org.apache.flink.api.common.JobID;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
     import org.apache.flink.runtime.jobgraph.JobVertexID;
    -import org.apache.flink.util.Preconditions;
    +import org.apache.flink.util.FileUtils;
     
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
     import javax.annotation.Nonnull;
    +import javax.annotation.concurrent.GuardedBy;
     
    +import java.io.File;
    +import java.io.IOException;
     import java.util.HashMap;
     import java.util.Map;
    +import java.util.concurrent.Executor;
     
     /**
    - * This class holds the all {@link TaskLocalStateStore} objects for a task executor (manager).
    - *
    - * TODO: this still still work in progress and partially still acts as a placeholder.
    + * This class holds the all {@link TaskLocalStateStoreImpl} objects for a task executor (manager).
      */
     public class TaskExecutorLocalStateStoresManager {
     
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskExecutorLocalStateStoresManager.class);
    +
     	/**
     	 * This map holds all local state stores for tasks running on the task manager / executor that own the instance of
    -	 * this.
    +	 * this. Maps from allocation id to all the subtask's local state stores.
     	 */
    -	private final Map<JobID, Map<JobVertexSubtaskKey, TaskLocalStateStore>> taskStateManagers;
    +	@GuardedBy("lock")
    +	private final Map<AllocationID, Map<JobVertexSubtaskKey, TaskLocalStateStoreImpl>> taskStateStoresByAllocationID;
    +
    +	/** The configured mode for local recovery on this task manager. */
    +	private final LocalRecoveryConfig.LocalRecoveryMode localRecoveryMode;
    +
    +	/** This is the root directory for all local state of this task manager / executor. */
    +	private final File[] localStateRootDirectories;
    +
    +	/** Executor that runs the discarding of released state objects. */
    +	private final Executor discardExecutor;
    +
    +	/** Guarding lock for taskStateStoresByAllocationID and closed-flag. */
    +	private final Object lock;
    +
    +	private final Thread shutdownHook;
     
    -	public TaskExecutorLocalStateStoresManager() {
    -		this.taskStateManagers = new HashMap<>();
    +	@GuardedBy("lock")
    +	private boolean closed;
    +
    +	public TaskExecutorLocalStateStoresManager(
    +		@Nonnull LocalRecoveryConfig.LocalRecoveryMode localRecoveryMode,
    +		@Nonnull File[] localStateRootDirectories,
    +		@Nonnull Executor discardExecutor) {
    +
    +		this.taskStateStoresByAllocationID = new HashMap<>();
    +		this.localRecoveryMode = localRecoveryMode;
    +		this.localStateRootDirectories = localStateRootDirectories;
    +		this.discardExecutor = discardExecutor;
    +		this.lock = new Object();
    +		this.closed = false;
    +
    +		for (File localStateRecoveryRootDir : localStateRootDirectories) {
    +			if (!localStateRecoveryRootDir.exists()) {
    +
    +				if (!localStateRecoveryRootDir.mkdirs()) {
    +					throw new IllegalStateException("Could not create root directory for local recovery: " +
    +						localStateRecoveryRootDir);
    +				}
    +			}
    +		}
    +
    +		// install a shutdown hook
    +		this.shutdownHook = new Thread("TaskExecutorLocalStateStoresManager shutdown hook") {
    +			@Override
    +			public void run() {
    +				shutdown();
    +			}
    +		};
    +		try {
    +			Runtime.getRuntime().addShutdownHook(this.shutdownHook);
    +		} catch (IllegalStateException e) {
    +			// race, JVM is in shutdown already, we can safely ignore this
    +			LOG.debug("Unable to add shutdown hook for TaskExecutorLocalStateStoresManager, shutdown already in progress", e);
    +		} catch (Throwable t) {
    +			LOG.warn("Error while adding shutdown hook for TaskExecutorLocalStateStoresManager", t);
    +		}
     	}
     
    -	public TaskLocalStateStore localStateStoreForTask(
    -		JobID jobId,
    -		JobVertexID jobVertexID,
    -		int subtaskIndex) {
    +	@Nonnull
    +	public TaskLocalStateStore localStateStoreForSubtask(
    +		@Nonnull JobID jobId,
    +		@Nonnull AllocationID allocationID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex) {
    +
    +		synchronized (lock) {
    +
    +			if (closed) {
    +				throw new IllegalStateException("TaskExecutorLocalStateStoresManager is already closed and cannot " +
    +					"register a new TaskLocalStateStore.");
    +			}
    +
    +			final Map<JobVertexSubtaskKey, TaskLocalStateStoreImpl> taskStateManagers =
    +				this.taskStateStoresByAllocationID.computeIfAbsent(allocationID, k -> new HashMap<>());
    +
    +			final JobVertexSubtaskKey taskKey = new JobVertexSubtaskKey(jobVertexID, subtaskIndex);
     
    -		Preconditions.checkNotNull(jobId);
    -		final JobVertexSubtaskKey taskKey = new JobVertexSubtaskKey(jobVertexID, subtaskIndex);
    +			// create the allocation base dirs, one inside each root dir.
    +			File[] allocationBaseDirectories = allocationBaseDirectories(allocationID);
     
    -		final Map<JobVertexSubtaskKey, TaskLocalStateStore> taskStateManagers =
    -			this.taskStateManagers.computeIfAbsent(jobId, k -> new HashMap<>());
    +			LocalRecoveryDirectoryProviderImpl directoryProvider = new LocalRecoveryDirectoryProviderImpl(
    +				allocationBaseDirectories,
    +				jobId,
    +				allocationID,
    +				jobVertexID,
    +				subtaskIndex);
     
    -		return taskStateManagers.computeIfAbsent(
    -			taskKey, k -> new TaskLocalStateStore(jobId, jobVertexID, subtaskIndex));
    +			LocalRecoveryConfig localRecoveryConfig = new LocalRecoveryConfig(
    +				localRecoveryMode,
    +				directoryProvider);
    +
    +			return taskStateManagers.computeIfAbsent(
    +				taskKey,
    +				k -> new TaskLocalStateStoreImpl(
    +					jobId,
    +					allocationID,
    +					jobVertexID,
    +					subtaskIndex,
    +					localRecoveryConfig,
    +					discardExecutor));
    +		}
     	}
     
    -	public void releaseJob(JobID jobID) {
    +	public void releaseLocalStateForAllocationId(@Nonnull AllocationID allocationID) {
     
    --- End diff --
    
    In `TaskExecutorLocalStateStoresManager` I only found this `releaseLocalStateForAllocationId ()` which releases resources according to `allocationID`. Is there missing a function for releasing resources according to `jobID`? I think it should be used when a job is cancelled. 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168177861
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskStateManagerImplTest.java ---
    @@ -25,17 +25,24 @@
     import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
     import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
     import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
     import org.apache.flink.runtime.jobgraph.OperatorID;
     import org.apache.flink.runtime.taskmanager.CheckpointResponder;
     import org.apache.flink.runtime.taskmanager.TestCheckpointResponder;
     
     import org.junit.Assert;
     import org.junit.Test;
    +import org.junit.rules.TemporaryFolder;
    +
    +import java.io.IOException;
     
     import static org.mockito.Mockito.mock;
     
     public class TaskStateManagerImplTest {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170035520
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalRecoveryConfig.java ---
    @@ -0,0 +1,120 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.io.ParseException;
    +import org.apache.flink.configuration.CheckpointingOptions;
    +import org.apache.flink.configuration.Configuration;
    +
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnull;
    +
    +/**
    + * This class encapsulates the completed configuration for local recovery, i.e. the root
    + * directories into which all file-based snapshots can be written and the general mode for the local recover feature.
    + */
    +public class LocalRecoveryConfig {
    +
    +	/**
    +	 * Enum over modes of local recovery:
    +	 * <p><ul>
    +	 * <li>DISABLED: disables local recovery.
    +	 * <li>ENABLE_FILE_BASED: enables local recovery in a variant that is based on local files.
    +	 * </ul>
    +	 */
    +	public enum LocalRecoveryMode {
    +		DISABLED,
    +		ENABLE_FILE_BASED;
    +
    +		/**
    +		 * Attempts to parses the given string to a {@link LocalRecoveryMode}.
    +		 *
    +		 * @param input string to parse
    +		 * @return the {@link LocalRecoveryMode} that is represented by the input string.
    +		 * @throws ParseException if the string cannot be parsed to a {@link LocalRecoveryMode}.
    +		 */
    +		@Nonnull
    +		public static LocalRecoveryMode fromString(String input) throws ParseException {
    +			if (input != null) {
    +				String trimInput = input.trim();
    +				if (DISABLED.toString().equalsIgnoreCase(trimInput)) {
    +					return DISABLED;
    +				} else if (ENABLE_FILE_BASED.toString().equalsIgnoreCase(trimInput)) {
    +					return ENABLE_FILE_BASED;
    +				}
    +			}
    +			throw new ParseException("Cannot parse input to LocalRecoveryMode: " + input);
    +		}
    --- End diff --
    
    Why not simply using `LocalRecoveryMode.valueOf()`?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168777749
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PrioritizedOperatorSubtaskState.java ---
    @@ -0,0 +1,265 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.checkpoint;
    +
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.Preconditions;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.util.ArrayList;
    +import java.util.Collections;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.function.BiFunction;
    +
    +/**
    + * This class is a wrapper over multiple alternative {@link OperatorSubtaskState} that are (partial) substitutes for
    + * each other and imposes a priority ordering over all alternatives for the different states which define an order in
    + * which the operator should attempt to restore the state from them. One OperatorSubtaskState is considered as the
    + * "ground truth" about which state should be represented. Alternatives may be complete or partial substitutes for
    + * the "ground truth" with a higher priority (if they had a lower alternative, they would not really be alternatives).
    + * Substitution is determined on a per-sub-state basis.
    + */
    +public class PrioritizedOperatorSubtaskState {
    +
    +	/** Singleton instance for an empty, non-restored operator state. */
    +	private static final PrioritizedOperatorSubtaskState EMPTY_NON_RESTORED_INSTANCE =
    +		new PrioritizedOperatorSubtaskState(new OperatorSubtaskState(), Collections.emptyList(), false);
    +
    +	/** List of prioritized snapshot alternatives for managed operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedManagedOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for raw operator state. */
    +	private final List<StateObjectCollection<OperatorStateHandle>> prioritizedRawOperatorState;
    +
    +	/** List of prioritized snapshot alternatives for managed keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedManagedKeyedState;
    +
    +	/** List of prioritized snapshot alternatives for raw keyed state. */
    +	private final List<StateObjectCollection<KeyedStateHandle>> prioritizedRawKeyedState;
    +
    +	/** Signal flag if this represents state for a restored operator. */
    +	private final boolean restored;
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority) {
    +		this(jobManagerState, alternativesByPriority, true);
    +	}
    +
    +	public PrioritizedOperatorSubtaskState(
    +		@Nonnull OperatorSubtaskState jobManagerState,
    +		@Nonnull List<OperatorSubtaskState> alternativesByPriority,
    +		boolean restored) {
    +
    +		Preconditions.checkNotNull(jobManagerState, "Job manager state is null.");
    +		int size = Preconditions.checkNotNull(alternativesByPriority, "Alternative states are null.").size();
    +
    +		this.restored = restored;
    +
    +		List<StateObjectCollection<OperatorStateHandle>> managedOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> managedKeyedAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<OperatorStateHandle>> rawOperatorAlternatives = new ArrayList<>(size);
    +		List<StateObjectCollection<KeyedStateHandle>> rawKeyedAlternatives = new ArrayList<>(size);
    +
    +		for (OperatorSubtaskState subtaskState : alternativesByPriority) {
    +
    +			if (subtaskState != null) {
    +				managedKeyedAlternatives.add(subtaskState.getManagedKeyedState());
    +				rawKeyedAlternatives.add(subtaskState.getRawKeyedState());
    +				managedOperatorAlternatives.add(subtaskState.getManagedOperatorState());
    +				rawOperatorAlternatives.add(subtaskState.getRawOperatorState());
    +			}
    +		}
    +
    +		// Key-groups should match.
    +		BiFunction<KeyedStateHandle, KeyedStateHandle, Boolean> keyedStateApprover =
    +			(ref, alt) -> ref.getKeyGroupRange().equals(alt.getKeyGroupRange());
    +
    --- End diff --
    
    In theory, we could be less strict, yes. But in practice it doesn't matter because rescaling is not supported right now. See my other comment.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160959867
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamOperatorSnapshotRestoreTest.java ---
    @@ -98,7 +101,26 @@ public Integer getKey(Integer value) throws Exception {
     				TypeInformation.of(Integer.class),
     				MAX_PARALLELISM,
     				1 /* num subtasks */,
    -				0 /* subtask index */);
    +				0 /* subtask index */) {
    +
    +			@Override
    +			protected StreamTaskStateManager createStreamTaskStateManager(
    +				Environment env,
    +				StateBackend stateBackend,
    +				ProcessingTimeService processingTimeService) {
    +
    +				return new StreamTaskStateManagerImpl(env, stateBackend, processingTimeService) {
    +					@Override
    +					protected <K> InternalTimeServiceManager<?, K> internalTimeServiceManager(
    +						AbstractKeyedStateBackend<K> keyedStatedBackend,
    +						KeyContext keyContext,
    +						Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
    +
    +						return null;
    --- End diff --
    
    Does it make sense to return null in this method? Isn't this some kind of hack or a walkaround some underlying issue?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162583331
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java ---
    @@ -669,31 +676,14 @@ else if (current == ExecutionState.CANCELING) {
     			Environment env = new RuntimeEnvironment(
     				jobId, vertexId, executionId, executionConfig, taskInfo,
     				jobConfiguration, taskConfiguration, userCodeClassLoader,
    -				memoryManager, ioManager, broadcastVariableManager,
    +				memoryManager, ioManager, broadcastVariableManager, taskStateManager,
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168777737
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -266,11 +305,16 @@ public RocksDBKeyedStateBackend(
     		RocksIterator iterator = db.newIterator(columnInfo.f0);
     		iterator.seekToFirst();
     
    -		Iterable<K> iterable = () -> new RocksIteratorWrapper<>(iterator, state, keySerializer, keyGroupPrefixBytes);
    +		Iterable<K> iterable = () -> new RocksIteratorToJavaIteratorAdapter<>(iterator, state, keySerializer, keyGroupPrefixBytes);
     		Stream<K> targetStream = StreamSupport.stream(iterable.spliterator(), false);
    --- End diff --
    
    Checked. In `HeapKeyedStateBackend` the implementation is 
    ```
    	@Override
    	public <N> Stream<K> getKeys(String state, N namespace) {
    		if (!stateTables.containsKey(state)) {
    			return Stream.empty();
    		}
    		StateTable<K, N, ?> table = (StateTable<K, N, ?>) stateTables.get(state);
    		return table.getKeys(namespace);
    	}
    ```
    So I think `namespace` is useful ...


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159866648
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java ---
    @@ -669,31 +676,14 @@ else if (current == ExecutionState.CANCELING) {
     			Environment env = new RuntimeEnvironment(
     				jobId, vertexId, executionId, executionConfig, taskInfo,
     				jobConfiguration, taskConfiguration, userCodeClassLoader,
    -				memoryManager, ioManager, broadcastVariableManager,
    +				memoryManager, ioManager, broadcastVariableManager, taskStateManager,
     				accumulatorRegistry, kvStateRegistry, inputSplitProvider,
     				distributedCacheEntries, writers, inputGates,
     				checkpointResponder, taskManagerConfig, metrics, this);
     
     			// let the task code create its readers and writers
     			invokable.setEnvironment(env);
     
    -			// the very last thing before the actual execution starts running is to inject
    -			// the state into the task. the state is non-empty if this is an execution
    -			// of a task that failed but had backuped state from a checkpoint
    -
    -			if (null != taskRestore && taskRestore.getTaskStateSnapshot() != null) {
    -				if (invokable instanceof StatefulTask) {
    -					StatefulTask op = (StatefulTask) invokable;
    -					op.setInitialState(taskRestore.getTaskStateSnapshot());
    --- End diff --
    
    grrr, this cryptic `op` name forced me to look into the source code to check whether this is is instance of `StatefullTask` or not :/ could you rename it to something that at least is not an abbreviation?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168180080
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/LocalStateForwardingTest.java ---
    @@ -156,30 +157,38 @@ public void acknowledgeCheckpoint(
     			}
     		};
     
    -		TaskLocalStateStore taskLocalStateStore = new TaskLocalStateStore(jobID, jobVertexID, subtaskIdx) {
    -			@Override
    -			public void storeLocalState(
    -				@Nonnull CheckpointMetaData checkpointMetaData,
    -				@Nullable TaskStateSnapshot localState) {
    -
    -				Assert.assertEquals(tm, localState);
    -				tmReported.set(true);
    -			}
    -		};
    +		TemporaryFolder temporaryFolder = new TemporaryFolder();
     
    -		TaskStateManagerImpl taskStateManager =
    -			new TaskStateManagerImpl(
    -				jobID,
    -				executionAttemptID,
    -				taskLocalStateStore,
    -				null,
    -				checkpointResponder);
    -
    -		taskStateManager.reportTaskStateSnapshots(
    -			checkpointMetaData,
    -			checkpointMetrics,
    -			jm,
    -			tm);
    +		try {
    +			TaskLocalStateStore taskLocalStateStore =
    +				new TaskLocalStateStore(jobID, jobVertexID, subtaskIdx, temporaryFolder.newFolder()) {
    +					@Override
    +					public void storeLocalState(
    +						@Nonnull CheckpointMetaData checkpointMetaData,
    +						@Nullable TaskStateSnapshot localState) {
    +
    +						Assert.assertEquals(tm, localState);
    +						tmReported.set(true);
    +					}
    +				};
    +
    +			TaskStateManagerImpl taskStateManager =
    +				new TaskStateManagerImpl(
    +					jobID,
    +					executionAttemptID,
    +					taskLocalStateStore,
    +					null,
    +					checkpointResponder);
    +
    +			taskStateManager.reportTaskStateSnapshots(
    +				checkpointMetaData,
    +				checkpointMetrics,
    +				jm,
    +				tm);
    +		} catch (Exception ex) {
    +			temporaryFolder.delete();
    +			throw new RuntimeException(ex);
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168528301
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -19,92 +19,224 @@
     package org.apache.flink.runtime.state;
     
     import org.apache.flink.api.common.JobID;
    -import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
     import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
     import org.apache.flink.runtime.jobgraph.JobVertexID;
    -import org.apache.flink.util.ExceptionUtils;
     
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
     import javax.annotation.Nonnull;
     import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
     
     import java.io.File;
    -import java.util.HashMap;
    +import java.util.Arrays;
    +import java.util.Iterator;
     import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
     
     /**
      * This class will service as a task-manager-level local storage for local checkpointed state. The purpose is to provide
      * access to a state that is stored locally for a faster recovery compared to the state that is stored remotely in a
      * stable store DFS. For now, this storage is only complementary to the stable storage and local state is typically
      * lost in case of machine failures. In such cases (and others), client code of this class must fall back to using the
      * slower but highly available store.
    - *
    - * TODO this is currently a placeholder / mock that still must be implemented!
      */
     public class TaskLocalStateStore {
     
    -	/** */
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStore.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	private static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    +
    +	/** JobID from the owning subtask. */
     	private final JobID jobID;
     
    -	/** */
    +	/** JobVertexID of the owning subtask. */
     	private final JobVertexID jobVertexID;
     
    -	/** */
    +	/** Subtask index of the owning subtask. */
     	private final int subtaskIndex;
     
    -	/** */
    -	private final Map<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    -
     	/** The root directories for all local state of this {@link TaskLocalStateStore}. */
     	private final File[] localStateRootDirectories;
     
    +	/** Executor that runs the discarding of released state objects. */
    +	private final Executor discardExecutor;
    +
    +	/** Lock for synchronisation on the storage map and the discarded status. */
    +	private final Object lock;
    +
    +	/** Status flag if this store was already discarded. */
    +	@GuardedBy("lock")
    +	private boolean discarded;
    +
    +	/** Maps checkpoint ids to local TaskStateSnapshots. */
    +	@GuardedBy("lock")
    +	private final SortedMap<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
     	public TaskLocalStateStore(
    -		JobID jobID,
    -		JobVertexID jobVertexID,
    -		int subtaskIndex,
    -		File[] localStateRootDirectories) {
    +		@Nonnull JobID jobID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex,
    +		@Nonnull File[] localStateRootDirectories,
    +		@Nonnull Executor discardExecutor) {
     
     		this.jobID = jobID;
     		this.jobVertexID = jobVertexID;
     		this.subtaskIndex = subtaskIndex;
    -		this.storedTaskStateByCheckpointID = new HashMap<>();
     		this.localStateRootDirectories = localStateRootDirectories;
    +		this.discardExecutor = discardExecutor;
    +		this.lock = new Object();
    +		this.storedTaskStateByCheckpointID = new TreeMap<>();
    +		this.discarded = false;
     	}
     
    +	@Nonnull
     	protected String createSubtaskPath() {
     		return jobID + File.separator + jobVertexID + File.separator + subtaskIndex;
     	}
     
    +	/**
    +	 * Stores the local state for the given checkpoint id.
    +	 *
    +	 * @param checkpointId id for the checkpoint that created the local state that will be stored.
    +	 * @param localState the local state to store.
    +	 */
     	public void storeLocalState(
    -		@Nonnull CheckpointMetaData checkpointMetaData,
    +		@Nonnegative long checkpointId,
     		@Nullable TaskStateSnapshot localState) {
     
    -		TaskStateSnapshot previous =
    -			storedTaskStateByCheckpointID.put(checkpointMetaData.getCheckpointId(), localState);
    +		if (localState == null) {
    +			localState = NULL_DUMMY;
    +		}
     
    -		if (previous != null) {
    -			throw new IllegalStateException("Found previously registered local state for checkpoint with id " +
    -				checkpointMetaData.getCheckpointId() + "! This indicated a problem.");
    +		LOG.info("Storing local state for checkpoint {}.", checkpointId);
    +		LOG.debug("Local state for checkpoint {} is {}.", checkpointId, localState);
    +
    +		synchronized (lock) {
    +			if (discarded) {
    +				// we ignore late stores and simply discard the state.
    +				discardStateObject(localState, checkpointId);
    +			} else {
    +				TaskStateSnapshot previous =
    +					storedTaskStateByCheckpointID.put(checkpointId, localState);
    +
    +				if (previous != null) {
    +					// this should never happen.
    +					discardStateObject(previous, checkpointId);
    +					throw new IllegalStateException("Found previously registered local state for checkpoint with id " +
    +						checkpointId + "! This indicated a problem.");
    +				}
    +
    +				// prune history.
    +				while (storedTaskStateByCheckpointID.size() > MAX_RETAINED_SNAPSHOTS) {
    +					Long removeCheckpointID = storedTaskStateByCheckpointID.firstKey();
    +					TaskStateSnapshot snapshot =
    +						storedTaskStateByCheckpointID.remove(storedTaskStateByCheckpointID.firstKey());
    --- End diff --
    
    Deleted in later commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168193098
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -336,1697 +386,1982 @@ public int getKeyGroupPrefixBytes() {
     	 * @param streamFactory The factory that we can use for writing our state to streams.
     	 * @param checkpointOptions Options for how to perform this checkpoint.
     	 * @return Future to the state handle of the snapshot data.
    -	 * @throws Exception
    +	 * @throws Exception indicating a problem in the synchronous part of the checkpoint.
     	 */
     	@Override
    -	public RunnableFuture<KeyedStateHandle> snapshot(
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot(
     		final long checkpointId,
     		final long timestamp,
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.",
    -						checkpointTimestamp);
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
     		try {
    -			snapshotOperation.takeSnapshot();
    -		} catch (Exception e) {
    -			snapshotOperation.stop();
    -			snapshotOperation.releaseResources(true);
    -			throw e;
    -		}
    -
    -		return new FutureTask<KeyedStateHandle>(
    -			new Callable<KeyedStateHandle>() {
    -				@Override
    -				public KeyedStateHandle call() throws Exception {
    -					return snapshotOperation.materializeSnapshot();
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else {
    +				KeyedStateHandle firstStateHandle = restoreState.iterator().next();
    +				if (firstStateHandle instanceof IncrementalKeyedStateHandle
    +					|| firstStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +					restoreOperation.restore(restoreState);
    +				} else {
    +					RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +					restoreOperation.doRestore(restoreState);
     				}
     			}
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    -
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    -			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<KeyedStateHandle> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at {}. Returning null.", timestamp);
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<KeyedStateHandle> ioCallable =
    -			new AbstractAsyncCallableWithResources<KeyedStateHandle>() {
    -
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				@Override
    -				public KeyGroupsStateHandle performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in FIRST position, see
    +		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, columnOptions));
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					return snapshotOperation.getSnapshotResultStateHandle();
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot ({}, synchronous part) in thread {} took {} ms.",
    -				streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    -
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
     
    -						assert (!hasMetaDataFollowsFlag(previousKey));
    +		private final RocksDBKeyedStateBackend<T> stateBackend;
     
    -						//set signal in first key byte that meta data will follow in the stream after this k/v pair
    -						if (mergeIterator.isNewKeyGroup() || mergeIterator.isNewKeyValueState()) {
    +		private RocksDBIncrementalRestoreOperation(RocksDBKeyedStateBackend<T> stateBackend) {
    +			this.stateBackend = stateBackend;
    +		}
     
    -							//be cooperative and check for interruption from time to time in the hot loop
    -							checkInterrupted();
    +		/**
    +		 * Root method that branches for different implementations of {@link KeyedStateHandle}.
    +		 */
    +		void restore(Collection<KeyedStateHandle> restoreStateHandles) throws Exception {
     
    -							setMetaDataFollowsFlagInKey(previousKey);
    -						}
    +			boolean hasExtraKeys = (restoreStateHandles.size() > 1 ||
    +				!Objects.equals(restoreStateHandles.iterator().next().getKeyGroupRange(), stateBackend.keyGroupRange));
     
    -						writeKeyValuePair(previousKey, previousValue, kgOutView);
    +			if (hasExtraKeys) {
    +				stateBackend.createDB();
    +			}
     
    -						//write meta data if we have to
    -						if (mergeIterator.isNewKeyGroup()) {
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -							// this will just close the outer stream
    -							kgOutStream.close();
    -							//begin new key-group
    -							keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -							//write the kev-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -							kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						} else if (mergeIterator.isNewKeyValueState()) {
    -							//write the k/v-state
    -							//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -							kgOutView.writeShort(mergeIterator.kvStateId());
    -						}
    +			for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
     
    -						//request next k/v pair
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    -					}
    +				if (rawStateHandle instanceof IncrementalKeyedStateHandle) {
    +					restoreInstance((IncrementalKeyedStateHandle) rawStateHandle, hasExtraKeys);
    +				} else if (rawStateHandle instanceof IncrementalLocalKeyedStateHandle) {
    +					Preconditions.checkState(!hasExtraKeys, "Cannot recover from local state after rescaling.");
    +					restoreInstance((IncrementalLocalKeyedStateHandle) rawStateHandle);
    +				} else {
    +					throw new IllegalStateException("Unexpected state handle type, " +
    +						"expected " + IncrementalKeyedStateHandle.class +
    +						", but found " + rawStateHandle.getClass());
     				}
    +			}
    +		}
     
    -				//epilogue: write last key-group
    -				if (previousKey != null) {
    -					assert (!hasMetaDataFollowsFlag(previousKey));
    -					setMetaDataFollowsFlagInKey(previousKey);
    -					writeKeyValuePair(previousKey, previousValue, kgOutView);
    -					//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -					kgOutView.writeShort(END_OF_KEY_GROUP_MARK);
    -					// this will just close the outer stream
    -					kgOutStream.close();
    -					kgOutStream = null;
    -				}
    +		/**
    +		 * Recovery from remote incremental state.
    +		 */
    +		private void restoreInstance(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			boolean hasExtraKeys) throws Exception {
    +
    +			// read state data
    +			Path temporaryRestoreInstancePath = new Path(
    +				stateBackend.instanceBasePath.getAbsolutePath(),
    +				UUID.randomUUID().toString());
    +
    +			try {
    +
    +				transferAllStateDataToDirectory(restoreStateHandle, temporaryRestoreInstancePath);
     
    +				// read meta data
    +				List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +					readMetaData(restoreStateHandle.getMetaStateHandle());
    +
    +				List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +					createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
    +
    +				if (hasExtraKeys) {
    +					restoreKeyGroupsShardWithTemporaryHelperInstance(
    +						temporaryRestoreInstancePath,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				} else {
    +
    +					// since we transferred all remote state to a local directory, we can use the same code as for
    +					// local recovery.
    +					IncrementalLocalKeyedStateHandle localKeyedStateHandle = new IncrementalLocalKeyedStateHandle(
    +						restoreStateHandle.getBackendIdentifier(),
    +						restoreStateHandle.getCheckpointId(),
    +						new DirectoryStateHandle(temporaryRestoreInstancePath),
    +						restoreStateHandle.getKeyGroupRange(),
    +						restoreStateHandle.getMetaStateHandle(),
    +						restoreStateHandle.getSharedState().keySet());
    +
    +					restoreLocalStateIntoFullInstance(
    +						localKeyedStateHandle,
    +						columnFamilyDescriptors,
    +						stateMetaInfoSnapshots);
    +				}
     			} finally {
    -				// this will just close the outer stream
    -				IOUtils.closeQuietly(kgOutStream);
    +				FileSystem restoreFileSystem = temporaryRestoreInstancePath.getFileSystem();
    +				if (restoreFileSystem.exists(temporaryRestoreInstancePath)) {
    +					restoreFileSystem.delete(temporaryRestoreInstancePath, true);
    +				}
     			}
     		}
     
    -		private void writeKeyValuePair(byte[] key, byte[] value, DataOutputView out) throws IOException {
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(key, out);
    -			BytePrimitiveArraySerializer.INSTANCE.serialize(value, out);
    -		}
    +		/**
    +		 * Recovery from local incremental state.
    +		 */
    +		private void restoreInstance(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
    +			// read meta data
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots =
    +				readMetaData(localKeyedStateHandle.getMetaDataState());
     
    -		static void setMetaDataFollowsFlagInKey(byte[] key) {
    -			key[0] |= FIRST_BIT_IN_BYTE_MASK;
    -		}
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots);
     
    -		static void clearMetaDataFollowsFlag(byte[] key) {
    -			key[0] &= (~RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    +			restoreLocalStateIntoFullInstance(
    +				localKeyedStateHandle,
    +				columnFamilyDescriptors,
    +				stateMetaInfoSnapshots);
     		}
     
    -		static boolean hasMetaDataFollowsFlag(byte[] key) {
    -			return 0 != (key[0] & RocksDBFullSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
    -		}
    +		/**
    +		 * This method recreates and registers all {@link ColumnFamilyDescriptor} from Flink's state meta data snapshot.
    +		 */
    +		private List<ColumnFamilyDescriptor> createAndRegisterColumnFamilyDescriptors(
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) {
     
    -		private static void checkInterrupted() throws InterruptedException {
    -			if (Thread.currentThread().isInterrupted()) {
    -				throw new InterruptedException("RocksDB snapshot interrupted.");
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +				new ArrayList<>(1 + stateMetaInfoSnapshots.size());
    +
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot : stateMetaInfoSnapshots) {
    +
    +				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +					stateMetaInfoSnapshot.getName().getBytes(ConfigConstants.DEFAULT_CHARSET),
    +					stateBackend.columnOptions);
    +
    +				columnFamilyDescriptors.add(columnFamilyDescriptor);
    +				stateBackend.restoredKvStateMetaInfos.put(stateMetaInfoSnapshot.getName(), stateMetaInfoSnapshot);
     			}
    +			return columnFamilyDescriptors;
     		}
    -	}
     
    -	private static final class RocksDBIncrementalSnapshotOperation<K> {
    +		/**
    +		 * This method implements the core of the restore logic that unifies how local and remote state are recovered.
    +		 */
    +		private void restoreLocalStateIntoFullInstance(
    +			IncrementalLocalKeyedStateHandle restoreStateHandle,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
    +			// pick up again the old backend id, so the we can reference existing state
    +			stateBackend.backendUID = restoreStateHandle.getBackendIdentifier();
    +
    +			LOG.debug("Restoring keyed backend uid in operator {} from incremental snapshot to {}.",
    +				stateBackend.operatorIdentifier, stateBackend.backendUID);
    +
    +			// create hard links in the instance directory
    +			if (!stateBackend.instanceRocksDBPath.mkdirs()) {
    +				throw new IOException("Could not create RocksDB data directory.");
    +			}
     
    -		/** The backend which we snapshot. */
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    +			Path restoreSourcePath = restoreStateHandle.getDirectoryStateHandle().getDirectory();
    +			restoreInstanceDirectoryFromPath(restoreSourcePath);
     
    -		/** Stream factory that creates the outpus streams to DFS. */
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -		/** Id for the current checkpoint. */
    -		private final long checkpointId;
    +			stateBackend.db = stateBackend.openDB(
    +				stateBackend.instanceRocksDBPath.getAbsolutePath(),
    +				columnFamilyDescriptors, columnFamilyHandles);
     
    -		/** Timestamp for the current checkpoint. */
    -		private final long checkpointTimestamp;
    +			// extract and store the default column family which is located at the first index
    +			stateBackend.defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -		/** All sst files that were part of the last previously completed checkpoint. */
    -		private Set<StateHandleID> baseSstFiles;
    +			for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +				RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -		/** The state meta data. */
    -		private final List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots = new ArrayList<>();
    +				ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +				RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +					new RegisteredKeyedBackendStateMetaInfo<>(
    +						stateMetaInfoSnapshot.getStateType(),
    +						stateMetaInfoSnapshot.getName(),
    +						stateMetaInfoSnapshot.getNamespaceSerializer(),
    +						stateMetaInfoSnapshot.getStateSerializer());
     
    -		private FileSystem backupFileSystem;
    -		private Path backupPath;
    +				stateBackend.kvStateInformation.put(
    +					stateMetaInfoSnapshot.getName(),
    +					new Tuple2<>(columnFamilyHandle, stateMetaInfo));
    +			}
     
    -		// Registry for all opened i/o streams
    -		private final CloseableRegistry closeableRegistry = new CloseableRegistry();
    +			// use the restore sst files as the base for succeeding checkpoints
    +			synchronized (stateBackend.materializedSstFiles) {
    +				stateBackend.materializedSstFiles.put(
    +					restoreStateHandle.getCheckpointId(),
    +					restoreStateHandle.getSharedStateHandleIDs());
    +			}
     
    -		// new sst files since the last completed checkpoint
    -		private final Map<StateHandleID, StreamStateHandle> sstFiles = new HashMap<>();
    +			stateBackend.lastCompletedCheckpointId = restoreStateHandle.getCheckpointId();
    +		}
     
    -		// handles to the misc files in the current snapshot
    -		private final Map<StateHandleID, StreamStateHandle> miscFiles = new HashMap<>();
    +		/**
    +		 * This recreates the new working directory of the recovered RocksDB instance and links/copies the contents from
    +		 * a local state.
    +		 */
    +		private void restoreInstanceDirectoryFromPath(Path source) throws IOException {
     
    -		// This lease protects from concurrent disposal of the native rocksdb instance.
    -		private final ResourceGuard.Lease dbLease;
    +			FileSystem fileSystem = source.getFileSystem();
     
    -		private StreamStateHandle metaStateHandle = null;
    +			final FileStatus[] fileStatuses = fileSystem.listStatus(source);
     
    -		private RocksDBIncrementalSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			long checkpointId,
    -			long checkpointTimestamp) throws IOException {
    +			if (fileStatuses == null) {
    +				throw new IOException("Cannot list file statues. Directory " + source + " does not exist.");
    +			}
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimestamp = checkpointTimestamp;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    +			for (FileStatus fileStatus : fileStatuses) {
    +				final Path filePath = fileStatus.getPath();
    +				final String fileName = filePath.getName();
    +				File restoreFile = new File(source.getPath(), fileName);
    +				File targetFile = new File(stateBackend.instanceRocksDBPath.getPath(), fileName);
    +				if (fileName.endsWith(SST_FILE_SUFFIX)) {
    +					// hardlink'ing the immutable sst-files.
    +					Files.createLink(targetFile.toPath(), restoreFile.toPath());
    +				} else {
    +					// true copy for all other files.
    +					Files.copy(restoreFile.toPath(), targetFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    +				}
    +			}
     		}
     
    -		private StreamStateHandle materializeStateData(Path filePath) throws Exception {
    +		/**
    +		 * Reads Flink's state meta data file from the state handle.
    +		 */
    +		private List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> readMetaData(
    +			StreamStateHandle metaStateHandle) throws Exception {
    +
     			FSDataInputStream inputStream = null;
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
     
     			try {
    -				final byte[] buffer = new byte[8 * 1024];
    -
    -				FileSystem backupFileSystem = backupPath.getFileSystem();
    -				inputStream = backupFileSystem.open(filePath);
    -				closeableRegistry.registerCloseable(inputStream);
    -
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
    -				closeableRegistry.registerCloseable(outputStream);
    -
    -				while (true) {
    -					int numBytes = inputStream.read(buffer);
    +				inputStream = metaStateHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -					if (numBytes == -1) {
    -						break;
    -					}
    +				KeyedBackendSerializationProxy<T> serializationProxy =
    +					new KeyedBackendSerializationProxy<>(stateBackend.userCodeClassLoader);
    +				DataInputView in = new DataInputViewStreamWrapper(inputStream);
    +				serializationProxy.read(in);
     
    -					outputStream.write(buffer, 0, numBytes);
    -				}
    +				// check for key serializer compatibility; this also reconfigures the
    +				// key serializer to be compatible, if it is required and is possible
    +				if (CompatibilityUtil.resolveCompatibilityResult(
    +					serializationProxy.getKeySerializer(),
    +					UnloadableDummyTypeSerializer.class,
    +					serializationProxy.getKeySerializerConfigSnapshot(),
    +					stateBackend.keySerializer)
    +					.isRequiresMigration()) {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    +					// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +					throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +						"Aborting now since state migration is currently not available");
     				}
    -				return result;
     
    +				return serializationProxy.getStateMetaInfoSnapshots();
     			} finally {
    -				if (inputStream != null && closeableRegistry.unregisterCloseable(inputStream)) {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
     					inputStream.close();
     				}
    -
    -				if (outputStream != null && closeableRegistry.unregisterCloseable(outputStream)) {
    -					outputStream.close();
    -				}
     			}
     		}
     
    -		private StreamStateHandle materializeMetaData() throws Exception {
    -			CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;
    -
    -			try {
    -				outputStream = checkpointStreamFactory
    -					.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);
    -				closeableRegistry.registerCloseable(outputStream);
    +		private void transferAllStateDataToDirectory(
    +			IncrementalKeyedStateHandle restoreStateHandle,
    +			Path dest) throws IOException {
     
    -				//no need for compression scheme support because sst-files are already compressed
    -				KeyedBackendSerializationProxy<K> serializationProxy =
    -					new KeyedBackendSerializationProxy<>(
    -						stateBackend.keySerializer,
    -						stateMetaInfoSnapshots,
    -						false);
    +			final Map<StateHandleID, StreamStateHandle> sstFiles =
    +				restoreStateHandle.getSharedState();
    +			final Map<StateHandleID, StreamStateHandle> miscFiles =
    +				restoreStateHandle.getPrivateState();
     
    -				DataOutputView out = new DataOutputViewStreamWrapper(outputStream);
    +			transferAllDataFromStateHandles(sstFiles, dest);
    +			transferAllDataFromStateHandles(miscFiles, dest);
    +		}
     
    -				serializationProxy.write(out);
    +		/**
    +		 * Copies all the files from the given stream state handles to the given path, renaming the files w.r.t. their
    +		 * {@link StateHandleID}.
    +		 */
    +		private void transferAllDataFromStateHandles(
    +			Map<StateHandleID, StreamStateHandle> stateHandleMap,
    +			Path restoreInstancePath) throws IOException {
     
    -				StreamStateHandle result = null;
    -				if (closeableRegistry.unregisterCloseable(outputStream)) {
    -					result = outputStream.closeAndGetHandle();
    -					outputStream = null;
    -				}
    -				return result;
    -			} finally {
    -				if (outputStream != null) {
    -					if (closeableRegistry.unregisterCloseable(outputStream)) {
    -						outputStream.close();
    -					}
    -				}
    +			for (Map.Entry<StateHandleID, StreamStateHandle> entry : stateHandleMap.entrySet()) {
    +				StateHandleID stateHandleID = entry.getKey();
    +				StreamStateHandle remoteFileHandle = entry.getValue();
    +				copyStateDataHandleData(new Path(restoreInstancePath, stateHandleID.toString()), remoteFileHandle);
     			}
     		}
     
    -		void takeSnapshot() throws Exception {
    +		/**
    +		 * Copies the file from a single state handle to the given path.
    +		 */
    +		private void copyStateDataHandleData(
    +			Path restoreFilePath,
    +			StreamStateHandle remoteFileHandle) throws IOException {
     
    -			final long lastCompletedCheckpoint;
    +			FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
     
    -			// use the last completed checkpoint as the comparison base.
    -			synchronized (stateBackend.materializedSstFiles) {
    -				lastCompletedCheckpoint = stateBackend.lastCompletedCheckpointId;
    -				baseSstFiles = stateBackend.materializedSstFiles.get(lastCompletedCheckpoint);
    -			}
    +			FSDataInputStream inputStream = null;
    +			FSDataOutputStream outputStream = null;
     
    -			LOG.trace("Taking incremental snapshot for checkpoint {}. Snapshot is based on last completed checkpoint {} " +
    -				"assuming the following (shared) files as base: {}.", checkpointId, lastCompletedCheckpoint, baseSstFiles);
    +			try {
    +				inputStream = remoteFileHandle.openInputStream();
    +				stateBackend.cancelStreamRegistry.registerCloseable(inputStream);
     
    -			// save meta data
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> stateMetaInfoEntry
    -				: stateBackend.kvStateInformation.entrySet()) {
    -				stateMetaInfoSnapshots.add(stateMetaInfoEntry.getValue().f1.snapshot());
    -			}
    +				outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
    +				stateBackend.cancelStreamRegistry.registerCloseable(outputStream);
     
    -			// save state data
    -			backupPath = new Path(stateBackend.instanceBasePath.getAbsolutePath(), "chk-" + checkpointId);
    +				byte[] buffer = new byte[8 * 1024];
    +				while (true) {
    +					int numBytes = inputStream.read(buffer);
    +					if (numBytes == -1) {
    +						break;
    +					}
     
    -			LOG.trace("Local RocksDB checkpoint goes to backup path {}.", backupPath);
    +					outputStream.write(buffer, 0, numBytes);
    +				}
    +			} finally {
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(inputStream)) {
    +					inputStream.close();
    +				}
     
    -			backupFileSystem = backupPath.getFileSystem();
    -			if (backupFileSystem.exists(backupPath)) {
    -				throw new IllegalStateException("Unexpected existence of the backup directory.");
    +				if (stateBackend.cancelStreamRegistry.unregisterCloseable(outputStream)) {
    +					outputStream.close();
    +				}
     			}
    -
    -			// create hard links of living files in the checkpoint path
    -			Checkpoint checkpoint = Checkpoint.create(stateBackend.db);
    -			checkpoint.createCheckpoint(backupPath.getPath());
     		}
     
    -		KeyedStateHandle materializeSnapshot() throws Exception {
    +		/**
    +		 * In case of 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 restoreKeyGroupsShardWithTemporaryHelperInstance(
    +			Path restoreInstancePath,
    +			List<ColumnFamilyDescriptor> columnFamilyDescriptors,
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> stateMetaInfoSnapshots) throws Exception {
     
    -			stateBackend.cancelStreamRegistry.registerCloseable(closeableRegistry);
    +			List<ColumnFamilyHandle> columnFamilyHandles =
    +				new ArrayList<>(1 + columnFamilyDescriptors.size());
     
    -			// write meta data
    -			metaStateHandle = materializeMetaData();
    +			try (RocksDB restoreDb = stateBackend.openDB(
    +				restoreInstancePath.getPath(),
    +				columnFamilyDescriptors,
    +				columnFamilyHandles)) {
     
    -			// write state data
    -			Preconditions.checkState(backupFileSystem.exists(backupPath));
    +				final ColumnFamilyHandle defaultColumnFamily = columnFamilyHandles.remove(0);
     
    -			FileStatus[] fileStatuses = backupFileSystem.listStatus(backupPath);
    -			if (fileStatuses != null) {
    -				for (FileStatus fileStatus : fileStatuses) {
    -					final Path filePath = fileStatus.getPath();
    -					final String fileName = filePath.getName();
    -					final StateHandleID stateHandleID = new StateHandleID(fileName);
    +				Preconditions.checkState(columnFamilyHandles.size() == columnFamilyDescriptors.size());
     
    -					if (fileName.endsWith(SST_FILE_SUFFIX)) {
    -						final boolean existsAlready =
    -							baseSstFiles != null && baseSstFiles.contains(stateHandleID);
    +				try {
    +					for (int i = 0; i < columnFamilyDescriptors.size(); ++i) {
    +						ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
    +						ColumnFamilyDescriptor columnFamilyDescriptor = columnFamilyDescriptors.get(i);
    +						RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> stateMetaInfoSnapshot = stateMetaInfoSnapshots.get(i);
     
    -						if (existsAlready) {
    -							// we introduce a placeholder state handle, that is replaced with the
    -							// original from the shared state registry (created from a previous checkpoint)
    -							sstFiles.put(
    -								stateHandleID,
    -								new PlaceholderStreamStateHandle());
    -						} else {
    -							sstFiles.put(stateHandleID, materializeStateData(filePath));
    -						}
    -					} else {
    -						StreamStateHandle fileHandle = materializeStateData(filePath);
    -						miscFiles.put(stateHandleID, fileHandle);
    -					}
    -				}
    -			}
    +						Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredStateMetaInfoEntry =
    +							stateBackend.kvStateInformation.get(stateMetaInfoSnapshot.getName());
     
    -			synchronized (stateBackend.materializedSstFiles) {
    -				stateBackend.materializedSstFiles.put(checkpointId, sstFiles.keySet());
    -			}
    +						if (null == registeredStateMetaInfoEntry) {
     
    -			return new IncrementalKeyedStateHandle(
    -				stateBackend.backendUID,
    -				stateBackend.keyGroupRange,
    -				checkpointId,
    -				sstFiles,
    -				miscFiles,
    -				metaStateHandle);
    -		}
    +							RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +								new RegisteredKeyedBackendStateMetaInfo<>(
    +									stateMetaInfoSnapshot.getStateType(),
    +									stateMetaInfoSnapshot.getName(),
    +									stateMetaInfoSnapshot.getNamespaceSerializer(),
    +									stateMetaInfoSnapshot.getStateSerializer());
     
    -		void stop() {
    +							registeredStateMetaInfoEntry =
    +								new Tuple2<>(
    +									stateBackend.db.createColumnFamily(columnFamilyDescriptor),
    +									stateMetaInfo);
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Could not properly close io streams.", e);
    -				}
    -			}
    -		}
    +							stateBackend.kvStateInformation.put(
    +								stateMetaInfoSnapshot.getName(),
    +								registeredStateMetaInfoEntry);
    +						}
     
    -		void releaseResources(boolean canceled) {
    +						ColumnFamilyHandle targetColumnFamilyHandle = registeredStateMetaInfoEntry.f0;
     
    -			dbLease.close();
    +						try (RocksIterator iterator = restoreDb.newIterator(columnFamilyHandle)) {
     
    -			if (stateBackend.cancelStreamRegistry.unregisterCloseable(closeableRegistry)) {
    -				try {
    -					closeableRegistry.close();
    -				} catch (IOException e) {
    -					LOG.warn("Exception on closing registry.", e);
    -				}
    -			}
    +							int startKeyGroup = stateBackend.getKeyGroupRange().getStartKeyGroup();
    +							byte[] startKeyGroupPrefixBytes = new byte[stateBackend.keyGroupPrefixBytes];
    +							for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +								startKeyGroupPrefixBytes[j] = (byte) (startKeyGroup >>> ((stateBackend.keyGroupPrefixBytes - j - 1) * Byte.SIZE));
    +							}
     
    -			if (backupPath != null) {
    -				try {
    -					if (backupFileSystem.exists(backupPath)) {
    +							iterator.seek(startKeyGroupPrefixBytes);
     
    -						LOG.trace("Deleting local RocksDB backup path {}.", backupPath);
    -						backupFileSystem.delete(backupPath, true);
    -					}
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly delete the checkpoint directory.", e);
    -				}
    -			}
    +							while (iterator.isValid()) {
     
    -			if (canceled) {
    -				Collection<StateObject> statesToDiscard =
    -					new ArrayList<>(1 + miscFiles.size() + sstFiles.size());
    +								int keyGroup = 0;
    +								for (int j = 0; j < stateBackend.keyGroupPrefixBytes; ++j) {
    +									keyGroup = (keyGroup << Byte.SIZE) + iterator.key()[j];
    +								}
     
    -				statesToDiscard.add(metaStateHandle);
    -				statesToDiscard.addAll(miscFiles.values());
    -				statesToDiscard.addAll(sstFiles.values());
    +								if (stateBackend.keyGroupRange.contains(keyGroup)) {
    +									stateBackend.db.put(targetColumnFamilyHandle,
    +										iterator.key(), iterator.value());
    +								}
     
    -				try {
    -					StateUtil.bestEffortDiscardAllStateObjects(statesToDiscard);
    -				} catch (Exception e) {
    -					LOG.warn("Could not properly discard states.", e);
    +								iterator.next();
    +							}
    +						} // releases native iterator resources
    +					}
    +				} finally {
    +
    +					//release native tmp db column family resources
    +					IOUtils.closeQuietly(defaultColumnFamily);
    +
    +					for (ColumnFamilyHandle flinkColumnFamilyHandle : columnFamilyHandles) {
    +						IOUtils.closeQuietly(flinkColumnFamilyHandle);
    +					}
     				}
    -			}
    +			} // releases native tmp db resources
     		}
     	}
     
    -	@Override
    -	public void restore(Collection<KeyedStateHandle> restoreState) throws Exception {
    -		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
    +	// ------------------------------------------------------------------------
    +	//  State factories
    +	// ------------------------------------------------------------------------
     
    -		if (LOG.isDebugEnabled()) {
    -			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
    -		}
    +	/**
    +	 * Creates a column family handle for use with a k/v state. When restoring from a snapshot
    +	 * we don't restore the individual k/v states, just the global RocksDB database and the
    +	 * list of column families. When a k/v state is first requested we check here whether we
    +	 * already have a column family for that and return it or create a new one if it doesn't exist.
    +	 *
    +	 * <p>This also checks whether the {@link StateDescriptor} for a state matches the one
    +	 * that we checkpointed, i.e. is already in the map of column families.
    +	 */
    +	@SuppressWarnings("rawtypes, unchecked")
    +	protected <N, S> ColumnFamilyHandle getColumnFamily(
    +		StateDescriptor<?, S> descriptor, TypeSerializer<N> namespaceSerializer) throws IOException, StateMigrationException {
     
    -		// clear all meta data
    -		kvStateInformation.clear();
    -		restoredKvStateMetaInfos.clear();
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> stateInfo =
    +			kvStateInformation.get(descriptor.getName());
     
    -		try {
    -			if (restoreState == null || restoreState.isEmpty()) {
    -				createDB();
    -			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    -				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    -				restoreOperation.restore(restoreState);
    -			} else {
    -				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    -				restoreOperation.doRestore(restoreState);
    +		RegisteredKeyedBackendStateMetaInfo<N, S> newMetaInfo = new RegisteredKeyedBackendStateMetaInfo<>(
    +			descriptor.getType(),
    +			descriptor.getName(),
    +			namespaceSerializer,
    +			descriptor.getSerializer());
    +
    +		if (stateInfo != null) {
    +			// TODO with eager registration in place, these checks should be moved to restore()
    +
    +			RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S> restoredMetaInfo =
    +				(RegisteredKeyedBackendStateMetaInfo.Snapshot<N, S>) restoredKvStateMetaInfos.get(descriptor.getName());
    +
    +			Preconditions.checkState(
    +				Objects.equals(newMetaInfo.getName(), restoredMetaInfo.getName()),
    +				"Incompatible state names. " +
    +					"Was [" + restoredMetaInfo.getName() + "], " +
    +					"registered with [" + newMetaInfo.getName() + "].");
    +
    +			if (!Objects.equals(newMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)
    +				&& !Objects.equals(restoredMetaInfo.getStateType(), StateDescriptor.Type.UNKNOWN)) {
    +
    +				Preconditions.checkState(
    +					newMetaInfo.getStateType() == restoredMetaInfo.getStateType(),
    +					"Incompatible state types. " +
    +						"Was [" + restoredMetaInfo.getStateType() + "], " +
    +						"registered with [" + newMetaInfo.getStateType() + "].");
     			}
    -		} catch (Exception ex) {
    -			dispose();
    -			throw ex;
    +
    +			// check compatibility results to determine if state migration is required
    +			CompatibilityResult<N> namespaceCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getNamespaceSerializer(),
    +				null,
    +				restoredMetaInfo.getNamespaceSerializerConfigSnapshot(),
    +				newMetaInfo.getNamespaceSerializer());
    +
    +			CompatibilityResult<S> stateCompatibility = CompatibilityUtil.resolveCompatibilityResult(
    +				restoredMetaInfo.getStateSerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				restoredMetaInfo.getStateSerializerConfigSnapshot(),
    +				newMetaInfo.getStateSerializer());
    +
    +			if (namespaceCompatibility.isRequiresMigration() || stateCompatibility.isRequiresMigration()) {
    +				// TODO state migration currently isn't possible.
    +				throw new StateMigrationException("State migration isn't supported, yet.");
    +			} else {
    +				stateInfo.f1 = newMetaInfo;
    +				return stateInfo.f0;
    +			}
    +		}
    +
    +		byte[] nameBytes = descriptor.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
    +		Preconditions.checkState(!Arrays.equals(RocksDB.DEFAULT_COLUMN_FAMILY, nameBytes),
    +			"The chosen state name 'default' collides with the name of the default column family!");
    +
    +		ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(nameBytes, columnOptions);
    +
    +		final ColumnFamilyHandle columnFamily;
    +
    +		try {
    +			columnFamily = db.createColumnFamily(columnDescriptor);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error creating ColumnFamilyHandle.", e);
     		}
    +
    +		Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<N, S>> tuple =
    +			new Tuple2<>(columnFamily, newMetaInfo);
    +		Map rawAccess = kvStateInformation;
    +		rawAccess.put(descriptor.getName(), tuple);
    +		return columnFamily;
     	}
     
     	@Override
    -	public void notifyCheckpointComplete(long completedCheckpointId) {
    -		synchronized (materializedSstFiles) {
    -			if (completedCheckpointId < lastCompletedCheckpointId) {
    -				return;
    -			}
    +	protected <N, T> InternalValueState<N, T> createValueState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ValueStateDescriptor<T> stateDesc) throws Exception {
     
    -			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -			lastCompletedCheckpointId = completedCheckpointId;
    -		}
    +		return new RocksDBValueState<>(columnFamily, namespaceSerializer,  stateDesc, this);
     	}
     
    -	private void createDB() throws IOException {
    -		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    -		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    -		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	@Override
    +	protected <N, T> InternalListState<N, T> createListState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ListStateDescriptor<T> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
    -	private RocksDB openDB(
    -		String path,
    -		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    -		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
    +	@Override
    +	protected <N, T> InternalReducingState<N, T> createReducingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		ReducingStateDescriptor<T> stateDesc) throws Exception {
     
    -		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    -			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
    +		return new RocksDBReducingState<>(columnFamily, namespaceSerializer,  stateDesc, this);
    +	}
     
    -		// we add the required descriptor for the default CF in last position.
    -		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
    +	@Override
    +	protected <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception {
     
    -		RocksDB dbRef;
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +		return new RocksDBAggregatingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
     
    -		try {
    -			dbRef = RocksDB.open(
    -				Preconditions.checkNotNull(dbOptions),
    -				Preconditions.checkNotNull(path),
    -				columnFamilyDescriptors,
    -				stateColumnFamilyHandles);
    -		} catch (RocksDBException e) {
    -			throw new IOException("Error while opening RocksDB instance.", e);
    -		}
    +	@Override
    +	protected <N, T, ACC> InternalFoldingState<N, T, ACC> createFoldingState(
    +		TypeSerializer<N> namespaceSerializer,
    +		FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
     
    -		// requested + default CF
    -		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    -			"Not all requested column family handles have been created");
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
     
    -		return dbRef;
    +		return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
    +	}
    +
    +	@Override
    +	protected <N, UK, UV> InternalMapState<N, UK, UV> createMapState(
    +		TypeSerializer<N> namespaceSerializer,
    +		MapStateDescriptor<UK, UV> stateDesc) throws Exception {
    +
    +		ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc, namespaceSerializer);
    +
    +		return new RocksDBMapState<>(columnFamily, namespaceSerializer, stateDesc, this);
     	}
     
     	/**
    -	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
    +	 * Only visible for testing, DO NOT USE.
     	 */
    -	static final class RocksDBFullRestoreOperation<K> {
    +	public File getInstanceBasePath() {
    +		return instanceBasePath;
    +	}
     
    -		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
    +	@Override
    +	public boolean supportsAsynchronousSnapshots() {
    +		return true;
    +	}
     
    -		/** Current key-groups state handle from which we restore key-groups. */
    -		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    -		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    -		private FSDataInputStream currentStateHandleInStream;
    -		/** Current data input view that wraps currentStateHandleInStream. */
    -		private DataInputView currentStateHandleInView;
    -		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    -		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    -		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    -		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
    +	@VisibleForTesting
    +	@SuppressWarnings("unchecked")
    +	@Override
    +	public int numStateEntries() {
    +		int count = 0;
     
    -		/**
    -		 * Creates a restore operation object for the given state backend instance.
    -		 *
    -		 * @param rocksDBKeyedStateBackend the state backend into which we restore
    -		 */
    -		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    -			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
    +		for (Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> column : kvStateInformation.values()) {
    +			try (RocksIterator rocksIterator = db.newIterator(column.f0)) {
    +				rocksIterator.seekToFirst();
    +
    +				while (rocksIterator.isValid()) {
    +					count++;
    +					rocksIterator.next();
    +				}
    +			}
    +		}
    +
    +		return count;
    +	}
    +
    +
    +
    +	/**
    +	 * Iterator that merges multiple RocksDB iterators to partition all states into contiguous key-groups.
    +	 * The resulting iteration sequence is ordered by (key-group, kv-state).
    +	 */
    +	@VisibleForTesting
    +	static final class RocksDBMergeIterator implements AutoCloseable {
    +
    +		private final PriorityQueue<MergeIterator> heap;
    +		private final int keyGroupPrefixByteCount;
    +		private boolean newKeyGroup;
    +		private boolean newKVState;
    +		private boolean valid;
    +
    +		private MergeIterator currentSubIterator;
    +
    +		private static final List<Comparator<MergeIterator>> COMPARATORS;
    +
    +		static {
    +			int maxBytes = 4;
    +			COMPARATORS = new ArrayList<>(maxBytes);
    +			for (int i = 0; i < maxBytes; ++i) {
    +				final int currentBytes = i;
    +				COMPARATORS.add(new Comparator<MergeIterator>() {
    +					@Override
    +					public int compare(MergeIterator o1, MergeIterator o2) {
    +						int arrayCmpRes = compareKeyGroupsForByteArrays(
    +							o1.currentKey, o2.currentKey, currentBytes);
    +						return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
    +					}
    +				});
    +			}
    +		}
    +
    +		RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) {
    +			Preconditions.checkNotNull(kvStateIterators);
    +			this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
    +
    +			Comparator<MergeIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount);
    +
    +			if (kvStateIterators.size() > 0) {
    +				PriorityQueue<MergeIterator> iteratorPriorityQueue =
    +					new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
    +
    +				for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
    +					final RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
    +					rocksIterator.seekToFirst();
    +					if (rocksIterator.isValid()) {
    +						iteratorPriorityQueue.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
    +					} else {
    +						IOUtils.closeQuietly(rocksIterator);
    +					}
    +				}
    +
    +				kvStateIterators.clear();
    +
    +				this.heap = iteratorPriorityQueue;
    +				this.valid = !heap.isEmpty();
    +				this.currentSubIterator = heap.poll();
    +			} else {
    +				// creating a PriorityQueue of size 0 results in an exception.
    +				this.heap = null;
    +				this.valid = false;
    +			}
    +
    +			this.newKeyGroup = true;
    +			this.newKVState = true;
     		}
     
     		/**
    -		 * Restores all key-groups data that is referenced by the passed state handles.
    -		 *
    -		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
    +		 * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
    +		 * calls to {@link #next()}.
     		 */
    -		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    -			throws IOException, StateMigrationException, RocksDBException {
    +		public void next() {
    +			newKeyGroup = false;
    +			newKVState = false;
     
    -			rocksDBKeyedStateBackend.createDB();
    +			final RocksIterator rocksIterator = currentSubIterator.getIterator();
    +			rocksIterator.next();
     
    -			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    -				if (keyedStateHandle != null) {
    +			byte[] oldKey = currentSubIterator.getCurrentKey();
    +			if (rocksIterator.isValid()) {
    +				currentSubIterator.currentKey = rocksIterator.key();
     
    -					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    -						throw new IllegalStateException("Unexpected state handle type, " +
    -							"expected: " + KeyGroupsStateHandle.class +
    -							", but found: " + keyedStateHandle.getClass());
    -					}
    -					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    -					restoreKeyGroupsInStateHandle();
    +				if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
    +					heap.offer(currentSubIterator);
    +					currentSubIterator 
    --- End diff --
    
    This should also be remove when canceled is true.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168510973
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamTaskStateInitializerImpl.java ---
    @@ -223,155 +223,110 @@ public StreamOperatorStateContext streamOperatorStateContext(
     
     	protected OperatorStateBackend operatorStateBackend(
     		String operatorIdentifierText,
    -		OperatorSubtaskState operatorSubtaskStateFromJobManager,
    +		PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
     		CloseableRegistry backendCloseableRegistry) throws Exception {
     
    -		//TODO search in local state for a local recovery opportunity.
    +		BackendRestorerProcedure<OperatorStateBackend, OperatorStateHandle> backendRestorer =
    +			new BackendRestorerProcedure<>(
    +				() -> stateBackend.createOperatorStateBackend(environment, operatorIdentifierText),
    +				backendCloseableRegistry);
     
    -		return createOperatorStateBackendFromJobManagerState(
    -			operatorIdentifierText,
    -			operatorSubtaskStateFromJobManager,
    -			backendCloseableRegistry);
    +		return backendRestorer.createAndRestore(prioritizedOperatorSubtaskStates.getPrioritizedManagedOperatorState());
     	}
     
     	protected <K> AbstractKeyedStateBackend<K> keyedStatedBackend(
     		TypeSerializer<K> keySerializer,
     		String operatorIdentifierText,
    -		OperatorSubtaskState operatorSubtaskStateFromJobManager,
    +		PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
     		CloseableRegistry backendCloseableRegistry) throws Exception {
     
     		if (keySerializer == null) {
     			return null;
     		}
     
    -		//TODO search in local state for a local recovery opportunity.
    +		TaskInfo taskInfo = environment.getTaskInfo();
    +
    +		final KeyGroupRange keyGroupRange = KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(
    +			taskInfo.getMaxNumberOfParallelSubtasks(),
    +			taskInfo.getNumberOfParallelSubtasks(),
    +			taskInfo.getIndexOfThisSubtask());
     
    -		return createKeyedStatedBackendFromJobManagerState(
    -			keySerializer,
    -			operatorIdentifierText,
    -			operatorSubtaskStateFromJobManager,
    -			backendCloseableRegistry);
    +		BackendRestorerProcedure<AbstractKeyedStateBackend<K>, KeyedStateHandle> backendRestorer =
    +			new BackendRestorerProcedure<>(
    +				() -> stateBackend.createKeyedStateBackend(
    +					environment,
    +					environment.getJobID(),
    +					operatorIdentifierText,
    +					keySerializer,
    +					taskInfo.getMaxNumberOfParallelSubtasks(),
    +					keyGroupRange,
    +					environment.getTaskKvStateRegistry()),
    +				backendCloseableRegistry);
    +
    +		return backendRestorer.createAndRestore(prioritizedOperatorSubtaskStates.getPrioritizedManagedKeyedState());
     	}
     
     	protected CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs(
    -		OperatorSubtaskState operatorSubtaskStateFromJobManager) {
    +		Iterator<StateObjectCollection<OperatorStateHandle>> restoreStateAlternatives) {
     
    -		if (operatorSubtaskStateFromJobManager != null) {
    +		if (restoreStateAlternatives.hasNext()) {
     
     			final CloseableRegistry closeableRegistry = new CloseableRegistry();
     
    -			Collection<OperatorStateHandle> rawOperatorState =
    -				operatorSubtaskStateFromJobManager.getRawOperatorState();
    -
    -			return new CloseableIterable<StatePartitionStreamProvider>() {
    -				@Override
    -				public void close() throws IOException {
    -					closeableRegistry.close();
    -				}
    -
    -				@Nonnull
    -				@Override
    -				public Iterator<StatePartitionStreamProvider> iterator() {
    -					return new OperatorStateStreamIterator(
    -						DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME,
    -						rawOperatorState.iterator(), closeableRegistry);
    -				}
    -			};
    -		}
    -
    -		return CloseableIterable.empty();
    -	}
    -
    -	protected CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs(
    -		OperatorSubtaskState operatorSubtaskStateFromJobManager) {
    +			Collection<OperatorStateHandle> rawOperatorState = restoreStateAlternatives.next();
    +			// TODO currently this does not support local state recovery, so we expect there is only one handle.
    +			Preconditions.checkState(!restoreStateAlternatives.hasNext());
    --- End diff --
    
    Maybe we could add a proper failure message here.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168169548
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotResult.java ---
    @@ -0,0 +1,83 @@
    +package org.apache.flink.runtime.state;/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162585043
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamTaskStateManagerImplTest.java ---
    @@ -0,0 +1,323 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.api.common.typeutils.TypeSerializer;
    +import org.apache.flink.api.common.typeutils.base.IntSerializer;
    +import org.apache.flink.core.fs.CloseableRegistry;
    +import org.apache.flink.runtime.checkpoint.JobManagerTaskRestore;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.checkpoint.savepoint.CheckpointTestUtils;
    +import org.apache.flink.runtime.execution.Environment;
    +import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
    +import org.apache.flink.runtime.query.TaskKvStateRegistry;
    +import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
    +import org.apache.flink.runtime.state.CheckpointStreamFactory;
    +import org.apache.flink.runtime.state.KeyGroupRange;
    +import org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.OperatorStateBackend;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.StateBackend;
    +import org.apache.flink.runtime.state.StatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.TaskStateManager;
    +import org.apache.flink.runtime.state.TaskStateManagerImplTest;
    +import org.apache.flink.runtime.state.memory.MemoryStateBackend;
    +import org.apache.flink.runtime.taskmanager.TestCheckpointResponder;
    +import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService;
    +import org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService;
    +import org.apache.flink.util.CloseableIterable;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import javax.annotation.Nullable;
    +
    +import java.io.Closeable;
    +import java.io.IOException;
    +import java.util.Collections;
    +import java.util.Random;
    +
    +import static org.mockito.Matchers.any;
    +import static org.mockito.Matchers.anyInt;
    +import static org.mockito.Matchers.eq;
    +import static org.mockito.Mockito.mock;
    +import static org.mockito.Mockito.spy;
    +import static org.mockito.Mockito.verify;
    +import static org.mockito.Mockito.when;
    +
    +/**
    + * Test for {@link StreamTaskStateManagerImpl}.
    + */
    +public class StreamTaskStateManagerImplTest {
    +
    +	@Test
    +	public void testNoRestore() throws Exception {
    +
    +		MemoryStateBackend stateBackend = spy(new MemoryStateBackend(1024));
    +
    +		// No job manager provided state to restore
    +		StreamTaskStateManager 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,
    +			typeSerializer,
    +			closeableRegistry);
    +
    +		verify(stateBackend).createKeyedStateBackend(
    +			any(Environment.class),
    +			any(JobID.class),
    +			any(String.class),
    +			eq(typeSerializer),
    +			anyInt(),
    +			any(KeyGroupRange.class),
    +			any(TaskKvStateRegistry.class));
    +
    +		verify(stateBackend).createOperatorStateBackend(
    +			any(Environment.class),
    +			any(String.class));
    +
    +		verify(stateBackend).createStreamFactory(
    +			any(JobID.class),
    +			any(String.class));
    +
    +		OperatorStateBackend operatorStateBackend = stateContext.operatorStateBackend();
    +		AbstractKeyedStateBackend<?> keyedStateBackend = stateContext.keyedStateBackend();
    +		InternalTimeServiceManager<?, ?> timeServiceManager = stateContext.internalTimerServiceManager();
    +		CheckpointStreamFactory streamFactory = stateContext.checkpointStreamFactory();
    +		CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = stateContext.rawKeyedStateInputs();
    +		CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = stateContext.rawOperatorStateInputs();
    +
    +		Assert.assertEquals(false, stateContext.isRestored());
    +		Assert.assertNotNull(operatorStateBackend);
    --- End diff --
    
    This check is sufficient for this test.


---

[GitHub] flink issue #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on the issue:

    https://github.com/apache/flink/pull/5239
  
    Thanks for going through the general design @StephanEwen ! As we discussed, I agree with your first point. For the second point about RocksDB, this PR already contains an optimized way to deal with incremental local checkpoints that we did not discuss in our review, because I thought it is too much of a low level detail.
    It does not work with duplicating streams. Instead, I introduced a state handle type for a local directory. In fact, I mapped the previous incremental recovery from DFS state also to this new handle type: dfs state is first downloaded and then it also simply becomes a local directory state handle. From there, both incremental recovery paths are identical.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168179062
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TaskStateManagerImplTest.java ---
    @@ -97,6 +104,39 @@ public void testStateReportingAndRetrieving() {
     		Assert.assertNull(taskStateManager.operatorStates(operatorID_3));
     	}
     
    +	/**
    +	 * This tests if the {@link TaskStateManager} properly returns the the subtask local state dir from the
    +	 * corresponding {@link TaskLocalStateStore}.
    +	 */
    +	@Test
    +	public void testForwardingSubtaskLocalStateBaseDirFromLocalStateStore() throws IOException {
    +		JobID jobID = new JobID(42L, 43L);
    +		JobVertexID jobVertexID = new JobVertexID(12L, 34L);
    +		ExecutionAttemptID executionAttemptID = new ExecutionAttemptID(23L, 24L);
    +		TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();
    +
    +		TemporaryFolder tmpFolder = new TemporaryFolder();
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168530422
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TestTaskStateManager.java ---
    @@ -118,12 +121,38 @@ public void reportTaskStateSnapshots(
     		}
     	}
     
    +	@Nonnull
     	@Override
    -	public OperatorSubtaskState operatorStates(OperatorID operatorID) {
    -		TaskStateSnapshot taskStateSnapshot = getLastJobManagerTaskStateSnapshot();
    -		return taskStateSnapshot != null ? taskStateSnapshot.getSubtaskStateByOperatorID(operatorID) : null;
    +	public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {
    +		TaskStateSnapshot jmTaskStateSnapshot = getLastJobManagerTaskStateSnapshot();
    +		TaskStateSnapshot tmTaskStateSnapshot = getLastTaskManagerTaskStateSnapshot();
    +
    +		OperatorSubtaskState jmOpState = null;
    +		List<OperatorSubtaskState> tmStateCollection = null;
    +
    +		if (jmTaskStateSnapshot != null) {
    +			jmOpState = jmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
    +		}
    +
    +		if (tmTaskStateSnapshot != null) {
    +			OperatorSubtaskState tmOpState = tmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
    +			if (tmOpState != null) {
    +				tmStateCollection = Collections.singletonList(tmOpState);
    +			}
    +		}
    +
    +		if (jmOpState == null) {
    +			jmOpState = new OperatorSubtaskState();
    +		}
    +
    +		if (tmStateCollection == null) {
    +			tmStateCollection = Collections.emptyList();
    +		}
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168479432
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -44,11 +44,21 @@
     	private final Map<JobID, Map<JobVertexSubtaskKey, TaskLocalStateStore>> taskStateManagers;
     
     	/** This is the root directory for all local state of this task manager / executor. */
    -	private final File localStateRootDirectory;
    +	private final File[] localStateRootDirectories;
     
    -	public TaskExecutorLocalStateStoresManager(File localStateRootDirectory) {
    +	public TaskExecutorLocalStateStoresManager(File[] localStateRootDirectories) {
     		this.taskStateManagers = new HashMap<>();
    -		this.localStateRootDirectory = Preconditions.checkNotNull(localStateRootDirectory);
    +		this.localStateRootDirectories = Preconditions.checkNotNull(localStateRootDirectories);
    +
    +		for (File localStateRecoveryRootDir : localStateRootDirectories) {
    +			if (!localStateRecoveryRootDir.exists()) {
    +
    +				if (!localStateRecoveryRootDir.mkdirs()) {
    +					throw new IllegalStateException("Could not create root directory for local recovery: " +
    --- End diff --
    
    Isn't this more of an `IOException`. This would make it a bit more explicit what can happen if one instantiates the `TaskExecutorLocalStateStoresManager`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165349012
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -363,1691 +372,1780 @@ public int getKeyGroupPrefixBytes() {
     		final CheckpointStreamFactory streamFactory,
     		CheckpointOptions checkpointOptions) throws Exception {
     
    -		if (checkpointOptions.getCheckpointType() != CheckpointOptions.CheckpointType.SAVEPOINT &&
    -			enableIncrementalCheckpointing) {
    -			return snapshotIncrementally(checkpointId, timestamp, streamFactory);
    -		} else {
    -			return snapshotFully(checkpointId, timestamp, streamFactory);
    -		}
    +		return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     	}
     
    -	private RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotIncrementally(
    -		final long checkpointId,
    -		final long checkpointTimestamp,
    -		final CheckpointStreamFactory checkpointStreamFactory) throws Exception {
    -
    -		if (db == null) {
    -			throw new IOException("RocksDB closed.");
    -		}
    +	@Override
    +	public void restore(StateObjectCollection<KeyedStateHandle> restoreState) throws Exception {
    +		LOG.info("Initializing RocksDB keyed state backend from snapshot.");
     
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at " +
    -					checkpointTimestamp + " . Returning null.");
    -			}
    -			return DoneFuture.nullValue();
    +		if (LOG.isDebugEnabled()) {
    +			LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
     		}
     
    -		final RocksDBIncrementalSnapshotOperation<K> snapshotOperation =
    -			new RocksDBIncrementalSnapshotOperation<>(
    -				this,
    -				checkpointStreamFactory,
    -				checkpointId,
    -				checkpointTimestamp);
    -
    -		snapshotOperation.takeSnapshot();
    -
    -		return new FutureTask<SnapshotResult<KeyedStateHandle>>(
    -			() -> snapshotOperation.materializeSnapshot()
    -		) {
    -			@Override
    -			public boolean cancel(boolean mayInterruptIfRunning) {
    -				snapshotOperation.stop();
    -				return super.cancel(mayInterruptIfRunning);
    -			}
    +		// clear all meta data
    +		kvStateInformation.clear();
    +		restoredKvStateMetaInfos.clear();
     
    -			@Override
    -			protected void done() {
    -				snapshotOperation.releaseResources(isCancelled());
    +		try {
    +			if (restoreState == null || restoreState.isEmpty()) {
    +				createDB();
    +			} else if (restoreState.iterator().next() instanceof IncrementalKeyedStateHandle) {
    +				RocksDBIncrementalRestoreOperation<K> restoreOperation = new RocksDBIncrementalRestoreOperation<>(this);
    +				restoreOperation.restore(restoreState);
    +			} else {
    +				RocksDBFullRestoreOperation<K> restoreOperation = new RocksDBFullRestoreOperation<>(this);
    +				restoreOperation.doRestore(restoreState);
     			}
    -		};
    +		} catch (Exception ex) {
    +			dispose();
    +			throw ex;
    +		}
     	}
     
    -	private RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotFully(
    -		final long checkpointId,
    -		final long timestamp,
    -		final CheckpointStreamFactory streamFactory) throws Exception {
    -
    -		long startTime = System.currentTimeMillis();
    -		final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
    -
    -		final RocksDBFullSnapshotOperation<K> snapshotOperation;
    -
    -		if (kvStateInformation.isEmpty()) {
    -			if (LOG.isDebugEnabled()) {
    -				LOG.debug("Asynchronous RocksDB snapshot performed on empty keyed state at " + timestamp +
    -					" . Returning null.");
    -			}
    +	@Override
    +	public void notifyCheckpointComplete(long completedCheckpointId) {
     
    -			return DoneFuture.nullValue();
    +		if (!enableIncrementalCheckpointing) {
    +			return;
     		}
     
    -		snapshotOperation = new RocksDBFullSnapshotOperation<>(this, streamFactory, snapshotCloseableRegistry);
    -		snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
    -
    -		// implementation of the async IO operation, based on FutureTask
    -		AbstractAsyncCallableWithResources<SnapshotResult<KeyedStateHandle>> ioCallable =
    -			new AbstractAsyncCallableWithResources<SnapshotResult<KeyedStateHandle>>() {
    +		synchronized (materializedSstFiles) {
     
    -				@Override
    -				protected void acquireResources() throws Exception {
    -					cancelStreamRegistry.registerCloseable(snapshotCloseableRegistry);
    -					snapshotOperation.openCheckpointStream();
    -				}
    +			if (completedCheckpointId < lastCompletedCheckpointId) {
    +				return;
    +			}
     
    -				@Override
    -				protected void releaseResources() throws Exception {
    -					closeLocalRegistry();
    -					releaseSnapshotOperationResources();
    -				}
    +			materializedSstFiles.keySet().removeIf(checkpointId -> checkpointId < completedCheckpointId);
     
    -				private void releaseSnapshotOperationResources() {
    -					// hold the db lock while operation on the db to guard us against async db disposal
    -					snapshotOperation.releaseSnapshotResources();
    -				}
    +			lastCompletedCheckpointId = completedCheckpointId;
    +		}
    +	}
     
    -				@Override
    -				protected void stopOperation() throws Exception {
    -					closeLocalRegistry();
    -				}
    +	private void createDB() throws IOException {
    +		List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    +		this.db = openDB(instanceRocksDBPath.getAbsolutePath(), Collections.emptyList(), columnFamilyHandles);
    +		this.defaultColumnFamily = columnFamilyHandles.get(0);
    +	}
     
    -				private void closeLocalRegistry() {
    -					if (cancelStreamRegistry.unregisterCloseable(snapshotCloseableRegistry)) {
    -						try {
    -							snapshotCloseableRegistry.close();
    -						} catch (Exception ex) {
    -							LOG.warn("Error closing local registry", ex);
    -						}
    -					}
    -				}
    +	private RocksDB openDB(
    +		String path,
    +		List<ColumnFamilyDescriptor> stateColumnFamilyDescriptors,
    +		List<ColumnFamilyHandle> stateColumnFamilyHandles) throws IOException {
     
    -				@Override
    -				public SnapshotResult<KeyedStateHandle> performOperation() throws Exception {
    -					long startTime = System.currentTimeMillis();
    +		List<ColumnFamilyDescriptor> columnFamilyDescriptors =
    +			new ArrayList<>(1 + stateColumnFamilyDescriptors.size());
     
    -					if (isStopped()) {
    -						throw new IOException("RocksDB closed.");
    -					}
    +		columnFamilyDescriptors.addAll(stateColumnFamilyDescriptors);
     
    -					snapshotOperation.writeDBSnapshot();
    +		// we add the required descriptor for the default CF in last position.
    +		columnFamilyDescriptors.add(new ColumnFamilyDescriptor(DEFAULT_COLUMN_FAMILY_NAME_BYTES, columnOptions));
     
    -					LOG.info("Asynchronous RocksDB snapshot ({}, asynchronous part) in thread {} took {} ms.",
    -						streamFactory, Thread.currentThread(), (System.currentTimeMillis() - startTime));
    +		RocksDB dbRef;
     
    -					KeyGroupsStateHandle stateHandle = snapshotOperation.getSnapshotResultStateHandle();
    -					return new SnapshotResult<>(stateHandle, null);
    -				}
    -			};
    +		try {
    +			dbRef = RocksDB.open(
    +				Preconditions.checkNotNull(dbOptions),
    +				Preconditions.checkNotNull(path),
    +				columnFamilyDescriptors,
    +				stateColumnFamilyHandles);
    +		} catch (RocksDBException e) {
    +			throw new IOException("Error while opening RocksDB instance.", e);
    +		}
     
    -		LOG.info("Asynchronous RocksDB snapshot (" + streamFactory + ", synchronous part) in thread " +
    -			Thread.currentThread() + " took " + (System.currentTimeMillis() - startTime) + " ms.");
    +		// requested + default CF
    +		Preconditions.checkState(1 + stateColumnFamilyDescriptors.size() == stateColumnFamilyHandles.size(),
    +			"Not all requested column family handles have been created");
     
    -		return AsyncStoppableTaskWithCallback.from(ioCallable);
    +		return dbRef;
     	}
     
     	/**
    -	 * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a full snapshot.
     	 */
    -	static final class RocksDBFullSnapshotOperation<K> {
    +	private static final class RocksDBFullRestoreOperation<K> {
     
    -		static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
    -		static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
    -
    -		private final RocksDBKeyedStateBackend<K> stateBackend;
    -		private final KeyGroupRangeOffsets keyGroupRangeOffsets;
    -		private final CheckpointStreamFactory checkpointStreamFactory;
    -		private final CloseableRegistry snapshotCloseableRegistry;
    -		private final ResourceGuard.Lease dbLease;
    -
    -		private long checkpointId;
    -		private long checkpointTimeStamp;
    -
    -		private Snapshot snapshot;
    -		private ReadOptions readOptions;
    -		private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
    -
    -		private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
    -		private DataOutputView outputView;
    -
    -		RocksDBFullSnapshotOperation(
    -			RocksDBKeyedStateBackend<K> stateBackend,
    -			CheckpointStreamFactory checkpointStreamFactory,
    -			CloseableRegistry registry) throws IOException {
    +		private final RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend;
     
    -			this.stateBackend = stateBackend;
    -			this.checkpointStreamFactory = checkpointStreamFactory;
    -			this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
    -			this.snapshotCloseableRegistry = registry;
    -			this.dbLease = this.stateBackend.rocksDBResourceGuard.acquireResource();
    -		}
    +		/** Current key-groups state handle from which we restore key-groups. */
    +		private KeyGroupsStateHandle currentKeyGroupsStateHandle;
    +		/** Current input stream we obtained from currentKeyGroupsStateHandle. */
    +		private FSDataInputStream currentStateHandleInStream;
    +		/** Current data input view that wraps currentStateHandleInStream. */
    +		private DataInputView currentStateHandleInView;
    +		/** Current list of ColumnFamilyHandles for all column families we restore from currentKeyGroupsStateHandle. */
    +		private List<ColumnFamilyHandle> currentStateHandleKVStateColumnFamilies;
    +		/** The compression decorator that was used for writing the state, as determined by the meta data. */
    +		private StreamCompressionDecorator keygroupStreamCompressionDecorator;
     
     		/**
    -		 * 1) Create a snapshot object from RocksDB.
    +		 * Creates a restore operation object for the given state backend instance.
     		 *
    -		 * @param checkpointId id of the checkpoint for which we take the snapshot
    -		 * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
    +		 * @param rocksDBKeyedStateBackend the state backend into which we restore
     		 */
    -		public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) {
    -			Preconditions.checkArgument(snapshot == null, "Only one ongoing snapshot allowed!");
    -			this.kvStateIterators = new ArrayList<>(stateBackend.kvStateInformation.size());
    -			this.checkpointId = checkpointId;
    -			this.checkpointTimeStamp = checkpointTimeStamp;
    -			this.snapshot = stateBackend.db.getSnapshot();
    +		public RocksDBFullRestoreOperation(RocksDBKeyedStateBackend<K> rocksDBKeyedStateBackend) {
    +			this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
     		}
     
     		/**
    -		 * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
    +		 * Restores all key-groups data that is referenced by the passed state handles.
     		 *
    -		 * @throws Exception
    +		 * @param keyedStateHandles List of all key groups state handles that shall be restored.
     		 */
    -		public void openCheckpointStream() throws Exception {
    -			Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
    -			outStream = checkpointStreamFactory.createCheckpointStateOutputStream(checkpointId, checkpointTimeStamp);
    -			snapshotCloseableRegistry.registerCloseable(outStream);
    -			outputView = new DataOutputViewStreamWrapper(outStream);
    -		}
    +		public void doRestore(Collection<KeyedStateHandle> keyedStateHandles)
    +			throws IOException, StateMigrationException, RocksDBException {
     
    -		/**
    -		 * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
    -		 *
    -		 * @throws IOException
    -		 */
    -		public void writeDBSnapshot() throws IOException, InterruptedException {
    +			rocksDBKeyedStateBackend.createDB();
     
    -			if (null == snapshot) {
    -				throw new IOException("No snapshot available. Might be released due to cancellation.");
    -			}
    +			for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
    +				if (keyedStateHandle != null) {
     
    -			Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
    -			writeKVStateMetaData();
    -			writeKVStateData();
    +					if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
    +						throw new IllegalStateException("Unexpected state handle type, " +
    +							"expected: " + KeyGroupsStateHandle.class +
    +							", but found: " + keyedStateHandle.getClass());
    +					}
    +					this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
    +					restoreKeyGroupsInStateHandle();
    +				}
    +			}
     		}
     
     		/**
    -		 * 4) Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
    -		 *
    -		 * @return state handle to the completed snapshot
    +		 * Restore one key groups state handle.
     		 */
    -		public KeyGroupsStateHandle getSnapshotResultStateHandle() throws IOException {
    -
    -			if (snapshotCloseableRegistry.unregisterCloseable(outStream)) {
    -
    -				StreamStateHandle stateHandle = outStream.closeAndGetHandle();
    -				outStream = null;
    -
    -				if (stateHandle != null) {
    -					return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
    +		private void restoreKeyGroupsInStateHandle()
    +			throws IOException, StateMigrationException, RocksDBException {
    +			try {
    +				currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
    +				rocksDBKeyedStateBackend.cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
    +				currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
    +				restoreKVStateMetaData();
    +				restoreKVStateData();
    +			} finally {
    +				if (rocksDBKeyedStateBackend.cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
    +					IOUtils.closeQuietly(currentStateHandleInStream);
     				}
     			}
    -			return null;
     		}
     
     		/**
    -		 * 5) Release the snapshot object for RocksDB and clean up.
    +		 * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws ClassNotFoundException
    +		 * @throws RocksDBException
     		 */
    -		public void releaseSnapshotResources() {
    +		private void restoreKVStateMetaData() throws IOException, StateMigrationException, RocksDBException {
     
    -			outStream = null;
    +			KeyedBackendSerializationProxy<K> serializationProxy =
    +				new KeyedBackendSerializationProxy<>(rocksDBKeyedStateBackend.userCodeClassLoader);
     
    -			if (null != kvStateIterators) {
    -				for (Tuple2<RocksIterator, Integer> kvStateIterator : kvStateIterators) {
    -					IOUtils.closeQuietly(kvStateIterator.f0);
    -				}
    -				kvStateIterators = null;
    -			}
    +			serializationProxy.read(currentStateHandleInView);
     
    -			if (null != snapshot) {
    -				if (null != stateBackend.db) {
    -					stateBackend.db.releaseSnapshot(snapshot);
    -				}
    -				IOUtils.closeQuietly(snapshot);
    -				snapshot = null;
    -			}
    +			// check for key serializer compatibility; this also reconfigures the
    +			// key serializer to be compatible, if it is required and is possible
    +			if (CompatibilityUtil.resolveCompatibilityResult(
    +				serializationProxy.getKeySerializer(),
    +				UnloadableDummyTypeSerializer.class,
    +				serializationProxy.getKeySerializerConfigSnapshot(),
    +				rocksDBKeyedStateBackend.keySerializer)
    +				.isRequiresMigration()) {
     
    -			if (null != readOptions) {
    -				IOUtils.closeQuietly(readOptions);
    -				readOptions = null;
    +				// TODO replace with state migration; note that key hash codes need to remain the same after migration
    +				throw new StateMigrationException("The new key serializer is not compatible to read previous keys. " +
    +					"Aborting now since state migration is currently not available");
     			}
     
    -			this.dbLease.close();
    -		}
    -
    -		private void writeKVStateMetaData() throws IOException {
    +			this.keygroupStreamCompressionDecorator = serializationProxy.isUsingKeyGroupCompression() ?
    +				SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE;
     
    -			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> metaInfoSnapshots =
    -				new ArrayList<>(stateBackend.kvStateInformation.size());
    +			List<RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?>> restoredMetaInfos =
    +				serializationProxy.getStateMetaInfoSnapshots();
    +			currentStateHandleKVStateColumnFamilies = new ArrayList<>(restoredMetaInfos.size());
    +			//rocksDBKeyedStateBackend.restoredKvStateMetaInfos = new HashMap<>(restoredMetaInfos.size());
     
    -			int kvStateId = 0;
    -			for (Map.Entry<String, Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>>> column :
    -				stateBackend.kvStateInformation.entrySet()) {
    +			for (RegisteredKeyedBackendStateMetaInfo.Snapshot<?, ?> restoredMetaInfo : restoredMetaInfos) {
     
    -				metaInfoSnapshots.add(column.getValue().f1.snapshot());
    +				Tuple2<ColumnFamilyHandle, RegisteredKeyedBackendStateMetaInfo<?, ?>> registeredColumn =
    +					rocksDBKeyedStateBackend.kvStateInformation.get(restoredMetaInfo.getName());
     
    -				//retrieve iterator for this k/v states
    -				readOptions = new ReadOptions();
    -				readOptions.setSnapshot(snapshot);
    +				if (registeredColumn == null) {
    +					byte[] nameBytes = restoredMetaInfo.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
     
    -				kvStateIterators.add(
    -					new Tuple2<>(stateBackend.db.newIterator(column.getValue().f0, readOptions), kvStateId));
    +					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
    +						nameBytes,
    +						rocksDBKeyedStateBackend.columnOptions);
     
    -				++kvStateId;
    -			}
    +					RegisteredKeyedBackendStateMetaInfo<?, ?> stateMetaInfo =
    +						new RegisteredKeyedBackendStateMetaInfo<>(
    +							restoredMetaInfo.getStateType(),
    +							restoredMetaInfo.getName(),
    +							restoredMetaInfo.getNamespaceSerializer(),
    +							restoredMetaInfo.getStateSerializer());
     
    -			KeyedBackendSerializationProxy<K> serializationProxy =
    -				new KeyedBackendSerializationProxy<>(
    -					stateBackend.getKeySerializer(),
    -					metaInfoSnapshots,
    -					!Objects.equals(UncompressedStreamCompressionDecorator.INSTANCE, stateBackend.keyGroupCompressionDecorator));
    +					rocksDBKeyedStateBackend.restoredKvStateMetaInfos.put(restoredMetaInfo.getName(), restoredMetaInfo);
     
    -			serializationProxy.write(outputView);
    -		}
    +					ColumnFamilyHandle columnFamily = rocksDBKeyedStateBackend.db.createColumnFamily(columnFamilyDescriptor);
     
    -		private void writeKVStateData() throws IOException, InterruptedException {
    +					registeredColumn = new Tuple2<>(columnFamily, stateMetaInfo);
    +					rocksDBKeyedStateBackend.kvStateInformation.put(stateMetaInfo.getName(), registeredColumn);
     
    -			byte[] previousKey = null;
    -			byte[] previousValue = null;
    -			OutputStream kgOutStream = null;
    -			DataOutputView kgOutView = null;
    +				} else {
    +					// TODO with eager state registration in place, check here for serializer migration strategies
    +				}
    +				currentStateHandleKVStateColumnFamilies.add(registeredColumn.f0);
    +			}
    +		}
     
    -			try {
    -				// Here we transfer ownership of RocksIterators to the RocksDBMergeIterator
    -				try (RocksDBMergeIterator mergeIterator = new RocksDBMergeIterator(
    -					kvStateIterators, stateBackend.keyGroupPrefixBytes)) {
    +		/**
    +		 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state handle.
    +		 *
    +		 * @throws IOException
    +		 * @throws RocksDBException
    +		 */
    +		private void restoreKVStateData() throws IOException, RocksDBException {
    +			//for all key-groups in the current state handle...
    +			for (Tuple2<Integer, Long> keyGroupOffset : currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
    +				int keyGroup = keyGroupOffset.f0;
     
    -					// handover complete, null out to prevent double close
    -					kvStateIterators = null;
    +				// Check that restored key groups all belong to the backend
    +				Preconditions.checkState(rocksDBKeyedStateBackend.getKeyGroupRange().contains(keyGroup),
    +					"The key group must belong to the backend");
     
    -					//preamble: setup with first key-group as our lookahead
    -					if (mergeIterator.isValid()) {
    -						//begin first key-group by recording the offset
    -						keyGroupRangeOffsets.setKeyGroupOffset(mergeIterator.keyGroup(), outStream.getPos());
    -						//write the k/v-state id as metadata
    -						kgOutStream = stateBackend.keyGroupCompressionDecorator.decorateWithCompression(outStream);
    -						kgOutView = new DataOutputViewStreamWrapper(kgOutStream);
    +				long offset = keyGroupOffset.f1;
    +				//not empty key-group?
    +				if (0L != offset) {
    +					currentStateHandleInStream.seek(offset);
    +					try (InputStream compressedKgIn = keygroupStreamCompressionDecorator.decorateWithCompression(currentStateHandleInStream)) {
    +						DataInputViewStreamWrapper compressedKgInputView = new DataInputViewStreamWrapper(compressedKgIn);
     						//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    -						kgOutView.writeShort(mergeIterator.kvStateId());
    -						previousKey = mergeIterator.key();
    -						previousValue = mergeIterator.value();
    -						mergeIterator.next();
    +						int kvStateId = compressedKgInputView.readShort();
    +						ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +						//insert all k/v pairs into DB
    +						boolean keyGroupHasMoreKeys = true;
    +						while (keyGroupHasMoreKeys) {
    +							byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(compressedKgInputView);
    +							if (RocksDBFullSnapshotOperation.hasMetaDataFollowsFlag(key)) {
    +								//clear the signal bit in the key to make it ready for insertion again
    +								RocksDBFullSnapshotOperation.clearMetaDataFollowsFlag(key);
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +								//TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
    +								kvStateId = RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK
    +									& compressedKgInputView.readShort();
    +								if (RocksDBFullSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
    +									keyGroupHasMoreKeys = false;
    +								} else {
    +									handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
    +								}
    +							} else {
    +								rocksDBKeyedStateBackend.db.put(handle, key, value);
    +							}
    +						}
     					}
    +				}
    +			}
    +		}
    +	}
     
    -					//main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
    -					while (mergeIterator.isValid()) {
    +	/**
    +	 * Encapsulates the process of restoring a RocksDBKeyedStateBackend from an incremental snapshot.
    +	 */
    +	private static class RocksDBIncrementalRestoreOperation<T> {
    --- End diff --
    
    Would it actually make sense to give these classes their own files? The `RocksDBKeyedStateBackend` file is 2100 lines long which makes it quite hard to handle.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168189852
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStoreImpl.java ---
    @@ -0,0 +1,298 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
    +
    +import java.io.File;
    +import java.io.IOException;
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.HashMap;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
    +
    +/**
    + * Main implementation of a {@link TaskLocalStateStore}.
    + */
    +public class TaskLocalStateStoreImpl implements TaskLocalStateStore {
    +
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStoreImpl.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	@VisibleForTesting
    +	static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    --- End diff --
    
    - What is `MAX_RETAINED_SNAPSHOTS` used for?  I found all the checkpoints that satisfied `checkpointId < confirmedCheckpointId ` will be removed in `confirmCheckpoint()`.
    - Is it's value can be configured?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165317571
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java ---
    @@ -916,57 +916,78 @@ private void handleExecutionException(Exception e) {
     				CheckpointingOperation.AsynCheckpointState.COMPLETED,
     				CheckpointingOperation.AsynCheckpointState.RUNNING);
     
    -			try {
    -				cleanup();
    -			} catch (Exception cleanupException) {
    -				e.addSuppressed(cleanupException);
    -			}
    +			if (asyncCheckpointState.compareAndSet(
    +				CheckpointingOperation.AsynCheckpointState.RUNNING,
    +				CheckpointingOperation.AsynCheckpointState.DISCARDED)) {
    --- End diff --
    
    This could be `asyncCheckpointState.compareAndSet(currentAsyncCheckpointState, DISCARDED)`. Then we could remove the compare and set above this if condition.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r167356422
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    --- End diff --
    
    to recover~y~


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165107698
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotResult.java ---
    @@ -0,0 +1,83 @@
    +package org.apache.flink.runtime.state;/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +import org.apache.flink.util.ExceptionUtils;
    +
    +/**
    + *
    + */
    +public class SnapshotResult<T extends StateObject> implements StateObject {
    --- End diff --
    
    I'm wondering whether it wouldn't make sense to have a `SnapshotResult` and a subclass `SnapshotResultWithLocalSnapshot`. The latter has an additional field for the `taskLocalSnapshot`. That way it's much clearer how to initialize these structures. Ideally the arguments are both not null.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168742457
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStoreImpl.java ---
    @@ -0,0 +1,298 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
    +
    +import java.io.File;
    +import java.io.IOException;
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.HashMap;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
    +
    +/**
    + * Main implementation of a {@link TaskLocalStateStore}.
    + */
    +public class TaskLocalStateStoreImpl implements TaskLocalStateStore {
    +
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStoreImpl.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	@VisibleForTesting
    +	static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    +
    +	/** JobID from the owning subtask. */
    +	@Nonnull
    +	private final JobID jobID;
    +
    +	/** AllocationID of the owning slot. */
    +	@Nonnull
    +	private final AllocationID allocationID;
    +
    +	/** JobVertexID of the owning subtask. */
    +	@Nonnull
    +	private final JobVertexID jobVertexID;
    +
    +	/** Subtask index of the owning subtask. */
    +	@Nonnegative
    +	private final int subtaskIndex;
    +
    +	/** The configured mode for local recovery. */
    +	@Nonnull
    +	private final LocalRecoveryConfig localRecoveryConfig;
    +
    +	/** Executor that runs the discarding of released state objects. */
    +	@Nonnull
    +	private final Executor discardExecutor;
    +
    +	/** Lock for synchronisation on the storage map and the discarded status. */
    +	@Nonnull
    +	private final Object lock;
    +
    +	/** Status flag if this store was already discarded. */
    +	@GuardedBy("lock")
    +	private boolean discarded;
    +
    +	/** Maps checkpoint ids to local TaskStateSnapshots. */
    +	@Nonnull
    +	@GuardedBy("lock")
    +	private final SortedMap<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
    +	public TaskLocalStateStoreImpl(
    +		@Nonnull JobID jobID,
    +		@Nonnull AllocationID allocationID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex,
    +		@Nonnull LocalRecoveryConfig localRecoveryConfig,
    +		@Nonnull Executor discardExecutor) {
    +
    +		this.jobID = jobID;
    +		this.allocationID = allocationID;
    +		this.jobVertexID = jobVertexID;
    +		this.subtaskIndex = subtaskIndex;
    +		this.discardExecutor = discardExecutor;
    +		this.lock = new Object();
    +		this.storedTaskStateByCheckpointID = new TreeMap<>();
    +		this.discarded = false;
    +		this.localRecoveryConfig = localRecoveryConfig;
    +	}
    +
    +	@Override
    +	public void storeLocalState(
    +		@Nonnegative long checkpointId,
    +		@Nullable TaskStateSnapshot localState) {
    +
    +		if (localState == null) {
    +			localState = NULL_DUMMY;
    +		}
    +
    +		LOG.info("Storing local state for checkpoint {}.", checkpointId);
    +		LOG.debug("Local state for checkpoint {} is {}.", checkpointId, localState);
    +
    +		Map<Long, TaskStateSnapshot> toDiscard = new HashMap<>(MAX_RETAINED_SNAPSHOTS);
    +
    +		synchronized (lock) {
    +			if (discarded) {
    +				// we ignore late stores and simply discard the state.
    +				toDiscard.put(checkpointId, localState);
    +			} else {
    +				TaskStateSnapshot previous =
    +					storedTaskStateByCheckpointID.put(checkpointId, localState);
    +
    +				if (previous != null) {
    +					toDiscard.put(checkpointId, previous);
    +				}
    +
    +				// remove from history.
    +				while (storedTaskStateByCheckpointID.size() > MAX_RETAINED_SNAPSHOTS) {
    +					Long removeCheckpointID = storedTaskStateByCheckpointID.firstKey();
    --- End diff --
    
    Aha, I think simply removed this pruning will also lead to problematic.... Imagine that, if checkpoint always completed in this subtask but failed on global, this will lead to resources leak. I think this can be fixed by introduce a scheduled `canceller` for clear these `pending completed checkpoint` just like what we have done in `CheckpointCoordinator`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165343021
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/LocalStateForwardingTest.java ---
    @@ -156,30 +157,38 @@ public void acknowledgeCheckpoint(
     			}
     		};
     
    -		TaskLocalStateStore taskLocalStateStore = new TaskLocalStateStore(jobID, jobVertexID, subtaskIdx) {
    -			@Override
    -			public void storeLocalState(
    -				@Nonnull CheckpointMetaData checkpointMetaData,
    -				@Nullable TaskStateSnapshot localState) {
    -
    -				Assert.assertEquals(tm, localState);
    -				tmReported.set(true);
    -			}
    -		};
    +		TemporaryFolder temporaryFolder = new TemporaryFolder();
     
    -		TaskStateManagerImpl taskStateManager =
    -			new TaskStateManagerImpl(
    -				jobID,
    -				executionAttemptID,
    -				taskLocalStateStore,
    -				null,
    -				checkpointResponder);
    -
    -		taskStateManager.reportTaskStateSnapshots(
    -			checkpointMetaData,
    -			checkpointMetrics,
    -			jm,
    -			tm);
    +		try {
    +			TaskLocalStateStore taskLocalStateStore =
    +				new TaskLocalStateStore(jobID, jobVertexID, subtaskIdx, temporaryFolder.newFolder()) {
    +					@Override
    +					public void storeLocalState(
    +						@Nonnull CheckpointMetaData checkpointMetaData,
    +						@Nullable TaskStateSnapshot localState) {
    +
    +						Assert.assertEquals(tm, localState);
    +						tmReported.set(true);
    +					}
    +				};
    +
    +			TaskStateManagerImpl taskStateManager =
    +				new TaskStateManagerImpl(
    +					jobID,
    +					executionAttemptID,
    +					taskLocalStateStore,
    +					null,
    +					checkpointResponder);
    +
    +			taskStateManager.reportTaskStateSnapshots(
    +				checkpointMetaData,
    +				checkpointMetrics,
    +				jm,
    +				tm);
    +		} catch (Exception ex) {
    +			temporaryFolder.delete();
    +			throw new RuntimeException(ex);
    --- End diff --
    
    Why do we throw a `RuntimeException`?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168490613
  
    --- Diff: flink-core/src/main/java/org/apache/flink/util/ThrowingSupplier.java ---
    @@ -0,0 +1,38 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.util;
    +
    +/**
    + * This interface is basically Java's {@link java.util.function.Supplier} interface enhanced with the ability to throw
    + * an exception.
    + *
    + * @param <OUT> type of the supplied result.
    + * @param <T>   type of {@link Throwable}.
    + */
    +@FunctionalInterface
    +public interface ThrowingSupplier<OUT, T extends Throwable> {
    +
    +	/**
    +	 * Gets a result.
    +	 *
    +	 * @return a result.
    +	 * @throws T on errors.
    +	 */
    +	OUT get() throws T, Exception;
    +}
    --- End diff --
    
    Sure, I guess this class did not yet exists when I implement mine.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162583016
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java ---
    @@ -145,8 +145,8 @@
     
     	private volatile Throwable failureCause;          // once assigned, never changes
     
    -	/** The handle to the state that the task gets on restore */
    -	private volatile TaskStateSnapshot taskState;
    +	/** Information to restore the task onr recovery, such as checkpoint id and task state snapshot */
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168499407
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java ---
    @@ -727,19 +729,17 @@ static void resetRocksDBLoadedFlag() throws Exception {
     	/**
     	 * This class encapsulates the configuration for local recovery of this backend.
     	 */
    -	public static final class LocalRecoveryConfig implements Serializable {
    +	public static final class LocalRecoveryConfig extends LocalRecoveryConfigBase {
     
    -		private static final long serialVersionUID = 1L;
    --- End diff --
    
    Changed in later commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168475165
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamWithResultProvider.java ---
    @@ -0,0 +1,175 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.state.filesystem.FixFileFsStateOutputStream;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import java.io.Closeable;
    +import java.io.File;
    +import java.io.IOException;
    +
    +/**
    + * Interface that provides access to a CheckpointStateOutputStream and a method to provide the {@link SnapshotResult}.
    + * This abstracts from different ways that a result is obtained from checkpoint output streams.
    + */
    +public interface CheckpointStreamWithResultProvider extends Closeable {
    +
    +	SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +		KeyGroupRangeOffsets kgOffs) throws IOException;
    +
    +	CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream();
    +
    +	@Override
    +	default void close() throws IOException {
    +		CheckpointStreamFactory.CheckpointStateOutputStream outputStream = getCheckpointOutputStream();
    +		if (outputStream != null) {
    +			outputStream.close();
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that only creates the
    +	 * primary/remote/jm-owned state.
    +	 */
    +	class PrimaryStreamOnly implements CheckpointStreamWithResultProvider {
    +
    +		private final CheckpointStreamFactory.CheckpointStateOutputStream outputStream;
    +
    +		public PrimaryStreamOnly(CheckpointStreamFactory.CheckpointStateOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets kgOffs) throws IOException {
    +
    +			StreamStateHandle streamStateHandle = outputStream.closeAndGetHandle();
    +			KeyGroupsStateHandle primaryStateHandle = new KeyGroupsStateHandle(kgOffs, streamStateHandle);
    +			return new SnapshotResult<>(primaryStateHandle, null);
    +		}
    +
    +		@Override
    +		public CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that creates both, the
    +	 * primary/remote/jm-owned state and the secondary/local/tm-owned state.
    +	 */
    +	class PrimaryAndSecondaryStream implements CheckpointStreamWithResultProvider {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		private final DuplicatingCheckpointOutputStream outputStream;
    +
    +		public PrimaryAndSecondaryStream(DuplicatingCheckpointOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets offsets) throws IOException {
    +
    +			final StreamStateHandle primaryStreamStateHandle;
    +
    +			try {
    +				primaryStreamStateHandle = outputStream.closeAndGetPrimaryHandle();
    +			} catch (IOException primaryEx) {
    +				try {
    +					outputStream.close();
    +				} catch (IOException closeEx) {
    +					primaryEx = ExceptionUtils.firstOrSuppressed(closeEx, primaryEx);
    +				}
    +				throw primaryEx;
    +			}
    +
    +			StreamStateHandle secondaryStreamStateHandle = null;
    +
    +			try {
    +				secondaryStreamStateHandle = outputStream.closeAndGetSecondaryHandle();
    +			} catch (IOException secondaryEx) {
    +				LOG.warn("Exception from secondary/local checkpoint stream.", secondaryEx);
    +			}
    +
    +			if (primaryStreamStateHandle != null) {
    +
    +				final KeyGroupsStateHandle primaryKeyGroupsStateHandle =
    +					new KeyGroupsStateHandle(offsets, primaryStreamStateHandle);
    +
    +				final KeyGroupsStateHandle secondaryKeyGroupsStateHandle = secondaryStreamStateHandle != null ?
    +					new KeyGroupsStateHandle(offsets, secondaryStreamStateHandle) :
    +					null;
    +
    +				return new SnapshotResult<>(primaryKeyGroupsStateHandle, secondaryKeyGroupsStateHandle);
    +			}
    +
    +			return null;
    +		}
    +
    +		@Override
    +		public DuplicatingCheckpointOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	class Factory {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		public CheckpointStreamWithResultProvider create(
    +			long checkpointId,
    +			long timestamp,
    +			CheckpointStreamFactory primaryStreamFactory,
    +			LocalRecoveryDirectoryProvider secondaryStreamDirProvider) throws Exception {
    +
    +			CheckpointStreamFactory.CheckpointStateOutputStream primaryOut =
    +				primaryStreamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
    +
    +			CheckpointStreamFactory.CheckpointStateOutputStream secondaryOut;
    +
    +			if (secondaryStreamDirProvider != null) {
    +				try {
    +					File currentLocalRecoveryRoot = secondaryStreamDirProvider.nextRootDirectory();
    +					File outFile = new File(
    +						currentLocalRecoveryRoot,
    +						secondaryStreamDirProvider.specificFileForCheckpointId(checkpointId));
    +					Path outPath = new Path(outFile.toURI());
    +					secondaryOut = new FixFileFsStateOutputStream(outPath.getFileSystem(), outPath);
    +					final DuplicatingCheckpointOutputStream effectiveOutStream =
    +						new DuplicatingCheckpointOutputStream(primaryOut, secondaryOut);
    +
    +					return new CheckpointStreamWithResultProvider.PrimaryAndSecondaryStream(effectiveOutStream);
    --- End diff --
    
    Wouldn't it make more sense to pass in the primary and secondary stream given the method name? Internally one would then create the `DuplicatingCheckpointOutputStream`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r167357375
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    +
    +- RocksDBStateBackend: task-local recovery is supported for keyed state. For *full checkpoints*, state is duplicated to a local file. This can introduce additional write costs
    +and occupies local disk space. For *incremental snapshots*, the local state is based on RocksDB's native checkpointing mechanism. This mechanism is also used as the first step
    --- End diff --
    
    you mean 'incremental checkpoint' rather than 'incremental snapshot'? I haven't heard of the latter term yet 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168498267
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/LocalRecoveryDirectoryProviderTest.java ---
    @@ -0,0 +1,95 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.junit.After;
    +import org.junit.Assert;
    +import org.junit.Before;
    +import org.junit.Test;
    +import org.junit.rules.TemporaryFolder;
    +
    +import java.io.File;
    +import java.io.IOException;
    +
    +public class LocalRecoveryDirectoryProviderTest {
    +
    +	private static final String SPECIFIC_PATH = "specificPath";
    +
    +	private TemporaryFolder tmpFolder;
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168474584
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamWithResultProvider.java ---
    @@ -0,0 +1,175 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.state.filesystem.FixFileFsStateOutputStream;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import java.io.Closeable;
    +import java.io.File;
    +import java.io.IOException;
    +
    +/**
    + * Interface that provides access to a CheckpointStateOutputStream and a method to provide the {@link SnapshotResult}.
    + * This abstracts from different ways that a result is obtained from checkpoint output streams.
    + */
    +public interface CheckpointStreamWithResultProvider extends Closeable {
    +
    +	SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +		KeyGroupRangeOffsets kgOffs) throws IOException;
    +
    +	CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream();
    +
    +	@Override
    +	default void close() throws IOException {
    +		CheckpointStreamFactory.CheckpointStateOutputStream outputStream = getCheckpointOutputStream();
    +		if (outputStream != null) {
    +			outputStream.close();
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that only creates the
    +	 * primary/remote/jm-owned state.
    +	 */
    +	class PrimaryStreamOnly implements CheckpointStreamWithResultProvider {
    +
    +		private final CheckpointStreamFactory.CheckpointStateOutputStream outputStream;
    +
    +		public PrimaryStreamOnly(CheckpointStreamFactory.CheckpointStateOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets kgOffs) throws IOException {
    +
    +			StreamStateHandle streamStateHandle = outputStream.closeAndGetHandle();
    +			KeyGroupsStateHandle primaryStateHandle = new KeyGroupsStateHandle(kgOffs, streamStateHandle);
    +			return new SnapshotResult<>(primaryStateHandle, null);
    +		}
    +
    +		@Override
    +		public CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that creates both, the
    +	 * primary/remote/jm-owned state and the secondary/local/tm-owned state.
    +	 */
    +	class PrimaryAndSecondaryStream implements CheckpointStreamWithResultProvider {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		private final DuplicatingCheckpointOutputStream outputStream;
    +
    +		public PrimaryAndSecondaryStream(DuplicatingCheckpointOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets offsets) throws IOException {
    +
    +			final StreamStateHandle primaryStreamStateHandle;
    +
    +			try {
    +				primaryStreamStateHandle = outputStream.closeAndGetPrimaryHandle();
    +			} catch (IOException primaryEx) {
    +				try {
    +					outputStream.close();
    +				} catch (IOException closeEx) {
    +					primaryEx = ExceptionUtils.firstOrSuppressed(closeEx, primaryEx);
    +				}
    +				throw primaryEx;
    +			}
    +
    +			StreamStateHandle secondaryStreamStateHandle = null;
    +
    +			try {
    +				secondaryStreamStateHandle = outputStream.closeAndGetSecondaryHandle();
    +			} catch (IOException secondaryEx) {
    +				LOG.warn("Exception from secondary/local checkpoint stream.", secondaryEx);
    +			}
    +
    +			if (primaryStreamStateHandle != null) {
    +
    +				final KeyGroupsStateHandle primaryKeyGroupsStateHandle =
    +					new KeyGroupsStateHandle(offsets, primaryStreamStateHandle);
    +
    +				final KeyGroupsStateHandle secondaryKeyGroupsStateHandle = secondaryStreamStateHandle != null ?
    +					new KeyGroupsStateHandle(offsets, secondaryStreamStateHandle) :
    +					null;
    +
    +				return new SnapshotResult<>(primaryKeyGroupsStateHandle, secondaryKeyGroupsStateHandle);
    +			}
    +
    +			return null;
    +		}
    +
    +		@Override
    +		public DuplicatingCheckpointOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	class Factory {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		public CheckpointStreamWithResultProvider create(
    +			long checkpointId,
    +			long timestamp,
    +			CheckpointStreamFactory primaryStreamFactory,
    +			LocalRecoveryDirectoryProvider secondaryStreamDirProvider) throws Exception {
    +
    +			CheckpointStreamFactory.CheckpointStateOutputStream primaryOut =
    +				primaryStreamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
    +
    +			CheckpointStreamFactory.CheckpointStateOutputStream secondaryOut;
    --- End diff --
    
    Wrong scope of variable.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165275719
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/StateSnapshotContextSynchronousImpl.java ---
    @@ -109,14 +109,20 @@ public OperatorStateCheckpointOutputStream getRawOperatorStateOutput() throws Ex
     		return operatorStateCheckpointOutputStream;
     	}
     
    -	public RunnableFuture<KeyedStateHandle> getKeyedStateStreamFuture() throws IOException {
    -		KeyGroupsStateHandle keyGroupsStateHandle = closeAndUnregisterStreamToObtainStateHandle(keyedStateCheckpointOutputStream);
    -		return new DoneFuture<KeyedStateHandle>(keyGroupsStateHandle);
    +	public RunnableFuture<SnapshotResult<KeyedStateHandle>> getKeyedStateStreamFuture() throws IOException {
    +		return getGenericStateStreamFuture(keyedStateCheckpointOutputStream);
     	}
     
    -	public RunnableFuture<OperatorStateHandle> getOperatorStateStreamFuture() throws IOException {
    -		OperatorStateHandle operatorStateHandle = closeAndUnregisterStreamToObtainStateHandle(operatorStateCheckpointOutputStream);
    -		return new DoneFuture<>(operatorStateHandle);
    +	public RunnableFuture<SnapshotResult<OperatorStateHandle>> getOperatorStateStreamFuture() throws IOException {
    +		return getGenericStateStreamFuture(operatorStateCheckpointOutputStream);
    +	}
    +
    +	private <T extends StateObject> RunnableFuture<SnapshotResult<T>> getGenericStateStreamFuture(
    +		NonClosingCheckpointOutputStream<? extends T> stream) throws IOException {
    +		T operatorStateHandle = (T) closeAndUnregisterStreamToObtainStateHandle(stream);
    --- End diff --
    
    This cast seems a bit fishy to me. I think it should not be necessary if the generics are applied correctly. A way to solve it would be `T extends StreamStateHandle` and `RunnableFuture<? extends SnapshotResult<? extends KeyedStateHandle>> getKeyedStateStreamFuture()`


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162584582
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StateInitializationContextImplTest.java ---
    @@ -124,14 +140,69 @@ public void setUp() throws Exception {
     			operatorStateHandles.add(operatorStateHandle);
     		}
     
    +		OperatorSubtaskState operatorSubtaskState = new OperatorSubtaskState(
    +			Collections.emptyList(),
    +			operatorStateHandles,
    +			Collections.emptyList(),
    +			keyedStateHandles);
    +
    +		OperatorID operatorID = new OperatorID();
    +		TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    +		taskStateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
    +
    +		JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(0L, taskStateSnapshot);
    +
    +		TaskStateManager manager = new TaskStateManagerImpl(
    +			new JobID(),
    +			new ExecutionAttemptID(),
    +			mock(TaskLocalStateStore.class),
    +			jobManagerTaskRestore,
    +			mock(CheckpointResponder.class));
    +
    +		DummyEnvironment environment = new DummyEnvironment(
    +			"test",
    +			1,
    +			0,
    +			prev);
    +
    +		environment.setTaskStateManager(manager);
    +
    +		StateBackend stateBackend = new MemoryStateBackend(1024);
    +		StreamTaskStateManager streamTaskStateManager = new StreamTaskStateManagerImpl(
    +			environment,
    +			stateBackend,
    +			mock(ProcessingTimeService.class)) {
    +
    +			@Override
    +			protected <K> InternalTimeServiceManager<?, K> internalTimeServiceManager(
    +				AbstractKeyedStateBackend<K> keyedStatedBackend,
    +				KeyContext keyContext,
    +				Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
    +
    +				// We do not initialize a timer service manager here, because it would already consume the raw keyed
    +				// state as part of initialization. For the purpose of this test, we want an unconsumed raw keyed
    +				// stream.
    +				return null;
    +			}
    +		};
    +
    +		AbstractStreamOperator<?> mockOperator = mock(AbstractStreamOperator.class);
    +		when(mockOperator.getOperatorID()).thenReturn(operatorID);
    +
    +		StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(
    +			mockOperator,
    +			// notice that this essentially disables the previous test of the keyed stream because it was and is always
    +			// consumed by the timer service.
    +			mock(TypeSerializer.class),
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165315625
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java ---
    @@ -859,56 +861,77 @@ public void run() {
     				if (asyncCheckpointState.compareAndSet(CheckpointingOperation.AsynCheckpointState.RUNNING,
     					CheckpointingOperation.AsynCheckpointState.COMPLETED)) {
     
    -					TaskStateSnapshot acknowledgedState = hasState ? taskOperatorSubtaskStates : null;
    -
    -					TaskStateManager taskStateManager = owner.getEnvironment().getTaskStateManager();
    -
    -					// we signal stateless tasks by reporting null, so that there are no attempts to assign empty state
    -					// to stateless tasks on restore. This enables simple job modifications that only concern
    -					// stateless without the need to assign them uids to match their (always empty) states.
    -					taskStateManager.reportTaskStateSnapshot(
    -						checkpointMetaData,
    -						checkpointMetrics,
    -						acknowledgedState);
    -
    -					LOG.debug("{} - finished asynchronous part of checkpoint {}. Asynchronous duration: {} ms",
    -						owner.getName(), checkpointMetaData.getCheckpointId(), asyncDurationMillis);
    -
    -					LOG.trace("{} - reported the following states in snapshot for checkpoint {}: {}.",
    -						owner.getName(), checkpointMetaData.getCheckpointId(), acknowledgedState);
    +					reportCompletedSnapshotStates(
    +						jobManagerTaskOperatorSubtaskStates,
    +						localTaskOperatorSubtaskStates,
    +						asyncDurationMillis);
     
     				} else {
     					LOG.debug("{} - asynchronous part of checkpoint {} could not be completed because it was closed before.",
     						owner.getName(),
     						checkpointMetaData.getCheckpointId());
     				}
     			} catch (Exception e) {
    -				// the state is completed if an exception occurred in the acknowledgeCheckpoint call
    -				// in order to clean up, we have to set it to RUNNING again.
    -				asyncCheckpointState.compareAndSet(
    -					CheckpointingOperation.AsynCheckpointState.COMPLETED,
    -					CheckpointingOperation.AsynCheckpointState.RUNNING);
    -
    -				try {
    -					cleanup();
    -				} catch (Exception cleanupException) {
    -					e.addSuppressed(cleanupException);
    -				}
    -
    -				Exception checkpointException = new Exception(
    -					"Could not materialize checkpoint " + checkpointId + " for operator " +
    -						owner.getName() + '.',
    -					e);
    -
    -				owner.asynchronousCheckpointExceptionHandler.tryHandleCheckpointException(
    -					checkpointMetaData,
    -					checkpointException);
    +				handleExecutionException(e);
     			} finally {
     				owner.cancelables.unregisterCloseable(this);
     				FileSystemSafetyNet.closeSafetyNetAndGuardedResourcesForThread();
     			}
     		}
     
    +		private void reportCompletedSnapshotStates(
    +			TaskStateSnapshot acknowledgedTaskStateSnapshot,
    +			TaskStateSnapshot localTaskStateSnapshot,
    +			long asyncDurationMillis) {
    +
    +			TaskStateManager taskStateManager = owner.getEnvironment().getTaskStateManager();
    +
    +			boolean hasAckState = acknowledgedTaskStateSnapshot.hasState();
    +			boolean hasLocalState = localTaskStateSnapshot.hasState();
    +
    +			Preconditions.checkState(hasAckState || !hasLocalState,
    +				"Found cached state but no corresponding primary state is reported to the job " +
    +					"manager. This indicates a problem.");
    +
    +			// we signal stateless tasks by reporting null, so that there are no attempts to assign empty state
    +			// to stateless tasks on restore. This enables simple job modifications that only concern
    +			// stateless without the need to assign them uids to match their (always empty) states.
    +			taskStateManager.reportTaskStateSnapshots(
    +				checkpointMetaData,
    +				checkpointMetrics,
    +				hasAckState ? acknowledgedTaskStateSnapshot : null,
    +				hasLocalState ? localTaskStateSnapshot : null);
    +
    +			LOG.debug("{} - finished asynchronous part of checkpoint {}. Asynchronous duration: {} ms",
    +				owner.getName(), checkpointMetaData.getCheckpointId(), asyncDurationMillis);
    +
    +			LOG.trace("{} - reported the following states in snapshot for checkpoint {}: {}.",
    +				owner.getName(), checkpointMetaData.getCheckpointId(), acknowledgedTaskStateSnapshot);
    +		}
    +
    +		private void handleExecutionException(Exception e) {
    +			// the state is completed if an exception occurred in the acknowledgeCheckpoint call
    --- End diff --
    
    comment needs to be adapted. 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165318246
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTerminationTest.java ---
    @@ -227,21 +229,13 @@ protected void cleanup() throws Exception {
     			// has been stopped
     			CLEANUP_LATCH.trigger();
     
    -			// wait until handle async exception has been called to proceed with the termination of the
    -			// StreamTask
    -			HANDLE_ASYNC_EXCEPTION_LATCH.await();
    --- End diff --
    
    This handle should be deleted since it is no longer needed.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170260996
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalRecoveryConfig.java ---
    @@ -0,0 +1,120 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.io.ParseException;
    +import org.apache.flink.configuration.CheckpointingOptions;
    +import org.apache.flink.configuration.Configuration;
    +
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnull;
    +
    +/**
    + * This class encapsulates the completed configuration for local recovery, i.e. the root
    + * directories into which all file-based snapshots can be written and the general mode for the local recover feature.
    + */
    +public class LocalRecoveryConfig {
    +
    +	/**
    +	 * Enum over modes of local recovery:
    +	 * <p><ul>
    +	 * <li>DISABLED: disables local recovery.
    +	 * <li>ENABLE_FILE_BASED: enables local recovery in a variant that is based on local files.
    +	 * </ul>
    +	 */
    +	public enum LocalRecoveryMode {
    +		DISABLED,
    +		ENABLE_FILE_BASED;
    +
    +		/**
    +		 * Attempts to parses the given string to a {@link LocalRecoveryMode}.
    +		 *
    +		 * @param input string to parse
    +		 * @return the {@link LocalRecoveryMode} that is represented by the input string.
    +		 * @throws ParseException if the string cannot be parsed to a {@link LocalRecoveryMode}.
    +		 */
    +		@Nonnull
    +		public static LocalRecoveryMode fromString(String input) throws ParseException {
    +			if (input != null) {
    +				String trimInput = input.trim();
    +				if (DISABLED.toString().equalsIgnoreCase(trimInput)) {
    +					return DISABLED;
    +				} else if (ENABLE_FILE_BASED.toString().equalsIgnoreCase(trimInput)) {
    +					return ENABLE_FILE_BASED;
    +				}
    +			}
    +			throw new ParseException("Cannot parse input to LocalRecoveryMode: " + input);
    +		}
    +
    +		/**
    +		 * Extracts the {@link LocalRecoveryMode} from the given configuration. Defaults to LocalRecoveryMode.DISABLED
    +		 * if no configuration value is specified or parsing the value resulted in an exception.
    +		 *
    +		 * @param configuration the configuration that specifies the value for the local recovery mode.
    +		 * @return the local recovery mode as found in the config, or LocalRecoveryMode.DISABLED if no mode was
    +		 * configured or the specified mode could not be parsed.
    +		 */
    +		@Nonnull
    +		public static LocalRecoveryMode fromConfig(@Nonnull Configuration configuration) {
    +			String localRecoveryConfString = configuration.getString(CheckpointingOptions.LOCAL_RECOVERY);
    +			try {
    +				return LocalRecoveryConfig.LocalRecoveryMode.fromString(localRecoveryConfString);
    +			} catch (ParseException ex) {
    +				LoggerFactory.getLogger(LocalRecoveryConfig.class).warn(
    +					"Exception while parsing configuration of local recovery mode. Local recovery will be disabled.",
    +					ex);
    +				return LocalRecoveryConfig.LocalRecoveryMode.DISABLED;
    +			}
    +		}
    +	}
    +
    +	/** The local recovery mode. */
    +	@Nonnull
    +	private final LocalRecoveryMode localRecoveryMode;
    --- End diff --
    
    This is a bit tricky, I could see that sometimes the stream abstraction might not be enough for all (optimized) types of  local snapshotting, and a directory is more powerful in that respect. For the time being, I think we should keep it like this and adjust if we have more concrete requirements that call for a different generalization.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168488712
  
    --- Diff: flink-runtime/src/main/scala/org/apache/flink/runtime/taskmanager/TaskManager.scala ---
    @@ -120,15 +120,16 @@ import scala.language.postfixOps
      *      requires a clean JVM.
      */
     class TaskManager(
    -    protected val config: TaskManagerConfiguration,
    -    protected val resourceID: ResourceID,
    -    protected val location: TaskManagerLocation,
    -    protected val memoryManager: MemoryManager,
    -    protected val ioManager: IOManager,
    -    protected val network: NetworkEnvironment,
    -    protected val numberOfSlots: Int,
    -    protected val highAvailabilityServices: HighAvailabilityServices,
    -    protected val taskManagerMetricGroup: TaskManagerMetricGroup)
    +                   protected val config: TaskManagerConfiguration,
    +                   protected val resourceID: ResourceID,
    +                   protected val location: TaskManagerLocation,
    +                   protected val memoryManager: MemoryManager,
    +                   protected val ioManager: IOManager,
    +                   protected val network: NetworkEnvironment,
    +                   protected val taskManagerLocalStateStoresManager: TaskExecutorLocalStateStoresManager,
    +                   protected val numberOfSlots: Int,
    +                   protected val highAvailabilityServices: HighAvailabilityServices,
    +                   protected val taskManagerMetricGroup: TaskManagerMetricGroup)
    --- End diff --
    
    Please revert indentation here.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170242003
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileBasedStateOutputStream.java ---
    @@ -0,0 +1,155 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state.filesystem;
    +
    +import org.apache.flink.core.fs.FSDataOutputStream;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.FileSystem.WriteMode;
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.state.CheckpointStreamFactory.CheckpointStateOutputStream;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.IOException;
    +
    +import static org.apache.flink.util.Preconditions.checkNotNull;
    +
    +/**
    + * A {@link CheckpointStateOutputStream} that writes into a specified file and
    + * returns a {@link FileStateHandle} upon closing.
    + *
    + * <p>Unlike the {@link org.apache.flink.runtime.state.filesystem.FsCheckpointStreamFactory.FsCheckpointStateOutputStream},
    + * this stream does not have a threshold below which it returns a memory byte stream handle,
    + * and does not create random files, but writes to a specified file.
    + */
    +public final class FileBasedStateOutputStream extends CheckpointStateOutputStream {
    +
    +	private static final Logger LOG = LoggerFactory.getLogger(FileBasedStateOutputStream.class);
    +
    +	// ------------------------------------------------------------------------
    +
    +	private final FSDataOutputStream out;
    +
    +	private final Path path;
    +
    +	private final FileSystem fileSystem;
    +
    +	private volatile boolean closed;
    +
    +
    +	public FileBasedStateOutputStream(FileSystem fileSystem, Path path) throws IOException {
    +		this.fileSystem = checkNotNull(fileSystem);
    +		this.path = checkNotNull(path);
    +
    +		this.out = fileSystem.create(path, WriteMode.NO_OVERWRITE);
    +	}
    +
    +	// ------------------------------------------------------------------------
    +	//  I/O
    +	// ------------------------------------------------------------------------
    +
    +	@Override
    +	public final void write(int b) throws IOException {
    +		out.write(b);
    --- End diff --
    
    I found this is inconsistency with `FsCheckpointStateOutputStream`, this means that invoking `DuplicatingCheckpointOutputStream.write(int b).` will broken the position align. Is this intended?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r167356160
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    --- End diff --
    
    life cycle**s**
    
    primary cop**ies**


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168545164
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/BackendRestorerProcedure.java ---
    @@ -0,0 +1,143 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.core.fs.CloseableRegistry;
    +import org.apache.flink.runtime.checkpoint.StateObjectCollection;
    +import org.apache.flink.runtime.state.Snapshotable;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.util.Disposable;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +import org.apache.flink.util.ThrowingSupplier;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.Closeable;
    +import java.io.IOException;
    +import java.util.Iterator;
    +
    +/**
    + * This class implements the logic that creates (and potentially restores) a state backend. The restore logic
    + * considers multiple, prioritized options of snapshots to restore from, where all of the options should recreate
    + * the same state for the backend. When we fail to restore from the snapshot with the highest priority (typically
    + * the "fastest" to restore), we fallback to the next snapshot with the next highest priority. We also take care
    + * of cleaning up from failed restore attempts. We only reattempt when the problem occurs during the restore call
    + * and will only stop after all snapshot alternatives are exhausted and all failed.
    + *
    + * @param <T> type of the restored backend.
    + * @param <S> type of the supplied snapshots from which the backend restores.
    + */
    +public class BackendRestorerProcedure<
    +	T extends Closeable & Disposable & Snapshotable<?, StateObjectCollection<S>>,
    +	S extends StateObject> {
    +
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(BackendRestorerProcedure.class);
    +
    +	/** Factory for new, fresh backends without state. */
    +	private final ThrowingSupplier<T, Exception> instanceSupplier;
    +
    +	/** This registry is used so that recovery can participate in the task lifecycle, i.e. can be canceled. */
    +	private final CloseableRegistry backendCloseableRegistry;
    +
    +	/**
    +	 * Creates a new backend restorer using the given backend supplier and the closeable registry.
    +	 *
    +	 * @param instanceSupplier factory function for new, empty backend instances.
    +	 * @param backendCloseableRegistry registry to allow participation in task lifecycle, e.g. react to cancel.
    +	 */
    +	public BackendRestorerProcedure(
    +		@Nonnull ThrowingSupplier<T, Exception> instanceSupplier,
    +		@Nonnull CloseableRegistry backendCloseableRegistry) {
    +
    +		this.instanceSupplier = Preconditions.checkNotNull(instanceSupplier);
    +		this.backendCloseableRegistry = Preconditions.checkNotNull(backendCloseableRegistry);
    +	}
    +
    +	/**
    +	 * Creates a new state backend and restores it from the provided set of state snapshot alternatives.
    +	 *
    +	 * @param restoreOptions iterator over a prioritized set of state snapshot alternatives for recovery.
    +	 * @return the created (and restored) state backend.
    +	 * @throws Exception if the backend could not be created or restored.
    +	 */
    +	public @Nonnull T createAndRestore(@Nonnull Iterator<StateObjectCollection<S>> restoreOptions) throws Exception {
    +
    +		StateObjectCollection<S> restoreState = null;
    +
    +		boolean retry;
    +
    +		while (true) {
    --- End diff --
    
    It is not quiet that easy, because the restore method must always be called at least once and also the log message might differ between iterations (the "will retry" part). But I agree it could be simplified.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169898668
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/SnapshotDirectory.java ---
    @@ -0,0 +1,172 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.FileStatus;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.io.IOException;
    +import java.util.concurrent.atomic.AtomicReference;
    +
    +/**
    + * This class represents a directory that is the target for a state snapshot. This class provides some method that
    + * simplify resource management when dealing with such directories, e.g. it can produce a {@link DirectoryStateHandle}
    + * when the snapshot is completed and disposal considers whether or not a snapshot was already completed. For a
    + * completed snapshot, the ownership for cleanup is transferred to the created directory state handle. For incomplete
    + * snapshots, calling {@link #deleteIfIncompleteSnapshot()} will delete the underlying directory resource.
    + */
    +public class SnapshotDirectory {
    +
    +	/**
    +	 * Lifecycle stages of a snapshot directory.
    +	 */
    +	enum State {
    +		ONGOING, COMPLETED, DELETED
    +	}
    +
    +	/** This path describes the underlying directory for the snapshot. */
    +	@Nonnull
    +	private final Path directory;
    +
    +	/** The filesystem that contains the snapshot directory. */
    +	@Nonnull
    +	private final FileSystem fileSystem;
    --- End diff --
    
    In this class, I would keep it because it is used in several operations and obtaining the `FileSystem` from `Path` looks like it might not be a cheap operation. So this is more like a cache, or in case that you already have the `FileSystem` when constructing the object.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r159881015
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/MultiStreamStateHandle.java ---
    @@ -1,104 +0,0 @@
    -/*
    --- End diff --
    
    That was a dead code, right? Please extract removal of this (and `MultiStreamStateHandleTest`) to separate commit.
      


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165341071
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerServices.java ---
    @@ -197,7 +201,13 @@ public static TaskManagerServices fromConfiguration(
     		final JobManagerTable jobManagerTable = new JobManagerTable();
     
     		final JobLeaderService jobLeaderService = new JobLeaderService(taskManagerLocation);
    -		final TaskExecutorLocalStateStoresManager taskStateManager = new TaskExecutorLocalStateStoresManager();
    +
    +		final File taskExecutorLocalStateRootDir =
    +			new File(taskManagerServicesConfiguration.getLocalStateRootDir(), LOCAL_STATE_SUB_DIRECTORY_ROOT);
    --- End diff --
    
    Are we giving the `TaskExecutorLocalStateStoresManager` a unique directory? What happens if multiple TMs run on the same machine? Will they use different directories?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by sihuazhou <gi...@git.apache.org>.
Github user sihuazhou commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168184247
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStoreImpl.java ---
    @@ -0,0 +1,298 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.annotation.VisibleForTesting;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.core.fs.FileSystem;
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +import javax.annotation.concurrent.GuardedBy;
    +
    +import java.io.File;
    +import java.io.IOException;
    +import java.util.ArrayList;
    +import java.util.Collection;
    +import java.util.HashMap;
    +import java.util.Iterator;
    +import java.util.List;
    +import java.util.Map;
    +import java.util.SortedMap;
    +import java.util.TreeMap;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
    +
    +/**
    + * Main implementation of a {@link TaskLocalStateStore}.
    + */
    +public class TaskLocalStateStoreImpl implements TaskLocalStateStore {
    +
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskLocalStateStoreImpl.class);
    +
    +	/** Maximum number of retained snapshots. */
    +	@VisibleForTesting
    +	static final int MAX_RETAINED_SNAPSHOTS = 5;
    +
    +	/** Dummy value to use instead of null to satisfy {@link ConcurrentHashMap}. */
    +	private final TaskStateSnapshot NULL_DUMMY = new TaskStateSnapshot();
    +
    +	/** JobID from the owning subtask. */
    +	@Nonnull
    +	private final JobID jobID;
    +
    +	/** AllocationID of the owning slot. */
    +	@Nonnull
    +	private final AllocationID allocationID;
    +
    +	/** JobVertexID of the owning subtask. */
    +	@Nonnull
    +	private final JobVertexID jobVertexID;
    +
    +	/** Subtask index of the owning subtask. */
    +	@Nonnegative
    +	private final int subtaskIndex;
    +
    +	/** The configured mode for local recovery. */
    +	@Nonnull
    +	private final LocalRecoveryConfig localRecoveryConfig;
    +
    +	/** Executor that runs the discarding of released state objects. */
    +	@Nonnull
    +	private final Executor discardExecutor;
    +
    +	/** Lock for synchronisation on the storage map and the discarded status. */
    +	@Nonnull
    +	private final Object lock;
    +
    +	/** Status flag if this store was already discarded. */
    +	@GuardedBy("lock")
    +	private boolean discarded;
    +
    +	/** Maps checkpoint ids to local TaskStateSnapshots. */
    +	@Nonnull
    +	@GuardedBy("lock")
    +	private final SortedMap<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
    +
    +	public TaskLocalStateStoreImpl(
    +		@Nonnull JobID jobID,
    +		@Nonnull AllocationID allocationID,
    +		@Nonnull JobVertexID jobVertexID,
    +		@Nonnegative int subtaskIndex,
    +		@Nonnull LocalRecoveryConfig localRecoveryConfig,
    +		@Nonnull Executor discardExecutor) {
    +
    +		this.jobID = jobID;
    +		this.allocationID = allocationID;
    +		this.jobVertexID = jobVertexID;
    +		this.subtaskIndex = subtaskIndex;
    +		this.discardExecutor = discardExecutor;
    +		this.lock = new Object();
    +		this.storedTaskStateByCheckpointID = new TreeMap<>();
    +		this.discarded = false;
    +		this.localRecoveryConfig = localRecoveryConfig;
    +	}
    +
    +	@Override
    +	public void storeLocalState(
    +		@Nonnegative long checkpointId,
    +		@Nullable TaskStateSnapshot localState) {
    +
    +		if (localState == null) {
    +			localState = NULL_DUMMY;
    +		}
    +
    +		LOG.info("Storing local state for checkpoint {}.", checkpointId);
    +		LOG.debug("Local state for checkpoint {} is {}.", checkpointId, localState);
    +
    +		Map<Long, TaskStateSnapshot> toDiscard = new HashMap<>(MAX_RETAINED_SNAPSHOTS);
    +
    +		synchronized (lock) {
    +			if (discarded) {
    +				// we ignore late stores and simply discard the state.
    +				toDiscard.put(checkpointId, localState);
    +			} else {
    +				TaskStateSnapshot previous =
    +					storedTaskStateByCheckpointID.put(checkpointId, localState);
    +
    +				if (previous != null) {
    +					toDiscard.put(checkpointId, previous);
    +				}
    +
    +				// remove from history.
    +				while (storedTaskStateByCheckpointID.size() > MAX_RETAINED_SNAPSHOTS) {
    +					Long removeCheckpointID = storedTaskStateByCheckpointID.firstKey();
    --- End diff --
    
    I found this can be problematic, even this method has been invoked but it doesn't mean the checkpoint is truly completed. If we perform remove local task state here we may removed the `really completed local checkpoint data`. Image the situation when `max concurrent checkpoint number = MAX_RETAINED_SNAPSHOTS + 1`...


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168486303
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamWithResultProvider.java ---
    @@ -0,0 +1,175 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.core.fs.Path;
    +import org.apache.flink.runtime.state.filesystem.FixFileFsStateOutputStream;
    +import org.apache.flink.util.ExceptionUtils;
    +import org.apache.flink.util.Preconditions;
    +
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import java.io.Closeable;
    +import java.io.File;
    +import java.io.IOException;
    +
    +/**
    + * Interface that provides access to a CheckpointStateOutputStream and a method to provide the {@link SnapshotResult}.
    + * This abstracts from different ways that a result is obtained from checkpoint output streams.
    + */
    +public interface CheckpointStreamWithResultProvider extends Closeable {
    +
    +	SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +		KeyGroupRangeOffsets kgOffs) throws IOException;
    +
    +	CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream();
    +
    +	@Override
    +	default void close() throws IOException {
    +		CheckpointStreamFactory.CheckpointStateOutputStream outputStream = getCheckpointOutputStream();
    +		if (outputStream != null) {
    +			outputStream.close();
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that only creates the
    +	 * primary/remote/jm-owned state.
    +	 */
    +	class PrimaryStreamOnly implements CheckpointStreamWithResultProvider {
    +
    +		private final CheckpointStreamFactory.CheckpointStateOutputStream outputStream;
    +
    +		public PrimaryStreamOnly(CheckpointStreamFactory.CheckpointStateOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets kgOffs) throws IOException {
    +
    +			StreamStateHandle streamStateHandle = outputStream.closeAndGetHandle();
    +			KeyGroupsStateHandle primaryStateHandle = new KeyGroupsStateHandle(kgOffs, streamStateHandle);
    +			return new SnapshotResult<>(primaryStateHandle, null);
    +		}
    +
    +		@Override
    +		public CheckpointStreamFactory.CheckpointStateOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	/**
    +	 * Implementation of {@link CheckpointStreamWithResultProvider} that creates both, the
    +	 * primary/remote/jm-owned state and the secondary/local/tm-owned state.
    +	 */
    +	class PrimaryAndSecondaryStream implements CheckpointStreamWithResultProvider {
    +
    +		private static final Logger LOG = LoggerFactory.getLogger(PrimaryAndSecondaryStream.class);
    +
    +		private final DuplicatingCheckpointOutputStream outputStream;
    +
    +		public PrimaryAndSecondaryStream(DuplicatingCheckpointOutputStream outputStream) {
    +			this.outputStream = Preconditions.checkNotNull(outputStream);
    +		}
    +
    +		@Override
    +		public SnapshotResult<KeyedStateHandle> closeAndFinalizeCheckpointStreamResult(
    +			KeyGroupRangeOffsets offsets) throws IOException {
    +
    +			final StreamStateHandle primaryStreamStateHandle;
    +
    +			try {
    +				primaryStreamStateHandle = outputStream.closeAndGetPrimaryHandle();
    +			} catch (IOException primaryEx) {
    +				try {
    +					outputStream.close();
    +				} catch (IOException closeEx) {
    +					primaryEx = ExceptionUtils.firstOrSuppressed(closeEx, primaryEx);
    +				}
    +				throw primaryEx;
    +			}
    +
    +			StreamStateHandle secondaryStreamStateHandle = null;
    +
    +			try {
    +				secondaryStreamStateHandle = outputStream.closeAndGetSecondaryHandle();
    +			} catch (IOException secondaryEx) {
    +				LOG.warn("Exception from secondary/local checkpoint stream.", secondaryEx);
    +			}
    +
    +			if (primaryStreamStateHandle != null) {
    +
    +				final KeyGroupsStateHandle primaryKeyGroupsStateHandle =
    +					new KeyGroupsStateHandle(offsets, primaryStreamStateHandle);
    +
    +				final KeyGroupsStateHandle secondaryKeyGroupsStateHandle = secondaryStreamStateHandle != null ?
    +					new KeyGroupsStateHandle(offsets, secondaryStreamStateHandle) :
    +					null;
    +
    +				return new SnapshotResult<>(primaryKeyGroupsStateHandle, secondaryKeyGroupsStateHandle);
    +			}
    +
    +			return null;
    +		}
    +
    +		@Override
    +		public DuplicatingCheckpointOutputStream getCheckpointOutputStream() {
    +			return outputStream;
    +		}
    +	}
    +
    +	class Factory {
    --- End diff --
    
    For what do we need this class? Wouldn't two static methods `createSimpleStream` and `createDuplicatedStream` be enough?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StephanEwen <gi...@git.apache.org>.
Github user StephanEwen commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160492829
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamFactory.java ---
    @@ -43,9 +44,9 @@ CheckpointStateOutputStream createCheckpointStateOutputStream(
     	 * Closes the stream factory, releasing all internal resources, but does not delete any
     	 * persistent checkpoint data.
     	 *
    -	 * @throws Exception Exceptions can be forwarded and will be logged by the system
    +	 * @throws IOException Exceptions can be forwarded and will be logged by the system
     	 */
    -	void close() throws Exception;
    +	void close() throws IOException;
    --- End diff --
    
    I think this method is actually removed in the latest master.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168480072
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -55,25 +54,24 @@
     	/** */
     	private final Map<Long, TaskStateSnapshot> storedTaskStateByCheckpointID;
     
    -	/** This is the base directory for all local state of the subtask that owns this {@link TaskLocalStateStore}. */
    -	private final File subtaskLocalStateBaseDirectory;
    +	/** The root directories for all local state of this {@link TaskLocalStateStore}. */
    +	private final File[] localStateRootDirectories;
     
     	public TaskLocalStateStore(
     		JobID jobID,
     		JobVertexID jobVertexID,
     		int subtaskIndex,
    -		File localStateRootDir) {
    +		File[] localStateRootDirectories) {
     
     		this.jobID = jobID;
     		this.jobVertexID = jobVertexID;
     		this.subtaskIndex = subtaskIndex;
     		this.storedTaskStateByCheckpointID = new HashMap<>();
    -		this.subtaskLocalStateBaseDirectory =
    -			new File(localStateRootDir, createSubtaskPath(jobID, jobVertexID, subtaskIndex));
    +		this.localStateRootDirectories = localStateRootDirectories;
     	}
     
    -	static String createSubtaskPath(JobID jobID, JobVertexID jobVertexID, int subtaskIndex) {
    -		return "jid-" + jobID + "_vtx-" + jobVertexID + "_sti-" + subtaskIndex;
    +	protected String createSubtaskPath() {
    +		return jobID + File.separator + jobVertexID + File.separator + subtaskIndex;
    --- End diff --
    
    Maybe using `Paths.get` for that?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r162584949
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamTaskStateManagerImplTest.java ---
    @@ -0,0 +1,323 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.api.operators;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.api.common.typeutils.TypeSerializer;
    +import org.apache.flink.api.common.typeutils.base.IntSerializer;
    +import org.apache.flink.core.fs.CloseableRegistry;
    +import org.apache.flink.runtime.checkpoint.JobManagerTaskRestore;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.checkpoint.savepoint.CheckpointTestUtils;
    +import org.apache.flink.runtime.execution.Environment;
    +import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
    +import org.apache.flink.runtime.query.TaskKvStateRegistry;
    +import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
    +import org.apache.flink.runtime.state.CheckpointStreamFactory;
    +import org.apache.flink.runtime.state.KeyGroupRange;
    +import org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.OperatorStateBackend;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.StateBackend;
    +import org.apache.flink.runtime.state.StatePartitionStreamProvider;
    +import org.apache.flink.runtime.state.TaskStateManager;
    +import org.apache.flink.runtime.state.TaskStateManagerImplTest;
    +import org.apache.flink.runtime.state.memory.MemoryStateBackend;
    +import org.apache.flink.runtime.taskmanager.TestCheckpointResponder;
    +import org.apache.flink.streaming.runtime.tasks.ProcessingTimeService;
    +import org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService;
    +import org.apache.flink.util.CloseableIterable;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import javax.annotation.Nullable;
    +
    +import java.io.Closeable;
    +import java.io.IOException;
    +import java.util.Collections;
    +import java.util.Random;
    +
    +import static org.mockito.Matchers.any;
    +import static org.mockito.Matchers.anyInt;
    +import static org.mockito.Matchers.eq;
    +import static org.mockito.Mockito.mock;
    +import static org.mockito.Mockito.spy;
    +import static org.mockito.Mockito.verify;
    +import static org.mockito.Mockito.when;
    +
    +/**
    + * Test for {@link StreamTaskStateManagerImpl}.
    + */
    +public class StreamTaskStateManagerImplTest {
    +
    +	@Test
    +	public void testNoRestore() throws Exception {
    +
    +		MemoryStateBackend stateBackend = spy(new MemoryStateBackend(1024));
    +
    +		// No job manager provided state to restore
    +		StreamTaskStateManager 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,
    +			typeSerializer,
    +			closeableRegistry);
    +
    +		verify(stateBackend).createKeyedStateBackend(
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168498608
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -20,35 +20,47 @@
     
     import org.apache.flink.annotation.VisibleForTesting;
     import org.apache.flink.api.common.JobID;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
     import org.apache.flink.runtime.jobgraph.JobVertexID;
    -import org.apache.flink.util.ExceptionUtils;
    -import org.apache.flink.util.Preconditions;
     
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
     import javax.annotation.Nonnull;
     
     import java.io.File;
    -import java.util.HashMap;
     import java.util.Map;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
     
     /**
      * This class holds the all {@link TaskLocalStateStore} objects for a task executor (manager).
    - *
    - * TODO: this still still work in progress and partially still acts as a placeholder.
      */
     public class TaskExecutorLocalStateStoresManager {
     
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskExecutorLocalStateStoresManager.class);
    +
     	/**
     	 * This map holds all local state stores for tasks running on the task manager / executor that own the instance of
    -	 * this.
    +	 * this. Maps from allocation id to all the subtask's local state stores.
     	 */
    -	private final Map<JobID, Map<JobVertexSubtaskKey, TaskLocalStateStore>> taskStateStoresByJobID;
    +	private final Map<AllocationID, Map<JobVertexSubtaskKey, TaskLocalStateStore>> taskStateStoresByAllocationID;
    --- End diff --
    
    Maybe we should change the type of the field to `ConcurrentHashMap` that way it is clearer that we need a thread safe structure here.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165326694
  
    --- Diff: flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/LocalStateForwardingTest.java ---
    @@ -0,0 +1,213 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.streaming.runtime.tasks;
    +
    +import org.apache.flink.api.common.ExecutionConfig;
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.configuration.Configuration;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
    +import org.apache.flink.runtime.checkpoint.CheckpointMetrics;
    +import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
    +import org.apache.flink.runtime.checkpoint.StateObjectCollection;
    +import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
    +import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +import org.apache.flink.runtime.jobgraph.OperatorID;
    +import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
    +import org.apache.flink.runtime.state.DoneFuture;
    +import org.apache.flink.runtime.state.KeyedStateHandle;
    +import org.apache.flink.runtime.state.OperatorStateHandle;
    +import org.apache.flink.runtime.state.SnapshotResult;
    +import org.apache.flink.runtime.state.StateObject;
    +import org.apache.flink.runtime.state.TaskLocalStateStore;
    +import org.apache.flink.runtime.state.TaskStateManagerImpl;
    +import org.apache.flink.runtime.state.TestTaskStateManager;
    +import org.apache.flink.runtime.taskmanager.TestCheckpointResponder;
    +import org.apache.flink.streaming.api.operators.OperatorSnapshotFutures;
    +
    +import org.junit.Assert;
    +import org.junit.Test;
    +
    +import javax.annotation.Nonnull;
    +import javax.annotation.Nullable;
    +
    +import java.util.HashMap;
    +import java.util.Map;
    +import java.util.concurrent.Future;
    +import java.util.concurrent.RunnableFuture;
    +import java.util.concurrent.atomic.AtomicBoolean;
    +
    +import static org.mockito.Mockito.mock;
    +
    +/**
    + * Test for forwarding of state reporting to and from {@link org.apache.flink.runtime.state.TaskStateManager}.
    + */
    +public class LocalStateForwardingTest {
    +
    +	/**
    +	 * This tests the forwarding of jm and tm-local state from the futures reported by the backends, through the
    +	 * async checkpointing thread to the {@link org.apache.flink.runtime.state.TaskStateManager}.
    +	 */
    +	@Test
    +	public void testForwardingFromSnapshotToTaskStateManager() throws Exception {
    +
    +		TestTaskStateManager taskStateManager = new TestTaskStateManager();
    +
    +		StreamMockEnvironment streamMockEnvironment = new StreamMockEnvironment(
    +			new Configuration(),
    +			new Configuration(),
    +			new ExecutionConfig(),
    +			1024*1024,
    +			new MockInputSplitProvider(),
    +			0,
    +			taskStateManager);
    +
    +		StreamTask testStreamTask = new StreamTaskTest.NoOpStreamTask(streamMockEnvironment);
    +		CheckpointMetaData checkpointMetaData = new CheckpointMetaData(0L, 0L);
    +		CheckpointMetrics checkpointMetrics = new CheckpointMetrics();
    +
    +		Map<OperatorID, OperatorSnapshotFutures> snapshots = new HashMap<>(1);
    +		OperatorSnapshotFutures osFuture = new OperatorSnapshotFutures();
    +
    +		osFuture.setKeyedStateManagedFuture(createSnapshotResult(KeyedStateHandle.class));
    +		osFuture.setKeyedStateRawFuture(createSnapshotResult(KeyedStateHandle.class));
    +		osFuture.setOperatorStateManagedFuture(createSnapshotResult(OperatorStateHandle.class));
    +		osFuture.setOperatorStateRawFuture(createSnapshotResult(OperatorStateHandle.class));
    +
    +		OperatorID operatorID = new OperatorID();
    +		snapshots.put(operatorID, osFuture);
    +
    +		StreamTask.AsyncCheckpointRunnable checkpointRunnable =
    +			new StreamTask.AsyncCheckpointRunnable(
    +				testStreamTask,
    +				snapshots,
    +				checkpointMetaData,
    +				checkpointMetrics,
    +				0L);
    +
    +		checkpointRunnable.run();
    +
    +		TaskStateSnapshot lastJobManagerTaskStateSnapshot = taskStateManager.getLastJobManagerTaskStateSnapshot();
    +		TaskStateSnapshot lastTaskManagerTaskStateSnapshot = taskStateManager.getLastTaskManagerTaskStateSnapshot();
    +
    +		OperatorSubtaskState jmState =
    +			lastJobManagerTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
    +
    +		OperatorSubtaskState tmState =
    +			lastTaskManagerTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
    +
    +		performCheck(osFuture.getKeyedStateManagedFuture(), jmState.getManagedKeyedState(), tmState.getManagedKeyedState());
    +		performCheck(osFuture.getKeyedStateRawFuture(), jmState.getRawKeyedState(), tmState.getRawKeyedState());
    +		performCheck(osFuture.getOperatorStateManagedFuture(), jmState.getManagedOperatorState(), tmState.getManagedOperatorState());
    +		performCheck(osFuture.getOperatorStateRawFuture(), jmState.getRawOperatorState(), tmState.getRawOperatorState());
    +	}
    +
    +	/**
    +	 * This tests that state that was reported to the {@link org.apache.flink.runtime.state.TaskStateManager} is also
    +	 * reported to {@link org.apache.flink.runtime.taskmanager.CheckpointResponder} and {@link TaskLocalStateStore}.
    +	 */
    +	@Test
    +	public void testForwardingFromTaskStateManagerToResponderAndTaskLocalStateStore() {
    +
    +		final JobID jobID = new JobID();
    +		final ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
    +		final CheckpointMetaData checkpointMetaData = new CheckpointMetaData(42L, 4711L);
    +		final CheckpointMetrics checkpointMetrics = new CheckpointMetrics();
    +		final int subtaskIdx = 42;
    +		JobVertexID jobVertexID = new JobVertexID();
    +
    +		TaskStateSnapshot jm = new TaskStateSnapshot();
    +		TaskStateSnapshot tm = new TaskStateSnapshot();
    +
    +		final AtomicBoolean jmReported = new AtomicBoolean(false);
    +		final AtomicBoolean tmReported = new AtomicBoolean(false);
    +
    +		TestCheckpointResponder checkpointResponder = new TestCheckpointResponder() {
    +
    +			@Override
    +			public void acknowledgeCheckpoint(
    +				JobID lJobID,
    +				ExecutionAttemptID lExecutionAttemptID,
    +				long lCheckpointId,
    +				CheckpointMetrics lCheckpointMetrics,
    +				TaskStateSnapshot lSubtaskState) {
    +
    +				Assert.assertEquals(jobID, lJobID);
    +				Assert.assertEquals(executionAttemptID, lExecutionAttemptID);
    +				Assert.assertEquals(checkpointMetaData.getCheckpointId(), lCheckpointId);
    +				Assert.assertEquals(checkpointMetrics, lCheckpointMetrics);
    +				jmReported.set(true);
    +			}
    +		};
    +
    +		TaskLocalStateStore taskLocalStateStore = new TaskLocalStateStore(jobID, jobVertexID, subtaskIdx) {
    +			@Override
    +			public void storeLocalState(
    +				@Nonnull CheckpointMetaData checkpointMetaData,
    +				@Nullable TaskStateSnapshot localState) {
    +
    +				Assert.assertEquals(tm, localState);
    +				tmReported.set(true);
    +			}
    +		};
    +
    +		TaskStateManagerImpl taskStateManager =
    +			new TaskStateManagerImpl(
    +				jobID,
    +				executionAttemptID,
    +				taskLocalStateStore,
    +				null,
    +				checkpointResponder);
    +
    +		taskStateManager.reportTaskStateSnapshots(
    +			checkpointMetaData,
    +			checkpointMetrics,
    +			jm,
    +			tm);
    +
    +		Assert.assertTrue("Reporting for JM state was not called.", jmReported.get());
    +		Assert.assertTrue("Reporting for TM state was not called.", tmReported.get());
    +	}
    +
    +	private static <T extends StateObject> void performCheck(
    +		Future<SnapshotResult<T>> cc,
    +		StateObjectCollection<T> jm,
    +		StateObjectCollection<T> tm) {
    --- End diff --
    
    Naming of variables could be a bit more expressive.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r167357486
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    +
    +- RocksDBStateBackend: task-local recovery is supported for keyed state. For *full checkpoints*, state is duplicated to a local file. This can introduce additional write costs
    +and occupies local disk space. For *incremental snapshots*, the local state is based on RocksDB's native checkpointing mechanism. This mechanism is also used as the first step
    +to create the primary copy, which means that in this case no additional cost is introduced for create the secondary copy. We simply keep the native checkpoint directory around
    --- End diff --
    
    creat**ing**


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170229683
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalRecoveryConfig.java ---
    @@ -0,0 +1,120 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.io.ParseException;
    +import org.apache.flink.configuration.CheckpointingOptions;
    +import org.apache.flink.configuration.Configuration;
    +
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnull;
    +
    +/**
    + * This class encapsulates the completed configuration for local recovery, i.e. the root
    + * directories into which all file-based snapshots can be written and the general mode for the local recover feature.
    + */
    +public class LocalRecoveryConfig {
    +
    +	/**
    +	 * Enum over modes of local recovery:
    +	 * <p><ul>
    +	 * <li>DISABLED: disables local recovery.
    +	 * <li>ENABLE_FILE_BASED: enables local recovery in a variant that is based on local files.
    +	 * </ul>
    +	 */
    +	public enum LocalRecoveryMode {
    +		DISABLED,
    +		ENABLE_FILE_BASED;
    +
    +		/**
    +		 * Attempts to parses the given string to a {@link LocalRecoveryMode}.
    +		 *
    +		 * @param input string to parse
    +		 * @return the {@link LocalRecoveryMode} that is represented by the input string.
    +		 * @throws ParseException if the string cannot be parsed to a {@link LocalRecoveryMode}.
    +		 */
    +		@Nonnull
    +		public static LocalRecoveryMode fromString(String input) throws ParseException {
    +			if (input != null) {
    +				String trimInput = input.trim();
    +				if (DISABLED.toString().equalsIgnoreCase(trimInput)) {
    +					return DISABLED;
    +				} else if (ENABLE_FILE_BASED.toString().equalsIgnoreCase(trimInput)) {
    +					return ENABLE_FILE_BASED;
    +				}
    +			}
    +			throw new ParseException("Cannot parse input to LocalRecoveryMode: " + input);
    +		}
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168501090
  
    --- Diff: flink-runtime/src/main/scala/org/apache/flink/runtime/taskmanager/TaskManager.scala ---
    @@ -1195,22 +1202,15 @@ class TaskManager(
               config.getTimeout().getSize(),
               config.getTimeout().getUnit()))
     
    -      //TODO: make configurable, this is just the future fallback case,
    -      // integrate with ConfigConstants.TASK_MANAGER_LOCAL_STATE_ROOT_DIR_KEY
    -      // TODO: wire this so that the manager survives the end of the task
    -      //TODO create "localState" sub-directories!!!
    -      val taskExecutorLocalStateStoresManager =
    -        new TaskExecutorLocalStateStoresManager(ioManager.getSpillingDirectories)
    -
    -      val localStateStore = taskExecutorLocalStateStoresManager.localStateStoreForTask(
    +      val taskLocalStateStore = taskManagerLocalStateStoresManager.localStateStoreForTask(
             jobInformation.getJobId,
             taskInformation.getJobVertexId,
             tdd.getSubtaskIndex)
     
    -      val slotStateManager = new TaskStateManagerImpl(
    +      val taskLocalStateManager = new TaskStateManagerImpl(
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170270552
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskLocalStateStore.java ---
    @@ -166,20 +190,38 @@ public TaskStateSnapshot retrieveLocalState(long checkpointID) {
     	 * Disposes the state of all local snapshots managed by this object.
     	 */
     	public void dispose() {
    +
    +		Collection<Map.Entry<Long, TaskStateSnapshot>> statesCopy;
    +
     		synchronized (lock) {
    -			for (Map.Entry<Long, TaskStateSnapshot> entry : storedTaskStateByCheckpointID.entrySet()) {
    -				discardStateObject(entry.getValue(), entry.getKey());
    -			}
     			discarded = true;
    +			statesCopy = new ArrayList<>(storedTaskStateByCheckpointID.entrySet());
     		}
    +
    +		discardExecutor.execute(() -> {
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168169207
  
    --- Diff: flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java ---
    @@ -859,56 +861,77 @@ public void run() {
     				if (asyncCheckpointState.compareAndSet(CheckpointingOperation.AsynCheckpointState.RUNNING,
     					CheckpointingOperation.AsynCheckpointState.COMPLETED)) {
     
    -					TaskStateSnapshot acknowledgedState = hasState ? taskOperatorSubtaskStates : null;
    -
    -					TaskStateManager taskStateManager = owner.getEnvironment().getTaskStateManager();
    -
    -					// we signal stateless tasks by reporting null, so that there are no attempts to assign empty state
    -					// to stateless tasks on restore. This enables simple job modifications that only concern
    -					// stateless without the need to assign them uids to match their (always empty) states.
    -					taskStateManager.reportTaskStateSnapshot(
    -						checkpointMetaData,
    -						checkpointMetrics,
    -						acknowledgedState);
    -
    -					LOG.debug("{} - finished asynchronous part of checkpoint {}. Asynchronous duration: {} ms",
    -						owner.getName(), checkpointMetaData.getCheckpointId(), asyncDurationMillis);
    -
    -					LOG.trace("{} - reported the following states in snapshot for checkpoint {}: {}.",
    -						owner.getName(), checkpointMetaData.getCheckpointId(), acknowledgedState);
    +					reportCompletedSnapshotStates(
    +						jobManagerTaskOperatorSubtaskStates,
    +						localTaskOperatorSubtaskStates,
    +						asyncDurationMillis);
     
     				} else {
     					LOG.debug("{} - asynchronous part of checkpoint {} could not be completed because it was closed before.",
     						owner.getName(),
     						checkpointMetaData.getCheckpointId());
     				}
     			} catch (Exception e) {
    -				// the state is completed if an exception occurred in the acknowledgeCheckpoint call
    -				// in order to clean up, we have to set it to RUNNING again.
    -				asyncCheckpointState.compareAndSet(
    -					CheckpointingOperation.AsynCheckpointState.COMPLETED,
    -					CheckpointingOperation.AsynCheckpointState.RUNNING);
    -
    -				try {
    -					cleanup();
    -				} catch (Exception cleanupException) {
    -					e.addSuppressed(cleanupException);
    -				}
    -
    -				Exception checkpointException = new Exception(
    -					"Could not materialize checkpoint " + checkpointId + " for operator " +
    -						owner.getName() + '.',
    -					e);
    -
    -				owner.asynchronousCheckpointExceptionHandler.tryHandleCheckpointException(
    -					checkpointMetaData,
    -					checkpointException);
    +				handleExecutionException(e);
     			} finally {
     				owner.cancelables.unregisterCloseable(this);
     				FileSystemSafetyNet.closeSafetyNetAndGuardedResourcesForThread();
     			}
     		}
     
    +		private void reportCompletedSnapshotStates(
    +			TaskStateSnapshot acknowledgedTaskStateSnapshot,
    +			TaskStateSnapshot localTaskStateSnapshot,
    +			long asyncDurationMillis) {
    +
    +			TaskStateManager taskStateManager = owner.getEnvironment().getTaskStateManager();
    +
    +			boolean hasAckState = acknowledgedTaskStateSnapshot.hasState();
    +			boolean hasLocalState = localTaskStateSnapshot.hasState();
    +
    +			Preconditions.checkState(hasAckState || !hasLocalState,
    +				"Found cached state but no corresponding primary state is reported to the job " +
    +					"manager. This indicates a problem.");
    +
    +			// we signal stateless tasks by reporting null, so that there are no attempts to assign empty state
    +			// to stateless tasks on restore. This enables simple job modifications that only concern
    +			// stateless without the need to assign them uids to match their (always empty) states.
    +			taskStateManager.reportTaskStateSnapshots(
    +				checkpointMetaData,
    +				checkpointMetrics,
    +				hasAckState ? acknowledgedTaskStateSnapshot : null,
    +				hasLocalState ? localTaskStateSnapshot : null);
    +
    +			LOG.debug("{} - finished asynchronous part of checkpoint {}. Asynchronous duration: {} ms",
    +				owner.getName(), checkpointMetaData.getCheckpointId(), asyncDurationMillis);
    +
    +			LOG.trace("{} - reported the following states in snapshot for checkpoint {}: {}.",
    +				owner.getName(), checkpointMetaData.getCheckpointId(), acknowledgedTaskStateSnapshot);
    +		}
    +
    +		private void handleExecutionException(Exception e) {
    +			// the state is completed if an exception occurred in the acknowledgeCheckpoint call
    --- End diff --
    
    👍 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by StefanRRichter <gi...@git.apache.org>.
Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168524862
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java ---
    @@ -525,19 +526,17 @@ public String toString() {
     	/**
     	 * This class encapsulates the configuration for local recovery of this backend.
     	 */
    -	public static final class LocalRecoveryConfig implements Serializable {
    +	public static final class LocalRecoveryConfig extends LocalRecoveryConfigBase {
     
    -		private static final long serialVersionUID = 1L;
    --- End diff --
    
    changed in later commit.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r167354757
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    --- End diff --
    
    long recovery time**~s~**


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by pnowojski <gi...@git.apache.org>.
Github user pnowojski commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r160923874
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -0,0 +1,123 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.JobID;
    +import org.apache.flink.runtime.jobgraph.JobVertexID;
    +import org.apache.flink.util.Preconditions;
    +
    +import javax.annotation.Nonnull;
    +
    +import java.util.HashMap;
    +import java.util.Map;
    +
    +/**
    + * This class holds the all {@link TaskLocalStateStore} objects for a task executor (manager).
    + *
    + * TODO: this still still work in progress and partially still acts as a placeholder.
    + */
    +public class TaskExecutorLocalStateStoresManager {
    --- End diff --
    
    Rename to just `LocalStateStoresManager`? `TaskExecutor` part kind of hardcodes the usage in the class name?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by bowenli86 <gi...@git.apache.org>.
Github user bowenli86 commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r167389639
  
    --- Diff: docs/ops/state/large_state_tuning.md ---
    @@ -234,4 +234,97 @@ Compression can be activated through the `ExecutionConfig`:
     **Notice:** The compression option has no impact on incremental snapshots, because they are using RocksDB's internal
     format which is always using snappy compression out of the box.
     
    +## Task-Local Recovery
    +
    +### Motivation
    +
    +In Flink's checkpointing, each task produces a snapshot of its state that is then written to a distributed store. Each task acknowledges
    +a successful write of the state to the job manager by sending a handle that describes the location of the state in the distributed store.
    +The job manager, in turn, collects the handles from all tasks and bundles them into a checkpoint object.
    +
    +In case of recovery, the job manager opens the latest checkpoint object and sends the handles back to the corresponding tasks, which can
    +then restore their state from the distributed storage. Using a distributed storage to store state has two important advantages. First, the storage
    +is fault tolerant and second, all state in the distributed store is accessible to all nodes and can be easily redistributed (e.g. for rescaling).
    +
    +However, using a remote distributed store has also one big disadvantage: all tasks must read their state from a remote location, over the network.
    +In many scenarios, recovery could reschedule failed tasks to the same task manager as in the previous run (of course there are exceptions like machine
    +failures), but we still have to read remote state. This can result in *long recovery times for large states*, even if there was only a small failure on
    +a single machine.
    +
    +### Approach
    +
    +Task-local state recovery targets exactly this problem of long recovery times and the main idea is the following: for every checkpoint, we do not
    +only write task states to the distributed storage, but also keep *a secondary copy of the state snapshot in a storage that is local to the task*
    +(e.g. on local disk or in memory). Notice that the primary store for snapshots must still be the distributed store, because local storage does not
    +ensure durability under node failures abd also does not provide access for other nodes to redistribute state, this functionality still requires the
    +primary copy.
    +
    +However, for each task that can be rescheduled to the previous location for recovery, we can restore state from the secondary, local
    +copy and avoid the costs of reading the state remotely. Given that *many failures are not node failures and node failures typically only affect one
    +or very few nodes at a time*, it is very likely that in a recovery most tasks can return to their previous location and find their local state intact.
    +This is what makes local recovery effective in reducing recovery time.
    +
    +Please note that this can come at some additional costs per checkpoint for creating and storing the secondary local state copy, depending on the
    +chosen state backend and checkpointing strategy. For example, in most cases the implementation will simply duplicate the writes to the distributed
    +store to a local file.
    +
    +<img src="../../fig/local_recovery.png" class="center" width="80%" alt="Illustration of checkpointing with task-local recovery."/>
    +
    +### Relationship of primary (distributed store) and secondary (task-local) state snapshots
    +
    +Task-local state is always considered a secondary copy, the ground truth of the checkpoint state is the primary copy in the distributed store. This
    +has implications for problems with local state during checkpointing and recovery:
    +
    +- For checkpointing, the *primary copy must be successful* and a failure to produce the *secondary, local copy will not fail* the checkpoint. A checkpoint
    +will fail if the primary copy could not be created, even if the secondary copy was successfully created.
    +
    +- Only the primary copy is acknowledged and managed by the job manager, secondary copies are owned by task managers and their life cycle can be
    +independent from their primary copy. For example, it is possible to retain a history of the 3 latest checkpoints as primary copies and only keep
    +the task-local state of the latest checkpoint.
    +
    +- For recovery, Flink will always *attempt to restore from task-local state first*, if a matching secondary copy is available. If any problem occurs during
    +the recovery from the secondary copy, Flink will *transparently retry to recovery the task from the primary copy*. Recovery only fails, if primary
    +and the (optional) secondary copy failed. In this case, depending on the configuration Flink could still fall back to an older checkpoint.
    +
    +- It is possible that the task-local copy contains only parts of the full task state (e.g. exception while writing one local file). In this case,
    +Flink will first try to recover local parts locally, non-local state is restored from the primary copy. Primary state must always be complete and is
    +a *superset of the task-local state*.
    +
    +- Task-local state can have a different format than the primary state, they are not required to be byte identical. For example, it could even possible that
    +the task-local state is an in-memory consisting of heap objects, and not stored in any files.
    +
    +- If a task manager is lost, the local state from all its task is lost.
    +
    +### Configuring task-local recovery
    +
    +Task-local recovery is *deactivated by default* and can be activated through Flink's configuration with the key `state.backend.local-recovery` as specified
    +in `CheckpointingOptions.LOCAL_RECOVERY`. Users have currently two choices:
    +
    +- `DISABLED`: Local recovery is disabled (default).
    +- `ENABLE_FILE_BASED`: Local recovery is activated, based on writing a secondary copy of the task state on local disk.
    +
    +### Details on task-local recovery for different state backends
    +
    +***Limitation**: Currently, task-local recovery only covers keyed state backends. Keyed state is typically by far the largest part of the state. In the near future, we will
    +also cover operator state and timers.*
    +
    +The following state backends can support task-local recovery.
    +
    +- FsStateBackend: task-local recovery is supported for keyed state. The implementation will duplicate the state to a local file. This can introduce additional write costs
    +and occupies local disk space. In the future, we might also offer an implementation that keeps task-local state in memory.
    +
    +- RocksDBStateBackend: task-local recovery is supported for keyed state. For *full checkpoints*, state is duplicated to a local file. This can introduce additional write costs
    +and occupies local disk space. For *incremental snapshots*, the local state is based on RocksDB's native checkpointing mechanism. This mechanism is also used as the first step
    +to create the primary copy, which means that in this case no additional cost is introduced for create the secondary copy. We simply keep the native checkpoint directory around
    +instead of deleting it after uploading to the distributed store. This local copy can share active files with the working directory of RocksDB (via hard links), so for active
    +files also no additional disk space is consumed for task-local recovery with incremental snapshots.
    +
    +### Allocation-preserving scheduling
    +
    +Task-local recovery assumes allocation-preserving task scheduling under failures, which was introduced as part of FLIP-6 and works as follows. Each task remembers its previous
    +allocation and *requests the exact same slot* to restart in recovery. If this slot is not available, the task will request a *new, fresh slot* from the resource manager. This way,
    +if a task manager is no longer available, a task that cannot return to its previous location *will not drive other recovering tasks out of their previous slots*. Our reasoning is
    +that the previous slot can only disappear when a task manager is no longer available, and in this case *some* task has to request a new slot anyways. With our scheduling strategy
    --- End diff --
    
    some task**s** **have** to 


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r169724637
  
    --- Diff: flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java ---
    @@ -2123,9 +2162,21 @@ void takeSnapshot() throws Exception {
     				checkpointId,
     				sstFiles,
     				miscFiles,
    -				metaStateHandle);
    +				metaStateHandle.getJobManagerOwnedSnapshot());
    +
    +			DirectoryStateHandle directoryStateHandle = localBackupDirectory.completeSnapshotAndGetHandle();
    +			StreamStateHandle taskLocalSnapshotMetaDataStateHandle = metaStateHandle.getTaskLocalSnapshot();
    +			IncrementalLocalKeyedStateHandle directoryKeyedStateHandle =
    +				directoryStateHandle != null && taskLocalSnapshotMetaDataStateHandle != null ?
    +					new IncrementalLocalKeyedStateHandle(
    +						stateBackend.backendUID,
    +						checkpointId,
    +						directoryStateHandle,
    +						stateBackend.keyGroupRange,
    +						taskLocalSnapshotMetaDataStateHandle) :
    +					null;
     
    -			return new SnapshotResult<>(incrementalKeyedStateHandle, null);
    +			return new SnapshotResult<>(incrementalKeyedStateHandle, directoryKeyedStateHandle);
    --- End diff --
    
    I guess you changed the creation of the `SnapshotResult` in a fixup commit in your local branch, right? Otherwise we might refactor this in order to get rid of the many `nulls` in the lines above.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168501126
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/TaskExecutorLocalStateStoresManager.java ---
    @@ -20,35 +20,47 @@
     
     import org.apache.flink.annotation.VisibleForTesting;
     import org.apache.flink.api.common.JobID;
    +import org.apache.flink.runtime.clusterframework.types.AllocationID;
     import org.apache.flink.runtime.jobgraph.JobVertexID;
    -import org.apache.flink.util.ExceptionUtils;
    -import org.apache.flink.util.Preconditions;
     
    +import org.slf4j.Logger;
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnegative;
     import javax.annotation.Nonnull;
     
     import java.io.File;
    -import java.util.HashMap;
     import java.util.Map;
    +import java.util.concurrent.ConcurrentHashMap;
    +import java.util.concurrent.Executor;
     
     /**
      * This class holds the all {@link TaskLocalStateStore} objects for a task executor (manager).
    - *
    - * TODO: this still still work in progress and partially still acts as a placeholder.
      */
     public class TaskExecutorLocalStateStoresManager {
     
    +	/** Logger for this class. */
    +	private static final Logger LOG = LoggerFactory.getLogger(TaskExecutorLocalStateStoresManager.class);
    +
     	/**
     	 * This map holds all local state stores for tasks running on the task manager / executor that own the instance of
    -	 * this.
    +	 * this. Maps from allocation id to all the subtask's local state stores.
     	 */
    -	private final Map<JobID, Map<JobVertexSubtaskKey, TaskLocalStateStore>> taskStateStoresByJobID;
    +	private final Map<AllocationID, Map<JobVertexSubtaskKey, TaskLocalStateStore>> taskStateStoresByAllocationID;
     
     	/** This is the root directory for all local state of this task manager / executor. */
     	private final File[] localStateRootDirectories;
     
    -	public TaskExecutorLocalStateStoresManager(File[] localStateRootDirectories) {
    -		this.taskStateStoresByJobID = new HashMap<>();
    -		this.localStateRootDirectories = Preconditions.checkNotNull(localStateRootDirectories);
    +	/** Executor that runs the discarding of released state objects. */
    +	private final Executor discardExecutor;
    +
    +	public TaskExecutorLocalStateStoresManager(
    +		@Nonnull File[] localStateRootDirectories,
    +		@Nonnull Executor discardExecutor) {
    --- End diff --
    
    Are these `@Nonnull` annotations enforced? Meaning are we executing `FindBugs` which fails if on passes `null` to the `TaskExecutorLocalStateStoresManager`?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168472392
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java ---
    @@ -507,17 +509,78 @@ public boolean supportsAsynchronousSnapshots() {
     		return localRecoveryConfig;
     	}
     
    +	private interface SnapshotStrategySynchronicityBehavior<K> {
    +
    +		default void finalizeSnapshotBeforeReturnHook(Runnable runnable) {
    +
    +		}
    +
    +		default void logOperationCompleted(CheckpointStreamFactory streamFactory, long startTime) {
    +
    +		}
    +
    +		boolean isAsynchronous();
    +
    +		<N, V> StateTable<K, N, V> newStateTable(RegisteredKeyedBackendStateMetaInfo<N, V> newMetaInfo);
    --- End diff --
    
    I think this method should be part of the `SnapshotStrategy` interface. That way we would make the `SnapshotStrategySynchronicityBehavior` sub classes independent of the `HeapKeyedStateBackend`.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170037177
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalRecoveryConfig.java ---
    @@ -0,0 +1,120 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + * http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.apache.flink.api.common.io.ParseException;
    +import org.apache.flink.configuration.CheckpointingOptions;
    +import org.apache.flink.configuration.Configuration;
    +
    +import org.slf4j.LoggerFactory;
    +
    +import javax.annotation.Nonnull;
    +
    +/**
    + * This class encapsulates the completed configuration for local recovery, i.e. the root
    + * directories into which all file-based snapshots can be written and the general mode for the local recover feature.
    + */
    +public class LocalRecoveryConfig {
    +
    +	/**
    +	 * Enum over modes of local recovery:
    +	 * <p><ul>
    +	 * <li>DISABLED: disables local recovery.
    +	 * <li>ENABLE_FILE_BASED: enables local recovery in a variant that is based on local files.
    +	 * </ul>
    +	 */
    +	public enum LocalRecoveryMode {
    +		DISABLED,
    +		ENABLE_FILE_BASED;
    +
    +		/**
    +		 * Attempts to parses the given string to a {@link LocalRecoveryMode}.
    +		 *
    +		 * @param input string to parse
    +		 * @return the {@link LocalRecoveryMode} that is represented by the input string.
    +		 * @throws ParseException if the string cannot be parsed to a {@link LocalRecoveryMode}.
    +		 */
    +		@Nonnull
    +		public static LocalRecoveryMode fromString(String input) throws ParseException {
    +			if (input != null) {
    +				String trimInput = input.trim();
    +				if (DISABLED.toString().equalsIgnoreCase(trimInput)) {
    +					return DISABLED;
    +				} else if (ENABLE_FILE_BASED.toString().equalsIgnoreCase(trimInput)) {
    +					return ENABLE_FILE_BASED;
    +				}
    +			}
    +			throw new ParseException("Cannot parse input to LocalRecoveryMode: " + input);
    +		}
    +
    +		/**
    +		 * Extracts the {@link LocalRecoveryMode} from the given configuration. Defaults to LocalRecoveryMode.DISABLED
    +		 * if no configuration value is specified or parsing the value resulted in an exception.
    +		 *
    +		 * @param configuration the configuration that specifies the value for the local recovery mode.
    +		 * @return the local recovery mode as found in the config, or LocalRecoveryMode.DISABLED if no mode was
    +		 * configured or the specified mode could not be parsed.
    +		 */
    +		@Nonnull
    +		public static LocalRecoveryMode fromConfig(@Nonnull Configuration configuration) {
    +			String localRecoveryConfString = configuration.getString(CheckpointingOptions.LOCAL_RECOVERY);
    +			try {
    +				return LocalRecoveryConfig.LocalRecoveryMode.fromString(localRecoveryConfString);
    +			} catch (ParseException ex) {
    +				LoggerFactory.getLogger(LocalRecoveryConfig.class).warn(
    +					"Exception while parsing configuration of local recovery mode. Local recovery will be disabled.",
    +					ex);
    +				return LocalRecoveryConfig.LocalRecoveryMode.DISABLED;
    +			}
    +		}
    +	}
    +
    +	/** The local recovery mode. */
    +	@Nonnull
    +	private final LocalRecoveryMode localRecoveryMode;
    --- End diff --
    
    This is more of a general question than a review comment: How will it look in the future if we want to support a non file based local recovery? E.g. writing some things to a local DB, for example. Then we would not need the `LocalRecoveryDirectoryProvider` and we would have to adapt how we create the duplicating streams. So I'm wondering whether we should not have something like a `LocalStateStreamFactory` instead. For file based recovery it could wrap the `LocalRecoveryDirectoryProvider` to select the next file to write into. What are the plans for that in the future?


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r170033424
  
    --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/taskexecutor/TaskManagerServicesConfiguration.java ---
    @@ -193,7 +193,12 @@ public static TaskManagerServicesConfiguration fromConfiguration(
     		}
     
     		final String[] tmpDirs = ConfigurationUtils.parseTempDirectories(configuration);
    -		final String[] localStateRootDir = ConfigurationUtils.parseLocalStateDirectories(configuration);
    +		String[] localStateRootDir = ConfigurationUtils.parseLocalStateDirectories(configuration);
    --- End diff --
    
    Could also be `final`


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r165321013
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/TestTaskStateManager.java ---
    @@ -67,8 +69,11 @@ public TestTaskStateManager(
     		this.jobId = jobId;
     		this.executionAttemptID = executionAttemptID;
     		this.checkpointResponder = checkpointResponder;
    -		this.taskStateSnapshotsByCheckpointId = new HashMap<>();
    +		this.jobManagerTaskStateSnapshotsByCheckpointId = new HashMap<>();
    +		this.taskManagerTaskStateSnapshotsByCheckpointId = new HashMap<>();
     		this.reportedCheckpointId = -1L;
    +		this.subtaskLocalStateBaseDirectory =
    +			new File(System.getProperty("java.io.tmpdir"), "testLocalState_" + UUID.randomUUID());
    --- End diff --
    
    Shouldn't this come from something like a `TemporaryFolder` such that the directory get's cleaned up in all cases. Consequently, maybe we should pass in the `subtaskLocalStateBaseDirectory` from the test using this class.


---

[GitHub] flink pull request #5239: [FLINK-8360] Implement task-local state recovery

Posted by tillrohrmann <gi...@git.apache.org>.
Github user tillrohrmann commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5239#discussion_r168482439
  
    --- Diff: flink-runtime/src/test/java/org/apache/flink/runtime/state/LocalRecoveryDirectoryProviderTest.java ---
    @@ -0,0 +1,95 @@
    +/*
    + * Licensed to the Apache Software Foundation (ASF) under one
    + * or more contributor license agreements.  See the NOTICE file
    + * distributed with this work for additional information
    + * regarding copyright ownership.  The ASF licenses this file
    + * to you under the Apache License, Version 2.0 (the
    + * "License"); you may not use this file except in compliance
    + * with the License.  You may obtain a copy of the License at
    + *
    + *     http://www.apache.org/licenses/LICENSE-2.0
    + *
    + * Unless required by applicable law or agreed to in writing, software
    + * distributed under the License is distributed on an "AS IS" BASIS,
    + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    + * See the License for the specific language governing permissions and
    + * limitations under the License.
    + */
    +
    +package org.apache.flink.runtime.state;
    +
    +import org.junit.After;
    +import org.junit.Assert;
    +import org.junit.Before;
    +import org.junit.Test;
    +import org.junit.rules.TemporaryFolder;
    +
    +import java.io.File;
    +import java.io.IOException;
    +
    +public class LocalRecoveryDirectoryProviderTest {
    --- End diff --
    
    `TestLogger` missing.


---