This is an automated email from the ASF dual-hosted git repository.
dwysakowicz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 93c726bc90609d31f07494344ccaaed4d2c7e7c5
Author: Dawid Wysakowicz <dwysakow...@apache.org>
AuthorDate: Wed Jan 13 17:20:29 2021 +0100
[FLINK-20978] Extract common restoring logic from RocksDB
---
.../restore/FullSnapshotRestoreOperation.java | 323 +++++++++++++++++++++
.../flink/runtime/state/restore/KeyGroup.java | 41 +++
.../flink/runtime/state/restore/KeyGroupEntry.java | 47 +++
.../state/restore/SavepointRestoreResult.java | 44 +++
.../runtime/state/restore/ThrowingIterator.java | 52 ++++
.../state/restore/RocksDBFullRestoreOperation.java | 175 +++--------
6 files changed, 549 insertions(+), 133 deletions(-)
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/FullSnapshotRestoreOperation.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/FullSnapshotRestoreOperation.java
new file mode 100644
index 0000000..8147ea5
--- /dev/null
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/FullSnapshotRestoreOperation.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.runtime.state.restore;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.common.typeutils.TypeSerializerSchemaCompatibility;
+import
org.apache.flink.api.common.typeutils.base.array.BytePrimitiveArraySerializer;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.memory.DataInputView;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.KeyedBackendSerializationProxy;
+import org.apache.flink.runtime.state.KeyedStateHandle;
+import org.apache.flink.runtime.state.RestoreOperation;
+import org.apache.flink.runtime.state.SnappyStreamCompressionDecorator;
+import org.apache.flink.runtime.state.StateSerializerProvider;
+import org.apache.flink.runtime.state.StreamCompressionDecorator;
+import org.apache.flink.runtime.state.UncompressedStreamCompressionDecorator;
+import org.apache.flink.util.Preconditions;
+import org.apache.flink.util.StateMigrationException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.Nonnull;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+import static
org.apache.flink.runtime.state.FullSnapshotUtil.END_OF_KEY_GROUP_MARK;
+import static
org.apache.flink.runtime.state.FullSnapshotUtil.clearMetaDataFollowsFlag;
+import static
org.apache.flink.runtime.state.FullSnapshotUtil.hasMetaDataFollowsFlag;
+import static
org.apache.flink.runtime.state.StateUtil.unexpectedStateHandleException;
+
+/**
+ * The procedure of restoring state from a savepoint written with the unified
binary format. All
+ * state backends should support restoring from that format.
+ *
+ * <p>The format was adopted from the RocksDB state backend. It is as follows:
+ *
+ * <pre>
+ *
+-------------------------------------------------------------+
+ * | Keyed Backend Meta Information
|
+ * Meta Information
+--------------------+-----+----------------------------------+
+ * | State Meta 0 | ... | State Meta M
|
+ *
+-------------------------------------------------------------+
+ * | State ID (short)
|
+ * State 0
+--------------------+-----+----------------------------------+
+ * | State (K,V) pair 0 | ... | State (K,V) pair X (flipped
MSB) |
+ *
+--------------------+-----+----------------------------------+
+ * | State ID (short)
|
+ * State 1
+--------------------+-----+----------------------------------+
+ * | State (K,V) pair 0 | ... | State (K,V) pair X (flipped
MSB) |
+ *
+--------------------+-----+----------------------------------+
+ * | State ID (short)
|
+ *
+--------------------+-----+----------------------------------+
+ * State M | State (K,V) pair 0 | ... | State (K,V) pair X (flipped
MSB) |
+ *
+--------------------+-----+----------------------------------+
+ * | END_OF_KEY_GROUP_MARK (0xFFFF)
|
+ *
+-------------------------------------------------------------+
+ * </pre>
+ *
+ * <p>Additionally the format of the (K,V) pairs differs slightly depending on
the type of the state
+ * object:
+ *
+ * <pre>
+ * +------------------+-------------------------------------+
+ * | ValueState | [CompositeKey(KG, K, NS), SV] |
+ * |------------------+-------------------------------------+
+ * | ListState | [CompositeKey(KG, K, NS), SV] |
+ * +------------------+-------------------------------------+
+ * | MapState | [CompositeKey(KG, K, NS) :: UK, UV] |
+ * +------------------+-------------------------------------+
+ * | AggregatingState | [CompositeKey(KG, K, NS), SV] |
+ * +------------------+-------------------------------------+
+ * | ReducingState | [CompositeKey(KG, K, NS), SV] |
+ * | +-------------------------------------+
+ * | | [CompositeKey(KG, K, NS), SV] |
+ * +------------------+-------------------------------------+
+ * | Timers | [KG :: TS :: K :: NS, (empty)] |
+ * +------------------+-------------------------------------+
+ * </pre>
+ *
+ * <p>For detailed information see FLIP-41:
https://cwiki.apache.org/confluence/x/VJDiBg
+ *
+ * @param <K> The data type of the key.
+ */
+@Internal
+public class FullSnapshotRestoreOperation<K>
+ implements RestoreOperation<ThrowingIterator<SavepointRestoreResult>> {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(FullSnapshotRestoreOperation.class);
+
+ private final KeyGroupRange keyGroupRange;
+ private final ClassLoader userCodeClassLoader;
+ private final Collection<KeyedStateHandle> restoreStateHandles;
+ private final StateSerializerProvider<K> keySerializerProvider;
+
+ private boolean isKeySerializerCompatibilityChecked;
+
+ public FullSnapshotRestoreOperation(
+ KeyGroupRange keyGroupRange,
+ ClassLoader userCodeClassLoader,
+ Collection<KeyedStateHandle> restoreStateHandles,
+ StateSerializerProvider<K> keySerializerProvider) {
+ this.keyGroupRange = keyGroupRange;
+ this.userCodeClassLoader = userCodeClassLoader;
+ this.restoreStateHandles =
+
restoreStateHandles.stream().filter(Objects::nonNull).collect(Collectors.toList());
+ this.keySerializerProvider = keySerializerProvider;
+ }
+
+ @Override
+ public ThrowingIterator<SavepointRestoreResult> restore()
+ throws IOException, StateMigrationException {
+ return new ThrowingIterator<SavepointRestoreResult>() {
+
+ private final Iterator<KeyedStateHandle> keyedStateHandlesIterator
=
+ restoreStateHandles.iterator();
+
+ @Override
+ public boolean hasNext() {
+ return keyedStateHandlesIterator.hasNext();
+ }
+
+ @Override
+ public SavepointRestoreResult next() throws IOException,
StateMigrationException {
+ KeyedStateHandle keyedStateHandle =
keyedStateHandlesIterator.next();
+ if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
+ throw unexpectedStateHandleException(
+ KeyGroupsStateHandle.class,
keyedStateHandle.getClass());
+ }
+ KeyGroupsStateHandle groupsStateHandle =
(KeyGroupsStateHandle) keyedStateHandle;
+ return restoreKeyGroupsInStateHandle(groupsStateHandle);
+ }
+
+ @Override
+ public void close() {}
+ };
+ }
+
+ private SavepointRestoreResult restoreKeyGroupsInStateHandle(
+ @Nonnull KeyGroupsStateHandle keyedStateHandle)
+ throws IOException, StateMigrationException {
+ FSDataInputStream currentStateHandleInStream =
keyedStateHandle.openInputStream();
+ KeyedBackendSerializationProxy<K> serializationProxy =
+ readMetaData(new
DataInputViewStreamWrapper(currentStateHandleInStream));
+ KeyGroupsIterator groupsIterator =
+ new KeyGroupsIterator(
+ keyGroupRange,
+ keyedStateHandle,
+ currentStateHandleInStream,
+ serializationProxy.isUsingKeyGroupCompression()
+ ? SnappyStreamCompressionDecorator.INSTANCE
+ :
UncompressedStreamCompressionDecorator.INSTANCE);
+
+ return new SavepointRestoreResult(
+ serializationProxy.getStateMetaInfoSnapshots(),
groupsIterator);
+ }
+
+ private KeyedBackendSerializationProxy<K> readMetaData(DataInputView
dataInputView)
+ throws IOException, StateMigrationException {
+ // isSerializerPresenceRequired flag is set to false, since for the
RocksDB state backend,
+ // deserialization of state happens lazily during runtime; we depend
on the fact
+ // that the new serializer for states could be compatible, and
therefore the restore can
+ // continue
+ // without old serializers required to be present.
+ KeyedBackendSerializationProxy<K> serializationProxy =
+ new KeyedBackendSerializationProxy<>(userCodeClassLoader);
+ serializationProxy.read(dataInputView);
+ if (!isKeySerializerCompatibilityChecked) {
+ // fetch current serializer now because if it is incompatible, we
can't access
+ // it anymore to improve the error message
+ TypeSerializer<K> currentSerializer =
keySerializerProvider.currentSchemaSerializer();
+ // check for key serializer compatibility; this also reconfigures
the
+ // key serializer to be compatible, if it is required and is
possible
+ TypeSerializerSchemaCompatibility<K> keySerializerSchemaCompat =
+
keySerializerProvider.setPreviousSerializerSnapshotForRestoredState(
+ serializationProxy.getKeySerializerSnapshot());
+ if (keySerializerSchemaCompat.isCompatibleAfterMigration()
+ || keySerializerSchemaCompat.isIncompatible()) {
+ throw new StateMigrationException(
+ "The new key serializer ("
+ + currentSerializer
+ + ") must be compatible with the previous key
serializer ("
+ +
keySerializerProvider.previousSchemaSerializer()
+ + ").");
+ }
+
+ isKeySerializerCompatibilityChecked = true;
+ }
+
+ return serializationProxy;
+ }
+
+ private static class KeyGroupsIterator implements
ThrowingIterator<KeyGroup> {
+ @Nonnull private final KeyGroupRange keyGroupRange;
+ @Nonnull private final Iterator<Tuple2<Integer, Long>> keyGroups;
+ @Nonnull private final FSDataInputStream currentStateHandleInStream;
+ @Nonnull private final StreamCompressionDecorator
keygroupStreamCompressionDecorator;
+ @Nonnull private final KeyGroupsStateHandle
currentKeyGroupsStateHandle;
+
+ private KeyGroupsIterator(
+ @Nonnull KeyGroupRange keyGroupRange,
+ @Nonnull KeyGroupsStateHandle currentKeyGroupsStateHandle,
+ @Nonnull FSDataInputStream currentStateHandleInStream,
+ @Nonnull StreamCompressionDecorator
keygroupStreamCompressionDecorator) {
+ this.keyGroupRange = keyGroupRange;
+ this.keyGroups =
currentKeyGroupsStateHandle.getGroupRangeOffsets().iterator();
+ this.currentStateHandleInStream = currentStateHandleInStream;
+ this.keygroupStreamCompressionDecorator =
keygroupStreamCompressionDecorator;
+ this.currentKeyGroupsStateHandle = currentKeyGroupsStateHandle;
+ LOG.info("Starting to restore from state handle: {}.",
currentKeyGroupsStateHandle);
+ }
+
+ public boolean hasNext() {
+ return keyGroups.hasNext();
+ }
+
+ public KeyGroup next() throws IOException {
+ Tuple2<Integer, Long> keyGroupOffset = keyGroups.next();
+ int keyGroup = keyGroupOffset.f0;
+
+ // Check that restored key groups all belong to the backend
+ Preconditions.checkState(
+ keyGroupRange.contains(keyGroup), "The key group must
belong to the backend");
+
+ long offset = keyGroupOffset.f1;
+ if (0L != offset) {
+ currentStateHandleInStream.seek(offset);
+ InputStream compressedKgIn =
+
keygroupStreamCompressionDecorator.decorateWithCompression(
+ currentStateHandleInStream);
+ DataInputViewStreamWrapper compressedKgInputView =
+ new DataInputViewStreamWrapper(compressedKgIn);
+ return new KeyGroup(keyGroup, new
KeyGroupEntriesIterator(compressedKgInputView));
+ } else {
+ return new KeyGroup(keyGroup, new KeyGroupEntriesIterator());
+ }
+ }
+
+ @Override
+ public void close() throws IOException {
+ currentStateHandleInStream.close();
+ LOG.info("Finished restoring from state handle: {}.",
currentKeyGroupsStateHandle);
+ }
+ }
+
+ private static class KeyGroupEntriesIterator implements
ThrowingIterator<KeyGroupEntry> {
+ private final DataInputViewStreamWrapper kgInputView;
+ private Integer currentKvStateId;
+
+ private KeyGroupEntriesIterator(@Nonnull DataInputViewStreamWrapper
kgInputView)
+ throws IOException {
+ this.kgInputView = kgInputView;
+ this.currentKvStateId = END_OF_KEY_GROUP_MARK &
kgInputView.readShort();
+ }
+
+ // creates an empty iterator
+ private KeyGroupEntriesIterator() {
+ this.kgInputView = null;
+ this.currentKvStateId = null;
+ }
+
+ public boolean hasNext() {
+ return currentKvStateId != null;
+ }
+
+ public KeyGroupEntry next() throws IOException {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+
+ byte[] key =
BytePrimitiveArraySerializer.INSTANCE.deserialize(kgInputView);
+ byte[] value =
BytePrimitiveArraySerializer.INSTANCE.deserialize(kgInputView);
+ final int entryStateId = currentKvStateId;
+ if (hasMetaDataFollowsFlag(key)) {
+ // clear the signal bit in the key to make it ready for
insertion
+ // again
+ clearMetaDataFollowsFlag(key);
+ // TODO this could be aware of keyGroupPrefixBytes and write
only
+ // one byte if possible
+ currentKvStateId = END_OF_KEY_GROUP_MARK &
kgInputView.readShort();
+ if (END_OF_KEY_GROUP_MARK == currentKvStateId) {
+ currentKvStateId = null;
+ }
+ }
+
+ return new KeyGroupEntry(entryStateId, key, value);
+ }
+
+ @Override
+ public void close() throws IOException {
+ if (kgInputView != null) {
+ kgInputView.close();
+ }
+ }
+ }
+}
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/KeyGroup.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/KeyGroup.java
new file mode 100644
index 0000000..f321ca5
--- /dev/null
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/KeyGroup.java
@@ -0,0 +1,41 @@
+/*
+ * 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.restore;
+
+import org.apache.flink.annotation.Internal;
+
+/** Part of a savepoint representing data for a single key group. */
+@Internal
+public class KeyGroup {
+ private final int keyGroupId;
+ private final ThrowingIterator<KeyGroupEntry> keyGroupEntries;
+
+ KeyGroup(int keyGroupId, ThrowingIterator<KeyGroupEntry> keyGroupEntries) {
+ this.keyGroupId = keyGroupId;
+ this.keyGroupEntries = keyGroupEntries;
+ }
+
+ public int getKeyGroupId() {
+ return keyGroupId;
+ }
+
+ public ThrowingIterator<KeyGroupEntry> getKeyGroupEntries() {
+ return keyGroupEntries;
+ }
+}
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/KeyGroupEntry.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/KeyGroupEntry.java
new file mode 100644
index 0000000..afc2c72
--- /dev/null
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/KeyGroupEntry.java
@@ -0,0 +1,47 @@
+/*
+ * 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.restore;
+
+import org.apache.flink.annotation.Internal;
+
+/** Part of a savepoint representing data for a single state entry in a key
group. */
+@Internal
+public class KeyGroupEntry {
+ private final int kvStateId;
+ private final byte[] key;
+ private final byte[] value;
+
+ KeyGroupEntry(int kvStateId, byte[] key, byte[] value) {
+ this.kvStateId = kvStateId;
+ this.key = key;
+ this.value = value;
+ }
+
+ public int getKvStateId() {
+ return kvStateId;
+ }
+
+ public byte[] getKey() {
+ return key;
+ }
+
+ public byte[] getValue() {
+ return value;
+ }
+}
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/SavepointRestoreResult.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/SavepointRestoreResult.java
new file mode 100644
index 0000000..2f46ee5
--- /dev/null
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/SavepointRestoreResult.java
@@ -0,0 +1,44 @@
+/*
+ * 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.restore;
+
+import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot;
+
+import java.util.List;
+
+/** A result from restoring a single {@link
org.apache.flink.runtime.state.KeyedStateHandle}. */
+public class SavepointRestoreResult {
+ private final List<StateMetaInfoSnapshot> stateMetaInfoSnapshots;
+ private final ThrowingIterator<KeyGroup> restoredKeyGroups;
+
+ public SavepointRestoreResult(
+ List<StateMetaInfoSnapshot> stateMetaInfoSnapshots,
+ ThrowingIterator<KeyGroup> groupsIterator) {
+ this.stateMetaInfoSnapshots = stateMetaInfoSnapshots;
+ this.restoredKeyGroups = groupsIterator;
+ }
+
+ public List<StateMetaInfoSnapshot> getStateMetaInfoSnapshots() {
+ return stateMetaInfoSnapshots;
+ }
+
+ public ThrowingIterator<KeyGroup> getRestoredKeyGroups() {
+ return restoredKeyGroups;
+ }
+}
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/ThrowingIterator.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/ThrowingIterator.java
new file mode 100644
index 0000000..f2bd986
--- /dev/null
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/state/restore/ThrowingIterator.java
@@ -0,0 +1,52 @@
+/*
+ * 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.restore;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.util.StateMigrationException;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.NoSuchElementException;
+
+/**
+ * Modified {@link java.util.Iterator} interface that lets the implementation
throw an {@link
+ * IOException}.
+ *
+ * @param <E> the type of elements returned by this iterator
+ * @see java.util.Iterator
+ */
+@Internal
+public interface ThrowingIterator<E> extends Closeable {
+ /**
+ * Returns {@code true} if the iteration has more elements. (In other
words, returns {@code
+ * true} if {@link #next} would return an element rather than throwing an
exception.)
+ *
+ * @return {@code true} if the iteration has more elements
+ */
+ boolean hasNext();
+
+ /**
+ * Returns the next element in the iteration.
+ *
+ * @return the next element in the iteration
+ * @throws NoSuchElementException if the iteration has no more elements
+ */
+ E next() throws IOException, StateMigrationException;
+}
diff --git
a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/restore/RocksDBFullRestoreOperation.java
b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/restore/RocksDBFullRestoreOperation.java
index d6c066b..6c98d9b 100644
---
a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/restore/RocksDBFullRestoreOperation.java
+++
b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/restore/RocksDBFullRestoreOperation.java
@@ -18,28 +18,21 @@
package org.apache.flink.contrib.streaming.state.restore;
-import
org.apache.flink.api.common.typeutils.base.array.BytePrimitiveArraySerializer;
-import org.apache.flink.api.java.tuple.Tuple2;
import
org.apache.flink.contrib.streaming.state.RocksDBKeyedStateBackend.RocksDbKvStateInfo;
import org.apache.flink.contrib.streaming.state.RocksDBNativeMetricOptions;
import org.apache.flink.contrib.streaming.state.RocksDBWriteBatchWrapper;
import
org.apache.flink.contrib.streaming.state.ttl.RocksDbTtlCompactFiltersManager;
import org.apache.flink.core.fs.CloseableRegistry;
-import org.apache.flink.core.fs.FSDataInputStream;
-import org.apache.flink.core.memory.DataInputView;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.state.KeyGroupRange;
-import org.apache.flink.runtime.state.KeyGroupsStateHandle;
-import org.apache.flink.runtime.state.KeyedBackendSerializationProxy;
import org.apache.flink.runtime.state.KeyedStateHandle;
-import org.apache.flink.runtime.state.SnappyStreamCompressionDecorator;
import org.apache.flink.runtime.state.StateSerializerProvider;
-import org.apache.flink.runtime.state.StreamCompressionDecorator;
-import org.apache.flink.runtime.state.UncompressedStreamCompressionDecorator;
import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot;
-import org.apache.flink.util.IOUtils;
-import org.apache.flink.util.Preconditions;
+import org.apache.flink.runtime.state.restore.FullSnapshotRestoreOperation;
+import org.apache.flink.runtime.state.restore.KeyGroup;
+import org.apache.flink.runtime.state.restore.KeyGroupEntry;
+import org.apache.flink.runtime.state.restore.SavepointRestoreResult;
+import org.apache.flink.runtime.state.restore.ThrowingIterator;
import org.apache.flink.util.StateMigrationException;
import org.rocksdb.ColumnFamilyHandle;
@@ -52,38 +45,17 @@ import javax.annotation.Nonnull;
import java.io.File;
import java.io.IOException;
-import java.io.InputStream;
-import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
+import java.util.stream.Collectors;
-import static
org.apache.flink.runtime.state.FullSnapshotUtil.END_OF_KEY_GROUP_MARK;
-import static
org.apache.flink.runtime.state.FullSnapshotUtil.clearMetaDataFollowsFlag;
-import static
org.apache.flink.runtime.state.FullSnapshotUtil.hasMetaDataFollowsFlag;
-import static
org.apache.flink.runtime.state.StateUtil.unexpectedStateHandleException;
import static org.apache.flink.util.Preconditions.checkArgument;
/** Encapsulates the process of restoring a RocksDB instance from a full
snapshot. */
public class RocksDBFullRestoreOperation<K> extends
AbstractRocksDBRestoreOperation<K> {
- /** 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;
-
+ private final FullSnapshotRestoreOperation<K> savepointRestoreOperation;
/** Write batch size used in {@link RocksDBWriteBatchWrapper}. */
private final long writeBatchSize;
@@ -124,6 +96,12 @@ public class RocksDBFullRestoreOperation<K> extends
AbstractRocksDBRestoreOperat
writeBufferManagerCapacity);
checkArgument(writeBatchSize >= 0, "Write batch size have to be no
negative.");
this.writeBatchSize = writeBatchSize;
+ this.savepointRestoreOperation =
+ new FullSnapshotRestoreOperation<>(
+ keyGroupRange,
+ userCodeClassLoader,
+ restoreStateHandles,
+ keySerializerProvider);
}
/** Restores all key-groups data that is referenced by the passed state
handles. */
@@ -131,60 +109,32 @@ public class RocksDBFullRestoreOperation<K> extends
AbstractRocksDBRestoreOperat
public RocksDBRestoreResult restore()
throws IOException, StateMigrationException, RocksDBException {
openDB();
- for (KeyedStateHandle keyedStateHandle : restoreStateHandles) {
- if (keyedStateHandle != null) {
-
- if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
- throw unexpectedStateHandleException(
- KeyGroupsStateHandle.class,
keyedStateHandle.getClass());
- }
- this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle)
keyedStateHandle;
- restoreKeyGroupsInStateHandle();
+ try (ThrowingIterator<SavepointRestoreResult> restore =
+ savepointRestoreOperation.restore()) {
+ while (restore.hasNext()) {
+ applyRestoreResult(restore.next());
}
}
return new RocksDBRestoreResult(
this.db, defaultColumnFamilyHandle, nativeMetricMonitor, -1,
null, null);
}
- /** Restore one key groups state handle. */
- private void restoreKeyGroupsInStateHandle()
- throws IOException, StateMigrationException, RocksDBException {
- try {
- logger.info("Starting to restore from state handle: {}.",
currentKeyGroupsStateHandle);
- currentStateHandleInStream =
currentKeyGroupsStateHandle.openInputStream();
- cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
- currentStateHandleInView = new
DataInputViewStreamWrapper(currentStateHandleInStream);
- restoreKVStateMetaData();
- restoreKVStateData();
- logger.info("Finished restoring from state handle: {}.",
currentKeyGroupsStateHandle);
- } finally {
- if
(cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
- IOUtils.closeQuietly(currentStateHandleInStream);
- }
- }
- }
-
- /**
- * Restore the KV-state / ColumnFamily meta data for all key-groups
referenced by the current
- * state handle.
- */
- private void restoreKVStateMetaData() throws IOException,
StateMigrationException {
- KeyedBackendSerializationProxy<K> serializationProxy =
- readMetaData(currentStateHandleInView);
-
- this.keygroupStreamCompressionDecorator =
- serializationProxy.isUsingKeyGroupCompression()
- ? SnappyStreamCompressionDecorator.INSTANCE
- : UncompressedStreamCompressionDecorator.INSTANCE;
-
+ private void applyRestoreResult(SavepointRestoreResult
savepointRestoreResult)
+ throws IOException, RocksDBException, StateMigrationException {
List<StateMetaInfoSnapshot> restoredMetaInfos =
- serializationProxy.getStateMetaInfoSnapshots();
- currentStateHandleKVStateColumnFamilies = new
ArrayList<>(restoredMetaInfos.size());
-
- for (StateMetaInfoSnapshot restoredMetaInfo : restoredMetaInfos) {
- RocksDbKvStateInfo registeredStateCFHandle =
- getOrRegisterStateColumnFamilyHandle(null,
restoredMetaInfo);
-
currentStateHandleKVStateColumnFamilies.add(registeredStateCFHandle.columnFamilyHandle);
+ savepointRestoreResult.getStateMetaInfoSnapshots();
+ List<ColumnFamilyHandle> columnFamilyHandles =
+ restoredMetaInfos.stream()
+ .map(
+ stateMetaInfoSnapshot -> {
+ RocksDbKvStateInfo registeredStateCFHandle
=
+
getOrRegisterStateColumnFamilyHandle(
+ null,
stateMetaInfoSnapshot);
+ return
registeredStateCFHandle.columnFamilyHandle;
+ })
+ .collect(Collectors.toList());
+ try (ThrowingIterator<KeyGroup> keyGroups =
savepointRestoreResult.getRestoredKeyGroups()) {
+ restoreKVStateData(keyGroups, columnFamilyHandles);
}
}
@@ -192,60 +142,19 @@ public class RocksDBFullRestoreOperation<K> extends
AbstractRocksDBRestoreOperat
* Restore the KV-state / ColumnFamily data for all key-groups referenced
by the current state
* handle.
*/
- private void restoreKVStateData() throws IOException, RocksDBException {
+ private void restoreKVStateData(
+ ThrowingIterator<KeyGroup> keyGroups, List<ColumnFamilyHandle>
columnFamilies)
+ throws IOException, RocksDBException, StateMigrationException {
// for all key-groups in the current state handle...
try (RocksDBWriteBatchWrapper writeBatchWrapper =
new RocksDBWriteBatchWrapper(db, writeBatchSize)) {
- for (Tuple2<Integer, Long> keyGroupOffset :
- currentKeyGroupsStateHandle.getGroupRangeOffsets()) {
- int keyGroup = keyGroupOffset.f0;
-
- // Check that restored key groups all belong to the backend
- Preconditions.checkState(
- keyGroupRange.contains(keyGroup),
- "The key group must belong to the backend");
-
- 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
- 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 (hasMetaDataFollowsFlag(key)) {
- // clear the signal bit in the key to make it
ready for insertion
- // again
- clearMetaDataFollowsFlag(key);
- writeBatchWrapper.put(handle, key, value);
- // TODO this could be aware of
keyGroupPrefixBytes and write only
- // one byte if possible
- kvStateId =
- END_OF_KEY_GROUP_MARK &
compressedKgInputView.readShort();
- if (END_OF_KEY_GROUP_MARK == kvStateId) {
- keyGroupHasMoreKeys = false;
- } else {
- handle =
currentStateHandleKVStateColumnFamilies.get(kvStateId);
- }
- } else {
- writeBatchWrapper.put(handle, key, value);
- }
- }
+ while (keyGroups.hasNext()) {
+ KeyGroup keyGroup = keyGroups.next();
+ try (ThrowingIterator<KeyGroupEntry> groupEntries =
keyGroup.getKeyGroupEntries()) {
+ while (groupEntries.hasNext()) {
+ KeyGroupEntry groupEntry = groupEntries.next();
+ ColumnFamilyHandle handle =
columnFamilies.get(groupEntry.getKvStateId());
+ writeBatchWrapper.put(handle, groupEntry.getKey(),
groupEntry.getValue());
}
}
}
