You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by al...@apache.org on 2016/08/31 17:28:19 UTC
[01/27] flink git commit: [FLINK-4380] Introduce KeyGroupAssigner and
Max-Parallelism Parameter
Repository: flink
Updated Branches:
refs/heads/master 7cd9bb5f1 -> bdf9f86c5
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-tests/src/test/java/org/apache/flink/test/streaming/api/StreamingOperatorsITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/streaming/api/StreamingOperatorsITCase.java b/flink-tests/src/test/java/org/apache/flink/test/streaming/api/StreamingOperatorsITCase.java
index 5d99de4..3b98d33 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/streaming/api/StreamingOperatorsITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/streaming/api/StreamingOperatorsITCase.java
@@ -78,16 +78,20 @@ public class StreamingOperatorsITCase extends StreamingMultipleProgramsTestBase
final int numKeys = 2;
StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+
DataStream<Tuple2<Integer, Integer>> sourceStream = env.addSource(new TupleSource(numElements, numKeys));
SplitStream<Tuple2<Integer, Integer>> splittedResult = sourceStream
.keyBy(0)
.fold(0, new FoldFunction<Tuple2<Integer, Integer>, Integer>() {
+ private static final long serialVersionUID = 4875723041825726082L;
+
@Override
public Integer fold(Integer accumulator, Tuple2<Integer, Integer> value) throws Exception {
return accumulator + value.f1;
}
}).map(new RichMapFunction<Integer, Tuple2<Integer, Integer>>() {
+ private static final long serialVersionUID = 8538355101606319744L;
int key = -1;
@Override
public Tuple2<Integer, Integer> map(Integer value) throws Exception {
@@ -97,6 +101,8 @@ public class StreamingOperatorsITCase extends StreamingMultipleProgramsTestBase
return new Tuple2<>(key, value);
}
}).split(new OutputSelector<Tuple2<Integer, Integer>>() {
+ private static final long serialVersionUID = -8439325199163362470L;
+
@Override
public Iterable<String> select(Tuple2<Integer, Integer> value) {
List<String> output = new ArrayList<>();
@@ -107,6 +113,8 @@ public class StreamingOperatorsITCase extends StreamingMultipleProgramsTestBase
});
splittedResult.select("0").map(new MapFunction<Tuple2<Integer,Integer>, Integer>() {
+ private static final long serialVersionUID = 2114608668010092995L;
+
@Override
public Integer map(Tuple2<Integer, Integer> value) throws Exception {
return value.f1;
@@ -114,6 +122,8 @@ public class StreamingOperatorsITCase extends StreamingMultipleProgramsTestBase
}).writeAsText(resultPath1, FileSystem.WriteMode.OVERWRITE);
splittedResult.select("1").map(new MapFunction<Tuple2<Integer, Integer>, Integer>() {
+ private static final long serialVersionUID = 5631104389744681308L;
+
@Override
public Integer map(Tuple2<Integer, Integer> value) throws Exception {
return value.f1;
@@ -149,6 +159,7 @@ public class StreamingOperatorsITCase extends StreamingMultipleProgramsTestBase
final int numElements = 10;
StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+
DataStream<Tuple2<Integer, NonSerializable>> input = env.addSource(new NonSerializableTupleSource(numElements));
input
@@ -156,12 +167,16 @@ public class StreamingOperatorsITCase extends StreamingMultipleProgramsTestBase
.fold(
new NonSerializable(42),
new FoldFunction<Tuple2<Integer, NonSerializable>, NonSerializable>() {
+ private static final long serialVersionUID = 2705497830143608897L;
+
@Override
public NonSerializable fold(NonSerializable accumulator, Tuple2<Integer, NonSerializable> value) throws Exception {
return new NonSerializable(accumulator.value + value.f1.value);
}
})
.map(new MapFunction<NonSerializable, Integer>() {
+ private static final long serialVersionUID = 6906984044674568945L;
+
@Override
public Integer map(NonSerializable value) throws Exception {
return value.value;
@@ -192,6 +207,7 @@ public class StreamingOperatorsITCase extends StreamingMultipleProgramsTestBase
}
private static class NonSerializableTupleSource implements SourceFunction<Tuple2<Integer, NonSerializable>> {
+ private static final long serialVersionUID = 3949171986015451520L;
private final int numElements;
public NonSerializableTupleSource(int numElements) {
@@ -212,6 +228,7 @@ public class StreamingOperatorsITCase extends StreamingMultipleProgramsTestBase
private static class TupleSource implements SourceFunction<Tuple2<Integer, Integer>> {
+ private static final long serialVersionUID = -8110466235852024821L;
private final int numElements;
private final int numKeys;
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/DataStreamPojoITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/DataStreamPojoITCase.java b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/DataStreamPojoITCase.java
index c345b37..cc8b699 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/DataStreamPojoITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/DataStreamPojoITCase.java
@@ -59,6 +59,7 @@ public class DataStreamPojoITCase extends StreamingMultipleProgramsTestBase {
.sum("sum")
.keyBy("aaa", "abc", "wxyz")
.flatMap(new FlatMapFunction<Data, Data>() {
+ private static final long serialVersionUID = 788865239171396315L;
Data[] first = new Data[3];
@Override
public void flatMap(Data value, Collector<Data> out) throws Exception {
@@ -105,6 +106,7 @@ public class DataStreamPojoITCase extends StreamingMultipleProgramsTestBase {
.sum("sum")
.keyBy("aaa", "stats.count")
.flatMap(new FlatMapFunction<Data, Data>() {
+ private static final long serialVersionUID = -3678267280397950258L;
Data[] first = new Data[3];
@Override
public void flatMap(Data value, Collector<Data> out) throws Exception {
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/TimestampITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/TimestampITCase.java b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/TimestampITCase.java
index d69c140..d693aaa 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/TimestampITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/TimestampITCase.java
@@ -46,6 +46,7 @@ import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.test.util.ForkableFlinkMiniCluster;
import org.apache.flink.util.TestLogger;
+import org.apache.flink.util.TestLogger;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.Before;
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-tests/src/test/java/org/apache/flink/test/web/WebFrontendITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/web/WebFrontendITCase.java b/flink-tests/src/test/java/org/apache/flink/test/web/WebFrontendITCase.java
index 032c8fe..fc90994 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/web/WebFrontendITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/web/WebFrontendITCase.java
@@ -36,6 +36,7 @@ import org.apache.flink.runtime.webmonitor.testutils.HttpTestClient;
import org.apache.flink.test.util.ForkableFlinkMiniCluster;
import org.apache.flink.test.util.TestBaseUtils;
+import org.apache.flink.util.TestLogger;
import org.junit.BeforeClass;
import org.junit.Test;
@@ -56,7 +57,7 @@ import static org.junit.Assert.fail;
import static org.apache.flink.test.util.TestBaseUtils.getFromHTTP;
-public class WebFrontendITCase {
+public class WebFrontendITCase extends TestLogger {
private static final int NUM_TASK_MANAGERS = 2;
private static final int NUM_SLOTS = 4;
[05/27] flink git commit: [FLINK-3761] Refactor RocksDB Backend/Make
Key-Group Aware
Posted by al...@apache.org.
[FLINK-3761] Refactor RocksDB Backend/Make Key-Group Aware
This change makes the RocksDB backend key-group aware by building on the
changes in the previous commit.
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/addd0842
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/addd0842
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/addd0842
Branch: refs/heads/master
Commit: addd0842f9d74b05fc44374c6682b6e697603939
Parents: 4809f53
Author: Stefan Richter <s....@data-artisans.com>
Authored: Wed Aug 17 14:50:18 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:10:01 2016 +0200
----------------------------------------------------------------------
.../streaming/state/AbstractRocksDBState.java | 98 ++-
.../streaming/state/RocksDBFoldingState.java | 25 +-
.../state/RocksDBKeyedStateBackend.java | 786 ++++++++++++++++++-
.../streaming/state/RocksDBListState.java | 22 +-
.../streaming/state/RocksDBReducingState.java | 24 +-
.../streaming/state/RocksDBStateBackend.java | 24 +-
.../streaming/state/RocksDBValueState.java | 27 +-
.../state/RocksDBAsyncKVSnapshotTest.java | 330 --------
.../state/RocksDBAsyncSnapshotTest.java | 397 ++++++++++
.../state/RocksDBMergeIteratorTest.java | 140 ++++
.../state/RocksDBStateBackendConfigTest.java | 690 ++++++++--------
.../io/async/AbstractAsyncIOCallable.java | 157 ++++
.../runtime/io/async/AsyncDoneCallback.java | 31 +
.../flink/runtime/io/async/AsyncStoppable.java | 47 ++
.../async/AsyncStoppableTaskWithCallback.java | 55 ++
.../io/async/StoppableCallbackCallable.java | 30 +
.../memory/MemCheckpointStreamFactory.java | 6 +-
.../streaming/runtime/tasks/StreamTask.java | 18 +
18 files changed, 2136 insertions(+), 771 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
index 710f506..cbc2757 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
@@ -20,8 +20,11 @@ package org.apache.flink.contrib.streaming.state;
import org.apache.flink.api.common.state.State;
import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
import org.apache.flink.core.memory.DataOutputView;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
import org.apache.flink.runtime.state.KvState;
import org.apache.flink.util.Preconditions;
import org.rocksdb.ColumnFamilyHandle;
@@ -30,7 +33,6 @@ import org.rocksdb.WriteOptions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import java.io.ByteArrayOutputStream;
import java.io.IOException;
/**
@@ -56,7 +58,7 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta
private N currentNamespace;
/** Backend that holds the actual RocksDB instance where we store state */
- protected RocksDBKeyedStateBackend backend;
+ protected RocksDBKeyedStateBackend<K> backend;
/** The column family of this particular instance of state */
protected ColumnFamilyHandle columnFamily;
@@ -69,14 +71,20 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta
*/
private final WriteOptions writeOptions;
+ protected final ByteArrayOutputStreamWithPos keySerializationStream;
+ protected final DataOutputView keySerializationDateDataOutputView;
+
+ private final boolean ambiguousKeyPossible;
+
/**
* Creates a new RocksDB backed state.
* @param namespaceSerializer The serializer for the namespace.
*/
- protected AbstractRocksDBState(ColumnFamilyHandle columnFamily,
+ protected AbstractRocksDBState(
+ ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
SD stateDesc,
- RocksDBKeyedStateBackend backend) {
+ RocksDBKeyedStateBackend<K> backend) {
this.namespaceSerializer = namespaceSerializer;
this.backend = backend;
@@ -85,31 +93,27 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta
writeOptions = new WriteOptions();
writeOptions.setDisableWAL(true);
-
this.stateDesc = Preconditions.checkNotNull(stateDesc, "State Descriptor");
+
+ this.keySerializationStream = new ByteArrayOutputStreamWithPos(128);
+ this.keySerializationDateDataOutputView = new DataOutputViewStreamWrapper(keySerializationStream);
+ this.ambiguousKeyPossible = (backend.getKeySerializer().getLength() < 0)
+ && (namespaceSerializer.getLength() < 0);
}
// ------------------------------------------------------------------------
@Override
public void clear() {
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);
try {
- writeKeyAndNamespace(out);
- byte[] key = baos.toByteArray();
+ writeCurrentKeyWithGroupAndNamespace();
+ byte[] key = keySerializationStream.toByteArray();
backend.db.remove(columnFamily, writeOptions, key);
} catch (IOException|RocksDBException e) {
throw new RuntimeException("Error while removing entry from RocksDB", e);
}
}
- protected void writeKeyAndNamespace(DataOutputView out) throws IOException {
- backend.getKeySerializer().serialize(backend.getCurrentKey(), out);
- out.writeByte(42);
- namespaceSerializer.serialize(currentNamespace, out);
- }
-
@Override
public void setCurrentNamespace(N namespace) {
this.currentNamespace = Preconditions.checkNotNull(namespace, "Namespace");
@@ -118,17 +122,67 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta
@Override
@SuppressWarnings("unchecked")
public byte[] getSerializedValue(byte[] serializedKeyAndNamespace) throws Exception {
- // Serialized key and namespace is expected to be of the same format
- // as writeKeyAndNamespace()
Preconditions.checkNotNull(serializedKeyAndNamespace, "Serialized key and namespace");
- byte[] value = backend.db.get(columnFamily, serializedKeyAndNamespace);
+ //TODO make KvStateRequestSerializer key-group aware to save this round trip and key-group computation
+ Tuple2<K, N> des = KvStateRequestSerializer.<K, N>deserializeKeyAndNamespace(
+ serializedKeyAndNamespace,
+ backend.getKeySerializer(),
+ namespaceSerializer);
+
+ int keyGroup = backend.getKeyGroupAssigner().getKeyGroupIndex(des.f0);
+ writeKeyWithGroupAndNamespace(keyGroup, des.f0, des.f1);
+ return backend.db.get(columnFamily, keySerializationStream.toByteArray());
+
+ }
+
+ protected void writeCurrentKeyWithGroupAndNamespace() throws IOException {
+ writeKeyWithGroupAndNamespace(backend.getCurrentKeyGroupIndex(), backend.getCurrentKey(), currentNamespace);
+ }
+
+ protected void writeKeyWithGroupAndNamespace(int keyGroup, K key, N namespace) throws IOException {
+ keySerializationStream.reset();
+ writeKeyGroup(keyGroup);
+ writeKey(key);
+ writeNameSpace(namespace);
+ }
+
+ private void writeKeyGroup(int keyGroup) throws IOException {
+ for (int i = backend.getKeyGroupPrefixBytes(); --i >= 0;) {
+ keySerializationDateDataOutputView.writeByte(keyGroup >>> (i << 3));
+ }
+ }
+
+ private void writeKey(K key) throws IOException {
+ //write key
+ int beforeWrite = (int) keySerializationStream.getPosition();
+ backend.getKeySerializer().serialize(key, keySerializationDateDataOutputView);
+
+ if (ambiguousKeyPossible) {
+ //write size of key
+ writeLengthFrom(beforeWrite);
+ }
+ }
+
+ private void writeNameSpace(N namespace) throws IOException {
+ int beforeWrite = (int) keySerializationStream.getPosition();
+ namespaceSerializer.serialize(namespace, keySerializationDateDataOutputView);
- if (value != null) {
- return value;
- } else {
- return null;
+ if (ambiguousKeyPossible) {
+ //write length of namespace
+ writeLengthFrom(beforeWrite);
}
}
+ private void writeLengthFrom(int fromPosition) throws IOException {
+ int length = (int) (keySerializationStream.getPosition() - fromPosition);
+ writeVariableIntBytes(length);
+ }
+
+ private void writeVariableIntBytes(int value) throws IOException {
+ do {
+ keySerializationDateDataOutputView.writeByte(value);
+ value >>>= 8;
+ } while (value != 0);
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java
index 8c0799b..3018f7b 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java
@@ -29,7 +29,6 @@ import org.rocksdb.RocksDBException;
import org.rocksdb.WriteOptions;
import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
import java.io.IOException;
/**
@@ -66,7 +65,7 @@ public class RocksDBFoldingState<K, N, T, ACC>
public RocksDBFoldingState(ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
FoldingStateDescriptor<T, ACC> stateDesc,
- RocksDBKeyedStateBackend backend) {
+ RocksDBKeyedStateBackend<K> backend) {
super(columnFamily, namespaceSerializer, stateDesc, backend);
@@ -79,11 +78,9 @@ public class RocksDBFoldingState<K, N, T, ACC>
@Override
public ACC get() {
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);
try {
- writeKeyAndNamespace(out);
- byte[] key = baos.toByteArray();
+ writeCurrentKeyWithGroupAndNamespace();
+ byte[] key = keySerializationStream.toByteArray();
byte[] valueBytes = backend.db.get(columnFamily, key);
if (valueBytes == null) {
return null;
@@ -96,23 +93,21 @@ public class RocksDBFoldingState<K, N, T, ACC>
@Override
public void add(T value) throws IOException {
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);
try {
- writeKeyAndNamespace(out);
- byte[] key = baos.toByteArray();
+ writeCurrentKeyWithGroupAndNamespace();
+ byte[] key = keySerializationStream.toByteArray();
byte[] valueBytes = backend.db.get(columnFamily, key);
-
+ DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(keySerializationStream);
if (valueBytes == null) {
- baos.reset();
+ keySerializationStream.reset();
valueSerializer.serialize(foldFunction.fold(stateDesc.getDefaultValue(), value), out);
- backend.db.put(columnFamily, writeOptions, key, baos.toByteArray());
+ backend.db.put(columnFamily, writeOptions, key, keySerializationStream.toByteArray());
} else {
ACC oldValue = valueSerializer.deserialize(new DataInputViewStreamWrapper(new ByteArrayInputStream(valueBytes)));
ACC newValue = foldFunction.fold(oldValue, value);
- baos.reset();
+ keySerializationStream.reset();
valueSerializer.serialize(newValue, out);
- backend.db.put(columnFamily, writeOptions, key, baos.toByteArray());
+ backend.db.put(columnFamily, writeOptions, key, keySerializationStream.toByteArray());
}
} catch (Exception e) {
throw new RuntimeException("Error while adding data to RocksDB", e);
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
index 63f1fa2..a1634b2 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
@@ -21,6 +21,7 @@ import org.apache.commons.io.FileUtils;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.state.FoldingState;
import org.apache.flink.api.common.state.FoldingStateDescriptor;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ReducingState;
@@ -29,30 +30,49 @@ import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
+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.core.memory.DataOutputView;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.runtime.io.async.AbstractAsyncIOCallable;
+import org.apache.flink.runtime.io.async.AsyncStoppableTaskWithCallback;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.DoneFuture;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyGroupRangeOffsets;
import org.apache.flink.runtime.state.KeyGroupsStateHandle;
import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.util.SerializableObject;
+import org.apache.flink.util.Preconditions;
import org.rocksdb.ColumnFamilyDescriptor;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.DBOptions;
+import org.rocksdb.ReadOptions;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
+import org.rocksdb.RocksIterator;
+import org.rocksdb.Snapshot;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import javax.annotation.concurrent.GuardedBy;
import java.io.File;
import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
import java.util.ArrayList;
+import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import java.util.concurrent.Future;
+import java.util.PriorityQueue;
+import java.util.concurrent.RunnableFuture;
/**
* A {@link KeyedStateBackend} that stores its state in {@code RocksDB} and will serialize state to
@@ -80,25 +100,29 @@ public class RocksDBKeyedStateBackend<K> extends KeyedStateBackend<K> {
private final File instanceRocksDBPath;
/**
+ * Lock for protecting cleanup of the RocksDB db. We acquire this when doing asynchronous
+ * checkpoints and when disposing the db. Otherwise, the asynchronous snapshot might try
+ * iterating over a disposed db.
+ */
+ private final SerializableObject dbDisposeLock = new SerializableObject();
+
+ /**
* Our RocksDB data base, this is used by the actual subclasses of {@link AbstractRocksDBState}
* to store state. The different k/v states that we have don't each have their own RocksDB
* instance. They all write to this instance but to their own column family.
*/
+ @GuardedBy("dbDisposeLock")
protected volatile RocksDB db;
/**
- * Lock for protecting cleanup of the RocksDB db. We acquire this when doing asynchronous
- * checkpoints and when disposing the db. Otherwise, the asynchronous snapshot might try
- * iterating over a disposed db.
- */
- private final SerializableObject dbCleanupLock = new SerializableObject();
-
- /**
* Information about the k/v states as we create them. This is used to retrieve the
* column family that is used for a state and also for sanity checks when restoring.
*/
private Map<String, Tuple2<ColumnFamilyHandle, StateDescriptor>> kvStateInformation;
+ /** Number of bytes required to prefix the key groups. */
+ private final int keyGroupPrefixBytes;
+
public RocksDBKeyedStateBackend(
JobID jobId,
String operatorIdentifier,
@@ -108,7 +132,7 @@ public class RocksDBKeyedStateBackend<K> extends KeyedStateBackend<K> {
TaskKvStateRegistry kvStateRegistry,
TypeSerializer<K> keySerializer,
KeyGroupAssigner<K> keyGroupAssigner,
- KeyGroupRange keyGroupRange
+ KeyGroupRange keyGroupRange
) throws Exception {
super(kvStateRegistry, keySerializer, keyGroupAssigner, keyGroupRange);
@@ -147,35 +171,543 @@ public class RocksDBKeyedStateBackend<K> extends KeyedStateBackend<K> {
} catch (RocksDBException e) {
throw new RuntimeException("Error while opening RocksDB instance.", e);
}
-
+ keyGroupPrefixBytes = getNumberOfKeyGroups() > (Byte.MAX_VALUE + 1) ? 2 : 1;
kvStateInformation = new HashMap<>();
}
+ public RocksDBKeyedStateBackend(
+ JobID jobId,
+ String operatorIdentifier,
+ File instanceBasePath,
+ DBOptions dbOptions,
+ ColumnFamilyOptions columnFamilyOptions,
+ TaskKvStateRegistry kvStateRegistry,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ List<KeyGroupsStateHandle> restoreState
+ ) throws Exception {
+ this(
+ jobId,
+ operatorIdentifier,
+ instanceBasePath,
+ dbOptions,
+ columnFamilyOptions,
+ kvStateRegistry,
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange);
+
+ LOG.info("Initializing RocksDB keyed state backend from snapshot.");
+
+ if (LOG.isDebugEnabled()) {
+ LOG.debug("Restoring snapshot from state handles: {}.", restoreState);
+ }
+
+ RocksDBRestoreOperation restoreOperation = new RocksDBRestoreOperation(this);
+ restoreOperation.doRestore(restoreState);
+ }
+
+ /**
+ * @see java.io.Closeable
+ *
+ * Should only be called by one thread.
+ *
+ * @throws Exception
+ */
@Override
public void close() throws Exception {
super.close();
- // we have to lock because we might have an asynchronous checkpoint going on
- synchronized (dbCleanupLock) {
- if (db != null) {
- for (Tuple2<ColumnFamilyHandle, StateDescriptor> column : kvStateInformation.values()) {
- column.f0.dispose();
- }
+ final RocksDB cleanupRockDBReference;
- db.dispose();
- db = null;
+ // Acquire the log on dbDisposeLock, so that no ongoing snapshots access the db during cleanup
+ synchronized (dbDisposeLock) {
+ // IMPORTANT: null reference to signal potential async checkpoint workers that the db was disposed, as
+ // working on the disposed object results in SEGFAULTS. Other code has to check field #db for null
+ // and access it in a synchronized block that locks on #dbDisposeLock.
+ cleanupRockDBReference = db;
+ db = null;
+ }
+
+ // Dispose decoupled db
+ if (cleanupRockDBReference != null) {
+ for (Tuple2<ColumnFamilyHandle, StateDescriptor> column : kvStateInformation.values()) {
+ column.f0.dispose();
}
+ cleanupRockDBReference.dispose();
}
FileUtils.deleteDirectory(instanceBasePath);
}
+ public int getKeyGroupPrefixBytes() {
+ return keyGroupPrefixBytes;
+ }
+
+ /**
+ * Triggers an asynchronous snapshot of the keyed state backend from RocksDB. This snapshot can be canceled and
+ * is also stopped when the backend is closed through {@link #close()}. For each backend, this method must always
+ * be called by the same thread.
+ *
+ * @param checkpointId The Id of the checkpoint.
+ * @param timestamp The timestamp of the checkpoint.
+ * @param streamFactory The factory that we can use for writing our state to streams.
+ *
+ * @return Future to the state handle of the snapshot data.
+ * @throws Exception
+ */
@Override
- public Future<KeyGroupsStateHandle> snapshot(
- long checkpointId,
- long timestamp,
- CheckpointStreamFactory streamFactory) throws Exception {
- throw new RuntimeException("Not implemented.");
+ public RunnableFuture<KeyGroupsStateHandle> snapshot(
+ final long checkpointId,
+ final long timestamp,
+ final CheckpointStreamFactory streamFactory) throws Exception {
+
+ long startTime = System.currentTimeMillis();
+
+ if (kvStateInformation.isEmpty()) {
+ LOG.info("Asynchronous RocksDB snapshot performed on empty keyed state at " + timestamp + " . Returning null.");
+ return new DoneFuture<>(null);
+ }
+
+ final RocksDBSnapshotOperation snapshotOperation = new RocksDBSnapshotOperation(this, streamFactory);
+ // hold the db lock while operation on the db to guard us against async db disposal
+ synchronized (dbDisposeLock) {
+ if (db != null) {
+ snapshotOperation.takeDBSnapShot(checkpointId, timestamp);
+ } else {
+ throw new IOException("RocksDB closed.");
+ }
+ }
+
+ // implementation of the async IO operation, based on FutureTask
+ AbstractAsyncIOCallable<KeyGroupsStateHandle, CheckpointStreamFactory.CheckpointStateOutputStream> ioCallable =
+ new AbstractAsyncIOCallable<KeyGroupsStateHandle, CheckpointStreamFactory.CheckpointStateOutputStream>() {
+
+ @Override
+ public CheckpointStreamFactory.CheckpointStateOutputStream openIOHandle() throws Exception {
+ snapshotOperation.openCheckpointStream();
+ return snapshotOperation.getOutStream();
+ }
+
+ @Override
+ public KeyGroupsStateHandle performOperation() throws Exception {
+ long startTime = System.currentTimeMillis();
+ try {
+ // hold the db lock while operation on the db to guard us against async db disposal
+ synchronized (dbDisposeLock) {
+ if (db != null) {
+ snapshotOperation.writeDBSnapshot();
+ } else {
+ throw new IOException("RocksDB closed.");
+ }
+ }
+
+ } finally {
+ snapshotOperation.closeCheckpointStream();
+ }
+
+ LOG.info("Asynchronous RocksDB snapshot (" + streamFactory + ", asynchronous part) in thread " +
+ Thread.currentThread() + " took " + (System.currentTimeMillis() - startTime) + " ms.");
+
+ return snapshotOperation.getSnapshotResultStateHandle();
+ }
+
+ @Override
+ public void done() {
+ // hold the db lock while operation on the db to guard us against async db disposal
+ synchronized (dbDisposeLock) {
+ if (db != null) {
+ snapshotOperation.releaseDBSnapshot();
+ }
+ }
+ }
+ };
+
+ LOG.info("Asynchronous RocksDB snapshot (" + streamFactory + ", synchronous part) in thread " +
+ Thread.currentThread() + " took " + (System.currentTimeMillis() - startTime) + " ms.");
+
+ return AsyncStoppableTaskWithCallback.from(ioCallable);
+ }
+
+ /**
+ * Encapsulates the process to perform a snapshot of a RocksDBKeyedStateBackend.
+ */
+ static final class RocksDBSnapshotOperation {
+
+ static final int FIRST_BIT_IN_BYTE_MASK = 0x80;
+ static final int END_OF_KEY_GROUP_MARK = 0xFFFF;
+
+ private final RocksDBKeyedStateBackend<?> stateBackend;
+ private final KeyGroupRangeOffsets keyGroupRangeOffsets;
+ private final CheckpointStreamFactory checkpointStreamFactory;
+
+ private long checkpointId;
+ private long checkpointTimeStamp;
+
+ private Snapshot snapshot;
+ private CheckpointStreamFactory.CheckpointStateOutputStream outStream;
+ private DataOutputView outputView;
+ private List<Tuple2<RocksIterator, Integer>> kvStateIterators;
+ private KeyGroupsStateHandle snapshotResultStateHandle;
+
+
+
+ public RocksDBSnapshotOperation(
+ RocksDBKeyedStateBackend<?> stateBackend,
+ CheckpointStreamFactory checkpointStreamFactory) {
+
+ this.stateBackend = stateBackend;
+ this.checkpointStreamFactory = checkpointStreamFactory;
+ this.keyGroupRangeOffsets = new KeyGroupRangeOffsets(stateBackend.keyGroupRange);
+ }
+
+ /**
+ * 1) Create a snapshot object from RocksDB.
+ *
+ * @param checkpointId id of the checkpoint for which we take the snapshot
+ * @param checkpointTimeStamp timestamp of the checkpoint for which we take the snapshot
+ */
+ public void takeDBSnapShot(long checkpointId, long checkpointTimeStamp) throws IOException {
+ 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();
+ }
+
+ /**
+ * 2) Open CheckpointStateOutputStream through the checkpointStreamFactory into which we will write.
+ *
+ * @throws Exception
+ */
+ public void openCheckpointStream() throws Exception {
+ Preconditions.checkArgument(outStream == null, "Output stream for snapshot is already set.");
+ outStream = checkpointStreamFactory.
+ createCheckpointStateOutputStream(checkpointId, checkpointTimeStamp);
+ outputView = new DataOutputViewStreamWrapper(outStream);
+ }
+
+ /**
+ * 3) Write the actual data from RocksDB from the time we took the snapshot object in (1).
+ *
+ * @return
+ * @throws IOException
+ */
+ public void writeDBSnapshot() throws IOException, InterruptedException {
+ Preconditions.checkNotNull(snapshot, "No ongoing snapshot to write.");
+ Preconditions.checkNotNull(outStream, "No output stream to write snapshot.");
+ writeKVStateMetaData();
+ writeKVStateData();
+ }
+
+ /**
+ * 4) Close the CheckpointStateOutputStream after writing and receive a state handle.
+ *
+ * @throws IOException
+ */
+ public void closeCheckpointStream() throws IOException {
+ if(outStream != null) {
+ snapshotResultStateHandle = closeSnapshotStreamAndGetHandle();
+ }
+ }
+
+ /**
+ * 5) Release the snapshot object for RocksDB and clean up.
+ *
+ */
+ public void releaseDBSnapshot() {
+ Preconditions.checkNotNull(snapshot, "No ongoing snapshot to release.");
+ stateBackend.db.releaseSnapshot(snapshot);
+ snapshot = null;
+ outStream = null;
+ outputView = null;
+ kvStateIterators = null;
+ }
+
+ /**
+ * Returns the current CheckpointStateOutputStream (when it was opened and not yet closed) into which we write
+ * the state snapshot.
+ *
+ * @return the current CheckpointStateOutputStream
+ */
+ public CheckpointStreamFactory.CheckpointStateOutputStream getOutStream() {
+ return outStream;
+ }
+
+ /**
+ * Returns a state handle to the snapshot after the snapshot procedure is completed and null before.
+ *
+ * @return state handle to the completed snapshot
+ */
+ public KeyGroupsStateHandle getSnapshotResultStateHandle() {
+ return snapshotResultStateHandle;
+ }
+
+ private void writeKVStateMetaData() throws IOException, InterruptedException {
+ //write number of k/v states
+ outputView.writeInt(stateBackend.kvStateInformation.size());
+
+ int kvStateId = 0;
+ //iterate all column families, where each column family holds one k/v state, to write the metadata
+ for (Map.Entry<String, Tuple2<ColumnFamilyHandle, StateDescriptor>> column : stateBackend.kvStateInformation.entrySet()) {
+
+ //be cooperative and check for interruption from time to time in the hot loop
+ checkInterrupted();
+
+ //write StateDescriptor for this k/v state
+ ObjectOutputStream ooOut = new ObjectOutputStream(outStream);
+ ooOut.writeObject(column.getValue().f1);
+ //retrieve iterator for this k/v states
+ ReadOptions readOptions = new ReadOptions();
+ readOptions.setSnapshot(snapshot);
+ RocksIterator iterator = stateBackend.db.newIterator(column.getValue().f0, readOptions);
+ kvStateIterators.add(new Tuple2<RocksIterator, Integer>(iterator, kvStateId));
+ ++kvStateId;
+ }
+ }
+
+ private void writeKVStateData() throws IOException, InterruptedException {
+
+ RocksDBMergeIterator iterator = new RocksDBMergeIterator(kvStateIterators, stateBackend.keyGroupPrefixBytes);
+
+ byte[] previousKey = null;
+ byte[] previousValue = null;
+
+ //preamble: setup with first key-group as our lookahead
+ if (iterator.isValid()) {
+ //begin first key-group by recording the offset
+ keyGroupRangeOffsets.setKeyGroupOffset(iterator.keyGroup(), outStream.getPos());
+ //write the k/v-state id as metadata
+ //TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
+ outputView.writeShort(iterator.kvStateId());
+ previousKey = iterator.key();
+ previousValue = iterator.value();
+ iterator.next();
+ }
+
+ //main loop: write k/v pairs ordered by (key-group, kv-state), thereby tracking key-group offsets.
+ while (iterator.isValid()) {
+
+ assert (!hasMetaDataFollowsFlag(previousKey));
+
+ //set signal in first key byte that meta data will follow in the stream after this k/v pair
+ if (iterator.isNewKeyGroup() || iterator.isNewKeyValueState()) {
+
+ //be cooperative and check for interruption from time to time in the hot loop
+ checkInterrupted();
+
+ setMetaDataFollowsFlagInKey(previousKey);
+ }
+
+ writeKeyValuePair(previousKey, previousValue);
+
+ //write meta data if we have to
+ if (iterator.isNewKeyGroup()) {
+ //
+ //TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
+ outputView.writeShort(END_OF_KEY_GROUP_MARK);
+ //begin new key-group
+ keyGroupRangeOffsets.setKeyGroupOffset(iterator.keyGroup(), outStream.getPos());
+ //write the kev-state
+ //TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
+ outputView.writeShort(iterator.kvStateId());
+ } else if (iterator.isNewKeyValueState()) {
+ //write the k/v-state
+ //TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
+ outputView.writeShort(iterator.kvStateId());
+ }
+
+ //request next k/v pair
+ previousKey = iterator.key();
+ previousValue = iterator.value();
+ iterator.next();
+ }
+
+ //epilogue: write last key-group
+ if (previousKey != null) {
+ assert (!hasMetaDataFollowsFlag(previousKey));
+ setMetaDataFollowsFlagInKey(previousKey);
+ writeKeyValuePair(previousKey, previousValue);
+ //TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
+ outputView.writeShort(END_OF_KEY_GROUP_MARK);
+ }
+ }
+
+ private KeyGroupsStateHandle closeSnapshotStreamAndGetHandle() throws IOException {
+ StreamStateHandle stateHandle = outStream.closeAndGetHandle();
+ outStream = null;
+ if (stateHandle != null) {
+ return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
+ } else {
+ throw new IOException("Output stream returned null on close.");
+ }
+ }
+
+ private void writeKeyValuePair(byte[] key, byte[] value) throws IOException {
+ BytePrimitiveArraySerializer.INSTANCE.serialize(key, outputView);
+ BytePrimitiveArraySerializer.INSTANCE.serialize(value, outputView);
+ }
+
+ static void setMetaDataFollowsFlagInKey(byte[] key) {
+ key[0] |= FIRST_BIT_IN_BYTE_MASK;
+ }
+
+ static void clearMetaDataFollowsFlag(byte[] key) {
+ key[0] &= (~RocksDBSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
+ }
+
+ static boolean hasMetaDataFollowsFlag(byte[] key) {
+ return 0 != (key[0] & RocksDBSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
+ }
+
+ private static void checkInterrupted() throws InterruptedException {
+ if(Thread.currentThread().isInterrupted()) {
+ throw new InterruptedException("Snapshot canceled.");
+ }
+ }
+ }
+
+ /**
+ * Encapsulates the process of restoring a RocksDBKeyedStateBackend from a snapshot.
+ */
+ static final class RocksDBRestoreOperation {
+
+ private final RocksDBKeyedStateBackend<?> rocksDBKeyedStateBackend;
+
+ /** 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;
+
+ /**
+ * Creates a restore operation object for the given state backend instance.
+ *
+ * @param rocksDBKeyedStateBackend the state backend into which we restore
+ */
+ public RocksDBRestoreOperation(RocksDBKeyedStateBackend<?> rocksDBKeyedStateBackend) {
+ this.rocksDBKeyedStateBackend = Preconditions.checkNotNull(rocksDBKeyedStateBackend);
+ }
+
+ /**
+ * Restores all key-groups data that is referenced by the passed state handles.
+ *
+ * @param keyGroupsStateHandles List of all key groups state handles that shall be restored.
+ * @throws IOException
+ * @throws ClassNotFoundException
+ * @throws RocksDBException
+ */
+ public void doRestore(List<KeyGroupsStateHandle> keyGroupsStateHandles)
+ throws IOException, ClassNotFoundException, RocksDBException {
+
+ for (KeyGroupsStateHandle keyGroupsStateHandle : keyGroupsStateHandles) {
+ if (keyGroupsStateHandle != null) {
+ this.currentKeyGroupsStateHandle = keyGroupsStateHandle;
+ restoreKeyGroupsInStateHandle();
+ }
+ }
+ }
+
+ /**
+ * Restore one key groups state handle
+ *
+ * @throws IOException
+ * @throws RocksDBException
+ * @throws ClassNotFoundException
+ */
+ private void restoreKeyGroupsInStateHandle()
+ throws IOException, RocksDBException, ClassNotFoundException {
+ try {
+ currentStateHandleInStream = currentKeyGroupsStateHandle.getStateHandle().openInputStream();
+ currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
+ restoreKVStateMetaData();
+ restoreKVStateData();
+ } finally {
+ if(currentStateHandleInStream != null) {
+ currentStateHandleInStream.close();
+ }
+ }
+
+ }
+
+ /**
+ * Restore the KV-state / ColumnFamily meta data for all key-groups referenced by the current state handle
+ *
+ * @throws IOException
+ * @throws ClassNotFoundException
+ * @throws RocksDBException
+ */
+ private void restoreKVStateMetaData() throws IOException, ClassNotFoundException, RocksDBException {
+ //read number of k/v states
+ int numColumns = currentStateHandleInView.readInt();
+
+ //those two lists are aligned and should later have the same size!
+ currentStateHandleKVStateColumnFamilies = new ArrayList<>(numColumns);
+
+ //restore the empty columns for the k/v states through the metadata
+ for (int i = 0; i < numColumns; i++) {
+ ObjectInputStream ooIn = new ObjectInputStream(currentStateHandleInStream);
+ StateDescriptor stateDescriptor = (StateDescriptor) ooIn.readObject();
+ Tuple2<ColumnFamilyHandle, StateDescriptor> columnFamily = rocksDBKeyedStateBackend.
+ kvStateInformation.get(stateDescriptor.getName());
+
+ if(null == columnFamily) {
+ ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(
+ stateDescriptor.getName().getBytes(), rocksDBKeyedStateBackend.columnOptions);
+
+ columnFamily = new Tuple2<>(rocksDBKeyedStateBackend.db.
+ createColumnFamily(columnFamilyDescriptor), stateDescriptor);
+ rocksDBKeyedStateBackend.kvStateInformation.put(stateDescriptor.getName(), columnFamily);
+ }
+
+ currentStateHandleKVStateColumnFamilies.add(columnFamily.f0);
+ }
+ }
+
+ /**
+ * 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()) {
+ long offset = keyGroupOffset.f1;
+ //not empty key-group?
+ if (0L != offset) {
+ currentStateHandleInStream.seek(offset);
+ boolean keyGroupHasMoreKeys = true;
+ //TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
+ int kvStateId = currentStateHandleInView.readShort();
+ ColumnFamilyHandle handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
+ //insert all k/v pairs into DB
+ while (keyGroupHasMoreKeys) {
+ byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(currentStateHandleInView);
+ byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(currentStateHandleInView);
+ if (RocksDBSnapshotOperation.hasMetaDataFollowsFlag(key)) {
+ //clear the signal bit in the key to make it ready for insertion again
+ RocksDBSnapshotOperation.clearMetaDataFollowsFlag(key);
+ rocksDBKeyedStateBackend.db.put(handle, key, value);
+ //TODO this could be aware of keyGroupPrefixBytes and write only one byte if possible
+ kvStateId = RocksDBSnapshotOperation.END_OF_KEY_GROUP_MARK
+ & currentStateHandleInView.readShort();
+ if (RocksDBSnapshotOperation.END_OF_KEY_GROUP_MARK == kvStateId) {
+ keyGroupHasMoreKeys = false;
+ } else {
+ handle = currentStateHandleKVStateColumnFamilies.get(kvStateId);
+ }
+ } else {
+ rocksDBKeyedStateBackend.db.put(handle, key, value);
+ }
+ }
+ }
+ }
+ }
}
// ------------------------------------------------------------------------
@@ -197,12 +729,14 @@ public class RocksDBKeyedStateBackend<K> extends KeyedStateBackend<K> {
if (stateInfo != null) {
if (!stateInfo.f1.equals(descriptor)) {
- throw new RuntimeException("Trying to access state using wrong StateDescriptor, was " + stateInfo.f1 + " trying access with " + descriptor);
+ throw new RuntimeException("Trying to access state using wrong StateDescriptor, was " + stateInfo.f1 +
+ " trying access with " + descriptor);
}
return stateInfo.f0;
}
- ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(descriptor.getName().getBytes(), columnOptions);
+ ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(
+ descriptor.getName().getBytes(), columnOptions);
try {
ColumnFamilyHandle columnFamily = db.createColumnFamily(columnDescriptor);
@@ -248,4 +782,206 @@ public class RocksDBKeyedStateBackend<K> extends KeyedStateBackend<K> {
return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
}
+
+ /**
+ * Wraps a RocksDB iterator to cache it's current key and assign an id for the key/value state to the iterator.
+ * Used by #MergeIterator.
+ */
+ static final class MergeIterator {
+
+ /**
+ *
+ * @param iterator The #RocksIterator to wrap .
+ * @param kvStateId Id of the K/V state to which this iterator belongs.
+ */
+ public MergeIterator(RocksIterator iterator, int kvStateId) {
+ this.iterator = Preconditions.checkNotNull(iterator);
+ this.currentKey = iterator.key();
+ this.kvStateId = kvStateId;
+ }
+
+ private byte[] currentKey;
+ private final RocksIterator iterator;
+ private final int kvStateId;
+
+ public byte[] getCurrentKey() {
+ return currentKey;
+ }
+
+ public void setCurrentKey(byte[] currentKey) {
+ this.currentKey = currentKey;
+ }
+
+ public RocksIterator getIterator() {
+ return iterator;
+ }
+
+ public int getKvStateId() {
+ return kvStateId;
+ }
+ }
+
+ /**
+ * 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).
+ */
+ static final class RocksDBMergeIterator {
+
+ private final PriorityQueue<MergeIterator> heap;
+ private final int keyGroupPrefixByteCount;
+ private boolean newKeyGroup;
+ private boolean newKVState;
+ private boolean valid;
+
+ private MergeIterator currentSubIterator;
+
+ RocksDBMergeIterator(List<Tuple2<RocksIterator, Integer>> kvStateIterators, final int keyGroupPrefixByteCount) throws IOException {
+ Preconditions.checkNotNull(kvStateIterators);
+ this.keyGroupPrefixByteCount = keyGroupPrefixByteCount;
+
+ Comparator<MergeIterator> iteratorComparator = new Comparator<MergeIterator>() {
+ @Override
+ public int compare(MergeIterator o1, MergeIterator o2) {
+ int arrayCmpRes = compareKeyGroupsForByteArrays(
+ o1.currentKey, o2.currentKey, keyGroupPrefixByteCount);
+ return arrayCmpRes == 0 ? o1.getKvStateId() - o2.getKvStateId() : arrayCmpRes;
+ }
+ };
+
+ if (kvStateIterators.size() > 0) {
+ this.heap = new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);
+
+ for (Tuple2<RocksIterator, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
+ RocksIterator rocksIterator = rocksIteratorWithKVStateId.f0;
+ rocksIterator.seekToFirst();
+ if (rocksIterator.isValid()) {
+ heap.offer(new MergeIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
+ }
+ }
+ 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;
+ }
+
+ /**
+ * Advance the iterator. Should only be called if {@link #isValid()} returned true. Valid can only chance after
+ * calls to {@link #next()}.
+ */
+ public void next() {
+ newKeyGroup = false;
+ newKVState = false;
+
+ final RocksIterator rocksIterator = currentSubIterator.getIterator();
+ rocksIterator.next();
+
+ byte[] oldKey = currentSubIterator.getCurrentKey();
+ if (rocksIterator.isValid()) {
+ currentSubIterator.currentKey = rocksIterator.key();
+
+ if (isDifferentKeyGroup(oldKey, currentSubIterator.getCurrentKey())) {
+ heap.offer(currentSubIterator);
+ currentSubIterator = heap.poll();
+ newKVState = currentSubIterator.getIterator() != rocksIterator;
+ detectNewKeyGroup(oldKey);
+ }
+ } else if (heap.isEmpty()) {
+ valid = false;
+ } else {
+ currentSubIterator = heap.poll();
+ newKVState = true;
+ detectNewKeyGroup(oldKey);
+ }
+
+ }
+
+ private boolean isDifferentKeyGroup(byte[] a, byte[] b) {
+ return 0 != compareKeyGroupsForByteArrays(a, b, keyGroupPrefixByteCount);
+ }
+
+ private void detectNewKeyGroup(byte[] oldKey) {
+ if (isDifferentKeyGroup(oldKey, currentSubIterator.currentKey)) {
+ newKeyGroup = true;
+ }
+ }
+
+ /**
+ * 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;
+ }
+
+ public byte[] key() {
+ return currentSubIterator.getCurrentKey();
+ }
+
+ public byte[] value() {
+ return currentSubIterator.getIterator().value();
+ }
+
+ /**
+ * 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();
+ }
+
+ /**
+ * Indicates if current key starts a new k/v-state, i.e. belong to a different k/v-state than it's predecessor.
+ * @return true iff the current key belong to a different k/v-state than it's predecessor.
+ */
+ public boolean isNewKeyValueState() {
+ return newKVState;
+ }
+
+ /**
+ * Indicates if current key starts a new key-group, i.e. belong to a different key-group than it's predecessor.
+ * @return true iff the current key belong to a different key-group than it's predecessor.
+ */
+ public boolean isNewKeyGroup() {
+ return newKeyGroup;
+ }
+
+ /**
+ * Check if the iterator is still valid. Getters like {@link #key()}, {@link #value()}, etc. as well as
+ * {@link #next()} should only be called if valid returned true. Should be checked after each call to
+ * {@link #next()} before accessing iterator state.
+ * @return True iff this iterator is valid.
+ */
+ public boolean isValid() {
+ return valid;
+ }
+
+ private static int compareKeyGroupsForByteArrays(byte[] a, byte[] b, int len) {
+ for (int i = 0; i < len; ++i) {
+ int diff = (a[i] & 0xFF) - (b[i] & 0xFF);
+ if (diff != 0) {
+ return diff;
+ }
+ }
+ return 0;
+ }
+ }
+
+ /**
+ * Only visible for testing, DO NOT USE.
+ */
+ public File getInstanceBasePath() {
+ return instanceBasePath;
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java
index d8f937b..beea81a 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java
@@ -28,7 +28,6 @@ import org.rocksdb.RocksDBException;
import org.rocksdb.WriteOptions;
import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
@@ -67,7 +66,7 @@ public class RocksDBListState<K, N, V>
public RocksDBListState(ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
ListStateDescriptor<V> stateDesc,
- RocksDBKeyedStateBackend backend) {
+ RocksDBKeyedStateBackend<K> backend) {
super(columnFamily, namespaceSerializer, stateDesc, backend);
this.valueSerializer = stateDesc.getSerializer();
@@ -78,11 +77,9 @@ public class RocksDBListState<K, N, V>
@Override
public Iterable<V> get() {
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);
try {
- writeKeyAndNamespace(out);
- byte[] key = baos.toByteArray();
+ writeCurrentKeyWithGroupAndNamespace();
+ byte[] key = keySerializationStream.toByteArray();
byte[] valueBytes = backend.db.get(columnFamily, key);
if (valueBytes == null) {
@@ -107,16 +104,13 @@ public class RocksDBListState<K, N, V>
@Override
public void add(V value) throws IOException {
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);
try {
- writeKeyAndNamespace(out);
- byte[] key = baos.toByteArray();
-
- baos.reset();
-
+ writeCurrentKeyWithGroupAndNamespace();
+ byte[] key = keySerializationStream.toByteArray();
+ keySerializationStream.reset();
+ DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(keySerializationStream);
valueSerializer.serialize(value, out);
- backend.db.merge(columnFamily, writeOptions, key, baos.toByteArray());
+ backend.db.merge(columnFamily, writeOptions, key, keySerializationStream.toByteArray());
} catch (Exception e) {
throw new RuntimeException("Error while adding data to RocksDB", e);
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java
index 15ae493..068c051 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java
@@ -29,7 +29,6 @@ import org.rocksdb.RocksDBException;
import org.rocksdb.WriteOptions;
import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
import java.io.IOException;
/**
@@ -65,7 +64,7 @@ public class RocksDBReducingState<K, N, V>
public RocksDBReducingState(ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
ReducingStateDescriptor<V> stateDesc,
- RocksDBKeyedStateBackend backend) {
+ RocksDBKeyedStateBackend<K> backend) {
super(columnFamily, namespaceSerializer, stateDesc, backend);
this.valueSerializer = stateDesc.getSerializer();
@@ -77,11 +76,9 @@ public class RocksDBReducingState<K, N, V>
@Override
public V get() {
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);
try {
- writeKeyAndNamespace(out);
- byte[] key = baos.toByteArray();
+ writeCurrentKeyWithGroupAndNamespace();
+ byte[] key = keySerializationStream.toByteArray();
byte[] valueBytes = backend.db.get(columnFamily, key);
if (valueBytes == null) {
return null;
@@ -94,23 +91,22 @@ public class RocksDBReducingState<K, N, V>
@Override
public void add(V value) throws IOException {
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);
try {
- writeKeyAndNamespace(out);
- byte[] key = baos.toByteArray();
+ writeCurrentKeyWithGroupAndNamespace();
+ byte[] key = keySerializationStream.toByteArray();
byte[] valueBytes = backend.db.get(columnFamily, key);
+ DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(keySerializationStream);
if (valueBytes == null) {
- baos.reset();
+ keySerializationStream.reset();
valueSerializer.serialize(value, out);
- backend.db.put(columnFamily, writeOptions, key, baos.toByteArray());
+ backend.db.put(columnFamily, writeOptions, key, keySerializationStream.toByteArray());
} else {
V oldValue = valueSerializer.deserialize(new DataInputViewStreamWrapper(new ByteArrayInputStream(valueBytes)));
V newValue = reduceFunction.reduce(oldValue, value);
- baos.reset();
+ keySerializationStream.reset();
valueSerializer.serialize(newValue, out);
- backend.db.put(columnFamily, writeOptions, key, baos.toByteArray());
+ backend.db.put(columnFamily, writeOptions, key, keySerializationStream.toByteArray());
}
} catch (Exception e) {
throw new RuntimeException("Error while adding data to RocksDB", e);
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
index 62b71d9..f950751 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
@@ -221,12 +221,6 @@ public class RocksDBStateBackend extends AbstractStateBackend {
@Override
public CheckpointStreamFactory createStreamFactory(JobID jobId,
String operatorIdentifier) throws IOException {
- return null;
- }
-
- if (fullyAsyncBackup) {
- return performFullyAsyncSnapshot(checkpointId, timestamp);
- } else {
return checkpointStreamBackend.createStreamFactory(jobId, operatorIdentifier);
}
@@ -261,10 +255,24 @@ public class RocksDBStateBackend extends AbstractStateBackend {
String operatorIdentifier,
TypeSerializer<K> keySerializer,
KeyGroupAssigner<K> keyGroupAssigner,
- KeyGroupRange keyGroupRange,
+ KeyGroupRange keyGroupRange,
List<KeyGroupsStateHandle> restoredState,
TaskKvStateRegistry kvStateRegistry) throws Exception {
- throw new RuntimeException("Not implemented.");
+
+ lazyInitializeForJob(env, operatorIdentifier);
+
+ File instanceBasePath = new File(getDbPath(), UUID.randomUUID().toString());
+ return new RocksDBKeyedStateBackend<>(
+ jobID,
+ operatorIdentifier,
+ instanceBasePath,
+ getDbOptions(),
+ getColumnOptions(),
+ kvStateRegistry,
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange,
+ restoredState);
}
// ------------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java
index b9c0e83..9563ed8 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java
@@ -24,13 +24,11 @@ import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.core.memory.DataInputViewStreamWrapper;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.util.Preconditions;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.RocksDBException;
import org.rocksdb.WriteOptions;
import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
import java.io.IOException;
/**
@@ -63,7 +61,7 @@ public class RocksDBValueState<K, N, V>
public RocksDBValueState(ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
ValueStateDescriptor<V> stateDesc,
- RocksDBKeyedStateBackend backend) {
+ RocksDBKeyedStateBackend<K> backend) {
super(columnFamily, namespaceSerializer, stateDesc, backend);
this.valueSerializer = stateDesc.getSerializer();
@@ -74,11 +72,9 @@ public class RocksDBValueState<K, N, V>
@Override
public V value() {
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);
try {
- writeKeyAndNamespace(out);
- byte[] key = baos.toByteArray();
+ writeCurrentKeyWithGroupAndNamespace();
+ byte[] key = keySerializationStream.toByteArray();
byte[] valueBytes = backend.db.get(columnFamily, key);
if (valueBytes == null) {
return stateDesc.getDefaultValue();
@@ -95,14 +91,13 @@ public class RocksDBValueState<K, N, V>
clear();
return;
}
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);
+ DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(keySerializationStream);
try {
- writeKeyAndNamespace(out);
- byte[] key = baos.toByteArray();
- baos.reset();
+ writeCurrentKeyWithGroupAndNamespace();
+ byte[] key = keySerializationStream.toByteArray();
+ keySerializationStream.reset();
valueSerializer.serialize(value, out);
- backend.db.put(columnFamily, writeOptions, key, baos.toByteArray());
+ backend.db.put(columnFamily, writeOptions, key, keySerializationStream.toByteArray());
} catch (Exception e) {
throw new RuntimeException("Error while adding data to RocksDB", e);
}
@@ -110,11 +105,7 @@ public class RocksDBValueState<K, N, V>
@Override
public byte[] getSerializedValue(byte[] serializedKeyAndNamespace) throws Exception {
- // Serialized key and namespace is expected to be of the same format
- // as writeKeyAndNamespace()
- Preconditions.checkNotNull(serializedKeyAndNamespace, "Serialized key and namespace");
-
- byte[] value = backend.db.get(columnFamily, serializedKeyAndNamespace);
+ byte[] value = super.getSerializedValue(serializedKeyAndNamespace);
if (value != null) {
return value;
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
deleted file mode 100644
index 0e35b60..0000000
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
+++ /dev/null
@@ -1,330 +0,0 @@
-/*
- * 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.contrib.streaming.state;
-
-import org.apache.flink.api.common.functions.MapFunction;
-import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
-import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
-import org.apache.flink.api.common.typeutils.base.StringSerializer;
-import org.apache.flink.api.java.functions.KeySelector;
-import org.apache.flink.configuration.ConfigConstants;
-import org.apache.flink.core.testutils.OneShotLatch;
-import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
-import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
-import org.apache.flink.runtime.state.VoidNamespace;
-import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-import org.apache.flink.runtime.state.ChainedStateHandle;
-import org.apache.flink.runtime.state.KeyGroupsStateHandle;
-import org.apache.flink.runtime.state.StreamStateHandle;
-import org.apache.flink.runtime.state.memory.MemoryStateBackend;
-import org.apache.flink.streaming.api.graph.StreamConfig;
-import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
-import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
-import org.apache.flink.streaming.api.watermark.Watermark;
-import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
-import org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness;
-import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
-import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.util.OperatingSystem;
-import org.apache.hadoop.conf.Configuration;
-import org.apache.hadoop.fs.FileSystem;
-import org.apache.hadoop.fs.LocalFileSystem;
-import org.junit.Assume;
-import org.junit.Before;
-import org.junit.Ignore;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.powermock.api.mockito.PowerMockito;
-import org.powermock.core.classloader.annotations.PowerMockIgnore;
-import org.powermock.core.classloader.annotations.PrepareForTest;
-import org.powermock.modules.junit4.PowerMockRunner;
-
-import java.io.File;
-import java.lang.reflect.Field;
-import java.net.URI;
-import java.util.List;
-import java.util.UUID;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-
-/**
- * Tests for asynchronous RocksDB Key/Value state checkpoints.
- */
-@RunWith(PowerMockRunner.class)
-@PrepareForTest({ResultPartitionWriter.class, FileSystem.class})
-@PowerMockIgnore({"javax.management.*", "com.sun.jndi.*"})
-@SuppressWarnings("serial")
-public class RocksDBAsyncKVSnapshotTest {
-
- @Before
- public void checkOperatingSystem() {
- Assume.assumeTrue("This test can't run successfully on Windows.", !OperatingSystem.isWindows());
- }
-
- /**
- * This ensures that asynchronous state handles are actually materialized asynchonously.
- *
- * <p>We use latches to block at various stages and see if the code still continues through
- * the parts that are not asynchronous. If the checkpoint is not done asynchronously the
- * test will simply lock forever.
- */
- @Test
- public void testAsyncCheckpoints() throws Exception {
- LocalFileSystem localFS = new LocalFileSystem();
- localFS.initialize(new URI("file:///"), new Configuration());
- PowerMockito.stub(PowerMockito.method(FileSystem.class, "get", URI.class, Configuration.class)).toReturn(localFS);
-
- final OneShotLatch delayCheckpointLatch = new OneShotLatch();
- final OneShotLatch ensureCheckpointLatch = new OneShotLatch();
-
- final OneInputStreamTask<String, String> task = new OneInputStreamTask<>();
-
- final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
-
- testHarness.configureForKeyedStream(new KeySelector<String, String>() {
- @Override
- public String getKey(String value) throws Exception {
- return value;
- }
- }, BasicTypeInfo.STRING_TYPE_INFO);
-
- StreamConfig streamConfig = testHarness.getStreamConfig();
-
- File dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state");
- File chkDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "snapshots");
-
- RocksDBStateBackend backend = new RocksDBStateBackend(chkDir.getAbsoluteFile().toURI(), new MemoryStateBackend());
- backend.setDbStoragePath(dbDir.getAbsolutePath());
-
- streamConfig.setStateBackend(backend);
-
- streamConfig.setStreamOperator(new AsyncCheckpointOperator());
-
- StreamMockEnvironment mockEnv = new StreamMockEnvironment(
- testHarness.jobConfig,
- testHarness.taskConfig,
- testHarness.memorySize,
- new MockInputSplitProvider(),
- testHarness.bufferSize) {
-
- @Override
- public void acknowledgeCheckpoint(long checkpointId) {
- super.acknowledgeCheckpoint(checkpointId);
- }
-
- @Override
- public void acknowledgeCheckpoint(long checkpointId,
- ChainedStateHandle<StreamStateHandle> chainedStateHandle,
- List<KeyGroupsStateHandle> keyGroupStateHandles) {
- super.acknowledgeCheckpoint(checkpointId, chainedStateHandle, keyGroupStateHandles);
-
- // block on the latch, to verify that triggerCheckpoint returns below,
- // even though the async checkpoint would not finish
- try {
- delayCheckpointLatch.await();
- } catch (InterruptedException e) {
- e.printStackTrace();
- }
-
-
- // should be only one k/v state
-
- assertEquals(1, keyGroupStateHandles.size());
-
- // we now know that the checkpoint went through
- ensureCheckpointLatch.trigger();
- }
- };
-
- testHarness.invoke(mockEnv);
-
- // wait for the task to be running
- for (Field field: StreamTask.class.getDeclaredFields()) {
- if (field.getName().equals("isRunning")) {
- field.setAccessible(true);
- while (!field.getBoolean(task)) {
- Thread.sleep(10);
- }
-
- }
- }
-
- testHarness.processElement(new StreamRecord<>("Wohoo", 0));
-
- task.triggerCheckpoint(42, 17);
-
- // now we allow the checkpoint
- delayCheckpointLatch.trigger();
-
- // wait for the checkpoint to go through
- ensureCheckpointLatch.await();
-
- testHarness.endInput();
- testHarness.waitForTaskCompletion();
- }
-
- /**
- * This ensures that asynchronous state handles are actually materialized asynchonously.
- *
- * <p>We use latches to block at various stages and see if the code still continues through
- * the parts that are not asynchronous. If the checkpoint is not done asynchronously the
- * test will simply lock forever.
- */
- @Test
- public void testFullyAsyncCheckpoints() throws Exception {
- LocalFileSystem localFS = new LocalFileSystem();
- localFS.initialize(new URI("file:///"), new Configuration());
- PowerMockito.stub(PowerMockito.method(FileSystem.class, "get", URI.class, Configuration.class)).toReturn(localFS);
-
- final OneShotLatch delayCheckpointLatch = new OneShotLatch();
- final OneShotLatch ensureCheckpointLatch = new OneShotLatch();
-
- final OneInputStreamTask<String, String> task = new OneInputStreamTask<>();
-
- final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
-
- testHarness.configureForKeyedStream(new KeySelector<String, String>() {
- @Override
- public String getKey(String value) throws Exception {
- return value;
- }
- }, BasicTypeInfo.STRING_TYPE_INFO);
-
- StreamConfig streamConfig = testHarness.getStreamConfig();
-
- File dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state");
- File chkDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "snapshots");
-
- RocksDBStateBackend backend = new RocksDBStateBackend(chkDir.getAbsoluteFile().toURI(), new MemoryStateBackend());
- backend.setDbStoragePath(dbDir.getAbsolutePath());
-// backend.enableFullyAsyncSnapshots();
-
- streamConfig.setStateBackend(backend);
-
- streamConfig.setStreamOperator(new AsyncCheckpointOperator());
-
- StreamMockEnvironment mockEnv = new StreamMockEnvironment(
- testHarness.jobConfig,
- testHarness.taskConfig,
- testHarness.memorySize,
- new MockInputSplitProvider(),
- testHarness.bufferSize) {
-
- @Override
- public void acknowledgeCheckpoint(long checkpointId) {
- super.acknowledgeCheckpoint(checkpointId);
- }
-
- @Override
- public void acknowledgeCheckpoint(long checkpointId,
- ChainedStateHandle<StreamStateHandle> chainedStateHandle,
- List<KeyGroupsStateHandle> keyGroupStateHandles) {
- super.acknowledgeCheckpoint(checkpointId, chainedStateHandle, keyGroupStateHandles);
-
- // block on the latch, to verify that triggerCheckpoint returns below,
- // even though the async checkpoint would not finish
- try {
- delayCheckpointLatch.await();
- } catch (InterruptedException e) {
- e.printStackTrace();
- }
-
- // should be only one k/v state
- assertEquals(1, keyGroupStateHandles.size());
-
- // we now know that the checkpoint went through
- ensureCheckpointLatch.trigger();
- }
- };
-
- testHarness.invoke(mockEnv);
-
- // wait for the task to be running
- for (Field field: StreamTask.class.getDeclaredFields()) {
- if (field.getName().equals("isRunning")) {
- field.setAccessible(true);
- while (!field.getBoolean(task)) {
- Thread.sleep(10);
- }
-
- }
- }
-
- testHarness.processElement(new StreamRecord<>("Wohoo", 0));
-
- task.triggerCheckpoint(42, 17);
-
- // now we allow the checkpoint
- delayCheckpointLatch.trigger();
-
- // wait for the checkpoint to go through
- ensureCheckpointLatch.await();
-
- testHarness.endInput();
- testHarness.waitForTaskCompletion();
- }
-
-
- // ------------------------------------------------------------------------
-
- public static class AsyncCheckpointOperator
- extends AbstractStreamOperator<String>
- implements OneInputStreamOperator<String, String> {
-
- @Override
- public void open() throws Exception {
- super.open();
-
- // also get the state in open, this way we are sure that it was created before
- // we trigger the test checkpoint
- ValueState<String> state = getPartitionedState(
- VoidNamespace.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- new ValueStateDescriptor<>("count",
- StringSerializer.INSTANCE, "hello"));
-
- }
-
- @Override
- public void processElement(StreamRecord<String> element) throws Exception {
- // we also don't care
-
- ValueState<String> state = getPartitionedState(
- VoidNamespace.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- new ValueStateDescriptor<>("count",
- StringSerializer.INSTANCE, "hello"));
-
- state.update(element.getValue());
- }
-
- @Override
- public void processWatermark(Watermark mark) throws Exception {
- // not interested
- }
- }
-
- public static class DummyMapFunction<T> implements MapFunction<T, T> {
- @Override
- public T map(T value) { return value; }
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
new file mode 100644
index 0000000..624905c
--- /dev/null
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
@@ -0,0 +1,397 @@
+/*
+ * 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.contrib.streaming.state;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.functions.MapFunction;
+import org.apache.flink.api.common.state.ValueState;
+import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.api.common.typeutils.base.StringSerializer;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.ConfigConstants;
+import org.apache.flink.core.testutils.OneShotLatch;
+import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.VoidNamespace;
+import org.apache.flink.runtime.state.VoidNamespaceSerializer;
+import org.apache.flink.runtime.state.memory.MemCheckpointStreamFactory;
+import org.apache.flink.runtime.state.memory.MemoryStateBackend;
+import org.apache.flink.streaming.api.graph.StreamConfig;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.runtime.tasks.AsynchronousException;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
+import org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness;
+import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
+import org.apache.flink.util.OperatingSystem;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.LocalFileSystem;
+import org.junit.Assert;
+import org.junit.Assume;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.mockito.PowerMockito;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
+
+import java.io.File;
+import java.io.IOException;
+import java.lang.reflect.Field;
+import java.net.URI;
+import java.util.Arrays;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.CancellationException;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * Tests for asynchronous RocksDB Key/Value state checkpoints.
+ */
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({ResultPartitionWriter.class, FileSystem.class})
+@PowerMockIgnore({"javax.management.*", "com.sun.jndi.*"})
+@SuppressWarnings("serial")
+public class RocksDBAsyncSnapshotTest {
+
+ @Before
+ public void checkOperatingSystem() {
+ Assume.assumeTrue("This test can't run successfully on Windows.", !OperatingSystem.isWindows());
+ }
+
+ /**
+ * This ensures that asynchronous state handles are actually materialized asynchonously.
+ *
+ * <p>We use latches to block at various stages and see if the code still continues through
+ * the parts that are not asynchronous. If the checkpoint is not done asynchronously the
+ * test will simply lock forever.
+ */
+ @Test
+ public void testFullyAsyncSnapshot() throws Exception {
+
+ LocalFileSystem localFS = new LocalFileSystem();
+ localFS.initialize(new URI("file:///"), new Configuration());
+ PowerMockito.stub(PowerMockito.method(FileSystem.class, "get", URI.class, Configuration.class)).toReturn(localFS);
+
+ final OneInputStreamTask<String, String> task = new OneInputStreamTask<>();
+
+ final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
+
+ testHarness.configureForKeyedStream(new KeySelector<String, String>() {
+ @Override
+ public String getKey(String value) throws Exception {
+ return value;
+ }
+ }, BasicTypeInfo.STRING_TYPE_INFO);
+
+ StreamConfig streamConfig = testHarness.getStreamConfig();
+
+ File dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state");
+ File chkDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "snapshots");
+
+ RocksDBStateBackend backend = new RocksDBStateBackend(chkDir.getAbsoluteFile().toURI(), new MemoryStateBackend());
+ backend.setDbStoragePath(dbDir.getAbsolutePath());
+
+ streamConfig.setStateBackend(backend);
+
+ streamConfig.setStreamOperator(new AsyncCheckpointOperator());
+
+ final OneShotLatch delayCheckpointLatch = new OneShotLatch();
+ final OneShotLatch ensureCheckpointLatch = new OneShotLatch();
+
+ StreamMockEnvironment mockEnv = new StreamMockEnvironment(
+ testHarness.jobConfig,
+ testHarness.taskConfig,
+ testHarness.memorySize,
+ new MockInputSplitProvider(),
+ testHarness.bufferSize) {
+
+ @Override
+ public void acknowledgeCheckpoint(long checkpointId) {
+ super.acknowledgeCheckpoint(checkpointId);
+ }
+
+ @Override
+ public void acknowledgeCheckpoint(long checkpointId,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
+ List<KeyGroupsStateHandle> keyGroupStateHandles) {
+ super.acknowledgeCheckpoint(checkpointId, chainedStateHandle, keyGroupStateHandles);
+
+ // block on the latch, to verify that triggerCheckpoint returns below,
+ // even though the async checkpoint would not finish
+ try {
+ delayCheckpointLatch.await();
+ } catch (InterruptedException e) {
+ e.printStackTrace();
+ }
+
+ // should be only one k/v state
+ assertEquals(1, keyGroupStateHandles.size());
+
+ // we now know that the checkpoint went through
+ ensureCheckpointLatch.trigger();
+ }
+ };
+
+ testHarness.invoke(mockEnv);
+
+ // wait for the task to be running
+ for (Field field: StreamTask.class.getDeclaredFields()) {
+ if (field.getName().equals("isRunning")) {
+ field.setAccessible(true);
+ while (!field.getBoolean(task)) {
+ Thread.sleep(10);
+ }
+
+ }
+ }
+
+ testHarness.processElement(new StreamRecord<>("Wohoo", 0));
+
+ task.triggerCheckpoint(42, 17);
+
+ // now we allow the checkpoint
+ delayCheckpointLatch.trigger();
+
+ // wait for the checkpoint to go through
+ ensureCheckpointLatch.await();
+
+ testHarness.endInput();
+ testHarness.waitForTaskCompletion();
+ }
+
+ /**
+ * This tests ensures that canceling of asynchronous snapshots works as expected and does not block.
+ * @throws Exception
+ */
+ @Test
+ public void testCancelFullyAsyncCheckpoints() throws Exception {
+ LocalFileSystem localFS = new LocalFileSystem();
+ localFS.initialize(new URI("file:///"), new Configuration());
+ PowerMockito.stub(PowerMockito.method(FileSystem.class, "get", URI.class, Configuration.class)).toReturn(localFS);
+
+ final OneInputStreamTask<String, String> task = new OneInputStreamTask<>();
+
+ //ensure that the async threads complete before invoke method of the tasks returns.
+ task.setThreadPoolTerminationTimeout(Long.MAX_VALUE);
+
+ final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
+
+ testHarness.configureForKeyedStream(new KeySelector<String, String>() {
+ @Override
+ public String getKey(String value) throws Exception {
+ return value;
+ }
+ }, BasicTypeInfo.STRING_TYPE_INFO);
+
+ StreamConfig streamConfig = testHarness.getStreamConfig();
+
+ File dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state");
+ File chkDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "snapshots");
+
+ BlockingStreamMemoryStateBackend memoryStateBackend = new BlockingStreamMemoryStateBackend();
+
+ RocksDBStateBackend backend = new RocksDBStateBackend(chkDir.getAbsoluteFile().toURI(), memoryStateBackend);
+ backend.setDbStoragePath(dbDir.getAbsolutePath());
+
+ streamConfig.setStateBackend(backend);
+
+ streamConfig.setStreamOperator(new AsyncCheckpointOperator());
+
+ StreamMockEnvironment mockEnv = new StreamMockEnvironment(
+ testHarness.jobConfig,
+ testHarness.taskConfig,
+ testHarness.memorySize,
+ new MockInputSplitProvider(),
+ testHarness.bufferSize);
+
+ testHarness.invoke(mockEnv);
+
+ // wait for the task to be running
+ for (Field field: StreamTask.class.getDeclaredFields()) {
+ if (field.getName().equals("isRunning")) {
+ field.setAccessible(true);
+ while (!field.getBoolean(task)) {
+ Thread.sleep(10);
+ }
+
+ }
+ }
+
+ testHarness.processElement(new StreamRecord<>("Wohoo", 0));
+
+ task.triggerCheckpoint(42, 17);
+
+ BlockingStreamMemoryStateBackend.waitFirstWriteLatch.await();
+ task.cancel();
+
+ BlockingStreamMemoryStateBackend.unblockCancelLatch.trigger();
+
+ testHarness.endInput();
+ try {
+ testHarness.waitForTaskCompletion();
+ Assert.fail("Operation completed. Cancel failed.");
+ } catch (Exception expected) {
+ // we expect the exception from canceling snapshots
+ Throwable cause = expected.getCause();
+ if(cause instanceof AsynchronousException) {
+ AsynchronousException asynchronousException = (AsynchronousException) cause;
+ cause = asynchronousException.getCause();
+ Assert.assertTrue("Unexpected Exception: " + cause,
+ cause instanceof CancellationException //future canceled
+ || cause instanceof InterruptedException); //thread interrupted
+
+ } else {
+ Assert.fail();
+ }
+ }
+ }
+
+ @Test
+ public void testConsistentSnapshotSerializationFlagsAndMasks() {
+
+ Assert.assertEquals(0xFFFF, RocksDBKeyedStateBackend.RocksDBSnapshotOperation.END_OF_KEY_GROUP_MARK);
+ Assert.assertEquals(0x80, RocksDBKeyedStateBackend.RocksDBSnapshotOperation.FIRST_BIT_IN_BYTE_MASK);
+
+ byte[] expectedKey = new byte[] {42, 42};
+ byte[] modKey = expectedKey.clone();
+
+ Assert.assertFalse(RocksDBKeyedStateBackend.RocksDBSnapshotOperation.hasMetaDataFollowsFlag(modKey));
+
+ RocksDBKeyedStateBackend.RocksDBSnapshotOperation.setMetaDataFollowsFlagInKey(modKey);
+ Assert.assertTrue(RocksDBKeyedStateBackend.RocksDBSnapshotOperation.hasMetaDataFollowsFlag(modKey));
+
+ RocksDBKeyedStateBackend.RocksDBSnapshotOperation.clearMetaDataFollowsFlag(modKey);
+ Assert.assertFalse(RocksDBKeyedStateBackend.RocksDBSnapshotOperation.hasMetaDataFollowsFlag(modKey));
+
+ Assert.assertTrue(Arrays.equals(expectedKey, modKey));
+ }
+
+ // ------------------------------------------------------------------------
+
+ /**
+ * Creates us a CheckpointStateOutputStream that blocks write ops on a latch to delay writing of snapshots.
+ */
+ static class BlockingStreamMemoryStateBackend extends MemoryStateBackend {
+
+ public static OneShotLatch waitFirstWriteLatch = new OneShotLatch();
+
+ public static OneShotLatch unblockCancelLatch = new OneShotLatch();
+
+ volatile boolean closed = false;
+
+ @Override
+ public CheckpointStreamFactory createStreamFactory(JobID jobId, String operatorIdentifier) throws IOException {
+ return new MemCheckpointStreamFactory(4 * 1024 * 1024) {
+ @Override
+ public CheckpointStateOutputStream createCheckpointStateOutputStream(long checkpointID, long timestamp) throws Exception {
+
+ return new MemoryCheckpointOutputStream(4 * 1024 * 1024) {
+ @Override
+ public void write(int b) throws IOException {
+ waitFirstWriteLatch.trigger();
+ try {
+ unblockCancelLatch.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ if(closed) {
+ throw new IOException("Stream closed.");
+ }
+ super.write(b);
+ }
+
+ @Override
+ public void write(byte[] b, int off, int len) throws IOException {
+ waitFirstWriteLatch.trigger();
+ try {
+ unblockCancelLatch.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ if(closed) {
+ throw new IOException("Stream closed.");
+ }
+ super.write(b, off, len);
+ }
+
+ @Override
+ public void close() {
+ closed = true;
+ super.close();
+ }
+ };
+ }
+ };
+ }
+ }
+
+ public static class AsyncCheckpointOperator
+ extends AbstractStreamOperator<String>
+ implements OneInputStreamOperator<String, String> {
+
+ @Override
+ public void open() throws Exception {
+ super.open();
+
+ // also get the state in open, this way we are sure that it was created before
+ // we trigger the test checkpoint
+ ValueState<String> state = getPartitionedState(
+ VoidNamespace.INSTANCE,
+ VoidNamespaceSerializer.INSTANCE,
+ new ValueStateDescriptor<>("count",
+ StringSerializer.INSTANCE, "hello"));
+
+ }
+
+ @Override
+ public void processElement(StreamRecord<String> element) throws Exception {
+ // we also don't care
+
+ ValueState<String> state = getPartitionedState(
+ VoidNamespace.INSTANCE,
+ VoidNamespaceSerializer.INSTANCE,
+ new ValueStateDescriptor<>("count",
+ StringSerializer.INSTANCE, "hello"));
+
+ state.update(element.getValue());
+ }
+
+ @Override
+ public void processWatermark(Watermark mark) throws Exception {
+ // not interested
+ }
+ }
+
+ public static class DummyMapFunction<T> implements MapFunction<T, T> {
+ @Override
+ public T map(T value) { return value; }
+ }
+}
[08/27] flink git commit: [FLINK-3761] Refactor State Backends/Make
Keyed State Key-Group Aware
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java
index d653f73..30d91b6 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java
@@ -31,18 +31,23 @@ import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.runtime.jobgraph.JobVertexID;
-import org.apache.flink.runtime.query.KvStateID;
+import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.query.KvStateRegistry;
import org.apache.flink.runtime.query.KvStateServerAddress;
import org.apache.flink.runtime.query.netty.message.KvStateRequestResult;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
import org.apache.flink.runtime.query.netty.message.KvStateRequestType;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-import org.apache.flink.runtime.state.memory.MemValueState;
+import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.junit.AfterClass;
import org.junit.Test;
@@ -52,6 +57,7 @@ import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
public class KvStateServerTest {
@@ -84,26 +90,37 @@ public class KvStateServerTest {
KvStateServerAddress serverAddress = server.getAddress();
- // Register state
- MemValueState<Integer, VoidNamespace, Integer> kvState = new MemValueState<>(
+ AbstractStateBackend abstractBackend = new MemoryStateBackend();
+ DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
+ dummyEnv.setKvStateRegistry(registry);
+ KeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(
+ dummyEnv,
+ new JobID(),
+ "test_op",
IntSerializer.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null));
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ registry.createTaskRegistry(new JobID(), new JobVertexID()));
- KvStateID kvStateId = registry.registerKvState(
- new JobID(),
- new JobVertexID(),
- 0,
- "vanilla",
- kvState);
+ final KvStateServerHandlerTest.TestRegistryListener registryListener =
+ new KvStateServerHandlerTest.TestRegistryListener();
+
+ registry.registerListener(registryListener);
+
+ ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null);
+ desc.setQueryable("vanilla");
+
+ ValueState<Integer> state = backend.getPartitionedState(
+ VoidNamespace.INSTANCE,
+ VoidNamespaceSerializer.INSTANCE,
+ desc);
// Update KvState
int expectedValue = 712828289;
int key = 99812822;
- kvState.setCurrentKey(key);
- kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
- kvState.update(expectedValue);
+ backend.setCurrentKey(key);
+ state.update(expectedValue);
// Request
byte[] serializedKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(
@@ -128,10 +145,12 @@ public class KvStateServerTest {
.sync().channel();
long requestId = Integer.MAX_VALUE + 182828L;
+
+ assertTrue(registryListener.registrationName.equals("vanilla"));
ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(
channel.alloc(),
requestId,
- kvStateId,
+ registryListener.kvStateId,
serializedKeyAndNamespace);
channel.writeAndFlush(request);
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java
index 04fa089..bc0b9c3 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java
@@ -19,6 +19,7 @@
package org.apache.flink.runtime.state;
import org.apache.commons.io.FileUtils;
+import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.core.fs.Path;
@@ -29,7 +30,9 @@ import org.apache.flink.runtime.state.filesystem.FileStateHandle;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+import org.junit.Rule;
import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
import java.io.File;
import java.io.IOException;
@@ -42,17 +45,13 @@ import static org.junit.Assert.*;
public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
- private File stateDir;
+ @Rule
+ public TemporaryFolder tempFolder = new TemporaryFolder();
@Override
protected FsStateBackend getStateBackend() throws Exception {
- stateDir = new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString());
- return new FsStateBackend(localFileUri(stateDir));
- }
-
- @Override
- protected void cleanup() throws Exception {
- deleteDirectorySilently(stateDir);
+ File checkpointPath = tempFolder.newFolder();
+ return new FsStateBackend(localFileUri(checkpointPath));
}
// disable these because the verification does not work for this state backend
@@ -69,66 +68,19 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
public void testReducingStateRestoreWithWrongSerializers() {}
@Test
- public void testSetupAndSerialization() {
- File tempDir = new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString());
- try {
- final String backendDir = localFileUri(tempDir);
- FsStateBackend originalBackend = new FsStateBackend(backendDir);
-
- assertFalse(originalBackend.isInitialized());
- assertEquals(new URI(backendDir), originalBackend.getBasePath().toUri());
- assertNull(originalBackend.getCheckpointDirectory());
-
- // serialize / copy the backend
- FsStateBackend backend = CommonTestUtils.createCopySerializable(originalBackend);
- assertFalse(backend.isInitialized());
- assertEquals(new URI(backendDir), backend.getBasePath().toUri());
- assertNull(backend.getCheckpointDirectory());
-
- // no file operations should be possible right now
- try {
- FsStateBackend.FsCheckpointStateOutputStream out = backend.createCheckpointStateOutputStream(
- 2L,
- System.currentTimeMillis());
-
- out.write(1);
- out.closeAndGetHandle();
- fail("should fail with an exception");
- } catch (IllegalStateException e) {
- // supreme!
- }
-
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test-op", IntSerializer.INSTANCE);
- assertNotNull(backend.getCheckpointDirectory());
-
- File checkpointDir = new File(backend.getCheckpointDirectory().toUri().getPath());
- assertTrue(checkpointDir.exists());
- assertTrue(isDirectoryEmpty(checkpointDir));
-
- backend.disposeAllStateForCurrentJob();
- assertNull(backend.getCheckpointDirectory());
+ public void testStateOutputStream() throws IOException {
+ File basePath = tempFolder.newFolder().getAbsoluteFile();
- assertTrue(isDirectoryEmpty(tempDir));
- }
- catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
- }
- finally {
- deleteDirectorySilently(tempDir);
- }
- }
-
- @Test
- public void testStateOutputStream() {
- File tempDir = new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString());
try {
// the state backend has a very low in-mem state threshold (15 bytes)
- FsStateBackend backend = CommonTestUtils.createCopySerializable(new FsStateBackend(tempDir.toURI(), 15));
+ FsStateBackend backend = CommonTestUtils.createCopySerializable(new FsStateBackend(basePath.toURI(), 15));
+ JobID jobId = new JobID();
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test-op", IntSerializer.INSTANCE);
+ // we know how FsCheckpointStreamFactory is implemented so we know where it
+ // will store checkpoints
+ File checkpointPath = new File(basePath.getAbsolutePath(), jobId.toString());
- File checkpointDir = new File(backend.getCheckpointDirectory().toUri().getPath());
+ CheckpointStreamFactory streamFactory = backend.createStreamFactory(jobId, "test_op");
byte[] state1 = new byte[1274673];
byte[] state2 = new byte[1];
@@ -143,12 +95,14 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
long checkpointId = 97231523452L;
- FsStateBackend.FsCheckpointStateOutputStream stream1 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
- FsStateBackend.FsCheckpointStateOutputStream stream2 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
- FsStateBackend.FsCheckpointStateOutputStream stream3 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
+ CheckpointStreamFactory.CheckpointStateOutputStream stream1 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
+
+ CheckpointStreamFactory.CheckpointStateOutputStream stream2 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
+
+ CheckpointStreamFactory.CheckpointStateOutputStream stream3 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
stream1.write(state1);
stream2.write(state2);
@@ -160,15 +114,15 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
// use with try-with-resources
StreamStateHandle handle4;
- try (AbstractStateBackend.CheckpointStateOutputStream stream4 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis())) {
+ try (CheckpointStreamFactory.CheckpointStateOutputStream stream4 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis())) {
stream4.write(state4);
handle4 = stream4.closeAndGetHandle();
}
// close before accessing handle
- AbstractStateBackend.CheckpointStateOutputStream stream5 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
+ CheckpointStreamFactory.CheckpointStateOutputStream stream5 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
stream5.write(state4);
stream5.close();
try {
@@ -180,7 +134,7 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
validateBytesInStream(handle1.openInputStream(), state1);
handle1.discardState();
- assertFalse(isDirectoryEmpty(checkpointDir));
+ assertFalse(isDirectoryEmpty(basePath));
ensureLocalFileDeleted(handle1.getFilePath());
validateBytesInStream(handle2.openInputStream(), state2);
@@ -191,15 +145,12 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
validateBytesInStream(handle4.openInputStream(), state4);
handle4.discardState();
- assertTrue(isDirectoryEmpty(checkpointDir));
+ assertTrue(isDirectoryEmpty(checkpointPath));
}
catch (Exception e) {
e.printStackTrace();
fail(e.getMessage());
}
- finally {
- deleteDirectorySilently(tempDir);
- }
}
// ------------------------------------------------------------------------
@@ -253,8 +204,7 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
@Test
public void testConcurrentMapIfQueryable() throws Exception {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- StateBackendTestBase.testConcurrentMapIfQueryable(backend);
+ super.testConcurrentMapIfQueryable();
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java
index 940b337..944938b 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.state;
+import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
@@ -40,9 +41,6 @@ public class MemoryStateBackendTest extends StateBackendTestBase<MemoryStateBack
return new MemoryStateBackend();
}
- @Override
- protected void cleanup() throws Exception { }
-
// disable these because the verification does not work for this state backend
@Override
@Test
@@ -60,15 +58,15 @@ public class MemoryStateBackendTest extends StateBackendTestBase<MemoryStateBack
public void testOversizedState() {
try {
MemoryStateBackend backend = new MemoryStateBackend(10);
+ CheckpointStreamFactory streamFactory = backend.createStreamFactory(new JobID(), "test_op");
HashMap<String, Integer> state = new HashMap<>();
state.put("hey there", 2);
state.put("the crazy brown fox stumbles over a sentence that does not contain every letter", 77);
try {
- AbstractStateBackend.CheckpointStateOutputStream outStream = backend.createCheckpointStateOutputStream(
- 12,
- 459);
+ CheckpointStreamFactory.CheckpointStateOutputStream outStream =
+ streamFactory.createCheckpointStateOutputStream(12, 459);
ObjectOutputStream oos = new ObjectOutputStream(outStream);
oos.writeObject(state);
@@ -93,12 +91,13 @@ public class MemoryStateBackendTest extends StateBackendTestBase<MemoryStateBack
public void testStateStream() {
try {
MemoryStateBackend backend = new MemoryStateBackend();
+ CheckpointStreamFactory streamFactory = backend.createStreamFactory(new JobID(), "test_op");
HashMap<String, Integer> state = new HashMap<>();
state.put("hey there", 2);
state.put("the crazy brown fox stumbles over a sentence that does not contain every letter", 77);
- AbstractStateBackend.CheckpointStateOutputStream os = backend.createCheckpointStateOutputStream(1, 2);
+ CheckpointStreamFactory.CheckpointStateOutputStream os = streamFactory.createCheckpointStateOutputStream(1, 2);
ObjectOutputStream oos = new ObjectOutputStream(os);
oos.writeObject(state);
oos.flush();
@@ -121,12 +120,13 @@ public class MemoryStateBackendTest extends StateBackendTestBase<MemoryStateBack
public void testOversizedStateStream() {
try {
MemoryStateBackend backend = new MemoryStateBackend(10);
+ CheckpointStreamFactory streamFactory = backend.createStreamFactory(new JobID(), "test_op");
HashMap<String, Integer> state = new HashMap<>();
state.put("hey there", 2);
state.put("the crazy brown fox stumbles over a sentence that does not contain every letter", 77);
- AbstractStateBackend.CheckpointStateOutputStream os = backend.createCheckpointStateOutputStream(1, 2);
+ CheckpointStreamFactory.CheckpointStateOutputStream os = streamFactory.createCheckpointStateOutputStream(1, 2);
ObjectOutputStream oos = new ObjectOutputStream(os);
try {
@@ -147,7 +147,6 @@ public class MemoryStateBackendTest extends StateBackendTestBase<MemoryStateBack
@Test
public void testConcurrentMapIfQueryable() throws Exception {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- StateBackendTestBase.testConcurrentMapIfQueryable(backend);
+ super.testConcurrentMapIfQueryable();
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
index 834c35c..f094bd5 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
@@ -21,10 +21,12 @@ package org.apache.flink.runtime.state;
import com.google.common.base.Joiner;
import org.apache.commons.io.output.ByteArrayOutputStream;
import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.functions.FoldFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.state.FoldingState;
import org.apache.flink.api.common.state.FoldingStateDescriptor;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ReducingState;
@@ -37,21 +39,25 @@ import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.api.common.typeutils.base.LongSerializer;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.runtime.checkpoint.CheckpointCoordinator;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.query.KvStateID;
import org.apache.flink.runtime.query.KvStateRegistry;
import org.apache.flink.runtime.query.KvStateRegistryListener;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
+import org.apache.flink.runtime.state.heap.AbstractHeapState;
+import org.apache.flink.runtime.state.heap.StateTable;
import org.apache.flink.types.IntValue;
-import org.junit.After;
-import org.junit.Before;
import org.junit.Test;
import java.util.Collections;
-import java.util.HashMap;
import java.util.List;
-import java.util.Map;
+import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Future;
+import java.util.concurrent.RunnableFuture;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
@@ -71,27 +77,77 @@ import static org.mockito.Mockito.verify;
@SuppressWarnings("serial")
public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
- protected B backend;
-
protected abstract B getStateBackend() throws Exception;
- protected abstract void cleanup() throws Exception;
+ protected CheckpointStreamFactory createStreamFactory() throws Exception {
+ return getStateBackend().createStreamFactory(new JobID(), "test_op");
+ }
+
+ protected <K> KeyedStateBackend<K> createKeyedBackend(TypeSerializer<K> keySerializer) throws Exception {
+ return createKeyedBackend(keySerializer, new DummyEnvironment("test", 1, 0));
+ }
+
+ protected <K> KeyedStateBackend<K> createKeyedBackend(TypeSerializer<K> keySerializer, Environment env) throws Exception {
+ return createKeyedBackend(
+ keySerializer,
+ new HashKeyGroupAssigner<K>(10),
+ new KeyGroupRange(0, 9),
+ env);
+ }
+
+ protected <K> KeyedStateBackend<K> createKeyedBackend(
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ Environment env) throws Exception {
+ return getStateBackend().createKeyedStateBackend(
+ env,
+ new JobID(),
+ "test_op",
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange,
+ env.getTaskKvStateRegistry());
+ }
- @Before
- public void setup() throws Exception {
- this.backend = getStateBackend();
+ protected <K> KeyedStateBackend<K> restoreKeyedBackend(TypeSerializer<K> keySerializer, KeyGroupsStateHandle state) throws Exception {
+ return restoreKeyedBackend(keySerializer, state, new DummyEnvironment("test", 1, 0));
}
- @After
- public void teardown() throws Exception {
- this.backend.discardState();
- cleanup();
+ protected <K> KeyedStateBackend<K> restoreKeyedBackend(
+ TypeSerializer<K> keySerializer,
+ KeyGroupsStateHandle state,
+ Environment env) throws Exception {
+ return restoreKeyedBackend(
+ keySerializer,
+ new HashKeyGroupAssigner<K>(10),
+ new KeyGroupRange(0, 9),
+ Collections.singletonList(state),
+ env);
+ }
+
+ protected <K> KeyedStateBackend<K> restoreKeyedBackend(
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ List<KeyGroupsStateHandle> state,
+ Environment env) throws Exception {
+ return getStateBackend().restoreKeyedStateBackend(
+ env,
+ new JobID(),
+ "test_op",
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange,
+ state,
+ env.getTaskKvStateRegistry());
}
@Test
@SuppressWarnings("unchecked")
public void testValueState() throws Exception {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class, null);
kvId.initializeSerializerUnlessSet(new ExecutionConfig());
@@ -102,7 +158,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
ValueState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> kvState = (KvState<Integer, VoidNamespace, ?, ?, B>) state;
+ KvState<VoidNamespace> kvState = (KvState<VoidNamespace>) state;
// some modifications to the state
backend.setCurrentKey(1);
@@ -118,13 +174,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("1", getSerializedValue(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
// draw a snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot1 = backend.snapshotPartitionedState(682375462378L, 2);
-
- for (String key: snapshot1.keySet()) {
- if (snapshot1.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot1.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot1.get(key)).materialize());
- }
- }
+ KeyGroupsStateHandle snapshot1 = runSnapshot(backend.snapshot(682375462378L, 2, streamFactory));
// make some more modifications
backend.setCurrentKey(1);
@@ -135,13 +185,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
state.update("u3");
// draw another snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot2 = backend.snapshotPartitionedState(682375462379L, 4);
-
- for (String key: snapshot2.keySet()) {
- if (snapshot2.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot2.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot2.get(key)).materialize());
- }
- }
+ KeyGroupsStateHandle snapshot2 = runSnapshot(backend.snapshot(682375462379L, 4, streamFactory));
// validate the original state
backend.setCurrentKey(1);
@@ -154,18 +198,14 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("u3", state.value());
assertEquals("u3", getSerializedValue(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.discardState();
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ backend.close();
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
-
- for (String key: snapshot1.keySet()) {
- snapshot1.get(key).discardState();
- }
+ snapshot1.discardState();
ValueState<String> restored1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> restoredKvState1 = (KvState<Integer, VoidNamespace, ?, ?, B>) restored1;
+ KvState<VoidNamespace> restoredKvState1 = (KvState<VoidNamespace>) restored1;
backend.setCurrentKey(1);
assertEquals("1", restored1.value());
@@ -174,18 +214,14 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("2", restored1.value());
assertEquals("2", getSerializedValue(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.discardState();
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ backend.close();
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot2);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot2);
-
- for (String key: snapshot2.keySet()) {
- snapshot2.get(key).discardState();
- }
+ snapshot2.discardState();
ValueState<String> restored2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> restoredKvState2 = (KvState<Integer, VoidNamespace, ?, ?, B>) restored2;
+ KvState<VoidNamespace> restoredKvState2 = (KvState<VoidNamespace>) restored2;
backend.setCurrentKey(1);
assertEquals("u1", restored2.value());
@@ -196,6 +232,68 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
backend.setCurrentKey(3);
assertEquals("u3", restored2.value());
assertEquals("u3", getSerializedValue(restoredKvState2, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
+
+ backend.close();
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ public void testMultipleValueStates() throws Exception {
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+
+ KeyedStateBackend<Integer> backend = createKeyedBackend(
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ new DummyEnvironment("test_op", 1, 0));
+
+ ValueStateDescriptor<String> desc1 = new ValueStateDescriptor<>("a-string", StringSerializer.INSTANCE, null);
+ ValueStateDescriptor<Integer> desc2 = new ValueStateDescriptor<>("an-integer", IntSerializer.INSTANCE, null);
+
+ desc1.initializeSerializerUnlessSet(new ExecutionConfig());
+ desc2.initializeSerializerUnlessSet(new ExecutionConfig());
+
+ ValueState<String> state1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc1);
+ ValueState<Integer> state2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc2);
+
+ // some modifications to the state
+ backend.setCurrentKey(1);
+ assertNull(state1.value());
+ assertNull(state2.value());
+ state1.update("1");
+
+ // state2 should still have nothing
+ assertEquals("1", state1.value());
+ assertNull(state2.value());
+ state2.update(13);
+
+ // both have some state now
+ assertEquals("1", state1.value());
+ assertEquals(13, (int) state2.value());
+
+ // draw a snapshot
+ KeyGroupsStateHandle snapshot1 = runSnapshot(backend.snapshot(682375462378L, 2, streamFactory));
+
+ backend.close();
+ backend = restoreKeyedBackend(
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ Collections.singletonList(snapshot1),
+ new DummyEnvironment("test_op", 1, 0));
+
+ snapshot1.discardState();
+
+ backend.setCurrentKey(1);
+
+ state1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc1);
+ state2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc2);
+
+ // verify that they are still the same
+ assertEquals("1", state1.value());
+ assertEquals(13, (int) state2.value());
+
+ backend.close();
}
/**
@@ -217,7 +315,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
// alrighty
}
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
ValueStateDescriptor<Long> kvId = new ValueStateDescriptor<>("id", LongSerializer.INSTANCE, 42L);
kvId.initializeSerializerUnlessSet(new ExecutionConfig());
@@ -246,31 +345,24 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals(42L, (long) state.value());
// draw a snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot1 = backend.snapshotPartitionedState(682375462378L, 2);
-
- for (String key: snapshot1.keySet()) {
- if (snapshot1.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot1.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot1.get(key)).materialize());
- }
- }
-
- backend.discardState();
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ KeyGroupsStateHandle snapshot1 = runSnapshot(backend.snapshot(682375462378L, 2, streamFactory));
- backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
+ backend.close();
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
- for (String key: snapshot1.keySet()) {
- snapshot1.get(key).discardState();
- }
+ snapshot1.discardState();
backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
+
+ backend.close();
}
@Test
@SuppressWarnings("unchecked,rawtypes")
public void testListState() {
try {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
ListStateDescriptor<String> kvId = new ListStateDescriptor<>("id", String.class);
kvId.initializeSerializerUnlessSet(new ExecutionConfig());
@@ -281,7 +373,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
ListState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> kvState = (KvState<Integer, VoidNamespace, ?, ?, B>) state;
+ KvState<VoidNamespace> kvState = (KvState<VoidNamespace>) state;
Joiner joiner = Joiner.on(",");
// some modifications to the state
@@ -298,13 +390,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("1", joiner.join(getSerializedList(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
// draw a snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot1 = backend.snapshotPartitionedState(682375462378L, 2);
-
- for (String key: snapshot1.keySet()) {
- if (snapshot1.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot1.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot1.get(key)).materialize());
- }
- }
+ KeyGroupsStateHandle snapshot1 = runSnapshot(backend.snapshot(682375462378L, 2, streamFactory));
// make some more modifications
backend.setCurrentKey(1);
@@ -315,13 +401,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
state.add("u3");
// draw another snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot2 = backend.snapshotPartitionedState(682375462379L, 4);
-
- for (String key: snapshot2.keySet()) {
- if (snapshot2.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot2.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot2.get(key)).materialize());
- }
- }
+ KeyGroupsStateHandle snapshot2 = runSnapshot(backend.snapshot(682375462379L, 4, streamFactory));
// validate the original state
backend.setCurrentKey(1);
@@ -334,19 +414,14 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("u3", joiner.join(state.get()));
assertEquals("u3", joiner.join(getSerializedList(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
- backend.discardState();
-
+ backend.close();
// restore the first snapshot and validate it
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
-
- for (String key: snapshot1.keySet()) {
- snapshot1.get(key).discardState();
- }
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
+ snapshot1.discardState();
ListState<String> restored1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> restoredKvState1 = (KvState<Integer, VoidNamespace, ?, ?, B>) restored1;
+ KvState<VoidNamespace> restoredKvState1 = (KvState<VoidNamespace>) restored1;
backend.setCurrentKey(1);
assertEquals("1", joiner.join(restored1.get()));
@@ -355,19 +430,14 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("2", joiner.join(restored1.get()));
assertEquals("2", joiner.join(getSerializedList(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
- backend.discardState();
-
+ backend.close();
// restore the second snapshot and validate it
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot2);
-
- for (String key: snapshot2.keySet()) {
- snapshot2.get(key).discardState();
- }
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot2);
+ snapshot2.discardState();
ListState<String> restored2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> restoredKvState2 = (KvState<Integer, VoidNamespace, ?, ?, B>) restored2;
+ KvState<VoidNamespace> restoredKvState2 = (KvState<VoidNamespace>) restored2;
backend.setCurrentKey(1);
assertEquals("1,u1", joiner.join(restored2.get()));
@@ -378,6 +448,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
backend.setCurrentKey(3);
assertEquals("u3", joiner.join(restored2.get()));
assertEquals("u3", joiner.join(getSerializedList(restoredKvState2, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
+
+ backend.close();
}
catch (Exception e) {
e.printStackTrace();
@@ -389,7 +461,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
@SuppressWarnings("unchecked")
public void testReducingState() {
try {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
ReducingStateDescriptor<String> kvId = new ReducingStateDescriptor<>("id", new AppendingReduce(), String.class);
kvId.initializeSerializerUnlessSet(new ExecutionConfig());
@@ -400,7 +473,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
ReducingState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> kvState = (KvState<Integer, VoidNamespace, ?, ?, B>) state;
+ KvState<VoidNamespace> kvState = (KvState<VoidNamespace>) state;
// some modifications to the state
backend.setCurrentKey(1);
@@ -416,13 +489,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("1", getSerializedValue(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
// draw a snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot1 = backend.snapshotPartitionedState(682375462378L, 2);
-
- for (String key: snapshot1.keySet()) {
- if (snapshot1.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot1.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot1.get(key)).materialize());
- }
- }
+ KeyGroupsStateHandle snapshot1 = runSnapshot(backend.snapshot(682375462378L, 2, streamFactory));
// make some more modifications
backend.setCurrentKey(1);
@@ -433,13 +500,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
state.add("u3");
// draw another snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot2 = backend.snapshotPartitionedState(682375462379L, 4);
-
- for (String key: snapshot2.keySet()) {
- if (snapshot2.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot2.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot2.get(key)).materialize());
- }
- }
+ KeyGroupsStateHandle snapshot2 = runSnapshot(backend.snapshot(682375462379L, 4, streamFactory));
// validate the original state
backend.setCurrentKey(1);
@@ -452,19 +513,14 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("u3", state.get());
assertEquals("u3", getSerializedValue(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.discardState();
-
+ backend.close();
// restore the first snapshot and validate it
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
-
- for (String key: snapshot1.keySet()) {
- snapshot1.get(key).discardState();
- }
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
+ snapshot1.discardState();
ReducingState<String> restored1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> restoredKvState1 = (KvState<Integer, VoidNamespace, ?, ?, B>) restored1;
+ KvState<VoidNamespace> restoredKvState1 = (KvState<VoidNamespace>) restored1;
backend.setCurrentKey(1);
assertEquals("1", restored1.get());
@@ -473,19 +529,14 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("2", restored1.get());
assertEquals("2", getSerializedValue(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.discardState();
-
+ backend.close();
// restore the second snapshot and validate it
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot2);
-
- for (String key: snapshot2.keySet()) {
- snapshot2.get(key).discardState();
- }
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot2);
+ snapshot2.discardState();
ReducingState<String> restored2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> restoredKvState2 = (KvState<Integer, VoidNamespace, ?, ?, B>) restored2;
+ KvState<VoidNamespace> restoredKvState2 = (KvState<VoidNamespace>) restored2;
backend.setCurrentKey(1);
assertEquals("1,u1", restored2.get());
@@ -496,6 +547,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
backend.setCurrentKey(3);
assertEquals("u3", restored2.get());
assertEquals("u3", getSerializedValue(restoredKvState2, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
+
+ backend.close();
}
catch (Exception e) {
e.printStackTrace();
@@ -507,7 +560,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
@SuppressWarnings("unchecked,rawtypes")
public void testFoldingState() {
try {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
FoldingStateDescriptor<Integer, String> kvId = new FoldingStateDescriptor<>("id",
"Fold-Initial:",
@@ -521,7 +575,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
FoldingState<Integer, String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> kvState = (KvState<Integer, VoidNamespace, ?, ?, B>) state;
+ KvState<VoidNamespace> kvState = (KvState<VoidNamespace>) state;
// some modifications to the state
backend.setCurrentKey(1);
@@ -537,13 +591,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("Fold-Initial:,1", getSerializedValue(kvState, 1, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
// draw a snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot1 = backend.snapshotPartitionedState(682375462378L, 2);
-
- for (String key: snapshot1.keySet()) {
- if (snapshot1.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot1.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot1.get(key)).materialize());
- }
- }
+ KeyGroupsStateHandle snapshot1 = runSnapshot(backend.snapshot(682375462378L, 2, streamFactory));
// make some more modifications
backend.setCurrentKey(1);
@@ -555,13 +603,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
state.add(103);
// draw another snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot2 = backend.snapshotPartitionedState(682375462379L, 4);
-
- for (String key: snapshot2.keySet()) {
- if (snapshot2.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot2.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot2.get(key)).materialize());
- }
- }
+ KeyGroupsStateHandle snapshot2 = runSnapshot(backend.snapshot(682375462379L, 4, streamFactory));
// validate the original state
backend.setCurrentKey(1);
@@ -574,19 +616,14 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("Fold-Initial:,103", state.get());
assertEquals("Fold-Initial:,103", getSerializedValue(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.discardState();
-
+ backend.close();
// restore the first snapshot and validate it
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
-
- for (String key: snapshot1.keySet()) {
- snapshot1.get(key).discardState();
- }
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
+ snapshot1.discardState();
FoldingState<Integer, String> restored1 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> restoredKvState1 = (KvState<Integer, VoidNamespace, ?, ?, B>) restored1;
+ KvState<VoidNamespace> restoredKvState1 = (KvState<VoidNamespace>) restored1;
backend.setCurrentKey(1);
assertEquals("Fold-Initial:,1", restored1.get());
@@ -595,20 +632,15 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("Fold-Initial:,2", restored1.get());
assertEquals("Fold-Initial:,2", getSerializedValue(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.discardState();
-
+ backend.close();
// restore the second snapshot and validate it
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot2);
-
- for (String key: snapshot2.keySet()) {
- snapshot2.get(key).discardState();
- }
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot2);
+ snapshot1.discardState();
@SuppressWarnings("unchecked")
FoldingState<Integer, String> restored2 = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@SuppressWarnings("unchecked")
- KvState<Integer, VoidNamespace, ?, ?, B> restoredKvState2 = (KvState<Integer, VoidNamespace, ?, ?, B>) restored2;
+ KvState<VoidNamespace> restoredKvState2 = (KvState<VoidNamespace>) restored2;
backend.setCurrentKey(1);
assertEquals("Fold-Initial:,101", restored2.get());
@@ -619,6 +651,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
backend.setCurrentKey(3);
assertEquals("Fold-Initial:,103", restored2.get());
assertEquals("Fold-Initial:,103", getSerializedValue(restoredKvState2, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
+
+ backend.close();
}
catch (Exception e) {
e.printStackTrace();
@@ -626,17 +660,115 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
}
}
+ /**
+ * This test verifies that state is correctly assigned to key groups and that restore
+ * restores the relevant key groups in the backend.
+ *
+ * <p>We have ten key groups. Initially, one backend is responsible for all ten key groups.
+ * Then we snapshot, split up the state and restore in to backends where each is responsible
+ * for five key groups. Then we make sure that the state is only available in the correct
+ * backend.
+ * @throws Exception
+ */
+ @Test
+ public void testKeyGroupSnapshotRestore() throws Exception {
+ final int MAX_PARALLELISM = 10;
+
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+
+ HashKeyGroupAssigner<Integer> keyGroupAssigner = new HashKeyGroupAssigner<>(10);
+
+ KeyedStateBackend<Integer> backend = createKeyedBackend(
+ IntSerializer.INSTANCE,
+ keyGroupAssigner,
+ new KeyGroupRange(0, MAX_PARALLELISM - 1),
+ new DummyEnvironment("test", 1, 0));
+
+ ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class, null);
+ kvId.initializeSerializerUnlessSet(new ExecutionConfig());
+
+ ValueState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
+
+ // keys that fall into the first half/second half of the key groups, respectively
+ int keyInFirstHalf = 17;
+ int keyInSecondHalf = 42;
+ Random rand = new Random(0);
+
+ // for each key, determine into which half of the key-group space they fall
+ int firstKeyHalf = keyGroupAssigner.getKeyGroupIndex(keyInFirstHalf) * 2 / MAX_PARALLELISM;
+ int secondKeyHalf = keyGroupAssigner.getKeyGroupIndex(keyInSecondHalf) * 2 / MAX_PARALLELISM;
+
+ while (firstKeyHalf == secondKeyHalf) {
+ keyInSecondHalf = rand.nextInt();
+ secondKeyHalf = keyGroupAssigner.getKeyGroupIndex(keyInSecondHalf) * 2 / MAX_PARALLELISM;
+ }
+
+ backend.setCurrentKey(keyInFirstHalf);
+ state.update("ShouldBeInFirstHalf");
+
+ backend.setCurrentKey(keyInSecondHalf);
+ state.update("ShouldBeInSecondHalf");
+
+
+ KeyGroupsStateHandle snapshot = runSnapshot(backend.snapshot(0, 0, streamFactory));
+
+ List<KeyGroupsStateHandle> firstHalfKeyGroupStates = CheckpointCoordinator.getKeyGroupsStateHandles(
+ Collections.singletonList(snapshot),
+ new KeyGroupRange(0, 4));
+
+ List<KeyGroupsStateHandle> secondHalfKeyGroupStates = CheckpointCoordinator.getKeyGroupsStateHandles(
+ Collections.singletonList(snapshot),
+ new KeyGroupRange(5, 9));
+
+ backend.close();
+
+ // backend for the first half of the key group range
+ KeyedStateBackend<Integer> firstHalfBackend = restoreKeyedBackend(
+ IntSerializer.INSTANCE,
+ keyGroupAssigner,
+ new KeyGroupRange(0, 4),
+ firstHalfKeyGroupStates,
+ new DummyEnvironment("test", 1, 0));
+
+ // backend for the second half of the key group range
+ KeyedStateBackend<Integer> secondHalfBackend = restoreKeyedBackend(
+ IntSerializer.INSTANCE,
+ keyGroupAssigner,
+ new KeyGroupRange(5, 9),
+ secondHalfKeyGroupStates,
+ new DummyEnvironment("test", 1, 0));
+
+
+ ValueState<String> firstHalfState = firstHalfBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
+
+ firstHalfBackend.setCurrentKey(keyInFirstHalf);
+ assertTrue(firstHalfState.value().equals("ShouldBeInFirstHalf"));
+
+ firstHalfBackend.setCurrentKey(keyInSecondHalf);
+ assertTrue(firstHalfState.value() == null);
+
+ ValueState<String> secondHalfState = secondHalfBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
+
+ secondHalfBackend.setCurrentKey(keyInFirstHalf);
+ assertTrue(secondHalfState.value() == null);
+
+ secondHalfBackend.setCurrentKey(keyInSecondHalf);
+ assertTrue(secondHalfState.value().equals("ShouldBeInSecondHalf"));
+
+ firstHalfBackend.close();
+ secondHalfBackend.close();
+ }
+
@Test
@SuppressWarnings("unchecked")
public void testValueStateRestoreWithWrongSerializers() {
try {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0),
- "test_op",
- IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class, null);
kvId.initializeSerializerUnlessSet(new ExecutionConfig());
-
+
ValueState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
backend.setCurrentKey(1);
@@ -645,23 +777,12 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
state.update("2");
// draw a snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot1 = backend.snapshotPartitionedState(682375462378L, 2);
-
- for (String key: snapshot1.keySet()) {
- if (snapshot1.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot1.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot1.get(key)).materialize());
- }
- }
-
- backend.discardState();
+ KeyGroupsStateHandle snapshot1 = runSnapshot(backend.snapshot(682375462378L, 2, streamFactory));
+ backend.close();
// restore the first snapshot and validate it
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
-
- for (String key: snapshot1.keySet()) {
- snapshot1.get(key).discardState();
- }
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
+ snapshot1.discardState();
@SuppressWarnings("unchecked")
TypeSerializer<String> fakeStringSerializer =
@@ -683,6 +804,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
} catch (Exception e) {
fail("wrong exception " + e);
}
+ backend.close();
}
catch (Exception e) {
e.printStackTrace();
@@ -694,7 +816,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
@SuppressWarnings("unchecked")
public void testListStateRestoreWithWrongSerializers() {
try {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
ListStateDescriptor<String> kvId = new ListStateDescriptor<>("id", String.class);
ListState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);
@@ -705,23 +828,12 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
state.add("2");
// draw a snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot1 = backend.snapshotPartitionedState(682375462378L, 2);
-
- for (String key: snapshot1.keySet()) {
- if (snapshot1.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot1.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot1.get(key)).materialize());
- }
- }
-
- backend.discardState();
+ KeyGroupsStateHandle snapshot1 = runSnapshot(backend.snapshot(682375462378L, 2, streamFactory));
+ backend.close();
// restore the first snapshot and validate it
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
-
- for (String key: snapshot1.keySet()) {
- snapshot1.get(key).discardState();
- }
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
+ snapshot1.discardState();
@SuppressWarnings("unchecked")
TypeSerializer<String> fakeStringSerializer =
@@ -743,6 +855,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
} catch (Exception e) {
fail("wrong exception " + e);
}
+ backend.close();
}
catch (Exception e) {
e.printStackTrace();
@@ -754,7 +867,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
@SuppressWarnings("unchecked")
public void testReducingStateRestoreWithWrongSerializers() {
try {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
ReducingStateDescriptor<String> kvId = new ReducingStateDescriptor<>("id",
new AppendingReduce(),
@@ -767,23 +881,12 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
state.add("2");
// draw a snapshot
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot1 = backend.snapshotPartitionedState(682375462378L, 2);
-
- for (String key: snapshot1.keySet()) {
- if (snapshot1.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot1.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot1.get(key)).materialize());
- }
- }
-
- backend.discardState();
+ KeyGroupsStateHandle snapshot1 = runSnapshot(backend.snapshot(682375462378L, 2, streamFactory));
+ backend.close();
// restore the first snapshot and validate it
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
-
- for (String key: snapshot1.keySet()) {
- snapshot1.get(key).discardState();
- }
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot1);
+ snapshot1.discardState();
@SuppressWarnings("unchecked")
TypeSerializer<String> fakeStringSerializer =
@@ -805,6 +908,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
} catch (Exception e) {
fail("wrong exception " + e);
}
+ backend.close();
}
catch (Exception e) {
e.printStackTrace();
@@ -814,7 +918,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
@Test
public void testCopyDefaultValue() throws Exception {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
ValueStateDescriptor<IntValue> kvId = new ValueStateDescriptor<>("id", IntValue.class, new IntValue(-1));
kvId.initializeSerializerUnlessSet(new ExecutionConfig());
@@ -831,6 +935,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertNotNull(default2);
assertEquals(default1, default2);
assertFalse(default1 == default2);
+
+ backend.close();
}
/**
@@ -840,7 +946,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
*/
@Test
public void testRequireNonNullNamespace() throws Exception {
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
ValueStateDescriptor<IntValue> kvId = new ValueStateDescriptor<>("id", IntValue.class, new IntValue(-1));
kvId.initializeSerializerUnlessSet(new ExecutionConfig());
@@ -862,6 +968,8 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
fail("Did not throw expected NullPointerException");
} catch (NullPointerException ignored) {
}
+
+ backend.close();
}
/**
@@ -869,7 +977,13 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
* flag and create concurrent variants for internal state structures.
*/
@SuppressWarnings("unchecked")
- protected static <B extends AbstractStateBackend> void testConcurrentMapIfQueryable(B backend) throws Exception {
+ protected void testConcurrentMapIfQueryable() throws Exception {
+ KeyedStateBackend<Integer> backend = createKeyedBackend(
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ new DummyEnvironment("test_op", 1, 0));
+
{
// ValueState
ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>(
@@ -884,20 +998,19 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
VoidNamespaceSerializer.INSTANCE,
desc);
- KvState<Integer, VoidNamespace, ?, ?, ?> kvState = (KvState<Integer, VoidNamespace, ?, ?, ?>) state;
+ KvState<VoidNamespace> kvState = (KvState<VoidNamespace>) state;
assertTrue(kvState instanceof AbstractHeapState);
- Map<VoidNamespace, Map<Integer, ?>> stateMap = ((AbstractHeapState) kvState).getStateMap();
- assertTrue(stateMap instanceof ConcurrentHashMap);
-
kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
- kvState.setCurrentKey(1);
+ backend.setCurrentKey(1);
state.update(121818273);
- Map<Integer, ?> namespaceMap = stateMap.get(VoidNamespace.INSTANCE);
+ int keyGroupIndex = new HashKeyGroupAssigner<>(1).getKeyGroupIndex(1);
+ StateTable stateTable = ((AbstractHeapState) kvState).getStateTable();
+ assertNotNull("State not set", stateTable.get(keyGroupIndex));
+ assertTrue(stateTable.get(keyGroupIndex) instanceof ConcurrentHashMap);
+ assertTrue(stateTable.get(keyGroupIndex).get(VoidNamespace.INSTANCE) instanceof ConcurrentHashMap);
- assertNotNull("Value not set", namespaceMap);
- assertTrue(namespaceMap instanceof ConcurrentHashMap);
}
{
@@ -911,20 +1024,18 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
VoidNamespaceSerializer.INSTANCE,
desc);
- KvState<Integer, VoidNamespace, ?, ?, ?> kvState = (KvState<Integer, VoidNamespace, ?, ?, ?>) state;
+ KvState<VoidNamespace> kvState = (KvState<VoidNamespace>) state;
assertTrue(kvState instanceof AbstractHeapState);
- Map<VoidNamespace, Map<Integer, ?>> stateMap = ((AbstractHeapState) kvState).getStateMap();
- assertTrue(stateMap instanceof ConcurrentHashMap);
-
kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
- kvState.setCurrentKey(1);
+ backend.setCurrentKey(1);
state.add(121818273);
- Map<Integer, ?> namespaceMap = stateMap.get(VoidNamespace.INSTANCE);
-
- assertNotNull("List not set", namespaceMap);
- assertTrue(namespaceMap instanceof ConcurrentHashMap);
+ int keyGroupIndex = new HashKeyGroupAssigner<>(1).getKeyGroupIndex(1);
+ StateTable stateTable = ((AbstractHeapState) kvState).getStateTable();
+ assertNotNull("State not set", stateTable.get(keyGroupIndex));
+ assertTrue(stateTable.get(keyGroupIndex) instanceof ConcurrentHashMap);
+ assertTrue(stateTable.get(keyGroupIndex).get(VoidNamespace.INSTANCE) instanceof ConcurrentHashMap);
}
{
@@ -944,20 +1055,18 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
VoidNamespaceSerializer.INSTANCE,
desc);
- KvState<Integer, VoidNamespace, ?, ?, ?> kvState = (KvState<Integer, VoidNamespace, ?, ?, ?>) state;
+ KvState<VoidNamespace> kvState = (KvState<VoidNamespace>) state;
assertTrue(kvState instanceof AbstractHeapState);
- Map<VoidNamespace, Map<Integer, ?>> stateMap = ((AbstractHeapState) kvState).getStateMap();
- assertTrue(stateMap instanceof ConcurrentHashMap);
-
kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
- kvState.setCurrentKey(1);
+ backend.setCurrentKey(1);
state.add(121818273);
- Map<Integer, ?> namespaceMap = stateMap.get(VoidNamespace.INSTANCE);
-
- assertNotNull("List not set", namespaceMap);
- assertTrue(namespaceMap instanceof ConcurrentHashMap);
+ int keyGroupIndex = new HashKeyGroupAssigner<>(1).getKeyGroupIndex(1);
+ StateTable stateTable = ((AbstractHeapState) kvState).getStateTable();
+ assertNotNull("State not set", stateTable.get(keyGroupIndex));
+ assertTrue(stateTable.get(keyGroupIndex) instanceof ConcurrentHashMap);
+ assertTrue(stateTable.get(keyGroupIndex).get(VoidNamespace.INSTANCE) instanceof ConcurrentHashMap);
}
{
@@ -977,21 +1086,21 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
VoidNamespaceSerializer.INSTANCE,
desc);
- KvState<Integer, VoidNamespace, ?, ?, ?> kvState = (KvState<Integer, VoidNamespace, ?, ?, ?>) state;
+ KvState<VoidNamespace> kvState = (KvState<VoidNamespace>) state;
assertTrue(kvState instanceof AbstractHeapState);
- Map<VoidNamespace, Map<Integer, ?>> stateMap = ((AbstractHeapState) kvState).getStateMap();
- assertTrue(stateMap instanceof ConcurrentHashMap);
-
kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
- kvState.setCurrentKey(1);
+ backend.setCurrentKey(1);
state.add(121818273);
- Map<Integer, ?> namespaceMap = stateMap.get(VoidNamespace.INSTANCE);
-
- assertNotNull("List not set", namespaceMap);
- assertTrue(namespaceMap instanceof ConcurrentHashMap);
+ int keyGroupIndex = new HashKeyGroupAssigner<>(1).getKeyGroupIndex(1);
+ StateTable stateTable = ((AbstractHeapState) kvState).getStateTable();
+ assertNotNull("State not set", stateTable.get(keyGroupIndex));
+ assertTrue(stateTable.get(keyGroupIndex) instanceof ConcurrentHashMap);
+ assertTrue(stateTable.get(keyGroupIndex).get(VoidNamespace.INSTANCE) instanceof ConcurrentHashMap);
}
+
+ backend.close();
}
/**
@@ -1002,11 +1111,12 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
DummyEnvironment env = new DummyEnvironment("test", 1, 0);
KvStateRegistry registry = env.getKvStateRegistry();
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE, env);
+
KvStateRegistryListener listener = mock(KvStateRegistryListener.class);
registry.registerListener(listener);
- backend.initializeForJob(env, "test_op", IntSerializer.INSTANCE);
-
ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>(
"test",
IntSerializer.INSTANCE,
@@ -1020,25 +1130,16 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
eq(env.getJobID()), eq(env.getJobVertexId()), eq(0), eq("banana"), any(KvStateID.class));
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot = backend
- .snapshotPartitionedState(682375462379L, 4);
-
- for (String key: snapshot.keySet()) {
- if (snapshot.get(key) instanceof AsynchronousKvStateSnapshot) {
- snapshot.put(key, ((AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) snapshot.get(key)).materialize());
- }
- }
+ KeyGroupsStateHandle snapshot = runSnapshot(backend.snapshot(682375462379L, 4, streamFactory));
- // Verify unregistered
- backend.discardState();
+ backend.close();
verify(listener, times(1)).notifyKvStateUnregistered(
eq(env.getJobID()), eq(env.getJobVertexId()), eq(0), eq("banana"));
-
+ backend.close();
// Initialize again
- backend.initializeForJob(env, "test_op", IntSerializer.INSTANCE);
-
- backend.injectKeyValueStateSnapshots((HashMap) snapshot);
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot, env);
+ snapshot.discardState();
backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
@@ -1046,6 +1147,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
verify(listener, times(2)).notifyKvStateRegistered(
eq(env.getJobID()), eq(env.getJobVertexId()), eq(0), eq("banana"), any(KvStateID.class));
+ backend.close();
}
@@ -1093,7 +1195,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
* if it is not null.
*/
private static <V, K, N> V getSerializedValue(
- KvState<K, N, ?, ?, ?> kvState,
+ KvState<N> kvState,
K key,
TypeSerializer<K> keySerializer,
N namespace,
@@ -1117,7 +1219,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
* if it is not null.
*/
private static <V, K, N> List<V> getSerializedList(
- KvState<K, N, ?, ?, ?> kvState,
+ KvState<N> kvState,
K key,
TypeSerializer<K> keySerializer,
N namespace,
@@ -1135,4 +1237,12 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
return KvStateRequestSerializer.deserializeList(serializedValue, valueSerializer);
}
}
+
+ private KeyGroupsStateHandle runSnapshot(RunnableFuture<KeyGroupsStateHandle> snapshotRunnableFuture) throws Exception {
+ if(!snapshotRunnableFuture.isDone()) {
+ Thread runner = new Thread(snapshotRunnableFuture);
+ runner.start();
+ }
+ return snapshotRunnableFuture.get();
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStateOutputStreamTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStateOutputStreamTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStateOutputStreamTest.java
index 1d45115..a6a555d 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStateOutputStreamTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStateOutputStreamTest.java
@@ -34,29 +34,26 @@ import java.util.Random;
import static org.junit.Assert.*;
public class FsCheckpointStateOutputStreamTest {
-
+
/** The temp dir, obtained in a platform neutral way */
private static final Path TEMP_DIR_PATH = new Path(new File(System.getProperty("java.io.tmpdir")).toURI());
-
-
+
+
@Test(expected = IllegalArgumentException.class)
public void testWrongParameters() {
// this should fail
- new FsStateBackend.FsCheckpointStateOutputStream(
+ new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), 4000, 5000);
}
@Test
public void testEmptyState() throws Exception {
- AbstractStateBackend.CheckpointStateOutputStream stream = new FsStateBackend.FsCheckpointStateOutputStream(
- TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), 1024, 512);
+ FsCheckpointStreamFactory.CheckpointStateOutputStream stream =
+ new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), 1024, 512);
StreamStateHandle handle = stream.closeAndGetHandle();
- assertTrue(handle instanceof ByteStreamStateHandle);
-
- InputStream inStream = handle.openInputStream();
- assertEquals(-1, inStream.read());
+ assertTrue(handle == null);
}
@Test
@@ -73,17 +70,17 @@ public class FsCheckpointStateOutputStreamTest {
public void testStateAboveMemThreshold() throws Exception {
runTest(576446, 259, 17, true);
}
-
+
@Test
public void testZeroThreshold() throws Exception {
runTest(16678, 4096, 0, true);
}
-
+
private void runTest(int numBytes, int bufferSize, int threshold, boolean expectFile) throws Exception {
- AbstractStateBackend.CheckpointStateOutputStream stream =
- new FsStateBackend.FsCheckpointStateOutputStream(
+ FsCheckpointStreamFactory.CheckpointStateOutputStream stream =
+ new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), bufferSize, threshold);
-
+
Random rnd = new Random();
byte[] original = new byte[numBytes];
byte[] bytes = new byte[original.length];
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
index 429fc6b..ab4ca3b 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
@@ -159,7 +159,7 @@ public class TaskAsyncCallTest {
TaskDeploymentDescriptor tdd = new TaskDeploymentDescriptor(
new JobID(), "Job Name", new JobVertexID(), new ExecutionAttemptID(),
new SerializedValue<>(new ExecutionConfig()),
- "Test Task", 0, 1, 0,
+ "Test Task", 1, 0, 1, 0,
new Configuration(), new Configuration(),
CheckpointsInOrderInvokable.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java
index ce88c09..54cd7c6 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskManagerTest.java
@@ -162,7 +162,7 @@ public class TaskManagerTest extends TestLogger {
final SerializedValue<ExecutionConfig> executionConfig = new SerializedValue<>(new ExecutionConfig());
final TaskDeploymentDescriptor tdd = new TaskDeploymentDescriptor(jid, "TestJob", vid, eid, executionConfig,
- "TestTask", 2, 7, 0, new Configuration(), new Configuration(),
+ "TestTask", 7, 2, 7, 0, new Configuration(), new Configuration(),
TestInvokableCorrect.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.<InputGateDeploymentDescriptor>emptyList(),
@@ -265,7 +265,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd1 = new TaskDeploymentDescriptor(
jid1, "TestJob1", vid1, eid1,
new SerializedValue<>(new ExecutionConfig()),
- "TestTask1", 1, 5, 0,
+ "TestTask1", 5, 1, 5, 0,
new Configuration(), new Configuration(), TestInvokableBlockingCancelable.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.<InputGateDeploymentDescriptor>emptyList(),
@@ -274,7 +274,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd2 = new TaskDeploymentDescriptor(
jid2, "TestJob2", vid2, eid2,
new SerializedValue<>(new ExecutionConfig()),
- "TestTask2", 2, 7, 0,
+ "TestTask2", 7, 2, 7, 0,
new Configuration(), new Configuration(), TestInvokableBlockingCancelable.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.<InputGateDeploymentDescriptor>emptyList(),
@@ -403,13 +403,13 @@ public class TaskManagerTest extends TestLogger {
final SerializedValue<ExecutionConfig> executionConfig = new SerializedValue<>(new ExecutionConfig());
final TaskDeploymentDescriptor tdd1 = new TaskDeploymentDescriptor(jid1, "TestJob", vid1, eid1, executionConfig,
- "TestTask1", 1, 5, 0, new Configuration(), new Configuration(), StoppableInvokable.class.getName(),
+ "TestTask1", 5, 1, 5, 0, new Configuration(), new Configuration(), StoppableInvokable.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.<InputGateDeploymentDescriptor>emptyList(),
new ArrayList<BlobKey>(), Collections.<URL>emptyList(), 0);
final TaskDeploymentDescriptor tdd2 = new TaskDeploymentDescriptor(jid2, "TestJob", vid2, eid2, executionConfig,
- "TestTask2", 2, 7, 0, new Configuration(), new Configuration(), TestInvokableBlockingCancelable.class.getName(),
+ "TestTask2", 7, 2, 7, 0, new Configuration(), new Configuration(), TestInvokableBlockingCancelable.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.<InputGateDeploymentDescriptor>emptyList(),
new ArrayList<BlobKey>(), Collections.<URL>emptyList(), 0);
@@ -531,7 +531,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd1 = new TaskDeploymentDescriptor(
jid, "TestJob", vid1, eid1,
new SerializedValue<>(new ExecutionConfig()),
- "Sender", 0, 1, 0,
+ "Sender", 1, 0, 1, 0,
new Configuration(), new Configuration(), Tasks.Sender.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.<InputGateDeploymentDescriptor>emptyList(),
@@ -540,7 +540,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd2 = new TaskDeploymentDescriptor(
jid, "TestJob", vid2, eid2,
new SerializedValue<>(new ExecutionConfig()),
- "Receiver", 2, 7, 0,
+ "Receiver", 7, 2, 7, 0,
new Configuration(), new Configuration(), Tasks.Receiver.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.<InputGateDeploymentDescriptor>emptyList(),
@@ -636,7 +636,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd1 = new TaskDeploymentDescriptor(
jid, "TestJob", vid1, eid1,
new SerializedValue<>(new ExecutionConfig()),
- "Sender", 0, 1, 0,
+ "Sender", 1, 0, 1, 0,
new Configuration(), new Configuration(), Tasks.Sender.class.getName(),
irpdd, Collections.<InputGateDeploymentDescriptor>emptyList(), new ArrayList<BlobKey>(),
Collections.<URL>emptyList(), 0);
@@ -644,7 +644,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd2 = new TaskDeploymentDescriptor(
jid, "TestJob", vid2, eid2,
new SerializedValue<>(new ExecutionConfig()),
- "Receiver", 2, 7, 0,
+ "Receiver", 7, 2, 7, 0,
new Configuration(), new Configuration(), Tasks.Receiver.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.singletonList(ircdd),
@@ -781,7 +781,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd1 = new TaskDeploymentDescriptor(
jid, "TestJob", vid1, eid1,
new SerializedValue<>(new ExecutionConfig()),
- "Sender", 0, 1, 0,
+ "Sender", 1, 0, 1, 0,
new Configuration(), new Configuration(), Tasks.Sender.class.getName(),
irpdd, Collections.<InputGateDeploymentDescriptor>emptyList(),
new ArrayList<BlobKey>(), Collections.<URL>emptyList(), 0);
@@ -789,7 +789,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd2 = new TaskDeploymentDescriptor(
jid, "TestJob", vid2, eid2,
new SerializedValue<>(new ExecutionConfig()),
- "Receiver", 2, 7, 0,
+ "Receiver", 7, 2, 7, 0,
new Configuration(), new Configuration(), Tasks.BlockingReceiver.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
Collections.singletonList(ircdd),
@@ -929,7 +929,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd = new TaskDeploymentDescriptor(
jid, "TestJob", vid, eid,
new SerializedValue<>(new ExecutionConfig()),
- "Receiver", 0, 1, 0,
+ "Receiver", 1, 0, 1, 0,
new Configuration(), new Configuration(),
Tasks.AgnosticReceiver.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
@@ -1025,7 +1025,7 @@ public class TaskManagerTest extends TestLogger {
final TaskDeploymentDescriptor tdd = new TaskDeploymentDescriptor(
jid, "TestJob", vid, eid,
new SerializedValue<>(new ExecutionConfig()),
- "Receiver", 0, 1, 0,
+ "Receiver", 1, 0, 1, 0,
new Configuration(), new Configuration(),
Tasks.AgnosticReceiver.class.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
@@ -1104,6 +1104,7 @@ public class TaskManagerTest extends TestLogger {
new ExecutionAttemptID(),
new SerializedValue<>(new ExecutionConfig()),
"Task",
+ 1,
0,
1,
0,
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java
index 2f8e3db..f145b48 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskTest.java
@@ -639,7 +639,7 @@ public class TaskTest {
return new TaskDeploymentDescriptor(
new JobID(), "Test Job", new JobVertexID(), new ExecutionAttemptID(),
execConfig,
- "Test Task", 0, 1, 0,
+ "Test Task", 1, 0, 1, 0,
new Configuration(), new Configuration(),
invokable.getName(),
Collections.<ResultPartitionDeploymentDescriptor>emptyList(),
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java
index 6d67560..0c0b81a 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java
@@ -125,7 +125,7 @@ public class ContinuousFileReaderOperator<OUT, S extends Serializable> extends A
}
@Override
- public void dispose() {
+ public void dispose() throws Exception {
super.dispose();
// first try to cancel it properly and
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
index 81c5c48..f9f26e9 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
@@ -152,7 +152,6 @@ public class StreamGraphGenerator {
if (maxParallelism <= 0) {
maxParallelism = transform.getParallelism();
-
/**
* TODO: Remove once the parallelism settings works properly in Flink (FLINK-3885)
* Currently, the parallelism will be set to 1 on the JobManager iff it encounters
[17/27] flink git commit: [FLINK-4381] Refactor State to Prepare For
Key-Group State Backends
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStore.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStore.java b/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStore.java
index dea3452..d62b13e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStore.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStore.java
@@ -22,7 +22,7 @@ import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.api.BackgroundCallback;
import org.apache.curator.utils.ZKPaths;
import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
import org.apache.flink.util.InstantiationUtil;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
@@ -40,9 +40,10 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
/**
* State handles backed by ZooKeeper.
*
- * <p>Added state is persisted via {@link StateHandle}s, which in turn are written to
- * ZooKeeper. This level of indirection is necessary to keep the amount of data in ZooKeeper
- * small. ZooKeeper is build for data in the KB range whereas state can grow to multiple MBs.
+ * <p>Added state is persisted via {@link RetrievableStateHandle RetrievableStateHandles},
+ * which in turn are written to ZooKeeper. This level of indirection is necessary to keep the
+ * amount of data in ZooKeeper small. ZooKeeper is build for data in the KB range whereas
+ * state can grow to multiple MBs.
*
* <p>State modifications require some care, because it is possible that certain failures bring
* the state handle backend and ZooKeeper out of sync.
@@ -72,7 +73,7 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
/** Curator ZooKeeper client */
private final CuratorFramework client;
- private final StateStorageHelper<T> storage;
+ private final RetrievableStateStorageHelper<T> storage;
/**
* Creates a {@link ZooKeeperStateHandleStore}.
@@ -84,7 +85,7 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
*/
public ZooKeeperStateHandleStore(
CuratorFramework client,
- StateStorageHelper storage) throws IOException {
+ RetrievableStateStorageHelper<T> storage) throws IOException {
this.client = checkNotNull(client, "Curator client");
this.storage = checkNotNull(storage, "State storage");
@@ -94,9 +95,9 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
* Creates a state handle and stores it in ZooKeeper with create mode {@link
* CreateMode#PERSISTENT}.
*
- * @see #add(String, Serializable, CreateMode)
+ * @see #add(String, T, CreateMode)
*/
- public StateHandle<T> add(String pathInZooKeeper, T state) throws Exception {
+ public RetrievableStateHandle<T> add(String pathInZooKeeper, T state) throws Exception {
return add(pathInZooKeeper, state, CreateMode.PERSISTENT);
}
@@ -111,39 +112,39 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
* start with a '/')
* @param state State to be added
* @param createMode The create mode for the new path in ZooKeeper
- * @return Created {@link StateHandle}
+ *
+ * @return The Created {@link RetrievableStateHandle}.
* @throws Exception If a ZooKeeper or state handle operation fails
*/
- public StateHandle<T> add(
+ public RetrievableStateHandle<T> add(
String pathInZooKeeper,
T state,
CreateMode createMode) throws Exception {
checkNotNull(pathInZooKeeper, "Path in ZooKeeper");
checkNotNull(state, "State");
- StateHandle<T> stateHandle = storage.store(state);
+ RetrievableStateHandle<T> storeHandle = storage.store(state);
boolean success = false;
try {
// Serialize the state handle. This writes the state to the backend.
- byte[] serializedStateHandle = InstantiationUtil.serializeObject(stateHandle);
+ byte[] serializedStoreHandle = InstantiationUtil.serializeObject(storeHandle);
// Write state handle (not the actual state) to ZooKeeper. This is expected to be
// smaller than the state itself. This level of indirection makes sure that data in
// ZooKeeper is small, because ZooKeeper is designed for data in the KB range, but
// the state can be larger.
- client.create().withMode(createMode).forPath(pathInZooKeeper, serializedStateHandle);
+ client.create().withMode(createMode).forPath(pathInZooKeeper, serializedStoreHandle);
success = true;
-
- return stateHandle;
+ return storeHandle;
}
finally {
if (!success) {
// Cleanup the state handle if it was not written to ZooKeeper.
- if (stateHandle != null) {
- stateHandle.discardState();
+ if (storeHandle != null) {
+ storeHandle.discardState();
}
}
}
@@ -161,31 +162,29 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
checkNotNull(pathInZooKeeper, "Path in ZooKeeper");
checkNotNull(state, "State");
- StateHandle<T> oldStateHandle = get(pathInZooKeeper);
+ RetrievableStateHandle<T> oldStateHandle = get(pathInZooKeeper);
- StateHandle<T> stateHandle = storage.store(state);
+ RetrievableStateHandle<T> newStateHandle = storage.store(state);
boolean success = false;
try {
// Serialize the new state handle. This writes the state to the backend.
- byte[] serializedStateHandle = InstantiationUtil.serializeObject(stateHandle);
+ byte[] serializedStateHandle = InstantiationUtil.serializeObject(newStateHandle);
// Replace state handle in ZooKeeper.
client.setData()
.withVersion(expectedVersion)
.forPath(pathInZooKeeper, serializedStateHandle);
-
success = true;
- }
- finally {
- if (success) {
+ } finally {
+ if(success) {
oldStateHandle.discardState();
- }
- else {
- stateHandle.discardState();
+ } else {
+ newStateHandle.discardState();
}
}
+
}
/**
@@ -216,13 +215,11 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
* @throws Exception If a ZooKeeper or state handle operation fails
*/
@SuppressWarnings("unchecked")
- public StateHandle<T> get(String pathInZooKeeper) throws Exception {
+ public RetrievableStateHandle<T> get(String pathInZooKeeper) throws Exception {
checkNotNull(pathInZooKeeper, "Path in ZooKeeper");
byte[] data = client.getData().forPath(pathInZooKeeper);
-
- return (StateHandle<T>) InstantiationUtil
- .deserializeObject(data, ClassLoader.getSystemClassLoader());
+ return InstantiationUtil.deserializeObject(data);
}
/**
@@ -234,8 +231,8 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
* @throws Exception If a ZooKeeper or state handle operation fails
*/
@SuppressWarnings("unchecked")
- public List<Tuple2<StateHandle<T>, String>> getAll() throws Exception {
- final List<Tuple2<StateHandle<T>, String>> stateHandles = new ArrayList<>();
+ public List<Tuple2<RetrievableStateHandle<T>, String>> getAll() throws Exception {
+ final List<Tuple2<RetrievableStateHandle<T>, String>> stateHandles = new ArrayList<>();
boolean success = false;
@@ -254,7 +251,7 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
path = "/" + path;
try {
- final StateHandle<T> stateHandle = get(path);
+ final RetrievableStateHandle<T> stateHandle = get(path);
stateHandles.add(new Tuple2<>(stateHandle, path));
} catch (KeeperException.NoNodeException ignored) {
// Concurrent deletion, retry
@@ -272,6 +269,7 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
return stateHandles;
}
+
/**
* Gets all available state handles from ZooKeeper sorted by name (ascending).
*
@@ -281,8 +279,8 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
* @throws Exception If a ZooKeeper or state handle operation fails
*/
@SuppressWarnings("unchecked")
- public List<Tuple2<StateHandle<T>, String>> getAllSortedByName() throws Exception {
- final List<Tuple2<StateHandle<T>, String>> stateHandles = new ArrayList<>();
+ public List<Tuple2<RetrievableStateHandle<T>, String>> getAllSortedByName() throws Exception {
+ final List<Tuple2<RetrievableStateHandle<T>, String>> stateHandles = new ArrayList<>();
boolean success = false;
@@ -303,7 +301,7 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
path = "/" + path;
try {
- final StateHandle<T> stateHandle = get(path);
+ final RetrievableStateHandle<T> stateHandle = get(path);
stateHandles.add(new Tuple2<>(stateHandle, path));
} catch (KeeperException.NoNodeException ignored) {
// Concurrent deletion, retry
@@ -364,7 +362,7 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
public void removeAndDiscardState(String pathInZooKeeper) throws Exception {
checkNotNull(pathInZooKeeper, "Path in ZooKeeper");
- StateHandle<T> stateHandle = get(pathInZooKeeper);
+ RetrievableStateHandle<T> stateHandle = get(pathInZooKeeper);
// Delete the state handle from ZooKeeper first
client.delete().deletingChildrenIfNeeded().forPath(pathInZooKeeper);
@@ -381,7 +379,7 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
* @throws Exception If a ZooKeeper or state handle operation fails
*/
public void removeAndDiscardAllState() throws Exception {
- final List<Tuple2<StateHandle<T>, String>> allStateHandles = getAll();
+ final List<Tuple2<RetrievableStateHandle<T>, String>> allStateHandles = getAll();
ZKPaths.deleteChildren(
client.getZookeeperClient().getZooKeeper(),
@@ -389,7 +387,7 @@ public class ZooKeeperStateHandleStore<T extends Serializable> {
false);
// Discard the state handles only after they have been successfully deleted from ZooKeeper.
- for (Tuple2<StateHandle<T>, String> stateHandleAndPath : allStateHandles) {
+ for (Tuple2<RetrievableStateHandle<T>, String> stateHandleAndPath : allStateHandles) {
stateHandleAndPath.f0.discardState();
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/filesystem/FileSystemStateStorageHelper.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/filesystem/FileSystemStateStorageHelper.java b/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/filesystem/FileSystemStateStorageHelper.java
index 6692ef0..a534b40 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/filesystem/FileSystemStateStorageHelper.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/filesystem/FileSystemStateStorageHelper.java
@@ -21,22 +21,22 @@ package org.apache.flink.runtime.zookeeper.filesystem;
import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.core.fs.Path;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.state.filesystem.FileSerializableStateHandle;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
+import org.apache.flink.runtime.state.RetrievableStreamStateHandle;
+import org.apache.flink.runtime.zookeeper.RetrievableStateStorageHelper;
import org.apache.flink.util.FileUtils;
import org.apache.flink.util.Preconditions;
-import org.apache.flink.runtime.zookeeper.StateStorageHelper;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.io.Serializable;
/**
- * {@link StateStorageHelper} implementation which stores the state in the given filesystem path.
+ * {@link RetrievableStateStorageHelper} implementation which stores the state in the given filesystem path.
*
- * @param <T>
+ * @param <T> The type of the data that can be stored by this storage helper.
*/
-public class FileSystemStateStorageHelper<T extends Serializable> implements StateStorageHelper<T> {
+public class FileSystemStateStorageHelper<T extends Serializable> implements RetrievableStateStorageHelper<T> {
private final Path rootPath;
@@ -56,7 +56,7 @@ public class FileSystemStateStorageHelper<T extends Serializable> implements Sta
}
@Override
- public StateHandle<T> store(T state) throws Exception {
+ public RetrievableStateHandle<T> store(T state) throws Exception {
Exception latestException = null;
for (int attempt = 0; attempt < 10; attempt++) {
@@ -73,8 +73,7 @@ public class FileSystemStateStorageHelper<T extends Serializable> implements Sta
try(ObjectOutputStream os = new ObjectOutputStream(outStream)) {
os.writeObject(state);
}
-
- return new FileSerializableStateHandle<>(filePath);
+ return new RetrievableStreamStateHandle<T>(filePath);
}
throw new Exception("Could not open output stream for state backend", latestException);
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
index 356f1a9..407fa01 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
@@ -738,29 +738,18 @@ class JobManager(
sender() ! TriggerSavepointFailure(jobId, new IllegalArgumentException("Unknown job."))
}
- case DisposeSavepoint(savepointPath, blobKeys) =>
+ case DisposeSavepoint(savepointPath) =>
val senderRef = sender()
future {
try {
log.info(s"Disposing savepoint at '$savepointPath'.")
- if (blobKeys.isDefined) {
- // We don't need a real ID here for the library cache manager
- val jid = new JobID()
+ val savepoint = savepointStore.loadSavepoint(savepointPath)
- try {
- libraryCacheManager.registerJob(jid, blobKeys.get, java.util.Collections.emptyList())
- val classLoader = libraryCacheManager.getClassLoader(jid)
+ log.debug(s"$savepoint")
- // Discard with user code loader
- savepointStore.disposeSavepoint(savepointPath, classLoader)
- } finally {
- libraryCacheManager.unregisterJob(jid)
- }
- } else {
- // Discard with system class loader
- savepointStore.disposeSavepoint(savepointPath, getClass.getClassLoader)
- }
+ // Dispose the savepoint
+ savepointStore.disposeSavepoint(savepointPath)
senderRef ! DisposeSavepointSuccess
} catch {
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala
index 40c4dcf..5e2b547 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/messages/JobManagerMessages.scala
@@ -490,13 +490,9 @@ object JobManagerMessages {
* Disposes a savepoint.
*
* @param savepointPath The path of the savepoint to dispose.
- * @param blobKeys BLOB keys if a user program JAR was uploaded for disposal.
- * This is required when we dispose state which contains
- * custom state instances (e.g. reducing state, rocksDB state).
*/
case class DisposeSavepoint(
- savepointPath: String,
- blobKeys: Option[java.util.List[BlobKey]] = None)
+ savepointPath: String)
extends RequiresLeaderSessionID
/** Response after a successful savepoint dispose. */
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java
index 1816fc9..5416292 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java
@@ -28,14 +28,19 @@ import org.apache.flink.runtime.executiongraph.ExecutionJobVertex;
import org.apache.flink.runtime.executiongraph.ExecutionVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.messages.checkpoint.AcknowledgeCheckpoint;
-import org.apache.flink.runtime.state.LocalStateHandle;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
import org.apache.flink.runtime.util.SerializableObject;
-import org.apache.flink.util.SerializedValue;
import org.junit.Test;
import org.mockito.Mockito;
+import java.util.ArrayList;
+import java.util.Collections;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import static org.junit.Assert.assertEquals;
@@ -55,8 +60,11 @@ public class CheckpointStateRestoreTest {
@Test
public void testSetState() {
try {
- final SerializedValue<StateHandle<?>> serializedState = new SerializedValue<StateHandle<?>>(
- new LocalStateHandle<SerializableObject>(new SerializableObject()));
+
+ final ChainedStateHandle<StreamStateHandle> serializedState = CheckpointCoordinatorTest.generateChainedStateHandle(new SerializableObject());
+ KeyGroupRange keyGroupRange = KeyGroupRange.of(0,0);
+ List<SerializableObject> testStates = Collections.singletonList(new SerializableObject());
+ final List<KeyGroupsStateHandle> serializedKeyGroupStates = CheckpointCoordinatorTest.generateKeyGroupState(keyGroupRange, testStates);
final JobID jid = new JobID();
final JobVertexID statefulId = new JobVertexID();
@@ -106,9 +114,9 @@ public class CheckpointStateRestoreTest {
PendingCheckpoint pending = coord.getPendingCheckpoints().values().iterator().next();
final long checkpointId = pending.getCheckpointId();
- coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec1.getAttemptId(), checkpointId, serializedState, 0));
- coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec2.getAttemptId(), checkpointId, serializedState, 0));
- coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec3.getAttemptId(), checkpointId, serializedState, 0));
+ coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec1.getAttemptId(), checkpointId, serializedState, serializedKeyGroupStates));
+ coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec2.getAttemptId(), checkpointId, serializedState, serializedKeyGroupStates));
+ coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec3.getAttemptId(), checkpointId, serializedState, serializedKeyGroupStates));
coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statelessExec1.getAttemptId(), checkpointId));
coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statelessExec2.getAttemptId(), checkpointId));
@@ -119,11 +127,11 @@ public class CheckpointStateRestoreTest {
coord.restoreLatestCheckpointedState(map, true, false);
// verify that each stateful vertex got the state
- verify(statefulExec1, times(1)).setInitialState(Mockito.eq(serializedState), Mockito.<Map<Integer, SerializedValue<StateHandle<?>>>>any());
- verify(statefulExec2, times(1)).setInitialState(Mockito.eq(serializedState), Mockito.<Map<Integer, SerializedValue<StateHandle<?>>>>any());
- verify(statefulExec3, times(1)).setInitialState(Mockito.eq(serializedState), Mockito.<Map<Integer, SerializedValue<StateHandle<?>>>>any());
- verify(statelessExec1, times(0)).setInitialState(Mockito.<SerializedValue<StateHandle<?>>>any(), Mockito.<Map<Integer, SerializedValue<StateHandle<?>>>>any());
- verify(statelessExec2, times(0)).setInitialState(Mockito.<SerializedValue<StateHandle<?>>>any(), Mockito.<Map<Integer, SerializedValue<StateHandle<?>>>>any());
+ verify(statefulExec1, times(1)).setInitialState(Mockito.eq(serializedState), Mockito.<List<KeyGroupsStateHandle>>any());
+ verify(statefulExec2, times(1)).setInitialState(Mockito.eq(serializedState), Mockito.<List<KeyGroupsStateHandle>>any());
+ verify(statefulExec3, times(1)).setInitialState(Mockito.eq(serializedState), Mockito.<List<KeyGroupsStateHandle>>any());
+ verify(statelessExec1, times(0)).setInitialState(Mockito.<ChainedStateHandle<StreamStateHandle>>any(), Mockito.<List<KeyGroupsStateHandle>>any());
+ verify(statelessExec2, times(0)).setInitialState(Mockito.<ChainedStateHandle<StreamStateHandle>>any(), Mockito.<List<KeyGroupsStateHandle>>any());
}
catch (Exception e) {
e.printStackTrace();
@@ -134,8 +142,10 @@ public class CheckpointStateRestoreTest {
@Test
public void testStateOnlyPartiallyAvailable() {
try {
- final SerializedValue<StateHandle<?>> serializedState = new SerializedValue<StateHandle<?>>(
- new LocalStateHandle<SerializableObject>(new SerializableObject()));
+ final ChainedStateHandle<StreamStateHandle> serializedState = CheckpointCoordinatorTest.generateChainedStateHandle(new SerializableObject());
+ KeyGroupRange keyGroupRange = KeyGroupRange.of(0,0);
+ List<SerializableObject> testStates = Collections.singletonList(new SerializableObject());
+ final List<KeyGroupsStateHandle> serializedKeyGroupStates = CheckpointCoordinatorTest.generateKeyGroupState(keyGroupRange, testStates);
final JobID jid = new JobID();
final JobVertexID statefulId = new JobVertexID();
@@ -186,9 +196,9 @@ public class CheckpointStateRestoreTest {
final long checkpointId = pending.getCheckpointId();
// the difference to the test "testSetState" is that one stateful subtask does not report state
- coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec1.getAttemptId(), checkpointId, serializedState, 0));
+ coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec1.getAttemptId(), checkpointId, serializedState, serializedKeyGroupStates));
coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec2.getAttemptId(), checkpointId));
- coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec3.getAttemptId(), checkpointId, serializedState, 0));
+ coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statefulExec3.getAttemptId(), checkpointId, serializedState, serializedKeyGroupStates));
coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statelessExec1.getAttemptId(), checkpointId));
coord.receiveAcknowledgeMessage(new AcknowledgeCheckpoint(jid, statelessExec2.getAttemptId(), checkpointId));
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointStoreTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointStoreTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointStoreTest.java
index 634e177..6182ffd 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointStoreTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointStoreTest.java
@@ -21,8 +21,9 @@ package org.apache.flink.runtime.checkpoint;
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.messages.CheckpointMessagesTest;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
import org.apache.flink.util.TestLogger;
import org.junit.Test;
@@ -107,8 +108,6 @@ public abstract class CompletedCheckpointStoreTest extends TestLogger {
// The ZooKeeper implementation discards asynchronously
expected[i - 1].awaitDiscard();
assertTrue(expected[i - 1].isDiscarded());
- assertEquals(userClassLoader, expected[i - 1].getDiscardClassLoader());
-
assertEquals(1, checkpoints.getNumberOfRetainedCheckpoints());
}
}
@@ -183,7 +182,6 @@ public abstract class CompletedCheckpointStoreTest extends TestLogger {
// The ZooKeeper implementation discards asynchronously
checkpoint.awaitDiscard();
assertTrue(checkpoint.isDiscarded());
- assertEquals(userClassLoader, checkpoint.getDiscardClassLoader());
}
}
@@ -199,14 +197,14 @@ public abstract class CompletedCheckpointStoreTest extends TestLogger {
JobVertexID jvid = new JobVertexID();
Map<JobVertexID, TaskState> taskGroupStates = new HashMap<>();
- TaskState taskState = new TaskState(jvid, numberOfStates);
+ TaskState taskState = new TaskState(jvid, numberOfStates, numberOfStates);
taskGroupStates.put(jvid, taskState);
for (int i = 0; i < numberOfStates; i++) {
- SerializedValue<StateHandle<?>> stateHandle = new SerializedValue<StateHandle<?>>(
+ ChainedStateHandle<StreamStateHandle> stateHandle = CheckpointCoordinatorTest.generateChainedStateHandle(
new CheckpointMessagesTest.MyHandle());
- taskState.putState(i, new SubtaskState(stateHandle, 0, 0));
+ taskState.putState(i, new SubtaskState(stateHandle, 0));
}
return new TestCompletedCheckpoint(new JobID(), id, 0, taskGroupStates);
@@ -230,8 +228,6 @@ public abstract class CompletedCheckpointStoreTest extends TestLogger {
// Latch for test variants which discard asynchronously
private transient final CountDownLatch discardLatch = new CountDownLatch(1);
- private transient ClassLoader discardClassLoader;
-
public TestCompletedCheckpoint(
JobID jobId,
long checkpointId,
@@ -242,11 +238,10 @@ public abstract class CompletedCheckpointStoreTest extends TestLogger {
}
@Override
- public void discard(ClassLoader userClassLoader) throws Exception {
- super.discard(userClassLoader);
+ public void discardState() throws Exception {
+ super.discardState();
if (!isDiscarded) {
- this.discardClassLoader = userClassLoader;
this.isDiscarded = true;
if (discardLatch != null) {
@@ -265,10 +260,6 @@ public abstract class CompletedCheckpointStoreTest extends TestLogger {
}
}
- public ClassLoader getDiscardClassLoader() {
- return discardClassLoader;
- }
-
@Override
public boolean equals(Object o) {
if (this == o) return true;
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointTest.java
index 90a6836..9b04244 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CompletedCheckpointTest.java
@@ -45,14 +45,14 @@ public class CompletedCheckpointTest {
// Verify discard call is forwarded to state
CompletedCheckpoint checkpoint = new CompletedCheckpoint(new JobID(), 0, 0, 1, taskStates, true);
- checkpoint.discard(ClassLoader.getSystemClassLoader());
- verify(state, times(1)).discard(Matchers.any(ClassLoader.class));
+ checkpoint.discardState();
+ verify(state, times(1)).discardState();
Mockito.reset(state);
// Verify discard call is not forwarded to state
checkpoint = new CompletedCheckpoint(new JobID(), 0, 0, 1, taskStates, false);
- checkpoint.discard(ClassLoader.getSystemClassLoader());
- verify(state, times(0)).discard(Matchers.any(ClassLoader.class));
+ checkpoint.discardState();
+ verify(state, times(0)).discardState();
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingCheckpointTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingCheckpointTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingCheckpointTest.java
index d235e61..fd4e02d 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingCheckpointTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingCheckpointTest.java
@@ -66,7 +66,7 @@ public class PendingCheckpointTest {
setTaskState(pending, state);
pending.abortDeclined();
- verify(state, times(1)).discard(Matchers.any(ClassLoader.class));
+ verify(state, times(1)).discardState();
// Abort error
Mockito.reset(state);
@@ -75,7 +75,7 @@ public class PendingCheckpointTest {
setTaskState(pending, state);
pending.abortError(new Exception("Expected Test Exception"));
- verify(state, times(1)).discard(Matchers.any(ClassLoader.class));
+ verify(state, times(1)).discardState();
// Abort expired
Mockito.reset(state);
@@ -84,7 +84,7 @@ public class PendingCheckpointTest {
setTaskState(pending, state);
pending.abortExpired();
- verify(state, times(1)).discard(Matchers.any(ClassLoader.class));
+ verify(state, times(1)).discardState();
// Abort subsumed
Mockito.reset(state);
@@ -93,7 +93,7 @@ public class PendingCheckpointTest {
setTaskState(pending, state);
pending.abortSubsumed();
- verify(state, times(1)).discard(Matchers.any(ClassLoader.class));
+ verify(state, times(1)).discardState();
}
/**
@@ -106,21 +106,20 @@ public class PendingCheckpointTest {
PendingCheckpoint pending = createPendingCheckpoint();
PendingCheckpointTest.setTaskState(pending, state);
- pending.acknowledgeTask(ATTEMPT_ID, null, 0, null);
+ pending.acknowledgeTask(ATTEMPT_ID, null, null);
CompletedCheckpoint checkpoint = pending.finalizeCheckpoint();
// Does discard state
- checkpoint.discard(ClassLoader.getSystemClassLoader());
- verify(state, times(1)).discard(Matchers.any(ClassLoader.class));
+ checkpoint.discardState();
+ verify(state, times(1)).discardState();
}
// ------------------------------------------------------------------------
private static PendingCheckpoint createPendingCheckpoint() {
- ClassLoader classLoader = ClassLoader.getSystemClassLoader();
Map<ExecutionAttemptID, ExecutionVertex> ackTasks = new HashMap<>(ACK_TASKS);
- return new PendingCheckpoint(new JobID(), 0, 1, ackTasks, classLoader);
+ return new PendingCheckpoint(new JobID(), 0, 1, ackTasks);
}
@SuppressWarnings("unchecked")
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingSavepointTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingSavepointTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingSavepointTest.java
index 6ae6e1c..7258545 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingSavepointTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/PendingSavepointTest.java
@@ -71,7 +71,7 @@ public class PendingSavepointTest {
PendingCheckpointTest.setTaskState(pending, state);
pending.abortDeclined();
- verify(state, times(1)).discard(Matchers.any(ClassLoader.class));
+ verify(state, times(1)).discardState();
// Abort error
Mockito.reset(state);
@@ -81,7 +81,7 @@ public class PendingSavepointTest {
Future<String> future = pending.getCompletionFuture();
pending.abortError(new Exception("Expected Test Exception"));
- verify(state, times(1)).discard(Matchers.any(ClassLoader.class));
+ verify(state, times(1)).discardState();
assertTrue(future.failed().isCompleted());
// Abort expired
@@ -92,7 +92,7 @@ public class PendingSavepointTest {
future = pending.getCompletionFuture();
pending.abortExpired();
- verify(state, times(1)).discard(Matchers.any(ClassLoader.class));
+ verify(state, times(1)).discardState();
assertTrue(future.failed().isCompleted());
// Abort subsumed
@@ -117,13 +117,13 @@ public class PendingSavepointTest {
Future<String> future = pending.getCompletionFuture();
- pending.acknowledgeTask(ATTEMPT_ID, null, 0, null);
+ pending.acknowledgeTask(ATTEMPT_ID, null, null);
CompletedCheckpoint checkpoint = pending.finalizeCheckpoint();
// Does _NOT_ discard state
- checkpoint.discard(ClassLoader.getSystemClassLoader());
- verify(state, times(0)).discard(Matchers.any(ClassLoader.class));
+ checkpoint.discardState();
+ verify(state, times(0)).discardState();
// Future is completed
String path = Await.result(future, Duration.Zero());
@@ -133,9 +133,8 @@ public class PendingSavepointTest {
// ------------------------------------------------------------------------
private static PendingSavepoint createPendingSavepoint() {
- ClassLoader classLoader = ClassLoader.getSystemClassLoader();
Map<ExecutionAttemptID, ExecutionVertex> ackTasks = new HashMap<>(ACK_TASKS);
- return new PendingSavepoint(new JobID(), 0, 1, ackTasks, classLoader, new HeapSavepointStore());
+ return new PendingSavepoint(new JobID(), 0, 1, ackTasks, new HeapSavepointStore());
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStoreITCase.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStoreITCase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStoreITCase.java
index 380ba2c..f273797 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStoreITCase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStoreITCase.java
@@ -19,9 +19,8 @@
package org.apache.flink.runtime.checkpoint;
import org.apache.curator.framework.CuratorFramework;
-import org.apache.flink.runtime.state.LocalStateHandle;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.zookeeper.StateStorageHelper;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
+import org.apache.flink.runtime.zookeeper.RetrievableStateStorageHelper;
import org.apache.flink.runtime.zookeeper.ZooKeeperTestEnvironment;
import org.junit.AfterClass;
import org.junit.Before;
@@ -29,6 +28,8 @@ import org.junit.Test;
import scala.concurrent.duration.Deadline;
import scala.concurrent.duration.FiniteDuration;
+import java.io.IOException;
+import java.io.Serializable;
import java.util.concurrent.TimeUnit;
import static org.junit.Assert.assertEquals;
@@ -61,10 +62,10 @@ public class ZooKeeperCompletedCheckpointStoreITCase extends CompletedCheckpoint
int maxNumberOfCheckpointsToRetain, ClassLoader userLoader) throws Exception {
return new ZooKeeperCompletedCheckpointStore(maxNumberOfCheckpointsToRetain, userLoader,
- ZooKeeper.createClient(), CheckpointsPath, new StateStorageHelper<CompletedCheckpoint>() {
+ ZooKeeper.createClient(), CheckpointsPath, new RetrievableStateStorageHelper<CompletedCheckpoint>() {
@Override
- public StateHandle<CompletedCheckpoint> store(CompletedCheckpoint state) throws Exception {
- return new LocalStateHandle<>(state);
+ public RetrievableStateHandle<CompletedCheckpoint> store(CompletedCheckpoint state) throws Exception {
+ return new HeapRetrievableStateHandle<CompletedCheckpoint>(state);
}
});
}
@@ -160,4 +161,35 @@ public class ZooKeeperCompletedCheckpointStoreITCase extends CompletedCheckpoint
CompletedCheckpoint recovered = store.getLatestCheckpoint();
assertEquals(checkpoint, recovered);
}
+
+ static class HeapRetrievableStateHandle<T extends Serializable> implements RetrievableStateHandle<T> {
+
+ private static final long serialVersionUID = -268548467968932L;
+
+ public HeapRetrievableStateHandle(T state) {
+ this.state = state;
+ }
+
+ private T state;
+
+ @Override
+ public T retrieveState() throws Exception {
+ return state;
+ }
+
+ @Override
+ public void discardState() throws Exception {
+ state = null;
+ }
+
+ @Override
+ public long getStateSize() throws Exception {
+ return 0;
+ }
+
+ @Override
+ public void close() throws IOException {
+
+ }
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStoreTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStoreTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStoreTest.java
index 6b8c651..3e2de80 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStoreTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStoreTest.java
@@ -58,7 +58,7 @@ public class FsSavepointStoreTest {
assertEquals(0, tmp.getRoot().listFiles().length);
// Store
- SavepointV0 stored = new SavepointV0(1929292, SavepointV0Test.createTaskStates(4, 24));
+ SavepointV1 stored = new SavepointV1(1929292, SavepointV1Test.createTaskStates(4, 24));
String path = store.storeSavepoint(stored);
assertEquals(1, tmp.getRoot().listFiles().length);
@@ -67,7 +67,7 @@ public class FsSavepointStoreTest {
assertEquals(stored, loaded);
// Dispose
- store.disposeSavepoint(path, ClassLoader.getSystemClassLoader());
+ store.disposeSavepoint(path);
assertEquals(0, tmp.getRoot().listFiles().length);
}
@@ -122,7 +122,7 @@ public class FsSavepointStoreTest {
assertEquals(1, tmp.getRoot().listFiles().length);
// Savepoint v0
- Savepoint savepoint = new SavepointV0(checkpointId, SavepointV0Test.createTaskStates(4, 32));
+ Savepoint savepoint = new SavepointV1(checkpointId, SavepointV1Test.createTaskStates(4, 32));
String pathSavepoint = store.storeSavepoint(savepoint);
assertEquals(2, tmp.getRoot().listFiles().length);
@@ -208,7 +208,7 @@ public class FsSavepointStoreTest {
}
@Override
- public void dispose(ClassLoader classLoader) {
+ public void dispose() {
}
@Override
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java
index 6a85195..d703bd6 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java
@@ -65,7 +65,7 @@ public class SavepointLoaderTest {
true);
// Store savepoint
- SavepointV0 savepoint = new SavepointV0(stored.getCheckpointID(), taskStates.values());
+ SavepointV1 savepoint = new SavepointV1(stored.getCheckpointID(), taskStates.values());
SavepointStore store = new HeapSavepointStore();
String path = store.storeSavepoint(savepoint);
@@ -84,8 +84,8 @@ public class SavepointLoaderTest {
assertEquals(stored.getCheckpointID(), loaded.getCheckpointID());
// The loaded checkpoint should not discard state when its discarded
- loaded.discard(ClassLoader.getSystemClassLoader());
- verify(state, times(0)).discard(any(ClassLoader.class));
+ loaded.discardState();
+ verify(state, times(0)).discardState();
// 2) Load and validate: parallelism mismatch
when(vertex.getParallelism()).thenReturn(222);
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0SerializerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0SerializerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0SerializerTest.java
deleted file mode 100644
index b656d90..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0SerializerTest.java
+++ /dev/null
@@ -1,52 +0,0 @@
-/*
- * 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.savepoint;
-
-import org.apache.commons.io.output.ByteArrayOutputStream;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
-import org.junit.Test;
-
-import java.io.ByteArrayInputStream;
-
-import static org.junit.Assert.assertEquals;
-
-public class SavepointV0SerializerTest {
-
- /**
- * Test serialization of {@link SavepointV0} instance.
- */
- @Test
- public void testSerializeDeserializeV1() throws Exception {
- SavepointV0 expected = new SavepointV0(123123, SavepointV0Test.createTaskStates(8, 32));
-
- SavepointV0Serializer serializer = SavepointV0Serializer.INSTANCE;
-
- // Serialize
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- serializer.serialize(expected, new DataOutputViewStreamWrapper(baos));
- byte[] bytes = baos.toByteArray();
-
- // Deserialize
- ByteArrayInputStream bais = new ByteArrayInputStream(bytes);
- Savepoint actual = serializer.deserialize(new DataInputViewStreamWrapper(bais));
-
- assertEquals(expected, actual);
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0Test.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0Test.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0Test.java
deleted file mode 100644
index 4d72c42..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0Test.java
+++ /dev/null
@@ -1,81 +0,0 @@
-/*
- * 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.savepoint;
-
-import org.apache.flink.runtime.checkpoint.SubtaskState;
-import org.apache.flink.runtime.checkpoint.TaskState;
-import org.apache.flink.runtime.jobgraph.JobVertexID;
-import org.apache.flink.runtime.messages.CheckpointMessagesTest;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
-import org.junit.Test;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-import java.util.concurrent.ThreadLocalRandom;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-
-public class SavepointV0Test {
-
- /**
- * Simple test of savepoint methods.
- */
- @Test
- public void testSavepointV0() throws Exception {
- long checkpointId = ThreadLocalRandom.current().nextLong(Integer.MAX_VALUE);
- int numTaskStates = 4;
- int numSubtaskStates = 16;
-
- Collection<TaskState> expected = createTaskStates(numTaskStates, numSubtaskStates);
-
- SavepointV0 savepoint = new SavepointV0(checkpointId, expected);
-
- assertEquals(SavepointV0.VERSION, savepoint.getVersion());
- assertEquals(checkpointId, savepoint.getCheckpointId());
- assertEquals(expected, savepoint.getTaskStates());
-
- assertFalse(savepoint.getTaskStates().isEmpty());
- savepoint.dispose(ClassLoader.getSystemClassLoader());
- assertTrue(savepoint.getTaskStates().isEmpty());
- }
-
- static Collection<TaskState> createTaskStates(int numTaskStates, int numSubtaskStates) throws IOException {
- List<TaskState> taskStates = new ArrayList<>(numTaskStates);
-
- for (int i = 0; i < numTaskStates; i++) {
- TaskState taskState = new TaskState(new JobVertexID(), numSubtaskStates);
- for (int j = 0; j < numSubtaskStates; j++) {
- SerializedValue<StateHandle<?>> stateHandle = new SerializedValue<StateHandle<?>>(
- new CheckpointMessagesTest.MyHandle());
-
- taskState.putState(i, new SubtaskState(stateHandle, 0, 0));
- }
-
- taskStates.add(taskState);
- }
-
- return taskStates;
- }
-
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1SerializerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1SerializerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1SerializerTest.java
new file mode 100644
index 0000000..bad836b
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1SerializerTest.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.checkpoint.savepoint;
+
+import org.apache.commons.io.output.ByteArrayOutputStream;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.junit.Test;
+
+import java.io.ByteArrayInputStream;
+
+import static org.junit.Assert.assertEquals;
+
+public class SavepointV1SerializerTest {
+
+ /**
+ * Test serialization of {@link SavepointV1} instance.
+ */
+ @Test
+ public void testSerializeDeserializeV1() throws Exception {
+ SavepointV1 expected = new SavepointV1(123123, SavepointV1Test.createTaskStates(8, 32));
+
+ SavepointV1Serializer serializer = SavepointV1Serializer.INSTANCE;
+
+ // Serialize
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ serializer.serialize(expected, new DataOutputViewStreamWrapper(baos));
+ byte[] bytes = baos.toByteArray();
+
+ // Deserialize
+ ByteArrayInputStream bais = new ByteArrayInputStream(bytes);
+ Savepoint actual = serializer.deserialize(new DataInputViewStreamWrapper(bais));
+
+ assertEquals(expected, actual);
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1Test.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1Test.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1Test.java
new file mode 100644
index 0000000..ef10032
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1Test.java
@@ -0,0 +1,88 @@
+/*
+ * 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.savepoint;
+
+import org.apache.flink.runtime.checkpoint.SubtaskState;
+import org.apache.flink.runtime.checkpoint.TaskState;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupRangeOffsets;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class SavepointV1Test {
+
+ /**
+ * Simple test of savepoint methods.
+ */
+ @Test
+ public void testSavepointV1() throws Exception {
+ long checkpointId = ThreadLocalRandom.current().nextLong(Integer.MAX_VALUE);
+ int numTaskStates = 4;
+ int numSubtaskStates = 16;
+
+ Collection<TaskState> expected = createTaskStates(numTaskStates, numSubtaskStates);
+
+ SavepointV1 savepoint = new SavepointV1(checkpointId, expected);
+
+ assertEquals(SavepointV1.VERSION, savepoint.getVersion());
+ assertEquals(checkpointId, savepoint.getCheckpointId());
+ assertEquals(expected, savepoint.getTaskStates());
+
+ assertFalse(savepoint.getTaskStates().isEmpty());
+ savepoint.dispose();
+ assertTrue(savepoint.getTaskStates().isEmpty());
+ }
+
+ static Collection<TaskState> createTaskStates(int numTaskStates, int numSubtaskStates) throws IOException {
+ List<TaskState> taskStates = new ArrayList<>(numTaskStates);
+
+ for (int i = 0; i < numTaskStates; i++) {
+ TaskState taskState = new TaskState(new JobVertexID(), numSubtaskStates, numSubtaskStates);
+ for (int j = 0; j < numSubtaskStates; j++) {
+ StreamStateHandle stateHandle = new ByteStreamStateHandle("Hello".getBytes());
+ taskState.putState(i, new SubtaskState(
+ new ChainedStateHandle<>(Collections.singletonList(stateHandle)), 0));
+ }
+
+ taskState.putKeyedState(
+ 0,
+ new KeyGroupsStateHandle(
+ new KeyGroupRangeOffsets(1,1, new long[] {42}), new ByteStreamStateHandle("Hello".getBytes())));
+
+ taskStates.add(taskState);
+ }
+
+ return taskStates;
+ }
+
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTrackerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTrackerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTrackerTest.java
index 12bbf82..c513e26 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTrackerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTrackerTest.java
@@ -19,17 +19,18 @@
package org.apache.flink.runtime.checkpoint.stats;
import org.apache.flink.api.common.JobID;
+import org.apache.flink.core.fs.Path;
import org.apache.flink.metrics.groups.UnregisteredMetricsGroup;
import org.apache.flink.runtime.checkpoint.CompletedCheckpoint;
import org.apache.flink.runtime.checkpoint.SubtaskState;
import org.apache.flink.runtime.checkpoint.TaskState;
import org.apache.flink.runtime.executiongraph.ExecutionJobVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.filesystem.FileStateHandle;
import org.junit.Test;
-import java.io.IOException;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.Arrays;
@@ -51,7 +52,7 @@ import static org.mockito.Mockito.when;
public class SimpleCheckpointStatsTrackerTest {
private static final Random RAND = new Random();
-
+
@Test
public void testNoCompletedCheckpointYet() throws Exception {
CheckpointStatsTracker tracker = new SimpleCheckpointStatsTracker(
@@ -154,7 +155,7 @@ public class SimpleCheckpointStatsTrackerTest {
private static void verifyJobStats(
CheckpointStatsTracker tracker,
int historySize,
- CompletedCheckpoint[] checkpoints) {
+ CompletedCheckpoint[] checkpoints) throws Exception {
assertTrue(tracker.getJobStats().isDefined());
JobCheckpointStats jobStats = tracker.getJobStats().get();
@@ -275,14 +276,15 @@ public class SimpleCheckpointStatsTrackerTest {
}
private static CompletedCheckpoint[] generateRandomCheckpoints(
- int numCheckpoints) throws IOException {
+ int numCheckpoints) throws Exception {
// Config
JobID jobId = new JobID();
int minNumOperators = 4;
int maxNumOperators = 32;
- int minParallelism = 4;
- int maxParallelism = 16;
+ int minOperatorParallelism = 4;
+ int maxOperatorParallelism = 16;
+ int maxParallelism = 32;
// Use yuge numbers here in order to test that summing up state sizes
// does not overflow. This was a bug in the initial version, because
@@ -299,7 +301,7 @@ public class SimpleCheckpointStatsTrackerTest {
for (int i = 0; i < numOperators; i++) {
operatorIds[i] = new JobVertexID();
- operatorParallelism[i] = RAND.nextInt(maxParallelism - minParallelism + 1) + minParallelism;
+ operatorParallelism[i] = RAND.nextInt(maxOperatorParallelism - minOperatorParallelism + 1) + minOperatorParallelism;
}
// Generate checkpoints
@@ -317,7 +319,7 @@ public class SimpleCheckpointStatsTrackerTest {
JobVertexID operatorId = operatorIds[operatorIndex];
int parallelism = operatorParallelism[operatorIndex];
- TaskState taskState = new TaskState(operatorId, parallelism);
+ TaskState taskState = new TaskState(operatorId, parallelism, maxParallelism);
taskGroupStates.put(operatorId, taskState);
@@ -328,9 +330,11 @@ public class SimpleCheckpointStatsTrackerTest {
completionDuration = duration;
}
+ final long proxySize = minStateSize + ((long) (RAND.nextDouble() * (maxStateSize - minStateSize)));
+ StreamStateHandle proxy = new StateHandleProxy(new Path(), proxySize);
+
SubtaskState subtaskState = new SubtaskState(
- new SerializedValue<StateHandle<?>>(null),
- minStateSize + ((long) (RAND.nextDouble() * (maxStateSize - minStateSize))),
+ new ChainedStateHandle<>(Arrays.asList(proxy)),
duration);
taskState.putState(subtaskIndex, subtaskState);
@@ -356,10 +360,32 @@ public class SimpleCheckpointStatsTrackerTest {
ExecutionJobVertex v = mock(ExecutionJobVertex.class);
when(v.getJobVertexId()).thenReturn(operatorId);
when(v.getParallelism()).thenReturn(parallelism);
-
+
jobVertices.add(v);
}
return jobVertices;
}
+
+ private static class StateHandleProxy extends FileStateHandle {
+
+ private static final long serialVersionUID = 35356735683568L;
+
+ public StateHandleProxy(Path filePath, long proxySize) {
+ super(filePath);
+ this.proxySize = proxySize;
+ }
+
+ private long proxySize;
+
+ @Override
+ public void discardState() throws Exception {
+
+ }
+
+ @Override
+ public long getStateSize() {
+ return proxySize;
+ }
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerHARecoveryTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerHARecoveryTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerHARecoveryTest.java
index 0e1c7c5..6b80c3d 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerHARecoveryTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/JobManagerHARecoveryTest.java
@@ -53,8 +53,11 @@ import org.apache.flink.runtime.leaderelection.TestingLeaderElectionService;
import org.apache.flink.runtime.leaderelection.TestingLeaderRetrievalService;
import org.apache.flink.runtime.leaderretrieval.LeaderRetrievalService;
import org.apache.flink.runtime.messages.JobManagerMessages;
-import org.apache.flink.runtime.state.LocalStateHandle;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.RetrievableStreamStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
import org.apache.flink.runtime.taskmanager.TaskManager;
import org.apache.flink.runtime.testingUtils.TestingJobManager;
import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages;
@@ -62,6 +65,7 @@ import org.apache.flink.runtime.testingUtils.TestingMessages;
import org.apache.flink.runtime.testingUtils.TestingTaskManager;
import org.apache.flink.runtime.testingUtils.TestingTaskManagerMessages;
import org.apache.flink.runtime.testingUtils.TestingUtils;
+import org.apache.flink.util.InstantiationUtil;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Rule;
@@ -141,38 +145,38 @@ public class JobManagerHARecoveryTest {
instanceManager.addInstanceListener(scheduler);
archive = system.actorOf(Props.create(
- MemoryArchivist.class,
- 10), "archive");
+ MemoryArchivist.class,
+ 10), "archive");
Props jobManagerProps = Props.create(
- TestingJobManager.class,
- flinkConfiguration,
- new ForkJoinPool(),
- instanceManager,
- scheduler,
- new BlobLibraryCacheManager(new BlobServer(flinkConfiguration), 3600000),
- archive,
- new FixedDelayRestartStrategy.FixedDelayRestartStrategyFactory(Int.MaxValue(), 100),
- timeout,
- myLeaderElectionService,
- mySubmittedJobGraphStore,
- checkpointStateFactory,
- new HeapSavepointStore(),
- jobRecoveryTimeout,
- Option.apply(null));
+ TestingJobManager.class,
+ flinkConfiguration,
+ new ForkJoinPool(),
+ instanceManager,
+ scheduler,
+ new BlobLibraryCacheManager(new BlobServer(flinkConfiguration), 3600000),
+ archive,
+ new FixedDelayRestartStrategy.FixedDelayRestartStrategyFactory(Int.MaxValue(), 100),
+ timeout,
+ myLeaderElectionService,
+ mySubmittedJobGraphStore,
+ checkpointStateFactory,
+ new HeapSavepointStore(),
+ jobRecoveryTimeout,
+ Option.apply(null));
jobManager = system.actorOf(jobManagerProps, "jobmanager");
ActorGateway gateway = new AkkaActorGateway(jobManager, leaderSessionID);
taskManager = TaskManager.startTaskManagerComponentsAndActor(
- flinkConfiguration,
- ResourceID.generate(),
- system,
- "localhost",
- Option.apply("taskmanager"),
- Option.apply((LeaderRetrievalService) myLeaderRetrievalService),
- true,
- TestingTaskManager.class);
+ flinkConfiguration,
+ ResourceID.generate(),
+ system,
+ "localhost",
+ Option.apply("taskmanager"),
+ Option.apply((LeaderRetrievalService) myLeaderRetrievalService),
+ true,
+ TestingTaskManager.class);
ActorGateway tmGateway = new AkkaActorGateway(taskManager, leaderSessionID);
@@ -199,12 +203,12 @@ public class JobManagerHARecoveryTest {
BlockingStatefulInvokable.initializeStaticHelpers(slots);
Future<Object> isLeader = gateway.ask(
- TestingJobManagerMessages.getNotifyWhenLeader(),
- deadline.timeLeft());
+ TestingJobManagerMessages.getNotifyWhenLeader(),
+ deadline.timeLeft());
Future<Object> isConnectedToJobManager = tmGateway.ask(
- new TestingTaskManagerMessages.NotifyWhenRegisteredAtJobManager(jobManager),
- deadline.timeLeft());
+ new TestingTaskManagerMessages.NotifyWhenRegisteredAtJobManager(jobManager),
+ deadline.timeLeft());
// tell jobManager that he's the leader
myLeaderElectionService.isLeader(leaderSessionID);
@@ -216,8 +220,8 @@ public class JobManagerHARecoveryTest {
// submit blocking job
Future<Object> jobSubmitted = gateway.ask(
- new JobManagerMessages.SubmitJob(jobGraph, ListeningBehaviour.DETACHED),
- deadline.timeLeft());
+ new JobManagerMessages.SubmitJob(jobGraph, ListeningBehaviour.DETACHED),
+ deadline.timeLeft());
Await.ready(jobSubmitted, deadline.timeLeft());
@@ -298,7 +302,7 @@ public class JobManagerHARecoveryTest {
public void addCheckpoint(CompletedCheckpoint checkpoint) throws Exception {
checkpoints.addLast(checkpoint);
if (checkpoints.size() > 1) {
- checkpoints.removeFirst().discard(ClassLoader.getSystemClassLoader());
+ checkpoints.removeFirst().discardState();
}
}
@@ -342,10 +346,12 @@ public class JobManagerHARecoveryTest {
}
@Override
- public void start() {}
+ public void start() {
+ }
@Override
- public void stop() {}
+ public void stop() {
+ }
@Override
public CompletedCheckpointStore createCheckpointStore(JobID jobId, ClassLoader userClassLoader) throws Exception {
@@ -408,7 +414,7 @@ public class JobManagerHARecoveryTest {
@Override
public void invoke() throws Exception {
- while(blocking) {
+ while (blocking) {
synchronized (lock) {
lock.wait();
}
@@ -424,7 +430,7 @@ public class JobManagerHARecoveryTest {
}
}
- public static class BlockingStatefulInvokable extends BlockingInvokable implements StatefulTask<StateHandle<Long>> {
+ public static class BlockingStatefulInvokable extends BlockingInvokable implements StatefulTask {
private static final int NUM_CHECKPOINTS_TO_COMPLETE = 5;
@@ -435,18 +441,28 @@ public class JobManagerHARecoveryTest {
private int completedCheckpoints = 0;
@Override
- public void setInitialState(StateHandle<Long> stateHandle) throws Exception {
+ public void setInitialState(ChainedStateHandle<StreamStateHandle> chainedState, List<KeyGroupsStateHandle> keyGroupsState) throws Exception {
int subtaskIndex = getIndexInSubtaskGroup();
if (subtaskIndex < recoveredStates.length) {
- recoveredStates[subtaskIndex] = stateHandle.getState(getUserCodeClassLoader());
+ recoveredStates[subtaskIndex] = InstantiationUtil.deserializeObject(chainedState.get(0).openInputStream());
}
}
@Override
public boolean triggerCheckpoint(long checkpointId, long timestamp) {
- StateHandle<Long> state = new LocalStateHandle<>(checkpointId);
- getEnvironment().acknowledgeCheckpoint(checkpointId, state);
- return true;
+ try {
+ ByteStreamStateHandle byteStreamStateHandle = new ByteStreamStateHandle(
+ InstantiationUtil.serializeObject(checkpointId));
+ RetrievableStreamStateHandle<Long> state = new RetrievableStreamStateHandle<Long>(byteStreamStateHandle);
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle = new ChainedStateHandle<StreamStateHandle>(Collections.singletonList(state));
+ getEnvironment().acknowledgeCheckpoint(
+ checkpointId,
+ chainedStateHandle,
+ Collections.<KeyGroupsStateHandle>emptyList());
+ return true;
+ } catch (Exception ex) {
+ throw new RuntimeException(ex);
+ }
}
@Override
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphsStoreITCase.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphsStoreITCase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphsStoreITCase.java
index 426dfba..6ef184d 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphsStoreITCase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphsStoreITCase.java
@@ -19,16 +19,17 @@
package org.apache.flink.runtime.jobmanager;
import akka.actor.ActorRef;
-import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.akka.ListeningBehaviour;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobgraph.JobVertex;
import org.apache.flink.runtime.jobmanager.SubmittedJobGraphStore.SubmittedJobGraphListener;
-import org.apache.flink.runtime.state.LocalStateHandle;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.zookeeper.StateStorageHelper;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
+import org.apache.flink.runtime.state.RetrievableStreamStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+import org.apache.flink.runtime.zookeeper.RetrievableStateStorageHelper;
import org.apache.flink.runtime.zookeeper.ZooKeeperTestEnvironment;
+import org.apache.flink.util.InstantiationUtil;
import org.apache.flink.util.TestLogger;
import org.junit.AfterClass;
import org.junit.Before;
@@ -36,6 +37,7 @@ import org.junit.Test;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
+import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.CountDownLatch;
@@ -57,10 +59,12 @@ public class ZooKeeperSubmittedJobGraphsStoreITCase extends TestLogger {
private final static ZooKeeperTestEnvironment ZooKeeper = new ZooKeeperTestEnvironment(1);
- private final static StateStorageHelper<SubmittedJobGraph> localStateStorage = new StateStorageHelper<SubmittedJobGraph>() {
+ private final static RetrievableStateStorageHelper<SubmittedJobGraph> localStateStorage = new RetrievableStateStorageHelper<SubmittedJobGraph>() {
@Override
- public StateHandle<SubmittedJobGraph> store(SubmittedJobGraph state) throws Exception {
- return new LocalStateHandle<>(state);
+ public RetrievableStateHandle<SubmittedJobGraph> store(SubmittedJobGraph state) throws IOException {
+ ByteStreamStateHandle byteStreamStateHandle = new ByteStreamStateHandle(
+ InstantiationUtil.serializeObject(state));
+ return new RetrievableStreamStateHandle<SubmittedJobGraph>(byteStreamStateHandle);
}
};
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java
index 73bf204..c6eb249 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java
@@ -18,32 +18,38 @@
package org.apache.flink.runtime.messages;
-import static org.junit.Assert.*;
-
import org.apache.flink.api.common.JobID;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.testutils.CommonTestUtils;
+import org.apache.flink.runtime.checkpoint.CheckpointCoordinatorTest;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
-import org.apache.flink.runtime.messages.checkpoint.NotifyCheckpointComplete;
import org.apache.flink.runtime.messages.checkpoint.AcknowledgeCheckpoint;
+import org.apache.flink.runtime.messages.checkpoint.NotifyCheckpointComplete;
import org.apache.flink.runtime.messages.checkpoint.TriggerCheckpoint;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.core.testutils.CommonTestUtils;
-import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.StateObject;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.junit.Test;
import java.io.IOException;
import java.io.Serializable;
+import java.util.Collections;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.fail;
public class CheckpointMessagesTest {
-
+
@Test
public void testTriggerAndConfirmCheckpoint() {
try {
NotifyCheckpointComplete cc = new NotifyCheckpointComplete(new JobID(), new ExecutionAttemptID(), 45287698767345L, 467L);
testSerializabilityEqualsHashCode(cc);
-
+
TriggerCheckpoint tc = new TriggerCheckpoint(new JobID(), new ExecutionAttemptID(), 347652734L, 7576752L);
testSerializabilityEqualsHashCode(tc);
-
+
}
catch (Exception e) {
e.printStackTrace();
@@ -55,35 +61,40 @@ public class CheckpointMessagesTest {
public void testConfirmTaskCheckpointed() {
try {
AcknowledgeCheckpoint noState = new AcknowledgeCheckpoint(
- new JobID(), new ExecutionAttemptID(), 569345L);
+ new JobID(), new ExecutionAttemptID(), 569345L);
+
+ KeyGroupRange keyGroupRange = KeyGroupRange.of(42,42);
AcknowledgeCheckpoint withState = new AcknowledgeCheckpoint(
- new JobID(), new ExecutionAttemptID(), 87658976143L,
- new SerializedValue<StateHandle<?>>(new MyHandle()), 0);
-
+ new JobID(),
+ new ExecutionAttemptID(),
+ 87658976143L,
+ CheckpointCoordinatorTest.generateChainedStateHandle(new MyHandle()),
+ CheckpointCoordinatorTest.generateKeyGroupState(
+ keyGroupRange, Collections.singletonList(new MyHandle())));
+
testSerializabilityEqualsHashCode(noState);
testSerializabilityEqualsHashCode(withState);
- }
- catch (Exception e) {
+ } catch (Exception e) {
e.printStackTrace();
fail(e.getMessage());
}
}
-
+
private static void testSerializabilityEqualsHashCode(Serializable o) throws IOException {
Object copy = CommonTestUtils.createCopySerializable(o);
+ System.out.println(o.getClass() +" "+copy.getClass());
assertEquals(o, copy);
assertEquals(o.hashCode(), copy.hashCode());
assertNotNull(o.toString());
assertNotNull(copy.toString());
}
-
- public static class MyHandle implements StateHandle<Serializable> {
+
+ public static class MyHandle implements StreamStateHandle {
private static final long serialVersionUID = 8128146204128728332L;
- @Override
- public Serializable getState(ClassLoader userCodeClassLoader) {
+ public Serializable get(ClassLoader userCodeClassLoader) {
return null;
}
@@ -107,5 +118,10 @@ public class CheckpointMessagesTest {
@Override
public void close() throws IOException {}
+
+ @Override
+ public FSDataInputStream openInputStream() throws Exception {
+ return null;
+ }
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java
index 87540bc..19317f9 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java
@@ -36,10 +36,13 @@ import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.query.KvStateRegistry;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
import java.util.Collections;
+import java.util.List;
import java.util.Map;
import java.util.concurrent.Future;
@@ -149,10 +152,15 @@ public class DummyEnvironment implements Environment {
}
@Override
- public void acknowledgeCheckpoint(long checkpointId) {}
+ public void acknowledgeCheckpoint(long checkpointId) {
+
+ }
@Override
- public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) {}
+ public void acknowledgeCheckpoint(long checkpointId,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
+ List<KeyGroupsStateHandle> keyGroupStateHandles) {
+ }
@Override
public void failExternally(Throwable cause) {
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
index 7b966c3..2c76399 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
@@ -45,7 +45,10 @@ import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.metrics.groups.TaskMetricGroup;
import org.apache.flink.runtime.query.KvStateRegistry;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
import org.apache.flink.types.Record;
import org.apache.flink.util.MutableObjectIterator;
@@ -96,9 +99,13 @@ public class MockEnvironment implements Environment {
private final int bufferSize;
public MockEnvironment(String taskName, long memorySize, MockInputSplitProvider inputSplitProvider, int bufferSize) {
+ this(taskName, memorySize, inputSplitProvider, bufferSize, new Configuration());
+ }
+
+ public MockEnvironment(String taskName, long memorySize, MockInputSplitProvider inputSplitProvider, int bufferSize, Configuration taskConfiguration) {
this.taskInfo = new TaskInfo(taskName, 0, 1, 0);
this.jobConfiguration = new Configuration();
- this.taskConfiguration = new Configuration();
+ this.taskConfiguration = taskConfiguration;
this.inputs = new LinkedList<InputGate>();
this.outputs = new LinkedList<ResultPartitionWriter>();
@@ -298,7 +305,9 @@ public class MockEnvironment implements Environment {
}
@Override
- public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) {
+ public void acknowledgeCheckpoint(long checkpointId,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
+ List<KeyGroupsStateHandle> keyGroupStateHandles) {
throw new UnsupportedOperationException();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/state/AbstractCloseableHandleTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/AbstractCloseableHandleTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/AbstractCloseableHandleTest.java
index ad3339a..40e1852 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/AbstractCloseableHandleTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/AbstractCloseableHandleTest.java
@@ -85,5 +85,15 @@ public class AbstractCloseableHandleTest {
private static final class CloseableHandle extends AbstractCloseableHandle {
private static final long serialVersionUID = 1L;
+
+ @Override
+ public void discardState() throws Exception {
+
+ }
+
+ @Override
+ public long getStateSize() throws Exception {
+ return 0;
+ }
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java
index 0f1c0f7..04fa089 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/FileStateBackendTest.java
@@ -24,7 +24,8 @@ import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.core.fs.Path;
import org.apache.flink.core.testutils.CommonTestUtils;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
-import org.apache.flink.runtime.state.filesystem.FileStreamStateHandle;
+
+import org.apache.flink.runtime.state.filesystem.FileStateHandle;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
@@ -86,7 +87,12 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
// no file operations should be possible right now
try {
- backend.checkpointStateSerializable("exception train rolling in", 2L, System.currentTimeMillis());
+ FsStateBackend.FsCheckpointStateOutputStream out = backend.createCheckpointStateOutputStream(
+ 2L,
+ System.currentTimeMillis());
+
+ out.write(1);
+ out.closeAndGetHandle();
fail("should fail with an exception");
} catch (IllegalStateException e) {
// supreme!
@@ -114,43 +120,6 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
}
@Test
- public void testSerializableState() {
- File tempDir = new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString());
- try {
- FsStateBackend backend = CommonTestUtils.createCopySerializable(new FsStateBackend(localFileUri(tempDir)));
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test-op", IntSerializer.INSTANCE);
-
- File checkpointDir = new File(backend.getCheckpointDirectory().toUri().getPath());
-
- String state1 = "dummy state";
- String state2 = "row row row your boat";
- Integer state3 = 42;
-
- StateHandle<String> handle1 = backend.checkpointStateSerializable(state1, 439568923746L, System.currentTimeMillis());
- StateHandle<String> handle2 = backend.checkpointStateSerializable(state2, 439568923746L, System.currentTimeMillis());
- StateHandle<Integer> handle3 = backend.checkpointStateSerializable(state3, 439568923746L, System.currentTimeMillis());
-
- assertEquals(state1, handle1.getState(getClass().getClassLoader()));
- handle1.discardState();
-
- assertEquals(state2, handle2.getState(getClass().getClassLoader()));
- handle2.discardState();
-
- assertEquals(state3, handle3.getState(getClass().getClassLoader()));
- handle3.discardState();
-
- assertTrue(isDirectoryEmpty(checkpointDir));
- }
- catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
- }
- finally {
- deleteDirectorySilently(tempDir);
- }
- }
-
- @Test
public void testStateOutputStream() {
File tempDir = new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString());
try {
@@ -185,16 +154,16 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
stream2.write(state2);
stream3.write(state3);
- FileStreamStateHandle handle1 = (FileStreamStateHandle) stream1.closeAndGetHandle();
+ FileStateHandle handle1 = (FileStateHandle) stream1.closeAndGetHandle();
ByteStreamStateHandle handle2 = (ByteStreamStateHandle) stream2.closeAndGetHandle();
ByteStreamStateHandle handle3 = (ByteStreamStateHandle) stream3.closeAndGetHandle();
// use with try-with-resources
- FileStreamStateHandle handle4;
+ StreamStateHandle handle4;
try (AbstractStateBackend.CheckpointStateOutputStream stream4 =
backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis())) {
stream4.write(state4);
- handle4 = (FileStreamStateHandle) stream4.closeAndGetHandle();
+ handle4 = stream4.closeAndGetHandle();
}
// close before accessing handle
@@ -209,18 +178,18 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
// uh-huh
}
- validateBytesInStream(handle1.getState(getClass().getClassLoader()), state1);
+ validateBytesInStream(handle1.openInputStream(), state1);
handle1.discardState();
assertFalse(isDirectoryEmpty(checkpointDir));
ensureLocalFileDeleted(handle1.getFilePath());
- validateBytesInStream(handle2.getState(getClass().getClassLoader()), state2);
+ validateBytesInStream(handle2.openInputStream(), state2);
handle2.discardState();
- validateBytesInStream(handle3.getState(getClass().getClassLoader()), state3);
+ validateBytesInStream(handle3.openInputStream(), state3);
handle3.discardState();
- validateBytesInStream(handle4.getState(getClass().getClassLoader()), state4);
+ validateBytesInStream(handle4.openInputStream(), state4);
handle4.discardState();
assertTrue(isDirectoryEmpty(checkpointDir));
}
[23/27] flink git commit: [FLINK-3755] Ignore QueryableStateITCase
Posted by al...@apache.org.
[FLINK-3755] Ignore QueryableStateITCase
This doesn't work yet because the state query machinery is not yet
properly aware of key-grouped state.
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/2b7a8d68
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/2b7a8d68
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/2b7a8d68
Branch: refs/heads/master
Commit: 2b7a8d6846d7e5f03819a9ac6cfbd42e9bd72476
Parents: addd084
Author: Aljoscha Krettek <al...@gmail.com>
Authored: Thu Aug 25 14:09:12 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:10:01 2016 +0200
----------------------------------------------------------------------
.../java/org/apache/flink/test/query/QueryableStateITCase.java | 2 ++
1 file changed, 2 insertions(+)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/2b7a8d68/flink-tests/src/test/java/org/apache/flink/test/query/QueryableStateITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/query/QueryableStateITCase.java b/flink-tests/src/test/java/org/apache/flink/test/query/QueryableStateITCase.java
index c31f4e5..40732df 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/query/QueryableStateITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/query/QueryableStateITCase.java
@@ -68,6 +68,7 @@ import org.apache.flink.util.Preconditions;
import org.apache.flink.util.TestLogger;
import org.junit.AfterClass;
import org.junit.BeforeClass;
+import org.junit.Ignore;
import org.junit.Test;
import scala.concurrent.Await;
import scala.concurrent.Future;
@@ -92,6 +93,7 @@ import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
+@Ignore
public class QueryableStateITCase extends TestLogger {
private final static FiniteDuration TEST_TIMEOUT = new FiniteDuration(100, TimeUnit.SECONDS);
[25/27] flink git commit: [hotfix] Improve Logging in
CheckpointCoordinator, StreamTask, State Code
Posted by al...@apache.org.
[hotfix] Improve Logging in CheckpointCoordinator, StreamTask, State Code
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/bdf9f86c
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/bdf9f86c
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/bdf9f86c
Branch: refs/heads/master
Commit: bdf9f86c51aa2122441adea6944ce46f393d729f
Parents: f7ef82b
Author: Stefan Richter <s....@data-artisans.com>
Authored: Wed Aug 31 11:19:13 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:10:02 2016 +0200
----------------------------------------------------------------------
.../flink/runtime/checkpoint/CheckpointCoordinator.java | 4 ++--
.../runtime/state/filesystem/FsCheckpointStreamFactory.java | 4 +++-
.../org/apache/flink/streaming/runtime/tasks/StreamTask.java | 7 +++++--
3 files changed, 10 insertions(+), 5 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/bdf9f86c/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
index 52f6d9a..3586d98 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
@@ -19,7 +19,6 @@
package org.apache.flink.runtime.checkpoint;
import akka.dispatch.Futures;
-
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.checkpoint.savepoint.SavepointStore;
import org.apache.flink.runtime.checkpoint.stats.CheckpointStatsTracker;
@@ -42,7 +41,6 @@ import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.util.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-
import scala.concurrent.Future;
import java.util.ArrayDeque;
@@ -778,6 +776,8 @@ public class CheckpointCoordinator {
}
}
+ LOG.info("Restoring from latest valid checkpoint: {}.", latest);
+
for (Map.Entry<JobVertexID, TaskState> taskGroupStateEntry: latest.getTaskStates().entrySet()) {
TaskState taskState = taskGroupStateEntry.getValue();
ExecutionJobVertex executionJobVertex = tasks.get(taskGroupStateEntry.getKey());
http://git-wip-us.apache.org/repos/asf/flink/blob/bdf9f86c/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStreamFactory.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStreamFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStreamFactory.java
index cc13a72..c027558 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStreamFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStreamFactory.java
@@ -97,7 +97,9 @@ public class FsCheckpointStreamFactory implements CheckpointStreamFactory {
Path dir = new Path(basePath, jobId.toString());
- LOG.info("Initializing file stream factory to URI {}.", dir);
+ if (LOG.isDebugEnabled()) {
+ LOG.debug("Initializing file stream factory to URI {}.", dir);
+ }
filesystem = basePath.getFileSystem();
filesystem.mkdirs(dir);
http://git-wip-us.apache.org/repos/asf/flink/blob/bdf9f86c/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
index bedc8fa..9c26509 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
@@ -970,10 +970,13 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
if (chainedStateHandles.isEmpty() && keyedStates.isEmpty()) {
owner.getEnvironment().acknowledgeCheckpoint(checkpointId);
} else {
- owner. getEnvironment().acknowledgeCheckpoint(checkpointId, chainedStateHandles, keyedStates);
+ owner.getEnvironment().acknowledgeCheckpoint(checkpointId, chainedStateHandles, keyedStates);
}
- LOG.debug("Finished asynchronous checkpoints for checkpoint {} on task {}", checkpointId, name);
+ if(LOG.isDebugEnabled()) {
+ LOG.debug("Finished asynchronous checkpoints for checkpoint {} on task {}. Returning handles on " +
+ "keyed states {}.", checkpointId, name, keyedStates);
+ }
}
catch (Exception e) {
if (owner.isRunning()) {
[18/27] flink git commit: [FLINK-4381] Refactor State to Prepare For
Key-Group State Backends
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChainedStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChainedStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChainedStateHandle.java
new file mode 100644
index 0000000..9b308a3
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/ChainedStateHandle.java
@@ -0,0 +1,131 @@
+/*
+ * 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.IOException;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Handle to the non-partitioned states for the operators in an operator chain.
+ */
+public class ChainedStateHandle<T extends StateObject> implements StateObject {
+
+ private static final long serialVersionUID = 1L;
+
+ /** The state handles for all operators in the chain */
+ private final List<? extends T> operatorStateHandles;
+
+ /**
+ * Wraps a list to the state handles for the operators in a chain. Individual state handles can be null.
+ *
+ * @param operatorStateHandles list with the state handles for the states in the operator chain.
+ */
+ public ChainedStateHandle(List<? extends T> operatorStateHandles) {
+ this.operatorStateHandles = Preconditions.checkNotNull(operatorStateHandles);
+ }
+
+ /**
+ * Check if there are any states handles present. Notice that this can be true even if {@link #getLength()} is
+ * greater than zero, because state handles can be null.
+ *
+ * @return true if there are no state handles for any operator.
+ */
+ public boolean isEmpty() {
+ for (T state : operatorStateHandles) {
+ if (state != null) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ /**
+ * Returns the length of the operator chain. This can be different from the number of operator state handles,
+ * because the some operators in the chain can have no state and thus their state handle can be null.
+ *
+ * @return length of the operator chain
+ */
+ public int getLength() {
+ return operatorStateHandles.size();
+ }
+
+ /**
+ * Get the state handle for a single operator in the operator chain by it's index.
+ *
+ * @param index the index in the operator chain
+ * @return state handle to the operator at the given position in the operator chain. can be null.
+ */
+ public T get(int index) {
+ return operatorStateHandles.get(index);
+ }
+
+ @Override
+ public void discardState() throws Exception {
+ StateUtil.bestEffortDiscardAllStateObjects(operatorStateHandles);
+ }
+
+ @Override
+ public long getStateSize() throws Exception {
+ long sumStateSize = 0;
+
+ if (operatorStateHandles != null) {
+ for (T state : operatorStateHandles) {
+ if (state != null) {
+ sumStateSize += state.getStateSize();
+ }
+ }
+ }
+
+ // State size as sum of all state sizes
+ return sumStateSize;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+
+ ChainedStateHandle<?> that = (ChainedStateHandle<?>) o;
+
+ return operatorStateHandles.equals(that.operatorStateHandles);
+
+ }
+
+ @Override
+ public int hashCode() {
+ return operatorStateHandles.hashCode();
+ }
+
+ public static <T extends StateObject> ChainedStateHandle<T> wrapSingleHandle(T stateHandleToWrap) {
+ return new ChainedStateHandle<T>(Collections.singletonList(stateHandleToWrap));
+ }
+
+ @Override
+ public void close() throws IOException {
+ StateUtil.bestEffortCloseAllStateObjects(operatorStateHandles);
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
index 280746d..9ee4b90 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
@@ -56,11 +56,11 @@ public class HashKeyGroupAssigner<K> implements KeyGroupAssigner<K> {
}
@Override
- public void setup(int numberKeyGroups) {
- Preconditions.checkArgument(numberKeyGroups > 0, "The number of key groups has to be " +
+ public void setup(int numberOfKeygroups) {
+ Preconditions.checkArgument(numberOfKeygroups > 0, "The number of key groups has to be " +
"greater than 0. Use setMaxParallelism() to specify the number of key " +
"groups.");
- this.numberKeyGroups = numberKeyGroups;
+ this.numberKeyGroups = numberOfKeygroups;
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
new file mode 100644
index 0000000..de42bdb
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
@@ -0,0 +1,217 @@
+/*
+ * 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.Serializable;
+import java.util.Iterator;
+
+/**
+ * This class defines a range of key-group indexes. Key-groups are the granularity into which the keyspace of a job
+ * is partitioned for keyed state-handling in state backends. The boundaries of the range are inclusive.
+ */
+public class KeyGroupRange implements Iterable<Integer>, Serializable {
+
+ /** The empty key-group */
+ public static final KeyGroupRange EMPTY_KEY_GROUP = new KeyGroupRange();
+
+ private final int startKeyGroup;
+ private final int endKeyGroup;
+
+ /**
+ * Empty KeyGroup Constructor
+ */
+ private KeyGroupRange() {
+ this.startKeyGroup = 0;
+ this.endKeyGroup = -1;
+ }
+
+ /**
+ * Defines the range [startKeyGroup, endKeyGroup]
+ *
+ * @param startKeyGroup start of the range (inclusive)
+ * @param endKeyGroup end of the range (inclusive)
+ */
+ public KeyGroupRange(int startKeyGroup, int endKeyGroup) {
+ Preconditions.checkArgument(startKeyGroup >= 0);
+ Preconditions.checkArgument(startKeyGroup <= endKeyGroup);
+ this.startKeyGroup = startKeyGroup;
+ this.endKeyGroup = endKeyGroup;
+ Preconditions.checkArgument(getNumberOfKeyGroups() >= 0, "Potential overflow detected.");
+ }
+
+
+ /**
+ * Checks whether or not a single key-group is contained in the range.
+ *
+ * @param keyGroup Key-group to check for inclusion.
+ * @return True, only if the key-group is in the range.
+ */
+ public boolean contains(int keyGroup) {
+ return keyGroup >= startKeyGroup && keyGroup <= endKeyGroup;
+ }
+
+ /**
+ * Create a range that represent the intersection between this range and the given range.
+ *
+ * @param other A KeyGroupRange to intersect.
+ * @return Key-group range that is the intersection between this and the given key-group range.
+ */
+ public KeyGroupRange getIntersection(KeyGroupRange other) {
+ int start = Math.max(startKeyGroup, other.startKeyGroup);
+ int end = Math.min(endKeyGroup, other.endKeyGroup);
+ return start <= end ? new KeyGroupRange(start, end) : EMPTY_KEY_GROUP;
+ }
+
+ /**
+ *
+ * @return The number of key-groups in the range
+ */
+ public int getNumberOfKeyGroups() {
+ return 1 + endKeyGroup - startKeyGroup;
+ }
+
+ /**
+ *
+ * @return The first key-group in the range.
+ */
+ public int getStartKeyGroup() {
+ return startKeyGroup;
+ }
+
+ /**
+ *
+ * @return The last key-group in the range.
+ */
+ public int getEndKeyGroup() {
+ return endKeyGroup;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof KeyGroupRange)) {
+ return false;
+ }
+
+ KeyGroupRange that = (KeyGroupRange) o;
+ return startKeyGroup == that.startKeyGroup && endKeyGroup == that.endKeyGroup;
+ }
+
+
+ @Override
+ public int hashCode() {
+ int result = startKeyGroup;
+ result = 31 * result + endKeyGroup;
+ return result;
+ }
+
+ @Override
+ public String toString() {
+ return "KeyGroupRange{" +
+ "startKeyGroup=" + startKeyGroup +
+ ", endKeyGroup=" + endKeyGroup +
+ '}';
+ }
+
+ @Override
+ public Iterator<Integer> iterator() {
+ return new KeyGroupIterator();
+ }
+
+ private final class KeyGroupIterator implements Iterator<Integer> {
+
+ public KeyGroupIterator() {
+ this.iteratorPos = 0;
+ }
+
+ private int iteratorPos;
+
+ @Override
+ public boolean hasNext() {
+ return iteratorPos < getNumberOfKeyGroups();
+ }
+
+ @Override
+ public Integer next() {
+ int rv = startKeyGroup + iteratorPos;
+ ++iteratorPos;
+ return rv;
+ }
+
+ @Override
+ public void remove() {
+ throw new UnsupportedOperationException("Unsupported by this iterator!");
+ }
+ }
+
+ /**
+ * Factory method that also handles creation of empty key-groups.
+ *
+ * @param startKeyGroup start of the range (inclusive)
+ * @param endKeyGroup end of the range (inclusive)
+ * @return the key-group from start to end or an empty key-group range.
+ */
+ public static KeyGroupRange of(int startKeyGroup, int endKeyGroup) {
+ return startKeyGroup <= endKeyGroup ? new KeyGroupRange(startKeyGroup, endKeyGroup) : EMPTY_KEY_GROUP;
+ }
+
+ /**
+ * Computes the range of key-groups that are assigned to a given operator under the given parallelism and maximum
+ * parallelism.
+ *
+ * IMPORTANT: maxParallelism must be <= Short.MAX_VALUE to avoid rounding problems in this method. If we ever want
+ * to go beyond this boundary, this method must perform arithmetic on long values.
+ *
+ * @param maxParallelism Maximal parallelism that the job was initially created with.
+ * @param parallelism The current parallelism under which the job runs. Must be <= maxParallelism.
+ * @param operatorIndex Id of a key-group. 0 <= keyGroupID < maxParallelism.
+ * @return
+ */
+ public static KeyGroupRange computeKeyGroupRangeForOperatorIndex(
+ int maxParallelism,
+ int parallelism,
+ int operatorIndex) {
+
+ int start = operatorIndex == 0 ? 0 : ((operatorIndex * maxParallelism - 1) / parallelism) + 1;
+ int end = ((operatorIndex + 1) * maxParallelism - 1) / parallelism;
+ return new KeyGroupRange(start, end);
+ }
+
+ /**
+ * Computes the index of the operator to which a key-group belongs under the given parallelism and maximum
+ * parallelism.
+ *
+ * IMPORTANT: maxParallelism must be <= Short.MAX_VALUE to avoid rounding problems in this method. If we ever want
+ * to go beyond this boundary, this method must perform arithmetic on long values.
+ *
+ * @param maxParallelism Maximal parallelism that the job was initially created with.
+ * 0 < parallelism <= maxParallelism <= Short.MAX_VALUE must hold.
+ * @param parallelism The current parallelism under which the job runs. Must be <= maxParallelism.
+ * @param keyGroupId Id of a key-group. 0 <= keyGroupID < maxParallelism.
+ * @return The index of the operator to which elements from the given key-group should be routed under the given
+ * parallelism and maxParallelism.
+ */
+ public static final int computeOperatorIndexForKeyGroup(int maxParallelism, int parallelism, int keyGroupId) {
+ return keyGroupId * parallelism / maxParallelism;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRangeOffsets.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRangeOffsets.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRangeOffsets.java
new file mode 100644
index 0000000..4f0a82b
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRangeOffsets.java
@@ -0,0 +1,203 @@
+/*
+ * 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.java.tuple.Tuple2;
+import org.apache.flink.util.Preconditions;
+
+import java.io.Serializable;
+import java.util.Arrays;
+import java.util.Iterator;
+
+
+/**
+ * This class combines a key-group range with offsets that correspond to the key-groups in the range.
+ */
+public class KeyGroupRangeOffsets implements Iterable<Tuple2<Integer, Long>> , Serializable {
+
+ /** the range of key-groups */
+ private final KeyGroupRange keyGroupRange;
+
+ /** the aligned array of offsets for the key-groups */
+ private final long[] offsets;
+
+ /**
+ * Creates key-group range with offsets from the given key-group range. The order of given offsets must be aligned
+ * with respect to the key-groups in the range.
+ *
+ * @param keyGroupRange The range of key-groups.
+ * @param offsets The aligned array of offsets for the given key-groups.
+ */
+ public KeyGroupRangeOffsets(KeyGroupRange keyGroupRange, long[] offsets) {
+ this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange);
+ this.offsets = Preconditions.checkNotNull(offsets);
+ Preconditions.checkArgument(offsets.length == keyGroupRange.getNumberOfKeyGroups());
+ }
+
+ /**
+ * Creates key-group range with offsets from the given start key-group to end key-group. The order of given offsets
+ * must be aligned with respect to the key-groups in the range.
+ *
+ * @param rangeStart Start key-group of the range (inclusive)
+ * @param rangeEnd End key-group of the range (inclusive)
+ * @param offsets The aligned array of offsets for the given key-groups.
+ */
+ public KeyGroupRangeOffsets(int rangeStart, int rangeEnd, long[] offsets) {
+ this(KeyGroupRange.of(rangeStart, rangeEnd), offsets);
+ }
+
+ /**
+ * Creates key-group range with offsets from the given start key-group to end key-group.
+ * All offsets are initially zero.
+ *
+ * @param rangeStart Start key-group of the range (inclusive)
+ * @param rangeEnd End key-group of the range (inclusive)
+ */
+ public KeyGroupRangeOffsets(int rangeStart, int rangeEnd) {
+ this(KeyGroupRange.of(rangeStart, rangeEnd));
+ }
+
+ /**
+ * Creates key-group range with offsets for the given key-group range, where all offsets are initially zero.
+ *
+ * @param keyGroupRange The range of key-groups.
+ */
+ public KeyGroupRangeOffsets(KeyGroupRange keyGroupRange) {
+ this(keyGroupRange, new long[keyGroupRange.getNumberOfKeyGroups()]);
+ }
+
+ /**
+ * Returns the offset for the given key-group. The key-group must be contained in the range.
+ *
+ * @param keyGroup Key-group for which we query the offset. Key-group must be contained in the range.
+ * @return The offset for the given key-group which must be contained in the range.
+ */
+ public long getKeyGroupOffset(int keyGroup) {
+ return offsets[computeKeyGroupIndex(keyGroup)];
+ }
+
+ /**
+ * Sets the offset for the given key-group. The key-group must be contained in the range.
+ *
+ * @param keyGroup Key-group for which we set the offset. Must be contained in the range.
+ * @param offset Offset for the key-group.
+ */
+ public void setKeyGroupOffset(int keyGroup, long offset) {
+ offsets[computeKeyGroupIndex(keyGroup)] = offset;
+ }
+
+ /**
+ * Returns a key-group range with offsets which is the intersection of the internal key-group range with the given
+ * key-group range.
+ *
+ * @param keyGroupRange Key-group range to intersect with the internal key-group range.
+ * @return The key-group range with offsets for the intersection of the internal key-group range with the given
+ * key-group range.
+ */
+ public KeyGroupRangeOffsets getIntersection(KeyGroupRange keyGroupRange) {
+ Preconditions.checkNotNull(keyGroupRange);
+ KeyGroupRange intersection = this.keyGroupRange.getIntersection(keyGroupRange);
+ long[] subOffsets = new long[intersection.getNumberOfKeyGroups()];
+ if(subOffsets.length > 0) {
+ System.arraycopy(
+ offsets,
+ computeKeyGroupIndex(intersection.getStartKeyGroup()),
+ subOffsets,
+ 0,
+ subOffsets.length);
+ }
+ return new KeyGroupRangeOffsets(intersection, subOffsets);
+ }
+
+ public KeyGroupRange getKeyGroupRange() {
+ return keyGroupRange;
+ }
+
+ @Override
+ public Iterator<Tuple2<Integer, Long>> iterator() {
+ return new KeyGroupOffsetsIterator();
+ }
+
+ private int computeKeyGroupIndex(int keyGroup) {
+ return keyGroup - keyGroupRange.getStartKeyGroup();
+ }
+
+ /**
+ * Iterator for the Key-group/Offset pairs.
+ */
+ private final class KeyGroupOffsetsIterator implements Iterator<Tuple2<Integer, Long>> {
+
+ public KeyGroupOffsetsIterator() {
+ this.keyGroupIterator = keyGroupRange.iterator();
+ }
+
+ private final Iterator<Integer> keyGroupIterator;
+
+ @Override
+ public boolean hasNext() {
+ return keyGroupIterator.hasNext();
+ }
+
+ @Override
+ public Tuple2<Integer, Long> next() {
+ Integer currentKeyGroup = keyGroupIterator.next();
+ Tuple2<Integer,Long> result = new Tuple2<>(
+ currentKeyGroup,
+ offsets[currentKeyGroup - keyGroupRange.getStartKeyGroup()]);
+ return result;
+ }
+
+ @Override
+ public void remove() {
+ throw new UnsupportedOperationException("Unsupported by this iterator!");
+ }
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof KeyGroupRangeOffsets)) {
+ return false;
+ }
+
+ KeyGroupRangeOffsets that = (KeyGroupRangeOffsets) o;
+
+ if (keyGroupRange != null ? !keyGroupRange.equals(that.keyGroupRange) : that.keyGroupRange != null) {
+ return false;
+ }
+ return Arrays.equals(offsets, that.offsets);
+ }
+
+ @Override
+ public int hashCode() {
+ int result = keyGroupRange != null ? keyGroupRange.hashCode() : 0;
+ result = 31 * result + Arrays.hashCode(offsets);
+ return result;
+ }
+
+ @Override
+ public String toString() {
+ return "KeyGroupRangeOffsets{" +
+ "keyGroupRange=" + keyGroupRange +
+ ", offsets=" + Arrays.toString(offsets) +
+ '}';
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupsStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupsStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupsStateHandle.java
new file mode 100644
index 0000000..0a36f92
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupsStateHandle.java
@@ -0,0 +1,163 @@
+/*
+ * 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.IOException;
+
+/**
+ * A handle to the partitioned stream operator state after it has been checkpointed. This state
+ * consists of a range of key group snapshots. A key group is subset of the available
+ * key space. The key groups are identified by their key group indices.
+ */
+public class KeyGroupsStateHandle implements StateObject {
+
+ private static final long serialVersionUID = -8070326169926626355L;
+
+ /** Range of key-groups with their respective offsets in the stream state */
+ private final KeyGroupRangeOffsets groupRangeOffsets;
+
+ /** Inner stream handle to the actual states of the key-groups in the range */
+ private final StreamStateHandle stateHandle;
+
+ /**
+ *
+ * @param groupRangeOffsets range of key-group ids that in the state of this handle
+ * @param streamStateHandle handle to the actual state of the key-groups
+ */
+ public KeyGroupsStateHandle(KeyGroupRangeOffsets groupRangeOffsets, StreamStateHandle streamStateHandle) {
+ Preconditions.checkNotNull(groupRangeOffsets);
+ Preconditions.checkNotNull(streamStateHandle);
+
+ this.groupRangeOffsets = groupRangeOffsets;
+ this.stateHandle = streamStateHandle;
+ }
+
+ /**
+ *
+ * @return iterable over the key-group range for the key-group state referenced by this handle
+ */
+ public Iterable<Integer> keyGroups() {
+ return groupRangeOffsets.getKeyGroupRange();
+ }
+
+
+ /**
+ *
+ * @param keyGroupId the id of a key-group
+ * @return true if the provided key-group id is contained in the key-group range of this handle
+ */
+ public boolean containsKeyGroup(int keyGroupId) {
+ return groupRangeOffsets.getKeyGroupRange().contains(keyGroupId);
+ }
+
+ /**
+ *
+ * @param keyGroupId the id of a key-group. the id must be contained in the range of this handle.
+ * @return offset to the position of data for the provided key-group in the stream referenced by this state handle
+ */
+ public long getOffsetForKeyGroup(int keyGroupId) {
+ return groupRangeOffsets.getKeyGroupOffset(keyGroupId);
+ }
+
+ /**
+ *
+ * @param keyGroupRange a key group range to intersect.
+ * @return key-group state over a range that is the intersection between this handle's key-group range and the
+ * provided key-group range.
+ */
+ public KeyGroupsStateHandle getKeyGroupIntersection(KeyGroupRange keyGroupRange) {
+ return new KeyGroupsStateHandle(groupRangeOffsets.getIntersection(keyGroupRange), stateHandle);
+ }
+
+ /**
+ *
+ * @return the internal key-group range to offsets metadata
+ */
+ public KeyGroupRangeOffsets getGroupRangeOffsets() {
+ return groupRangeOffsets;
+ }
+
+ /**
+ *
+ * @return number of key-groups in the key-group range of this handle
+ */
+ public int getNumberOfKeyGroups() {
+ return groupRangeOffsets.getKeyGroupRange().getNumberOfKeyGroups();
+ }
+
+ /**
+ *
+ * @return the inner stream state handle to the actual key-group states
+ */
+ public StreamStateHandle getStateHandle() {
+ return stateHandle;
+ }
+
+ @Override
+ public void discardState() throws Exception {
+ stateHandle.discardState();
+ }
+
+ @Override
+ public long getStateSize() throws Exception {
+ return stateHandle.getStateSize();
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+
+ if (!(o instanceof KeyGroupsStateHandle)) {
+ return false;
+ }
+
+ KeyGroupsStateHandle that = (KeyGroupsStateHandle) o;
+
+ if (!groupRangeOffsets.equals(that.groupRangeOffsets)) {
+ return false;
+ }
+ return stateHandle.equals(that.stateHandle);
+
+ }
+
+ @Override
+ public int hashCode() {
+ int result = groupRangeOffsets.hashCode();
+ result = 31 * result + stateHandle.hashCode();
+ return result;
+ }
+
+ @Override
+ public String toString() {
+ return "KeyGroupsStateHandle{" +
+ "groupRangeOffsets=" + groupRangeOffsets +
+ ", data=" + stateHandle +
+ '}';
+ }
+
+ @Override
+ public void close() throws IOException {
+ stateHandle.close();
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalStateHandle.java
deleted file mode 100644
index 4e7531f..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/LocalStateHandle.java
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * 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 java.io.IOException;
-import java.io.Serializable;
-
-/**
- * A StateHandle that includes the operator states directly.
- */
-public class LocalStateHandle<T extends Serializable> implements StateHandle<T> {
-
- private static final long serialVersionUID = 2093619217898039610L;
-
- private final T state;
-
- public LocalStateHandle(T state) {
- this.state = state;
- }
-
- @Override
- public T getState(ClassLoader userCodeClassLoader) {
- // The object has been deserialized correctly before
- return state;
- }
-
- @Override
- public void discardState() {}
-
- @Override
- public long getStateSize() {
- return 0;
- }
-
- @Override
- public void close() throws IOException {}
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStateHandle.java
new file mode 100644
index 0000000..d547624
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStateHandle.java
@@ -0,0 +1,32 @@
+/*
+ * 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 java.io.Serializable;
+
+/**
+ * Handle to state that can be read back again via {@link #retrieveState()}.
+ */
+public interface RetrievableStateHandle<T extends Serializable> extends StateObject {
+
+ /**
+ * Retrieves the object that was previously written to state.
+ */
+ T retrieveState() throws Exception;
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStreamStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStreamStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStreamStateHandle.java
new file mode 100644
index 0000000..e3538af
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStreamStateHandle.java
@@ -0,0 +1,82 @@
+/*
+ * 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.core.fs.Path;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.filesystem.FileStateHandle;
+import org.apache.flink.util.InstantiationUtil;
+import org.apache.flink.util.Preconditions;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.io.Serializable;
+
+/**
+ * Wrapper around a {@link StreamStateHandle} to make the referenced state object retrievable trough a simple get call.
+ * This implementation expects that the object was serialized through default serialization of Java's
+ * {@link java.io.ObjectOutputStream}.
+ *
+ * @param <T> type of the retrievable object which is stored under the wrapped stream handle
+ */
+public class RetrievableStreamStateHandle<T extends Serializable> implements
+ StreamStateHandle, RetrievableStateHandle<T>, Closeable {
+
+ private static final long serialVersionUID = 314567453677355L;
+ /** wrapped inner stream state handle from which we deserialize on retrieval */
+ private final StreamStateHandle wrappedStreamStateHandle;
+
+ public RetrievableStreamStateHandle(StreamStateHandle streamStateHandle) {
+ this.wrappedStreamStateHandle = Preconditions.checkNotNull(streamStateHandle);
+ }
+
+ public RetrievableStreamStateHandle(Path filePath) {
+ Preconditions.checkNotNull(filePath);
+ this.wrappedStreamStateHandle = new FileStateHandle(filePath);
+ }
+
+ @Override
+ public T retrieveState() throws Exception {
+ try (FSDataInputStream in = openInputStream()) {
+ return InstantiationUtil.deserializeObject(in);
+ }
+ }
+
+ @Override
+ public FSDataInputStream openInputStream() throws Exception {
+ return wrappedStreamStateHandle.openInputStream();
+ }
+
+ @Override
+ public void discardState() throws Exception {
+ wrappedStreamStateHandle.discardState();
+ }
+
+ @Override
+ public long getStateSize() throws Exception {
+ return wrappedStreamStateHandle.getStateSize();
+ }
+
+ @Override
+ public void close() throws IOException {
+ wrappedStreamStateHandle.close();
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateBackendFactory.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateBackendFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateBackendFactory.java
index f17eb6e..39e7ed2 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateBackendFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateBackendFactory.java
@@ -20,13 +20,15 @@ package org.apache.flink.runtime.state;
import org.apache.flink.configuration.Configuration;
+import java.io.Serializable;
+
/**
* A factory to create a specific state backend. The state backend creation gets a Configuration
* object that can be used to read further config values.
*
* @param <T> The type of the state backend created.
*/
-public interface StateBackendFactory<T extends AbstractStateBackend> {
+public interface StateBackendFactory<T extends AbstractStateBackend> extends Serializable {
/**
* Creates the state backend, optionally using the given configuration.
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateUtil.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateUtil.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateUtil.java
new file mode 100644
index 0000000..3c5157e
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateUtil.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 java.io.IOException;
+
+/**
+ * Helpers for {@link StateObject} related code.
+ */
+public class StateUtil {
+
+ private StateUtil() {
+ throw new AssertionError();
+ }
+
+ /**
+ * Iterates through the passed state handles and calls discardState() on each handle that is not null. All
+ * occurring exceptions are suppressed and collected until the iteration is over and emitted as a single exception.
+ *
+ * @param handlesToDiscard State handles to discard. Passed iterable is allowed to deliver null values.
+ * @throws Exception exception that is a collection of all suppressed exceptions that were caught during iteration
+ */
+ public static void bestEffortDiscardAllStateObjects(
+ Iterable<? extends StateObject> handlesToDiscard) throws Exception {
+
+ if (handlesToDiscard != null) {
+
+ Exception suppressedExceptions = null;
+
+ for (StateObject state : handlesToDiscard) {
+
+ if (state != null) {
+ try {
+ state.discardState();
+ } catch (Exception ex) {
+ //best effort to still cleanup other states and deliver exceptions in the end
+ if (suppressedExceptions == null) {
+ suppressedExceptions = new Exception(ex);
+ }
+ suppressedExceptions.addSuppressed(ex);
+ }
+ }
+ }
+
+ if (suppressedExceptions != null) {
+ throw suppressedExceptions;
+ }
+ }
+ }
+
+ /**
+ * Iterates through the passed state handles and calls discardState() on each handle that is not null. All
+ * occurring exceptions are suppressed and collected until the iteration is over and emitted as a single exception.
+ *
+ * @param handlesToDiscard State handles to discard. Passed iterable is allowed to deliver null values.
+ * @throws Exception exception that is a collection of all suppressed exceptions that were caught during iteration
+ */
+ public static void bestEffortCloseAllStateObjects(
+ Iterable<? extends StateObject> handlesToDiscard) throws IOException {
+
+ if (handlesToDiscard != null) {
+
+ IOException suppressedExceptions = null;
+
+ for (StateObject state : handlesToDiscard) {
+
+ if (state != null) {
+ try {
+ state.close();
+ } catch (Exception ex) {
+ //best effort to still cleanup other states and deliver exceptions in the end
+ if (suppressedExceptions == null) {
+ suppressedExceptions = new IOException(ex);
+ }
+ suppressedExceptions.addSuppressed(ex);
+ }
+ }
+ }
+
+ if (suppressedExceptions != null) {
+ throw suppressedExceptions;
+ }
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateUtils.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateUtils.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateUtils.java
deleted file mode 100644
index b130c70..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateUtils.java
+++ /dev/null
@@ -1,59 +0,0 @@
-/*
- * 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.jobgraph.tasks.StatefulTask;
-
-/**
- * A collection of utility methods for dealing with operator state.
- */
-public class StateUtils {
-
- /**
- * Utility method to define a common generic bound to be used for setting a
- * generic state handle on a generic state carrier.
- *
- * This has no impact on runtime, since internally, it performs unchecked
- * casts. The purpose is merely to allow the use of generic interfaces
- * without resorting to raw types, by giving the compiler a common type
- * bound.
- *
- * @param op
- * The state carrier operator.
- * @param state
- * The state handle.
- * @param <T>
- * Type bound for the
- */
- public static <T extends StateHandle<?>> void setOperatorState(StatefulTask<?> op, StateHandle<?> state)
- throws Exception {
-
- @SuppressWarnings("unchecked")
- StatefulTask<T> typedOp = (StatefulTask<T>) op;
- @SuppressWarnings("unchecked")
- T typedHandle = (T) state;
-
- typedOp.setInitialState(typedHandle);
- }
-
- // ------------------------------------------------------------------------
-
- /** Do not instantiate */
- private StateUtils() {}
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java
index 891243b..46e4299 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java
@@ -7,7 +7,7 @@
* "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
+ * 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,
@@ -18,19 +18,17 @@
package org.apache.flink.runtime.state;
-import java.io.InputStream;
-import java.io.Serializable;
+import org.apache.flink.core.fs.FSDataInputStream;
/**
- * A state handle that produces an input stream when resolved.
+ * A {@link StateObject} that represents state that was written to a stream. The data can be read
+ * back via {@link #openInputStream()}.
*/
-public interface StreamStateHandle extends StateHandle<InputStream> {
+public interface StreamStateHandle extends StateObject {
/**
- * Converts this stream state handle into a state handle that de-serializes
- * the stream into an object using Java's serialization mechanism.
- *
- * @return The state handle that automatically de-serializes.
+ * Returns an {@link FSDataInputStream} that can be used to read back the data that
+ * was previously written to the stream.
*/
- <T extends Serializable> StateHandle<T> toSerializableHandle();
+ FSDataInputStream openInputStream() throws Exception;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFileStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFileStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFileStateHandle.java
deleted file mode 100644
index 0585062..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFileStateHandle.java
+++ /dev/null
@@ -1,98 +0,0 @@
-/*
- * 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.FileSystem;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.runtime.state.AbstractCloseableHandle;
-import org.apache.flink.runtime.state.StateObject;
-
-import java.io.IOException;
-
-import static org.apache.flink.util.Preconditions.checkNotNull;
-
-/**
- * Base class for state that is stored in a file.
- */
-public abstract class AbstractFileStateHandle extends AbstractCloseableHandle implements StateObject {
-
- private static final long serialVersionUID = 350284443258002355L;
-
- /** The path to the file in the filesystem, fully describing the file system */
- private final Path filePath;
-
- /** Cached file system handle */
- private transient FileSystem fs;
-
- /**
- * Creates a new file state for the given file path.
- *
- * @param filePath The path to the file that stores the state.
- */
- protected AbstractFileStateHandle(Path filePath) {
- this.filePath = checkNotNull(filePath);
- }
-
- /**
- * Gets the path where this handle's state is stored.
- * @return The path where this handle's state is stored.
- */
- public Path getFilePath() {
- return filePath;
- }
-
- /**
- * Discard the state by deleting the file that stores the state. If the parent directory
- * of the state is empty after deleting the state file, it is also deleted.
- *
- * @throws Exception Thrown, if the file deletion (not the directory deletion) fails.
- */
- @Override
- public void discardState() throws Exception {
- getFileSystem().delete(filePath, false);
-
- // send a call to delete the checkpoint directory containing the file. This will
- // fail (and be ignored) when some files still exist
- try {
- getFileSystem().delete(filePath.getParent(), false);
- } catch (IOException ignored) {}
- }
-
- /**
- * Gets the file system that stores the file state.
- * @return The file system that stores the file state.
- * @throws IOException Thrown if the file system cannot be accessed.
- */
- protected FileSystem getFileSystem() throws IOException {
- if (fs == null) {
- fs = FileSystem.get(filePath.toUri());
- }
- return fs;
- }
-
- /**
- * Returns the file size in bytes.
- *
- * @return The file size in bytes.
- * @throws IOException Thrown if the file system cannot be accessed.
- */
- protected long getFileSize() throws IOException {
- return getFileSystem().getFileStatus(filePath).getLen();
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsStateSnapshot.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsStateSnapshot.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsStateSnapshot.java
index 0692541..51e8b5a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsStateSnapshot.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsStateSnapshot.java
@@ -38,8 +38,9 @@ import java.util.Map;
* @param <N> The type of the namespace in the snapshot state.
* @param <SV> The type of the state value.
*/
-public abstract class AbstractFsStateSnapshot<K, N, SV, S extends State, SD extends StateDescriptor<S, ?>>
- extends AbstractFileStateHandle implements KvStateSnapshot<K, N, S, SD, FsStateBackend> {
+public abstract class AbstractFsStateSnapshot<K, N, SV, S extends State, SD extends StateDescriptor<S, ?>>
+ extends FileStateHandle
+ implements KvStateSnapshot<K, N, S, SD, FsStateBackend> {
private static final long serialVersionUID = 1L;
@@ -132,7 +133,7 @@ public abstract class AbstractFsStateSnapshot<K, N, SV, S extends State, SD exte
* @throws IOException Thrown if the file system cannot be accessed.
*/
@Override
- public long getStateSize() throws IOException {
- return getFileSize();
+ public void discardState() throws Exception {
+ super.discardState();
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileSerializableStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileSerializableStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileSerializableStateHandle.java
deleted file mode 100644
index 34a1cb0..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileSerializableStateHandle.java
+++ /dev/null
@@ -1,72 +0,0 @@
-/*
- * 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.FSDataInputStream;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.InstantiationUtil;
-
-import java.io.IOException;
-import java.io.ObjectInputStream;
-import java.io.Serializable;
-
-/**
- * A state handle that points to state stored in a file via Java Serialization.
- *
- * @param <T> The type of state pointed to by the state handle.
- */
-public class FileSerializableStateHandle<T extends Serializable> extends AbstractFileStateHandle implements StateHandle<T> {
-
- private static final long serialVersionUID = -657631394290213622L;
-
- /**
- * Creates a new FileSerializableStateHandle pointing to state at the given file path.
- *
- * @param filePath The path to the file containing the checkpointed state.
- */
- public FileSerializableStateHandle(Path filePath) {
- super(filePath);
- }
-
- @Override
- @SuppressWarnings("unchecked")
- public T getState(ClassLoader classLoader) throws Exception {
- ensureNotClosed();
-
- try (FSDataInputStream inStream = getFileSystem().open(getFilePath())) {
- // make sure any deserialization can be aborted
- registerCloseable(inStream);
-
- ObjectInputStream ois = new InstantiationUtil.ClassLoaderObjectInputStream(inStream, classLoader);
- return (T) ois.readObject();
- }
- }
-
- /**
- * Returns the file size in bytes.
- *
- * @return The file size in bytes.
- * @throws IOException Thrown if the file system cannot be accessed.
- */
- @Override
- public long getStateSize() throws IOException {
- return getFileSize();
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStateHandle.java
new file mode 100644
index 0000000..871e56c
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStateHandle.java
@@ -0,0 +1,136 @@
+/*
+ * 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.FSDataInputStream;
+import org.apache.flink.core.fs.FileSystem;
+import org.apache.flink.core.fs.Path;
+import org.apache.flink.runtime.state.AbstractCloseableHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
+import java.io.IOException;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * {@link StreamStateHandle} for state that was written to a file stream. The written data is
+ * identifier by the file path. The state can be read again by calling {@link #openInputStream()}.
+ */
+public class FileStateHandle extends AbstractCloseableHandle implements StreamStateHandle {
+
+ private static final long serialVersionUID = 350284443258002355L;
+
+ /**
+ * The path to the file in the filesystem, fully describing the file system
+ */
+ private final Path filePath;
+
+ /**
+ * Cached file system handle
+ */
+ private transient FileSystem fs;
+
+ /**
+ * Creates a new file state for the given file path.
+ *
+ * @param filePath The path to the file that stores the state.
+ */
+ public FileStateHandle(Path filePath) {
+ this.filePath = requireNonNull(filePath);
+ }
+
+ /**
+ * Gets the path where this handle's state is stored.
+ *
+ * @return The path where this handle's state is stored.
+ */
+ public Path getFilePath() {
+ return filePath;
+ }
+
+ @Override
+ public FSDataInputStream openInputStream() throws Exception {
+ ensureNotClosed();
+ FSDataInputStream inputStream = getFileSystem().open(filePath);
+ registerCloseable(inputStream);
+ return inputStream;
+ }
+
+ /**
+ * Discard the state by deleting the file that stores the state. If the parent directory
+ * of the state is empty after deleting the state file, it is also deleted.
+ *
+ * @throws Exception Thrown, if the file deletion (not the directory deletion) fails.
+ */
+ @Override
+ public void discardState() throws Exception {
+ getFileSystem().delete(filePath, false);
+
+ // send a call to delete the checkpoint directory containing the file. This will
+ // fail (and be ignored) when some files still exist
+ try {
+ getFileSystem().delete(filePath.getParent(), false);
+ } catch (IOException ignored) {
+ }
+ }
+
+ /**
+ * Returns the file size in bytes.
+ *
+ * @return The file size in bytes.
+ * @throws IOException Thrown if the file system cannot be accessed.
+ */
+ @Override
+ public long getStateSize() throws IOException {
+ return getFileSystem().getFileStatus(filePath).getLen();
+ }
+
+ /**
+ * Gets the file system that stores the file state.
+ *
+ * @return The file system that stores the file state.
+ * @throws IOException Thrown if the file system cannot be accessed.
+ */
+ private FileSystem getFileSystem() throws IOException {
+ if (fs == null) {
+ fs = FileSystem.get(filePath.toUri());
+ }
+ return fs;
+ }
+
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof FileStateHandle)) {
+ return false;
+ }
+
+ FileStateHandle that = (FileStateHandle) o;
+ return filePath.equals(that.filePath);
+
+ }
+
+ @Override
+ public int hashCode() {
+ return filePath.hashCode();
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStreamStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStreamStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStreamStateHandle.java
deleted file mode 100644
index 5bfb4ee..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStreamStateHandle.java
+++ /dev/null
@@ -1,83 +0,0 @@
-/*
- * 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.Path;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.state.StreamStateHandle;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.Serializable;
-
-/**
- * A state handle that points to state in a file system, accessible as an input stream.
- */
-public class FileStreamStateHandle extends AbstractFileStateHandle implements StreamStateHandle {
-
- private static final long serialVersionUID = -6826990484549987311L;
-
- /**
- * Creates a new FileStreamStateHandle pointing to state at the given file path.
- *
- * @param filePath The path to the file containing the checkpointed state.
- */
- public FileStreamStateHandle(Path filePath) {
- super(filePath);
- }
-
- @Override
- public InputStream getState(ClassLoader userCodeClassLoader) throws Exception {
- ensureNotClosed();
-
- InputStream inStream = getFileSystem().open(getFilePath());
- // make sure the state handle is cancelable
- registerCloseable(inStream);
-
- return inStream;
- }
-
- /**
- * Returns the file size in bytes.
- *
- * @return The file size in bytes.
- * @throws IOException Thrown if the file system cannot be accessed.
- */
- @Override
- public long getStateSize() throws IOException {
- return getFileSize();
- }
-
- @Override
- public <T extends Serializable> StateHandle<T> toSerializableHandle() {
- FileSerializableStateHandle<T> handle = new FileSerializableStateHandle<>(getFilePath());
-
- // forward closed status
- if (isClosed()) {
- try {
- handle.close();
- } catch (IOException e) {
- // should not happen on a fresh handle, but forward anyways
- throw new RuntimeException(e);
- }
- }
-
- return handle;
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
index 61cf741..a3f4682 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
@@ -32,15 +32,12 @@ import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.core.fs.Path;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StateHandle;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
-import java.io.ObjectOutputStream;
-import java.io.Serializable;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Arrays;
@@ -294,24 +291,6 @@ public class FsStateBackend extends AbstractStateBackend {
}
@Override
- public <S extends Serializable> StateHandle<S> checkpointStateSerializable(
- S state, long checkpointID, long timestamp) throws Exception
- {
- checkFileSystemInitialized();
-
- Path checkpointDir = createCheckpointDirPath(checkpointID);
- int bufferSize = Math.max(DEFAULT_WRITE_BUFFER_SIZE, fileStateThreshold);
-
- FsCheckpointStateOutputStream stream =
- new FsCheckpointStateOutputStream(checkpointDir, filesystem, bufferSize, fileStateThreshold);
-
- try (ObjectOutputStream os = new ObjectOutputStream(stream)) {
- os.writeObject(state);
- return stream.closeAndGetHandle().toSerializableHandle();
- }
- }
-
- @Override
public FsCheckpointStateOutputStream createCheckpointStateOutputStream(long checkpointID, long timestamp) throws Exception {
checkFileSystemInitialized();
@@ -520,6 +499,11 @@ public class FsStateBackend extends AbstractStateBackend {
}
}
+ @Override
+ public void sync() throws IOException {
+ outStream.sync();
+ }
+
/**
* If the stream is only closed, we remove the produced file (cleanup through the auto close
* feature, for example). This method throws no exception if the deletion fails, but only
@@ -559,7 +543,7 @@ public class FsStateBackend extends AbstractStateBackend {
flush();
outStream.close();
closed = true;
- return new FileStreamStateHandle(statePath);
+ return new FileStateHandle(statePath);
}
}
else {
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java
index ba6de42..a42bec2 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java
@@ -18,28 +18,32 @@
package org.apache.flink.runtime.state.memory;
+import org.apache.flink.core.fs.FSDataInputStream;
import org.apache.flink.runtime.state.AbstractCloseableHandle;
-import org.apache.flink.runtime.state.StateHandle;
import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.util.InstantiationUtil;
+
+import org.apache.flink.util.Preconditions;
-import java.io.ByteArrayInputStream;
import java.io.IOException;
-import java.io.InputStream;
import java.io.Serializable;
+import java.util.Arrays;
/**
* A state handle that contains stream state in a byte array.
*/
-public final class ByteStreamStateHandle extends AbstractCloseableHandle implements StreamStateHandle {
+public class ByteStreamStateHandle extends AbstractCloseableHandle implements StreamStateHandle {
private static final long serialVersionUID = -5280226231200217594L;
-
- /** the state data */
- private final byte[] data;
+
+ /**
+ * the state data
+ */
+ protected final byte[] data;
/**
* Creates a new ByteStreamStateHandle containing the given data.
- *
+ *
* @param data The state data.
*/
public ByteStreamStateHandle(byte[] data) {
@@ -47,17 +51,39 @@ public final class ByteStreamStateHandle extends AbstractCloseableHandle impleme
}
@Override
- public InputStream getState(ClassLoader userCodeClassLoader) throws Exception {
+ public FSDataInputStream openInputStream() throws Exception {
ensureNotClosed();
- ByteArrayInputStream stream = new ByteArrayInputStream(data);
- registerCloseable(stream);
+ FSDataInputStream inputStream = new FSDataInputStream() {
+ int index = 0;
+
+ @Override
+ public void seek(long desired) throws IOException {
+ Preconditions.checkArgument(desired >= 0 && desired < Integer.MAX_VALUE);
+ index = (int) desired;
+ }
+
+ @Override
+ public long getPos() throws IOException {
+ return index;
+ }
+
+ @Override
+ public int read() throws IOException {
+ return index < data.length ? data[index++] & 0xFF : -1;
+ }
+ };
+ registerCloseable(inputStream);
+ return inputStream;
+ }
- return stream;
+ public byte[] getData() {
+ return data;
}
@Override
- public void discardState() {}
+ public void discardState() {
+ }
@Override
public long getStateSize() {
@@ -65,19 +91,27 @@ public final class ByteStreamStateHandle extends AbstractCloseableHandle impleme
}
@Override
- public <T extends Serializable> StateHandle<T> toSerializableHandle() {
- SerializedStateHandle<T> serializableHandle = new SerializedStateHandle<T>(data);
-
- // forward the closed status
- if (isClosed()) {
- try {
- serializableHandle.close();
- } catch (IOException e) {
- // should not happen on a fresh handle, but forward anyways
- throw new RuntimeException(e);
- }
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
}
+ if (!(o instanceof ByteStreamStateHandle)) {
+ return false;
+ }
+
+ ByteStreamStateHandle that = (ByteStreamStateHandle) o;
+ return Arrays.equals(data, that.data);
+
+ }
+
+ @Override
+ public int hashCode() {
+ int result = super.hashCode();
+ result = 31 * result + Arrays.hashCode(data);
+ return result;
+ }
- return serializableHandle;
+ public static StreamStateHandle fromSerializable(Serializable value) throws IOException {
+ return new ByteStreamStateHandle(InstantiationUtil.serializeObject(value));
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
index 7b9d21b..af84394 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
@@ -27,13 +27,12 @@ import org.apache.flink.api.common.state.ReducingStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.runtime.state.StateHandle;
import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.runtime.state.StreamStateHandle;
+
import java.io.ByteArrayOutputStream;
import java.io.IOException;
-import java.io.Serializable;
/**
* A {@link AbstractStateBackend} that stores all its data and checkpoints in memory and has no
@@ -104,27 +103,6 @@ public class MemoryStateBackend extends AbstractStateBackend {
return new MemFoldingState<>(keySerializer, namespaceSerializer, stateDesc);
}
- /**
- * Serialized the given state into bytes using Java serialization and creates a state handle that
- * can re-create that state.
- *
- * @param state The state to checkpoint.
- * @param checkpointID The ID of the checkpoint.
- * @param timestamp The timestamp of the checkpoint.
- * @param <S> The type of the state.
- *
- * @return A state handle that contains the given state serialized as bytes.
- * @throws Exception Thrown, if the serialization fails.
- */
- @Override
- public <S extends Serializable> StateHandle<S> checkpointStateSerializable(
- S state, long checkpointID, long timestamp) throws Exception
- {
- SerializedStateHandle<S> handle = new SerializedStateHandle<>(state);
- checkSize(handle.getSizeOfSerializedState(), maxStateSize);
- return new SerializedStateHandle<S>(state);
- }
-
@Override
public CheckpointStateOutputStream createCheckpointStateOutputStream(
long checkpointID, long timestamp) throws Exception
@@ -177,6 +155,14 @@ public class MemoryStateBackend extends AbstractStateBackend {
os.write(b, off, len);
}
+ @Override
+ public void flush() throws IOException {
+ os.flush();
+ }
+
+ @Override
+ public void sync() throws IOException { }
+
// --------------------------------------------------------------------
@Override
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/SerializedStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/SerializedStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/SerializedStateHandle.java
deleted file mode 100644
index 4420470..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/SerializedStateHandle.java
+++ /dev/null
@@ -1,87 +0,0 @@
-/*
- * 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.memory;
-
-import org.apache.flink.runtime.state.AbstractCloseableHandle;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.InstantiationUtil;
-
-import java.io.IOException;
-import java.io.Serializable;
-
-/**
- * A state handle that represents its state in serialized form as bytes.
- *
- * @param <T> The type of state represented by this state handle.
- */
-public class SerializedStateHandle<T extends Serializable> extends AbstractCloseableHandle implements StateHandle<T> {
-
- private static final long serialVersionUID = 4145685722538475769L;
-
- /** The serialized data */
- private final byte[] serializedData;
-
- /**
- * Creates a new serialized state handle, eagerly serializing the given state object.
- *
- * @param value The state object.
- * @throws IOException Thrown, if the serialization fails.
- */
- public SerializedStateHandle(T value) throws IOException {
- this.serializedData = value == null ? null : InstantiationUtil.serializeObject(value);
- }
-
- /**
- * Creates a new serialized state handle, based in the given already serialized data.
- *
- * @param serializedData The serialized data.
- */
- public SerializedStateHandle(byte[] serializedData) {
- this.serializedData = serializedData;
- }
-
- @Override
- public T getState(ClassLoader classLoader) throws Exception {
- if (classLoader == null) {
- throw new NullPointerException();
- }
-
- ensureNotClosed();
- return serializedData == null ? null : InstantiationUtil.<T>deserializeObject(serializedData, classLoader);
- }
-
- /**
- * Gets the size of the serialized state.
- * @return The size of the serialized state.
- */
- public int getSizeOfSerializedState() {
- return serializedData.length;
- }
-
- /**
- * Discarding heap-memory backed state is a no-op, so this method does nothing.
- */
- @Override
- public void discardState() {}
-
- @Override
- public long getStateSize() {
- return serializedData.length;
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java
index 6958784..d54826a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/RuntimeEnvironment.java
@@ -23,7 +23,6 @@ import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.TaskInfo;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.fs.Path;
-import org.apache.flink.runtime.metrics.groups.TaskMetricGroup;
import org.apache.flink.runtime.accumulators.AccumulatorRegistry;
import org.apache.flink.runtime.broadcast.BroadcastVariableManager;
import org.apache.flink.runtime.execution.Environment;
@@ -36,10 +35,13 @@ import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.messages.checkpoint.AcknowledgeCheckpoint;
+import org.apache.flink.runtime.metrics.groups.TaskMetricGroup;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import java.util.List;
import java.util.Map;
import java.util.concurrent.Future;
@@ -240,39 +242,21 @@ public class RuntimeEnvironment implements Environment {
@Override
public void acknowledgeCheckpoint(long checkpointId) {
- acknowledgeCheckpoint(checkpointId, null);
+ acknowledgeCheckpoint(checkpointId, null, null);
}
@Override
- public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) {
- // try and create a serialized version of the state handle
- SerializedValue<StateHandle<?>> serializedState;
- long stateSize;
-
- if (state == null) {
- serializedState = null;
- stateSize = 0;
- } else {
- try {
- serializedState = new SerializedValue<StateHandle<?>>(state);
- } catch (Exception e) {
- throw new RuntimeException("Failed to serialize state handle during checkpoint confirmation", e);
- }
-
- try {
- stateSize = state.getStateSize();
- }
- catch (Exception e) {
- throw new RuntimeException("Failed to fetch state handle size", e);
- }
- }
-
+ public void acknowledgeCheckpoint(
+ long checkpointId,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
+ List<KeyGroupsStateHandle> keyGroupStateHandles) {
+
AcknowledgeCheckpoint message = new AcknowledgeCheckpoint(
jobId,
executionId,
checkpointId,
- serializedState,
- stateSize);
+ chainedStateHandle,
+ keyGroupStateHandles);
jobManager.tell(message);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
index c98d512..73601c4 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/taskmanager/Task.java
@@ -58,10 +58,11 @@ import org.apache.flink.runtime.messages.TaskMessages.TaskInFinalState;
import org.apache.flink.runtime.messages.TaskMessages.UpdateTaskExecutionState;
import org.apache.flink.runtime.messages.checkpoint.DeclineCheckpoint;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.state.StateUtils;
-import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.util.SerializedValue;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -223,9 +224,17 @@ public class Task implements Runnable {
/** Serial executor for asynchronous calls (checkpoints, etc), lazily initialized */
private volatile ExecutorService asyncCallDispatcher;
- /** The handle to the state that the operator was initialized with. Will be set to null after the
- * initialization, to be memory friendly */
- private volatile SerializedValue<StateHandle<?>> operatorState;
+ /**
+ * The handle to the chained operator state that the task was initialized with. Will be set
+ * to null after the initialization, to be memory friendly.
+ */
+ private volatile ChainedStateHandle<StreamStateHandle> chainedOperatorState;
+
+ /**
+ * The handle to the key group state that the task was initialized with. Will be set
+ * to null after the initialization, to be memory friendly.
+ */
+ private volatile List<KeyGroupsStateHandle> keyGroupStates;
/** Initialized from the Flink configuration. May also be set at the ExecutionConfig */
private long taskCancellationInterval;
@@ -257,8 +266,9 @@ public class Task implements Runnable {
this.requiredJarFiles = checkNotNull(tdd.getRequiredJarFiles());
this.requiredClasspaths = checkNotNull(tdd.getRequiredClasspaths());
this.nameOfInvokableClass = checkNotNull(tdd.getInvokableClassName());
- this.operatorState = tdd.getOperatorState();
+ this.chainedOperatorState = tdd.getOperatorState();
this.serializedExecutionConfig = checkNotNull(tdd.getSerializedExecutionConfig());
+ this.keyGroupStates = tdd.getKeyGroupState();
this.taskCancellationInterval = jobConfiguration.getLong(
ConfigConstants.TASK_CANCELLATION_INTERVAL_MILLIS,
@@ -538,21 +548,11 @@ public class Task implements Runnable {
// 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
- // get our private reference onto the stack (be safe against concurrent changes)
- SerializedValue<StateHandle<?>> operatorState = this.operatorState;
-
- if (operatorState != null) {
+ if (chainedOperatorState != null || keyGroupStates != null) {
if (invokable instanceof StatefulTask) {
- try {
- StateHandle<?> state = operatorState.deserializeValue(userCodeClassLoader);
- StatefulTask<?> op = (StatefulTask<?>) invokable;
- StateUtils.setOperatorState(op, state);
- }
- catch (Exception e) {
- throw new RuntimeException("Failed to deserialize state handle and setup initial operator state.", e);
- }
- }
- else {
+ StatefulTask op = (StatefulTask) invokable;
+ op.setInitialState(chainedOperatorState, keyGroupStates);
+ } else {
throw new IllegalStateException("Found operator state for a non-stateful task invokable");
}
}
@@ -560,8 +560,8 @@ public class Task implements Runnable {
// be memory and GC friendly - since the code stays in invoke() for a potentially long time,
// we clear the reference to the state handle
//noinspection UnusedAssignment
- operatorState = null;
- this.operatorState = null;
+ this.chainedOperatorState = null;
+ this.keyGroupStates = null;
// ----------------------------------------------------------------
// actual task core work
@@ -936,7 +936,7 @@ public class Task implements Runnable {
if (invokable instanceof StatefulTask) {
// build a local closure
- final StatefulTask<?> statefulTask = (StatefulTask<?>) invokable;
+ final StatefulTask statefulTask = (StatefulTask) invokable;
final String taskName = taskNameWithSubtask;
Runnable runnable = new Runnable() {
@@ -977,7 +977,7 @@ public class Task implements Runnable {
if (invokable instanceof StatefulTask) {
// build a local closure
- final StatefulTask<?> statefulTask = (StatefulTask<?>) invokable;
+ final StatefulTask statefulTask = (StatefulTask) invokable;
final String taskName = taskNameWithSubtask;
Runnable runnable = new Runnable() {
@@ -1192,7 +1192,6 @@ public class Task implements Runnable {
// reason, we spawn a separate thread that repeatedly interrupts the user code until
// it exits
while (executer.isAlive()) {
-
// build the stack trace of where the thread is stuck, for the log
StringBuilder bld = new StringBuilder();
StackTraceElement[] stack = executer.getStackTrace();
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/util/ZooKeeperUtils.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/util/ZooKeeperUtils.java b/flink-runtime/src/main/java/org/apache/flink/runtime/util/ZooKeeperUtils.java
index b585fe6..91db564 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/util/ZooKeeperUtils.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/util/ZooKeeperUtils.java
@@ -34,7 +34,7 @@ import org.apache.flink.runtime.jobmanager.SubmittedJobGraph;
import org.apache.flink.runtime.jobmanager.ZooKeeperSubmittedJobGraphStore;
import org.apache.flink.runtime.leaderelection.ZooKeeperLeaderElectionService;
import org.apache.flink.runtime.leaderretrieval.ZooKeeperLeaderRetrievalService;
-import org.apache.flink.runtime.zookeeper.StateStorageHelper;
+import org.apache.flink.runtime.zookeeper.RetrievableStateStorageHelper;
import org.apache.flink.runtime.zookeeper.filesystem.FileSystemStateStorageHelper;
import org.apache.flink.util.ConfigurationUtil;
import org.slf4j.Logger;
@@ -228,7 +228,7 @@ public class ZooKeeperUtils {
checkNotNull(configuration, "Configuration");
- StateStorageHelper<SubmittedJobGraph> stateStorage = createFileSystemStateStorage(configuration, "submittedJobGraph");
+ RetrievableStateStorageHelper<SubmittedJobGraph> stateStorage = createFileSystemStateStorage(configuration, "submittedJobGraph");
// ZooKeeper submitted jobs root dir
String zooKeeperSubmittedJobsPath = ConfigurationUtil.getStringWithDeprecatedKeys(
@@ -266,7 +266,7 @@ public class ZooKeeperUtils {
ConfigConstants.DEFAULT_ZOOKEEPER_CHECKPOINTS_PATH,
ConfigConstants.ZOOKEEPER_CHECKPOINTS_PATH);
- StateStorageHelper<CompletedCheckpoint> stateStorage = createFileSystemStateStorage(
+ RetrievableStateStorageHelper<CompletedCheckpoint> stateStorage = createFileSystemStateStorage(
configuration,
"completedCheckpoint");
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/RetrievableStateStorageHelper.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/RetrievableStateStorageHelper.java b/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/RetrievableStateStorageHelper.java
new file mode 100644
index 0000000..1434f74
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/RetrievableStateStorageHelper.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.zookeeper;
+
+import org.apache.flink.runtime.state.RetrievableStateHandle;
+
+import java.io.Serializable;
+
+/**
+ * State storage helper which is used by {@link ZooKeeperStateHandleStore} to persiste state before
+ * the state handle is written to ZooKeeper.
+ *
+ * @param <T> The type of the data that can be stored by this storage helper.
+ */
+public interface RetrievableStateStorageHelper<T extends Serializable> {
+
+ /**
+ * Stores the given state and returns a state handle to it.
+ *
+ * @param state State to be stored
+ * @return State handle to the stored state
+ * @throws Exception
+ */
+ RetrievableStateHandle<T> store(T state) throws Exception;
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/StateStorageHelper.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/StateStorageHelper.java b/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/StateStorageHelper.java
deleted file mode 100644
index 36fb849..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/zookeeper/StateStorageHelper.java
+++ /dev/null
@@ -1,41 +0,0 @@
-/*
- * 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.zookeeper;
-
-import org.apache.flink.runtime.state.StateHandle;
-
-import java.io.Serializable;
-
-/**
- * State storage helper which is used by {@link ZooKeeperStateHandleStore} to persiste state before
- * the state handle is written to ZooKeeper.
- *
- * @param <T>
- */
-public interface StateStorageHelper<T extends Serializable> {
-
- /**
- * Stores the given state and returns a state handle to it.
- *
- * @param state State to be stored
- * @return State handle to the stored state
- * @throws Exception
- */
- StateHandle<T> store(T state) throws Exception;
-}
[10/27] flink git commit: [FLINK-3761] Refactor State Backends/Make
Keyed State Key-Group Aware
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStreamFactory.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStreamFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStreamFactory.java
new file mode 100644
index 0000000..cc13a72
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStreamFactory.java
@@ -0,0 +1,313 @@
+/*
+ * 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.api.common.JobID;
+import org.apache.flink.core.fs.FSDataOutputStream;
+import org.apache.flink.core.fs.FileSystem;
+import org.apache.flink.core.fs.Path;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.net.URI;
+import java.util.Arrays;
+import java.util.UUID;
+
+/**
+ * {@link org.apache.flink.runtime.state.CheckpointStreamFactory} that produces streams that
+ * write to a {@link FileSystem}.
+ *
+ * <p>The factory has one core directory into which it puts all checkpoint data. Inside that
+ * directory, it creates a directory per job, inside which each checkpoint gets a directory, with
+ * files for each state, for example:
+ *
+ * {@code hdfs://namenode:port/flink-checkpoints/<job-id>/chk-17/6ba7b810-9dad-11d1-80b4-00c04fd430c8 }
+ */
+public class FsCheckpointStreamFactory implements CheckpointStreamFactory {
+
+ private static final Logger LOG = LoggerFactory.getLogger(FsCheckpointStreamFactory.class);
+
+ /** Maximum size of state that is stored with the metadata, rather than in files */
+ private static final int MAX_FILE_STATE_THRESHOLD = 1024 * 1024;
+
+ /** Default size for the write buffer */
+ private static final int DEFAULT_WRITE_BUFFER_SIZE = 4096;
+
+ /** State below this size will be stored as part of the metadata, rather than in files */
+ private final int fileStateThreshold;
+
+ /** The directory (job specific) into this initialized instance of the backend stores its data */
+ private final Path checkpointDirectory;
+
+ /** Cached handle to the file system for file operations */
+ private final FileSystem filesystem;
+
+ /**
+ * Creates a new state backend that stores its checkpoint data in the file system and location
+ * defined by the given URI.
+ *
+ * <p>A file system for the file system scheme in the URI (e.g., 'file://', 'hdfs://', or 'S3://')
+ * must be accessible via {@link FileSystem#get(URI)}.
+ *
+ * <p>For a state backend targeting HDFS, this means that the URI must either specify the authority
+ * (host and port), or that the Hadoop configuration that describes that information must be in the
+ * classpath.
+ *
+ * @param checkpointDataUri The URI describing the filesystem (scheme and optionally authority),
+ * and the path to the checkpoint data directory.
+ * @param fileStateSizeThreshold State up to this size will be stored as part of the metadata,
+ * rather than in files
+ *
+ * @throws IOException Thrown, if no file system can be found for the scheme in the URI.
+ */
+ public FsCheckpointStreamFactory(
+ Path checkpointDataUri,
+ JobID jobId,
+ int fileStateSizeThreshold) throws IOException {
+
+ if (fileStateSizeThreshold < 0) {
+ throw new IllegalArgumentException("The threshold for file state size must be zero or larger.");
+ }
+ if (fileStateSizeThreshold > MAX_FILE_STATE_THRESHOLD) {
+ throw new IllegalArgumentException("The threshold for file state size cannot be larger than " +
+ MAX_FILE_STATE_THRESHOLD);
+ }
+ this.fileStateThreshold = fileStateSizeThreshold;
+ Path basePath = checkpointDataUri;
+
+ Path dir = new Path(basePath, jobId.toString());
+
+ LOG.info("Initializing file stream factory to URI {}.", dir);
+
+ filesystem = basePath.getFileSystem();
+ filesystem.mkdirs(dir);
+
+ checkpointDirectory = dir;
+ }
+
+ @Override
+ public void close() throws Exception {}
+
+ @Override
+ public FsCheckpointStateOutputStream createCheckpointStateOutputStream(long checkpointID, long timestamp) throws Exception {
+ checkFileSystemInitialized();
+
+ Path checkpointDir = createCheckpointDirPath(checkpointID);
+ int bufferSize = Math.max(DEFAULT_WRITE_BUFFER_SIZE, fileStateThreshold);
+ return new FsCheckpointStateOutputStream(checkpointDir, filesystem, bufferSize, fileStateThreshold);
+ }
+
+ // ------------------------------------------------------------------------
+ // utilities
+ // ------------------------------------------------------------------------
+
+ private void checkFileSystemInitialized() throws IllegalStateException {
+ if (filesystem == null || checkpointDirectory == null) {
+ throw new IllegalStateException("filesystem has not been re-initialized after deserialization");
+ }
+ }
+
+ private Path createCheckpointDirPath(long checkpointID) {
+ return new Path(checkpointDirectory, "chk-" + checkpointID);
+ }
+
+ @Override
+ public String toString() {
+ return "File Stream Factory @ " + checkpointDirectory;
+ }
+
+ /**
+ * A {@link CheckpointStreamFactory.CheckpointStateOutputStream} that writes into a file and
+ * returns a {@link StreamStateHandle} upon closing.
+ */
+ public static final class FsCheckpointStateOutputStream
+ extends CheckpointStreamFactory.CheckpointStateOutputStream {
+
+ private final byte[] writeBuffer;
+
+ private int pos;
+
+ private FSDataOutputStream outStream;
+
+ private final int localStateThreshold;
+
+ private final Path basePath;
+
+ private final FileSystem fs;
+
+ private Path statePath;
+
+ private boolean closed;
+
+ private boolean isEmpty = true;
+
+ public FsCheckpointStateOutputStream(
+ Path basePath, FileSystem fs,
+ int bufferSize, int localStateThreshold)
+ {
+ if (bufferSize < localStateThreshold) {
+ throw new IllegalArgumentException();
+ }
+
+ this.basePath = basePath;
+ this.fs = fs;
+ this.writeBuffer = new byte[bufferSize];
+ this.localStateThreshold = localStateThreshold;
+ }
+
+
+ @Override
+ public void write(int b) throws IOException {
+ if (pos >= writeBuffer.length) {
+ flush();
+ }
+ writeBuffer[pos++] = (byte) b;
+
+ isEmpty = false;
+ }
+
+ @Override
+ public void write(byte[] b, int off, int len) throws IOException {
+ if (len < writeBuffer.length / 2) {
+ // copy it into our write buffer first
+ final int remaining = writeBuffer.length - pos;
+ if (len > remaining) {
+ // copy as much as fits
+ System.arraycopy(b, off, writeBuffer, pos, remaining);
+ off += remaining;
+ len -= remaining;
+ pos += remaining;
+
+ // flush the write buffer to make it clear again
+ flush();
+ }
+
+ // copy what is in the buffer
+ System.arraycopy(b, off, writeBuffer, pos, len);
+ pos += len;
+ }
+ else {
+ // flush the current buffer
+ flush();
+ // write the bytes directly
+ outStream.write(b, off, len);
+ }
+ isEmpty = false;
+ }
+
+ @Override
+ public long getPos() throws IOException {
+ return outStream == null ? pos : outStream.getPos();
+ }
+
+ @Override
+ public void flush() throws IOException {
+ if (!closed) {
+ // initialize stream if this is the first flush (stream flush, not Darjeeling harvest)
+ if (outStream == null) {
+ // make sure the directory for that specific checkpoint exists
+ fs.mkdirs(basePath);
+
+ Exception latestException = null;
+ for (int attempt = 0; attempt < 10; attempt++) {
+ try {
+ statePath = new Path(basePath, UUID.randomUUID().toString());
+ outStream = fs.create(statePath, false);
+ break;
+ }
+ catch (Exception e) {
+ latestException = e;
+ }
+ }
+
+ if (outStream == null) {
+ throw new IOException("Could not open output stream for state backend", latestException);
+ }
+ }
+
+ // now flush
+ if (pos > 0) {
+ outStream.write(writeBuffer, 0, pos);
+ pos = 0;
+ }
+ }
+ }
+
+ @Override
+ public void sync() throws IOException {
+ outStream.sync();
+ }
+
+ /**
+ * If the stream is only closed, we remove the produced file (cleanup through the auto close
+ * feature, for example). This method throws no exception if the deletion fails, but only
+ * logs the error.
+ */
+ @Override
+ public void close() {
+ if (!closed) {
+ closed = true;
+ if (outStream != null) {
+ try {
+ outStream.close();
+ fs.delete(statePath, false);
+
+ // attempt to delete the parent (will fail and be ignored if the parent has more files)
+ try {
+ fs.delete(basePath, false);
+ } catch (IOException ignored) {}
+ }
+ catch (Exception e) {
+ LOG.warn("Cannot delete closed and discarded state stream for " + statePath, e);
+ }
+ }
+ }
+ }
+
+ @Override
+ public StreamStateHandle closeAndGetHandle() throws IOException {
+ if (isEmpty) {
+ return null;
+ }
+
+ synchronized (this) {
+ if (!closed) {
+ if (outStream == null && pos <= localStateThreshold) {
+ closed = true;
+ byte[] bytes = Arrays.copyOf(writeBuffer, pos);
+ return new ByteStreamStateHandle(bytes);
+ }
+ else {
+ flush();
+ outStream.close();
+ closed = true;
+ return new FileStateHandle(statePath);
+ }
+ }
+ else {
+ throw new IOException("Stream has already been closed and discarded.");
+ }
+ }
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsFoldingState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsFoldingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsFoldingState.java
deleted file mode 100644
index 2fbbdc9..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsFoldingState.java
+++ /dev/null
@@ -1,161 +0,0 @@
-/*
- * 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.api.common.functions.FoldFunction;
-import org.apache.flink.api.common.state.FoldingState;
-import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.util.Preconditions;
-
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Heap-backed partitioned {@link FoldingState} that is
- * snapshotted into files.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <T> The type of the values that can be folded into the state.
- * @param <ACC> The type of the value in the folding state.
- */
-public class FsFoldingState<K, N, T, ACC>
- extends AbstractFsState<K, N, ACC, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>>
- implements FoldingState<T, ACC> {
-
- private final FoldFunction<T, ACC> foldFunction;
-
- /**
- * Creates a new and empty partitioned state.
- *
- * @param backend The file system state backend backing snapshots of this state
- * @param keySerializer The serializer for the key.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateDesc The state identifier for the state. This contains name
- * and can create a default state value.
- */
- public FsFoldingState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- FoldingStateDescriptor<T, ACC> stateDesc) {
- super(backend, keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc);
- this.foldFunction = stateDesc.getFoldFunction();
- }
-
- /**
- * Creates a new key/value state with the given state contents.
- * This method is used to re-create key/value state with existing data, for example from
- * a snapshot.
- *
- * @param backend The file system state backend backing snapshots of this state
- * @param keySerializer The serializer for the key.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateDesc The state identifier for the state. This contains name
- * and can create a default state value.
- * @param state The map of key/value pairs to initialize the state with.
- */
- public FsFoldingState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- FoldingStateDescriptor<T, ACC> stateDesc,
- HashMap<N, Map<K, ACC>> state) {
- super(backend, keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc, state);
- this.foldFunction = stateDesc.getFoldFunction();
- }
-
- @Override
- public ACC get() {
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = state.get(currentNamespace);
- }
- if (currentNSState != null) {
- Preconditions.checkState(currentKey != null, "No key set");
- return currentNSState.get(currentKey);
- } else {
- return null;
- }
- }
-
- @Override
- public void add(T value) throws IOException {
- Preconditions.checkState(currentKey != null, "No key set");
-
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = createNewNamespaceMap();
- state.put(currentNamespace, currentNSState);
- }
-
- ACC currentValue = currentNSState.get(currentKey);
- try {
- if (currentValue == null) {
- currentNSState.put(currentKey, foldFunction.fold(stateDesc.getDefaultValue(), value));
- } else {
- currentNSState.put(currentKey, foldFunction.fold(currentValue, value));
-
- }
- } catch (Exception e) {
- throw new RuntimeException("Could not add value to folding state.", e);
- }
- }
-
- @Override
- public KvStateSnapshot<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>, FsStateBackend> createHeapSnapshot(Path filePath) {
- return new Snapshot<>(getKeySerializer(), getNamespaceSerializer(), stateSerializer, stateDesc, filePath);
- }
-
- @Override
- public byte[] getSerializedValue(K key, N namespace) throws Exception {
- Preconditions.checkNotNull(key, "Key");
- Preconditions.checkNotNull(namespace, "Namespace");
-
- Map<K, ACC> stateByKey = state.get(namespace);
-
- if (stateByKey != null) {
- return KvStateRequestSerializer.serializeValue(stateByKey.get(key), stateDesc.getSerializer());
- } else {
- return null;
- }
- }
-
- public static class Snapshot<K, N, T, ACC> extends AbstractFsStateSnapshot<K, N, ACC, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>> {
- private static final long serialVersionUID = 1L;
-
- public Snapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<ACC> stateSerializer,
- FoldingStateDescriptor<T, ACC> stateDescs,
- Path filePath) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDescs, filePath);
- }
-
- @Override
- public KvState<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>, FsStateBackend> createFsState(FsStateBackend backend, HashMap<N, Map<K, ACC>> stateMap) {
- return new FsFoldingState<>(backend, keySerializer, namespaceSerializer, stateDesc, stateMap);
- }
- }
-
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsListState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsListState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsListState.java
deleted file mode 100644
index dbef900..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsListState.java
+++ /dev/null
@@ -1,149 +0,0 @@
-/*
- * 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.api.common.state.ListState;
-import org.apache.flink.api.common.state.ListStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.runtime.state.ArrayListSerializer;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.util.Preconditions;
-
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Heap-backed partitioned {@link org.apache.flink.api.common.state.ListState} that is snapshotted
- * into files.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <V> The type of the value.
- */
-public class FsListState<K, N, V>
- extends AbstractFsState<K, N, ArrayList<V>, ListState<V>, ListStateDescriptor<V>>
- implements ListState<V> {
-
- /**
- * Creates a new and empty partitioned state.
- *
- * @param keySerializer The serializer for the key.
- * @param stateDesc The state identifier for the state. This contains name
- * and can create a default state value.
- * @param backend The file system state backend backing snapshots of this state
- */
- public FsListState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ListStateDescriptor<V> stateDesc) {
- super(backend, keySerializer, namespaceSerializer, new ArrayListSerializer<>(stateDesc.getSerializer()), stateDesc);
- }
-
- /**
- * Creates a new key/value state with the given state contents.
- * This method is used to re-create key/value state with existing data, for example from
- * a snapshot.
- *
- * @param keySerializer The serializer for the key.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateDesc The state identifier for the state. This contains name
- * and can create a default state value.
- * @param state The map of key/value pairs to initialize the state with.
- * @param backend The file system state backend backing snapshots of this state
- */
- public FsListState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ListStateDescriptor<V> stateDesc,
- HashMap<N, Map<K, ArrayList<V>>> state) {
- super(backend, keySerializer, namespaceSerializer, new ArrayListSerializer<>(stateDesc.getSerializer()), stateDesc, state);
- }
-
-
- @Override
- public Iterable<V> get() {
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = state.get(currentNamespace);
- }
- if (currentNSState != null) {
- Preconditions.checkState(currentKey != null, "No key set");
- return currentNSState.get(currentKey);
- } else {
- return null;
- }
- }
-
- @Override
- public void add(V value) {
- Preconditions.checkState(currentKey != null, "No key set");
-
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = createNewNamespaceMap();
- state.put(currentNamespace, currentNSState);
- }
-
- ArrayList<V> list = currentNSState.get(currentKey);
- if (list == null) {
- list = new ArrayList<>();
- currentNSState.put(currentKey, list);
- }
- list.add(value);
- }
-
- @Override
- public KvStateSnapshot<K, N, ListState<V>, ListStateDescriptor<V>, FsStateBackend> createHeapSnapshot(Path filePath) {
- return new Snapshot<>(getKeySerializer(), getNamespaceSerializer(), new ArrayListSerializer<>(stateDesc.getSerializer()), stateDesc, filePath);
- }
-
- @Override
- public byte[] getSerializedValue(K key, N namespace) throws Exception {
- Preconditions.checkNotNull(key, "Key");
- Preconditions.checkNotNull(namespace, "Namespace");
-
- Map<K, ArrayList<V>> stateByKey = state.get(namespace);
- if (stateByKey != null) {
- return KvStateRequestSerializer.serializeList(stateByKey.get(key), stateDesc.getSerializer());
- } else {
- return null;
- }
- }
-
- public static class Snapshot<K, N, V> extends AbstractFsStateSnapshot<K, N, ArrayList<V>, ListState<V>, ListStateDescriptor<V>> {
- private static final long serialVersionUID = 1L;
-
- public Snapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<ArrayList<V>> stateSerializer,
- ListStateDescriptor<V> stateDescs,
- Path filePath) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDescs, filePath);
- }
-
- @Override
- public KvState<K, N, ListState<V>, ListStateDescriptor<V>, FsStateBackend> createFsState(FsStateBackend backend, HashMap<N, Map<K, ArrayList<V>>> stateMap) {
- return new FsListState<>(backend, keySerializer, namespaceSerializer, stateDesc, stateMap);
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsReducingState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsReducingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsReducingState.java
deleted file mode 100644
index bb389d9..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsReducingState.java
+++ /dev/null
@@ -1,165 +0,0 @@
-/*
- * 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.api.common.functions.ReduceFunction;
-import org.apache.flink.api.common.state.ReducingState;
-import org.apache.flink.api.common.state.ReducingStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.util.Preconditions;
-
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Heap-backed partitioned {@link org.apache.flink.api.common.state.ReducingState} that is
- * snapshotted into files.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <V> The type of the value.
- */
-public class FsReducingState<K, N, V>
- extends AbstractFsState<K, N, V, ReducingState<V>, ReducingStateDescriptor<V>>
- implements ReducingState<V> {
-
- private final ReduceFunction<V> reduceFunction;
-
- /**
- * Creates a new and empty partitioned state.
- *
- * @param backend The file system state backend backing snapshots of this state
- * @param keySerializer The serializer for the key.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateDesc The state identifier for the state. This contains name
- * and can create a default state value.
- */
- public FsReducingState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ReducingStateDescriptor<V> stateDesc) {
- super(backend, keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc);
- this.reduceFunction = stateDesc.getReduceFunction();
- }
-
- /**
- * Creates a new key/value state with the given state contents.
- * This method is used to re-create key/value state with existing data, for example from
- * a snapshot.
- *
- * @param backend The file system state backend backing snapshots of this state
- * @param keySerializer The serializer for the key.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateDesc The state identifier for the state. This contains name
-* and can create a default state value.
- * @param state The map of key/value pairs to initialize the state with.
- */
- public FsReducingState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ReducingStateDescriptor<V> stateDesc,
- HashMap<N, Map<K, V>> state) {
- super(backend, keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc, state);
- this.reduceFunction = stateDesc.getReduceFunction();
- }
-
-
- @Override
- public V get() {
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = state.get(currentNamespace);
- }
- if (currentNSState != null) {
- Preconditions.checkState(currentKey != null, "No key set");
- return currentNSState.get(currentKey);
- }
- return null;
- }
-
- @Override
- public void add(V value) throws IOException {
- Preconditions.checkState(currentKey != null, "No key set");
-
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = createNewNamespaceMap();
- state.put(currentNamespace, currentNSState);
- }
-// currentKeyState.merge(currentNamespace, value, new BiFunction<V, V, V>() {
-// @Override
-// public V apply(V v, V v2) {
-// try {
-// return reduceFunction.reduce(v, v2);
-// } catch (Exception e) {
-// return null;
-// }
-// }
-// });
- V currentValue = currentNSState.get(currentKey);
- if (currentValue == null) {
- currentNSState.put(currentKey, value);
- } else {
- try {
- currentNSState.put(currentKey, reduceFunction.reduce(currentValue, value));
- } catch (Exception e) {
- throw new RuntimeException("Could not add value to reducing state.", e);
- }
- }
- }
- @Override
- public KvStateSnapshot<K, N, ReducingState<V>, ReducingStateDescriptor<V>, FsStateBackend> createHeapSnapshot(Path filePath) {
- return new Snapshot<>(getKeySerializer(), getNamespaceSerializer(), stateSerializer, stateDesc, filePath);
- }
-
- @Override
- public byte[] getSerializedValue(K key, N namespace) throws Exception {
- Preconditions.checkNotNull(key, "Key");
- Preconditions.checkNotNull(namespace, "Namespace");
-
- Map<K, V> stateByKey = state.get(namespace);
- if (stateByKey != null) {
- return KvStateRequestSerializer.serializeValue(stateByKey.get(key), stateDesc.getSerializer());
- } else {
- return null;
- }
- }
-
- public static class Snapshot<K, N, V> extends AbstractFsStateSnapshot<K, N, V, ReducingState<V>, ReducingStateDescriptor<V>> {
- private static final long serialVersionUID = 1L;
-
- public Snapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<V> stateSerializer,
- ReducingStateDescriptor<V> stateDescs,
- Path filePath) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDescs, filePath);
- }
-
- @Override
- public KvState<K, N, ReducingState<V>, ReducingStateDescriptor<V>, FsStateBackend> createFsState(FsStateBackend backend, HashMap<N, Map<K, V>> stateMap) {
- return new FsReducingState<>(backend, keySerializer, namespaceSerializer, stateDesc, stateMap);
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
index a3f4682..5495244 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
@@ -18,30 +18,26 @@
package org.apache.flink.runtime.state.filesystem;
-import org.apache.flink.api.common.state.FoldingState;
-import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.state.ListState;
-import org.apache.flink.api.common.state.ListStateDescriptor;
-import org.apache.flink.api.common.state.ReducingState;
-import org.apache.flink.api.common.state.ReducingStateDescriptor;
-import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.core.fs.Path;
import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StreamStateHandle;
-import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.heap.HeapKeyedStateBackend;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
-import java.util.Arrays;
-import java.util.UUID;
+import java.util.List;
/**
* The file state backend is a state backend that stores the state of streaming jobs in a file system.
@@ -63,12 +59,8 @@ public class FsStateBackend extends AbstractStateBackend {
public static final int DEFAULT_FILE_STATE_THRESHOLD = 1024;
/** Maximum size of state that is stored with the metadata, rather than in files */
- public static final int MAX_FILE_STATE_THRESHOLD = 1024 * 1024;
+ private static final int MAX_FILE_STATE_THRESHOLD = 1024 * 1024;
- /** Default size for the write buffer */
- private static final int DEFAULT_WRITE_BUFFER_SIZE = 4096;
-
-
/** The path to the directory for the checkpoint data, including the file system
* description via scheme and optional authority */
private final Path basePath;
@@ -76,13 +68,6 @@ public class FsStateBackend extends AbstractStateBackend {
/** State below this size will be stored as part of the metadata, rather than in files */
private final int fileStateThreshold;
- /** The directory (job specific) into this initialized instance of the backend stores its data */
- private transient Path checkpointDirectory;
-
- /** Cached handle to the file system for file operations */
- private transient FileSystem filesystem;
-
-
/**
* Creates a new state backend that stores its checkpoint data in the file system and location
* defined by the given URI.
@@ -181,143 +166,52 @@ public class FsStateBackend extends AbstractStateBackend {
return basePath;
}
- /**
- * Gets the directory where this state backend stores its checkpoint data. Will be null if
- * the state backend has not been initialized.
- *
- * @return The directory where this state backend stores its checkpoint data.
- */
- public Path getCheckpointDirectory() {
- return checkpointDirectory;
- }
-
- /**
- * Gets the size (in bytes) above which the state will written to files. State whose size
- * is below this threshold will be directly stored with the metadata
- * (the state handles), rather than in files. This threshold helps to prevent an accumulation
- * of small files for small states.
- *
- * @return The threshold (in bytes) above which state is written to files.
- */
- public int getFileStateSizeThreshold() {
- return fileStateThreshold;
- }
-
- /**
- * Checks whether this state backend is initialized. Note that initialization does not carry
- * across serialization. After each serialization, the state backend needs to be initialized.
- *
- * @return True, if the file state backend has been initialized, false otherwise.
- */
- public boolean isInitialized() {
- return filesystem != null && checkpointDirectory != null;
- }
-
- /**
- * Gets the file system handle for the file system that stores the state for this backend.
- *
- * @return This backend's file system handle.
- */
- public FileSystem getFileSystem() {
- if (filesystem != null) {
- return filesystem;
- }
- else {
- throw new IllegalStateException("State backend has not been initialized.");
- }
- }
-
// ------------------------------------------------------------------------
// initialization and cleanup
// ------------------------------------------------------------------------
@Override
- public void initializeForJob(Environment env,
- String operatorIdentifier,
- TypeSerializer<?> keySerializer) throws Exception {
- super.initializeForJob(env, operatorIdentifier, keySerializer);
-
- Path dir = new Path(basePath, env.getJobID().toString());
-
- LOG.info("Initializing file state backend to URI " + dir);
-
- filesystem = basePath.getFileSystem();
- filesystem.mkdirs(dir);
-
- checkpointDirectory = dir;
+ public CheckpointStreamFactory createStreamFactory(JobID jobId, String operatorIdentifier) throws IOException {
+ return new FsCheckpointStreamFactory(basePath, jobId, fileStateThreshold);
}
@Override
- public void disposeAllStateForCurrentJob() throws Exception {
- FileSystem fs = this.filesystem;
- Path dir = this.checkpointDirectory;
-
- if (fs != null && dir != null) {
- this.filesystem = null;
- this.checkpointDirectory = null;
- fs.delete(dir, true);
- }
- else {
- throw new IllegalStateException("state backend has not been initialized");
- }
+ public <K> KeyedStateBackend<K> createKeyedStateBackend(
+ Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ TaskKvStateRegistry kvStateRegistry) throws Exception {
+ return new HeapKeyedStateBackend<>(
+ kvStateRegistry,
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange);
}
@Override
- public void close() throws Exception {}
-
- // ------------------------------------------------------------------------
- // state backend operations
- // ------------------------------------------------------------------------
-
- @Override
- public <N, V> ValueState<V> createValueState(TypeSerializer<N> namespaceSerializer, ValueStateDescriptor<V> stateDesc) throws Exception {
- return new FsValueState<>(this, keySerializer, namespaceSerializer, stateDesc);
- }
-
- @Override
- public <N, T> ListState<T> createListState(TypeSerializer<N> namespaceSerializer, ListStateDescriptor<T> stateDesc) throws Exception {
- return new FsListState<>(this, keySerializer, namespaceSerializer, stateDesc);
- }
-
- @Override
- public <N, T> ReducingState<T> createReducingState(TypeSerializer<N> namespaceSerializer, ReducingStateDescriptor<T> stateDesc) throws Exception {
- return new FsReducingState<>(this, keySerializer, namespaceSerializer, stateDesc);
- }
-
- @Override
- protected <N, T, ACC> FoldingState<T, ACC> createFoldingState(TypeSerializer<N> namespaceSerializer,
- FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
- return new FsFoldingState<>(this, keySerializer, namespaceSerializer, stateDesc);
- }
-
- @Override
- public FsCheckpointStateOutputStream createCheckpointStateOutputStream(long checkpointID, long timestamp) throws Exception {
- checkFileSystemInitialized();
-
- Path checkpointDir = createCheckpointDirPath(checkpointID);
- int bufferSize = Math.max(DEFAULT_WRITE_BUFFER_SIZE, fileStateThreshold);
- return new FsCheckpointStateOutputStream(checkpointDir, filesystem, bufferSize, fileStateThreshold);
- }
-
- // ------------------------------------------------------------------------
- // utilities
- // ------------------------------------------------------------------------
-
- private void checkFileSystemInitialized() throws IllegalStateException {
- if (filesystem == null || checkpointDirectory == null) {
- throw new IllegalStateException("filesystem has not been re-initialized after deserialization");
- }
- }
-
- private Path createCheckpointDirPath(long checkpointID) {
- return new Path(checkpointDirectory, "chk-" + checkpointID);
+ public <K> KeyedStateBackend<K> restoreKeyedStateBackend(
+ Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ List<KeyGroupsStateHandle> restoredState,
+ TaskKvStateRegistry kvStateRegistry) throws Exception {
+ return new HeapKeyedStateBackend<>(
+ kvStateRegistry,
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange,
+ restoredState);
}
@Override
public String toString() {
- return checkpointDirectory == null ?
- "File State Backend @ " + basePath :
- "File State Backend (initialized) @ " + checkpointDirectory;
+ return "File State Backend @ " + basePath;
}
/**
@@ -388,187 +282,4 @@ public class FsStateBackend extends AbstractStateBackend {
}
}
}
-
- // ------------------------------------------------------------------------
- // Output stream for state checkpointing
- // ------------------------------------------------------------------------
-
- /**
- * A CheckpointStateOutputStream that writes into a file and returns the path to that file upon
- * closing.
- */
- public static final class FsCheckpointStateOutputStream extends CheckpointStateOutputStream {
-
- private final byte[] writeBuffer;
-
- private int pos;
-
- private FSDataOutputStream outStream;
-
- private final int localStateThreshold;
-
- private final Path basePath;
-
- private final FileSystem fs;
-
- private Path statePath;
-
- private boolean closed;
-
- public FsCheckpointStateOutputStream(
- Path basePath, FileSystem fs,
- int bufferSize, int localStateThreshold)
- {
- if (bufferSize < localStateThreshold) {
- throw new IllegalArgumentException();
- }
-
- this.basePath = basePath;
- this.fs = fs;
- this.writeBuffer = new byte[bufferSize];
- this.localStateThreshold = localStateThreshold;
- }
-
-
- @Override
- public void write(int b) throws IOException {
- if (pos >= writeBuffer.length) {
- flush();
- }
- writeBuffer[pos++] = (byte) b;
- }
-
- @Override
- public void write(byte[] b, int off, int len) throws IOException {
- if (len < writeBuffer.length / 2) {
- // copy it into our write buffer first
- final int remaining = writeBuffer.length - pos;
- if (len > remaining) {
- // copy as much as fits
- System.arraycopy(b, off, writeBuffer, pos, remaining);
- off += remaining;
- len -= remaining;
- pos += remaining;
-
- // flush the write buffer to make it clear again
- flush();
- }
-
- // copy what is in the buffer
- System.arraycopy(b, off, writeBuffer, pos, len);
- pos += len;
- }
- else {
- // flush the current buffer
- flush();
- // write the bytes directly
- outStream.write(b, off, len);
- }
- }
-
- @Override
- public void flush() throws IOException {
- if (!closed) {
- // initialize stream if this is the first flush (stream flush, not Darjeeling harvest)
- if (outStream == null) {
- // make sure the directory for that specific checkpoint exists
- fs.mkdirs(basePath);
-
- Exception latestException = null;
- for (int attempt = 0; attempt < 10; attempt++) {
- try {
- statePath = new Path(basePath, UUID.randomUUID().toString());
- outStream = fs.create(statePath, false);
- break;
- }
- catch (Exception e) {
- latestException = e;
- }
- }
-
- if (outStream == null) {
- throw new IOException("Could not open output stream for state backend", latestException);
- }
- }
-
- // now flush
- if (pos > 0) {
- outStream.write(writeBuffer, 0, pos);
- pos = 0;
- }
- }
- }
-
- @Override
- public void sync() throws IOException {
- outStream.sync();
- }
-
- /**
- * If the stream is only closed, we remove the produced file (cleanup through the auto close
- * feature, for example). This method throws no exception if the deletion fails, but only
- * logs the error.
- */
- @Override
- public void close() {
- if (!closed) {
- closed = true;
- if (outStream != null) {
- try {
- outStream.close();
- fs.delete(statePath, false);
-
- // attempt to delete the parent (will fail and be ignored if the parent has more files)
- try {
- fs.delete(basePath, false);
- } catch (IOException ignored) {}
- }
- catch (Exception e) {
- LOG.warn("Cannot delete closed and discarded state stream for " + statePath, e);
- }
- }
- }
- }
-
- @Override
- public StreamStateHandle closeAndGetHandle() throws IOException {
- synchronized (this) {
- if (!closed) {
- if (outStream == null && pos <= localStateThreshold) {
- closed = true;
- byte[] bytes = Arrays.copyOf(writeBuffer, pos);
- return new ByteStreamStateHandle(bytes);
- }
- else {
- flush();
- outStream.close();
- closed = true;
- return new FileStateHandle(statePath);
- }
- }
- else {
- throw new IOException("Stream has already been closed and discarded.");
- }
- }
- }
-
- /**
- * Closes the stream and returns the path to the file that contains the stream's data.
- * @return The path to the file that contains the stream's data.
- * @throws IOException Thrown if the stream cannot be successfully closed.
- */
- public Path closeAndGetPath() throws IOException {
- synchronized (this) {
- if (!closed) {
- closed = true;
- flush();
- outStream.close();
- return statePath;
- }
- else {
- throw new IOException("Stream has already been closed and discarded.");
- }
- }
- }
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsValueState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsValueState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsValueState.java
deleted file mode 100644
index 698bc1f..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsValueState.java
+++ /dev/null
@@ -1,148 +0,0 @@
-/*
- * 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.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.util.Preconditions;
-
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Heap-backed partitioned {@link org.apache.flink.api.common.state.ValueState} that is snapshotted
- * into files.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <V> The type of the value.
- */
-public class FsValueState<K, N, V>
- extends AbstractFsState<K, N, V, ValueState<V>, ValueStateDescriptor<V>>
- implements ValueState<V> {
-
- /**
- * Creates a new and empty key/value state.
- *
- * @param keySerializer The serializer for the key.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateDesc The state identifier for the state. This contains name
- * and can create a default state value.
- * @param backend The file system state backend backing snapshots of this state
- */
- public FsValueState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ValueStateDescriptor<V> stateDesc) {
- super(backend, keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc);
- }
-
- /**
- * Creates a new key/value state with the given state contents.
- * This method is used to re-create key/value state with existing data, for example from
- * a snapshot.
- *
- * @param keySerializer The serializer for the key.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateDesc The state identifier for the state. This contains name
- * and can create a default state value.
- * @param state The map of key/value pairs to initialize the state with.
- * @param backend The file system state backend backing snapshots of this state
- */
- public FsValueState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ValueStateDescriptor<V> stateDesc,
- HashMap<N, Map<K, V>> state) {
- super(backend, keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc, state);
- }
-
- @Override
- public V value() {
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = state.get(currentNamespace);
- }
- if (currentNSState != null) {
- Preconditions.checkState(currentKey != null, "No key set");
- V value = currentNSState.get(currentKey);
- return value != null ? value : stateDesc.getDefaultValue();
- }
- return stateDesc.getDefaultValue();
- }
-
- @Override
- public void update(V value) {
- Preconditions.checkState(currentKey != null, "No key set");
-
- if (value == null) {
- clear();
- return;
- }
-
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = createNewNamespaceMap();
- state.put(currentNamespace, currentNSState);
- }
-
- currentNSState.put(currentKey, value);
- }
-
- @Override
- public KvStateSnapshot<K, N, ValueState<V>, ValueStateDescriptor<V>, FsStateBackend> createHeapSnapshot(Path filePath) {
- return new Snapshot<>(getKeySerializer(), getNamespaceSerializer(), stateSerializer, stateDesc, filePath);
- }
-
- @Override
- public byte[] getSerializedValue(K key, N namespace) throws Exception {
- Preconditions.checkNotNull(key, "Key");
- Preconditions.checkNotNull(namespace, "Namespace");
-
- Map<K, V> stateByKey = state.get(namespace);
- V value = stateByKey != null ? stateByKey.get(key) : stateDesc.getDefaultValue();
- if (value != null) {
- return KvStateRequestSerializer.serializeValue(value, stateDesc.getSerializer());
- } else {
- return KvStateRequestSerializer.serializeValue(stateDesc.getDefaultValue(), stateDesc.getSerializer());
- }
- }
-
- public static class Snapshot<K, N, V> extends AbstractFsStateSnapshot<K, N, V, ValueState<V>, ValueStateDescriptor<V>> {
- private static final long serialVersionUID = 1L;
-
- public Snapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<V> stateSerializer,
- ValueStateDescriptor<V> stateDescs,
- Path filePath) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDescs, filePath);
- }
-
- @Override
- public KvState<K, N, ValueState<V>, ValueStateDescriptor<V>, FsStateBackend> createFsState(FsStateBackend backend, HashMap<N, Map<K, V>> stateMap) {
- return new FsValueState<>(backend, keySerializer, namespaceSerializer, stateDesc, stateMap);
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapState.java
new file mode 100644
index 0000000..9863c93
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapState.java
@@ -0,0 +1,187 @@
+/*
+ * 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.heap;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.State;
+import org.apache.flink.api.common.state.StateDescriptor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.KvState;
+import org.apache.flink.runtime.state.heap.StateTable;
+import org.apache.flink.util.Preconditions;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Base class for partitioned {@link ListState} implementations that are backed by a regular
+ * heap hash map. The concrete implementations define how the state is checkpointed.
+ *
+ * @param <K> The type of the key.
+ * @param <N> The type of the namespace.
+ * @param <SV> The type of the values in the state.
+ * @param <S> The type of State
+ * @param <SD> The type of StateDescriptor for the State S
+ */
+public abstract class AbstractHeapState<K, N, SV, S extends State, SD extends StateDescriptor<S, ?>>
+ implements KvState<N>, State {
+
+ /** Map containing the actual key/value pairs */
+ protected final StateTable<K, N, SV> stateTable;
+
+ /** This holds the name of the state and can create an initial default value for the state. */
+ protected final SD stateDesc;
+
+ /** The current namespace, which the access methods will refer to. */
+ protected N currentNamespace = null;
+
+ protected final KeyedStateBackend<K> backend;
+
+ protected final TypeSerializer<K> keySerializer;
+
+ protected final TypeSerializer<N> namespaceSerializer;
+
+ /**
+ * Creates a new key/value state for the given hash map of key/value pairs.
+ *
+ * @param backend The state backend backing that created this state.
+ * @param stateDesc The state identifier for the state. This contains name
+ * and can create a default state value.
+ * @param stateTable The state tab;e to use in this kev/value state. May contain initial state.
+ */
+ protected AbstractHeapState(
+ KeyedStateBackend<K> backend,
+ SD stateDesc,
+ StateTable<K, N, SV> stateTable,
+ TypeSerializer<K> keySerializer,
+ TypeSerializer<N> namespaceSerializer) {
+
+ Preconditions.checkNotNull(stateTable, "State table must not be null.");
+
+ this.backend = backend;
+ this.stateDesc = stateDesc;
+ this.stateTable = stateTable;
+ this.keySerializer = keySerializer;
+ this.namespaceSerializer = namespaceSerializer;
+ }
+
+ // ------------------------------------------------------------------------
+
+ @Override
+ public final void clear() {
+ Preconditions.checkState(currentNamespace != null, "No namespace set.");
+ Preconditions.checkState(backend.getCurrentKey() != null, "No key set.");
+
+ Map<N, Map<K, SV>> namespaceMap =
+ stateTable.get(backend.getCurrentKeyGroupIndex());
+
+ if (namespaceMap == null) {
+ return;
+ }
+
+ Map<K, SV> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ return;
+ }
+
+ SV removed = keyedMap.remove(backend.getCurrentKey());
+
+ if (removed == null) {
+ return;
+ }
+
+ if (!keyedMap.isEmpty()) {
+ return;
+ }
+
+ namespaceMap.remove(currentNamespace);
+ }
+
+ @Override
+ public final void setCurrentNamespace(N namespace) {
+ this.currentNamespace = Preconditions.checkNotNull(namespace, "Namespace must not be null.");
+ }
+
+ @Override
+ public byte[] getSerializedValue(byte[] serializedKeyAndNamespace) throws Exception {
+ Preconditions.checkNotNull(serializedKeyAndNamespace, "Serialized key and namespace");
+
+ Tuple2<K, N> keyAndNamespace = KvStateRequestSerializer.deserializeKeyAndNamespace(
+ serializedKeyAndNamespace, keySerializer, namespaceSerializer);
+
+ return getSerializedValue(keyAndNamespace.f0, keyAndNamespace.f1);
+ }
+
+ public byte[] getSerializedValue(K key, N namespace) throws Exception {
+ Preconditions.checkState(namespace != null, "No namespace given.");
+ Preconditions.checkState(key != null, "No key given.");
+
+ Map<N, Map<K, SV>> namespaceMap =
+ stateTable.get(backend.getKeyGroupAssigner().getKeyGroupIndex(key));
+
+ if (namespaceMap == null) {
+ return null;
+ }
+
+ Map<K, SV> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ return null;
+ }
+
+ SV result = keyedMap.get(key);
+
+ if (result == null) {
+ return null;
+ }
+
+ @SuppressWarnings("unchecked,rawtypes")
+ TypeSerializer serializer = stateDesc.getSerializer();
+
+ return KvStateRequestSerializer.serializeValue(result, serializer);
+ }
+
+ /**
+ * Creates a new map for use in Heap based state.
+ *
+ * <p>If the state queryable ({@link StateDescriptor#isQueryable()}, this
+ * will create a concurrent hash map instead of a regular one.
+ *
+ * @return A new namespace map.
+ */
+ protected <MK, MV> Map<MK, MV> createNewMap() {
+ if (stateDesc.isQueryable()) {
+ return new ConcurrentHashMap<>();
+ } else {
+ return new HashMap<>();
+ }
+ }
+
+ /**
+ * This should only be used for testing.
+ */
+ public StateTable<K, N, SV> getStateTable() {
+ return stateTable;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapFoldingState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapFoldingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapFoldingState.java
new file mode 100644
index 0000000..1679122
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapFoldingState.java
@@ -0,0 +1,124 @@
+/*
+ * 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.heap;
+
+import org.apache.flink.api.common.functions.FoldFunction;
+import org.apache.flink.api.common.state.FoldingState;
+import org.apache.flink.api.common.state.FoldingStateDescriptor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.util.Preconditions;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * Heap-backed partitioned {@link FoldingState} that is
+ * snapshotted into files.
+ *
+ * @param <K> The type of the key.
+ * @param <N> The type of the namespace.
+ * @param <T> The type of the values that can be folded into the state.
+ * @param <ACC> The type of the value in the folding state.
+ */
+public class HeapFoldingState<K, N, T, ACC>
+ extends AbstractHeapState<K, N, ACC, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>>
+ implements FoldingState<T, ACC> {
+
+ private final FoldFunction<T, ACC> foldFunction;
+
+ /**
+ * Creates a new key/value state for the given hash map of key/value pairs.
+ *
+ * @param backend The state backend backing that created this state.
+ * @param stateDesc The state identifier for the state. This contains name
+ * and can create a default state value.
+ * @param stateTable The state tab;e to use in this kev/value state. May contain initial state.
+ */
+ public HeapFoldingState(
+ KeyedStateBackend<K> backend,
+ FoldingStateDescriptor<T, ACC> stateDesc,
+ StateTable<K, N, ACC> stateTable,
+ TypeSerializer<K> keySerializer,
+ TypeSerializer<N> namespaceSerializer) {
+ super(backend, stateDesc, stateTable, keySerializer, namespaceSerializer);
+ this.foldFunction = stateDesc.getFoldFunction();
+ }
+
+ @Override
+ public ACC get() {
+ Preconditions.checkState(currentNamespace != null, "No namespace set.");
+ Preconditions.checkState(backend.getCurrentKey() != null, "No key set.");
+
+ Map<N, Map<K, ACC>> namespaceMap =
+ stateTable.get(backend.getCurrentKeyGroupIndex());
+
+ if (namespaceMap == null) {
+ return null;
+ }
+
+ Map<K, ACC> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ return null;
+ }
+
+ return keyedMap.get(backend.<K>getCurrentKey());
+ }
+
+ @Override
+ public void add(T value) throws IOException {
+ Preconditions.checkState(currentNamespace != null, "No namespace set.");
+ Preconditions.checkState(backend.getCurrentKey() != null, "No key set.");
+
+ if (value == null) {
+ clear();
+ return;
+ }
+
+ Map<N, Map<K, ACC>> namespaceMap =
+ stateTable.get(backend.getCurrentKeyGroupIndex());
+
+ if (namespaceMap == null) {
+ namespaceMap = createNewMap();
+ stateTable.set(backend.getCurrentKeyGroupIndex(), namespaceMap);
+ }
+
+ Map<K, ACC> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ keyedMap = createNewMap();
+ namespaceMap.put(currentNamespace, keyedMap);
+ }
+
+ ACC currentValue = keyedMap.get(backend.<K>getCurrentKey());
+
+ try {
+
+ if (currentValue == null) {
+ keyedMap.put(backend.<K>getCurrentKey(),
+ foldFunction.fold(stateDesc.getDefaultValue(), value));
+ } else {
+ keyedMap.put(backend.<K>getCurrentKey(), foldFunction.fold(currentValue, value));
+ }
+ } catch (Exception e) {
+ throw new RuntimeException("Could not add value to folding state.", e);
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
new file mode 100644
index 0000000..fcb4bef
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
@@ -0,0 +1,328 @@
+/*
+ * 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.heap;
+
+import org.apache.flink.api.common.state.FoldingState;
+import org.apache.flink.api.common.state.FoldingStateDescriptor;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.state.ReducingState;
+import org.apache.flink.api.common.state.ReducingStateDescriptor;
+import org.apache.flink.api.common.state.ValueState;
+import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.runtime.query.TaskKvStateRegistry;
+import org.apache.flink.runtime.state.ArrayListSerializer;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.DoneFuture;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeOffsets;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.util.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.RunnableFuture;
+
+/**
+ * A {@link KeyedStateBackend} that keeps state on the Java Heap and will serialize state to
+ * streams provided by a {@link org.apache.flink.runtime.state.CheckpointStreamFactory} upon
+ * checkpointing.
+ *
+ * @param <K> The key by which state is keyed.
+ */
+public class HeapKeyedStateBackend<K> extends KeyedStateBackend<K> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(HeapKeyedStateBackend.class);
+
+ /**
+ * Map of state tables that stores all state of key/value states. We store it centrally so
+ * that we can easily checkpoint/restore it.
+ *
+ * <p>The actual parameters of StateTable are {@code StateTable<NamespaceT, Map<KeyT, StateT>>}
+ * but we can't put them here because different key/value states with different types and
+ * namespace types share this central list of tables.
+ */
+ private final Map<String, StateTable<K, ?, ?>> stateTables = new HashMap<>();
+
+ public HeapKeyedStateBackend(
+ TaskKvStateRegistry kvStateRegistry,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange) {
+
+ super(kvStateRegistry, keySerializer, keyGroupAssigner, keyGroupRange);
+
+ LOG.info("Initializing heap keyed state backend with stream factory.");
+ }
+
+ public HeapKeyedStateBackend(TaskKvStateRegistry kvStateRegistry,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ List<KeyGroupsStateHandle> restoredState) throws Exception {
+ super(kvStateRegistry, keySerializer, keyGroupAssigner, keyGroupRange);
+
+ LOG.info("Initializing heap keyed state backend from snapshot.");
+
+ if (LOG.isDebugEnabled()) {
+ LOG.debug("Restoring snapshot from state handles: {}.", restoredState);
+ }
+
+ restorePartitionedState(restoredState);
+ }
+
+ // ------------------------------------------------------------------------
+ // state backend operations
+ // ------------------------------------------------------------------------
+
+ @Override
+ public <N, V> ValueState<V> createValueState(TypeSerializer<N> namespaceSerializer, ValueStateDescriptor<V> stateDesc) throws Exception {
+ @SuppressWarnings("unchecked,rawtypes")
+ StateTable<K, N, V> stateTable = (StateTable) stateTables.get(stateDesc.getName());
+
+
+ if (stateTable == null) {
+ stateTable = new StateTable<>(stateDesc.getSerializer(), namespaceSerializer, keyGroupRange);
+ stateTables.put(stateDesc.getName(), stateTable);
+ }
+
+ return new HeapValueState<>(this, stateDesc, stateTable, keySerializer, namespaceSerializer);
+ }
+
+ @Override
+ public <N, T> ListState<T> createListState(TypeSerializer<N> namespaceSerializer, ListStateDescriptor<T> stateDesc) throws Exception {
+ @SuppressWarnings("unchecked,rawtypes")
+ StateTable<K, N, ArrayList<T>> stateTable = (StateTable) stateTables.get(stateDesc.getName());
+
+ if (stateTable == null) {
+ stateTable = new StateTable<>(new ArrayListSerializer<>(stateDesc.getSerializer()), namespaceSerializer, keyGroupRange);
+ stateTables.put(stateDesc.getName(), stateTable);
+ }
+
+ return new HeapListState<>(this, stateDesc, stateTable, keySerializer, namespaceSerializer);
+ }
+
+ @Override
+ public <N, T> ReducingState<T> createReducingState(TypeSerializer<N> namespaceSerializer, ReducingStateDescriptor<T> stateDesc) throws Exception {
+ @SuppressWarnings("unchecked,rawtypes")
+ StateTable<K, N, T> stateTable = (StateTable) stateTables.get(stateDesc.getName());
+
+
+ if (stateTable == null) {
+ stateTable = new StateTable<>(stateDesc.getSerializer(), namespaceSerializer, keyGroupRange);
+ stateTables.put(stateDesc.getName(), stateTable);
+ }
+
+ return new HeapReducingState<>(this, stateDesc, stateTable, keySerializer, namespaceSerializer);
+ }
+
+ @Override
+ protected <N, T, ACC> FoldingState<T, ACC> createFoldingState(TypeSerializer<N> namespaceSerializer, FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
+ @SuppressWarnings("unchecked,rawtypes")
+ StateTable<K, N, ACC> stateTable = (StateTable) stateTables.get(stateDesc.getName());
+
+ if (stateTable == null) {
+ stateTable = new StateTable<>(stateDesc.getSerializer(), namespaceSerializer, keyGroupRange);
+ stateTables.put(stateDesc.getName(), stateTable);
+ }
+
+ return new HeapFoldingState<>(this, stateDesc, stateTable, keySerializer, namespaceSerializer);
+ }
+
+ @Override
+ @SuppressWarnings("rawtypes,unchecked")
+ public RunnableFuture<KeyGroupsStateHandle> snapshot(
+ long checkpointId,
+ long timestamp,
+ CheckpointStreamFactory streamFactory) throws Exception {
+
+ CheckpointStreamFactory.CheckpointStateOutputStream stream =
+ streamFactory.createCheckpointStateOutputStream(
+ checkpointId,
+ timestamp);
+
+ if (stateTables.isEmpty()) {
+ return new DoneFuture<>(null);
+ }
+
+ DataOutputViewStreamWrapper outView = new DataOutputViewStreamWrapper(stream);
+
+ Preconditions.checkState(stateTables.size() <= Short.MAX_VALUE,
+ "Too many KV-States: " + stateTables.size() +
+ ". Currently at most " + Short.MAX_VALUE + " states are supported");
+
+ outView.writeShort(stateTables.size());
+
+ Map<String, Integer> kVStateToId = new HashMap<>(stateTables.size());
+
+ for (Map.Entry<String, StateTable<K, ?, ?>> kvState : stateTables.entrySet()) {
+
+ outView.writeUTF(kvState.getKey());
+
+ TypeSerializer namespaceSerializer = kvState.getValue().getNamespaceSerializer();
+ TypeSerializer stateSerializer = kvState.getValue().getStateSerializer();
+
+ ObjectOutputStream oos = new ObjectOutputStream(outView);
+ oos.writeObject(namespaceSerializer);
+ oos.writeObject(stateSerializer);
+ oos.flush();
+
+ kVStateToId.put(kvState.getKey(), kVStateToId.size());
+ }
+
+ int offsetCounter = 0;
+ long[] keyGroupRangeOffsets = new long[keyGroupRange.getNumberOfKeyGroups()];
+
+ for (int keyGroupIndex = keyGroupRange.getStartKeyGroup(); keyGroupIndex <= keyGroupRange.getEndKeyGroup(); keyGroupIndex++) {
+ keyGroupRangeOffsets[offsetCounter++] = stream.getPos();
+ outView.writeInt(keyGroupIndex);
+
+ for (Map.Entry<String, StateTable<K, ?, ?>> kvState : stateTables.entrySet()) {
+
+ outView.writeShort(kVStateToId.get(kvState.getKey()));
+
+ TypeSerializer namespaceSerializer = kvState.getValue().getNamespaceSerializer();
+ TypeSerializer stateSerializer = kvState.getValue().getStateSerializer();
+
+ // Map<NamespaceT, Map<KeyT, StateT>>
+ Map<?, ? extends Map<K, ?>> namespaceMap = kvState.getValue().get(keyGroupIndex);
+ if (namespaceMap == null) {
+ outView.writeByte(0);
+ continue;
+ }
+
+ outView.writeByte(1);
+
+ // number of namespaces
+ outView.writeInt(namespaceMap.size());
+ for (Map.Entry<?, ? extends Map<K, ?>> namespace : namespaceMap.entrySet()) {
+ namespaceSerializer.serialize(namespace.getKey(), outView);
+
+ Map<K, ?> entryMap = namespace.getValue();
+
+ // number of entries
+ outView.writeInt(entryMap.size());
+ for (Map.Entry<K, ?> entry : entryMap.entrySet()) {
+ keySerializer.serialize(entry.getKey(), outView);
+ stateSerializer.serialize(entry.getValue(), outView);
+ }
+ }
+ }
+ outView.flush();
+ }
+
+ StreamStateHandle streamStateHandle = stream.closeAndGetHandle();
+
+ KeyGroupRangeOffsets offsets = new KeyGroupRangeOffsets(keyGroupRange, keyGroupRangeOffsets);
+ final KeyGroupsStateHandle keyGroupsStateHandle = new KeyGroupsStateHandle(offsets, streamStateHandle);
+
+ return new DoneFuture(keyGroupsStateHandle);
+ }
+
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ public void restorePartitionedState(List<KeyGroupsStateHandle> state) throws Exception {
+
+ for (KeyGroupsStateHandle keyGroupsHandle : state) {
+
+ if(keyGroupsHandle == null) {
+ continue;
+ }
+
+ FSDataInputStream fsDataInputStream = keyGroupsHandle.getStateHandle().openInputStream();
+ DataInputViewStreamWrapper inView = new DataInputViewStreamWrapper(fsDataInputStream);
+
+ int numKvStates = inView.readShort();
+
+ Map<Integer, String> kvStatesById = new HashMap<>(numKvStates);
+
+ for (int i = 0; i < numKvStates; ++i) {
+ String stateName = inView.readUTF();
+
+ ObjectInputStream ois = new ObjectInputStream(inView);
+
+ TypeSerializer namespaceSerializer = (TypeSerializer) ois.readObject();
+ TypeSerializer stateSerializer = (TypeSerializer) ois.readObject();
+ StateTable<K, ?, ?> stateTable = new StateTable(stateSerializer,
+ namespaceSerializer,
+ keyGroupRange);
+ stateTables.put(stateName, stateTable);
+ kvStatesById.put(i, stateName);
+ }
+
+ for (int keyGroupIndex = keyGroupRange.getStartKeyGroup(); keyGroupIndex <= keyGroupRange.getEndKeyGroup(); keyGroupIndex++) {
+ long offset = keyGroupsHandle.getOffsetForKeyGroup(keyGroupIndex);
+ fsDataInputStream.seek(offset);
+
+ int writtenKeyGroupIndex = inView.readInt();
+ assert writtenKeyGroupIndex == keyGroupIndex;
+
+ for (int i = 0; i < numKvStates; i++) {
+ int kvStateId = inView.readShort();
+
+ byte isPresent = inView.readByte();
+ if (isPresent == 0) {
+ continue;
+ }
+
+ StateTable<K, ?, ?> stateTable = stateTables.get(kvStatesById.get(kvStateId));
+ Preconditions.checkNotNull(stateTable);
+
+ TypeSerializer namespaceSerializer = stateTable.getNamespaceSerializer();
+ TypeSerializer stateSerializer = stateTable.getStateSerializer();
+
+ Map namespaceMap = new HashMap<>();
+ stateTable.set(keyGroupIndex, namespaceMap);
+
+ int numNamespaces = inView.readInt();
+ for (int k = 0; k < numNamespaces; k++) {
+ Object namespace = namespaceSerializer.deserialize(inView);
+ Map entryMap = new HashMap<>();
+ namespaceMap.put(namespace, entryMap);
+
+ int numEntries = inView.readInt();
+ for (int l = 0; l < numEntries; l++) {
+ Object key = keySerializer.deserialize(inView);
+ Object value = stateSerializer.deserialize(inView);
+ entryMap.put(key, value);
+ }
+ }
+ }
+ }
+ }
+ }
+
+ @Override
+ public String toString() {
+ return "HeapKeyedStateBackend";
+ }
+
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapListState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapListState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapListState.java
new file mode 100644
index 0000000..4c65c25
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapListState.java
@@ -0,0 +1,156 @@
+/*
+ * 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.heap;
+
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.util.Preconditions;
+
+import java.io.ByteArrayOutputStream;
+import java.util.ArrayList;
+import java.util.Map;
+
+/**
+ * Heap-backed partitioned {@link org.apache.flink.api.common.state.ListState} that is snapshotted
+ * into files.
+ *
+ * @param <K> The type of the key.
+ * @param <N> The type of the namespace.
+ * @param <V> The type of the value.
+ */
+public class HeapListState<K, N, V>
+ extends AbstractHeapState<K, N, ArrayList<V>, ListState<V>, ListStateDescriptor<V>>
+ implements ListState<V> {
+
+ /**
+ * Creates a new key/value state for the given hash map of key/value pairs.
+ *
+ * @param backend The state backend backing that created this state.
+ * @param stateDesc The state identifier for the state. This contains name
+ * and can create a default state value.
+ * @param stateTable The state tab;e to use in this kev/value state. May contain initial state.
+ */
+ public HeapListState(
+ KeyedStateBackend<K> backend,
+ ListStateDescriptor<V> stateDesc,
+ StateTable<K, N, ArrayList<V>> stateTable,
+ TypeSerializer<K> keySerializer,
+ TypeSerializer<N> namespaceSerializer) {
+ super(backend, stateDesc, stateTable, keySerializer, namespaceSerializer);
+ }
+
+ @Override
+ public Iterable<V> get() {
+ Preconditions.checkState(currentNamespace != null, "No namespace set.");
+ Preconditions.checkState(backend.getCurrentKey() != null, "No key set.");
+
+ Map<N, Map<K, ArrayList<V>>> namespaceMap =
+ stateTable.get(backend.getCurrentKeyGroupIndex());
+
+ if (namespaceMap == null) {
+ return null;
+ }
+
+ Map<K, ArrayList<V>> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ return null;
+ }
+
+ return keyedMap.get(backend.<K>getCurrentKey());
+ }
+
+ @Override
+ public void add(V value) {
+ Preconditions.checkState(currentNamespace != null, "No namespace set.");
+ Preconditions.checkState(backend.getCurrentKey() != null, "No key set.");
+
+ if (value == null) {
+ clear();
+ return;
+ }
+
+ Map<N, Map<K, ArrayList<V>>> namespaceMap =
+ stateTable.get(backend.getCurrentKeyGroupIndex());
+
+ if (namespaceMap == null) {
+ namespaceMap = createNewMap();
+ stateTable.set(backend.getCurrentKeyGroupIndex(), namespaceMap);
+ }
+
+ Map<K, ArrayList<V>> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ keyedMap = createNewMap();
+ namespaceMap.put(currentNamespace, keyedMap);
+ }
+
+ ArrayList<V> list = keyedMap.get(backend.<K>getCurrentKey());
+
+ if (list == null) {
+ list = new ArrayList<>();
+ keyedMap.put(backend.<K>getCurrentKey(), list);
+ }
+ list.add(value);
+ }
+
+ @Override
+ public byte[] getSerializedValue(K key, N namespace) throws Exception {
+ Preconditions.checkState(namespace != null, "No namespace given.");
+ Preconditions.checkState(key != null, "No key given.");
+
+ Map<N, Map<K, ArrayList<V>>> namespaceMap =
+ stateTable.get(backend.getKeyGroupAssigner().getKeyGroupIndex(key));
+
+ if (namespaceMap == null) {
+ return null;
+ }
+
+ Map<K, ArrayList<V>> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ return null;
+ }
+
+ ArrayList<V> result = keyedMap.get(key);
+
+ if (result == null) {
+ return null;
+ }
+
+ TypeSerializer<V> serializer = stateDesc.getSerializer();
+
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ DataOutputViewStreamWrapper view = new DataOutputViewStreamWrapper(baos);
+
+ // write the same as RocksDB writes lists, with one ',' separator
+ for (int i = 0; i < result.size(); i++) {
+ serializer.serialize(result.get(i), view);
+ if (i < result.size() -1) {
+ view.writeByte(',');
+ }
+ }
+ view.flush();
+
+ return baos.toByteArray();
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapReducingState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapReducingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapReducingState.java
new file mode 100644
index 0000000..37aa812
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapReducingState.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.heap;
+
+import org.apache.flink.api.common.functions.ReduceFunction;
+import org.apache.flink.api.common.state.ReducingState;
+import org.apache.flink.api.common.state.ReducingStateDescriptor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.util.Preconditions;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * Heap-backed partitioned {@link org.apache.flink.api.common.state.ReducingState} that is
+ * snapshotted into files.
+ *
+ * @param <K> The type of the key.
+ * @param <N> The type of the namespace.
+ * @param <V> The type of the value.
+ */
+public class HeapReducingState<K, N, V>
+ extends AbstractHeapState<K, N, V, ReducingState<V>, ReducingStateDescriptor<V>>
+ implements ReducingState<V> {
+
+ private final ReduceFunction<V> reduceFunction;
+
+ /**
+ * Creates a new key/value state for the given hash map of key/value pairs.
+ *
+ * @param backend The state backend backing that created this state.
+ * @param stateDesc The state identifier for the state. This contains name
+ * and can create a default state value.
+ * @param stateTable The state tab;e to use in this kev/value state. May contain initial state.
+ */
+ public HeapReducingState(
+ KeyedStateBackend<K> backend,
+ ReducingStateDescriptor<V> stateDesc,
+ StateTable<K, N, V> stateTable,
+ TypeSerializer<K> keySerializer,
+ TypeSerializer<N> namespaceSerializer) {
+ super(backend, stateDesc, stateTable, keySerializer, namespaceSerializer);
+ this.reduceFunction = stateDesc.getReduceFunction();
+ }
+
+ @Override
+ public V get() {
+ Preconditions.checkState(currentNamespace != null, "No namespace set.");
+ Preconditions.checkState(backend.getCurrentKey() != null, "No key set.");
+
+ Map<N, Map<K, V>> namespaceMap =
+ stateTable.get(backend.getCurrentKeyGroupIndex());
+
+ if (namespaceMap == null) {
+ return null;
+ }
+
+ Map<K, V> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ return null;
+ }
+
+ return keyedMap.get(backend.<K>getCurrentKey());
+ }
+
+ @Override
+ public void add(V value) throws IOException {
+ Preconditions.checkState(currentNamespace != null, "No namespace set.");
+ Preconditions.checkState(backend.getCurrentKey() != null, "No key set.");
+
+ if (value == null) {
+ clear();
+ return;
+ }
+
+ Map<N, Map<K, V>> namespaceMap =
+ stateTable.get(backend.getCurrentKeyGroupIndex());
+
+ if (namespaceMap == null) {
+ namespaceMap = createNewMap();
+ stateTable.set(backend.getCurrentKeyGroupIndex(), namespaceMap);
+ }
+
+ Map<K, V> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ keyedMap = createNewMap();
+ namespaceMap.put(currentNamespace, keyedMap);
+ }
+
+ V currentValue = keyedMap.put(backend.<K>getCurrentKey(), value);
+
+ if (currentValue == null) {
+ // we're good, just added the new value
+ } else {
+ V reducedValue = null;
+ try {
+ reducedValue = reduceFunction.reduce(currentValue, value);
+ } catch (Exception e) {
+ throw new RuntimeException("Could not add value to reducing state.", e);
+ }
+ keyedMap.put(backend.<K>getCurrentKey(), reducedValue);
+ }
+ }
+}
[02/27] flink git commit: [FLINK-4380] Introduce KeyGroupAssigner and
Max-Parallelism Parameter
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/DataStreamTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/DataStreamTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/DataStreamTest.java
index 5a86c5c..17bea68 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/DataStreamTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/DataStreamTest.java
@@ -52,7 +52,7 @@ import org.apache.flink.streaming.api.windowing.triggers.PurgingTrigger;
import org.apache.flink.streaming.api.windowing.windows.GlobalWindow;
import org.apache.flink.streaming.runtime.partitioner.BroadcastPartitioner;
import org.apache.flink.streaming.runtime.partitioner.CustomPartitionerWrapper;
-import org.apache.flink.streaming.runtime.partitioner.HashPartitioner;
+import org.apache.flink.streaming.runtime.partitioner.KeyGroupStreamPartitioner;
import org.apache.flink.streaming.runtime.partitioner.ForwardPartitioner;
import org.apache.flink.streaming.runtime.partitioner.GlobalPartitioner;
import org.apache.flink.streaming.runtime.partitioner.RebalancePartitioner;
@@ -672,7 +672,7 @@ public class DataStreamTest {
assertNotNull(env.getStreamGraph().getStreamNode(sink2.getTransformation().getId()).getStateKeySerializer());
assertNotNull(env.getStreamGraph().getStreamNode(sink2.getTransformation().getId()).getStateKeySerializer());
assertEquals(key1, env.getStreamGraph().getStreamNode(sink2.getTransformation().getId()).getStatePartitioner1());
- assertTrue(env.getStreamGraph().getStreamNode(sink2.getTransformation().getId()).getInEdges().get(0).getPartitioner() instanceof HashPartitioner);
+ assertTrue(env.getStreamGraph().getStreamNode(sink2.getTransformation().getId()).getInEdges().get(0).getPartitioner() instanceof KeyGroupStreamPartitioner);
KeySelector<Long, Long> key2 = new KeySelector<Long, Long>() {
@@ -688,7 +688,7 @@ public class DataStreamTest {
assertTrue(env.getStreamGraph().getStreamNode(sink3.getTransformation().getId()).getStatePartitioner1() != null);
assertEquals(key2, env.getStreamGraph().getStreamNode(sink3.getTransformation().getId()).getStatePartitioner1());
- assertTrue(env.getStreamGraph().getStreamNode(sink3.getTransformation().getId()).getInEdges().get(0).getPartitioner() instanceof HashPartitioner);
+ assertTrue(env.getStreamGraph().getStreamNode(sink3.getTransformation().getId()).getInEdges().get(0).getPartitioner() instanceof KeyGroupStreamPartitioner);
}
@Test
@@ -783,7 +783,7 @@ public class DataStreamTest {
private static boolean isPartitioned(List<StreamEdge> edges) {
boolean result = true;
for (StreamEdge edge: edges) {
- if (!(edge.getPartitioner() instanceof HashPartitioner)) {
+ if (!(edge.getPartitioner() instanceof KeyGroupStreamPartitioner)) {
result = false;
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/RestartStrategyTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/RestartStrategyTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/RestartStrategyTest.java
index c57bea7..d6fcd61 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/RestartStrategyTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/RestartStrategyTest.java
@@ -22,11 +22,11 @@ import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.graph.StreamGraph;
-
+import org.apache.flink.util.TestLogger;
import org.junit.Assert;
import org.junit.Test;
-public class RestartStrategyTest {
+public class RestartStrategyTest extends TestLogger {
/**
* Tests that in a streaming use case where checkpointing is enabled, a
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/SlotAllocationTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/SlotAllocationTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/SlotAllocationTest.java
index bab43fa..d873771 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/SlotAllocationTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/SlotAllocationTest.java
@@ -28,6 +28,7 @@ import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoMapFunction;
+import org.apache.flink.util.TestLogger;
import org.junit.Test;
/**
@@ -37,7 +38,7 @@ import org.junit.Test;
* resource groups/slot sharing groups.
*/
@SuppressWarnings("serial")
-public class SlotAllocationTest {
+public class SlotAllocationTest extends TestLogger {
@Test
public void testTwoPipelines() {
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
index a4ee18e..06d381f 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
@@ -19,13 +19,17 @@
package org.apache.flink.streaming.api.graph;
import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
+import org.apache.flink.streaming.api.functions.co.CoMapFunction;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.Output;
@@ -34,8 +38,10 @@ import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.partitioner.BroadcastPartitioner;
import org.apache.flink.streaming.runtime.partitioner.GlobalPartitioner;
+import org.apache.flink.streaming.runtime.partitioner.KeyGroupStreamPartitioner;
import org.apache.flink.streaming.runtime.partitioner.RebalancePartitioner;
import org.apache.flink.streaming.runtime.partitioner.ShufflePartitioner;
+import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
import org.apache.flink.streaming.util.EvenOddOutputSelector;
@@ -236,6 +242,207 @@ public class StreamGraphGeneratorTest {
assertEquals(BasicTypeInfo.INT_TYPE_INFO, outputTypeConfigurableOperation.getTypeInformation());
}
+ /**
+ * Tests that the KeyGroupStreamPartitioner are properly set up with the correct value of
+ * maximum parallelism.
+ */
+ @Test
+ public void testSetupOfKeyGroupPartitioner() {
+ int maxParallelism = 42;
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.getConfig().setMaxParallelism(maxParallelism);
+
+ DataStream<Integer> source = env.fromElements(1, 2, 3);
+
+ DataStream<Integer> keyedResult = source.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 9205556348021992189L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap());
+
+ keyedResult.addSink(new DiscardingSink<Integer>());
+
+ StreamGraph graph = env.getStreamGraph();
+
+ StreamNode keyedResultNode = graph.getStreamNode(keyedResult.getId());
+
+ StreamPartitioner<?> streamPartitioner = keyedResultNode.getInEdges().get(0).getPartitioner();
+
+ HashKeyGroupAssigner<?> hashKeyGroupAssigner = extractHashKeyGroupAssigner(streamPartitioner);
+
+ assertEquals(maxParallelism, hashKeyGroupAssigner.getNumberKeyGroups());
+ }
+
+ /**
+ * Tests that the global and operator-wide max parallelism setting is respected
+ */
+ @Test
+ public void testMaxParallelismForwarding() {
+ int globalMaxParallelism = 42;
+ int keyedResult2MaxParallelism = 17;
+
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.getConfig().setMaxParallelism(globalMaxParallelism);
+
+ DataStream<Integer> source = env.fromElements(1, 2, 3);
+
+ DataStream<Integer> keyedResult1 = source.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 9205556348021992189L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap());
+
+ DataStream<Integer> keyedResult2 = keyedResult1.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 1250168178707154838L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap()).setMaxParallelism(keyedResult2MaxParallelism);
+
+ keyedResult2.addSink(new DiscardingSink<Integer>());
+
+ StreamGraph graph = env.getStreamGraph();
+
+ StreamNode keyedResult1Node = graph.getStreamNode(keyedResult1.getId());
+ StreamNode keyedResult2Node = graph.getStreamNode(keyedResult2.getId());
+
+ assertEquals(globalMaxParallelism, keyedResult1Node.getMaxParallelism());
+ assertEquals(keyedResult2MaxParallelism, keyedResult2Node.getMaxParallelism());
+ }
+
+ /**
+ * Tests that the max parallelism is automatically set to the parallelism if it has not been
+ * specified.
+ */
+ @Test
+ public void testAutoMaxParallelism() {
+ int globalParallelism = 42;
+ int mapParallelism = 17;
+ int maxParallelism = 21;
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(globalParallelism);
+
+ DataStream<Integer> source = env.fromElements(1, 2, 3);
+
+ DataStream<Integer> keyedResult1 = source.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 9205556348021992189L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap());
+
+ DataStream<Integer> keyedResult2 = keyedResult1.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 1250168178707154838L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap()).setParallelism(mapParallelism);
+
+ DataStream<Integer> keyedResult3 = keyedResult2.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 1250168178707154838L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap()).setMaxParallelism(maxParallelism);
+
+ DataStream<Integer> keyedResult4 = keyedResult3.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 1250168178707154838L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap()).setMaxParallelism(maxParallelism).setParallelism(mapParallelism);
+
+ keyedResult4.addSink(new DiscardingSink<Integer>());
+
+ StreamGraph graph = env.getStreamGraph();
+
+ StreamNode keyedResult1Node = graph.getStreamNode(keyedResult1.getId());
+ StreamNode keyedResult2Node = graph.getStreamNode(keyedResult2.getId());
+ StreamNode keyedResult3Node = graph.getStreamNode(keyedResult3.getId());
+ StreamNode keyedResult4Node = graph.getStreamNode(keyedResult4.getId());
+
+ assertEquals(globalParallelism, keyedResult1Node.getMaxParallelism());
+ assertEquals(mapParallelism, keyedResult2Node.getMaxParallelism());
+ assertEquals(maxParallelism, keyedResult3Node.getMaxParallelism());
+ assertEquals(maxParallelism, keyedResult4Node.getMaxParallelism());
+ }
+
+ /**
+ * Tests that the max parallelism and the key group partitioner is properly set for connected
+ * streams.
+ */
+ @Test
+ public void testMaxParallelismWithConnectedKeyedStream() {
+ int maxParallelism = 42;
+
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ DataStream<Integer> input1 = env.fromElements(1, 2, 3, 4).setMaxParallelism(128);
+ DataStream<Integer> input2 = env.fromElements(1, 2, 3, 4).setMaxParallelism(129);
+
+ env.getConfig().setMaxParallelism(maxParallelism);
+
+ DataStream<Integer> keyedResult = input1.connect(input2).keyBy(
+ new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = -6908614081449363419L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ },
+ new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 3195683453223164931L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntCoMap());
+
+ keyedResult.addSink(new DiscardingSink<Integer>());
+
+ StreamGraph graph = env.getStreamGraph();
+
+ StreamNode keyedResultNode = graph.getStreamNode(keyedResult.getId());
+
+ StreamPartitioner<?> streamPartitioner1 = keyedResultNode.getInEdges().get(0).getPartitioner();
+ StreamPartitioner<?> streamPartitioner2 = keyedResultNode.getInEdges().get(1).getPartitioner();
+
+ HashKeyGroupAssigner<?> hashKeyGroupAssigner1 = extractHashKeyGroupAssigner(streamPartitioner1);
+ assertEquals(maxParallelism, hashKeyGroupAssigner1.getNumberKeyGroups());
+
+ HashKeyGroupAssigner<?> hashKeyGroupAssigner2 = extractHashKeyGroupAssigner(streamPartitioner2);
+ assertEquals(maxParallelism, hashKeyGroupAssigner2.getNumberKeyGroups());
+ }
+
+ private HashKeyGroupAssigner<?> extractHashKeyGroupAssigner(StreamPartitioner<?> streamPartitioner) {
+ assertTrue(streamPartitioner instanceof KeyGroupStreamPartitioner);
+
+ KeyGroupStreamPartitioner<?, ?> keyGroupStreamPartitioner = (KeyGroupStreamPartitioner<?, ?>) streamPartitioner;
+
+ KeyGroupAssigner<?> keyGroupAssigner = keyGroupStreamPartitioner.getKeyGroupAssigner();
+
+ assertTrue(keyGroupAssigner instanceof HashKeyGroupAssigner);
+
+ return (HashKeyGroupAssigner<?>) keyGroupAssigner;
+ }
+
private static class OutputTypeConfigurableOperationWithTwoInputs
extends AbstractStreamOperator<Integer>
implements TwoInputStreamOperator<Integer, Integer, Integer>, OutputTypeConfigurable<Integer> {
@@ -297,4 +504,17 @@ public class StreamGraphGeneratorTest {
}
}
+ static class NoOpIntCoMap implements CoMapFunction<Integer, Integer, Integer> {
+ private static final long serialVersionUID = 1886595528149124270L;
+
+ public Integer map1(Integer value) throws Exception {
+ return value;
+ }
+
+ public Integer map2(Integer value) throws Exception {
+ return value;
+ }
+
+ };
+
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java
index 7f94aa0..277fab4 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java
@@ -18,24 +18,33 @@
package org.apache.flink.streaming.api.graph;
import java.io.IOException;
+import java.util.List;
import java.util.Random;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.functions.MapFunction;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
+import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.jobgraph.JobVertex;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
+import org.apache.flink.streaming.util.NoOpIntMap;
import org.apache.flink.util.InstantiationUtil;
import org.apache.flink.util.SerializedValue;
+import org.apache.flink.util.TestLogger;
import org.junit.Test;
import static org.junit.Assert.*;
@SuppressWarnings("serial")
-public class StreamingJobGraphGeneratorTest {
+public class StreamingJobGraphGeneratorTest extends TestLogger {
@Test
public void testExecutionConfigSerialization() throws IOException, ClassNotFoundException {
@@ -114,6 +123,8 @@ public class StreamingJobGraphGeneratorTest {
DataStream<Tuple2<String, String>> input = env
.fromElements("a", "b", "c", "d", "e", "f")
.map(new MapFunction<String, Tuple2<String, String>>() {
+ private static final long serialVersionUID = 471891682418382583L;
+
@Override
public Tuple2<String, String> map(String value) {
return new Tuple2<>(value, value);
@@ -124,6 +135,8 @@ public class StreamingJobGraphGeneratorTest {
.keyBy(0)
.map(new MapFunction<Tuple2<String, String>, Tuple2<String, String>>() {
+ private static final long serialVersionUID = 3583760206245136188L;
+
@Override
public Tuple2<String, String> map(Tuple2<String, String> value) {
return value;
@@ -131,6 +144,8 @@ public class StreamingJobGraphGeneratorTest {
});
result.addSink(new SinkFunction<Tuple2<String, String>>() {
+ private static final long serialVersionUID = -5614849094269539342L;
+
@Override
public void invoke(Tuple2<String, String> value) {}
});
@@ -145,4 +160,203 @@ public class StreamingJobGraphGeneratorTest {
assertEquals(1, jobGraph.getVerticesAsArray()[0].getParallelism());
assertEquals(1, jobGraph.getVerticesAsArray()[1].getParallelism());
}
+
+ /**
+ * Tests that the KeyGroupAssigner is properly set in the {@link StreamConfig} if the max
+ * parallelism is set for the whole job.
+ */
+ @Test
+ public void testKeyGroupAssignerProperlySet() {
+ int maxParallelism = 42;
+
+ final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.getConfig().setMaxParallelism(maxParallelism);
+
+ DataStream<Integer> input = env.fromElements(1, 2, 3);
+
+ DataStream<Integer> keyedResult = input.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 350461576474507944L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap());
+
+ keyedResult.addSink(new DiscardingSink<Integer>());
+
+ JobGraph jobGraph = env.getStreamGraph().getJobGraph();
+
+ List<JobVertex> jobVertices = jobGraph.getVerticesSortedTopologicallyFromSources();
+
+ assertEquals(maxParallelism, jobVertices.get(1).getMaxParallelism());
+
+ HashKeyGroupAssigner<Integer> hashKeyGroupAssigner = extractHashKeyGroupAssigner(jobVertices.get(1));
+
+ assertEquals(maxParallelism, hashKeyGroupAssigner.getNumberKeyGroups());
+ }
+
+ /**
+ * Tests that the key group assigner for the keyed streams in the stream config is properly
+ * initialized with the max parallelism value if there is no max parallelism defined for the
+ * whole job.
+ */
+ @Test
+ public void testKeyGroupAssignerProperlySetAutoMaxParallelism() {
+ int globalParallelism = 42;
+ int mapParallelism = 17;
+ int maxParallelism = 43;
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(globalParallelism);
+
+ DataStream<Integer> source = env.fromElements(1, 2, 3);
+
+ DataStream<Integer> keyedResult1 = source.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 9205556348021992189L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap());
+
+ DataStream<Integer> keyedResult2 = keyedResult1.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 1250168178707154838L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap()).setParallelism(mapParallelism);
+
+ DataStream<Integer> keyedResult3 = keyedResult2.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 1250168178707154838L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap()).setMaxParallelism(maxParallelism);
+
+ DataStream<Integer> keyedResult4 = keyedResult3.keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 1250168178707154838L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new NoOpIntMap()).setMaxParallelism(maxParallelism).setParallelism(mapParallelism);
+
+ keyedResult4.addSink(new DiscardingSink<Integer>());
+
+ JobGraph jobGraph = env.getStreamGraph().getJobGraph();
+
+ List<JobVertex> vertices = jobGraph.getVerticesSortedTopologicallyFromSources();
+
+ JobVertex keyedResultJV1 = vertices.get(1);
+ JobVertex keyedResultJV2 = vertices.get(2);
+ JobVertex keyedResultJV3 = vertices.get(3);
+ JobVertex keyedResultJV4 = vertices.get(4);
+
+ HashKeyGroupAssigner<Integer> hashKeyGroupAssigner1 = extractHashKeyGroupAssigner(keyedResultJV1);
+ HashKeyGroupAssigner<Integer> hashKeyGroupAssigner2 = extractHashKeyGroupAssigner(keyedResultJV2);
+ HashKeyGroupAssigner<Integer> hashKeyGroupAssigner3 = extractHashKeyGroupAssigner(keyedResultJV3);
+ HashKeyGroupAssigner<Integer> hashKeyGroupAssigner4 = extractHashKeyGroupAssigner(keyedResultJV4);
+
+ assertEquals(globalParallelism, hashKeyGroupAssigner1.getNumberKeyGroups());
+ assertEquals(mapParallelism, hashKeyGroupAssigner2.getNumberKeyGroups());
+ assertEquals(maxParallelism, hashKeyGroupAssigner3.getNumberKeyGroups());
+ assertEquals(maxParallelism, hashKeyGroupAssigner4.getNumberKeyGroups());
+ }
+
+ /**
+ * Tests that the {@link KeyGroupAssigner} is properly set in the {@link StreamConfig} for
+ * connected streams.
+ */
+ @Test
+ public void testMaxParallelismWithConnectedKeyedStream() {
+ int maxParallelism = 42;
+
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ DataStream<Integer> input1 = env.fromElements(1, 2, 3, 4).setMaxParallelism(128).name("input1");
+ DataStream<Integer> input2 = env.fromElements(1, 2, 3, 4).setMaxParallelism(129).name("input2");
+
+ env.getConfig().setMaxParallelism(maxParallelism);
+
+ DataStream<Integer> keyedResult = input1.connect(input2).keyBy(
+ new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = -6908614081449363419L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ },
+ new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = 3195683453223164931L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ }).map(new StreamGraphGeneratorTest.NoOpIntCoMap());
+
+ keyedResult.addSink(new DiscardingSink<Integer>());
+
+ JobGraph jobGraph = env.getStreamGraph().getJobGraph();
+
+ List<JobVertex> jobVertices = jobGraph.getVerticesSortedTopologicallyFromSources();
+
+ JobVertex input1JV = jobVertices.get(0);
+ JobVertex input2JV = jobVertices.get(1);
+ JobVertex connectedJV = jobVertices.get(2);
+
+ // disambiguate the partial order of the inputs
+ if (input1JV.getName().equals("Source: input1")) {
+ assertEquals(128, input1JV.getMaxParallelism());
+ assertEquals(129, input2JV.getMaxParallelism());
+ } else {
+ assertEquals(128, input2JV.getMaxParallelism());
+ assertEquals(129, input1JV.getMaxParallelism());
+ }
+
+ assertEquals(maxParallelism, connectedJV.getMaxParallelism());
+
+ HashKeyGroupAssigner<Integer> hashKeyGroupAssigner = extractHashKeyGroupAssigner(connectedJV);
+
+ assertEquals(maxParallelism, hashKeyGroupAssigner.getNumberKeyGroups());
+ }
+
+ /**
+ * Tests that the {@link JobGraph} creation fails if the parallelism is greater than the max
+ * parallelism.
+ */
+ @Test(expected=IllegalStateException.class)
+ public void testFailureOfJobJobCreationIfParallelismGreaterThanMaxParallelism() {
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.getConfig().setMaxParallelism(42);
+
+ DataStream<Integer> input = env.fromElements(1, 2, 3, 4);
+
+ DataStream<Integer> result = input.map(new NoOpIntMap()).setParallelism(43);
+
+ result.addSink(new DiscardingSink<Integer>());
+
+ env.getStreamGraph().getJobGraph();
+
+ fail("The JobGraph should not have been created because the parallelism is greater than " +
+ "the max parallelism.");
+ }
+
+ private HashKeyGroupAssigner<Integer> extractHashKeyGroupAssigner(JobVertex jobVertex) {
+ Configuration config = jobVertex.getConfiguration();
+
+ StreamConfig streamConfig = new StreamConfig(config);
+
+ KeyGroupAssigner<Integer> keyGroupAssigner = streamConfig.getKeyGroupAssigner(getClass().getClassLoader());
+
+ assertTrue(keyGroupAssigner instanceof HashKeyGroupAssigner);
+
+ return (HashKeyGroupAssigner<Integer>) keyGroupAssigner;
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/graph/StreamingJobGraphGeneratorNodeHashTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/graph/StreamingJobGraphGeneratorNodeHashTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/graph/StreamingJobGraphGeneratorNodeHashTest.java
index bcf621a..340981b 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/graph/StreamingJobGraphGeneratorNodeHashTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/graph/StreamingJobGraphGeneratorNodeHashTest.java
@@ -27,12 +27,11 @@ import org.apache.flink.runtime.jobgraph.JobVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
-import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.functions.source.ParallelSourceFunction;
import org.apache.flink.streaming.api.graph.StreamGraph;
import org.apache.flink.streaming.api.graph.StreamNode;
-
+import org.apache.flink.util.TestLogger;
import org.junit.Test;
import java.util.HashMap;
@@ -52,7 +51,7 @@ import static org.junit.Assert.assertTrue;
* {@link JobGraph} instances.
*/
@SuppressWarnings("serial")
-public class StreamingJobGraphGeneratorNodeHashTest {
+public class StreamingJobGraphGeneratorNodeHashTest extends TestLogger {
// ------------------------------------------------------------------------
// Deterministic hash assignment
@@ -126,53 +125,6 @@ public class StreamingJobGraphGeneratorNodeHashTest {
}
/**
- * Verifies that parallelism affects the node hash.
- */
- @Test
- public void testNodeHashParallelism() throws Exception {
- StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironment();
- env.disableOperatorChaining();
-
- env.addSource(new NoOpSourceFunction(), "src").setParallelism(4)
- .addSink(new DiscardingSink<String>()).name("sink").setParallelism(4);
-
- JobGraph jobGraph = env.getStreamGraph().getJobGraph();
-
- Map<JobVertexID, String> ids = rememberIds(jobGraph);
-
- // Change parallelism of source
- env = StreamExecutionEnvironment.createLocalEnvironment();
- env.disableOperatorChaining();
-
- env.addSource(new NoOpSourceFunction(), "src").setParallelism(8)
- .addSink(new DiscardingSink<String>()).name("sink").setParallelism(4);
-
- jobGraph = env.getStreamGraph().getJobGraph();
-
- verifyIdsNotEqual(jobGraph, ids);
-
- // Change parallelism of sink
- env = StreamExecutionEnvironment.createLocalEnvironment();
- env.disableOperatorChaining();
-
- env.addSource(new NoOpSourceFunction(), "src").setParallelism(4)
- .addSink(new DiscardingSink<String>()).name("sink").setParallelism(8);
-
- jobGraph = env.getStreamGraph().getJobGraph();
-
- // The source hash will should be the same
- JobVertex[] vertices = jobGraph.getVerticesAsArray();
- if (vertices[0].isInputVertex()) {
- assertTrue(ids.containsKey(vertices[0].getID()));
- assertFalse(ids.containsKey(vertices[1].getID()));
- }
- else {
- assertTrue(ids.containsKey(vertices[1].getID()));
- assertFalse(ids.containsKey(vertices[0].getID()));
- }
- }
-
- /**
* Tests that there are no collisions with two identical sources.
*
* <pre>
@@ -516,6 +468,8 @@ public class StreamingJobGraphGeneratorNodeHashTest {
private static class NoOpSourceFunction implements ParallelSourceFunction<String> {
+ private static final long serialVersionUID = -5459224792698512636L;
+
@Override
public void run(SourceContext<String> ctx) throws Exception {
}
@@ -527,6 +481,8 @@ public class StreamingJobGraphGeneratorNodeHashTest {
private static class NoOpSinkFunction implements SinkFunction<String> {
+ private static final long serialVersionUID = -5654199886203297279L;
+
@Override
public void invoke(String value) throws Exception {
}
@@ -534,6 +490,8 @@ public class StreamingJobGraphGeneratorNodeHashTest {
private static class NoOpMapFunction implements MapFunction<String, String> {
+ private static final long serialVersionUID = 6584823409744624276L;
+
@Override
public String map(String value) throws Exception {
return value;
@@ -542,6 +500,8 @@ public class StreamingJobGraphGeneratorNodeHashTest {
private static class NoOpFilterFunction implements FilterFunction<String> {
+ private static final long serialVersionUID = 500005424900187476L;
+
@Override
public boolean filter(String value) throws Exception {
return true;
@@ -550,6 +510,8 @@ public class StreamingJobGraphGeneratorNodeHashTest {
private static class NoOpKeySelector implements KeySelector<String, String> {
+ private static final long serialVersionUID = -96127515593422991L;
+
@Override
public String getKey(String value) throws Exception {
return value;
@@ -557,6 +519,8 @@ public class StreamingJobGraphGeneratorNodeHashTest {
}
private static class NoOpReduceFunction implements ReduceFunction<String> {
+ private static final long serialVersionUID = -8775747640749256372L;
+
@Override
public String reduce(String value1, String value2) throws Exception {
return value1;
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AllWindowTranslationTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AllWindowTranslationTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AllWindowTranslationTest.java
index ebe6bea..7ac9e13 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AllWindowTranslationTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AllWindowTranslationTest.java
@@ -253,6 +253,8 @@ public class AllWindowTranslationTest {
try {
windowedStream.fold("", new FoldFunction<String, String>() {
+ private static final long serialVersionUID = -8722899157560218917L;
+
@Override
public String fold(String accumulator, String value) throws Exception {
return accumulator;
@@ -278,6 +280,8 @@ public class AllWindowTranslationTest {
try {
windowedStream.trigger(new Trigger<String, TimeWindow>() {
+ private static final long serialVersionUID = 8360971631424870421L;
+
@Override
public TriggerResult onElement(String element,
long timestamp,
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowTranslationTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowTranslationTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowTranslationTest.java
index 39d89cf..2707108 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowTranslationTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowTranslationTest.java
@@ -76,6 +76,8 @@ public class WindowTranslationTest {
.keyBy(0)
.window(SlidingEventTimeWindows.of(Time.of(1, TimeUnit.SECONDS), Time.of(100, TimeUnit.MILLISECONDS)))
.reduce(new RichReduceFunction<Tuple2<String, Integer>>() {
+ private static final long serialVersionUID = -6448847205314995812L;
+
@Override
public Tuple2<String, Integer> reduce(Tuple2<String, Integer> value1,
Tuple2<String, Integer> value2) throws Exception {
@@ -242,6 +244,8 @@ public class WindowTranslationTest {
WindowedStream<String, String, TimeWindow> windowedStream = env.fromElements("Hello", "Ciao")
.keyBy(new KeySelector<String, String>() {
+ private static final long serialVersionUID = -3298887124448443076L;
+
@Override
public String getKey(String value) throws Exception {
return value;
@@ -251,6 +255,8 @@ public class WindowTranslationTest {
try {
windowedStream.fold("", new FoldFunction<String, String>() {
+ private static final long serialVersionUID = -4567902917104921706L;
+
@Override
public String fold(String accumulator, String value) throws Exception {
return accumulator;
@@ -273,6 +279,8 @@ public class WindowTranslationTest {
WindowedStream<String, String, TimeWindow> windowedStream = env.fromElements("Hello", "Ciao")
.keyBy(new KeySelector<String, String>() {
+ private static final long serialVersionUID = 598309916882894293L;
+
@Override
public String getKey(String value) throws Exception {
return value;
@@ -282,6 +290,8 @@ public class WindowTranslationTest {
try {
windowedStream.trigger(new Trigger<String, TimeWindow>() {
+ private static final long serialVersionUID = 6558046711583024443L;
+
@Override
public TriggerResult onElement(String element,
long timestamp,
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/HashPartitionerTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/HashPartitionerTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/HashPartitionerTest.java
deleted file mode 100644
index 6dbf932..0000000
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/HashPartitionerTest.java
+++ /dev/null
@@ -1,71 +0,0 @@
-/*
- * 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.partitioner;
-
-import static org.junit.Assert.assertArrayEquals;
-import static org.junit.Assert.assertEquals;
-
-import org.apache.flink.api.java.functions.KeySelector;
-import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.runtime.plugable.SerializationDelegate;
-import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.junit.Before;
-import org.junit.Test;
-
-public class HashPartitionerTest {
-
- private HashPartitioner<Tuple2<String, Integer>> hashPartitioner;
- private StreamRecord<Tuple2<String, Integer>> streamRecord1 = new StreamRecord<Tuple2<String, Integer>>(new Tuple2<String, Integer>("test", 0));
- private StreamRecord<Tuple2<String, Integer>> streamRecord2 = new StreamRecord<Tuple2<String, Integer>>(new Tuple2<String, Integer>("test", 42));
- private SerializationDelegate<StreamRecord<Tuple2<String, Integer>>> sd1 = new SerializationDelegate<StreamRecord<Tuple2<String, Integer>>>(null);
- private SerializationDelegate<StreamRecord<Tuple2<String, Integer>>> sd2 = new SerializationDelegate<StreamRecord<Tuple2<String, Integer>>>(null);
-
- @Before
- public void setPartitioner() {
- hashPartitioner = new HashPartitioner<Tuple2<String, Integer>>(new KeySelector<Tuple2<String, Integer>, String>() {
-
- private static final long serialVersionUID = 1L;
-
- @Override
- public String getKey(Tuple2<String, Integer> value) throws Exception {
- return value.getField(0);
- }
- });
- }
-
- @Test
- public void testSelectChannelsLength() {
- sd1.setInstance(streamRecord1);
- assertEquals(1, hashPartitioner.selectChannels(sd1, 1).length);
- assertEquals(1, hashPartitioner.selectChannels(sd1, 2).length);
- assertEquals(1, hashPartitioner.selectChannels(sd1, 1024).length);
- }
-
- @Test
- public void testSelectChannelsGrouping() {
- sd1.setInstance(streamRecord1);
- sd2.setInstance(streamRecord2);
-
- assertArrayEquals(hashPartitioner.selectChannels(sd1, 1),
- hashPartitioner.selectChannels(sd2, 1));
- assertArrayEquals(hashPartitioner.selectChannels(sd1, 2),
- hashPartitioner.selectChannels(sd2, 2));
- assertArrayEquals(hashPartitioner.selectChannels(sd1, 1024),
- hashPartitioner.selectChannels(sd2, 1024));
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitionerTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitionerTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitionerTest.java
new file mode 100644
index 0000000..6fbf35e
--- /dev/null
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitionerTest.java
@@ -0,0 +1,74 @@
+/*
+ * 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.partitioner;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.plugable.SerializationDelegate;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.util.TestLogger;
+import org.junit.Before;
+import org.junit.Test;
+
+public class KeyGroupStreamPartitionerTest extends TestLogger {
+
+ private KeyGroupStreamPartitioner<Tuple2<String, Integer>, String> keyGroupPartitioner;
+ private StreamRecord<Tuple2<String, Integer>> streamRecord1 = new StreamRecord<Tuple2<String, Integer>>(new Tuple2<String, Integer>("test", 0));
+ private StreamRecord<Tuple2<String, Integer>> streamRecord2 = new StreamRecord<Tuple2<String, Integer>>(new Tuple2<String, Integer>("test", 42));
+ private SerializationDelegate<StreamRecord<Tuple2<String, Integer>>> sd1 = new SerializationDelegate<StreamRecord<Tuple2<String, Integer>>>(null);
+ private SerializationDelegate<StreamRecord<Tuple2<String, Integer>>> sd2 = new SerializationDelegate<StreamRecord<Tuple2<String, Integer>>>(null);
+
+ @Before
+ public void setPartitioner() {
+ keyGroupPartitioner = new KeyGroupStreamPartitioner<Tuple2<String, Integer>, String>(new KeySelector<Tuple2<String, Integer>, String>() {
+
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public String getKey(Tuple2<String, Integer> value) throws Exception {
+ return value.getField(0);
+ }
+ },
+ new HashKeyGroupAssigner<String>(1024));
+ }
+
+ @Test
+ public void testSelectChannelsLength() {
+ sd1.setInstance(streamRecord1);
+ assertEquals(1, keyGroupPartitioner.selectChannels(sd1, 1).length);
+ assertEquals(1, keyGroupPartitioner.selectChannels(sd1, 2).length);
+ assertEquals(1, keyGroupPartitioner.selectChannels(sd1, 1024).length);
+ }
+
+ @Test
+ public void testSelectChannelsGrouping() {
+ sd1.setInstance(streamRecord1);
+ sd2.setInstance(streamRecord2);
+
+ assertArrayEquals(keyGroupPartitioner.selectChannels(sd1, 1),
+ keyGroupPartitioner.selectChannels(sd2, 1));
+ assertArrayEquals(keyGroupPartitioner.selectChannels(sd1, 2),
+ keyGroupPartitioner.selectChannels(sd2, 2));
+ assertArrayEquals(keyGroupPartitioner.selectChannels(sd1, 1024),
+ keyGroupPartitioner.selectChannels(sd2, 1024));
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/RescalePartitionerTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/RescalePartitionerTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/RescalePartitionerTest.java
index 8c7360a..37ea68a 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/RescalePartitionerTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/RescalePartitionerTest.java
@@ -94,6 +94,8 @@ public class RescalePartitionerTest extends TestLogger {
// get input data
DataStream<String> text = env.addSource(new ParallelSourceFunction<String>() {
+ private static final long serialVersionUID = 7772338606389180774L;
+
@Override
public void run(SourceContext<String> ctx) throws Exception {
@@ -108,6 +110,8 @@ public class RescalePartitionerTest extends TestLogger {
DataStream<Tuple2<String, Integer>> counts = text
.rescale()
.flatMap(new FlatMapFunction<String, Tuple2<String, Integer>>() {
+ private static final long serialVersionUID = -5255930322161596829L;
+
@Override
public void flatMap(String value,
Collector<Tuple2<String, Integer>> out) throws Exception {
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java
index 145edc2..5f73e25 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java
@@ -23,6 +23,7 @@ import org.apache.flink.api.java.ClosureCleaner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.runtime.event.AbstractEvent;
import org.apache.flink.runtime.io.network.partition.consumer.StreamTestSingleInputGate;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import java.io.IOException;
@@ -105,6 +106,7 @@ public class OneInputStreamTaskTestHarness<IN, OUT> extends StreamTaskTestHarnes
ClosureCleaner.clean(keySelector, false);
streamConfig.setStatePartitioner(0, keySelector);
streamConfig.setStateKeySerializer(keyType.createSerializer(executionConfig));
+ streamConfig.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(10));
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
index bcd8a5f..3d9d50f 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
@@ -19,8 +19,11 @@
package org.apache.flink.streaming.runtime.tasks;
import akka.actor.ActorRef;
+
+import akka.dispatch.Futures;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.blob.BlobKey;
import org.apache.flink.runtime.broadcast.BroadcastVariableManager;
@@ -40,9 +43,11 @@ import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.jobgraph.tasks.AbstractInvokable;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.metrics.groups.TaskMetricGroup;
+import org.apache.flink.runtime.messages.TaskMessages;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.taskmanager.Task;
import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
+import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.api.operators.Output;
@@ -51,17 +56,27 @@ import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.SerializedValue;
import org.junit.Test;
+
+import scala.concurrent.Await;
import scala.concurrent.ExecutionContext;
import scala.concurrent.Future;
+import scala.concurrent.duration.Deadline;
import scala.concurrent.duration.FiniteDuration;
+import scala.concurrent.impl.Promise;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.net.URL;
import java.util.Collections;
+import java.util.Comparator;
+import java.util.PriorityQueue;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
@@ -72,55 +87,140 @@ import static org.mockito.Mockito.when;
public class StreamTaskTest {
- /**
+ /**
* This test checks that cancel calls that are issued before the operator is
* instantiated still lead to proper canceling.
*/
@Test
- public void testEarlyCanceling() {
- try {
- StreamConfig cfg = new StreamConfig(new Configuration());
- cfg.setStreamOperator(new SlowlyDeserializingOperator());
-
- Task task = createTask(SourceStreamTask.class, cfg);
- task.startTaskThread();
-
- // wait until the task thread reached state RUNNING
- while (task.getExecutionState() == ExecutionState.CREATED ||
- task.getExecutionState() == ExecutionState.DEPLOYING)
- {
- Thread.sleep(5);
- }
-
- // make sure the task is really running
- if (task.getExecutionState() != ExecutionState.RUNNING) {
- fail("Task entered state " + task.getExecutionState() + " with error "
- + ExceptionUtils.stringifyException(task.getFailureCause()));
- }
-
- // send a cancel. because the operator takes a long time to deserialize, this should
- // hit the task before the operator is deserialized
- task.cancelExecution();
-
- // the task should reach state canceled eventually
- assertTrue(task.getExecutionState() == ExecutionState.CANCELING ||
- task.getExecutionState() == ExecutionState.CANCELED);
-
- task.getExecutingThread().join(60000);
-
- assertFalse("Task did not cancel", task.getExecutingThread().isAlive());
- assertEquals(ExecutionState.CANCELED, task.getExecutionState());
- }
- catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
+ public void testEarlyCanceling() throws Exception {
+ Deadline deadline = new FiniteDuration(2, TimeUnit.MINUTES).fromNow();
+ StreamConfig cfg = new StreamConfig(new Configuration());
+ cfg.setStreamOperator(new SlowlyDeserializingOperator());
+ cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);
+
+ Task task = createTask(SourceStreamTask.class, cfg);
+
+ ExecutionStateListener executionStateListener = new ExecutionStateListener();
+
+ task.registerExecutionListener(executionStateListener);
+ task.startTaskThread();
+
+ Future<ExecutionState> running = executionStateListener.notifyWhenExecutionState(ExecutionState.RUNNING);
+
+ // wait until the task thread reached state RUNNING
+ ExecutionState executionState = Await.result(running, deadline.timeLeft());
+
+ // make sure the task is really running
+ if (executionState != ExecutionState.RUNNING) {
+ fail("Task entered state " + task.getExecutionState() + " with error "
+ + ExceptionUtils.stringifyException(task.getFailureCause()));
}
+
+ // send a cancel. because the operator takes a long time to deserialize, this should
+ // hit the task before the operator is deserialized
+ task.cancelExecution();
+
+ Future<ExecutionState> canceling = executionStateListener.notifyWhenExecutionState(ExecutionState.CANCELING);
+
+ executionState = Await.result(canceling, deadline.timeLeft());
+
+ // the task should reach state canceled eventually
+ assertTrue(executionState == ExecutionState.CANCELING ||
+ executionState == ExecutionState.CANCELED);
+
+ task.getExecutingThread().join(deadline.timeLeft().toMillis());
+
+ assertFalse("Task did not cancel", task.getExecutingThread().isAlive());
+ assertEquals(ExecutionState.CANCELED, task.getExecutionState());
}
+
// ------------------------------------------------------------------------
// Test Utilities
// ------------------------------------------------------------------------
+ private static class ExecutionStateListener implements ActorGateway {
+
+ private static final long serialVersionUID = 8926442805035692182L;
+
+ ExecutionState executionState = null;
+
+ PriorityQueue<Tuple2<ExecutionState, Promise<ExecutionState>>> priorityQueue = new PriorityQueue<>(
+ 1,
+ new Comparator<Tuple2<ExecutionState, Promise<ExecutionState>>>() {
+ @Override
+ public int compare(Tuple2<ExecutionState, Promise<ExecutionState>> o1, Tuple2<ExecutionState, Promise<ExecutionState>> o2) {
+ return o1.f0.ordinal() - o2.f0.ordinal();
+ }
+ });
+
+ public Future<ExecutionState> notifyWhenExecutionState(ExecutionState executionState) {
+ synchronized (priorityQueue) {
+ if (this.executionState != null && this.executionState.ordinal() >= executionState.ordinal()) {
+ return Futures.successful(executionState);
+ } else {
+ Promise<ExecutionState> promise = new Promise.DefaultPromise<ExecutionState>();
+
+ priorityQueue.offer(Tuple2.of(executionState, promise));
+
+ return promise.future();
+ }
+ }
+ }
+
+ @Override
+ public Future<Object> ask(Object message, FiniteDuration timeout) {
+ return null;
+ }
+
+ @Override
+ public void tell(Object message) {
+ this.tell(message, null);
+ }
+
+ @Override
+ public void tell(Object message, ActorGateway sender) {
+ if (message instanceof TaskMessages.UpdateTaskExecutionState) {
+ TaskMessages.UpdateTaskExecutionState updateTaskExecutionState = (TaskMessages.UpdateTaskExecutionState) message;
+
+ synchronized (priorityQueue) {
+ this.executionState = updateTaskExecutionState.taskExecutionState().getExecutionState();
+
+ while (!priorityQueue.isEmpty() && priorityQueue.peek().f0.ordinal() <= this.executionState.ordinal()) {
+ Promise<ExecutionState> promise = priorityQueue.poll().f1;
+
+ promise.success(this.executionState);
+ }
+ }
+ }
+ }
+
+ @Override
+ public void forward(Object message, ActorGateway sender) {
+
+ }
+
+ @Override
+ public Future<Object> retry(Object message, int numberRetries, FiniteDuration timeout, ExecutionContext executionContext) {
+ return null;
+ }
+
+ @Override
+ public String path() {
+ return null;
+ }
+
+ @Override
+ public ActorRef actor() {
+ return null;
+ }
+
+ @Override
+ public UUID leaderSessionID() {
+ return null;
+ }
+ }
+
private Task createTask(Class<? extends AbstractInvokable> invokable, StreamConfig taskConfig) throws Exception {
LibraryCacheManager libCache = mock(LibraryCacheManager.class);
when(libCache.getClassLoader(any(JobID.class))).thenReturn(getClass().getClassLoader());
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTestHarness.java
index 00e95b9..cb10c5c 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTestHarness.java
@@ -186,23 +186,58 @@ public class StreamTaskTestHarness<OUT> {
taskThread.start();
}
+ /**
+ * Waits for the task completion.
+ *
+ * @throws Exception
+ */
public void waitForTaskCompletion() throws Exception {
+ waitForTaskCompletion(Long.MAX_VALUE);
+ }
+
+ /**
+ * Waits for the task completion. If this does not happen within the timeout, then a
+ * TimeoutException is thrown.
+ *
+ * @param timeout Timeout for the task completion
+ * @throws Exception
+ */
+ public void waitForTaskCompletion(long timeout) throws Exception {
if (taskThread == null) {
throw new IllegalStateException("Task thread was not started.");
}
- taskThread.join();
+ taskThread.join(timeout);
if (taskThread.getError() != null) {
throw new Exception("error in task", taskThread.getError());
}
}
+ /**
+ * Waits for the task to be running.
+ *
+ * @throws Exception
+ */
public void waitForTaskRunning() throws Exception {
+ waitForTaskRunning(Long.MAX_VALUE);
+ }
+
+ /**
+ * Waits fro the task to be running. If this does not happen within the timeout, then a
+ * TimeoutException is thrown.
+ *
+ * @param timeout Timeout for the task to be running.
+ * @throws Exception
+ */
+ public void waitForTaskRunning(long timeout) throws Exception {
if (taskThread == null) {
throw new IllegalStateException("Task thread was not started.");
}
else {
if (taskThread.task instanceof StreamTask) {
+ long base = System.currentTimeMillis();
+ long now = 0;
+
StreamTask<?, ?> streamTask = (StreamTask<?, ?>) taskThread.task;
while (!streamTask.isRunning()) {
Thread.sleep(100);
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/resources/log4j-test.properties
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/resources/log4j-test.properties b/flink-streaming-java/src/test/resources/log4j-test.properties
index 0b686e5..881dc06 100644
--- a/flink-streaming-java/src/test/resources/log4j-test.properties
+++ b/flink-streaming-java/src/test/resources/log4j-test.properties
@@ -24,4 +24,4 @@ log4j.appender.A1=org.apache.log4j.ConsoleAppender
# A1 uses PatternLayout.
log4j.appender.A1.layout=org.apache.log4j.PatternLayout
-log4j.appender.A1.layout.ConversionPattern=%-4r [%t] %-5p %c %x - %m%n
\ No newline at end of file
+log4j.appender.A1.layout.ConversionPattern=%-4r [%t] %-5p %c %x - %m%n
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/DataStream.scala
----------------------------------------------------------------------
diff --git a/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/DataStream.scala b/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/DataStream.scala
index 8693834..4fe73e9 100644
--- a/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/DataStream.scala
+++ b/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/DataStream.scala
@@ -133,6 +133,17 @@ class DataStream[T](stream: JavaStream[T]) {
this
}
+ def setMaxParallelism(maxParallelism: Int): DataStream[T] = {
+ stream match {
+ case ds: SingleOutputStreamOperator[T] => ds.setMaxParallelism(maxParallelism)
+ case _ =>
+ throw new UnsupportedOperationException("Operator " + stream + " cannot set the maximum" +
+ "paralllelism")
+ }
+
+ this
+ }
+
/**
* Gets the name of the current data stream. This name is
* used by the visualization and logging during runtime.
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/StreamExecutionEnvironment.scala
----------------------------------------------------------------------
diff --git a/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/StreamExecutionEnvironment.scala b/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/StreamExecutionEnvironment.scala
index 9cb36a5..2e432ba 100644
--- a/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/StreamExecutionEnvironment.scala
+++ b/flink-streaming-scala/src/main/scala/org/apache/flink/streaming/api/scala/StreamExecutionEnvironment.scala
@@ -59,12 +59,30 @@ class StreamExecutionEnvironment(javaEnv: JavaEnv) {
}
/**
+ * Sets the maximum degree of parallelism defined for the program.
+ * The maximum degree of parallelism specifies the upper limit for dynamic scaling. It also
+ * defines the number of key groups used for partitioned state.
+ **/
+ def setMaxParallelism(maxParallelism: Int): Unit = {
+ javaEnv.setMaxParallelism(maxParallelism)
+ }
+
+ /**
* Returns the default parallelism for this execution environment. Note that this
* value can be overridden by individual operations using [[DataStream#setParallelism(int)]]
*/
def getParallelism = javaEnv.getParallelism
/**
+ * Returns the maximum degree of parallelism defined for the program.
+ *
+ * The maximum degree of parallelism specifies the upper limit for dynamic scaling. It also
+ * defines the number of key groups used for partitioned state.
+ *
+ */
+ def getMaxParallelism = javaEnv.getMaxParallelism
+
+ /**
* Sets the maximum time frequency (milliseconds) for the flushing of the
* output buffers. By default the output buffers flush frequently to provide
* low latency and to aid smooth developer experience. Setting the parameter
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/DataStreamTest.scala
----------------------------------------------------------------------
diff --git a/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/DataStreamTest.scala b/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/DataStreamTest.scala
index 16fcfc3..b73eae8 100644
--- a/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/DataStreamTest.scala
+++ b/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/DataStreamTest.scala
@@ -512,7 +512,7 @@ class DataStreamTest extends StreamingMultipleProgramsTestBase {
private def isPartitioned(edges: java.util.List[StreamEdge]): Boolean = {
import scala.collection.JavaConverters._
- edges.asScala.forall( _.getPartitioner.isInstanceOf[HashPartitioner[_]])
+ edges.asScala.forall( _.getPartitioner.isInstanceOf[KeyGroupStreamPartitioner[_, _]])
}
private def isCustomPartitioned(edges: java.util.List[StreamEdge]): Boolean = {
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeAllWindowCheckpointingITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeAllWindowCheckpointingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeAllWindowCheckpointingITCase.java
index 6faee45..163fb42 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeAllWindowCheckpointingITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeAllWindowCheckpointingITCase.java
@@ -162,7 +162,7 @@ public class EventTimeAllWindowCheckpointingITCase extends TestLogger {
env.setParallelism(PARALLELISM);
env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);
env.enableCheckpointing(100);
- env.setRestartStrategy(RestartStrategies.fixedDelayRestart(3, 0));
+ env.setRestartStrategy(RestartStrategies.fixedDelayRestart(1, 0));
env.getConfig().disableSysoutLogging();
env
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
new file mode 100644
index 0000000..0de2a75
--- /dev/null
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
@@ -0,0 +1,683 @@
+/*
+ * 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.test.checkpointing;
+
+import io.netty.util.internal.ConcurrentSet;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.functions.RichFlatMapFunction;
+import org.apache.flink.api.common.restartstrategy.RestartStrategies;
+import org.apache.flink.api.common.state.ValueState;
+import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.configuration.ConfigConstants;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.runtime.checkpoint.savepoint.SavepointStoreFactory;
+import org.apache.flink.runtime.client.JobExecutionException;
+import org.apache.flink.runtime.execution.SuppressRestartsException;
+import org.apache.flink.runtime.instance.ActorGateway;
+import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.messages.JobManagerMessages;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.filesystem.FsStateBackendFactory;
+import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages;
+import org.apache.flink.streaming.api.checkpoint.Checkpointed;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
+import org.apache.flink.streaming.api.functions.sink.SinkFunction;
+import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
+import org.apache.flink.test.util.ForkableFlinkMiniCluster;
+import org.apache.flink.util.Collector;
+import org.apache.flink.util.TestLogger;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import scala.concurrent.Await;
+import scala.concurrent.Future;
+import scala.concurrent.duration.Deadline;
+import scala.concurrent.duration.FiniteDuration;
+
+import java.io.File;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+public class RescalingITCase extends TestLogger {
+
+ private static int numTaskManagers = 2;
+ private static int slotsPerTaskManager = 2;
+ private static int numSlots = numTaskManagers * slotsPerTaskManager;
+
+ private static ForkableFlinkMiniCluster cluster;
+
+ @ClassRule
+ public static TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ @BeforeClass
+ public static void setup() throws Exception {
+ Configuration config = new Configuration();
+ config.setInteger(ConfigConstants.LOCAL_NUMBER_TASK_MANAGER, numTaskManagers);
+ config.setInteger(ConfigConstants.TASK_MANAGER_NUM_TASK_SLOTS, slotsPerTaskManager);
+
+ final File checkpointDir = temporaryFolder.newFolder();
+ final File savepointDir = temporaryFolder.newFolder();
+
+ config.setString(ConfigConstants.STATE_BACKEND, "filesystem");
+ config.setString(FsStateBackendFactory.CHECKPOINT_DIRECTORY_URI_CONF_KEY, checkpointDir.toURI().toString());
+ config.setString(SavepointStoreFactory.SAVEPOINT_BACKEND_KEY, "filesystem");
+ config.setString(SavepointStoreFactory.SAVEPOINT_DIRECTORY_KEY, savepointDir.toURI().toString());
+
+ cluster = new ForkableFlinkMiniCluster(config);
+ cluster.start();
+ }
+
+ @AfterClass
+ public static void teardown() {
+ if (cluster != null) {
+ cluster.shutdown();
+ }
+ }
+
+ /**
+ * Tests that a a job with purely partitioned state can be restarted from a savepoint
+ * with a different parallelism.
+ */
+ @Test
+ public void testSavepointRescalingWithPartitionedState() throws Exception {
+ int numberKeys = 42;
+ int numberElements = 1000;
+ int numberElements2 = 500;
+ int parallelism = numSlots / 2;
+ int parallelism2 = numSlots;
+ int maxParallelism = 13;
+
+ FiniteDuration timeout = new FiniteDuration(3, TimeUnit.MINUTES);
+ Deadline deadline = timeout.fromNow();
+
+ ActorGateway jobManager = null;
+ JobID jobID = null;
+
+ try {
+ jobManager = cluster.getLeaderGateway(deadline.timeLeft());
+
+ JobGraph jobGraph = createPartitionedStateJobGraph(parallelism, maxParallelism, numberKeys, numberElements, false, 100);
+
+ jobID = jobGraph.getJobID();
+
+ cluster.submitJobDetached(jobGraph);
+
+ // wait til the sources have emitted numberElements for each key and completed a checkpoint
+ SubtaskIndexFlatMapper.workCompletedLatch.await(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);
+
+ // verify the current state
+
+ Set<Tuple2<Integer, Integer>> actualResult = CollectionSink.getElementsSet();
+
+ Set<Tuple2<Integer, Integer>> expectedResult = new HashSet<>();
+
+ HashKeyGroupAssigner<Integer> keyGroupAssigner = new HashKeyGroupAssigner<>(maxParallelism);
+
+ for (int key = 0; key < numberKeys; key++) {
+ int keyGroupIndex = keyGroupAssigner.getKeyGroupIndex(key);
+
+ expectedResult.add(Tuple2.of(keyGroupIndex % parallelism, numberElements * key));
+ }
+
+ assertEquals(expectedResult, actualResult);
+
+ // clear the CollectionSink set for the restarted job
+ CollectionSink.clearElementsSet();
+
+ Future<Object> savepointPathFuture = jobManager.ask(new JobManagerMessages.TriggerSavepoint(jobID), deadline.timeLeft());
+
+ final String savepointPath = ((JobManagerMessages.TriggerSavepointSuccess)
+ Await.result(savepointPathFuture, deadline.timeLeft())).savepointPath();
+
+ Future<Object> jobRemovedFuture = jobManager.ask(new TestingJobManagerMessages.NotifyWhenJobRemoved(jobID), deadline.timeLeft());
+
+ Future<Object> cancellationResponseFuture = jobManager.ask(new JobManagerMessages.CancelJob(jobID), deadline.timeLeft());
+
+ Object cancellationResponse = Await.result(cancellationResponseFuture, deadline.timeLeft());
+
+ assertTrue(cancellationResponse instanceof JobManagerMessages.CancellationSuccess);
+
+ Await.ready(jobRemovedFuture, deadline.timeLeft());
+
+ jobID = null;
+
+ JobGraph scaledJobGraph = createPartitionedStateJobGraph(parallelism2, maxParallelism, numberKeys, numberElements2, true, 100);
+
+ scaledJobGraph.setSavepointPath(savepointPath);
+
+ jobID = scaledJobGraph.getJobID();
+
+ cluster.submitJobAndWait(scaledJobGraph, false);
+
+ jobID = null;
+
+ Set<Tuple2<Integer, Integer>> actualResult2 = CollectionSink.getElementsSet();
+
+ Set<Tuple2<Integer, Integer>> expectedResult2 = new HashSet<>();
+
+ HashKeyGroupAssigner<Integer> keyGroupAssigner2 = new HashKeyGroupAssigner<>(maxParallelism);
+
+ for (int key = 0; key < numberKeys; key++) {
+ int keyGroupIndex = keyGroupAssigner2.getKeyGroupIndex(key);
+ expectedResult2.add(Tuple2.of(keyGroupIndex % parallelism2, key * (numberElements + numberElements2)));
+ }
+
+ assertEquals(expectedResult2, actualResult2);
+
+ } finally {
+ // clear the CollectionSink set for the restarted job
+ CollectionSink.clearElementsSet();
+
+ // clear any left overs from a possibly failed job
+ if (jobID != null && jobManager != null) {
+ Future<Object> jobRemovedFuture = jobManager.ask(new TestingJobManagerMessages.NotifyWhenJobRemoved(jobID), timeout);
+
+ try {
+ Await.ready(jobRemovedFuture, timeout);
+ } catch (TimeoutException | InterruptedException ie) {
+ fail("Failed while cleaning up the cluster.");
+ }
+ }
+ }
+ }
+
+ /**
+ * Tests that a job cannot be restarted from a savepoint with a different parallelism if the
+ * rescaled operator has non-partitioned state.
+ *
+ * @throws Exception
+ */
+ @Test
+ public void testSavepointRescalingFailureWithNonPartitionedState() throws Exception {
+ int parallelism = numSlots / 2;
+ int parallelism2 = numSlots;
+ int maxParallelism = 13;
+
+ FiniteDuration timeout = new FiniteDuration(3, TimeUnit.MINUTES);
+ Deadline deadline = timeout.fromNow();
+
+ JobID jobID = null;
+ ActorGateway jobManager = null;
+
+ try {
+ jobManager = cluster.getLeaderGateway(deadline.timeLeft());
+
+ JobGraph jobGraph = createNonPartitionedStateJobGraph(parallelism, maxParallelism, 500);
+
+ jobID = jobGraph.getJobID();
+
+ cluster.submitJobDetached(jobGraph);
+
+ Future<Object> allTasksRunning = jobManager.ask(new TestingJobManagerMessages.WaitForAllVerticesToBeRunning(jobID), deadline.timeLeft());
+
+ Await.ready(allTasksRunning, deadline.timeLeft());
+
+ Future<Object> savepointPathFuture = jobManager.ask(new JobManagerMessages.TriggerSavepoint(jobID), deadline.timeLeft());
+
+ Object savepointResponse = Await.result(savepointPathFuture, deadline.timeLeft());
+
+ assertTrue(savepointResponse instanceof JobManagerMessages.TriggerSavepointSuccess);
+
+ final String savepointPath = ((JobManagerMessages.TriggerSavepointSuccess)savepointResponse).savepointPath();
+
+ Future<Object> jobRemovedFuture = jobManager.ask(new TestingJobManagerMessages.NotifyWhenJobRemoved(jobID), deadline.timeLeft());
+
+ Future<Object> cancellationResponseFuture = jobManager.ask(new JobManagerMessages.CancelJob(jobID), deadline.timeLeft());
+
+ Object cancellationResponse = Await.result(cancellationResponseFuture, deadline.timeLeft());
+
+ assertTrue(cancellationResponse instanceof JobManagerMessages.CancellationSuccess);
+
+ Await.ready(jobRemovedFuture, deadline.timeLeft());
+
+ // job successfully removed
+ jobID = null;
+
+ JobGraph scaledJobGraph = createNonPartitionedStateJobGraph(parallelism2, maxParallelism, 500);
+
+ scaledJobGraph.setSavepointPath(savepointPath);
+
+ jobID = scaledJobGraph.getJobID();
+
+ cluster.submitJobAndWait(scaledJobGraph, false);
+
+ jobID = null;
+
+ } catch (JobExecutionException exception) {
+ if (exception.getCause() instanceof SuppressRestartsException) {
+ SuppressRestartsException suppressRestartsException = (SuppressRestartsException) exception.getCause();
+
+ if (suppressRestartsException.getCause() instanceof IllegalStateException) {
+ // we expect a IllegalStateException wrapped in a SuppressRestartsException wrapped
+ // in a JobExecutionException, because the job containing non-partitioned state
+ // is being rescaled
+ } else {
+ throw exception;
+ }
+ } else {
+ throw exception;
+ }
+ } finally {
+ // clear any left overs from a possibly failed job
+ if (jobID != null && jobManager != null) {
+ Future<Object> jobRemovedFuture = jobManager.ask(new TestingJobManagerMessages.NotifyWhenJobRemoved(jobID), timeout);
+
+ try {
+ Await.ready(jobRemovedFuture, timeout);
+ } catch (TimeoutException | InterruptedException ie) {
+ fail("Failed while cleaning up the cluster.");
+ }
+ }
+ }
+ }
+
+ /**
+ * Tests that a job with non partitioned state can be restarted from a savepoint with a
+ * different parallelism if the operator with non-partitioned state are not rescaled.
+ *
+ * @throws Exception
+ */
+ @Test
+ public void testSavepointRescalingWithPartiallyNonPartitionedState() throws Exception {
+ int numberKeys = 42;
+ int numberElements = 1000;
+ int numberElements2 = 500;
+ int parallelism = numSlots / 2;
+ int parallelism2 = numSlots;
+ int maxParallelism = 13;
+
+ FiniteDuration timeout = new FiniteDuration(3, TimeUnit.MINUTES);
+ Deadline deadline = timeout.fromNow();
+
+ ActorGateway jobManager = null;
+ JobID jobID = null;
+
+ try {
+ jobManager = cluster.getLeaderGateway(deadline.timeLeft());
+
+ JobGraph jobGraph = createPartitionedNonPartitionedStateJobGraph(
+ parallelism,
+ maxParallelism,
+ parallelism,
+ numberKeys,
+ numberElements,
+ false,
+ 100);
+
+ jobID = jobGraph.getJobID();
+
+ cluster.submitJobDetached(jobGraph);
+
+ // wait til the sources have emitted numberElements for each key and completed a checkpoint
+ SubtaskIndexFlatMapper.workCompletedLatch.await(deadline.timeLeft().toMillis(), TimeUnit.MILLISECONDS);
+
+ // verify the current state
+
+ Set<Tuple2<Integer, Integer>> actualResult = CollectionSink.getElementsSet();
+
+ Set<Tuple2<Integer, Integer>> expectedResult = new HashSet<>();
+
+ HashKeyGroupAssigner<Integer> keyGroupAssigner = new HashKeyGroupAssigner<>(maxParallelism);
+
+ for (int key = 0; key < numberKeys; key++) {
+ int keyGroupIndex = keyGroupAssigner.getKeyGroupIndex(key);
+
+ expectedResult.add(Tuple2.of(keyGroupIndex % parallelism, numberElements * key));
+ }
+
+ assertEquals(expectedResult, actualResult);
+
+ // clear the CollectionSink set for the restarted job
+ CollectionSink.clearElementsSet();
+
+ Future<Object> savepointPathFuture = jobManager.ask(new JobManagerMessages.TriggerSavepoint(jobID), deadline.timeLeft());
+
+ final String savepointPath = ((JobManagerMessages.TriggerSavepointSuccess)
+ Await.result(savepointPathFuture, deadline.timeLeft())).savepointPath();
+
+ Future<Object> jobRemovedFuture = jobManager.ask(new TestingJobManagerMessages.NotifyWhenJobRemoved(jobID), deadline.timeLeft());
+
+ Future<Object> cancellationResponseFuture = jobManager.ask(new JobManagerMessages.CancelJob(jobID), deadline.timeLeft());
+
+ Object cancellationResponse = Await.result(cancellationResponseFuture, deadline.timeLeft());
+
+ assertTrue(cancellationResponse instanceof JobManagerMessages.CancellationSuccess);
+
+ Await.ready(jobRemovedFuture, deadline.timeLeft());
+
+ jobID = null;
+
+ JobGraph scaledJobGraph = createPartitionedNonPartitionedStateJobGraph(
+ parallelism2,
+ maxParallelism,
+ parallelism,
+ numberKeys,
+ numberElements + numberElements2,
+ true,
+ 100);
+
+ scaledJobGraph.setSavepointPath(savepointPath);
+
+ jobID = scaledJobGraph.getJobID();
+
+ cluster.submitJobAndWait(scaledJobGraph, false);
+
+ jobID = null;
+
+ Set<Tuple2<Integer, Integer>> actualResult2 = CollectionSink.getElementsSet();
+
+ Set<Tuple2<Integer, Integer>> expectedResult2 = new HashSet<>();
+
+ HashKeyGroupAssigner<Integer> keyGroupAssigner2 = new HashKeyGroupAssigner<>(maxParallelism);
+
+ for (int key = 0; key < numberKeys; key++) {
+ int keyGroupIndex = keyGroupAssigner2.getKeyGroupIndex(key);
+ expectedResult2.add(Tuple2.of(keyGroupIndex % parallelism2, key * (numberElements + numberElements2)));
+ }
+
+ assertEquals(expectedResult2, actualResult2);
+
+ } finally {
+ // clear the CollectionSink set for the restarted job
+ CollectionSink.clearElementsSet();
+
+ // clear any left overs from a possibly failed job
+ if (jobID != null && jobManager != null) {
+ Future<Object> jobRemovedFuture = jobManager.ask(new TestingJobManagerMessages.NotifyWhenJobRemoved(jobID), timeout);
+
+ try {
+ Await.ready(jobRemovedFuture, timeout);
+ } catch (TimeoutException | InterruptedException ie) {
+ fail("Failed while cleaning up the cluster.");
+ }
+ }
+ }
+ }
+
+ private static JobGraph createNonPartitionedStateJobGraph(int parallelism, int maxParallelism, long checkpointInterval) {
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(parallelism);
+ env.getConfig().setMaxParallelism(maxParallelism);
+ env.enableCheckpointing(checkpointInterval);
+ env.setRestartStrategy(RestartStrategies.noRestart());
+
+ DataStream<Integer> input = env.addSource(new NonPartitionedStateSource());
+
+ input.addSink(new DiscardingSink<Integer>());
+
+ return env.getStreamGraph().getJobGraph();
+ }
+
+ private static JobGraph createPartitionedStateJobGraph(
+ int parallelism,
+ int maxParallelism,
+ int numberKeys,
+ int numberElements,
+ boolean terminateAfterEmission,
+ int checkpointingInterval) {
+
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(parallelism);
+ env.getConfig().setMaxParallelism(maxParallelism);
+ env.enableCheckpointing(checkpointingInterval);
+ env.setRestartStrategy(RestartStrategies.noRestart());
+
+ DataStream<Integer> input = env.addSource(new SubtaskIndexSource(
+ numberKeys,
+ numberElements,
+ terminateAfterEmission))
+ .keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = -7952298871120320940L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ });
+
+ SubtaskIndexFlatMapper.workCompletedLatch = new CountDownLatch(numberKeys);
+
+ DataStream<Tuple2<Integer, Integer>> result = input.flatMap(new SubtaskIndexFlatMapper(numberElements));
+
+ result.addSink(new CollectionSink());
+
+ return env.getStreamGraph().getJobGraph();
+ }
+
+ private static JobGraph createPartitionedNonPartitionedStateJobGraph(
+ int parallelism,
+ int maxParallelism,
+ int fixedParallelism,
+ int numberKeys,
+ int numberElements,
+ boolean terminateAfterEmission,
+ int checkpointingInterval) {
+
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(parallelism);
+ env.getConfig().setMaxParallelism(maxParallelism);
+ env.enableCheckpointing(checkpointingInterval);
+ env.setRestartStrategy(RestartStrategies.noRestart());
+
+ DataStream<Integer> input = env.addSource(new SubtaskIndexNonPartitionedStateSource(
+ numberKeys,
+ numberElements,
+ terminateAfterEmission))
+ .setParallelism(fixedParallelism)
+ .keyBy(new KeySelector<Integer, Integer>() {
+ private static final long serialVersionUID = -7952298871120320940L;
+
+ @Override
+ public Integer getKey(Integer value) throws Exception {
+ return value;
+ }
+ });
+
+ SubtaskIndexFlatMapper.workCompletedLatch = new CountDownLatch(numberKeys);
+
+ DataStream<Tuple2<Integer, Integer>> result = input.flatMap(new SubtaskIndexFlatMapper(numberElements));
+
+ result.addSink(new CollectionSink());
+
+ return env.getStreamGraph().getJobGraph();
+ }
+
+ private static class SubtaskIndexSource
+ extends RichParallelSourceFunction<Integer> {
+
+ private static final long serialVersionUID = -400066323594122516L;
+
+ private final int numberKeys;
+ private final int numberElements;
+ private final boolean terminateAfterEmission;
+
+ protected int counter = 0;
+
+ private boolean running = true;
+
+ SubtaskIndexSource(
+ int numberKeys,
+ int numberElements,
+ boolean terminateAfterEmission) {
+
+ this.numberKeys = numberKeys;
+ this.numberElements = numberElements;
+ this.terminateAfterEmission = terminateAfterEmission;
+ }
+
+ @Override
+ public void run(SourceContext<Integer> ctx) throws Exception {
+ final Object lock = ctx.getCheckpointLock();
+ final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
+
+ while (running) {
+
+ if (counter < numberElements) {
+ synchronized (lock) {
+ for (int value = subtaskIndex;
+ value < numberKeys;
+ value += getRuntimeContext().getNumberOfParallelSubtasks()) {
+
+ ctx.collect(value);
+ }
+
+ counter++;
+ }
+ } else {
+ if (terminateAfterEmission) {
+ running = false;
+ } else {
+ Thread.sleep(100);
+ }
+ }
+ }
+ }
+
+ @Override
+ public void cancel() {
+ running = false;
+ }
+ }
+
+ private static class SubtaskIndexNonPartitionedStateSource extends SubtaskIndexSource implements Checkpointed<Integer> {
+
+ private static final long serialVersionUID = 8388073059042040203L;
+
+ SubtaskIndexNonPartitionedStateSource(int numberKeys, int numberElements, boolean terminateAfterEmission) {
+ super(numberKeys, numberElements, terminateAfterEmission);
+ }
+
+ @Override
+ public Integer snapshotState(long checkpointId, long checkpointTimestamp) throws Exception {
+ return counter;
+ }
+
+ @Override
+ public void restoreState(Integer state) throws Exception {
+ counter = state;
+ }
+ }
+
+ private static class SubtaskIndexFlatMapper extends RichFlatMapFunction<Integer, Tuple2<Integer, Integer>> {
+
+ private static final long serialVersionUID = 5273172591283191348L;
+
+ private static volatile CountDownLatch workCompletedLatch = new CountDownLatch(1);
+
+ private transient ValueState<Integer> counter;
+ private transient ValueState<Integer> sum;
+
+ private final int numberElements;
+
+ SubtaskIndexFlatMapper(int numberElements) {
+ this.numberElements = numberElements;
+ }
+
+ @Override
+ public void open(Configuration configuration) {
+ counter = getRuntimeContext().getState(new ValueStateDescriptor<Integer>("counter", Integer.class, 0));
+ sum = getRuntimeContext().getState(new ValueStateDescriptor<Integer>("sum", Integer.class, 0));
+ }
+
+ @Override
+ public void flatMap(Integer value, Collector<Tuple2<Integer, Integer>> out) throws Exception {
+ int count = counter.value() + 1;
+ counter.update(count);
+
+ int s = sum.value() + value;
+ sum.update(s);
+
+ if (count % numberElements == 0) {
+ out.collect(Tuple2.of(getRuntimeContext().getIndexOfThisSubtask(), s));
+ workCompletedLatch.countDown();
+ }
+ }
+ }
+
+ private static class CollectionSink<IN> implements SinkFunction<IN> {
+
+ private static ConcurrentSet<Object> elements = new ConcurrentSet<Object>();
+
+ private static final long serialVersionUID = -1652452958040267745L;
+
+ public static <IN> Set<IN> getElementsSet() {
+ return (Set<IN>) elements;
+ }
+
+ public static void clearElementsSet() {
+ elements.clear();
+ }
+
+ @Override
+ public void invoke(IN value) throws Exception {
+ elements.add(value);
+ }
+ }
+
+ private static class NonPartitionedStateSource extends RichParallelSourceFunction<Integer> implements Checkpointed<Integer> {
+
+ private static final long serialVersionUID = -8108185918123186841L;
+
+ private int counter = 0;
+ private boolean running = true;
+
+ @Override
+ public Integer snapshotState(long checkpointId, long checkpointTimestamp) throws Exception {
+ return counter;
+ }
+
+ @Override
+ public void restoreState(Integer state) throws Exception {
+ counter = state;
+ }
+
+ @Override
+ public void run(SourceContext<Integer> ctx) throws Exception {
+ final Object lock = ctx.getCheckpointLock();
+
+ while (running) {
+ synchronized (lock) {
+ counter++;
+
+ ctx.collect(counter * getRuntimeContext().getIndexOfThisSubtask());
+ }
+
+ Thread.sleep(100);
+ }
+ }
+
+ @Override
+ public void cancel() {
+ running = true;
+ }
+ }
+}
[11/27] flink git commit: [FLINK-3761] Refactor State Backends/Make
Keyed State Key-Group Aware
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/query/KvStateRegistry.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/query/KvStateRegistry.java b/flink-runtime/src/main/java/org/apache/flink/runtime/query/KvStateRegistry.java
index e09b868..5213fe9 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/query/KvStateRegistry.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/query/KvStateRegistry.java
@@ -40,7 +40,7 @@ import java.util.concurrent.atomic.AtomicReference;
public class KvStateRegistry {
/** All registered KvState instances. */
- private final ConcurrentHashMap<KvStateID, KvState<?, ?, ?, ?, ?>> registeredKvStates =
+ private final ConcurrentHashMap<KvStateID, KvState<?>> registeredKvStates =
new ConcurrentHashMap<>();
/** Registry listener to be notified on registration/unregistration. */
@@ -83,7 +83,7 @@ public class KvStateRegistry {
JobVertexID jobVertexId,
int keyGroupIndex,
String registrationName,
- KvState<?, ?, ?, ?, ?> kvState) {
+ KvState<?> kvState) {
KvStateID kvStateId = new KvStateID();
@@ -136,7 +136,7 @@ public class KvStateRegistry {
* @param kvStateId KvStateID to identify the KvState instance
* @return KvState instance identified by the KvStateID or <code>null</code>
*/
- public KvState<?, ?, ?, ?, ?> getKvState(KvStateID kvStateId) {
+ public KvState<?> getKvState(KvStateID kvStateId) {
return registeredKvStates.get(kvStateId);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/query/TaskKvStateRegistry.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/query/TaskKvStateRegistry.java b/flink-runtime/src/main/java/org/apache/flink/runtime/query/TaskKvStateRegistry.java
index 15f0160..b5c09aa 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/query/TaskKvStateRegistry.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/query/TaskKvStateRegistry.java
@@ -58,7 +58,7 @@ public class TaskKvStateRegistry {
* descriptor used to create the KvState instance)
* @param kvState The
*/
- public void registerKvState(int keyGroupIndex, String registrationName, KvState<?, ?, ?, ?, ?> kvState) {
+ public void registerKvState(int keyGroupIndex, String registrationName, KvState<?> kvState) {
KvStateID kvStateId = registry.registerKvState(jobId, jobVertexId, keyGroupIndex, registrationName, kvState);
registeredKvStates.add(new KvStateInfo(keyGroupIndex, registrationName, kvStateId));
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/query/netty/KvStateServerHandler.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/query/netty/KvStateServerHandler.java b/flink-runtime/src/main/java/org/apache/flink/runtime/query/netty/KvStateServerHandler.java
index 47f2ad6..8201708 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/query/netty/KvStateServerHandler.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/query/netty/KvStateServerHandler.java
@@ -103,7 +103,7 @@ class KvStateServerHandler extends ChannelInboundHandlerAdapter {
stats.reportRequest();
- KvState<?, ?, ?, ?, ?> kvState = registry.getKvState(request.getKvStateId());
+ KvState<?> kvState = registry.getKvState(request.getKvStateId());
if (kvState != null) {
// Execute actual query async, because it is possibly
@@ -186,7 +186,7 @@ class KvStateServerHandler extends ChannelInboundHandlerAdapter {
private final KvStateRequest request;
- private final KvState<?, ?, ?, ?, ?> kvState;
+ private final KvState<?> kvState;
private final KvStateRequestStats stats;
@@ -195,7 +195,7 @@ class KvStateServerHandler extends ChannelInboundHandlerAdapter {
public AsyncKvStateQueryTask(
ChannelHandlerContext ctx,
KvStateRequest request,
- KvState<?, ?, ?, ?, ?> kvState,
+ KvState<?> kvState,
KvStateRequestStats stats) {
this.ctx = Objects.requireNonNull(ctx, "Channel handler context");
@@ -238,6 +238,8 @@ class KvStateServerHandler extends ChannelInboundHandlerAdapter {
success = true;
} else {
+ kvState.getSerializedValue(serializedKeyAndNamespace);
+
// No data for the key/namespace. This is considered to be
// a failure.
ByteBuf unknownKey = KvStateRequestSerializer.serializeKvStateRequestFailure(
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractHeapState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractHeapState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractHeapState.java
deleted file mode 100644
index 6fa4575..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractHeapState.java
+++ /dev/null
@@ -1,220 +0,0 @@
-/*
- * 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.state.ListState;
-import org.apache.flink.api.common.state.State;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.util.Preconditions;
-
-import java.util.HashMap;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-
-/**
- * Base class for partitioned {@link ListState} implementations that are backed by a regular
- * heap hash map. The concrete implementations define how the state is checkpointed.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <SV> The type of the values in the state.
- * @param <S> The type of State
- * @param <SD> The type of StateDescriptor for the State S
- * @param <Backend> The type of the backend that snapshots this key/value state.
- */
-public abstract class AbstractHeapState<K, N, SV, S extends State, SD extends StateDescriptor<S, ?>, Backend extends AbstractStateBackend>
- implements KvState<K, N, S, SD, Backend>, State {
-
- /** Map containing the actual key/value pairs */
- protected final Map<N, Map<K, SV>> state;
-
- /** Serializer for the state value. The state value could be a List<V>, for example. */
- protected final TypeSerializer<SV> stateSerializer;
-
- /** The serializer for the keys */
- protected final TypeSerializer<K> keySerializer;
-
- /** The serializer for the namespace */
- protected final TypeSerializer<N> namespaceSerializer;
-
- /** This holds the name of the state and can create an initial default value for the state. */
- protected final SD stateDesc;
-
- /** The current key, which the next value methods will refer to */
- protected K currentKey;
-
- /** The current namespace, which the access methods will refer to. */
- protected N currentNamespace = null;
-
- /** Cache the state map for the current key. */
- protected Map<K, SV> currentNSState;
-
- /**
- * Creates a new empty key/value state.
- *
- * @param keySerializer The serializer for the keys.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateDesc The state identifier for the state. This contains name
- * and can create a default state value.
- */
- protected AbstractHeapState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<SV> stateSerializer,
- SD stateDesc) {
- this(keySerializer, namespaceSerializer, stateSerializer, stateDesc, new HashMap<N, Map<K, SV>>());
- }
-
- /**
- * Creates a new key/value state for the given hash map of key/value pairs.
- *
- * @param keySerializer The serializer for the keys.
- * @param stateDesc The state identifier for the state. This contains name
- * and can create a default state value.
- * @param state The state map to use in this kev/value state. May contain initial state.
- */
- protected AbstractHeapState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<SV> stateSerializer,
- SD stateDesc,
- Map<N, Map<K, SV>> state) {
-
- Preconditions.checkNotNull(state, "State map");
-
- // Make sure that the state map supports concurrent read access for
- // queries. See also #createNewNamespaceMap for the namespace maps.
- if (stateDesc.isQueryable()) {
- this.state = new ConcurrentHashMap<>(state);
- } else {
- this.state = state;
- }
-
- this.keySerializer = Preconditions.checkNotNull(keySerializer);
- this.namespaceSerializer = Preconditions.checkNotNull(namespaceSerializer);
- this.stateSerializer = stateSerializer;
- this.stateDesc = stateDesc;
- }
-
- // ------------------------------------------------------------------------
-
- @Override
- public final void clear() {
- if (currentNSState != null) {
- currentNSState.remove(currentKey);
- if (currentNSState.isEmpty()) {
- state.remove(currentNamespace);
- currentNSState = null;
- }
- }
- }
-
- @Override
- public final void setCurrentKey(K currentKey) {
- this.currentKey = Preconditions.checkNotNull(currentKey, "Key");
- }
-
- @Override
- public final void setCurrentNamespace(N namespace) {
- if (namespace != null && namespace.equals(this.currentNamespace)) {
- return;
- }
- this.currentNamespace = Preconditions.checkNotNull(namespace, "Namespace");
- this.currentNSState = state.get(currentNamespace);
- }
-
- @Override
- public byte[] getSerializedValue(byte[] serializedKeyAndNamespace) throws Exception {
- Preconditions.checkNotNull(serializedKeyAndNamespace, "Serialized key and namespace");
-
- Tuple2<K, N> keyAndNamespace = KvStateRequestSerializer.deserializeKeyAndNamespace(
- serializedKeyAndNamespace, keySerializer, namespaceSerializer);
-
- return getSerializedValue(keyAndNamespace.f0, keyAndNamespace.f1);
- }
-
- protected abstract byte[] getSerializedValue(K key, N namespace) throws Exception;
-
- /**
- * Returns the number of all state pairs in this state, across namespaces.
- */
- protected final int size() {
- int size = 0;
- for (Map<K, SV> namespace: state.values()) {
- size += namespace.size();
- }
- return size;
- }
-
- @Override
- public void dispose() {
- state.clear();
- }
-
- @Override
- public SD getStateDescriptor() {
- return stateDesc;
- }
-
- /**
- * Gets the serializer for the keys.
- *
- * @return The serializer for the keys.
- */
- public final TypeSerializer<K> getKeySerializer() {
- return keySerializer;
- }
-
- /**
- * Gets the serializer for the namespace.
- *
- * @return The serializer for the namespace.
- */
- public final TypeSerializer<N> getNamespaceSerializer() {
- return namespaceSerializer;
- }
-
- /**
- * Creates a new namespace map.
- *
- * <p>If the state queryable ({@link StateDescriptor#isQueryable()}, this
- * will create a concurrent hash map instead of a regular one.
- *
- * @return A new namespace map.
- */
- protected Map<K, SV> createNewNamespaceMap() {
- if (stateDesc.isQueryable()) {
- return new ConcurrentHashMap<>();
- } else {
- return new HashMap<>();
- }
- }
-
- // ------------------------------------------------------------------------
-
- /**
- * Returns the internal state map for testing.
- *
- * @return The internal state map
- */
- Map<N, Map<K, SV>> getStateMap() {
- return state;
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
index b2cde22..e6093a8 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
@@ -18,417 +18,57 @@
package org.apache.flink.runtime.state;
-import org.apache.flink.api.common.ExecutionConfig;
-import org.apache.flink.api.common.functions.ReduceFunction;
-import org.apache.flink.api.common.state.FoldingState;
-import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.state.ListState;
-import org.apache.flink.api.common.state.ListStateDescriptor;
-import org.apache.flink.api.common.state.MergingState;
-import org.apache.flink.api.common.state.ReducingState;
-import org.apache.flink.api.common.state.ReducingStateDescriptor;
-import org.apache.flink.api.common.state.State;
-import org.apache.flink.api.common.state.StateBackend;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
-import org.apache.flink.util.Preconditions;
import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.HashMap;
import java.util.List;
-import java.util.Map;
/**
* A state backend defines how state is stored and snapshotted during checkpoints.
*/
public abstract class AbstractStateBackend implements java.io.Serializable {
-
- private static final long serialVersionUID = 4620413814639220247L;
-
- protected transient TypeSerializer<?> keySerializer;
-
- protected transient ClassLoader userCodeClassLoader;
-
- protected transient Object currentKey;
-
- /** For efficient access in setCurrentKey() */
- private transient KvState<?, ?, ?, ?, ?>[] keyValueStates;
-
- /** So that we can give out state when the user uses the same key. */
- protected transient HashMap<String, KvState<?, ?, ?, ?, ?>> keyValueStatesByName;
-
- /** For caching the last accessed partitioned state */
- private transient String lastName;
-
- @SuppressWarnings("rawtypes")
- private transient KvState lastState;
-
- /** KvStateRegistry helper for this task */
- protected transient TaskKvStateRegistry kvStateRegistry;
-
- /** Key group index of this state backend */
- protected transient int keyGroupIndex;
-
- // ------------------------------------------------------------------------
- // initialization and cleanup
- // ------------------------------------------------------------------------
-
- /**
- * This method is called by the task upon deployment to initialize the state backend for
- * data for a specific job.
- *
- * @param env The {@link Environment} of the task that instantiated the state backend
- * @param operatorIdentifier Unique identifier for naming states created by this backend
- * @throws Exception Overwritten versions of this method may throw exceptions, in which
- * case the job that uses the state backend is considered failed during
- * deployment.
- */
- public void initializeForJob(
- Environment env,
- String operatorIdentifier,
- TypeSerializer<?> keySerializer) throws Exception {
-
- this.userCodeClassLoader = env.getUserClassLoader();
- this.keySerializer = keySerializer;
-
- this.keyGroupIndex = env.getTaskInfo().getIndexOfThisSubtask();
- this.kvStateRegistry = env.getTaskKvStateRegistry();
- }
-
- /**
- * Disposes all state associated with the current job.
- *
- * @throws Exception Exceptions may occur during disposal of the state and should be forwarded.
- */
- public abstract void disposeAllStateForCurrentJob() throws Exception;
-
- /**
- * Closes the state backend, 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
- */
- public abstract void close() throws Exception;
-
- public void discardState() throws Exception {
- if (kvStateRegistry != null) {
- kvStateRegistry.unregisterAll();
- }
-
- lastName = null;
- lastState = null;
- if (keyValueStates != null) {
- for (KvState<?, ?, ?, ?, ?> state : keyValueStates) {
- state.dispose();
- }
- }
- keyValueStates = null;
- keyValueStatesByName = null;
- }
-
- // ------------------------------------------------------------------------
- // key/value state
- // ------------------------------------------------------------------------
-
- /**
- * Creates and returns a new {@link ValueState}.
- *
- * @param namespaceSerializer TypeSerializer for the state namespace.
- * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
- *
- * @param <N> The type of the namespace.
- * @param <T> The type of the value that the {@code ValueState} can store.
- */
- protected abstract <N, T> ValueState<T> createValueState(TypeSerializer<N> namespaceSerializer, ValueStateDescriptor<T> stateDesc) throws Exception;
-
- /**
- * Creates and returns a new {@link ListState}.
- *
- * @param namespaceSerializer TypeSerializer for the state namespace.
- * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
- *
- * @param <N> The type of the namespace.
- * @param <T> The type of the values that the {@code ListState} can store.
- */
- protected abstract <N, T> ListState<T> createListState(TypeSerializer<N> namespaceSerializer, ListStateDescriptor<T> stateDesc) throws Exception;
-
- /**
- * Creates and returns a new {@link ReducingState}.
- *
- * @param namespaceSerializer TypeSerializer for the state namespace.
- * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
- *
- * @param <N> The type of the namespace.
- * @param <T> The type of the values that the {@code ListState} can store.
- */
- protected abstract <N, T> ReducingState<T> createReducingState(TypeSerializer<N> namespaceSerializer, ReducingStateDescriptor<T> stateDesc) throws Exception;
-
- /**
- * Creates and returns a new {@link FoldingState}.
- *
- * @param namespaceSerializer TypeSerializer for the state namespace.
- * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
- *
- * @param <N> The type of the namespace.
- * @param <T> Type of the values folded into the state
- * @param <ACC> Type of the value in the state *
- */
- protected abstract <N, T, ACC> FoldingState<T, ACC> createFoldingState(TypeSerializer<N> namespaceSerializer, FoldingStateDescriptor<T, ACC> stateDesc) throws Exception;
-
- /**
- * Sets the current key that is used for partitioned state.
- * @param currentKey The current key.
- */
- @SuppressWarnings({"unchecked", "rawtypes"})
- public void setCurrentKey(Object currentKey) {
- this.currentKey = Preconditions.checkNotNull(currentKey, "Key");
- if (keyValueStates != null) {
- for (KvState kv : keyValueStates) {
- kv.setCurrentKey(currentKey);
- }
- }
- }
-
- public Object getCurrentKey() {
- return currentKey;
- }
+ private static final long serialVersionUID = 4620415814639230247L;
/**
- * Creates or retrieves a partitioned state backed by this state backend.
- *
- * @param stateDescriptor The state identifier for the state. This contains name
- * and can create a default state value.
-
- * @param <N> The type of the namespace.
- * @param <S> The type of the state.
+ * Creates a {@link CheckpointStreamFactory} that can be used to create streams
+ * that should end up in a checkpoint.
*
- * @return A new key/value state backed by this backend.
- *
- * @throws Exception Exceptions may occur during initialization of the state and should be forwarded.
+ * @param jobId The {@link JobID} of the job for which we are creating checkpoint streams.
+ * @param operatorIdentifier An identifier of the operator for which we create streams.
*/
- @SuppressWarnings({"rawtypes", "unchecked"})
- public <N, S extends State> S getPartitionedState(final N namespace, final TypeSerializer<N> namespaceSerializer, final StateDescriptor<S, ?> stateDescriptor) throws Exception {
- Preconditions.checkNotNull(namespace, "Namespace");
- Preconditions.checkNotNull(namespaceSerializer, "Namespace serializer");
-
- if (keySerializer == null) {
- throw new RuntimeException("State key serializer has not been configured in the config. " +
- "This operation cannot use partitioned state.");
- }
-
- if (!stateDescriptor.isSerializerInitialized()) {
- stateDescriptor.initializeSerializerUnlessSet(new ExecutionConfig());
- }
-
- if (keyValueStatesByName == null) {
- keyValueStatesByName = new HashMap<>();
- }
-
- if (lastName != null && lastName.equals(stateDescriptor.getName())) {
- lastState.setCurrentNamespace(namespace);
- return (S) lastState;
- }
-
- KvState<?, ?, ?, ?, ?> previous = keyValueStatesByName.get(stateDescriptor.getName());
- if (previous != null) {
- lastState = previous;
- lastState.setCurrentNamespace(namespace);
- lastName = stateDescriptor.getName();
- return (S) previous;
- }
-
- // create a new blank key/value state
- S state = stateDescriptor.bind(new StateBackend() {
- @Override
- public <T> ValueState<T> createValueState(ValueStateDescriptor<T> stateDesc) throws Exception {
- return AbstractStateBackend.this.createValueState(namespaceSerializer, stateDesc);
- }
-
- @Override
- public <T> ListState<T> createListState(ListStateDescriptor<T> stateDesc) throws Exception {
- return AbstractStateBackend.this.createListState(namespaceSerializer, stateDesc);
- }
-
- @Override
- public <T> ReducingState<T> createReducingState(ReducingStateDescriptor<T> stateDesc) throws Exception {
- return AbstractStateBackend.this.createReducingState(namespaceSerializer, stateDesc);
- }
-
- @Override
- public <T, ACC> FoldingState<T, ACC> createFoldingState(FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
- return AbstractStateBackend.this.createFoldingState(namespaceSerializer, stateDesc);
- }
-
- });
-
- KvState kvState = (KvState) state;
-
- keyValueStatesByName.put(stateDescriptor.getName(), kvState);
- keyValueStates = keyValueStatesByName.values().toArray(new KvState[keyValueStatesByName.size()]);
-
- lastName = stateDescriptor.getName();
- lastState = kvState;
-
- if (currentKey != null) {
- kvState.setCurrentKey(currentKey);
- }
-
- kvState.setCurrentNamespace(namespace);
-
- // Publish queryable state
- if (stateDescriptor.isQueryable()) {
- if (kvStateRegistry == null) {
- throw new IllegalStateException("State backend has not been initialized for job.");
- }
-
- String name = stateDescriptor.getQueryableStateName();
- kvStateRegistry.registerKvState(keyGroupIndex, name, kvState);
- }
-
- return state;
- }
-
- @SuppressWarnings("unchecked,rawtypes")
- public <N, S extends MergingState<?, ?>> void mergePartitionedStates(final N target, Collection<N> sources, final TypeSerializer<N> namespaceSerializer, final StateDescriptor<S, ?> stateDescriptor) throws Exception {
- if (stateDescriptor instanceof ReducingStateDescriptor) {
- ReducingStateDescriptor reducingStateDescriptor = (ReducingStateDescriptor) stateDescriptor;
- ReduceFunction reduceFn = reducingStateDescriptor.getReduceFunction();
- ReducingState state = (ReducingState) getPartitionedState(target, namespaceSerializer, stateDescriptor);
- KvState kvState = (KvState) state;
- Object result = null;
- for (N source: sources) {
- kvState.setCurrentNamespace(source);
- Object sourceValue = state.get();
- if (result == null) {
- result = state.get();
- } else if (sourceValue != null) {
- result = reduceFn.reduce(result, sourceValue);
- }
- state.clear();
- }
- kvState.setCurrentNamespace(target);
- if (result != null) {
- state.add(result);
- }
- } else if (stateDescriptor instanceof ListStateDescriptor) {
- ListState<Object> state = (ListState) getPartitionedState(target, namespaceSerializer, stateDescriptor);
- KvState kvState = (KvState) state;
- List<Object> result = new ArrayList<>();
- for (N source: sources) {
- kvState.setCurrentNamespace(source);
- Iterable<Object> sourceValue = state.get();
- if (sourceValue != null) {
- for (Object o : sourceValue) {
- result.add(o);
- }
- }
- state.clear();
- }
- kvState.setCurrentNamespace(target);
- for (Object o : result) {
- state.add(o);
- }
- } else {
- throw new RuntimeException("Cannot merge states for " + stateDescriptor);
- }
- }
-
- public HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshotPartitionedState(long checkpointId, long timestamp) throws Exception {
- if (keyValueStates != null) {
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshots = new HashMap<>(keyValueStatesByName.size());
-
- for (Map.Entry<String, KvState<?, ?, ?, ?, ?>> entry : keyValueStatesByName.entrySet()) {
- KvStateSnapshot<?, ?, ?, ?, ?> snapshot = entry.getValue().snapshot(checkpointId, timestamp);
- snapshots.put(entry.getKey(), snapshot);
- }
- return snapshots;
- }
-
- return null;
- }
-
- public void notifyOfCompletedCheckpoint(long checkpointId) throws Exception {
- // We check whether the KvStates require notifications
- if (keyValueStates != null) {
- for (KvState<?, ?, ?, ?, ?> kvstate : keyValueStates) {
- if (kvstate instanceof CheckpointListener) {
- ((CheckpointListener) kvstate).notifyCheckpointComplete(checkpointId);
- }
- }
- }
- }
-
- /**
- * Injects K/V state snapshots for lazy restore.
- * @param keyValueStateSnapshots The Map of snapshots
- */
- @SuppressWarnings("unchecked,rawtypes")
- public void injectKeyValueStateSnapshots(HashMap<String, KvStateSnapshot> keyValueStateSnapshots) throws Exception {
- if (keyValueStateSnapshots != null) {
- if (keyValueStatesByName == null) {
- keyValueStatesByName = new HashMap<>();
- }
-
- for (Map.Entry<String, KvStateSnapshot> state : keyValueStateSnapshots.entrySet()) {
- KvState kvState = state.getValue().restoreState(this,
- keySerializer,
- userCodeClassLoader);
- keyValueStatesByName.put(state.getKey(), kvState);
-
- try {
- // Publish queryable state
- StateDescriptor stateDesc = kvState.getStateDescriptor();
- if (stateDesc.isQueryable()) {
- String queryableStateName = stateDesc.getQueryableStateName();
- kvStateRegistry.registerKvState(keyGroupIndex, queryableStateName, kvState);
- }
- } catch (Throwable ignored) {
- }
- }
- keyValueStates = keyValueStatesByName.values().toArray(new KvState[keyValueStatesByName.size()]);
- }
- }
-
- // ------------------------------------------------------------------------
- // storing state for a checkpoint
- // ------------------------------------------------------------------------
+ public abstract CheckpointStreamFactory createStreamFactory(
+ JobID jobId,
+ String operatorIdentifier) throws IOException;
/**
- * Creates an output stream that writes into the state of the given checkpoint. When the stream
- * is closes, it returns a state handle that can retrieve the state back.
- *
- * @param checkpointID The ID of the checkpoint.
- * @param timestamp The timestamp of the checkpoint.
- * @return An output stream that writes state for the given checkpoint.
- *
- * @throws Exception Exceptions may occur while creating the stream and should be forwarded.
+ * Creates a new {@link KeyedStateBackend} that is responsible for keeping keyed state
+ * and can be checkpointed to checkpoint streams.
*/
- public abstract CheckpointStateOutputStream createCheckpointStateOutputStream(
- long checkpointID, long timestamp) throws Exception;
-
- // ------------------------------------------------------------------------
- // Checkpoint state output stream
- // ------------------------------------------------------------------------
+ public abstract <K> KeyedStateBackend<K> createKeyedStateBackend(
+ Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ TaskKvStateRegistry kvStateRegistry) throws Exception;
/**
- * A dedicated output stream that produces a {@link StreamStateHandle} when closed.
+ * Creates a new {@link KeyedStateBackend} that restores its state from the given list
+ * {@link KeyGroupsStateHandle KeyGroupStateHandles}.
*/
- public static abstract class CheckpointStateOutputStream extends FSDataOutputStream {
+ public abstract <K> KeyedStateBackend<K> restoreKeyedStateBackend(
+ Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ List<KeyGroupsStateHandle> restoredState,
+ TaskKvStateRegistry kvStateRegistry) throws Exception;
- /**
- * Closes the stream and gets a state handle that can create an input stream
- * producing the data written to this stream.
- *
- * @return A state handle that can create an input stream producing the data written to this stream.
- * @throws IOException Thrown, if the stream cannot be closed.
- */
- public abstract StreamStateHandle closeAndGetHandle() throws IOException;
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousKvStateSnapshot.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousKvStateSnapshot.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousKvStateSnapshot.java
deleted file mode 100644
index c2fc8a4..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousKvStateSnapshot.java
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * 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.state.State;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-
-import java.io.IOException;
-
-/**
- * {@link KvStateSnapshot} that asynchronously materializes the state that it represents. Instead
- * of representing a materialized handle to state this would normally hold the (immutable) state
- * internally and materializes it when {@link #materialize()} is called.
- *
- * @param <K> The type of the key
- * @param <N> The type of the namespace
- * @param <S> The type of the {@link State}
- * @param <SD> The type of the {@link StateDescriptor}
- * @param <Backend> The type of the backend that can restore the state from this snapshot.
- */
-public abstract class AsynchronousKvStateSnapshot<K, N, S extends State, SD extends StateDescriptor<S, ?>, Backend extends AbstractStateBackend> implements KvStateSnapshot<K, N, S, SD, Backend> {
- private static final long serialVersionUID = 1L;
-
- /**
- * Materializes the state held by this {@code AsynchronousKvStateSnapshot}.
- */
- public abstract KvStateSnapshot<K, N, S, SD, Backend> materialize() throws Exception;
-
- @Override
- public final KvState<K, N, S, SD, Backend> restoreState(
- Backend stateBackend,
- TypeSerializer<K> keySerializer,
- ClassLoader classLoader) throws Exception {
- throw new RuntimeException("This should never be called and probably points to a bug.");
- }
-
- @Override
- public void discardState() throws Exception {
- throw new RuntimeException("This should never be called and probably points to a bug.");
- }
-
- @Override
- public long getStateSize() throws Exception {
- throw new RuntimeException("This should never be called and probably points to a bug.");
- }
-
- @Override
- public void close() throws IOException {
- throw new RuntimeException("This should never be called and probably points to a bug.");
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamFactory.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamFactory.java
new file mode 100644
index 0000000..199a856
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/CheckpointStreamFactory.java
@@ -0,0 +1,67 @@
+/*
+ * 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.FSDataOutputStream;
+
+import java.io.IOException;
+import java.io.OutputStream;
+
+public interface CheckpointStreamFactory {
+
+ /**
+ * Creates an new {@link CheckpointStateOutputStream}. When the stream
+ * is closed, it returns a state handle that can retrieve the state back.
+ *
+ * @param checkpointID The ID of the checkpoint.
+ * @param timestamp The timestamp of the checkpoint.
+ *
+ * @return An output stream that writes state for the given checkpoint.
+ *
+ * @throws Exception Exceptions may occur while creating the stream and should be forwarded.
+ */
+ CheckpointStateOutputStream createCheckpointStateOutputStream(
+ long checkpointID,
+ long timestamp) throws Exception;
+
+ /**
+ * 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
+ */
+ void close() throws Exception;
+
+ /**
+ * A dedicated output stream that produces a {@link StreamStateHandle} when closed.
+ *
+ * <p>Note: This is an abstract class and not an interface because {@link OutputStream}
+ * is an abstract class.
+ */
+ abstract class CheckpointStateOutputStream extends FSDataOutputStream {
+
+ /**
+ * Closes the stream and gets a state handle that can create an input stream
+ * producing the data written to this stream.
+ *
+ * @return A state handle that can create an input stream producing the data written to this stream.
+ * @throws IOException Thrown, if the stream cannot be closed.
+ */
+ public abstract StreamStateHandle closeAndGetHandle() throws IOException;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/DoneFuture.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/DoneFuture.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/DoneFuture.java
new file mode 100644
index 0000000..777ab69
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/DoneFuture.java
@@ -0,0 +1,70 @@
+/*
+ * 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 java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.RunnableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+/**
+ * A {@link Future} that is always done and will just yield the object that was given at creation
+ * time.
+ *
+ * @param <T> The type of object in this {@code Future}.
+ */
+public class DoneFuture<T> implements RunnableFuture<T> {
+ private final T keyGroupsStateHandle;
+
+ public DoneFuture(T keyGroupsStateHandle) {
+ this.keyGroupsStateHandle = keyGroupsStateHandle;
+ }
+
+ @Override
+ public boolean cancel(boolean mayInterruptIfRunning) {
+ return false;
+ }
+
+ @Override
+ public boolean isCancelled() {
+ return false;
+ }
+
+ @Override
+ public boolean isDone() {
+ return true;
+ }
+
+ @Override
+ public T get() throws InterruptedException, ExecutionException {
+ return keyGroupsStateHandle;
+ }
+
+ @Override
+ public T get(
+ long timeout,
+ TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
+ return get();
+ }
+
+ @Override
+ public void run() {
+
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericFoldingState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericFoldingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericFoldingState.java
index e13ac98..ee2d86d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericFoldingState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericFoldingState.java
@@ -20,25 +20,18 @@ package org.apache.flink.runtime.state;
import org.apache.flink.api.common.functions.FoldFunction;
import org.apache.flink.api.common.state.FoldingState;
-import org.apache.flink.api.common.state.FoldingStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-
-import java.io.IOException;
/**
* Generic implementation of {@link FoldingState} based on a wrapped {@link ValueState}.
*
- * @param <K> The type of the key.
* @param <N> The type of the namespace.
* @param <T> The type of the values that can be folded into the state.
* @param <ACC> The type of the value in the folding state.
- * @param <Backend> The type of {@link AbstractStateBackend} that manages this {@code KvState}.
* @param <W> Generic type that extends both the underlying {@code ValueState} and {@code KvState}.
*/
-public class GenericFoldingState<K, N, T, ACC, Backend extends AbstractStateBackend, W extends ValueState<ACC> & KvState<K, N, ValueState<ACC>, ValueStateDescriptor<ACC>, Backend>>
- implements FoldingState<T, ACC>, KvState<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>, Backend> {
+public class GenericFoldingState<N, T, ACC, W extends ValueState<ACC> & KvState<N>>
+ implements FoldingState<T, ACC>, KvState<N> {
private final W wrappedState;
private final FoldFunction<T, ACC> foldFunction;
@@ -60,11 +53,6 @@ public class GenericFoldingState<K, N, T, ACC, Backend extends AbstractStateBack
}
@Override
- public void setCurrentKey(K key) {
- wrappedState.setCurrentKey(key);
- }
-
- @Override
public void setCurrentNamespace(N namespace) {
wrappedState.setCurrentNamespace(namespace);
}
@@ -75,26 +63,6 @@ public class GenericFoldingState<K, N, T, ACC, Backend extends AbstractStateBack
}
@Override
- public KvStateSnapshot<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>, Backend> snapshot(
- long checkpointId,
- long timestamp) throws Exception {
- KvStateSnapshot<K, N, ValueState<ACC>, ValueStateDescriptor<ACC>, Backend> wrappedSnapshot = wrappedState.snapshot(
- checkpointId,
- timestamp);
- return new Snapshot<>(wrappedSnapshot, foldFunction);
- }
-
- @Override
- public void dispose() {
- wrappedState.dispose();
- }
-
- @Override
- public FoldingStateDescriptor<T, ACC> getStateDescriptor() {
- throw new UnsupportedOperationException("Not supported by generic state type");
- }
-
- @Override
public ACC get() throws Exception {
return wrappedState.value();
}
@@ -109,42 +77,4 @@ public class GenericFoldingState<K, N, T, ACC, Backend extends AbstractStateBack
public void clear() {
wrappedState.clear();
}
-
- private static class Snapshot<K, N, T, ACC, Backend extends AbstractStateBackend> implements KvStateSnapshot<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>, Backend> {
- private static final long serialVersionUID = 1L;
-
- private final KvStateSnapshot<K, N, ValueState<ACC>, ValueStateDescriptor<ACC>, Backend> wrappedSnapshot;
-
- private final FoldFunction<T, ACC> foldFunction;
-
- public Snapshot(KvStateSnapshot<K, N, ValueState<ACC>, ValueStateDescriptor<ACC>, Backend> wrappedSnapshot,
- FoldFunction<T, ACC> foldFunction) {
- this.wrappedSnapshot = wrappedSnapshot;
- this.foldFunction = foldFunction;
- }
-
- @Override
- @SuppressWarnings("unchecked")
- public KvState<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>, Backend> restoreState(
- Backend stateBackend,
- TypeSerializer<K> keySerializer,
- ClassLoader classLoader) throws Exception {
- return new GenericFoldingState((ValueState<ACC>) wrappedSnapshot.restoreState(stateBackend, keySerializer, classLoader), foldFunction);
- }
-
- @Override
- public void discardState() throws Exception {
- wrappedSnapshot.discardState();
- }
-
- @Override
- public long getStateSize() throws Exception {
- return wrappedSnapshot.getStateSize();
- }
-
- @Override
- public void close() throws IOException {
- wrappedSnapshot.close();
- }
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericListState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericListState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericListState.java
index 45460b4..ba81837 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericListState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericListState.java
@@ -19,25 +19,19 @@
package org.apache.flink.runtime.state;
import org.apache.flink.api.common.state.ListState;
-import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import java.io.IOException;
import java.util.ArrayList;
/**
* Generic implementation of {@link ListState} based on a wrapped {@link ValueState}.
*
- * @param <K> The type of the key.
* @param <N> The type of the namespace.
* @param <T> The type of the values stored in this {@code ListState}.
- * @param <Backend> The type of {@link AbstractStateBackend} that manages this {@code KvState}.
* @param <W> Generic type that extends both the underlying {@code ValueState} and {@code KvState}.
*/
-public class GenericListState<K, N, T, Backend extends AbstractStateBackend, W extends ValueState<ArrayList<T>> & KvState<K, N, ValueState<ArrayList<T>>, ValueStateDescriptor<ArrayList<T>>, Backend>>
- implements ListState<T>, KvState<K, N, ListState<T>, ListStateDescriptor<T>, Backend> {
+public class GenericListState<N, T, W extends ValueState<ArrayList<T>> & KvState<N>>
+ implements ListState<T>, KvState<N> {
private final W wrappedState;
@@ -56,11 +50,6 @@ public class GenericListState<K, N, T, Backend extends AbstractStateBackend, W e
}
@Override
- public void setCurrentKey(K key) {
- wrappedState.setCurrentKey(key);
- }
-
- @Override
public void setCurrentNamespace(N namespace) {
wrappedState.setCurrentNamespace(namespace);
}
@@ -71,26 +60,6 @@ public class GenericListState<K, N, T, Backend extends AbstractStateBackend, W e
}
@Override
- public KvStateSnapshot<K, N, ListState<T>, ListStateDescriptor<T>, Backend> snapshot(
- long checkpointId,
- long timestamp) throws Exception {
- KvStateSnapshot<K, N, ValueState<ArrayList<T>>, ValueStateDescriptor<ArrayList<T>>, Backend> wrappedSnapshot = wrappedState.snapshot(
- checkpointId,
- timestamp);
- return new Snapshot<>(wrappedSnapshot);
- }
-
- @Override
- public void dispose() {
- wrappedState.dispose();
- }
-
- @Override
- public ListStateDescriptor<T> getStateDescriptor() {
- throw new UnsupportedOperationException("Not supported by generic state type");
- }
-
- @Override
public Iterable<T> get() throws Exception {
return wrappedState.value();
}
@@ -112,38 +81,4 @@ public class GenericListState<K, N, T, Backend extends AbstractStateBackend, W e
public void clear() {
wrappedState.clear();
}
-
- private static class Snapshot<K, N, T, Backend extends AbstractStateBackend> implements KvStateSnapshot<K, N, ListState<T>, ListStateDescriptor<T>, Backend> {
- private static final long serialVersionUID = 1L;
-
- private final KvStateSnapshot<K, N, ValueState<ArrayList<T>>, ValueStateDescriptor<ArrayList<T>>, Backend> wrappedSnapshot;
-
- public Snapshot(KvStateSnapshot<K, N, ValueState<ArrayList<T>>, ValueStateDescriptor<ArrayList<T>>, Backend> wrappedSnapshot) {
- this.wrappedSnapshot = wrappedSnapshot;
- }
-
- @Override
- @SuppressWarnings("unchecked")
- public KvState<K, N, ListState<T>, ListStateDescriptor<T>, Backend> restoreState(
- Backend stateBackend,
- TypeSerializer<K> keySerializer,
- ClassLoader classLoader) throws Exception {
- return new GenericListState((ValueState<T>) wrappedSnapshot.restoreState(stateBackend, keySerializer, classLoader));
- }
-
- @Override
- public void discardState() throws Exception {
- wrappedSnapshot.discardState();
- }
-
- @Override
- public long getStateSize() throws Exception {
- return wrappedSnapshot.getStateSize();
- }
-
- @Override
- public void close() throws IOException {
- wrappedSnapshot.close();
- }
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericReducingState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericReducingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericReducingState.java
index e4bb279..214231e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericReducingState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/GenericReducingState.java
@@ -20,24 +20,17 @@ package org.apache.flink.runtime.state;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.state.ReducingState;
-import org.apache.flink.api.common.state.ReducingStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-
-import java.io.IOException;
/**
* Generic implementation of {@link ReducingState} based on a wrapped {@link ValueState}.
*
- * @param <K> The type of the key.
* @param <N> The type of the namespace.
* @param <T> The type of the values stored in this {@code ReducingState}.
- * @param <Backend> The type of {@link AbstractStateBackend} that manages this {@code KvState}.
* @param <W> Generic type that extends both the underlying {@code ValueState} and {@code KvState}.
*/
-public class GenericReducingState<K, N, T, Backend extends AbstractStateBackend, W extends ValueState<T> & KvState<K, N, ValueState<T>, ValueStateDescriptor<T>, Backend>>
- implements ReducingState<T>, KvState<K, N, ReducingState<T>, ReducingStateDescriptor<T>, Backend> {
+public class GenericReducingState<N, T, W extends ValueState<T> & KvState<N>>
+ implements ReducingState<T>, KvState<N> {
private final W wrappedState;
private final ReduceFunction<T> reduceFunction;
@@ -59,11 +52,6 @@ public class GenericReducingState<K, N, T, Backend extends AbstractStateBackend,
}
@Override
- public void setCurrentKey(K key) {
- wrappedState.setCurrentKey(key);
- }
-
- @Override
public void setCurrentNamespace(N namespace) {
wrappedState.setCurrentNamespace(namespace);
}
@@ -74,26 +62,6 @@ public class GenericReducingState<K, N, T, Backend extends AbstractStateBackend,
}
@Override
- public KvStateSnapshot<K, N, ReducingState<T>, ReducingStateDescriptor<T>, Backend> snapshot(
- long checkpointId,
- long timestamp) throws Exception {
- KvStateSnapshot<K, N, ValueState<T>, ValueStateDescriptor<T>, Backend> wrappedSnapshot = wrappedState.snapshot(
- checkpointId,
- timestamp);
- return new Snapshot<>(wrappedSnapshot, reduceFunction);
- }
-
- @Override
- public void dispose() {
- wrappedState.dispose();
- }
-
- @Override
- public ReducingStateDescriptor<T> getStateDescriptor() {
- throw new UnsupportedOperationException("Not supported by generic state type");
- }
-
- @Override
public T get() throws Exception {
return wrappedState.value();
}
@@ -112,42 +80,4 @@ public class GenericReducingState<K, N, T, Backend extends AbstractStateBackend,
public void clear() {
wrappedState.clear();
}
-
- private static class Snapshot<K, N, T, Backend extends AbstractStateBackend> implements KvStateSnapshot<K, N, ReducingState<T>, ReducingStateDescriptor<T>, Backend> {
- private static final long serialVersionUID = 1L;
-
- private final KvStateSnapshot<K, N, ValueState<T>, ValueStateDescriptor<T>, Backend> wrappedSnapshot;
-
- private final ReduceFunction<T> reduceFunction;
-
- public Snapshot(KvStateSnapshot<K, N, ValueState<T>, ValueStateDescriptor<T>, Backend> wrappedSnapshot,
- ReduceFunction<T> reduceFunction) {
- this.wrappedSnapshot = wrappedSnapshot;
- this.reduceFunction = reduceFunction;
- }
-
- @Override
- @SuppressWarnings("unchecked")
- public KvState<K, N, ReducingState<T>, ReducingStateDescriptor<T>, Backend> restoreState(
- Backend stateBackend,
- TypeSerializer<K> keySerializer,
- ClassLoader classLoader) throws Exception {
- return new GenericReducingState((ValueState<T>) wrappedSnapshot.restoreState(stateBackend, keySerializer, classLoader), reduceFunction);
- }
-
- @Override
- public void discardState() throws Exception {
- wrappedSnapshot.discardState();
- }
-
- @Override
- public long getStateSize() throws Exception {
- return wrappedSnapshot.getStateSize();
- }
-
- @Override
- public void close() throws IOException {
- wrappedSnapshot.close();
- }
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
index de42bdb..9e74036 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
@@ -191,6 +191,9 @@ public class KeyGroupRange implements Iterable<Integer>, Serializable {
int maxParallelism,
int parallelism,
int operatorIndex) {
+ Preconditions.checkArgument(parallelism > 0, "Parallelism must not be smaller than zero.");
+ Preconditions.checkArgument(maxParallelism >= parallelism, "Maximum parallelism must not be smaller than parallelism.");
+ Preconditions.checkArgument(maxParallelism <= Short.MAX_VALUE, "Maximum parallelism must be smaller than Short.MAX_VALUE.");
int start = operatorIndex == 0 ? 0 : ((operatorIndex * maxParallelism - 1) / parallelism) + 1;
int end = ((operatorIndex + 1) * maxParallelism - 1) / parallelism;
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
new file mode 100644
index 0000000..2d1d25c
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
@@ -0,0 +1,340 @@
+/*
+ * 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.ExecutionConfig;
+import org.apache.flink.api.common.functions.ReduceFunction;
+import org.apache.flink.api.common.state.FoldingState;
+import org.apache.flink.api.common.state.FoldingStateDescriptor;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.state.MergingState;
+import org.apache.flink.api.common.state.ReducingState;
+import org.apache.flink.api.common.state.ReducingStateDescriptor;
+import org.apache.flink.api.common.state.State;
+import org.apache.flink.api.common.state.StateBackend;
+import org.apache.flink.api.common.state.StateDescriptor;
+import org.apache.flink.api.common.state.ValueState;
+import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.runtime.query.TaskKvStateRegistry;
+import org.apache.flink.util.Preconditions;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.concurrent.RunnableFuture;
+
+/**
+ * A keyed state backend is responsible for managing keyed state. The state can be checkpointed
+ * to streams using {@link #snapshot(long, long, CheckpointStreamFactory)}.
+ *
+ * @param <K> The key by which state is keyed.
+ */
+public abstract class KeyedStateBackend<K> {
+
+ /** {@link TypeSerializer} for our key. */
+ protected final TypeSerializer<K> keySerializer;
+
+ /** The currently active key. */
+ protected K currentKey;
+
+ /** The key group of the currently active key */
+ private int currentKeyGroup;
+
+ /** So that we can give out state when the user uses the same key. */
+ protected HashMap<String, KvState<?>> keyValueStatesByName;
+
+ /** For caching the last accessed partitioned state */
+ private String lastName;
+
+ @SuppressWarnings("rawtypes")
+ private KvState lastState;
+
+ /** KeyGroupAssigner which determines the key group for each keys */
+ protected final KeyGroupAssigner<K> keyGroupAssigner;
+
+ /** Range of key-groups for which this backend is responsible */
+ protected final KeyGroupRange keyGroupRange;
+
+ /** KvStateRegistry helper for this task */
+ protected final TaskKvStateRegistry kvStateRegistry;
+
+ public KeyedStateBackend(
+ TaskKvStateRegistry kvStateRegistry,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange) {
+
+ this.kvStateRegistry = Preconditions.checkNotNull(kvStateRegistry);
+ this.keySerializer = Preconditions.checkNotNull(keySerializer);
+ this.keyGroupAssigner = Preconditions.checkNotNull(keyGroupAssigner);
+ this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange);
+ }
+
+ /**
+ * Closes the state backend, 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
+ */
+ public void close() throws Exception {
+ if (kvStateRegistry != null) {
+ kvStateRegistry.unregisterAll();
+ }
+
+ lastName = null;
+ lastState = null;
+ keyValueStatesByName = null;
+ }
+
+ /**
+ * Creates and returns a new {@link ValueState}.
+ *
+ * @param namespaceSerializer TypeSerializer for the state namespace.
+ * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
+ *
+ * @param <N> The type of the namespace.
+ * @param <T> The type of the value that the {@code ValueState} can store.
+ */
+ protected abstract <N, T> ValueState<T> createValueState(TypeSerializer<N> namespaceSerializer, ValueStateDescriptor<T> stateDesc) throws Exception;
+
+ /**
+ * Creates and returns a new {@link ListState}.
+ *
+ * @param namespaceSerializer TypeSerializer for the state namespace.
+ * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
+ *
+ * @param <N> The type of the namespace.
+ * @param <T> The type of the values that the {@code ListState} can store.
+ */
+ protected abstract <N, T> ListState<T> createListState(TypeSerializer<N> namespaceSerializer, ListStateDescriptor<T> stateDesc) throws Exception;
+
+ /**
+ * Creates and returns a new {@link ReducingState}.
+ *
+ * @param namespaceSerializer TypeSerializer for the state namespace.
+ * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
+ *
+ * @param <N> The type of the namespace.
+ * @param <T> The type of the values that the {@code ListState} can store.
+ */
+ protected abstract <N, T> ReducingState<T> createReducingState(TypeSerializer<N> namespaceSerializer, ReducingStateDescriptor<T> stateDesc) throws Exception;
+
+ /**
+ * Creates and returns a new {@link FoldingState}.
+ *
+ * @param namespaceSerializer TypeSerializer for the state namespace.
+ * @param stateDesc The {@code StateDescriptor} that contains the name of the state.
+ *
+ * @param <N> The type of the namespace.
+ * @param <T> Type of the values folded into the state
+ * @param <ACC> Type of the value in the state *
+ */
+ protected abstract <N, T, ACC> FoldingState<T, ACC> createFoldingState(TypeSerializer<N> namespaceSerializer, FoldingStateDescriptor<T, ACC> stateDesc) throws Exception;
+
+ /**
+ * Sets the current key that is used for partitioned state.
+ * @param newKey The new current key.
+ */
+ public void setCurrentKey(K newKey) {
+ this.currentKey = newKey;
+ this.currentKeyGroup = keyGroupAssigner.getKeyGroupIndex(newKey);
+ }
+
+ /**
+ * {@link TypeSerializer} for the state backend key type.
+ */
+ public TypeSerializer<K> getKeySerializer() {
+ return keySerializer;
+ }
+
+ /**
+ * Used by states to access the current key.
+ */
+ public K getCurrentKey() {
+ return currentKey;
+ }
+
+ public int getCurrentKeyGroupIndex() {
+ return currentKeyGroup;
+ }
+
+ public int getNumberOfKeyGroups() {
+ return keyGroupAssigner.getNumberKeyGroups();
+ }
+
+ public KeyGroupAssigner<K> getKeyGroupAssigner() {
+ return keyGroupAssigner;
+ }
+
+ /**
+ * Creates or retrieves a partitioned state backed by this state backend.
+ *
+ * @param stateDescriptor The state identifier for the state. This contains name
+ * and can create a default state value.
+
+ * @param <N> The type of the namespace.
+ * @param <S> The type of the state.
+ *
+ * @return A new key/value state backed by this backend.
+ *
+ * @throws Exception Exceptions may occur during initialization of the state and should be forwarded.
+ */
+ @SuppressWarnings({"rawtypes", "unchecked"})
+ public <N, S extends State> S getPartitionedState(final N namespace, final TypeSerializer<N> namespaceSerializer, final StateDescriptor<S, ?> stateDescriptor) throws Exception {
+ Preconditions.checkNotNull(namespace, "Namespace");
+ Preconditions.checkNotNull(namespaceSerializer, "Namespace serializer");
+
+ if (keySerializer == null) {
+ throw new RuntimeException("State key serializer has not been configured in the config. " +
+ "This operation cannot use partitioned state.");
+ }
+
+ if (!stateDescriptor.isSerializerInitialized()) {
+ stateDescriptor.initializeSerializerUnlessSet(new ExecutionConfig());
+ }
+
+ if (keyValueStatesByName == null) {
+ keyValueStatesByName = new HashMap<>();
+ }
+
+ if (lastName != null && lastName.equals(stateDescriptor.getName())) {
+ lastState.setCurrentNamespace(namespace);
+ return (S) lastState;
+ }
+
+ KvState<?> previous = keyValueStatesByName.get(stateDescriptor.getName());
+ if (previous != null) {
+ lastState = previous;
+ lastState.setCurrentNamespace(namespace);
+ lastName = stateDescriptor.getName();
+ return (S) previous;
+ }
+
+ // create a new blank key/value state
+ S state = stateDescriptor.bind(new StateBackend() {
+ @Override
+ public <T> ValueState<T> createValueState(ValueStateDescriptor<T> stateDesc) throws Exception {
+ return KeyedStateBackend.this.createValueState(namespaceSerializer, stateDesc);
+ }
+
+ @Override
+ public <T> ListState<T> createListState(ListStateDescriptor<T> stateDesc) throws Exception {
+ return KeyedStateBackend.this.createListState(namespaceSerializer, stateDesc);
+ }
+
+ @Override
+ public <T> ReducingState<T> createReducingState(ReducingStateDescriptor<T> stateDesc) throws Exception {
+ return KeyedStateBackend.this.createReducingState(namespaceSerializer, stateDesc);
+ }
+
+ @Override
+ public <T, ACC> FoldingState<T, ACC> createFoldingState(FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
+ return KeyedStateBackend.this.createFoldingState(namespaceSerializer, stateDesc);
+ }
+
+ });
+
+ KvState kvState = (KvState) state;
+
+ keyValueStatesByName.put(stateDescriptor.getName(), kvState);
+
+ lastName = stateDescriptor.getName();
+ lastState = kvState;
+
+ kvState.setCurrentNamespace(namespace);
+
+ // Publish queryable state
+ if (stateDescriptor.isQueryable()) {
+ if (kvStateRegistry == null) {
+ throw new IllegalStateException("State backend has not been initialized for job.");
+ }
+
+ String name = stateDescriptor.getQueryableStateName();
+ // TODO: deal with key group indices here
+ kvStateRegistry.registerKvState(0, name, kvState);
+ }
+
+ return state;
+ }
+
+ @SuppressWarnings("unchecked,rawtypes")
+ public <N, S extends MergingState<?, ?>> void mergePartitionedStates(final N target, Collection<N> sources, final TypeSerializer<N> namespaceSerializer, final StateDescriptor<S, ?> stateDescriptor) throws Exception {
+ if (stateDescriptor instanceof ReducingStateDescriptor) {
+ ReducingStateDescriptor reducingStateDescriptor = (ReducingStateDescriptor) stateDescriptor;
+ ReduceFunction reduceFn = reducingStateDescriptor.getReduceFunction();
+ ReducingState state = (ReducingState) getPartitionedState(target, namespaceSerializer, stateDescriptor);
+ KvState kvState = (KvState) state;
+ Object result = null;
+ for (N source: sources) {
+ kvState.setCurrentNamespace(source);
+ Object sourceValue = state.get();
+ if (result == null) {
+ result = state.get();
+ } else if (sourceValue != null) {
+ result = reduceFn.reduce(result, sourceValue);
+ }
+ state.clear();
+ }
+ kvState.setCurrentNamespace(target);
+ if (result != null) {
+ state.add(result);
+ }
+ } else if (stateDescriptor instanceof ListStateDescriptor) {
+ ListState<Object> state = (ListState) getPartitionedState(target, namespaceSerializer, stateDescriptor);
+ KvState kvState = (KvState) state;
+ List<Object> result = new ArrayList<>();
+ for (N source: sources) {
+ kvState.setCurrentNamespace(source);
+ Iterable<Object> sourceValue = state.get();
+ if (sourceValue != null) {
+ for (Object o : sourceValue) {
+ result.add(o);
+ }
+ }
+ state.clear();
+ }
+ kvState.setCurrentNamespace(target);
+ for (Object o : result) {
+ state.add(o);
+ }
+ } else {
+ throw new RuntimeException("Cannot merge states for " + stateDescriptor);
+ }
+ }
+
+ /**
+ * Snapshots the keyed state by writing it to streams that are provided by a
+ * {@link CheckpointStreamFactory}.
+ *
+ * @param checkpointId The ID of the checkpoint.
+ * @param timestamp The timestamp of the checkpoint.
+ * @param streamFactory The factory that we can use for writing our state to streams.
+ *
+ * @return A future that will yield a {@link KeyGroupsStateHandle} with the index and
+ * written key group state stream.
+ */
+ public abstract RunnableFuture<KeyGroupsStateHandle> snapshot(
+ long checkpointId,
+ long timestamp,
+ CheckpointStreamFactory streamFactory) throws Exception;
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/KvState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KvState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KvState.java
index a8aa872..aded79f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KvState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KvState.java
@@ -18,9 +18,6 @@
package org.apache.flink.runtime.state;
-import org.apache.flink.api.common.state.State;
-import org.apache.flink.api.common.state.StateDescriptor;
-
/**
* Key/Value state implementation for user-defined state. The state is backed by a state
* backend, which typically follows one of the following patterns: Either the state is stored
@@ -29,20 +26,9 @@ import org.apache.flink.api.common.state.StateDescriptor;
* by an external key/value store as the state backend, and checkpoints merely record the
* metadata of what is considered part of the checkpoint.
*
- * @param <K> The type of the key.
* @param <N> The type of the namespace.
- * @param <S> The type of {@link State} this {@code KvState} holds.
- * @param <SD> The type of the {@link StateDescriptor} for state {@code S}.
- * @param <Backend> The type of {@link AbstractStateBackend} that manages this {@code KvState}.
*/
-public interface KvState<K, N, S extends State, SD extends StateDescriptor<S, ?>, Backend extends AbstractStateBackend> {
-
- /**
- * Sets the current key, which will be used when using the state access methods.
- *
- * @param key The key.
- */
- void setCurrentKey(K key);
+public interface KvState<N> {
/**
* Sets the current namespace, which will be used when using the state access methods.
@@ -63,27 +49,4 @@ public interface KvState<K, N, S extends State, SD extends StateDescriptor<S, ?>
* @throws Exception Exceptions during serialization are forwarded
*/
byte[] getSerializedValue(byte[] serializedKeyAndNamespace) throws Exception;
-
- /**
- * Creates a snapshot of this state.
- *
- * @param checkpointId The ID of the checkpoint for which the snapshot should be created.
- * @param timestamp The timestamp of the checkpoint.
- * @return A snapshot handle for this key/value state.
- *
- * @throws Exception Exceptions during snapshotting the state should be forwarded, so the system
- * can react to failed snapshots.
- */
- KvStateSnapshot<K, N, S, SD, Backend> snapshot(long checkpointId, long timestamp) throws Exception;
-
- /**
- * Disposes the key/value state, releasing all occupied resources.
- */
- void dispose();
-
- /**
- * Returns the state descriptor from which the KvState instance was created.
- */
- SD getStateDescriptor();
-
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/KvStateSnapshot.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KvStateSnapshot.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KvStateSnapshot.java
deleted file mode 100644
index 5654845..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KvStateSnapshot.java
+++ /dev/null
@@ -1,61 +0,0 @@
-/*
- * 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.state.State;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-
-/**
- * This class represents a snapshot of the {@link KvState}, taken for a checkpoint. Where exactly
- * the snapshot stores the snapshot data (in this object, in an external data store, etc) depends
- * on the actual implementation. This snapshot defines merely how to restore the state and
- * how to discard the state.
- *
- * <p>One possible implementation is that this snapshot simply contains a copy of the key/value map.
- *
- * <p>Another possible implementation for this snapshot is that the key/value map is serialized into
- * a file and this snapshot object contains a pointer to that file.
- *
- * @param <K> The type of the key
- * @param <N> The type of the namespace
- * @param <S> The type of the {@link State}
- * @param <SD> The type of the {@link StateDescriptor}
- * @param <Backend> The type of the backend that can restore the state from this snapshot.
- */
-public interface KvStateSnapshot<K, N, S extends State, SD extends StateDescriptor<S, ?>, Backend extends AbstractStateBackend>
- extends StateObject {
-
- /**
- * Loads the key/value state back from this snapshot.
- *
- * @param stateBackend The state backend that created this snapshot and can restore the key/value state
- * from this snapshot.
- * @param keySerializer The serializer for the keys.
- * @param classLoader The class loader for user-defined types.
- *
- * @return An instance of the key/value state loaded from this snapshot.
- *
- * @throws Exception Exceptions can occur during the state loading and are forwarded.
- */
- KvState<K, N, S, SD, Backend> restoreState(
- Backend stateBackend,
- TypeSerializer<K> keySerializer,
- ClassLoader classLoader) throws Exception;
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStreamStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStreamStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStreamStateHandle.java
index e3538af..c6fd02c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStreamStateHandle.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/RetrievableStreamStateHandle.java
@@ -61,7 +61,7 @@ public class RetrievableStreamStateHandle<T extends Serializable> implements
}
@Override
- public FSDataInputStream openInputStream() throws Exception {
+ public FSDataInputStream openInputStream() throws IOException {
return wrappedStreamStateHandle.openInputStream();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateObject.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateObject.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateObject.java
index a43a2c5..47103c1 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateObject.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StateObject.java
@@ -20,9 +20,8 @@ package org.apache.flink.runtime.state;
/**
* Base of all types that represent checkpointed state. Specializations are for
- * example {@link StateHandle StateHandles} (directly resolve to state) and
- * {@link KvStateSnapshot key/value state snapshots}.
- *
+ * example {@link StateHandle StateHandles} (directly resolve to state).
+ *
* <p>State objects define how to:
* <ul>
* <li><b>Discard State</b>: The {@link #discardState()} method defines how state is permanently
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java
index 46e4299..e792e62 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/StreamStateHandle.java
@@ -20,6 +20,8 @@ package org.apache.flink.runtime.state;
import org.apache.flink.core.fs.FSDataInputStream;
+import java.io.IOException;
+
/**
* A {@link StateObject} that represents state that was written to a stream. The data can be read
* back via {@link #openInputStream()}.
@@ -30,5 +32,5 @@ public interface StreamStateHandle extends StateObject {
* Returns an {@link FSDataInputStream} that can be used to read back the data that
* was previously written to the stream.
*/
- FSDataInputStream openInputStream() throws Exception;
+ FSDataInputStream openInputStream() throws IOException;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsState.java
deleted file mode 100644
index 3cae629..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsState.java
+++ /dev/null
@@ -1,95 +0,0 @@
-/*
- * 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.api.common.state.ListState;
-import org.apache.flink.api.common.state.State;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
-import org.apache.flink.runtime.state.AbstractHeapState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-
-import java.io.DataOutputStream;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Base class for partitioned {@link ListState} implementations that are backed by a regular
- * heap hash map. The concrete implementations define how the state is checkpointed.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <SV> The type of the values in the state.
- * @param <S> The type of State
- * @param <SD> The type of StateDescriptor for the State S
- */
-public abstract class AbstractFsState<K, N, SV, S extends State, SD extends StateDescriptor<S, ?>>
- extends AbstractHeapState<K, N, SV, S, SD, FsStateBackend> {
-
- /** The file system state backend backing snapshots of this state */
- private final FsStateBackend backend;
-
- public AbstractFsState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<SV> stateSerializer,
- SD stateDesc) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDesc);
- this.backend = backend;
- }
-
- public AbstractFsState(FsStateBackend backend,
- TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<SV> stateSerializer,
- SD stateDesc,
- HashMap<N, Map<K, SV>> state) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDesc, state);
- this.backend = backend;
- }
-
- public abstract KvStateSnapshot<K, N, S, SD, FsStateBackend> createHeapSnapshot(Path filePath);
-
- @Override
- public KvStateSnapshot<K, N, S, SD, FsStateBackend> snapshot(long checkpointId, long timestamp) throws Exception {
-
- try (FsStateBackend.FsCheckpointStateOutputStream out = backend.createCheckpointStateOutputStream(checkpointId, timestamp)) {
-
- // serialize the state to the output stream
- DataOutputViewStreamWrapper outView = new DataOutputViewStreamWrapper(new DataOutputStream(out));
- outView.writeInt(state.size());
- for (Map.Entry<N, Map<K, SV>> namespaceState: state.entrySet()) {
- N namespace = namespaceState.getKey();
- namespaceSerializer.serialize(namespace, outView);
- outView.writeInt(namespaceState.getValue().size());
- for (Map.Entry<K, SV> entry: namespaceState.getValue().entrySet()) {
- keySerializer.serialize(entry.getKey(), outView);
- stateSerializer.serialize(entry.getValue(), outView);
- }
- }
- outView.flush();
-
- // create a handle to the state
-// return new FsHeapValueStateSnapshot<>(getKeySerializer(), getNamespaceSerializer(), stateDesc, out.closeAndGetPath());
- return createHeapSnapshot(out.closeAndGetPath());
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsStateSnapshot.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsStateSnapshot.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsStateSnapshot.java
deleted file mode 100644
index 51e8b5a..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/AbstractFsStateSnapshot.java
+++ /dev/null
@@ -1,139 +0,0 @@
-/*
- * 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.api.common.state.State;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.core.fs.FSDataInputStream;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * A snapshot of a heap key/value state stored in a file.
- *
- * @param <K> The type of the key in the snapshot state.
- * @param <N> The type of the namespace in the snapshot state.
- * @param <SV> The type of the state value.
- */
-public abstract class AbstractFsStateSnapshot<K, N, SV, S extends State, SD extends StateDescriptor<S, ?>>
- extends FileStateHandle
- implements KvStateSnapshot<K, N, S, SD, FsStateBackend> {
-
- private static final long serialVersionUID = 1L;
-
- /** Key Serializer */
- protected final TypeSerializer<K> keySerializer;
-
- /** Namespace Serializer */
- protected final TypeSerializer<N> namespaceSerializer;
-
- /** Serializer for the state value */
- protected final TypeSerializer<SV> stateSerializer;
-
- /** StateDescriptor, for sanity checks */
- protected final SD stateDesc;
-
- /**
- * Creates a new state snapshot with data in the file system.
- *
- * @param keySerializer The serializer for the keys.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateSerializer The serializer for the elements in the state HashMap
- * @param stateDesc The state identifier
- * @param filePath The path where the snapshot data is stored.
- */
- public AbstractFsStateSnapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<SV> stateSerializer,
- SD stateDesc,
- Path filePath) {
- super(filePath);
- this.stateDesc = stateDesc;
- this.keySerializer = keySerializer;
- this.stateSerializer = stateSerializer;
- this.namespaceSerializer = namespaceSerializer;
-
- }
-
- public abstract KvState<K, N, S, SD, FsStateBackend> createFsState(FsStateBackend backend, HashMap<N, Map<K, SV>> stateMap);
-
- @Override
- public KvState<K, N, S, SD, FsStateBackend> restoreState(
- FsStateBackend stateBackend,
- final TypeSerializer<K> keySerializer,
- ClassLoader classLoader) throws Exception {
-
- // validity checks
- if (!this.keySerializer.equals(keySerializer)) {
- throw new IllegalArgumentException(
- "Cannot restore the state from the snapshot with the given serializers. " +
- "State (K/V) was serialized with " +
- "(" + this.keySerializer + ") " +
- "now is (" + keySerializer + ")");
- }
-
- // state restore
- ensureNotClosed();
-
- try (FSDataInputStream inStream = stateBackend.getFileSystem().open(getFilePath())) {
- // make sure the in-progress restore from the handle can be closed
- registerCloseable(inStream);
-
- DataInputViewStreamWrapper inView = new DataInputViewStreamWrapper(inStream);
-
- final int numKeys = inView.readInt();
- HashMap<N, Map<K, SV>> stateMap = new HashMap<>(numKeys);
-
- for (int i = 0; i < numKeys; i++) {
- N namespace = namespaceSerializer.deserialize(inView);
- final int numValues = inView.readInt();
- Map<K, SV> namespaceMap = new HashMap<>(numValues);
- stateMap.put(namespace, namespaceMap);
- for (int j = 0; j < numValues; j++) {
- K key = keySerializer.deserialize(inView);
- SV value = stateSerializer.deserialize(inView);
- namespaceMap.put(key, value);
- }
- }
-
- return createFsState(stateBackend, stateMap);
- }
- catch (Exception e) {
- throw new Exception("Failed to restore state from file system", e);
- }
- }
-
- /**
- * Returns the file size in bytes.
- *
- * @return The file size in bytes.
- * @throws IOException Thrown if the file system cannot be accessed.
- */
- @Override
- public void discardState() throws Exception {
- super.discardState();
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStateHandle.java
index 871e56c..5ae751b 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStateHandle.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FileStateHandle.java
@@ -65,7 +65,7 @@ public class FileStateHandle extends AbstractCloseableHandle implements StreamSt
}
@Override
- public FSDataInputStream openInputStream() throws Exception {
+ public FSDataInputStream openInputStream() throws IOException {
ensureNotClosed();
FSDataInputStream inputStream = getFileSystem().open(filePath);
registerCloseable(inputStream);
[21/27] flink git commit: [FLINK-4380] Remove KeyGroupAssigner in
favor of static method/Have default max. parallelism at 128
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java
index 277fab4..6259598 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGeneratorTest.java
@@ -23,13 +23,11 @@ import java.util.Random;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.functions.MapFunction;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobgraph.JobVertex;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
@@ -161,202 +159,4 @@ public class StreamingJobGraphGeneratorTest extends TestLogger {
assertEquals(1, jobGraph.getVerticesAsArray()[1].getParallelism());
}
- /**
- * Tests that the KeyGroupAssigner is properly set in the {@link StreamConfig} if the max
- * parallelism is set for the whole job.
- */
- @Test
- public void testKeyGroupAssignerProperlySet() {
- int maxParallelism = 42;
-
- final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
- env.getConfig().setMaxParallelism(maxParallelism);
-
- DataStream<Integer> input = env.fromElements(1, 2, 3);
-
- DataStream<Integer> keyedResult = input.keyBy(new KeySelector<Integer, Integer>() {
- private static final long serialVersionUID = 350461576474507944L;
-
- @Override
- public Integer getKey(Integer value) throws Exception {
- return value;
- }
- }).map(new NoOpIntMap());
-
- keyedResult.addSink(new DiscardingSink<Integer>());
-
- JobGraph jobGraph = env.getStreamGraph().getJobGraph();
-
- List<JobVertex> jobVertices = jobGraph.getVerticesSortedTopologicallyFromSources();
-
- assertEquals(maxParallelism, jobVertices.get(1).getMaxParallelism());
-
- HashKeyGroupAssigner<Integer> hashKeyGroupAssigner = extractHashKeyGroupAssigner(jobVertices.get(1));
-
- assertEquals(maxParallelism, hashKeyGroupAssigner.getNumberKeyGroups());
- }
-
- /**
- * Tests that the key group assigner for the keyed streams in the stream config is properly
- * initialized with the max parallelism value if there is no max parallelism defined for the
- * whole job.
- */
- @Test
- public void testKeyGroupAssignerProperlySetAutoMaxParallelism() {
- int globalParallelism = 42;
- int mapParallelism = 17;
- int maxParallelism = 43;
- StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
- env.setParallelism(globalParallelism);
-
- DataStream<Integer> source = env.fromElements(1, 2, 3);
-
- DataStream<Integer> keyedResult1 = source.keyBy(new KeySelector<Integer, Integer>() {
- private static final long serialVersionUID = 9205556348021992189L;
-
- @Override
- public Integer getKey(Integer value) throws Exception {
- return value;
- }
- }).map(new NoOpIntMap());
-
- DataStream<Integer> keyedResult2 = keyedResult1.keyBy(new KeySelector<Integer, Integer>() {
- private static final long serialVersionUID = 1250168178707154838L;
-
- @Override
- public Integer getKey(Integer value) throws Exception {
- return value;
- }
- }).map(new NoOpIntMap()).setParallelism(mapParallelism);
-
- DataStream<Integer> keyedResult3 = keyedResult2.keyBy(new KeySelector<Integer, Integer>() {
- private static final long serialVersionUID = 1250168178707154838L;
-
- @Override
- public Integer getKey(Integer value) throws Exception {
- return value;
- }
- }).map(new NoOpIntMap()).setMaxParallelism(maxParallelism);
-
- DataStream<Integer> keyedResult4 = keyedResult3.keyBy(new KeySelector<Integer, Integer>() {
- private static final long serialVersionUID = 1250168178707154838L;
-
- @Override
- public Integer getKey(Integer value) throws Exception {
- return value;
- }
- }).map(new NoOpIntMap()).setMaxParallelism(maxParallelism).setParallelism(mapParallelism);
-
- keyedResult4.addSink(new DiscardingSink<Integer>());
-
- JobGraph jobGraph = env.getStreamGraph().getJobGraph();
-
- List<JobVertex> vertices = jobGraph.getVerticesSortedTopologicallyFromSources();
-
- JobVertex keyedResultJV1 = vertices.get(1);
- JobVertex keyedResultJV2 = vertices.get(2);
- JobVertex keyedResultJV3 = vertices.get(3);
- JobVertex keyedResultJV4 = vertices.get(4);
-
- HashKeyGroupAssigner<Integer> hashKeyGroupAssigner1 = extractHashKeyGroupAssigner(keyedResultJV1);
- HashKeyGroupAssigner<Integer> hashKeyGroupAssigner2 = extractHashKeyGroupAssigner(keyedResultJV2);
- HashKeyGroupAssigner<Integer> hashKeyGroupAssigner3 = extractHashKeyGroupAssigner(keyedResultJV3);
- HashKeyGroupAssigner<Integer> hashKeyGroupAssigner4 = extractHashKeyGroupAssigner(keyedResultJV4);
-
- assertEquals(globalParallelism, hashKeyGroupAssigner1.getNumberKeyGroups());
- assertEquals(mapParallelism, hashKeyGroupAssigner2.getNumberKeyGroups());
- assertEquals(maxParallelism, hashKeyGroupAssigner3.getNumberKeyGroups());
- assertEquals(maxParallelism, hashKeyGroupAssigner4.getNumberKeyGroups());
- }
-
- /**
- * Tests that the {@link KeyGroupAssigner} is properly set in the {@link StreamConfig} for
- * connected streams.
- */
- @Test
- public void testMaxParallelismWithConnectedKeyedStream() {
- int maxParallelism = 42;
-
- StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
- DataStream<Integer> input1 = env.fromElements(1, 2, 3, 4).setMaxParallelism(128).name("input1");
- DataStream<Integer> input2 = env.fromElements(1, 2, 3, 4).setMaxParallelism(129).name("input2");
-
- env.getConfig().setMaxParallelism(maxParallelism);
-
- DataStream<Integer> keyedResult = input1.connect(input2).keyBy(
- new KeySelector<Integer, Integer>() {
- private static final long serialVersionUID = -6908614081449363419L;
-
- @Override
- public Integer getKey(Integer value) throws Exception {
- return value;
- }
- },
- new KeySelector<Integer, Integer>() {
- private static final long serialVersionUID = 3195683453223164931L;
-
- @Override
- public Integer getKey(Integer value) throws Exception {
- return value;
- }
- }).map(new StreamGraphGeneratorTest.NoOpIntCoMap());
-
- keyedResult.addSink(new DiscardingSink<Integer>());
-
- JobGraph jobGraph = env.getStreamGraph().getJobGraph();
-
- List<JobVertex> jobVertices = jobGraph.getVerticesSortedTopologicallyFromSources();
-
- JobVertex input1JV = jobVertices.get(0);
- JobVertex input2JV = jobVertices.get(1);
- JobVertex connectedJV = jobVertices.get(2);
-
- // disambiguate the partial order of the inputs
- if (input1JV.getName().equals("Source: input1")) {
- assertEquals(128, input1JV.getMaxParallelism());
- assertEquals(129, input2JV.getMaxParallelism());
- } else {
- assertEquals(128, input2JV.getMaxParallelism());
- assertEquals(129, input1JV.getMaxParallelism());
- }
-
- assertEquals(maxParallelism, connectedJV.getMaxParallelism());
-
- HashKeyGroupAssigner<Integer> hashKeyGroupAssigner = extractHashKeyGroupAssigner(connectedJV);
-
- assertEquals(maxParallelism, hashKeyGroupAssigner.getNumberKeyGroups());
- }
-
- /**
- * Tests that the {@link JobGraph} creation fails if the parallelism is greater than the max
- * parallelism.
- */
- @Test(expected=IllegalStateException.class)
- public void testFailureOfJobJobCreationIfParallelismGreaterThanMaxParallelism() {
- StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
- env.getConfig().setMaxParallelism(42);
-
- DataStream<Integer> input = env.fromElements(1, 2, 3, 4);
-
- DataStream<Integer> result = input.map(new NoOpIntMap()).setParallelism(43);
-
- result.addSink(new DiscardingSink<Integer>());
-
- env.getStreamGraph().getJobGraph();
-
- fail("The JobGraph should not have been created because the parallelism is greater than " +
- "the max parallelism.");
- }
-
- private HashKeyGroupAssigner<Integer> extractHashKeyGroupAssigner(JobVertex jobVertex) {
- Configuration config = jobVertex.getConfiguration();
-
- StreamConfig streamConfig = new StreamConfig(config);
-
- KeyGroupAssigner<Integer> keyGroupAssigner = streamConfig.getKeyGroupAssigner(getClass().getClassLoader());
-
- assertTrue(keyGroupAssigner instanceof HashKeyGroupAssigner);
-
- return (HashKeyGroupAssigner<Integer>) keyGroupAssigner;
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
index d3b7ff9..fe09788 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
@@ -30,19 +30,16 @@ import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
-import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.java.typeutils.runtime.kryo.KryoSerializer;
import org.apache.flink.core.fs.Path;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.query.KvStateRegistry;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-import org.apache.flink.runtime.state.heap.HeapListState;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.junit.Test;
import org.mockito.invocation.InvocationOnMock;
@@ -196,7 +193,7 @@ public class StreamingRuntimeContextTest {
new JobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ 1,
new KeyGroupRange(0, 0),
new KvStateRegistry().createTaskRegistry(new JobID(),
new JobVertexID()));
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
index 4e7e4d0..59bfe6f 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
@@ -22,7 +22,6 @@ import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.accumulators.Accumulator;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.functions.RichReduceFunction;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
@@ -37,9 +36,6 @@ import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.memory.MemoryStateBackend;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
import org.apache.flink.streaming.api.graph.StreamConfig;
@@ -52,7 +48,6 @@ import org.apache.flink.streaming.runtime.tasks.TestTimeServiceProvider;
import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.junit.After;
-import org.junit.Ignore;
import org.junit.Test;
import org.mockito.invocation.InvocationOnMock;
@@ -212,7 +207,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
op = new AggregatingProcessingTimeWindowOperator<>(mockFunction, mockKeySelector,
StringSerializer.INSTANCE, StringSerializer.INSTANCE, 5000, 1000);
- op.setup(mockTask, createTaskConfig(mockKeySelector, StringSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), mockOut);
+ op.setup(mockTask, createTaskConfig(mockKeySelector, StringSerializer.INSTANCE, 10), mockOut);
op.open();
assertTrue(op.getNextSlideTime() % 1000 == 0);
assertTrue(op.getNextEvaluationTime() % 1000 == 0);
@@ -264,7 +259,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
final Object lock = new Object();
final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
- op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, 10), out);
op.open();
final int numElements = 1000;
@@ -322,7 +317,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
windowSize, windowSize);
- op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, 10), out);
op.open();
final int numWindows = 10;
@@ -389,7 +384,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
150, 50);
- op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, 10), out);
op.open();
final int numElements = 1000;
@@ -458,7 +453,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
sumFunction, fieldOneSelector,
IntSerializer.INSTANCE, tupleSerializer, 150, 50);
- op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, 10), out);
op.open();
synchronized (lock) {
@@ -520,7 +515,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
hundredYears, hundredYears);
- op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, 10), out);
op.open();
for (int i = 0; i < 100; i++) {
@@ -973,7 +968,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
return mockTask;
}
- private static StreamConfig createTaskConfig(KeySelector<?, ?> partitioner, TypeSerializer<?> keySerializer, KeyGroupAssigner<?> keyGroupAssigner) {
+ private static StreamConfig createTaskConfig(KeySelector<?, ?> partitioner, TypeSerializer<?> keySerializer, int numberOfKeGroups) {
StreamConfig cfg = new StreamConfig(new Configuration());
return cfg;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitionerTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitionerTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitionerTest.java
index 6fbf35e..4ca7449 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitionerTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitionerTest.java
@@ -23,7 +23,6 @@ import static org.junit.Assert.assertEquals;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.runtime.plugable.SerializationDelegate;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.util.TestLogger;
import org.junit.Before;
@@ -48,7 +47,7 @@ public class KeyGroupStreamPartitionerTest extends TestLogger {
return value.getField(0);
}
},
- new HashKeyGroupAssigner<String>(1024));
+ 1024);
}
@Test
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java
index 5f73e25..5573a53 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTestHarness.java
@@ -23,7 +23,6 @@ import org.apache.flink.api.java.ClosureCleaner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.runtime.event.AbstractEvent;
import org.apache.flink.runtime.io.network.partition.consumer.StreamTestSingleInputGate;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import java.io.IOException;
@@ -106,7 +105,7 @@ public class OneInputStreamTaskTestHarness<IN, OUT> extends StreamTaskTestHarnes
ClosureCleaner.clean(keySelector, false);
streamConfig.setStatePartitioner(0, keySelector);
streamConfig.setStateKeySerializer(keyType.createSerializer(executionConfig));
- streamConfig.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(10));
+ streamConfig.setNumberOfKeyGroups(10);
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
index 5594193..03f50f9 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
@@ -19,14 +19,12 @@ package org.apache.flink.streaming.util;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.ClosureCleaner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.core.fs.FSDataInputStream;
import org.apache.flink.runtime.state.CheckpointStreamFactory;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyGroupsStateHandle;
import org.apache.flink.runtime.state.KeyedStateBackend;
@@ -70,7 +68,7 @@ public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
ClosureCleaner.clean(keySelector, false);
config.setStatePartitioner(0, keySelector);
config.setStateKeySerializer(keyType.createSerializer(executionConfig));
- config.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(MAX_PARALLELISM));
+ config.setNumberOfKeyGroups(MAX_PARALLELISM);
setupMockTaskCreateKeyedBackend();
}
@@ -84,7 +82,7 @@ public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
ClosureCleaner.clean(keySelector, false);
config.setStatePartitioner(0, keySelector);
config.setStateKeySerializer(keyType.createSerializer(executionConfig));
- config.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(MAX_PARALLELISM));
+ config.setNumberOfKeyGroups(MAX_PARALLELISM);
setupMockTaskCreateKeyedBackend();
}
@@ -99,7 +97,7 @@ public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
ClosureCleaner.clean(keySelector, false);
config.setStatePartitioner(0, keySelector);
config.setStateKeySerializer(keyType.createSerializer(executionConfig));
- config.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(MAX_PARALLELISM));
+ config.setNumberOfKeyGroups(MAX_PARALLELISM);
setupMockTaskCreateKeyedBackend();
}
@@ -112,7 +110,7 @@ public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
public KeyedStateBackend answer(InvocationOnMock invocationOnMock) throws Throwable {
final TypeSerializer keySerializer = (TypeSerializer) invocationOnMock.getArguments()[0];
- final KeyGroupAssigner keyGroupAssigner = (KeyGroupAssigner) invocationOnMock.getArguments()[1];
+ final int numberOfKeyGroups = (Integer) invocationOnMock.getArguments()[1];
final KeyGroupRange keyGroupRange = (KeyGroupRange) invocationOnMock.getArguments()[2];
if (restoredKeyedState == null) {
@@ -121,7 +119,7 @@ public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
new JobID(),
"test_op",
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange,
mockTask.getEnvironment().getTaskKvStateRegistry());
return keyedStateBackend;
@@ -131,7 +129,7 @@ public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
new JobID(),
"test_op",
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange,
Collections.singletonList(restoredKeyedState),
mockTask.getEnvironment().getTaskKvStateRegistry());
@@ -139,7 +137,7 @@ public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
return keyedStateBackend;
}
}
- }).when(mockTask).createKeyedStateBackend(any(TypeSerializer.class), any(KeyGroupAssigner.class), any(KeyGroupRange.class));
+ }).when(mockTask).createKeyedStateBackend(any(TypeSerializer.class), anyInt(), any(KeyGroupRange.class));
} catch (Exception e) {
throw new RuntimeException(e.getMessage(), e);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
index 78e05b7..15074a7 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
@@ -19,7 +19,6 @@ package org.apache.flink.streaming.util;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.typeutils.TypeExtractor;
import org.apache.flink.configuration.Configuration;
@@ -29,8 +28,6 @@ import org.apache.flink.runtime.operators.testutils.MockEnvironment;
import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.runtime.state.CheckpointStreamFactory;
-import org.apache.flink.runtime.state.KeyGroupRange;
-import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.streaming.api.graph.StreamConfig;
@@ -47,7 +44,6 @@ import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import java.util.Collection;
-import java.util.Collections;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Executors;
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
index 39f3086..82dbd1f 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
@@ -34,8 +34,7 @@ import org.apache.flink.runtime.execution.SuppressRestartsException;
import org.apache.flink.runtime.instance.ActorGateway;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.messages.JobManagerMessages;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
-import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.state.filesystem.FsStateBackendFactory;
import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages;
import org.apache.flink.streaming.api.checkpoint.Checkpointed;
@@ -141,12 +140,10 @@ public class RescalingITCase extends TestLogger {
Set<Tuple2<Integer, Integer>> expectedResult = new HashSet<>();
- HashKeyGroupAssigner<Integer> keyGroupAssigner = new HashKeyGroupAssigner<>(maxParallelism);
-
for (int key = 0; key < numberKeys; key++) {
- int keyGroupIndex = keyGroupAssigner.getKeyGroupIndex(key);
+ int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);
- expectedResult.add(Tuple2.of(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, keyGroupIndex), numberElements * key));
+ expectedResult.add(Tuple2.of(KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, keyGroupIndex), numberElements * key));
}
assertEquals(expectedResult, actualResult);
@@ -185,11 +182,9 @@ public class RescalingITCase extends TestLogger {
Set<Tuple2<Integer, Integer>> expectedResult2 = new HashSet<>();
- HashKeyGroupAssigner<Integer> keyGroupAssigner2 = new HashKeyGroupAssigner<>(maxParallelism);
-
for (int key = 0; key < numberKeys; key++) {
- int keyGroupIndex = keyGroupAssigner2.getKeyGroupIndex(key);
- expectedResult2.add(Tuple2.of(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism2, keyGroupIndex), key * (numberElements + numberElements2)));
+ int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);
+ expectedResult2.add(Tuple2.of(KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(maxParallelism, parallelism2, keyGroupIndex), key * (numberElements + numberElements2)));
}
assertEquals(expectedResult2, actualResult2);
@@ -347,12 +342,10 @@ public class RescalingITCase extends TestLogger {
Set<Tuple2<Integer, Integer>> expectedResult = new HashSet<>();
- HashKeyGroupAssigner<Integer> keyGroupAssigner = new HashKeyGroupAssigner<>(maxParallelism);
-
for (int key = 0; key < numberKeys; key++) {
- int keyGroupIndex = keyGroupAssigner.getKeyGroupIndex(key);
+ int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);
- expectedResult.add(Tuple2.of(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, keyGroupIndex) , numberElements * key));
+ expectedResult.add(Tuple2.of(KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, keyGroupIndex) , numberElements * key));
}
assertEquals(expectedResult, actualResult);
@@ -398,11 +391,9 @@ public class RescalingITCase extends TestLogger {
Set<Tuple2<Integer, Integer>> expectedResult2 = new HashSet<>();
- HashKeyGroupAssigner<Integer> keyGroupAssigner2 = new HashKeyGroupAssigner<>(maxParallelism);
-
for (int key = 0; key < numberKeys; key++) {
- int keyGroupIndex = keyGroupAssigner2.getKeyGroupIndex(key);
- expectedResult2.add(Tuple2.of(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism2, keyGroupIndex), key * (numberElements + numberElements2)));
+ int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(key, maxParallelism);
+ expectedResult2.add(Tuple2.of(KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(maxParallelism, parallelism2, keyGroupIndex), key * (numberElements + numberElements2)));
}
assertEquals(expectedResult2, actualResult2);
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
index d4dd475..694f006 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
@@ -21,7 +21,6 @@ package org.apache.flink.test.streaming.runtime;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
@@ -104,7 +103,7 @@ public class StateBackendITCase extends StreamingMultipleProgramsTestBase {
JobID jobID,
String operatorIdentifier,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
TaskKvStateRegistry kvStateRegistry) throws Exception {
throw new SuccessException();
@@ -115,7 +114,7 @@ public class StateBackendITCase extends StreamingMultipleProgramsTestBase {
JobID jobID,
String operatorIdentifier,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
List<KeyGroupsStateHandle> restoredState,
TaskKvStateRegistry kvStateRegistry) throws Exception {
[20/27] flink git commit: [FLINK-4381] Refactor State to Prepare For
Key-Group State Backends
Posted by al...@apache.org.
[FLINK-4381] Refactor State to Prepare For Key-Group State Backends
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/847ead01
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/847ead01
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/847ead01
Branch: refs/heads/master
Commit: 847ead01f2f0aaf318b2b1ba8501bc697d245900
Parents: ec975aa
Author: Stefan Richter <s....@data-artisans.com>
Authored: Thu Aug 11 11:59:07 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:10:01 2016 +0200
----------------------------------------------------------------------
.../org/apache/flink/client/CliFrontend.java | 3 +-
.../flink/client/CliFrontendSavepointTest.java | 12 +-
.../streaming/state/RocksDBStateBackend.java | 67 ++-
.../state/RocksDBAsyncKVSnapshotTest.java | 36 +-
.../storm/tests/StormFieldsGroupingITCase.java | 26 +-
.../api/common/state/KeyGroupAssigner.java | 10 +-
.../apache/flink/util/InstantiationUtil.java | 34 +-
.../flink/hdfstests/FileStateBackendTest.java | 58 +--
.../operator/AbstractCEPPatternOperator.java | 44 +-
.../AbstractKeyedCEPPatternOperator.java | 25 +-
.../flink/cep/operator/CEPOperatorTest.java | 26 +-
.../store/ZooKeeperMesosWorkerStore.java | 16 +-
.../checkpoint/CheckpointCoordinator.java | 103 ++--
.../runtime/checkpoint/CompletedCheckpoint.java | 56 ++-
.../flink/runtime/checkpoint/KeyGroupState.java | 87 ----
.../runtime/checkpoint/PendingCheckpoint.java | 67 ++-
.../runtime/checkpoint/PendingSavepoint.java | 7 +-
.../StandaloneCompletedCheckpointStore.java | 16 +-
.../flink/runtime/checkpoint/SubtaskState.java | 56 ++-
.../flink/runtime/checkpoint/TaskState.java | 151 +++---
.../ZooKeeperCompletedCheckpointStore.java | 46 +-
.../checkpoint/savepoint/FsSavepointStore.java | 6 +-
.../savepoint/HeapSavepointStore.java | 7 +-
.../runtime/checkpoint/savepoint/Savepoint.java | 8 +-
.../checkpoint/savepoint/SavepointLoader.java | 4 +-
.../savepoint/SavepointSerializers.java | 7 +-
.../checkpoint/savepoint/SavepointStore.java | 7 +-
.../checkpoint/savepoint/SavepointV0.java | 97 ----
.../savepoint/SavepointV0Serializer.java | 186 --------
.../checkpoint/savepoint/SavepointV1.java | 95 ++++
.../savepoint/SavepointV1Serializer.java | 232 +++++++++
.../stats/SimpleCheckpointStatsTracker.java | 15 +-
.../deployment/TaskDeploymentDescriptor.java | 22 +-
.../flink/runtime/execution/Environment.java | 17 +-
.../flink/runtime/executiongraph/Execution.java | 41 +-
.../runtime/executiongraph/ExecutionVertex.java | 32 +-
.../flink/runtime/jobgraph/JobVertex.java | 2 +-
.../runtime/jobgraph/tasks/StatefulTask.java | 14 +-
.../ZooKeeperSubmittedJobGraphStore.java | 20 +-
.../checkpoint/AcknowledgeCheckpoint.java | 53 ++-
.../runtime/state/AbstractCloseableHandle.java | 23 +-
.../runtime/state/AbstractStateBackend.java | 105 +----
.../runtime/state/AsynchronousStateHandle.java | 43 --
.../flink/runtime/state/ChainedStateHandle.java | 131 ++++++
.../runtime/state/HashKeyGroupAssigner.java | 6 +-
.../flink/runtime/state/KeyGroupRange.java | 217 +++++++++
.../runtime/state/KeyGroupRangeOffsets.java | 203 ++++++++
.../runtime/state/KeyGroupsStateHandle.java | 163 +++++++
.../flink/runtime/state/LocalStateHandle.java | 53 ---
.../runtime/state/RetrievableStateHandle.java | 32 ++
.../state/RetrievableStreamStateHandle.java | 82 ++++
.../runtime/state/StateBackendFactory.java | 4 +-
.../apache/flink/runtime/state/StateUtil.java | 101 ++++
.../apache/flink/runtime/state/StateUtils.java | 59 ---
.../flink/runtime/state/StreamStateHandle.java | 18 +-
.../filesystem/AbstractFileStateHandle.java | 98 ----
.../filesystem/AbstractFsStateSnapshot.java | 9 +-
.../filesystem/FileSerializableStateHandle.java | 72 ---
.../state/filesystem/FileStateHandle.java | 136 ++++++
.../state/filesystem/FileStreamStateHandle.java | 83 ----
.../state/filesystem/FsStateBackend.java | 28 +-
.../state/memory/ByteStreamStateHandle.java | 84 +++-
.../state/memory/MemoryStateBackend.java | 32 +-
.../state/memory/SerializedStateHandle.java | 87 ----
.../runtime/taskmanager/RuntimeEnvironment.java | 42 +-
.../apache/flink/runtime/taskmanager/Task.java | 51 +-
.../flink/runtime/util/ZooKeeperUtils.java | 6 +-
.../RetrievableStateStorageHelper.java | 41 ++
.../runtime/zookeeper/StateStorageHelper.java | 41 --
.../zookeeper/ZooKeeperStateHandleStore.java | 78 ++-
.../FileSystemStateStorageHelper.java | 17 +-
.../flink/runtime/jobmanager/JobManager.scala | 21 +-
.../runtime/messages/JobManagerMessages.scala | 6 +-
.../checkpoint/CheckpointStateRestoreTest.java | 44 +-
.../CompletedCheckpointStoreTest.java | 25 +-
.../checkpoint/CompletedCheckpointTest.java | 8 +-
.../checkpoint/PendingCheckpointTest.java | 17 +-
.../checkpoint/PendingSavepointTest.java | 15 +-
...ZooKeeperCompletedCheckpointStoreITCase.java | 44 +-
.../savepoint/FsSavepointStoreTest.java | 8 +-
.../savepoint/SavepointLoaderTest.java | 6 +-
.../savepoint/SavepointV0SerializerTest.java | 52 --
.../checkpoint/savepoint/SavepointV0Test.java | 81 ----
.../savepoint/SavepointV1SerializerTest.java | 52 ++
.../checkpoint/savepoint/SavepointV1Test.java | 88 ++++
.../stats/SimpleCheckpointStatsTrackerTest.java | 52 +-
.../jobmanager/JobManagerHARecoveryTest.java | 102 ++--
.../ZooKeeperSubmittedJobGraphsStoreITCase.java | 18 +-
.../messages/CheckpointMessagesTest.java | 56 ++-
.../operators/testutils/DummyEnvironment.java | 14 +-
.../operators/testutils/MockEnvironment.java | 15 +-
.../state/AbstractCloseableHandleTest.java | 10 +
.../runtime/state/FileStateBackendTest.java | 61 +--
.../FsCheckpointStateOutputStreamTest.java | 128 -----
.../runtime/state/KeyGroupRangeOffsetTest.java | 136 ++++++
.../flink/runtime/state/KeyGroupRangeTest.java | 101 ++++
.../runtime/state/MemoryStateBackendTest.java | 33 +-
.../runtime/state/StateBackendTestBase.java | 28 +-
.../FsCheckpointStateOutputStreamTest.java | 128 +++++
.../runtime/taskmanager/TaskAsyncCallTest.java | 16 +-
.../ZooKeeperStateHandleStoreITCase.java | 62 +--
.../cassandra/CassandraConnectorITCase.java | 18 +-
.../fs/bucketing/BucketingSinkTest.java | 6 +-
.../source/ContinuousFileReaderOperator.java | 33 +-
.../api/operators/AbstractStreamOperator.java | 98 ++--
.../operators/AbstractUdfStreamOperator.java | 67 ++-
.../streaming/api/operators/StreamOperator.java | 16 +-
.../operators/GenericWriteAheadSink.java | 89 ++--
...ractAlignedProcessingTimeWindowOperator.java | 49 +-
.../operators/windowing/WindowOperator.java | 32 +-
.../partitioner/KeyGroupStreamPartitioner.java | 7 +-
.../streaming/runtime/tasks/StreamTask.java | 363 +++++++-------
.../runtime/tasks/StreamTaskState.java | 185 --------
.../runtime/tasks/StreamTaskStateList.java | 123 -----
.../operators/GenericWriteAheadSinkTest.java | 68 ++-
.../operators/StreamSourceOperatorTest.java | 2 +
.../operators/WriteAheadSinkTestBase.java | 163 +++----
...AlignedProcessingTimeWindowOperatorTest.java | 117 +++--
...AlignedProcessingTimeWindowOperatorTest.java | 79 ++--
.../operators/windowing/WindowOperatorTest.java | 39 +-
.../windowing/WindowingTestHarnessTest.java | 6 +-
.../tasks/InterruptSensitiveRestoreTest.java | 63 ++-
.../runtime/tasks/StreamMockEnvironment.java | 9 +-
.../tasks/StreamTaskAsyncCheckpointTest.java | 470 +++++++++----------
.../flink/streaming/util/MockContext.java | 78 +--
.../util/OneInputStreamOperatorTestHarness.java | 74 ++-
.../streaming/util/WindowingTestHarness.java | 12 +-
.../EventTimeWindowCheckpointingITCase.java | 8 +-
.../test/checkpointing/RescalingITCase.java | 10 +-
.../test/checkpointing/SavepointITCase.java | 44 +-
.../test/classloading/ClassLoaderITCase.java | 2 +-
.../state/StateHandleSerializationTest.java | 8 +-
.../streaming/runtime/StateBackendITCase.java | 10 -
133 files changed, 4096 insertions(+), 3635 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java
----------------------------------------------------------------------
diff --git a/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java b/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java
index c90bc29..c3f1354 100644
--- a/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java
+++ b/flink-clients/src/main/java/org/apache/flink/client/CliFrontend.java
@@ -72,7 +72,6 @@ import org.apache.flink.util.Preconditions;
import org.apache.flink.util.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import scala.Option;
import scala.concurrent.Await;
import scala.concurrent.Future;
import scala.concurrent.duration.FiniteDuration;
@@ -727,7 +726,7 @@ public class CliFrontend {
logAndSysout("Disposing savepoint '" + savepointPath + "'.");
}
- Object msg = new DisposeSavepoint(savepointPath, Option.apply(blobKeys));
+ Object msg = new DisposeSavepoint(savepointPath);
Future<Object> response = jobManager.ask(msg, clientTimeout);
Object result;
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-clients/src/test/java/org/apache/flink/client/CliFrontendSavepointTest.java
----------------------------------------------------------------------
diff --git a/flink-clients/src/test/java/org/apache/flink/client/CliFrontendSavepointTest.java b/flink-clients/src/test/java/org/apache/flink/client/CliFrontendSavepointTest.java
index d9e7075..f9c7e8c 100644
--- a/flink-clients/src/test/java/org/apache/flink/client/CliFrontendSavepointTest.java
+++ b/flink-clients/src/test/java/org/apache/flink/client/CliFrontendSavepointTest.java
@@ -212,7 +212,7 @@ public class CliFrontendSavepointTest {
Promise<Object> triggerResponse = new scala.concurrent.impl.Promise.DefaultPromise<>();
when(jobManager.ask(
- Mockito.eq(new DisposeSavepoint(savepointPath, Option.<List<BlobKey>>empty())),
+ Mockito.eq(new DisposeSavepoint(savepointPath)),
any(FiniteDuration.class))).thenReturn(triggerResponse.future());
triggerResponse.success(getDisposeSavepointSuccess());
@@ -225,7 +225,7 @@ public class CliFrontendSavepointTest {
assertEquals(0, returnCode);
verify(jobManager, times(1)).ask(
- Mockito.eq(new DisposeSavepoint(savepointPath, Option.<List<BlobKey>>empty())),
+ Mockito.eq(new DisposeSavepoint(savepointPath)),
any(FiniteDuration.class));
String outMsg = buffer.toString();
@@ -307,7 +307,7 @@ public class CliFrontendSavepointTest {
Promise<Object> triggerResponse = new scala.concurrent.impl.Promise.DefaultPromise<>();
when(jobManager.ask(
- Mockito.eq(new DisposeSavepoint(savepointPath, Option.<List<BlobKey>>empty())),
+ Mockito.eq(new DisposeSavepoint(savepointPath)),
any(FiniteDuration.class)))
.thenReturn(triggerResponse.future());
@@ -323,7 +323,7 @@ public class CliFrontendSavepointTest {
assertTrue(returnCode != 0);
verify(jobManager, times(1)).ask(
- Mockito.eq(new DisposeSavepoint(savepointPath, Option.<List<BlobKey>>empty())),
+ Mockito.eq(new DisposeSavepoint(savepointPath)),
any(FiniteDuration.class));
assertTrue(buffer.toString().contains("expectedTestException"));
@@ -344,7 +344,7 @@ public class CliFrontendSavepointTest {
Promise<Object> triggerResponse = new scala.concurrent.impl.Promise.DefaultPromise<>();
when(jobManager.ask(
- Mockito.eq(new DisposeSavepoint(savepointPath, Option.<List<BlobKey>>empty())),
+ Mockito.eq(new DisposeSavepoint(savepointPath)),
any(FiniteDuration.class)))
.thenReturn(triggerResponse.future());
@@ -358,7 +358,7 @@ public class CliFrontendSavepointTest {
assertTrue(returnCode != 0);
verify(jobManager, times(1)).ask(
- Mockito.eq(new DisposeSavepoint(savepointPath, Option.<List<BlobKey>>empty())),
+ Mockito.eq(new DisposeSavepoint(savepointPath)),
any(FiniteDuration.class));
String errMsg = buffer.toString();
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
index 2c9a5d2..74276c0 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
@@ -17,22 +17,7 @@
package org.apache.flink.contrib.streaming.state;
-import java.io.EOFException;
-import java.io.File;
-import java.io.IOException;
-import java.io.ObjectInputStream;
-import java.io.ObjectOutputStream;
-import java.io.Serializable;
-import java.net.URI;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Random;
-import java.util.UUID;
-
import org.apache.commons.io.FileUtils;
-
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.state.FoldingState;
import org.apache.flink.api.common.state.FoldingStateDescriptor;
@@ -40,6 +25,7 @@ import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ReducingState;
import org.apache.flink.api.common.state.ReducingStateDescriptor;
+import org.apache.flink.api.common.state.StateBackend;
import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
@@ -49,21 +35,21 @@ import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.typeutils.runtime.DataInputViewStream;
import org.apache.flink.core.fs.Path;
import org.apache.flink.core.memory.DataInputView;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputView;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.fs.hdfs.HadoopFileSystem;
import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.runtime.state.AsynchronousKvStateSnapshot;
import org.apache.flink.runtime.state.KvState;
import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.api.common.state.StateBackend;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.runtime.util.SerializableObject;
import org.apache.flink.streaming.util.HDFSCopyFromLocal;
import org.apache.flink.streaming.util.HDFSCopyToLocal;
-
import org.apache.hadoop.fs.FileSystem;
-
import org.rocksdb.BackupEngine;
import org.rocksdb.BackupableDBOptions;
import org.rocksdb.ColumnFamilyDescriptor;
@@ -76,10 +62,22 @@ import org.rocksdb.RestoreOptions;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.RocksIterator;
-
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.io.EOFException;
+import java.io.File;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.net.URI;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.UUID;
+
import static java.util.Objects.requireNonNull;
/**
@@ -312,9 +310,9 @@ public class RocksDBStateBackend extends AbstractStateBackend {
}
@Override
- public void dispose() {
- super.dispose();
- nonPartitionedStateBackend.dispose();
+ public void discardState() throws Exception {
+ super.discardState();
+ nonPartitionedStateBackend.discardState();
// we have to lock because we might have an asynchronous checkpoint going on
synchronized (dbCleanupLock) {
@@ -569,7 +567,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
private void restoreFromFullyAsyncSnapshot(FinalFullyAsyncSnapshot snapshot) throws Exception {
- DataInputView inputView = snapshot.stateHandle.getState(userCodeClassLoader);
+ DataInputView inputView = new DataInputViewStreamWrapper(snapshot.stateHandle.openInputStream());
// clear k/v state information before filling it
kvStateInformation.clear();
@@ -729,8 +727,8 @@ public class RocksDBStateBackend extends AbstractStateBackend {
try {
long startTime = System.currentTimeMillis();
- CheckpointStateOutputView outputView = backend.createCheckpointStateOutputView(checkpointId, startTime);
-
+ CheckpointStateOutputStream outputStream = backend.createCheckpointStateOutputStream(checkpointId, startTime);
+ DataOutputView outputView = new DataOutputViewStreamWrapper(outputStream);
outputView.writeInt(columnFamilies.size());
// we don't know how many key/value pairs there are in each column family.
@@ -743,7 +741,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
outputView.writeByte(count);
- ObjectOutputStream ooOut = new ObjectOutputStream(outputView);
+ ObjectOutputStream ooOut = new ObjectOutputStream(outputStream);
ooOut.writeObject(column.getValue().f1);
ooOut.flush();
@@ -774,7 +772,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
}
}
- StateHandle<DataInputView> stateHandle = outputView.closeAndGetHandle();
+ StreamStateHandle stateHandle = outputStream.closeAndGetHandle();
long endTime = System.currentTimeMillis();
LOG.info("Fully asynchronous RocksDB materialization to " + backupUri + " (asynchronous part) took " + (endTime - startTime) + " ms.");
@@ -798,14 +796,14 @@ public class RocksDBStateBackend extends AbstractStateBackend {
private static class FinalFullyAsyncSnapshot implements KvStateSnapshot<Object, Object, ValueState<Object>, ValueStateDescriptor<Object>, RocksDBStateBackend> {
private static final long serialVersionUID = 1L;
- final StateHandle<DataInputView> stateHandle;
+ final StreamStateHandle stateHandle;
final long checkpointId;
/**
* Creates a new snapshot from the given state parameters.
*/
- private FinalFullyAsyncSnapshot(StateHandle<DataInputView> stateHandle, long checkpointId) {
- this.stateHandle = requireNonNull(stateHandle);
+ private FinalFullyAsyncSnapshot(StreamStateHandle stateHandle, long checkpointId) {
+ this.stateHandle = stateHandle;
this.checkpointId = checkpointId;
}
@@ -929,13 +927,6 @@ public class RocksDBStateBackend extends AbstractStateBackend {
return nonPartitionedStateBackend.createCheckpointStateOutputStream(checkpointID, timestamp);
}
- @Override
- public <S extends Serializable> StateHandle<S> checkpointStateSerializable(
- S state, long checkpointID, long timestamp) throws Exception {
-
- return nonPartitionedStateBackend.checkpointStateSerializable(state, checkpointID, timestamp);
- }
-
// ------------------------------------------------------------------------
// Parameters
// ------------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
index 7118cf6..d720c6d 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
@@ -28,9 +28,11 @@ import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.core.testutils.OneShotLatch;
import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
-import org.apache.flink.runtime.state.StateHandle;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
@@ -41,14 +43,13 @@ import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
import org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness;
import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskStateList;
import org.apache.flink.util.OperatingSystem;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.LocalFileSystem;
import org.junit.Assume;
import org.junit.Before;
+import org.junit.Ignore;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.powermock.api.mockito.PowerMockito;
@@ -59,6 +60,7 @@ import org.powermock.modules.junit4.PowerMockRunner;
import java.io.File;
import java.lang.reflect.Field;
import java.net.URI;
+import java.util.List;
import java.util.UUID;
import static org.junit.Assert.assertEquals;
@@ -86,6 +88,7 @@ public class RocksDBAsyncKVSnapshotTest {
* test will simply lock forever.
*/
@Test
+ @Ignore
public void testAsyncCheckpoints() throws Exception {
LocalFileSystem localFS = new LocalFileSystem();
localFS.initialize(new URI("file:///"), new Configuration());
@@ -130,8 +133,10 @@ public class RocksDBAsyncKVSnapshotTest {
}
@Override
- public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) {
- super.acknowledgeCheckpoint(checkpointId, state);
+ public void acknowledgeCheckpoint(long checkpointId,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
+ List<KeyGroupsStateHandle> keyGroupStateHandles) {
+ super.acknowledgeCheckpoint(checkpointId, chainedStateHandle, keyGroupStateHandles);
// block on the latch, to verify that triggerCheckpoint returns below,
// even though the async checkpoint would not finish
@@ -141,12 +146,10 @@ public class RocksDBAsyncKVSnapshotTest {
e.printStackTrace();
}
- assertTrue(state instanceof StreamTaskStateList);
- StreamTaskStateList stateList = (StreamTaskStateList) state;
// should be only one k/v state
- StreamTaskState taskState = stateList.getState(this.getUserClassLoader())[0];
- assertEquals(1, taskState.getKvStates().size());
+
+ assertEquals(1, keyGroupStateHandles.size());
// we now know that the checkpoint went through
ensureCheckpointLatch.trigger();
@@ -188,6 +191,7 @@ public class RocksDBAsyncKVSnapshotTest {
* test will simply lock forever.
*/
@Test
+ @Ignore
public void testFullyAsyncCheckpoints() throws Exception {
LocalFileSystem localFS = new LocalFileSystem();
localFS.initialize(new URI("file:///"), new Configuration());
@@ -233,8 +237,10 @@ public class RocksDBAsyncKVSnapshotTest {
}
@Override
- public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) {
- super.acknowledgeCheckpoint(checkpointId, state);
+ public void acknowledgeCheckpoint(long checkpointId,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
+ List<KeyGroupsStateHandle> keyGroupStateHandles) {
+ super.acknowledgeCheckpoint(checkpointId, chainedStateHandle, keyGroupStateHandles);
// block on the latch, to verify that triggerCheckpoint returns below,
// even though the async checkpoint would not finish
@@ -244,12 +250,8 @@ public class RocksDBAsyncKVSnapshotTest {
e.printStackTrace();
}
- assertTrue(state instanceof StreamTaskStateList);
- StreamTaskStateList stateList = (StreamTaskStateList) state;
-
// should be only one k/v state
- StreamTaskState taskState = stateList.getState(this.getUserClassLoader())[0];
- assertEquals(1, taskState.getKvStates().size());
+ assertEquals(1, keyGroupStateHandles.size());
// we now know that the checkpoint went through
ensureCheckpointLatch.trigger();
@@ -322,7 +324,7 @@ public class RocksDBAsyncKVSnapshotTest {
// not interested
}
}
-
+
public static class DummyMapFunction<T> implements MapFunction<T, T> {
@Override
public T map(T value) { return value; }
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java b/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java
index b43b24d..eb04038 100644
--- a/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java
+++ b/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java
@@ -20,8 +20,6 @@ package org.apache.flink.storm.tests;
import backtype.storm.Config;
import backtype.storm.topology.TopologyBuilder;
import backtype.storm.tuple.Fields;
-
-import org.apache.flink.util.MathUtils;
import org.apache.flink.storm.api.FlinkLocalCluster;
import org.apache.flink.storm.api.FlinkTopology;
import org.apache.flink.storm.tests.operators.FiniteRandomSpout;
@@ -29,6 +27,13 @@ import org.apache.flink.storm.tests.operators.TaskIdBolt;
import org.apache.flink.storm.util.BoltFileSink;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.util.StreamingProgramTestBase;
+import org.apache.flink.util.MathUtils;
+import org.junit.Assert;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
/**
* This test relies on the hash function used by the {@link DataStream#keyBy}, which is
@@ -49,9 +54,20 @@ public class StormFieldsGroupingITCase extends StreamingProgramTestBase {
@Override
protected void postSubmit() throws Exception {
- compareResultsByLinesInMemory("4> -1155484576\n" + "3> 1033096058\n" + "3> -1930858313\n" +
- "4> 1431162155\n" + "3> -1557280266\n" + "4> -1728529858\n" + "3> 1654374947\n" +
- "3> -65105105\n" + "3> -518907128\n" + "4> -252332814\n", this.resultPath);
+ List<String> expectedResults = Arrays.asList(
+ "-1155484576", "1033096058", "-1930858313", "1431162155", "-1557280266", "-1728529858", "1654374947",
+ "-65105105", "-518907128", "-252332814");
+
+ List<String> actualResults = new ArrayList<>();
+ readAllResultLines(actualResults, resultPath, new String[0], false);
+
+ Assert.assertEquals(expectedResults.size(),actualResults.size());
+ Collections.sort(actualResults);
+ Collections.sort(expectedResults);
+ for(int i=0; i< actualResults.size(); ++i) {
+ //compare against actual results with removed prefex (as it depends e.g. on the hash function used)
+ Assert.assertEquals(expectedResults.get(i), actualResults.get(i).substring(3));
+ }
}
@Override
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java b/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java
index 23463e1..bb0691e 100644
--- a/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java
+++ b/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java
@@ -41,7 +41,13 @@ public interface KeyGroupAssigner<K> extends Serializable {
/**
* Setups the key group assigner with the maximum parallelism (= number of key groups).
*
- * @param maxParallelism Maximum parallelism (= number of key groups)
+ * @param numberOfKeygroups Maximum parallelism (= number of key groups)
*/
- void setup(int maxParallelism);
+ void setup(int numberOfKeygroups);
+
+ /**
+ *
+ * @return configured maximum parallelism
+ */
+ int getNumberKeyGroups();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-core/src/main/java/org/apache/flink/util/InstantiationUtil.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/util/InstantiationUtil.java b/flink-core/src/main/java/org/apache/flink/util/InstantiationUtil.java
index e94b254..b1dddae 100644
--- a/flink-core/src/main/java/org/apache/flink/util/InstantiationUtil.java
+++ b/flink-core/src/main/java/org/apache/flink/util/InstantiationUtil.java
@@ -32,6 +32,7 @@ import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.ObjectStreamClass;
+import java.io.OutputStream;
import java.io.Serializable;
import java.lang.reflect.Constructor;
import java.lang.reflect.Modifier;
@@ -294,15 +295,46 @@ public final class InstantiationUtil {
Thread.currentThread().setContextClassLoader(old);
}
}
-
+
+ @SuppressWarnings("unchecked")
+ public static <T> T deserializeObject(InputStream in, ClassLoader cl) throws IOException, ClassNotFoundException {
+ final ClassLoader old = Thread.currentThread().getContextClassLoader();
+ try (ObjectInputStream oois = new ClassLoaderObjectInputStream(in, cl)) {
+ Thread.currentThread().setContextClassLoader(cl);
+ return (T) oois.readObject();
+ }
+ finally {
+ Thread.currentThread().setContextClassLoader(old);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ public static <T> T deserializeObject(byte[] bytes) throws IOException, ClassNotFoundException {
+ ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(bytes);
+ return deserializeObject(byteArrayInputStream);
+ }
+
+ @SuppressWarnings("unchecked")
+ public static <T> T deserializeObject(InputStream in) throws IOException, ClassNotFoundException {
+ ObjectInputStream objectInputStream = new ObjectInputStream(in);
+ return (T) objectInputStream.readObject();
+ }
+
public static byte[] serializeObject(Object o) throws IOException {
try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
ObjectOutputStream oos = new ObjectOutputStream(baos)) {
oos.writeObject(o);
+ oos.flush();
return baos.toByteArray();
}
}
+ public static void serializeObject(OutputStream out, Object o) throws IOException {
+ ObjectOutputStream oos = new ObjectOutputStream(out);
+ oos.writeObject(o);
+ oos.flush();
+ }
+
/**
* Clones the given serializable object using Java serialization.
*
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java
----------------------------------------------------------------------
diff --git a/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java b/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java
index 4ec94f6..fd7bf5d 100644
--- a/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java
+++ b/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java
@@ -26,13 +26,11 @@ import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.core.fs.Path;
import org.apache.flink.core.testutils.CommonTestUtils;
import org.apache.flink.runtime.state.StateBackendTestBase;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.state.filesystem.FileStreamStateHandle;
+import org.apache.flink.runtime.state.filesystem.FileStateHandle;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
import org.apache.hadoop.conf.Configuration;
@@ -152,7 +150,12 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
// no file operations should be possible right now
try {
- backend.checkpointStateSerializable("exception train rolling in", 2L, System.currentTimeMillis());
+ FsStateBackend.FsCheckpointStateOutputStream out = backend.createCheckpointStateOutputStream(
+ 2L,
+ System.currentTimeMillis());
+
+ out.write(1);
+ out.closeAndGetHandle();
fail("should fail with an exception");
} catch (IllegalStateException e) {
// supreme!
@@ -177,39 +180,6 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
}
@Test
- public void testSerializableState() {
- try {
- FsStateBackend backend = CommonTestUtils.createCopySerializable(new FsStateBackend(randomHdfsFileUri(), 40));
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "dummy", IntSerializer.INSTANCE);
-
- Path checkpointDir = backend.getCheckpointDirectory();
-
- String state1 = "dummy state";
- String state2 = "row row row your boat";
- Integer state3 = 42;
-
- StateHandle<String> handle1 = backend.checkpointStateSerializable(state1, 439568923746L, System.currentTimeMillis());
- StateHandle<String> handle2 = backend.checkpointStateSerializable(state2, 439568923746L, System.currentTimeMillis());
- StateHandle<Integer> handle3 = backend.checkpointStateSerializable(state3, 439568923746L, System.currentTimeMillis());
-
- assertEquals(state1, handle1.getState(getClass().getClassLoader()));
- handle1.discardState();
-
- assertEquals(state2, handle2.getState(getClass().getClassLoader()));
- handle2.discardState();
-
- assertEquals(state3, handle3.getState(getClass().getClassLoader()));
- handle3.discardState();
-
- assertTrue(isDirectoryEmpty(checkpointDir));
- }
- catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
- }
- }
-
- @Test
public void testStateOutputStream() {
try {
FsStateBackend backend = CommonTestUtils.createCopySerializable(new FsStateBackend(randomHdfsFileUri(), 15));
@@ -241,16 +211,16 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
stream2.write(state2);
stream3.write(state3);
- FileStreamStateHandle handle1 = (FileStreamStateHandle) stream1.closeAndGetHandle();
+ FileStateHandle handle1 = (FileStateHandle) stream1.closeAndGetHandle();
ByteStreamStateHandle handle2 = (ByteStreamStateHandle) stream2.closeAndGetHandle();
ByteStreamStateHandle handle3 = (ByteStreamStateHandle) stream3.closeAndGetHandle();
// use with try-with-resources
- StreamStateHandle handle4;
+ FileStateHandle handle4;
try (AbstractStateBackend.CheckpointStateOutputStream stream4 =
backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis())) {
stream4.write(state4);
- handle4 = stream4.closeAndGetHandle();
+ handle4 = (FileStateHandle) stream4.closeAndGetHandle();
}
// close before accessing handle
@@ -265,18 +235,18 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
// uh-huh
}
- validateBytesInStream(handle1.getState(getClass().getClassLoader()), state1);
+ validateBytesInStream(handle1.openInputStream(), state1);
handle1.discardState();
assertFalse(isDirectoryEmpty(checkpointDir));
ensureFileDeleted(handle1.getFilePath());
- validateBytesInStream(handle2.getState(getClass().getClassLoader()), state2);
+ validateBytesInStream(handle2.openInputStream(), state2);
handle2.discardState();
- validateBytesInStream(handle3.getState(getClass().getClassLoader()), state3);
+ validateBytesInStream(handle3.openInputStream(), state3);
handle3.discardState();
- validateBytesInStream(handle4.getState(getClass().getClassLoader()), state4);
+ validateBytesInStream(handle4.openInputStream(), state4);
handle4.discardState();
assertTrue(isDirectoryEmpty(checkpointDir));
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractCEPPatternOperator.java
----------------------------------------------------------------------
diff --git a/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractCEPPatternOperator.java b/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractCEPPatternOperator.java
index c26eaeb..64ffa2a 100644
--- a/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractCEPPatternOperator.java
+++ b/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractCEPPatternOperator.java
@@ -21,17 +21,16 @@ package org.apache.flink.cep.operator;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.cep.nfa.NFA;
import org.apache.flink.cep.nfa.compiler.NFACompiler;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.core.memory.DataInputViewStreamWrapper;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.MultiplexingStreamRecordSerializer;
import org.apache.flink.streaming.runtime.streamrecord.StreamElement;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import java.io.IOException;
-import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.PriorityQueue;
@@ -104,51 +103,34 @@ abstract public class AbstractCEPPatternOperator<IN, OUT> extends AbstractCEPBas
}
@Override
- public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
- StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
-
- final AbstractStateBackend.CheckpointStateOutputStream os = this.getStateBackend().createCheckpointStateOutputStream(
- checkpointId,
- timestamp);
-
- final ObjectOutputStream oos = new ObjectOutputStream(os);
- final AbstractStateBackend.CheckpointStateOutputView ov = new AbstractStateBackend.CheckpointStateOutputView(os);
+ public void snapshotState(FSDataOutputStream out, long checkpointId, long timestamp) throws Exception {
+ super.snapshotState(out, checkpointId, timestamp);
+ final ObjectOutputStream oos = new ObjectOutputStream(out);
oos.writeObject(nfa);
-
- ov.writeInt(priorityQueue.size());
+ oos.writeInt(priorityQueue.size());
for (StreamRecord<IN> streamRecord: priorityQueue) {
- streamRecordSerializer.serialize(streamRecord, ov);
+ streamRecordSerializer.serialize(streamRecord, new DataOutputViewStreamWrapper(oos));
}
-
- taskState.setOperatorState(os.closeAndGetHandle());
-
- return taskState;
+ oos.flush();
}
@Override
@SuppressWarnings("unchecked")
- public void restoreState(StreamTaskState state) throws Exception {
+ public void restoreState(FSDataInputStream state) throws Exception {
super.restoreState(state);
-
- StreamStateHandle stream = (StreamStateHandle)state.getOperatorState();
-
- final InputStream is = stream.getState(getUserCodeClassloader());
- final ObjectInputStream ois = new ObjectInputStream(is);
- final DataInputViewStreamWrapper div = new DataInputViewStreamWrapper(is);
+ final ObjectInputStream ois = new ObjectInputStream(state);
nfa = (NFA<IN>)ois.readObject();
- int numberPriorityQueueEntries = div.readInt();
+ int numberPriorityQueueEntries = ois.readInt();
priorityQueue = new PriorityQueue<StreamRecord<IN>>(numberPriorityQueueEntries, new StreamRecordComparator<IN>());
for (int i = 0; i <numberPriorityQueueEntries; i++) {
- StreamElement streamElement = streamRecordSerializer.deserialize(div);
+ StreamElement streamElement = streamRecordSerializer.deserialize(new DataInputViewStreamWrapper(ois));
priorityQueue.offer(streamElement.<IN>asRecord());
}
-
- div.close();
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java
----------------------------------------------------------------------
diff --git a/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java b/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java
index 206be47..e3f924c 100644
--- a/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java
+++ b/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java
@@ -24,14 +24,15 @@ import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.cep.nfa.NFA;
import org.apache.flink.cep.nfa.compiler.NFACompiler;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.core.memory.DataInputView;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
import org.apache.flink.core.memory.DataOutputView;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.MultiplexingStreamRecordSerializer;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import java.io.IOException;
import java.io.Serializable;
@@ -183,30 +184,22 @@ abstract public class AbstractKeyedCEPPatternOperator<IN, KEY, OUT> extends Abst
}
@Override
- public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
- StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
-
- AbstractStateBackend.CheckpointStateOutputView ov = getStateBackend().createCheckpointStateOutputView(checkpointId, timestamp);
+ public void snapshotState(FSDataOutputStream out, long checkpointId, long timestamp) throws Exception {
+ super.snapshotState(out, checkpointId, timestamp);
+ DataOutputView ov = new DataOutputViewStreamWrapper(out);
ov.writeInt(keys.size());
for (KEY key: keys) {
keySerializer.serialize(key, ov);
}
-
- taskState.setOperatorState(ov.closeAndGetHandle());
-
- return taskState;
}
@Override
- public void restoreState(StreamTaskState state) throws Exception {
+ public void restoreState(FSDataInputStream state) throws Exception {
super.restoreState(state);
- @SuppressWarnings("unchecked")
- StateHandle<DataInputView> stateHandle = (StateHandle<DataInputView>) state.getOperatorState();
-
- DataInputView inputView = stateHandle.getState(getUserCodeClassloader());
+ DataInputView inputView = new DataInputViewStreamWrapper(state);
if (keys == null) {
keys = new HashSet<>();
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java b/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
index 56c4161..54c1477 100644
--- a/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
+++ b/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
@@ -28,11 +28,11 @@ import org.apache.flink.cep.nfa.NFA;
import org.apache.flink.cep.nfa.compiler.NFACompiler;
import org.apache.flink.cep.pattern.Pattern;
import org.apache.flink.contrib.streaming.state.RocksDBStateBackend;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.util.TestLogger;
import org.junit.Rule;
@@ -135,7 +135,7 @@ public class CEPOperatorTest extends TestLogger {
harness.processElement(new StreamRecord<Event>(new Event(42, "foobar", 1.0), 2));
// simulate snapshot/restore with some elements in internal sorting queue
- StreamTaskState snapshot = harness.snapshot(0, 0);
+ StreamStateHandle snapshot = harness.snapshot(0, 0);
harness = new OneInputStreamOperatorTestHarness<>(
new CEPPatternOperator<>(
@@ -144,7 +144,7 @@ public class CEPOperatorTest extends TestLogger {
new NFAFactory()));
harness.setup();
- harness.restore(snapshot, 1);
+ harness.restore(snapshot);
harness.open();
harness.processWatermark(new Watermark(Long.MIN_VALUE));
@@ -156,7 +156,7 @@ public class CEPOperatorTest extends TestLogger {
harness.processWatermark(new Watermark(2));
// simulate snapshot/restore with empty element queue but NFA state
- StreamTaskState snapshot2 = harness.snapshot(1, 1);
+ StreamStateHandle snapshot2 = harness.snapshot(1, 1);
harness = new OneInputStreamOperatorTestHarness<>(
new CEPPatternOperator<>(
@@ -165,7 +165,7 @@ public class CEPOperatorTest extends TestLogger {
new NFAFactory()));
harness.setup();
- harness.restore(snapshot2, 2);
+ harness.restore(snapshot2);
harness.open();
harness.processElement(new StreamRecord<Event>(middleEvent, 3));
@@ -226,7 +226,7 @@ public class CEPOperatorTest extends TestLogger {
harness.processElement(new StreamRecord<Event>(new Event(42, "foobar", 1.0), 2));
// simulate snapshot/restore with some elements in internal sorting queue
- StreamTaskState snapshot = harness.snapshot(0, 0);
+ StreamStateHandle snapshot = harness.snapshot(0, 0);
harness = new OneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
@@ -238,7 +238,7 @@ public class CEPOperatorTest extends TestLogger {
harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
harness.setup();
- harness.restore(snapshot, 1);
+ harness.restore(snapshot);
harness.open();
harness.processWatermark(new Watermark(Long.MIN_VALUE));
@@ -250,7 +250,7 @@ public class CEPOperatorTest extends TestLogger {
harness.processWatermark(new Watermark(2));
// simulate snapshot/restore with empty element queue but NFA state
- StreamTaskState snapshot2 = harness.snapshot(1, 1);
+ StreamStateHandle snapshot2 = harness.snapshot(1, 1);
harness = new OneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
@@ -262,7 +262,7 @@ public class CEPOperatorTest extends TestLogger {
harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
harness.setup();
- harness.restore(snapshot2, 2);
+ harness.restore(snapshot2);
harness.open();
harness.processElement(new StreamRecord<Event>(middleEvent, 3));
@@ -330,7 +330,7 @@ public class CEPOperatorTest extends TestLogger {
harness.processElement(new StreamRecord<Event>(new Event(42, "foobar", 1.0), 2));
// simulate snapshot/restore with some elements in internal sorting queue
- StreamTaskState snapshot = harness.snapshot(0, 0);
+ StreamStateHandle snapshot = harness.snapshot(0, 0);
harness = new OneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
@@ -346,7 +346,7 @@ public class CEPOperatorTest extends TestLogger {
harness.setStateBackend(rocksDBStateBackend);
harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
harness.setup();
- harness.restore(snapshot, 1);
+ harness.restore(snapshot);
harness.open();
harness.processWatermark(new Watermark(Long.MIN_VALUE));
@@ -358,7 +358,7 @@ public class CEPOperatorTest extends TestLogger {
harness.processWatermark(new Watermark(2));
// simulate snapshot/restore with empty element queue but NFA state
- StreamTaskState snapshot2 = harness.snapshot(1, 1);
+ StreamStateHandle snapshot2 = harness.snapshot(1, 1);
harness = new OneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
@@ -374,7 +374,7 @@ public class CEPOperatorTest extends TestLogger {
harness.setStateBackend(rocksDBStateBackend);
harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
harness.setup();
- harness.restore(snapshot2, 2);
+ harness.restore(snapshot2);
harness.open();
harness.processElement(new StreamRecord<Event>(middleEvent, 3));
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-mesos/src/main/java/org/apache/flink/mesos/runtime/clusterframework/store/ZooKeeperMesosWorkerStore.java
----------------------------------------------------------------------
diff --git a/flink-mesos/src/main/java/org/apache/flink/mesos/runtime/clusterframework/store/ZooKeeperMesosWorkerStore.java b/flink-mesos/src/main/java/org/apache/flink/mesos/runtime/clusterframework/store/ZooKeeperMesosWorkerStore.java
index c5cef8e..551852e 100644
--- a/flink-mesos/src/main/java/org/apache/flink/mesos/runtime/clusterframework/store/ZooKeeperMesosWorkerStore.java
+++ b/flink-mesos/src/main/java/org/apache/flink/mesos/runtime/clusterframework/store/ZooKeeperMesosWorkerStore.java
@@ -25,9 +25,9 @@ import org.apache.curator.framework.recipes.shared.VersionedValue;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
import org.apache.flink.runtime.util.ZooKeeperUtils;
-import org.apache.flink.runtime.zookeeper.StateStorageHelper;
+import org.apache.flink.runtime.zookeeper.RetrievableStateStorageHelper;
import org.apache.flink.runtime.zookeeper.ZooKeeperStateHandleStore;
import org.apache.mesos.Protos;
import org.apache.zookeeper.KeeperException;
@@ -74,7 +74,7 @@ public class ZooKeeperMesosWorkerStore implements MesosWorkerStore {
ZooKeeperMesosWorkerStore(
CuratorFramework client,
String storePath,
- StateStorageHelper<MesosWorkerStore.Worker> stateStorage
+ RetrievableStateStorageHelper<Worker> stateStorage
) throws Exception {
checkNotNull(storePath, "storePath");
checkNotNull(stateStorage, "stateStorage");
@@ -100,7 +100,7 @@ public class ZooKeeperMesosWorkerStore implements MesosWorkerStore {
// using late-binding as a workaround for shaded curator dependency of flink-runtime.
this.workersInZooKeeper = ZooKeeperStateHandleStore.class
- .getConstructor(CuratorFramework.class, StateStorageHelper.class)
+ .getConstructor(CuratorFramework.class, RetrievableStateStorageHelper.class)
.newInstance(storeFacade, stateStorage);
}
@@ -203,12 +203,12 @@ public class ZooKeeperMesosWorkerStore implements MesosWorkerStore {
synchronized (startStopLock) {
verifyIsRunning();
- List<Tuple2<StateHandle<MesosWorkerStore.Worker>, String>> handles = workersInZooKeeper.getAll();
+ List<Tuple2<RetrievableStateHandle<Worker>, String>> handles = workersInZooKeeper.getAll();
if(handles.size() != 0) {
List<MesosWorkerStore.Worker> workers = new ArrayList<>(handles.size());
- for (Tuple2<StateHandle<MesosWorkerStore.Worker>, String> handle : handles) {
- Worker worker = handle.f0.getState(ClassLoader.getSystemClassLoader());
+ for (Tuple2<RetrievableStateHandle<Worker>, String> handle : handles) {
+ Worker worker = handle.f0.retrieveState();
workers.add(worker);
}
@@ -288,7 +288,7 @@ public class ZooKeeperMesosWorkerStore implements MesosWorkerStore {
checkNotNull(configuration, "Configuration");
- StateStorageHelper<MesosWorkerStore.Worker> stateStorage =
+ RetrievableStateStorageHelper<MesosWorkerStore.Worker> stateStorage =
ZooKeeperUtils.createFileSystemStateStorage(configuration, "mesosWorkerStore");
String zooKeeperMesosWorkerStorePath = configuration.getString(
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
index 24fca5f..e78e203 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
@@ -34,9 +34,11 @@ import org.apache.flink.runtime.messages.checkpoint.AcknowledgeCheckpoint;
import org.apache.flink.runtime.messages.checkpoint.DeclineCheckpoint;
import org.apache.flink.runtime.messages.checkpoint.NotifyCheckpointComplete;
import org.apache.flink.runtime.messages.checkpoint.TriggerCheckpoint;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
-
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.util.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -45,12 +47,10 @@ import scala.concurrent.Future;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.HashMap;
-import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
-import java.util.Set;
import java.util.Timer;
import java.util.TimerTask;
@@ -65,7 +65,7 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
*/
public class CheckpointCoordinator {
- static final Logger LOG = LoggerFactory.getLogger(CheckpointCoordinator.class);
+ protected static final Logger LOG = LoggerFactory.getLogger(CheckpointCoordinator.class);
/** The number of recent checkpoints whose IDs are remembered */
private static final int NUM_GHOST_CHECKPOINT_IDS = 16;
@@ -402,9 +402,9 @@ public class CheckpointCoordinator {
return new CheckpointTriggerResult(CheckpointDeclineReason.EXCEPTION);
}
- final PendingCheckpoint checkpoint = props.isSavepoint() ?
- new PendingSavepoint(job, checkpointID, timestamp, ackTasks, userClassLoader, savepointStore) :
- new PendingCheckpoint(job, checkpointID, timestamp, ackTasks, userClassLoader);
+ final PendingCheckpoint checkpoint = props.isSavepoint() ?
+ new PendingSavepoint(job, checkpointID, timestamp, ackTasks, savepointStore) :
+ new PendingCheckpoint(job, checkpointID, timestamp, ackTasks);
// schedule the timer that will clean up the expired checkpoints
TimerTask canceller = new TimerTask() {
@@ -632,10 +632,8 @@ public class CheckpointCoordinator {
if (checkpoint.acknowledgeTask(
message.getTaskExecutionId(),
- message.getState(),
- message.getStateSize(),
- null)) { // TODO: Give KV-state to the acknowledgeTask method
-
+ message.getStateHandle(),
+ message.getKeyGroupsStateHandle())) {
if (checkpoint.isFullyAcknowledged()) {
completed = checkpoint.finalizeCheckpoint();
@@ -783,32 +781,60 @@ public class CheckpointCoordinator {
ExecutionJobVertex executionJobVertex = tasks.get(taskGroupStateEntry.getKey());
if (executionJobVertex != null) {
- // check that we only restore the state if the parallelism has not been changed
- if (taskState.getParallelism() != executionJobVertex.getParallelism()) {
- throw new RuntimeException("Cannot restore the latest checkpoint because " +
- "the parallelism changed. The operator" + executionJobVertex.getJobVertexId() +
+ // check that the number of key groups have not changed
+ if (taskState.getMaxParallelism() != executionJobVertex.getMaxParallelism()) {
+ throw new IllegalStateException("The maximum parallelism (" +
+ taskState.getMaxParallelism() + ") with which the latest " +
+ "checkpoint of the execution job vertex " + executionJobVertex +
+ " has been taken and the current maximum parallelism (" +
+ executionJobVertex.getMaxParallelism() + ") changed. This " +
+ "is currently not supported.");
+ }
+
+
+ boolean hasNonPartitionedState = taskState.hasNonPartitionedState();
+
+ if (hasNonPartitionedState && taskState.getParallelism() != executionJobVertex.getParallelism()) {
+ throw new IllegalStateException("Cannot restore the latest checkpoint because " +
+ "the operator " + executionJobVertex.getJobVertexId() + " has non-partitioned " +
+ "state and its parallelism changed. The operator" + executionJobVertex.getJobVertexId() +
" has parallelism " + executionJobVertex.getParallelism() + " whereas the corresponding" +
"state object has a parallelism of " + taskState.getParallelism());
}
+ List<KeyGroupRange> keyGroupPartitions = createKeyGroupPartitions(
+ executionJobVertex.getMaxParallelism(),
+ executionJobVertex.getParallelism());
+
int counter = 0;
+ for (int i = 0; i < executionJobVertex.getParallelism(); i++) {
+ ChainedStateHandle<StreamStateHandle> state = null;
- List<Set<Integer>> keyGroupPartitions = createKeyGroupPartitions(executionJobVertex.getMaxParallelism(), executionJobVertex.getParallelism());
+ if (hasNonPartitionedState) {
+ SubtaskState subtaskState = taskState.getState(i);
- for (int i = 0; i < executionJobVertex.getParallelism(); i++) {
- SubtaskState subtaskState = taskState.getState(i);
- SerializedValue<StateHandle<?>> state = null;
+ if (subtaskState != null) {
+ // count the number of executions for which we set a state
+ counter++;
+ state = subtaskState.getChainedStateHandle();
+ }
+ }
+
+ KeyGroupRange subtaskKeyGroupIds = keyGroupPartitions.get(i);
+ List<KeyGroupsStateHandle> subtaskKeyGroupStates = new ArrayList<>();
- if (subtaskState != null) {
- // count the number of executions for which we set a state
- counter++;
- state = subtaskState.getState();
+ for (KeyGroupsStateHandle storedKeyGroup : taskState.getKeyGroupStates()) {
+ KeyGroupsStateHandle intersection = storedKeyGroup.getKeyGroupIntersection(subtaskKeyGroupIds);
+ if(intersection.getNumberOfKeyGroups() > 0) {
+ subtaskKeyGroupStates.add(intersection);
+ }
}
- Map<Integer, SerializedValue<StateHandle<?>>> kvStateForTaskMap = taskState.getUnwrappedKvStates(keyGroupPartitions.get(i));
+ Execution currentExecutionAttempt = executionJobVertex
+ .getTaskVertices()[i]
+ .getCurrentExecutionAttempt();
- Execution currentExecutionAttempt = executionJobVertex.getTaskVertices()[i].getCurrentExecutionAttempt();
- currentExecutionAttempt.setInitialState(state, kvStateForTaskMap);
+ currentExecutionAttempt.setInitialState(state, subtaskKeyGroupStates);
}
if (allOrNothingState && counter > 0 && counter < executionJobVertex.getParallelism()) {
@@ -830,23 +856,20 @@ public class CheckpointCoordinator {
* the set of key groups which is assigned to the same task. Each set of the returned list
* constitutes a key group partition.
*
+ * <b>IMPORTANT</b>: The assignment of key groups to partitions has to be in sync with the
+ * KeyGroupStreamPartitioner.
+ *
* @param numberKeyGroups Number of available key groups (indexed from 0 to numberKeyGroups - 1)
* @param parallelism Parallelism to generate the key group partitioning for
* @return List of key group partitions
*/
- protected List<Set<Integer>> createKeyGroupPartitions(int numberKeyGroups, int parallelism) {
- ArrayList<Set<Integer>> result = new ArrayList<>(parallelism);
-
- for (int p = 0; p < parallelism; p++) {
- HashSet<Integer> keyGroupPartition = new HashSet<>();
-
- for (int k = p; k < numberKeyGroups; k += parallelism) {
- keyGroupPartition.add(k);
- }
-
- result.add(keyGroupPartition);
+ public static List<KeyGroupRange> createKeyGroupPartitions(int numberKeyGroups, int parallelism) {
+ Preconditions.checkArgument(numberKeyGroups >= parallelism);
+ List<KeyGroupRange> result = new ArrayList<>(parallelism);
+ int start = 0;
+ for (int i = 0; i < parallelism; ++i) {
+ result.add(KeyGroupRange.computeKeyGroupRangeForOperatorIndex(numberKeyGroups, parallelism, i));
}
-
return result;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CompletedCheckpoint.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CompletedCheckpoint.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CompletedCheckpoint.java
index d8c6c7d..e412006 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CompletedCheckpoint.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CompletedCheckpoint.java
@@ -20,9 +20,12 @@ package org.apache.flink.runtime.checkpoint;
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.state.StateObject;
+import org.apache.flink.runtime.state.StateUtil;
-import java.io.Serializable;
+import java.io.IOException;
import java.util.Map;
+import java.util.Objects;
import static org.apache.flink.util.Preconditions.checkArgument;
import static org.apache.flink.util.Preconditions.checkNotNull;
@@ -31,12 +34,12 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
* A successful checkpoint describes a checkpoint after all required tasks acknowledged it (with their state)
* and that is considered completed.
*/
-public class CompletedCheckpoint implements Serializable {
+public class CompletedCheckpoint implements StateObject {
private static final long serialVersionUID = -8360248179615702014L;
private final JobID job;
-
+
private final long checkpointID;
/** The timestamp when the checkpoint was triggered. */
@@ -92,11 +95,24 @@ public class CompletedCheckpoint implements Serializable {
return duration;
}
- public long getStateSize() {
+ @Override
+ public void discardState() throws Exception {
+ if (deleteStateWhenDisposed) {
+
+ try {
+ StateUtil.bestEffortDiscardAllStateObjects(taskStates.values());
+ } finally {
+ taskStates.clear();
+ }
+ }
+ }
+
+ @Override
+ public long getStateSize() throws Exception {
long result = 0L;
for (TaskState taskState : taskStates.values()) {
- result += taskState.getStateSize();
+ result += taskState.getStateSize();
}
return result;
@@ -112,20 +128,34 @@ public class CompletedCheckpoint implements Serializable {
// --------------------------------------------------------------------------------------------
- public void discard(ClassLoader userClassLoader) throws Exception {
- if (deleteStateWhenDisposed) {
- for (TaskState state: taskStates.values()) {
- state.discard(userClassLoader);
- }
+ @Override
+ public boolean equals(Object obj) {
+ if (obj instanceof CompletedCheckpoint) {
+ CompletedCheckpoint other = (CompletedCheckpoint) obj;
+
+ return job.equals(other.job) && checkpointID == other.checkpointID &&
+ timestamp == other.timestamp && duration == other.duration &&
+ taskStates.equals(other.taskStates);
+ } else {
+ return false;
}
-
- taskStates.clear();
}
- // --------------------------------------------------------------------------------------------
+ @Override
+ public int hashCode() {
+ return (int) (this.checkpointID ^ this.checkpointID >>> 32) +
+ 31 * ((int) (this.timestamp ^ this.timestamp >>> 32) +
+ 31 * ((int) (this.duration ^ this.duration >>> 32) +
+ 31 * Objects.hash(job, taskStates)));
+ }
@Override
public String toString() {
return String.format("Checkpoint %d @ %d for %s", checkpointID, timestamp, job);
}
+
+ @Override
+ public void close() throws IOException {
+ StateUtil.bestEffortCloseAllStateObjects(taskStates.values());
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/KeyGroupState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/KeyGroupState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/KeyGroupState.java
deleted file mode 100644
index eb358b6..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/KeyGroupState.java
+++ /dev/null
@@ -1,87 +0,0 @@
-/*
- * 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.StateHandle;
-import org.apache.flink.util.SerializedValue;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.Serializable;
-
-/**
- * Simple container class which contains the serialized state handle for a key group.
- *
- * The key group state handle is kept in serialized form because it can contain user code classes
- * which might not be available on the JobManager.
- */
-public class KeyGroupState implements Serializable {
- private static final long serialVersionUID = -5926696455438467634L;
-
- private static final Logger LOG = LoggerFactory.getLogger(KeyGroupState.class);
-
- private final SerializedValue<StateHandle<?>> keyGroupState;
-
- private final long stateSize;
-
- private final long duration;
-
- public KeyGroupState(SerializedValue<StateHandle<?>> keyGroupState, long stateSize, long duration) {
- this.keyGroupState = keyGroupState;
-
- this.stateSize = stateSize;
-
- this.duration = duration;
- }
-
- public SerializedValue<StateHandle<?>> getKeyGroupState() {
- return keyGroupState;
- }
-
- public long getDuration() {
- return duration;
- }
-
- public long getStateSize() {
- return stateSize;
- }
-
- public void discard(ClassLoader classLoader) throws Exception {
- keyGroupState.deserializeValue(classLoader).discardState();
- }
-
- @Override
- public boolean equals(Object obj) {
- if (obj instanceof KeyGroupState) {
- KeyGroupState other = (KeyGroupState) obj;
-
- return keyGroupState.equals(other.keyGroupState) && stateSize == other.stateSize &&
- duration == other.duration;
- } else {
- return false;
- }
- }
-
- @Override
- public int hashCode() {
- return (int) (this.stateSize ^ this.stateSize >>> 32) +
- 31 * ((int) (this.duration ^ this.duration >>> 32) +
- 31 * keyGroupState.hashCode());
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingCheckpoint.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingCheckpoint.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingCheckpoint.java
index ab3fe0a..d499a5a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingCheckpoint.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingCheckpoint.java
@@ -18,15 +18,18 @@
package org.apache.flink.runtime.checkpoint;
-import java.util.HashMap;
-import java.util.Map;
-
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
import org.apache.flink.runtime.executiongraph.ExecutionVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StateUtil;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
import static org.apache.flink.util.Preconditions.checkArgument;
import static org.apache.flink.util.Preconditions.checkNotNull;
@@ -53,8 +56,6 @@ public class PendingCheckpoint {
private final Map<ExecutionAttemptID, ExecutionVertex> notYetAcknowledgedTasks;
- private final ClassLoader userCodeClassLoader;
-
private final boolean disposeWhenSubsumed;
private int numAcknowledgedTasks;
@@ -67,9 +68,8 @@ public class PendingCheckpoint {
JobID jobId,
long checkpointId,
long checkpointTimestamp,
- Map<ExecutionAttemptID, ExecutionVertex> verticesToConfirm,
- ClassLoader userCodeClassLoader) {
- this(jobId, checkpointId, checkpointTimestamp, verticesToConfirm, userCodeClassLoader, true);
+ Map<ExecutionAttemptID, ExecutionVertex> verticesToConfirm) {
+ this(jobId, checkpointId, checkpointTimestamp, verticesToConfirm, true);
}
PendingCheckpoint(
@@ -77,14 +77,12 @@ public class PendingCheckpoint {
long checkpointId,
long checkpointTimestamp,
Map<ExecutionAttemptID, ExecutionVertex> verticesToConfirm,
- ClassLoader userCodeClassLoader,
boolean disposeWhenSubsumed)
{
this.jobId = checkNotNull(jobId);
this.checkpointId = checkpointId;
this.checkpointTimestamp = checkpointTimestamp;
this.notYetAcknowledgedTasks = checkNotNull(verticesToConfirm);
- this.userCodeClassLoader = checkNotNull(userCodeClassLoader);
this.disposeWhenSubsumed = disposeWhenSubsumed;
this.taskStates = new HashMap<>();
@@ -92,6 +90,8 @@ public class PendingCheckpoint {
"Checkpoint needs at least one vertex that commits the checkpoint");
}
+ // --------------------------------------------------------------------------------------------
+
// ------------------------------------------------------------------------
// Properties
// ------------------------------------------------------------------------
@@ -167,10 +167,9 @@ public class PendingCheckpoint {
}
public boolean acknowledgeTask(
- ExecutionAttemptID attemptID,
- SerializedValue<StateHandle<?>> state,
- long stateSize,
- Map<Integer, SerializedValue<StateHandle<?>>> kvState) {
+ ExecutionAttemptID attemptID,
+ ChainedStateHandle<StreamStateHandle> state,
+ List<KeyGroupsStateHandle> keyGroupsState) {
synchronized (lock) {
if (discarded) {
@@ -179,7 +178,7 @@ public class PendingCheckpoint {
ExecutionVertex vertex = notYetAcknowledgedTasks.remove(attemptID);
if (vertex != null) {
- if (state != null || kvState != null) {
+ if (state != null || keyGroupsState != null) {
JobVertexID jobVertexID = vertex.getJobvertexId();
@@ -188,33 +187,23 @@ public class PendingCheckpoint {
if (taskStates.containsKey(jobVertexID)) {
taskState = taskStates.get(jobVertexID);
} else {
- taskState = new TaskState(jobVertexID, vertex.getTotalNumberOfParallelSubtasks());
+ taskState = new TaskState(jobVertexID, vertex.getTotalNumberOfParallelSubtasks(), vertex.getMaxParallelism());
taskStates.put(jobVertexID, taskState);
}
- long timestamp = System.currentTimeMillis() - checkpointTimestamp;
+ long duration = System.currentTimeMillis() - checkpointTimestamp;
if (state != null) {
taskState.putState(
vertex.getParallelSubtaskIndex(),
- new SubtaskState(
- state,
- stateSize,
- timestamp
- )
- );
+ new SubtaskState(state, duration));
}
- if (kvState != null) {
- for (Map.Entry<Integer, SerializedValue<StateHandle<?>>> entry : kvState.entrySet()) {
- taskState.putKvState(
- entry.getKey(),
- new KeyGroupState(
- entry.getValue(),
- 0L,
- timestamp
- ));
- }
+ // currently a checkpoint can only contain keyed state
+ // for the head operator
+ if (keyGroupsState != null && !keyGroupsState.isEmpty()) {
+ KeyGroupsStateHandle keyGroupsStateHandle = keyGroupsState.get(0);
+ taskState.putKeyedState(vertex.getParallelSubtaskIndex(), keyGroupsStateHandle);
}
}
numAcknowledgedTasks++;
@@ -258,13 +247,11 @@ public class PendingCheckpoint {
protected void dispose(boolean releaseState) throws Exception {
synchronized (lock) {
- discarded = true;
- numAcknowledgedTasks = -1;
try {
+ discarded = true;
+ numAcknowledgedTasks = -1;
if (releaseState) {
- for (TaskState taskState : taskStates.values()) {
- taskState.discard(userCodeClassLoader);
- }
+ StateUtil.bestEffortDiscardAllStateObjects(taskStates.values());
}
} finally {
taskStates.clear();
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingSavepoint.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingSavepoint.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingSavepoint.java
index 460ff8e..0bb6a91 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingSavepoint.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/PendingSavepoint.java
@@ -21,7 +21,7 @@ package org.apache.flink.runtime.checkpoint;
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.checkpoint.savepoint.Savepoint;
import org.apache.flink.runtime.checkpoint.savepoint.SavepointStore;
-import org.apache.flink.runtime.checkpoint.savepoint.SavepointV0;
+import org.apache.flink.runtime.checkpoint.savepoint.SavepointV1;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
import org.apache.flink.runtime.executiongraph.ExecutionVertex;
import org.apache.flink.util.ExceptionUtils;
@@ -53,10 +53,9 @@ public class PendingSavepoint extends PendingCheckpoint {
long checkpointId,
long checkpointTimestamp,
Map<ExecutionAttemptID, ExecutionVertex> verticesToConfirm,
- ClassLoader userCodeClassLoader,
SavepointStore store)
{
- super(jobId, checkpointId, checkpointTimestamp, verticesToConfirm, userCodeClassLoader, false);
+ super(jobId, checkpointId, checkpointTimestamp, verticesToConfirm, false);
this.store = checkNotNull(store);
this.onCompletionPromise = new scala.concurrent.impl.Promise.DefaultPromise<>();
@@ -77,7 +76,7 @@ public class PendingSavepoint extends PendingCheckpoint {
// now store the checkpoint externally as a savepoint
try {
- Savepoint savepoint = new SavepointV0(
+ Savepoint savepoint = new SavepointV1(
completedCheckpoint.getCheckpointID(),
completedCheckpoint.getTaskStates().values());
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StandaloneCompletedCheckpointStore.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StandaloneCompletedCheckpointStore.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StandaloneCompletedCheckpointStore.java
index a35ca77..5e03988 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StandaloneCompletedCheckpointStore.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/StandaloneCompletedCheckpointStore.java
@@ -19,13 +19,13 @@
package org.apache.flink.runtime.checkpoint;
import org.apache.flink.runtime.jobmanager.HighAvailabilityMode;
+import org.apache.flink.runtime.state.StateUtil;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.List;
import static org.apache.flink.util.Preconditions.checkArgument;
-import static org.apache.flink.util.Preconditions.checkNotNull;
/**
* {@link CompletedCheckpointStore} for JobManagers running in {@link HighAvailabilityMode#NONE}.
@@ -35,9 +35,6 @@ public class StandaloneCompletedCheckpointStore implements CompletedCheckpointSt
/** The maximum number of checkpoints to retain (at least 1). */
private final int maxNumberOfCheckpointsToRetain;
- /** User class loader for discarding {@link CompletedCheckpoint} instances. */
- private final ClassLoader userClassLoader;
-
/** The completed checkpoints. */
private final ArrayDeque<CompletedCheckpoint> checkpoints;
@@ -56,7 +53,6 @@ public class StandaloneCompletedCheckpointStore implements CompletedCheckpointSt
checkArgument(maxNumberOfCheckpointsToRetain >= 1, "Must retain at least one checkpoint.");
this.maxNumberOfCheckpointsToRetain = maxNumberOfCheckpointsToRetain;
- this.userClassLoader = checkNotNull(userClassLoader, "User class loader");
this.checkpoints = new ArrayDeque<>(maxNumberOfCheckpointsToRetain + 1);
}
@@ -70,7 +66,7 @@ public class StandaloneCompletedCheckpointStore implements CompletedCheckpointSt
public void addCheckpoint(CompletedCheckpoint checkpoint) throws Exception {
checkpoints.addLast(checkpoint);
if (checkpoints.size() > maxNumberOfCheckpointsToRetain) {
- checkpoints.removeFirst().discard(userClassLoader);
+ checkpoints.removeFirst().discardState();
}
}
@@ -91,11 +87,11 @@ public class StandaloneCompletedCheckpointStore implements CompletedCheckpointSt
@Override
public void shutdown() throws Exception {
- for (CompletedCheckpoint checkpoint : checkpoints) {
- checkpoint.discard(userClassLoader);
+ try {
+ StateUtil.bestEffortDiscardAllStateObjects(checkpoints);
+ } finally {
+ checkpoints.clear();
}
-
- checkpoints.clear();
}
@Override
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/SubtaskState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/SubtaskState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/SubtaskState.java
index f326eaa..9beb233 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/SubtaskState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/SubtaskState.java
@@ -18,32 +18,28 @@
package org.apache.flink.runtime.checkpoint;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.StateObject;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import java.io.Serializable;
+import java.io.IOException;
import static org.apache.flink.util.Preconditions.checkNotNull;
/**
- * Simple bean to describe the state belonging to a parallel operator. It is part of the
+ * Container for the chained state of one parallel subtask of an operator/task. This is part of the
* {@link TaskState}.
- *
- * The state itself is kept in serialized form, since the checkpoint coordinator itself
- * is never looking at it anyways and only sends it back out in case of a recovery.
- * Furthermore, the state may involve user-defined classes that are not accessible without
- * the respective classloader.
*/
-public class SubtaskState implements Serializable {
+public class SubtaskState implements StateObject {
private static final long serialVersionUID = -2394696997971923995L;
private static final Logger LOG = LoggerFactory.getLogger(SubtaskState.class);
/** The state of the parallel operator */
- private final SerializedValue<StateHandle<?>> state;
+ private final ChainedStateHandle<StreamStateHandle> chainedStateHandle;
/**
* The state size. This is also part of the deserialized state handle.
@@ -52,27 +48,29 @@ public class SubtaskState implements Serializable {
*/
private final long stateSize;
- /** The duration of the acknowledged (ack timestamp - trigger timestamp). */
+ /** The duration of the checkpoint (ack timestamp - trigger timestamp). */
private final long duration;
public SubtaskState(
- SerializedValue<StateHandle<?>> state,
- long stateSize,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
long duration) {
- this.state = checkNotNull(state, "State");
- // Sanity check and don't fail checkpoint because of this.
- this.stateSize = stateSize >= 0 ? stateSize : 0;
-
+ this.chainedStateHandle = checkNotNull(chainedStateHandle, "State");
this.duration = duration;
+ try {
+ stateSize = chainedStateHandle.getStateSize();
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to get state size.", e);
+ }
}
// --------------------------------------------------------------------------------------------
- public SerializedValue<StateHandle<?>> getState() {
- return state;
+ public ChainedStateHandle<StreamStateHandle> getChainedStateHandle() {
+ return chainedStateHandle;
}
+ @Override
public long getStateSize() {
return stateSize;
}
@@ -81,8 +79,9 @@ public class SubtaskState implements Serializable {
return duration;
}
- public void discard(ClassLoader userClassLoader) throws Exception {
- state.deserializeValue(userClassLoader).discardState();
+ @Override
+ public void discardState() throws Exception {
+ chainedStateHandle.discardState();
}
// --------------------------------------------------------------------------------------------
@@ -94,7 +93,7 @@ public class SubtaskState implements Serializable {
}
else if (o instanceof SubtaskState) {
SubtaskState that = (SubtaskState) o;
- return this.state.equals(that.state) && stateSize == that.stateSize &&
+ return this.chainedStateHandle.equals(that.chainedStateHandle) && stateSize == that.stateSize &&
duration == that.duration;
}
else {
@@ -106,11 +105,18 @@ public class SubtaskState implements Serializable {
public int hashCode() {
return (int) (this.stateSize ^ this.stateSize >>> 32) +
31 * ((int) (this.duration ^ this.duration >>> 32) +
- 31 * state.hashCode());
+ 31 * chainedStateHandle.hashCode());
}
@Override
public String toString() {
- return String.format("SubtaskState(Size: %d, Duration: %d, State: %s)", stateSize, duration, state);
+ return String.format("SubtaskState(Size: %d, Duration: %d, State: %s)", stateSize, duration, chainedStateHandle);
}
+
+ @Override
+ public void close() throws IOException {
+ chainedStateHandle.close();
+ }
+
+
}
[03/27] flink git commit: [FLINK-4380] Introduce KeyGroupAssigner and
Max-Parallelism Parameter
Posted by al...@apache.org.
[FLINK-4380] Introduce KeyGroupAssigner and Max-Parallelism Parameter
This introduces a new KeySelector that assigns keys to key groups and
also adds the max parallelism parameter throughout all API levels.
This also adds tests for the newly introduced features.
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/ec975aab
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/ec975aab
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/ec975aab
Branch: refs/heads/master
Commit: ec975aaba79449bd93020f296b05ea509ea57bdc
Parents: 7cd9bb5
Author: Till Rohrmann <tr...@apache.org>
Authored: Thu Jul 28 15:08:24 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:04:31 2016 +0200
----------------------------------------------------------------------
docs/dev/api_concepts.md | 2 +
.../flink-statebackend-rocksdb/pom.xml | 7 +
.../src/test/resources/log4j-test.properties | 2 +-
.../src/test/resources/log4j.properties | 2 +-
.../flink/api/common/ExecutionConfig.java | 51 +-
.../api/common/state/KeyGroupAssigner.java | 47 ++
.../flink/api/common/state/StateDescriptor.java | 1 +
.../checkpoint/CheckpointCoordinator.java | 6 +-
.../InputChannelDeploymentDescriptor.java | 1 +
.../InputGateDeploymentDescriptor.java | 1 +
.../deployment/ResultPartitionLocation.java | 1 +
.../runtime/executiongraph/ExecutionGraph.java | 2 -
.../executiongraph/ExecutionJobVertex.java | 35 +-
.../runtime/executiongraph/ExecutionVertex.java | 4 +
.../flink/runtime/jobgraph/JobVertex.java | 30 +-
.../runtime/state/HashKeyGroupAssigner.java | 66 ++
.../flink/runtime/jobmanager/JobManager.scala | 1 -
.../checkpoint/CheckpointCoordinatorTest.java | 17 -
.../checkpoint/CheckpointStateRestoreTest.java | 5 +-
...ExecutionGraphCheckpointCoordinatorTest.java | 1 -
.../scheduler/SchedulerTestUtils.java | 2 +
.../testutils/DummyEnvironment.java.orig | 185 -----
.../streaming/api/datastream/KeyedStream.java | 12 +-
.../datastream/SingleOutputStreamOperator.java | 19 +
.../environment/StreamExecutionEnvironment.java | 42 +-
.../flink/streaming/api/graph/StreamConfig.java | 49 +-
.../flink/streaming/api/graph/StreamGraph.java | 33 +-
.../api/graph/StreamGraphGenerator.java | 43 +-
.../flink/streaming/api/graph/StreamNode.java | 25 +
.../api/graph/StreamingJobGraphGenerator.java | 31 +-
.../transformations/StreamTransformation.java | 24 +
.../ConfigurableStreamPartitioner.java | 39 ++
.../runtime/partitioner/HashPartitioner.java | 66 --
.../partitioner/KeyGroupStreamPartitioner.java | 82 +++
.../streaming/api/AggregationFunctionTest.java | 8 +-
.../flink/streaming/api/DataStreamTest.java | 8 +-
.../streaming/api/RestartStrategyTest.java | 4 +-
.../streaming/api/graph/SlotAllocationTest.java | 3 +-
.../api/graph/StreamGraphGeneratorTest.java | 220 ++++++
.../graph/StreamingJobGraphGeneratorTest.java | 216 +++++-
.../StreamingJobGraphGeneratorNodeHashTest.java | 64 +-
.../windowing/AllWindowTranslationTest.java | 4 +
.../windowing/WindowTranslationTest.java | 10 +
.../partitioner/HashPartitionerTest.java | 71 --
.../KeyGroupStreamPartitionerTest.java | 74 ++
.../partitioner/RescalePartitionerTest.java | 4 +
.../tasks/OneInputStreamTaskTestHarness.java | 2 +
.../streaming/runtime/tasks/StreamTaskTest.java | 176 +++--
.../runtime/tasks/StreamTaskTestHarness.java | 37 +-
.../src/test/resources/log4j-test.properties | 2 +-
.../flink/streaming/api/scala/DataStream.scala | 11 +
.../api/scala/StreamExecutionEnvironment.scala | 18 +
.../streaming/api/scala/DataStreamTest.scala | 2 +-
.../EventTimeAllWindowCheckpointingITCase.java | 2 +-
.../test/checkpointing/RescalingITCase.java | 683 +++++++++++++++++++
.../streaming/api/StreamingOperatorsITCase.java | 17 +
.../streaming/runtime/DataStreamPojoITCase.java | 2 +
.../test/streaming/runtime/TimestampITCase.java | 1 +
.../flink/test/web/WebFrontendITCase.java | 3 +-
59 files changed, 2074 insertions(+), 502 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/docs/dev/api_concepts.md
----------------------------------------------------------------------
diff --git a/docs/dev/api_concepts.md b/docs/dev/api_concepts.md
index 468085a..b0c034c 100644
--- a/docs/dev/api_concepts.md
+++ b/docs/dev/api_concepts.md
@@ -960,6 +960,8 @@ With the closure cleaner disabled, it might happen that an anonymous user functi
- `getParallelism()` / `setParallelism(int parallelism)` Set the default parallelism for the job.
+- `getMaxParallelism()` / `setMaxParallelism(int parallelism)` Set the default maximum parallelism for the job. This setting determines the maximum degree of parallelism and specifies the upper limit for dynamic scaling.
+
- `getNumberOfExecutionRetries()` / `setNumberOfExecutionRetries(int numberOfExecutionRetries)` Sets the number of times that failed tasks are re-executed. A value of zero effectively disables fault tolerance. A value of `-1` indicates that the system default value (as defined in the configuration) should be used.
- `getExecutionRetryDelay()` / `setExecutionRetryDelay(long executionRetryDelay)` Sets the delay in milliseconds that the system waits after a job has failed, before re-executing it. The delay starts after all tasks have been successfully been stopped on the TaskManagers, and once the delay is past, the tasks are re-started. This parameter is useful to delay re-execution in order to let certain time-out related failures surface fully (like broken connections that have not fully timed out), before attempting a re-execution and immediately failing again due to the same problem. This parameter only has an effect if the number of execution re-tries is one or more.
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-contrib/flink-statebackend-rocksdb/pom.xml
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/pom.xml b/flink-contrib/flink-statebackend-rocksdb/pom.xml
index efd673c..314b83f 100644
--- a/flink-contrib/flink-statebackend-rocksdb/pom.xml
+++ b/flink-contrib/flink-statebackend-rocksdb/pom.xml
@@ -69,6 +69,13 @@ under the License.
<dependency>
<groupId>org.apache.flink</groupId>
+ <artifactId>flink-core</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
<artifactId>flink-streaming-java_2.10</artifactId>
<version>${project.version}</version>
<type>test-jar</type>
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j-test.properties
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j-test.properties b/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j-test.properties
index 0b686e5..881dc06 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j-test.properties
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j-test.properties
@@ -24,4 +24,4 @@ log4j.appender.A1=org.apache.log4j.ConsoleAppender
# A1 uses PatternLayout.
log4j.appender.A1.layout=org.apache.log4j.PatternLayout
-log4j.appender.A1.layout.ConversionPattern=%-4r [%t] %-5p %c %x - %m%n
\ No newline at end of file
+log4j.appender.A1.layout.ConversionPattern=%-4r [%t] %-5p %c %x - %m%n
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j.properties
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j.properties b/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j.properties
index ed2bbcb..4a30c6f 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j.properties
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/resources/log4j.properties
@@ -24,4 +24,4 @@ log4j.rootLogger=OFF, console
log4j.appender.console=org.apache.log4j.ConsoleAppender
log4j.appender.console.target = System.err
log4j.appender.console.layout=org.apache.log4j.PatternLayout
-log4j.appender.console.layout.ConversionPattern=%d{HH:mm:ss,SSS} %-5p %-60c %x - %m%n
\ No newline at end of file
+log4j.appender.console.layout.ConversionPattern=%d{HH:mm:ss,SSS} %-5p %-60c %x - %m%n
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-core/src/main/java/org/apache/flink/api/common/ExecutionConfig.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/api/common/ExecutionConfig.java b/flink-core/src/main/java/org/apache/flink/api/common/ExecutionConfig.java
index 5b69794..81ee930 100644
--- a/flink-core/src/main/java/org/apache/flink/api/common/ExecutionConfig.java
+++ b/flink-core/src/main/java/org/apache/flink/api/common/ExecutionConfig.java
@@ -74,6 +74,13 @@ public class ExecutionConfig implements Serializable {
*/
public static final int PARALLELISM_DEFAULT = -1;
+ /**
+ * The flag value indicating an unknown or unset parallelism. This value is
+ * not a valid parallelism and indicates that the parallelism should remain
+ * unchanged.
+ */
+ public static final int PARALLELISM_UNKNOWN = -2;
+
private static final long DEFAULT_RESTART_DELAY = 10000L;
// --------------------------------------------------------------------------------------------
@@ -86,6 +93,13 @@ public class ExecutionConfig implements Serializable {
private int parallelism = PARALLELISM_DEFAULT;
/**
+ * The program wide maximum parallelism used for operators which haven't specified a maximum
+ * parallelism. The maximum parallelism specifies the upper limit for dynamic scaling and the
+ * number of key groups used for partitioned state.
+ */
+ private int maxParallelism = -1;
+
+ /**
* @deprecated Should no longer be used because it is subsumed by RestartStrategyConfiguration
*/
@Deprecated
@@ -219,12 +233,41 @@ public class ExecutionConfig implements Serializable {
* @param parallelism The parallelism to use
*/
public ExecutionConfig setParallelism(int parallelism) {
- Preconditions.checkArgument(parallelism > 0 || parallelism == PARALLELISM_DEFAULT,
- "The parallelism of an operator must be at least 1.");
+ if (parallelism != PARALLELISM_UNKNOWN) {
+ if (parallelism < 1 && parallelism != PARALLELISM_DEFAULT) {
+ throw new IllegalArgumentException(
+ "Parallelism must be at least one, or ExecutionConfig.PARALLELISM_DEFAULT (use system default).");
+ }
+ this.parallelism = parallelism;
+ }
+ return this;
+ }
- this.parallelism = parallelism;
+ /**
+ * Gets the maximum degree of parallelism defined for the program.
+ *
+ * The maximum degree of parallelism specifies the upper limit for dynamic scaling. It also
+ * defines the number of key groups used for partitioned state.
+ *
+ * @return Maximum degree of parallelism
+ */
+ @PublicEvolving
+ public int getMaxParallelism() {
+ return maxParallelism;
+ }
- return this;
+ /**
+ * Sets the maximum degree of parallelism defined for the program.
+ *
+ * The maximum degree of parallelism specifies the upper limit for dynamic scaling. It also
+ * defines the number of key groups used for partitioned state.
+ *
+ * @param maxParallelism Maximum degree of parallelism to be used for the program.
+ */
+ @PublicEvolving
+ public void setMaxParallelism(int maxParallelism) {
+ Preconditions.checkArgument(maxParallelism > 0, "The maximum parallelism must be greater than 0.");
+ this.maxParallelism = maxParallelism;
}
/**
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java b/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java
new file mode 100644
index 0000000..23463e1
--- /dev/null
+++ b/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.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.api.common.state;
+
+import org.apache.flink.annotation.Internal;
+
+import java.io.Serializable;
+
+/**
+ * Assigns a key to a key group index. A key group is the smallest unit of partitioned state
+ * which is assigned to an operator. An operator can be assigned multiple key groups.
+ *
+ * @param <K> Type of the key
+ */
+@Internal
+public interface KeyGroupAssigner<K> extends Serializable {
+ /**
+ * Calculates the key group index for the given key.
+ *
+ * @param key Key to be used
+ * @return Key group index for the given key
+ */
+ int getKeyGroupIndex(K key);
+
+ /**
+ * Setups the key group assigner with the maximum parallelism (= number of key groups).
+ *
+ * @param maxParallelism Maximum parallelism (= number of key groups)
+ */
+ void setup(int maxParallelism);
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-core/src/main/java/org/apache/flink/api/common/state/StateDescriptor.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/api/common/state/StateDescriptor.java b/flink-core/src/main/java/org/apache/flink/api/common/state/StateDescriptor.java
index d99f4de..483c954 100644
--- a/flink-core/src/main/java/org/apache/flink/api/common/state/StateDescriptor.java
+++ b/flink-core/src/main/java/org/apache/flink/api/common/state/StateDescriptor.java
@@ -44,6 +44,7 @@ import static java.util.Objects.requireNonNull;
* <p>Subclasses must correctly implement {@link #equals(Object)} and {@link #hashCode()}.
*
* @param <S> The type of the State objects created from this {@code StateDescriptor}.
+ * @param <T> The type of the value of the state object described by this state descriptor.
*/
@PublicEvolving
public abstract class StateDescriptor<S extends State, T> implements Serializable {
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
index 3619f48..24fca5f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
@@ -155,8 +155,6 @@ public class CheckpointCoordinator {
/** Helper for tracking checkpoint statistics */
private final CheckpointStatsTracker statsTracker;
- private final int numberKeyGroups;
-
// --------------------------------------------------------------------------------------------
public CheckpointCoordinator(
@@ -165,7 +163,6 @@ public class CheckpointCoordinator {
long checkpointTimeout,
long minPauseBetweenCheckpoints,
int maxConcurrentCheckpointAttempts,
- int numberKeyGroups,
ExecutionVertex[] tasksToTrigger,
ExecutionVertex[] tasksToWaitFor,
ExecutionVertex[] tasksToCommitTo,
@@ -202,7 +199,6 @@ public class CheckpointCoordinator {
this.recentPendingCheckpoints = new ArrayDeque<>(NUM_GHOST_CHECKPOINT_IDS);
this.userClassLoader = checkNotNull(userClassLoader);
this.statsTracker = checkNotNull(statsTracker);
- this.numberKeyGroups = numberKeyGroups;
this.timer = new Timer("Checkpoint Timer", true);
@@ -797,7 +793,7 @@ public class CheckpointCoordinator {
int counter = 0;
- List<Set<Integer>> keyGroupPartitions = createKeyGroupPartitions(numberKeyGroups, executionJobVertex.getParallelism());
+ List<Set<Integer>> keyGroupPartitions = createKeyGroupPartitions(executionJobVertex.getMaxParallelism(), executionJobVertex.getParallelism());
for (int i = 0; i < executionJobVertex.getParallelism(); i++) {
SubtaskState subtaskState = taskState.getState(i);
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputChannelDeploymentDescriptor.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputChannelDeploymentDescriptor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputChannelDeploymentDescriptor.java
index e00a480..f31febb 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputChannelDeploymentDescriptor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputChannelDeploymentDescriptor.java
@@ -48,6 +48,7 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
*/
public class InputChannelDeploymentDescriptor implements Serializable {
+ private static final long serialVersionUID = 373711381640454080L;
private static Logger LOG = LoggerFactory.getLogger(InputChannelDeploymentDescriptor.class);
/** The ID of the partition the input channel is going to consume. */
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputGateDeploymentDescriptor.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputGateDeploymentDescriptor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputGateDeploymentDescriptor.java
index ec4bd40..dde1ed7 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputGateDeploymentDescriptor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/InputGateDeploymentDescriptor.java
@@ -38,6 +38,7 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
*/
public class InputGateDeploymentDescriptor implements Serializable {
+ private static final long serialVersionUID = -7143441863165366704L;
/**
* The ID of the consumed intermediate result. Each input gate consumes partitions of the
* intermediate result specified by this ID. This ID also identifies the input gate at the
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionLocation.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionLocation.java b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionLocation.java
index ca63e6b..895bea0 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionLocation.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/ResultPartitionLocation.java
@@ -47,6 +47,7 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
*/
public class ResultPartitionLocation implements Serializable {
+ private static final long serialVersionUID = -6354238166937194463L;
/** The type of location for the result partition. */
private final LocationType locationType;
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java
index d7e40a3..92cab41 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionGraph.java
@@ -340,7 +340,6 @@ public class ExecutionGraph {
long checkpointTimeout,
long minPauseBetweenCheckpoints,
int maxConcurrentCheckpoints,
- int numberKeyGroups,
List<ExecutionJobVertex> verticesToTrigger,
List<ExecutionJobVertex> verticesToWaitFor,
List<ExecutionJobVertex> verticesToCommitTo,
@@ -373,7 +372,6 @@ public class ExecutionGraph {
checkpointTimeout,
minPauseBetweenCheckpoints,
maxConcurrentCheckpoints,
- numberKeyGroups,
tasksToTrigger,
tasksToWaitFor,
tasksToCommitTo,
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionJobVertex.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionJobVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionJobVertex.java
index 6272151..d3dc8fe 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionJobVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionJobVertex.java
@@ -40,7 +40,7 @@ import org.apache.flink.runtime.jobmanager.scheduler.NoResourceAvailableExceptio
import org.apache.flink.runtime.jobmanager.scheduler.Scheduler;
import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup;
import org.apache.flink.runtime.util.SerializableObject;
-
+import org.apache.flink.util.Preconditions;
import org.slf4j.Logger;
import scala.concurrent.duration.FiniteDuration;
@@ -69,6 +69,8 @@ public class ExecutionJobVertex {
private final List<IntermediateResult> inputs;
private final int parallelism;
+
+ private final int maxParallelism;
private final boolean[] finishedSubtasks;
@@ -81,16 +83,23 @@ public class ExecutionJobVertex {
private final InputSplit[] inputSplits;
private InputSplitAssigner splitAssigner;
+
+ public ExecutionJobVertex(
+ ExecutionGraph graph,
+ JobVertex jobVertex,
+ int defaultParallelism,
+ FiniteDuration timeout) throws JobException {
- public ExecutionJobVertex(ExecutionGraph graph, JobVertex jobVertex,
- int defaultParallelism, FiniteDuration timeout) throws JobException {
this(graph, jobVertex, defaultParallelism, timeout, System.currentTimeMillis());
}
- public ExecutionJobVertex(ExecutionGraph graph, JobVertex jobVertex,
- int defaultParallelism, FiniteDuration timeout, long createTimestamp)
- throws JobException
- {
+ public ExecutionJobVertex(
+ ExecutionGraph graph,
+ JobVertex jobVertex,
+ int defaultParallelism,
+ FiniteDuration timeout,
+ long createTimestamp) throws JobException {
+
if (graph == null || jobVertex == null) {
throw new NullPointerException();
}
@@ -102,6 +111,14 @@ public class ExecutionJobVertex {
int numTaskVertices = vertexParallelism > 0 ? vertexParallelism : defaultParallelism;
this.parallelism = numTaskVertices;
+
+ int maxParallelism = jobVertex.getMaxParallelism();
+
+ Preconditions.checkArgument(maxParallelism >= parallelism, "The maximum parallelism (" +
+ maxParallelism + ") must be greater or equal than the parallelism (" + parallelism +
+ ").");
+ this.maxParallelism = maxParallelism;
+
this.taskVertices = new ExecutionVertex[numTaskVertices];
this.inputs = new ArrayList<IntermediateResult>(jobVertex.getInputs().size());
@@ -177,6 +194,10 @@ public class ExecutionJobVertex {
return parallelism;
}
+ public int getMaxParallelism() {
+ return maxParallelism;
+ }
+
public JobID getJobId() {
return graph.getJobID();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
index 2495316..b15f851 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
@@ -187,6 +187,10 @@ public class ExecutionVertex {
return this.subTaskIndex;
}
+ public int getMaxParallelism() {
+ return this.jobVertex.getMaxParallelism();
+ }
+
public int getNumberOfInputs() {
return this.inputEdges.length;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
index 379a42a..4786388 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
@@ -18,9 +18,6 @@
package org.apache.flink.runtime.jobgraph;
-import java.util.ArrayList;
-import java.util.List;
-
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.io.InputSplitSource;
@@ -31,6 +28,9 @@ import org.apache.flink.runtime.jobmanager.scheduler.CoLocationGroup;
import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup;
import org.apache.flink.util.Preconditions;
+import java.util.ArrayList;
+import java.util.List;
+
/**
* The base class for job vertexes.
*/
@@ -57,6 +57,9 @@ public class JobVertex implements java.io.Serializable {
/** Number of subtasks to split this task into at runtime.*/
private int parallelism = ExecutionConfig.PARALLELISM_DEFAULT;
+ /** Maximum number of subtasks to split this taks into a runtime. */
+ private int maxParallelism = Short.MAX_VALUE;
+
/** Custom configuration passed to the assigned task at runtime. */
private Configuration configuration;
@@ -234,6 +237,27 @@ public class JobVertex implements java.io.Serializable {
this.parallelism = parallelism;
}
+ /**
+ * Gets the maximum parallelism for the task.
+ *
+ * @return The maximum parallelism for the task.
+ */
+ public int getMaxParallelism() {
+ return maxParallelism;
+ }
+
+ /**
+ * Sets the maximum parallelism for the task.
+ *
+ * @param maxParallelism The maximum parallelism to be set.
+ */
+ public void setMaxParallelism(int maxParallelism) {
+ org.apache.flink.util.Preconditions.checkArgument(
+ maxParallelism > 0 && maxParallelism <= Short.MAX_VALUE, "The max parallelism must be at least 1.");
+
+ this.maxParallelism = maxParallelism;
+ }
+
public InputSplitSource<?> getInputSplitSource() {
return inputSplitSource;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
new file mode 100644
index 0000000..280746d
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
@@ -0,0 +1,66 @@
+/*
+ * 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.state.KeyGroupAssigner;
+import org.apache.flink.util.MathUtils;
+import org.apache.flink.util.Preconditions;
+
+/**
+ * Hash based key group assigner. The assigner assigns each key to a key group using the hash value
+ * of the key.
+ *
+ * @param <K> Type of the key
+ */
+public class HashKeyGroupAssigner<K> implements KeyGroupAssigner<K> {
+ private static final long serialVersionUID = -6319826921798945448L;
+
+ private static final int UNDEFINED_NUMBER_KEY_GROUPS = Integer.MIN_VALUE;
+
+ private int numberKeyGroups;
+
+ public HashKeyGroupAssigner() {
+ this(UNDEFINED_NUMBER_KEY_GROUPS);
+ }
+
+ public HashKeyGroupAssigner(int numberKeyGroups) {
+ Preconditions.checkArgument(numberKeyGroups > 0 || numberKeyGroups == UNDEFINED_NUMBER_KEY_GROUPS,
+ "The number of key groups has to be greater than 0 or undefined. Use " +
+ "setMaxParallelism() to specify the number of key groups.");
+ this.numberKeyGroups = numberKeyGroups;
+ }
+
+ public int getNumberKeyGroups() {
+ return numberKeyGroups;
+ }
+
+ @Override
+ public int getKeyGroupIndex(K key) {
+ return MathUtils.murmurHash(key.hashCode()) % numberKeyGroups;
+ }
+
+ @Override
+ public void setup(int numberKeyGroups) {
+ Preconditions.checkArgument(numberKeyGroups > 0, "The number of key groups has to be " +
+ "greater than 0. Use setMaxParallelism() to specify the number of key " +
+ "groups.");
+
+ this.numberKeyGroups = numberKeyGroups;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala b/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
index b706a1a..356f1a9 100644
--- a/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
+++ b/flink-runtime/src/main/scala/org/apache/flink/runtime/jobmanager/JobManager.scala
@@ -1271,7 +1271,6 @@ class JobManager(
snapshotSettings.getCheckpointTimeout,
snapshotSettings.getMinPauseBetweenCheckpoints,
snapshotSettings.getMaxConcurrentCheckpoints,
- parallelism,
triggerVertices,
ackVertices,
confirmVertices,
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
index 09c53d6..50330fa 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
@@ -86,7 +86,6 @@ public class CheckpointCoordinatorTest {
600000,
600000,
0, Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { triggerVertex1, triggerVertex2 },
new ExecutionVertex[] { ackVertex1, ackVertex2 },
new ExecutionVertex[] {},
@@ -140,7 +139,6 @@ public class CheckpointCoordinatorTest {
600000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { triggerVertex1, triggerVertex2 },
new ExecutionVertex[] { ackVertex1, ackVertex2 },
new ExecutionVertex[] {},
@@ -192,7 +190,6 @@ public class CheckpointCoordinatorTest {
600000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { triggerVertex1, triggerVertex2 },
new ExecutionVertex[] { ackVertex1, ackVertex2 },
new ExecutionVertex[] {},
@@ -245,7 +242,6 @@ public class CheckpointCoordinatorTest {
600000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
@@ -371,7 +367,6 @@ public class CheckpointCoordinatorTest {
600000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
@@ -493,7 +488,6 @@ public class CheckpointCoordinatorTest {
600000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
@@ -645,7 +639,6 @@ public class CheckpointCoordinatorTest {
600000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { triggerVertex1, triggerVertex2 },
new ExecutionVertex[] { ackVertex1, ackVertex2, ackVertex3 },
new ExecutionVertex[] { commitVertex },
@@ -782,7 +775,6 @@ public class CheckpointCoordinatorTest {
600000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { triggerVertex1, triggerVertex2 },
new ExecutionVertex[] { ackVertex1, ackVertex2, ackVertex3 },
new ExecutionVertex[] { commitVertex },
@@ -905,7 +897,6 @@ public class CheckpointCoordinatorTest {
200,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { triggerVertex },
new ExecutionVertex[] { ackVertex1, ackVertex2 },
new ExecutionVertex[] { commitVertex },
@@ -975,7 +966,6 @@ public class CheckpointCoordinatorTest {
200000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { triggerVertex },
new ExecutionVertex[] { ackVertex1, ackVertex2 },
new ExecutionVertex[] { commitVertex },
@@ -1056,7 +1046,6 @@ public class CheckpointCoordinatorTest {
200000, // timeout is very long (200 s)
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { triggerVertex },
new ExecutionVertex[] { ackVertex },
new ExecutionVertex[] { commitVertex },
@@ -1149,7 +1138,6 @@ public class CheckpointCoordinatorTest {
200000, // timeout is very long (200 s)
500, // 500ms delay between checkpoints
10,
- 42,
new ExecutionVertex[] { vertex1 },
new ExecutionVertex[] { vertex1 },
new ExecutionVertex[] { vertex1 },
@@ -1235,7 +1223,6 @@ public class CheckpointCoordinatorTest {
600000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
@@ -1374,7 +1361,6 @@ public class CheckpointCoordinatorTest {
600000,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
new ExecutionVertex[] { vertex1, vertex2 },
@@ -1462,7 +1448,6 @@ public class CheckpointCoordinatorTest {
200000, // timeout is very long (200 s)
0L, // no extra delay
maxConcurrentAttempts,
- 42,
new ExecutionVertex[] { triggerVertex },
new ExecutionVertex[] { ackVertex },
new ExecutionVertex[] { commitVertex }, cl, new StandaloneCheckpointIDCounter
@@ -1534,7 +1519,6 @@ public class CheckpointCoordinatorTest {
200000, // timeout is very long (200 s)
0L, // no extra delay
maxConcurrentAttempts, // max two concurrent checkpoints
- 42,
new ExecutionVertex[] { triggerVertex },
new ExecutionVertex[] { ackVertex },
new ExecutionVertex[] { commitVertex }, cl, new StandaloneCheckpointIDCounter
@@ -1615,7 +1599,6 @@ public class CheckpointCoordinatorTest {
200000, // timeout is very long (200 s)
0L, // no extra delay
2, // max two concurrent checkpoints
- 42,
new ExecutionVertex[] { triggerVertex },
new ExecutionVertex[] { ackVertex },
new ExecutionVertex[] { commitVertex }, cl, new StandaloneCheckpointIDCounter(),
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java
index 061059a..1816fc9 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointStateRestoreTest.java
@@ -90,7 +90,6 @@ public class CheckpointStateRestoreTest {
200000L,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { stateful1, stateful2, stateful3, stateless1, stateless2 },
new ExecutionVertex[] { stateful1, stateful2, stateful3, stateless1, stateless2 },
new ExecutionVertex[0],
@@ -170,7 +169,6 @@ public class CheckpointStateRestoreTest {
200000L,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { stateful1, stateful2, stateful3, stateless1, stateless2 },
new ExecutionVertex[] { stateful1, stateful2, stateful3, stateless1, stateless2 },
new ExecutionVertex[0],
@@ -221,7 +219,6 @@ public class CheckpointStateRestoreTest {
200000L,
0,
Integer.MAX_VALUE,
- 42,
new ExecutionVertex[] { mock(ExecutionVertex.class) },
new ExecutionVertex[] { mock(ExecutionVertex.class) },
new ExecutionVertex[0], cl,
@@ -263,12 +260,14 @@ public class CheckpointStateRestoreTest {
when(mock.getParallelSubtaskIndex()).thenReturn(subtask);
when(mock.getCurrentExecutionAttempt()).thenReturn(execution);
when(mock.getTotalNumberOfParallelSubtasks()).thenReturn(parallelism);
+ when(mock.getMaxParallelism()).thenReturn(parallelism);
return mock;
}
private ExecutionJobVertex mockExecutionJobVertex(JobVertexID id, ExecutionVertex[] vertices) {
ExecutionJobVertex vertex = mock(ExecutionJobVertex.class);
when(vertex.getParallelism()).thenReturn(vertices.length);
+ when(vertex.getMaxParallelism()).thenReturn(vertices.length);
when(vertex.getJobVertexId()).thenReturn(id);
when(vertex.getTaskVertices()).thenReturn(vertices);
return vertex;
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ExecutionGraphCheckpointCoordinatorTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ExecutionGraphCheckpointCoordinatorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ExecutionGraphCheckpointCoordinatorTest.java
index 49a9449..1de7098 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ExecutionGraphCheckpointCoordinatorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/ExecutionGraphCheckpointCoordinatorTest.java
@@ -112,7 +112,6 @@ public class ExecutionGraphCheckpointCoordinatorTest {
100,
100,
1,
- 42,
Collections.<ExecutionJobVertex>emptyList(),
Collections.<ExecutionJobVertex>emptyList(),
Collections.<ExecutionJobVertex>emptyList(),
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java
index ed6d1ee..983d6e6 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/jobmanager/scheduler/SchedulerTestUtils.java
@@ -104,6 +104,7 @@ public class SchedulerTestUtils {
when(vertex.getJobvertexId()).thenReturn(jid);
when(vertex.getParallelSubtaskIndex()).thenReturn(taskIndex);
when(vertex.getTotalNumberOfParallelSubtasks()).thenReturn(numTasks);
+ when(vertex.getMaxParallelism()).thenReturn(numTasks);
when(vertex.toString()).thenReturn("TEST-VERTEX");
when(vertex.getSimpleName()).thenReturn("TEST-VERTEX");
@@ -121,6 +122,7 @@ public class SchedulerTestUtils {
when(vertex.getJobvertexId()).thenReturn(jid);
when(vertex.getParallelSubtaskIndex()).thenReturn(taskIndex);
when(vertex.getTotalNumberOfParallelSubtasks()).thenReturn(numTasks);
+ when(vertex.getMaxParallelism()).thenReturn(numTasks);
when(vertex.toString()).thenReturn("TEST-VERTEX");
Execution execution = mock(Execution.class);
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java.orig
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java.orig b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java.orig
deleted file mode 100644
index 393ee4c..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java.orig
+++ /dev/null
@@ -1,185 +0,0 @@
-/*
- * 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.operators.testutils;
-
-import org.apache.flink.api.common.ExecutionConfig;
-import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.TaskInfo;
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.core.fs.Path;
-import org.apache.flink.metrics.groups.TaskMetricGroup;
-import org.apache.flink.runtime.accumulators.AccumulatorRegistry;
-import org.apache.flink.runtime.broadcast.BroadcastVariableManager;
-import org.apache.flink.runtime.execution.Environment;
-import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
-import org.apache.flink.runtime.io.disk.iomanager.IOManager;
-import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
-import org.apache.flink.runtime.io.network.partition.consumer.InputGate;
-import org.apache.flink.runtime.jobgraph.JobVertexID;
-import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider;
-import org.apache.flink.runtime.memory.MemoryManager;
-import org.apache.flink.runtime.query.KvStateRegistry;
-import org.apache.flink.runtime.query.TaskKvStateRegistry;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
-
-import java.util.Collections;
-import java.util.Map;
-import java.util.concurrent.Future;
-
-public class DummyEnvironment implements Environment {
-
- private final JobID jobId = new JobID();
- private final JobVertexID jobVertexId = new JobVertexID();
- private final ExecutionAttemptID executionId = new ExecutionAttemptID();
- private final ExecutionConfig executionConfig = new ExecutionConfig();
-<<<<<<< 9a73dbc71b83080b7deccc62b8b6ffa9f102e847
- private final TaskInfo taskInfo;
-=======
- private final KvStateRegistry kvStateRegistry = new KvStateRegistry();
- private final TaskKvStateRegistry taskKvStateRegistry;
->>>>>>> [FLINK-3779] [runtime] Add KvStateRegistry for queryable KvState
-
- public DummyEnvironment(String taskName, int numSubTasks, int subTaskIndex) {
- this.taskInfo = new TaskInfo(taskName, subTaskIndex, numSubTasks, 0);
-
- this.taskKvStateRegistry = kvStateRegistry.createTaskRegistry(jobId, jobVertexId);
- }
-
- public KvStateRegistry getKvStateRegistry() {
- return kvStateRegistry;
- }
-
- @Override
- public ExecutionConfig getExecutionConfig() {
- return executionConfig;
- }
-
- @Override
- public JobID getJobID() {
- return jobId;
- }
-
- @Override
- public JobVertexID getJobVertexId() {
- return jobVertexId;
- }
-
- @Override
- public ExecutionAttemptID getExecutionId() {
- return executionId;
- }
-
- @Override
- public Configuration getTaskConfiguration() {
- return new Configuration();
- }
-
- @Override
- public TaskManagerRuntimeInfo getTaskManagerInfo() {
- return null;
- }
-
- @Override
- public TaskMetricGroup getMetricGroup() {
- return new UnregisteredTaskMetricsGroup();
- }
-
- @Override
- public Configuration getJobConfiguration() {
- return new Configuration();
- }
-
- @Override
- public TaskInfo getTaskInfo() {
- return taskInfo;
- }
-
- @Override
- public InputSplitProvider getInputSplitProvider() {
- return null;
- }
-
- @Override
- public IOManager getIOManager() {
- return null;
- }
-
- @Override
- public MemoryManager getMemoryManager() {
- return null;
- }
-
- @Override
- public ClassLoader getUserClassLoader() {
- return getClass().getClassLoader();
- }
-
- @Override
- public Map<String, Future<Path>> getDistributedCacheEntries() {
- return Collections.emptyMap();
- }
-
- @Override
- public BroadcastVariableManager getBroadcastVariableManager() {
- return null;
- }
-
- @Override
- public AccumulatorRegistry getAccumulatorRegistry() {
- return null;
- }
-
- @Override
-<<<<<<< 9a73dbc71b83080b7deccc62b8b6ffa9f102e847
- public void acknowledgeCheckpoint(long checkpointId) {}
-=======
- public TaskKvStateRegistry getTaskKvStateRegistry() {
- return taskKvStateRegistry;
- }
-
- @Override
- public void acknowledgeCheckpoint(long checkpointId) {
- }
->>>>>>> [FLINK-3779] [runtime] Add KvStateRegistry for queryable KvState
-
- @Override
- public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) {}
-
- @Override
- public ResultPartitionWriter getWriter(int index) {
- return null;
- }
-
- @Override
- public ResultPartitionWriter[] getAllWriters() {
- return null;
- }
-
- @Override
- public InputGate getInputGate(int index) {
- return null;
- }
-
- @Override
- public InputGate[] getAllInputGates() {
- return null;
- }
-
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java
index 6998890..1fb34b8 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java
@@ -30,6 +30,7 @@ import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.Utils;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.typeutils.TypeExtractor;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.functions.aggregation.AggregationFunction;
import org.apache.flink.streaming.api.functions.aggregation.ComparableAggregator;
@@ -55,7 +56,7 @@ import org.apache.flink.streaming.api.windowing.triggers.PurgingTrigger;
import org.apache.flink.streaming.api.windowing.windows.GlobalWindow;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.api.windowing.windows.Window;
-import org.apache.flink.streaming.runtime.partitioner.HashPartitioner;
+import org.apache.flink.streaming.runtime.partitioner.KeyGroupStreamPartitioner;
import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner;
import java.util.UUID;
@@ -105,8 +106,13 @@ public class KeyedStream<T, KEY> extends DataStream<T> {
* Function for determining state partitions
*/
public KeyedStream(DataStream<T> dataStream, KeySelector<T, KEY> keySelector, TypeInformation<KEY> keyType) {
- super(dataStream.getExecutionEnvironment(), new PartitionTransformation<>(
- dataStream.getTransformation(), new HashPartitioner<>(keySelector)));
+ super(
+ dataStream.getExecutionEnvironment(),
+ new PartitionTransformation<>(
+ dataStream.getTransformation(),
+ new KeyGroupStreamPartitioner<>(
+ keySelector,
+ new HashKeyGroupAssigner<KEY>())));
this.keySelector = keySelector;
this.keyType = keyType;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/SingleOutputStreamOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/SingleOutputStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/SingleOutputStreamOperator.java
index 02ea219..614f19b 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/SingleOutputStreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/SingleOutputStreamOperator.java
@@ -28,6 +28,7 @@ import org.apache.flink.streaming.api.operators.ChainingStrategy;
import org.apache.flink.streaming.api.transformations.PartitionTransformation;
import org.apache.flink.streaming.api.transformations.StreamTransformation;
import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner;
+import org.apache.flink.util.Preconditions;
import static java.util.Objects.requireNonNull;
@@ -106,6 +107,24 @@ public class SingleOutputStreamOperator<T> extends DataStream<T> {
}
/**
+ * Sets the maximum parallelism of this operator.
+ *
+ * The maximum parallelism specifies the upper bound for dynamic scaling. It also defines the
+ * number of key groups used for partitioned state.
+ *
+ * @param maxParallelism Maximum parallelism
+ * @return The operator with set maximum parallelism
+ */
+ @PublicEvolving
+ public SingleOutputStreamOperator<T> setMaxParallelism(int maxParallelism) {
+ Preconditions.checkArgument(maxParallelism > 0, "The maximum parallelism must be greater than 0.");
+
+ transformation.setMaxParallelism(maxParallelism);
+
+ return this;
+ }
+
+ /**
* Sets the parallelism of this operator to one.
* And mark this operator cannot set a non-1 degree of parallelism.
*
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/environment/StreamExecutionEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/environment/StreamExecutionEnvironment.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/environment/StreamExecutionEnvironment.java
index ead9564..78aab97 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/environment/StreamExecutionEnvironment.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/environment/StreamExecutionEnvironment.java
@@ -18,16 +18,16 @@
package org.apache.flink.streaming.api.environment;
import com.esotericsoftware.kryo.Serializer;
-
-import org.apache.flink.annotation.PublicEvolving;
import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.Public;
+import org.apache.flink.annotation.PublicEvolving;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.InvalidProgramException;
import org.apache.flink.api.common.JobExecutionResult;
import org.apache.flink.api.common.functions.InvalidTypesException;
import org.apache.flink.api.common.functions.StoppableFunction;
import org.apache.flink.api.common.io.FileInputFormat;
+import org.apache.flink.api.common.io.FilePathFilter;
import org.apache.flink.api.common.io.InputFormat;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.ValueState;
@@ -46,22 +46,22 @@ import org.apache.flink.client.program.OptimizerPlanEnvironment;
import org.apache.flink.client.program.PreviewPlanEnvironment;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.fs.Path;
+import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
-import org.apache.flink.streaming.api.functions.source.FileMonitoringFunction;
-import org.apache.flink.api.common.io.FilePathFilter;
-import org.apache.flink.streaming.api.functions.source.FileReadFunction;
import org.apache.flink.streaming.api.functions.source.ContinuousFileMonitoringFunction;
import org.apache.flink.streaming.api.functions.source.ContinuousFileReaderOperator;
-import org.apache.flink.streaming.api.functions.source.InputFormatSourceFunction;
+import org.apache.flink.streaming.api.functions.source.FileMonitoringFunction;
+import org.apache.flink.streaming.api.functions.source.FileProcessingMode;
+import org.apache.flink.streaming.api.functions.source.FileReadFunction;
import org.apache.flink.streaming.api.functions.source.FromElementsFunction;
import org.apache.flink.streaming.api.functions.source.FromIteratorFunction;
import org.apache.flink.streaming.api.functions.source.FromSplittableIteratorFunction;
+import org.apache.flink.streaming.api.functions.source.InputFormatSourceFunction;
import org.apache.flink.streaming.api.functions.source.ParallelSourceFunction;
-import org.apache.flink.streaming.api.functions.source.FileProcessingMode;
import org.apache.flink.streaming.api.functions.source.SocketTextStreamFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.functions.source.StatefulSequenceSource;
@@ -69,7 +69,6 @@ import org.apache.flink.streaming.api.graph.StreamGraph;
import org.apache.flink.streaming.api.graph.StreamGraphGenerator;
import org.apache.flink.streaming.api.operators.StoppableStreamSource;
import org.apache.flink.streaming.api.operators.StreamSource;
-import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.streaming.api.transformations.StreamTransformation;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.SplittableIterator;
@@ -168,6 +167,21 @@ public abstract class StreamExecutionEnvironment {
}
/**
+ * Sets the maximum degree of parallelism defined for the program.
+ *
+ * The maximum degree of parallelism specifies the upper limit for dynamic scaling. It also
+ * defines the number of key groups used for partitioned state.
+ *
+ * @param maxParallelism Maximum degree of parallelism to be used for the program., with 0 < maxParallelism <= 2^15
+ */
+ public StreamExecutionEnvironment setMaxParallelism(int maxParallelism) {
+ Preconditions.checkArgument(maxParallelism > 0 && maxParallelism <= (1 << 15),
+ "maxParallelism is out of bounds 0 < maxParallelism <= 2^15. Found: " + maxParallelism);
+ config.setMaxParallelism(maxParallelism);
+ return this;
+ }
+
+ /**
* Gets the parallelism with which operation are executed by default.
* Operations can individually override this value to use a specific
* parallelism.
@@ -180,6 +194,18 @@ public abstract class StreamExecutionEnvironment {
}
/**
+ * Gets the maximum degree of parallelism defined for the program.
+ *
+ * The maximum degree of parallelism specifies the upper limit for dynamic scaling. It also
+ * defines the number of key groups used for partitioned state.
+ *
+ * @return Maximum degree of parallelism
+ */
+ public int getMaxParallelism() {
+ return config.getMaxParallelism();
+ }
+
+ /**
* Sets the maximum time frequency (milliseconds) for the flushing of the
* output buffers. By default the output buffers flush frequently to provide
* low latency and to aid smooth developer experience. Setting the parameter
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java
index 783b3e2..1a92ba4 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java
@@ -26,6 +26,7 @@ import java.util.List;
import java.util.Map;
import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
@@ -74,6 +75,10 @@ public class StreamConfig implements Serializable {
private static final String STATE_BACKEND = "statebackend";
private static final String STATE_PARTITIONER = "statePartitioner";
+
+ /** key for the {@link KeyGroupAssigner} for key to key group index mappings */
+ private static final String KEY_GROUP_ASSIGNER = "keyGroupAssigner";
+
private static final String STATE_KEY_SERIALIZER = "statekeyser";
private static final String TIME_CHARACTERISTIC = "timechar";
@@ -402,10 +407,12 @@ public class StreamConfig implements Serializable {
// ------------------------------------------------------------------------
public void setStateBackend(AbstractStateBackend backend) {
- try {
- InstantiationUtil.writeObjectToConfig(backend, this.config, STATE_BACKEND);
- } catch (Exception e) {
- throw new StreamTaskException("Could not serialize stateHandle provider.", e);
+ if (backend != null) {
+ try {
+ InstantiationUtil.writeObjectToConfig(backend, this.config, STATE_BACKEND);
+ } catch (Exception e) {
+ throw new StreamTaskException("Could not serialize stateHandle provider.", e);
+ }
}
}
@@ -416,6 +423,10 @@ public class StreamConfig implements Serializable {
throw new StreamTaskException("Could not instantiate statehandle provider.", e);
}
}
+
+ public byte[] getSerializedStateBackend() {
+ return this.config.getBytes(STATE_BACKEND, null);
+ }
public void setStatePartitioner(int input, KeySelector<?, ?> partitioner) {
try {
@@ -432,6 +443,34 @@ public class StreamConfig implements Serializable {
throw new StreamTaskException("Could not instantiate state partitioner.", e);
}
}
+
+ /**
+ * Sets the {@link KeyGroupAssigner} to be used for the current {@link StreamOperator}.
+ *
+ * @param keyGroupAssigner Key group assigner to be used
+ */
+ public void setKeyGroupAssigner(KeyGroupAssigner<?> keyGroupAssigner) {
+ try {
+ InstantiationUtil.writeObjectToConfig(keyGroupAssigner, this.config, KEY_GROUP_ASSIGNER);
+ } catch (Exception e) {
+ throw new StreamTaskException("Could not serialize virtual state partitioner.", e);
+ }
+ }
+
+ /**
+ * Gets the {@link KeyGroupAssigner} for the {@link StreamOperator}.
+ *
+ * @param classLoader Classloader to be used for the deserialization
+ * @param <K> Type of the keys to be assigned to key groups
+ * @return Key group assigner
+ */
+ public <K> KeyGroupAssigner<K> getKeyGroupAssigner(ClassLoader classLoader) {
+ try {
+ return InstantiationUtil.readObjectFromConfig(this.config, KEY_GROUP_ASSIGNER, classLoader);
+ } catch (Exception e) {
+ throw new StreamTaskException("Could not instantiate virtual state partitioner.", e);
+ }
+ }
public void setStateKeySerializer(TypeSerializer<?> serializer) {
try {
@@ -448,6 +487,8 @@ public class StreamConfig implements Serializable {
throw new StreamTaskException("Could not instantiate state key serializer from task config.", e);
}
}
+
+
// ------------------------------------------------------------------------
// Miscellansous
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraph.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraph.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraph.java
index 4be2874..c946e98 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraph.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraph.java
@@ -50,6 +50,7 @@ import org.apache.flink.streaming.api.operators.StreamOperator;
import org.apache.flink.streaming.api.operators.StreamSource;
import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.streaming.runtime.partitioner.ConfigurableStreamPartitioner;
import org.apache.flink.streaming.runtime.partitioner.ForwardPartitioner;
import org.apache.flink.streaming.runtime.partitioner.RebalancePartitioner;
import org.apache.flink.streaming.runtime.partitioner.StreamPartitioner;
@@ -357,6 +358,18 @@ public class StreamGraph extends StreamingPlan {
if (partitioner == null) {
partitioner = virtuaPartitionNodes.get(virtualId).f1;
}
+
+ if (partitioner instanceof ConfigurableStreamPartitioner) {
+ StreamNode downstreamNode = getStreamNode(downStreamVertexID);
+
+ ConfigurableStreamPartitioner configurableStreamPartitioner = (ConfigurableStreamPartitioner) partitioner;
+
+ // Configure the partitioner with the max parallelism. This is necessary if the
+ // partitioner has been created before the maximum parallelism has been set. The
+ // maximum parallelism is necessary for the key group mapping.
+ configurableStreamPartitioner.configure(downstreamNode.getMaxParallelism());
+ }
+
addEdgeInternal(upStreamVertexID, downStreamVertexID, typeNumber, partitioner, outputNames);
} else {
StreamNode upstreamNode = getStreamNode(upStreamVertexID);
@@ -407,6 +420,12 @@ public class StreamGraph extends StreamingPlan {
}
}
+ public void setMaxParallelism(int vertexID, int maxParallelism) {
+ if (getStreamNode(vertexID) != null) {
+ getStreamNode(vertexID).setMaxParallelism(maxParallelism);
+ }
+ }
+
public void setOneInputStateKey(Integer vertexID, KeySelector<?, ?> keySelector, TypeSerializer<?> keySerializer) {
StreamNode node = getStreamNode(vertexID);
node.setStatePartitioner1(keySelector);
@@ -514,7 +533,13 @@ public class StreamGraph extends StreamingPlan {
return vertexIDtoLoopTimeout.get(vertexID);
}
- public Tuple2<StreamNode, StreamNode> createIterationSourceAndSink(int loopId, int sourceId, int sinkId, long timeout, int parallelism) {
+ public Tuple2<StreamNode, StreamNode> createIterationSourceAndSink(
+ int loopId,
+ int sourceId,
+ int sinkId,
+ long timeout,
+ int parallelism,
+ int maxParallelism) {
StreamNode source = this.addNode(sourceId,
null,
StreamIterationHead.class,
@@ -522,6 +547,7 @@ public class StreamGraph extends StreamingPlan {
"IterationSource-" + loopId);
sources.add(source.getId());
setParallelism(source.getId(), parallelism);
+ setMaxParallelism(source.getId(), maxParallelism);
StreamNode sink = this.addNode(sinkId,
null,
@@ -530,6 +556,7 @@ public class StreamGraph extends StreamingPlan {
"IterationSink-" + loopId);
sinks.add(sink.getId());
setParallelism(sink.getId(), parallelism);
+ setMaxParallelism(sink.getId(), parallelism);
iterationSourceSinkPairs.add(new Tuple2<>(source, sink));
@@ -604,8 +631,4 @@ public class StreamGraph extends StreamingPlan {
}
}
}
-
- public static enum ResourceStrategy {
- DEFAULT, ISOLATE, NEWGROUP
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
index de80e25..81c5c48 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
@@ -144,6 +144,31 @@ public class StreamGraphGenerator {
LOG.debug("Transforming " + transform);
+ if (transform.getMaxParallelism() <= 0) {
+ // if the max parallelism hasn't been set, then first use the job wide max parallelism
+ // from theExecutionConfig. If this value has not been specified either, then use the
+ // parallelism of the operator.
+ int maxParallelism = env.getConfig().getMaxParallelism();
+
+ if (maxParallelism <= 0) {
+ maxParallelism = transform.getParallelism();
+
+ /**
+ * TODO: Remove once the parallelism settings works properly in Flink (FLINK-3885)
+ * Currently, the parallelism will be set to 1 on the JobManager iff it encounters
+ * a negative parallelism value. We need to know this for the
+ * KeyGroupStreamPartitioner on the client-side. Thus, we already set the value to
+ * 1 here.
+ */
+ if (maxParallelism <= 0) {
+ transform.setParallelism(1);
+ maxParallelism = 1;
+ }
+ }
+
+ transform.setMaxParallelism(maxParallelism);
+ }
+
// call at least once to trigger exceptions about MissingTypeInfo
transform.getOutputType();
@@ -309,11 +334,12 @@ public class StreamGraphGenerator {
// create the fake iteration source/sink pair
Tuple2<StreamNode, StreamNode> itSourceAndSink = streamGraph.createIterationSourceAndSink(
- iterate.getId(),
- getNewIterationNodeId(),
- getNewIterationNodeId(),
- iterate.getWaitTime(),
- iterate.getParallelism());
+ iterate.getId(),
+ getNewIterationNodeId(),
+ getNewIterationNodeId(),
+ iterate.getWaitTime(),
+ iterate.getParallelism(),
+ iterate.getMaxParallelism());
StreamNode itSource = itSourceAndSink.f0;
StreamNode itSink = itSourceAndSink.f1;
@@ -377,7 +403,8 @@ public class StreamGraphGenerator {
getNewIterationNodeId(),
getNewIterationNodeId(),
coIterate.getWaitTime(),
- coIterate.getParallelism());
+ coIterate.getParallelism(),
+ coIterate.getMaxParallelism());
StreamNode itSource = itSourceAndSink.f0;
StreamNode itSink = itSourceAndSink.f1;
@@ -430,6 +457,7 @@ public class StreamGraphGenerator {
streamGraph.setInputFormat(source.getId(), fs.getFormat());
}
streamGraph.setParallelism(source.getId(), source.getParallelism());
+ streamGraph.setMaxParallelism(source.getId(), source.getMaxParallelism());
return Collections.singleton(source.getId());
}
@@ -450,6 +478,7 @@ public class StreamGraphGenerator {
"Sink: " + sink.getName());
streamGraph.setParallelism(sink.getId(), sink.getParallelism());
+ streamGraph.setMaxParallelism(sink.getId(), sink.getMaxParallelism());
for (Integer inputId: inputIds) {
streamGraph.addEdge(inputId,
@@ -498,6 +527,7 @@ public class StreamGraphGenerator {
}
streamGraph.setParallelism(transform.getId(), transform.getParallelism());
+ streamGraph.setMaxParallelism(transform.getId(), transform.getMaxParallelism());
for (Integer inputId: inputIds) {
streamGraph.addEdge(inputId, transform.getId(), 0);
@@ -545,6 +575,7 @@ public class StreamGraphGenerator {
streamGraph.setParallelism(transform.getId(), transform.getParallelism());
+ streamGraph.setMaxParallelism(transform.getId(), transform.getMaxParallelism());
for (Integer inputId: inputIds1) {
streamGraph.addEdge(inputId,
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamNode.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamNode.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamNode.java
index 430760b..9051891 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamNode.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamNode.java
@@ -30,6 +30,7 @@ import org.apache.flink.runtime.jobgraph.tasks.AbstractInvokable;
import org.apache.flink.streaming.api.collector.selector.OutputSelector;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.operators.StreamOperator;
+import org.apache.flink.util.Preconditions;
/**
* Class representing the operators in the streaming programs, with all their properties.
@@ -43,6 +44,11 @@ public class StreamNode implements Serializable {
private final int id;
private Integer parallelism = null;
+ /**
+ * Maximum parallelism for this stream node. The maximum parallelism is the upper limit for
+ * dynamic scaling and the number of key groups used for partitioned state.
+ */
+ private int maxParallelism;
private Long bufferTimeout = null;
private final String operatorName;
private String slotSharingGroup;
@@ -141,6 +147,25 @@ public class StreamNode implements Serializable {
this.parallelism = parallelism;
}
+ /**
+ * Get the maximum parallelism for this stream node.
+ *
+ * @return Maximum parallelism
+ */
+ int getMaxParallelism() {
+ return maxParallelism;
+ }
+
+ /**
+ * Set the maximum parallelism for this stream node.
+ *
+ * @param maxParallelism Maximum parallelism to be set
+ */
+ void setMaxParallelism(int maxParallelism) {
+ Preconditions.checkArgument(maxParallelism > 0, "The maximum parallelism must be at least 1.");
+ this.maxParallelism = maxParallelism;
+ }
+
public Long getBufferTimeout() {
return bufferTimeout != null ? bufferTimeout : env.getBufferTimeout();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java
index 28d982c..76fdaca 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java
@@ -26,6 +26,7 @@ import org.apache.commons.lang3.StringUtils;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.operators.util.UserCodeObjectWrapper;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
@@ -40,6 +41,7 @@ import org.apache.flink.runtime.jobgraph.tasks.JobSnapshottingSettings;
import org.apache.flink.runtime.jobmanager.scheduler.CoLocationGroup;
import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup;
import org.apache.flink.runtime.operators.util.TaskConfig;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.environment.CheckpointConfig;
import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator;
@@ -289,6 +291,21 @@ public class StreamingJobGraphGenerator {
if (parallelism > 0) {
jobVertex.setParallelism(parallelism);
+ } else {
+ parallelism = jobVertex.getParallelism();
+ }
+
+ int maxParallelism = streamNode.getMaxParallelism();
+
+ // the maximum parallelism specifies the upper bound for the parallelism
+ if (parallelism > maxParallelism) {
+ // the parallelism should always be smaller or equal than the max parallelism
+ throw new IllegalStateException("The maximum parallelism (" + maxParallelism + ") of " +
+ "the stream node " + streamNode + " is smaller than the parallelism (" +
+ parallelism + "). Increase the maximum parallelism or decrease the parallelism of" +
+ "this operator.");
+ } else {
+ jobVertex.setMaxParallelism(streamNode.getMaxParallelism());
}
if (LOG.isDebugEnabled()) {
@@ -325,7 +342,7 @@ public class StreamingJobGraphGenerator {
config.setTimeCharacteristic(streamGraph.getEnvironment().getStreamTimeCharacteristic());
final CheckpointConfig ceckpointCfg = streamGraph.getCheckpointConfig();
-
+
config.setStateBackend(streamGraph.getStateBackend());
config.setCheckpointingEnabled(ceckpointCfg.isCheckpointingEnabled());
if (ceckpointCfg.isCheckpointingEnabled()) {
@@ -339,7 +356,15 @@ public class StreamingJobGraphGenerator {
config.setStatePartitioner(0, vertex.getStatePartitioner1());
config.setStatePartitioner(1, vertex.getStatePartitioner2());
config.setStateKeySerializer(vertex.getStateKeySerializer());
-
+
+ // only set the key group assigner if the vertex uses partitioned state (= KeyedStream).
+ if (vertex.getStatePartitioner1() != null) {
+ // the key group assigner has to know the number of key groups (= maxParallelism)
+ KeyGroupAssigner<Object> keyGroupAssigner = new HashKeyGroupAssigner<Object>(vertex.getMaxParallelism());
+
+ config.setKeyGroupAssigner(keyGroupAssigner);
+ }
+
Class<? extends AbstractInvokable> vertexClass = vertex.getJobVertexClass();
if (vertexClass.equals(StreamIterationHead.class)
@@ -725,8 +750,6 @@ public class StreamingJobGraphGenerator {
// stream graph.
hasher.putInt(id);
- hasher.putInt(node.getParallelism());
-
if (node.getOperator() instanceof AbstractUdfStreamOperator) {
String udfClassName = ((AbstractUdfStreamOperator<?, ?>) node.getOperator())
.getUserFunction().getClass().getName();
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/transformations/StreamTransformation.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/transformations/StreamTransformation.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/transformations/StreamTransformation.java
index 1d2a1bb..e674619 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/transformations/StreamTransformation.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/transformations/StreamTransformation.java
@@ -117,6 +117,12 @@ public abstract class StreamTransformation<T> {
private int parallelism;
/**
+ * The maximum parallelism for this stream transformation. It defines the upper limit for
+ * dynamic scaling and the number of key groups used for partitioned state.
+ */
+ private int maxParallelism = -1;
+
+ /**
* User-specified ID for this transformation. This is used to assign the
* same operator ID across job restarts. There is also the automatically
* generated {@link #id}, which is assigned from a static counter. That
@@ -181,6 +187,24 @@ public abstract class StreamTransformation<T> {
}
/**
+ * Gets the maximum parallelism for this stream transformation.
+ *
+ * @return Maximum parallelism of this transformation.
+ */
+ public int getMaxParallelism() {
+ return maxParallelism;
+ }
+
+ /**
+ * Sets the maximum parallelism for this stream transformation.
+ *
+ * @param maxParallelism Maximum parallelism for this stream transformation.
+ */
+ public void setMaxParallelism(int maxParallelism) {
+ this.maxParallelism = maxParallelism;
+ }
+
+ /**
* Sets an ID for this {@link StreamTransformation}.
*
* <p>The specified ID is used to assign the same operator ID across job
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/ConfigurableStreamPartitioner.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/ConfigurableStreamPartitioner.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/ConfigurableStreamPartitioner.java
new file mode 100644
index 0000000..c59c88a
--- /dev/null
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/ConfigurableStreamPartitioner.java
@@ -0,0 +1,39 @@
+/*
+ * 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.partitioner;
+
+/**
+ * Interface for {@link StreamPartitioner} which have to be configured with the maximum parallelism
+ * of the stream transformation. The configure method is called by the StreamGraph when adding
+ * internal edges.
+ *
+ * This interface is required since the stream partitioners are instantiated eagerly. Due to that
+ * the maximum parallelism might not have been determined and needs to be set at a stage when the
+ * maximum parallelism could have been determined.
+ */
+public interface ConfigurableStreamPartitioner {
+
+ /**
+ * Configure the {@link StreamPartitioner} with the maximum parallelism of the down stream
+ * operator.
+ *
+ * @param maxParallelism Maximum parallelism of the down stream operator.
+ */
+ void configure(int maxParallelism);
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/HashPartitioner.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/HashPartitioner.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/HashPartitioner.java
deleted file mode 100644
index 3c93fb7..0000000
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/HashPartitioner.java
+++ /dev/null
@@ -1,66 +0,0 @@
-/*
- * 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.partitioner;
-
-import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.java.functions.KeySelector;
-import org.apache.flink.runtime.plugable.SerializationDelegate;
-import org.apache.flink.util.MathUtils;
-import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-
-/**
- * Partitioner selects the target channel based on the hash value of a key from a
- * {@link KeySelector}.
- *
- * @param <T> Type of the elements in the Stream being partitioned
- */
-@Internal
-public class HashPartitioner<T> extends StreamPartitioner<T> {
- private static final long serialVersionUID = 1L;
-
- private int[] returnArray = new int[1];
- KeySelector<T, ?> keySelector;
-
- public HashPartitioner(KeySelector<T, ?> keySelector) {
- this.keySelector = keySelector;
- }
-
- @Override
- public int[] selectChannels(SerializationDelegate<StreamRecord<T>> record,
- int numberOfOutputChannels) {
- Object key;
- try {
- key = keySelector.getKey(record.getInstance().getValue());
- } catch (Exception e) {
- throw new RuntimeException("Could not extract key from " + record.getInstance().getValue(), e);
- }
- returnArray[0] = MathUtils.murmurHash(key.hashCode()) % numberOfOutputChannels;
-
- return returnArray;
- }
-
- @Override
- public StreamPartitioner<T> copy() {
- return this;
- }
-
- @Override
- public String toString() {
- return "HASH";
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
new file mode 100644
index 0000000..5bcf41b
--- /dev/null
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
@@ -0,0 +1,82 @@
+/*
+ * 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.partitioner;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.plugable.SerializationDelegate;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.util.Preconditions;
+
+/**
+ * Partitioner selects the target channel based on the key group index. The key group
+ * index is derived from the key of the elements using the {@link KeyGroupAssigner}.
+ *
+ * @param <T> Type of the elements in the Stream being partitioned
+ */
+@Internal
+public class KeyGroupStreamPartitioner<T, K> extends StreamPartitioner<T> implements ConfigurableStreamPartitioner {
+ private static final long serialVersionUID = 1L;
+
+ private final int[] returnArray = new int[1];
+
+ private final KeySelector<T, K> keySelector;
+
+ private final KeyGroupAssigner<K> keyGroupAssigner;
+
+ public KeyGroupStreamPartitioner(KeySelector<T, K> keySelector, KeyGroupAssigner<K> keyGroupAssigner) {
+ this.keySelector = Preconditions.checkNotNull(keySelector);
+ this.keyGroupAssigner = Preconditions.checkNotNull(keyGroupAssigner);
+ }
+
+ public KeyGroupAssigner<K> getKeyGroupAssigner() {
+ return keyGroupAssigner;
+ }
+
+ @Override
+ public int[] selectChannels(
+ SerializationDelegate<StreamRecord<T>> record,
+ int numberOfOutputChannels) {
+
+ K key;
+ try {
+ key = keySelector.getKey(record.getInstance().getValue());
+ } catch (Exception e) {
+ throw new RuntimeException("Could not extract key from " + record.getInstance().getValue(), e);
+ }
+ returnArray[0] = keyGroupAssigner.getKeyGroupIndex(key) % numberOfOutputChannels;
+
+ return returnArray;
+ }
+
+ @Override
+ public StreamPartitioner<T> copy() {
+ return this;
+ }
+
+ @Override
+ public String toString() {
+ return "HASH";
+ }
+
+ @Override
+ public void configure(int maxParallelism) {
+ keyGroupAssigner.setup(maxParallelism);
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/ec975aab/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/AggregationFunctionTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/AggregationFunctionTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/AggregationFunctionTest.java
index d0617d0..1183306 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/AggregationFunctionTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/AggregationFunctionTest.java
@@ -46,7 +46,7 @@ import org.junit.Test;
public class AggregationFunctionTest {
@Test
- public void groupSumIntegerTest() {
+ public void groupSumIntegerTest() throws Exception {
// preparing expected outputs
List<Tuple2<Integer, Integer>> expectedGroupSumList = new ArrayList<>();
@@ -115,7 +115,7 @@ public class AggregationFunctionTest {
}
@Test
- public void pojoGroupSumIntegerTest() {
+ public void pojoGroupSumIntegerTest() throws Exception {
// preparing expected outputs
List<MyPojo> expectedGroupSumList = new ArrayList<>();
@@ -183,7 +183,7 @@ public class AggregationFunctionTest {
}
@Test
- public void minMaxByTest() {
+ public void minMaxByTest() throws Exception {
// Tuples are grouped on field 0, aggregated on field 1
// preparing expected outputs
@@ -250,7 +250,7 @@ public class AggregationFunctionTest {
}
@Test
- public void pojoMinMaxByTest() {
+ public void pojoMinMaxByTest() throws Exception {
// Pojos are grouped on field 0, aggregated on field 1
// preparing expected outputs
[26/27] flink git commit: [FLINK-3755] Fix variety of test problems
cause by Keyed-State Refactoring
Posted by al...@apache.org.
[FLINK-3755] Fix variety of test problems cause by Keyed-State Refactoring
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/f44b57cc
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/f44b57cc
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/f44b57cc
Branch: refs/heads/master
Commit: f44b57ccf8f088f2ad4c1f10f479ed62be17eb8b
Parents: 6d43061
Author: Stefan Richter <s....@data-artisans.com>
Authored: Mon Aug 29 16:10:15 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:10:02 2016 +0200
----------------------------------------------------------------------
.../state/RocksDBAsyncSnapshotTest.java | 58 ++++++++++++--------
.../flink/cep/operator/CEPOperatorTest.java | 6 ++
.../flink/runtime/executiongraph/Execution.java | 4 +-
.../checkpoint/CheckpointCoordinatorTest.java | 18 ------
.../runtime/state/StateBackendTestBase.java | 20 +++----
.../streaming/runtime/tasks/StreamTask.java | 8 ++-
.../api/graph/StreamGraphGeneratorTest.java | 5 +-
.../KeyedOneInputStreamOperatorTestHarness.java | 7 +++
.../api/scala/StreamingOperatorsITCase.scala | 1 +
.../WindowCheckpointingITCase.java | 14 +++--
.../test/streaming/runtime/IterateITCase.java | 1 +
.../translation/CustomPartitioningTest.scala | 5 +-
.../DeltaIterationTranslationTest.scala | 1 +
13 files changed, 83 insertions(+), 65 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
index 624905c..d5b9b46 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
@@ -69,6 +69,8 @@ import java.util.Arrays;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CancellationException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
import static org.junit.Assert.assertEquals;
@@ -173,10 +175,10 @@ public class RocksDBAsyncSnapshotTest {
}
}
- testHarness.processElement(new StreamRecord<>("Wohoo", 0));
-
task.triggerCheckpoint(42, 17);
+ testHarness.processElement(new StreamRecord<>("Wohoo", 0));
+
// now we allow the checkpoint
delayCheckpointLatch.trigger();
@@ -184,7 +186,13 @@ public class RocksDBAsyncSnapshotTest {
ensureCheckpointLatch.await();
testHarness.endInput();
+
+ ExecutorService threadPool = task.getAsyncOperationsThreadPool();
+ threadPool.shutdown();
+ Assert.assertTrue(threadPool.awaitTermination(60_000, TimeUnit.MILLISECONDS));
+
testHarness.waitForTaskCompletion();
+ task.checkTimerException();
}
/**
@@ -199,9 +207,6 @@ public class RocksDBAsyncSnapshotTest {
final OneInputStreamTask<String, String> task = new OneInputStreamTask<>();
- //ensure that the async threads complete before invoke method of the tasks returns.
- task.setThreadPoolTerminationTimeout(Long.MAX_VALUE);
-
final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
testHarness.configureForKeyedStream(new KeySelector<String, String>() {
@@ -232,6 +237,9 @@ public class RocksDBAsyncSnapshotTest {
new MockInputSplitProvider(),
testHarness.bufferSize);
+ BlockingStreamMemoryStateBackend.waitFirstWriteLatch = new OneShotLatch();
+ BlockingStreamMemoryStateBackend.unblockCancelLatch = new OneShotLatch();
+
testHarness.invoke(mockEnv);
// wait for the task to be running
@@ -241,36 +249,40 @@ public class RocksDBAsyncSnapshotTest {
while (!field.getBoolean(task)) {
Thread.sleep(10);
}
-
}
}
- testHarness.processElement(new StreamRecord<>("Wohoo", 0));
-
task.triggerCheckpoint(42, 17);
-
+ testHarness.processElement(new StreamRecord<>("Wohoo", 0));
BlockingStreamMemoryStateBackend.waitFirstWriteLatch.await();
task.cancel();
-
BlockingStreamMemoryStateBackend.unblockCancelLatch.trigger();
-
testHarness.endInput();
try {
+
+ ExecutorService threadPool = task.getAsyncOperationsThreadPool();
+ threadPool.shutdown();
+ Assert.assertTrue(threadPool.awaitTermination(60_000, TimeUnit.MILLISECONDS));
testHarness.waitForTaskCompletion();
+ task.checkTimerException();
+
Assert.fail("Operation completed. Cancel failed.");
} catch (Exception expected) {
- // we expect the exception from canceling snapshots
- Throwable cause = expected.getCause();
- if(cause instanceof AsynchronousException) {
- AsynchronousException asynchronousException = (AsynchronousException) cause;
- cause = asynchronousException.getCause();
- Assert.assertTrue("Unexpected Exception: " + cause,
- cause instanceof CancellationException //future canceled
- || cause instanceof InterruptedException); //thread interrupted
+ AsynchronousException asynchronousException = null;
+ if (expected instanceof AsynchronousException) {
+ asynchronousException = (AsynchronousException) expected;
+ } else if (expected.getCause() instanceof AsynchronousException) {
+ asynchronousException = (AsynchronousException) expected.getCause();
} else {
- Assert.fail();
+ Assert.fail("Unexpected exception: " + expected);
}
+
+ // we expect the exception from canceling snapshots
+ Throwable innerCause = asynchronousException.getCause();
+ Assert.assertTrue("Unexpected inner cause: " + innerCause,
+ innerCause instanceof CancellationException //future canceled
+ || innerCause instanceof InterruptedException); //thread interrupted
}
}
@@ -301,11 +313,11 @@ public class RocksDBAsyncSnapshotTest {
*/
static class BlockingStreamMemoryStateBackend extends MemoryStateBackend {
- public static OneShotLatch waitFirstWriteLatch = new OneShotLatch();
+ public static volatile OneShotLatch waitFirstWriteLatch = null;
- public static OneShotLatch unblockCancelLatch = new OneShotLatch();
+ public static volatile OneShotLatch unblockCancelLatch = null;
- volatile boolean closed = false;
+ private volatile boolean closed = false;
@Override
public CheckpointStreamFactory createStreamFactory(JobID jobId, String operatorIdentifier) throws IOException {
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java b/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
index 52a02d1..1fd8de8 100644
--- a/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
+++ b/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
@@ -136,6 +136,7 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with some elements in internal sorting queue
StreamStateHandle snapshot = harness.snapshot(0, 0);
+ harness.close();
harness = new OneInputStreamOperatorTestHarness<>(
new CEPPatternOperator<>(
@@ -157,6 +158,7 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with empty element queue but NFA state
StreamStateHandle snapshot2 = harness.snapshot(1, 1);
+ harness.close();
harness = new OneInputStreamOperatorTestHarness<>(
new CEPPatternOperator<>(
@@ -227,6 +229,7 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with some elements in internal sorting queue
StreamStateHandle snapshot = harness.snapshot(0, 0);
+ harness.close();
harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
@@ -252,6 +255,7 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with empty element queue but NFA state
StreamStateHandle snapshot2 = harness.snapshot(1, 1);
+ harness.close();
harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
@@ -334,6 +338,7 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with some elements in internal sorting queue
StreamStateHandle snapshot = harness.snapshot(0, 0);
+ harness.close();
harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
@@ -364,6 +369,7 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with empty element queue but NFA state
StreamStateHandle snapshot2 = harness.snapshot(1, 1);
+ harness.close();
harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
index 1981f5b..efddecc 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
@@ -387,8 +387,8 @@ public class Execution {
final TaskDeploymentDescriptor deployment = vertex.createDeploymentDescriptor(
attemptId,
slot,
- chainedStateHandle,
- keyGroupsStateHandles,
+ chainedStateHandle,
+ keyGroupsStateHandles,
attemptNumber);
// register this execution at the execution graph, to receive call backs
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
index 4972c51..bc61742 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
@@ -1669,7 +1669,6 @@ public class CheckpointCoordinatorTest {
200000,
0L,
1, // max one checkpoint at a time => should not affect savepoints
- 42,
new ExecutionVertex[] { vertex1 },
new ExecutionVertex[] { vertex1 },
new ExecutionVertex[] { vertex1 },
@@ -1721,7 +1720,6 @@ public class CheckpointCoordinatorTest {
200000,
100000000L, // very long min delay => should not affect savepoints
1,
- 42,
new ExecutionVertex[] { vertex1 },
new ExecutionVertex[] { vertex1 },
new ExecutionVertex[] { vertex1 },
@@ -1742,22 +1740,6 @@ public class CheckpointCoordinatorTest {
// Utilities
// ------------------------------------------------------------------------
- private static ExecutionVertex mockExecutionVertex(ExecutionAttemptID attemptID) {
- return mockExecutionVertex(attemptID, ExecutionState.RUNNING);
- }
-
- private static ExecutionVertex mockExecutionVertex(ExecutionAttemptID attemptID,
- ExecutionState state, ExecutionState ... successiveStates) {
- final Execution exec = mock(Execution.class);
- when(exec.getAttemptId()).thenReturn(attemptID);
- when(exec.getState()).thenReturn(state, successiveStates);
-
- ExecutionVertex vertex = mock(ExecutionVertex.class);
- when(vertex.getJobvertexId()).thenReturn(new JobVertexID());
- when(vertex.getCurrentExecutionAttempt()).thenReturn(exec);
-
- return vertex;
- }
/**
* Tests that the checkpointed partitioned and non-partitioned state is assigned properly to
* the {@link Execution} upon recovery.
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
index 5984aca..33ec182 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
@@ -1148,20 +1148,20 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
@Test
public void testEmptyStateCheckpointing() {
+
try {
- DummyEnvironment env = new DummyEnvironment("test", 1, 0);
- backend.initializeForJob(env, "test_op", IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = createStreamFactory();
+ KeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
+
+ ListStateDescriptor<String> kvId = new ListStateDescriptor<>("id", String.class);
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshot = backend
- .snapshotPartitionedState(682375462379L, 1);
-
+ // draw a snapshot
+ KeyGroupsStateHandle snapshot = runSnapshot(backend.snapshot(682375462379L, 1, streamFactory));
assertNull(snapshot);
- backend.dispose();
+ backend.close();
- // Make sure we can restore from empty state
- backend.initializeForJob(env, "test_op", IntSerializer.INSTANCE);
- backend.injectKeyValueStateSnapshots((HashMap) snapshot);
- backend.dispose();
+ backend = restoreKeyedBackend(IntSerializer.INSTANCE, snapshot);
+ backend.close();
}
catch (Exception e) {
e.printStackTrace();
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
index 701281b..bedc8fa 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
@@ -614,7 +614,6 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
private boolean performCheckpoint(final long checkpointId, final long timestamp) throws Exception {
LOG.debug("Starting checkpoint {} on task {}", checkpointId, getName());
-
synchronized (lock) {
if (isRunning) {
@@ -677,7 +676,6 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
synchronized (cancelables) {
cancelables.add(asyncCheckpointRunnable);
}
-
asyncOperationsThreadPool.submit(asyncCheckpointRunnable);
return true;
} else {
@@ -685,7 +683,11 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
}
}
}
-
+
+ public ExecutorService getAsyncOperationsThreadPool() {
+ return asyncOperationsThreadPool;
+ }
+
@Override
public void notifyCheckpointComplete(long checkpointId) throws Exception {
synchronized (lock) {
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
index 874274f..c93a439 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
@@ -22,6 +22,7 @@ import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
@@ -371,8 +372,8 @@ public class StreamGraphGeneratorTest {
StreamNode keyedResult3Node = graph.getStreamNode(keyedResult3.getId());
StreamNode keyedResult4Node = graph.getStreamNode(keyedResult4.getId());
- assertEquals(globalParallelism, keyedResult1Node.getMaxParallelism());
- assertEquals(mapParallelism, keyedResult2Node.getMaxParallelism());
+ assertEquals(KeyGroupRangeAssignment.DEFAULT_MAX_PARALLELISM, keyedResult1Node.getMaxParallelism());
+ assertEquals(KeyGroupRangeAssignment.DEFAULT_MAX_PARALLELISM, keyedResult2Node.getMaxParallelism());
assertEquals(maxParallelism, keyedResult3Node.getMaxParallelism());
assertEquals(maxParallelism, keyedResult4Node.getMaxParallelism());
}
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
index 03f50f9..7e86da0 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
@@ -113,6 +113,10 @@ public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
final int numberOfKeyGroups = (Integer) invocationOnMock.getArguments()[1];
final KeyGroupRange keyGroupRange = (KeyGroupRange) invocationOnMock.getArguments()[2];
+ if(keyedStateBackend != null) {
+ keyedStateBackend.close();
+ }
+
if (restoredKeyedState == null) {
keyedStateBackend = stateBackend.createKeyedStateBackend(
mockTask.getEnvironment(),
@@ -195,5 +199,8 @@ public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
*/
public void close() throws Exception {
super.close();
+ if(keyedStateBackend != null) {
+ keyedStateBackend.close();
+ }
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/StreamingOperatorsITCase.scala
----------------------------------------------------------------------
diff --git a/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/StreamingOperatorsITCase.scala b/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/StreamingOperatorsITCase.scala
index d353468..c57c29c 100644
--- a/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/StreamingOperatorsITCase.scala
+++ b/flink-streaming-scala/src/test/scala/org/apache/flink/streaming/api/scala/StreamingOperatorsITCase.scala
@@ -69,6 +69,7 @@ class StreamingOperatorsITCase extends ScalaStreamingMultipleProgramsTestBase {
val env = StreamExecutionEnvironment.getExecutionEnvironment
env.setParallelism(2)
+ env.getConfig.setMaxParallelism(2);
val sourceStream = env.addSource(new SourceFunction[(Int, Int)] {
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-tests/src/test/java/org/apache/flink/test/checkpointing/WindowCheckpointingITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/WindowCheckpointingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/WindowCheckpointingITCase.java
index 2d634de..2e6ce78 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/WindowCheckpointingITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/WindowCheckpointingITCase.java
@@ -337,7 +337,6 @@ public class WindowCheckpointingITCase extends TestLogger {
// we loop longer than we have elements, to permit delayed checkpoints
// to still cause a failure
while (running) {
-
if (!failedBefore) {
// delay a bit, if we have not failed before
Thread.sleep(1);
@@ -350,17 +349,15 @@ public class WindowCheckpointingITCase extends TestLogger {
}
if (numElementsEmitted < numElementsToEmit &&
- (failedBefore || numElementsEmitted <= failureAfterNumElements))
- {
+ (failedBefore || numElementsEmitted <= failureAfterNumElements)) {
// the function failed before, or we are in the elements before the failure
synchronized (ctx.getCheckpointLock()) {
int next = numElementsEmitted++;
ctx.collect(new Tuple2<Long, IntType>((long) next, new IntType(next)));
}
- }
- else {
+ } else {
// if our work is done, delay a bit to prevent busy waiting
- Thread.sleep(1);
+ Thread.sleep(10);
}
}
}
@@ -409,6 +406,7 @@ public class WindowCheckpointingITCase extends TestLogger {
public void open(Configuration parameters) throws Exception {
// this sink can only work with DOP 1
assertEquals(1, getRuntimeContext().getNumberOfParallelSubtasks());
+ checkSuccess();
}
@Override
@@ -423,6 +421,10 @@ public class WindowCheckpointingITCase extends TestLogger {
// check if we have seen all we expect
aggCount += value.f1.value;
+ checkSuccess();
+ }
+
+ private void checkSuccess() throws SuccessException {
if (aggCount >= elementCountExpected * countPerElementExpected) {
// we are done. validate
assertEquals(elementCountExpected, counts.size());
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/IterateITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/IterateITCase.java b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/IterateITCase.java
index 1fbebd0..e49f832 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/IterateITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/IterateITCase.java
@@ -524,6 +524,7 @@ public class IterateITCase extends StreamingMultipleProgramsTestBase {
try {
StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
env.setParallelism(DEFAULT_PARALLELISM - 1);
+ env.getConfig().setMaxParallelism(env.getParallelism());
KeySelector<Integer, Integer> key = new KeySelector<Integer, Integer>() {
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/CustomPartitioningTest.scala
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/CustomPartitioningTest.scala b/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/CustomPartitioningTest.scala
index 7ebf378..2ef5f01 100644
--- a/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/CustomPartitioningTest.scala
+++ b/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/CustomPartitioningTest.scala
@@ -39,6 +39,7 @@ class CustomPartitioningTest extends CompilerTestBase {
val env = ExecutionEnvironment.getExecutionEnvironment
env.setParallelism(parallelism)
+ env.getConfig.setMaxParallelism(parallelism);
val data = env.fromElements( (0,0) ).rebalance()
@@ -108,7 +109,8 @@ class CustomPartitioningTest extends CompilerTestBase {
val env = ExecutionEnvironment.getExecutionEnvironment
env.setParallelism(parallelism)
-
+ env.getConfig.setMaxParallelism(parallelism);
+
val data = env.fromElements(new Pojo()).rebalance()
data
@@ -179,6 +181,7 @@ class CustomPartitioningTest extends CompilerTestBase {
val env = ExecutionEnvironment.getExecutionEnvironment
env.setParallelism(parallelism)
+ env.getConfig.setMaxParallelism(parallelism);
val data = env.fromElements(new Pojo()).rebalance()
http://git-wip-us.apache.org/repos/asf/flink/blob/f44b57cc/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/DeltaIterationTranslationTest.scala
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/DeltaIterationTranslationTest.scala b/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/DeltaIterationTranslationTest.scala
index 3121d68..05294b9 100644
--- a/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/DeltaIterationTranslationTest.scala
+++ b/flink-tests/src/test/scala/org/apache/flink/api/scala/operators/translation/DeltaIterationTranslationTest.scala
@@ -52,6 +52,7 @@ class DeltaIterationTranslationTest {
val env = ExecutionEnvironment.getExecutionEnvironment
env.setParallelism(DEFAULT_PARALLELISM)
+ env.getConfig.setMaxParallelism(DEFAULT_PARALLELISM);
val initialSolutionSet = env.fromElements((3.44, 5L, "abc"))
val initialWorkSet = env.fromElements((1.23, "abc"))
[16/27] flink git commit: [FLINK-4381] Refactor State to Prepare For
Key-Group State Backends
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/state/FsCheckpointStateOutputStreamTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/FsCheckpointStateOutputStreamTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/FsCheckpointStateOutputStreamTest.java
deleted file mode 100644
index 5964b72..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/FsCheckpointStateOutputStreamTest.java
+++ /dev/null
@@ -1,128 +0,0 @@
-/*
- * 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 org.apache.flink.runtime.state.filesystem.FileStreamStateHandle;
-import org.apache.flink.runtime.state.filesystem.FsStateBackend;
-import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
-
-import org.junit.Test;
-
-import java.io.File;
-import java.io.InputStream;
-import java.util.Random;
-
-import static org.junit.Assert.*;
-
-public class FsCheckpointStateOutputStreamTest {
-
- /** The temp dir, obtained in a platform neutral way */
- private static final Path TEMP_DIR_PATH = new Path(new File(System.getProperty("java.io.tmpdir")).toURI());
-
-
- @Test(expected = IllegalArgumentException.class)
- public void testWrongParameters() {
- // this should fail
- new FsStateBackend.FsCheckpointStateOutputStream(
- TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), 4000, 5000);
- }
-
-
- @Test
- public void testEmptyState() throws Exception {
- AbstractStateBackend.CheckpointStateOutputStream stream = new FsStateBackend.FsCheckpointStateOutputStream(
- TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), 1024, 512);
-
- StreamStateHandle handle = stream.closeAndGetHandle();
- assertTrue(handle instanceof ByteStreamStateHandle);
-
- InputStream inStream = handle.getState(ClassLoader.getSystemClassLoader());
- assertEquals(-1, inStream.read());
- }
-
- @Test
- public void testStateBlowMemThreshold() throws Exception {
- runTest(222, 999, 512, false);
- }
-
- @Test
- public void testStateOneBufferAboveThreshold() throws Exception {
- runTest(896, 1024, 15, true);
- }
-
- @Test
- public void testStateAboveMemThreshold() throws Exception {
- runTest(576446, 259, 17, true);
- }
-
- @Test
- public void testZeroThreshold() throws Exception {
- runTest(16678, 4096, 0, true);
- }
-
- private void runTest(int numBytes, int bufferSize, int threshold, boolean expectFile) throws Exception {
- AbstractStateBackend.CheckpointStateOutputStream stream =
- new FsStateBackend.FsCheckpointStateOutputStream(
- TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), bufferSize, threshold);
-
- Random rnd = new Random();
- byte[] original = new byte[numBytes];
- byte[] bytes = new byte[original.length];
-
- rnd.nextBytes(original);
- System.arraycopy(original, 0, bytes, 0, original.length);
-
- // the test writes a mixture of writing individual bytes and byte arrays
- int pos = 0;
- while (pos < bytes.length) {
- boolean single = rnd.nextBoolean();
- if (single) {
- stream.write(bytes[pos++]);
- }
- else {
- int num = rnd.nextInt(Math.min(10, bytes.length - pos));
- stream.write(bytes, pos, num);
- pos += num;
- }
- }
-
- StreamStateHandle handle = stream.closeAndGetHandle();
- if (expectFile) {
- assertTrue(handle instanceof FileStreamStateHandle);
- } else {
- assertTrue(handle instanceof ByteStreamStateHandle);
- }
-
- // make sure the writing process did not alter the original byte array
- assertArrayEquals(original, bytes);
-
- InputStream inStream = handle.getState(ClassLoader.getSystemClassLoader());
- byte[] validation = new byte[bytes.length];
- int bytesRead = inStream.read(validation);
-
- assertEquals(numBytes, bytesRead);
- assertEquals(-1, inStream.read());
-
- assertArrayEquals(bytes, validation);
-
- handle.discardState();
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/state/KeyGroupRangeOffsetTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/KeyGroupRangeOffsetTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/KeyGroupRangeOffsetTest.java
new file mode 100644
index 0000000..95564cc
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/KeyGroupRangeOffsetTest.java
@@ -0,0 +1,136 @@
+/*
+ * 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.java.tuple.Tuple2;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.Arrays;
+
+public class KeyGroupRangeOffsetTest {
+
+ @Test
+ public void testKeyGroupIntersection() {
+ long[] offsets = new long[9];
+ for (int i = 0; i < offsets.length; ++i) {
+ offsets[i] = i;
+ }
+
+ int startKeyGroup = 2;
+
+ KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(KeyGroupRange.of(startKeyGroup, 10), offsets);
+ KeyGroupRangeOffsets intersection = keyGroupRangeOffsets.getIntersection(KeyGroupRange.of(3, 7));
+ KeyGroupRangeOffsets expected = new KeyGroupRangeOffsets(
+ KeyGroupRange.of(3, 7), Arrays.copyOfRange(offsets, 3 - startKeyGroup, 8 - startKeyGroup));
+ Assert.assertEquals(expected, intersection);
+
+ Assert.assertEquals(keyGroupRangeOffsets, keyGroupRangeOffsets.getIntersection(
+ keyGroupRangeOffsets.getKeyGroupRange()));
+
+ intersection = keyGroupRangeOffsets.getIntersection(KeyGroupRange.of(11, 13));
+ Assert.assertEquals(KeyGroupRange.EMPTY_KEY_GROUP, intersection.getKeyGroupRange());
+ Assert.assertFalse(intersection.iterator().hasNext());
+
+ intersection = keyGroupRangeOffsets.getIntersection(KeyGroupRange.of(5, 13));
+ expected = new KeyGroupRangeOffsets(KeyGroupRange.of(5, 10), Arrays.copyOfRange(
+ offsets, 5 - startKeyGroup, 11 - startKeyGroup));
+ Assert.assertEquals(expected, intersection);
+
+ intersection = keyGroupRangeOffsets.getIntersection(KeyGroupRange.of(0, 2));
+ expected = new KeyGroupRangeOffsets(KeyGroupRange.of(2, 2), Arrays.copyOfRange(
+ offsets, 2 - startKeyGroup, 3 - startKeyGroup));
+ Assert.assertEquals(intersection, intersection);
+ }
+
+ @Test
+ public void testKeyGroupRangeOffsetsBasics() {
+ testKeyGroupRangeOffsetsBasicsInternal(0, 0);
+ testKeyGroupRangeOffsetsBasicsInternal(0, 1);
+ testKeyGroupRangeOffsetsBasicsInternal(1, 2);
+ testKeyGroupRangeOffsetsBasicsInternal(42, 42);
+ testKeyGroupRangeOffsetsBasicsInternal(3, 7);
+ testKeyGroupRangeOffsetsBasicsInternal(0, Short.MAX_VALUE);
+ testKeyGroupRangeOffsetsBasicsInternal(Short.MAX_VALUE - 1, Short.MAX_VALUE);
+
+ try {
+ testKeyGroupRangeOffsetsBasicsInternal(-3, 2);
+ Assert.fail();
+ } catch (IllegalArgumentException ex) {
+ //expected
+ }
+
+ KeyGroupRangeOffsets testNoGivenOffsets = new KeyGroupRangeOffsets(3, 7);
+ for (int i = 3; i <= 7; ++i) {
+ testNoGivenOffsets.setKeyGroupOffset(i, i + 1);
+ }
+ for (int i = 3; i <= 7; ++i) {
+ Assert.assertEquals(i + 1, testNoGivenOffsets.getKeyGroupOffset(i));
+ }
+
+ }
+
+ private void testKeyGroupRangeOffsetsBasicsInternal(int startKeyGroup, int endKeyGroup) {
+
+ long[] offsets = new long[endKeyGroup - startKeyGroup + 1];
+ for (int i = 0; i < offsets.length; ++i) {
+ offsets[i] = i;
+ }
+
+ KeyGroupRangeOffsets keyGroupRange = new KeyGroupRangeOffsets(startKeyGroup, endKeyGroup, offsets);
+ KeyGroupRangeOffsets sameButDifferentConstr =
+ new KeyGroupRangeOffsets(KeyGroupRange.of(startKeyGroup, endKeyGroup), offsets);
+ Assert.assertEquals(keyGroupRange, sameButDifferentConstr);
+
+ int numberOfKeyGroup = keyGroupRange.getKeyGroupRange().getNumberOfKeyGroups();
+ Assert.assertEquals(Math.max(0, endKeyGroup - startKeyGroup + 1), numberOfKeyGroup);
+ if (numberOfKeyGroup > 0) {
+ Assert.assertEquals(startKeyGroup, keyGroupRange.getKeyGroupRange().getStartKeyGroup());
+ Assert.assertEquals(endKeyGroup, keyGroupRange.getKeyGroupRange().getEndKeyGroup());
+ int c = startKeyGroup;
+ for (Tuple2<Integer, Long> tuple : keyGroupRange) {
+ Assert.assertEquals(c, (int) tuple.f0);
+ Assert.assertTrue(keyGroupRange.getKeyGroupRange().contains(tuple.f0));
+ Assert.assertEquals((long) c - startKeyGroup, (long) tuple.f1);
+ ++c;
+ }
+
+ for (int i = startKeyGroup; i <= endKeyGroup; ++i) {
+ Assert.assertEquals(i - startKeyGroup, keyGroupRange.getKeyGroupOffset(i));
+ }
+
+ int newOffset = 42;
+ for (int i = startKeyGroup; i <= endKeyGroup; ++i) {
+ keyGroupRange.setKeyGroupOffset(i, newOffset);
+ ++newOffset;
+ }
+
+ for (int i = startKeyGroup; i <= endKeyGroup; ++i) {
+ Assert.assertEquals(42 + i - startKeyGroup, keyGroupRange.getKeyGroupOffset(i));
+ }
+
+ Assert.assertEquals(endKeyGroup + 1, c);
+ Assert.assertFalse(keyGroupRange.getKeyGroupRange().contains(startKeyGroup - 1));
+ Assert.assertFalse(keyGroupRange.getKeyGroupRange().contains(endKeyGroup + 1));
+ } else {
+ Assert.assertEquals(KeyGroupRange.EMPTY_KEY_GROUP, keyGroupRange);
+ }
+ }
+
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/state/KeyGroupRangeTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/KeyGroupRangeTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/KeyGroupRangeTest.java
new file mode 100644
index 0000000..ab0c327
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/KeyGroupRangeTest.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.junit.Assert;
+import org.junit.Test;
+
+public class KeyGroupRangeTest {
+
+ @Test
+ public void testKeyGroupIntersection() {
+ KeyGroupRange keyGroupRange1 = KeyGroupRange.of(0, 10);
+ KeyGroupRange keyGroupRange2 = KeyGroupRange.of(3, 7);
+ KeyGroupRange intersection = keyGroupRange1.getIntersection(keyGroupRange2);
+ Assert.assertEquals(3, intersection.getStartKeyGroup());
+ Assert.assertEquals(7, intersection.getEndKeyGroup());
+ Assert.assertEquals(intersection, keyGroupRange2.getIntersection(keyGroupRange1));
+
+ Assert.assertEquals(keyGroupRange1, keyGroupRange1.getIntersection(keyGroupRange1));
+
+ keyGroupRange1 = KeyGroupRange.of(0,5);
+ keyGroupRange2 = KeyGroupRange.of(6,10);
+ intersection =keyGroupRange1.getIntersection(keyGroupRange2);
+ Assert.assertEquals(KeyGroupRange.EMPTY_KEY_GROUP, intersection);
+ Assert.assertEquals(intersection, keyGroupRange2.getIntersection(keyGroupRange1));
+
+ keyGroupRange1 = KeyGroupRange.of(0, 10);
+ keyGroupRange2 = KeyGroupRange.of(5, 20);
+ intersection = keyGroupRange1.getIntersection(keyGroupRange2);
+ Assert.assertEquals(5, intersection.getStartKeyGroup());
+ Assert.assertEquals(10, intersection.getEndKeyGroup());
+ Assert.assertEquals(intersection, keyGroupRange2.getIntersection(keyGroupRange1));
+
+ keyGroupRange1 = KeyGroupRange.of(3, 12);
+ keyGroupRange2 = KeyGroupRange.of(0, 10);
+ intersection = keyGroupRange1.getIntersection(keyGroupRange2);
+ Assert.assertEquals(3, intersection.getStartKeyGroup());
+ Assert.assertEquals(10, intersection.getEndKeyGroup());
+ Assert.assertEquals(intersection, keyGroupRange2.getIntersection(keyGroupRange1));
+ }
+
+
+ @Test
+ public void testKeyGroupRangeBasics() {
+ testKeyGroupRangeBasicsInternal(0, 0);
+ testKeyGroupRangeBasicsInternal(0, 1);
+ testKeyGroupRangeBasicsInternal(1, 2);
+ testKeyGroupRangeBasicsInternal(42, 42);
+ testKeyGroupRangeBasicsInternal(3, 7);
+ testKeyGroupRangeBasicsInternal(0, Short.MAX_VALUE);
+ testKeyGroupRangeBasicsInternal(Short.MAX_VALUE - 1, Short.MAX_VALUE);
+
+ try {
+ testKeyGroupRangeBasicsInternal(-3, 2);
+ Assert.fail();
+ } catch (IllegalArgumentException ex) {
+ //expected
+ }
+
+ }
+
+ private void testKeyGroupRangeBasicsInternal(int startKeyGroup, int endKeyGroup) {
+ KeyGroupRange keyGroupRange = KeyGroupRange.of(startKeyGroup, endKeyGroup);
+ int numberOfKeyGroup = keyGroupRange.getNumberOfKeyGroups();
+ Assert.assertEquals(Math.max(0, endKeyGroup - startKeyGroup + 1), numberOfKeyGroup);
+ if (keyGroupRange.getNumberOfKeyGroups() > 0) {
+ Assert.assertEquals(startKeyGroup, keyGroupRange.getStartKeyGroup());
+ Assert.assertEquals(endKeyGroup, keyGroupRange.getEndKeyGroup());
+ int c = startKeyGroup;
+ for(int i : keyGroupRange) {
+ Assert.assertEquals(c, i);
+ Assert.assertTrue(keyGroupRange.contains(i));
+ ++c;
+ }
+
+ Assert.assertEquals(endKeyGroup + 1, c);
+ Assert.assertFalse(keyGroupRange.contains(startKeyGroup - 1));
+ Assert.assertFalse(keyGroupRange.contains(endKeyGroup + 1));
+ } else {
+ Assert.assertEquals(KeyGroupRange.EMPTY_KEY_GROUP, keyGroupRange);
+ }
+ }
+
+
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java
index d3b4dbc..940b337 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/MemoryStateBackendTest.java
@@ -57,37 +57,26 @@ public class MemoryStateBackendTest extends StateBackendTestBase<MemoryStateBack
public void testReducingStateRestoreWithWrongSerializers() {}
@Test
- public void testSerializableState() {
+ public void testOversizedState() {
try {
- MemoryStateBackend backend = new MemoryStateBackend();
+ MemoryStateBackend backend = new MemoryStateBackend(10);
HashMap<String, Integer> state = new HashMap<>();
state.put("hey there", 2);
state.put("the crazy brown fox stumbles over a sentence that does not contain every letter", 77);
- StateHandle<HashMap<String, Integer>> handle = backend.checkpointStateSerializable(state, 12, 459);
- assertNotNull(handle);
+ try {
+ AbstractStateBackend.CheckpointStateOutputStream outStream = backend.createCheckpointStateOutputStream(
+ 12,
+ 459);
- HashMap<String, Integer> restored = handle.getState(getClass().getClassLoader());
- assertEquals(state, restored);
- }
- catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
- }
- }
+ ObjectOutputStream oos = new ObjectOutputStream(outStream);
+ oos.writeObject(state);
- @Test
- public void testOversizedState() {
- try {
- MemoryStateBackend backend = new MemoryStateBackend(10);
+ oos.flush();
- HashMap<String, Integer> state = new HashMap<>();
- state.put("hey there", 2);
- state.put("the crazy brown fox stumbles over a sentence that does not contain every letter", 77);
+ outStream.closeAndGetHandle();
- try {
- backend.checkpointStateSerializable(state, 12, 459);
fail("this should cause an exception");
}
catch (IOException e) {
@@ -117,7 +106,7 @@ public class MemoryStateBackendTest extends StateBackendTestBase<MemoryStateBack
assertNotNull(handle);
- ObjectInputStream ois = new ObjectInputStream(handle.getState(getClass().getClassLoader()));
+ ObjectInputStream ois = new ObjectInputStream(handle.openInputStream());
assertEquals(state, ois.readObject());
assertTrue(ois.available() <= 0);
ois.close();
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
index 7b00b27..834c35c 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
@@ -84,7 +84,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
@After
public void teardown() throws Exception {
- this.backend.dispose();
+ this.backend.discardState();
cleanup();
}
@@ -154,7 +154,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("u3", state.value());
assertEquals("u3", getSerializedValue(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.dispose();
+ backend.discardState();
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
@@ -174,7 +174,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("2", restored1.value());
assertEquals("2", getSerializedValue(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.dispose();
+ backend.discardState();
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
backend.injectKeyValueStateSnapshots((HashMap) snapshot2);
@@ -254,7 +254,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
}
}
- backend.dispose();
+ backend.discardState();
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
backend.injectKeyValueStateSnapshots((HashMap) snapshot1);
@@ -334,7 +334,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("u3", joiner.join(state.get()));
assertEquals("u3", joiner.join(getSerializedList(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
- backend.dispose();
+ backend.discardState();
// restore the first snapshot and validate it
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
@@ -355,7 +355,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("2", joiner.join(restored1.get()));
assertEquals("2", joiner.join(getSerializedList(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer)));
- backend.dispose();
+ backend.discardState();
// restore the second snapshot and validate it
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
@@ -452,7 +452,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("u3", state.get());
assertEquals("u3", getSerializedValue(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.dispose();
+ backend.discardState();
// restore the first snapshot and validate it
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
@@ -473,7 +473,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("2", restored1.get());
assertEquals("2", getSerializedValue(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.dispose();
+ backend.discardState();
// restore the second snapshot and validate it
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
@@ -574,7 +574,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("Fold-Initial:,103", state.get());
assertEquals("Fold-Initial:,103", getSerializedValue(kvState, 3, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.dispose();
+ backend.discardState();
// restore the first snapshot and validate it
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
@@ -595,7 +595,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
assertEquals("Fold-Initial:,2", restored1.get());
assertEquals("Fold-Initial:,2", getSerializedValue(restoredKvState1, 2, keySerializer, VoidNamespace.INSTANCE, namespaceSerializer, valueSerializer));
- backend.dispose();
+ backend.discardState();
// restore the second snapshot and validate it
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
@@ -653,7 +653,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
}
}
- backend.dispose();
+ backend.discardState();
// restore the first snapshot and validate it
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
@@ -713,7 +713,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
}
}
- backend.dispose();
+ backend.discardState();
// restore the first snapshot and validate it
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
@@ -775,7 +775,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
}
}
- backend.dispose();
+ backend.discardState();
// restore the first snapshot and validate it
backend.initializeForJob(new DummyEnvironment("test", 1, 0), "test_op", IntSerializer.INSTANCE);
@@ -1030,7 +1030,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
}
// Verify unregistered
- backend.dispose();
+ backend.discardState();
verify(listener, times(1)).notifyKvStateUnregistered(
eq(env.getJobID()), eq(env.getJobVertexId()), eq(0), eq("banana"));
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStateOutputStreamTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStateOutputStreamTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStateOutputStreamTest.java
new file mode 100644
index 0000000..1d45115
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/filesystem/FsCheckpointStateOutputStreamTest.java
@@ -0,0 +1,128 @@
+/*
+ * 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.FileSystem;
+import org.apache.flink.core.fs.Path;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+import org.junit.Test;
+
+import java.io.DataInputStream;
+import java.io.File;
+import java.io.InputStream;
+import java.util.Random;
+
+import static org.junit.Assert.*;
+
+public class FsCheckpointStateOutputStreamTest {
+
+ /** The temp dir, obtained in a platform neutral way */
+ private static final Path TEMP_DIR_PATH = new Path(new File(System.getProperty("java.io.tmpdir")).toURI());
+
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testWrongParameters() {
+ // this should fail
+ new FsStateBackend.FsCheckpointStateOutputStream(
+ TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), 4000, 5000);
+ }
+
+
+ @Test
+ public void testEmptyState() throws Exception {
+ AbstractStateBackend.CheckpointStateOutputStream stream = new FsStateBackend.FsCheckpointStateOutputStream(
+ TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), 1024, 512);
+
+ StreamStateHandle handle = stream.closeAndGetHandle();
+ assertTrue(handle instanceof ByteStreamStateHandle);
+
+ InputStream inStream = handle.openInputStream();
+ assertEquals(-1, inStream.read());
+ }
+
+ @Test
+ public void testStateBlowMemThreshold() throws Exception {
+ runTest(222, 999, 512, false);
+ }
+
+ @Test
+ public void testStateOneBufferAboveThreshold() throws Exception {
+ runTest(896, 1024, 15, true);
+ }
+
+ @Test
+ public void testStateAboveMemThreshold() throws Exception {
+ runTest(576446, 259, 17, true);
+ }
+
+ @Test
+ public void testZeroThreshold() throws Exception {
+ runTest(16678, 4096, 0, true);
+ }
+
+ private void runTest(int numBytes, int bufferSize, int threshold, boolean expectFile) throws Exception {
+ AbstractStateBackend.CheckpointStateOutputStream stream =
+ new FsStateBackend.FsCheckpointStateOutputStream(
+ TEMP_DIR_PATH, FileSystem.getLocalFileSystem(), bufferSize, threshold);
+
+ Random rnd = new Random();
+ byte[] original = new byte[numBytes];
+ byte[] bytes = new byte[original.length];
+
+ rnd.nextBytes(original);
+ System.arraycopy(original, 0, bytes, 0, original.length);
+
+ // the test writes a mixture of writing individual bytes and byte arrays
+ int pos = 0;
+ while (pos < bytes.length) {
+ boolean single = rnd.nextBoolean();
+ if (single) {
+ stream.write(bytes[pos++]);
+ }
+ else {
+ int num = rnd.nextInt(Math.min(10, bytes.length - pos));
+ stream.write(bytes, pos, num);
+ pos += num;
+ }
+ }
+
+ StreamStateHandle handle = stream.closeAndGetHandle();
+ if (expectFile) {
+ assertTrue(handle instanceof FileStateHandle);
+ } else {
+ assertTrue(handle instanceof ByteStreamStateHandle);
+ }
+
+ // make sure the writing process did not alter the original byte array
+ assertArrayEquals(original, bytes);
+
+ InputStream inStream = handle.openInputStream();
+ byte[] validation = new byte[bytes.length];
+
+ DataInputStream dataInputStream = new DataInputStream(inStream);
+ dataInputStream.readFully(validation);
+
+ assertArrayEquals(bytes, validation);
+
+ handle.discardState();
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
index fc1a5df..429fc6b 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/taskmanager/TaskAsyncCallTest.java
@@ -44,16 +44,19 @@ import org.apache.flink.runtime.jobgraph.tasks.StatefulTask;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
import org.apache.flink.util.SerializedValue;
import org.junit.Before;
import org.junit.Test;
import scala.concurrent.duration.FiniteDuration;
-import java.io.Serializable;
import java.net.URL;
import java.util.Collections;
+import java.util.List;
import java.util.concurrent.TimeUnit;
import static org.junit.Assert.assertFalse;
@@ -179,8 +182,8 @@ public class TaskAsyncCallTest {
new TaskManagerRuntimeInfo("localhost", new Configuration(), System.getProperty("java.io.tmpdir")),
mock(TaskMetricGroup.class));
}
-
- public static class CheckpointsInOrderInvokable extends AbstractInvokable implements StatefulTask<StateHandle<Serializable>> {
+
+ public static class CheckpointsInOrderInvokable extends AbstractInvokable implements StatefulTask {
private volatile long lastCheckpointId = 0;
@@ -204,7 +207,10 @@ public class TaskAsyncCallTest {
}
@Override
- public void setInitialState(StateHandle<Serializable> stateHandle) throws Exception {}
+ public void setInitialState(ChainedStateHandle<StreamStateHandle> chainedState,
+ List<KeyGroupsStateHandle> keyGroupsState) throws Exception {
+
+ }
@Override
public boolean triggerCheckpoint(long checkpointId, long timestamp) {
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/test/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStoreITCase.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStoreITCase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStoreITCase.java
index 7505bfc..7e8868c 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStoreITCase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/zookeeper/ZooKeeperStateHandleStoreITCase.java
@@ -22,7 +22,7 @@ import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.api.BackgroundCallback;
import org.apache.curator.framework.api.CuratorEvent;
import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
import org.apache.flink.util.InstantiationUtil;
import org.apache.flink.util.TestLogger;
import org.apache.zookeeper.CreateMode;
@@ -97,7 +97,7 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
// Verify
// State handle created
assertEquals(1, store.getAll().size());
- assertEquals(state, store.get(pathInZooKeeper).getState(null));
+ assertEquals(state, store.get(pathInZooKeeper).retrieveState());
// Path created and is persistent
Stat stat = ZooKeeper.getClient().checkExists().forPath(pathInZooKeeper);
@@ -106,9 +106,9 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
// Data is equal
@SuppressWarnings("unchecked")
- Long actual = ((StateHandle<Long>) InstantiationUtil.deserializeObject(
+ Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject(
ZooKeeper.getClient().getData().forPath(pathInZooKeeper),
- ClassLoader.getSystemClassLoader())).getState(null);
+ ClassLoader.getSystemClassLoader())).retrieveState();
assertEquals(state, actual);
}
@@ -149,7 +149,7 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
// Verify
// State handle created
assertEquals(i + 1, store.getAll().size());
- assertEquals(state, longStateStorage.getStateHandles().get(i).getState(null));
+ assertEquals(state, longStateStorage.getStateHandles().get(i).retrieveState());
// Path created
Stat stat = ZooKeeper.getClient().checkExists().forPath(pathInZooKeeper);
@@ -166,9 +166,9 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
// Data is equal
@SuppressWarnings("unchecked")
- Long actual = ((StateHandle<Long>) InstantiationUtil.deserializeObject(
+ Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject(
ZooKeeper.getClient().getData().forPath(pathInZooKeeper),
- ClassLoader.getSystemClassLoader())).getState(null);
+ ClassLoader.getSystemClassLoader())).retrieveState();
assertEquals(state, actual);
}
@@ -218,7 +218,7 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
// Verify
// State handle created and discarded
assertEquals(1, stateHandleProvider.getStateHandles().size());
- assertEquals(state, stateHandleProvider.getStateHandles().get(0).getState(null));
+ assertEquals(state, stateHandleProvider.getStateHandles().get(0).retrieveState());
assertEquals(1, stateHandleProvider.getStateHandles().get(0).getNumberOfDiscardCalls());
}
@@ -245,8 +245,8 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
// Verify
// State handles created
assertEquals(2, stateHandleProvider.getStateHandles().size());
- assertEquals(initialState, stateHandleProvider.getStateHandles().get(0).getState(null));
- assertEquals(replaceState, stateHandleProvider.getStateHandles().get(1).getState(null));
+ assertEquals(initialState, stateHandleProvider.getStateHandles().get(0).retrieveState());
+ assertEquals(replaceState, stateHandleProvider.getStateHandles().get(1).retrieveState());
// Path created and is persistent
Stat stat = ZooKeeper.getClient().checkExists().forPath(pathInZooKeeper);
@@ -255,9 +255,9 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
// Data is equal
@SuppressWarnings("unchecked")
- Long actual = ((StateHandle<Long>) InstantiationUtil.deserializeObject(
+ Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject(
ZooKeeper.getClient().getData().forPath(pathInZooKeeper),
- ClassLoader.getSystemClassLoader())).getState(null);
+ ClassLoader.getSystemClassLoader())).retrieveState();
assertEquals(replaceState, actual);
}
@@ -267,7 +267,7 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
*/
@Test(expected = Exception.class)
public void testReplaceNonExistingPath() throws Exception {
- StateStorageHelper<Long> stateStorage = new LongStateStorage();
+ RetrievableStateStorageHelper<Long> stateStorage = new LongStateStorage();
ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
ZooKeeper.getClient(), stateStorage);
@@ -307,15 +307,15 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
// Verify
// State handle created and discarded
assertEquals(2, stateHandleProvider.getStateHandles().size());
- assertEquals(initialState, stateHandleProvider.getStateHandles().get(0).getState(null));
- assertEquals(replaceState, stateHandleProvider.getStateHandles().get(1).getState(null));
+ assertEquals(initialState, stateHandleProvider.getStateHandles().get(0).retrieveState());
+ assertEquals(replaceState, stateHandleProvider.getStateHandles().get(1).retrieveState());
assertEquals(1, stateHandleProvider.getStateHandles().get(1).getNumberOfDiscardCalls());
// Initial value
@SuppressWarnings("unchecked")
- Long actual = ((StateHandle<Long>) InstantiationUtil.deserializeObject(
+ Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject(
ZooKeeper.getClient().getData().forPath(pathInZooKeeper),
- ClassLoader.getSystemClassLoader())).getState(null);
+ ClassLoader.getSystemClassLoader())).retrieveState();
assertEquals(initialState, actual);
}
@@ -339,10 +339,10 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
assertEquals(-1, store.exists(pathInZooKeeper));
store.add(pathInZooKeeper, state);
- StateHandle<Long> actual = store.get(pathInZooKeeper);
+ RetrievableStateHandle<Long> actual = store.get(pathInZooKeeper);
// Verify
- assertEquals(state, actual.getState(null));
+ assertEquals(state, actual.retrieveState());
assertTrue(store.exists(pathInZooKeeper) >= 0);
}
@@ -384,8 +384,8 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
store.add(pathInZooKeeper, val, CreateMode.PERSISTENT_SEQUENTIAL);
}
- for (Tuple2<StateHandle<Long>, String> val : store.getAll()) {
- assertTrue(expected.remove(val.f0.getState(null)));
+ for (Tuple2<RetrievableStateHandle<Long>, String> val : store.getAll()) {
+ assertTrue(expected.remove(val.f0.retrieveState()));
}
assertEquals(0, expected.size());
}
@@ -412,11 +412,11 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
store.add(pathInZooKeeper, val, CreateMode.PERSISTENT_SEQUENTIAL);
}
- List<Tuple2<StateHandle<Long>, String>> actual = store.getAllSortedByName();
+ List<Tuple2<RetrievableStateHandle<Long>, String>> actual = store.getAllSortedByName();
assertEquals(expected.length, actual.size());
for (int i = 0; i < expected.length; i++) {
- assertEquals(expected[i], actual.get(i).f0.getState(null));
+ assertEquals(expected[i], actual.get(i).f0.retrieveState());
}
}
@@ -540,24 +540,24 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
// Simple test helpers
// ---------------------------------------------------------------------------------------------
- private static class LongStateStorage implements StateStorageHelper<Long> {
+ private static class LongStateStorage implements RetrievableStateStorageHelper<Long> {
- private final List<LongStateHandle> stateHandles = new ArrayList<>();
+ private final List<LongRetrievableStateHandle> stateHandles = new ArrayList<>();
@Override
- public StateHandle<Long> store(Long state) throws Exception {
- LongStateHandle stateHandle = new LongStateHandle(state);
+ public RetrievableStateHandle<Long> store(Long state) throws Exception {
+ LongRetrievableStateHandle stateHandle = new LongRetrievableStateHandle(state);
stateHandles.add(stateHandle);
return stateHandle;
}
- List<LongStateHandle> getStateHandles() {
+ List<LongRetrievableStateHandle> getStateHandles() {
return stateHandles;
}
}
- private static class LongStateHandle implements StateHandle<Long> {
+ private static class LongRetrievableStateHandle implements RetrievableStateHandle<Long> {
private static final long serialVersionUID = -3555329254423838912L;
@@ -565,12 +565,12 @@ public class ZooKeeperStateHandleStoreITCase extends TestLogger {
private int numberOfDiscardCalls;
- public LongStateHandle(Long state) {
+ public LongRetrievableStateHandle(Long state) {
this.state = state;
}
@Override
- public Long getState(ClassLoader ignored) throws Exception {
+ public Long retrieveState() throws Exception {
return state;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-connectors/flink-connector-cassandra/src/test/java/org/apache/flink/streaming/connectors/cassandra/CassandraConnectorITCase.java
----------------------------------------------------------------------
diff --git a/flink-streaming-connectors/flink-connector-cassandra/src/test/java/org/apache/flink/streaming/connectors/cassandra/CassandraConnectorITCase.java b/flink-streaming-connectors/flink-connector-cassandra/src/test/java/org/apache/flink/streaming/connectors/cassandra/CassandraConnectorITCase.java
index 9388818..4822f91 100644
--- a/flink-streaming-connectors/flink-connector-cassandra/src/test/java/org/apache/flink/streaming/connectors/cassandra/CassandraConnectorITCase.java
+++ b/flink-streaming-connectors/flink-connector-cassandra/src/test/java/org/apache/flink/streaming/connectors/cassandra/CassandraConnectorITCase.java
@@ -6,9 +6,9 @@
* 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.
@@ -46,11 +46,10 @@ import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.runtime.operators.WriteAheadSinkTestBase;
-import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
-import org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness;
+
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.TestStreamEnvironment;
import org.apache.flink.test.util.ForkableFlinkMiniCluster;
-
import org.apache.flink.test.util.TestEnvironment;
import org.junit.After;
import org.junit.AfterClass;
@@ -276,8 +275,7 @@ public class CassandraConnectorITCase extends WriteAheadSinkTestBase<Tuple3<Stri
@Override
protected void verifyResultsIdealCircumstances(
- OneInputStreamTaskTestHarness<Tuple3<String, Integer, Integer>, Tuple3<String, Integer, Integer>> harness,
- OneInputStreamTask<Tuple3<String, Integer, Integer>, Tuple3<String, Integer, Integer>> task,
+ OneInputStreamOperatorTestHarness<Tuple3<String, Integer, Integer>, Tuple3<String, Integer, Integer>> harness,
CassandraTupleWriteAheadSink<Tuple3<String, Integer, Integer>> sink) {
ResultSet result = session.execute(SELECT_DATA_QUERY);
@@ -294,8 +292,7 @@ public class CassandraConnectorITCase extends WriteAheadSinkTestBase<Tuple3<Stri
@Override
protected void verifyResultsDataPersistenceUponMissedNotify(
- OneInputStreamTaskTestHarness<Tuple3<String, Integer, Integer>, Tuple3<String, Integer, Integer>> harness,
- OneInputStreamTask<Tuple3<String, Integer, Integer>, Tuple3<String, Integer, Integer>> task,
+ OneInputStreamOperatorTestHarness<Tuple3<String, Integer, Integer>, Tuple3<String, Integer, Integer>> harness,
CassandraTupleWriteAheadSink<Tuple3<String, Integer, Integer>> sink) {
ResultSet result = session.execute(SELECT_DATA_QUERY);
@@ -312,8 +309,7 @@ public class CassandraConnectorITCase extends WriteAheadSinkTestBase<Tuple3<Stri
@Override
protected void verifyResultsDataDiscardingUponRestore(
- OneInputStreamTaskTestHarness<Tuple3<String, Integer, Integer>, Tuple3<String, Integer, Integer>> harness,
- OneInputStreamTask<Tuple3<String, Integer, Integer>, Tuple3<String, Integer, Integer>> task,
+ OneInputStreamOperatorTestHarness<Tuple3<String, Integer, Integer>, Tuple3<String, Integer, Integer>> harness,
CassandraTupleWriteAheadSink<Tuple3<String, Integer, Integer>> sink) {
ResultSet result = session.execute(SELECT_DATA_QUERY);
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-connectors/flink-connector-filesystem/src/test/java/org/apache/flink/streaming/connectors/fs/bucketing/BucketingSinkTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-connectors/flink-connector-filesystem/src/test/java/org/apache/flink/streaming/connectors/fs/bucketing/BucketingSinkTest.java b/flink-streaming-connectors/flink-connector-filesystem/src/test/java/org/apache/flink/streaming/connectors/fs/bucketing/BucketingSinkTest.java
index 828dcc9..e274fdd 100644
--- a/flink-streaming-connectors/flink-connector-filesystem/src/test/java/org/apache/flink/streaming/connectors/fs/bucketing/BucketingSinkTest.java
+++ b/flink-streaming-connectors/flink-connector-filesystem/src/test/java/org/apache/flink/streaming/connectors/fs/bucketing/BucketingSinkTest.java
@@ -28,13 +28,13 @@ import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.streaming.api.operators.StreamSink;
import org.apache.flink.streaming.connectors.fs.AvroKeyValueSinkWriter;
import org.apache.flink.streaming.connectors.fs.Clock;
import org.apache.flink.streaming.connectors.fs.SequenceFileWriter;
import org.apache.flink.streaming.connectors.fs.StringWriter;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import org.apache.flink.streaming.runtime.tasks.TestTimeServiceProvider;
import org.apache.flink.streaming.runtime.tasks.TimeServiceProvider;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
@@ -137,11 +137,11 @@ public class BucketingSinkTest {
// snapshot but don't call notify to simulate a notify that never
// arrives, the sink should move pending files in restore() in that case
- StreamTaskState snapshot1 = testHarness.snapshot(0, 0);
+ StreamStateHandle snapshot1 = testHarness.snapshot(0, 0);
testHarness = createTestSink(dataDir, clock);
testHarness.setup();
- testHarness.restore(snapshot1, 1);
+ testHarness.restore(snapshot1);
testHarness.open();
testHarness.processElement(new StreamRecord<>("Hello"));
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java
index fda5efd..6d67560 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/ContinuousFileReaderOperator.java
@@ -24,11 +24,10 @@ import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.core.fs.FileInputSplit;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
import org.apache.flink.metrics.Counter;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.OutputTypeConfigurable;
@@ -36,13 +35,11 @@ import org.apache.flink.streaming.api.operators.TimestampedCollector;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import org.apache.flink.util.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
-import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
@@ -376,14 +373,10 @@ public class ContinuousFileReaderOperator<OUT, S extends Serializable> extends A
// --------------------- Checkpointing --------------------------
@Override
- public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
- StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
-
- final AbstractStateBackend.CheckpointStateOutputStream os =
- this.getStateBackend().createCheckpointStateOutputStream(checkpointId, timestamp);
+ public void snapshotState(FSDataOutputStream os, long checkpointId, long timestamp) throws Exception {
+ super.snapshotState(os, checkpointId, timestamp);
final ObjectOutputStream oos = new ObjectOutputStream(os);
- final AbstractStateBackend.CheckpointStateOutputView ov = new AbstractStateBackend.CheckpointStateOutputView(os);
Tuple3<List<FileInputSplit>, FileInputSplit, S> readerState = this.reader.getReaderState();
List<FileInputSplit> pendingSplits = readerState.f0;
@@ -392,35 +385,28 @@ public class ContinuousFileReaderOperator<OUT, S extends Serializable> extends A
// write the current split
oos.writeObject(currSplit);
-
- // write the pending ones
- ov.writeInt(pendingSplits.size());
+ oos.writeInt(pendingSplits.size());
for (FileInputSplit split : pendingSplits) {
oos.writeObject(split);
}
// write the state of the reading channel
oos.writeObject(formatState);
- taskState.setOperatorState(os.closeAndGetHandle());
- return taskState;
+ oos.flush();
}
@Override
- public void restoreState(StreamTaskState state) throws Exception {
- super.restoreState(state);
-
- StreamStateHandle stream = (StreamStateHandle) state.getOperatorState();
+ public void restoreState(FSDataInputStream is) throws Exception {
+ super.restoreState(is);
- final InputStream is = stream.getState(getUserCodeClassloader());
final ObjectInputStream ois = new ObjectInputStream(is);
- final DataInputViewStreamWrapper div = new DataInputViewStreamWrapper(is);
// read the split that was being read
FileInputSplit currSplit = (FileInputSplit) ois.readObject();
// read the pending splits list
List<FileInputSplit> pendingSplits = new LinkedList<>();
- int noOfSplits = div.readInt();
+ int noOfSplits = ois.readInt();
for (int i = 0; i < noOfSplits; i++) {
FileInputSplit split = (FileInputSplit) ois.readObject();
pendingSplits.add(split);
@@ -435,6 +421,5 @@ public class ContinuousFileReaderOperator<OUT, S extends Serializable> extends A
"The reader state has already been initialized.");
this.readerState = new Tuple3<>(pendingSplits, currSplit, formatState);
- div.close();
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
index 15bb384..59ecd15 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
@@ -25,6 +25,9 @@ import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.metrics.Counter;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
+import org.apache.flink.runtime.state.AsynchronousKvStateSnapshot;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.runtime.state.KvStateSnapshot;
@@ -35,11 +38,14 @@ import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
+import org.apache.flink.util.InstantiationUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
import java.util.HashMap;
+import java.util.Set;
import java.util.concurrent.ScheduledFuture;
/**
@@ -91,7 +97,7 @@ public abstract class AbstractStreamOperator<OUT>
private transient KeySelector<?, ?> stateKeySelector2;
/** The state backend that stores the state and checkpoints for this task */
- private AbstractStateBackend stateBackend = null;
+ private transient AbstractStateBackend stateBackend;
protected MetricGroup metrics;
// ------------------------------------------------------------------------
@@ -164,7 +170,7 @@ public abstract class AbstractStreamOperator<OUT>
if (stateBackend != null) {
try {
stateBackend.close();
- stateBackend.dispose();
+ stateBackend.discardState();
} catch (Exception e) {
throw new RuntimeException("Error while closing/disposing state backend.", e);
}
@@ -176,30 +182,45 @@ public abstract class AbstractStreamOperator<OUT>
// ------------------------------------------------------------------------
@Override
- public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
- // here, we deal with key/value state snapshots
-
- StreamTaskState state = new StreamTaskState();
-
- if (stateBackend != null) {
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> partitionedSnapshots =
- stateBackend.snapshotPartitionedState(checkpointId, timestamp);
- if (partitionedSnapshots != null) {
- state.setKvStates(partitionedSnapshots);
+ public void snapshotState(FSDataOutputStream out,
+ long checkpointId,
+ long timestamp) throws Exception {
+
+ HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> keyedState =
+ stateBackend.snapshotPartitionedState(checkpointId,timestamp);
+
+ // Materialize asynchronous snapshots, if any
+ if (keyedState != null) {
+ Set<String> keys = keyedState.keySet();
+ for (String key: keys) {
+ if (keyedState.get(key) instanceof AsynchronousKvStateSnapshot) {
+ AsynchronousKvStateSnapshot<?, ?, ?, ?, ?> asyncHandle = (AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) keyedState.get(key);
+ keyedState.put(key, asyncHandle.materialize());
+ }
}
}
- return state;
+ byte[] serializedSnapshot = InstantiationUtil.serializeObject(keyedState);
+
+ DataOutputStream dos = new DataOutputStream(out);
+ dos.writeInt(serializedSnapshot.length);
+ dos.write(serializedSnapshot);
+
+ dos.flush();
+
}
-
+
@Override
- @SuppressWarnings("rawtypes,unchecked")
- public void restoreState(StreamTaskState state) throws Exception {
- // restore the key/value state. the actual restore happens lazily, when the function requests
- // the state again, because the restore method needs information provided by the user function
- if (stateBackend != null) {
- stateBackend.injectKeyValueStateSnapshots((HashMap)state.getKvStates());
- }
+ public void restoreState(FSDataInputStream in) throws Exception {
+ DataInputStream dis = new DataInputStream(in);
+ int size = dis.readInt();
+ byte[] serializedSnapshot = new byte[size];
+ dis.readFully(serializedSnapshot);
+
+ HashMap<String, KvStateSnapshot> keyedState =
+ InstantiationUtil.deserializeObject(serializedSnapshot, getUserCodeClassloader());
+
+ stateBackend.injectKeyValueStateSnapshots(keyedState);
}
@Override
@@ -249,8 +270,8 @@ public abstract class AbstractStreamOperator<OUT>
}
/**
- * Register a timer callback. At the specified time the provided {@link Triggerable} will
- * be invoked. This call is guaranteed to not happen concurrently with method calls on the operator.
+ * Register a timer callback. At the specified time the {@link Triggerable} will be invoked.
+ * This call is guaranteed to not happen concurrently with method calls on the operator.
*
* @param time The absolute time in milliseconds.
* @param target The target to be triggered.
@@ -281,11 +302,18 @@ public abstract class AbstractStreamOperator<OUT>
*/
@SuppressWarnings("unchecked")
protected <S extends State, N> S getPartitionedState(N namespace, TypeSerializer<N> namespaceSerializer, StateDescriptor<S, ?> stateDescriptor) throws Exception {
- return getStateBackend().getPartitionedState(namespace, (TypeSerializer<Object>) namespaceSerializer,
- stateDescriptor);
+ if (stateBackend != null) {
+ return stateBackend.getPartitionedState(
+ namespace,
+ namespaceSerializer,
+ stateDescriptor);
+ } else {
+ throw new RuntimeException("Cannot create partitioned state. The key grouped state " +
+ "backend has not been set. This indicates that the operator is not " +
+ "partitioned/keyed.");
+ }
}
-
@Override
@SuppressWarnings({"unchecked", "rawtypes"})
public void setKeyContextElement1(StreamRecord record) throws Exception {
@@ -300,17 +328,25 @@ public abstract class AbstractStreamOperator<OUT>
public void setKeyContextElement2(StreamRecord record) throws Exception {
if (stateKeySelector2 != null) {
Object key = ((KeySelector) stateKeySelector2).getKey(record.getValue());
- getStateBackend().setCurrentKey(key);
+
+ setKeyContext(key);
}
}
@SuppressWarnings({"unchecked", "rawtypes"})
public void setKeyContext(Object key) {
- if (stateKeySelector1 != null) {
- stateBackend.setCurrentKey(key);
+ if (stateBackend != null) {
+ try {
+ stateBackend.setCurrentKey(key);
+ } catch (Exception e) {
+ throw new RuntimeException("Exception occurred while setting the current key context.", e);
+ }
+ } else {
+ throw new RuntimeException("Could not set the current key context, because the " +
+ "AbstractStateBackend has not been initialized.");
}
}
-
+
// ------------------------------------------------------------------------
// Context and chaining properties
// ------------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java
index 1ddd934..b1bc531 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java
@@ -18,6 +18,8 @@
package org.apache.flink.streaming.api.operators;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
import java.io.Serializable;
import org.apache.flink.annotation.PublicEvolving;
@@ -26,14 +28,13 @@ import org.apache.flink.api.common.functions.Function;
import org.apache.flink.api.common.functions.util.FunctionUtils;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.Configuration;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.runtime.state.CheckpointListener;
import org.apache.flink.streaming.api.checkpoint.Checkpointed;
import org.apache.flink.streaming.api.graph.StreamConfig;
-import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import static java.util.Objects.requireNonNull;
@@ -117,8 +118,8 @@ public abstract class AbstractUdfStreamOperator<OUT, F extends Function> extends
// ------------------------------------------------------------------------
@Override
- public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
- StreamTaskState state = super.snapshotOperatorState(checkpointId, timestamp);
+ public void snapshotState(FSDataOutputStream out, long checkpointId, long timestamp) throws Exception {
+ super.snapshotState(out, checkpointId, timestamp);
if (userFunction instanceof Checkpointed) {
@SuppressWarnings("unchecked")
@@ -127,45 +128,39 @@ public abstract class AbstractUdfStreamOperator<OUT, F extends Function> extends
Serializable udfState;
try {
udfState = chkFunction.snapshotState(checkpointId, timestamp);
- }
- catch (Exception e) {
- throw new Exception("Failed to draw state snapshot from function: " + e.getMessage(), e);
- }
-
- if (udfState != null) {
- try {
- AbstractStateBackend stateBackend = getStateBackend();
- StateHandle<Serializable> handle =
- stateBackend.checkpointStateSerializable(udfState, checkpointId, timestamp);
- state.setFunctionState(handle);
- }
- catch (Exception e) {
- throw new Exception("Failed to add the state snapshot of the function to the checkpoint: "
- + e.getMessage(), e);
+ if (udfState != null) {
+ out.write(1);
+ ObjectOutputStream os = new ObjectOutputStream(out);
+ os.writeObject(udfState);
+ os.flush();
+ } else {
+ out.write(0);
}
+ } catch (Exception e) {
+ throw new Exception("Failed to draw state snapshot from function: " + e.getMessage(), e);
}
}
-
- return state;
}
@Override
- public void restoreState(StreamTaskState state) throws Exception {
- super.restoreState(state);
-
- StateHandle<Serializable> stateHandle = state.getFunctionState();
-
- if (userFunction instanceof Checkpointed && stateHandle != null) {
+ public void restoreState(FSDataInputStream in) throws Exception {
+ super.restoreState(in);
+
+ if (userFunction instanceof Checkpointed) {
@SuppressWarnings("unchecked")
Checkpointed<Serializable> chkFunction = (Checkpointed<Serializable>) userFunction;
-
- Serializable functionState = stateHandle.getState(getUserCodeClassloader());
- if (functionState != null) {
- try {
- chkFunction.restoreState(functionState);
- }
- catch (Exception e) {
- throw new Exception("Failed to restore state to function: " + e.getMessage(), e);
+
+ int hasUdfState = in.read();
+
+ if (hasUdfState == 1) {
+ ObjectInputStream ois = new ObjectInputStream(in);
+ Serializable functionState = (Serializable) ois.readObject();
+ if (functionState != null) {
+ try {
+ chkFunction.restoreState(functionState);
+ } catch (Exception e) {
+ throw new Exception("Failed to restore state to function: " + e.getMessage(), e);
+ }
}
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java
index 3e38165..3411a60 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java
@@ -21,10 +21,11 @@ import java.io.Serializable;
import org.apache.flink.annotation.PublicEvolving;
import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
/**
* Basic interface for stream operators. Implementers would implement one of
@@ -94,17 +95,15 @@ public interface StreamOperator<OUT> extends Serializable {
* (if the operator is stateful) and the key/value state (if it is being used and has been
* initialized).
*
+ * @param out The stream to which we have to write our state.
* @param checkpointId The ID of the checkpoint.
* @param timestamp The timestamp of the checkpoint.
*
- * @return The StreamTaskState object, possibly containing the snapshots for the
- * operator and key/value state.
- *
* @throws Exception Forwards exceptions that occur while drawing snapshots from the operator
* and the key/value state.
*/
- StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception;
-
+ void snapshotState(FSDataOutputStream out, long checkpointId, long timestamp) throws Exception;
+
/**
* Restores the operator state, if this operator's execution is recovering from a checkpoint.
* This method restores the operator state (if the operator is stateful) and the key/value state
@@ -113,13 +112,12 @@ public interface StreamOperator<OUT> extends Serializable {
* <p>This method is called after {@link #setup(StreamTask, StreamConfig, Output)}
* and before {@link #open()}.
*
- * @param state The state of operator that was snapshotted as part of checkpoint
- * from which the execution is restored.
+ * @param in The stream from which we have to restore our state.
*
* @throws Exception Exceptions during state restore should be forwarded, so that the system can
* properly react to failed state restore and fail the execution attempt.
*/
- void restoreState(StreamTaskState state) throws Exception;
+ void restoreState(FSDataInputStream in) throws Exception;
/**
* Called when the checkpoint with the given ID is completed and acknowledged on the JobManager.
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java
index 34c62fb..8d074cc 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java
@@ -19,17 +19,19 @@ package org.apache.flink.streaming.runtime.operators;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.core.memory.DataInputView;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.io.disk.InputViewIterator;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.util.ReusingMutableToRegularIteratorWrapper;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
-import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.InstantiationUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -54,7 +56,7 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
protected static final Logger LOG = LoggerFactory.getLogger(GenericWriteAheadSink.class);
private final CheckpointCommitter committer;
- private transient AbstractStateBackend.CheckpointStateOutputView out;
+ private transient AbstractStateBackend.CheckpointStateOutputStream out;
protected final TypeSerializer<IN> serializer;
private final String id;
@@ -89,7 +91,7 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
private void saveHandleInState(final long checkpointId, final long timestamp) throws Exception {
//only add handle if a new OperatorState was created since the last snapshot
if (out != null) {
- StateHandle<DataInputView> handle = out.closeAndGetHandle();
+ StreamStateHandle handle = out.closeAndGetHandle();
if (state.pendingHandles.containsKey(checkpointId)) {
//we already have a checkpoint stored for that ID that may have been partially written,
//so we discard this "alternate version" and use the stored checkpoint
@@ -102,18 +104,21 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
}
@Override
- public StreamTaskState snapshotOperatorState(final long checkpointId, final long timestamp) throws Exception {
- StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
+ public void snapshotState(FSDataOutputStream out,
+ long checkpointId,
+ long timestamp) throws Exception {
+ super.snapshotState(out, checkpointId, timestamp);
+
saveHandleInState(checkpointId, timestamp);
- taskState.setFunctionState(state);
- return taskState;
+
+ InstantiationUtil.serializeObject(out, state);
}
@Override
- public void restoreState(StreamTaskState state) throws Exception {
- super.restoreState(state);
- this.state = (ExactlyOnceState) state.getFunctionState();
- out = null;
+ public void restoreState(FSDataInputStream in) throws Exception {
+ super.restoreState(in);
+
+ this.state = InstantiationUtil.deserializeObject(in, getUserCodeClassloader());
}
private void cleanState() throws Exception {
@@ -142,9 +147,9 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
if (pastCheckpointId <= checkpointId) {
try {
if (!committer.isCheckpointCommitted(pastCheckpointId)) {
- Tuple2<Long, StateHandle<DataInputView>> handle = state.pendingHandles.get(pastCheckpointId);
- DataInputView in = handle.f1.getState(getUserCodeClassloader());
- boolean success = sendValues(new ReusingMutableToRegularIteratorWrapper<>(new InputViewIterator<>(in, serializer), serializer), handle.f0);
+ Tuple2<Long, StreamStateHandle> handle = state.pendingHandles.get(pastCheckpointId);
+ FSDataInputStream in = handle.f1.openInputStream();
+ boolean success = sendValues(new ReusingMutableToRegularIteratorWrapper<>(new InputViewIterator<>(new DataInputViewStreamWrapper(in), serializer), serializer), handle.f0);
if (success) { //if the sending has failed we will retry on the next notify
committer.commitCheckpoint(pastCheckpointId);
checkpointsToRemove.add(pastCheckpointId);
@@ -159,7 +164,7 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
}
}
for (Long toRemove : checkpointsToRemove) {
- Tuple2<Long, StateHandle<DataInputView>> handle = state.pendingHandles.get(toRemove);
+ Tuple2<Long, StreamStateHandle> handle = state.pendingHandles.get(toRemove);
state.pendingHandles.remove(toRemove);
handle.f1.discardState();
}
@@ -181,9 +186,9 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
IN value = element.getValue();
//generate initial operator state
if (out == null) {
- out = getStateBackend().createCheckpointStateOutputView(0, 0);
+ out = getStateBackend().createCheckpointStateOutputStream(0, 0);
}
- serializer.serialize(value, out);
+ serializer.serialize(value, new DataOutputViewStreamWrapper(out));
}
@Override
@@ -195,59 +200,21 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
* This state is used to keep a list of all StateHandles (essentially references to past OperatorStates) that were
* used since the last completed checkpoint.
**/
- public static class ExactlyOnceState implements StateHandle<Serializable> {
+ public static class ExactlyOnceState implements Serializable {
private static final long serialVersionUID = -3571063495273460743L;
- protected TreeMap<Long, Tuple2<Long, StateHandle<DataInputView>>> pendingHandles;
+ protected TreeMap<Long, Tuple2<Long, StreamStateHandle>> pendingHandles;
public ExactlyOnceState() {
pendingHandles = new TreeMap<>();
}
- @Override
- public TreeMap<Long, Tuple2<Long, StateHandle<DataInputView>>> getState(ClassLoader userCodeClassLoader) throws Exception {
+ public TreeMap<Long, Tuple2<Long, StreamStateHandle>> getState(ClassLoader userCodeClassLoader) throws Exception {
return pendingHandles;
}
@Override
- public void discardState() throws Exception {
- //we specifically want the state to survive failed jobs, so we don't discard anything
- }
-
- @Override
- public long getStateSize() throws Exception {
- int stateSize = 0;
- for (Tuple2<Long, StateHandle<DataInputView>> pair : pendingHandles.values()) {
- stateSize += pair.f1.getStateSize();
- }
- return stateSize;
- }
-
- @Override
- public void close() throws IOException {
- Throwable exception = null;
-
- for (Tuple2<Long, StateHandle<DataInputView>> pair : pendingHandles.values()) {
- StateHandle<DataInputView> handle = pair.f1;
- if (handle != null) {
- try {
- handle.close();
- }
- catch (Throwable t) {
- if (exception != null) {
- exception = t;
- }
- }
- }
- }
-
- if (exception != null) {
- ExceptionUtils.rethrowIOException(exception);
- }
- }
-
- @Override
public String toString() {
return this.pendingHandles.toString();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java
index fdc8117..2c95099 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java
@@ -24,18 +24,18 @@ import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.functions.Function;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.functions.KeySelector;
-import org.apache.flink.core.memory.DataInputView;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.util.MathUtils;
import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator;
import org.apache.flink.streaming.api.operators.TimestampedCollector;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
-import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import static java.util.Objects.requireNonNull;
@@ -244,36 +244,35 @@ public abstract class AbstractAlignedProcessingTimeWindowOperator<KEY, IN, OUT,
// ------------------------------------------------------------------------
@Override
- public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
- StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
-
+ public void snapshotState(FSDataOutputStream out, long checkpointId, long timestamp) throws Exception {
+ super.snapshotState(out, checkpointId, timestamp);
+
// we write the panes with the key/value maps into the stream, as well as when this state
// should have triggered and slided
- AbstractStateBackend.CheckpointStateOutputView out =
- getStateBackend().createCheckpointStateOutputView(checkpointId, timestamp);
- out.writeLong(nextEvaluationTime);
- out.writeLong(nextSlideTime);
- panes.writeToOutput(out, keySerializer, stateTypeSerializer);
-
- taskState.setOperatorState(out.closeAndGetHandle());
- return taskState;
+ DataOutputViewStreamWrapper outView = new DataOutputViewStreamWrapper(out);
+
+ outView.writeLong(nextEvaluationTime);
+ outView.writeLong(nextSlideTime);
+
+ panes.writeToOutput(outView, keySerializer, stateTypeSerializer);
+
+ outView.flush();
}
@Override
- public void restoreState(StreamTaskState taskState) throws Exception {
- super.restoreState(taskState);
+ public void restoreState(FSDataInputStream in) throws Exception {
+ super.restoreState(in);
- @SuppressWarnings("unchecked")
- StateHandle<DataInputView> inputState = (StateHandle<DataInputView>) taskState.getOperatorState();
- DataInputView in = inputState.getState(getUserCodeClassloader());
-
- final long nextEvaluationTime = in.readLong();
- final long nextSlideTime = in.readLong();
+ DataInputViewStreamWrapper inView = new DataInputViewStreamWrapper(in);
+
+ final long nextEvaluationTime = inView.readLong();
+ final long nextSlideTime = inView.readLong();
AbstractKeyedTimePanes<IN, KEY, STATE, OUT> panes = createPanes(keySelector, function);
- panes.readFromInput(in, keySerializer, stateTypeSerializer);
-
+
+ panes.readFromInput(inView, keySerializer, stateTypeSerializer);
+
restoredState = new RestoredState<>(panes, nextEvaluationTime, nextSlideTime);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java
index 12ed60e..dbdd660 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java
@@ -39,11 +39,13 @@ import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.typeutils.InputTypeConfigurable;
import org.apache.flink.api.java.typeutils.TypeExtractor;
import org.apache.flink.api.java.typeutils.runtime.TupleSerializer;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.core.memory.DataInputView;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
import org.apache.flink.core.memory.DataOutputView;
import org.apache.flink.metrics.MetricGroup;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator;
@@ -59,7 +61,6 @@ import org.apache.flink.streaming.api.windowing.windows.Window;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.operators.windowing.functions.InternalWindowFunction;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import org.apache.flink.util.Preconditions;
import java.io.IOException;
@@ -857,7 +858,7 @@ public class WindowOperator<K, IN, ACC, OUT, W extends Window>
@Override
@SuppressWarnings("unchecked")
- public StreamTaskState snapshotOperatorState(long checkpointId, long timestamp) throws Exception {
+ public void snapshotState(FSDataOutputStream out, long checkpointId, long timestamp) throws Exception {
if (mergingWindowsByKey != null) {
TupleSerializer<Tuple2<W, W>> tupleSerializer = new TupleSerializer<>((Class) Tuple2.class, new TypeSerializer[] {windowSerializer, windowSerializer} );
@@ -870,29 +871,16 @@ public class WindowOperator<K, IN, ACC, OUT, W extends Window>
}
}
- StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
+ snapshotTimers(new DataOutputViewStreamWrapper(out));
- AbstractStateBackend.CheckpointStateOutputView out =
- getStateBackend().createCheckpointStateOutputView(checkpointId, timestamp);
-
- snapshotTimers(out);
-
- taskState.setOperatorState(out.closeAndGetHandle());
-
- return taskState;
+ super.snapshotState(out, checkpointId, timestamp);
}
@Override
- public void restoreState(StreamTaskState taskState) throws Exception {
- super.restoreState(taskState);
-
- final ClassLoader userClassloader = getUserCodeClassloader();
-
- @SuppressWarnings("unchecked")
- StateHandle<DataInputView> inputState = (StateHandle<DataInputView>) taskState.getOperatorState();
- DataInputView in = inputState.getState(userClassloader);
+ public void restoreState(FSDataInputStream in) throws Exception {
+ restoreTimers(new DataInputViewStreamWrapper(in));
- restoreTimers(in);
+ super.restoreState(in);
}
private void restoreTimers(DataInputView in ) throws IOException {
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
index 5bcf41b..108a3ae 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
@@ -21,6 +21,7 @@ import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.runtime.plugable.SerializationDelegate;
+import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.util.Preconditions;
@@ -60,8 +61,10 @@ public class KeyGroupStreamPartitioner<T, K> extends StreamPartitioner<T> implem
} catch (Exception e) {
throw new RuntimeException("Could not extract key from " + record.getInstance().getValue(), e);
}
- returnArray[0] = keyGroupAssigner.getKeyGroupIndex(key) % numberOfOutputChannels;
-
+ returnArray[0] = KeyGroupRange.computeOperatorIndexForKeyGroup(
+ keyGroupAssigner.getNumberKeyGroups(),
+ numberOfOutputChannels,
+ keyGroupAssigner.getKeyGroupIndex(key));
return returnArray;
}
[09/27] flink git commit: [FLINK-3761] Refactor State Backends/Make
Keyed State Key-Group Aware
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapValueState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapValueState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapValueState.java
new file mode 100644
index 0000000..cccaacb
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapValueState.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.heap;
+
+import org.apache.flink.api.common.state.ValueState;
+import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.util.Preconditions;
+
+import java.util.Map;
+
+/**
+ * Heap-backed partitioned {@link org.apache.flink.api.common.state.ValueState} that is snapshotted
+ * into files.
+ *
+ * @param <K> The type of the key.
+ * @param <N> The type of the namespace.
+ * @param <V> The type of the value.
+ */
+public class HeapValueState<K, N, V>
+ extends AbstractHeapState<K, N, V, ValueState<V>, ValueStateDescriptor<V>>
+ implements ValueState<V> {
+
+ /**
+ * Creates a new key/value state for the given hash map of key/value pairs.
+ *
+ * @param backend The state backend backing that created this state.
+ * @param stateDesc The state identifier for the state. This contains name
+ * and can create a default state value.
+ * @param stateTable The state tab;e to use in this kev/value state. May contain initial state.
+ */
+ public HeapValueState(
+ KeyedStateBackend<K> backend,
+ ValueStateDescriptor<V> stateDesc,
+ StateTable<K, N, V> stateTable,
+ TypeSerializer<K> keySerializer,
+ TypeSerializer<N> namespaceSerializer) {
+ super(backend, stateDesc, stateTable, keySerializer, namespaceSerializer);
+ }
+
+ @Override
+ public V value() {
+ Preconditions.checkState(currentNamespace != null, "No namespace set.");
+ Preconditions.checkState(backend.getCurrentKey() != null, "No key set.");
+
+ Map<N, Map<K, V>> namespaceMap =
+ stateTable.get(backend.getCurrentKeyGroupIndex());
+
+ if (namespaceMap == null) {
+ return stateDesc.getDefaultValue();
+ }
+
+ Map<K, V> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ return stateDesc.getDefaultValue();
+ }
+
+ V result = keyedMap.get(backend.<K>getCurrentKey());
+
+ if (result == null) {
+ return stateDesc.getDefaultValue();
+ }
+
+ return result;
+ }
+
+ @Override
+ public void update(V value) {
+ Preconditions.checkState(currentNamespace != null, "No namespace set.");
+ Preconditions.checkState(backend.getCurrentKey() != null, "No key set.");
+
+ if (value == null) {
+ clear();
+ return;
+ }
+
+ Map<N, Map<K, V>> namespaceMap =
+ stateTable.get(backend.getCurrentKeyGroupIndex());
+
+ if (namespaceMap == null) {
+ namespaceMap = createNewMap();
+ stateTable.set(backend.getCurrentKeyGroupIndex(), namespaceMap);
+ }
+
+ Map<K, V> keyedMap = namespaceMap.get(currentNamespace);
+
+ if (keyedMap == null) {
+ keyedMap = createNewMap();
+ namespaceMap.put(currentNamespace, keyedMap);
+ }
+
+ keyedMap.put(backend.<K>getCurrentKey(), value);
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java
new file mode 100644
index 0000000..96e23d6
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/StateTable.java
@@ -0,0 +1,77 @@
+/*
+ * 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.heap;
+
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.runtime.state.KeyGroupRange;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+public class StateTable<K, N, ST> {
+
+ /** Serializer for the state value. The state value could be a List<V>, for example. */
+ protected final TypeSerializer<ST> stateSerializer;
+
+ /** The serializer for the namespace */
+ protected final TypeSerializer<N> namespaceSerializer;
+
+ /** Map for holding the actual state objects. */
+ private final List<Map<N, Map<K, ST>>> state;
+
+ protected final KeyGroupRange keyGroupRange;
+
+ public StateTable(
+ TypeSerializer<ST> stateSerializer,
+ TypeSerializer<N> namespaceSerializer,
+ KeyGroupRange keyGroupRange) {
+ this.stateSerializer = stateSerializer;
+ this.namespaceSerializer = namespaceSerializer;
+ this.keyGroupRange = keyGroupRange;
+
+ this.state = Arrays.asList((Map<N, Map<K, ST>>[]) new Map[keyGroupRange.getNumberOfKeyGroups()]);
+ }
+
+ private int indexToOffset(int index) {
+ return index - keyGroupRange.getStartKeyGroup();
+ }
+
+ public Map<N, Map<K, ST>> get(int index) {
+ return keyGroupRange.contains(index) ? state.get(indexToOffset(index)) : null;
+ }
+
+ public void set(int index, Map<N, Map<K, ST>> map) {
+ if (!keyGroupRange.contains(index)) {
+ throw new RuntimeException("Unexpected key group index. This indicates a bug.");
+ }
+ state.set(indexToOffset(index), map);
+ }
+
+ public TypeSerializer<ST> getStateSerializer() {
+ return stateSerializer;
+ }
+
+ public TypeSerializer<N> getNamespaceSerializer() {
+ return namespaceSerializer;
+ }
+
+ public List<Map<N, Map<K, ST>>> getState() {
+ return state;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/AbstractMemState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/AbstractMemState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/AbstractMemState.java
deleted file mode 100644
index cae673d..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/AbstractMemState.java
+++ /dev/null
@@ -1,82 +0,0 @@
-/*
- * 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.memory;
-
-import org.apache.flink.api.common.state.ListState;
-import org.apache.flink.api.common.state.State;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.runtime.state.AbstractHeapState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.runtime.util.DataOutputSerializer;
-
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Base class for partitioned {@link ListState} implementations that are backed by a regular
- * heap hash map. The concrete implementations define how the state is checkpointed.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <SV> The type of the values in the state.
- * @param <S> The type of State
- * @param <SD> The type of StateDescriptor for the State S
- */
-public abstract class AbstractMemState<K, N, SV, S extends State, SD extends StateDescriptor<S, ?>>
- extends AbstractHeapState<K, N, SV, S, SD, MemoryStateBackend> {
-
- public AbstractMemState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<SV> stateSerializer,
- SD stateDesc) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDesc);
- }
-
- public AbstractMemState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<SV> stateSerializer,
- SD stateDesc,
- HashMap<N, Map<K, SV>> state) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDesc, state);
- }
-
- public abstract KvStateSnapshot<K, N, S, SD, MemoryStateBackend> createHeapSnapshot(byte[] bytes);
-
- @Override
- public KvStateSnapshot<K, N, S, SD, MemoryStateBackend> snapshot(long checkpointId, long timestamp) throws Exception {
-
- DataOutputSerializer out = new DataOutputSerializer(Math.max(size() * 16, 16));
-
- out.writeInt(state.size());
- for (Map.Entry<N, Map<K, SV>> namespaceState: state.entrySet()) {
- N namespace = namespaceState.getKey();
- namespaceSerializer.serialize(namespace, out);
- out.writeInt(namespaceState.getValue().size());
- for (Map.Entry<K, SV> entry: namespaceState.getValue().entrySet()) {
- keySerializer.serialize(entry.getKey(), out);
- stateSerializer.serialize(entry.getValue(), out);
- }
- }
-
- byte[] bytes = out.getCopyOfBuffer();
-
- return createHeapSnapshot(bytes);
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/AbstractMemStateSnapshot.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/AbstractMemStateSnapshot.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/AbstractMemStateSnapshot.java
deleted file mode 100644
index e1b62d2..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/AbstractMemStateSnapshot.java
+++ /dev/null
@@ -1,144 +0,0 @@
-/*
- * 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.memory;
-
-import org.apache.flink.api.common.state.State;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.runtime.util.DataInputDeserializer;
-
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * A snapshot of a {@link MemValueState} for a checkpoint. The data is stored in a heap byte
- * array, in serialized form.
- *
- * @param <K> The type of the key in the snapshot state.
- * @param <N> The type of the namespace in the snapshot state.
- * @param <SV> The type of the value in the snapshot state.
- */
-public abstract class AbstractMemStateSnapshot<K, N, SV, S extends State, SD extends StateDescriptor<S, ?>>
- implements KvStateSnapshot<K, N, S, SD, MemoryStateBackend> {
-
- private static final long serialVersionUID = 1L;
-
- /** Key Serializer */
- protected final TypeSerializer<K> keySerializer;
-
- /** Namespace Serializer */
- protected final TypeSerializer<N> namespaceSerializer;
-
- /** Serializer for the state value */
- protected final TypeSerializer<SV> stateSerializer;
-
- /** StateDescriptor, for sanity checks */
- protected final SD stateDesc;
-
- /** The serialized data of the state key/value pairs */
- private final byte[] data;
-
- private transient boolean closed;
-
- /**
- * Creates a new heap memory state snapshot.
- *
- * @param keySerializer The serializer for the keys.
- * @param namespaceSerializer The serializer for the namespace.
- * @param stateSerializer The serializer for the elements in the state HashMap
- * @param stateDesc The state identifier
- * @param data The serialized data of the state key/value pairs
- */
- public AbstractMemStateSnapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<SV> stateSerializer,
- SD stateDesc,
- byte[] data) {
- this.keySerializer = keySerializer;
- this.namespaceSerializer = namespaceSerializer;
- this.stateSerializer = stateSerializer;
- this.stateDesc = stateDesc;
- this.data = data;
- }
-
- public abstract KvState<K, N, S, SD, MemoryStateBackend> createMemState(HashMap<N, Map<K, SV>> stateMap);
-
- @Override
- public KvState<K, N, S, SD, MemoryStateBackend> restoreState(
- MemoryStateBackend stateBackend,
- final TypeSerializer<K> keySerializer,
- ClassLoader classLoader) throws Exception {
-
- // validity checks
- if (!this.keySerializer.equals(keySerializer)) {
- throw new IllegalArgumentException(
- "Cannot restore the state from the snapshot with the given serializers. " +
- "State (K/V) was serialized with " +
- "(" + this.keySerializer + ") " +
- "now is (" + keySerializer + ")");
- }
-
- if (closed) {
- throw new IOException("snapshot has been closed");
- }
-
- // restore state
- DataInputDeserializer inView = new DataInputDeserializer(data, 0, data.length);
-
- final int numKeys = inView.readInt();
- HashMap<N, Map<K, SV>> stateMap = new HashMap<>(numKeys);
-
- for (int i = 0; i < numKeys && !closed; i++) {
- N namespace = namespaceSerializer.deserialize(inView);
- final int numValues = inView.readInt();
- Map<K, SV> namespaceMap = new HashMap<>(numValues);
- stateMap.put(namespace, namespaceMap);
- for (int j = 0; j < numValues; j++) {
- K key = keySerializer.deserialize(inView);
- SV value = stateSerializer.deserialize(inView);
- namespaceMap.put(key, value);
- }
- }
-
- if (closed) {
- throw new IOException("snapshot has been closed");
- }
-
- return createMemState(stateMap);
- }
-
- /**
- * Discarding the heap state is a no-op.
- */
- @Override
- public void discardState() {}
-
- @Override
- public long getStateSize() {
- return data.length;
- }
-
- @Override
- public void close() {
- closed = true;
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java
index a42bec2..b9ff255 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/ByteStreamStateHandle.java
@@ -51,7 +51,7 @@ public class ByteStreamStateHandle extends AbstractCloseableHandle implements St
}
@Override
- public FSDataInputStream openInputStream() throws Exception {
+ public FSDataInputStream openInputStream() throws IOException {
ensureNotClosed();
FSDataInputStream inputStream = new FSDataInputStream() {
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemCheckpointStreamFactory.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemCheckpointStreamFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemCheckpointStreamFactory.java
new file mode 100644
index 0000000..4801d85
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemCheckpointStreamFactory.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.runtime.state.memory;
+
+import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
+import java.io.IOException;
+
+/**
+ * {@link CheckpointStreamFactory} that produces streams that write to in-memory byte arrays.
+ */
+public class MemCheckpointStreamFactory implements CheckpointStreamFactory {
+
+ /** The maximal size that the snapshotted memory state may have */
+ private final int maxStateSize;
+
+ /**
+ * Creates a new in-memory stream factory that accepts states whose serialized forms are
+ * up to the given number of bytes.
+ *
+ * @param maxStateSize The maximal size of the serialized state
+ */
+ public MemCheckpointStreamFactory(int maxStateSize) {
+ this.maxStateSize = maxStateSize;
+ }
+
+ @Override
+ public void close() throws Exception {}
+
+ @Override
+ public CheckpointStateOutputStream createCheckpointStateOutputStream(
+ long checkpointID, long timestamp) throws Exception
+ {
+ return new MemoryCheckpointOutputStream(maxStateSize);
+ }
+
+ @Override
+ public String toString() {
+ return "In-Memory Stream Factory";
+ }
+
+ static void checkSize(int size, int maxSize) throws IOException {
+ if (size > maxSize) {
+ throw new IOException(
+ "Size of the state is larger than the maximum permitted memory-backed state. Size="
+ + size + " , maxSize=" + maxSize
+ + " . Consider using a different state backend, like the File System State backend.");
+ }
+ }
+
+
+
+ /**
+ * A {@code CheckpointStateOutputStream} that writes into a byte array.
+ */
+ public static final class MemoryCheckpointOutputStream extends CheckpointStateOutputStream {
+
+ private final ByteArrayOutputStreamWithPos os = new ByteArrayOutputStreamWithPos();
+
+ private final int maxSize;
+
+ private boolean closed;
+
+ boolean isEmpty = true;
+
+ public MemoryCheckpointOutputStream(int maxSize) {
+ this.maxSize = maxSize;
+ }
+
+ @Override
+ public void write(int b) {
+ os.write(b);
+ isEmpty = false;
+ }
+
+ @Override
+ public void write(byte[] b, int off, int len) {
+ os.write(b, off, len);
+ isEmpty = false;
+ }
+
+ @Override
+ public void flush() throws IOException {
+ os.flush();
+ }
+
+ @Override
+ public void sync() throws IOException { }
+
+ // --------------------------------------------------------------------
+
+ @Override
+ public void close() {
+ closed = true;
+ os.reset();
+ }
+
+ @Override
+ public StreamStateHandle closeAndGetHandle() throws IOException {
+ if (isEmpty) {
+ return null;
+ }
+ return new ByteStreamStateHandle(closeAndGetBytes());
+ }
+
+ @Override
+ public long getPos() throws IOException {
+ return os.getPosition();
+ }
+
+ /**
+ * Closes the stream and returns the byte array containing the stream's data.
+ * @return The byte array containing the stream's data.
+ * @throws IOException Thrown if the size of the data exceeds the maximal
+ */
+ public byte[] closeAndGetBytes() throws IOException {
+ if (!closed) {
+ checkSize(os.size(), maxSize);
+ byte[] bytes = os.toByteArray();
+ close();
+ return bytes;
+ }
+ else {
+ throw new IllegalStateException("stream has already been closed");
+ }
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemFoldingState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemFoldingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemFoldingState.java
deleted file mode 100644
index a4dec3b..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemFoldingState.java
+++ /dev/null
@@ -1,135 +0,0 @@
-/*
- * 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.memory;
-
-import org.apache.flink.api.common.functions.FoldFunction;
-import org.apache.flink.api.common.state.FoldingState;
-import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.util.Preconditions;
-
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Heap-backed partitioned {@link FoldingState} that is
- * snapshotted into a serialized memory copy.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <T> The type of the values that can be folded into the state.
- * @param <ACC> The type of the value in the folding state.
- */
-public class MemFoldingState<K, N, T, ACC>
- extends AbstractMemState<K, N, ACC, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>>
- implements FoldingState<T, ACC> {
-
- private final FoldFunction<T, ACC> foldFunction;
-
- public MemFoldingState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- FoldingStateDescriptor<T, ACC> stateDesc) {
- super(keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc);
- this.foldFunction = stateDesc.getFoldFunction();
- }
-
- public MemFoldingState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- FoldingStateDescriptor<T, ACC> stateDesc,
- HashMap<N, Map<K, ACC>> state) {
- super(keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc, state);
- this.foldFunction = stateDesc.getFoldFunction();
- }
-
- @Override
- public ACC get() {
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = state.get(currentNamespace);
- }
- if (currentNSState != null) {
- Preconditions.checkState(currentKey != null, "No key set");
- return currentNSState.get(currentKey);
- } else {
- return null;
- }
- }
-
- @Override
- public void add(T value) throws IOException {
- Preconditions.checkState(currentKey != null, "No key set");
-
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = createNewNamespaceMap();
- state.put(currentNamespace, currentNSState);
- }
-
- ACC currentValue = currentNSState.get(currentKey);
- try {
- if (currentValue == null) {
- currentNSState.put(currentKey, foldFunction.fold(stateDesc.getDefaultValue(), value));
- } else {
- currentNSState.put(currentKey, foldFunction.fold(currentValue, value));
-
- }
- } catch (Exception e) {
- throw new RuntimeException("Could not add value to folding state.", e);
- }
- }
-
- @Override
- public KvStateSnapshot<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>, MemoryStateBackend> createHeapSnapshot(byte[] bytes) {
- return new Snapshot<>(getKeySerializer(), getNamespaceSerializer(), stateSerializer, stateDesc, bytes);
- }
-
- @Override
- public byte[] getSerializedValue(K key, N namespace) throws Exception {
- Preconditions.checkNotNull(key, "Key");
- Preconditions.checkNotNull(namespace, "Namespace");
-
- Map<K, ACC> stateByKey = state.get(namespace);
-
- if (stateByKey != null) {
- return KvStateRequestSerializer.serializeValue(stateByKey.get(key), stateDesc.getSerializer());
- } else {
- return null;
- }
- }
-
- public static class Snapshot<K, N, T, ACC> extends AbstractMemStateSnapshot<K, N, ACC, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>> {
- private static final long serialVersionUID = 1L;
-
- public Snapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<ACC> stateSerializer,
- FoldingStateDescriptor<T, ACC> stateDescs, byte[] data) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDescs, data);
- }
-
- @Override
- public KvState<K, N, FoldingState<T, ACC>, FoldingStateDescriptor<T, ACC>, MemoryStateBackend> createMemState(HashMap<N, Map<K, ACC>> stateMap) {
- return new MemFoldingState<>(keySerializer, namespaceSerializer, stateDesc, stateMap);
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemListState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemListState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemListState.java
deleted file mode 100644
index 20b6eb5..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemListState.java
+++ /dev/null
@@ -1,120 +0,0 @@
-/*
- * 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.memory;
-
-import org.apache.flink.api.common.state.ListState;
-import org.apache.flink.api.common.state.ListStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.runtime.state.ArrayListSerializer;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.util.Preconditions;
-
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Heap-backed partitioned {@link org.apache.flink.api.common.state.ListState} that is snapshotted
- * into a serialized memory copy.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <V> The type of the values in the list state.
- */
-public class MemListState<K, N, V>
- extends AbstractMemState<K, N, ArrayList<V>, ListState<V>, ListStateDescriptor<V>>
- implements ListState<V> {
-
- public MemListState(TypeSerializer<K> keySerializer, TypeSerializer<N> namespaceSerializer, ListStateDescriptor<V> stateDesc) {
- super(keySerializer, namespaceSerializer, new ArrayListSerializer<>(stateDesc.getSerializer()), stateDesc);
- }
-
- public MemListState(TypeSerializer<K> keySerializer, TypeSerializer<N> namespaceSerializer, ListStateDescriptor<V> stateDesc, HashMap<N, Map<K, ArrayList<V>>> state) {
- super(keySerializer, namespaceSerializer, new ArrayListSerializer<>(stateDesc.getSerializer()), stateDesc, state);
- }
-
- @Override
- public Iterable<V> get() {
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = state.get(currentNamespace);
- }
- if (currentNSState != null) {
- Preconditions.checkState(currentKey != null, "No key set");
- return currentNSState.get(currentKey);
- } else {
- return null;
- }
- }
-
- @Override
- public void add(V value) {
- Preconditions.checkState(currentKey != null, "No key set");
-
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = createNewNamespaceMap();
- state.put(currentNamespace, currentNSState);
- }
-
- ArrayList<V> list = currentNSState.get(currentKey);
- if (list == null) {
- list = new ArrayList<>();
- currentNSState.put(currentKey, list);
- }
- list.add(value);
- }
-
- @Override
- public KvStateSnapshot<K, N, ListState<V>, ListStateDescriptor<V>, MemoryStateBackend> createHeapSnapshot(byte[] bytes) {
- return new Snapshot<>(getKeySerializer(), getNamespaceSerializer(), stateSerializer, stateDesc, bytes);
- }
-
- @Override
- public byte[] getSerializedValue(K key, N namespace) throws Exception {
- Preconditions.checkNotNull(key, "Key");
- Preconditions.checkNotNull(namespace, "Namespace");
-
- Map<K, ArrayList<V>> stateByKey = state.get(namespace);
- if (stateByKey != null) {
- return KvStateRequestSerializer.serializeList(stateByKey.get(key), stateDesc.getSerializer());
- } else {
- return null;
- }
- }
-
- public static class Snapshot<K, N, V> extends AbstractMemStateSnapshot<K, N, ArrayList<V>, ListState<V>, ListStateDescriptor<V>> {
- private static final long serialVersionUID = 1L;
-
- public Snapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<ArrayList<V>> stateSerializer,
- ListStateDescriptor<V> stateDescs, byte[] data) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDescs, data);
- }
-
- @Override
- public KvState<K, N, ListState<V>, ListStateDescriptor<V>, MemoryStateBackend> createMemState(HashMap<N, Map<K, ArrayList<V>>> stateMap) {
- return new MemListState<>(keySerializer, namespaceSerializer, stateDesc, stateMap);
- }
- }
-
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemReducingState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemReducingState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemReducingState.java
deleted file mode 100644
index 9a4c676..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemReducingState.java
+++ /dev/null
@@ -1,139 +0,0 @@
-/*
- * 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.memory;
-
-import org.apache.flink.api.common.functions.ReduceFunction;
-import org.apache.flink.api.common.state.ReducingState;
-import org.apache.flink.api.common.state.ReducingStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.util.Preconditions;
-
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Heap-backed partitioned {@link org.apache.flink.api.common.state.ReducingState} that is
- * snapshotted into a serialized memory copy.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <V> The type of the values in the list state.
- */
-public class MemReducingState<K, N, V>
- extends AbstractMemState<K, N, V, ReducingState<V>, ReducingStateDescriptor<V>>
- implements ReducingState<V> {
-
- private final ReduceFunction<V> reduceFunction;
-
- public MemReducingState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ReducingStateDescriptor<V> stateDesc) {
- super(keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc);
- this.reduceFunction = stateDesc.getReduceFunction();
- }
-
- public MemReducingState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ReducingStateDescriptor<V> stateDesc,
- HashMap<N, Map<K, V>> state) {
- super(keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc, state);
- this.reduceFunction = stateDesc.getReduceFunction();
- }
-
- @Override
- public V get() {
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = state.get(currentNamespace);
- }
- if (currentNSState != null) {
- Preconditions.checkState(currentKey != null, "No key set");
- return currentNSState.get(currentKey);
- }
- return null;
- }
-
- @Override
- public void add(V value) throws IOException {
- Preconditions.checkState(currentKey != null, "No key set");
-
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = createNewNamespaceMap();
- state.put(currentNamespace, currentNSState);
- }
-// currentKeyState.merge(currentNamespace, value, new BiFunction<V, V, V>() {
-// @Override
-// public V apply(V v, V v2) {
-// try {
-// return reduceFunction.reduce(v, v2);
-// } catch (Exception e) {
-// return null;
-// }
-// }
-// });
- V currentValue = currentNSState.get(currentKey);
- if (currentValue == null) {
- currentNSState.put(currentKey, value);
- } else {
- try {
- currentNSState.put(currentKey, reduceFunction.reduce(currentValue, value));
- } catch (Exception e) {
- throw new RuntimeException("Could not add value to reducing state.", e);
- }
- }
- }
-
- @Override
- public KvStateSnapshot<K, N, ReducingState<V>, ReducingStateDescriptor<V>, MemoryStateBackend> createHeapSnapshot(byte[] bytes) {
- return new Snapshot<>(getKeySerializer(), getNamespaceSerializer(), stateSerializer, stateDesc, bytes);
- }
-
- @Override
- public byte[] getSerializedValue(K key, N namespace) throws Exception {
- Preconditions.checkNotNull(key, "Key");
- Preconditions.checkNotNull(namespace, "Namespace");
-
- Map<K, V> stateByKey = state.get(namespace);
- if (stateByKey != null) {
- return KvStateRequestSerializer.serializeValue(stateByKey.get(key), stateDesc.getSerializer());
- } else {
- return null;
- }
- }
-
- public static class Snapshot<K, N, V> extends AbstractMemStateSnapshot<K, N, V, ReducingState<V>, ReducingStateDescriptor<V>> {
- private static final long serialVersionUID = 1L;
-
- public Snapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<V> stateSerializer,
- ReducingStateDescriptor<V> stateDescs, byte[] data) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDescs, data);
- }
-
- @Override
- public KvState<K, N, ReducingState<V>, ReducingStateDescriptor<V>, MemoryStateBackend> createMemState(HashMap<N, Map<K, V>> stateMap) {
- return new MemReducingState<>(keySerializer, namespaceSerializer, stateDesc, stateMap);
- }
- }}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemValueState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemValueState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemValueState.java
deleted file mode 100644
index c0e3779..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemValueState.java
+++ /dev/null
@@ -1,122 +0,0 @@
-/*
- * 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.memory;
-
-import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.util.Preconditions;
-
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Heap-backed key/value state that is snapshotted into a serialized memory copy.
- *
- * @param <K> The type of the key.
- * @param <N> The type of the namespace.
- * @param <V> The type of the value.
- */
-public class MemValueState<K, N, V>
- extends AbstractMemState<K, N, V, ValueState<V>, ValueStateDescriptor<V>>
- implements ValueState<V> {
-
- public MemValueState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ValueStateDescriptor<V> stateDesc) {
- super(keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc);
- }
-
- public MemValueState(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- ValueStateDescriptor<V> stateDesc,
- HashMap<N, Map<K, V>> state) {
- super(keySerializer, namespaceSerializer, stateDesc.getSerializer(), stateDesc, state);
- }
-
- @Override
- public V value() {
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = state.get(currentNamespace);
- }
- if (currentNSState != null) {
- Preconditions.checkState(currentKey != null, "No key set");
- V value = currentNSState.get(currentKey);
- return value != null ? value : stateDesc.getDefaultValue();
- }
- return stateDesc.getDefaultValue();
- }
-
- @Override
- public void update(V value) {
- Preconditions.checkState(currentKey != null, "No key set");
-
- if (value == null) {
- clear();
- return;
- }
-
- if (currentNSState == null) {
- Preconditions.checkState(currentNamespace != null, "No namespace set");
- currentNSState = createNewNamespaceMap();
- state.put(currentNamespace, currentNSState);
- }
-
- currentNSState.put(currentKey, value);
- }
-
- @Override
- public KvStateSnapshot<K, N, ValueState<V>, ValueStateDescriptor<V>, MemoryStateBackend> createHeapSnapshot(byte[] bytes) {
- return new Snapshot<>(getKeySerializer(), getNamespaceSerializer(), stateSerializer, stateDesc, bytes);
- }
-
- @Override
- public byte[] getSerializedValue(K key, N namespace) throws Exception {
- Preconditions.checkNotNull(key, "Key");
- Preconditions.checkNotNull(namespace, "Namespace");
-
- Map<K, V> stateByKey = state.get(namespace);
- V value = stateByKey != null ? stateByKey.get(key) : stateDesc.getDefaultValue();
- if (value != null) {
- return KvStateRequestSerializer.serializeValue(value, stateDesc.getSerializer());
- } else {
- return KvStateRequestSerializer.serializeValue(stateDesc.getDefaultValue(), stateDesc.getSerializer());
- }
- }
-
- public static class Snapshot<K, N, V> extends AbstractMemStateSnapshot<K, N, V, ValueState<V>, ValueStateDescriptor<V>> {
- private static final long serialVersionUID = 1L;
-
- public Snapshot(TypeSerializer<K> keySerializer,
- TypeSerializer<N> namespaceSerializer,
- TypeSerializer<V> stateSerializer,
- ValueStateDescriptor<V> stateDescs, byte[] data) {
- super(keySerializer, namespaceSerializer, stateSerializer, stateDescs, data);
- }
-
- @Override
- public KvState<K, N, ValueState<V>, ValueStateDescriptor<V>, MemoryStateBackend> createMemState(HashMap<N, Map<K, V>> stateMap) {
- return new MemValueState<>(keySerializer, namespaceSerializer, stateDesc, stateMap);
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
index af84394..654c367 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
@@ -18,21 +18,20 @@
package org.apache.flink.runtime.state.memory;
-import org.apache.flink.api.common.state.FoldingState;
-import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.state.ListState;
-import org.apache.flink.api.common.state.ListStateDescriptor;
-import org.apache.flink.api.common.state.ReducingState;
-import org.apache.flink.api.common.state.ReducingStateDescriptor;
-import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.heap.HeapKeyedStateBackend;
-
-import java.io.ByteArrayOutputStream;
import java.io.IOException;
+import java.util.List;
/**
* A {@link AbstractStateBackend} that stores all its data and checkpoints in memory and has no
@@ -67,142 +66,47 @@ public class MemoryStateBackend extends AbstractStateBackend {
this.maxStateSize = maxStateSize;
}
- // ------------------------------------------------------------------------
- // initialization and cleanup
- // ------------------------------------------------------------------------
-
- @Override
- public void disposeAllStateForCurrentJob() {
- // nothing to do here, GC will do it
- }
-
- @Override
- public void close() throws Exception {}
-
- // ------------------------------------------------------------------------
- // State backend operations
- // ------------------------------------------------------------------------
-
@Override
- public <N, V> ValueState<V> createValueState(TypeSerializer<N> namespaceSerializer, ValueStateDescriptor<V> stateDesc) throws Exception {
- return new MemValueState<>(keySerializer, namespaceSerializer, stateDesc);
- }
-
- @Override
- public <N, T> ListState<T> createListState(TypeSerializer<N> namespaceSerializer, ListStateDescriptor<T> stateDesc) throws Exception {
- return new MemListState<>(keySerializer, namespaceSerializer, stateDesc);
- }
-
- @Override
- public <N, T> ReducingState<T> createReducingState(TypeSerializer<N> namespaceSerializer, ReducingStateDescriptor<T> stateDesc) throws Exception {
- return new MemReducingState<>(keySerializer, namespaceSerializer, stateDesc);
+ public String toString() {
+ return "MemoryStateBackend (data in heap memory / checkpoints to JobManager)";
}
@Override
- public <N, T, ACC> FoldingState<T, ACC> createFoldingState(TypeSerializer<N> namespaceSerializer, FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
- return new MemFoldingState<>(keySerializer, namespaceSerializer, stateDesc);
+ public CheckpointStreamFactory createStreamFactory(JobID jobId, String operatorIdentifier) throws IOException {
+ return new MemCheckpointStreamFactory(maxStateSize);
}
@Override
- public CheckpointStateOutputStream createCheckpointStateOutputStream(
- long checkpointID, long timestamp) throws Exception
- {
- return new MemoryCheckpointOutputStream(maxStateSize);
+ public <K> KeyedStateBackend<K> createKeyedStateBackend(
+ Environment env, JobID jobID,
+ String operatorIdentifier, TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ TaskKvStateRegistry kvStateRegistry) throws IOException {
+
+ return new HeapKeyedStateBackend<>(
+ kvStateRegistry,
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange);
}
- // ------------------------------------------------------------------------
- // Utilities
- // ------------------------------------------------------------------------
-
@Override
- public String toString() {
- return "MemoryStateBackend (data in heap memory / checkpoints to JobManager)";
+ public <K> KeyedStateBackend<K> restoreKeyedStateBackend(
+ Environment env, JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ List<KeyGroupsStateHandle> restoredState,
+ TaskKvStateRegistry kvStateRegistry) throws Exception {
+
+ return new HeapKeyedStateBackend<>(
+ kvStateRegistry,
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange,
+ restoredState);
}
- static void checkSize(int size, int maxSize) throws IOException {
- if (size > maxSize) {
- throw new IOException(
- "Size of the state is larger than the maximum permitted memory-backed state. Size="
- + size + " , maxSize=" + maxSize
- + " . Consider using a different state backend, like the File System State backend.");
- }
- }
-
- // ------------------------------------------------------------------------
-
- /**
- * A CheckpointStateOutputStream that writes into a byte array.
- */
- public static final class MemoryCheckpointOutputStream extends CheckpointStateOutputStream {
-
- private final ByteArrayOutputStream os = new ByteArrayOutputStream();
-
- private final int maxSize;
-
- private boolean closed;
-
- public MemoryCheckpointOutputStream(int maxSize) {
- this.maxSize = maxSize;
- }
-
- @Override
- public void write(int b) {
- os.write(b);
- }
-
- @Override
- public void write(byte[] b, int off, int len) {
- os.write(b, off, len);
- }
-
- @Override
- public void flush() throws IOException {
- os.flush();
- }
-
- @Override
- public void sync() throws IOException { }
-
- // --------------------------------------------------------------------
-
- @Override
- public void close() {
- closed = true;
- os.reset();
- }
-
- @Override
- public StreamStateHandle closeAndGetHandle() throws IOException {
- return new ByteStreamStateHandle(closeAndGetBytes());
- }
-
- /**
- * Closes the stream and returns the byte array containing the stream's data.
- * @return The byte array containing the stream's data.
- * @throws IOException Thrown if the size of the data exceeds the maximal
- */
- public byte[] closeAndGetBytes() throws IOException {
- if (!closed) {
- checkSize(os.size(), maxSize);
- byte[] bytes = os.toByteArray();
- close();
- return bytes;
- }
- else {
- throw new IllegalStateException("stream has already been closed");
- }
- }
- }
-
- // ------------------------------------------------------------------------
- // Static default instance
- // ------------------------------------------------------------------------
-
- /**
- * Gets the default instance of this state backend, using the default maximal state size.
- * @return The default instance of this state backend.
- */
- public static MemoryStateBackend create() {
- return new MemoryStateBackend();
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java
index d703bd6..766531a 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoaderTest.java
@@ -87,14 +87,14 @@ public class SavepointLoaderTest {
loaded.discardState();
verify(state, times(0)).discardState();
- // 2) Load and validate: parallelism mismatch
- when(vertex.getParallelism()).thenReturn(222);
+ // 2) Load and validate: max parallelism mismatch
+ when(vertex.getMaxParallelism()).thenReturn(222);
try {
SavepointLoader.loadAndValidateSavepoint(jobId, tasks, store, path);
fail("Did not throw expected Exception");
} catch (IllegalStateException expected) {
- assertTrue(expected.getMessage().contains("Parallelism mismatch"));
+ assertTrue(expected.getMessage().contains("Max parallelism mismatch"));
}
// 3) Load and validate: missing vertex (this should be relaxed)
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorTest.java
index 56da9c8..39ea176 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptorTest.java
@@ -48,6 +48,7 @@ public class TaskDeploymentDescriptorTest {
final ExecutionAttemptID execId = new ExecutionAttemptID();
final String jobName = "job name";
final String taskName = "task name";
+ final int numberOfKeyGroups = 1;
final int indexInSubtaskGroup = 0;
final int currentNumberOfSubtasks = 1;
final int attemptNumber = 0;
@@ -61,7 +62,7 @@ public class TaskDeploymentDescriptorTest {
final SerializedValue<ExecutionConfig> executionConfig = new SerializedValue<>(new ExecutionConfig());
final TaskDeploymentDescriptor orig = new TaskDeploymentDescriptor(jobID, jobName, vertexID, execId,
- executionConfig, taskName, indexInSubtaskGroup, currentNumberOfSubtasks, attemptNumber,
+ executionConfig, taskName, numberOfKeyGroups, indexInSubtaskGroup, currentNumberOfSubtasks, attemptNumber,
jobConfiguration, taskConfiguration, invokableClass.getName(), producedResults, inputGates,
requiredJars, requiredClasspaths, 47);
@@ -76,6 +77,7 @@ public class TaskDeploymentDescriptorTest {
assertEquals(orig.getJobID(), copy.getJobID());
assertEquals(orig.getVertexID(), copy.getVertexID());
assertEquals(orig.getTaskName(), copy.getTaskName());
+ assertEquals(orig.getNumberOfKeyGroups(), copy.getNumberOfKeyGroups());
assertEquals(orig.getIndexInSubtaskGroup(), copy.getIndexInSubtaskGroup());
assertEquals(orig.getNumberOfSubtasks(), copy.getNumberOfSubtasks());
assertEquals(orig.getAttemptNumber(), copy.getAttemptNumber());
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java
index c6eb249..6a6ac64 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/messages/CheckpointMessagesTest.java
@@ -120,7 +120,7 @@ public class CheckpointMessagesTest {
public void close() throws IOException {}
@Override
- public FSDataInputStream openInputStream() throws Exception {
+ public FSDataInputStream openInputStream() throws IOException {
return null;
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/metrics/groups/TaskManagerGroupTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/metrics/groups/TaskManagerGroupTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/metrics/groups/TaskManagerGroupTest.java
index b80ff6a..b2c5dc7 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/metrics/groups/TaskManagerGroupTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/metrics/groups/TaskManagerGroupTest.java
@@ -77,7 +77,7 @@ public class TaskManagerGroupTest {
execution11,
new SerializedValue<>(new ExecutionConfig()),
"test",
- 17, 18, 0,
+ 18, 17, 18, 0,
new Configuration(), new Configuration(),
"",
new ArrayList<ResultPartitionDeploymentDescriptor>(),
@@ -92,7 +92,7 @@ public class TaskManagerGroupTest {
execution12,
new SerializedValue<>(new ExecutionConfig()),
"test",
- 13, 18, 1,
+ 18, 13, 18, 1,
new Configuration(), new Configuration(),
"",
new ArrayList<ResultPartitionDeploymentDescriptor>(),
@@ -107,7 +107,7 @@ public class TaskManagerGroupTest {
execution21,
new SerializedValue<>(new ExecutionConfig()),
"test",
- 7, 18, 2,
+ 18, 7, 18, 2,
new Configuration(), new Configuration(),
"",
new ArrayList<ResultPartitionDeploymentDescriptor>(),
@@ -122,7 +122,7 @@ public class TaskManagerGroupTest {
execution13,
new SerializedValue<>(new ExecutionConfig()),
"test",
- 0, 18, 0,
+ 18, 0, 18, 0,
new Configuration(), new Configuration(),
"",
new ArrayList<ResultPartitionDeploymentDescriptor>(),
@@ -193,7 +193,7 @@ public class TaskManagerGroupTest {
execution11,
new SerializedValue<>(new ExecutionConfig()),
"test",
- 17, 18, 0,
+ 18, 17, 18, 0,
new Configuration(), new Configuration(),
"",
new ArrayList<ResultPartitionDeploymentDescriptor>(),
@@ -208,7 +208,7 @@ public class TaskManagerGroupTest {
execution12,
new SerializedValue<>(new ExecutionConfig()),
"test",
- 13, 18, 1,
+ 18, 13, 18, 1,
new Configuration(), new Configuration(),
"",
new ArrayList<ResultPartitionDeploymentDescriptor>(),
@@ -223,7 +223,7 @@ public class TaskManagerGroupTest {
execution21,
new SerializedValue<>(new ExecutionConfig()),
"test",
- 7, 18, 1,
+ 18, 7, 18, 1,
new Configuration(), new Configuration(),
"",
new ArrayList<ResultPartitionDeploymentDescriptor>(),
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java
index 19317f9..4654507 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/DummyEnvironment.java
@@ -53,13 +53,14 @@ public class DummyEnvironment implements Environment {
private final ExecutionAttemptID executionId = new ExecutionAttemptID();
private final ExecutionConfig executionConfig = new ExecutionConfig();
private final TaskInfo taskInfo;
- private final KvStateRegistry kvStateRegistry = new KvStateRegistry();
- private final TaskKvStateRegistry taskKvStateRegistry;
+ private KvStateRegistry kvStateRegistry = new KvStateRegistry();
public DummyEnvironment(String taskName, int numSubTasks, int subTaskIndex) {
- this.taskInfo = new TaskInfo(taskName, subTaskIndex, numSubTasks, 0);
+ this.taskInfo = new TaskInfo(taskName, numSubTasks, subTaskIndex, numSubTasks, 0);
+ }
- this.taskKvStateRegistry = kvStateRegistry.createTaskRegistry(jobId, jobVertexId);
+ public void setKvStateRegistry(KvStateRegistry kvStateRegistry) {
+ this.kvStateRegistry = kvStateRegistry;
}
public KvStateRegistry getKvStateRegistry() {
@@ -148,7 +149,7 @@ public class DummyEnvironment implements Environment {
@Override
public TaskKvStateRegistry getTaskKvStateRegistry() {
- return taskKvStateRegistry;
+ return kvStateRegistry.createTaskRegistry(jobId, jobVertexId);
}
@Override
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
index 2c76399..e7bf6e1 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
@@ -99,11 +99,24 @@ public class MockEnvironment implements Environment {
private final int bufferSize;
public MockEnvironment(String taskName, long memorySize, MockInputSplitProvider inputSplitProvider, int bufferSize) {
- this(taskName, memorySize, inputSplitProvider, bufferSize, new Configuration());
+ this(taskName, memorySize, inputSplitProvider, bufferSize, new Configuration(), new ExecutionConfig());
}
- public MockEnvironment(String taskName, long memorySize, MockInputSplitProvider inputSplitProvider, int bufferSize, Configuration taskConfiguration) {
- this.taskInfo = new TaskInfo(taskName, 0, 1, 0);
+ public MockEnvironment(String taskName, long memorySize, MockInputSplitProvider inputSplitProvider, int bufferSize, Configuration taskConfiguration, ExecutionConfig executionConfig) {
+ this(taskName, memorySize, inputSplitProvider, bufferSize, taskConfiguration, executionConfig, 1, 1, 0);
+ }
+
+ public MockEnvironment(
+ String taskName,
+ long memorySize,
+ MockInputSplitProvider inputSplitProvider,
+ int bufferSize,
+ Configuration taskConfiguration,
+ ExecutionConfig executionConfig,
+ int maxParallelism,
+ int parallelism,
+ int subtaskIndex) {
+ this.taskInfo = new TaskInfo(taskName, maxParallelism, subtaskIndex, parallelism, 0);
this.jobConfiguration = new Configuration();
this.taskConfiguration = taskConfiguration;
this.inputs = new LinkedList<InputGate>();
@@ -111,7 +124,7 @@ public class MockEnvironment implements Environment {
this.memManager = new MemoryManager(memorySize, 1);
this.ioManager = new IOManagerAsync();
- this.executionConfig = new ExecutionConfig();
+ this.executionConfig = executionConfig;
this.inputSplitProvider = inputSplitProvider;
this.bufferSize = bufferSize;
@@ -121,6 +134,7 @@ public class MockEnvironment implements Environment {
this.kvStateRegistry = registry.createTaskRegistry(jobID, getJobVertexId());
}
+
public IteratorWrappingTestSingleInputGate<Record> addInput(MutableObjectIterator<Record> inputIterator) {
try {
final IteratorWrappingTestSingleInputGate<Record> reader = new IteratorWrappingTestSingleInputGate<Record>(bufferSize, Record.class, inputIterator);
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java
index 36f2b45..3380907 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java
@@ -26,14 +26,20 @@ import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.akka.AkkaUtils;
import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.query.netty.AtomicKvStateRequestStats;
import org.apache.flink.runtime.query.netty.KvStateClient;
import org.apache.flink.runtime.query.netty.KvStateServer;
import org.apache.flink.runtime.query.netty.UnknownKvStateID;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-import org.apache.flink.runtime.state.memory.MemValueState;
+import org.apache.flink.runtime.state.heap.HeapValueState;
+import org.apache.flink.runtime.state.heap.StateTable;
+import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.util.MathUtils;
import org.junit.AfterClass;
import org.junit.Test;
@@ -237,10 +243,22 @@ public class QueryableStateClientTest {
KvStateClient networkClient = null;
AtomicKvStateRequestStats networkClientStats = new AtomicKvStateRequestStats();
+ MemoryStateBackend backend = new MemoryStateBackend();
+ DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
+
+ KeyedStateBackend<Integer> keyedStateBackend = backend.createKeyedStateBackend(dummyEnv,
+ new JobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
+
+
try {
KvStateRegistry[] registries = new KvStateRegistry[numServers];
KvStateID[] kvStateIds = new KvStateID[numServers];
- List<MemValueState<Integer, VoidNamespace, Integer>> kvStates = new ArrayList<>();
+ List<HeapValueState<Integer, VoidNamespace, Integer>> kvStates = new ArrayList<>();
// Start the servers
for (int i = 0; i < numServers; i++) {
@@ -249,11 +267,14 @@ public class QueryableStateClientTest {
servers[i] = new KvStateServer(InetAddress.getLocalHost(), 0, 1, 1, registries[i], serverStats[i]);
servers[i].start();
+
// Register state
- MemValueState<Integer, VoidNamespace, Integer> kvState = new MemValueState<>(
+ HeapValueState<Integer, VoidNamespace, Integer> kvState = new HeapValueState<>(
+ keyedStateBackend,
+ new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null),
+ new StateTable<Integer, VoidNamespace, Integer>(IntSerializer.INSTANCE, VoidNamespaceSerializer.INSTANCE, new KeyGroupRange(0, 1)),
IntSerializer.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null));
+ VoidNamespaceSerializer.INSTANCE);
kvStates.add(kvState);
@@ -271,9 +292,9 @@ public class QueryableStateClientTest {
int targetKeyGroupIndex = MathUtils.murmurHash(key) % numServers;
expectedRequests[targetKeyGroupIndex]++;
- MemValueState<Integer, VoidNamespace, Integer> kvState = kvStates.get(targetKeyGroupIndex);
+ HeapValueState<Integer, VoidNamespace, Integer> kvState = kvStates.get(targetKeyGroupIndex);
- kvState.setCurrentKey(key);
+ keyedStateBackend.setCurrentKey(key);
kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
kvState.update(1337 + key);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java
index ac03f94..796481c 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java
@@ -30,18 +30,25 @@ import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.query.KvStateID;
import org.apache.flink.runtime.query.KvStateRegistry;
import org.apache.flink.runtime.query.KvStateServerAddress;
import org.apache.flink.runtime.query.netty.message.KvStateRequest;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
import org.apache.flink.runtime.query.netty.message.KvStateRequestType;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.KvState;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-import org.apache.flink.runtime.state.memory.MemValueState;
+import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.util.NetUtils;
import org.junit.AfterClass;
import org.junit.Test;
@@ -526,6 +533,20 @@ public class KvStateClientTest {
final int batchSize = 16;
+ AbstractStateBackend abstractBackend = new MemoryStateBackend();
+ KvStateRegistry dummyRegistry = new KvStateRegistry();
+ DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
+ dummyEnv.setKvStateRegistry(dummyRegistry);
+ KeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(
+ dummyEnv,
+ new JobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ dummyRegistry.createTaskRegistry(new JobID(), new JobVertexID()));
+
+
final FiniteDuration timeout = new FiniteDuration(10, TimeUnit.SECONDS);
AtomicKvStateRequestStats clientStats = new AtomicKvStateRequestStats();
@@ -542,11 +563,6 @@ public class KvStateClientTest {
ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null);
desc.setQueryable("any");
- MemValueState<Integer, VoidNamespace, Integer> kvState = new MemValueState<>(
- IntSerializer.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- desc);
-
// Create servers
KvStateRegistry[] registry = new KvStateRegistry[numServers];
AtomicKvStateRequestStats[] serverStats = new AtomicKvStateRequestStats[numServers];
@@ -565,10 +581,17 @@ public class KvStateClientTest {
server[i].start();
+ backend.setCurrentKey(1010 + i);
+
// Value per server
- kvState.setCurrentKey(1010 + i);
- kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
- kvState.update(201 + i);
+ ValueState<Integer> state = backend.getPartitionedState(VoidNamespace.INSTANCE,
+ VoidNamespaceSerializer.INSTANCE,
+ desc);
+
+ state.update(201 + i);
+
+ // we know it must be a KvStat but this is not exposed to the user via State
+ KvState<?> kvState = (KvState<?>) state;
// Register KvState (one state instance for all server)
ids[i] = registry[i].registerKvState(new JobID(), new JobVertexID(), 0, "any", kvState);
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java
index 6ad7ece..3d2e8b5 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java
@@ -24,21 +24,28 @@ import io.netty.channel.ChannelHandler;
import io.netty.channel.embedded.EmbeddedChannel;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.common.typeutils.base.array.BytePrimitiveArraySerializer;
import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.query.KvStateID;
import org.apache.flink.runtime.query.KvStateRegistry;
+import org.apache.flink.runtime.query.KvStateRegistryListener;
import org.apache.flink.runtime.query.netty.message.KvStateRequestFailure;
import org.apache.flink.runtime.query.netty.message.KvStateRequestResult;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
import org.apache.flink.runtime.query.netty.message.KvStateRequestType;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.KvState;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-import org.apache.flink.runtime.state.memory.MemValueState;
+import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.junit.AfterClass;
import org.junit.Test;
@@ -80,28 +87,34 @@ public class KvStateServerHandlerTest {
// Register state
ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null);
- desc.setQueryable("any");
+ desc.setQueryable("vanilla");
- MemValueState<Integer, VoidNamespace, Integer> kvState = new MemValueState<>(
+ AbstractStateBackend abstractBackend = new MemoryStateBackend();
+ DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
+ dummyEnv.setKvStateRegistry(registry);
+ KeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(
+ dummyEnv,
+ new JobID(),
+ "test_op",
IntSerializer.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- desc);
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
- KvStateID kvStateId = registry.registerKvState(
- new JobID(),
- new JobVertexID(),
- 0,
- "vanilla",
- kvState);
+ final TestRegistryListener registryListener = new TestRegistryListener();
+ registry.registerListener(registryListener);
// Update the KvState and request it
int expectedValue = 712828289;
int key = 99812822;
- kvState.setCurrentKey(key);
- kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
+ backend.setCurrentKey(key);
+ ValueState<Integer> state = backend.getPartitionedState(
+ VoidNamespace.INSTANCE,
+ VoidNamespaceSerializer.INSTANCE,
+ desc);
- kvState.update(expectedValue);
+ state.update(expectedValue);
byte[] serializedKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(
key,
@@ -110,10 +123,13 @@ public class KvStateServerHandlerTest {
VoidNamespaceSerializer.INSTANCE);
long requestId = Integer.MAX_VALUE + 182828L;
+
+ assertTrue(registryListener.registrationName.equals("vanilla"));
+
ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(
channel.alloc(),
requestId,
- kvStateId,
+ registryListener.kvStateId,
serializedKeyAndNamespace);
// Write the request and wait for the response
@@ -184,21 +200,26 @@ public class KvStateServerHandlerTest {
KvStateServerHandler handler = new KvStateServerHandler(registry, TEST_THREAD_POOL, stats);
EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
+ AbstractStateBackend abstractBackend = new MemoryStateBackend();
+ DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
+ dummyEnv.setKvStateRegistry(registry);
+ KeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(
+ dummyEnv,
+ new JobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
+
+ final TestRegistryListener registryListener = new TestRegistryListener();
+ registry.registerListener(registryListener);
+
// Register state
ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null);
- desc.setQueryable("any");
+ desc.setQueryable("vanilla");
- MemValueState<Integer, VoidNamespace, Integer> kvState = new MemValueState<>(
- IntSerializer.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- desc);
-
- KvStateID kvStateId = registry.registerKvState(
- new JobID(),
- new JobVertexID(),
- 0,
- "vanilla",
- kvState);
+ backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
byte[] serializedKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(
1238283,
@@ -207,10 +228,13 @@ public class KvStateServerHandlerTest {
VoidNamespaceSerializer.INSTANCE);
long requestId = Integer.MAX_VALUE + 22982L;
+
+ assertTrue(registryListener.registrationName.equals("vanilla"));
+
ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(
channel.alloc(),
requestId,
- kvStateId,
+ registryListener.kvStateId,
serializedKeyAndNamespace);
// Write the request and wait for the response
@@ -225,6 +249,8 @@ public class KvStateServerHandlerTest {
assertEquals(requestId, response.getRequestId());
+ System.out.println("RESPOINSE: " + response);
+
assertTrue("Did not respond with expected failure cause", response.getCause() instanceof UnknownKeyOrNamespace);
assertEquals(1, stats.getNumRequests());
@@ -244,7 +270,7 @@ public class KvStateServerHandlerTest {
EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
// Failing KvState
- KvState<?, ?, ?, ?, ?> kvState = mock(KvState.class);
+ KvState<?> kvState = mock(KvState.class);
when(kvState.getSerializedValue(any(byte[].class)))
.thenThrow(new RuntimeException("Expected test Exception"));
@@ -320,26 +346,33 @@ public class KvStateServerHandlerTest {
KvStateServerHandler handler = new KvStateServerHandler(registry, closedExecutor, stats);
EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
+ AbstractStateBackend abstractBackend = new MemoryStateBackend();
+ DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
+ dummyEnv.setKvStateRegistry(registry);
+ KeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(
+ dummyEnv,
+ new JobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
+
+ final TestRegistryListener registryListener = new TestRegistryListener();
+ registry.registerListener(registryListener);
+
// Register state
ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null);
- desc.setQueryable("any");
+ desc.setQueryable("vanilla");
- MemValueState<Integer, VoidNamespace, Integer> kvState = new MemValueState<>(
- IntSerializer.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- desc);
+ backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
- KvStateID kvStateId = registry.registerKvState(
- new JobID(),
- new JobVertexID(),
- 0,
- "vanilla",
- kvState);
+ assertTrue(registryListener.registrationName.equals("vanilla"));
ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(
channel.alloc(),
282872,
- kvStateId,
+ registryListener.kvStateId,
new byte[0]);
// Write the request and wait for the response
@@ -451,28 +484,35 @@ public class KvStateServerHandlerTest {
KvStateServerHandler handler = new KvStateServerHandler(registry, TEST_THREAD_POOL, stats);
EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
+ AbstractStateBackend abstractBackend = new MemoryStateBackend();
+ DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
+ dummyEnv.setKvStateRegistry(registry);
+ KeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(
+ dummyEnv,
+ new JobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
+
+ final TestRegistryListener registryListener = new TestRegistryListener();
+ registry.registerListener(registryListener);
+
// Register state
ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null);
- desc.setQueryable("any");
+ desc.setQueryable("vanilla");
- MemValueState<Integer, VoidNamespace, Integer> kvState = new MemValueState<>(
- IntSerializer.INSTANCE,
+ ValueState<Integer> state = backend.getPartitionedState(
+ VoidNamespace.INSTANCE,
VoidNamespaceSerializer.INSTANCE,
desc);
- KvStateID kvStateId = registry.registerKvState(
- new JobID(),
- new JobVertexID(),
- 0,
- "vanilla",
- kvState);
-
int key = 99812822;
// Update the KvState
- kvState.setCurrentKey(key);
- kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
- kvState.update(712828289);
+ backend.setCurrentKey(key);
+ state.update(712828289);
byte[] wrongKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(
"wrong-key-type",
@@ -486,10 +526,11 @@ public class KvStateServerHandlerTest {
"wrong-namespace-type",
StringSerializer.INSTANCE);
+ assertTrue(registryListener.registrationName.equals("vanilla"));
ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(
channel.alloc(),
182828,
- kvStateId,
+ registryListener.kvStateId,
wrongKeyAndNamespace);
// Write the request and wait for the response
@@ -508,7 +549,7 @@ public class KvStateServerHandlerTest {
request = KvStateRequestSerializer.serializeKvStateRequest(
channel.alloc(),
182829,
- kvStateId,
+ registryListener.kvStateId,
wrongNamespace);
// Write the request and wait for the response
@@ -538,22 +579,30 @@ public class KvStateServerHandlerTest {
KvStateServerHandler handler = new KvStateServerHandler(registry, TEST_THREAD_POOL, stats);
EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
+ AbstractStateBackend abstractBackend = new MemoryStateBackend();
+ DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
+ dummyEnv.setKvStateRegistry(registry);
+ KeyedStateBackend<Integer> backend = abstractBackend.createKeyedStateBackend(
+ dummyEnv,
+ new JobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
+
+ final TestRegistryListener registryListener = new TestRegistryListener();
+ registry.registerListener(registryListener);
+
// Register state
ValueStateDescriptor<byte[]> desc = new ValueStateDescriptor<>("any", BytePrimitiveArraySerializer.INSTANCE, null);
- desc.setQueryable("any");
+ desc.setQueryable("vanilla");
- MemValueState<Integer, VoidNamespace, byte[]> kvState = new MemValueState<>(
- IntSerializer.INSTANCE,
+ ValueState<byte[]> state = backend.getPartitionedState(
+ VoidNamespace.INSTANCE,
VoidNamespaceSerializer.INSTANCE,
desc);
- KvStateID kvStateId = registry.registerKvState(
- new JobID(),
- new JobVertexID(),
- 0,
- "vanilla",
- kvState);
-
// Update KvState
byte[] bytes = new byte[2 * channel.config().getWriteBufferHighWaterMark()];
@@ -563,9 +612,8 @@ public class KvStateServerHandlerTest {
}
int key = 99812822;
- kvState.setCurrentKey(key);
- kvState.setCurrentNamespace(VoidNamespace.INSTANCE);
- kvState.update(bytes);
+ backend.setCurrentKey(key);
+ state.update(bytes);
// Request
byte[] serializedKeyAndNamespace = KvStateRequestSerializer.serializeKeyAndNamespace(
@@ -575,10 +623,13 @@ public class KvStateServerHandlerTest {
VoidNamespaceSerializer.INSTANCE);
long requestId = Integer.MAX_VALUE + 182828L;
+
+ assertTrue(registryListener.registrationName.equals("vanilla"));
+
ByteBuf request = KvStateRequestSerializer.serializeKvStateRequest(
channel.alloc(),
requestId,
- kvStateId,
+ registryListener.kvStateId,
serializedKeyAndNamespace);
// Write the request and wait for the response
@@ -619,4 +670,35 @@ public class KvStateServerHandlerTest {
private ChannelHandler getFrameDecoder() {
return new LengthFieldBasedFrameDecoder(Integer.MAX_VALUE, 0, 4, 0, 4);
}
+
+ /**
+ * A listener that keeps the last updated KvState information so that a test
+ * can retrieve it.
+ */
+ static class TestRegistryListener implements KvStateRegistryListener {
+ volatile JobVertexID jobVertexID;
+ volatile int keyGroupIndex;
+ volatile String registrationName;
+ volatile KvStateID kvStateId;
+
+ @Override
+ public void notifyKvStateRegistered(JobID jobId,
+ JobVertexID jobVertexId,
+ int keyGroupIndex,
+ String registrationName,
+ KvStateID kvStateId) {
+ this.jobVertexID = jobVertexId;
+ this.keyGroupIndex = keyGroupIndex;
+ this.registrationName = registrationName;
+ this.kvStateId = kvStateId;
+ }
+
+ @Override
+ public void notifyKvStateUnregistered(JobID jobId,
+ JobVertexID jobVertexId,
+ int keyGroupIndex,
+ String registrationName) {
+
+ }
+ }
}
[04/27] flink git commit: [FLINK-3761] Refactor RocksDB Backend/Make
Key-Group Aware
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBMergeIteratorTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBMergeIteratorTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBMergeIteratorTest.java
new file mode 100644
index 0000000..1cb3b2b
--- /dev/null
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBMergeIteratorTest.java
@@ -0,0 +1,140 @@
+/*
+ * 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.contrib.streaming.state;
+
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
+import org.apache.flink.runtime.testutils.CommonTestUtils;
+import org.junit.Assert;
+import org.junit.Test;
+import org.rocksdb.ColumnFamilyDescriptor;
+import org.rocksdb.ColumnFamilyHandle;
+import org.rocksdb.RocksDB;
+import org.rocksdb.RocksIterator;
+
+import java.io.DataOutputStream;
+import java.io.File;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+
+public class RocksDBMergeIteratorTest {
+
+ private static final int NUM_KEY_VAL_STATES = 50;
+ private static final int MAX_NUM_KEYS = 20;
+
+ @Test
+ public void testEmptyMergeIterator() throws IOException {
+ RocksDBKeyedStateBackend.RocksDBMergeIterator emptyIterator =
+ new RocksDBKeyedStateBackend.RocksDBMergeIterator(Collections.EMPTY_LIST, 2);
+ Assert.assertFalse(emptyIterator.isValid());
+ }
+
+ @Test
+ public void testMergeIterator() throws Exception {
+ Assert.assertTrue(MAX_NUM_KEYS <= Byte.MAX_VALUE);
+
+ testMergeIterator(Byte.MAX_VALUE);
+ testMergeIterator(Short.MAX_VALUE);
+ }
+
+ public void testMergeIterator(int maxParallelism) throws Exception {
+ Random random = new Random(1234);
+
+ File tmpDir = CommonTestUtils.createTempDirectory();
+
+ RocksDB rocksDB = RocksDB.open(tmpDir.getAbsolutePath());
+ try {
+ List<Tuple2<RocksIterator, Integer>> rocksIteratorsWithKVStateId = new ArrayList<>();
+ List<Tuple2<ColumnFamilyHandle, Integer>> columnFamilyHandlesWithKeyCount = new ArrayList<>();
+
+ int totalKeysExpected = 0;
+
+ for (int c = 0; c < NUM_KEY_VAL_STATES; ++c) {
+ ColumnFamilyHandle handle = rocksDB.createColumnFamily(
+ new ColumnFamilyDescriptor(("column-" + c).getBytes()));
+
+ ByteArrayOutputStreamWithPos bos = new ByteArrayOutputStreamWithPos();
+ DataOutputStream dos = new DataOutputStream(bos);
+
+ int numKeys = random.nextInt(MAX_NUM_KEYS + 1);
+
+ for (int i = 0; i < numKeys; ++i) {
+ if (maxParallelism <= Byte.MAX_VALUE) {
+ dos.writeByte(i);
+ } else {
+ dos.writeShort(i);
+ }
+ dos.writeInt(i);
+ byte[] key = bos.toByteArray();
+ byte[] val = new byte[]{42};
+ rocksDB.put(handle, key, val);
+
+ bos.reset();
+ }
+ columnFamilyHandlesWithKeyCount.add(new Tuple2<>(handle, numKeys));
+ totalKeysExpected += numKeys;
+ }
+
+ int id = 0;
+ for (Tuple2<ColumnFamilyHandle, Integer> columnFamilyHandle : columnFamilyHandlesWithKeyCount) {
+ rocksIteratorsWithKVStateId.add(new Tuple2<>(rocksDB.newIterator(columnFamilyHandle.f0), id));
+ ++id;
+ }
+
+ RocksDBKeyedStateBackend.RocksDBMergeIterator mergeIterator = new RocksDBKeyedStateBackend.RocksDBMergeIterator(rocksIteratorsWithKVStateId, maxParallelism <= Byte.MAX_VALUE ? 1 : 2);
+
+ int prevKVState = -1;
+ int prevKey = -1;
+ int prevKeyGroup = -1;
+ int totalKeysActual = 0;
+
+ while (mergeIterator.isValid()) {
+ ByteBuffer bb = ByteBuffer.wrap(mergeIterator.key());
+
+ int keyGroup = maxParallelism > Byte.MAX_VALUE ? bb.getShort() : bb.get();
+ int key = bb.getInt();
+
+ Assert.assertTrue(keyGroup >= prevKeyGroup);
+ Assert.assertTrue(key >= prevKey);
+ Assert.assertEquals(prevKeyGroup != keyGroup, mergeIterator.isNewKeyGroup());
+ Assert.assertEquals(prevKVState != mergeIterator.kvStateId(), mergeIterator.isNewKeyValueState());
+
+ prevKeyGroup = keyGroup;
+ prevKVState = mergeIterator.kvStateId();
+
+ //System.out.println(keyGroup + " " + key + " " + mergeIterator.kvStateId());
+ mergeIterator.next();
+ ++totalKeysActual;
+ }
+
+ Assert.assertEquals(totalKeysExpected, totalKeysActual);
+
+ for (Tuple2<ColumnFamilyHandle, Integer> handleWithCount : columnFamilyHandlesWithKeyCount) {
+ rocksDB.dropColumnFamily(handleWithCount.f0);
+ }
+ } finally {
+ rocksDB.close();
+ }
+ }
+
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
index 6f4a983..acf6cb8 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
@@ -1,322 +1,368 @@
-///*
-// * 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.contrib.streaming.state;
-//
-//import org.apache.commons.io.FileUtils;
-//import org.apache.flink.api.common.JobID;
-//import org.apache.flink.api.common.TaskInfo;
-//import org.apache.flink.api.common.state.ValueStateDescriptor;
-//import org.apache.flink.api.common.typeutils.TypeSerializer;
-//import org.apache.flink.api.common.typeutils.base.IntSerializer;
-//import org.apache.flink.runtime.execution.Environment;
-//import org.apache.flink.runtime.io.disk.iomanager.IOManager;
-//import org.apache.flink.runtime.state.AbstractStateBackend;
-//
-//import org.apache.flink.runtime.state.VoidNamespace;
-//import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-//import org.apache.flink.util.OperatingSystem;
-//import org.junit.Assume;
-//import org.junit.Before;
-//import org.junit.Test;
-//
-//import org.rocksdb.ColumnFamilyOptions;
-//import org.rocksdb.CompactionStyle;
-//import org.rocksdb.DBOptions;
-//
-//import java.io.File;
-//import java.util.UUID;
-//
-//import static org.junit.Assert.*;
-//import static org.mockito.Mockito.*;
-//
-///**
-// * Tests for configuring the RocksDB State Backend
-// */
-//@SuppressWarnings("serial")
-//public class RocksDBStateBackendConfigTest {
-//
-// private static final String TEMP_URI = new File(System.getProperty("java.io.tmpdir")).toURI().toString();
-//
-// @Before
-// public void checkOperatingSystem() {
-// Assume.assumeTrue("This test can't run successfully on Windows.", !OperatingSystem.isWindows());
-// }
-//
-// // ------------------------------------------------------------------------
-// // RocksDB local file directory
-// // ------------------------------------------------------------------------
-//
-// @Test
-// public void testSetDbPath() throws Exception {
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-//
-// assertNull(rocksDbBackend.getDbStoragePaths());
-//
-// rocksDbBackend.setDbStoragePath("/abc/def");
-// assertArrayEquals(new String[] { "/abc/def" }, rocksDbBackend.getDbStoragePaths());
-//
-// rocksDbBackend.setDbStoragePath(null);
-// assertNull(rocksDbBackend.getDbStoragePaths());
-//
-// rocksDbBackend.setDbStoragePaths("/abc/def", "/uvw/xyz");
-// assertArrayEquals(new String[] { "/abc/def", "/uvw/xyz" }, rocksDbBackend.getDbStoragePaths());
-//
-// //noinspection NullArgumentToVariableArgMethod
-// rocksDbBackend.setDbStoragePaths(null);
-// assertNull(rocksDbBackend.getDbStoragePaths());
-// }
-//
-// @Test(expected = IllegalArgumentException.class)
-// public void testSetNullPaths() throws Exception {
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-// rocksDbBackend.setDbStoragePaths();
-// }
-//
-// @Test(expected = IllegalArgumentException.class)
-// public void testNonFileSchemePath() throws Exception {
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-// rocksDbBackend.setDbStoragePath("hdfs:///some/path/to/perdition");
-// }
-//
-// // ------------------------------------------------------------------------
-// // RocksDB local file automatic from temp directories
-// // ------------------------------------------------------------------------
-//
-// @Test
-// public void testUseTempDirectories() throws Exception {
-// File dir1 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
-// File dir2 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
-//
-// File[] tempDirs = new File[] { dir1, dir2 };
-//
-// try {
-// assertTrue(dir1.mkdirs());
-// assertTrue(dir2.mkdirs());
-//
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-// assertNull(rocksDbBackend.getDbStoragePaths());
-//
-// rocksDbBackend.initializeForJob(getMockEnvironment(tempDirs), "foobar", IntSerializer.INSTANCE);
-// assertArrayEquals(tempDirs, rocksDbBackend.getStoragePaths());
-// }
-// finally {
-// FileUtils.deleteDirectory(dir1);
-// FileUtils.deleteDirectory(dir2);
-// }
-// }
-//
-// // ------------------------------------------------------------------------
-// // RocksDB local file directory initialization
-// // ------------------------------------------------------------------------
-//
-// @Test
-// public void testFailWhenNoLocalStorageDir() throws Exception {
-// File targetDir = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
-// try {
-// assertTrue(targetDir.mkdirs());
-//
-// if (!targetDir.setWritable(false, false)) {
-// System.err.println("Cannot execute 'testFailWhenNoLocalStorageDir' because cannot mark directory non-writable");
-// return;
-// }
-//
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-// rocksDbBackend.setDbStoragePath(targetDir.getAbsolutePath());
-//
-// boolean hasFailure = false;
-// try {
-// rocksDbBackend.initializeForJob(getMockEnvironment(), "foobar", IntSerializer.INSTANCE);
-// }
-// catch (Exception e) {
-// assertTrue(e.getMessage().contains("No local storage directories available"));
-// assertTrue(e.getMessage().contains(targetDir.getAbsolutePath()));
-// hasFailure = true;
-// }
-// assertTrue("We must see a failure because no storaged directory is feasible.", hasFailure);
-// }
-// finally {
-// //noinspection ResultOfMethodCallIgnored
-// targetDir.setWritable(true, false);
-// FileUtils.deleteDirectory(targetDir);
-// }
-// }
-//
-// @Test
-// public void testContinueOnSomeDbDirectoriesMissing() throws Exception {
-// File targetDir1 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
-// File targetDir2 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
-//
-// try {
-// assertTrue(targetDir1.mkdirs());
-// assertTrue(targetDir2.mkdirs());
-//
-// if (!targetDir1.setWritable(false, false)) {
-// System.err.println("Cannot execute 'testContinueOnSomeDbDirectoriesMissing' because cannot mark directory non-writable");
-// return;
-// }
-//
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-// rocksDbBackend.setDbStoragePaths(targetDir1.getAbsolutePath(), targetDir2.getAbsolutePath());
-//
-// try {
-// rocksDbBackend.initializeForJob(getMockEnvironment(), "foobar", IntSerializer.INSTANCE);
-//
-// // actually get a state to see whether we can write to the storage directory
-// rocksDbBackend.getPartitionedState(
-// VoidNamespace.INSTANCE,
-// VoidNamespaceSerializer.INSTANCE,
-// new ValueStateDescriptor<>("test", String.class, ""));
-// }
-// catch (Exception e) {
-// e.printStackTrace();
-// fail("Backend initialization failed even though some paths were available");
-// }
-// } finally {
-// //noinspection ResultOfMethodCallIgnored
-// targetDir1.setWritable(true, false);
-// FileUtils.deleteDirectory(targetDir1);
-// FileUtils.deleteDirectory(targetDir2);
-// }
-// }
-//
-// // ------------------------------------------------------------------------
-// // RocksDB Options
-// // ------------------------------------------------------------------------
-//
-// @Test
-// public void testPredefinedOptions() throws Exception {
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-//
-// assertEquals(PredefinedOptions.DEFAULT, rocksDbBackend.getPredefinedOptions());
-//
-// rocksDbBackend.setPredefinedOptions(PredefinedOptions.SPINNING_DISK_OPTIMIZED);
-// assertEquals(PredefinedOptions.SPINNING_DISK_OPTIMIZED, rocksDbBackend.getPredefinedOptions());
-//
-// DBOptions opt1 = rocksDbBackend.getDbOptions();
-// DBOptions opt2 = rocksDbBackend.getDbOptions();
-//
-// assertEquals(opt1, opt2);
-//
-// ColumnFamilyOptions columnOpt1 = rocksDbBackend.getColumnOptions();
-// ColumnFamilyOptions columnOpt2 = rocksDbBackend.getColumnOptions();
-//
-// assertEquals(columnOpt1, columnOpt2);
-//
-// assertEquals(CompactionStyle.LEVEL, columnOpt1.compactionStyle());
-// }
-//
-// @Test
-// public void testOptionsFactory() throws Exception {
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-//
-// rocksDbBackend.setOptions(new OptionsFactory() {
-// @Override
-// public DBOptions createDBOptions(DBOptions currentOptions) {
-// return currentOptions;
-// }
-//
-// @Override
-// public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions) {
-// return currentOptions.setCompactionStyle(CompactionStyle.FIFO);
-// }
-// });
-//
-// assertNotNull(rocksDbBackend.getOptions());
-// assertEquals(CompactionStyle.FIFO, rocksDbBackend.getColumnOptions().compactionStyle());
-// }
-//
-// @Test
-// public void testPredefinedAndOptionsFactory() throws Exception {
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-//
-// assertEquals(PredefinedOptions.DEFAULT, rocksDbBackend.getPredefinedOptions());
-//
-// rocksDbBackend.setPredefinedOptions(PredefinedOptions.SPINNING_DISK_OPTIMIZED);
-// rocksDbBackend.setOptions(new OptionsFactory() {
-// @Override
-// public DBOptions createDBOptions(DBOptions currentOptions) {
-// return currentOptions;
-// }
-//
-// @Override
-// public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions) {
-// return currentOptions.setCompactionStyle(CompactionStyle.UNIVERSAL);
-// }
-// });
-//
-// assertEquals(PredefinedOptions.SPINNING_DISK_OPTIMIZED, rocksDbBackend.getPredefinedOptions());
-// assertNotNull(rocksDbBackend.getOptions());
-// assertEquals(CompactionStyle.UNIVERSAL, rocksDbBackend.getColumnOptions().compactionStyle());
-// }
-//
-// @Test
-// public void testPredefinedOptionsEnum() {
-// for (PredefinedOptions o : PredefinedOptions.values()) {
-// DBOptions opt = o.createDBOptions();
-// try {
-// assertNotNull(opt);
-// } finally {
-// opt.dispose();
-// }
-// }
-// }
-//
-// // ------------------------------------------------------------------------
-// // Contained Non-partitioned State Backend
-// // ------------------------------------------------------------------------
-//
-// @Test
-// public void testCallsForwardedToNonPartitionedBackend() throws Exception {
-// AbstractStateBackend nonPartBackend = mock(AbstractStateBackend.class);
-// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI, nonPartBackend);
-//
-// rocksDbBackend.initializeForJob(getMockEnvironment(), "foo", IntSerializer.INSTANCE);
-// verify(nonPartBackend, times(1)).initializeForJob(any(Environment.class), anyString(), any(TypeSerializer.class));
-//
-// rocksDbBackend.disposeAllStateForCurrentJob();
-// verify(nonPartBackend, times(1)).disposeAllStateForCurrentJob();
-//
-// rocksDbBackend.close();
-// verify(nonPartBackend, times(1)).close();
-// }
-//
-// // ------------------------------------------------------------------------
-// // Utilities
-// // ------------------------------------------------------------------------
-//
-// private static Environment getMockEnvironment() {
-// return getMockEnvironment(new File[] { new File(System.getProperty("java.io.tmpdir")) });
-// }
-//
-// private static Environment getMockEnvironment(File[] tempDirs) {
-// IOManager ioMan = mock(IOManager.class);
-// when(ioMan.getSpillingDirectories()).thenReturn(tempDirs);
-//
-// Environment env = mock(Environment.class);
-// when(env.getJobID()).thenReturn(new JobID());
-// when(env.getUserClassLoader()).thenReturn(RocksDBStateBackendConfigTest.class.getClassLoader());
-// when(env.getIOManager()).thenReturn(ioMan);
-//
-// TaskInfo taskInfo = mock(TaskInfo.class);
-// when(env.getTaskInfo()).thenReturn(taskInfo);
-//
-// when(taskInfo.getIndexOfThisSubtask()).thenReturn(0);
-// return env;
-// }
-//}
+/*
+ * 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.contrib.streaming.state;
+
+import org.apache.commons.io.FileUtils;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.TaskInfo;
+import org.apache.flink.api.common.typeutils.base.IntSerializer;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.io.disk.iomanager.IOManager;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.query.KvStateRegistry;
+
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.util.OperatingSystem;
+import org.junit.Assume;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+import org.junit.rules.TemporaryFolder;
+import org.rocksdb.ColumnFamilyOptions;
+import org.rocksdb.CompactionStyle;
+import org.rocksdb.DBOptions;
+
+import java.io.File;
+
+import static org.hamcrest.CoreMatchers.anyOf;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.startsWith;
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
+
+
+/**
+ * Tests for configuring the RocksDB State Backend
+ */
+@SuppressWarnings("serial")
+public class RocksDBStateBackendConfigTest {
+
+
+ @Rule
+ public TemporaryFolder tempFolder = new TemporaryFolder();
+
+ @Before
+ public void checkOperatingSystem() {
+ Assume.assumeTrue("This test can't run successfully on Windows.", !OperatingSystem.isWindows());
+ }
+
+ // ------------------------------------------------------------------------
+ // RocksDB local file directory
+ // ------------------------------------------------------------------------
+
+ @Test
+ public void testSetDbPath() throws Exception {
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ File testDir1 = tempFolder.newFolder();
+ File testDir2 = tempFolder.newFolder();
+
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
+
+ assertNull(rocksDbBackend.getDbStoragePaths());
+
+ rocksDbBackend.setDbStoragePath(testDir1.getAbsolutePath());
+ assertArrayEquals(new String[] { testDir1.getAbsolutePath() }, rocksDbBackend.getDbStoragePaths());
+
+ rocksDbBackend.setDbStoragePath(null);
+ assertNull(rocksDbBackend.getDbStoragePaths());
+
+ rocksDbBackend.setDbStoragePaths(testDir1.getAbsolutePath(), testDir2.getAbsolutePath());
+ assertArrayEquals(new String[] { testDir1.getAbsolutePath(), testDir2.getAbsolutePath() }, rocksDbBackend.getDbStoragePaths());
+
+ Environment env = getMockEnvironment(new File[] {});
+ RocksDBKeyedStateBackend<Integer> keyedBackend = (RocksDBKeyedStateBackend<Integer>) rocksDbBackend.createKeyedStateBackend(env,
+ env.getJobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ env.getTaskKvStateRegistry());
+
+
+ File instanceBasePath = keyedBackend.getInstanceBasePath();
+ assertThat(instanceBasePath.getAbsolutePath(), anyOf(startsWith(testDir1.getAbsolutePath()), startsWith(testDir2.getAbsolutePath())));
+
+ //noinspection NullArgumentToVariableArgMethod
+ rocksDbBackend.setDbStoragePaths(null);
+ assertNull(rocksDbBackend.getDbStoragePaths());
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testSetNullPaths() throws Exception {
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
+ rocksDbBackend.setDbStoragePaths();
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testNonFileSchemePath() throws Exception {
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
+ rocksDbBackend.setDbStoragePath("hdfs:///some/path/to/perdition");
+ }
+
+ // ------------------------------------------------------------------------
+ // RocksDB local file automatic from temp directories
+ // ------------------------------------------------------------------------
+
+ /**
+ * This tests whether the RocksDB backends uses the temp directories that are provided
+ * from the {@link Environment} when no db storage path is set.
+ *
+ * @throws Exception
+ */
+ @Test
+ public void testUseTempDirectories() throws Exception {
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
+
+ File dir1 = tempFolder.newFolder();
+ File dir2 = tempFolder.newFolder();
+
+ File[] tempDirs = new File[] { dir1, dir2 };
+
+ assertNull(rocksDbBackend.getDbStoragePaths());
+
+ Environment env = getMockEnvironment(tempDirs);
+ RocksDBKeyedStateBackend<Integer> keyedBackend = (RocksDBKeyedStateBackend<Integer>) rocksDbBackend.createKeyedStateBackend(env,
+ env.getJobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ env.getTaskKvStateRegistry());
+
+
+ File instanceBasePath = keyedBackend.getInstanceBasePath();
+ assertThat(instanceBasePath.getAbsolutePath(), anyOf(startsWith(dir1.getAbsolutePath()), startsWith(dir2.getAbsolutePath())));
+ }
+
+ // ------------------------------------------------------------------------
+ // RocksDB local file directory initialization
+ // ------------------------------------------------------------------------
+
+ @Test
+ public void testFailWhenNoLocalStorageDir() throws Exception {
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
+ File targetDir = tempFolder.newFolder();
+
+ try {
+ if (!targetDir.setWritable(false, false)) {
+ System.err.println("Cannot execute 'testFailWhenNoLocalStorageDir' because cannot mark directory non-writable");
+ return;
+ }
+
+ rocksDbBackend.setDbStoragePath(targetDir.getAbsolutePath());
+
+ boolean hasFailure = false;
+ try {
+ Environment env = getMockEnvironment();
+ rocksDbBackend.createKeyedStateBackend(
+ env,
+ env.getJobID(),
+ "foobar",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()));
+ }
+ catch (Exception e) {
+ assertTrue(e.getMessage().contains("No local storage directories available"));
+ assertTrue(e.getMessage().contains(targetDir.getAbsolutePath()));
+ hasFailure = true;
+ }
+ assertTrue("We must see a failure because no storaged directory is feasible.", hasFailure);
+ }
+ finally {
+ //noinspection ResultOfMethodCallIgnored
+ targetDir.setWritable(true, false);
+ FileUtils.deleteDirectory(targetDir);
+ }
+ }
+
+ @Test
+ public void testContinueOnSomeDbDirectoriesMissing() throws Exception {
+ File targetDir1 = tempFolder.newFolder();
+ File targetDir2 = tempFolder.newFolder();
+
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
+
+ try {
+
+ if (!targetDir1.setWritable(false, false)) {
+ System.err.println("Cannot execute 'testContinueOnSomeDbDirectoriesMissing' because cannot mark directory non-writable");
+ return;
+ }
+
+ rocksDbBackend.setDbStoragePaths(targetDir1.getAbsolutePath(), targetDir2.getAbsolutePath());
+
+ try {
+ Environment env = getMockEnvironment();
+ rocksDbBackend.createKeyedStateBackend(
+ env,
+ env.getJobID(),
+ "foobar",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()));
+ }
+ catch (Exception e) {
+ e.printStackTrace();
+ fail("Backend initialization failed even though some paths were available");
+ }
+ } finally {
+ //noinspection ResultOfMethodCallIgnored
+ targetDir1.setWritable(true, false);
+ FileUtils.deleteDirectory(targetDir1);
+ FileUtils.deleteDirectory(targetDir2);
+ }
+ }
+
+ // ------------------------------------------------------------------------
+ // RocksDB Options
+ // ------------------------------------------------------------------------
+
+ @Test
+ public void testPredefinedOptions() throws Exception {
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
+
+ assertEquals(PredefinedOptions.DEFAULT, rocksDbBackend.getPredefinedOptions());
+
+ rocksDbBackend.setPredefinedOptions(PredefinedOptions.SPINNING_DISK_OPTIMIZED);
+ assertEquals(PredefinedOptions.SPINNING_DISK_OPTIMIZED, rocksDbBackend.getPredefinedOptions());
+
+ DBOptions opt1 = rocksDbBackend.getDbOptions();
+ DBOptions opt2 = rocksDbBackend.getDbOptions();
+
+ assertEquals(opt1, opt2);
+
+ ColumnFamilyOptions columnOpt1 = rocksDbBackend.getColumnOptions();
+ ColumnFamilyOptions columnOpt2 = rocksDbBackend.getColumnOptions();
+
+ assertEquals(columnOpt1, columnOpt2);
+
+ assertEquals(CompactionStyle.LEVEL, columnOpt1.compactionStyle());
+ }
+
+ @Test
+ public void testOptionsFactory() throws Exception {
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
+
+ rocksDbBackend.setOptions(new OptionsFactory() {
+ @Override
+ public DBOptions createDBOptions(DBOptions currentOptions) {
+ return currentOptions;
+ }
+
+ @Override
+ public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions) {
+ return currentOptions.setCompactionStyle(CompactionStyle.FIFO);
+ }
+ });
+
+ assertNotNull(rocksDbBackend.getOptions());
+ assertEquals(CompactionStyle.FIFO, rocksDbBackend.getColumnOptions().compactionStyle());
+ }
+
+ @Test
+ public void testPredefinedAndOptionsFactory() throws Exception {
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);
+
+ assertEquals(PredefinedOptions.DEFAULT, rocksDbBackend.getPredefinedOptions());
+
+ rocksDbBackend.setPredefinedOptions(PredefinedOptions.SPINNING_DISK_OPTIMIZED);
+ rocksDbBackend.setOptions(new OptionsFactory() {
+ @Override
+ public DBOptions createDBOptions(DBOptions currentOptions) {
+ return currentOptions;
+ }
+
+ @Override
+ public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions) {
+ return currentOptions.setCompactionStyle(CompactionStyle.UNIVERSAL);
+ }
+ });
+
+ assertEquals(PredefinedOptions.SPINNING_DISK_OPTIMIZED, rocksDbBackend.getPredefinedOptions());
+ assertNotNull(rocksDbBackend.getOptions());
+ assertEquals(CompactionStyle.UNIVERSAL, rocksDbBackend.getColumnOptions().compactionStyle());
+ }
+
+ @Test
+ public void testPredefinedOptionsEnum() {
+ for (PredefinedOptions o : PredefinedOptions.values()) {
+ DBOptions opt = o.createDBOptions();
+ try {
+ assertNotNull(opt);
+ } finally {
+ opt.dispose();
+ }
+ }
+ }
+
+ // ------------------------------------------------------------------------
+ // Contained Non-partitioned State Backend
+ // ------------------------------------------------------------------------
+
+ @Test
+ public void testCallsForwardedToNonPartitionedBackend() throws Exception {
+ AbstractStateBackend nonPartBackend = mock(AbstractStateBackend.class);
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath, nonPartBackend);
+
+ Environment env = getMockEnvironment();
+ rocksDbBackend.createStreamFactory(env.getJobID(), "foobar");
+
+ verify(nonPartBackend, times(1)).createStreamFactory(any(JobID.class), anyString());
+ }
+
+ // ------------------------------------------------------------------------
+ // Utilities
+ // ------------------------------------------------------------------------
+
+ private static Environment getMockEnvironment() {
+ return getMockEnvironment(new File[] { new File(System.getProperty("java.io.tmpdir")) });
+ }
+
+ private static Environment getMockEnvironment(File[] tempDirs) {
+ IOManager ioMan = mock(IOManager.class);
+ when(ioMan.getSpillingDirectories()).thenReturn(tempDirs);
+
+ Environment env = mock(Environment.class);
+ when(env.getJobID()).thenReturn(new JobID());
+ when(env.getUserClassLoader()).thenReturn(RocksDBStateBackendConfigTest.class.getClassLoader());
+ when(env.getIOManager()).thenReturn(ioMan);
+ when(env.getTaskKvStateRegistry()).thenReturn(new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
+
+ TaskInfo taskInfo = mock(TaskInfo.class);
+ when(env.getTaskInfo()).thenReturn(taskInfo);
+
+ when(taskInfo.getIndexOfThisSubtask()).thenReturn(0);
+ return env;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AbstractAsyncIOCallable.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AbstractAsyncIOCallable.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AbstractAsyncIOCallable.java
new file mode 100644
index 0000000..989e868
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AbstractAsyncIOCallable.java
@@ -0,0 +1,157 @@
+/*
+ * 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.io.async;
+
+import java.io.Closeable;
+import java.io.IOException;
+
+/**
+ * The abstract class encapsulates the lifecycle and execution strategy for asynchronous IO operations
+ *
+ * @param <V> return type of the asynchronous call
+ * @param <D> type of the IO handle
+ */
+public abstract class AbstractAsyncIOCallable<V, D extends Closeable> implements StoppableCallbackCallable<V> {
+
+ private volatile boolean stopped;
+
+ /**
+ * Closable handle to IO, e.g. an InputStream
+ */
+ private volatile D ioHandle;
+
+ /**
+ * Stores exception that might happen during close
+ */
+ private volatile IOException stopException;
+
+
+ public AbstractAsyncIOCallable() {
+ this.stopped = false;
+ }
+
+ /**
+ * This method implements the strategy for the actual IO operation:
+ *
+ * 1) Open the IO handle
+ * 2) Perform IO operation
+ * 3) Close IO handle
+ *
+ * @return Result of the IO operation, e.g. a deserialized object.
+ * @throws Exception exception that happened during the call.
+ */
+ @Override
+ public V call() throws Exception {
+
+ synchronized (this) {
+ if (isStopped()) {
+ throw new IOException("Task was already stopped. No I/O handle opened.");
+ }
+
+ ioHandle = openIOHandle();
+ }
+
+ try {
+
+ return performOperation();
+
+ } finally {
+ closeIOHandle();
+ }
+
+ }
+
+ /**
+ * Open the IO Handle (e.g. a stream) on which the operation will be performed.
+ *
+ * @return the opened IO handle that implements #Closeable
+ * @throws Exception
+ */
+ protected abstract D openIOHandle() throws Exception;
+
+ /**
+ * Implements the actual IO operation on the opened IO handle.
+ *
+ * @return Result of the IO operation
+ * @throws Exception
+ */
+ protected abstract V performOperation() throws Exception;
+
+ /**
+ * Stops the I/O operation by closing the I/O handle. If an exception is thrown on close, it can be accessed via
+ * #getStopException().
+ */
+ @Override
+ public void stop() {
+ closeIOHandle();
+ }
+
+ private synchronized void closeIOHandle() {
+
+ if (!stopped) {
+ stopped = true;
+
+ final D handle = ioHandle;
+ if (handle != null) {
+ try {
+ handle.close();
+ } catch (IOException ex) {
+ stopException = ex;
+ }
+ }
+ }
+ }
+
+ /**
+ * Returns the IO handle.
+ * @return the IO handle
+ */
+ protected D getIoHandle() {
+ return ioHandle;
+ }
+
+ /**
+ * Optional callback that subclasses can implement. This is called when the callable method completed, e.g. because
+ * it finished or was stopped.
+ */
+ @Override
+ public void done() {
+ //optional callback hook
+ }
+
+ /**
+ * Check if the IO operation is stopped
+ *
+ * @return true if stop() was called
+ */
+ @Override
+ public boolean isStopped() {
+ return stopped;
+ }
+
+ /**
+ * Returns Exception that might happen on stop.
+ *
+ * @return Potential Exception that happened open stopping.
+ */
+ @Override
+ public IOException getStopException() {
+ return stopException;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncDoneCallback.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncDoneCallback.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncDoneCallback.java
new file mode 100644
index 0000000..13d9057
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncDoneCallback.java
@@ -0,0 +1,31 @@
+/*
+ * 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.io.async;
+
+/**
+ * Callback for an asynchronous operation that is called on termination
+ */
+public interface AsyncDoneCallback {
+
+ /**
+ * the callback
+ */
+ void done();
+
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncStoppable.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncStoppable.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncStoppable.java
new file mode 100644
index 0000000..560e56a
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncStoppable.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.io.async;
+
+import java.io.IOException;
+
+/**
+ * An asynchronous operation that can be stopped.
+ */
+public interface AsyncStoppable {
+
+ /**
+ * Stop the operation
+ */
+ void stop();
+
+ /**
+ * Check whether the operation is stopped
+ *
+ * @return true iff operation is stopped
+ */
+ boolean isStopped();
+
+ /**
+ * Delivers Exception that might happen during {@link #stop()}
+ *
+ * @return Exception that can happen during stop
+ */
+ IOException getStopException();
+
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncStoppableTaskWithCallback.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncStoppableTaskWithCallback.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncStoppableTaskWithCallback.java
new file mode 100644
index 0000000..8316e4f
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/AsyncStoppableTaskWithCallback.java
@@ -0,0 +1,55 @@
+/*
+ * 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.io.async;
+
+import org.apache.flink.util.Preconditions;
+
+import java.util.concurrent.FutureTask;
+
+/**
+ * @param <V> return type of the callable function
+ */
+public class AsyncStoppableTaskWithCallback<V> extends FutureTask<V> {
+
+ protected final StoppableCallbackCallable<V> stoppableCallbackCallable;
+
+ public AsyncStoppableTaskWithCallback(StoppableCallbackCallable<V> callable) {
+ super(Preconditions.checkNotNull(callable));
+ this.stoppableCallbackCallable = callable;
+ }
+
+ @Override
+ public boolean cancel(boolean mayInterruptIfRunning) {
+
+ if (mayInterruptIfRunning) {
+ stoppableCallbackCallable.stop();
+ }
+
+ return super.cancel(mayInterruptIfRunning);
+ }
+
+ @Override
+ protected void done() {
+ stoppableCallbackCallable.done();
+ }
+
+ public static <V> AsyncStoppableTaskWithCallback<V> from(StoppableCallbackCallable<V> callable) {
+ return new AsyncStoppableTaskWithCallback<>(callable);
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/StoppableCallbackCallable.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/StoppableCallbackCallable.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/StoppableCallbackCallable.java
new file mode 100644
index 0000000..d459316
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/async/StoppableCallbackCallable.java
@@ -0,0 +1,30 @@
+/*
+ * 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.io.async;
+
+import java.util.concurrent.Callable;
+
+/**
+ * A {@link Callable} that can be stopped and offers a callback on termination.
+ *
+ * @param <V> return value of the call operation.
+ */
+public interface StoppableCallbackCallable<V> extends Callable<V>, AsyncStoppable, AsyncDoneCallback {
+
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemCheckpointStreamFactory.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemCheckpointStreamFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemCheckpointStreamFactory.java
index 4801d85..6f0a814 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemCheckpointStreamFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemCheckpointStreamFactory.java
@@ -71,7 +71,7 @@ public class MemCheckpointStreamFactory implements CheckpointStreamFactory {
/**
* A {@code CheckpointStateOutputStream} that writes into a byte array.
*/
- public static final class MemoryCheckpointOutputStream extends CheckpointStateOutputStream {
+ public static class MemoryCheckpointOutputStream extends CheckpointStateOutputStream {
private final ByteArrayOutputStreamWithPos os = new ByteArrayOutputStreamWithPos();
@@ -86,13 +86,13 @@ public class MemCheckpointStreamFactory implements CheckpointStreamFactory {
}
@Override
- public void write(int b) {
+ public void write(int b) throws IOException {
os.write(b);
isEmpty = false;
}
@Override
- public void write(byte[] b, int off, int len) {
+ public void write(byte[] b, int off, int len) throws IOException {
os.write(b, off, len);
isEmpty = false;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/addd0842/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
index 02579aa..13f650c 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
@@ -66,6 +66,7 @@ import java.util.concurrent.Executors;
import java.util.concurrent.RunnableFuture;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
/**
* Base class for all streaming tasks. A task is the unit of local processing that is deployed
@@ -176,8 +177,12 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
private long lastCheckpointSize = 0;
+ /** Thread pool for async snapshot workers */
private ExecutorService asyncOperationsThreadPool;
+ /** Timeout to await the termination of the thread pool in milliseconds */
+ private long threadPoolTerminationTimeout = 0L;
+
// ------------------------------------------------------------------------
// Life cycle methods for specific implementations
// ------------------------------------------------------------------------
@@ -441,6 +446,10 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
if (!asyncOperationsThreadPool.isShutdown()) {
asyncOperationsThreadPool.shutdownNow();
}
+
+ if(threadPoolTerminationTimeout > 0L) {
+ asyncOperationsThreadPool.awaitTermination(threadPoolTerminationTimeout, TimeUnit.MILLISECONDS);
+ }
}
/**
@@ -861,6 +870,15 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
};
}
+ /**
+ * Sets a timeout for the async thread pool. Default should always be 0 to avoid blocking restarts of task.
+ *
+ * @param threadPoolTerminationTimeout timeout for the async thread pool in milliseconds
+ */
+ public void setThreadPoolTerminationTimeout(long threadPoolTerminationTimeout) {
+ this.threadPoolTerminationTimeout = threadPoolTerminationTimeout;
+ }
+
// ------------------------------------------------------------------------
/**
[27/27] flink git commit: [FLINK-3755] Extended
EventTimeWindowCheckpointITCase to test the boundaries of maxParallelism.
Posted by al...@apache.org.
[FLINK-3755] Extended EventTimeWindowCheckpointITCase to test the boundaries of maxParallelism.
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/f7ef82b3
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/f7ef82b3
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/f7ef82b3
Branch: refs/heads/master
Commit: f7ef82b38ce9361b7d705199d71ac684ba4a76c3
Parents: f44b57c
Author: Stefan Richter <s....@data-artisans.com>
Authored: Tue Aug 30 14:59:12 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:10:02 2016 +0200
----------------------------------------------------------------------
.../EventTimeWindowCheckpointingITCase.java | 12 +++++++++++-
1 file changed, 11 insertions(+), 1 deletion(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/f7ef82b3/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
index 9f8ab90..fa5339d 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
@@ -196,7 +196,16 @@ public class EventTimeWindowCheckpointingITCase extends TestLogger {
}
@Test
- public void testTumblingTimeWindowWithKVState() {
+ public void testTumblingTimeWindowWithKVStateMinMaxParallelism() {
+ doTestTumblingTimeWindowWithKVState(PARALLELISM);
+ }
+
+ @Test
+ public void testTumblingTimeWindowWithKVStateMaxMaxParallelism() {
+ doTestTumblingTimeWindowWithKVState(1 << 15);
+ }
+
+ public void doTestTumblingTimeWindowWithKVState(int maxParallelism) {
final int NUM_ELEMENTS_PER_KEY = 3000;
final int WINDOW_SIZE = 100;
final int NUM_KEYS = 100;
@@ -207,6 +216,7 @@ public class EventTimeWindowCheckpointingITCase extends TestLogger {
"localhost", cluster.getLeaderRPCPort());
env.setParallelism(PARALLELISM);
+ env.setMaxParallelism(maxParallelism);
env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime);
env.enableCheckpointing(100);
env.setRestartStrategy(RestartStrategies.fixedDelayRestart(3, 0));
[07/27] flink git commit: [FLINK-3761] Refactor State Backends/Make
Keyed State Key-Group Aware
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
index 59ecd15..718c0c7 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
@@ -24,13 +24,12 @@ import org.apache.flink.api.common.state.State;
import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.functions.KeySelector;
-import org.apache.flink.metrics.Counter;
import org.apache.flink.core.fs.FSDataInputStream;
import org.apache.flink.core.fs.FSDataOutputStream;
-import org.apache.flink.runtime.state.AsynchronousKvStateSnapshot;
+import org.apache.flink.metrics.Counter;
import org.apache.flink.metrics.MetricGroup;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.KvStateSnapshot;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
import org.apache.flink.streaming.api.graph.StreamConfig;
@@ -38,14 +37,9 @@ import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.util.InstantiationUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import java.io.DataInputStream;
-import java.io.DataOutputStream;
-import java.util.HashMap;
-import java.util.Set;
import java.util.concurrent.ScheduledFuture;
/**
@@ -96,9 +90,10 @@ public abstract class AbstractStreamOperator<OUT>
private transient KeySelector<?, ?> stateKeySelector1;
private transient KeySelector<?, ?> stateKeySelector2;
- /** The state backend that stores the state and checkpoints for this task */
- private transient AbstractStateBackend stateBackend;
- protected MetricGroup metrics;
+ /** Backend for keyed state. This might be empty if we're not on a keyed stream. */
+ private transient KeyedStateBackend<?> keyedStateBackend;
+
+ protected transient MetricGroup metrics;
// ------------------------------------------------------------------------
// Life Cycle
@@ -116,16 +111,6 @@ public abstract class AbstractStreamOperator<OUT>
stateKeySelector1 = config.getStatePartitioner(0, getUserCodeClassloader());
stateKeySelector2 = config.getStatePartitioner(1, getUserCodeClassloader());
-
- try {
- TypeSerializer<Object> keySerializer = config.getStateKeySerializer(getUserCodeClassloader());
- // if the keySerializer is null we still need to create the state backend
- // for the non-partitioned state features it provides, such as the state output streams
- String operatorIdentifier = getClass().getSimpleName() + "_" + config.getVertexID() + "_" + runtimeContext.getIndexOfThisSubtask();
- stateBackend = container.createStateBackend(operatorIdentifier, keySerializer);
- } catch (Exception e) {
- throw new RuntimeException("Could not initialize state backend. ", e);
- }
}
public MetricGroup getMetricGroup() {
@@ -141,7 +126,27 @@ public abstract class AbstractStreamOperator<OUT>
* @throws Exception An exception in this method causes the operator to fail.
*/
@Override
- public void open() throws Exception {}
+ public void open() throws Exception {
+ try {
+ TypeSerializer<Object> keySerializer = config.getStateKeySerializer(getUserCodeClassloader());
+ // create a keyed state backend if there is keyed state, as indicated by the presence of a key serializer
+ if (null != keySerializer) {
+ ExecutionConfig execConf = container.getEnvironment().getExecutionConfig();;
+
+ KeyGroupRange subTaskKeyGroupRange = KeyGroupRange.computeKeyGroupRangeForOperatorIndex(
+ container.getEnvironment().getTaskInfo().getNumberOfKeyGroups(),
+ container.getEnvironment().getTaskInfo().getNumberOfParallelSubtasks(),
+ container.getIndexInSubtaskGroup());
+
+ keyedStateBackend = container.createKeyedStateBackend(
+ keySerializer,
+ container.getConfiguration().getKeyGroupAssigner(getUserCodeClassloader()),
+ subTaskKeyGroupRange);
+ }
+ } catch (Exception e) {
+ throw new RuntimeException("Could not initialize keyed state backend.", e);
+ }
+ }
/**
* This method is called after all records have been added to the operators via the methods
@@ -166,69 +171,22 @@ public abstract class AbstractStreamOperator<OUT>
* that the operator has acquired.
*/
@Override
- public void dispose() {
- if (stateBackend != null) {
- try {
- stateBackend.close();
- stateBackend.discardState();
- } catch (Exception e) {
- throw new RuntimeException("Error while closing/disposing state backend.", e);
- }
+ public void dispose() throws Exception {
+ if (keyedStateBackend != null) {
+ keyedStateBackend.close();
}
}
-
- // ------------------------------------------------------------------------
- // Checkpointing
- // ------------------------------------------------------------------------
@Override
public void snapshotState(FSDataOutputStream out,
long checkpointId,
- long timestamp) throws Exception {
-
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> keyedState =
- stateBackend.snapshotPartitionedState(checkpointId,timestamp);
-
- // Materialize asynchronous snapshots, if any
- if (keyedState != null) {
- Set<String> keys = keyedState.keySet();
- for (String key: keys) {
- if (keyedState.get(key) instanceof AsynchronousKvStateSnapshot) {
- AsynchronousKvStateSnapshot<?, ?, ?, ?, ?> asyncHandle = (AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) keyedState.get(key);
- keyedState.put(key, asyncHandle.materialize());
- }
- }
- }
-
- byte[] serializedSnapshot = InstantiationUtil.serializeObject(keyedState);
-
- DataOutputStream dos = new DataOutputStream(out);
- dos.writeInt(serializedSnapshot.length);
- dos.write(serializedSnapshot);
-
- dos.flush();
-
- }
+ long timestamp) throws Exception {}
@Override
- public void restoreState(FSDataInputStream in) throws Exception {
- DataInputStream dis = new DataInputStream(in);
- int size = dis.readInt();
- byte[] serializedSnapshot = new byte[size];
- dis.readFully(serializedSnapshot);
-
- HashMap<String, KvStateSnapshot> keyedState =
- InstantiationUtil.deserializeObject(serializedSnapshot, getUserCodeClassloader());
+ public void restoreState(FSDataInputStream in) throws Exception {}
- stateBackend.injectKeyValueStateSnapshots(keyedState);
- }
-
@Override
- public void notifyOfCompletedCheckpoint(long checkpointId) throws Exception {
- if (stateBackend != null) {
- stateBackend.notifyOfCompletedCheckpoint(checkpointId);
- }
- }
+ public void notifyOfCompletedCheckpoint(long checkpointId) throws Exception {}
// ------------------------------------------------------------------------
// Properties and Services
@@ -265,13 +223,14 @@ public abstract class AbstractStreamOperator<OUT>
return runtimeContext;
}
- public AbstractStateBackend getStateBackend() {
- return stateBackend;
+ @SuppressWarnings("rawtypes, unchecked")
+ public <K> KeyedStateBackend<K> getStateBackend() {
+ return (KeyedStateBackend<K>) keyedStateBackend;
}
/**
- * Register a timer callback. At the specified time the {@link Triggerable} will be invoked.
- * This call is guaranteed to not happen concurrently with method calls on the operator.
+ * Register a timer callback. At the specified time the provided {@link Triggerable} will
+ * be invoked. This call is guaranteed to not happen concurrently with method calls on the operator.
*
* @param time The absolute time in milliseconds.
* @param target The target to be triggered.
@@ -291,7 +250,7 @@ public abstract class AbstractStreamOperator<OUT>
* @throws Exception Thrown, if the state backend cannot create the key/value state.
*/
protected <S extends State> S getPartitionedState(StateDescriptor<S, ?> stateDescriptor) throws Exception {
- return getStateBackend().getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor);
+ return getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor);
}
/**
@@ -302,13 +261,13 @@ public abstract class AbstractStreamOperator<OUT>
*/
@SuppressWarnings("unchecked")
protected <S extends State, N> S getPartitionedState(N namespace, TypeSerializer<N> namespaceSerializer, StateDescriptor<S, ?> stateDescriptor) throws Exception {
- if (stateBackend != null) {
- return stateBackend.getPartitionedState(
- namespace,
- namespaceSerializer,
- stateDescriptor);
+ if (keyedStateBackend != null) {
+ return keyedStateBackend.getPartitionedState(
+ namespace,
+ namespaceSerializer,
+ stateDescriptor);
} else {
- throw new RuntimeException("Cannot create partitioned state. The key grouped state " +
+ throw new RuntimeException("Cannot create partitioned state. The keyed state " +
"backend has not been set. This indicates that the operator is not " +
"partitioned/keyed.");
}
@@ -335,15 +294,16 @@ public abstract class AbstractStreamOperator<OUT>
@SuppressWarnings({"unchecked", "rawtypes"})
public void setKeyContext(Object key) {
- if (stateBackend != null) {
+ if (keyedStateBackend != null) {
try {
- stateBackend.setCurrentKey(key);
+ // need to work around type restrictions
+ @SuppressWarnings("unchecked,rawtypes")
+ KeyedStateBackend rawBackend = (KeyedStateBackend) keyedStateBackend;
+
+ rawBackend.setCurrentKey(key);
} catch (Exception e) {
throw new RuntimeException("Exception occurred while setting the current key context.", e);
}
- } else {
- throw new RuntimeException("Could not set the current key context, because the " +
- "AbstractStateBackend has not been initialized.");
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java
index b1bc531..6ac73e7 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractUdfStreamOperator.java
@@ -18,8 +18,6 @@
package org.apache.flink.streaming.api.operators;
-import java.io.ObjectInputStream;
-import java.io.ObjectOutputStream;
import java.io.Serializable;
import org.apache.flink.annotation.PublicEvolving;
@@ -35,6 +33,7 @@ import org.apache.flink.streaming.api.checkpoint.Checkpointed;
import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
+import org.apache.flink.util.InstantiationUtil;
import static java.util.Objects.requireNonNull;
@@ -49,7 +48,9 @@ import static java.util.Objects.requireNonNull;
* The type of the user function
*/
@PublicEvolving
-public abstract class AbstractUdfStreamOperator<OUT, F extends Function> extends AbstractStreamOperator<OUT> implements OutputTypeConfigurable<OUT> {
+public abstract class AbstractUdfStreamOperator<OUT, F extends Function>
+ extends AbstractStreamOperator<OUT>
+ implements OutputTypeConfigurable<OUT> {
private static final long serialVersionUID = 1L;
@@ -100,16 +101,11 @@ public abstract class AbstractUdfStreamOperator<OUT, F extends Function> extends
}
@Override
- public void dispose() {
+ public void dispose() throws Exception {
super.dispose();
if (!functionsClosed) {
functionsClosed = true;
- try {
- FunctionUtils.closeFunction(userFunction);
- }
- catch (Throwable t) {
- LOG.error("Exception while closing user function while failing or canceling task", t);
- }
+ FunctionUtils.closeFunction(userFunction);
}
}
@@ -130,9 +126,7 @@ public abstract class AbstractUdfStreamOperator<OUT, F extends Function> extends
udfState = chkFunction.snapshotState(checkpointId, timestamp);
if (udfState != null) {
out.write(1);
- ObjectOutputStream os = new ObjectOutputStream(out);
- os.writeObject(udfState);
- os.flush();
+ InstantiationUtil.serializeObject(out, udfState);
} else {
out.write(0);
}
@@ -153,8 +147,7 @@ public abstract class AbstractUdfStreamOperator<OUT, F extends Function> extends
int hasUdfState = in.read();
if (hasUdfState == 1) {
- ObjectInputStream ois = new ObjectInputStream(in);
- Serializable functionState = (Serializable) ois.readObject();
+ Serializable functionState = InstantiationUtil.deserializeObject(in, getUserCodeClassloader());
if (functionState != null) {
try {
chkFunction.restoreState(functionState);
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java
index 3411a60..f1e8160 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamOperator.java
@@ -20,9 +20,9 @@ package org.apache.flink.streaming.api.operators;
import java.io.Serializable;
import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.core.fs.FSDataInputStream;
-import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
@@ -84,7 +84,7 @@ public interface StreamOperator<OUT> extends Serializable {
* This method is expected to make a thorough effort to release all resources
* that the operator has acquired.
*/
- void dispose();
+ void dispose() throws Exception;
// ------------------------------------------------------------------------
// state snapshots
@@ -92,8 +92,7 @@ public interface StreamOperator<OUT> extends Serializable {
/**
* Called to draw a state snapshot from the operator. This method snapshots the operator state
- * (if the operator is stateful) and the key/value state (if it is being used and has been
- * initialized).
+ * (if the operator is stateful).
*
* @param out The stream to which we have to write our state.
* @param checkpointId The ID of the checkpoint.
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java
index 8d074cc..35d1108 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSink.java
@@ -24,7 +24,7 @@ import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.core.memory.DataInputViewStreamWrapper;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.io.disk.InputViewIterator;
-import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.util.ReusingMutableToRegularIteratorWrapper;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
@@ -56,9 +56,10 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
protected static final Logger LOG = LoggerFactory.getLogger(GenericWriteAheadSink.class);
private final CheckpointCommitter committer;
- private transient AbstractStateBackend.CheckpointStateOutputStream out;
+ private transient CheckpointStreamFactory.CheckpointStateOutputStream out;
protected final TypeSerializer<IN> serializer;
private final String id;
+ private transient CheckpointStreamFactory checkpointStreamFactory;
private ExactlyOnceState state = new ExactlyOnceState();
@@ -76,6 +77,8 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
committer.setOperatorSubtaskId(getRuntimeContext().getIndexOfThisSubtask());
committer.open();
cleanState();
+ checkpointStreamFactory =
+ getContainingTask().createCheckpointStreamFactory(this);
}
public void close() throws Exception {
@@ -184,9 +187,9 @@ public abstract class GenericWriteAheadSink<IN> extends AbstractStreamOperator<I
@Override
public void processElement(StreamRecord<IN> element) throws Exception {
IN value = element.getValue();
- //generate initial operator state
+ // generate initial operator state
if (out == null) {
- out = getStateBackend().createCheckpointStateOutputStream(0, 0);
+ out = checkpointStreamFactory.createCheckpointStateOutputStream(0, 0);
}
serializer.serialize(value, new DataOutputViewStreamWrapper(out));
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java
index 2c95099..e74dd87 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/AbstractAlignedProcessingTimeWindowOperator.java
@@ -125,7 +125,7 @@ public abstract class AbstractAlignedProcessingTimeWindowOperator<KEY, IN, OUT,
// decide when to first compute the window and when to slide it
// the values should align with the start of time (that is, the UNIX epoch, not the big bang)
- final long now = System.currentTimeMillis();
+ final long now = getRuntimeContext().getCurrentProcessingTime();
nextEvaluationTime = now + windowSlide - (now % windowSlide);
nextSlideTime = now + paneSize - (now % paneSize);
@@ -178,7 +178,7 @@ public abstract class AbstractAlignedProcessingTimeWindowOperator<KEY, IN, OUT,
}
@Override
- public void dispose() {
+ public void dispose() throws Exception {
super.dispose();
// acquire the lock during shutdown, to prevent trigger calls at the same time
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java
index dbdd660..25ec519 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperator.java
@@ -277,7 +277,7 @@ public class WindowOperator<K, IN, ACC, OUT, W extends Window>
}
@Override
- public void dispose() {
+ public void dispose() throws Exception {
super.dispose();
timestampedCollector = null;
watermarkTimers = null;
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
index 074257c..02579aa 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
@@ -19,6 +19,7 @@ package org.apache.flink.streaming.runtime.tasks;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.accumulators.Accumulator;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
@@ -30,7 +31,10 @@ import org.apache.flink.runtime.jobgraph.tasks.AbstractInvokable;
import org.apache.flink.runtime.jobgraph.tasks.StatefulTask;
import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.StateBackendFactory;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
@@ -40,7 +44,6 @@ import org.apache.flink.runtime.taskmanager.DispatcherThreadFactory;
import org.apache.flink.runtime.util.event.EventListener;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.graph.StreamConfig;
-import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.Output;
import org.apache.flink.streaming.api.operators.StreamOperator;
import org.apache.flink.streaming.runtime.io.RecordWriterOutput;
@@ -50,6 +53,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.Closeable;
+import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -57,6 +61,9 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.RunnableFuture;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
@@ -130,6 +137,12 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
/** The class loader used to load dynamic classes of a job */
private ClassLoader userClassLoader;
+ /** Our state backend. We use this to create checkpoint streams and a keyed state backend. */
+ private AbstractStateBackend stateBackend;
+
+ /** Keyed state backend for the head operator, if it is keyed. There can only ever be one. */
+ private KeyedStateBackend<?> keyedStateBackend;
+
/**
* The internal {@link TimeServiceProvider} used to define the current
* processing time (default = {@code System.currentTimeMillis()}) and
@@ -152,7 +165,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
private volatile AsynchronousException asyncException;
/** The currently active background materialization threads */
- private final Set<Closeable> cancelables = new HashSet<Closeable>();
+ private final Set<Closeable> cancelables = new HashSet<>();
/** Flag to mark the task "in operation", in which case check
* needs to be initialized to true, so that early cancel() before invoke() behaves correctly */
@@ -163,6 +176,8 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
private long lastCheckpointSize = 0;
+ private ExecutorService asyncOperationsThreadPool;
+
// ------------------------------------------------------------------------
// Life cycle methods for specific implementations
// ------------------------------------------------------------------------
@@ -205,9 +220,14 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
// -------- Initialize ---------
LOG.debug("Initializing {}", getName());
+ asyncOperationsThreadPool = Executors.newCachedThreadPool();
+
userClassLoader = getUserCodeClassLoader();
configuration = new StreamConfig(getTaskConfiguration());
+
+ stateBackend = createStateBackend();
+
accumulatorMap = getEnvironment().getAccumulatorRegistry().getUserMap();
// if the clock is not already set, then assign a default TimeServiceProvider
@@ -252,8 +272,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
// first order of business is to give operators back their state
restoreState();
lazyRestoreChainedOperatorState = null; // GC friendliness
- lazyRestoreKeyGroupStates = null; // GC friendliness
-
+
// we need to make sure that any triggers scheduled in open() cannot be
// executed before all operators are opened
synchronized (lock) {
@@ -292,6 +311,9 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
// still let the computation fail
tryDisposeAllOperators();
disposed = true;
+
+ // Don't forget to check and throw exceptions that happened in async thread one last time
+ checkTimerException();
}
finally {
// clean up everything we initialized
@@ -307,16 +329,17 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
LOG.error("Could not shut down timer service", t);
}
}
-
+
// stop all asynchronous checkpoint threads
try {
closeAllClosables();
+ shutdownAsyncThreads();
}
catch (Throwable t) {
// catch and log the exception to not replace the original exception
LOG.error("Could not shut down async checkpoint threads", t);
}
-
+
// release the output resources. this method should never fail.
if (operatorChain != null) {
operatorChain.releaseOutputs();
@@ -330,7 +353,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
// catch and log the exception to not replace the original exception
LOG.error("Error during cleanup of stream task", t);
}
-
+
// if the operators were not disposed before, do a hard dispose
if (!disposed) {
disposeAllOperators();
@@ -414,6 +437,12 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
}
}
+ private void shutdownAsyncThreads() throws Exception {
+ if (!asyncOperationsThreadPool.isShutdown()) {
+ asyncOperationsThreadPool.shutdownNow();
+ }
+ }
+
/**
* Execute the operator-specific {@link StreamOperator#dispose()} method in each
* of the operators in the chain of this {@link StreamTask}. </b> Disposing happens
@@ -558,69 +587,6 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
cancelables.remove(lazyRestoreChainedOperatorState);
}
}
-// if (lazyRestoreState != null || lazyRestoreKeyGroupStates != null) {
-//
-// LOG.info("Restoring checkpointed state to task {}", getName());
-//
-// final StreamOperator<?>[] allOperators = operatorChain.getAllOperators();
-// StreamOperatorNonPartitionedState[] nonPartitionedStates;
-//
-// final List<Map<Integer, PartitionedStateSnapshot>> keyGroupStates = new ArrayList<Map<Integer, PartitionedStateSnapshot>>(allOperators.length);
-//
-// for (int i = 0; i < allOperators.length; i++) {
-// keyGroupStates.add(new HashMap<Integer, PartitionedStateSnapshot>());
-// }
-//
-// if (lazyRestoreState != null) {
-// try {
-// nonPartitionedStates = lazyRestoreState.get(getUserCodeClassLoader());
-//
-// // be GC friendly
-// lazyRestoreState = null;
-// } catch (Exception e) {
-// throw new Exception("Could not restore checkpointed non-partitioned state.", e);
-// }
-// } else {
-// nonPartitionedStates = new StreamOperatorNonPartitionedState[allOperators.length];
-// }
-//
-// if (lazyRestoreKeyGroupStates != null) {
-// try {
-// // construct key groups state for operators
-// for (Map.Entry<Integer, ChainedStateHandle> lazyRestoreKeyGroupState : lazyRestoreKeyGroupStates.entrySet()) {
-// int keyGroupId = lazyRestoreKeyGroupState.getKey();
-//
-// Map<Integer, PartitionedStateSnapshot> chainedKeyGroupStates = lazyRestoreKeyGroupState.getValue().get(getUserCodeClassLoader());
-//
-// for (Map.Entry<Integer, PartitionedStateSnapshot> chainedKeyGroupState : chainedKeyGroupStates.entrySet()) {
-// int chainedIndex = chainedKeyGroupState.getKey();
-//
-// Map<Integer, PartitionedStateSnapshot> keyGroupState;
-//
-// keyGroupState = keyGroupStates.get(chainedIndex);
-// keyGroupState.put(keyGroupId, chainedKeyGroupState.getValue());
-// }
-// }
-//
-// lazyRestoreKeyGroupStates = null;
-//
-// } catch (Exception e) {
-// throw new Exception("Could not restore checkpointed partitioned state.", e);
-// }
-// }
-//
-// for (int i = 0; i < nonPartitionedStates.length; i++) {
-// StreamOperatorNonPartitionedState nonPartitionedState = nonPartitionedStates[i];
-// StreamOperator<?> operator = allOperators[i];
-// KeyGroupsStateHandle partitionedState = new KeyGroupsStateHandle(keyGroupStates.get(i));
-// StreamOperatorState operatorState = new StreamOperatorState(partitionedState, nonPartitionedState);
-//
-// if (operator != null) {
-// LOG.debug("Restore state of task {} in chain ({}).", i, getName());
-// operator.restoreState(operatorState, recoveryTimestamp);
-// }
-// }
-// }
}
@Override
@@ -658,8 +624,15 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
StreamOperator<?> operator = allOperators[i];
if (operator != null) {
- AbstractStateBackend.CheckpointStateOutputStream outStream =
- ((AbstractStreamOperator) operator).getStateBackend().createCheckpointStateOutputStream(checkpointId, timestamp);
+ CheckpointStreamFactory streamFactory =
+ stateBackend.createStreamFactory(
+ getEnvironment().getJobID(),
+ createOperatorIdentifier(
+ operator,
+ configuration.getVertexID()));
+
+ CheckpointStreamFactory.CheckpointStateOutputStream outStream =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
operator.snapshotState(outStream, checkpointId, timestamp);
@@ -667,22 +640,37 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
}
}
- if (!isRunning) {
- // Rethrow the cancel exception because some state backends could swallow
- // exceptions and seem to exit cleanly.
- throw new CancelTaskException();
+ RunnableFuture<KeyGroupsStateHandle> keyGroupsStateHandleFuture = null;
+
+ if (keyedStateBackend != null) {
+ CheckpointStreamFactory streamFactory = stateBackend.createStreamFactory(
+ getEnvironment().getJobID(),
+ createOperatorIdentifier(
+ headOperator,
+ configuration.getVertexID()));
+ keyGroupsStateHandleFuture = keyedStateBackend.snapshot(
+ checkpointId,
+ timestamp,
+ streamFactory);
}
- ChainedStateHandle<StreamStateHandle> states = new ChainedStateHandle<>(nonPartitionedStates);
- List<KeyGroupsStateHandle> keyedStates = Collections.<KeyGroupsStateHandle>emptyList();
+ ChainedStateHandle<StreamStateHandle> chainedStateHandles = new ChainedStateHandle<>(nonPartitionedStates);
+ LOG.debug("Finished synchronous checkpoints for checkpoint {} on task {}", checkpointId, getName());
- if (states.isEmpty() && keyedStates.isEmpty()) {
- getEnvironment().acknowledgeCheckpoint(checkpointId);
- } else {
- this.lastCheckpointSize = states.getStateSize();
- getEnvironment().acknowledgeCheckpoint(checkpointId, states, keyedStates);
+ AsyncCheckpointRunnable asyncCheckpointRunnable = new AsyncCheckpointRunnable(
+ "checkpoint-" + checkpointId + "-" + timestamp,
+ this,
+ cancelables,
+ chainedStateHandles,
+ keyGroupsStateHandleFuture,
+ checkpointId);
+
+ synchronized (cancelables) {
+ cancelables.add(asyncCheckpointRunnable);
}
+
+ asyncOperationsThreadPool.submit(asyncCheckpointRunnable);
return true;
} else {
return false;
@@ -712,7 +700,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
// State backend
// ------------------------------------------------------------------------
- public AbstractStateBackend createStateBackend(String operatorIdentifier, TypeSerializer<?> keySerializer) throws Exception {
+ private AbstractStateBackend createStateBackend() throws Exception {
AbstractStateBackend stateBackend = configuration.getStateBackend(getUserCodeClassLoader());
if (stateBackend != null) {
@@ -732,7 +720,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
switch (backendName) {
case "jobmanager":
LOG.info("State backend is set to heap memory (checkpoint to jobmanager)");
- stateBackend = MemoryStateBackend.create();
+ stateBackend = new MemoryStateBackend();
break;
case "filesystem":
@@ -760,10 +748,69 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
}
}
}
- stateBackend.initializeForJob(getEnvironment(), operatorIdentifier, keySerializer);
return stateBackend;
}
+ public <K> KeyedStateBackend<K> createKeyedStateBackend(
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange) throws Exception {
+
+ if (keyedStateBackend != null) {
+ throw new RuntimeException("The keyed state backend can only be created once.");
+ }
+
+ String operatorIdentifier = createOperatorIdentifier(
+ headOperator,
+ configuration.getVertexID());
+
+ if (lazyRestoreKeyGroupStates != null) {
+ keyedStateBackend = stateBackend.restoreKeyedStateBackend(
+ getEnvironment(),
+ getEnvironment().getJobID(),
+ operatorIdentifier,
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange,
+ lazyRestoreKeyGroupStates,
+ getEnvironment().getTaskKvStateRegistry());
+
+ lazyRestoreKeyGroupStates = null; // GC friendliness
+ } else {
+ keyedStateBackend = stateBackend.createKeyedStateBackend(
+ getEnvironment(),
+ getEnvironment().getJobID(),
+ operatorIdentifier,
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange,
+ getEnvironment().getTaskKvStateRegistry());
+ }
+
+ return (KeyedStateBackend<K>) keyedStateBackend;
+ }
+
+ /**
+ * This is only visible because
+ * {@link org.apache.flink.streaming.runtime.operators.GenericWriteAheadSink} uses the
+ * checkpoint stream factory to write write-ahead logs. <b>This should not be used for
+ * anything else.</b>
+ */
+ public CheckpointStreamFactory createCheckpointStreamFactory(StreamOperator operator) throws IOException {
+ return stateBackend.createStreamFactory(
+ getEnvironment().getJobID(),
+ createOperatorIdentifier(
+ operator,
+ configuration.getVertexID()));
+
+ }
+
+ private String createOperatorIdentifier(StreamOperator operator, int vertexId) {
+ return operator.getClass().getSimpleName() +
+ "_" + vertexId +
+ "_" + getEnvironment().getTaskInfo().getIndexOfThisSubtask();
+ }
+
/**
* Registers a timer.
*/
@@ -852,77 +899,83 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
// ------------------------------------------------------------------------
-// private static class AsyncCheckpointThread extends Thread implements Closeable {
-//
-// private final StreamTask<?, ?> owner;
-//
-// private final Set<Closeable> cancelables;
-//
-// private final StreamTaskState[] states;
-//
-// private final long checkpointId;
-//
-// AsyncCheckpointThread(String name, StreamTask<?, ?> owner, Set<Closeable> cancelables,
-// StreamTaskState[] states, long checkpointId) {
-// super(name);
-// setDaemon(true);
-//
-// this.owner = owner;
-// this.cancelables = cancelables;
-// this.states = states;
-// this.checkpointId = checkpointId;
-// }
-//
-// @Override
-// public void run() {
-// try {
-// for (StreamTaskState state : states) {
-// if (state != null) {
-// if (state.getFunctionState() instanceof AsynchronousStateHandle) {
-// AsynchronousStateHandle<Serializable> asyncState = (AsynchronousStateHandle<Serializable>) state.getFunctionState();
-// state.setFunctionState(asyncState.materialize());
-// }
-// if (state.getOperatorState() instanceof AsynchronousStateHandle) {
-// AsynchronousStateHandle<?> asyncState = (AsynchronousStateHandle<?>) state.getOperatorState();
-// state.setOperatorState(asyncState.materialize());
-// }
-// if (state.getKvStates() != null) {
-// Set<String> keys = state.getKvStates().keySet();
-// HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates = state.getKvStates();
-// for (String key: keys) {
-// if (kvStates.get(key) instanceof AsynchronousKvStateSnapshot) {
-// AsynchronousKvStateSnapshot<?, ?, ?, ?, ?> asyncHandle = (AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) kvStates.get(key);
-// kvStates.put(key, asyncHandle.materialize());
-// }
-// }
-// }
-//
-// }
-// }
-// StreamTaskStateList allStates = new StreamTaskStateList(states);
-// owner.lastCheckpointSize = allStates.getStateSize();
-// owner.getEnvironment().acknowledgeCheckpoint(checkpointId, allStates);
-//
-// LOG.debug("Finished asynchronous checkpoints for checkpoint {} on task {}", checkpointId, getName());
-// }
-// catch (Exception e) {
-// if (owner.isRunning()) {
-// LOG.error("Caught exception while materializing asynchronous checkpoints.", e);
-// }
-// if (owner.asyncException == null) {
-// owner.asyncException = new AsynchronousException(e);
-// }
-// }
-// finally {
-// synchronized (cancelables) {
-// cancelables.remove(this);
-// }
-// }
-// }
-//
-// @Override
-// public void close() {
-// interrupt();
-// }
-// }
+ private static class AsyncCheckpointRunnable implements Runnable, Closeable {
+
+ private final StreamTask<?, ?> owner;
+
+ private final Set<Closeable> cancelables;
+
+ private final ChainedStateHandle<StreamStateHandle> chainedStateHandles;
+
+ private final RunnableFuture<KeyGroupsStateHandle> keyGroupsStateHandleFuture;
+
+ private final long checkpointId;
+
+ private final String name;
+
+ AsyncCheckpointRunnable(
+ String name,
+ StreamTask<?, ?> owner,
+ Set<Closeable> cancelables,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandles,
+ RunnableFuture<KeyGroupsStateHandle> keyGroupsStateHandleFuture,
+ long checkpointId) {
+
+ this.name = name;
+ this.owner = owner;
+ this.cancelables = cancelables;
+ this.chainedStateHandles = chainedStateHandles;
+ this.keyGroupsStateHandleFuture = keyGroupsStateHandleFuture;
+ this.checkpointId = checkpointId;
+ }
+
+ @Override
+ public void run() {
+ try {
+
+ List<KeyGroupsStateHandle> keyedStates = Collections.emptyList();
+
+ if (keyGroupsStateHandleFuture != null) {
+
+ if (!keyGroupsStateHandleFuture.isDone()) {
+ //TODO this currently works because we only have one RunnableFuture
+ keyGroupsStateHandleFuture.run();
+ }
+
+ KeyGroupsStateHandle keyGroupsStateHandle = this.keyGroupsStateHandleFuture.get();
+ if (keyGroupsStateHandle != null) {
+ keyedStates = Arrays.asList(keyGroupsStateHandle);
+ }
+ }
+
+ if (chainedStateHandles.isEmpty() && keyedStates.isEmpty()) {
+ owner.getEnvironment().acknowledgeCheckpoint(checkpointId);
+ } else {
+ owner. getEnvironment().acknowledgeCheckpoint(checkpointId, chainedStateHandles, keyedStates);
+ }
+
+ LOG.debug("Finished asynchronous checkpoints for checkpoint {} on task {}", checkpointId, name);
+ }
+ catch (Exception e) {
+ if (owner.isRunning()) {
+ LOG.error("Caught exception while materializing asynchronous checkpoints.", e);
+ }
+ if (owner.asyncException == null) {
+ owner.asyncException = new AsynchronousException(e);
+ }
+ }
+ finally {
+ synchronized (cancelables) {
+ cancelables.remove(this);
+ }
+ }
+ }
+
+ @Override
+ public void close() {
+ if (keyGroupsStateHandleFuture != null) {
+ keyGroupsStateHandleFuture.cancel(true);
+ }
+ }
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedFoldTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedFoldTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedFoldTest.java
index f6e7e6b..68a2bb2 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedFoldTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedFoldTest.java
@@ -27,6 +27,7 @@ import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.TestHarnessUtil;
@@ -69,8 +70,8 @@ public class StreamGroupedFoldTest {
StreamGroupedFold<Integer, String, String> operator = new StreamGroupedFold<>(new MyFolder(), "100");
operator.setOutputType(BasicTypeInfo.STRING_TYPE_INFO, new ExecutionConfig());
- OneInputStreamOperatorTestHarness<Integer, String> testHarness = new OneInputStreamOperatorTestHarness<>(operator);
- testHarness.configureForKeyedStream(keySelector, BasicTypeInfo.STRING_TYPE_INFO);
+ OneInputStreamOperatorTestHarness<Integer, String> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, BasicTypeInfo.STRING_TYPE_INFO);
long initialTime = 0L;
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -107,10 +108,9 @@ public class StreamGroupedFoldTest {
new TestOpenCloseFoldFunction(), "init");
operator.setOutputType(BasicTypeInfo.STRING_TYPE_INFO, new ExecutionConfig());
- OneInputStreamOperatorTestHarness<Integer, String> testHarness = new OneInputStreamOperatorTestHarness<>(operator);
- testHarness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
-
-
+ OneInputStreamOperatorTestHarness<Integer, String> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, BasicTypeInfo.INT_TYPE_INFO);
+
long initialTime = 0L;
testHarness.open();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedReduceTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedReduceTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedReduceTest.java
index 6cb46c9..0f304a0 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedReduceTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamGroupedReduceTest.java
@@ -28,6 +28,7 @@ import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.TestHarnessUtil;
import org.junit.Assert;
@@ -52,8 +53,8 @@ public class StreamGroupedReduceTest {
StreamGroupedReduce<Integer> operator = new StreamGroupedReduce<>(new MyReducer(), IntSerializer.INSTANCE);
- OneInputStreamOperatorTestHarness<Integer, Integer> testHarness = new OneInputStreamOperatorTestHarness<>(operator);
- testHarness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
+ OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, BasicTypeInfo.INT_TYPE_INFO);
long initialTime = 0L;
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -84,8 +85,8 @@ public class StreamGroupedReduceTest {
StreamGroupedReduce<Integer> operator =
new StreamGroupedReduce<>(new TestOpenCloseReduceFunction(), IntSerializer.INSTANCE);
- OneInputStreamOperatorTestHarness<Integer, Integer> testHarness = new OneInputStreamOperatorTestHarness<>(operator);
- testHarness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
+ OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, BasicTypeInfo.INT_TYPE_INFO);
long initialTime = 0L;
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
index 3a88d94..d3b7ff9 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamingRuntimeContextTest.java
@@ -19,6 +19,7 @@
package org.apache.flink.streaming.api.operators;
import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.TaskInfo;
import org.apache.flink.api.common.accumulators.Accumulator;
import org.apache.flink.api.common.functions.ReduceFunction;
@@ -28,13 +29,21 @@ import org.apache.flink.api.common.state.ReducingStateDescriptor;
import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.java.typeutils.runtime.kryo.KryoSerializer;
import org.apache.flink.core.fs.Path;
import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
+import org.apache.flink.runtime.query.KvStateRegistry;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-import org.apache.flink.runtime.state.memory.MemListState;
+import org.apache.flink.runtime.state.heap.HeapListState;
+import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.junit.Test;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
@@ -181,11 +190,18 @@ public class StreamingRuntimeContextTest {
@Override
public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable {
ListStateDescriptor<String> descr =
- (ListStateDescriptor<String>) invocationOnMock.getArguments()[0];
- MemListState<String, VoidNamespace, String> listState = new MemListState<>(
- StringSerializer.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
- listState.setCurrentNamespace(VoidNamespace.INSTANCE);
- return listState;
+ (ListStateDescriptor<String>) invocationOnMock.getArguments()[0];
+ KeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend(
+ new DummyEnvironment("test_task", 1, 0),
+ new JobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ new HashKeyGroupAssigner<Integer>(1),
+ new KeyGroupRange(0, 0),
+ new KvStateRegistry().createTaskRegistry(new JobID(),
+ new JobVertexID()));
+ backend.setCurrentKey(0);
+ return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr);
}
});
@@ -196,7 +212,7 @@ public class StreamingRuntimeContextTest {
Environment env = mock(Environment.class);
when(env.getUserClassLoader()).thenReturn(StreamingRuntimeContextTest.class.getClassLoader());
when(env.getDistributedCacheEntries()).thenReturn(Collections.<String, Future<Path>>emptyMap());
- when(env.getTaskInfo()).thenReturn(new TaskInfo("test task", 0, 1, 1));
+ when(env.getTaskInfo()).thenReturn(new TaskInfo("test task", 1, 0, 1, 1));
return env;
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamOperatorChainingTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamOperatorChainingTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamOperatorChainingTest.java
index 3b201dc..f4ac5b2 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamOperatorChainingTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamOperatorChainingTest.java
@@ -333,21 +333,6 @@ public class StreamOperatorChainingTest {
when(mockTask.getEnvironment()).thenReturn(env);
when(mockTask.getExecutionConfig()).thenReturn(new ExecutionConfig().enableObjectReuse());
- try {
- doAnswer(new Answer<AbstractStateBackend>() {
- @Override
- public AbstractStateBackend answer(InvocationOnMock invocationOnMock) throws Throwable {
- final String operatorIdentifier = (String) invocationOnMock.getArguments()[0];
- final TypeSerializer<?> keySerializer = (TypeSerializer<?>) invocationOnMock.getArguments()[1];
- MemoryStateBackend backend = MemoryStateBackend.create();
- backend.initializeForJob(env, operatorIdentifier, keySerializer);
- return backend;
- }
- }).when(mockTask).createStateBackend(any(String.class), any(TypeSerializer.class));
- } catch (Exception e) {
- throw new RuntimeException(e.getMessage(), e);
- }
-
return mockTask;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java
index 2cb1809..40a6c79 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java
@@ -21,10 +21,9 @@ package org.apache.flink.streaming.runtime.operators.windowing;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.TaskInfo;
import org.apache.flink.api.common.accumulators.Accumulator;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.java.ClosureCleaner;
@@ -33,9 +32,6 @@ import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.operators.testutils.UnregisteredTaskMetricsGroup;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.memory.MemoryStateBackend;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
import org.apache.flink.streaming.api.functions.windowing.RichWindowFunction;
@@ -46,11 +42,12 @@ import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
+import org.apache.flink.streaming.runtime.tasks.TestTimeServiceProvider;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.util.Collector;
import org.junit.After;
-import org.junit.Ignore;
import org.junit.Test;
import org.mockito.invocation.InvocationOnMock;
@@ -79,7 +76,6 @@ import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@SuppressWarnings({"serial", "SynchronizationOnLocalVariableOrMethodParameter"})
-@Ignore
public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
@SuppressWarnings("unchecked")
@@ -203,7 +199,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
op = new AccumulatingProcessingTimeWindowOperator<>(mockFunction, mockKeySelector,
StringSerializer.INSTANCE, StringSerializer.INSTANCE, 5000, 1000);
- op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), mockOut);
+ op.setup(mockTask, new StreamConfig(new Configuration()), mockOut);
op.open();
assertTrue(op.getNextSlideTime() % 1000 == 0);
assertTrue(op.getNextEvaluationTime() % 1000 == 0);
@@ -255,7 +251,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, IntSerializer.INSTANCE,
windowSize, windowSize);
- op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
+ op.setup(mockTask, new StreamConfig(new Configuration()), out);
op.open();
final int numElements = 1000;
@@ -306,7 +302,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE, 150, 50);
- op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
+ op.setup(mockTask, new StreamConfig(new Configuration()), out);
op.open();
final int numElements = 1000;
@@ -367,7 +363,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE, 50, 50);
- op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
+ op.setup(mockTask, new StreamConfig(new Configuration()), out);
op.open();
synchronized (lock) {
@@ -423,7 +419,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE, 150, 50);
- op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
+ op.setup(mockTask, new StreamConfig(new Configuration()), out);
op.open();
synchronized (lock) {
@@ -457,67 +453,58 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
@Test
public void checkpointRestoreWithPendingWindowTumbling() {
- final ScheduledExecutorService timerService = Executors.newSingleThreadScheduledExecutor();
try {
final int windowSize = 200;
- final CollectingOutput<Integer> out = new CollectingOutput<>(windowSize);
- final Object lock = new Object();
- final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
- // tumbling window that triggers every 50 milliseconds
+ // tumbling window that triggers every 200 milliseconds
AccumulatingProcessingTimeWindowOperator<Integer, Integer, Integer> op =
new AccumulatingProcessingTimeWindowOperator<>(
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE,
windowSize, windowSize);
- OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
- new OneInputStreamOperatorTestHarness<>(op);
+ TestTimeServiceProvider timerService = new TestTimeServiceProvider();
+
+ OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
+ new OneInputStreamOperatorTestHarness<>(op, new ExecutionConfig(), timerService);
testHarness.setup();
testHarness.open();
+ timerService.setCurrentTime(0);
+
// inject some elements
final int numElementsFirst = 700;
final int numElements = 1000;
for (int i = 0; i < numElementsFirst; i++) {
- synchronized (lock) {
- op.processElement(new StreamRecord<Integer>(i));
- }
- Thread.sleep(1);
+ testHarness.processElement(new StreamRecord<>(i));
}
// draw a snapshot and dispose the window
- StreamStateHandle state;
- List<Integer> resultAtSnapshot;
- synchronized (lock) {
- int beforeSnapShot = out.getElements().size();
- state = testHarness.snapshot(1L, System.currentTimeMillis());
- resultAtSnapshot = new ArrayList<>(out.getElements());
- int afterSnapShot = out.getElements().size();
- assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
- assertTrue(afterSnapShot <= numElementsFirst);
- }
+ System.out.println("GOT: " + testHarness.getOutput());
+ int beforeSnapShot = testHarness.getOutput().size();
+ StreamStateHandle state = testHarness.snapshot(1L, System.currentTimeMillis());
+ List<Integer> resultAtSnapshot = extractFromStreamRecords(testHarness.getOutput());
+ int afterSnapShot = testHarness.getOutput().size();
+ assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
+ assertTrue(afterSnapShot <= numElementsFirst);
// inject some random elements, which should not show up in the state
for (int i = 0; i < 300; i++) {
- synchronized (lock) {
- op.processElement(new StreamRecord<Integer>(i + numElementsFirst));
- }
- Thread.sleep(1);
+ testHarness.processElement(new StreamRecord<>(i + numElementsFirst));
}
op.dispose();
// re-create the operator and restore the state
- final CollectingOutput<Integer> out2 = new CollectingOutput<>(windowSize);
op = new AccumulatingProcessingTimeWindowOperator<>(
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE,
windowSize, windowSize);
testHarness =
- new OneInputStreamOperatorTestHarness<>(op);
+ new OneInputStreamOperatorTestHarness<>(op, new ExecutionConfig(), timerService);
+
testHarness.setup();
testHarness.restore(state);
@@ -525,18 +512,16 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
// inject some more elements
for (int i = numElementsFirst; i < numElements; i++) {
- synchronized (lock) {
- op.processElement(new StreamRecord<Integer>(i));
- }
- Thread.sleep(1);
+ testHarness.processElement(new StreamRecord<>(i));
}
-
- out2.waitForNElements(numElements - resultAtSnapshot.size(), 60_000);
+ timerService.setCurrentTime(400);
// get and verify the result
- List<Integer> finalResult = new ArrayList<>(resultAtSnapshot);
- finalResult.addAll(out2.getElements());
+ List<Integer> finalResult = new ArrayList<>();
+ finalResult.addAll(resultAtSnapshot);
+ List<Integer> finalPartialResult = extractFromStreamRecords(testHarness.getOutput());
+ finalResult.addAll(finalPartialResult);
assertEquals(numElements, finalResult.size());
Collections.sort(finalResult);
@@ -548,22 +533,16 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
e.printStackTrace();
fail(e.getMessage());
}
- finally {
- timerService.shutdown();
- }
}
@Test
public void checkpointRestoreWithPendingWindowSliding() {
- final ScheduledExecutorService timerService = Executors.newSingleThreadScheduledExecutor();
try {
final int factor = 4;
final int windowSlide = 50;
final int windowSize = factor * windowSlide;
-
- final CollectingOutput<Integer> out = new CollectingOutput<>(windowSlide);
- final Object lock = new Object();
- final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
+
+ TestTimeServiceProvider timerService = new TestTimeServiceProvider();
// sliding window (200 msecs) every 50 msecs
AccumulatingProcessingTimeWindowOperator<Integer, Integer, Integer> op =
@@ -573,7 +552,9 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
windowSize, windowSlide);
OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
- new OneInputStreamOperatorTestHarness<>(op);
+ new OneInputStreamOperatorTestHarness<>(op, new ExecutionConfig(), timerService);
+
+ timerService.setCurrentTime(0);
testHarness.setup();
testHarness.open();
@@ -583,44 +564,32 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
final int numElementsFirst = 700;
for (int i = 0; i < numElementsFirst; i++) {
- synchronized (lock) {
- op.processElement(new StreamRecord<Integer>(i));
- }
- Thread.sleep(1);
+ testHarness.processElement(new StreamRecord<>(i));
}
// draw a snapshot
- StreamStateHandle state;
- List<Integer> resultAtSnapshot;
- synchronized (lock) {
- int beforeSnapShot = out.getElements().size();
- state = testHarness.snapshot(1L, System.currentTimeMillis());
- resultAtSnapshot = new ArrayList<>(out.getElements());
- int afterSnapShot = out.getElements().size();
- assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
- }
-
+ List<Integer> resultAtSnapshot = extractFromStreamRecords(testHarness.getOutput());
+ int beforeSnapShot = testHarness.getOutput().size();
+ StreamStateHandle state = testHarness.snapshot(1L, System.currentTimeMillis());
+ int afterSnapShot = testHarness.getOutput().size();
+ assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
+
assertTrue(resultAtSnapshot.size() <= factor * numElementsFirst);
// inject the remaining elements - these should not influence the snapshot
for (int i = numElementsFirst; i < numElements; i++) {
- synchronized (lock) {
- op.processElement(new StreamRecord<Integer>(i));
- }
- Thread.sleep(1);
+ testHarness.processElement(new StreamRecord<>(i));
}
op.dispose();
// re-create the operator and restore the state
- final CollectingOutput<Integer> out2 = new CollectingOutput<>(windowSlide);
op = new AccumulatingProcessingTimeWindowOperator<>(
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE,
windowSize, windowSlide);
- testHarness =
- new OneInputStreamOperatorTestHarness<>(op);
+ testHarness = new OneInputStreamOperatorTestHarness<>(op, new ExecutionConfig(), timerService);
testHarness.setup();
testHarness.restore(state);
@@ -629,29 +598,24 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
// inject again the remaining elements
for (int i = numElementsFirst; i < numElements; i++) {
- synchronized (lock) {
- op.processElement(new StreamRecord<Integer>(i));
- }
- Thread.sleep(1);
+ testHarness.processElement(new StreamRecord<>(i));
}
- // for a deterministic result, we need to wait until all pending triggers
- // have fired and emitted their results
- long deadline = System.currentTimeMillis() + 120000;
- do {
- Thread.sleep(20);
- }
- while (resultAtSnapshot.size() + out2.getElements().size() < factor * numElements
- && System.currentTimeMillis() < deadline);
+ timerService.setCurrentTime(50);
+ timerService.setCurrentTime(100);
+ timerService.setCurrentTime(150);
+ timerService.setCurrentTime(200);
+ timerService.setCurrentTime(250);
+ timerService.setCurrentTime(300);
+ timerService.setCurrentTime(350);
- synchronized (lock) {
- op.close();
- }
+ testHarness.close();
op.dispose();
// get and verify the result
List<Integer> finalResult = new ArrayList<>(resultAtSnapshot);
- finalResult.addAll(out2.getElements());
+ List<Integer> finalPartialResult = extractFromStreamRecords(testHarness.getOutput());
+ finalResult.addAll(finalPartialResult);
assertEquals(factor * numElements, finalResult.size());
Collections.sort(finalResult);
@@ -663,19 +627,12 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
e.printStackTrace();
fail(e.getMessage());
}
- finally {
- timerService.shutdown();
- }
}
@Test
public void testKeyValueStateInWindowFunction() {
- final ScheduledExecutorService timerService = Executors.newSingleThreadScheduledExecutor();
try {
- final CollectingOutput<Integer> out = new CollectingOutput<>(50);
- final Object lock = new Object();
- final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
-
+
StatefulFunction.globalCounts.clear();
// tumbling window that triggers every 20 milliseconds
@@ -684,26 +641,28 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
new StatefulFunction(), identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE, 50, 50);
- op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
- op.open();
+ TestTimeServiceProvider timerService = new TestTimeServiceProvider();
- synchronized (lock) {
- op.processElement(new StreamRecord<Integer>(1));
- op.processElement(new StreamRecord<Integer>(2));
- }
- out.waitForNElements(2, 60000);
+ OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(op, new ExecutionConfig(), timerService, identitySelector, BasicTypeInfo.INT_TYPE_INFO);
- synchronized (lock) {
- op.processElement(new StreamRecord<Integer>(1));
- op.processElement(new StreamRecord<Integer>(2));
- op.processElement(new StreamRecord<Integer>(1));
- op.processElement(new StreamRecord<Integer>(1));
- op.processElement(new StreamRecord<Integer>(2));
- op.processElement(new StreamRecord<Integer>(2));
- }
- out.waitForNElements(8, 60000);
+ testHarness.open();
- List<Integer> result = out.getElements();
+ timerService.setCurrentTime(0);
+
+ testHarness.processElement(new StreamRecord<>(1));
+ testHarness.processElement(new StreamRecord<>(2));
+
+ op.processElement(new StreamRecord<>(1));
+ op.processElement(new StreamRecord<>(2));
+ op.processElement(new StreamRecord<>(1));
+ op.processElement(new StreamRecord<>(1));
+ op.processElement(new StreamRecord<>(2));
+ op.processElement(new StreamRecord<>(2));
+
+ timerService.setCurrentTime(1000);
+
+ List<Integer> result = extractFromStreamRecords(testHarness.getOutput());
assertEquals(8, result.size());
Collections.sort(result);
@@ -712,18 +671,13 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
assertEquals(4, StatefulFunction.globalCounts.get(1).intValue());
assertEquals(4, StatefulFunction.globalCounts.get(2).intValue());
- synchronized (lock) {
- op.close();
- }
+ testHarness.close();
op.dispose();
}
catch (Exception e) {
e.printStackTrace();
fail(e.getMessage());
}
- finally {
- timerService.shutdown();
- }
}
// ------------------------------------------------------------------------
@@ -793,27 +747,11 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
when(mockTaskManagerRuntimeInfo.getConfiguration()).thenReturn(configuration);
final Environment env = mock(Environment.class);
- when(env.getTaskInfo()).thenReturn(new TaskInfo("Test task name", 0, 1, 0));
+ when(env.getTaskInfo()).thenReturn(new TaskInfo("Test task name", 1, 0, 1, 0));
when(env.getUserClassLoader()).thenReturn(AggregatingAlignedProcessingTimeWindowOperatorTest.class.getClassLoader());
when(env.getMetricGroup()).thenReturn(new UnregisteredTaskMetricsGroup());
when(task.getEnvironment()).thenReturn(env);
-
- try {
- doAnswer(new Answer<AbstractStateBackend>() {
- @Override
- public AbstractStateBackend answer(InvocationOnMock invocationOnMock) throws Throwable {
- final String operatorIdentifier = (String) invocationOnMock.getArguments()[0];
- final TypeSerializer<?> keySerializer = (TypeSerializer<?>) invocationOnMock.getArguments()[1];
- MemoryStateBackend backend = MemoryStateBackend.create();
- backend.initializeForJob(env, operatorIdentifier, keySerializer);
- return backend;
- }
- }).when(task).createStateBackend(any(String.class), any(TypeSerializer.class));
- } catch (Exception e) {
- e.printStackTrace();
- }
-
return task;
}
@@ -846,11 +784,14 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
return mockTask;
}
- private static StreamConfig createTaskConfig(KeySelector<?, ?> partitioner, TypeSerializer<?> keySerializer, KeyGroupAssigner<?> keyGroupAssigner) {
- StreamConfig cfg = new StreamConfig(new Configuration());
- cfg.setStatePartitioner(0, partitioner);
- cfg.setStateKeySerializer(keySerializer);
- cfg.setKeyGroupAssigner(keyGroupAssigner);
- return cfg;
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ private <T> List<T> extractFromStreamRecords(Iterable<Object> input) {
+ List<T> result = new ArrayList<>();
+ for (Object in : input) {
+ if (in instanceof StreamRecord) {
+ result.add((T) ((StreamRecord) in).getValue());
+ }
+ }
+ return result;
}
}
[15/27] flink git commit: [FLINK-4381] Refactor State to Prepare For
Key-Group State Backends
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
index 940f699..074257c 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
@@ -29,10 +29,10 @@ import org.apache.flink.runtime.io.network.api.CheckpointBarrier;
import org.apache.flink.runtime.jobgraph.tasks.AbstractInvokable;
import org.apache.flink.runtime.jobgraph.tasks.StatefulTask;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.AsynchronousKvStateSnapshot;
-import org.apache.flink.runtime.state.AsynchronousStateHandle;
-import org.apache.flink.runtime.state.KvStateSnapshot;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
import org.apache.flink.runtime.state.StateBackendFactory;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.runtime.state.filesystem.FsStateBackendFactory;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
@@ -40,19 +40,19 @@ import org.apache.flink.runtime.taskmanager.DispatcherThreadFactory;
import org.apache.flink.runtime.util.event.EventListener;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.graph.StreamConfig;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.Output;
import org.apache.flink.streaming.api.operators.StreamOperator;
import org.apache.flink.streaming.runtime.io.RecordWriterOutput;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.Closeable;
-import java.io.Serializable;
import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Arrays;
+import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
@@ -102,7 +102,7 @@ import java.util.concurrent.ScheduledThreadPoolExecutor;
@Internal
public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
extends AbstractInvokable
- implements StatefulTask<StreamTaskStateList> {
+ implements StatefulTask {
/** The thread group that holds all trigger timer threads */
public static final ThreadGroup TRIGGER_THREAD_GROUP = new ThreadGroup("Triggers");
@@ -140,8 +140,10 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
/** The map of user-defined accumulators of this task */
private Map<String, Accumulator<?, ?>> accumulatorMap;
- /** The state to be restored once the initialization is done */
- private StreamTaskStateList lazyRestoreState;
+ /** The chained operator state to be restored once the initialization is done */
+ private ChainedStateHandle<StreamStateHandle> lazyRestoreChainedOperatorState;
+
+ private List<KeyGroupsStateHandle> lazyRestoreKeyGroupStates;
/**
* This field is used to forward an exception that is caught in the timer thread or other
@@ -204,6 +206,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
LOG.debug("Initializing {}", getName());
userClassLoader = getUserCodeClassLoader();
+
configuration = new StreamConfig(getTaskConfiguration());
accumulatorMap = getEnvironment().getAccumulatorRegistry().getUserMap();
@@ -247,8 +250,9 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
LOG.debug("Invoking {}", getName());
// first order of business is to give operators back their state
- restoreState(lazyRestoreState);
- lazyRestoreState = null; // GC friendliness
+ restoreState();
+ lazyRestoreChainedOperatorState = null; // GC friendliness
+ lazyRestoreKeyGroupStates = null; // GC friendliness
// we need to make sure that any triggers scheduled in open() cannot be
// executed before all operators are opened
@@ -473,7 +477,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
}
}
- protected boolean isSerializingTimestamps() {
+ boolean isSerializingTimestamps() {
TimeCharacteristic tc = configuration.getTimeCharacteristic();
return tc == TimeCharacteristic.EventTime | tc == TimeCharacteristic.IngestionTime;
}
@@ -506,11 +510,11 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
return accumulatorMap;
}
- public Output<StreamRecord<OUT>> getHeadOutput() {
+ Output<StreamRecord<OUT>> getHeadOutput() {
return operatorChain.getChainEntryPoint();
}
- public RecordWriterOutput<?>[] getStreamOutputs() {
+ RecordWriterOutput<?>[] getStreamOutputs() {
return operatorChain.getStreamOutputs();
}
@@ -519,44 +523,104 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
// ------------------------------------------------------------------------
@Override
- public void setInitialState(StreamTaskStateList initialState) {
- lazyRestoreState = initialState;
+ public void setInitialState(ChainedStateHandle<StreamStateHandle> chainedState, List<KeyGroupsStateHandle> keyGroupsState) {
+ lazyRestoreChainedOperatorState = chainedState;
+ lazyRestoreKeyGroupStates = keyGroupsState;
}
- private void restoreState(StreamTaskStateList restoredState) throws Exception {
- if (restoredState != null) {
- LOG.info("Restoring checkpointed state to task {}", getName());
-
- synchronized (cancelables) {
- cancelables.add(restoredState);
- }
+ private void restoreState() throws Exception {
+ final StreamOperator<?>[] allOperators = operatorChain.getAllOperators();
- try {
- final StreamOperator<?>[] allOperators = operatorChain.getAllOperators();
- final StreamTaskState[] states = restoredState.getState(userClassLoader);
-
- for (int i = 0; i < states.length; i++) {
- StreamTaskState state = states[i];
- StreamOperator<?> operator = allOperators[i];
-
- if (state != null && operator != null) {
- LOG.debug("Task {} in chain ({}) has checkpointed state", i, getName());
- operator.restoreState(state);
+ try {
+ if (lazyRestoreChainedOperatorState != null) {
+
+ synchronized (cancelables) {
+ cancelables.add(lazyRestoreChainedOperatorState);
+ }
+
+ for (int i = 0; i < lazyRestoreChainedOperatorState.getLength(); i++) {
+ StreamStateHandle state = lazyRestoreChainedOperatorState.get(i);
+ if (state == null) {
+ continue;
}
- else if (operator != null) {
- LOG.debug("Task {} in chain ({}) does not have checkpointed state", i, getName());
+ if (state != null) {
+ StreamOperator<?> operator = allOperators[i];
+
+ if (operator != null) {
+ LOG.debug("Restore state of task {} in chain ({}).", i, getName());
+ operator.restoreState(state.openInputStream());
+ }
}
}
}
- catch (Exception e) {
- throw new Exception("Could not restore checkpointed state to operators and functions", e);
- }
- finally {
- synchronized (cancelables) {
- cancelables.remove(restoredState);
- }
+ } finally {
+ synchronized (cancelables) {
+ cancelables.remove(lazyRestoreChainedOperatorState);
}
}
+// if (lazyRestoreState != null || lazyRestoreKeyGroupStates != null) {
+//
+// LOG.info("Restoring checkpointed state to task {}", getName());
+//
+// final StreamOperator<?>[] allOperators = operatorChain.getAllOperators();
+// StreamOperatorNonPartitionedState[] nonPartitionedStates;
+//
+// final List<Map<Integer, PartitionedStateSnapshot>> keyGroupStates = new ArrayList<Map<Integer, PartitionedStateSnapshot>>(allOperators.length);
+//
+// for (int i = 0; i < allOperators.length; i++) {
+// keyGroupStates.add(new HashMap<Integer, PartitionedStateSnapshot>());
+// }
+//
+// if (lazyRestoreState != null) {
+// try {
+// nonPartitionedStates = lazyRestoreState.get(getUserCodeClassLoader());
+//
+// // be GC friendly
+// lazyRestoreState = null;
+// } catch (Exception e) {
+// throw new Exception("Could not restore checkpointed non-partitioned state.", e);
+// }
+// } else {
+// nonPartitionedStates = new StreamOperatorNonPartitionedState[allOperators.length];
+// }
+//
+// if (lazyRestoreKeyGroupStates != null) {
+// try {
+// // construct key groups state for operators
+// for (Map.Entry<Integer, ChainedStateHandle> lazyRestoreKeyGroupState : lazyRestoreKeyGroupStates.entrySet()) {
+// int keyGroupId = lazyRestoreKeyGroupState.getKey();
+//
+// Map<Integer, PartitionedStateSnapshot> chainedKeyGroupStates = lazyRestoreKeyGroupState.getValue().get(getUserCodeClassLoader());
+//
+// for (Map.Entry<Integer, PartitionedStateSnapshot> chainedKeyGroupState : chainedKeyGroupStates.entrySet()) {
+// int chainedIndex = chainedKeyGroupState.getKey();
+//
+// Map<Integer, PartitionedStateSnapshot> keyGroupState;
+//
+// keyGroupState = keyGroupStates.get(chainedIndex);
+// keyGroupState.put(keyGroupId, chainedKeyGroupState.getValue());
+// }
+// }
+//
+// lazyRestoreKeyGroupStates = null;
+//
+// } catch (Exception e) {
+// throw new Exception("Could not restore checkpointed partitioned state.", e);
+// }
+// }
+//
+// for (int i = 0; i < nonPartitionedStates.length; i++) {
+// StreamOperatorNonPartitionedState nonPartitionedState = nonPartitionedStates[i];
+// StreamOperator<?> operator = allOperators[i];
+// KeyGroupsStateHandle partitionedState = new KeyGroupsStateHandle(keyGroupStates.get(i));
+// StreamOperatorState operatorState = new StreamOperatorState(partitionedState, nonPartitionedState);
+//
+// if (operator != null) {
+// LOG.debug("Restore state of task {} in chain ({}).", i, getName());
+// operator.restoreState(operatorState, recoveryTimestamp);
+// }
+// }
+// }
}
@Override
@@ -574,7 +638,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
}
}
- protected boolean performCheckpoint(final long checkpointId, final long timestamp) throws Exception {
+ private boolean performCheckpoint(final long checkpointId, final long timestamp) throws Exception {
LOG.debug("Starting checkpoint {} on task {}", checkpointId, getName());
synchronized (lock) {
@@ -588,29 +652,18 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
// now draw the state snapshot
final StreamOperator<?>[] allOperators = operatorChain.getAllOperators();
- final StreamTaskState[] states = new StreamTaskState[allOperators.length];
+ final List<StreamStateHandle> nonPartitionedStates = Arrays.asList(new StreamStateHandle[allOperators.length]);
- boolean hasAsyncStates = false;
-
- for (int i = 0; i < states.length; i++) {
+ for (int i = 0; i < allOperators.length; i++) {
StreamOperator<?> operator = allOperators[i];
+
if (operator != null) {
- StreamTaskState state = operator.snapshotOperatorState(checkpointId, timestamp);
- if (state.getOperatorState() instanceof AsynchronousStateHandle) {
- hasAsyncStates = true;
- }
- if (state.getFunctionState() instanceof AsynchronousStateHandle) {
- hasAsyncStates = true;
- }
- if (state.getKvStates() != null) {
- for (KvStateSnapshot<?, ?, ?, ?, ?> kvSnapshot: state.getKvStates().values()) {
- if (kvSnapshot instanceof AsynchronousKvStateSnapshot) {
- hasAsyncStates = true;
- }
- }
- }
+ AbstractStateBackend.CheckpointStateOutputStream outStream =
+ ((AbstractStreamOperator) operator).getStateBackend().createCheckpointStateOutputStream(checkpointId, timestamp);
+
+ operator.snapshotState(outStream, checkpointId, timestamp);
- states[i] = state.isEmpty() ? null : state;
+ nonPartitionedStates.set(i, outStream.closeAndGetHandle());
}
}
@@ -620,24 +673,15 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
throw new CancelTaskException();
}
- StreamTaskStateList allStates = new StreamTaskStateList(states);
+ ChainedStateHandle<StreamStateHandle> states = new ChainedStateHandle<>(nonPartitionedStates);
+ List<KeyGroupsStateHandle> keyedStates = Collections.<KeyGroupsStateHandle>emptyList();
- if (allStates.isEmpty()) {
+
+ if (states.isEmpty() && keyedStates.isEmpty()) {
getEnvironment().acknowledgeCheckpoint(checkpointId);
- } else if (!hasAsyncStates) {
- this.lastCheckpointSize = allStates.getStateSize();
- getEnvironment().acknowledgeCheckpoint(checkpointId, allStates);
- } else {
- // start a Thread that does the asynchronous materialization and
- // then sends the checkpoint acknowledge
- String threadName = "Materialize checkpoint state " + checkpointId + " - " + getName();
- AsyncCheckpointThread checkpointThread = new AsyncCheckpointThread(
- threadName, this, cancelables, states, checkpointId);
-
- synchronized (cancelables) {
- cancelables.add(checkpointThread);
- }
- checkpointThread.start();
+ } else {
+ this.lastCheckpointSize = states.getStateSize();
+ getEnvironment().acknowledgeCheckpoint(checkpointId, states, keyedStates);
}
return true;
} else {
@@ -663,13 +707,13 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
}
}
}
-
+
// ------------------------------------------------------------------------
// State backend
// ------------------------------------------------------------------------
public AbstractStateBackend createStateBackend(String operatorIdentifier, TypeSerializer<?> keySerializer) throws Exception {
- AbstractStateBackend stateBackend = configuration.getStateBackend(userClassLoader);
+ AbstractStateBackend stateBackend = configuration.getStateBackend(getUserCodeClassLoader());
if (stateBackend != null) {
// backend has been configured on the environment
@@ -694,7 +738,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
case "filesystem":
FsStateBackend backend = new FsStateBackendFactory().createFromConfig(flinkConfig);
LOG.info("State backend is set to heap memory (checkpoints to filesystem \""
- + backend.getBasePath() + "\")");
+ + backend.getBasePath() + "\")");
stateBackend = backend;
break;
@@ -702,15 +746,15 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
try {
@SuppressWarnings("rawtypes")
Class<? extends StateBackendFactory> clazz =
- Class.forName(backendName, false, userClassLoader).asSubclass(StateBackendFactory.class);
+ Class.forName(backendName, false, getUserCodeClassLoader()).asSubclass(StateBackendFactory.class);
stateBackend = ((StateBackendFactory<?>) clazz.newInstance()).createFromConfig(flinkConfig);
} catch (ClassNotFoundException e) {
throw new IllegalConfigurationException("Cannot find configured state backend: " + backendName);
} catch (ClassCastException e) {
throw new IllegalConfigurationException("The class configured under '" +
- ConfigConstants.STATE_BACKEND + "' is not a valid state backend factory (" +
- backendName + ')');
+ ConfigConstants.STATE_BACKEND + "' is not a valid state backend factory (" +
+ backendName + ')');
} catch (Throwable t) {
throw new IllegalConfigurationException("Cannot create configured state backend", t);
}
@@ -718,7 +762,6 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
}
stateBackend.initializeForJob(getEnvironment(), operatorIdentifier, keySerializer);
return stateBackend;
-
}
/**
@@ -754,7 +797,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
return getName();
}
- protected final EventListener<CheckpointBarrier> getCheckpointBarrierListener() {
+ final EventListener<CheckpointBarrier> getCheckpointBarrierListener() {
return new EventListener<CheckpointBarrier>() {
@Override
public void onEvent(CheckpointBarrier barrier) {
@@ -809,77 +852,77 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
// ------------------------------------------------------------------------
- private static class AsyncCheckpointThread extends Thread implements Closeable {
-
- private final StreamTask<?, ?> owner;
-
- private final Set<Closeable> cancelables;
-
- private final StreamTaskState[] states;
-
- private final long checkpointId;
-
- AsyncCheckpointThread(String name, StreamTask<?, ?> owner, Set<Closeable> cancelables,
- StreamTaskState[] states, long checkpointId) {
- super(name);
- setDaemon(true);
-
- this.owner = owner;
- this.cancelables = cancelables;
- this.states = states;
- this.checkpointId = checkpointId;
- }
-
- @Override
- public void run() {
- try {
- for (StreamTaskState state : states) {
- if (state != null) {
- if (state.getFunctionState() instanceof AsynchronousStateHandle) {
- AsynchronousStateHandle<Serializable> asyncState = (AsynchronousStateHandle<Serializable>) state.getFunctionState();
- state.setFunctionState(asyncState.materialize());
- }
- if (state.getOperatorState() instanceof AsynchronousStateHandle) {
- AsynchronousStateHandle<?> asyncState = (AsynchronousStateHandle<?>) state.getOperatorState();
- state.setOperatorState(asyncState.materialize());
- }
- if (state.getKvStates() != null) {
- Set<String> keys = state.getKvStates().keySet();
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates = state.getKvStates();
- for (String key: keys) {
- if (kvStates.get(key) instanceof AsynchronousKvStateSnapshot) {
- AsynchronousKvStateSnapshot<?, ?, ?, ?, ?> asyncHandle = (AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) kvStates.get(key);
- kvStates.put(key, asyncHandle.materialize());
- }
- }
- }
-
- }
- }
- StreamTaskStateList allStates = new StreamTaskStateList(states);
- owner.lastCheckpointSize = allStates.getStateSize();
- owner.getEnvironment().acknowledgeCheckpoint(checkpointId, allStates);
-
- LOG.debug("Finished asynchronous checkpoints for checkpoint {} on task {}", checkpointId, getName());
- }
- catch (Exception e) {
- if (owner.isRunning()) {
- LOG.error("Caught exception while materializing asynchronous checkpoints.", e);
- }
- if (owner.asyncException == null) {
- owner.asyncException = new AsynchronousException(e);
- }
- }
- finally {
- synchronized (cancelables) {
- cancelables.remove(this);
- }
- }
- }
-
- @Override
- public void close() {
- interrupt();
- }
- }
+// private static class AsyncCheckpointThread extends Thread implements Closeable {
+//
+// private final StreamTask<?, ?> owner;
+//
+// private final Set<Closeable> cancelables;
+//
+// private final StreamTaskState[] states;
+//
+// private final long checkpointId;
+//
+// AsyncCheckpointThread(String name, StreamTask<?, ?> owner, Set<Closeable> cancelables,
+// StreamTaskState[] states, long checkpointId) {
+// super(name);
+// setDaemon(true);
+//
+// this.owner = owner;
+// this.cancelables = cancelables;
+// this.states = states;
+// this.checkpointId = checkpointId;
+// }
+//
+// @Override
+// public void run() {
+// try {
+// for (StreamTaskState state : states) {
+// if (state != null) {
+// if (state.getFunctionState() instanceof AsynchronousStateHandle) {
+// AsynchronousStateHandle<Serializable> asyncState = (AsynchronousStateHandle<Serializable>) state.getFunctionState();
+// state.setFunctionState(asyncState.materialize());
+// }
+// if (state.getOperatorState() instanceof AsynchronousStateHandle) {
+// AsynchronousStateHandle<?> asyncState = (AsynchronousStateHandle<?>) state.getOperatorState();
+// state.setOperatorState(asyncState.materialize());
+// }
+// if (state.getKvStates() != null) {
+// Set<String> keys = state.getKvStates().keySet();
+// HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates = state.getKvStates();
+// for (String key: keys) {
+// if (kvStates.get(key) instanceof AsynchronousKvStateSnapshot) {
+// AsynchronousKvStateSnapshot<?, ?, ?, ?, ?> asyncHandle = (AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) kvStates.get(key);
+// kvStates.put(key, asyncHandle.materialize());
+// }
+// }
+// }
+//
+// }
+// }
+// StreamTaskStateList allStates = new StreamTaskStateList(states);
+// owner.lastCheckpointSize = allStates.getStateSize();
+// owner.getEnvironment().acknowledgeCheckpoint(checkpointId, allStates);
+//
+// LOG.debug("Finished asynchronous checkpoints for checkpoint {} on task {}", checkpointId, getName());
+// }
+// catch (Exception e) {
+// if (owner.isRunning()) {
+// LOG.error("Caught exception while materializing asynchronous checkpoints.", e);
+// }
+// if (owner.asyncException == null) {
+// owner.asyncException = new AsynchronousException(e);
+// }
+// }
+// finally {
+// synchronized (cancelables) {
+// cancelables.remove(this);
+// }
+// }
+// }
+//
+// @Override
+// public void close() {
+// interrupt();
+// }
+// }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskState.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskState.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskState.java
deleted file mode 100644
index 925dd8c..0000000
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskState.java
+++ /dev/null
@@ -1,185 +0,0 @@
-/*
- * 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.annotation.Internal;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.util.ExceptionUtils;
-
-import java.io.Closeable;
-import java.io.IOException;
-import java.io.Serializable;
-import java.util.ConcurrentModificationException;
-import java.util.HashMap;
-import java.util.Iterator;
-
-/**
- * The state checkpointed by a {@link org.apache.flink.streaming.api.operators.AbstractStreamOperator}.
- * This state consists of any combination of those three:
- * <ul>
- * <li>The state of the stream operator, if it implements the Checkpointed interface.</li>
- * <li>The state of the user function, if it implements the Checkpointed interface.</li>
- * <li>The key/value state of the operator, if it executes on a KeyedDataStream.</li>
- * </ul>
- */
-@Internal
-public class StreamTaskState implements Serializable, Closeable {
-
- private static final long serialVersionUID = 1L;
-
- private StateHandle<?> operatorState;
-
- private StateHandle<Serializable> functionState;
-
- private HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates;
-
- // ------------------------------------------------------------------------
-
- public StateHandle<?> getOperatorState() {
- return operatorState;
- }
-
- public void setOperatorState(StateHandle<?> operatorState) {
- this.operatorState = operatorState;
- }
-
- public StateHandle<Serializable> getFunctionState() {
- return functionState;
- }
-
- public void setFunctionState(StateHandle<Serializable> functionState) {
- this.functionState = functionState;
- }
-
- public HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> getKvStates() {
- return kvStates;
- }
-
- public void setKvStates(HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates) {
- this.kvStates = kvStates;
- }
-
- // ------------------------------------------------------------------------
-
- /**
- * Checks if this state object actually contains any state, or if all of the state
- * fields are null.
- *
- * @return True, if all state is null, false if at least one state is not null.
- */
- public boolean isEmpty() {
- return operatorState == null & functionState == null & kvStates == null;
- }
-
- /**
- * Discards all the contained states and sets them to null.
- *
- * @throws Exception Forwards exceptions that occur when releasing the
- * state handles and snapshots.
- */
- public void discardState() throws Exception {
- StateHandle<?> operatorState = this.operatorState;
- StateHandle<?> functionState = this.functionState;
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates = this.kvStates;
-
- this.operatorState = null;
- this.functionState = null;
- this.kvStates = null;
-
- if (operatorState != null) {
- operatorState.discardState();
- }
- if (functionState != null) {
- functionState.discardState();
- }
- if (kvStates != null) {
- while (kvStates.size() > 0) {
- try {
- Iterator<KvStateSnapshot<?, ?, ?, ?, ?>> values = kvStates.values().iterator();
- while (values.hasNext()) {
- KvStateSnapshot<?, ?, ?, ?, ?> s = values.next();
- s.discardState();
- values.remove();
- }
- }
- catch (ConcurrentModificationException e) {
- // fall through the loop
- }
- }
- }
- }
-
- @Override
- public void close() throws IOException {
- StateHandle<?> operatorState = this.operatorState;
- StateHandle<?> functionState = this.functionState;
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates = this.kvStates;
-
- this.operatorState = null;
- this.functionState = null;
- this.kvStates = null;
-
- Throwable firstException = null;
-
- if (operatorState != null) {
- try {
- operatorState.close();
- } catch (Throwable t) {
- firstException = t;
- }
- }
-
- if (functionState != null) {
- try {
- functionState.close();
- } catch (Throwable t) {
- if (firstException == null) {
- firstException = t;
- }
- }
- }
-
- if (kvStates != null) {
- while (kvStates.size() > 0) {
- try {
- Iterator<KvStateSnapshot<?, ?, ?, ?, ?>> values = kvStates.values().iterator();
- while (values.hasNext()) {
- KvStateSnapshot<?, ?, ?, ?, ?> s = values.next();
- try {
- s.close();
- } catch (Throwable t) {
- if (firstException == null) {
- firstException = t;
- }
- }
- values.remove();
- }
- }
- catch (ConcurrentModificationException e) {
- // fall through the loop
- }
- }
- }
-
- if (firstException != null) {
- ExceptionUtils.rethrowIOException(firstException);
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskStateList.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskStateList.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskStateList.java
deleted file mode 100644
index ae85d86..0000000
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTaskStateList.java
+++ /dev/null
@@ -1,123 +0,0 @@
-/*
- * 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.annotation.Internal;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.ExceptionUtils;
-
-import java.io.IOException;
-import java.util.HashMap;
-
-/**
- * List of task states for a chain of streaming tasks.
- */
-@Internal
-public class StreamTaskStateList implements StateHandle<StreamTaskState[]> {
-
- private static final long serialVersionUID = 1L;
-
- /** The states for all operator */
- private final StreamTaskState[] states;
-
- public StreamTaskStateList(StreamTaskState[] states) throws Exception {
- this.states = states;
- }
-
- public boolean isEmpty() {
- for (StreamTaskState state : states) {
- if (state != null) {
- return false;
- }
- }
- return true;
- }
-
- @Override
- public StreamTaskState[] getState(ClassLoader userCodeClassLoader) {
- return states;
- }
-
- @Override
- public void discardState() throws Exception {
- for (StreamTaskState state : states) {
- if (state != null) {
- state.discardState();
- }
- }
- }
-
- @Override
- public long getStateSize() throws Exception {
- long sumStateSize = 0;
-
- if (states != null) {
- for (StreamTaskState state : states) {
- if (state != null) {
- StateHandle<?> operatorState = state.getOperatorState();
- StateHandle<?> functionState = state.getFunctionState();
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates = state.getKvStates();
-
- if (operatorState != null) {
- sumStateSize += operatorState.getStateSize();
- }
-
- if (functionState != null) {
- sumStateSize += functionState.getStateSize();
- }
-
- if (kvStates != null) {
- for (KvStateSnapshot<?, ?, ?, ?, ?> kvState : kvStates.values()) {
- if (kvState != null) {
- sumStateSize += kvState.getStateSize();
- }
- }
- }
- }
- }
- }
-
- // State size as sum of all state sizes
- return sumStateSize;
- }
-
- @Override
- public void close() throws IOException {
- if (states != null) {
- Throwable firstException = null;
-
- for (StreamTaskState state : states) {
- if (state != null) {
- try {
- state.close();
- } catch (Throwable t) {
- if (firstException == null) {
- firstException = t;
- }
- }
- }
- }
-
- if (firstException != null) {
- ExceptionUtils.rethrowIOException(firstException);
- }
- }
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSinkTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSinkTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSinkTest.java
index d5b8fa6..d1ba489 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSinkTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/GenericWriteAheadSinkTest.java
@@ -18,15 +18,12 @@
package org.apache.flink.streaming.runtime.operators;
import org.apache.flink.api.common.ExecutionConfig;
-import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple1;
import org.apache.flink.api.java.typeutils.TupleTypeInfo;
import org.apache.flink.api.java.typeutils.TypeExtractor;
import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
-import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
-import org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.junit.Assert;
import org.junit.Test;
import org.junit.runner.RunWith;
@@ -56,10 +53,11 @@ public class GenericWriteAheadSinkTest extends WriteAheadSinkTestBase<Tuple1<Int
return new Tuple1<>(counter);
}
+
@Override
protected void verifyResultsIdealCircumstances(
- OneInputStreamTaskTestHarness<Tuple1<Integer>, Tuple1<Integer>> harness,
- OneInputStreamTask<Tuple1<Integer>, Tuple1<Integer>> task, ListSink sink) {
+ OneInputStreamOperatorTestHarness<Tuple1<Integer>, Tuple1<Integer>> harness,
+ ListSink sink) {
ArrayList<Integer> list = new ArrayList<>();
for (int x = 1; x <= 60; x++) {
@@ -75,8 +73,8 @@ public class GenericWriteAheadSinkTest extends WriteAheadSinkTestBase<Tuple1<Int
@Override
protected void verifyResultsDataPersistenceUponMissedNotify(
- OneInputStreamTaskTestHarness<Tuple1<Integer>, Tuple1<Integer>> harness,
- OneInputStreamTask<Tuple1<Integer>, Tuple1<Integer>> task, ListSink sink) {
+ OneInputStreamOperatorTestHarness<Tuple1<Integer>, Tuple1<Integer>> harness,
+ ListSink sink) {
ArrayList<Integer> list = new ArrayList<>();
for (int x = 1; x <= 60; x++) {
@@ -92,8 +90,8 @@ public class GenericWriteAheadSinkTest extends WriteAheadSinkTestBase<Tuple1<Int
@Override
protected void verifyResultsDataDiscardingUponRestore(
- OneInputStreamTaskTestHarness<Tuple1<Integer>, Tuple1<Integer>> harness,
- OneInputStreamTask<Tuple1<Integer>, Tuple1<Integer>> task, ListSink sink) {
+ OneInputStreamOperatorTestHarness<Tuple1<Integer>, Tuple1<Integer>> harness,
+ ListSink sink) {
ArrayList<Integer> list = new ArrayList<>();
for (int x = 1; x <= 20; x++) {
@@ -116,62 +114,56 @@ public class GenericWriteAheadSinkTest extends WriteAheadSinkTestBase<Tuple1<Int
* and later retry of the affected checkpoints.
*/
public void testCommitterException() throws Exception {
- OperatorExposingTask<Tuple1<Integer>> task = createTask();
- TypeInformation<Tuple1<Integer>> info = createTypeInfo();
- OneInputStreamTaskTestHarness<Tuple1<Integer>, Tuple1<Integer>> testHarness = new OneInputStreamTaskTestHarness<>(task, 1, 1, info, info);
- StreamConfig streamConfig = testHarness.getStreamConfig();
- streamConfig.setCheckpointingEnabled(true);
- streamConfig.setStreamOperator(new ListSink2());
+
+ ListSink2 sink = new ListSink2();
+
+ OneInputStreamOperatorTestHarness<Tuple1<Integer>, Tuple1<Integer>> testHarness =
+ new OneInputStreamOperatorTestHarness<>(sink);
+
+ testHarness.open();
int elementCounter = 1;
- testHarness.invoke();
- testHarness.waitForTaskRunning();
-
for (int x = 0; x < 10; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 0)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- task.getOperator().snapshotOperatorState(0, 0);
- task.notifyCheckpointComplete(0);
-
+ testHarness.snapshot(0, 0);
+ testHarness.notifyOfCompletedCheckpoint(0);
+
//isCommitted should have failed, thus sendValues() should never have been called
- Assert.assertTrue(((ListSink2) task.getOperator()).values.size() == 0);
+ Assert.assertTrue(sink.values.size() == 0);
for (int x = 0; x < 10; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 1)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- task.getOperator().snapshotOperatorState(1, 0);
- task.notifyCheckpointComplete(1);
+ testHarness.snapshot(1, 0);
+ testHarness.notifyOfCompletedCheckpoint(1);
//previous CP should be retried, but will fail the CP commit. Second CP should be skipped.
- Assert.assertTrue(((ListSink2) task.getOperator()).values.size() == 10);
+ Assert.assertTrue(sink.values.size() == 10);
for (int x = 0; x < 10; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 2)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- task.getOperator().snapshotOperatorState(2, 0);
- task.notifyCheckpointComplete(2);
+ testHarness.snapshot(2, 0);
+ testHarness.notifyOfCompletedCheckpoint(2);
//all CP's should be retried and succeed; since one CP was written twice we have 2 * 10 + 10 + 10 = 40 values
- Assert.assertTrue(((ListSink2) task.getOperator()).values.size() == 40);
-
- testHarness.endInput();
- testHarness.waitForTaskCompletion();
+ Assert.assertTrue(sink.values.size() == 40);
}
/**
* Simple sink that stores all records in a public list.
*/
public static class ListSink extends GenericWriteAheadSink<Tuple1<Integer>> {
+ private static final long serialVersionUID = 1L;
+
public List<Integer> values = new ArrayList<>();
public ListSink() throws Exception {
@@ -188,6 +180,8 @@ public class GenericWriteAheadSinkTest extends WriteAheadSinkTestBase<Tuple1<Int
}
public static class SimpleCommitter extends CheckpointCommitter {
+ private static final long serialVersionUID = 1L;
+
private List<Long> checkpoints;
@Override
@@ -218,6 +212,8 @@ public class GenericWriteAheadSinkTest extends WriteAheadSinkTestBase<Tuple1<Int
* Simple sink that stores all records in a public list.
*/
public static class ListSink2 extends GenericWriteAheadSink<Tuple1<Integer>> {
+ private static final long serialVersionUID = 1L;
+
public List<Integer> values = new ArrayList<>();
public ListSink2() throws Exception {
@@ -234,6 +230,8 @@ public class GenericWriteAheadSinkTest extends WriteAheadSinkTestBase<Tuple1<Int
}
public static class FailingCommitter extends CheckpointCommitter {
+ private static final long serialVersionUID = 1L;
+
private List<Long> checkpoints;
private boolean failIsCommitted = true;
private boolean failCommit = true;
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamSourceOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamSourceOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamSourceOperatorTest.java
index 0e18166..9c06b49 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamSourceOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/StreamSourceOperatorTest.java
@@ -24,6 +24,7 @@ import org.apache.flink.api.common.functions.StoppableFunction;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
+import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.graph.StreamConfig;
@@ -224,6 +225,7 @@ public class StreamSourceOperatorTest {
executionConfig.setAutoWatermarkInterval(watermarkInterval);
StreamConfig cfg = new StreamConfig(new Configuration());
+ cfg.setStateBackend(new MemoryStateBackend());
cfg.setTimeCharacteristic(timeChar);
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/WriteAheadSinkTestBase.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/WriteAheadSinkTestBase.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/WriteAheadSinkTestBase.java
index 221d7da..b7203b5 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/WriteAheadSinkTestBase.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/WriteAheadSinkTestBase.java
@@ -19,40 +19,21 @@ package org.apache.flink.streaming.runtime.operators;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
-import org.apache.flink.streaming.api.graph.StreamConfig;
-import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.OneInputStreamTask;
-import org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.util.TestLogger;
-
import org.junit.Test;
import org.junit.runner.RunWith;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.io.ObjectInputStream;
-import java.io.ObjectOutputStream;
-import java.util.ArrayList;
-
@RunWith(PowerMockRunner.class)
@PrepareForTest(ResultPartitionWriter.class)
+@PowerMockIgnore("javax.management.*")
public abstract class WriteAheadSinkTestBase<IN, S extends GenericWriteAheadSink<IN>> extends TestLogger {
- protected static class OperatorExposingTask<INT> extends OneInputStreamTask<INT, INT> {
- public OneInputStreamOperator<INT, INT> getOperator() {
- return this.headOperator;
- }
- }
-
- protected OperatorExposingTask<IN> createTask() {
- return new OperatorExposingTask<>();
- }
-
protected abstract S createSink() throws Exception;
protected abstract TypeInformation<IN> createTypeInfo();
@@ -60,172 +41,134 @@ public abstract class WriteAheadSinkTestBase<IN, S extends GenericWriteAheadSink
protected abstract IN generateValue(int counter, int checkpointID);
protected abstract void verifyResultsIdealCircumstances(
- OneInputStreamTaskTestHarness<IN, IN> harness, OneInputStreamTask<IN, IN> task, S sink) throws Exception;
+ OneInputStreamOperatorTestHarness<IN, IN> harness, S sink) throws Exception;
protected abstract void verifyResultsDataPersistenceUponMissedNotify(
- OneInputStreamTaskTestHarness<IN, IN> harness, OneInputStreamTask<IN, IN> task, S sink) throws Exception;
+ OneInputStreamOperatorTestHarness<IN, IN> harness, S sink) throws Exception;
protected abstract void verifyResultsDataDiscardingUponRestore(
- OneInputStreamTaskTestHarness<IN, IN> harness, OneInputStreamTask<IN, IN> task, S sink) throws Exception;
+ OneInputStreamOperatorTestHarness<IN, IN> harness, S sink) throws Exception;
@Test
public void testIdealCircumstances() throws Exception {
- OperatorExposingTask<IN> task = createTask();
- TypeInformation<IN> info = createTypeInfo();
- OneInputStreamTaskTestHarness<IN, IN> testHarness = new OneInputStreamTaskTestHarness<>(task, 1, 1, info, info);
- StreamConfig streamConfig = testHarness.getStreamConfig();
- streamConfig.setCheckpointingEnabled(true);
- streamConfig.setStreamOperator(createSink());
+ S sink = createSink();
- int elementCounter = 1;
+ OneInputStreamOperatorTestHarness<IN, IN> testHarness =
+ new OneInputStreamOperatorTestHarness<>(sink);
- testHarness.invoke();
- testHarness.waitForTaskRunning();
+ testHarness.open();
- ArrayList<StreamTaskState> states = new ArrayList<>();
+ int elementCounter = 1;
+ int snapshotCount = 0;
for (int x = 0; x < 20; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 0)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- task.notifyCheckpointComplete(states.size() - 1);
+ testHarness.snapshot(snapshotCount++, 0);
+ testHarness.notifyOfCompletedCheckpoint(snapshotCount - 1);
for (int x = 0; x < 20; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 1)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- task.notifyCheckpointComplete(states.size() - 1);
+ testHarness.snapshot(snapshotCount++, 0);
+ testHarness.notifyOfCompletedCheckpoint(snapshotCount - 1);
for (int x = 0; x < 20; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 2)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
-
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- task.notifyCheckpointComplete(states.size() - 1);
- testHarness.endInput();
+ testHarness.snapshot(snapshotCount++, 0);
+ testHarness.notifyOfCompletedCheckpoint(snapshotCount - 1);
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- testHarness.waitForTaskCompletion();
-
- verifyResultsIdealCircumstances(testHarness, task, (S) task.getOperator());
+ verifyResultsIdealCircumstances(testHarness, sink);
}
@Test
public void testDataPersistenceUponMissedNotify() throws Exception {
S sink = createSink();
- OperatorExposingTask<IN> task = createTask();
- TypeInformation<IN> info = createTypeInfo();
- OneInputStreamTaskTestHarness<IN, IN> testHarness = new OneInputStreamTaskTestHarness<>(task, 1, 1, info, info);
- StreamConfig streamConfig = testHarness.getStreamConfig();
- streamConfig.setCheckpointingEnabled(true);
- streamConfig.setStreamOperator(sink);
- int elementCounter = 1;
+ OneInputStreamOperatorTestHarness<IN, IN> testHarness =
+ new OneInputStreamOperatorTestHarness<>(sink);
- testHarness.invoke();
- testHarness.waitForTaskRunning();
+ testHarness.open();
- ArrayList<StreamTaskState> states = new ArrayList<>();
+ int elementCounter = 1;
+ int snapshotCount = 0;
for (int x = 0; x < 20; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 0)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- task.notifyCheckpointComplete(states.size() - 1);
+
+ testHarness.snapshot(snapshotCount++, 0);
+ testHarness.notifyOfCompletedCheckpoint(snapshotCount - 1);
for (int x = 0; x < 20; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 1)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
+
+ testHarness.snapshot(snapshotCount++, 0);
for (int x = 0; x < 20; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 2)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- task.notifyCheckpointComplete(states.size() - 1);
- testHarness.endInput();
+ testHarness.snapshot(snapshotCount++, 0);
+ testHarness.notifyOfCompletedCheckpoint(snapshotCount - 1);
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- testHarness.waitForTaskCompletion();
-
- verifyResultsDataPersistenceUponMissedNotify(testHarness, task, (S) task.getOperator());
+ verifyResultsDataPersistenceUponMissedNotify(testHarness, sink);
}
@Test
public void testDataDiscardingUponRestore() throws Exception {
S sink = createSink();
- OperatorExposingTask<IN> task = createTask();
- TypeInformation<IN> info = createTypeInfo();
- OneInputStreamTaskTestHarness<IN, IN> testHarness = new OneInputStreamTaskTestHarness<>(task, 1, 1, info, info);
- StreamConfig streamConfig = testHarness.getStreamConfig();
- streamConfig.setCheckpointingEnabled(true);
- streamConfig.setStreamOperator(sink);
- int elementCounter = 1;
+ OneInputStreamOperatorTestHarness<IN, IN> testHarness =
+ new OneInputStreamOperatorTestHarness<>(sink);
- testHarness.invoke();
- testHarness.waitForTaskRunning();
+ testHarness.open();
- ArrayList<StreamTaskState> states = new ArrayList<>();
+ int elementCounter = 1;
+ int snapshotCount = 0;
for (int x = 0; x < 20; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 0)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- task.notifyCheckpointComplete(states.size() - 1);
+
+ StreamStateHandle latestSnapshot = testHarness.snapshot(snapshotCount++, 0);
+ testHarness.notifyOfCompletedCheckpoint(snapshotCount - 1);
for (int x = 0; x < 20; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 1)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
-
- task.getOperator().close();
- task.getOperator().open();
- task.getOperator().restoreState(states.get(states.size() - 1));
+ testHarness.close();
+
+ sink = createSink();
+
+ testHarness =new OneInputStreamOperatorTestHarness<>(sink);
+
+ testHarness.setup();
+ testHarness.restore(latestSnapshot);
+ testHarness.open();
for (int x = 0; x < 20; x++) {
testHarness.processElement(new StreamRecord<>(generateValue(elementCounter, 2)));
elementCounter++;
}
- testHarness.waitForInputProcessing();
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- task.notifyCheckpointComplete(states.size() - 1);
-
- testHarness.endInput();
-
- states.add(copyTaskState(task.getOperator().snapshotOperatorState(states.size(), 0)));
- testHarness.waitForTaskCompletion();
-
- verifyResultsDataDiscardingUponRestore(testHarness, task, (S) task.getOperator());
- }
- protected StreamTaskState copyTaskState(StreamTaskState toCopy) throws IOException, ClassNotFoundException {
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- ObjectOutputStream oos = new ObjectOutputStream(baos);
- oos.writeObject(toCopy);
+ testHarness.snapshot(snapshotCount++, 0);
+ testHarness.notifyOfCompletedCheckpoint(snapshotCount - 1);
- ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
- ObjectInputStream ois = new ObjectInputStream(bais);
- return (StreamTaskState) ois.readObject();
+ verifyResultsDataDiscardingUponRestore(testHarness, sink);
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java
index d2f8e05..2cb1809 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AccumulatingAlignedProcessingTimeWindowOperatorTest.java
@@ -21,6 +21,7 @@ package org.apache.flink.streaming.runtime.operators.windowing;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.TaskInfo;
import org.apache.flink.api.common.accumulators.Accumulator;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
@@ -28,11 +29,15 @@ import org.apache.flink.api.common.typeutils.base.IntSerializer;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.java.ClosureCleaner;
import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.operators.testutils.UnregisteredTaskMetricsGroup;
import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
import org.apache.flink.streaming.api.functions.windowing.RichWindowFunction;
import org.apache.flink.streaming.api.functions.windowing.WindowFunction;
import org.apache.flink.streaming.api.graph.StreamConfig;
@@ -41,10 +46,11 @@ import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.util.Collector;
import org.junit.After;
+import org.junit.Ignore;
import org.junit.Test;
import org.mockito.invocation.InvocationOnMock;
@@ -73,6 +79,7 @@ import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@SuppressWarnings({"serial", "SynchronizationOnLocalVariableOrMethodParameter"})
+@Ignore
public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
@SuppressWarnings("unchecked")
@@ -196,7 +203,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
op = new AccumulatingProcessingTimeWindowOperator<>(mockFunction, mockKeySelector,
StringSerializer.INSTANCE, StringSerializer.INSTANCE, 5000, 1000);
- op.setup(mockTask, new StreamConfig(new Configuration()), mockOut);
+ op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), mockOut);
op.open();
assertTrue(op.getNextSlideTime() % 1000 == 0);
assertTrue(op.getNextEvaluationTime() % 1000 == 0);
@@ -248,7 +255,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, IntSerializer.INSTANCE,
windowSize, windowSize);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
+ op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
final int numElements = 1000;
@@ -286,7 +293,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
}
@Test
- public void testSlidingWindow() {
+ public void testSlidingWindow() throws Exception {
final ScheduledExecutorService timerService = Executors.newSingleThreadScheduledExecutor();
try {
final CollectingOutput<Integer> out = new CollectingOutput<>(50);
@@ -299,7 +306,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE, 150, 50);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
+ op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
final int numElements = 1000;
@@ -340,12 +347,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
lastCount = 1;
}
}
- }
- catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
- }
- finally {
+ } finally {
timerService.shutdown();
}
}
@@ -365,7 +367,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE, 50, 50);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
+ op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
synchronized (lock) {
@@ -421,7 +423,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE, 150, 50);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
+ op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
synchronized (lock) {
@@ -468,9 +470,12 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
validatingIdentityFunction, identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE,
windowSize, windowSize);
+ OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
- op.open();
+ new OneInputStreamOperatorTestHarness<>(op);
+
+ testHarness.setup();
+ testHarness.open();
// inject some elements
final int numElementsFirst = 700;
@@ -483,11 +488,11 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
}
// draw a snapshot and dispose the window
- StreamTaskState state;
+ StreamStateHandle state;
List<Integer> resultAtSnapshot;
synchronized (lock) {
- int beforeSnapShot = out.getElements().size();
- state = op.snapshotOperatorState(1L, System.currentTimeMillis());
+ int beforeSnapShot = out.getElements().size();
+ state = testHarness.snapshot(1L, System.currentTimeMillis());
resultAtSnapshot = new ArrayList<>(out.getElements());
int afterSnapShot = out.getElements().size();
assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
@@ -511,9 +516,12 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, IntSerializer.INSTANCE,
windowSize, windowSize);
- op.setup(mockTask, new StreamConfig(new Configuration()), out2);
- op.restoreState(state);
- op.open();
+ testHarness =
+ new OneInputStreamOperatorTestHarness<>(op);
+
+ testHarness.setup();
+ testHarness.restore(state);
+ testHarness.open();
// inject some more elements
for (int i = numElementsFirst; i < numElements; i++) {
@@ -564,8 +572,11 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, IntSerializer.INSTANCE,
windowSize, windowSlide);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
- op.open();
+ OneInputStreamOperatorTestHarness<Integer, Integer> testHarness =
+ new OneInputStreamOperatorTestHarness<>(op);
+
+ testHarness.setup();
+ testHarness.open();
// inject some elements
final int numElements = 1000;
@@ -579,11 +590,11 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
}
// draw a snapshot
- StreamTaskState state;
+ StreamStateHandle state;
List<Integer> resultAtSnapshot;
synchronized (lock) {
int beforeSnapShot = out.getElements().size();
- state = op.snapshotOperatorState(1L, System.currentTimeMillis());
+ state = testHarness.snapshot(1L, System.currentTimeMillis());
resultAtSnapshot = new ArrayList<>(out.getElements());
int afterSnapShot = out.getElements().size();
assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
@@ -608,10 +619,13 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, IntSerializer.INSTANCE,
windowSize, windowSlide);
- op.setup(mockTask, new StreamConfig(new Configuration()), out2);
- op.restoreState(state);
- op.open();
-
+ testHarness =
+ new OneInputStreamOperatorTestHarness<>(op);
+
+ testHarness.setup();
+ testHarness.restore(state);
+ testHarness.open();
+
// inject again the remaining elements
for (int i = numElementsFirst; i < numElements; i++) {
@@ -670,7 +684,7 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
new StatefulFunction(), identitySelector,
IntSerializer.INSTANCE, IntSerializer.INSTANCE, 50, 50);
- op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE), out);
+ op.setup(mockTask, createTaskConfig(identitySelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
synchronized (lock) {
@@ -731,41 +745,13 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
}
// ------------------------------------------------------------------------
-
- private static class FailingFunction implements WindowFunction<Integer, Integer, Integer, TimeWindow> {
-
- private final int failAfterElements;
-
- private int numElements;
-
- FailingFunction(int failAfterElements) {
- this.failAfterElements = failAfterElements;
- }
-
- @Override
- public void apply(Integer integer,
- TimeWindow window,
- Iterable<Integer> values,
- Collector<Integer> out) throws Exception {
- for (Integer i : values) {
- out.collect(i);
- numElements++;
-
- if (numElements >= failAfterElements) {
- throw new Exception("Artificial Test Exception");
- }
- }
- }
- }
-
- // ------------------------------------------------------------------------
private static class StatefulFunction extends RichWindowFunction<Integer, Integer, Integer, TimeWindow> {
-
+
// we use a concurrent map here even though there is no concurrency, to
// get "volatile" style access to entries
static final Map<Integer, Integer> globalCounts = new ConcurrentHashMap<>();
-
+
private ValueState<Integer> state;
@Override
@@ -795,11 +781,17 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
// ------------------------------------------------------------------------
private static StreamTask<?, ?> createMockTask() {
+ Configuration configuration = new Configuration();
+ configuration.setString(ConfigConstants.STATE_BACKEND, "jobmanager");
+
StreamTask<?, ?> task = mock(StreamTask.class);
when(task.getAccumulatorMap()).thenReturn(new HashMap<String, Accumulator<?, ?>>());
when(task.getName()).thenReturn("Test task name");
when(task.getExecutionConfig()).thenReturn(new ExecutionConfig());
+ final TaskManagerRuntimeInfo mockTaskManagerRuntimeInfo = mock(TaskManagerRuntimeInfo.class);
+ when(mockTaskManagerRuntimeInfo.getConfiguration()).thenReturn(configuration);
+
final Environment env = mock(Environment.class);
when(env.getTaskInfo()).thenReturn(new TaskInfo("Test task name", 0, 1, 0));
when(env.getUserClassLoader()).thenReturn(AggregatingAlignedProcessingTimeWindowOperatorTest.class.getClassLoader());
@@ -853,11 +845,12 @@ public class AccumulatingAlignedProcessingTimeWindowOperatorTest {
return mockTask;
}
-
- private static StreamConfig createTaskConfig(KeySelector<?, ?> partitioner, TypeSerializer<?> keySerializer) {
+
+ private static StreamConfig createTaskConfig(KeySelector<?, ?> partitioner, TypeSerializer<?> keySerializer, KeyGroupAssigner<?> keyGroupAssigner) {
StreamConfig cfg = new StreamConfig(new Configuration());
cfg.setStatePartitioner(0, partitioner);
cfg.setStateKeySerializer(keySerializer);
+ cfg.setKeyGroupAssigner(keyGroupAssigner);
return cfg;
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
index 585eaa7..472dccb 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
@@ -22,6 +22,7 @@ import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.accumulators.Accumulator;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.functions.RichReduceFunction;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
@@ -32,19 +33,24 @@ import org.apache.flink.api.java.ClosureCleaner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.typeutils.TupleTypeInfo;
+import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.api.operators.Output;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.junit.After;
+import org.junit.Ignore;
import org.junit.Test;
import org.mockito.invocation.InvocationOnMock;
@@ -75,6 +81,7 @@ import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@SuppressWarnings({"serial", "SynchronizationOnLocalVariableOrMethodParameter"})
+@Ignore
public class AggregatingAlignedProcessingTimeWindowOperatorTest {
@SuppressWarnings("unchecked")
@@ -204,7 +211,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
op = new AggregatingProcessingTimeWindowOperator<>(mockFunction, mockKeySelector,
StringSerializer.INSTANCE, StringSerializer.INSTANCE, 5000, 1000);
- op.setup(mockTask, new StreamConfig(new Configuration()), mockOut);
+ op.setup(mockTask, createTaskConfig(mockKeySelector, StringSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), mockOut);
op.open();
assertTrue(op.getNextSlideTime() % 1000 == 0);
assertTrue(op.getNextEvaluationTime() % 1000 == 0);
@@ -255,8 +262,8 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
final Object lock = new Object();
final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
-
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
+
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
final int numElements = 1000;
@@ -313,8 +320,8 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
sumFunction, fieldOneSelector,
IntSerializer.INSTANCE, tupleSerializer,
windowSize, windowSize);
-
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
+
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
final int numWindows = 10;
@@ -381,7 +388,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
150, 50);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
final int numElements = 1000;
@@ -450,7 +457,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
sumFunction, fieldOneSelector,
IntSerializer.INSTANCE, tupleSerializer, 150, 50);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
synchronized (lock) {
@@ -512,12 +519,12 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
hundredYears, hundredYears);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
for (int i = 0; i < 100; i++) {
synchronized (lock) {
- StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(1, 1));
+ StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(1, 1));
op.setKeyContextElement1(next);
op.processElement(next);
}
@@ -560,8 +567,11 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
windowSize, windowSize);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
- op.open();
+ OneInputStreamOperatorTestHarness<Tuple2<Integer, Integer>, Tuple2<Integer, Integer>> testHarness =
+ new OneInputStreamOperatorTestHarness<>(op);
+
+ testHarness.setup();
+ testHarness.open();
// inject some elements
final int numElementsFirst = 700;
@@ -577,11 +587,11 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
}
// draw a snapshot and dispose the window
- StreamTaskState state;
+ StreamStateHandle state;
List<Tuple2<Integer, Integer>> resultAtSnapshot;
synchronized (lock) {
int beforeSnapShot = out.getElements().size();
- state = op.snapshotOperatorState(1L, System.currentTimeMillis());
+ state = testHarness.snapshot(1L, System.currentTimeMillis());
resultAtSnapshot = new ArrayList<>(out.getElements());
int afterSnapShot = out.getElements().size();
assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
@@ -608,9 +618,12 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
windowSize, windowSize);
- op.setup(mockTask, new StreamConfig(new Configuration()), out2);
- op.restoreState(state);
- op.open();
+ testHarness =
+ new OneInputStreamOperatorTestHarness<>(op);
+
+ testHarness.setup();
+ testHarness.restore(state);
+ testHarness.open();
// inject the remaining elements
for (int i = numElementsFirst; i < numElements; i++) {
@@ -668,8 +681,11 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
windowSize, windowSlide);
- op.setup(mockTask, new StreamConfig(new Configuration()), out);
- op.open();
+ OneInputStreamOperatorTestHarness<Tuple2<Integer, Integer>, Tuple2<Integer, Integer>> testHarness =
+ new OneInputStreamOperatorTestHarness<>(op);
+
+ testHarness.setup();
+ testHarness.open();
// inject some elements
final int numElements = 1000;
@@ -685,11 +701,11 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
}
// draw a snapshot
- StreamTaskState state;
+ StreamStateHandle state;
List<Tuple2<Integer, Integer>> resultAtSnapshot;
synchronized (lock) {
int beforeSnapShot = out.getElements().size();
- state = op.snapshotOperatorState(1L, System.currentTimeMillis());
+ state = testHarness.snapshot(1L, System.currentTimeMillis());
resultAtSnapshot = new ArrayList<>(out.getElements());
int afterSnapShot = out.getElements().size();
assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
@@ -716,10 +732,12 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
windowSize, windowSlide);
- op.setup(mockTask, new StreamConfig(new Configuration()), out2);
- op.restoreState(state);
- op.open();
+ testHarness =
+ new OneInputStreamOperatorTestHarness<>(op);
+ testHarness.setup();
+ testHarness.restore(state);
+ testHarness.open();
// inject again the remaining elements
for (int i = numElementsFirst; i < numElements; i++) {
@@ -782,7 +800,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
new StatefulFunction(), fieldOneSelector,
IntSerializer.INSTANCE, tupleSerializer, twoSeconds, twoSeconds);
- op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
// because the window interval is so large, everything should be in one window
@@ -837,7 +855,7 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
new StatefulFunction(), fieldOneSelector,
IntSerializer.INSTANCE, tupleSerializer, windowSize, windowSlide);
- op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE), out);
+ op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
op.open();
// because the window interval is so large, everything should be in one window
@@ -959,11 +977,17 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
// ------------------------------------------------------------------------
private static StreamTask<?, ?> createMockTask() {
+ Configuration configuration = new Configuration();
+ configuration.setString(ConfigConstants.STATE_BACKEND, "jobmanager");
+
StreamTask<?, ?> task = mock(StreamTask.class);
when(task.getAccumulatorMap()).thenReturn(new HashMap<String, Accumulator<?, ?>>());
when(task.getName()).thenReturn("Test task name");
when(task.getExecutionConfig()).thenReturn(new ExecutionConfig());
+ final TaskManagerRuntimeInfo mockTaskManagerRuntimeInfo = mock(TaskManagerRuntimeInfo.class);
+ when(mockTaskManagerRuntimeInfo.getConfiguration()).thenReturn(configuration);
+
final Environment env = new DummyEnvironment("Test task name", 1, 0);
when(task.getEnvironment()).thenReturn(env);
@@ -1014,10 +1038,11 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
return mockTask;
}
- private static StreamConfig createTaskConfig(KeySelector<?, ?> partitioner, TypeSerializer<?> keySerializer) {
+ private static StreamConfig createTaskConfig(KeySelector<?, ?> partitioner, TypeSerializer<?> keySerializer, KeyGroupAssigner<?> keyGroupAssigner) {
StreamConfig cfg = new StreamConfig(new Configuration());
cfg.setStatePartitioner(0, partitioner);
cfg.setStateKeySerializer(keySerializer);
+ cfg.setKeyGroupAssigner(keyGroupAssigner);
return cfg;
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java
index 2d7b615..cda6e1e 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java
@@ -32,6 +32,7 @@ import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.api.java.typeutils.TypeInfoParser;
import org.apache.flink.configuration.Configuration;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.streaming.api.datastream.WindowedStream;
import org.apache.flink.streaming.api.environment.LocalStreamEnvironment;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
@@ -62,12 +63,12 @@ import org.apache.flink.streaming.api.windowing.windows.Window;
import org.apache.flink.streaming.runtime.operators.windowing.functions.InternalIterableWindowFunction;
import org.apache.flink.streaming.runtime.operators.windowing.functions.InternalSingleValueWindowFunction;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import org.apache.flink.streaming.runtime.tasks.TestTimeServiceProvider;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.TestHarnessUtil;
import org.apache.flink.streaming.util.WindowingTestHarness;
import org.apache.flink.util.Collector;
+import org.apache.flink.util.TestLogger;
import org.junit.Assert;
import org.junit.Test;
@@ -82,7 +83,7 @@ import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
-public class WindowOperatorTest {
+public class WindowOperatorTest extends TestLogger {
// For counting if close() is called the correct number of times on the SumReducer
private static AtomicInteger closeCalled = new AtomicInteger(0);
@@ -123,10 +124,10 @@ public class WindowOperatorTest {
TestHarnessUtil.assertOutputEqualsSorted("Output was not correct.", expectedOutput, testHarness.getOutput(), new Tuple2ResultSortComparator());
// do a snapshot, close and restore again
- StreamTaskState snapshot = testHarness.snapshot(0L, 0L);
+ StreamStateHandle snapshot = testHarness.snapshot(0L, 0L);
testHarness.close();
testHarness.setup();
- testHarness.restore(snapshot, 10L);
+ testHarness.restore(snapshot);
testHarness.open();
testHarness.processWatermark(new Watermark(3999));
@@ -261,10 +262,10 @@ public class WindowOperatorTest {
TestHarnessUtil.assertOutputEqualsSorted("Output was not correct.", expectedOutput, testHarness.getOutput(), new Tuple2ResultSortComparator());
// do a snapshot, close and restore again
- StreamTaskState snapshot = testHarness.snapshot(0L, 0L);
+ StreamStateHandle snapshot = testHarness.snapshot(0L, 0L);
testHarness.close();
testHarness.setup();
- testHarness.restore(snapshot, 10L);
+ testHarness.restore(snapshot);
testHarness.open();
testHarness.processWatermark(new Watermark(2999));
@@ -414,10 +415,10 @@ public class WindowOperatorTest {
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key1", 2), 1000));
// do a snapshot, close and restore again
- StreamTaskState snapshot = testHarness.snapshot(0L, 0L);
+ StreamStateHandle snapshot = testHarness.snapshot(0L, 0L);
testHarness.close();
testHarness.setup();
- testHarness.restore(snapshot, 10L);
+ testHarness.restore(snapshot);
testHarness.open();
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key1", 3), 2500));
@@ -487,10 +488,10 @@ public class WindowOperatorTest {
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key2", 3), 2500));
// do a snapshot, close and restore again
- StreamTaskState snapshot = testHarness.snapshot(0L, 0L);
+ StreamStateHandle snapshot = testHarness.snapshot(0L, 0L);
testHarness.close();
testHarness.setup();
- testHarness.restore(snapshot, 10L);
+ testHarness.restore(snapshot);
testHarness.open();
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key1", 1), 10));
@@ -569,10 +570,10 @@ public class WindowOperatorTest {
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key1", 2), 1000));
// do a snapshot, close and restore again
- StreamTaskState snapshot = testHarness.snapshot(0L, 0L);
+ StreamStateHandle snapshot = testHarness.snapshot(0L, 0L);
testHarness.close();
testHarness.setup();
- testHarness.restore(snapshot, 10L);
+ testHarness.restore(snapshot);
testHarness.open();
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key1", 3), 2500));
@@ -604,7 +605,7 @@ public class WindowOperatorTest {
WindowedStream<String, String, TimeWindow> windowedStream = env.fromElements("Hello", "Ciao")
.keyBy(new KeySelector<String, String>() {
- private static final long serialVersionUID = 1L;
+ private static final long serialVersionUID = -887743259776124087L;
@Override
public String getKey(String value) throws Exception {
@@ -671,10 +672,10 @@ public class WindowOperatorTest {
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key2", 33), 1000));
// do a snapshot, close and restore again
- StreamTaskState snapshot = testHarness.snapshot(0L, 0L);
+ StreamStateHandle snapshot = testHarness.snapshot(0L, 0L);
testHarness.close();
testHarness.setup();
- testHarness.restore(snapshot, 10L);
+ testHarness.restore(snapshot);
testHarness.open();
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key2", 33), 2500));
@@ -835,10 +836,10 @@ public class WindowOperatorTest {
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key2", 1), 1999));
// do a snapshot, close and restore again
- StreamTaskState snapshot = testHarness.snapshot(0L, 0L);
+ StreamStateHandle snapshot = testHarness.snapshot(0L, 0L);
testHarness.close();
testHarness.setup();
- testHarness.restore(snapshot, 10L);
+ testHarness.restore(snapshot);
testHarness.open();
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key2", 1), 1000));
@@ -908,7 +909,7 @@ public class WindowOperatorTest {
operator.processingTimeTimerTimestamps.add(3L, 1);
- StreamTaskState snapshot = testHarness.snapshot(0, 0);
+ StreamStateHandle snapshot = testHarness.snapshot(0, 0);
WindowOperator<String, Tuple2<String, Integer>, Tuple2<String, Integer>, Tuple2<String, Integer>, TimeWindow> otherOperator = new WindowOperator<>(
SlidingEventTimeWindows.of(Time.of(WINDOW_SIZE, TimeUnit.SECONDS), Time.of(WINDOW_SLIDE, TimeUnit.SECONDS)),
@@ -927,7 +928,7 @@ public class WindowOperatorTest {
otherOperator.setInputType(inputType, new ExecutionConfig());
otherTestHarness.setup();
- otherTestHarness.restore(snapshot, 0);
+ otherTestHarness.restore(snapshot);
otherTestHarness.open();
Assert.assertEquals(operator.processingTimeTimers, otherOperator.processingTimeTimers);
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowingTestHarnessTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowingTestHarnessTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowingTestHarnessTest.java
index 7242e1c..58a7897 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowingTestHarnessTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowingTestHarnessTest.java
@@ -24,13 +24,13 @@ import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.typeutils.TypeInfoParser;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.assigners.TumblingProcessingTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.triggers.EventTimeTrigger;
import org.apache.flink.streaming.api.windowing.triggers.ProcessingTimeTrigger;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import org.apache.flink.streaming.util.WindowingTestHarness;
import org.junit.Test;
@@ -175,9 +175,9 @@ public class WindowingTestHarnessTest {
testHarness.compareActualToExpectedOutput("Output was not correct.");
// do a snapshot, close and restore again
- StreamTaskState snapshot = testHarness.snapshot(0L, 0L);
+ StreamStateHandle snapshot = testHarness.snapshot(0L, 0L);
testHarness.close();
- testHarness.restore(snapshot, 10L);
+ testHarness.restore(snapshot);
testHarness.processWatermark(2999);
[24/27] flink git commit: [FLINK-4380] Add tests for new
Key-Group/Max-Parallelism
Posted by al...@apache.org.
[FLINK-4380] Add tests for new Key-Group/Max-Parallelism
This tests the rescaling features in CheckpointCoordinator and
SavepointCoordinator.
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/516ad011
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/516ad011
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/516ad011
Branch: refs/heads/master
Commit: 516ad011865ca5beece273ca9b985e2861b3435a
Parents: 847ead0
Author: Till Rohrmann <tr...@apache.org>
Authored: Thu Aug 11 12:14:18 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:10:01 2016 +0200
----------------------------------------------------------------------
.../checkpoint/CheckpointCoordinatorTest.java | 733 ++++++++++++++++++-
.../runtime/tasks/OneInputStreamTaskTest.java | 280 ++++++-
.../test/checkpointing/RescalingITCase.java | 1 -
3 files changed, 1007 insertions(+), 7 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/516ad011/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
index 50330fa..495dced 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
@@ -18,28 +18,45 @@
package org.apache.flink.runtime.checkpoint;
+import com.google.common.collect.Iterables;
import org.apache.flink.api.common.JobID;
+import org.apache.flink.core.fs.FSDataInputStream;
import org.apache.flink.runtime.checkpoint.savepoint.HeapSavepointStore;
import org.apache.flink.runtime.checkpoint.stats.DisabledCheckpointStatsTracker;
import org.apache.flink.runtime.execution.ExecutionState;
import org.apache.flink.runtime.executiongraph.Execution;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
+import org.apache.flink.runtime.executiongraph.ExecutionJobVertex;
import org.apache.flink.runtime.executiongraph.ExecutionVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.messages.checkpoint.AcknowledgeCheckpoint;
import org.apache.flink.runtime.messages.checkpoint.DeclineCheckpoint;
import org.apache.flink.runtime.messages.checkpoint.NotifyCheckpointComplete;
import org.apache.flink.runtime.messages.checkpoint.TriggerCheckpoint;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeOffsets;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+import org.apache.flink.util.InstantiationUtil;
+import org.apache.flink.util.Preconditions;
+import org.junit.Assert;
import org.junit.Test;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
+import scala.concurrent.ExecutionContext;
import scala.concurrent.Future;
+import java.io.IOException;
import java.io.Serializable;
import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
@@ -47,12 +64,14 @@ import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -124,7 +143,7 @@ public class CheckpointCoordinatorTest {
final ExecutionAttemptID triggerAttemptID1 = new ExecutionAttemptID();
final ExecutionAttemptID triggerAttemptID2 = new ExecutionAttemptID();
ExecutionVertex triggerVertex1 = mockExecutionVertex(triggerAttemptID1);
- ExecutionVertex triggerVertex2 = mockExecutionVertex(triggerAttemptID2, ExecutionState.FINISHED);
+ ExecutionVertex triggerVertex2 = mockExecutionVertex(triggerAttemptID2, new JobVertexID(), 1, 1, ExecutionState.FINISHED);
// create some mock Execution vertices that need to ack the checkpoint
final ExecutionAttemptID ackAttemptID1 = new ExecutionAttemptID();
@@ -1529,7 +1548,7 @@ public class CheckpointCoordinatorTest {
coord.startCheckpointScheduler();
// after a while, there should be exactly as many checkpoints
- // as concurrently permitted
+ // as concurrently permitted
long now = System.currentTimeMillis();
long timeout = now + 60000;
long minDuration = now + 100;
@@ -1622,7 +1641,7 @@ public class CheckpointCoordinatorTest {
}
while (System.currentTimeMillis() < timeout &&
coord.getNumberOfPendingCheckpoints() == 0);
-
+
assertTrue(coord.getNumberOfPendingCheckpoints() > 0);
}
catch (Exception e) {
@@ -1738,4 +1757,712 @@ public class CheckpointCoordinatorTest {
return vertex;
}
+/**
+ * Tests that the checkpointed partitioned and non-partitioned state is assigned properly to
+ * the {@link Execution} upon recovery.
+ *
+ * @throws Exception
+ */
+ @Test
+ public void testRestoreLatestCheckpointedState() throws Exception {
+ final JobID jid = new JobID();
+ final long timestamp = System.currentTimeMillis();
+
+ final JobVertexID jobVertexID1 = new JobVertexID();
+ final JobVertexID jobVertexID2 = new JobVertexID();
+ int parallelism1 = 3;
+ int parallelism2 = 2;
+ int maxParallelism1 = 42;
+ int maxParallelism2 = 13;
+
+ final ExecutionJobVertex jobVertex1 = mockExecutionJobVertex(
+ jobVertexID1,
+ parallelism1,
+ maxParallelism1);
+ final ExecutionJobVertex jobVertex2 = mockExecutionJobVertex(
+ jobVertexID2,
+ parallelism2,
+ maxParallelism2);
+
+ List<ExecutionVertex> allExecutionVertices = new ArrayList<>(parallelism1 + parallelism2);
+
+ allExecutionVertices.addAll(Arrays.asList(jobVertex1.getTaskVertices()));
+ allExecutionVertices.addAll(Arrays.asList(jobVertex2.getTaskVertices()));
+
+ ExecutionVertex[] arrayExecutionVertices = allExecutionVertices.toArray(new ExecutionVertex[0]);
+
+ // set up the coordinator and validate the initial state
+ CheckpointCoordinator coord = new CheckpointCoordinator(
+ jid,
+ 600000,
+ 600000,
+ 0,
+ Integer.MAX_VALUE,
+ arrayExecutionVertices,
+ arrayExecutionVertices,
+ arrayExecutionVertices,
+ cl,
+ new StandaloneCheckpointIDCounter(),
+ new StandaloneCompletedCheckpointStore(1, cl),
+ new HeapSavepointStore(),
+ new DisabledCheckpointStatsTracker());
+
+ // trigger the checkpoint
+ coord.triggerCheckpoint(timestamp);
+
+ assertTrue(coord.getPendingCheckpoints().keySet().size() == 1);
+ long checkpointId = Iterables.getOnlyElement(coord.getPendingCheckpoints().keySet());
+
+ List<KeyGroupRange> keyGroupPartitions1 = coord.createKeyGroupPartitions(maxParallelism1, parallelism1);
+ List<KeyGroupRange> keyGroupPartitions2 = coord.createKeyGroupPartitions(maxParallelism2, parallelism2);
+
+ for (int index = 0; index < jobVertex1.getParallelism(); index++) {
+ ChainedStateHandle<StreamStateHandle> nonPartitionedState = generateStateForVertex(jobVertexID1, index);
+ List<KeyGroupsStateHandle> partitionedKeyGroupState = generateKeyGroupState(jobVertexID1, keyGroupPartitions1.get(index));
+
+ AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
+ jid,
+ jobVertex1.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
+ checkpointId,
+ nonPartitionedState,
+ partitionedKeyGroupState);
+
+ coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
+ }
+
+
+ for (int index = 0; index < jobVertex2.getParallelism(); index++) {
+ ChainedStateHandle<StreamStateHandle> nonPartitionedState = generateStateForVertex(jobVertexID2, index);
+ List<KeyGroupsStateHandle> partitionedKeyGroupState = generateKeyGroupState(jobVertexID2, keyGroupPartitions2.get(index));
+
+ AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
+ jid,
+ jobVertex2.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
+ checkpointId,
+ nonPartitionedState,
+ partitionedKeyGroupState);
+
+ coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
+ }
+
+ List<CompletedCheckpoint> completedCheckpoints = coord.getSuccessfulCheckpoints();
+
+ assertEquals(1, completedCheckpoints.size());
+
+ Map<JobVertexID, ExecutionJobVertex> tasks = new HashMap<>();
+
+ tasks.put(jobVertexID1, jobVertex1);
+ tasks.put(jobVertexID2, jobVertex2);
+
+ coord.restoreLatestCheckpointedState(tasks, true, true);
+
+ // verify the restored state
+ verifiyStateRestore(jobVertexID1, jobVertex1, keyGroupPartitions1);
+ verifiyStateRestore(jobVertexID2, jobVertex2, keyGroupPartitions2);
+ }
+
+ /**
+ * Tests that the checkpoint restoration fails if the max parallelism of the job vertices has
+ * changed.
+ *
+ * @throws Exception
+ */
+ @Test(expected=IllegalStateException.class)
+ public void testRestoreLatestCheckpointFailureWhenMaxParallelismChanges() throws Exception {
+ final JobID jid = new JobID();
+ final long timestamp = System.currentTimeMillis();
+
+ final JobVertexID jobVertexID1 = new JobVertexID();
+ final JobVertexID jobVertexID2 = new JobVertexID();
+ int parallelism1 = 3;
+ int parallelism2 = 2;
+ int maxParallelism1 = 42;
+ int maxParallelism2 = 13;
+
+ final ExecutionJobVertex jobVertex1 = mockExecutionJobVertex(
+ jobVertexID1,
+ parallelism1,
+ maxParallelism1);
+ final ExecutionJobVertex jobVertex2 = mockExecutionJobVertex(
+ jobVertexID2,
+ parallelism2,
+ maxParallelism2);
+
+ List<ExecutionVertex> allExecutionVertices = new ArrayList<>(parallelism1 + parallelism2);
+
+ allExecutionVertices.addAll(Arrays.asList(jobVertex1.getTaskVertices()));
+ allExecutionVertices.addAll(Arrays.asList(jobVertex2.getTaskVertices()));
+
+ ExecutionVertex[] arrayExecutionVertices = allExecutionVertices.toArray(new ExecutionVertex[0]);
+
+ // set up the coordinator and validate the initial state
+ CheckpointCoordinator coord = new CheckpointCoordinator(
+ jid,
+ 600000,
+ 600000,
+ 0,
+ Integer.MAX_VALUE,
+ arrayExecutionVertices,
+ arrayExecutionVertices,
+ arrayExecutionVertices,
+ cl,
+ new StandaloneCheckpointIDCounter(),
+ new StandaloneCompletedCheckpointStore(1, cl),
+ new HeapSavepointStore(),
+ new DisabledCheckpointStatsTracker());
+
+ // trigger the checkpoint
+ coord.triggerCheckpoint(timestamp);
+
+ assertTrue(coord.getPendingCheckpoints().keySet().size() == 1);
+ long checkpointId = Iterables.getOnlyElement(coord.getPendingCheckpoints().keySet());
+
+ List<KeyGroupRange> keyGroupPartitions1 = coord.createKeyGroupPartitions(maxParallelism1, parallelism1);
+ List<KeyGroupRange> keyGroupPartitions2 = coord.createKeyGroupPartitions(maxParallelism2, parallelism2);
+
+ for (int index = 0; index < jobVertex1.getParallelism(); index++) {
+ ChainedStateHandle<StreamStateHandle> valueSizeTuple = generateStateForVertex(jobVertexID1, index);
+ List<KeyGroupsStateHandle> keyGroupState = generateKeyGroupState(jobVertexID1, keyGroupPartitions1.get(index));
+
+ AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
+ jid,
+ jobVertex1.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
+ checkpointId,
+ valueSizeTuple,
+ keyGroupState);
+
+ coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
+ }
+
+
+ for (int index = 0; index < jobVertex2.getParallelism(); index++) {
+ ChainedStateHandle<StreamStateHandle> valueSizeTuple = generateStateForVertex(jobVertexID2, index);
+ List<KeyGroupsStateHandle> keyGroupState = generateKeyGroupState(jobVertexID2, keyGroupPartitions2.get(index));
+
+ AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
+ jid,
+ jobVertex2.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
+ checkpointId,
+ valueSizeTuple,
+ keyGroupState);
+
+ coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
+ }
+
+ List<CompletedCheckpoint> completedCheckpoints = coord.getSuccessfulCheckpoints();
+
+ assertEquals(1, completedCheckpoints.size());
+
+ Map<JobVertexID, ExecutionJobVertex> tasks = new HashMap<>();
+
+ int newMaxParallelism1 = 20;
+ int newMaxParallelism2 = 42;
+
+ final ExecutionJobVertex newJobVertex1 = mockExecutionJobVertex(
+ jobVertexID1,
+ parallelism1,
+ newMaxParallelism1);
+
+ final ExecutionJobVertex newJobVertex2 = mockExecutionJobVertex(
+ jobVertexID2,
+ parallelism2,
+ newMaxParallelism2);
+
+ tasks.put(jobVertexID1, newJobVertex1);
+ tasks.put(jobVertexID2, newJobVertex2);
+
+ coord.restoreLatestCheckpointedState(tasks, true, true);
+
+ fail("The restoration should have failed because the max parallelism changed.");
+ }
+
+ /**
+ * Tests that the checkpoint restoration fails if the parallelism of a job vertices with
+ * non-partitioned state has changed.
+ *
+ * @throws Exception
+ */
+ @Test(expected=IllegalStateException.class)
+ public void testRestoreLatestCheckpointFailureWhenParallelismChanges() throws Exception {
+ final JobID jid = new JobID();
+ final long timestamp = System.currentTimeMillis();
+
+ final JobVertexID jobVertexID1 = new JobVertexID();
+ final JobVertexID jobVertexID2 = new JobVertexID();
+ int parallelism1 = 3;
+ int parallelism2 = 2;
+ int maxParallelism1 = 42;
+ int maxParallelism2 = 13;
+
+ final ExecutionJobVertex jobVertex1 = mockExecutionJobVertex(
+ jobVertexID1,
+ parallelism1,
+ maxParallelism1);
+ final ExecutionJobVertex jobVertex2 = mockExecutionJobVertex(
+ jobVertexID2,
+ parallelism2,
+ maxParallelism2);
+
+ List<ExecutionVertex> allExecutionVertices = new ArrayList<>(parallelism1 + parallelism2);
+
+ allExecutionVertices.addAll(Arrays.asList(jobVertex1.getTaskVertices()));
+ allExecutionVertices.addAll(Arrays.asList(jobVertex2.getTaskVertices()));
+
+ ExecutionVertex[] arrayExecutionVertices = allExecutionVertices.toArray(new ExecutionVertex[0]);
+
+ // set up the coordinator and validate the initial state
+ CheckpointCoordinator coord = new CheckpointCoordinator(
+ jid,
+ 600000,
+ 600000,
+ 0,
+ Integer.MAX_VALUE,
+ arrayExecutionVertices,
+ arrayExecutionVertices,
+ arrayExecutionVertices,
+ cl,
+ new StandaloneCheckpointIDCounter(),
+ new StandaloneCompletedCheckpointStore(1, cl),
+ new HeapSavepointStore(),
+ new DisabledCheckpointStatsTracker());
+
+ // trigger the checkpoint
+ coord.triggerCheckpoint(timestamp);
+
+ assertTrue(coord.getPendingCheckpoints().keySet().size() == 1);
+ long checkpointId = Iterables.getOnlyElement(coord.getPendingCheckpoints().keySet());
+
+ List<KeyGroupRange> keyGroupPartitions1 = coord.createKeyGroupPartitions(maxParallelism1, parallelism1);
+ List<KeyGroupRange> keyGroupPartitions2 = coord.createKeyGroupPartitions(maxParallelism2, parallelism2);
+
+ for (int index = 0; index < jobVertex1.getParallelism(); index++) {
+ ChainedStateHandle<StreamStateHandle> valueSizeTuple = generateStateForVertex(jobVertexID1, index);
+ List<KeyGroupsStateHandle> keyGroupState = generateKeyGroupState(
+ jobVertexID1, keyGroupPartitions1.get(index));
+
+ AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
+ jid,
+ jobVertex1.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
+ checkpointId,
+ valueSizeTuple,
+ keyGroupState);
+
+ coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
+ }
+
+
+ for (int index = 0; index < jobVertex2.getParallelism(); index++) {
+
+ ChainedStateHandle<StreamStateHandle> state = generateStateForVertex(jobVertexID2, index);
+ List<KeyGroupsStateHandle> keyGroupState = generateKeyGroupState(
+ jobVertexID2, keyGroupPartitions2.get(index));
+
+ AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
+ jid,
+ jobVertex2.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
+ checkpointId,
+ state,
+ keyGroupState);
+
+ coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
+ }
+
+ List<CompletedCheckpoint> completedCheckpoints = coord.getSuccessfulCheckpoints();
+
+ assertEquals(1, completedCheckpoints.size());
+
+ Map<JobVertexID, ExecutionJobVertex> tasks = new HashMap<>();
+
+ int newParallelism1 = 4;
+ int newParallelism2 = 3;
+
+ final ExecutionJobVertex newJobVertex1 = mockExecutionJobVertex(
+ jobVertexID1,
+ newParallelism1,
+ maxParallelism1);
+
+ final ExecutionJobVertex newJobVertex2 = mockExecutionJobVertex(
+ jobVertexID2,
+ newParallelism2,
+ maxParallelism2);
+
+ tasks.put(jobVertexID1, newJobVertex1);
+ tasks.put(jobVertexID2, newJobVertex2);
+
+ coord.restoreLatestCheckpointedState(tasks, true, true);
+
+ fail("The restoration should have failed because the parallelism of an vertex with " +
+ "non-partitioned state changed.");
+ }
+
+ /**
+ * Tests the checkpoint restoration with changing parallelism of job vertex with partitioned
+ * state.
+ *
+ * @throws Exception
+ */
+ @Test
+ public void testRestoreLatestCheckpointedStateWithChangingParallelism() throws Exception {
+ final JobID jid = new JobID();
+ final long timestamp = System.currentTimeMillis();
+
+ final JobVertexID jobVertexID1 = new JobVertexID();
+ final JobVertexID jobVertexID2 = new JobVertexID();
+ int parallelism1 = 3;
+ int parallelism2 = 2;
+ int maxParallelism1 = 42;
+ int maxParallelism2 = 13;
+
+ final ExecutionJobVertex jobVertex1 = mockExecutionJobVertex(
+ jobVertexID1,
+ parallelism1,
+ maxParallelism1);
+ final ExecutionJobVertex jobVertex2 = mockExecutionJobVertex(
+ jobVertexID2,
+ parallelism2,
+ maxParallelism2);
+
+ List<ExecutionVertex> allExecutionVertices = new ArrayList<>(parallelism1 + parallelism2);
+
+ allExecutionVertices.addAll(Arrays.asList(jobVertex1.getTaskVertices()));
+ allExecutionVertices.addAll(Arrays.asList(jobVertex2.getTaskVertices()));
+
+ ExecutionVertex[] arrayExecutionVertices = allExecutionVertices.toArray(new ExecutionVertex[0]);
+
+ // set up the coordinator and validate the initial state
+ CheckpointCoordinator coord = new CheckpointCoordinator(
+ jid,
+ 600000,
+ 600000,
+ 0,
+ Integer.MAX_VALUE,
+ arrayExecutionVertices,
+ arrayExecutionVertices,
+ arrayExecutionVertices,
+ cl,
+ new StandaloneCheckpointIDCounter(),
+ new StandaloneCompletedCheckpointStore(1, cl),
+ new HeapSavepointStore(),
+ new DisabledCheckpointStatsTracker());
+
+ // trigger the checkpoint
+ coord.triggerCheckpoint(timestamp);
+
+ assertTrue(coord.getPendingCheckpoints().keySet().size() == 1);
+ long checkpointId = Iterables.getOnlyElement(coord.getPendingCheckpoints().keySet());
+
+ List<KeyGroupRange> keyGroupPartitions1 = coord.createKeyGroupPartitions(maxParallelism1, parallelism1);
+ List<KeyGroupRange> keyGroupPartitions2 = coord.createKeyGroupPartitions(maxParallelism2, parallelism2);
+
+ for (int index = 0; index < jobVertex1.getParallelism(); index++) {
+ ChainedStateHandle<StreamStateHandle> valueSizeTuple = generateStateForVertex(jobVertexID1, index);
+ List<KeyGroupsStateHandle> keyGroupState = generateKeyGroupState(jobVertexID1, keyGroupPartitions1.get(index));
+
+ AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
+ jid,
+ jobVertex1.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
+ checkpointId,
+ valueSizeTuple,
+ keyGroupState);
+
+ coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
+ }
+
+
+ for (int index = 0; index < jobVertex2.getParallelism(); index++) {
+ List<KeyGroupsStateHandle> keyGroupState = generateKeyGroupState(jobVertexID2, keyGroupPartitions2.get(index));
+
+ AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
+ jid,
+ jobVertex2.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
+ checkpointId,
+ null,
+ keyGroupState);
+
+ coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
+ }
+
+ List<CompletedCheckpoint> completedCheckpoints = coord.getSuccessfulCheckpoints();
+
+ assertEquals(1, completedCheckpoints.size());
+
+ Map<JobVertexID, ExecutionJobVertex> tasks = new HashMap<>();
+
+ int newParallelism2 = 13;
+
+ List<KeyGroupRange> newKeyGroupPartitions2 = coord.createKeyGroupPartitions(maxParallelism2, newParallelism2);
+
+ final ExecutionJobVertex newJobVertex1 = mockExecutionJobVertex(
+ jobVertexID1,
+ parallelism1,
+ maxParallelism1);
+
+ final ExecutionJobVertex newJobVertex2 = mockExecutionJobVertex(
+ jobVertexID2,
+ newParallelism2,
+ maxParallelism2);
+
+ tasks.put(jobVertexID1, newJobVertex1);
+ tasks.put(jobVertexID2, newJobVertex2);
+ coord.restoreLatestCheckpointedState(tasks, true, true);
+
+ // verify the restored state
+ verifiyStateRestore(jobVertexID1, newJobVertex1, keyGroupPartitions1);
+
+ for (int i = 0; i < newJobVertex2.getParallelism(); i++) {
+ List<KeyGroupsStateHandle> originalKeyGroupState = generateKeyGroupState(jobVertexID2, newKeyGroupPartitions2.get(i));
+
+ ChainedStateHandle<StreamStateHandle> operatorState = newJobVertex2.getTaskVertices()[i].getCurrentExecutionAttempt().getChainedStateHandle();
+ List<KeyGroupsStateHandle> keyGroupState = newJobVertex2.getTaskVertices()[i].getCurrentExecutionAttempt().getKeyGroupsStateHandles();
+
+ assertNull(operatorState);
+ comparePartitionedState(originalKeyGroupState, keyGroupState);
+ }
+ }
+
+ // ------------------------------------------------------------------------
+ // Utilities
+ // ------------------------------------------------------------------------
+
+ static void sendAckMessageToCoordinator(
+ CheckpointCoordinator coord,
+ long checkpointId, JobID jid,
+ ExecutionJobVertex jobVertex,
+ JobVertexID jobVertexID,
+ List<KeyGroupRange> keyGroupPartitions) throws Exception {
+
+ for (int index = 0; index < jobVertex.getParallelism(); index++) {
+ ChainedStateHandle<StreamStateHandle> state = generateStateForVertex(jobVertexID, index);
+ List<KeyGroupsStateHandle> keyGroupState = generateKeyGroupState(
+ jobVertexID,
+ keyGroupPartitions.get(index));
+
+ AcknowledgeCheckpoint acknowledgeCheckpoint = new AcknowledgeCheckpoint(
+ jid,
+ jobVertex.getTaskVertices()[index].getCurrentExecutionAttempt().getAttemptId(),
+ checkpointId,
+ state,
+ keyGroupState);
+
+ coord.receiveAcknowledgeMessage(acknowledgeCheckpoint);
+ }
+ }
+
+ public static List<KeyGroupsStateHandle> generateKeyGroupState(
+ JobVertexID jobVertexID,
+ KeyGroupRange keyGroupPartition) throws IOException {
+
+ KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(keyGroupPartition);
+ List<Integer> testStatesLists = new ArrayList<>(keyGroupPartition.getNumberOfKeyGroups());
+ int runningGroupsOffset = 0;
+ // generate state for one keygroup
+ for (int keyGroupIndex : keyGroupPartition) {
+ Random random = new Random(jobVertexID.hashCode() + keyGroupIndex);
+ int simulatedStateValue = random.nextInt();
+ testStatesLists.add(simulatedStateValue);
+ }
+
+ return generateKeyGroupState(keyGroupPartition, testStatesLists);
+ }
+
+ public static List<KeyGroupsStateHandle> generateKeyGroupState(KeyGroupRange keyGroupRange, List< ? extends Serializable> states) throws IOException {
+ Preconditions.checkArgument(keyGroupRange.getNumberOfKeyGroups() == states.size());
+
+ long[] offsets = new long[keyGroupRange.getNumberOfKeyGroups()];
+ List<byte[]> serializedGroupValues = new ArrayList<>(offsets.length);
+
+ KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(keyGroupRange, offsets);
+
+ int runningGroupsOffset = 0;
+ // generate test state for all keygroups
+ int idx = 0;
+ for (int keyGroup : keyGroupRange) {
+ keyGroupRangeOffsets.setKeyGroupOffset(keyGroup,runningGroupsOffset);
+ byte[] serializedValue = InstantiationUtil.serializeObject(states.get(idx));
+ runningGroupsOffset += serializedValue.length;
+ serializedGroupValues.add(serializedValue);
+ ++idx;
+ }
+
+ //write all generated values in a single byte array, which is index by groupOffsetsInFinalByteArray
+ byte[] allSerializedValuesConcatenated = new byte[runningGroupsOffset];
+ runningGroupsOffset = 0;
+ byte[] old = null;
+ for(byte[] serializedGroupValue : serializedGroupValues) {
+ System.arraycopy(
+ serializedGroupValue,
+ 0,
+ allSerializedValuesConcatenated,
+ runningGroupsOffset,
+ serializedGroupValue.length);
+ runningGroupsOffset += serializedGroupValue.length;
+ old = serializedGroupValue;
+ }
+
+ ByteStreamStateHandle allSerializedStatesHandle = new ByteStreamStateHandle(
+ allSerializedValuesConcatenated);
+ KeyGroupsStateHandle keyGroupsStateHandle = new KeyGroupsStateHandle(
+ keyGroupRangeOffsets,
+ allSerializedStatesHandle);
+ List<KeyGroupsStateHandle> keyGroupsStateHandleList = new ArrayList<>();
+ keyGroupsStateHandleList.add(keyGroupsStateHandle);
+ return keyGroupsStateHandleList;
+ }
+
+ public static ChainedStateHandle<StreamStateHandle> generateStateForVertex(
+ JobVertexID jobVertexID,
+ int index) throws IOException {
+
+ Random random = new Random(jobVertexID.hashCode() + index);
+ int value = random.nextInt();
+ return generateChainedStateHandle(value);
+ }
+
+ public static ChainedStateHandle<StreamStateHandle> generateChainedStateHandle(
+ Serializable value) throws IOException {
+ return ChainedStateHandle.wrapSingleHandle(ByteStreamStateHandle.fromSerializable(value));
+ }
+
+ public static ExecutionJobVertex mockExecutionJobVertex(
+ JobVertexID jobVertexID,
+ int parallelism,
+ int maxParallelism) {
+ final ExecutionJobVertex executionJobVertex = mock(ExecutionJobVertex.class);
+
+ ExecutionVertex[] executionVertices = new ExecutionVertex[parallelism];
+
+ for (int i = 0; i < parallelism; i++) {
+ executionVertices[i] = mockExecutionVertex(
+ new ExecutionAttemptID(),
+ jobVertexID,
+ parallelism,
+ maxParallelism,
+ ExecutionState.RUNNING);
+
+ when(executionVertices[i].getParallelSubtaskIndex()).thenReturn(i);
+ }
+
+ when(executionJobVertex.getJobVertexId()).thenReturn(jobVertexID);
+ when(executionJobVertex.getTaskVertices()).thenReturn(executionVertices);
+ when(executionJobVertex.getParallelism()).thenReturn(parallelism);
+ when(executionJobVertex.getMaxParallelism()).thenReturn(maxParallelism);
+
+ return executionJobVertex;
+ }
+
+ private static ExecutionVertex mockExecutionVertex(ExecutionAttemptID attemptID) {
+ return mockExecutionVertex(
+ attemptID,
+ new JobVertexID(),
+ 1,
+ 1,
+ ExecutionState.RUNNING);
+ }
+
+ private static ExecutionVertex mockExecutionVertex(
+ ExecutionAttemptID attemptID,
+ JobVertexID jobVertexID,
+ int parallelism,
+ int maxParallelism,
+ ExecutionState state,
+ ExecutionState ... successiveStates) {
+
+ ExecutionVertex vertex = mock(ExecutionVertex.class);
+
+ final Execution exec = spy(new Execution(
+ mock(ExecutionContext.class),
+ vertex,
+ 1,
+ 1L,
+ null
+ ));
+ when(exec.getAttemptId()).thenReturn(attemptID);
+ when(exec.getState()).thenReturn(state, successiveStates);
+
+ when(vertex.getJobvertexId()).thenReturn(jobVertexID);
+ when(vertex.getCurrentExecutionAttempt()).thenReturn(exec);
+ when(vertex.getTotalNumberOfParallelSubtasks()).thenReturn(parallelism);
+ when(vertex.getMaxParallelism()).thenReturn(maxParallelism);
+
+ return vertex;
+ }
+
+ public static void verifiyStateRestore(
+ JobVertexID jobVertexID, ExecutionJobVertex executionJobVertex,
+ List<KeyGroupRange> keyGroupPartitions) throws Exception {
+
+ for (int i = 0; i < executionJobVertex.getParallelism(); i++) {
+
+ ChainedStateHandle<StreamStateHandle> expectNonPartitionedState = generateStateForVertex(jobVertexID, i);
+ ChainedStateHandle<StreamStateHandle> actualNonPartitionedState = executionJobVertex.
+ getTaskVertices()[i].getCurrentExecutionAttempt().getChainedStateHandle();
+ assertEquals(expectNonPartitionedState.get(0), actualNonPartitionedState.get(0));
+
+ List<KeyGroupsStateHandle> expectPartitionedKeyGroupState = generateKeyGroupState(
+ jobVertexID,
+ keyGroupPartitions.get(i));
+ List<KeyGroupsStateHandle> actualPartitionedKeyGroupState = executionJobVertex.
+ getTaskVertices()[i].getCurrentExecutionAttempt().getKeyGroupsStateHandles();
+ comparePartitionedState(expectPartitionedKeyGroupState, actualPartitionedKeyGroupState);
+ }
+ }
+
+ public static void comparePartitionedState(
+ List<KeyGroupsStateHandle> expectPartitionedKeyGroupState,
+ List<KeyGroupsStateHandle> actualPartitionedKeyGroupState) throws Exception {
+
+ KeyGroupsStateHandle expectedHeadOpKeyGroupStateHandle = expectPartitionedKeyGroupState.get(0);
+ int expectedTotalKeyGroups = expectedHeadOpKeyGroupStateHandle.getNumberOfKeyGroups();
+ int actualTotalKeyGroups = 0;
+ for(KeyGroupsStateHandle keyGroupsStateHandle: actualPartitionedKeyGroupState) {
+ actualTotalKeyGroups += keyGroupsStateHandle.getNumberOfKeyGroups();
+ }
+
+ assertEquals(expectedTotalKeyGroups, actualTotalKeyGroups);
+
+ FSDataInputStream inputStream = expectedHeadOpKeyGroupStateHandle.getStateHandle().openInputStream();
+ for(int groupId : expectedHeadOpKeyGroupStateHandle.keyGroups()) {
+ long offset = expectedHeadOpKeyGroupStateHandle.getOffsetForKeyGroup(groupId);
+ inputStream.seek(offset);
+ int expectedKeyGroupState = InstantiationUtil.deserializeObject(inputStream);
+ for(KeyGroupsStateHandle oneActualKeyGroupStateHandle : actualPartitionedKeyGroupState) {
+ if (oneActualKeyGroupStateHandle.containsKeyGroup(groupId)) {
+ long actualOffset = oneActualKeyGroupStateHandle.getOffsetForKeyGroup(groupId);
+ FSDataInputStream actualInputStream = oneActualKeyGroupStateHandle.getStateHandle().openInputStream();
+ actualInputStream.seek(actualOffset);
+ int actualGroupState = InstantiationUtil.deserializeObject(actualInputStream);
+
+ assertEquals(expectedKeyGroupState, actualGroupState);
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testCreateKeyGroupPartitions() {
+ testCreateKeyGroupPartitions(1, 1);
+ testCreateKeyGroupPartitions(13, 1);
+ testCreateKeyGroupPartitions(13, 2);
+ testCreateKeyGroupPartitions(Short.MAX_VALUE, 1);
+ testCreateKeyGroupPartitions(Short.MAX_VALUE, 13);
+ testCreateKeyGroupPartitions(Short.MAX_VALUE, Short.MAX_VALUE);
+
+ Random r = new Random(1234);
+ for (int k = 0; k < 1000; ++k) {
+ int maxParallelism = 1 + r.nextInt(Short.MAX_VALUE - 1);
+ int parallelism = 1 + r.nextInt(maxParallelism);
+ testCreateKeyGroupPartitions(maxParallelism, parallelism);
+ }
+ }
+
+ private void testCreateKeyGroupPartitions(int maxParallelism, int parallelism) {
+ List<KeyGroupRange> ranges = CheckpointCoordinator.createKeyGroupPartitions(maxParallelism, parallelism);
+ for (int i = 0; i < maxParallelism; ++i) {
+ KeyGroupRange range = ranges.get(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, i));
+ if (!range.contains(i)) {
+ Assert.fail("Could not find expected key-group " + i + " in range " + range);
+ }
+ }
+ }
+
}
http://git-wip-us.apache.org/repos/asf/flink/blob/516ad011/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java
index 5fcc59e..f757943 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java
@@ -18,26 +18,54 @@
package org.apache.flink.streaming.runtime.tasks;
+import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.runtime.io.network.api.CheckpointBarrier;
import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
+import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.streaming.api.graph.StreamConfig;
+import org.apache.flink.streaming.api.graph.StreamEdge;
+import org.apache.flink.streaming.api.graph.StreamNode;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.StreamMap;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.util.TestHarnessUtil;
+import org.apache.flink.util.InstantiationUtil;
+import org.apache.flink.util.Preconditions;
+import org.apache.flink.util.TestLogger;
import org.junit.Assert;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.powermock.core.classloader.annotations.PowerMockIgnore;
import org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;
+import scala.concurrent.duration.Deadline;
+import scala.concurrent.duration.FiniteDuration;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
+import java.util.Random;
import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.TimeUnit;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
/**
* Tests for {@link OneInputStreamTask}.
@@ -51,7 +79,7 @@ import java.util.concurrent.ConcurrentLinkedQueue;
@RunWith(PowerMockRunner.class)
@PrepareForTest({ResultPartitionWriter.class})
@PowerMockIgnore({"javax.management.*", "com.sun.jndi.*"})
-public class OneInputStreamTaskTest {
+public class OneInputStreamTaskTest extends TestLogger {
/**
* This test verifies that open() and close() are correctly called. This test also verifies
@@ -82,7 +110,7 @@ public class OneInputStreamTaskTest {
testHarness.waitForTaskCompletion();
- Assert.assertTrue("RichFunction methods where not called.", TestOpenCloseMapFunction.closeCalled);
+ assertTrue("RichFunction methods where not called.", TestOpenCloseMapFunction.closeCalled);
TestHarnessUtil.assertOutputEquals("Output was not correct.",
expectedOutput,
@@ -165,7 +193,7 @@ public class OneInputStreamTaskTest {
testHarness.waitForTaskCompletion();
List<String> resultElements = TestHarnessUtil.getRawElementsFromOutput(testHarness.getOutput());
- Assert.assertEquals(2, resultElements.size());
+ assertEquals(2, resultElements.size());
}
/**
@@ -293,6 +321,252 @@ public class OneInputStreamTaskTest {
TestHarnessUtil.assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
}
+ /**
+ * Tests that the stream operator can snapshot and restore the operator state of chained
+ * operators
+ */
+ @Test
+ public void testSnapshottingAndRestoring() throws Exception {
+ final Deadline deadline = new FiniteDuration(2, TimeUnit.MINUTES).fromNow();
+ final OneInputStreamTask<String, String> streamTask = new OneInputStreamTask<String, String>();
+ final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<String, String>(streamTask, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
+ IdentityKeySelector<String> keySelector = new IdentityKeySelector<>();
+ testHarness.configureForKeyedStream(keySelector, BasicTypeInfo.STRING_TYPE_INFO);
+
+ long checkpointId = 1L;
+ long checkpointTimestamp = 1L;
+ long recoveryTimestamp = 3L;
+ long seed = 2L;
+ int numberChainedTasks = 11;
+
+ StreamConfig streamConfig = testHarness.getStreamConfig();
+
+ configureChainedTestingStreamOperator(streamConfig, numberChainedTasks, seed, recoveryTimestamp);
+
+ AcknowledgeStreamMockEnvironment env = new AcknowledgeStreamMockEnvironment(
+ testHarness.jobConfig,
+ testHarness.taskConfig,
+ testHarness.executionConfig,
+ testHarness.memorySize,
+ new MockInputSplitProvider(),
+ testHarness.bufferSize);
+
+ // reset number of restore calls
+ TestingStreamOperator.numberRestoreCalls = 0;
+
+ testHarness.invoke(env);
+ testHarness.waitForTaskRunning(deadline.timeLeft().toMillis());
+
+ streamTask.triggerCheckpoint(checkpointId, checkpointTimestamp);
+
+ testHarness.endInput();
+ testHarness.waitForTaskCompletion(deadline.timeLeft().toMillis());
+
+ // since no state was set, there shouldn't be restore calls
+ assertEquals(0, TestingStreamOperator.numberRestoreCalls);
+
+ assertEquals(checkpointId, env.getCheckpointId());
+
+ final OneInputStreamTask<String, String> restoredTask = new OneInputStreamTask<String, String>();
+ restoredTask.setInitialState(env.getState(), env.getKeyGroupStates());
+
+ final OneInputStreamTaskTestHarness<String, String> restoredTaskHarness = new OneInputStreamTaskTestHarness<String, String>(restoredTask, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
+ restoredTaskHarness.configureForKeyedStream(keySelector, BasicTypeInfo.STRING_TYPE_INFO);
+
+ StreamConfig restoredTaskStreamConfig = restoredTaskHarness.getStreamConfig();
+
+ configureChainedTestingStreamOperator(restoredTaskStreamConfig, numberChainedTasks, seed, recoveryTimestamp);
+
+ TestingStreamOperator.numberRestoreCalls = 0;
+
+ restoredTaskHarness.invoke();
+ restoredTaskHarness.endInput();
+ restoredTaskHarness.waitForTaskCompletion(deadline.timeLeft().toMillis());
+
+ // restore of every chained operator should have been called
+ assertEquals(numberChainedTasks, TestingStreamOperator.numberRestoreCalls);
+
+ TestingStreamOperator.numberRestoreCalls = 0;
+ }
+
+ //==============================================================================================
+ // Utility functions and classes
+ //==============================================================================================
+
+ private void configureChainedTestingStreamOperator(
+ StreamConfig streamConfig,
+ int numberChainedTasks,
+ long seed,
+ long recoveryTimestamp) {
+
+ Preconditions.checkArgument(numberChainedTasks >= 1, "The operator chain must at least " +
+ "contain one operator.");
+
+ Random random = new Random(seed);
+
+ TestingStreamOperator<Integer, Integer> previousOperator = new TestingStreamOperator<>(random.nextLong(), recoveryTimestamp);
+ streamConfig.setStreamOperator(previousOperator);
+
+ // create the chain of operators
+ Map<Integer, StreamConfig> chainedTaskConfigs = new HashMap<>(numberChainedTasks - 1);
+ List<StreamEdge> outputEdges = new ArrayList<>(numberChainedTasks - 1);
+
+ for (int chainedIndex = 1; chainedIndex < numberChainedTasks; chainedIndex++) {
+ TestingStreamOperator<Integer, Integer> chainedOperator = new TestingStreamOperator<>(random.nextLong(), recoveryTimestamp);
+ StreamConfig chainedConfig = new StreamConfig(new Configuration());
+ chainedConfig.setStreamOperator(chainedOperator);
+ chainedTaskConfigs.put(chainedIndex, chainedConfig);
+
+ StreamEdge outputEdge = new StreamEdge(
+ new StreamNode(
+ null,
+ chainedIndex - 1,
+ null,
+ null,
+ null,
+ null,
+ null
+ ),
+ new StreamNode(
+ null,
+ chainedIndex,
+ null,
+ null,
+ null,
+ null,
+ null
+ ),
+ 0,
+ Collections.<String>emptyList(),
+ null
+ );
+
+ outputEdges.add(outputEdge);
+ }
+
+ streamConfig.setChainedOutputs(outputEdges);
+ streamConfig.setTransitiveChainedTaskConfigs(chainedTaskConfigs);
+ }
+
+ private static class IdentityKeySelector<IN> implements KeySelector<IN, IN> {
+
+ private static final long serialVersionUID = -3555913664416688425L;
+
+ @Override
+ public IN getKey(IN value) throws Exception {
+ return value;
+ }
+ }
+
+ private static class AcknowledgeStreamMockEnvironment extends StreamMockEnvironment {
+ private long checkpointId;
+ private ChainedStateHandle<StreamStateHandle> state;
+ private List<KeyGroupsStateHandle> keyGroupStates;
+
+ public long getCheckpointId() {
+ return checkpointId;
+ }
+
+ public ChainedStateHandle<StreamStateHandle> getState() {
+ return state;
+ }
+
+ List<KeyGroupsStateHandle> getKeyGroupStates() {
+ List<KeyGroupsStateHandle> result = new ArrayList<>();
+ for (int i = 0; i < keyGroupStates.size(); i++) {
+ if (keyGroupStates.get(i) != null) {
+ result.add(keyGroupStates.get(i));
+ }
+ }
+ return result;
+ }
+
+ AcknowledgeStreamMockEnvironment(Configuration jobConfig, Configuration taskConfig,
+ ExecutionConfig executionConfig, long memorySize,
+ MockInputSplitProvider inputSplitProvider, int bufferSize) {
+ super(jobConfig, taskConfig, executionConfig, memorySize, inputSplitProvider, bufferSize);
+ }
+
+
+ @Override
+ public void acknowledgeCheckpoint(long checkpointId, ChainedStateHandle<StreamStateHandle> state,
+ List<KeyGroupsStateHandle> keyGroupStates) {
+ this.checkpointId = checkpointId;
+ this.state = state;
+ this.keyGroupStates = keyGroupStates;
+ }
+ }
+
+ private static class TestingStreamOperator<IN, OUT>
+ extends AbstractStreamOperator<OUT> implements OneInputStreamOperator<IN, OUT> {
+
+ private static final long serialVersionUID = 774614855940397174L;
+
+ public static int numberRestoreCalls = 0;
+
+ private final long seed;
+ private final long recoveryTimestamp;
+
+ private transient Random random;
+
+ TestingStreamOperator(long seed, long recoveryTimestamp) {
+ this.seed = seed;
+ this.recoveryTimestamp = recoveryTimestamp;
+ }
+
+ @Override
+ public void processElement(StreamRecord<IN> element) throws Exception {
+
+ }
+
+ @Override
+ public void processWatermark(Watermark mark) throws Exception {
+
+ }
+
+ @Override
+ public void snapshotState(FSDataOutputStream out, long checkpointId, long timestamp) throws Exception {
+ if (random == null) {
+ random = new Random(seed);
+ }
+
+ Serializable functionState = generateFunctionState();
+ Integer operatorState = generateOperatorState();
+
+ InstantiationUtil.serializeObject(out, functionState);
+ InstantiationUtil.serializeObject(out, operatorState);
+ }
+
+ @Override
+ public void restoreState(FSDataInputStream in) throws Exception {
+ numberRestoreCalls++;
+
+ if (random == null) {
+ random = new Random(seed);
+ }
+
+ assertEquals(this.recoveryTimestamp, recoveryTimestamp);
+
+ assertNotNull(in);
+
+ Serializable functionState= InstantiationUtil.deserializeObject(in);
+ Integer operatorState= InstantiationUtil.deserializeObject(in);
+
+ assertEquals(random.nextInt(), functionState);
+ assertEquals(random.nextInt(), (int) operatorState);
+ }
+
+
+ private Serializable generateFunctionState() {
+ return random.nextInt();
+ }
+
+ private Integer generateOperatorState() {
+ return random.nextInt();
+ }
+ }
+
+
// This must only be used in one test, otherwise the static fields will be changed
// by several tests concurrently
private static class TestOpenCloseMapFunction extends RichMapFunction<String, String> {
http://git-wip-us.apache.org/repos/asf/flink/blob/516ad011/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
index 8d1baeb..39f3086 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
@@ -352,7 +352,6 @@ public class RescalingITCase extends TestLogger {
for (int key = 0; key < numberKeys; key++) {
int keyGroupIndex = keyGroupAssigner.getKeyGroupIndex(key);
-// expectedResult.add(Tuple2.of(keyGroupIndex % parallelism, numberElements * key));
expectedResult.add(Tuple2.of(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, keyGroupIndex) , numberElements * key));
}
[13/27] flink git commit: [FLINK-3761] Refactor State Backends/Make
Keyed State Key-Group Aware
Posted by al...@apache.org.
[FLINK-3761] Refactor State Backends/Make Keyed State Key-Group Aware
The biggest change in this is that functionality that used to be in
AbstractStateBackend is now moved to CheckpointStreamFactory and
KeyedStateBackend. The former is responsible for providing streams that
can be used to checkpoint data while the latter is responsible for
keeping keyed state. A keyed backend can checkpoint the state that it
keeps by using a CheckpointStreamFactory.
This also refactors how asynchronous keyed state snapshots work. They
are not implemented using a Future/RunnableFuture.
Also, this changes the keyed state backends to be key-group aware and to
snapshot the state in key-groups with an index for restoring.
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/4809f536
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/4809f536
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/4809f536
Branch: refs/heads/master
Commit: 4809f5367b08a9734fc1bd4875be51a9f3bb65aa
Parents: 516ad01
Author: Aljoscha Krettek <al...@gmail.com>
Authored: Wed Aug 10 18:44:50 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:10:01 2016 +0200
----------------------------------------------------------------------
.../streaming/state/AbstractRocksDBState.java | 30 +-
.../streaming/state/RocksDBFoldingState.java | 2 +-
.../state/RocksDBKeyedStateBackend.java | 251 ++++++
.../streaming/state/RocksDBListState.java | 4 +-
.../streaming/state/RocksDBReducingState.java | 4 +-
.../streaming/state/RocksDBStateBackend.java | 884 +++----------------
.../streaming/state/RocksDBValueState.java | 2 +-
.../FullyAsyncRocksDBStateBackendTest.java | 65 --
.../state/RocksDBAsyncKVSnapshotTest.java | 4 +-
.../state/RocksDBStateBackendConfigTest.java | 644 +++++++-------
.../state/RocksDBStateBackendTest.java | 29 +-
.../storm/tests/StormFieldsGroupingITCase.java | 12 +-
.../flink/storm/wrappers/BoltWrapper.java | 2 +-
.../flink/storm/wrappers/BoltWrapperTest.java | 4 +-
.../org/apache/flink/api/common/TaskInfo.java | 13 +-
.../common/operators/CollectionExecutor.java | 8 +-
.../flink/core/fs/FSDataOutputStream.java | 2 +
.../core/fs/local/LocalDataOutputStream.java | 5 +
.../memory/ByteArrayOutputStreamWithPos.java | 281 ++++++
.../functions/util/RuntimeUDFContextTest.java | 2 +-
.../api/common/io/RichInputFormatTest.java | 2 +-
.../api/common/io/RichOutputFormatTest.java | 2 +-
.../operators/GenericDataSinkBaseTest.java | 2 +-
.../operators/GenericDataSourceBaseTest.java | 2 +-
.../base/FlatMapOperatorCollectionTest.java | 2 +-
.../base/InnerJoinOperatorBaseTest.java | 2 +-
.../common/operators/base/MapOperatorTest.java | 2 +-
.../base/PartitionMapOperatorTest.java | 2 +-
.../flink/hdfstests/FileStateBackendTest.java | 117 +--
.../base/CoGroupOperatorCollectionTest.java | 2 +-
.../operators/base/GroupReduceOperatorTest.java | 2 +-
.../base/InnerJoinOperatorBaseTest.java | 2 +-
.../operators/base/ReduceOperatorTest.java | 2 +-
.../main/java/org/apache/flink/cep/nfa/NFA.java | 2 +-
.../AbstractKeyedCEPPatternOperator.java | 2 +
.../flink/cep/operator/CEPOperatorTest.java | 62 +-
.../checkpoint/CheckpointCoordinator.java | 32 +-
.../deployment/TaskDeploymentDescriptor.java | 16 +-
.../runtime/executiongraph/ExecutionVertex.java | 1 +
.../runtime/fs/hdfs/HadoopDataOutputStream.java | 5 +
.../flink/runtime/query/KvStateRegistry.java | 6 +-
.../runtime/query/TaskKvStateRegistry.java | 2 +-
.../query/netty/KvStateServerHandler.java | 8 +-
.../flink/runtime/state/AbstractHeapState.java | 220 -----
.../runtime/state/AbstractStateBackend.java | 422 +--------
.../state/AsynchronousKvStateSnapshot.java | 68 --
.../runtime/state/CheckpointStreamFactory.java | 67 ++
.../apache/flink/runtime/state/DoneFuture.java | 70 ++
.../runtime/state/GenericFoldingState.java | 74 +-
.../flink/runtime/state/GenericListState.java | 69 +-
.../runtime/state/GenericReducingState.java | 74 +-
.../flink/runtime/state/KeyGroupRange.java | 3 +
.../flink/runtime/state/KeyedStateBackend.java | 340 +++++++
.../org/apache/flink/runtime/state/KvState.java | 39 +-
.../flink/runtime/state/KvStateSnapshot.java | 61 --
.../state/RetrievableStreamStateHandle.java | 2 +-
.../apache/flink/runtime/state/StateObject.java | 5 +-
.../flink/runtime/state/StreamStateHandle.java | 4 +-
.../state/filesystem/AbstractFsState.java | 95 --
.../filesystem/AbstractFsStateSnapshot.java | 139 ---
.../state/filesystem/FileStateHandle.java | 2 +-
.../filesystem/FsCheckpointStreamFactory.java | 313 +++++++
.../state/filesystem/FsFoldingState.java | 161 ----
.../runtime/state/filesystem/FsListState.java | 149 ----
.../state/filesystem/FsReducingState.java | 165 ----
.../state/filesystem/FsStateBackend.java | 371 +-------
.../runtime/state/filesystem/FsValueState.java | 148 ----
.../runtime/state/heap/AbstractHeapState.java | 187 ++++
.../runtime/state/heap/HeapFoldingState.java | 124 +++
.../state/heap/HeapKeyedStateBackend.java | 328 +++++++
.../flink/runtime/state/heap/HeapListState.java | 156 ++++
.../runtime/state/heap/HeapReducingState.java | 123 +++
.../runtime/state/heap/HeapValueState.java | 112 +++
.../flink/runtime/state/heap/StateTable.java | 77 ++
.../runtime/state/memory/AbstractMemState.java | 82 --
.../state/memory/AbstractMemStateSnapshot.java | 144 ---
.../state/memory/ByteStreamStateHandle.java | 2 +-
.../memory/MemCheckpointStreamFactory.java | 146 +++
.../runtime/state/memory/MemFoldingState.java | 135 ---
.../runtime/state/memory/MemListState.java | 120 ---
.../runtime/state/memory/MemReducingState.java | 139 ---
.../runtime/state/memory/MemValueState.java | 122 ---
.../state/memory/MemoryStateBackend.java | 178 +---
.../savepoint/SavepointLoaderTest.java | 6 +-
.../TaskDeploymentDescriptorTest.java | 4 +-
.../messages/CheckpointMessagesTest.java | 2 +-
.../metrics/groups/TaskManagerGroupTest.java | 14 +-
.../operators/testutils/DummyEnvironment.java | 11 +-
.../operators/testutils/MockEnvironment.java | 22 +-
.../runtime/query/QueryableStateClientTest.java | 35 +-
.../runtime/query/netty/KvStateClientTest.java | 41 +-
.../query/netty/KvStateServerHandlerTest.java | 222 +++--
.../runtime/query/netty/KvStateServerTest.java | 51 +-
.../runtime/state/FileStateBackendTest.java | 110 +--
.../runtime/state/MemoryStateBackendTest.java | 19 +-
.../runtime/state/StateBackendTestBase.java | 650 ++++++++------
.../FsCheckpointStateOutputStreamTest.java | 27 +-
.../runtime/taskmanager/TaskAsyncCallTest.java | 2 +-
.../runtime/taskmanager/TaskManagerTest.java | 27 +-
.../flink/runtime/taskmanager/TaskTest.java | 2 +-
.../source/ContinuousFileReaderOperator.java | 2 +-
.../api/graph/StreamGraphGenerator.java | 1 -
.../api/operators/AbstractStreamOperator.java | 144 ++-
.../operators/AbstractUdfStreamOperator.java | 23 +-
.../streaming/api/operators/StreamOperator.java | 7 +-
.../operators/GenericWriteAheadSink.java | 11 +-
...ractAlignedProcessingTimeWindowOperator.java | 4 +-
.../operators/windowing/WindowOperator.java | 2 +-
.../streaming/runtime/tasks/StreamTask.java | 371 ++++----
.../api/operators/StreamGroupedFoldTest.java | 12 +-
.../api/operators/StreamGroupedReduceTest.java | 9 +-
.../operators/StreamingRuntimeContextTest.java | 30 +-
.../operators/StreamOperatorChainingTest.java | 15 -
...AlignedProcessingTimeWindowOperatorTest.java | 235 ++---
...AlignedProcessingTimeWindowOperatorTest.java | 279 +++---
.../windowing/EvictingWindowOperatorTest.java | 13 +-
.../operators/windowing/WindowOperatorTest.java | 159 ++--
.../tasks/InterruptSensitiveRestoreTest.java | 4 +-
.../runtime/tasks/OneInputStreamTaskTest.java | 22 +-
.../runtime/tasks/StreamMockEnvironment.java | 2 +-
.../streaming/runtime/tasks/StreamTaskTest.java | 2 +-
.../KeyedOneInputStreamOperatorTestHarness.java | 201 +++++
.../flink/streaming/util/MockContext.java | 4 +-
.../util/OneInputStreamOperatorTestHarness.java | 63 +-
.../flink/streaming/util/TestHarnessUtil.java | 5 +-
.../streaming/util/WindowingTestHarness.java | 3 +-
.../EventTimeWindowCheckpointingITCase.java | 13 +-
.../test/classloading/ClassLoaderITCase.java | 2 +-
.../classloading/jar/CustomKvStateProgram.java | 37 +-
.../state/StateHandleSerializationTest.java | 14 -
.../streaming/runtime/StateBackendITCase.java | 70 +-
131 files changed, 4955 insertions(+), 5810 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
index 3c4a209..710f506 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
@@ -23,7 +23,6 @@ import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.core.memory.DataOutputView;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
import org.apache.flink.util.Preconditions;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.RocksDBException;
@@ -46,7 +45,7 @@ import java.io.IOException;
* @param <SD> The type of {@link StateDescriptor}.
*/
public abstract class AbstractRocksDBState<K, N, S extends State, SD extends StateDescriptor<S, V>, V>
- implements KvState<K, N, S, SD, RocksDBStateBackend>, State {
+ implements KvState<N>, State {
private static final Logger LOG = LoggerFactory.getLogger(AbstractRocksDBState.class);
@@ -57,7 +56,7 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta
private N currentNamespace;
/** Backend that holds the actual RocksDB instance where we store state */
- protected RocksDBStateBackend backend;
+ protected RocksDBKeyedStateBackend backend;
/** The column family of this particular instance of state */
protected ColumnFamilyHandle columnFamily;
@@ -77,7 +76,7 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta
protected AbstractRocksDBState(ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
SD stateDesc,
- RocksDBStateBackend backend) {
+ RocksDBKeyedStateBackend backend) {
this.namespaceSerializer = namespaceSerializer;
this.backend = backend;
@@ -106,7 +105,7 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta
}
protected void writeKeyAndNamespace(DataOutputView out) throws IOException {
- backend.keySerializer().serialize(backend.currentKey(), out);
+ backend.getKeySerializer().serialize(backend.getCurrentKey(), out);
out.writeByte(42);
namespaceSerializer.serialize(currentNamespace, out);
}
@@ -117,27 +116,6 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta
}
@Override
- public void dispose() {
- // ignore because we don't hold any state ourselves
- }
-
- @Override
- public SD getStateDescriptor() {
- return stateDesc;
- }
-
- @Override
- public void setCurrentKey(K key) {
- // ignore because we don't hold any state ourselves
- }
-
- @Override
- public KvStateSnapshot<K, N, S, SD, RocksDBStateBackend> snapshot(long checkpointId,
- long timestamp) throws Exception {
- throw new RuntimeException("Should not be called. Backups happen in RocksDBStateBackend.");
- }
-
- @Override
@SuppressWarnings("unchecked")
public byte[] getSerializedValue(byte[] serializedKeyAndNamespace) throws Exception {
// Serialized key and namespace is expected to be of the same format
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java
index f1cf409..8c0799b 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBFoldingState.java
@@ -66,7 +66,7 @@ public class RocksDBFoldingState<K, N, T, ACC>
public RocksDBFoldingState(ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
FoldingStateDescriptor<T, ACC> stateDesc,
- RocksDBStateBackend backend) {
+ RocksDBKeyedStateBackend backend) {
super(columnFamily, namespaceSerializer, stateDesc, backend);
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
new file mode 100644
index 0000000..63f1fa2
--- /dev/null
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
@@ -0,0 +1,251 @@
+/*
+ * 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.contrib.streaming.state;
+
+import org.apache.commons.io.FileUtils;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.FoldingState;
+import org.apache.flink.api.common.state.FoldingStateDescriptor;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.state.ReducingState;
+import org.apache.flink.api.common.state.ReducingStateDescriptor;
+import org.apache.flink.api.common.state.StateDescriptor;
+import org.apache.flink.api.common.state.ValueState;
+import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.query.TaskKvStateRegistry;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeOffsets;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.util.SerializableObject;
+import org.rocksdb.ColumnFamilyDescriptor;
+import org.rocksdb.ColumnFamilyHandle;
+import org.rocksdb.ColumnFamilyOptions;
+import org.rocksdb.DBOptions;
+import org.rocksdb.RocksDB;
+import org.rocksdb.RocksDBException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Future;
+
+/**
+ * A {@link KeyedStateBackend} that stores its state in {@code RocksDB} and will serialize state to
+ * streams provided by a {@link org.apache.flink.runtime.state.CheckpointStreamFactory} upon
+ * checkpointing. This state backend can store very large state that exceeds memory and spills
+ * to disk.
+ */
+public class RocksDBKeyedStateBackend<K> extends KeyedStateBackend<K> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(RocksDBKeyedStateBackend.class);
+
+ /** Operator identifier that is used to uniqueify the RocksDB storage path. */
+ private final String operatorIdentifier;
+
+ /** JobID for uniquifying backup paths. */
+ private final JobID jobId;
+
+ /** The options from the options factory, cached */
+ private final ColumnFamilyOptions columnOptions;
+
+ /** Path where this configured instance stores its data directory */
+ private final File instanceBasePath;
+
+ /** Path where this configured instance stores its RocksDB data base */
+ private final File instanceRocksDBPath;
+
+ /**
+ * Our RocksDB data base, this is used by the actual subclasses of {@link AbstractRocksDBState}
+ * to store state. The different k/v states that we have don't each have their own RocksDB
+ * instance. They all write to this instance but to their own column family.
+ */
+ protected volatile RocksDB db;
+
+ /**
+ * Lock for protecting cleanup of the RocksDB db. We acquire this when doing asynchronous
+ * checkpoints and when disposing the db. Otherwise, the asynchronous snapshot might try
+ * iterating over a disposed db.
+ */
+ private final SerializableObject dbCleanupLock = new SerializableObject();
+
+ /**
+ * Information about the k/v states as we create them. This is used to retrieve the
+ * column family that is used for a state and also for sanity checks when restoring.
+ */
+ private Map<String, Tuple2<ColumnFamilyHandle, StateDescriptor>> kvStateInformation;
+
+ public RocksDBKeyedStateBackend(
+ JobID jobId,
+ String operatorIdentifier,
+ File instanceBasePath,
+ DBOptions dbOptions,
+ ColumnFamilyOptions columnFamilyOptions,
+ TaskKvStateRegistry kvStateRegistry,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange
+ ) throws Exception {
+
+ super(kvStateRegistry, keySerializer, keyGroupAssigner, keyGroupRange);
+
+ this.operatorIdentifier = operatorIdentifier;
+ this.jobId = jobId;
+ this.columnOptions = columnFamilyOptions;
+
+ this.instanceBasePath = instanceBasePath;
+ this.instanceRocksDBPath = new File(instanceBasePath, "db");
+
+ RocksDB.loadLibrary();
+
+ if (!instanceBasePath.exists()) {
+ if (!instanceBasePath.mkdirs()) {
+ throw new RuntimeException("Could not create RocksDB data directory.");
+ }
+ }
+
+ // clean it, this will remove the last part of the path but RocksDB will recreate it
+ try {
+ if (instanceRocksDBPath.exists()) {
+ LOG.warn("Deleting already existing db directory {}.", instanceRocksDBPath);
+ FileUtils.deleteDirectory(instanceRocksDBPath);
+ }
+ } catch (IOException e) {
+ throw new RuntimeException("Error cleaning RocksDB data directory.", e);
+ }
+
+ List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>(1);
+ // RocksDB seems to need this...
+ columnFamilyDescriptors.add(new ColumnFamilyDescriptor("default".getBytes()));
+ List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
+ try {
+ db = RocksDB.open(dbOptions, instanceRocksDBPath.getAbsolutePath(), columnFamilyDescriptors, columnFamilyHandles);
+ } catch (RocksDBException e) {
+ throw new RuntimeException("Error while opening RocksDB instance.", e);
+ }
+
+ kvStateInformation = new HashMap<>();
+ }
+
+ @Override
+ public void close() throws Exception {
+ super.close();
+
+ // we have to lock because we might have an asynchronous checkpoint going on
+ synchronized (dbCleanupLock) {
+ if (db != null) {
+ for (Tuple2<ColumnFamilyHandle, StateDescriptor> column : kvStateInformation.values()) {
+ column.f0.dispose();
+ }
+
+ db.dispose();
+ db = null;
+ }
+ }
+
+ FileUtils.deleteDirectory(instanceBasePath);
+ }
+
+ @Override
+ public Future<KeyGroupsStateHandle> snapshot(
+ long checkpointId,
+ long timestamp,
+ CheckpointStreamFactory streamFactory) throws Exception {
+ throw new RuntimeException("Not implemented.");
+ }
+
+ // ------------------------------------------------------------------------
+ // State factories
+ // ------------------------------------------------------------------------
+
+ /**
+ * 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.
+ */
+ protected ColumnFamilyHandle getColumnFamily(StateDescriptor descriptor) {
+
+ Tuple2<ColumnFamilyHandle, StateDescriptor> stateInfo = kvStateInformation.get(descriptor.getName());
+
+ if (stateInfo != null) {
+ if (!stateInfo.f1.equals(descriptor)) {
+ throw new RuntimeException("Trying to access state using wrong StateDescriptor, was " + stateInfo.f1 + " trying access with " + descriptor);
+ }
+ return stateInfo.f0;
+ }
+
+ ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(descriptor.getName().getBytes(), columnOptions);
+
+ try {
+ ColumnFamilyHandle columnFamily = db.createColumnFamily(columnDescriptor);
+ kvStateInformation.put(descriptor.getName(), new Tuple2<>(columnFamily, descriptor));
+ return columnFamily;
+ } catch (RocksDBException e) {
+ throw new RuntimeException("Error creating ColumnFamilyHandle.", e);
+ }
+ }
+
+ @Override
+ protected <N, T> ValueState<T> createValueState(TypeSerializer<N> namespaceSerializer,
+ ValueStateDescriptor<T> stateDesc) throws Exception {
+
+ ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc);
+
+ return new RocksDBValueState<>(columnFamily, namespaceSerializer, stateDesc, this);
+ }
+
+ @Override
+ protected <N, T> ListState<T> createListState(TypeSerializer<N> namespaceSerializer,
+ ListStateDescriptor<T> stateDesc) throws Exception {
+
+ ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc);
+
+ return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
+ }
+
+ @Override
+ protected <N, T> ReducingState<T> createReducingState(TypeSerializer<N> namespaceSerializer,
+ ReducingStateDescriptor<T> stateDesc) throws Exception {
+
+ ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc);
+
+ return new RocksDBReducingState<>(columnFamily, namespaceSerializer, stateDesc, this);
+ }
+
+ @Override
+ protected <N, T, ACC> FoldingState<T, ACC> createFoldingState(TypeSerializer<N> namespaceSerializer,
+ FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
+
+ ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc);
+
+ return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java
index ff1038e..d8f937b 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBListState.java
@@ -67,8 +67,8 @@ public class RocksDBListState<K, N, V>
public RocksDBListState(ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
ListStateDescriptor<V> stateDesc,
- RocksDBStateBackend backend) {
-
+ RocksDBKeyedStateBackend backend) {
+
super(columnFamily, namespaceSerializer, stateDesc, backend);
this.valueSerializer = stateDesc.getSerializer();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java
index efa2931..15ae493 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBReducingState.java
@@ -65,8 +65,8 @@ public class RocksDBReducingState<K, N, V>
public RocksDBReducingState(ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
ReducingStateDescriptor<V> stateDesc,
- RocksDBStateBackend backend) {
-
+ RocksDBKeyedStateBackend backend) {
+
super(columnFamily, namespaceSerializer, stateDesc, backend);
this.valueSerializer = stateDesc.getSerializer();
this.reduceFunction = stateDesc.getReduceFunction();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
index 74276c0..62b71d9 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
@@ -17,64 +17,31 @@
package org.apache.flink.contrib.streaming.state;
-import org.apache.commons.io.FileUtils;
import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.state.FoldingState;
-import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.state.ListState;
-import org.apache.flink.api.common.state.ListStateDescriptor;
-import org.apache.flink.api.common.state.ReducingState;
-import org.apache.flink.api.common.state.ReducingStateDescriptor;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.StateBackend;
-import org.apache.flink.api.common.state.StateDescriptor;
-import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.api.common.typeutils.base.array.BytePrimitiveArraySerializer;
-import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.api.java.typeutils.runtime.DataInputViewStream;
import org.apache.flink.core.fs.Path;
-import org.apache.flink.core.memory.DataInputView;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
-import org.apache.flink.core.memory.DataOutputView;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.execution.Environment;
-import org.apache.flink.runtime.fs.hdfs.HadoopFileSystem;
+import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.AsynchronousKvStateSnapshot;
-import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
-import org.apache.flink.runtime.util.SerializableObject;
-import org.apache.flink.streaming.util.HDFSCopyFromLocal;
-import org.apache.flink.streaming.util.HDFSCopyToLocal;
-import org.apache.hadoop.fs.FileSystem;
-import org.rocksdb.BackupEngine;
-import org.rocksdb.BackupableDBOptions;
-import org.rocksdb.ColumnFamilyDescriptor;
-import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.DBOptions;
-import org.rocksdb.Env;
-import org.rocksdb.ReadOptions;
-import org.rocksdb.RestoreOptions;
-import org.rocksdb.RocksDB;
-import org.rocksdb.RocksDBException;
-import org.rocksdb.RocksIterator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import java.io.EOFException;
import java.io.File;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.net.URI;
import java.util.ArrayList;
-import java.util.HashMap;
import java.util.List;
-import java.util.Map;
import java.util.Random;
import java.util.UUID;
@@ -83,12 +50,12 @@ import static java.util.Objects.requireNonNull;
/**
* A {@link StateBackend} that stores its state in {@code RocksDB}. This state backend can
* store very large state that exceeds memory and spills to disk.
- *
+ *
* <p>All key/value state (including windows) is stored in the key/value index of RocksDB.
* For persistence against loss of machines, checkpoints take a snapshot of the
* RocksDB database, and persist that snapshot in a file system (by default) or
* another configurable state backend.
- *
+ *
* <p>The behavior of the RocksDB instances can be parametrized by setting RocksDB Options
* using the methods {@link #setPredefinedOptions(PredefinedOptions)} and
* {@link #setOptions(OptionsFactory)}.
@@ -101,15 +68,9 @@ public class RocksDBStateBackend extends AbstractStateBackend {
// ------------------------------------------------------------------------
// Static configuration values
// ------------------------------------------------------------------------
-
- /** The checkpoint directory that we copy the RocksDB backups to. */
- private final Path checkpointDirectory;
-
- /** The state backend that stores the non-partitioned state */
- private final AbstractStateBackend nonPartitionedStateBackend;
- /** Whether we do snapshots fully asynchronous */
- private boolean fullyAsyncBackup = false;
+ /** The state backend that we use for creating checkpoint streams. */
+ private final AbstractStateBackend checkpointStreamBackend;
/** Operator identifier that is used to uniqueify the RocksDB storage path. */
private String operatorIdentifier;
@@ -118,66 +79,35 @@ public class RocksDBStateBackend extends AbstractStateBackend {
private JobID jobId;
// DB storage directories
-
+
/** Base paths for RocksDB directory, as configured. May be null. */
private Path[] configuredDbBasePaths;
/** Base paths for RocksDB directory, as initialized */
private File[] initializedDbBasePaths;
-
+
private int nextDirectory;
-
+
// RocksDB options
-
+
/** The pre-configured option settings */
private PredefinedOptions predefinedOptions = PredefinedOptions.DEFAULT;
-
+
/** The options factory to create the RocksDB options in the cluster */
private OptionsFactory optionsFactory;
-
+
/** The options from the options factory, cached */
private transient DBOptions dbOptions;
private transient ColumnFamilyOptions columnOptions;
- // ------------------------------------------------------------------------
- // Per operator values that are set in initializerForJob
- // ------------------------------------------------------------------------
-
- /** Path where this configured instance stores its data directory */
- private transient File instanceBasePath;
-
- /** Path where this configured instance stores its RocksDB data base */
- private transient File instanceRocksDBPath;
-
- /** Base path where this configured instance stores checkpoints */
- private transient String instanceCheckpointPath;
-
- /**
- * Our RocksDB data base, this is used by the actual subclasses of {@link AbstractRocksDBState}
- * to store state. The different k/v states that we have don't each have their own RocksDB
- * instance. They all write to this instance but to their own column family.
- */
- protected volatile transient RocksDB db;
-
- /**
- * Lock for protecting cleanup of the RocksDB db. We acquire this when doing asynchronous
- * checkpoints and when disposing the db. Otherwise, the asynchronous snapshot might try
- * iterating over a disposed db.
- */
- private final SerializableObject dbCleanupLock = new SerializableObject();
+ /** Whether we already lazily initialized our local storage directories. */
+ private transient boolean isInitialized = false;
- /**
- * Information about the k/v states as we create them. This is used to retrieve the
- * column family that is used for a state and also for sanity checks when restoring.
- */
- private Map<String, Tuple2<ColumnFamilyHandle, StateDescriptor>> kvStateInformation;
-
- // ------------------------------------------------------------------------
/**
* Creates a new {@code RocksDBStateBackend} that stores its checkpoint data in the
* file system and location defined by the given URI.
- *
+ *
* <p>A state backend that stores checkpoints in HDFS or S3 must specify the file system
* host and port in the URI, or have the Hadoop configuration that describes the file system
* (host / high-availability group / possibly credentials) either referenced from the Flink
@@ -205,38 +135,42 @@ public class RocksDBStateBackend extends AbstractStateBackend {
public RocksDBStateBackend(URI checkpointDataUri) throws IOException {
// creating the FsStateBackend automatically sanity checks the URI
FsStateBackend fsStateBackend = new FsStateBackend(checkpointDataUri);
-
- this.nonPartitionedStateBackend = fsStateBackend;
- this.checkpointDirectory = fsStateBackend.getBasePath();
+
+ this.checkpointStreamBackend = fsStateBackend;
}
public RocksDBStateBackend(String checkpointDataUri, AbstractStateBackend nonPartitionedStateBackend) throws IOException {
this(new Path(checkpointDataUri).toUri(), nonPartitionedStateBackend);
}
-
- public RocksDBStateBackend(URI checkpointDataUri, AbstractStateBackend nonPartitionedStateBackend) throws IOException {
- this.nonPartitionedStateBackend = requireNonNull(nonPartitionedStateBackend);
- this.checkpointDirectory = FsStateBackend.validateAndNormalizeUri(checkpointDataUri);
+
+ public RocksDBStateBackend(URI checkpointDataUri, AbstractStateBackend checkpointStreamBackend) throws IOException {
+ this.checkpointStreamBackend = requireNonNull(checkpointStreamBackend);
+ }
+
+ private void writeObject(ObjectOutputStream oos) throws IOException {
+ oos.defaultWriteObject();
}
+ private void readObject(ObjectInputStream ois) throws IOException, ClassNotFoundException {
+ ois.defaultReadObject();
+ isInitialized = false;
+ }
// ------------------------------------------------------------------------
// State backend methods
// ------------------------------------------------------------------------
-
- @Override
- public void initializeForJob(
- Environment env,
- String operatorIdentifier,
- TypeSerializer<?> keySerializer) throws Exception {
-
- super.initializeForJob(env, operatorIdentifier, keySerializer);
- this.nonPartitionedStateBackend.initializeForJob(env, operatorIdentifier, keySerializer);
-
+ private void lazyInitializeForJob(
+ Environment env,
+ String operatorIdentifier) throws Exception {
+
+ if (isInitialized) {
+ return;
+ }
+
this.operatorIdentifier = operatorIdentifier.replace(" ", "");
this.jobId = env.getJobID();
-
+
// initialize the paths where the local RocksDB files should be stored
if (configuredDbBasePaths == null) {
// initialize from the temp directories
@@ -245,7 +179,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
else {
List<File> dirs = new ArrayList<>(configuredDbBasePaths.length);
String errorMessage = "";
-
+
for (Path path : configuredDbBasePaths) {
File f = new File(path.toUri().getPath());
File testDir = new File(f, UUID.randomUUID().toString());
@@ -259,672 +193,78 @@ public class RocksDBStateBackend extends AbstractStateBackend {
}
testDir.delete();
}
-
+
if (dirs.isEmpty()) {
throw new Exception("No local storage directories available. " + errorMessage);
} else {
initializedDbBasePaths = dirs.toArray(new File[dirs.size()]);
}
}
-
- nextDirectory = new Random().nextInt(initializedDbBasePaths.length);
-
- instanceBasePath = new File(getDbPath("dummy_state"), UUID.randomUUID().toString());
- instanceCheckpointPath = getCheckpointPath("dummy_state");
- instanceRocksDBPath = new File(instanceBasePath, "db");
-
- RocksDB.loadLibrary();
-
- if (!instanceBasePath.exists()) {
- if (!instanceBasePath.mkdirs()) {
- throw new RuntimeException("Could not create RocksDB data directory.");
- }
- }
-
- // clean it, this will remove the last part of the path but RocksDB will recreate it
- try {
- if (instanceRocksDBPath.exists()) {
- LOG.warn("Deleting already existing db directory {}.", instanceRocksDBPath);
- FileUtils.deleteDirectory(instanceRocksDBPath);
- }
- } catch (IOException e) {
- throw new RuntimeException("Error cleaning RocksDB data directory.", e);
- }
-
- List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>(1);
- // RocksDB seems to need this...
- columnFamilyDescriptors.add(new ColumnFamilyDescriptor("default".getBytes()));
- List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
- try {
- db = RocksDB.open(getDbOptions(), instanceRocksDBPath.getAbsolutePath(), columnFamilyDescriptors, columnFamilyHandles);
- } catch (RocksDBException e) {
- throw new RuntimeException("Error while opening RocksDB instance.", e);
- }
-
- kvStateInformation = new HashMap<>();
- }
-
- @Override
- public void disposeAllStateForCurrentJob() throws Exception {
- nonPartitionedStateBackend.disposeAllStateForCurrentJob();
- }
-
- @Override
- public void discardState() throws Exception {
- super.discardState();
- nonPartitionedStateBackend.discardState();
-
- // we have to lock because we might have an asynchronous checkpoint going on
- synchronized (dbCleanupLock) {
- if (db != null) {
- if (this.dbOptions != null) {
- this.dbOptions.dispose();
- this.dbOptions = null;
- }
-
- for (Tuple2<ColumnFamilyHandle, StateDescriptor> column : kvStateInformation.values()) {
- column.f0.dispose();
- }
-
- db.dispose();
- db = null;
- }
- }
- }
- @Override
- public void close() throws Exception {
- nonPartitionedStateBackend.close();
-
- // we have to lock because we might have an asynchronous checkpoint going on
- synchronized (dbCleanupLock) {
- if (db != null) {
- if (this.dbOptions != null) {
- this.dbOptions.dispose();
- this.dbOptions = null;
- }
-
- for (Tuple2<ColumnFamilyHandle, StateDescriptor> column : kvStateInformation.values()) {
- column.f0.dispose();
- }
-
- db.dispose();
- db = null;
- }
- }
- }
+ nextDirectory = new Random().nextInt(initializedDbBasePaths.length);
- private File getDbPath(String stateName) {
- return new File(new File(new File(getNextStoragePath(), jobId.toString()), operatorIdentifier), stateName);
+ isInitialized = true;
}
- private String getCheckpointPath(String stateName) {
- return checkpointDirectory + "/" + jobId.toString() + "/" + operatorIdentifier + "/" + stateName;
+ private File getDbPath() {
+ return new File(new File(getNextStoragePath(), jobId.toString()), operatorIdentifier);
}
private File getNextStoragePath() {
int ni = nextDirectory + 1;
ni = ni >= initializedDbBasePaths.length ? 0 : ni;
nextDirectory = ni;
-
- return initializedDbBasePaths[ni];
- }
- /**
- * Visible for tests.
- */
- public File[] getStoragePaths() {
- return initializedDbBasePaths;
+ return initializedDbBasePaths[ni];
}
- // ------------------------------------------------------------------------
- // Snapshot and restore
- // ------------------------------------------------------------------------
-
@Override
- public HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> snapshotPartitionedState(long checkpointId, long timestamp) throws Exception {
- if (keyValueStatesByName == null || keyValueStatesByName.size() == 0) {
+ public CheckpointStreamFactory createStreamFactory(JobID jobId,
+ String operatorIdentifier) throws IOException {
return null;
}
if (fullyAsyncBackup) {
return performFullyAsyncSnapshot(checkpointId, timestamp);
} else {
- return performSemiAsyncSnapshot(checkpointId, timestamp);
- }
- }
-
- /**
- * Performs a checkpoint by using the RocksDB backup feature to backup to a directory.
- * This backup is the asynchronously copied to the final checkpoint location.
- */
- private HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> performSemiAsyncSnapshot(long checkpointId, long timestamp) throws Exception {
- // We don't snapshot individual k/v states since everything is stored in a central
- // RocksDB data base. Create a dummy KvStateSnapshot that holds the information about
- // that checkpoint. We use the in injectKeyValueStateSnapshots to restore.
-
- final File localBackupPath = new File(instanceBasePath, "local-chk-" + checkpointId);
- final URI backupUri = new URI(instanceCheckpointPath + "/chk-" + checkpointId);
-
- if (!localBackupPath.exists()) {
- if (!localBackupPath.mkdirs()) {
- throw new RuntimeException("Could not create local backup path " + localBackupPath);
- }
- }
-
- long startTime = System.currentTimeMillis();
-
- BackupableDBOptions backupOptions = new BackupableDBOptions(localBackupPath.getAbsolutePath());
- // we disabled the WAL
- backupOptions.setBackupLogFiles(false);
- // no need to sync since we use the backup only as intermediate data before writing to FileSystem snapshot
- backupOptions.setSync(false);
-
- try (BackupEngine backupEngine = BackupEngine.open(Env.getDefault(), backupOptions)) {
- // wait before flush with "true"
- backupEngine.createNewBackup(db, true);
- }
-
- long endTime = System.currentTimeMillis();
- LOG.info("RocksDB (" + instanceRocksDBPath + ") backup (synchronous part) took " + (endTime - startTime) + " ms.");
-
- // draw a copy in case it get's changed while performing the async snapshot
- List<StateDescriptor> kvStateInformationCopy = new ArrayList<>();
- for (Tuple2<ColumnFamilyHandle, StateDescriptor> state: kvStateInformation.values()) {
- kvStateInformationCopy.add(state.f1);
- }
- SemiAsyncSnapshot dummySnapshot = new SemiAsyncSnapshot(localBackupPath,
- backupUri,
- kvStateInformationCopy,
- checkpointId);
-
-
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> result = new HashMap<>();
- result.put("dummy_state", dummySnapshot);
- return result;
- }
-
- /**
- * Performs a checkpoint by drawing a {@link org.rocksdb.Snapshot} from RocksDB and then
- * iterating over all key/value pairs in RocksDB to store them in the final checkpoint
- * location. The only synchronous part is the drawing of the {@code Snapshot} which
- * is essentially free.
- */
- private HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> performFullyAsyncSnapshot(long checkpointId, long timestamp) throws Exception {
- // we draw a snapshot from RocksDB then iterate over all keys at that point
- // and store them in the backup location
-
- final URI backupUri = new URI(instanceCheckpointPath + "/chk-" + checkpointId);
-
- long startTime = System.currentTimeMillis();
-
- org.rocksdb.Snapshot snapshot = db.getSnapshot();
-
- long endTime = System.currentTimeMillis();
- LOG.info("Fully asynchronous RocksDB (" + instanceRocksDBPath + ") backup (synchronous part) took " + (endTime - startTime) + " ms.");
-
- // draw a copy in case it get's changed while performing the async snapshot
- Map<String, Tuple2<ColumnFamilyHandle, StateDescriptor>> columnFamiliesCopy = new HashMap<>();
- columnFamiliesCopy.putAll(kvStateInformation);
- FullyAsyncSnapshot dummySnapshot = new FullyAsyncSnapshot(snapshot,
- this,
- backupUri,
- columnFamiliesCopy,
- checkpointId);
-
-
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> result = new HashMap<>();
- result.put("dummy_state", dummySnapshot);
- return result;
+ return checkpointStreamBackend.createStreamFactory(jobId, operatorIdentifier);
}
@Override
- public final void injectKeyValueStateSnapshots(HashMap<String, KvStateSnapshot> keyValueStateSnapshots) throws Exception {
- if (keyValueStateSnapshots == null) {
- return;
- }
-
- KvStateSnapshot dummyState = keyValueStateSnapshots.get("dummy_state");
- if (dummyState instanceof FinalSemiAsyncSnapshot) {
- restoreFromSemiAsyncSnapshot((FinalSemiAsyncSnapshot) dummyState);
- } else if (dummyState instanceof FinalFullyAsyncSnapshot) {
- restoreFromFullyAsyncSnapshot((FinalFullyAsyncSnapshot) dummyState);
- } else {
- throw new RuntimeException("Unknown RocksDB snapshot: " + dummyState);
- }
- }
-
- private void restoreFromSemiAsyncSnapshot(FinalSemiAsyncSnapshot snapshot) throws Exception {
- // This does mostly the same work as initializeForJob, we remove the existing RocksDB
- // directory and create a new one from the backup.
- // This must be refactored. The StateBackend should either be initialized from
- // scratch or from a snapshot.
-
- if (!instanceBasePath.exists()) {
- if (!instanceBasePath.mkdirs()) {
- throw new RuntimeException("Could not create RocksDB data directory.");
- }
- }
-
- db.dispose();
-
- // clean it, this will remove the last part of the path but RocksDB will recreate it
- try {
- if (instanceRocksDBPath.exists()) {
- LOG.warn("Deleting already existing db directory {}.", instanceRocksDBPath);
- FileUtils.deleteDirectory(instanceRocksDBPath);
- }
- } catch (IOException e) {
- throw new RuntimeException("Error cleaning RocksDB data directory.", e);
- }
-
- final File localBackupPath = new File(instanceBasePath, "chk-" + snapshot.checkpointId);
-
- if (localBackupPath.exists()) {
- try {
- LOG.warn("Deleting already existing local backup directory {}.", localBackupPath);
- FileUtils.deleteDirectory(localBackupPath);
- } catch (IOException e) {
- throw new RuntimeException("Error cleaning RocksDB local backup directory.", e);
- }
- }
-
- HDFSCopyToLocal.copyToLocal(snapshot.backupUri, instanceBasePath);
-
- try (BackupEngine backupEngine = BackupEngine.open(Env.getDefault(), new BackupableDBOptions(localBackupPath.getAbsolutePath()))) {
- backupEngine.restoreDbFromLatestBackup(instanceRocksDBPath.getAbsolutePath(), instanceRocksDBPath.getAbsolutePath(), new RestoreOptions(true));
- } catch (RocksDBException|IllegalArgumentException e) {
- throw new RuntimeException("Error while restoring RocksDB state from " + localBackupPath, e);
- } finally {
- try {
- FileUtils.deleteDirectory(localBackupPath);
- } catch (IOException e) {
- LOG.error("Error cleaning up local restore directory " + localBackupPath, e);
- }
- }
-
-
- List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>(snapshot.stateDescriptors.size());
- for (StateDescriptor stateDescriptor: snapshot.stateDescriptors) {
- columnFamilyDescriptors.add(new ColumnFamilyDescriptor(stateDescriptor.getName().getBytes(), getColumnOptions()));
- }
-
- // RocksDB seems to need this...
- columnFamilyDescriptors.add(new ColumnFamilyDescriptor("default".getBytes()));
- List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(snapshot.stateDescriptors.size());
- try {
-
- db = RocksDB.open(getDbOptions(), instanceRocksDBPath.getAbsolutePath(), columnFamilyDescriptors, columnFamilyHandles);
- this.kvStateInformation = new HashMap<>();
- for (int i = 0; i < snapshot.stateDescriptors.size(); i++) {
- this.kvStateInformation.put(snapshot.stateDescriptors.get(i).getName(), new Tuple2<>(columnFamilyHandles.get(i), snapshot.stateDescriptors.get(i)));
- }
-
- } catch (RocksDBException e) {
- throw new RuntimeException("Error while opening RocksDB instance.", e);
- }
- }
-
- private void restoreFromFullyAsyncSnapshot(FinalFullyAsyncSnapshot snapshot) throws Exception {
-
- DataInputView inputView = new DataInputViewStreamWrapper(snapshot.stateHandle.openInputStream());
-
- // clear k/v state information before filling it
- kvStateInformation.clear();
-
- // first get the column family mapping
- int numColumns = inputView.readInt();
- Map<Byte, StateDescriptor> columnFamilyMapping = new HashMap<>(numColumns);
- for (int i = 0; i < numColumns; i++) {
- byte mappingByte = inputView.readByte();
-
- ObjectInputStream ooIn = new ObjectInputStream(new DataInputViewStream(inputView));
- StateDescriptor stateDescriptor = (StateDescriptor) ooIn.readObject();
-
- columnFamilyMapping.put(mappingByte, stateDescriptor);
-
- // this will fill in the k/v state information
- getColumnFamily(stateDescriptor);
- }
-
- // try and read until EOF
- try {
- // the EOFException will get us out of this...
- while (true) {
- byte mappingByte = inputView.readByte();
- ColumnFamilyHandle handle = getColumnFamily(columnFamilyMapping.get(mappingByte));
- byte[] key = BytePrimitiveArraySerializer.INSTANCE.deserialize(inputView);
- byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(inputView);
- db.put(handle, key, value);
- }
- } catch (EOFException e) {
- // expected
- }
- }
-
- // ------------------------------------------------------------------------
- // Semi-asynchronous Backup Classes
- // ------------------------------------------------------------------------
-
- /**
- * Upon snapshotting the RocksDB backup is created synchronously. The asynchronous part is
- * copying the backup to a (possibly) remote filesystem. This is done in {@link #materialize()}.
- */
- private static class SemiAsyncSnapshot extends AsynchronousKvStateSnapshot<Object, Object, ValueState<Object>, ValueStateDescriptor<Object>, RocksDBStateBackend> {
- private static final long serialVersionUID = 1L;
- private final File localBackupPath;
- private final URI backupUri;
- private final List<StateDescriptor> stateDescriptors;
- private final long checkpointId;
-
- private SemiAsyncSnapshot(File localBackupPath,
- URI backupUri,
- List<StateDescriptor> columnFamilies,
- long checkpointId) {
- this.localBackupPath = localBackupPath;
- this.backupUri = backupUri;
- this.stateDescriptors = columnFamilies;
- this.checkpointId = checkpointId;
- }
-
- @Override
- public KvStateSnapshot<Object, Object, ValueState<Object>, ValueStateDescriptor<Object>, RocksDBStateBackend> materialize() throws Exception {
- try {
- long startTime = System.currentTimeMillis();
- HDFSCopyFromLocal.copyFromLocal(localBackupPath, backupUri);
- long endTime = System.currentTimeMillis();
- LOG.info("RocksDB materialization from " + localBackupPath + " to " + backupUri + " (asynchronous part) took " + (endTime - startTime) + " ms.");
- return new FinalSemiAsyncSnapshot(backupUri, checkpointId, stateDescriptors);
- } catch (Exception e) {
- FileSystem fs = FileSystem.get(backupUri, HadoopFileSystem.getHadoopConfiguration());
- fs.delete(new org.apache.hadoop.fs.Path(backupUri), true);
- throw e;
- } finally {
- FileUtils.deleteQuietly(localBackupPath);
- }
- }
- }
-
- /**
- * Dummy {@link KvStateSnapshot} that holds the state of our one RocksDB data base. This
- * also stores the column families that we had at the time of the snapshot so that we can
- * restore these. This results from {@link SemiAsyncSnapshot}.
- */
- private static class FinalSemiAsyncSnapshot implements KvStateSnapshot<Object, Object, ValueState<Object>, ValueStateDescriptor<Object>, RocksDBStateBackend> {
- private static final long serialVersionUID = 1L;
-
- final URI backupUri;
- final long checkpointId;
- private final List<StateDescriptor> stateDescriptors;
-
- /**
- * Creates a new snapshot from the given state parameters.
- */
- private FinalSemiAsyncSnapshot(URI backupUri, long checkpointId, List<StateDescriptor> stateDescriptors) {
- this.backupUri = backupUri;
- this.checkpointId = checkpointId;
- this.stateDescriptors = stateDescriptors;
- }
-
- @Override
- public final KvState<Object, Object, ValueState<Object>, ValueStateDescriptor<Object>, RocksDBStateBackend> restoreState(
- RocksDBStateBackend stateBackend,
- TypeSerializer<Object> keySerializer,
- ClassLoader classLoader) throws Exception {
- throw new RuntimeException("Should never happen.");
- }
-
- @Override
- public final void discardState() throws Exception {
- FileSystem fs = FileSystem.get(backupUri, HadoopFileSystem.getHadoopConfiguration());
- fs.delete(new org.apache.hadoop.fs.Path(backupUri), true);
- }
-
- @Override
- public final long getStateSize() throws Exception {
- FileSystem fs = FileSystem.get(backupUri, HadoopFileSystem.getHadoopConfiguration());
- return fs.getContentSummary(new org.apache.hadoop.fs.Path(backupUri)).getLength();
- }
-
- @Override
- public void close() throws IOException {
- // cannot do much here
- }
- }
-
- // ------------------------------------------------------------------------
- // Fully asynchronous Backup Classes
- // ------------------------------------------------------------------------
-
- /**
- * This does the snapshot using a RocksDB snapshot and an iterator over all keys
- * at the point of that snapshot.
- */
- private class FullyAsyncSnapshot extends AsynchronousKvStateSnapshot<Object, Object, ValueState<Object>, ValueStateDescriptor<Object>, RocksDBStateBackend> {
- private static final long serialVersionUID = 1L;
-
- private transient org.rocksdb.Snapshot snapshot;
- private transient AbstractStateBackend backend;
-
- private final URI backupUri;
- private final Map<String, Tuple2<ColumnFamilyHandle, StateDescriptor>> columnFamilies;
- private final long checkpointId;
-
- private FullyAsyncSnapshot(org.rocksdb.Snapshot snapshot,
- AbstractStateBackend backend,
- URI backupUri,
- Map<String, Tuple2<ColumnFamilyHandle, StateDescriptor>> columnFamilies,
- long checkpointId) {
- this.snapshot = snapshot;
- this.backend = backend;
- this.backupUri = backupUri;
- this.columnFamilies = columnFamilies;
- this.checkpointId = checkpointId;
- }
-
- @Override
- public KvStateSnapshot<Object, Object, ValueState<Object>, ValueStateDescriptor<Object>, RocksDBStateBackend> materialize() throws Exception {
- try {
- long startTime = System.currentTimeMillis();
-
- CheckpointStateOutputStream outputStream = backend.createCheckpointStateOutputStream(checkpointId, startTime);
- DataOutputView outputView = new DataOutputViewStreamWrapper(outputStream);
- outputView.writeInt(columnFamilies.size());
-
- // we don't know how many key/value pairs there are in each column family.
- // We prefix every written element with a byte that signifies to which
- // column family it belongs, this way we can restore the column families
- byte count = 0;
- Map<String, Byte> columnFamilyMapping = new HashMap<>();
- for (Map.Entry<String, Tuple2<ColumnFamilyHandle, StateDescriptor>> column: columnFamilies.entrySet()) {
- columnFamilyMapping.put(column.getKey(), count);
-
- outputView.writeByte(count);
-
- ObjectOutputStream ooOut = new ObjectOutputStream(outputStream);
- ooOut.writeObject(column.getValue().f1);
- ooOut.flush();
-
- count++;
- }
-
- ReadOptions readOptions = new ReadOptions();
- readOptions.setSnapshot(snapshot);
-
- for (Map.Entry<String, Tuple2<ColumnFamilyHandle, StateDescriptor>> column: columnFamilies.entrySet()) {
- byte columnByte = columnFamilyMapping.get(column.getKey());
-
- synchronized (dbCleanupLock) {
- if (db == null) {
- throw new RuntimeException("RocksDB instance was disposed. This happens " +
- "when we are in the middle of a checkpoint and the job fails.");
- }
- RocksIterator iterator = db.newIterator(column.getValue().f0, readOptions);
- iterator.seekToFirst();
- while (iterator.isValid()) {
- outputView.writeByte(columnByte);
- BytePrimitiveArraySerializer.INSTANCE.serialize(iterator.key(),
- outputView);
- BytePrimitiveArraySerializer.INSTANCE.serialize(iterator.value(),
- outputView);
- iterator.next();
- }
- }
- }
-
- StreamStateHandle stateHandle = outputStream.closeAndGetHandle();
-
- long endTime = System.currentTimeMillis();
- LOG.info("Fully asynchronous RocksDB materialization to " + backupUri + " (asynchronous part) took " + (endTime - startTime) + " ms.");
- return new FinalFullyAsyncSnapshot(stateHandle, checkpointId);
- } finally {
- synchronized (dbCleanupLock) {
- if (db != null) {
- db.releaseSnapshot(snapshot);
- }
- }
- snapshot = null;
- }
- }
-
- }
-
- /**
- * Dummy {@link KvStateSnapshot} that holds the state of our one RocksDB data base. This
- * results from {@link FullyAsyncSnapshot}.
- */
- private static class FinalFullyAsyncSnapshot implements KvStateSnapshot<Object, Object, ValueState<Object>, ValueStateDescriptor<Object>, RocksDBStateBackend> {
- private static final long serialVersionUID = 1L;
-
- final StreamStateHandle stateHandle;
- final long checkpointId;
-
- /**
- * Creates a new snapshot from the given state parameters.
- */
- private FinalFullyAsyncSnapshot(StreamStateHandle stateHandle, long checkpointId) {
- this.stateHandle = stateHandle;
- this.checkpointId = checkpointId;
- }
-
- @Override
- public final KvState<Object, Object, ValueState<Object>, ValueStateDescriptor<Object>, RocksDBStateBackend> restoreState(
- RocksDBStateBackend stateBackend,
- TypeSerializer<Object> keySerializer,
- ClassLoader classLoader) throws Exception {
- throw new RuntimeException("Should never happen.");
- }
-
- @Override
- public final void discardState() throws Exception {
- stateHandle.discardState();
- }
-
- @Override
- public final long getStateSize() throws Exception {
- return stateHandle.getStateSize();
- }
-
- @Override
- public void close() throws IOException {
- stateHandle.close();
- }
- }
-
- // ------------------------------------------------------------------------
- // State factories
- // ------------------------------------------------------------------------
-
- /**
- * 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.
- */
- protected ColumnFamilyHandle getColumnFamily(StateDescriptor descriptor) {
-
- Tuple2<ColumnFamilyHandle, StateDescriptor> stateInfo = kvStateInformation.get(descriptor.getName());
-
- if (stateInfo != null) {
- if (!stateInfo.f1.equals(descriptor)) {
- throw new RuntimeException("Trying to access state using wrong StateDescriptor, was " + stateInfo.f1 + " trying access with " + descriptor);
- }
- return stateInfo.f0;
- }
-
- ColumnFamilyDescriptor columnDescriptor = new ColumnFamilyDescriptor(descriptor.getName().getBytes(), getColumnOptions());
-
- try {
- ColumnFamilyHandle columnFamily = db.createColumnFamily(columnDescriptor);
- kvStateInformation.put(descriptor.getName(), new Tuple2<>(columnFamily, descriptor));
- return columnFamily;
- } catch (RocksDBException e) {
- throw new RuntimeException("Error creating ColumnFamilyHandle.", e);
- }
- }
-
- /**
- * Used by k/v states to access the current key.
- */
- public Object currentKey() {
- return currentKey;
- }
-
- /**
- * Used by k/v states to access the key serializer.
- */
- public TypeSerializer keySerializer() {
- return keySerializer;
- }
-
- @Override
- protected <N, T> ValueState<T> createValueState(TypeSerializer<N> namespaceSerializer,
- ValueStateDescriptor<T> stateDesc) throws Exception {
-
- ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc);
-
- return new RocksDBValueState<>(columnFamily, namespaceSerializer, stateDesc, this);
- }
-
- @Override
- protected <N, T> ListState<T> createListState(TypeSerializer<N> namespaceSerializer,
- ListStateDescriptor<T> stateDesc) throws Exception {
-
- ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc);
-
- return new RocksDBListState<>(columnFamily, namespaceSerializer, stateDesc, this);
- }
-
- @Override
- protected <N, T> ReducingState<T> createReducingState(TypeSerializer<N> namespaceSerializer,
- ReducingStateDescriptor<T> stateDesc) throws Exception {
-
- ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc);
-
- return new RocksDBReducingState<>(columnFamily, namespaceSerializer, stateDesc, this);
- }
+ public <K> KeyedStateBackend<K> createKeyedStateBackend(
+ Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ TaskKvStateRegistry kvStateRegistry) throws Exception {
- @Override
- protected <N, T, ACC> FoldingState<T, ACC> createFoldingState(TypeSerializer<N> namespaceSerializer,
- FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
+ lazyInitializeForJob(env, operatorIdentifier);
- ColumnFamilyHandle columnFamily = getColumnFamily(stateDesc);
+ File instanceBasePath = new File(getDbPath(), UUID.randomUUID().toString());
- return new RocksDBFoldingState<>(columnFamily, namespaceSerializer, stateDesc, this);
+ return new RocksDBKeyedStateBackend<>(
+ jobID,
+ operatorIdentifier,
+ instanceBasePath,
+ getDbOptions(),
+ getColumnOptions(),
+ kvStateRegistry,
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange);
}
- // ------------------------------------------------------------------------
- // Non-partitioned state
- // ------------------------------------------------------------------------
-
@Override
- public CheckpointStateOutputStream createCheckpointStateOutputStream(
- long checkpointID, long timestamp) throws Exception {
-
- return nonPartitionedStateBackend.createCheckpointStateOutputStream(checkpointID, timestamp);
+ public <K> KeyedStateBackend<K> restoreKeyedStateBackend(Environment env, JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ List<KeyGroupsStateHandle> restoredState,
+ TaskKvStateRegistry kvStateRegistry) throws Exception {
+ throw new RuntimeException("Not implemented.");
}
// ------------------------------------------------------------------------
@@ -932,35 +272,13 @@ public class RocksDBStateBackend extends AbstractStateBackend {
// ------------------------------------------------------------------------
/**
- * Enables fully asynchronous snapshotting of the partitioned state held in RocksDB.
- *
- * <p>By default, this is disabled. This means that RocksDB state is copied in a synchronous
- * step, during which normal processing of elements pauses, followed by an asynchronous step
- * of copying the RocksDB backup to the final checkpoint location. Fully asynchronous
- * snapshots take longer (linear time requirement with respect to number of unique keys)
- * but normal processing of elements is not paused.
- */
- public void enableFullyAsyncSnapshots() {
- this.fullyAsyncBackup = true;
- }
-
- /**
- * Disables fully asynchronous snapshotting of the partitioned state held in RocksDB.
- *
- * <p>By default, this is disabled.
- */
- public void disableFullyAsyncSnapshots() {
- this.fullyAsyncBackup = false;
- }
-
- /**
* Sets the path where the RocksDB local database files should be stored on the local
* file system. Setting this path overrides the default behavior, where the
* files are stored across the configured temp directories.
- *
+ *
* <p>Passing {@code null} to this function restores the default behavior, where the configured
* temp directories will be used.
- *
+ *
* @param path The path where the local RocksDB database files are stored.
*/
public void setDbStoragePath(String path) {
@@ -971,44 +289,44 @@ public class RocksDBStateBackend extends AbstractStateBackend {
* Sets the paths across which the local RocksDB database files are distributed on the local
* file system. Setting these paths overrides the default behavior, where the
* files are stored across the configured temp directories.
- *
+ *
* <p>Each distinct state will be stored in one path, but when the state backend creates
* multiple states, they will store their files on different paths.
- *
+ *
* <p>Passing {@code null} to this function restores the default behavior, where the configured
* temp directories will be used.
- *
- * @param paths The paths across which the local RocksDB database files will be spread.
+ *
+ * @param paths The paths across which the local RocksDB database files will be spread.
*/
public void setDbStoragePaths(String... paths) {
if (paths == null) {
configuredDbBasePaths = null;
- }
+ }
else if (paths.length == 0) {
throw new IllegalArgumentException("empty paths");
}
else {
Path[] pp = new Path[paths.length];
-
+
for (int i = 0; i < paths.length; i++) {
if (paths[i] == null) {
throw new IllegalArgumentException("null path");
}
-
+
pp[i] = new Path(paths[i]);
String scheme = pp[i].toUri().getScheme();
if (scheme != null && !scheme.equalsIgnoreCase("file")) {
throw new IllegalArgumentException("Path " + paths[i] + " has a non local scheme");
}
}
-
+
configuredDbBasePaths = pp;
}
}
/**
- *
- * @return The configured DB storage paths, or null, if none were configured.
+ *
+ * @return The configured DB storage paths, or null, if none were configured.
*/
public String[] getDbStoragePaths() {
if (configuredDbBasePaths == null) {
@@ -1021,18 +339,18 @@ public class RocksDBStateBackend extends AbstractStateBackend {
return paths;
}
}
-
+
// ------------------------------------------------------------------------
// Parametrize with RocksDB Options
// ------------------------------------------------------------------------
/**
* Sets the predefined options for RocksDB.
- *
+ *
* <p>If a user-defined options factory is set (via {@link #setOptions(OptionsFactory)}),
* then the options from the factory are applied on top of the here specified
* predefined options.
- *
+ *
* @param options The options to set (must not be null).
*/
public void setPredefinedOptions(PredefinedOptions options) {
@@ -1043,10 +361,10 @@ public class RocksDBStateBackend extends AbstractStateBackend {
* Gets the currently set predefined options for RocksDB.
* The default options (if nothing was set via {@link #setPredefinedOptions(PredefinedOptions)})
* are {@link PredefinedOptions#DEFAULT}.
- *
+ *
* <p>If a user-defined options factory is set (via {@link #setOptions(OptionsFactory)}),
* then the options from the factory are applied on top of the predefined options.
- *
+ *
* @return The currently set predefined options for RocksDB.
*/
public PredefinedOptions getPredefinedOptions() {
@@ -1057,13 +375,13 @@ public class RocksDBStateBackend extends AbstractStateBackend {
* Sets {@link org.rocksdb.Options} for the RocksDB instances.
* Because the options are not serializable and hold native code references,
* they must be specified through a factory.
- *
- * <p>The options created by the factory here are applied on top of the pre-defined
+ *
+ * <p>The options created by the factory here are applied on top of the pre-defined
* options profile selected via {@link #setPredefinedOptions(PredefinedOptions)}.
* If the pre-defined options profile is the default
* ({@link PredefinedOptions#DEFAULT}), then the factory fully controls the RocksDB
* options.
- *
+ *
* @param optionsFactory The options factory that lazily creates the RocksDB options.
*/
public void setOptions(OptionsFactory optionsFactory) {
@@ -1072,7 +390,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
/**
* Gets the options factory that lazily creates the RocksDB options.
- *
+ *
* @return The options factory.
*/
public OptionsFactory getOptions() {
@@ -1091,7 +409,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
if (optionsFactory != null) {
opt = optionsFactory.createDBOptions(opt);
}
-
+
// add necessary default options
opt = opt.setCreateIfMissing(true);
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java
index 62bc366..b9c0e83 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBValueState.java
@@ -63,7 +63,7 @@ public class RocksDBValueState<K, N, V>
public RocksDBValueState(ColumnFamilyHandle columnFamily,
TypeSerializer<N> namespaceSerializer,
ValueStateDescriptor<V> stateDesc,
- RocksDBStateBackend backend) {
+ RocksDBKeyedStateBackend backend) {
super(columnFamily, namespaceSerializer, stateDesc, backend);
this.valueSerializer = stateDesc.getSerializer();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/FullyAsyncRocksDBStateBackendTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/FullyAsyncRocksDBStateBackendTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/FullyAsyncRocksDBStateBackendTest.java
deleted file mode 100644
index 7861542..0000000
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/FullyAsyncRocksDBStateBackendTest.java
+++ /dev/null
@@ -1,65 +0,0 @@
-/*
- * 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.contrib.streaming.state;
-
-import org.apache.commons.io.FileUtils;
-import org.apache.flink.configuration.ConfigConstants;
-import org.apache.flink.runtime.state.StateBackendTestBase;
-import org.apache.flink.runtime.state.memory.MemoryStateBackend;
-import org.apache.flink.util.OperatingSystem;
-import org.junit.Assume;
-import org.junit.Before;
-
-import java.io.File;
-import java.io.IOException;
-import java.util.UUID;
-
-/**
- * Tests for the partitioned state part of {@link RocksDBStateBackend} with fully asynchronous
- * checkpointing enabled.
- */
-public class FullyAsyncRocksDBStateBackendTest extends StateBackendTestBase<RocksDBStateBackend> {
-
- private File dbDir;
- private File chkDir;
-
- @Before
- public void checkOperatingSystem() {
- Assume.assumeTrue("This test can't run successfully on Windows.", !OperatingSystem.isWindows());
- }
-
- @Override
- protected RocksDBStateBackend getStateBackend() throws IOException {
- dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state");
- chkDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "snapshots");
-
- RocksDBStateBackend backend = new RocksDBStateBackend(chkDir.getAbsoluteFile().toURI(), new MemoryStateBackend());
- backend.setDbStoragePath(dbDir.getAbsolutePath());
- backend.enableFullyAsyncSnapshots();
- return backend;
- }
-
- @Override
- protected void cleanup() {
- try {
- FileUtils.deleteDirectory(dbDir);
- FileUtils.deleteDirectory(chkDir);
- } catch (IOException ignore) {}
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
index d720c6d..0e35b60 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncKVSnapshotTest.java
@@ -88,7 +88,6 @@ public class RocksDBAsyncKVSnapshotTest {
* test will simply lock forever.
*/
@Test
- @Ignore
public void testAsyncCheckpoints() throws Exception {
LocalFileSystem localFS = new LocalFileSystem();
localFS.initialize(new URI("file:///"), new Configuration());
@@ -191,7 +190,6 @@ public class RocksDBAsyncKVSnapshotTest {
* test will simply lock forever.
*/
@Test
- @Ignore
public void testFullyAsyncCheckpoints() throws Exception {
LocalFileSystem localFS = new LocalFileSystem();
localFS.initialize(new URI("file:///"), new Configuration());
@@ -218,7 +216,7 @@ public class RocksDBAsyncKVSnapshotTest {
RocksDBStateBackend backend = new RocksDBStateBackend(chkDir.getAbsoluteFile().toURI(), new MemoryStateBackend());
backend.setDbStoragePath(dbDir.getAbsolutePath());
- backend.enableFullyAsyncSnapshots();
+// backend.enableFullyAsyncSnapshots();
streamConfig.setStateBackend(backend);
[12/27] flink git commit: [FLINK-3761] Refactor State Backends/Make
Keyed State Key-Group Aware
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
index 657c57e..6f4a983 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
@@ -1,322 +1,322 @@
-/*
- * 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.contrib.streaming.state;
-
-import org.apache.commons.io.FileUtils;
-import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.TaskInfo;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.api.common.typeutils.base.IntSerializer;
-import org.apache.flink.runtime.execution.Environment;
-import org.apache.flink.runtime.io.disk.iomanager.IOManager;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-
-import org.apache.flink.runtime.state.VoidNamespace;
-import org.apache.flink.runtime.state.VoidNamespaceSerializer;
-import org.apache.flink.util.OperatingSystem;
-import org.junit.Assume;
-import org.junit.Before;
-import org.junit.Test;
-
-import org.rocksdb.ColumnFamilyOptions;
-import org.rocksdb.CompactionStyle;
-import org.rocksdb.DBOptions;
-
-import java.io.File;
-import java.util.UUID;
-
-import static org.junit.Assert.*;
-import static org.mockito.Mockito.*;
-
-/**
- * Tests for configuring the RocksDB State Backend
- */
-@SuppressWarnings("serial")
-public class RocksDBStateBackendConfigTest {
-
- private static final String TEMP_URI = new File(System.getProperty("java.io.tmpdir")).toURI().toString();
-
- @Before
- public void checkOperatingSystem() {
- Assume.assumeTrue("This test can't run successfully on Windows.", !OperatingSystem.isWindows());
- }
-
- // ------------------------------------------------------------------------
- // RocksDB local file directory
- // ------------------------------------------------------------------------
-
- @Test
- public void testSetDbPath() throws Exception {
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-
- assertNull(rocksDbBackend.getDbStoragePaths());
-
- rocksDbBackend.setDbStoragePath("/abc/def");
- assertArrayEquals(new String[] { "/abc/def" }, rocksDbBackend.getDbStoragePaths());
-
- rocksDbBackend.setDbStoragePath(null);
- assertNull(rocksDbBackend.getDbStoragePaths());
-
- rocksDbBackend.setDbStoragePaths("/abc/def", "/uvw/xyz");
- assertArrayEquals(new String[] { "/abc/def", "/uvw/xyz" }, rocksDbBackend.getDbStoragePaths());
-
- //noinspection NullArgumentToVariableArgMethod
- rocksDbBackend.setDbStoragePaths(null);
- assertNull(rocksDbBackend.getDbStoragePaths());
- }
-
- @Test(expected = IllegalArgumentException.class)
- public void testSetNullPaths() throws Exception {
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
- rocksDbBackend.setDbStoragePaths();
- }
-
- @Test(expected = IllegalArgumentException.class)
- public void testNonFileSchemePath() throws Exception {
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
- rocksDbBackend.setDbStoragePath("hdfs:///some/path/to/perdition");
- }
-
- // ------------------------------------------------------------------------
- // RocksDB local file automatic from temp directories
- // ------------------------------------------------------------------------
-
- @Test
- public void testUseTempDirectories() throws Exception {
- File dir1 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
- File dir2 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
-
- File[] tempDirs = new File[] { dir1, dir2 };
-
- try {
- assertTrue(dir1.mkdirs());
- assertTrue(dir2.mkdirs());
-
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
- assertNull(rocksDbBackend.getDbStoragePaths());
-
- rocksDbBackend.initializeForJob(getMockEnvironment(tempDirs), "foobar", IntSerializer.INSTANCE);
- assertArrayEquals(tempDirs, rocksDbBackend.getStoragePaths());
- }
- finally {
- FileUtils.deleteDirectory(dir1);
- FileUtils.deleteDirectory(dir2);
- }
- }
-
- // ------------------------------------------------------------------------
- // RocksDB local file directory initialization
- // ------------------------------------------------------------------------
-
- @Test
- public void testFailWhenNoLocalStorageDir() throws Exception {
- File targetDir = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
- try {
- assertTrue(targetDir.mkdirs());
-
- if (!targetDir.setWritable(false, false)) {
- System.err.println("Cannot execute 'testFailWhenNoLocalStorageDir' because cannot mark directory non-writable");
- return;
- }
-
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
- rocksDbBackend.setDbStoragePath(targetDir.getAbsolutePath());
-
- boolean hasFailure = false;
- try {
- rocksDbBackend.initializeForJob(getMockEnvironment(), "foobar", IntSerializer.INSTANCE);
- }
- catch (Exception e) {
- assertTrue(e.getMessage().contains("No local storage directories available"));
- assertTrue(e.getMessage().contains(targetDir.getAbsolutePath()));
- hasFailure = true;
- }
- assertTrue("We must see a failure because no storaged directory is feasible.", hasFailure);
- }
- finally {
- //noinspection ResultOfMethodCallIgnored
- targetDir.setWritable(true, false);
- FileUtils.deleteDirectory(targetDir);
- }
- }
-
- @Test
- public void testContinueOnSomeDbDirectoriesMissing() throws Exception {
- File targetDir1 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
- File targetDir2 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
-
- try {
- assertTrue(targetDir1.mkdirs());
- assertTrue(targetDir2.mkdirs());
-
- if (!targetDir1.setWritable(false, false)) {
- System.err.println("Cannot execute 'testContinueOnSomeDbDirectoriesMissing' because cannot mark directory non-writable");
- return;
- }
-
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
- rocksDbBackend.setDbStoragePaths(targetDir1.getAbsolutePath(), targetDir2.getAbsolutePath());
-
- try {
- rocksDbBackend.initializeForJob(getMockEnvironment(), "foobar", IntSerializer.INSTANCE);
-
- // actually get a state to see whether we can write to the storage directory
- rocksDbBackend.getPartitionedState(
- VoidNamespace.INSTANCE,
- VoidNamespaceSerializer.INSTANCE,
- new ValueStateDescriptor<>("test", String.class, ""));
- }
- catch (Exception e) {
- e.printStackTrace();
- fail("Backend initialization failed even though some paths were available");
- }
- } finally {
- //noinspection ResultOfMethodCallIgnored
- targetDir1.setWritable(true, false);
- FileUtils.deleteDirectory(targetDir1);
- FileUtils.deleteDirectory(targetDir2);
- }
- }
-
- // ------------------------------------------------------------------------
- // RocksDB Options
- // ------------------------------------------------------------------------
-
- @Test
- public void testPredefinedOptions() throws Exception {
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-
- assertEquals(PredefinedOptions.DEFAULT, rocksDbBackend.getPredefinedOptions());
-
- rocksDbBackend.setPredefinedOptions(PredefinedOptions.SPINNING_DISK_OPTIMIZED);
- assertEquals(PredefinedOptions.SPINNING_DISK_OPTIMIZED, rocksDbBackend.getPredefinedOptions());
-
- DBOptions opt1 = rocksDbBackend.getDbOptions();
- DBOptions opt2 = rocksDbBackend.getDbOptions();
-
- assertEquals(opt1, opt2);
-
- ColumnFamilyOptions columnOpt1 = rocksDbBackend.getColumnOptions();
- ColumnFamilyOptions columnOpt2 = rocksDbBackend.getColumnOptions();
-
- assertEquals(columnOpt1, columnOpt2);
-
- assertEquals(CompactionStyle.LEVEL, columnOpt1.compactionStyle());
- }
-
- @Test
- public void testOptionsFactory() throws Exception {
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-
- rocksDbBackend.setOptions(new OptionsFactory() {
- @Override
- public DBOptions createDBOptions(DBOptions currentOptions) {
- return currentOptions;
- }
-
- @Override
- public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions) {
- return currentOptions.setCompactionStyle(CompactionStyle.FIFO);
- }
- });
-
- assertNotNull(rocksDbBackend.getOptions());
- assertEquals(CompactionStyle.FIFO, rocksDbBackend.getColumnOptions().compactionStyle());
- }
-
- @Test
- public void testPredefinedAndOptionsFactory() throws Exception {
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
-
- assertEquals(PredefinedOptions.DEFAULT, rocksDbBackend.getPredefinedOptions());
-
- rocksDbBackend.setPredefinedOptions(PredefinedOptions.SPINNING_DISK_OPTIMIZED);
- rocksDbBackend.setOptions(new OptionsFactory() {
- @Override
- public DBOptions createDBOptions(DBOptions currentOptions) {
- return currentOptions;
- }
-
- @Override
- public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions) {
- return currentOptions.setCompactionStyle(CompactionStyle.UNIVERSAL);
- }
- });
-
- assertEquals(PredefinedOptions.SPINNING_DISK_OPTIMIZED, rocksDbBackend.getPredefinedOptions());
- assertNotNull(rocksDbBackend.getOptions());
- assertEquals(CompactionStyle.UNIVERSAL, rocksDbBackend.getColumnOptions().compactionStyle());
- }
-
- @Test
- public void testPredefinedOptionsEnum() {
- for (PredefinedOptions o : PredefinedOptions.values()) {
- DBOptions opt = o.createDBOptions();
- try {
- assertNotNull(opt);
- } finally {
- opt.dispose();
- }
- }
- }
-
- // ------------------------------------------------------------------------
- // Contained Non-partitioned State Backend
- // ------------------------------------------------------------------------
-
- @Test
- public void testCallsForwardedToNonPartitionedBackend() throws Exception {
- AbstractStateBackend nonPartBackend = mock(AbstractStateBackend.class);
- RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI, nonPartBackend);
-
- rocksDbBackend.initializeForJob(getMockEnvironment(), "foo", IntSerializer.INSTANCE);
- verify(nonPartBackend, times(1)).initializeForJob(any(Environment.class), anyString(), any(TypeSerializer.class));
-
- rocksDbBackend.disposeAllStateForCurrentJob();
- verify(nonPartBackend, times(1)).disposeAllStateForCurrentJob();
-
- rocksDbBackend.close();
- verify(nonPartBackend, times(1)).close();
- }
-
- // ------------------------------------------------------------------------
- // Utilities
- // ------------------------------------------------------------------------
-
- private static Environment getMockEnvironment() {
- return getMockEnvironment(new File[] { new File(System.getProperty("java.io.tmpdir")) });
- }
-
- private static Environment getMockEnvironment(File[] tempDirs) {
- IOManager ioMan = mock(IOManager.class);
- when(ioMan.getSpillingDirectories()).thenReturn(tempDirs);
-
- Environment env = mock(Environment.class);
- when(env.getJobID()).thenReturn(new JobID());
- when(env.getUserClassLoader()).thenReturn(RocksDBStateBackendConfigTest.class.getClassLoader());
- when(env.getIOManager()).thenReturn(ioMan);
-
- TaskInfo taskInfo = mock(TaskInfo.class);
- when(env.getTaskInfo()).thenReturn(taskInfo);
-
- when(taskInfo.getIndexOfThisSubtask()).thenReturn(0);
- return env;
- }
-}
+///*
+// * 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.contrib.streaming.state;
+//
+//import org.apache.commons.io.FileUtils;
+//import org.apache.flink.api.common.JobID;
+//import org.apache.flink.api.common.TaskInfo;
+//import org.apache.flink.api.common.state.ValueStateDescriptor;
+//import org.apache.flink.api.common.typeutils.TypeSerializer;
+//import org.apache.flink.api.common.typeutils.base.IntSerializer;
+//import org.apache.flink.runtime.execution.Environment;
+//import org.apache.flink.runtime.io.disk.iomanager.IOManager;
+//import org.apache.flink.runtime.state.AbstractStateBackend;
+//
+//import org.apache.flink.runtime.state.VoidNamespace;
+//import org.apache.flink.runtime.state.VoidNamespaceSerializer;
+//import org.apache.flink.util.OperatingSystem;
+//import org.junit.Assume;
+//import org.junit.Before;
+//import org.junit.Test;
+//
+//import org.rocksdb.ColumnFamilyOptions;
+//import org.rocksdb.CompactionStyle;
+//import org.rocksdb.DBOptions;
+//
+//import java.io.File;
+//import java.util.UUID;
+//
+//import static org.junit.Assert.*;
+//import static org.mockito.Mockito.*;
+//
+///**
+// * Tests for configuring the RocksDB State Backend
+// */
+//@SuppressWarnings("serial")
+//public class RocksDBStateBackendConfigTest {
+//
+// private static final String TEMP_URI = new File(System.getProperty("java.io.tmpdir")).toURI().toString();
+//
+// @Before
+// public void checkOperatingSystem() {
+// Assume.assumeTrue("This test can't run successfully on Windows.", !OperatingSystem.isWindows());
+// }
+//
+// // ------------------------------------------------------------------------
+// // RocksDB local file directory
+// // ------------------------------------------------------------------------
+//
+// @Test
+// public void testSetDbPath() throws Exception {
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
+//
+// assertNull(rocksDbBackend.getDbStoragePaths());
+//
+// rocksDbBackend.setDbStoragePath("/abc/def");
+// assertArrayEquals(new String[] { "/abc/def" }, rocksDbBackend.getDbStoragePaths());
+//
+// rocksDbBackend.setDbStoragePath(null);
+// assertNull(rocksDbBackend.getDbStoragePaths());
+//
+// rocksDbBackend.setDbStoragePaths("/abc/def", "/uvw/xyz");
+// assertArrayEquals(new String[] { "/abc/def", "/uvw/xyz" }, rocksDbBackend.getDbStoragePaths());
+//
+// //noinspection NullArgumentToVariableArgMethod
+// rocksDbBackend.setDbStoragePaths(null);
+// assertNull(rocksDbBackend.getDbStoragePaths());
+// }
+//
+// @Test(expected = IllegalArgumentException.class)
+// public void testSetNullPaths() throws Exception {
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
+// rocksDbBackend.setDbStoragePaths();
+// }
+//
+// @Test(expected = IllegalArgumentException.class)
+// public void testNonFileSchemePath() throws Exception {
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
+// rocksDbBackend.setDbStoragePath("hdfs:///some/path/to/perdition");
+// }
+//
+// // ------------------------------------------------------------------------
+// // RocksDB local file automatic from temp directories
+// // ------------------------------------------------------------------------
+//
+// @Test
+// public void testUseTempDirectories() throws Exception {
+// File dir1 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
+// File dir2 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
+//
+// File[] tempDirs = new File[] { dir1, dir2 };
+//
+// try {
+// assertTrue(dir1.mkdirs());
+// assertTrue(dir2.mkdirs());
+//
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
+// assertNull(rocksDbBackend.getDbStoragePaths());
+//
+// rocksDbBackend.initializeForJob(getMockEnvironment(tempDirs), "foobar", IntSerializer.INSTANCE);
+// assertArrayEquals(tempDirs, rocksDbBackend.getStoragePaths());
+// }
+// finally {
+// FileUtils.deleteDirectory(dir1);
+// FileUtils.deleteDirectory(dir2);
+// }
+// }
+//
+// // ------------------------------------------------------------------------
+// // RocksDB local file directory initialization
+// // ------------------------------------------------------------------------
+//
+// @Test
+// public void testFailWhenNoLocalStorageDir() throws Exception {
+// File targetDir = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
+// try {
+// assertTrue(targetDir.mkdirs());
+//
+// if (!targetDir.setWritable(false, false)) {
+// System.err.println("Cannot execute 'testFailWhenNoLocalStorageDir' because cannot mark directory non-writable");
+// return;
+// }
+//
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
+// rocksDbBackend.setDbStoragePath(targetDir.getAbsolutePath());
+//
+// boolean hasFailure = false;
+// try {
+// rocksDbBackend.initializeForJob(getMockEnvironment(), "foobar", IntSerializer.INSTANCE);
+// }
+// catch (Exception e) {
+// assertTrue(e.getMessage().contains("No local storage directories available"));
+// assertTrue(e.getMessage().contains(targetDir.getAbsolutePath()));
+// hasFailure = true;
+// }
+// assertTrue("We must see a failure because no storaged directory is feasible.", hasFailure);
+// }
+// finally {
+// //noinspection ResultOfMethodCallIgnored
+// targetDir.setWritable(true, false);
+// FileUtils.deleteDirectory(targetDir);
+// }
+// }
+//
+// @Test
+// public void testContinueOnSomeDbDirectoriesMissing() throws Exception {
+// File targetDir1 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
+// File targetDir2 = new File(System.getProperty("java.io.tmpdir"), UUID.randomUUID().toString());
+//
+// try {
+// assertTrue(targetDir1.mkdirs());
+// assertTrue(targetDir2.mkdirs());
+//
+// if (!targetDir1.setWritable(false, false)) {
+// System.err.println("Cannot execute 'testContinueOnSomeDbDirectoriesMissing' because cannot mark directory non-writable");
+// return;
+// }
+//
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
+// rocksDbBackend.setDbStoragePaths(targetDir1.getAbsolutePath(), targetDir2.getAbsolutePath());
+//
+// try {
+// rocksDbBackend.initializeForJob(getMockEnvironment(), "foobar", IntSerializer.INSTANCE);
+//
+// // actually get a state to see whether we can write to the storage directory
+// rocksDbBackend.getPartitionedState(
+// VoidNamespace.INSTANCE,
+// VoidNamespaceSerializer.INSTANCE,
+// new ValueStateDescriptor<>("test", String.class, ""));
+// }
+// catch (Exception e) {
+// e.printStackTrace();
+// fail("Backend initialization failed even though some paths were available");
+// }
+// } finally {
+// //noinspection ResultOfMethodCallIgnored
+// targetDir1.setWritable(true, false);
+// FileUtils.deleteDirectory(targetDir1);
+// FileUtils.deleteDirectory(targetDir2);
+// }
+// }
+//
+// // ------------------------------------------------------------------------
+// // RocksDB Options
+// // ------------------------------------------------------------------------
+//
+// @Test
+// public void testPredefinedOptions() throws Exception {
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
+//
+// assertEquals(PredefinedOptions.DEFAULT, rocksDbBackend.getPredefinedOptions());
+//
+// rocksDbBackend.setPredefinedOptions(PredefinedOptions.SPINNING_DISK_OPTIMIZED);
+// assertEquals(PredefinedOptions.SPINNING_DISK_OPTIMIZED, rocksDbBackend.getPredefinedOptions());
+//
+// DBOptions opt1 = rocksDbBackend.getDbOptions();
+// DBOptions opt2 = rocksDbBackend.getDbOptions();
+//
+// assertEquals(opt1, opt2);
+//
+// ColumnFamilyOptions columnOpt1 = rocksDbBackend.getColumnOptions();
+// ColumnFamilyOptions columnOpt2 = rocksDbBackend.getColumnOptions();
+//
+// assertEquals(columnOpt1, columnOpt2);
+//
+// assertEquals(CompactionStyle.LEVEL, columnOpt1.compactionStyle());
+// }
+//
+// @Test
+// public void testOptionsFactory() throws Exception {
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
+//
+// rocksDbBackend.setOptions(new OptionsFactory() {
+// @Override
+// public DBOptions createDBOptions(DBOptions currentOptions) {
+// return currentOptions;
+// }
+//
+// @Override
+// public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions) {
+// return currentOptions.setCompactionStyle(CompactionStyle.FIFO);
+// }
+// });
+//
+// assertNotNull(rocksDbBackend.getOptions());
+// assertEquals(CompactionStyle.FIFO, rocksDbBackend.getColumnOptions().compactionStyle());
+// }
+//
+// @Test
+// public void testPredefinedAndOptionsFactory() throws Exception {
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI);
+//
+// assertEquals(PredefinedOptions.DEFAULT, rocksDbBackend.getPredefinedOptions());
+//
+// rocksDbBackend.setPredefinedOptions(PredefinedOptions.SPINNING_DISK_OPTIMIZED);
+// rocksDbBackend.setOptions(new OptionsFactory() {
+// @Override
+// public DBOptions createDBOptions(DBOptions currentOptions) {
+// return currentOptions;
+// }
+//
+// @Override
+// public ColumnFamilyOptions createColumnOptions(ColumnFamilyOptions currentOptions) {
+// return currentOptions.setCompactionStyle(CompactionStyle.UNIVERSAL);
+// }
+// });
+//
+// assertEquals(PredefinedOptions.SPINNING_DISK_OPTIMIZED, rocksDbBackend.getPredefinedOptions());
+// assertNotNull(rocksDbBackend.getOptions());
+// assertEquals(CompactionStyle.UNIVERSAL, rocksDbBackend.getColumnOptions().compactionStyle());
+// }
+//
+// @Test
+// public void testPredefinedOptionsEnum() {
+// for (PredefinedOptions o : PredefinedOptions.values()) {
+// DBOptions opt = o.createDBOptions();
+// try {
+// assertNotNull(opt);
+// } finally {
+// opt.dispose();
+// }
+// }
+// }
+//
+// // ------------------------------------------------------------------------
+// // Contained Non-partitioned State Backend
+// // ------------------------------------------------------------------------
+//
+// @Test
+// public void testCallsForwardedToNonPartitionedBackend() throws Exception {
+// AbstractStateBackend nonPartBackend = mock(AbstractStateBackend.class);
+// RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(TEMP_URI, nonPartBackend);
+//
+// rocksDbBackend.initializeForJob(getMockEnvironment(), "foo", IntSerializer.INSTANCE);
+// verify(nonPartBackend, times(1)).initializeForJob(any(Environment.class), anyString(), any(TypeSerializer.class));
+//
+// rocksDbBackend.disposeAllStateForCurrentJob();
+// verify(nonPartBackend, times(1)).disposeAllStateForCurrentJob();
+//
+// rocksDbBackend.close();
+// verify(nonPartBackend, times(1)).close();
+// }
+//
+// // ------------------------------------------------------------------------
+// // Utilities
+// // ------------------------------------------------------------------------
+//
+// private static Environment getMockEnvironment() {
+// return getMockEnvironment(new File[] { new File(System.getProperty("java.io.tmpdir")) });
+// }
+//
+// private static Environment getMockEnvironment(File[] tempDirs) {
+// IOManager ioMan = mock(IOManager.class);
+// when(ioMan.getSpillingDirectories()).thenReturn(tempDirs);
+//
+// Environment env = mock(Environment.class);
+// when(env.getJobID()).thenReturn(new JobID());
+// when(env.getUserClassLoader()).thenReturn(RocksDBStateBackendConfigTest.class.getClassLoader());
+// when(env.getIOManager()).thenReturn(ioMan);
+//
+// TaskInfo taskInfo = mock(TaskInfo.class);
+// when(env.getTaskInfo()).thenReturn(taskInfo);
+//
+// when(taskInfo.getIndexOfThisSubtask()).thenReturn(0);
+// return env;
+// }
+//}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java
index 57f906e..9222f0b 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendTest.java
@@ -18,25 +18,23 @@
package org.apache.flink.contrib.streaming.state;
-import org.apache.commons.io.FileUtils;
-import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.runtime.state.StateBackendTestBase;
-import org.apache.flink.runtime.state.memory.MemoryStateBackend;
+import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.util.OperatingSystem;
import org.junit.Assume;
import org.junit.Before;
+import org.junit.Rule;
+import org.junit.rules.TemporaryFolder;
-import java.io.File;
import java.io.IOException;
-import java.util.UUID;
/**
* Tests for the partitioned state part of {@link RocksDBStateBackend}.
*/
public class RocksDBStateBackendTest extends StateBackendTestBase<RocksDBStateBackend> {
- private File dbDir;
- private File chkDir;
+ @Rule
+ public TemporaryFolder tempFolder = new TemporaryFolder();
@Before
public void checkOperatingSystem() {
@@ -45,19 +43,10 @@ public class RocksDBStateBackendTest extends StateBackendTestBase<RocksDBStateBa
@Override
protected RocksDBStateBackend getStateBackend() throws IOException {
- dbDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "state");
- chkDir = new File(new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()), "snapshots");
-
- RocksDBStateBackend backend = new RocksDBStateBackend(chkDir.getAbsoluteFile().toURI(), new MemoryStateBackend());
- backend.setDbStoragePath(dbDir.getAbsolutePath());
+ String dbPath = tempFolder.newFolder().getAbsolutePath();
+ String checkpointPath = tempFolder.newFolder().toURI().toString();
+ RocksDBStateBackend backend = new RocksDBStateBackend(checkpointPath, new FsStateBackend(checkpointPath));
+ backend.setDbStoragePath(dbPath);
return backend;
}
-
- @Override
- protected void cleanup() {
- try {
- FileUtils.deleteDirectory(dbDir);
- FileUtils.deleteDirectory(chkDir);
- } catch (IOException ignore) {}
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java b/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java
index eb04038..5df1337 100644
--- a/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java
+++ b/flink-contrib/flink-storm-examples/src/test/java/org/apache/flink/storm/tests/StormFieldsGroupingITCase.java
@@ -61,12 +61,22 @@ public class StormFieldsGroupingITCase extends StreamingProgramTestBase {
List<String> actualResults = new ArrayList<>();
readAllResultLines(actualResults, resultPath, new String[0], false);
+ //remove potential operator id prefix
+ for(int i = 0; i < actualResults.size(); ++i) {
+ String s = actualResults.get(i);
+ if(s.contains(">")) {
+ s = s.substring(s.indexOf(">") + 2);
+ actualResults.set(i, s);
+ }
+ }
+
Assert.assertEquals(expectedResults.size(),actualResults.size());
Collections.sort(actualResults);
Collections.sort(expectedResults);
+ System.out.println(actualResults);
for(int i=0; i< actualResults.size(); ++i) {
//compare against actual results with removed prefex (as it depends e.g. on the hash function used)
- Assert.assertEquals(expectedResults.get(i), actualResults.get(i).substring(3));
+ Assert.assertEquals(expectedResults.get(i), actualResults.get(i));
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-storm/src/main/java/org/apache/flink/storm/wrappers/BoltWrapper.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-storm/src/main/java/org/apache/flink/storm/wrappers/BoltWrapper.java b/flink-contrib/flink-storm/src/main/java/org/apache/flink/storm/wrappers/BoltWrapper.java
index 6e316e7..d59ff04 100644
--- a/flink-contrib/flink-storm/src/main/java/org/apache/flink/storm/wrappers/BoltWrapper.java
+++ b/flink-contrib/flink-storm/src/main/java/org/apache/flink/storm/wrappers/BoltWrapper.java
@@ -294,7 +294,7 @@ public class BoltWrapper<IN, OUT> extends AbstractStreamOperator<OUT> implements
}
@Override
- public void dispose() {
+ public void dispose() throws Exception {
super.dispose();
this.bolt.cleanup();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-contrib/flink-storm/src/test/java/org/apache/flink/storm/wrappers/BoltWrapperTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-storm/src/test/java/org/apache/flink/storm/wrappers/BoltWrapperTest.java b/flink-contrib/flink-storm/src/test/java/org/apache/flink/storm/wrappers/BoltWrapperTest.java
index 2ebb917..c15b5f6 100644
--- a/flink-contrib/flink-storm/src/test/java/org/apache/flink/storm/wrappers/BoltWrapperTest.java
+++ b/flink-contrib/flink-storm/src/test/java/org/apache/flink/storm/wrappers/BoltWrapperTest.java
@@ -327,7 +327,7 @@ public class BoltWrapperTest extends AbstractTest {
private static final class TestBolt implements IRichBolt {
private static final long serialVersionUID = 7278692872260138758L;
- private OutputCollector collector;
+ private transient OutputCollector collector;
@SuppressWarnings("rawtypes")
@Override
@@ -366,7 +366,7 @@ public class BoltWrapperTest extends AbstractTest {
public static StreamTask<?, ?> createMockStreamTask(ExecutionConfig execConfig) {
Environment env = mock(Environment.class);
- when(env.getTaskInfo()).thenReturn(new TaskInfo("Mock Task", 0, 1, 0));
+ when(env.getTaskInfo()).thenReturn(new TaskInfo("Mock Task", 1, 0, 1, 0));
when(env.getUserClassLoader()).thenReturn(BoltWrapperTest.class.getClassLoader());
when(env.getMetricGroup()).thenReturn(new UnregisteredTaskMetricsGroup());
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/main/java/org/apache/flink/api/common/TaskInfo.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/api/common/TaskInfo.java b/flink-core/src/main/java/org/apache/flink/api/common/TaskInfo.java
index ac87e74..5627ca8 100644
--- a/flink-core/src/main/java/org/apache/flink/api/common/TaskInfo.java
+++ b/flink-core/src/main/java/org/apache/flink/api/common/TaskInfo.java
@@ -31,16 +31,20 @@ public class TaskInfo {
private final String taskName;
private final String taskNameWithSubtasks;
+ private final int numberOfKeyGroups;
private final int indexOfSubtask;
private final int numberOfParallelSubtasks;
private final int attemptNumber;
- public TaskInfo(String taskName, int indexOfSubtask, int numberOfParallelSubtasks, int attemptNumber) {
+ public TaskInfo(String taskName, int numberOfKeyGroups, int indexOfSubtask, int numberOfParallelSubtasks, int attemptNumber) {
checkArgument(indexOfSubtask >= 0, "Task index must be a non-negative number.");
+ checkArgument(numberOfKeyGroups >= 1, "Max parallelism must be a positive number.");
+ checkArgument(numberOfKeyGroups >= numberOfParallelSubtasks, "Max parallelism must be >= than parallelism.");
checkArgument(numberOfParallelSubtasks >= 1, "Parallelism must be a positive number.");
checkArgument(indexOfSubtask < numberOfParallelSubtasks, "Task index must be less than parallelism.");
checkArgument(attemptNumber >= 0, "Attempt number must be a non-negative number.");
this.taskName = checkNotNull(taskName, "Task Name must not be null.");
+ this.numberOfKeyGroups = numberOfKeyGroups;
this.indexOfSubtask = indexOfSubtask;
this.numberOfParallelSubtasks = numberOfParallelSubtasks;
this.attemptNumber = attemptNumber;
@@ -57,6 +61,13 @@ public class TaskInfo {
}
/**
+ * Gets the number of key groups aka the max parallelism aka the max number of subtasks.
+ */
+ public int getNumberOfKeyGroups() {
+ return numberOfKeyGroups;
+ }
+
+ /**
* Gets the number of this parallel subtask. The numbering starts from 0 and goes up to
* parallelism-1 (parallelism as returned by {@link #getNumberOfParallelSubtasks()}).
*
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/main/java/org/apache/flink/api/common/operators/CollectionExecutor.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/api/common/operators/CollectionExecutor.java b/flink-core/src/main/java/org/apache/flink/api/common/operators/CollectionExecutor.java
index 913b205..d9240fe 100644
--- a/flink-core/src/main/java/org/apache/flink/api/common/operators/CollectionExecutor.java
+++ b/flink-core/src/main/java/org/apache/flink/api/common/operators/CollectionExecutor.java
@@ -183,7 +183,7 @@ public class CollectionExecutor {
GenericDataSinkBase<IN> typedSink = (GenericDataSinkBase<IN>) sink;
// build the runtime context and compute broadcast variables, if necessary
- TaskInfo taskInfo = new TaskInfo(typedSink.getName(), 0, 1, 0);
+ TaskInfo taskInfo = new TaskInfo(typedSink.getName(), 1, 0, 1, 0);
RuntimeUDFContext ctx;
MetricGroup metrics = new UnregisteredMetricsGroup();
@@ -203,7 +203,7 @@ public class CollectionExecutor {
@SuppressWarnings("unchecked")
GenericDataSourceBase<OUT, ?> typedSource = (GenericDataSourceBase<OUT, ?>) source;
// build the runtime context and compute broadcast variables, if necessary
- TaskInfo taskInfo = new TaskInfo(typedSource.getName(), 0, 1, 0);
+ TaskInfo taskInfo = new TaskInfo(typedSource.getName(), 1, 0, 1, 0);
RuntimeUDFContext ctx;
@@ -230,7 +230,7 @@ public class CollectionExecutor {
SingleInputOperator<IN, OUT, ?> typedOp = (SingleInputOperator<IN, OUT, ?>) operator;
// build the runtime context and compute broadcast variables, if necessary
- TaskInfo taskInfo = new TaskInfo(typedOp.getName(), 0, 1, 0);
+ TaskInfo taskInfo = new TaskInfo(typedOp.getName(), 1, 0, 1, 0);
RuntimeUDFContext ctx;
MetricGroup metrics = new UnregisteredMetricsGroup();
@@ -270,7 +270,7 @@ public class CollectionExecutor {
DualInputOperator<IN1, IN2, OUT, ?> typedOp = (DualInputOperator<IN1, IN2, OUT, ?>) operator;
// build the runtime context and compute broadcast variables, if necessary
- TaskInfo taskInfo = new TaskInfo(typedOp.getName(), 0, 1, 0);
+ TaskInfo taskInfo = new TaskInfo(typedOp.getName(), 1, 0, 1, 0);
RuntimeUDFContext ctx;
MetricGroup metrics = new UnregisteredMetricsGroup();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/main/java/org/apache/flink/core/fs/FSDataOutputStream.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/core/fs/FSDataOutputStream.java b/flink-core/src/main/java/org/apache/flink/core/fs/FSDataOutputStream.java
index a6becf7..0318d1f 100644
--- a/flink-core/src/main/java/org/apache/flink/core/fs/FSDataOutputStream.java
+++ b/flink-core/src/main/java/org/apache/flink/core/fs/FSDataOutputStream.java
@@ -29,6 +29,8 @@ import java.io.OutputStream;
@Public
public abstract class FSDataOutputStream extends OutputStream {
+ public abstract long getPos() throws IOException;
+
public abstract void flush() throws IOException;
public abstract void sync() throws IOException;
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/main/java/org/apache/flink/core/fs/local/LocalDataOutputStream.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/core/fs/local/LocalDataOutputStream.java b/flink-core/src/main/java/org/apache/flink/core/fs/local/LocalDataOutputStream.java
index 54ec8dd..c3b793d 100644
--- a/flink-core/src/main/java/org/apache/flink/core/fs/local/LocalDataOutputStream.java
+++ b/flink-core/src/main/java/org/apache/flink/core/fs/local/LocalDataOutputStream.java
@@ -90,4 +90,9 @@ public class LocalDataOutputStream extends FSDataOutputStream {
public void sync() throws IOException {
fos.getFD().sync();
}
+
+ @Override
+ public long getPos() throws IOException {
+ return fos.getChannel().position();
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/main/java/org/apache/flink/core/memory/ByteArrayOutputStreamWithPos.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/core/memory/ByteArrayOutputStreamWithPos.java b/flink-core/src/main/java/org/apache/flink/core/memory/ByteArrayOutputStreamWithPos.java
new file mode 100644
index 0000000..285e016
--- /dev/null
+++ b/flink-core/src/main/java/org/apache/flink/core/memory/ByteArrayOutputStreamWithPos.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.core.memory;
+
+import org.apache.flink.util.Preconditions;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.io.UnsupportedEncodingException;
+import java.util.Arrays;
+
+/**
+ * Un-synchronized copy of Java's ByteArrayOutputStream that also exposes the current position.
+ */
+public class ByteArrayOutputStreamWithPos extends OutputStream {
+
+ /**
+ * The buffer where data is stored.
+ */
+ protected byte[] buf;
+
+ /**
+ * The number of valid bytes in the buffer.
+ */
+ protected int count;
+
+ /**
+ * Creates a new byte array output stream. The buffer capacity is
+ * initially 32 bytes, though its size increases if necessary.
+ */
+ public ByteArrayOutputStreamWithPos() {
+ this(32);
+ }
+
+ /**
+ * Creates a new byte array output stream, with a buffer capacity of
+ * the specified size, in bytes.
+ *
+ * @param size the initial size.
+ * @throws IllegalArgumentException if size is negative.
+ */
+ public ByteArrayOutputStreamWithPos(int size) {
+ if (size < 0) {
+ throw new IllegalArgumentException("Negative initial size: "
+ + size);
+ }
+ buf = new byte[size];
+ }
+
+ /**
+ * Increases the capacity if necessary to ensure that it can hold
+ * at least the number of elements specified by the minimum
+ * capacity argument.
+ *
+ * @param minCapacity the desired minimum capacity
+ * @throws OutOfMemoryError if {@code minCapacity < 0}. This is
+ * interpreted as a request for the unsatisfiably large capacity
+ * {@code (long) Integer.MAX_VALUE + (minCapacity - Integer.MAX_VALUE)}.
+ */
+ private void ensureCapacity(int minCapacity) {
+ // overflow-conscious code
+ if (minCapacity - buf.length > 0) {
+ grow(minCapacity);
+ }
+ }
+
+ /**
+ * Increases the capacity to ensure that it can hold at least the
+ * number of elements specified by the minimum capacity argument.
+ *
+ * @param minCapacity the desired minimum capacity
+ */
+ private void grow(int minCapacity) {
+ // overflow-conscious code
+ int oldCapacity = buf.length;
+ int newCapacity = oldCapacity << 1;
+ if (newCapacity - minCapacity < 0) {
+ newCapacity = minCapacity;
+ }
+ if (newCapacity < 0) {
+ if (minCapacity < 0) { // overflow
+ throw new OutOfMemoryError();
+ }
+ newCapacity = Integer.MAX_VALUE;
+ }
+ buf = Arrays.copyOf(buf, newCapacity);
+ }
+
+ /**
+ * Writes the specified byte to this byte array output stream.
+ *
+ * @param b the byte to be written.
+ */
+ public void write(int b) {
+ ensureCapacity(count + 1);
+ buf[count] = (byte) b;
+ count += 1;
+ }
+
+ /**
+ * Writes <code>len</code> bytes from the specified byte array
+ * starting at offset <code>off</code> to this byte array output stream.
+ *
+ * @param b the data.
+ * @param off the start offset in the data.
+ * @param len the number of bytes to write.
+ */
+ public void write(byte[] b, int off, int len) {
+ if ((off < 0) || (off > b.length) || (len < 0) ||
+ ((off + len) - b.length > 0)) {
+ throw new IndexOutOfBoundsException();
+ }
+ ensureCapacity(count + len);
+ System.arraycopy(b, off, buf, count, len);
+ count += len;
+ }
+
+ /**
+ * Writes the complete contents of this byte array output stream to
+ * the specified output stream argument, as if by calling the output
+ * stream's write method using <code>out.write(buf, 0, count)</code>.
+ *
+ * @param out the output stream to which to write the data.
+ * @throws IOException if an I/O error occurs.
+ */
+ public void writeTo(OutputStream out) throws IOException {
+ out.write(buf, 0, count);
+ }
+
+ /**
+ * Resets the <code>count</code> field of this byte array output
+ * stream to zero, so that all currently accumulated output in the
+ * output stream is discarded. The output stream can be used again,
+ * reusing the already allocated buffer space.
+ *
+ * @see java.io.ByteArrayInputStream#count
+ */
+ public void reset() {
+ count = 0;
+ }
+
+ /**
+ * Creates a newly allocated byte array. Its size is the current
+ * size of this output stream and the valid contents of the buffer
+ * have been copied into it.
+ *
+ * @return the current contents of this output stream, as a byte array.
+ * @see java.io.ByteArrayOutputStream#size()
+ */
+ public byte toByteArray()[] {
+ return Arrays.copyOf(buf, count);
+ }
+
+ /**
+ * Returns the current size of the buffer.
+ *
+ * @return the value of the <code>count</code> field, which is the number
+ * of valid bytes in this output stream.
+ * @see java.io.ByteArrayOutputStream#count
+ */
+ public int size() {
+ return count;
+ }
+
+ /**
+ * Converts the buffer's contents into a string decoding bytes using the
+ * platform's default character set. The length of the new <tt>String</tt>
+ * is a function of the character set, and hence may not be equal to the
+ * size of the buffer.
+ * <p>
+ * <p> This method always replaces malformed-input and unmappable-character
+ * sequences with the default replacement string for the platform's
+ * default character set. The {@linkplain java.nio.charset.CharsetDecoder}
+ * class should be used when more control over the decoding process is
+ * required.
+ *
+ * @return String decoded from the buffer's contents.
+ * @since JDK1.1
+ */
+ public String toString() {
+ return new String(buf, 0, count);
+ }
+
+ /**
+ * Converts the buffer's contents into a string by decoding the bytes using
+ * the named {@link java.nio.charset.Charset charset}. The length of the new
+ * <tt>String</tt> is a function of the charset, and hence may not be equal
+ * to the length of the byte array.
+ * <p>
+ * <p> This method always replaces malformed-input and unmappable-character
+ * sequences with this charset's default replacement string. The {@link
+ * java.nio.charset.CharsetDecoder} class should be used when more control
+ * over the decoding process is required.
+ *
+ * @param charsetName the name of a supported
+ * {@link java.nio.charset.Charset charset}
+ * @return String decoded from the buffer's contents.
+ * @throws UnsupportedEncodingException If the named charset is not supported
+ * @since JDK1.1
+ */
+ public String toString(String charsetName)
+ throws UnsupportedEncodingException {
+ return new String(buf, 0, count, charsetName);
+ }
+
+ /**
+ * Creates a newly allocated string. Its size is the current size of
+ * the output stream and the valid contents of the buffer have been
+ * copied into it. Each character <i>c</i> in the resulting string is
+ * constructed from the corresponding element <i>b</i> in the byte
+ * array such that:
+ * <blockquote><pre>
+ * c == (char)(((hibyte & 0xff) << 8) | (b & 0xff))
+ * </pre></blockquote>
+ *
+ * @param hibyte the high byte of each resulting Unicode character.
+ * @return the current contents of the output stream, as a string.
+ * @see java.io.ByteArrayOutputStream#size()
+ * @see java.io.ByteArrayOutputStream#toString(String)
+ * @see java.io.ByteArrayOutputStream#toString()
+ * @deprecated This method does not properly convert bytes into characters.
+ * As of JDK 1.1, the preferred way to do this is via the
+ * <code>toString(String enc)</code> method, which takes an encoding-name
+ * argument, or the <code>toString()</code> method, which uses the
+ * platform's default character encoding.
+ */
+ @Deprecated
+ public String toString(int hibyte) {
+ return new String(buf, hibyte, 0, count);
+ }
+
+ /**
+ * Closing a <tt>ByteArrayOutputStream</tt> has no effect. The methods in
+ * this class can be called after the stream has been closed without
+ * generating an <tt>IOException</tt>.
+ */
+ public void close() throws IOException {
+ }
+
+ /**
+ * Returns the read/write offset position for the stream.
+ * @return the current position in the stream.
+ */
+ public int getPosition() {
+ return count;
+ }
+
+ /**
+ * Sets the read/write offset position for the stream.
+ *
+ * @param position the position to which the offset in the stream shall be set. Must be < getEndPosition
+ */
+ public void setPosition(int position) {
+ Preconditions.checkArgument(position < getEndPosition(), "Position out of bounds.");
+ count = position;
+ }
+
+ /**
+ * Returns the size of the internal buffer, which is the current end position for all setPosition calls.
+ * @return size of the internal buffer
+ */
+ public int getEndPosition() {
+ return buf.length;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/test/java/org/apache/flink/api/common/functions/util/RuntimeUDFContextTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/api/common/functions/util/RuntimeUDFContextTest.java b/flink-core/src/test/java/org/apache/flink/api/common/functions/util/RuntimeUDFContextTest.java
index 4cd2a64..7c5878d 100644
--- a/flink-core/src/test/java/org/apache/flink/api/common/functions/util/RuntimeUDFContextTest.java
+++ b/flink-core/src/test/java/org/apache/flink/api/common/functions/util/RuntimeUDFContextTest.java
@@ -38,7 +38,7 @@ import org.junit.Test;
public class RuntimeUDFContextTest {
- private final TaskInfo taskInfo = new TaskInfo("test name", 1, 3, 0);
+ private final TaskInfo taskInfo = new TaskInfo("test name", 3, 1, 3, 0);
@Test
public void testBroadcastVariableNotFound() {
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/test/java/org/apache/flink/api/common/io/RichInputFormatTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/api/common/io/RichInputFormatTest.java b/flink-core/src/test/java/org/apache/flink/api/common/io/RichInputFormatTest.java
index c3cbb58..fc3fb1a 100644
--- a/flink-core/src/test/java/org/apache/flink/api/common/io/RichInputFormatTest.java
+++ b/flink-core/src/test/java/org/apache/flink/api/common/io/RichInputFormatTest.java
@@ -41,7 +41,7 @@ public class RichInputFormatTest {
@Test
public void testCheckRuntimeContextAccess() {
final SerializedInputFormat<Value> inputFormat = new SerializedInputFormat<Value>();
- final TaskInfo taskInfo = new TaskInfo("test name", 1, 3, 0);
+ final TaskInfo taskInfo = new TaskInfo("test name", 3, 1, 3, 0);
inputFormat.setRuntimeContext(
new RuntimeUDFContext(
taskInfo, getClass().getClassLoader(), new ExecutionConfig(),
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/test/java/org/apache/flink/api/common/io/RichOutputFormatTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/api/common/io/RichOutputFormatTest.java b/flink-core/src/test/java/org/apache/flink/api/common/io/RichOutputFormatTest.java
index 4c303a6..95f8497 100644
--- a/flink-core/src/test/java/org/apache/flink/api/common/io/RichOutputFormatTest.java
+++ b/flink-core/src/test/java/org/apache/flink/api/common/io/RichOutputFormatTest.java
@@ -41,7 +41,7 @@ public class RichOutputFormatTest {
@Test
public void testCheckRuntimeContextAccess() {
final SerializedOutputFormat<Value> inputFormat = new SerializedOutputFormat<Value>();
- final TaskInfo taskInfo = new TaskInfo("test name", 1, 3, 0);
+ final TaskInfo taskInfo = new TaskInfo("test name", 3, 1, 3, 0);
inputFormat.setRuntimeContext(new RuntimeUDFContext(
taskInfo, getClass().getClassLoader(), new ExecutionConfig(),
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSinkBaseTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSinkBaseTest.java b/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSinkBaseTest.java
index 71bb102..b952c58 100644
--- a/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSinkBaseTest.java
+++ b/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSinkBaseTest.java
@@ -93,7 +93,7 @@ public class GenericDataSinkBaseTest implements java.io.Serializable {
ExecutionConfig executionConfig = new ExecutionConfig();
final HashMap<String, Accumulator<?, ?>> accumulatorMap = new HashMap<String, Accumulator<?, ?>>();
final HashMap<String, Future<Path>> cpTasks = new HashMap<>();
- final TaskInfo taskInfo = new TaskInfo("test_sink", 0, 1, 0);
+ final TaskInfo taskInfo = new TaskInfo("test_sink", 1, 0, 1, 0);
executionConfig.disableObjectReuse();
in.reset();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSourceBaseTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSourceBaseTest.java b/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSourceBaseTest.java
index 2dabe48..9a2b877 100644
--- a/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSourceBaseTest.java
+++ b/flink-core/src/test/java/org/apache/flink/api/common/operators/GenericDataSourceBaseTest.java
@@ -79,7 +79,7 @@ public class GenericDataSourceBaseTest implements java.io.Serializable {
final HashMap<String, Accumulator<?, ?>> accumulatorMap = new HashMap<String, Accumulator<?, ?>>();
final HashMap<String, Future<Path>> cpTasks = new HashMap<>();
- final TaskInfo taskInfo = new TaskInfo("test_source", 0, 1, 0);
+ final TaskInfo taskInfo = new TaskInfo("test_source", 1, 0, 1, 0);
ExecutionConfig executionConfig = new ExecutionConfig();
executionConfig.disableObjectReuse();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/test/java/org/apache/flink/api/common/operators/base/FlatMapOperatorCollectionTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/api/common/operators/base/FlatMapOperatorCollectionTest.java b/flink-core/src/test/java/org/apache/flink/api/common/operators/base/FlatMapOperatorCollectionTest.java
index f125c4b..232f510 100644
--- a/flink-core/src/test/java/org/apache/flink/api/common/operators/base/FlatMapOperatorCollectionTest.java
+++ b/flink-core/src/test/java/org/apache/flink/api/common/operators/base/FlatMapOperatorCollectionTest.java
@@ -77,7 +77,7 @@ public class FlatMapOperatorCollectionTest implements Serializable {
} else {
executionConfig.enableObjectReuse();
}
- final TaskInfo taskInfo = new TaskInfo("Test UDF", 0, 4, 0);
+ final TaskInfo taskInfo = new TaskInfo("Test UDF", 4, 0, 4, 0);
// run on collections
final List<String> result = getTestFlatMapOperator(udf)
.executeOnCollections(input,
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java b/flink-core/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java
index 8befcb9..72f2f2e 100644
--- a/flink-core/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java
+++ b/flink-core/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java
@@ -121,7 +121,7 @@ public class InnerJoinOperatorBaseTest implements Serializable {
try {
- final TaskInfo taskInfo = new TaskInfo(taskName, 0, 1, 0);
+ final TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
final HashMap<String, Accumulator<?, ?>> accumulatorMap = new HashMap<String, Accumulator<?, ?>>();
final HashMap<String, Future<Path>> cpTasks = new HashMap<>();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/test/java/org/apache/flink/api/common/operators/base/MapOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/api/common/operators/base/MapOperatorTest.java b/flink-core/src/test/java/org/apache/flink/api/common/operators/base/MapOperatorTest.java
index d79e2a5..1b2af7c 100644
--- a/flink-core/src/test/java/org/apache/flink/api/common/operators/base/MapOperatorTest.java
+++ b/flink-core/src/test/java/org/apache/flink/api/common/operators/base/MapOperatorTest.java
@@ -112,7 +112,7 @@ public class MapOperatorTest implements java.io.Serializable {
List<String> input = new ArrayList<String>(asList("1", "2", "3", "4", "5", "6"));
final HashMap<String, Accumulator<?, ?>> accumulatorMap = new HashMap<String, Accumulator<?, ?>>();
final HashMap<String, Future<Path>> cpTasks = new HashMap<>();
- final TaskInfo taskInfo = new TaskInfo(taskName, 0, 1, 0);
+ final TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
ExecutionConfig executionConfig = new ExecutionConfig();
executionConfig.disableObjectReuse();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-core/src/test/java/org/apache/flink/api/common/operators/base/PartitionMapOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/api/common/operators/base/PartitionMapOperatorTest.java b/flink-core/src/test/java/org/apache/flink/api/common/operators/base/PartitionMapOperatorTest.java
index 83c194a..e709152 100644
--- a/flink-core/src/test/java/org/apache/flink/api/common/operators/base/PartitionMapOperatorTest.java
+++ b/flink-core/src/test/java/org/apache/flink/api/common/operators/base/PartitionMapOperatorTest.java
@@ -83,7 +83,7 @@ public class PartitionMapOperatorTest implements java.io.Serializable {
List<String> input = new ArrayList<String>(asList("1", "2", "3", "4", "5", "6"));
- final TaskInfo taskInfo = new TaskInfo(taskName, 0, 1, 0);
+ final TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
ExecutionConfig executionConfig = new ExecutionConfig();
executionConfig.disableObjectReuse();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java
----------------------------------------------------------------------
diff --git a/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java b/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java
index fd7bf5d..df40998 100644
--- a/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java
+++ b/flink-fs-tests/src/test/java/org/apache/flink/hdfstests/FileStateBackendTest.java
@@ -20,17 +20,16 @@ package org.apache.flink.hdfstests;
import org.apache.commons.io.FileUtils;
+import org.apache.flink.api.common.JobID;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.core.fs.FileStatus;
import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.core.fs.Path;
import org.apache.flink.core.testutils.CommonTestUtils;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
import org.apache.flink.runtime.state.StateBackendTestBase;
import org.apache.flink.runtime.state.filesystem.FileStateHandle;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
-import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
-import org.apache.flink.api.common.typeutils.base.IntSerializer;
-import org.apache.flink.runtime.state.AbstractStateBackend;
import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
import org.apache.hadoop.conf.Configuration;
@@ -51,25 +50,23 @@ import java.util.UUID;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
-
+
private static File TEMP_DIR;
-
+
private static String HDFS_ROOT_URI;
-
+
private static MiniDFSCluster HDFS_CLUSTER;
-
+
private static FileSystem FS;
-
+
// ------------------------------------------------------------------------
// startup / shutdown
// ------------------------------------------------------------------------
-
+
@BeforeClass
public static void createHDFS() {
try {
@@ -82,7 +79,7 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
HDFS_ROOT_URI = "hdfs://" + HDFS_CLUSTER.getURI().getHost() + ":"
+ HDFS_CLUSTER.getNameNodePort() + "/";
-
+
FS = FileSystem.get(new URI(HDFS_ROOT_URI));
}
catch (Exception e) {
@@ -109,11 +106,6 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
}
- @Override
- protected void cleanup() throws Exception {
- FileSystem.get(stateBaseURI).delete(new Path(stateBaseURI), true);
- }
-
// ------------------------------------------------------------------------
// Tests
// ------------------------------------------------------------------------
@@ -132,60 +124,19 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
public void testReducingStateRestoreWithWrongSerializers() {}
@Test
- public void testSetupAndSerialization() {
- try {
- URI baseUri = new URI(HDFS_ROOT_URI + UUID.randomUUID().toString());
-
- FsStateBackend originalBackend = new FsStateBackend(baseUri);
-
- assertFalse(originalBackend.isInitialized());
- assertEquals(baseUri, originalBackend.getBasePath().toUri());
- assertNull(originalBackend.getCheckpointDirectory());
-
- // serialize / copy the backend
- FsStateBackend backend = CommonTestUtils.createCopySerializable(originalBackend);
- assertFalse(backend.isInitialized());
- assertEquals(baseUri, backend.getBasePath().toUri());
- assertNull(backend.getCheckpointDirectory());
-
- // no file operations should be possible right now
- try {
- FsStateBackend.FsCheckpointStateOutputStream out = backend.createCheckpointStateOutputStream(
- 2L,
- System.currentTimeMillis());
-
- out.write(1);
- out.closeAndGetHandle();
- fail("should fail with an exception");
- } catch (IllegalStateException e) {
- // supreme!
- }
-
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "dummy", IntSerializer.INSTANCE);
- assertNotNull(backend.getCheckpointDirectory());
-
- Path checkpointDir = backend.getCheckpointDirectory();
- assertTrue(FS.exists(checkpointDir));
- assertTrue(isDirectoryEmpty(checkpointDir));
+ public void testStateOutputStream() {
+ URI basePath = randomHdfsFileUri();
- backend.disposeAllStateForCurrentJob();
- assertNull(backend.getCheckpointDirectory());
+ try {
+ FsStateBackend backend = CommonTestUtils.createCopySerializable(new FsStateBackend(basePath, 15));
+ JobID jobId = new JobID();
- assertTrue(isDirectoryEmpty(baseUri));
- }
- catch (Exception e) {
- e.printStackTrace();
- fail(e.getMessage());
- }
- }
- @Test
- public void testStateOutputStream() {
- try {
- FsStateBackend backend = CommonTestUtils.createCopySerializable(new FsStateBackend(randomHdfsFileUri(), 15));
- backend.initializeForJob(new DummyEnvironment("test", 1, 0), "dummy", IntSerializer.INSTANCE);
+ CheckpointStreamFactory streamFactory = backend.createStreamFactory(jobId, "test_op");
- Path checkpointDir = backend.getCheckpointDirectory();
+ // we know how FsCheckpointStreamFactory is implemented so we know where it
+ // will store checkpoints
+ Path checkpointPath = new Path(new Path(basePath), jobId.toString());
byte[] state1 = new byte[1274673];
byte[] state2 = new byte[1];
@@ -200,12 +151,12 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
long checkpointId = 97231523452L;
- FsStateBackend.FsCheckpointStateOutputStream stream1 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
- FsStateBackend.FsCheckpointStateOutputStream stream2 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
- FsStateBackend.FsCheckpointStateOutputStream stream3 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
+ CheckpointStreamFactory.CheckpointStateOutputStream stream1 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
+ CheckpointStreamFactory.CheckpointStateOutputStream stream2 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
+ CheckpointStreamFactory.CheckpointStateOutputStream stream3 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
stream1.write(state1);
stream2.write(state2);
@@ -217,15 +168,15 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
// use with try-with-resources
FileStateHandle handle4;
- try (AbstractStateBackend.CheckpointStateOutputStream stream4 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis())) {
+ try (CheckpointStreamFactory.CheckpointStateOutputStream stream4 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis())) {
stream4.write(state4);
handle4 = (FileStateHandle) stream4.closeAndGetHandle();
}
// close before accessing handle
- AbstractStateBackend.CheckpointStateOutputStream stream5 =
- backend.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
+ CheckpointStreamFactory.CheckpointStateOutputStream stream5 =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, System.currentTimeMillis());
stream5.write(state4);
stream5.close();
try {
@@ -237,7 +188,7 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
validateBytesInStream(handle1.openInputStream(), state1);
handle1.discardState();
- assertFalse(isDirectoryEmpty(checkpointDir));
+ assertFalse(isDirectoryEmpty(checkpointPath));
ensureFileDeleted(handle1.getFilePath());
validateBytesInStream(handle2.openInputStream(), state2);
@@ -248,7 +199,7 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
validateBytesInStream(handle4.openInputStream(), state4);
handle4.discardState();
- assertTrue(isDirectoryEmpty(checkpointDir));
+ assertTrue(isDirectoryEmpty(checkpointPath));
}
catch (Exception e) {
e.printStackTrace();
@@ -270,7 +221,7 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
private static boolean isDirectoryEmpty(URI directory) {
return isDirectoryEmpty(new Path(directory));
}
-
+
private static boolean isDirectoryEmpty(Path directory) {
try {
FileStatus[] nested = FS.listStatus(directory);
@@ -293,14 +244,14 @@ public class FileStateBackendTest extends StateBackendTestBase<FsStateBackend> {
private static void validateBytesInStream(InputStream is, byte[] data) throws IOException {
byte[] holder = new byte[data.length];
-
+
int pos = 0;
int read;
while (pos < holder.length && (read = is.read(holder, pos, holder.length - pos)) != -1) {
pos += read;
}
-
- assertEquals("not enough data", holder.length, pos);
+
+ assertEquals("not enough data", holder.length, pos);
assertEquals("too much data", -1, is.read());
assertArrayEquals("wrong data", data, holder);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-java/src/test/java/org/apache/flink/api/common/operators/base/CoGroupOperatorCollectionTest.java
----------------------------------------------------------------------
diff --git a/flink-java/src/test/java/org/apache/flink/api/common/operators/base/CoGroupOperatorCollectionTest.java b/flink-java/src/test/java/org/apache/flink/api/common/operators/base/CoGroupOperatorCollectionTest.java
index 2682584..a4426e0 100644
--- a/flink-java/src/test/java/org/apache/flink/api/common/operators/base/CoGroupOperatorCollectionTest.java
+++ b/flink-java/src/test/java/org/apache/flink/api/common/operators/base/CoGroupOperatorCollectionTest.java
@@ -77,7 +77,7 @@ public class CoGroupOperatorCollectionTest implements Serializable {
ExecutionConfig executionConfig = new ExecutionConfig();
final HashMap<String, Accumulator<?, ?>> accumulators = new HashMap<String, Accumulator<?, ?>>();
final HashMap<String, Future<Path>> cpTasks = new HashMap<>();
- final TaskInfo taskInfo = new TaskInfo("Test UDF", 0, 4, 0);
+ final TaskInfo taskInfo = new TaskInfo("Test UDF", 4, 0, 4, 0);
final RuntimeContext ctx = new RuntimeUDFContext(
taskInfo, null, executionConfig, cpTasks, accumulators, new UnregisteredMetricsGroup());
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-java/src/test/java/org/apache/flink/api/common/operators/base/GroupReduceOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-java/src/test/java/org/apache/flink/api/common/operators/base/GroupReduceOperatorTest.java b/flink-java/src/test/java/org/apache/flink/api/common/operators/base/GroupReduceOperatorTest.java
index c5a247a..d0784a8 100644
--- a/flink-java/src/test/java/org/apache/flink/api/common/operators/base/GroupReduceOperatorTest.java
+++ b/flink-java/src/test/java/org/apache/flink/api/common/operators/base/GroupReduceOperatorTest.java
@@ -165,7 +165,7 @@ public class GroupReduceOperatorTest implements java.io.Serializable {
Integer>("foo", 3), new Tuple2<String, Integer>("bar", 2), new Tuple2<String,
Integer>("bar", 4)));
- final TaskInfo taskInfo = new TaskInfo(taskName, 0, 1, 0);
+ final TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
ExecutionConfig executionConfig = new ExecutionConfig();
executionConfig.disableObjectReuse();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-java/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java
----------------------------------------------------------------------
diff --git a/flink-java/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java b/flink-java/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java
index 89574a8..ef33ac0 100644
--- a/flink-java/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java
+++ b/flink-java/src/test/java/org/apache/flink/api/common/operators/base/InnerJoinOperatorBaseTest.java
@@ -107,7 +107,7 @@ public class InnerJoinOperatorBaseTest implements Serializable {
));
try {
- final TaskInfo taskInfo = new TaskInfo("op", 0, 1, 0);
+ final TaskInfo taskInfo = new TaskInfo("op", 1, 0, 1, 0);
ExecutionConfig executionConfig = new ExecutionConfig();
executionConfig.disableObjectReuse();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-java/src/test/java/org/apache/flink/api/common/operators/base/ReduceOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-java/src/test/java/org/apache/flink/api/common/operators/base/ReduceOperatorTest.java b/flink-java/src/test/java/org/apache/flink/api/common/operators/base/ReduceOperatorTest.java
index 150854d..9427d6f 100644
--- a/flink-java/src/test/java/org/apache/flink/api/common/operators/base/ReduceOperatorTest.java
+++ b/flink-java/src/test/java/org/apache/flink/api/common/operators/base/ReduceOperatorTest.java
@@ -145,7 +145,7 @@ public class ReduceOperatorTest implements java.io.Serializable {
Integer>("foo", 3), new Tuple2<String, Integer>("bar", 2), new Tuple2<String,
Integer>("bar", 4)));
- final TaskInfo taskInfo = new TaskInfo(taskName, 0, 1, 0);
+ final TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
ExecutionConfig executionConfig = new ExecutionConfig();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/nfa/NFA.java
----------------------------------------------------------------------
diff --git a/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/nfa/NFA.java b/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/nfa/NFA.java
index 624db0d..5ac638e 100644
--- a/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/nfa/NFA.java
+++ b/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/nfa/NFA.java
@@ -520,7 +520,7 @@ public class NFA<T> implements Serializable {
public void serialize(NFA<T> record, DataOutputView target) throws IOException {
ObjectOutputStream oos = new ObjectOutputStream(new DataOutputViewStream(target));
oos.writeObject(record);
- oos.close();
+ oos.flush();
}
@Override
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java
----------------------------------------------------------------------
diff --git a/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java b/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java
index e3f924c..09773a2 100644
--- a/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java
+++ b/flink-libraries/flink-cep/src/main/java/org/apache/flink/cep/operator/AbstractKeyedCEPPatternOperator.java
@@ -92,6 +92,8 @@ abstract public class AbstractKeyedCEPPatternOperator<IN, KEY, OUT> extends Abst
@Override
@SuppressWarnings("unchecked")
public void open() throws Exception {
+ super.open();
+
if (keys == null) {
keys = new HashSet<>();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java b/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
index 54c1477..52a02d1 100644
--- a/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
+++ b/flink-libraries/flink-cep/src/test/java/org/apache/flink/cep/operator/CEPOperatorTest.java
@@ -33,6 +33,7 @@ import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.util.TestLogger;
import org.junit.Rule;
@@ -83,16 +84,15 @@ public class CEPOperatorTest extends TestLogger {
}
};
- OneInputStreamOperatorTestHarness<Event, Map<String, Event>> harness = new OneInputStreamOperatorTestHarness<>(
- new KeyedCEPPatternOperator<>(
- Event.createTypeSerializer(),
- false,
+ OneInputStreamOperatorTestHarness<Event, Map<String, Event>> harness = new KeyedOneInputStreamOperatorTestHarness<>(
+ new KeyedCEPPatternOperator<>(
+ Event.createTypeSerializer(),
+ false,
+ keySelector,
+ IntSerializer.INSTANCE,
+ new NFAFactory()),
keySelector,
- IntSerializer.INSTANCE,
- new NFAFactory())
- );
-
- harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
+ BasicTypeInfo.INT_TYPE_INFO);
harness.open();
@@ -206,15 +206,15 @@ public class CEPOperatorTest extends TestLogger {
}
};
- OneInputStreamOperatorTestHarness<Event, Map<String, Event>> harness = new OneInputStreamOperatorTestHarness<>(
+ OneInputStreamOperatorTestHarness<Event, Map<String, Event>> harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
Event.createTypeSerializer(),
false,
keySelector,
IntSerializer.INSTANCE,
- new NFAFactory()));
-
- harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
+ new NFAFactory()),
+ keySelector,
+ BasicTypeInfo.INT_TYPE_INFO);
harness.open();
@@ -228,15 +228,16 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with some elements in internal sorting queue
StreamStateHandle snapshot = harness.snapshot(0, 0);
- harness = new OneInputStreamOperatorTestHarness<>(
+ harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
Event.createTypeSerializer(),
false,
keySelector,
IntSerializer.INSTANCE,
- new NFAFactory()));
+ new NFAFactory()),
+ keySelector,
+ BasicTypeInfo.INT_TYPE_INFO);
- harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
harness.setup();
harness.restore(snapshot);
harness.open();
@@ -252,15 +253,16 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with empty element queue but NFA state
StreamStateHandle snapshot2 = harness.snapshot(1, 1);
- harness = new OneInputStreamOperatorTestHarness<>(
+ harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
Event.createTypeSerializer(),
false,
keySelector,
IntSerializer.INSTANCE,
- new NFAFactory()));
+ new NFAFactory()),
+ keySelector,
+ BasicTypeInfo.INT_TYPE_INFO);
- harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
harness.setup();
harness.restore(snapshot2);
harness.open();
@@ -309,16 +311,17 @@ public class CEPOperatorTest extends TestLogger {
}
};
- OneInputStreamOperatorTestHarness<Event, Map<String, Event>> harness = new OneInputStreamOperatorTestHarness<>(
+ OneInputStreamOperatorTestHarness<Event, Map<String, Event>> harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
Event.createTypeSerializer(),
false,
keySelector,
IntSerializer.INSTANCE,
- new NFAFactory()));
+ new NFAFactory()),
+ keySelector,
+ BasicTypeInfo.INT_TYPE_INFO);
harness.setStateBackend(rocksDBStateBackend);
- harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
harness.open();
@@ -332,19 +335,21 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with some elements in internal sorting queue
StreamStateHandle snapshot = harness.snapshot(0, 0);
- harness = new OneInputStreamOperatorTestHarness<>(
+ harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
Event.createTypeSerializer(),
false,
keySelector,
IntSerializer.INSTANCE,
- new NFAFactory()));
+ new NFAFactory()),
+ keySelector,
+ BasicTypeInfo.INT_TYPE_INFO);
rocksDBStateBackend =
new RocksDBStateBackend(rocksDbBackups, new MemoryStateBackend());
rocksDBStateBackend.setDbStoragePath(rocksDbPath);
harness.setStateBackend(rocksDBStateBackend);
- harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
+
harness.setup();
harness.restore(snapshot);
harness.open();
@@ -360,19 +365,20 @@ public class CEPOperatorTest extends TestLogger {
// simulate snapshot/restore with empty element queue but NFA state
StreamStateHandle snapshot2 = harness.snapshot(1, 1);
- harness = new OneInputStreamOperatorTestHarness<>(
+ harness = new KeyedOneInputStreamOperatorTestHarness<>(
new KeyedCEPPatternOperator<>(
Event.createTypeSerializer(),
false,
keySelector,
IntSerializer.INSTANCE,
- new NFAFactory()));
+ new NFAFactory()),
+ keySelector,
+ BasicTypeInfo.INT_TYPE_INFO);
rocksDBStateBackend =
new RocksDBStateBackend(rocksDbBackups, new MemoryStateBackend());
rocksDBStateBackend.setDbStoragePath(rocksDbPath);
harness.setStateBackend(rocksDBStateBackend);
- harness.configureForKeyedStream(keySelector, BasicTypeInfo.INT_TYPE_INFO);
harness.setup();
harness.restore(snapshot2);
harness.open();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
index e78e203..e751e08 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
@@ -46,6 +46,7 @@ import scala.concurrent.Future;
import java.util.ArrayDeque;
import java.util.ArrayList;
+import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
@@ -821,14 +822,10 @@ public class CheckpointCoordinator {
}
KeyGroupRange subtaskKeyGroupIds = keyGroupPartitions.get(i);
- List<KeyGroupsStateHandle> subtaskKeyGroupStates = new ArrayList<>();
- for (KeyGroupsStateHandle storedKeyGroup : taskState.getKeyGroupStates()) {
- KeyGroupsStateHandle intersection = storedKeyGroup.getKeyGroupIntersection(subtaskKeyGroupIds);
- if(intersection.getNumberOfKeyGroups() > 0) {
- subtaskKeyGroupStates.add(intersection);
- }
- }
+ List<KeyGroupsStateHandle> subtaskKeyGroupStates = getKeyGroupsStateHandles(
+ taskState.getKeyGroupStates(),
+ subtaskKeyGroupIds);
Execution currentExecutionAttempt = executionJobVertex
.getTaskVertices()[i]
@@ -852,6 +849,27 @@ public class CheckpointCoordinator {
}
/**
+ * Determine the subset of {@link KeyGroupsStateHandle KeyGroupsStateHandles} with correct
+ * key group index for the given subtask {@link KeyGroupRange}.
+ *
+ * <p>This is publicly visible to be used in tests.
+ */
+ public static List<KeyGroupsStateHandle> getKeyGroupsStateHandles(
+ Collection<KeyGroupsStateHandle> allKeyGroupsHandles,
+ KeyGroupRange subtaskKeyGroupIds) {
+
+ List<KeyGroupsStateHandle> subtaskKeyGroupStates = new ArrayList<>();
+
+ for (KeyGroupsStateHandle storedKeyGroup : allKeyGroupsHandles) {
+ KeyGroupsStateHandle intersection = storedKeyGroup.getKeyGroupIntersection(subtaskKeyGroupIds);
+ if(intersection.getNumberOfKeyGroups() > 0) {
+ subtaskKeyGroupStates.add(intersection);
+ }
+ }
+ return subtaskKeyGroupStates;
+ }
+
+ /**
* Groups the available set of key groups into key group partitions. A key group partition is
* the set of key groups which is assigned to the same task. Each set of the returned list
* constitutes a key group partition.
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
index 8849e93..ca976e4 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
@@ -59,6 +59,9 @@ public final class TaskDeploymentDescriptor implements Serializable {
/** The task's name. */
private final String taskName;
+ /** The number of key groups aka the max parallelism aka the max number of subtasks. */
+ private final int numberOfKeyGroups;
+
/** The task's index in the subtask group. */
private final int indexInSubtaskGroup;
@@ -110,6 +113,7 @@ public final class TaskDeploymentDescriptor implements Serializable {
ExecutionAttemptID executionId,
SerializedValue<ExecutionConfig> serializedExecutionConfig,
String taskName,
+ int numberOfKeyGroups,
int indexInSubtaskGroup,
int numberOfSubtasks,
int attemptNumber,
@@ -135,6 +139,7 @@ public final class TaskDeploymentDescriptor implements Serializable {
this.executionId = checkNotNull(executionId);
this.serializedExecutionConfig = checkNotNull(serializedExecutionConfig);
this.taskName = checkNotNull(taskName);
+ this.numberOfKeyGroups = numberOfKeyGroups;
this.indexInSubtaskGroup = indexInSubtaskGroup;
this.numberOfSubtasks = numberOfSubtasks;
this.attemptNumber = attemptNumber;
@@ -157,6 +162,7 @@ public final class TaskDeploymentDescriptor implements Serializable {
ExecutionAttemptID executionId,
SerializedValue<ExecutionConfig> serializedExecutionConfig,
String taskName,
+ int numberOfKeyGroups,
int indexInSubtaskGroup,
int numberOfSubtasks,
int attemptNumber,
@@ -176,6 +182,7 @@ public final class TaskDeploymentDescriptor implements Serializable {
executionId,
serializedExecutionConfig,
taskName,
+ numberOfKeyGroups,
indexInSubtaskGroup,
numberOfSubtasks,
attemptNumber,
@@ -227,6 +234,13 @@ public final class TaskDeploymentDescriptor implements Serializable {
}
/**
+ * Returns the task's number of key groups.
+ */
+ public int getNumberOfKeyGroups() {
+ return numberOfKeyGroups;
+ }
+
+ /**
* Returns the task's index in the subtask group.
*
* @return the task's index in the subtask group
@@ -253,7 +267,7 @@ public final class TaskDeploymentDescriptor implements Serializable {
* Returns the {@link TaskInfo} object for the subtask
*/
public TaskInfo getTaskInfo() {
- return new TaskInfo(taskName, indexInSubtaskGroup, numberOfSubtasks, attemptNumber);
+ return new TaskInfo(taskName, numberOfKeyGroups, indexInSubtaskGroup, numberOfSubtasks, attemptNumber);
}
/**
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
index f3a8b6d..b215394 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
@@ -682,6 +682,7 @@ public class ExecutionVertex {
executionId,
serializedConfig,
getTaskName(),
+ getMaxParallelism(),
subTaskIndex,
getTotalNumberOfParallelSubtasks(),
attemptNumber,
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-runtime/src/main/java/org/apache/flink/runtime/fs/hdfs/HadoopDataOutputStream.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/fs/hdfs/HadoopDataOutputStream.java b/flink-runtime/src/main/java/org/apache/flink/runtime/fs/hdfs/HadoopDataOutputStream.java
index b0ff4b3..d6fbc19 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/fs/hdfs/HadoopDataOutputStream.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/fs/hdfs/HadoopDataOutputStream.java
@@ -51,6 +51,11 @@ public class HadoopDataOutputStream extends FSDataOutputStream {
}
@Override
+ public long getPos() throws IOException {
+ return fdos.getPos();
+ }
+
+ @Override
public void flush() throws IOException {
if (HFLUSH_METHOD != null) {
try {
[22/27] flink git commit: [FLINK-4380] Remove KeyGroupAssigner in
favor of static method/Have default max. parallelism at 128
Posted by al...@apache.org.
[FLINK-4380] Remove KeyGroupAssigner in favor of static method/Have default max. parallelism at 128
Project: http://git-wip-us.apache.org/repos/asf/flink/repo
Commit: http://git-wip-us.apache.org/repos/asf/flink/commit/6d430618
Tree: http://git-wip-us.apache.org/repos/asf/flink/tree/6d430618
Diff: http://git-wip-us.apache.org/repos/asf/flink/diff/6d430618
Branch: refs/heads/master
Commit: 6d4306186e09be6f2557600ed7a853c33d3ae6bd
Parents: 2b7a8d6
Author: Stefan Richter <s....@data-artisans.com>
Authored: Mon Aug 29 11:53:22 2016 +0200
Committer: Aljoscha Krettek <al...@gmail.com>
Committed: Wed Aug 31 19:10:01 2016 +0200
----------------------------------------------------------------------
.../streaming/state/AbstractRocksDBState.java | 3 +-
.../state/RocksDBKeyedStateBackend.java | 9 +-
.../streaming/state/RocksDBStateBackend.java | 9 +-
.../state/RocksDBStateBackendConfigTest.java | 9 +-
.../api/common/state/KeyGroupAssigner.java | 53 -----
.../java/org/apache/flink/util/MathUtils.java | 15 ++
.../org/apache/flink/util/MathUtilTest.java | 31 +++
.../checkpoint/CheckpointCoordinator.java | 3 +-
.../flink/runtime/jobgraph/JobVertex.java | 3 +-
.../runtime/state/AbstractStateBackend.java | 5 +-
.../runtime/state/HashKeyGroupAssigner.java | 66 ------
.../flink/runtime/state/KeyGroupRange.java | 42 ----
.../runtime/state/KeyGroupRangeAssignment.java | 97 +++++++++
.../flink/runtime/state/KeyedStateBackend.java | 17 +-
.../state/filesystem/FsStateBackend.java | 9 +-
.../runtime/state/heap/AbstractHeapState.java | 4 +-
.../state/heap/HeapKeyedStateBackend.java | 9 +-
.../flink/runtime/state/heap/HeapListState.java | 3 +-
.../state/memory/MemoryStateBackend.java | 9 +-
.../checkpoint/CheckpointCoordinatorTest.java | 3 +-
.../runtime/query/QueryableStateClientTest.java | 4 +-
.../runtime/query/netty/KvStateClientTest.java | 5 +-
.../query/netty/KvStateServerHandlerTest.java | 16 +-
.../runtime/query/netty/KvStateServerTest.java | 6 +-
.../runtime/state/StateBackendTestBase.java | 49 ++---
.../streaming/api/datastream/KeyedStream.java | 5 +-
.../flink/streaming/api/graph/StreamConfig.java | 22 +-
.../api/graph/StreamGraphGenerator.java | 24 ++-
.../api/graph/StreamingJobGraphGenerator.java | 9 +-
.../api/operators/AbstractStreamOperator.java | 5 +-
.../partitioner/KeyGroupStreamPartitioner.java | 24 +--
.../streaming/runtime/tasks/StreamTask.java | 7 +-
.../api/graph/StreamGraphGeneratorTest.java | 26 +--
.../graph/StreamingJobGraphGeneratorTest.java | 200 -------------------
.../operators/StreamingRuntimeContextTest.java | 5 +-
...AlignedProcessingTimeWindowOperatorTest.java | 19 +-
.../KeyGroupStreamPartitionerTest.java | 3 +-
.../tasks/OneInputStreamTaskTestHarness.java | 3 +-
.../KeyedOneInputStreamOperatorTestHarness.java | 16 +-
.../util/OneInputStreamOperatorTestHarness.java | 4 -
.../test/checkpointing/RescalingITCase.java | 27 +--
.../streaming/runtime/StateBackendITCase.java | 5 +-
42 files changed, 297 insertions(+), 586 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
index cbc2757..e878ad5 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/AbstractRocksDBState.java
@@ -25,6 +25,7 @@ import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
import org.apache.flink.core.memory.DataOutputView;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.state.KvState;
import org.apache.flink.util.Preconditions;
import org.rocksdb.ColumnFamilyHandle;
@@ -130,7 +131,7 @@ public abstract class AbstractRocksDBState<K, N, S extends State, SD extends Sta
backend.getKeySerializer(),
namespaceSerializer);
- int keyGroup = backend.getKeyGroupAssigner().getKeyGroupIndex(des.f0);
+ int keyGroup = KeyGroupRangeAssignment.assignToKeyGroup(des.f0, backend.getNumberOfKeyGroups());
writeKeyWithGroupAndNamespace(keyGroup, des.f0, des.f1);
return backend.db.get(columnFamily, keySerializationStream.toByteArray());
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
index a1634b2..177c09f 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
@@ -21,7 +21,6 @@ import org.apache.commons.io.FileUtils;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.state.FoldingState;
import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ReducingState;
@@ -131,11 +130,11 @@ public class RocksDBKeyedStateBackend<K> extends KeyedStateBackend<K> {
ColumnFamilyOptions columnFamilyOptions,
TaskKvStateRegistry kvStateRegistry,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange
) throws Exception {
- super(kvStateRegistry, keySerializer, keyGroupAssigner, keyGroupRange);
+ super(kvStateRegistry, keySerializer, numberOfKeyGroups, keyGroupRange);
this.operatorIdentifier = operatorIdentifier;
this.jobId = jobId;
@@ -183,7 +182,7 @@ public class RocksDBKeyedStateBackend<K> extends KeyedStateBackend<K> {
ColumnFamilyOptions columnFamilyOptions,
TaskKvStateRegistry kvStateRegistry,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
List<KeyGroupsStateHandle> restoreState
) throws Exception {
@@ -195,7 +194,7 @@ public class RocksDBKeyedStateBackend<K> extends KeyedStateBackend<K> {
columnFamilyOptions,
kvStateRegistry,
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange);
LOG.info("Initializing RocksDB keyed state backend from snapshot.");
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
index f950751..0fdbd5f 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
@@ -18,7 +18,6 @@
package org.apache.flink.contrib.streaming.state;
import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.StateBackend;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.core.fs.Path;
@@ -230,7 +229,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
JobID jobID,
String operatorIdentifier,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
TaskKvStateRegistry kvStateRegistry) throws Exception {
@@ -246,7 +245,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
getColumnOptions(),
kvStateRegistry,
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange);
}
@@ -254,7 +253,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
public <K> KeyedStateBackend<K> restoreKeyedStateBackend(Environment env, JobID jobID,
String operatorIdentifier,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
List<KeyGroupsStateHandle> restoredState,
TaskKvStateRegistry kvStateRegistry) throws Exception {
@@ -270,7 +269,7 @@ public class RocksDBStateBackend extends AbstractStateBackend {
getColumnOptions(),
kvStateRegistry,
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange,
restoredState);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
----------------------------------------------------------------------
diff --git a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
index acf6cb8..3b851be 100644
--- a/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
+++ b/flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
@@ -28,7 +28,6 @@ import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.query.KvStateRegistry;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.util.OperatingSystem;
import org.junit.Assume;
@@ -93,7 +92,7 @@ public class RocksDBStateBackendConfigTest {
env.getJobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ 1,
new KeyGroupRange(0, 0),
env.getTaskKvStateRegistry());
@@ -147,7 +146,7 @@ public class RocksDBStateBackendConfigTest {
env.getJobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ 1,
new KeyGroupRange(0, 0),
env.getTaskKvStateRegistry());
@@ -182,7 +181,7 @@ public class RocksDBStateBackendConfigTest {
env.getJobID(),
"foobar",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ 1,
new KeyGroupRange(0, 0),
new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()));
}
@@ -224,7 +223,7 @@ public class RocksDBStateBackendConfigTest {
env.getJobID(),
"foobar",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ 1,
new KeyGroupRange(0, 0),
new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()));
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java b/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java
deleted file mode 100644
index bb0691e..0000000
--- a/flink-core/src/main/java/org/apache/flink/api/common/state/KeyGroupAssigner.java
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * 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.api.common.state;
-
-import org.apache.flink.annotation.Internal;
-
-import java.io.Serializable;
-
-/**
- * Assigns a key to a key group index. A key group is the smallest unit of partitioned state
- * which is assigned to an operator. An operator can be assigned multiple key groups.
- *
- * @param <K> Type of the key
- */
-@Internal
-public interface KeyGroupAssigner<K> extends Serializable {
- /**
- * Calculates the key group index for the given key.
- *
- * @param key Key to be used
- * @return Key group index for the given key
- */
- int getKeyGroupIndex(K key);
-
- /**
- * Setups the key group assigner with the maximum parallelism (= number of key groups).
- *
- * @param numberOfKeygroups Maximum parallelism (= number of key groups)
- */
- void setup(int numberOfKeygroups);
-
- /**
- *
- * @return configured maximum parallelism
- */
- int getNumberKeyGroups();
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-core/src/main/java/org/apache/flink/util/MathUtils.java
----------------------------------------------------------------------
diff --git a/flink-core/src/main/java/org/apache/flink/util/MathUtils.java b/flink-core/src/main/java/org/apache/flink/util/MathUtils.java
index f40c83a..49056cc 100644
--- a/flink-core/src/main/java/org/apache/flink/util/MathUtils.java
+++ b/flink-core/src/main/java/org/apache/flink/util/MathUtils.java
@@ -155,6 +155,21 @@ public final class MathUtils {
}
}
+ /**
+ * Round the given number to the next power of two
+ * @param x number to round
+ * @return x rounded up to the next power of two
+ */
+ public static int roundUpToPowerOfTwo(int x) {
+ x = x - 1;
+ x |= x >> 1;
+ x |= x >> 2;
+ x |= x >> 4;
+ x |= x >> 8;
+ x |= x >> 16;
+ return x + 1;
+ }
+
// ============================================================================================
/**
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-core/src/test/java/org/apache/flink/util/MathUtilTest.java
----------------------------------------------------------------------
diff --git a/flink-core/src/test/java/org/apache/flink/util/MathUtilTest.java b/flink-core/src/test/java/org/apache/flink/util/MathUtilTest.java
index 7917a7b..c98b7fc 100644
--- a/flink-core/src/test/java/org/apache/flink/util/MathUtilTest.java
+++ b/flink-core/src/test/java/org/apache/flink/util/MathUtilTest.java
@@ -81,6 +81,37 @@ public class MathUtilTest {
}
@Test
+ public void testRoundUpToPowerOf2() {
+ assertEquals(0, MathUtils.roundUpToPowerOfTwo(0));
+ assertEquals(1, MathUtils.roundUpToPowerOfTwo(1));
+ assertEquals(2, MathUtils.roundUpToPowerOfTwo(2));
+ assertEquals(4, MathUtils.roundUpToPowerOfTwo(3));
+ assertEquals(4, MathUtils.roundUpToPowerOfTwo(4));
+ assertEquals(8, MathUtils.roundUpToPowerOfTwo(5));
+ assertEquals(8, MathUtils.roundUpToPowerOfTwo(6));
+ assertEquals(8, MathUtils.roundUpToPowerOfTwo(7));
+ assertEquals(8, MathUtils.roundUpToPowerOfTwo(8));
+ assertEquals(16, MathUtils.roundUpToPowerOfTwo(9));
+ assertEquals(16, MathUtils.roundUpToPowerOfTwo(15));
+ assertEquals(16, MathUtils.roundUpToPowerOfTwo(16));
+ assertEquals(32, MathUtils.roundUpToPowerOfTwo(17));
+ assertEquals(32, MathUtils.roundUpToPowerOfTwo(31));
+ assertEquals(32, MathUtils.roundUpToPowerOfTwo(32));
+ assertEquals(64, MathUtils.roundUpToPowerOfTwo(33));
+ assertEquals(64, MathUtils.roundUpToPowerOfTwo(42));
+ assertEquals(64, MathUtils.roundUpToPowerOfTwo(63));
+ assertEquals(64, MathUtils.roundUpToPowerOfTwo(64));
+ assertEquals(128, MathUtils.roundUpToPowerOfTwo(125));
+ assertEquals(32768, MathUtils.roundUpToPowerOfTwo(25654));
+ assertEquals(67108864, MathUtils.roundUpToPowerOfTwo(34366363));
+ assertEquals(67108864, MathUtils.roundUpToPowerOfTwo(67108863));
+ assertEquals(67108864, MathUtils.roundUpToPowerOfTwo(67108864));
+ assertEquals(0x40000000, MathUtils.roundUpToPowerOfTwo(0x3FFFFFFE));
+ assertEquals(0x40000000, MathUtils.roundUpToPowerOfTwo(0x3FFFFFFF));
+ assertEquals(0x40000000, MathUtils.roundUpToPowerOfTwo(0x40000000));
+ }
+
+ @Test
public void testPowerOfTwo() {
assertTrue(MathUtils.isPowerOf2(1));
assertTrue(MathUtils.isPowerOf2(2));
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
index e751e08..52f6d9a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java
@@ -36,6 +36,7 @@ import org.apache.flink.runtime.messages.checkpoint.NotifyCheckpointComplete;
import org.apache.flink.runtime.messages.checkpoint.TriggerCheckpoint;
import org.apache.flink.runtime.state.ChainedStateHandle;
import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.state.KeyGroupsStateHandle;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.util.Preconditions;
@@ -886,7 +887,7 @@ public class CheckpointCoordinator {
List<KeyGroupRange> result = new ArrayList<>(parallelism);
int start = 0;
for (int i = 0; i < parallelism; ++i) {
- result.add(KeyGroupRange.computeKeyGroupRangeForOperatorIndex(numberKeyGroups, parallelism, i));
+ result.add(KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(numberKeyGroups, parallelism, i));
}
return result;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
index a623295..8ddc9f5 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
@@ -253,7 +253,8 @@ public class JobVertex implements java.io.Serializable {
*/
public void setMaxParallelism(int maxParallelism) {
org.apache.flink.util.Preconditions.checkArgument(
- maxParallelism > 0 && maxParallelism <= Short.MAX_VALUE, "The max parallelism must be at least 1.");
+ maxParallelism > 0 && maxParallelism <= (1 << 15),
+ "The max parallelism must be at least 1 and smaller than 2^15.");
this.maxParallelism = maxParallelism;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
index e6093a8..0d2bf45 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
@@ -19,7 +19,6 @@
package org.apache.flink.runtime.state;
import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
@@ -53,7 +52,7 @@ public abstract class AbstractStateBackend implements java.io.Serializable {
JobID jobID,
String operatorIdentifier,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
TaskKvStateRegistry kvStateRegistry) throws Exception;
@@ -66,7 +65,7 @@ public abstract class AbstractStateBackend implements java.io.Serializable {
JobID jobID,
String operatorIdentifier,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
List<KeyGroupsStateHandle> restoredState,
TaskKvStateRegistry kvStateRegistry) throws Exception;
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
deleted file mode 100644
index 9ee4b90..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/HashKeyGroupAssigner.java
+++ /dev/null
@@ -1,66 +0,0 @@
-/*
- * 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.state.KeyGroupAssigner;
-import org.apache.flink.util.MathUtils;
-import org.apache.flink.util.Preconditions;
-
-/**
- * Hash based key group assigner. The assigner assigns each key to a key group using the hash value
- * of the key.
- *
- * @param <K> Type of the key
- */
-public class HashKeyGroupAssigner<K> implements KeyGroupAssigner<K> {
- private static final long serialVersionUID = -6319826921798945448L;
-
- private static final int UNDEFINED_NUMBER_KEY_GROUPS = Integer.MIN_VALUE;
-
- private int numberKeyGroups;
-
- public HashKeyGroupAssigner() {
- this(UNDEFINED_NUMBER_KEY_GROUPS);
- }
-
- public HashKeyGroupAssigner(int numberKeyGroups) {
- Preconditions.checkArgument(numberKeyGroups > 0 || numberKeyGroups == UNDEFINED_NUMBER_KEY_GROUPS,
- "The number of key groups has to be greater than 0 or undefined. Use " +
- "setMaxParallelism() to specify the number of key groups.");
- this.numberKeyGroups = numberKeyGroups;
- }
-
- public int getNumberKeyGroups() {
- return numberKeyGroups;
- }
-
- @Override
- public int getKeyGroupIndex(K key) {
- return MathUtils.murmurHash(key.hashCode()) % numberKeyGroups;
- }
-
- @Override
- public void setup(int numberOfKeygroups) {
- Preconditions.checkArgument(numberOfKeygroups > 0, "The number of key groups has to be " +
- "greater than 0. Use setMaxParallelism() to specify the number of key " +
- "groups.");
-
- this.numberKeyGroups = numberOfKeygroups;
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
index 9e74036..3a9d3d0 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRange.java
@@ -175,46 +175,4 @@ public class KeyGroupRange implements Iterable<Integer>, Serializable {
return startKeyGroup <= endKeyGroup ? new KeyGroupRange(startKeyGroup, endKeyGroup) : EMPTY_KEY_GROUP;
}
- /**
- * Computes the range of key-groups that are assigned to a given operator under the given parallelism and maximum
- * parallelism.
- *
- * IMPORTANT: maxParallelism must be <= Short.MAX_VALUE to avoid rounding problems in this method. If we ever want
- * to go beyond this boundary, this method must perform arithmetic on long values.
- *
- * @param maxParallelism Maximal parallelism that the job was initially created with.
- * @param parallelism The current parallelism under which the job runs. Must be <= maxParallelism.
- * @param operatorIndex Id of a key-group. 0 <= keyGroupID < maxParallelism.
- * @return
- */
- public static KeyGroupRange computeKeyGroupRangeForOperatorIndex(
- int maxParallelism,
- int parallelism,
- int operatorIndex) {
- Preconditions.checkArgument(parallelism > 0, "Parallelism must not be smaller than zero.");
- Preconditions.checkArgument(maxParallelism >= parallelism, "Maximum parallelism must not be smaller than parallelism.");
- Preconditions.checkArgument(maxParallelism <= Short.MAX_VALUE, "Maximum parallelism must be smaller than Short.MAX_VALUE.");
-
- int start = operatorIndex == 0 ? 0 : ((operatorIndex * maxParallelism - 1) / parallelism) + 1;
- int end = ((operatorIndex + 1) * maxParallelism - 1) / parallelism;
- return new KeyGroupRange(start, end);
- }
-
- /**
- * Computes the index of the operator to which a key-group belongs under the given parallelism and maximum
- * parallelism.
- *
- * IMPORTANT: maxParallelism must be <= Short.MAX_VALUE to avoid rounding problems in this method. If we ever want
- * to go beyond this boundary, this method must perform arithmetic on long values.
- *
- * @param maxParallelism Maximal parallelism that the job was initially created with.
- * 0 < parallelism <= maxParallelism <= Short.MAX_VALUE must hold.
- * @param parallelism The current parallelism under which the job runs. Must be <= maxParallelism.
- * @param keyGroupId Id of a key-group. 0 <= keyGroupID < maxParallelism.
- * @return The index of the operator to which elements from the given key-group should be routed under the given
- * parallelism and maxParallelism.
- */
- public static final int computeOperatorIndexForKeyGroup(int maxParallelism, int parallelism, int keyGroupId) {
- return keyGroupId * parallelism / maxParallelism;
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRangeAssignment.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRangeAssignment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRangeAssignment.java
new file mode 100644
index 0000000..eceb6f4
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyGroupRangeAssignment.java
@@ -0,0 +1,97 @@
+/*
+ * 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.MathUtils;
+import org.apache.flink.util.Preconditions;
+
+public final class KeyGroupRangeAssignment {
+
+ public static final int DEFAULT_MAX_PARALLELISM = 128;
+
+ private KeyGroupRangeAssignment() {
+ throw new AssertionError();
+ }
+
+ /**
+ * Assigns the given key to a parallel operator index.
+ *
+ * @param key the key to assign
+ * @param maxParallelism the maximum supported parallelism, aka the number of key-groups.
+ * @param parallelism the current parallelism of the operator
+ * @return the index of the parallel operator to which the given key should be routed.
+ */
+ public static int assignKeyToParallelOperator(Object key, int maxParallelism, int parallelism) {
+ return computeOperatorIndexForKeyGroup(maxParallelism, parallelism, assignToKeyGroup(key, maxParallelism));
+ }
+
+ /**
+ * Assigns the given key to a key-group index.
+ *
+ * @param key the key to assign
+ * @param maxParallelism the maximum supported parallelism, aka the number of key-groups.
+ * @return the key-group to which the given key is assigned
+ */
+ public static final int assignToKeyGroup(Object key, int maxParallelism) {
+ return MathUtils.murmurHash(key.hashCode()) % maxParallelism;
+ }
+
+ /**
+ * Computes the range of key-groups that are assigned to a given operator under the given parallelism and maximum
+ * parallelism.
+ *
+ * IMPORTANT: maxParallelism must be <= Short.MAX_VALUE to avoid rounding problems in this method. If we ever want
+ * to go beyond this boundary, this method must perform arithmetic on long values.
+ *
+ * @param maxParallelism Maximal parallelism that the job was initially created with.
+ * @param parallelism The current parallelism under which the job runs. Must be <= maxParallelism.
+ * @param operatorIndex Id of a key-group. 0 <= keyGroupID < maxParallelism.
+ * @return
+ */
+ public static KeyGroupRange computeKeyGroupRangeForOperatorIndex(
+ int maxParallelism,
+ int parallelism,
+ int operatorIndex) {
+ Preconditions.checkArgument(parallelism > 0, "Parallelism must not be smaller than zero.");
+ Preconditions.checkArgument(maxParallelism >= parallelism, "Maximum parallelism must not be smaller than parallelism.");
+ Preconditions.checkArgument(maxParallelism <= (1 << 15), "Maximum parallelism must be smaller than 2^15.");
+
+ int start = operatorIndex == 0 ? 0 : ((operatorIndex * maxParallelism - 1) / parallelism) + 1;
+ int end = ((operatorIndex + 1) * maxParallelism - 1) / parallelism;
+ return new KeyGroupRange(start, end);
+ }
+
+ /**
+ * Computes the index of the operator to which a key-group belongs under the given parallelism and maximum
+ * parallelism.
+ *
+ * IMPORTANT: maxParallelism must be <= Short.MAX_VALUE to avoid rounding problems in this method. If we ever want
+ * to go beyond this boundary, this method must perform arithmetic on long values.
+ *
+ * @param maxParallelism Maximal parallelism that the job was initially created with.
+ * 0 < parallelism <= maxParallelism <= Short.MAX_VALUE must hold.
+ * @param parallelism The current parallelism under which the job runs. Must be <= maxParallelism.
+ * @param keyGroupId Id of a key-group. 0 <= keyGroupID < maxParallelism.
+ * @return The index of the operator to which elements from the given key-group should be routed under the given
+ * parallelism and maxParallelism.
+ */
+ public static int computeOperatorIndexForKeyGroup(int maxParallelism, int parallelism, int keyGroupId) {
+ return keyGroupId * parallelism / maxParallelism;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
index 2d1d25c..bf9018e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/KeyedStateBackend.java
@@ -22,7 +22,6 @@ import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.state.FoldingState;
import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.MergingState;
@@ -69,8 +68,8 @@ public abstract class KeyedStateBackend<K> {
@SuppressWarnings("rawtypes")
private KvState lastState;
- /** KeyGroupAssigner which determines the key group for each keys */
- protected final KeyGroupAssigner<K> keyGroupAssigner;
+ /** The number of key-groups aka max parallelism */
+ protected final int numberOfKeyGroups;
/** Range of key-groups for which this backend is responsible */
protected final KeyGroupRange keyGroupRange;
@@ -81,12 +80,12 @@ public abstract class KeyedStateBackend<K> {
public KeyedStateBackend(
TaskKvStateRegistry kvStateRegistry,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange) {
this.kvStateRegistry = Preconditions.checkNotNull(kvStateRegistry);
this.keySerializer = Preconditions.checkNotNull(keySerializer);
- this.keyGroupAssigner = Preconditions.checkNotNull(keyGroupAssigner);
+ this.numberOfKeyGroups = Preconditions.checkNotNull(numberOfKeyGroups);
this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange);
}
@@ -157,7 +156,7 @@ public abstract class KeyedStateBackend<K> {
*/
public void setCurrentKey(K newKey) {
this.currentKey = newKey;
- this.currentKeyGroup = keyGroupAssigner.getKeyGroupIndex(newKey);
+ this.currentKeyGroup = KeyGroupRangeAssignment.assignToKeyGroup(newKey, numberOfKeyGroups);
}
/**
@@ -179,11 +178,7 @@ public abstract class KeyedStateBackend<K> {
}
public int getNumberOfKeyGroups() {
- return keyGroupAssigner.getNumberKeyGroups();
- }
-
- public KeyGroupAssigner<K> getKeyGroupAssigner() {
- return keyGroupAssigner;
+ return numberOfKeyGroups;
}
/**
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
index 5495244..6d92a4d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/filesystem/FsStateBackend.java
@@ -19,7 +19,6 @@
package org.apache.flink.runtime.state.filesystem;
import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.core.fs.Path;
@@ -181,13 +180,13 @@ public class FsStateBackend extends AbstractStateBackend {
JobID jobID,
String operatorIdentifier,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
TaskKvStateRegistry kvStateRegistry) throws Exception {
return new HeapKeyedStateBackend<>(
kvStateRegistry,
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange);
}
@@ -197,14 +196,14 @@ public class FsStateBackend extends AbstractStateBackend {
JobID jobID,
String operatorIdentifier,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
List<KeyGroupsStateHandle> restoredState,
TaskKvStateRegistry kvStateRegistry) throws Exception {
return new HeapKeyedStateBackend<>(
kvStateRegistry,
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange,
restoredState);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapState.java
index 9863c93..18d1bc7 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapState.java
@@ -24,9 +24,9 @@ import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.heap.StateTable;
import org.apache.flink.util.Preconditions;
import java.util.HashMap;
@@ -138,7 +138,7 @@ public abstract class AbstractHeapState<K, N, SV, S extends State, SD extends St
Preconditions.checkState(key != null, "No key given.");
Map<N, Map<K, SV>> namespaceMap =
- stateTable.get(backend.getKeyGroupAssigner().getKeyGroupIndex(key));
+ stateTable.get(KeyGroupRangeAssignment.assignToKeyGroup(key, backend.getNumberOfKeyGroups()));
if (namespaceMap == null) {
return null;
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
index fcb4bef..8d13941 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java
@@ -20,7 +20,6 @@ package org.apache.flink.runtime.state.heap;
import org.apache.flink.api.common.state.FoldingState;
import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ReducingState;
@@ -76,20 +75,20 @@ public class HeapKeyedStateBackend<K> extends KeyedStateBackend<K> {
public HeapKeyedStateBackend(
TaskKvStateRegistry kvStateRegistry,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange) {
- super(kvStateRegistry, keySerializer, keyGroupAssigner, keyGroupRange);
+ super(kvStateRegistry, keySerializer, numberOfKeyGroups, keyGroupRange);
LOG.info("Initializing heap keyed state backend with stream factory.");
}
public HeapKeyedStateBackend(TaskKvStateRegistry kvStateRegistry,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
List<KeyGroupsStateHandle> restoredState) throws Exception {
- super(kvStateRegistry, keySerializer, keyGroupAssigner, keyGroupRange);
+ super(kvStateRegistry, keySerializer, numberOfKeyGroups, keyGroupRange);
LOG.info("Initializing heap keyed state backend from snapshot.");
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapListState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapListState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapListState.java
index 4c65c25..9552325 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapListState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapListState.java
@@ -22,6 +22,7 @@ import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.util.Preconditions;
@@ -119,7 +120,7 @@ public class HeapListState<K, N, V>
Preconditions.checkState(key != null, "No key given.");
Map<N, Map<K, ArrayList<V>>> namespaceMap =
- stateTable.get(backend.getKeyGroupAssigner().getKeyGroupIndex(key));
+ stateTable.get(KeyGroupRangeAssignment.assignToKeyGroup(key, backend.getNumberOfKeyGroups()));
if (namespaceMap == null) {
return null;
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
index 654c367..179dfe7 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/memory/MemoryStateBackend.java
@@ -19,7 +19,6 @@
package org.apache.flink.runtime.state.memory;
import org.apache.flink.api.common.JobID;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
@@ -80,14 +79,14 @@ public class MemoryStateBackend extends AbstractStateBackend {
public <K> KeyedStateBackend<K> createKeyedStateBackend(
Environment env, JobID jobID,
String operatorIdentifier, TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
TaskKvStateRegistry kvStateRegistry) throws IOException {
return new HeapKeyedStateBackend<>(
kvStateRegistry,
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange);
}
@@ -96,7 +95,7 @@ public class MemoryStateBackend extends AbstractStateBackend {
Environment env, JobID jobID,
String operatorIdentifier,
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
List<KeyGroupsStateHandle> restoredState,
TaskKvStateRegistry kvStateRegistry) throws Exception {
@@ -104,7 +103,7 @@ public class MemoryStateBackend extends AbstractStateBackend {
return new HeapKeyedStateBackend<>(
kvStateRegistry,
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange,
restoredState);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
index 495dced..4972c51 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorTest.java
@@ -35,6 +35,7 @@ import org.apache.flink.runtime.messages.checkpoint.NotifyCheckpointComplete;
import org.apache.flink.runtime.messages.checkpoint.TriggerCheckpoint;
import org.apache.flink.runtime.state.ChainedStateHandle;
import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.state.KeyGroupRangeOffsets;
import org.apache.flink.runtime.state.KeyGroupsStateHandle;
import org.apache.flink.runtime.state.StreamStateHandle;
@@ -2458,7 +2459,7 @@ public class CheckpointCoordinatorTest {
private void testCreateKeyGroupPartitions(int maxParallelism, int parallelism) {
List<KeyGroupRange> ranges = CheckpointCoordinator.createKeyGroupPartitions(maxParallelism, parallelism);
for (int i = 0; i < maxParallelism; ++i) {
- KeyGroupRange range = ranges.get(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, i));
+ KeyGroupRange range = ranges.get(KeyGroupRangeAssignment.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, i));
if (!range.contains(i)) {
Assert.fail("Could not find expected key-group " + i + " in range " + range);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java
index 3380907..405f962 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/query/QueryableStateClientTest.java
@@ -32,7 +32,6 @@ import org.apache.flink.runtime.query.netty.KvStateClient;
import org.apache.flink.runtime.query.netty.KvStateServer;
import org.apache.flink.runtime.query.netty.UnknownKvStateID;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.VoidNamespace;
@@ -232,6 +231,7 @@ public class QueryableStateClientTest {
// Config
int numServers = 2;
int numKeys = 1024;
+ int numKeyGroups = 1;
JobID jobId = new JobID();
JobVertexID jobVertexId = new JobVertexID();
@@ -250,7 +250,7 @@ public class QueryableStateClientTest {
new JobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ numKeyGroups,
new KeyGroupRange(0, 0),
new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java
index 796481c..f785174 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateClientTest.java
@@ -42,7 +42,6 @@ import org.apache.flink.runtime.query.netty.message.KvStateRequest;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
import org.apache.flink.runtime.query.netty.message.KvStateRequestType;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.KvState;
@@ -533,6 +532,8 @@ public class KvStateClientTest {
final int batchSize = 16;
+ final int numKeyGroups = 1;
+
AbstractStateBackend abstractBackend = new MemoryStateBackend();
KvStateRegistry dummyRegistry = new KvStateRegistry();
DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
@@ -542,7 +543,7 @@ public class KvStateClientTest {
new JobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ numKeyGroups,
new KeyGroupRange(0, 0),
dummyRegistry.createTaskRegistry(new JobID(), new JobVertexID()));
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java
index 3d2e8b5..52c807f 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerHandlerTest.java
@@ -39,7 +39,6 @@ import org.apache.flink.runtime.query.netty.message.KvStateRequestResult;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
import org.apache.flink.runtime.query.netty.message.KvStateRequestType;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.KvState;
@@ -89,6 +88,7 @@ public class KvStateServerHandlerTest {
ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE, null);
desc.setQueryable("vanilla");
+ int numKeyGroups =1;
AbstractStateBackend abstractBackend = new MemoryStateBackend();
DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
dummyEnv.setKvStateRegistry(registry);
@@ -97,7 +97,7 @@ public class KvStateServerHandlerTest {
new JobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ numKeyGroups,
new KeyGroupRange(0, 0),
registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
@@ -200,6 +200,7 @@ public class KvStateServerHandlerTest {
KvStateServerHandler handler = new KvStateServerHandler(registry, TEST_THREAD_POOL, stats);
EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
+ int numKeyGroups = 1;
AbstractStateBackend abstractBackend = new MemoryStateBackend();
DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
dummyEnv.setKvStateRegistry(registry);
@@ -208,7 +209,7 @@ public class KvStateServerHandlerTest {
new JobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ numKeyGroups,
new KeyGroupRange(0, 0),
registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
@@ -346,6 +347,7 @@ public class KvStateServerHandlerTest {
KvStateServerHandler handler = new KvStateServerHandler(registry, closedExecutor, stats);
EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
+ int numKeyGroups = 1;
AbstractStateBackend abstractBackend = new MemoryStateBackend();
DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
dummyEnv.setKvStateRegistry(registry);
@@ -354,7 +356,7 @@ public class KvStateServerHandlerTest {
new JobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ numKeyGroups,
new KeyGroupRange(0, 0),
registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
@@ -484,6 +486,7 @@ public class KvStateServerHandlerTest {
KvStateServerHandler handler = new KvStateServerHandler(registry, TEST_THREAD_POOL, stats);
EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
+ int numKeyGroups = 1;
AbstractStateBackend abstractBackend = new MemoryStateBackend();
DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
dummyEnv.setKvStateRegistry(registry);
@@ -492,7 +495,7 @@ public class KvStateServerHandlerTest {
new JobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ numKeyGroups,
new KeyGroupRange(0, 0),
registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
@@ -579,6 +582,7 @@ public class KvStateServerHandlerTest {
KvStateServerHandler handler = new KvStateServerHandler(registry, TEST_THREAD_POOL, stats);
EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);
+ int numKeyGroups = 1;
AbstractStateBackend abstractBackend = new MemoryStateBackend();
DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
dummyEnv.setKvStateRegistry(registry);
@@ -587,7 +591,7 @@ public class KvStateServerHandlerTest {
new JobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ numKeyGroups,
new KeyGroupRange(0, 0),
registry.createTaskRegistry(dummyEnv.getJobID(), dummyEnv.getJobVertexId()));
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java
index 30d91b6..e92fb10 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/query/netty/KvStateServerTest.java
@@ -42,7 +42,6 @@ import org.apache.flink.runtime.query.netty.message.KvStateRequestResult;
import org.apache.flink.runtime.query.netty.message.KvStateRequestSerializer;
import org.apache.flink.runtime.query.netty.message.KvStateRequestType;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.VoidNamespace;
@@ -80,7 +79,6 @@ public class KvStateServerTest {
public void testSimpleRequest() throws Exception {
KvStateServer server = null;
Bootstrap bootstrap = null;
-
try {
KvStateRegistry registry = new KvStateRegistry();
KvStateRequestStats stats = new AtomicKvStateRequestStats();
@@ -89,7 +87,7 @@ public class KvStateServerTest {
server.start();
KvStateServerAddress serverAddress = server.getAddress();
-
+ int numKeyGroups = 1;
AbstractStateBackend abstractBackend = new MemoryStateBackend();
DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
dummyEnv.setKvStateRegistry(registry);
@@ -98,7 +96,7 @@ public class KvStateServerTest {
new JobID(),
"test_op",
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ numKeyGroups,
new KeyGroupRange(0, 0),
registry.createTaskRegistry(new JobID(), new JobVertexID()));
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
index f094bd5..5984aca 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/StateBackendTestBase.java
@@ -26,7 +26,6 @@ import org.apache.flink.api.common.functions.FoldFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.state.FoldingState;
import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ReducingState;
@@ -41,7 +40,6 @@ import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.checkpoint.CheckpointCoordinator;
import org.apache.flink.runtime.execution.Environment;
-import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.query.KvStateID;
import org.apache.flink.runtime.query.KvStateRegistry;
@@ -56,7 +54,6 @@ import java.util.Collections;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.Future;
import java.util.concurrent.RunnableFuture;
import static org.junit.Assert.assertEquals;
@@ -90,14 +87,14 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
protected <K> KeyedStateBackend<K> createKeyedBackend(TypeSerializer<K> keySerializer, Environment env) throws Exception {
return createKeyedBackend(
keySerializer,
- new HashKeyGroupAssigner<K>(10),
+ 10,
new KeyGroupRange(0, 9),
env);
}
protected <K> KeyedStateBackend<K> createKeyedBackend(
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
Environment env) throws Exception {
return getStateBackend().createKeyedStateBackend(
@@ -105,7 +102,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
new JobID(),
"test_op",
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange,
env.getTaskKvStateRegistry());
}
@@ -120,7 +117,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
Environment env) throws Exception {
return restoreKeyedBackend(
keySerializer,
- new HashKeyGroupAssigner<K>(10),
+ 10,
new KeyGroupRange(0, 9),
Collections.singletonList(state),
env);
@@ -128,7 +125,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
protected <K> KeyedStateBackend<K> restoreKeyedBackend(
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange,
List<KeyGroupsStateHandle> state,
Environment env) throws Exception {
@@ -137,7 +134,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
new JobID(),
"test_op",
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange,
state,
env.getTaskKvStateRegistry());
@@ -243,7 +240,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
KeyedStateBackend<Integer> backend = createKeyedBackend(
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ 1,
new KeyGroupRange(0, 0),
new DummyEnvironment("test_op", 1, 0));
@@ -277,7 +274,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
backend.close();
backend = restoreKeyedBackend(
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ 1,
new KeyGroupRange(0, 0),
Collections.singletonList(snapshot1),
new DummyEnvironment("test_op", 1, 0));
@@ -675,12 +672,9 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
final int MAX_PARALLELISM = 10;
CheckpointStreamFactory streamFactory = createStreamFactory();
-
- HashKeyGroupAssigner<Integer> keyGroupAssigner = new HashKeyGroupAssigner<>(10);
-
KeyedStateBackend<Integer> backend = createKeyedBackend(
IntSerializer.INSTANCE,
- keyGroupAssigner,
+ MAX_PARALLELISM,
new KeyGroupRange(0, MAX_PARALLELISM - 1),
new DummyEnvironment("test", 1, 0));
@@ -695,12 +689,12 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
Random rand = new Random(0);
// for each key, determine into which half of the key-group space they fall
- int firstKeyHalf = keyGroupAssigner.getKeyGroupIndex(keyInFirstHalf) * 2 / MAX_PARALLELISM;
- int secondKeyHalf = keyGroupAssigner.getKeyGroupIndex(keyInSecondHalf) * 2 / MAX_PARALLELISM;
+ int firstKeyHalf = KeyGroupRangeAssignment.assignKeyToParallelOperator(keyInFirstHalf, MAX_PARALLELISM, 2);
+ int secondKeyHalf = KeyGroupRangeAssignment.assignKeyToParallelOperator(keyInFirstHalf, MAX_PARALLELISM, 2);
while (firstKeyHalf == secondKeyHalf) {
keyInSecondHalf = rand.nextInt();
- secondKeyHalf = keyGroupAssigner.getKeyGroupIndex(keyInSecondHalf) * 2 / MAX_PARALLELISM;
+ secondKeyHalf = KeyGroupRangeAssignment.assignKeyToParallelOperator(keyInSecondHalf, MAX_PARALLELISM, 2);
}
backend.setCurrentKey(keyInFirstHalf);
@@ -714,18 +708,18 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
List<KeyGroupsStateHandle> firstHalfKeyGroupStates = CheckpointCoordinator.getKeyGroupsStateHandles(
Collections.singletonList(snapshot),
- new KeyGroupRange(0, 4));
+ KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(MAX_PARALLELISM, 2, 0));
List<KeyGroupsStateHandle> secondHalfKeyGroupStates = CheckpointCoordinator.getKeyGroupsStateHandles(
Collections.singletonList(snapshot),
- new KeyGroupRange(5, 9));
+ KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(MAX_PARALLELISM, 2, 1));
backend.close();
// backend for the first half of the key group range
KeyedStateBackend<Integer> firstHalfBackend = restoreKeyedBackend(
IntSerializer.INSTANCE,
- keyGroupAssigner,
+ MAX_PARALLELISM,
new KeyGroupRange(0, 4),
firstHalfKeyGroupStates,
new DummyEnvironment("test", 1, 0));
@@ -733,7 +727,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
// backend for the second half of the key group range
KeyedStateBackend<Integer> secondHalfBackend = restoreKeyedBackend(
IntSerializer.INSTANCE,
- keyGroupAssigner,
+ MAX_PARALLELISM,
new KeyGroupRange(5, 9),
secondHalfKeyGroupStates,
new DummyEnvironment("test", 1, 0));
@@ -978,9 +972,10 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
*/
@SuppressWarnings("unchecked")
protected void testConcurrentMapIfQueryable() throws Exception {
+ final int numberOfKeyGroups = 1;
KeyedStateBackend<Integer> backend = createKeyedBackend(
IntSerializer.INSTANCE,
- new HashKeyGroupAssigner<Integer>(1),
+ numberOfKeyGroups,
new KeyGroupRange(0, 0),
new DummyEnvironment("test_op", 1, 0));
@@ -1005,7 +1000,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
backend.setCurrentKey(1);
state.update(121818273);
- int keyGroupIndex = new HashKeyGroupAssigner<>(1).getKeyGroupIndex(1);
+ int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(1, numberOfKeyGroups);
StateTable stateTable = ((AbstractHeapState) kvState).getStateTable();
assertNotNull("State not set", stateTable.get(keyGroupIndex));
assertTrue(stateTable.get(keyGroupIndex) instanceof ConcurrentHashMap);
@@ -1031,7 +1026,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
backend.setCurrentKey(1);
state.add(121818273);
- int keyGroupIndex = new HashKeyGroupAssigner<>(1).getKeyGroupIndex(1);
+ int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(1, numberOfKeyGroups);
StateTable stateTable = ((AbstractHeapState) kvState).getStateTable();
assertNotNull("State not set", stateTable.get(keyGroupIndex));
assertTrue(stateTable.get(keyGroupIndex) instanceof ConcurrentHashMap);
@@ -1062,7 +1057,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
backend.setCurrentKey(1);
state.add(121818273);
- int keyGroupIndex = new HashKeyGroupAssigner<>(1).getKeyGroupIndex(1);
+ int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(1, numberOfKeyGroups);
StateTable stateTable = ((AbstractHeapState) kvState).getStateTable();
assertNotNull("State not set", stateTable.get(keyGroupIndex));
assertTrue(stateTable.get(keyGroupIndex) instanceof ConcurrentHashMap);
@@ -1093,7 +1088,7 @@ public abstract class StateBackendTestBase<B extends AbstractStateBackend> {
backend.setCurrentKey(1);
state.add(121818273);
- int keyGroupIndex = new HashKeyGroupAssigner<>(1).getKeyGroupIndex(1);
+ int keyGroupIndex = KeyGroupRangeAssignment.assignToKeyGroup(1, numberOfKeyGroups);
StateTable stateTable = ((AbstractHeapState) kvState).getStateTable();
assertNotNull("State not set", stateTable.get(keyGroupIndex));
assertTrue(stateTable.get(keyGroupIndex) instanceof ConcurrentHashMap);
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java
index 1fb34b8..af907e3 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/datastream/KeyedStream.java
@@ -30,7 +30,7 @@ import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.Utils;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.typeutils.TypeExtractor;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.functions.aggregation.AggregationFunction;
import org.apache.flink.streaming.api.functions.aggregation.ComparableAggregator;
@@ -111,8 +111,7 @@ public class KeyedStream<T, KEY> extends DataStream<T> {
new PartitionTransformation<>(
dataStream.getTransformation(),
new KeyGroupStreamPartitioner<>(
- keySelector,
- new HashKeyGroupAssigner<KEY>())));
+ keySelector, KeyGroupRangeAssignment.DEFAULT_MAX_PARALLELISM)));
this.keySelector = keySelector;
this.keyType = keyType;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java
index 1a92ba4..8e807db 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java
@@ -26,7 +26,6 @@ import java.util.List;
import java.util.Map;
import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
@@ -76,8 +75,7 @@ public class StreamConfig implements Serializable {
private static final String STATE_BACKEND = "statebackend";
private static final String STATE_PARTITIONER = "statePartitioner";
- /** key for the {@link KeyGroupAssigner} for key to key group index mappings */
- private static final String KEY_GROUP_ASSIGNER = "keyGroupAssigner";
+ private static final String NUMBER_OF_KEY_GROUPS = "numberOfKeyGroups";
private static final String STATE_KEY_SERIALIZER = "statekeyser";
@@ -445,28 +443,26 @@ public class StreamConfig implements Serializable {
}
/**
- * Sets the {@link KeyGroupAssigner} to be used for the current {@link StreamOperator}.
+ * Sets the number of key-groups to be used for the current {@link StreamOperator}.
*
- * @param keyGroupAssigner Key group assigner to be used
+ * @param numberOfKeyGroups Number of key-groups to be used
*/
- public void setKeyGroupAssigner(KeyGroupAssigner<?> keyGroupAssigner) {
+ public void setNumberOfKeyGroups(int numberOfKeyGroups) {
try {
- InstantiationUtil.writeObjectToConfig(keyGroupAssigner, this.config, KEY_GROUP_ASSIGNER);
+ InstantiationUtil.writeObjectToConfig(numberOfKeyGroups, this.config, NUMBER_OF_KEY_GROUPS);
} catch (Exception e) {
throw new StreamTaskException("Could not serialize virtual state partitioner.", e);
}
}
/**
- * Gets the {@link KeyGroupAssigner} for the {@link StreamOperator}.
+ * Gets the number of key-groups for the {@link StreamOperator}.
*
- * @param classLoader Classloader to be used for the deserialization
- * @param <K> Type of the keys to be assigned to key groups
- * @return Key group assigner
+ * @return the number of key-groups
*/
- public <K> KeyGroupAssigner<K> getKeyGroupAssigner(ClassLoader classLoader) {
+ public Integer getNumberOfKeyGroups(ClassLoader cl) {
try {
- return InstantiationUtil.readObjectFromConfig(this.config, KEY_GROUP_ASSIGNER, classLoader);
+ return InstantiationUtil.readObjectFromConfig(this.config, NUMBER_OF_KEY_GROUPS, cl);
} catch (Exception e) {
throw new StreamTaskException("Could not instantiate virtual state partitioner.", e);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
index f9f26e9..506b664 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamGraphGenerator.java
@@ -20,6 +20,7 @@ package org.apache.flink.streaming.api.graph;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.InputFormatSourceFunction;
import org.apache.flink.streaming.api.transformations.CoFeedbackTransformation;
@@ -33,6 +34,7 @@ import org.apache.flink.streaming.api.transformations.SplitTransformation;
import org.apache.flink.streaming.api.transformations.StreamTransformation;
import org.apache.flink.streaming.api.transformations.TwoInputTransformation;
import org.apache.flink.streaming.api.transformations.UnionTransformation;
+import org.apache.flink.util.MathUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -145,24 +147,24 @@ public class StreamGraphGenerator {
LOG.debug("Transforming " + transform);
if (transform.getMaxParallelism() <= 0) {
+
// if the max parallelism hasn't been set, then first use the job wide max parallelism
// from theExecutionConfig. If this value has not been specified either, then use the
// parallelism of the operator.
int maxParallelism = env.getConfig().getMaxParallelism();
if (maxParallelism <= 0) {
- maxParallelism = transform.getParallelism();
- /**
- * TODO: Remove once the parallelism settings works properly in Flink (FLINK-3885)
- * Currently, the parallelism will be set to 1 on the JobManager iff it encounters
- * a negative parallelism value. We need to know this for the
- * KeyGroupStreamPartitioner on the client-side. Thus, we already set the value to
- * 1 here.
- */
- if (maxParallelism <= 0) {
- transform.setParallelism(1);
- maxParallelism = 1;
+
+ int parallelism = transform.getParallelism();
+
+ if(parallelism <= 0) {
+ parallelism = 1;
+ transform.setParallelism(parallelism);
}
+
+ maxParallelism = Math.max(
+ MathUtils.roundUpToPowerOfTwo(parallelism + (parallelism / 2)),
+ KeyGroupRangeAssignment.DEFAULT_MAX_PARALLELISM);
}
transform.setMaxParallelism(maxParallelism);
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java
index 76fdaca..a024895 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamingJobGraphGenerator.java
@@ -26,7 +26,6 @@ import org.apache.commons.lang3.StringUtils;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.operators.util.UserCodeObjectWrapper;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
@@ -41,7 +40,6 @@ import org.apache.flink.runtime.jobgraph.tasks.JobSnapshottingSettings;
import org.apache.flink.runtime.jobmanager.scheduler.CoLocationGroup;
import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup;
import org.apache.flink.runtime.operators.util.TaskConfig;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.environment.CheckpointConfig;
import org.apache.flink.streaming.api.operators.AbstractUdfStreamOperator;
@@ -357,12 +355,9 @@ public class StreamingJobGraphGenerator {
config.setStatePartitioner(1, vertex.getStatePartitioner2());
config.setStateKeySerializer(vertex.getStateKeySerializer());
- // only set the key group assigner if the vertex uses partitioned state (= KeyedStream).
+ // only set the max parallelism if the vertex uses partitioned state (= KeyedStream).
if (vertex.getStatePartitioner1() != null) {
- // the key group assigner has to know the number of key groups (= maxParallelism)
- KeyGroupAssigner<Object> keyGroupAssigner = new HashKeyGroupAssigner<Object>(vertex.getMaxParallelism());
-
- config.setKeyGroupAssigner(keyGroupAssigner);
+ config.setNumberOfKeyGroups(vertex.getMaxParallelism());
}
Class<? extends AbstractInvokable> vertexClass = vertex.getJobVertexClass();
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
index 718c0c7..71296e3 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java
@@ -29,6 +29,7 @@ import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.metrics.Counter;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.VoidNamespace;
import org.apache.flink.runtime.state.VoidNamespaceSerializer;
@@ -133,14 +134,14 @@ public abstract class AbstractStreamOperator<OUT>
if (null != keySerializer) {
ExecutionConfig execConf = container.getEnvironment().getExecutionConfig();;
- KeyGroupRange subTaskKeyGroupRange = KeyGroupRange.computeKeyGroupRangeForOperatorIndex(
+ KeyGroupRange subTaskKeyGroupRange = KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(
container.getEnvironment().getTaskInfo().getNumberOfKeyGroups(),
container.getEnvironment().getTaskInfo().getNumberOfParallelSubtasks(),
container.getIndexInSubtaskGroup());
keyedStateBackend = container.createKeyedStateBackend(
keySerializer,
- container.getConfiguration().getKeyGroupAssigner(getUserCodeClassloader()),
+ container.getConfiguration().getNumberOfKeyGroups(getUserCodeClassloader()),
subTaskKeyGroupRange);
}
} catch (Exception e) {
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
index 108a3ae..256fee1 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/partitioner/KeyGroupStreamPartitioner.java
@@ -18,16 +18,14 @@
package org.apache.flink.streaming.runtime.partitioner;
import org.apache.flink.annotation.Internal;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.runtime.plugable.SerializationDelegate;
-import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.util.Preconditions;
/**
- * Partitioner selects the target channel based on the key group index. The key group
- * index is derived from the key of the elements using the {@link KeyGroupAssigner}.
+ * Partitioner selects the target channel based on the key group index.
*
* @param <T> Type of the elements in the Stream being partitioned
*/
@@ -39,15 +37,16 @@ public class KeyGroupStreamPartitioner<T, K> extends StreamPartitioner<T> implem
private final KeySelector<T, K> keySelector;
- private final KeyGroupAssigner<K> keyGroupAssigner;
+ private int maxParallelism;
- public KeyGroupStreamPartitioner(KeySelector<T, K> keySelector, KeyGroupAssigner<K> keyGroupAssigner) {
+ public KeyGroupStreamPartitioner(KeySelector<T, K> keySelector, int maxParallelism) {
+ Preconditions.checkArgument(maxParallelism > 0, "Number of key-groups must be > 0!");
this.keySelector = Preconditions.checkNotNull(keySelector);
- this.keyGroupAssigner = Preconditions.checkNotNull(keyGroupAssigner);
+ this.maxParallelism = maxParallelism;
}
- public KeyGroupAssigner<K> getKeyGroupAssigner() {
- return keyGroupAssigner;
+ public int getMaxParallelism() {
+ return maxParallelism;
}
@Override
@@ -61,10 +60,7 @@ public class KeyGroupStreamPartitioner<T, K> extends StreamPartitioner<T> implem
} catch (Exception e) {
throw new RuntimeException("Could not extract key from " + record.getInstance().getValue(), e);
}
- returnArray[0] = KeyGroupRange.computeOperatorIndexForKeyGroup(
- keyGroupAssigner.getNumberKeyGroups(),
- numberOfOutputChannels,
- keyGroupAssigner.getKeyGroupIndex(key));
+ returnArray[0] = KeyGroupRangeAssignment.assignKeyToParallelOperator(key, maxParallelism, numberOfOutputChannels);
return returnArray;
}
@@ -80,6 +76,6 @@ public class KeyGroupStreamPartitioner<T, K> extends StreamPartitioner<T> implem
@Override
public void configure(int maxParallelism) {
- keyGroupAssigner.setup(maxParallelism);
+ this.maxParallelism = maxParallelism;
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
index 13f650c..701281b 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/runtime/tasks/StreamTask.java
@@ -19,7 +19,6 @@ package org.apache.flink.streaming.runtime.tasks;
import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.accumulators.Accumulator;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
@@ -762,7 +761,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
public <K> KeyedStateBackend<K> createKeyedStateBackend(
TypeSerializer<K> keySerializer,
- KeyGroupAssigner<K> keyGroupAssigner,
+ int numberOfKeyGroups,
KeyGroupRange keyGroupRange) throws Exception {
if (keyedStateBackend != null) {
@@ -779,7 +778,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
getEnvironment().getJobID(),
operatorIdentifier,
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange,
lazyRestoreKeyGroupStates,
getEnvironment().getTaskKvStateRegistry());
@@ -791,7 +790,7 @@ public abstract class StreamTask<OUT, Operator extends StreamOperator<OUT>>
getEnvironment().getJobID(),
operatorIdentifier,
keySerializer,
- keyGroupAssigner,
+ numberOfKeyGroups,
keyGroupRange,
getEnvironment().getTaskKvStateRegistry());
}
http://git-wip-us.apache.org/repos/asf/flink/blob/6d430618/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
index 06d381f..874274f 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/graph/StreamGraphGeneratorTest.java
@@ -19,11 +19,9 @@
package org.apache.flink.streaming.api.graph;
import org.apache.flink.api.common.ExecutionConfig;
-import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
@@ -270,10 +268,6 @@ public class StreamGraphGeneratorTest {
StreamNode keyedResultNode = graph.getStreamNode(keyedResult.getId());
StreamPartitioner<?> streamPartitioner = keyedResultNode.getInEdges().get(0).getPartitioner();
-
- HashKeyGroupAssigner<?> hashKeyGroupAssigner = extractHashKeyGroupAssigner(streamPartitioner);
-
- assertEquals(maxParallelism, hashKeyGroupAssigner.getNumberKeyGroups());
}
/**
@@ -384,7 +378,7 @@ public class StreamGraphGeneratorTest {
}
/**
- * Tests that the max parallelism and the key group partitioner is properly set for connected
+ * Tests that the max parallelism is properly set for connected
* streams.
*/
@Test
@@ -423,24 +417,6 @@ public class StreamGraphGeneratorTest {
StreamPartitioner<?> streamPartitioner1 = keyedResultNode.getInEdges().get(0).getPartitioner();
StreamPartitioner<?> streamPartitioner2 = keyedResultNode.getInEdges().get(1).getPartitioner();
-
- HashKeyGroupAssigner<?> hashKeyGroupAssigner1 = extractHashKeyGroupAssigner(streamPartitioner1);
- assertEquals(maxParallelism, hashKeyGroupAssigner1.getNumberKeyGroups());
-
- HashKeyGroupAssigner<?> hashKeyGroupAssigner2 = extractHashKeyGroupAssigner(streamPartitioner2);
- assertEquals(maxParallelism, hashKeyGroupAssigner2.getNumberKeyGroups());
- }
-
- private HashKeyGroupAssigner<?> extractHashKeyGroupAssigner(StreamPartitioner<?> streamPartitioner) {
- assertTrue(streamPartitioner instanceof KeyGroupStreamPartitioner);
-
- KeyGroupStreamPartitioner<?, ?> keyGroupStreamPartitioner = (KeyGroupStreamPartitioner<?, ?>) streamPartitioner;
-
- KeyGroupAssigner<?> keyGroupAssigner = keyGroupStreamPartitioner.getKeyGroupAssigner();
-
- assertTrue(keyGroupAssigner instanceof HashKeyGroupAssigner);
-
- return (HashKeyGroupAssigner<?>) keyGroupAssigner;
}
private static class OutputTypeConfigurableOperationWithTwoInputs
[06/27] flink git commit: [FLINK-3761] Refactor State Backends/Make
Keyed State Key-Group Aware
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
index 472dccb..4e7e4d0 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/AggregatingAlignedProcessingTimeWindowOperatorTest.java
@@ -48,6 +48,8 @@ import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
+import org.apache.flink.streaming.runtime.tasks.TestTimeServiceProvider;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.junit.After;
import org.junit.Ignore;
@@ -81,7 +83,6 @@ import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@SuppressWarnings({"serial", "SynchronizationOnLocalVariableOrMethodParameter"})
-@Ignore
public class AggregatingAlignedProcessingTimeWindowOperatorTest {
@SuppressWarnings("unchecked")
@@ -553,12 +554,10 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
@Test
public void checkpointRestoreWithPendingWindowTumbling() {
- final ScheduledExecutorService timerService = Executors.newSingleThreadScheduledExecutor();
try {
final int windowSize = 200;
- final CollectingOutput<Tuple2<Integer, Integer>> out = new CollectingOutput<>(windowSize);
- final Object lock = new Object();
- final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
+
+ TestTimeServiceProvider timerService = new TestTimeServiceProvider();
// tumbling window that triggers every 50 milliseconds
AggregatingProcessingTimeWindowOperator<Integer, Tuple2<Integer, Integer>> op =
@@ -568,7 +567,9 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
windowSize, windowSize);
OneInputStreamOperatorTestHarness<Tuple2<Integer, Integer>, Tuple2<Integer, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(op);
+ new OneInputStreamOperatorTestHarness<>(op, new ExecutionConfig(), timerService);
+
+ timerService.setCurrentTime(0);
testHarness.setup();
testHarness.open();
@@ -578,48 +579,34 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
final int numElements = 1000;
for (int i = 0; i < numElementsFirst; i++) {
- synchronized (lock) {
- StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
- op.setKeyContextElement1(next);
- op.processElement(next);
- }
- Thread.sleep(1);
+ StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
+ testHarness.processElement(next);
}
- // draw a snapshot and dispose the window
- StreamStateHandle state;
- List<Tuple2<Integer, Integer>> resultAtSnapshot;
- synchronized (lock) {
- int beforeSnapShot = out.getElements().size();
- state = testHarness.snapshot(1L, System.currentTimeMillis());
- resultAtSnapshot = new ArrayList<>(out.getElements());
- int afterSnapShot = out.getElements().size();
- assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
- }
-
+ // draw a snapshot
+ List<Tuple2<Integer, Integer>> resultAtSnapshot = extractFromStreamRecords(testHarness.getOutput());
+ int beforeSnapShot = resultAtSnapshot.size();
+ StreamStateHandle state = testHarness.snapshot(1L, System.currentTimeMillis());
+ int afterSnapShot = testHarness.getOutput().size();
+ assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
+
assertTrue(resultAtSnapshot.size() <= numElementsFirst);
// inject some random elements, which should not show up in the state
for (int i = numElementsFirst; i < numElements; i++) {
- synchronized (lock) {
- StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
- op.setKeyContextElement1(next);
- op.processElement(next);
- }
- Thread.sleep(1);
+ StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
+ testHarness.processElement(next);
}
op.dispose();
// re-create the operator and restore the state
- final CollectingOutput<Tuple2<Integer, Integer>> out2 = new CollectingOutput<>(windowSize);
op = new AggregatingProcessingTimeWindowOperator<>(
sumFunction, fieldOneSelector,
IntSerializer.INSTANCE, tupleSerializer,
windowSize, windowSize);
- testHarness =
- new OneInputStreamOperatorTestHarness<>(op);
+ testHarness = new OneInputStreamOperatorTestHarness<>(op, new ExecutionConfig(), timerService);
testHarness.setup();
testHarness.restore(state);
@@ -627,24 +614,19 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
// inject the remaining elements
for (int i = numElementsFirst; i < numElements; i++) {
- synchronized (lock) {
- StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
- op.setKeyContextElement1(next);
- op.processElement(next);
- }
- Thread.sleep(1);
+ StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
+ testHarness.processElement(next);
}
- out2.waitForNElements(numElements - resultAtSnapshot.size(), 60_000);
+ timerService.setCurrentTime(200);
// get and verify the result
List<Tuple2<Integer, Integer>> finalResult = new ArrayList<>(resultAtSnapshot);
- finalResult.addAll(out2.getElements());
+ List<Tuple2<Integer, Integer>> partialFinalResult = extractFromStreamRecords(testHarness.getOutput());
+ finalResult.addAll(partialFinalResult);
assertEquals(numElements, finalResult.size());
- synchronized (lock) {
- op.close();
- }
+ testHarness.close();
op.dispose();
Collections.sort(finalResult, tupleComparator);
@@ -657,22 +639,16 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
e.printStackTrace();
fail(e.getMessage());
}
- finally {
- timerService.shutdown();
- }
}
@Test
public void checkpointRestoreWithPendingWindowSliding() {
- final ScheduledExecutorService timerService = Executors.newSingleThreadScheduledExecutor();
try {
final int factor = 4;
final int windowSlide = 50;
final int windowSize = factor * windowSlide;
- final CollectingOutput<Tuple2<Integer, Integer>> out = new CollectingOutput<>(windowSlide);
- final Object lock = new Object();
- final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
+ TestTimeServiceProvider timerService = new TestTimeServiceProvider();
// sliding window (200 msecs) every 50 msecs
AggregatingProcessingTimeWindowOperator<Integer, Tuple2<Integer, Integer>> op =
@@ -681,8 +657,10 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
IntSerializer.INSTANCE, tupleSerializer,
windowSize, windowSlide);
+ timerService.setCurrentTime(0);
+
OneInputStreamOperatorTestHarness<Tuple2<Integer, Integer>, Tuple2<Integer, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(op);
+ new OneInputStreamOperatorTestHarness<>(op, new ExecutionConfig(), timerService);
testHarness.setup();
testHarness.open();
@@ -692,48 +670,34 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
final int numElementsFirst = 700;
for (int i = 0; i < numElementsFirst; i++) {
- synchronized (lock) {
- StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
- op.setKeyContextElement1(next);
- op.processElement(next);
- }
- Thread.sleep(1);
+ StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
+ testHarness.processElement(next);
}
// draw a snapshot
- StreamStateHandle state;
- List<Tuple2<Integer, Integer>> resultAtSnapshot;
- synchronized (lock) {
- int beforeSnapShot = out.getElements().size();
- state = testHarness.snapshot(1L, System.currentTimeMillis());
- resultAtSnapshot = new ArrayList<>(out.getElements());
- int afterSnapShot = out.getElements().size();
- assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
- }
+ List<Tuple2<Integer, Integer>> resultAtSnapshot = extractFromStreamRecords(testHarness.getOutput());
+ int beforeSnapShot = resultAtSnapshot.size();
+ StreamStateHandle state = testHarness.snapshot(1L, System.currentTimeMillis());
+ int afterSnapShot = testHarness.getOutput().size();
+ assertEquals("operator performed computation during snapshot", beforeSnapShot, afterSnapShot);
assertTrue(resultAtSnapshot.size() <= factor * numElementsFirst);
// inject the remaining elements - these should not influence the snapshot
for (int i = numElementsFirst; i < numElements; i++) {
- synchronized (lock) {
- StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
- op.setKeyContextElement1(next);
- op.processElement(next);
- }
- Thread.sleep(1);
+ StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
+ testHarness.processElement(next);
}
op.dispose();
// re-create the operator and restore the state
- final CollectingOutput<Tuple2<Integer, Integer>> out2 = new CollectingOutput<>(windowSlide);
op = new AggregatingProcessingTimeWindowOperator<>(
sumFunction, fieldOneSelector,
IntSerializer.INSTANCE, tupleSerializer,
windowSize, windowSlide);
- testHarness =
- new OneInputStreamOperatorTestHarness<>(op);
+ testHarness = new OneInputStreamOperatorTestHarness<>(op, new ExecutionConfig(), timerService);
testHarness.setup();
testHarness.restore(state);
@@ -741,32 +705,27 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
// inject again the remaining elements
for (int i = numElementsFirst; i < numElements; i++) {
- synchronized (lock) {
- StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
- op.setKeyContextElement1(next);
- op.processElement(next);
- }
- Thread.sleep(1);
+ StreamRecord<Tuple2<Integer, Integer>> next = new StreamRecord<>(new Tuple2<>(i, i));
+ testHarness.processElement(next);
}
- // for a deterministic result, we need to wait until all pending triggers
- // have fired and emitted their results
- long deadline = System.currentTimeMillis() + 120000;
- do {
- Thread.sleep(20);
- }
- while (resultAtSnapshot.size() + out2.getElements().size() < factor * numElements
- && System.currentTimeMillis() < deadline);
+ timerService.setCurrentTime(50);
+ timerService.setCurrentTime(100);
+ timerService.setCurrentTime(150);
+ timerService.setCurrentTime(200);
+ timerService.setCurrentTime(250);
+ timerService.setCurrentTime(300);
+ timerService.setCurrentTime(350);
+ timerService.setCurrentTime(400);
- synchronized (lock) {
- op.close();
- }
+ testHarness.close();
op.dispose();
// get and verify the result
List<Tuple2<Integer, Integer>> finalResult = new ArrayList<>(resultAtSnapshot);
- finalResult.addAll(out2.getElements());
- assertEquals(factor * numElements, finalResult.size());
+ List<Tuple2<Integer, Integer>> partialFinalResult = extractFromStreamRecords(testHarness.getOutput());
+ finalResult.addAll(partialFinalResult);
+ assertEquals(numElements * factor, finalResult.size());
Collections.sort(finalResult, tupleComparator);
for (int i = 0; i < factor * numElements; i++) {
@@ -778,20 +737,14 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
e.printStackTrace();
fail(e.getMessage());
}
- finally {
- timerService.shutdown();
- }
}
@Test
public void testKeyValueStateInWindowFunctionTumbling() {
- final ScheduledExecutorService timerService = Executors.newSingleThreadScheduledExecutor();
try {
final long twoSeconds = 2000;
- final CollectingOutput<Tuple2<Integer, Integer>> out = new CollectingOutput<>();
- final Object lock = new Object();
- final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
+ TestTimeServiceProvider timerService = new TestTimeServiceProvider();
StatefulFunction.globalCounts.clear();
@@ -800,53 +753,52 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
new StatefulFunction(), fieldOneSelector,
IntSerializer.INSTANCE, tupleSerializer, twoSeconds, twoSeconds);
- op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
- op.open();
+ KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, Integer>, Tuple2<Integer, Integer>> testHarness = new KeyedOneInputStreamOperatorTestHarness<>(
+ op,
+ new ExecutionConfig(),
+ timerService,
+ fieldOneSelector,
+ BasicTypeInfo.INT_TYPE_INFO);
+
+ timerService.setCurrentTime(0);
+ testHarness.open();
// because the window interval is so large, everything should be in one window
// and aggregate into one value per key
- synchronized (lock) {
- for (int i = 0; i < 10; i++) {
- StreamRecord<Tuple2<Integer, Integer>> next1 = new StreamRecord<>(new Tuple2<>(1, i));
- op.setKeyContextElement1(next1);
- op.processElement(next1);
-
- StreamRecord<Tuple2<Integer, Integer>> next2 = new StreamRecord<>(new Tuple2<>(2, i));
- op.setKeyContextElement1(next2);
- op.processElement(next2);
- }
- }
+ for (int i = 0; i < 10; i++) {
+ StreamRecord<Tuple2<Integer, Integer>> next1 = new StreamRecord<>(new Tuple2<>(1, i));
+ testHarness.processElement(next1);
- while (StatefulFunction.globalCounts.get(1) < 10 ||
- StatefulFunction.globalCounts.get(2) < 10)
- {
- Thread.sleep(50);
+ StreamRecord<Tuple2<Integer, Integer>> next2 = new StreamRecord<>(new Tuple2<>(2, i));
+ testHarness.processElement(next2);
}
-
- op.close();
+
+ timerService.setCurrentTime(1000);
+
+ int count1 = StatefulFunction.globalCounts.get(1);
+ int count2 = StatefulFunction.globalCounts.get(2);
+
+ assertTrue(count1 >= 2 && count1 <= 2 * 10);
+ assertEquals(count1, count2);
+
+ testHarness.close();
op.dispose();
}
catch (Exception e) {
e.printStackTrace();
fail(e.getMessage());
}
- finally {
- timerService.shutdown();
- }
}
@Test
public void testKeyValueStateInWindowFunctionSliding() {
- final ScheduledExecutorService timerService = Executors.newSingleThreadScheduledExecutor();
try {
final int factor = 2;
final int windowSlide = 50;
final int windowSize = factor * windowSlide;
-
- final CollectingOutput<Tuple2<Integer, Integer>> out = new CollectingOutput<>();
- final Object lock = new Object();
- final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
+
+ TestTimeServiceProvider timerService = new TestTimeServiceProvider();
StatefulFunction.globalCounts.clear();
@@ -855,8 +807,15 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
new StatefulFunction(), fieldOneSelector,
IntSerializer.INSTANCE, tupleSerializer, windowSize, windowSlide);
- op.setup(mockTask, createTaskConfig(fieldOneSelector, IntSerializer.INSTANCE, new HashKeyGroupAssigner<Object>(10)), out);
- op.open();
+ timerService.setCurrentTime(0);
+ KeyedOneInputStreamOperatorTestHarness<Integer, Tuple2<Integer, Integer>, Tuple2<Integer, Integer>> testHarness = new KeyedOneInputStreamOperatorTestHarness<>(
+ op,
+ new ExecutionConfig(),
+ timerService,
+ fieldOneSelector,
+ BasicTypeInfo.INT_TYPE_INFO);
+
+ testHarness.open();
// because the window interval is so large, everything should be in one window
// and aggregate into one value per key
@@ -870,25 +829,19 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
StreamRecord<Tuple2<Integer, Integer>> next3 = new StreamRecord<>(new Tuple2<>(1, i));
StreamRecord<Tuple2<Integer, Integer>> next4 = new StreamRecord<>(new Tuple2<>(2, i));
- // because we do not release the lock between elements, they end up in the same windows
- synchronized (lock) {
- op.setKeyContextElement1(next1);
- op.processElement(next1);
- op.setKeyContextElement1(next2);
- op.processElement(next2);
- op.setKeyContextElement1(next3);
- op.processElement(next3);
- op.setKeyContextElement1(next4);
- op.processElement(next4);
- }
-
- Thread.sleep(1);
- }
-
- synchronized (lock) {
- op.close();
+ testHarness.processElement(next1);
+ testHarness.processElement(next2);
+ testHarness.processElement(next3);
+ testHarness.processElement(next4);
}
+ timerService.setCurrentTime(50);
+ timerService.setCurrentTime(100);
+ timerService.setCurrentTime(150);
+ timerService.setCurrentTime(200);
+
+ testHarness.close();
+
int count1 = StatefulFunction.globalCounts.get(1);
int count2 = StatefulFunction.globalCounts.get(2);
@@ -901,9 +854,6 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
e.printStackTrace();
fail(e.getMessage());
}
- finally {
- timerService.shutdown();
- }
}
// ------------------------------------------------------------------------
@@ -991,21 +941,6 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
final Environment env = new DummyEnvironment("Test task name", 1, 0);
when(task.getEnvironment()).thenReturn(env);
- try {
- doAnswer(new Answer<AbstractStateBackend>() {
- @Override
- public AbstractStateBackend answer(InvocationOnMock invocationOnMock) throws Throwable {
- final String operatorIdentifier = (String) invocationOnMock.getArguments()[0];
- final TypeSerializer<?> keySerializer = (TypeSerializer<?>) invocationOnMock.getArguments()[1];
- MemoryStateBackend backend = MemoryStateBackend.create();
- backend.initializeForJob(env, operatorIdentifier, keySerializer);
- return backend;
- }
- }).when(task).createStateBackend(any(String.class), any(TypeSerializer.class));
- } catch (Exception e) {
- e.printStackTrace();
- }
-
return task;
}
@@ -1040,9 +975,17 @@ public class AggregatingAlignedProcessingTimeWindowOperatorTest {
private static StreamConfig createTaskConfig(KeySelector<?, ?> partitioner, TypeSerializer<?> keySerializer, KeyGroupAssigner<?> keyGroupAssigner) {
StreamConfig cfg = new StreamConfig(new Configuration());
- cfg.setStatePartitioner(0, partitioner);
- cfg.setStateKeySerializer(keySerializer);
- cfg.setKeyGroupAssigner(keyGroupAssigner);
return cfg;
}
+
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ private <T> List<T> extractFromStreamRecords(Iterable<Object> input) {
+ List<T> result = new ArrayList<>();
+ for (Object in : input) {
+ if (in instanceof StreamRecord) {
+ result.add((T) ((StreamRecord) in).getValue());
+ }
+ }
+ return result;
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/EvictingWindowOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/EvictingWindowOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/EvictingWindowOperatorTest.java
index dc71440..681a334 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/EvictingWindowOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/EvictingWindowOperatorTest.java
@@ -41,6 +41,7 @@ import org.apache.flink.streaming.api.windowing.windows.Window;
import org.apache.flink.streaming.runtime.operators.windowing.functions.InternalIterableWindowFunction;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecordSerializer;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.TestHarnessUtil;
import org.apache.flink.util.Collector;
@@ -85,9 +86,7 @@ public class EvictingWindowOperatorTest {
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
long initialTime = 0L;
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -157,9 +156,7 @@ public class EvictingWindowOperatorTest {
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
long initialTime = 0L;
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -227,9 +224,7 @@ public class EvictingWindowOperatorTest {
operator.setInputType(inputType, new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
long initialTime = 0L;
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java
index cda6e1e..67a6f55 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/operators/windowing/WindowOperatorTest.java
@@ -64,6 +64,7 @@ import org.apache.flink.streaming.runtime.operators.windowing.functions.Internal
import org.apache.flink.streaming.runtime.operators.windowing.functions.InternalSingleValueWindowFunction;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.TestTimeServiceProvider;
+import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.TestHarnessUtil;
import org.apache.flink.streaming.util.WindowingTestHarness;
@@ -182,9 +183,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(inputType, new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
testHarness.setup();
testHarness.open();
@@ -220,9 +219,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(inputType, new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
testHarness.open();
@@ -323,9 +320,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
testHarness.open();
@@ -359,9 +354,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
testHarness.open();
@@ -398,9 +391,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -474,9 +465,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -552,9 +541,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -659,9 +646,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -721,10 +706,8 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
-
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
testHarness.open();
@@ -814,10 +797,8 @@ public class WindowOperatorTest extends TestLogger {
"Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
-
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
testHarness.open();
@@ -837,17 +818,39 @@ public class WindowOperatorTest extends TestLogger {
// do a snapshot, close and restore again
StreamStateHandle snapshot = testHarness.snapshot(0L, 0L);
+
testHarness.close();
+
+ ConcurrentLinkedQueue<Object> outputBeforeClose = testHarness.getOutput();
+
+ stateDesc = new ReducingStateDescriptor<>("window-contents",
+ new SumReducer(),
+ inputType.createSerializer(new ExecutionConfig()));
+
+ operator = new WindowOperator<>(
+ GlobalWindows.create(),
+ new GlobalWindow.Serializer(),
+ new TupleKeySelector(),
+ BasicTypeInfo.STRING_TYPE_INFO.createSerializer(new ExecutionConfig()),
+ stateDesc,
+ new InternalSingleValueWindowFunction<>(new PassThroughWindowFunction<String, GlobalWindow, Tuple2<String, Integer>>()),
+ PurgingTrigger.of(CountTrigger.of(WINDOW_SIZE)),
+ 0);
+
+ testHarness = new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+
+ operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse(
+ "Tuple2<String, Integer>"), new ExecutionConfig());
+
testHarness.setup();
testHarness.restore(snapshot);
testHarness.open();
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key2", 1), 1000));
-
expectedOutput.add(new StreamRecord<>(new Tuple2<>("key2", 4), Long.MAX_VALUE));
- TestHarnessUtil.assertOutputEqualsSorted("Output was not correct.", expectedOutput, testHarness.getOutput(), new Tuple2ResultSortComparator());
+ TestHarnessUtil.assertOutputEqualsSorted("Output was not correct.", expectedOutput, Iterables.concat(outputBeforeClose, testHarness.getOutput()), new Tuple2ResultSortComparator());
testHarness.processElement(new StreamRecord<>(new Tuple2<>("key1", 1), 10999));
@@ -858,7 +861,10 @@ public class WindowOperatorTest extends TestLogger {
expectedOutput.add(new StreamRecord<>(new Tuple2<>("key1", 4), Long.MAX_VALUE));
expectedOutput.add(new StreamRecord<>(new Tuple2<>("key2", 4), Long.MAX_VALUE));
- TestHarnessUtil.assertOutputEqualsSorted("Output was not correct.", expectedOutput, testHarness.getOutput(), new Tuple2ResultSortComparator());
+ System.out.println("BEFORE GOT: " + outputBeforeClose);
+ System.out.println("GOT: " + testHarness.getOutput());
+
+ TestHarnessUtil.assertOutputEqualsSorted("Output was not correct.", expectedOutput, Iterables.concat(outputBeforeClose, testHarness.getOutput()), new Tuple2ResultSortComparator());
testHarness.close();
}
@@ -887,9 +893,7 @@ public class WindowOperatorTest extends TestLogger {
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -922,9 +926,8 @@ public class WindowOperatorTest extends TestLogger {
0);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> otherTestHarness =
- new OneInputStreamOperatorTestHarness<>(otherOperator);
+ new KeyedOneInputStreamOperatorTestHarness<>(otherOperator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
- otherTestHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
otherOperator.setInputType(inputType, new ExecutionConfig());
otherTestHarness.setup();
@@ -959,9 +962,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator, new ExecutionConfig(), testTimeProvider);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new ExecutionConfig(), testTimeProvider, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -1021,9 +1022,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator, new ExecutionConfig(), testTimeProvider);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new ExecutionConfig(), testTimeProvider, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -1096,9 +1095,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator, new ExecutionConfig(), testTimeProvider);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new ExecutionConfig(), testTimeProvider, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
@@ -1163,9 +1160,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -1226,9 +1221,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -1295,9 +1288,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -1358,9 +1349,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -1438,9 +1427,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
testHarness.open();
@@ -1532,9 +1519,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
testHarness.open();
@@ -1620,9 +1605,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
testHarness.open();
@@ -1708,9 +1691,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
testHarness.open();
@@ -1805,9 +1786,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
testHarness.open();
@@ -1894,9 +1873,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
testHarness.open();
@@ -1981,9 +1958,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, String> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -2038,9 +2013,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -2087,9 +2060,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -2144,9 +2115,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -2192,9 +2161,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
@@ -2241,9 +2208,7 @@ public class WindowOperatorTest extends TestLogger {
operator.setInputType(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), new ExecutionConfig());
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple3<String, Long, Long>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
testHarness.open();
@@ -2294,9 +2259,7 @@ public class WindowOperatorTest extends TestLogger {
LATENESS);
OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Tuple2<String, Integer>> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, new TupleKeySelector(), BasicTypeInfo.STRING_TYPE_INFO); ;
operator.setInputType(inputType, new ExecutionConfig());
testHarness.open();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java
index 1d07bdd..f8b4063 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java
@@ -129,7 +129,7 @@ public class InterruptSensitiveRestoreTest {
new ExecutionAttemptID(),
new SerializedValue<>(new ExecutionConfig()),
"test task name",
- 0, 1, 0,
+ 1, 0, 1, 0,
new Configuration(),
taskConfig,
SourceStreamTask.class.getName(),
@@ -170,7 +170,7 @@ public class InterruptSensitiveRestoreTest {
private static class InterruptLockingStateHandle extends AbstractCloseableHandle implements StreamStateHandle {
@Override
- public FSDataInputStream openInputStream() throws Exception {
+ public FSDataInputStream openInputStream() throws IOException {
ensureNotClosed();
FSDataInputStream is = new FSDataInputStream() {
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java
index f757943..7ef0080 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/OneInputStreamTaskTest.java
@@ -26,6 +26,7 @@ import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.fs.FSDataInputStream;
import org.apache.flink.core.fs.FSDataOutputStream;
+import org.apache.flink.core.testutils.OneShotLatch;
import org.apache.flink.runtime.io.network.api.CheckpointBarrier;
import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
@@ -359,14 +360,16 @@ public class OneInputStreamTaskTest extends TestLogger {
streamTask.triggerCheckpoint(checkpointId, checkpointTimestamp);
- testHarness.endInput();
- testHarness.waitForTaskCompletion(deadline.timeLeft().toMillis());
-
// since no state was set, there shouldn't be restore calls
assertEquals(0, TestingStreamOperator.numberRestoreCalls);
+ env.getCheckpointLatch().await();
+
assertEquals(checkpointId, env.getCheckpointId());
+ testHarness.endInput();
+ testHarness.waitForTaskCompletion(deadline.timeLeft().toMillis());
+
final OneInputStreamTask<String, String> restoredTask = new OneInputStreamTask<String, String>();
restoredTask.setInitialState(env.getState(), env.getKeyGroupStates());
@@ -459,9 +462,11 @@ public class OneInputStreamTaskTest extends TestLogger {
}
private static class AcknowledgeStreamMockEnvironment extends StreamMockEnvironment {
- private long checkpointId;
- private ChainedStateHandle<StreamStateHandle> state;
- private List<KeyGroupsStateHandle> keyGroupStates;
+ private volatile long checkpointId;
+ private volatile ChainedStateHandle<StreamStateHandle> state;
+ private volatile List<KeyGroupsStateHandle> keyGroupStates;
+
+ private final OneShotLatch checkpointLatch = new OneShotLatch();
public long getCheckpointId() {
return checkpointId;
@@ -494,6 +499,11 @@ public class OneInputStreamTaskTest extends TestLogger {
this.checkpointId = checkpointId;
this.state = state;
this.keyGroupStates = keyGroupStates;
+ checkpointLatch.trigger();
+ }
+
+ public OneShotLatch getCheckpointLatch() {
+ return checkpointLatch;
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java
index 5e82569..0901b32 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java
@@ -103,7 +103,7 @@ public class StreamMockEnvironment implements Environment {
public StreamMockEnvironment(Configuration jobConfig, Configuration taskConfig, ExecutionConfig executionConfig,
long memorySize, MockInputSplitProvider inputSplitProvider, int bufferSize) {
- this.taskInfo = new TaskInfo("", 0, 1, 0);
+ this.taskInfo = new TaskInfo("", 1, 0, 1, 0);
this.jobConfiguration = jobConfig;
this.taskConfiguration = taskConfig;
this.inputs = new LinkedList<InputGate>();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
index 3d9d50f..408b5b1 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskTest.java
@@ -237,7 +237,7 @@ public class StreamTaskTest {
TaskDeploymentDescriptor tdd = new TaskDeploymentDescriptor(
new JobID(), "Job Name", new JobVertexID(), new ExecutionAttemptID(),
new SerializedValue<>(new ExecutionConfig()),
- "Test Task", 0, 1, 0,
+ "Test Task", 1, 0, 1, 0,
new Configuration(),
taskConfig.getConfiguration(),
invokable.getName(),
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
new file mode 100644
index 0000000..5594193
--- /dev/null
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/KeyedOneInputStreamOperatorTestHarness.java
@@ -0,0 +1,201 @@
+/**
+ * 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.util;
+
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.java.ClosureCleaner;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.KeyedStateBackend;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.memory.MemoryStateBackend;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.runtime.tasks.TimeServiceProvider;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.util.Collections;
+import java.util.concurrent.RunnableFuture;
+
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.*;
+
+/**
+ * Extension of {@link OneInputStreamOperatorTestHarness} that allows the operator to get
+ * a {@link KeyedStateBackend}.
+ *
+ */
+public class KeyedOneInputStreamOperatorTestHarness<K, IN, OUT>
+ extends OneInputStreamOperatorTestHarness<IN, OUT> {
+
+ // in case the operator creates one we store it here so that we
+ // can snapshot its state
+ private KeyedStateBackend<?> keyedStateBackend = null;
+
+ // when we restore we keep the state here so that we can call restore
+ // when the operator requests the keyed state backend
+ private KeyGroupsStateHandle restoredKeyedState = null;
+
+ public KeyedOneInputStreamOperatorTestHarness(
+ OneInputStreamOperator<IN, OUT> operator,
+ final KeySelector<IN, K> keySelector,
+ TypeInformation<K> keyType) {
+ super(operator);
+
+ ClosureCleaner.clean(keySelector, false);
+ config.setStatePartitioner(0, keySelector);
+ config.setStateKeySerializer(keyType.createSerializer(executionConfig));
+ config.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(MAX_PARALLELISM));
+
+ setupMockTaskCreateKeyedBackend();
+ }
+
+ public KeyedOneInputStreamOperatorTestHarness(OneInputStreamOperator<IN, OUT> operator,
+ ExecutionConfig executionConfig,
+ KeySelector<IN, K> keySelector,
+ TypeInformation<K> keyType) {
+ super(operator, executionConfig);
+
+ ClosureCleaner.clean(keySelector, false);
+ config.setStatePartitioner(0, keySelector);
+ config.setStateKeySerializer(keyType.createSerializer(executionConfig));
+ config.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(MAX_PARALLELISM));
+
+ setupMockTaskCreateKeyedBackend();
+ }
+
+ public KeyedOneInputStreamOperatorTestHarness(OneInputStreamOperator<IN, OUT> operator,
+ ExecutionConfig executionConfig,
+ TimeServiceProvider testTimeProvider,
+ KeySelector<IN, K> keySelector,
+ TypeInformation<K> keyType) {
+ super(operator, executionConfig, testTimeProvider);
+
+ ClosureCleaner.clean(keySelector, false);
+ config.setStatePartitioner(0, keySelector);
+ config.setStateKeySerializer(keyType.createSerializer(executionConfig));
+ config.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(MAX_PARALLELISM));
+
+ setupMockTaskCreateKeyedBackend();
+ }
+
+ private void setupMockTaskCreateKeyedBackend() {
+
+ try {
+ doAnswer(new Answer<KeyedStateBackend>() {
+ @Override
+ public KeyedStateBackend answer(InvocationOnMock invocationOnMock) throws Throwable {
+
+ final TypeSerializer keySerializer = (TypeSerializer) invocationOnMock.getArguments()[0];
+ final KeyGroupAssigner keyGroupAssigner = (KeyGroupAssigner) invocationOnMock.getArguments()[1];
+ final KeyGroupRange keyGroupRange = (KeyGroupRange) invocationOnMock.getArguments()[2];
+
+ if (restoredKeyedState == null) {
+ keyedStateBackend = stateBackend.createKeyedStateBackend(
+ mockTask.getEnvironment(),
+ new JobID(),
+ "test_op",
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange,
+ mockTask.getEnvironment().getTaskKvStateRegistry());
+ return keyedStateBackend;
+ } else {
+ keyedStateBackend = stateBackend.restoreKeyedStateBackend(
+ mockTask.getEnvironment(),
+ new JobID(),
+ "test_op",
+ keySerializer,
+ keyGroupAssigner,
+ keyGroupRange,
+ Collections.singletonList(restoredKeyedState),
+ mockTask.getEnvironment().getTaskKvStateRegistry());
+ restoredKeyedState = null;
+ return keyedStateBackend;
+ }
+ }
+ }).when(mockTask).createKeyedStateBackend(any(TypeSerializer.class), any(KeyGroupAssigner.class), any(KeyGroupRange.class));
+ } catch (Exception e) {
+ throw new RuntimeException(e.getMessage(), e);
+ }
+ }
+
+ /**
+ * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#snapshotState(org.apache.flink.core.fs.FSDataOutputStream, long, long)} ()}
+ */
+ @Override
+ public StreamStateHandle snapshot(long checkpointId, long timestamp) throws Exception {
+ // simply use an in-memory handle
+ MemoryStateBackend backend = new MemoryStateBackend();
+
+ CheckpointStreamFactory streamFactory = backend.createStreamFactory(new JobID(), "test_op");
+ CheckpointStreamFactory.CheckpointStateOutputStream outStream =
+ streamFactory.createCheckpointStateOutputStream(checkpointId, timestamp);
+
+ operator.snapshotState(outStream, checkpointId, timestamp);
+
+ if (keyedStateBackend != null) {
+ RunnableFuture<KeyGroupsStateHandle> keyedSnapshotRunnable = keyedStateBackend.snapshot(checkpointId,
+ timestamp,
+ streamFactory);
+ if(!keyedSnapshotRunnable.isDone()) {
+ Thread runner = new Thread(keyedSnapshotRunnable);
+ runner.start();
+ }
+ outStream.write(1);
+ ObjectOutputStream oos = new ObjectOutputStream(outStream);
+ oos.writeObject(keyedSnapshotRunnable.get());
+ oos.flush();
+ } else {
+ outStream.write(0);
+ }
+ return outStream.closeAndGetHandle();
+ }
+
+ /**
+ * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#restoreState(org.apache.flink.core.fs.FSDataInputStream)} ()}
+ */
+ @Override
+ public void restore(StreamStateHandle snapshot) throws Exception {
+ FSDataInputStream inStream = snapshot.openInputStream();
+ operator.restoreState(inStream);
+
+ byte keyedStatePresent = (byte) inStream.read();
+ if (keyedStatePresent == 1) {
+ ObjectInputStream ois = new ObjectInputStream(inStream);
+ this.restoredKeyedState = (KeyGroupsStateHandle) ois.readObject();
+ }
+ }
+
+ /**
+ * Calls close and dispose on the operator.
+ */
+ public void close() throws Exception {
+ super.close();
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java
index bc255ff..65ed43d 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java
@@ -78,9 +78,7 @@ public class MockContext<IN, OUT> {
KeySelector<IN, KEY> keySelector, TypeInformation<KEY> keyType) throws Exception {
OneInputStreamOperatorTestHarness<IN, OUT> testHarness =
- new OneInputStreamOperatorTestHarness<>(operator);
-
- testHarness.configureForKeyedStream(keySelector, keyType);
+ new KeyedOneInputStreamOperatorTestHarness<>(operator, keySelector, keyType);
testHarness.setup();
testHarness.open();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
index 6cb46d6..78e05b7 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
@@ -18,22 +18,25 @@
package org.apache.flink.streaming.util;
import org.apache.flink.api.common.ExecutionConfig;
-import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.api.java.ClosureCleaner;
-import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.typeutils.TypeExtractor;
import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.operators.testutils.MockEnvironment;
import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.Output;
+import org.apache.flink.streaming.api.operators.StreamOperator;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
@@ -44,6 +47,7 @@ import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import java.util.Collection;
+import java.util.Collections;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Executors;
@@ -63,6 +67,8 @@ import static org.mockito.Mockito.when;
*/
public class OneInputStreamOperatorTestHarness<IN, OUT> {
+ protected static final int MAX_PARALLELISM = 10;
+
final OneInputStreamOperator<IN, OUT> operator;
final ConcurrentLinkedQueue<Object> outputList;
@@ -78,7 +84,7 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
StreamTask<?, ?> mockTask;
// use this as default for tests
- private AbstractStateBackend stateBackend = new MemoryStateBackend();
+ AbstractStateBackend stateBackend = new MemoryStateBackend();
/**
* Whether setup() was called on the operator. This is reset when calling close().
@@ -108,7 +114,7 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
this.executionConfig = executionConfig;
this.checkpointLock = new Object();
- final Environment env = new MockEnvironment("MockTwoInputTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024, underlyingConfig);
+ final Environment env = new MockEnvironment("MockTwoInputTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024, underlyingConfig, executionConfig, MAX_PARALLELISM, 1, 0);
mockTask = mock(StreamTask.class);
timeServiceProvider = testTimeProvider;
@@ -120,21 +126,6 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
when(mockTask.getExecutionConfig()).thenReturn(executionConfig);
when(mockTask.getUserCodeClassLoader()).thenReturn(this.getClass().getClassLoader());
- try {
- doAnswer(new Answer<AbstractStateBackend>() {
- @Override
- public AbstractStateBackend answer(InvocationOnMock invocationOnMock) throws Throwable {
- final String operatorIdentifier = (String) invocationOnMock.getArguments()[0];
- final TypeSerializer<?> keySerializer = (TypeSerializer<?>) invocationOnMock.getArguments()[1];
- OneInputStreamOperatorTestHarness.this.stateBackend.disposeAllStateForCurrentJob();
- OneInputStreamOperatorTestHarness.this.stateBackend.initializeForJob(env, operatorIdentifier, keySerializer);
- return OneInputStreamOperatorTestHarness.this.stateBackend;
- }
- }).when(mockTask).createStateBackend(any(String.class), any(TypeSerializer.class));
- } catch (Exception e) {
- throw new RuntimeException(e.getMessage(), e);
- }
-
doAnswer(new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
@@ -153,6 +144,20 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
return timeServiceProvider.getCurrentProcessingTime();
}
}).when(mockTask).getCurrentProcessingTime();
+
+ try {
+ doAnswer(new Answer<CheckpointStreamFactory>() {
+ @Override
+ public CheckpointStreamFactory answer(InvocationOnMock invocationOnMock) throws Throwable {
+
+ final StreamOperator operator = (StreamOperator) invocationOnMock.getArguments()[0];
+ return stateBackend.createStreamFactory(new JobID(), operator.getClass().getSimpleName());
+ }
+ }).when(mockTask).createCheckpointStreamFactory(any(StreamOperator.class));
+ } catch (Exception e) {
+ throw new RuntimeException(e.getMessage(), e);
+ }
+
}
public void setStateBackend(AbstractStateBackend stateBackend) {
@@ -167,13 +172,6 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
return this.mockTask.getEnvironment();
}
- public <K> void configureForKeyedStream(KeySelector<IN, K> keySelector, TypeInformation<K> keyType) {
- ClosureCleaner.clean(keySelector, false);
- config.setStatePartitioner(0, keySelector);
- config.setStateKeySerializer(keyType.createSerializer(executionConfig));
- config.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(10));
- }
-
/**
* Get all the output from the task. This contains StreamRecords and Events interleaved. Use
* {@link org.apache.flink.streaming.util.TestHarnessUtil#getStreamRecordsFromOutput(java.util.List)}
@@ -205,13 +203,12 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
}
/**
- * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#snapshotState(org.apache.flink.core.fs.FSDataOutputStream, long, long)} ()}
+ * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#snapshotState(FSDataOutputStream, long, long)} ()}
*/
public StreamStateHandle snapshot(long checkpointId, long timestamp) throws Exception {
- // simply use an in-memory handle
- MemoryStateBackend backend = new MemoryStateBackend();
- AbstractStateBackend.CheckpointStateOutputStream outStream =
- backend.createCheckpointStateOutputStream(checkpointId, timestamp);
+ CheckpointStreamFactory.CheckpointStateOutputStream outStream = stateBackend.createStreamFactory(
+ new JobID(),
+ "test_op").createCheckpointStateOutputStream(checkpointId, timestamp);
operator.snapshotState(outStream, checkpointId, timestamp);
return outStream.closeAndGetHandle();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/TestHarnessUtil.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/TestHarnessUtil.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/TestHarnessUtil.java
index bc4074f..58e8c6b 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/TestHarnessUtil.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/TestHarnessUtil.java
@@ -17,6 +17,7 @@
*/
package org.apache.flink.streaming.util;
+import com.google.common.collect.Iterables;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.junit.Assert;
@@ -76,8 +77,8 @@ public class TestHarnessUtil {
/**
* Compare the two queues containing operator/task output by converting them to an array first.
*/
- public static void assertOutputEqualsSorted(String message, Queue<Object> expected, Queue<Object> actual, Comparator<Object> comparator) {
- assertEquals(expected.size(), actual.size());
+ public static void assertOutputEqualsSorted(String message, Iterable<Object> expected, Iterable<Object> actual, Comparator<Object> comparator) {
+ assertEquals(Iterables.size(expected), Iterables.size(actual));
// first, compare only watermarks, their position should be deterministic
Iterator<Object> exIt = expected.iterator();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java
index 779436a..af1f3fa 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java
@@ -85,8 +85,7 @@ public class WindowingTestHarness<K, IN, W extends Window> {
operator.setInputType(inputType, executionConfig);
timeServiceProvider = new TestTimeServiceProvider();
- testHarness = new OneInputStreamOperatorTestHarness<>(operator, executionConfig, timeServiceProvider);
- testHarness.configureForKeyedStream(keySelector, keyType);
+ testHarness = new KeyedOneInputStreamOperatorTestHarness<>(operator, executionConfig, timeServiceProvider, keySelector, keyType);
}
/**
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
index 97c8339..9f8ab90 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
@@ -73,7 +73,6 @@ import static org.junit.Assert.*;
@RunWith(Parameterized.class)
public class EventTimeWindowCheckpointingITCase extends TestLogger {
-
private static final int MAX_MEM_STATE_SIZE = 10 * 1024 * 1024;
private static final int PARALLELISM = 4;
@@ -118,20 +117,11 @@ public class EventTimeWindowCheckpointingITCase extends TestLogger {
this.stateBackend = new FsStateBackend("file://" + backups);
break;
}
- case ROCKSDB: {
- String rocksDb = tempFolder.newFolder().getAbsolutePath();
- String rocksDbBackups = tempFolder.newFolder().toURI().toString();
- RocksDBStateBackend rdb = new RocksDBStateBackend(rocksDbBackups, new MemoryStateBackend(MAX_MEM_STATE_SIZE));
- rdb.setDbStoragePath(rocksDb);
- this.stateBackend = rdb;
- break;
- }
case ROCKSDB_FULLY_ASYNC: {
String rocksDb = tempFolder.newFolder().getAbsolutePath();
String rocksDbBackups = tempFolder.newFolder().toURI().toString();
RocksDBStateBackend rdb = new RocksDBStateBackend(rocksDbBackups, new MemoryStateBackend(MAX_MEM_STATE_SIZE));
rdb.setDbStoragePath(rocksDb);
- rdb.enableFullyAsyncSnapshots();
this.stateBackend = rdb;
break;
}
@@ -774,14 +764,13 @@ public class EventTimeWindowCheckpointingITCase extends TestLogger {
return Arrays.asList(new Object[][] {
{StateBackendEnum.MEM},
{StateBackendEnum.FILE},
- {StateBackendEnum.ROCKSDB},
{StateBackendEnum.ROCKSDB_FULLY_ASYNC}
}
);
}
private enum StateBackendEnum {
- MEM, FILE, ROCKSDB, ROCKSDB_FULLY_ASYNC
+ MEM, FILE, ROCKSDB_FULLY_ASYNC
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java b/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java
index 3bc3cf5..8b56d3d 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java
@@ -310,7 +310,7 @@ public class ClassLoaderITCase extends TestLogger {
// Success :-)
LOG.info("Disposed savepoint at " + savepointPath);
} else if (disposeResponse instanceof DisposeSavepointFailure) {
- throw new IllegalStateException("Failed to dispose savepoint");
+ throw new IllegalStateException("Failed to dispose savepoint " + disposeResponse);
} else {
throw new IllegalStateException("Unexpected response to DisposeSavepoint");
}
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-tests/src/test/java/org/apache/flink/test/classloading/jar/CustomKvStateProgram.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/classloading/jar/CustomKvStateProgram.java b/flink-tests/src/test/java/org/apache/flink/test/classloading/jar/CustomKvStateProgram.java
index c6a4c7f..8de4797 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/classloading/jar/CustomKvStateProgram.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/classloading/jar/CustomKvStateProgram.java
@@ -18,11 +18,13 @@
package org.apache.flink.test.classloading.jar;
+import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ReducingState;
import org.apache.flink.api.common.state.ReducingStateDescriptor;
import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.streaming.api.datastream.DataStream;
@@ -56,14 +58,23 @@ public class CustomKvStateProgram {
env.setStateBackend(new FsStateBackend(checkpointPath));
DataStream<Integer> source = env.addSource(new InfiniteIntegerSource());
- source.keyBy(new KeySelector<Integer, Integer>() {
- private static final long serialVersionUID = -9044152404048903826L;
-
- @Override
- public Integer getKey(Integer value) throws Exception {
- return ThreadLocalRandom.current().nextInt(parallelism);
- }
- }).flatMap(new ReducingStateFlatMap()).writeAsText(outputPath);
+ source
+ .map(new MapFunction<Integer, Tuple2<Integer, Integer>>() {
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public Tuple2<Integer, Integer> map(Integer value) throws Exception {
+ return new Tuple2<>(ThreadLocalRandom.current().nextInt(parallelism), value);
+ }
+ })
+ .keyBy(new KeySelector<Tuple2<Integer,Integer>, Integer>() {
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public Integer getKey(Tuple2<Integer, Integer> value) throws Exception {
+ return value.f0;
+ }
+ }).flatMap(new ReducingStateFlatMap()).writeAsText(outputPath);
env.execute();
}
@@ -88,10 +99,10 @@ public class CustomKvStateProgram {
}
}
- private static class ReducingStateFlatMap extends RichFlatMapFunction<Integer, Integer> {
+ private static class ReducingStateFlatMap extends RichFlatMapFunction<Tuple2<Integer, Integer>, Integer> {
private static final long serialVersionUID = -5939722892793950253L;
- private ReducingState<Integer> kvState;
+ private transient ReducingState<Integer> kvState;
@Override
public void open(Configuration parameters) throws Exception {
@@ -106,11 +117,13 @@ public class CustomKvStateProgram {
@Override
- public void flatMap(Integer value, Collector<Integer> out) throws Exception {
- kvState.add(value);
+ public void flatMap(Tuple2<Integer, Integer> value, Collector<Integer> out) throws Exception {
+ kvState.add(value.f1);
}
private static class ReduceSum implements ReduceFunction<Integer> {
+ private static final long serialVersionUID = 1L;
+
@Override
public Integer reduce(Integer value1, Integer value2) throws Exception {
return value1 + value2;
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java b/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java
index b3e2137..77a9b2e 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java
@@ -18,8 +18,6 @@
package org.apache.flink.test.state;
-import org.apache.flink.runtime.state.KvStateSnapshot;
-
import org.apache.flink.runtime.state.StateObject;
import org.junit.Test;
@@ -51,18 +49,6 @@ public class StateHandleSerializationTest {
for (Class<?> clazz : stateHandleImplementations) {
validataSerialVersionUID(clazz);
}
-
- // check all key/value snapshots
-
- @SuppressWarnings("unchecked")
- Set<Class<?>> kvStateSnapshotImplementations = (Set<Class<?>>) (Set<?>)
- reflections.getSubTypesOf(KvStateSnapshot.class);
-
- System.out.println(kvStateSnapshotImplementations);
-
- for (Class<?> clazz : kvStateSnapshotImplementations) {
- validataSerialVersionUID(clazz);
- }
}
catch (Exception e) {
e.printStackTrace();
http://git-wip-us.apache.org/repos/asf/flink/blob/4809f536/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
index 41455cf..d4dd475 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
@@ -18,26 +18,28 @@
package org.apache.flink.test.streaming.runtime;
+import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
-import org.apache.flink.api.common.state.FoldingState;
-import org.apache.flink.api.common.state.FoldingStateDescriptor;
-import org.apache.flink.api.common.state.ListState;
-import org.apache.flink.api.common.state.ListStateDescriptor;
-import org.apache.flink.api.common.state.ReducingState;
-import org.apache.flink.api.common.state.ReducingStateDescriptor;
-import org.apache.flink.api.common.state.ValueState;
-import org.apache.flink.api.common.state.ValueStateDescriptor;
+import org.apache.flink.api.common.state.KeyGroupAssigner;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.client.JobExecutionException;
import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.util.StreamingMultipleProgramsTestBase;
import org.junit.Test;
+import java.io.IOException;
+import java.util.List;
+
import static org.junit.Assert.fail;
public class StateBackendITCase extends StreamingMultipleProgramsTestBase {
@@ -74,7 +76,7 @@ public class StateBackendITCase extends StreamingMultipleProgramsTestBase {
}
})
.print();
-
+
try {
see.execute();
fail();
@@ -89,49 +91,39 @@ public class StateBackendITCase extends StreamingMultipleProgramsTestBase {
public static class FailingStateBackend extends AbstractStateBackend {
-
private static final long serialVersionUID = 1L;
@Override
- public void initializeForJob(Environment env, String operatorIdentifier, TypeSerializer<?> keySerializer) throws Exception {
+ public CheckpointStreamFactory createStreamFactory(JobID jobId,
+ String operatorIdentifier) throws IOException {
throw new SuccessException();
}
@Override
- public void disposeAllStateForCurrentJob() throws Exception {}
-
- @Override
- public void close() throws Exception {}
-
- @Override
- protected <N, T> ValueState<T> createValueState(TypeSerializer<N> namespaceSerializer, ValueStateDescriptor<T> stateDesc) throws Exception {
- return null;
- }
-
- @Override
- protected <N, T> ListState<T> createListState(TypeSerializer<N> namespaceSerializer, ListStateDescriptor<T> stateDesc) throws Exception {
- return null;
- }
-
- @Override
- protected <N, T> ReducingState<T> createReducingState(TypeSerializer<N> namespaceSerializer, ReducingStateDescriptor<T> stateDesc) throws Exception {
- return null;
- }
-
- @Override
- protected <N, T, ACC> FoldingState<T, ACC> createFoldingState(TypeSerializer<N> namespaceSerializer,
- FoldingStateDescriptor<T, ACC> stateDesc) throws Exception {
- return null;
+ public <K> KeyedStateBackend<K> createKeyedStateBackend(Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ TaskKvStateRegistry kvStateRegistry) throws Exception {
+ throw new SuccessException();
}
@Override
- public CheckpointStateOutputStream createCheckpointStateOutputStream(long checkpointID,
- long timestamp) throws Exception {
- return null;
+ public <K> KeyedStateBackend<K> restoreKeyedStateBackend(Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ KeyGroupAssigner<K> keyGroupAssigner,
+ KeyGroupRange keyGroupRange,
+ List<KeyGroupsStateHandle> restoredState,
+ TaskKvStateRegistry kvStateRegistry) throws Exception {
+ throw new SuccessException();
}
}
- static final class SuccessException extends Exception {
+ static final class SuccessException extends IOException {
private static final long serialVersionUID = -9218191172606739598L;
}
[19/27] flink git commit: [FLINK-4381] Refactor State to Prepare For
Key-Group State Backends
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskState.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskState.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskState.java
index ac4503d..9025090 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskState.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskState.java
@@ -18,46 +18,56 @@
package org.apache.flink.runtime.checkpoint;
+import com.google.common.collect.Iterables;
import org.apache.flink.runtime.jobgraph.JobVertexID;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StateObject;
+import org.apache.flink.runtime.state.StateUtil;
+import org.apache.flink.util.Preconditions;
-import java.io.Serializable;
+import java.io.IOException;
import java.util.Collection;
+import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
-import java.util.Set;
/**
- * Simple container class which contains the task state and key-value state handles for the sub
+ * Simple container class which contains the task state and key-group state handles for the sub
* tasks of a {@link org.apache.flink.runtime.jobgraph.JobVertex}.
*
- * This class basically groups all tasks and key groups belonging to the same job vertex together.
+ * This class basically groups all non-partitioned state and key-group state belonging to the same job vertex together.
*/
-public class TaskState implements Serializable {
+public class TaskState implements StateObject {
private static final long serialVersionUID = -4845578005863201810L;
private final JobVertexID jobVertexID;
- /** Map of task states which can be accessed by their sub task index */
+ /** handles to non-partitioned states, subtaskindex -> subtaskstate */
private final Map<Integer, SubtaskState> subtaskStates;
- /** Map of key-value states which can be accessed by their key group index */
- private final Map<Integer, KeyGroupState> kvStates;
+ /** handles to partitioned states, subtaskindex -> keyed state */
+ private final Map<Integer, KeyGroupsStateHandle> keyGroupsStateHandles;
- /** Parallelism of the operator when it was checkpointed */
+ /** parallelism of the operator when it was checkpointed */
private final int parallelism;
- public TaskState(JobVertexID jobVertexID, int parallelism) {
- this.jobVertexID = jobVertexID;
+ /** maximum parallelism of the operator when the job was first created */
+ private final int maxParallelism;
- this.subtaskStates = new HashMap<>(parallelism);
+ public TaskState(JobVertexID jobVertexID, int parallelism, int maxParallelism) {
+ Preconditions.checkArgument(
+ parallelism <= maxParallelism,
+ "Parallelism " + parallelism + " is not smaller or equal to max parallelism " + maxParallelism + ".");
- this.kvStates = new HashMap<>();
+ this.jobVertexID = jobVertexID;
+ //preallocate lists of the required size, so that we can randomly set values to indexes
+ this.subtaskStates = new HashMap<>(parallelism);
+ this.keyGroupsStateHandles = new HashMap<>(parallelism);
this.parallelism = parallelism;
+ this.maxParallelism = maxParallelism;
}
public JobVertexID getJobVertexID() {
@@ -65,6 +75,8 @@ public class TaskState implements Serializable {
}
public void putState(int subtaskIndex, SubtaskState subtaskState) {
+ Preconditions.checkNotNull(subtaskState);
+
if (subtaskIndex < 0 || subtaskIndex >= parallelism) {
throw new IndexOutOfBoundsException("The given sub task index " + subtaskIndex +
" exceeds the maximum number of sub tasks " + subtaskStates.size());
@@ -73,31 +85,38 @@ public class TaskState implements Serializable {
}
}
- public SubtaskState getState(int subtaskIndex) {
+ public void putKeyedState(int subtaskIndex, KeyGroupsStateHandle keyGroupsStateHandle) {
+ Preconditions.checkNotNull(keyGroupsStateHandle);
+
if (subtaskIndex < 0 || subtaskIndex >= parallelism) {
throw new IndexOutOfBoundsException("The given sub task index " + subtaskIndex +
- " exceeds the maximum number of sub tasks " + subtaskStates.size());
+ " exceeds the maximum number of sub tasks " + subtaskStates.size());
} else {
- return subtaskStates.get(subtaskIndex);
+ keyGroupsStateHandles.put(subtaskIndex, keyGroupsStateHandle);
}
}
- public Collection<SubtaskState> getStates() {
- return subtaskStates.values();
- }
-
- public long getStateSize() {
- long result = 0L;
- for (SubtaskState subtaskState : subtaskStates.values()) {
- result += subtaskState.getStateSize();
+ public SubtaskState getState(int subtaskIndex) {
+ if (subtaskIndex < 0 || subtaskIndex >= parallelism) {
+ throw new IndexOutOfBoundsException("The given sub task index " + subtaskIndex +
+ " exceeds the maximum number of sub tasks " + subtaskStates.size());
+ } else {
+ return subtaskStates.get(subtaskIndex);
}
+ }
- for (KeyGroupState keyGroupState : kvStates.values()) {
- result += keyGroupState.getStateSize();
+ public KeyGroupsStateHandle getKeyGroupState(int subtaskIndex) {
+ if (subtaskIndex < 0 || subtaskIndex >= parallelism) {
+ throw new IndexOutOfBoundsException("The given sub task index " + subtaskIndex +
+ " exceeds the maximum number of sub tasks " + keyGroupsStateHandles.size());
+ } else {
+ return keyGroupsStateHandles.get(subtaskIndex);
}
+ }
- return result;
+ public Collection<SubtaskState> getStates() {
+ return subtaskStates.values();
}
public int getNumberCollectedStates() {
@@ -108,48 +127,44 @@ public class TaskState implements Serializable {
return parallelism;
}
- public void putKvState(int keyGroupId, KeyGroupState keyGroupState) {
- kvStates.put(keyGroupId, keyGroupState);
+ public int getMaxParallelism() {
+ return maxParallelism;
}
- public KeyGroupState getKvState(int keyGroupId) {
- return kvStates.get(keyGroupId);
+ public Collection<KeyGroupsStateHandle> getKeyGroupStates() {
+ return keyGroupsStateHandles.values();
}
- /**
- * Retrieve the set of key-value state key groups specified by the given key group partition set.
- * The key groups are returned as a map where the key group index maps to the serialized state
- * handle of the key group.
- *
- * @param keyGroupPartition Set of key group indices
- * @return Map of serialized key group state handles indexed by their key group index.
- */
- public Map<Integer, SerializedValue<StateHandle<?>>> getUnwrappedKvStates(Set<Integer> keyGroupPartition) {
- HashMap<Integer, SerializedValue<StateHandle<?>>> result = new HashMap<>(keyGroupPartition.size());
-
- for (Integer keyGroupId : keyGroupPartition) {
- KeyGroupState keyGroupState = kvStates.get(keyGroupId);
-
- if (keyGroupState != null) {
- result.put(keyGroupId, kvStates.get(keyGroupId).getKeyGroupState());
+ public boolean hasNonPartitionedState() {
+ for(SubtaskState sts : subtaskStates.values()) {
+ if (sts != null && !sts.getChainedStateHandle().isEmpty()) {
+ return true;
}
}
-
- return result;
+ return false;
}
- public int getNumberCollectedKvStates() {
- return kvStates.size();
+ @Override
+ public void discardState() throws Exception {
+ StateUtil.bestEffortDiscardAllStateObjects(
+ Iterables.concat(subtaskStates.values(), keyGroupsStateHandles.values()));
}
- public void discard(ClassLoader classLoader) throws Exception {
- for (SubtaskState subtaskState : subtaskStates.values()) {
- subtaskState.discard(classLoader);
- }
- for (KeyGroupState keyGroupState : kvStates.values()) {
- keyGroupState.discard(classLoader);
+ @Override
+ public long getStateSize() throws Exception {
+ long result = 0L;
+
+ for (int i = 0; i < parallelism; i++) {
+ if (subtaskStates.get(i) != null) {
+ result += subtaskStates.get(i).getStateSize();
+ }
+ if (keyGroupsStateHandles.get(i) != null) {
+ result += keyGroupsStateHandles.get(i).getStateSize();
+ }
}
+
+ return result;
}
@Override
@@ -158,7 +173,7 @@ public class TaskState implements Serializable {
TaskState other = (TaskState) obj;
return jobVertexID.equals(other.jobVertexID) && parallelism == other.parallelism &&
- subtaskStates.equals(other.subtaskStates) && kvStates.equals(other.kvStates);
+ subtaskStates.equals(other.subtaskStates) && keyGroupsStateHandles.equals(other.keyGroupsStateHandles);
} else {
return false;
}
@@ -166,6 +181,20 @@ public class TaskState implements Serializable {
@Override
public int hashCode() {
- return parallelism + 31 * Objects.hash(jobVertexID, subtaskStates, kvStates);
+ return parallelism + 31 * Objects.hash(jobVertexID, subtaskStates, keyGroupsStateHandles);
+ }
+
+ @Override
+ public void close() throws IOException {
+ StateUtil.bestEffortCloseAllStateObjects(
+ Iterables.concat(subtaskStates.values(), keyGroupsStateHandles.values()));
+ }
+
+ public Map<Integer, SubtaskState> getSubtaskStates() {
+ return Collections.unmodifiableMap(subtaskStates);
+ }
+
+ public Map<Integer, KeyGroupsStateHandle> getKeyGroupsStateHandles() {
+ return Collections.unmodifiableMap(keyGroupsStateHandles);
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStore.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStore.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStore.java
index 376ef70..b826d9f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStore.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/ZooKeeperCompletedCheckpointStore.java
@@ -24,9 +24,9 @@ import org.apache.curator.framework.api.CuratorEvent;
import org.apache.curator.framework.api.CuratorEventType;
import org.apache.curator.utils.ZKPaths;
import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
+import org.apache.flink.runtime.zookeeper.RetrievableStateStorageHelper;
import org.apache.flink.runtime.jobmanager.HighAvailabilityMode;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.zookeeper.StateStorageHelper;
import org.apache.flink.runtime.zookeeper.ZooKeeperStateHandleStore;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -79,7 +79,7 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
private final ClassLoader userClassLoader;
/** Local completed checkpoints. */
- private final ArrayDeque<Tuple2<StateHandle<CompletedCheckpoint>, String>> checkpointStateHandles;
+ private final ArrayDeque<Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String>> checkpointStateHandles;
/**
* Creates a {@link ZooKeeperCompletedCheckpointStore} instance.
@@ -101,7 +101,7 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
ClassLoader userClassLoader,
CuratorFramework client,
String checkpointsPath,
- StateStorageHelper<CompletedCheckpoint> stateStorage) throws Exception {
+ RetrievableStateStorageHelper<CompletedCheckpoint> stateStorage) throws Exception {
checkArgument(maxNumberOfCheckpointsToRetain >= 1, "Must retain at least one checkpoint.");
checkNotNull(stateStorage, "State storage");
@@ -143,7 +143,7 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
checkpointStateHandles.clear();
// Get all there is first
- List<Tuple2<StateHandle<CompletedCheckpoint>, String>> initialCheckpoints;
+ List<Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String>> initialCheckpoints;
while (true) {
try {
initialCheckpoints = checkpointsInZooKeeper.getAllSortedByName();
@@ -161,10 +161,10 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
if (numberOfInitialCheckpoints > 0) {
// Take the last one. This is the latest checkpoints, because path names are strictly
// increasing (checkpoint ID).
- Tuple2<StateHandle<CompletedCheckpoint>, String> latest = initialCheckpoints
+ Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String> latest = initialCheckpoints
.get(numberOfInitialCheckpoints - 1);
- CompletedCheckpoint latestCheckpoint = latest.f0.getState(userClassLoader);
+ CompletedCheckpoint latestCheckpoint = latest.f0.retrieveState();
checkpointStateHandles.add(latest);
@@ -193,7 +193,8 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
// First add the new one. If it fails, we don't want to loose existing data.
String path = String.format("/%s", checkpoint.getCheckpointID());
- final StateHandle<CompletedCheckpoint> stateHandle = checkpointsInZooKeeper.add(path, checkpoint);
+ final RetrievableStateHandle<CompletedCheckpoint> stateHandle =
+ checkpointsInZooKeeper.add(path, checkpoint);
checkpointStateHandles.addLast(new Tuple2<>(stateHandle, path));
@@ -211,7 +212,7 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
return null;
}
else {
- return checkpointStateHandles.getLast().f0.getState(userClassLoader);
+ return checkpointStateHandles.getLast().f0.retrieveState();
}
}
@@ -219,8 +220,8 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
public List<CompletedCheckpoint> getAllCheckpoints() throws Exception {
List<CompletedCheckpoint> checkpoints = new ArrayList<>(checkpointStateHandles.size());
- for (Tuple2<StateHandle<CompletedCheckpoint>, String> stateHandle : checkpointStateHandles) {
- checkpoints.add(stateHandle.f0.getState(userClassLoader));
+ for (Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String> stateHandle : checkpointStateHandles) {
+ checkpoints.add(stateHandle.f0.retrieveState());
}
return checkpoints;
@@ -235,7 +236,7 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
public void shutdown() throws Exception {
LOG.info("Shutting down");
- for (Tuple2<StateHandle<CompletedCheckpoint>, String> checkpoint : checkpointStateHandles) {
+ for (Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String> checkpoint : checkpointStateHandles) {
try {
removeFromZooKeeperAndDiscardCheckpoint(checkpoint);
}
@@ -264,7 +265,7 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
* Removes the state handle from ZooKeeper, discards the checkpoints, and the state handle.
*/
private void removeFromZooKeeperAndDiscardCheckpoint(
- final Tuple2<StateHandle<CompletedCheckpoint>, String> stateHandleAndPath) throws Exception {
+ final Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String> stateHandleAndPath) throws Exception {
final BackgroundCallback callback = new BackgroundCallback() {
@Override
@@ -273,16 +274,15 @@ public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointSto
if (event.getType() == CuratorEventType.DELETE) {
if (event.getResultCode() == 0) {
// The checkpoint
- CompletedCheckpoint checkpoint = stateHandleAndPath
- .f0.getState(userClassLoader);
-
- checkpoint.discard(userClassLoader);
-
- // Discard the state handle
- stateHandleAndPath.f0.discardState();
-
- // Discard the checkpoint
- LOG.debug("Discarded " + checkpoint);
+ try {
+ CompletedCheckpoint checkpoint = stateHandleAndPath.f0.retrieveState();
+ checkpoint.discardState();
+ // Discard the checkpoint
+ LOG.debug("Discarded " + checkpoint);
+ } finally {
+ // Discard the state handle
+ stateHandleAndPath.f0.discardState();
+ }
}
else {
throw new IllegalStateException("Unexpected result code " +
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStore.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStore.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStore.java
index 6efc01e..49f51be 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStore.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/FsSavepointStore.java
@@ -142,13 +142,13 @@ public class FsSavepointStore implements SavepointStore {
}
@Override
- public void disposeSavepoint(String path, ClassLoader classLoader) throws Exception {
+ public void disposeSavepoint(String path) throws Exception {
Preconditions.checkNotNull(path, "Path");
- Preconditions.checkNotNull(classLoader, "Class loader");
try {
Savepoint savepoint = loadSavepoint(path);
- savepoint.dispose(classLoader);
+ LOG.info("Disposing savepoint: " + path);
+ savepoint.dispose();
Path filePath = new Path(path);
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/HeapSavepointStore.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/HeapSavepointStore.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/HeapSavepointStore.java
index cf30f5f..2cf8f31 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/HeapSavepointStore.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/HeapSavepointStore.java
@@ -108,9 +108,8 @@ public class HeapSavepointStore implements SavepointStore {
}
@Override
- public void disposeSavepoint(String path, ClassLoader classLoader) throws Exception {
+ public void disposeSavepoint(String path) throws Exception {
Preconditions.checkNotNull(path, "Path");
- Preconditions.checkNotNull(classLoader, "Class loader");
Savepoint savepoint;
synchronized (shutDownLock) {
@@ -118,7 +117,7 @@ public class HeapSavepointStore implements SavepointStore {
}
if (savepoint != null) {
- savepoint.dispose(classLoader);
+ savepoint.dispose();
} else {
throw new IllegalArgumentException("Invalid path '" + path + "'.");
}
@@ -131,7 +130,7 @@ public class HeapSavepointStore implements SavepointStore {
// available at this point.
for (Savepoint savepoint : savepoints.values()) {
try {
- savepoint.dispose(ClassLoader.getSystemClassLoader());
+ savepoint.dispose();
} catch (Throwable t) {
LOG.warn("Failed to dispose savepoint " + savepoint.getCheckpointId(), t);
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/Savepoint.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/Savepoint.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/Savepoint.java
index 7823ab9..643f14c 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/Savepoint.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/Savepoint.java
@@ -68,13 +68,7 @@ public interface Savepoint {
/**
* Disposes the savepoint.
- *
- * <p>The class loader is needed, because savepoints can currently point to
- * arbitrary snapshot {@link org.apache.flink.runtime.state.StateHandle}
- * instances, which need the user code class loader for deserialization.
- *
- * @param classLoader Class loader for disposal
*/
- void dispose(ClassLoader classLoader) throws Exception;
+ void dispose() throws Exception;
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoader.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoader.java
index 1be7a58..47917b4 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoader.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointLoader.java
@@ -60,12 +60,12 @@ public class SavepointLoader {
ExecutionJobVertex executionJobVertex = tasks.get(taskState.getJobVertexID());
if (executionJobVertex != null) {
- if (executionJobVertex.getParallelism() == taskState.getParallelism()) {
+ if (executionJobVertex.getMaxParallelism() == taskState.getMaxParallelism()) {
taskStates.put(taskState.getJobVertexID(), taskState);
}
else {
String msg = String.format("Failed to rollback to savepoint %s. " +
- "Parallelism mismatch between savepoint state and new program. " +
+ "Max parallelism mismatch between savepoint state and new program. " +
"Cannot map operator %s with parallelism %d to new program with " +
"parallelism %d. This indicates that the program has been changed " +
"in a non-compatible way after the savepoint.",
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointSerializers.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointSerializers.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointSerializers.java
index d06f3d0..20b3d89 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointSerializers.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointSerializers.java
@@ -28,10 +28,13 @@ import java.util.Map;
*/
public class SavepointSerializers {
+
+ private static final int SAVEPOINT_VERSION_0 = 0;
private static final Map<Integer, SavepointSerializer<?>> SERIALIZERS = new HashMap<>(1);
static {
- SERIALIZERS.put(SavepointV0.VERSION, SavepointV0Serializer.INSTANCE);
+ SERIALIZERS.put(SAVEPOINT_VERSION_0, null);
+ SERIALIZERS.put(SavepointV1.VERSION, SavepointV1Serializer.INSTANCE);
}
/**
@@ -66,7 +69,7 @@ public class SavepointSerializers {
if (serializer != null) {
return serializer;
} else {
- throw new IllegalArgumentException("Unknown savepoint version " + version + ".");
+ throw new IllegalArgumentException("Cannot restore savepoint version " + version + ".");
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointStore.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointStore.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointStore.java
index 71fcb34..68b88d2 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointStore.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointStore.java
@@ -48,15 +48,10 @@ public interface SavepointStore {
/**
* Disposes the savepoint at the specified path.
*
- * <p>The class loader is needed, because savepoints can currently point to
- * arbitrary snapshot {@link org.apache.flink.runtime.state.StateHandle}
- * instances, which need the user code class loader for deserialization.
- *
* @param path Path of savepoint to dispose
- * @param classLoader Class loader for disposal
* @throws Exception Failures during diposal are forwarded
*/
- void disposeSavepoint(String path, ClassLoader classLoader) throws Exception;
+ void disposeSavepoint(String path) throws Exception;
/**
* Shut downs the savepoint store.
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0.java
deleted file mode 100644
index d60d80e..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0.java
+++ /dev/null
@@ -1,97 +0,0 @@
-/*
- * 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.savepoint;
-
-import org.apache.flink.runtime.checkpoint.TaskState;
-import org.apache.flink.util.Preconditions;
-
-import java.util.ArrayList;
-import java.util.Collection;
-
-/**
- * Savepoint version 0.
- *
- * <p>This format was introduced with Flink 1.1.0.
- */
-public class SavepointV0 implements Savepoint {
-
- /** The savepoint version. */
- public static final int VERSION = 0;
-
- /** The checkpoint ID */
- private final long checkpointId;
-
- /** The task states */
- private final Collection<TaskState> taskStates = new ArrayList();
-
- public SavepointV0(long checkpointId, Collection<TaskState> taskStates) {
- this.checkpointId = checkpointId;
- this.taskStates.addAll(taskStates);
- }
-
- @Override
- public int getVersion() {
- return VERSION;
- }
-
- @Override
- public long getCheckpointId() {
- return checkpointId;
- }
-
- @Override
- public Collection<TaskState> getTaskStates() {
- return taskStates;
- }
-
- @Override
- public void dispose(ClassLoader classLoader) throws Exception {
- Preconditions.checkNotNull(classLoader, "Class loader");
- for (TaskState taskState : taskStates) {
- taskState.discard(classLoader);
- }
- taskStates.clear();
- }
-
- @Override
- public String toString() {
- return "Savepoint(version=" + VERSION + ")";
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
-
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
-
- SavepointV0 that = (SavepointV0) o;
- return checkpointId == that.checkpointId && getTaskStates().equals(that.getTaskStates());
- }
-
- @Override
- public int hashCode() {
- int result = (int) (checkpointId ^ (checkpointId >>> 32));
- result = 31 * result + taskStates.hashCode();
- return result;
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0Serializer.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0Serializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0Serializer.java
deleted file mode 100644
index e82b85f..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV0Serializer.java
+++ /dev/null
@@ -1,186 +0,0 @@
-/*
- * 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.savepoint;
-
-import org.apache.flink.runtime.checkpoint.KeyGroupState;
-import org.apache.flink.runtime.checkpoint.SubtaskState;
-import org.apache.flink.runtime.checkpoint.TaskState;
-import org.apache.flink.runtime.jobgraph.JobVertexID;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
-
-import java.io.DataInputStream;
-import java.io.DataOutputStream;
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-
-/**
- * Serializer for {@link SavepointV0} instances.
- *
- * <p>In contrast to previous savepoint versions, this serializer makes sure
- * that no default Java serialization is used for serialization. Therefore, we
- * don't rely on any involved Java classes to stay the same.
- */
-class SavepointV0Serializer implements SavepointSerializer<SavepointV0> {
-
- public static final SavepointV0Serializer INSTANCE = new SavepointV0Serializer();
-
- private SavepointV0Serializer() {
- }
-
- @Override
- public void serialize(SavepointV0 savepoint, DataOutputStream dos) throws IOException {
- dos.writeLong(savepoint.getCheckpointId());
-
- Collection<TaskState> taskStates = savepoint.getTaskStates();
- dos.writeInt(taskStates.size());
-
- for (TaskState taskState : savepoint.getTaskStates()) {
- // Vertex ID
- dos.writeLong(taskState.getJobVertexID().getLowerPart());
- dos.writeLong(taskState.getJobVertexID().getUpperPart());
-
- // Parallelism
- int parallelism = taskState.getParallelism();
- dos.writeInt(parallelism);
-
- // Sub task states
- dos.writeInt(taskState.getNumberCollectedStates());
-
- for (int i = 0; i < parallelism; i++) {
- SubtaskState subtaskState = taskState.getState(i);
-
- if (subtaskState != null) {
- dos.writeInt(i);
-
- SerializedValue<?> serializedValue = subtaskState.getState();
- if (serializedValue == null) {
- dos.writeInt(-1); // null
- } else {
- byte[] serialized = serializedValue.getByteArray();
- dos.writeInt(serialized.length);
- dos.write(serialized, 0, serialized.length);
- }
-
- dos.writeLong(subtaskState.getStateSize());
- dos.writeLong(subtaskState.getDuration());
- }
- }
-
- // Key group states
- dos.writeInt(taskState.getNumberCollectedKvStates());
-
- for (int i = 0; i < parallelism; i++) {
- KeyGroupState keyGroupState = taskState.getKvState(i);
-
- if (keyGroupState != null) {
- dos.write(i);
-
- SerializedValue<?> serializedValue = keyGroupState.getKeyGroupState();
- if (serializedValue == null) {
- dos.writeInt(-1); // null
- } else {
- byte[] serialized = serializedValue.getByteArray();
- dos.writeInt(serialized.length);
- dos.write(serialized, 0, serialized.length);
- }
-
- dos.writeLong(keyGroupState.getStateSize());
- dos.writeLong(keyGroupState.getDuration());
- }
- }
- }
- }
-
- @Override
- public SavepointV0 deserialize(DataInputStream dis) throws IOException {
- long checkpointId = dis.readLong();
-
- // Task states
- int numTaskStates = dis.readInt();
- List<TaskState> taskStates = new ArrayList<>(numTaskStates);
-
- for (int i = 0; i < numTaskStates; i++) {
- JobVertexID jobVertexId = new JobVertexID(dis.readLong(), dis.readLong());
- int parallelism = dis.readInt();
-
- // Add task state
- TaskState taskState = new TaskState(jobVertexId, parallelism);
- taskStates.add(taskState);
-
- // Sub task states
- int numSubTaskStates = dis.readInt();
- for (int j = 0; j < numSubTaskStates; j++) {
- int subtaskIndex = dis.readInt();
-
- int length = dis.readInt();
-
- SerializedValue<StateHandle<?>> serializedValue;
- if (length == -1) {
- serializedValue = new SerializedValue<>(null);
- } else {
- byte[] serializedData = new byte[length];
- dis.readFully(serializedData, 0, length);
- serializedValue = SerializedValue.fromBytes(serializedData);
- }
-
- long stateSize = dis.readLong();
- long duration = dis.readLong();
-
- SubtaskState subtaskState = new SubtaskState(
- serializedValue,
- stateSize,
- duration);
-
- taskState.putState(subtaskIndex, subtaskState);
- }
-
- // Key group states
- int numKvStates = dis.readInt();
- for (int j = 0; j < numKvStates; j++) {
- int keyGroupIndex = dis.readInt();
-
- int length = dis.readInt();
-
- SerializedValue<StateHandle<?>> serializedValue;
- if (length == -1) {
- serializedValue = new SerializedValue<>(null);
- } else {
- byte[] serializedData = new byte[length];
- dis.readFully(serializedData, 0, length);
- serializedValue = SerializedValue.fromBytes(serializedData);
- }
-
- long stateSize = dis.readLong();
- long duration = dis.readLong();
-
- KeyGroupState keyGroupState = new KeyGroupState(
- serializedValue,
- stateSize,
- duration);
-
- taskState.putKvState(keyGroupIndex, keyGroupState);
- }
- }
-
- return new SavepointV0(checkpointId, taskStates);
- }
-}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1.java
new file mode 100644
index 0000000..5976bbf
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1.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.checkpoint.savepoint;
+
+import org.apache.flink.runtime.checkpoint.TaskState;
+import org.apache.flink.util.Preconditions;
+
+import java.util.Collection;
+
+/**
+ * Savepoint version 0.
+ *
+ * <p>This format was introduced with Flink 1.1.0.
+ */
+public class SavepointV1 implements Savepoint {
+
+ /** The savepoint version. */
+ public static final int VERSION = 1;
+
+ /** The checkpoint ID */
+ private final long checkpointId;
+
+ /** The task states */
+ private final Collection<TaskState> taskStates;
+
+ public SavepointV1(long checkpointId, Collection<TaskState> taskStates) {
+ this.checkpointId = checkpointId;
+ this.taskStates = Preconditions.checkNotNull(taskStates, "Task States");
+ }
+
+ @Override
+ public int getVersion() {
+ return VERSION;
+ }
+
+ @Override
+ public long getCheckpointId() {
+ return checkpointId;
+ }
+
+ @Override
+ public Collection<TaskState> getTaskStates() {
+ return taskStates;
+ }
+
+ @Override
+ public void dispose() throws Exception {
+ for (TaskState taskState : taskStates) {
+ taskState.discardState();
+ }
+ taskStates.clear();
+ }
+
+ @Override
+ public String toString() {
+ return "Savepoint(version=" + VERSION + ")";
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+
+ SavepointV1 that = (SavepointV1) o;
+ return checkpointId == that.checkpointId && getTaskStates().equals(that.getTaskStates());
+ }
+
+ @Override
+ public int hashCode() {
+ int result = (int) (checkpointId ^ (checkpointId >>> 32));
+ result = 31 * result + taskStates.hashCode();
+ return result;
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1Serializer.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1Serializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1Serializer.java
new file mode 100644
index 0000000..8e05b81
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/savepoint/SavepointV1Serializer.java
@@ -0,0 +1,232 @@
+/*
+ * 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.savepoint;
+
+import org.apache.flink.core.fs.Path;
+import org.apache.flink.runtime.checkpoint.SubtaskState;
+import org.apache.flink.runtime.checkpoint.TaskState;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeOffsets;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.filesystem.FileStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Serializer for {@link SavepointV1} instances.
+ * <p>
+ * <p>In contrast to previous savepoint versions, this serializer makes sure
+ * that no default Java serialization is used for serialization. Therefore, we
+ * don't rely on any involved Java classes to stay the same.
+ */
+class SavepointV1Serializer implements SavepointSerializer<SavepointV1> {
+
+ private static final byte NULL_HANDLE = 0;
+ private static final byte BYTE_STREAM_STATE_HANDLE = 1;
+ private static final byte FILE_STREAM_STATE_HANDLE = 2;
+ private static final byte KEY_GROUPS_HANDLE = 3;
+
+
+ public static final SavepointV1Serializer INSTANCE = new SavepointV1Serializer();
+
+ private SavepointV1Serializer() {
+ }
+
+ @Override
+ public void serialize(SavepointV1 savepoint, DataOutputStream dos) throws IOException {
+ try {
+ dos.writeLong(savepoint.getCheckpointId());
+
+ Collection<TaskState> taskStates = savepoint.getTaskStates();
+ dos.writeInt(taskStates.size());
+
+ for (TaskState taskState : savepoint.getTaskStates()) {
+ // Vertex ID
+ dos.writeLong(taskState.getJobVertexID().getLowerPart());
+ dos.writeLong(taskState.getJobVertexID().getUpperPart());
+
+ // Parallelism
+ int parallelism = taskState.getParallelism();
+ dos.writeInt(parallelism);
+ dos.writeInt(taskState.getMaxParallelism());
+
+ // Sub task states
+ Map<Integer, SubtaskState> subtaskStateMap = taskState.getSubtaskStates();
+ dos.writeInt(subtaskStateMap.size());
+ for (Map.Entry<Integer, SubtaskState> entry : subtaskStateMap.entrySet()) {
+ dos.writeInt(entry.getKey());
+
+ SubtaskState subtaskState = entry.getValue();
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle = subtaskState.getChainedStateHandle();
+ dos.writeInt(chainedStateHandle.getLength());
+ for (int j = 0; j < chainedStateHandle.getLength(); ++j) {
+ StreamStateHandle stateHandle = chainedStateHandle.get(j);
+ serializeStreamStateHandle(stateHandle, dos);
+ }
+
+ dos.writeLong(subtaskState.getDuration());
+ }
+
+
+ Map<Integer, KeyGroupsStateHandle> keyGroupsStateHandles = taskState.getKeyGroupsStateHandles();
+ dos.writeInt(keyGroupsStateHandles.size());
+ for (Map.Entry<Integer, KeyGroupsStateHandle> entry : keyGroupsStateHandles.entrySet()) {
+ dos.writeInt(entry.getKey());
+ serializeKeyGroupStateHandle(entry.getValue(), dos);
+ }
+
+ }
+ } catch (Exception e) {
+ throw new IOException(e);
+ }
+ }
+
+ @Override
+ public SavepointV1 deserialize(DataInputStream dis) throws IOException {
+ long checkpointId = dis.readLong();
+
+ // Task states
+ int numTaskStates = dis.readInt();
+ List<TaskState> taskStates = new ArrayList<>(numTaskStates);
+
+ for (int i = 0; i < numTaskStates; i++) {
+ JobVertexID jobVertexId = new JobVertexID(dis.readLong(), dis.readLong());
+ int parallelism = dis.readInt();
+ int maxParallelism = dis.readInt();
+
+ // Add task state
+ TaskState taskState = new TaskState(jobVertexId, parallelism, maxParallelism);
+ taskStates.add(taskState);
+
+ // Sub task states
+ int numSubTaskStates = dis.readInt();
+
+ for (int j = 0; j < numSubTaskStates; j++) {
+ int subtaskIndex = dis.readInt();
+ int chainedStateHandleSize = dis.readInt();
+ List<StreamStateHandle> streamStateHandleList = new ArrayList<>(chainedStateHandleSize);
+ for (int k = 0; k < chainedStateHandleSize; ++k) {
+ StreamStateHandle streamStateHandle = deserializeStreamStateHandle(dis);
+ streamStateHandleList.add(streamStateHandle);
+ }
+
+ long duration = dis.readLong();
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle = new ChainedStateHandle<>(streamStateHandleList);
+ SubtaskState subtaskState = new SubtaskState(chainedStateHandle, duration);
+ taskState.putState(subtaskIndex, subtaskState);
+ }
+
+ // Key group states
+ int numKeyGroupStates = dis.readInt();
+ for (int j = 0; j < numKeyGroupStates; j++) {
+ int keyGroupIndex = dis.readInt();
+
+ KeyGroupsStateHandle keyGroupsStateHandle = deserializeKeyGroupStateHandle(dis);
+ if (keyGroupsStateHandle != null) {
+ taskState.putKeyedState(keyGroupIndex, keyGroupsStateHandle);
+ }
+ }
+ }
+
+ return new SavepointV1(checkpointId, taskStates);
+ }
+
+ public static void serializeKeyGroupStateHandle(KeyGroupsStateHandle stateHandle, DataOutputStream dos) throws IOException {
+ if (stateHandle != null) {
+ dos.writeByte(KEY_GROUPS_HANDLE);
+ dos.writeInt(stateHandle.getGroupRangeOffsets().getKeyGroupRange().getStartKeyGroup());
+ dos.writeInt(stateHandle.getNumberOfKeyGroups());
+ for (int keyGroup : stateHandle.keyGroups()) {
+ dos.writeLong(stateHandle.getOffsetForKeyGroup(keyGroup));
+ }
+ serializeStreamStateHandle(stateHandle.getStateHandle(), dos);
+ } else {
+ dos.writeByte(NULL_HANDLE);
+ }
+ }
+
+ public static KeyGroupsStateHandle deserializeKeyGroupStateHandle(DataInputStream dis) throws IOException {
+ int type = dis.readByte();
+ if (NULL_HANDLE == type) {
+ return null;
+ } else {
+ int startKeyGroup = dis.readInt();
+ int numKeyGroups = dis.readInt();
+ KeyGroupRange keyGroupRange = KeyGroupRange.of(startKeyGroup, startKeyGroup + numKeyGroups - 1);
+ long[] offsets = new long[numKeyGroups];
+ for (int i = 0; i < numKeyGroups; ++i) {
+ offsets[i] = dis.readLong();
+ }
+ KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(keyGroupRange, offsets);
+ StreamStateHandle stateHandle = deserializeStreamStateHandle(dis);
+ return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
+ }
+ }
+
+ public static void serializeStreamStateHandle(StreamStateHandle stateHandle, DataOutputStream dos) throws IOException {
+
+ if (stateHandle == null) {
+ dos.writeByte(NULL_HANDLE);
+
+ } else if (stateHandle instanceof FileStateHandle) {
+ dos.writeByte(FILE_STREAM_STATE_HANDLE);
+ FileStateHandle fileStateHandle = (FileStateHandle) stateHandle;
+ dos.writeUTF(fileStateHandle.getFilePath().toString());
+
+ } else if (stateHandle instanceof ByteStreamStateHandle) {
+ dos.writeByte(BYTE_STREAM_STATE_HANDLE);
+ ByteStreamStateHandle byteStreamStateHandle = (ByteStreamStateHandle) stateHandle;
+ byte[] internalData = byteStreamStateHandle.getData();
+ dos.writeInt(internalData.length);
+ dos.write(byteStreamStateHandle.getData());
+
+ } else {
+ throw new IOException("Unknown implementation of StreamStateHandle: " + stateHandle.getClass());
+ }
+
+ dos.flush();
+ }
+
+ public static StreamStateHandle deserializeStreamStateHandle(DataInputStream dis) throws IOException {
+ int type = dis.read();
+ if (NULL_HANDLE == type) {
+ return null;
+ } else if (FILE_STREAM_STATE_HANDLE == type) {
+ String pathString = dis.readUTF();
+ return new FileStateHandle(new Path(pathString));
+ } else if (BYTE_STREAM_STATE_HANDLE == type) {
+ int numBytes = dis.readInt();
+ byte[] data = new byte[numBytes];
+ dis.read(data);
+ return new ByteStreamStateHandle(data);
+ } else {
+ throw new IOException("Unknown implementation of StreamStateHandle, code: " + type);
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTracker.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTracker.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTracker.java
index 9d47457..2217fd4 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTracker.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/stats/SimpleCheckpointStatsTracker.java
@@ -22,8 +22,8 @@ import org.apache.flink.metrics.Gauge;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.checkpoint.CompletedCheckpoint;
import org.apache.flink.runtime.checkpoint.SubtaskState;
-import org.apache.flink.runtime.executiongraph.ExecutionJobVertex;
import org.apache.flink.runtime.checkpoint.TaskState;
+import org.apache.flink.runtime.executiongraph.ExecutionJobVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import scala.Option;
@@ -140,7 +140,12 @@ public class SimpleCheckpointStatsTracker implements CheckpointStatsTracker {
}
synchronized (statsLock) {
- long overallStateSize = checkpoint.getStateSize();
+ long overallStateSize;
+ try {
+ overallStateSize = checkpoint.getStateSize();
+ } catch (Exception ex) {
+ throw new RuntimeException(ex);
+ }
// Operator stats
Map<JobVertexID, long[][]> statsForSubTasks = new HashMap<>();
@@ -421,7 +426,11 @@ public class SimpleCheckpointStatsTracker implements CheckpointStatsTracker {
private class CheckpointSizeGauge implements Gauge<Long> {
@Override
public Long getValue() {
- return latestCompletedCheckpoint == null ? -1 : latestCompletedCheckpoint.getStateSize();
+ try {
+ return latestCompletedCheckpoint == null ? -1 : latestCompletedCheckpoint.getStateSize();
+ } catch (Exception ex) {
+ throw new RuntimeException(ex);
+ }
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
index 60fb45c..8849e93 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/deployment/TaskDeploymentDescriptor.java
@@ -25,7 +25,9 @@ import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.blob.BlobKey;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
import org.apache.flink.runtime.jobgraph.JobVertexID;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.util.SerializedValue;
import java.io.Serializable;
@@ -33,6 +35,7 @@ import java.net.URL;
import java.util.Collection;
import java.util.List;
+
import static org.apache.flink.util.Preconditions.checkArgument;
import static org.apache.flink.util.Preconditions.checkNotNull;
@@ -88,7 +91,11 @@ public final class TaskDeploymentDescriptor implements Serializable {
/** The list of classpaths required to run this task. */
private final List<URL> requiredClasspaths;
- private final SerializedValue<StateHandle<?>> operatorState;
+ /** Handle to the non-partitioned state of the operator chain */
+ private final ChainedStateHandle<StreamStateHandle> operatorState;
+
+ /** Handle to the key-grouped state of the head operator in the chain */
+ private final List<KeyGroupsStateHandle> keyGroupState;
/** The execution configuration (see {@link ExecutionConfig}) related to the specific job. */
private final SerializedValue<ExecutionConfig> serializedExecutionConfig;
@@ -114,7 +121,8 @@ public final class TaskDeploymentDescriptor implements Serializable {
List<BlobKey> requiredJarFiles,
List<URL> requiredClasspaths,
int targetSlotNumber,
- SerializedValue<StateHandle<?>> operatorState) {
+ ChainedStateHandle<StreamStateHandle> operatorState,
+ List<KeyGroupsStateHandle> keyGroupState) {
checkArgument(indexInSubtaskGroup >= 0);
checkArgument(numberOfSubtasks > indexInSubtaskGroup);
@@ -139,6 +147,7 @@ public final class TaskDeploymentDescriptor implements Serializable {
this.requiredClasspaths = checkNotNull(requiredClasspaths);
this.targetSlotNumber = targetSlotNumber;
this.operatorState = operatorState;
+ this.keyGroupState = keyGroupState;
}
public TaskDeploymentDescriptor(
@@ -178,6 +187,7 @@ public final class TaskDeploymentDescriptor implements Serializable {
requiredJarFiles,
requiredClasspaths,
targetSlotNumber,
+ null,
null);
}
@@ -316,7 +326,11 @@ public final class TaskDeploymentDescriptor implements Serializable {
return strBuilder.toString();
}
- public SerializedValue<StateHandle<?>> getOperatorState() {
+ public ChainedStateHandle<StreamStateHandle> getOperatorState() {
return operatorState;
}
+
+ public List<KeyGroupsStateHandle> getKeyGroupState() {
+ return keyGroupState;
+ }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java b/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java
index 2f158fd..1eee9d4 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/execution/Environment.java
@@ -35,9 +35,12 @@ import org.apache.flink.runtime.jobgraph.tasks.InputSplitProvider;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.KvState;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
+import java.util.List;
import java.util.Map;
import java.util.concurrent.Future;
@@ -166,14 +169,18 @@ public interface Environment {
void acknowledgeCheckpoint(long checkpointId);
/**
- * Confirms that the invokable has successfully completed all steps it needed to
- * to for the checkpoint with the give checkpoint-ID. This method does include
+ * Confirms that the invokable has successfully completed all required steps for
+ * the checkpoint with the give checkpoint-ID. This method does include
* the given state in the checkpoint.
*
* @param checkpointId The ID of the checkpoint.
- * @param state A handle to the state to be included in the checkpoint.
+ * @param chainedStateHandle Handle for the chained operator state
+ * @param keyGroupStateHandles Handles for key group state
*/
- void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state);
+ void acknowledgeCheckpoint(
+ long checkpointId,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
+ List<KeyGroupsStateHandle> keyGroupStateHandles);
/**
* Marks task execution failed for an external reason (a reason other than the task code itself
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
index 5bab780..1981f5b 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/Execution.java
@@ -46,9 +46,10 @@ import org.apache.flink.runtime.jobmanager.scheduler.SlotAllocationFutureAction;
import org.apache.flink.runtime.jobmanager.scheduler.SlotSharingGroup;
import org.apache.flink.runtime.messages.Messages;
import org.apache.flink.runtime.messages.TaskMessages.TaskOperationResult;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.util.SerializableObject;
-import org.apache.flink.util.SerializedValue;
import org.apache.flink.util.ExceptionUtils;
import org.slf4j.Logger;
@@ -134,8 +135,10 @@ public class Execution {
private volatile InstanceConnectionInfo assignedResourceLocation; // for the archived execution
- /** The state with which the execution attempt should start */
- private SerializedValue<StateHandle<?>> operatorState;
+ private ChainedStateHandle<StreamStateHandle> chainedStateHandle;
+
+ private List<KeyGroupsStateHandle> keyGroupsStateHandles;
+
/** The execution context which is used to execute futures. */
private ExecutionContext executionContext;
@@ -215,6 +218,14 @@ public class Execution {
return this.stateTimestamps[state.ordinal()];
}
+ public ChainedStateHandle<StreamStateHandle> getChainedStateHandle() {
+ return chainedStateHandle;
+ }
+
+ public List<KeyGroupsStateHandle> getKeyGroupsStateHandles() {
+ return keyGroupsStateHandles;
+ }
+
public boolean isFinished() {
return state.isTerminal();
}
@@ -234,19 +245,22 @@ public class Execution {
partialInputChannelDeploymentDescriptors = null;
}
+ /**
+ * Sets the initial state for the execution. The serialized state is then shipped via the
+ * {@link TaskDeploymentDescriptor} to the TaskManagers.
+ *
+ * @param chainedStateHandle Chained operator state
+ * @param keyGroupsStateHandles Key-group state (= partitioned state)
+ */
public void setInitialState(
- SerializedValue<StateHandle<?>> initialState,
- Map<Integer, SerializedValue<StateHandle<?>>> initialKvState) {
-
- if (initialKvState != null && initialKvState.size() > 0) {
- throw new UnsupportedOperationException("Error: inconsistent handling of key/value state snapshots");
- }
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
+ List<KeyGroupsStateHandle> keyGroupsStateHandles) {
if (state != ExecutionState.CREATED) {
throw new IllegalArgumentException("Can only assign operator state when execution attempt is in CREATED");
}
-
- this.operatorState = initialState;
+ this.chainedStateHandle = chainedStateHandle;
+ this.keyGroupsStateHandles = keyGroupsStateHandles;
}
// --------------------------------------------------------------------------------------------
@@ -373,7 +387,8 @@ public class Execution {
final TaskDeploymentDescriptor deployment = vertex.createDeploymentDescriptor(
attemptId,
slot,
- operatorState,
+ chainedStateHandle,
+ keyGroupsStateHandles,
attemptNumber);
// register this execution at the execution graph, to receive call backs
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
index b15f851..f3a8b6d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/executiongraph/ExecutionVertex.java
@@ -42,9 +42,11 @@ import org.apache.flink.runtime.jobmanager.scheduler.CoLocationConstraint;
import org.apache.flink.runtime.jobmanager.scheduler.CoLocationGroup;
import org.apache.flink.runtime.jobmanager.scheduler.NoResourceAvailableException;
import org.apache.flink.runtime.jobmanager.scheduler.Scheduler;
-import org.apache.flink.runtime.state.StateHandle;
import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.slf4j.Logger;
import scala.concurrent.duration.FiniteDuration;
@@ -183,14 +185,14 @@ public class ExecutionVertex {
return this.jobVertex.getParallelism();
}
- public int getParallelSubtaskIndex() {
- return this.subTaskIndex;
- }
-
public int getMaxParallelism() {
return this.jobVertex.getMaxParallelism();
}
+ public int getParallelSubtaskIndex() {
+ return this.subTaskIndex;
+ }
+
public int getNumberOfInputs() {
return this.inputEdges.length;
}
@@ -229,7 +231,7 @@ public class ExecutionVertex {
public InstanceConnectionInfo getCurrentAssignedResourceLocation() {
return currentExecution.getAssignedResourceLocation();
}
-
+
public Execution getPriorExecutionAttempt(int attemptNumber) {
if (attemptNumber >= 0 && attemptNumber < priorExecutions.size()) {
return priorExecutions.get(attemptNumber);
@@ -238,7 +240,7 @@ public class ExecutionVertex {
throw new IllegalArgumentException("attempt does not exist");
}
}
-
+
public ExecutionGraph getExecutionGraph() {
return this.jobVertex.getGraph();
}
@@ -275,7 +277,7 @@ public class ExecutionVertex {
this.inputEdges[inputNumber] = edges;
// add the consumers to the source
- // for now (until the receiver initiated handshake is in place), we need to register the
+ // for now (until the receiver initiated handshake is in place), we need to register the
// edges as the execution graph
for (ExecutionEdge ee : edges) {
ee.getSource().addConsumer(ee, consumerNumber);
@@ -486,11 +488,11 @@ public class ExecutionVertex {
ExecutionAttemptID attemptID,
ActorGateway sender) {
Execution exec = getCurrentExecutionAttempt();
-
+
// check that this is for the correct execution attempt
if (exec != null && exec.getAttemptId().equals(attemptID)) {
SimpleSlot slot = exec.getAssignedResource();
-
+
// send only if we actually have a target
if (slot != null) {
ActorGateway gateway = slot.getInstance().getActorGateway();
@@ -517,7 +519,7 @@ public class ExecutionVertex {
return false;
}
}
-
+
/**
* Schedules or updates the consumer tasks of the result partition with the given ID.
*/
@@ -633,13 +635,14 @@ public class ExecutionVertex {
/**
* Creates a task deployment descriptor to deploy a subtask to the given target slot.
- *
+ *
* TODO: This should actually be in the EXECUTION
*/
TaskDeploymentDescriptor createDeploymentDescriptor(
ExecutionAttemptID executionId,
SimpleSlot targetSlot,
- SerializedValue<StateHandle<?>> operatorState,
+ ChainedStateHandle<StreamStateHandle> operatorState,
+ List<KeyGroupsStateHandle> keyGroupStates,
int attemptNumber) {
// Produced intermediate results
@@ -690,7 +693,8 @@ public class ExecutionVertex {
jarFiles,
classpaths,
targetSlot.getRoot().getSlotNumber(),
- operatorState);
+ operatorState,
+ keyGroupStates);
}
// --------------------------------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
index 4786388..a623295 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/JobVertex.java
@@ -249,7 +249,7 @@ public class JobVertex implements java.io.Serializable {
/**
* Sets the maximum parallelism for the task.
*
- * @param maxParallelism The maximum parallelism to be set.
+ * @param maxParallelism The maximum parallelism to be set. must be between 1 and Short.MAX_VALUE.
*/
public void setMaxParallelism(int maxParallelism) {
org.apache.flink.util.Preconditions.checkArgument(
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/tasks/StatefulTask.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/tasks/StatefulTask.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/tasks/StatefulTask.java
index f8bba1a..cab7ed6 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/tasks/StatefulTask.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobgraph/tasks/StatefulTask.java
@@ -18,21 +18,27 @@
package org.apache.flink.runtime.jobgraph.tasks;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
+
+import java.util.List;
/**
* This interface must be implemented by any invokable that has recoverable state and participates
* in checkpointing.
*/
-public interface StatefulTask<T extends StateHandle<?>> {
+public interface StatefulTask {
/**
* Sets the initial state of the operator, upon recovery. The initial state is typically
* a snapshot of the state from a previous execution.
*
- * @param stateHandle The handle to the state.
+ * @param chainedState Handle for the chained operator states.
+ * @param keyGroupsState Handle for key group states.
*/
- void setInitialState(T stateHandle) throws Exception;
+ void setInitialState(ChainedStateHandle<StreamStateHandle> chainedState, List<KeyGroupsStateHandle> keyGroupsState) throws Exception;
/**
* This method is either called directly and asynchronously by the checkpoint
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphStore.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphStore.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphStore.java
index 7f7c5fe..ec05f1e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphStore.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmanager/ZooKeeperSubmittedJobGraphStore.java
@@ -25,8 +25,8 @@ import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.apache.curator.utils.ZKPaths;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.java.tuple.Tuple2;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.runtime.zookeeper.StateStorageHelper;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
+import org.apache.flink.runtime.zookeeper.RetrievableStateStorageHelper;
import org.apache.flink.runtime.zookeeper.ZooKeeperStateHandleStore;
import org.apache.zookeeper.KeeperException;
import org.slf4j.Logger;
@@ -98,7 +98,7 @@ public class ZooKeeperSubmittedJobGraphStore implements SubmittedJobGraphStore {
public ZooKeeperSubmittedJobGraphStore(
CuratorFramework client,
String currentJobsPath,
- StateStorageHelper<SubmittedJobGraph> stateStorage) throws Exception {
+ RetrievableStateStorageHelper<SubmittedJobGraph> stateStorage) throws Exception {
checkNotNull(currentJobsPath, "Current jobs path");
checkNotNull(stateStorage, "State storage");
@@ -153,7 +153,7 @@ public class ZooKeeperSubmittedJobGraphStore implements SubmittedJobGraphStore {
synchronized (cacheLock) {
verifyIsRunning();
- List<Tuple2<StateHandle<SubmittedJobGraph>, String>> submitted;
+ List<Tuple2<RetrievableStateHandle<SubmittedJobGraph>, String>> submitted;
while (true) {
try {
@@ -168,10 +168,8 @@ public class ZooKeeperSubmittedJobGraphStore implements SubmittedJobGraphStore {
if (submitted.size() != 0) {
List<SubmittedJobGraph> jobGraphs = new ArrayList<>(submitted.size());
- for (Tuple2<StateHandle<SubmittedJobGraph>, String> jobStateHandle : submitted) {
- SubmittedJobGraph jobGraph = jobStateHandle
- .f0.getState(ClassLoader.getSystemClassLoader());
-
+ for (Tuple2<RetrievableStateHandle<SubmittedJobGraph>, String> jobStateHandle : submitted) {
+ SubmittedJobGraph jobGraph = jobStateHandle.f0.retrieveState();
addedJobGraphs.add(jobGraph.getJobId());
jobGraphs.add(jobGraph);
@@ -196,11 +194,7 @@ public class ZooKeeperSubmittedJobGraphStore implements SubmittedJobGraphStore {
verifyIsRunning();
try {
- StateHandle<SubmittedJobGraph> jobStateHandle = jobGraphsInZooKeeper.get(path);
-
- SubmittedJobGraph jobGraph = jobStateHandle
- .getState(ClassLoader.getSystemClassLoader());
-
+ SubmittedJobGraph jobGraph = jobGraphsInZooKeeper.get(path).retrieveState();
addedJobGraphs.add(jobGraph.getJobId());
LOG.info("Recovered {}.", jobGraph);
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/messages/checkpoint/AcknowledgeCheckpoint.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/messages/checkpoint/AcknowledgeCheckpoint.java b/flink-runtime/src/main/java/org/apache/flink/runtime/messages/checkpoint/AcknowledgeCheckpoint.java
index a4d438b..0c56603 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/messages/checkpoint/AcknowledgeCheckpoint.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/messages/checkpoint/AcknowledgeCheckpoint.java
@@ -20,51 +20,50 @@ package org.apache.flink.runtime.messages.checkpoint;
import org.apache.flink.api.common.JobID;
import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.util.SerializedValue;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+
+import java.util.List;
/**
* This message is sent from the {@link org.apache.flink.runtime.taskmanager.TaskManager} to the
* {@link org.apache.flink.runtime.jobmanager.JobManager} to signal that the checkpoint of an
* individual task is completed.
*
- * This message may carry the handle to the task's state.
+ * <p>This message may carry the handle to the task's chained operator state and the key group
+ * state.
*/
public class AcknowledgeCheckpoint extends AbstractCheckpointMessage implements java.io.Serializable {
private static final long serialVersionUID = -7606214777192401493L;
- private final SerializedValue<StateHandle<?>> state;
+ private final ChainedStateHandle<StreamStateHandle> stateHandle;
- /**
- * The state size. This is an optimization in order to not deserialize the
- * state handle at the checkpoint coordinator when gathering stats about
- * the checkpoints.
- */
- private final long stateSize;
+ private final List<KeyGroupsStateHandle> keyGroupsStateHandle;
public AcknowledgeCheckpoint(JobID job, ExecutionAttemptID taskExecutionId, long checkpointId) {
- this(job, taskExecutionId, checkpointId, null, 0);
+ this(job, taskExecutionId, checkpointId, null, null);
}
public AcknowledgeCheckpoint(
- JobID job,
- ExecutionAttemptID taskExecutionId,
- long checkpointId,
- SerializedValue<StateHandle<?>> state,
- long stateSize) {
+ JobID job,
+ ExecutionAttemptID taskExecutionId,
+ long checkpointId,
+ ChainedStateHandle<StreamStateHandle> state,
+ List<KeyGroupsStateHandle> keyGroupStateAndSizes) {
super(job, taskExecutionId, checkpointId);
- this.state = state;
- this.stateSize = stateSize;
+ this.stateHandle = state;
+ this.keyGroupsStateHandle = keyGroupStateAndSizes;
}
- public SerializedValue<StateHandle<?>> getState() {
- return state;
+ public ChainedStateHandle<StreamStateHandle> getStateHandle() {
+ return stateHandle;
}
- public long getStateSize() {
- return stateSize;
+ public List<KeyGroupsStateHandle> getKeyGroupsStateHandle() {
+ return keyGroupsStateHandle;
}
// --------------------------------------------------------------------------------------------
@@ -76,8 +75,10 @@ public class AcknowledgeCheckpoint extends AbstractCheckpointMessage implements
}
else if (o instanceof AcknowledgeCheckpoint) {
AcknowledgeCheckpoint that = (AcknowledgeCheckpoint) o;
- return super.equals(o) && (this.state == null ? that.state == null :
- (that.state != null && this.state.equals(that.state)));
+ return super.equals(o) &&
+ (this.stateHandle == null ? that.stateHandle == null : (that.stateHandle != null && this.stateHandle.equals(that.stateHandle))) &&
+ (this.keyGroupsStateHandle == null ? that.keyGroupsStateHandle == null : (that.keyGroupsStateHandle != null && this.keyGroupsStateHandle.equals(that.keyGroupsStateHandle)));
+
}
else {
return false;
@@ -86,7 +87,7 @@ public class AcknowledgeCheckpoint extends AbstractCheckpointMessage implements
@Override
public String toString() {
- return String.format("Confirm Task Checkpoint %d for (%s/%s) - state=%s",
- getCheckpointId(), getJob(), getTaskExecutionId(), state);
+ return String.format("Confirm Task Checkpoint %d for (%s/%s) - state=%s keyGroupState=%s",
+ getCheckpointId(), getJob(), getTaskExecutionId(), stateHandle, keyGroupsStateHandle);
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractCloseableHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractCloseableHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractCloseableHandle.java
index 609158d..5966c95 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractCloseableHandle.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractCloseableHandle.java
@@ -20,27 +20,26 @@ package org.apache.flink.runtime.state;
import java.io.Closeable;
import java.io.IOException;
-import java.io.Serializable;
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
/**
* A simple base for closable handles.
- *
+ *
* Offers to register a stream (or other closable object) that close calls are delegated to if
* the handle is closed or was already closed.
*/
-public abstract class AbstractCloseableHandle implements Closeable, Serializable {
+public abstract class AbstractCloseableHandle implements Closeable, StateObject {
/** Serial Version UID must be constant to maintain format compatibility */
private static final long serialVersionUID = 1L;
/** To atomically update the "closable" field without needing to add a member class like "AtomicBoolean */
- private static final AtomicIntegerFieldUpdater<AbstractCloseableHandle> CLOSER =
+ private static final AtomicIntegerFieldUpdater<AbstractCloseableHandle> CLOSER =
AtomicIntegerFieldUpdater.newUpdater(AbstractCloseableHandle.class, "isClosed");
// ------------------------------------------------------------------------
- /** The closeable to close if this handle is closed late */
+ /** The closeable to close if this handle is closed late */
private transient volatile Closeable toClose;
/** Flag to remember if this handle was already closed */
@@ -53,7 +52,7 @@ public abstract class AbstractCloseableHandle implements Closeable, Serializable
if (toClose == null) {
return;
}
-
+
// NOTE: The order of operations matters here:
// (1) first setting the closeable
// (2) checking the flag.
@@ -73,16 +72,16 @@ public abstract class AbstractCloseableHandle implements Closeable, Serializable
/**
* Closes the handle.
- *
+ *
* <p>If a "Closeable" has been registered via {@link #registerCloseable(Closeable)},
* then this will be closes.
- *
+ *
* <p>If any "Closeable" will be registered via {@link #registerCloseable(Closeable)} in the future,
* it will immediately be closed and that method will throw an exception.
- *
+ *
* @throws IOException Exceptions occurring while closing an already registered {@code Closeable}
* are forwarded.
- *
+ *
* @see #registerCloseable(Closeable)
*/
@Override
@@ -106,7 +105,7 @@ public abstract class AbstractCloseableHandle implements Closeable, Serializable
/**
* Checks whether this handle has been closed.
- *
+ *
* @return True is the handle is closed, false otherwise.
*/
public boolean isClosed() {
@@ -116,7 +115,7 @@ public abstract class AbstractCloseableHandle implements Closeable, Serializable
/**
* This method checks whether the handle is closed and throws an exception if it is closed.
* If the handle is not closed, this method does nothing.
- *
+ *
* @throws IOException Thrown, if the handle has been closed.
*/
public void ensureNotClosed() throws IOException {
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
index 6fc9475..b2cde22 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractStateBackend.java
@@ -33,17 +33,12 @@ import org.apache.flink.api.common.state.StateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
-import org.apache.flink.core.memory.DataInputView;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
-import org.apache.flink.core.memory.DataOutputView;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.core.fs.FSDataOutputStream;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.util.Preconditions;
import java.io.IOException;
-import java.io.OutputStream;
-import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
@@ -122,7 +117,7 @@ public abstract class AbstractStateBackend implements java.io.Serializable {
*/
public abstract void close() throws Exception;
- public void dispose() {
+ public void discardState() throws Exception {
if (kvStateRegistry != null) {
kvStateRegistry.unregisterAll();
}
@@ -418,37 +413,6 @@ public abstract class AbstractStateBackend implements java.io.Serializable {
public abstract CheckpointStateOutputStream createCheckpointStateOutputStream(
long checkpointID, long timestamp) throws Exception;
- /**
- * Creates a {@link DataOutputView} stream that writes into the state of the given checkpoint.
- * When the stream is closes, it returns a state handle that can retrieve the state back.
- *
- * @param checkpointID The ID of the checkpoint.
- * @param timestamp The timestamp of the checkpoint.
- * @return An DataOutputView stream that writes state for the given checkpoint.
- *
- * @throws Exception Exceptions may occur while creating the stream and should be forwarded.
- */
- public CheckpointStateOutputView createCheckpointStateOutputView(
- long checkpointID, long timestamp) throws Exception {
- return new CheckpointStateOutputView(createCheckpointStateOutputStream(checkpointID, timestamp));
- }
-
- /**
- * Writes the given state into the checkpoint, and returns a handle that can retrieve the state back.
- *
- * @param state The state to be checkpointed.
- * @param checkpointID The ID of the checkpoint.
- * @param timestamp The timestamp of the checkpoint.
- * @param <S> The type of the state.
- *
- * @return A state handle that can retrieve the checkpoined state.
- *
- * @throws Exception Exceptions may occur during serialization / storing the state and should be forwarded.
- */
- public abstract <S extends Serializable> StateHandle<S> checkpointStateSerializable(
- S state, long checkpointID, long timestamp) throws Exception;
-
-
// ------------------------------------------------------------------------
// Checkpoint state output stream
// ------------------------------------------------------------------------
@@ -456,7 +420,7 @@ public abstract class AbstractStateBackend implements java.io.Serializable {
/**
* A dedicated output stream that produces a {@link StreamStateHandle} when closed.
*/
- public static abstract class CheckpointStateOutputStream extends OutputStream {
+ public static abstract class CheckpointStateOutputStream extends FSDataOutputStream {
/**
* Closes the stream and gets a state handle that can create an input stream
@@ -467,67 +431,4 @@ public abstract class AbstractStateBackend implements java.io.Serializable {
*/
public abstract StreamStateHandle closeAndGetHandle() throws IOException;
}
-
- /**
- * A dedicated DataOutputView stream that produces a {@code StateHandle<DataInputView>} when closed.
- */
- public static final class CheckpointStateOutputView extends DataOutputViewStreamWrapper {
-
- private final CheckpointStateOutputStream out;
-
- public CheckpointStateOutputView(CheckpointStateOutputStream out) {
- super(out);
- this.out = out;
- }
-
- /**
- * Closes the stream and gets a state handle that can create a DataInputView.
- * producing the data written to this stream.
- *
- * @return A state handle that can create an input stream producing the data written to this stream.
- * @throws IOException Thrown, if the stream cannot be closed.
- */
- public StateHandle<DataInputView> closeAndGetHandle() throws IOException {
- return new DataInputViewHandle(out.closeAndGetHandle());
- }
-
- @Override
- public void close() throws IOException {
- out.close();
- }
- }
-
- /**
- * Simple state handle that resolved a {@link DataInputView} from a StreamStateHandle.
- */
- private static final class DataInputViewHandle implements StateHandle<DataInputView> {
-
- private static final long serialVersionUID = 2891559813513532079L;
-
- private final StreamStateHandle stream;
-
- private DataInputViewHandle(StreamStateHandle stream) {
- this.stream = stream;
- }
-
- @Override
- public DataInputView getState(ClassLoader userCodeClassLoader) throws Exception {
- return new DataInputViewStreamWrapper(stream.getState(userCodeClassLoader));
- }
-
- @Override
- public void discardState() throws Exception {
- stream.discardState();
- }
-
- @Override
- public long getStateSize() throws Exception {
- return stream.getStateSize();
- }
-
- @Override
- public void close() throws IOException {
- stream.close();
- }
- }
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousStateHandle.java
----------------------------------------------------------------------
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousStateHandle.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousStateHandle.java
deleted file mode 100644
index fee1efe..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/AsynchronousStateHandle.java
+++ /dev/null
@@ -1,43 +0,0 @@
-/*
- * 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;
-
-/**
- * {@link StateHandle} that can asynchronously materialize the state that it represents. Instead
- * of representing a materialized handle to state this would normally hold the (immutable) state
- * internally and can materialize it if requested.
- */
-public abstract class AsynchronousStateHandle<T> implements StateHandle<T> {
- private static final long serialVersionUID = 1L;
-
- /**
- * Materializes the state held by this {@code AsynchronousStateHandle}.
- */
- public abstract StateHandle<T> materialize() throws Exception;
-
- @Override
- public final T getState(ClassLoader userCodeClassLoader) throws Exception {
- throw new UnsupportedOperationException("This must not be called. This is likely an internal bug.");
- }
-
- @Override
- public final void discardState() throws Exception {
- throw new UnsupportedOperationException("This must not be called. This is likely an internal bug.");
- }
-}
[14/27] flink git commit: [FLINK-4381] Refactor State to Prepare For
Key-Group State Backends
Posted by al...@apache.org.
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java
index f9698a8..1d07bdd 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/InterruptSensitiveRestoreTest.java
@@ -21,6 +21,7 @@ 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.core.fs.FSDataInputStream;
import org.apache.flink.core.testutils.OneShotLatch;
import org.apache.flink.runtime.blob.BlobKey;
import org.apache.flink.runtime.broadcast.BroadcastVariableManager;
@@ -38,7 +39,11 @@ import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.operators.testutils.UnregisteredTaskMetricsGroup;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.AbstractCloseableHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.RetrievableStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.taskmanager.Task;
import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
import org.apache.flink.runtime.util.EnvironmentInformation;
@@ -54,10 +59,12 @@ import org.junit.Test;
import scala.concurrent.duration.FiniteDuration;
+import java.io.EOFException;
import java.io.IOException;
import java.io.Serializable;
import java.net.URL;
import java.util.Collections;
+import java.util.List;
import java.util.concurrent.TimeUnit;
import static org.junit.Assert.*;
@@ -82,12 +89,9 @@ public class InterruptSensitiveRestoreTest {
cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);
cfg.setStreamOperator(new StreamSource<>(new TestSource()));
- StateHandle<Serializable> lockingHandle = new InterruptLockingStateHandle();
- StreamTaskState opState = new StreamTaskState();
- opState.setFunctionState(lockingHandle);
- StreamTaskStateList taskState = new StreamTaskStateList(new StreamTaskState[] { opState });
+ StreamStateHandle lockingHandle = new InterruptLockingStateHandle();
- TaskDeploymentDescriptor tdd = createTaskDeploymentDescriptor(taskConfig, taskState);
+ TaskDeploymentDescriptor tdd = createTaskDeploymentDescriptor(taskConfig, lockingHandle);
Task task = createTask(tdd);
// start the task and wait until it is in "restore"
@@ -113,7 +117,10 @@ public class InterruptSensitiveRestoreTest {
private static TaskDeploymentDescriptor createTaskDeploymentDescriptor(
Configuration taskConfig,
- StateHandle<?> state) throws IOException {
+ StreamStateHandle state) throws IOException {
+
+ ChainedStateHandle<StreamStateHandle> operatorState = new ChainedStateHandle<>(Collections.singletonList(state));
+ List<KeyGroupsStateHandle> keyGroupState = Collections.emptyList();
return new TaskDeploymentDescriptor(
new JobID(),
@@ -131,7 +138,8 @@ public class InterruptSensitiveRestoreTest {
Collections.<BlobKey>emptyList(),
Collections.<URL>emptyList(),
0,
- new SerializedValue<StateHandle<?>>(state));
+ operatorState,
+ keyGroupState);
}
private static Task createTask(TaskDeploymentDescriptor tdd) throws IOException {
@@ -159,14 +167,34 @@ public class InterruptSensitiveRestoreTest {
// ------------------------------------------------------------------------
@SuppressWarnings("serial")
- private static class InterruptLockingStateHandle implements StateHandle<Serializable> {
+ private static class InterruptLockingStateHandle extends AbstractCloseableHandle implements StreamStateHandle {
- private transient volatile boolean closed;
-
@Override
- public Serializable getState(ClassLoader userCodeClassLoader) {
+ public FSDataInputStream openInputStream() throws Exception {
+ ensureNotClosed();
+ FSDataInputStream is = new FSDataInputStream() {
+
+ @Override
+ public void seek(long desired) throws IOException {
+ }
+
+ @Override
+ public long getPos() throws IOException {
+ return 0;
+ }
+
+ @Override
+ public int read() throws IOException {
+ block();
+ throw new EOFException();
+ }
+ };
+ registerCloseable(is);
+ return is;
+ }
+
+ private void block() {
IN_RESTORE_LATCH.trigger();
-
// this mimics what happens in the HDFS client code.
// an interrupt on a waiting object leads to an infinite loop
try {
@@ -175,7 +203,7 @@ public class InterruptSensitiveRestoreTest {
}
}
catch (InterruptedException e) {
- while (!closed) {
+ while (!isClosed()) {
try {
synchronized (this) {
wait();
@@ -183,8 +211,6 @@ public class InterruptSensitiveRestoreTest {
} catch (InterruptedException ignored) {}
}
}
-
- return new SerializableObject();
}
@Override
@@ -194,11 +220,6 @@ public class InterruptSensitiveRestoreTest {
public long getStateSize() throws Exception {
return 0;
}
-
- @Override
- public void close() throws IOException {
- closed = true;
- }
}
// ------------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java
index 05b8e8c..5e82569 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamMockEnvironment.java
@@ -49,7 +49,9 @@ import org.apache.flink.runtime.plugable.DeserializationDelegate;
import org.apache.flink.runtime.plugable.NonReusingDeserializationDelegate;
import org.apache.flink.runtime.query.KvStateRegistry;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.ChainedStateHandle;
+import org.apache.flink.runtime.state.KeyGroupsStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
import org.mockito.invocation.InvocationOnMock;
@@ -310,7 +312,10 @@ public class StreamMockEnvironment implements Environment {
}
@Override
- public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) {
+ public void acknowledgeCheckpoint(long checkpointId,
+ ChainedStateHandle<StreamStateHandle> chainedStateHandle,
+ List<KeyGroupsStateHandle> keyGroupStateHandles) {
+
}
@Override
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java
index cfaeaad..66bc237 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/tasks/StreamTaskAsyncCheckpointTest.java
@@ -1,236 +1,234 @@
-/*
- * 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.functions.MapFunction;
-import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
-import org.apache.flink.core.testutils.OneShotLatch;
-import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
-import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
-import org.apache.flink.runtime.state.AsynchronousStateHandle;
-import org.apache.flink.runtime.state.StateHandle;
-import org.apache.flink.streaming.api.graph.StreamConfig;
-import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
-import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
-import org.apache.flink.streaming.api.watermark.Watermark;
-import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.powermock.core.classloader.annotations.PowerMockIgnore;
-import org.powermock.core.classloader.annotations.PrepareForTest;
-import org.powermock.modules.junit4.PowerMockRunner;
-
-import java.io.IOException;
-import java.lang.reflect.Field;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-
-/**
- * Tests for asynchronous checkpoints.
- */
-@RunWith(PowerMockRunner.class)
-@PrepareForTest(ResultPartitionWriter.class)
-@PowerMockIgnore({"javax.management.*", "com.sun.jndi.*"})
-@SuppressWarnings("serial")
-public class StreamTaskAsyncCheckpointTest {
-
- /**
- * This ensures that asynchronous state handles are actually materialized asynchonously.
- *
- * <p>We use latches to block at various stages and see if the code still continues through
- * the parts that are not asynchronous. If the checkpoint is not done asynchronously the
- * test will simply lock forever.
- * @throws Exception
- */
- @Test
- public void testAsyncCheckpoints() throws Exception {
- final OneShotLatch delayCheckpointLatch = new OneShotLatch();
- final OneShotLatch ensureCheckpointLatch = new OneShotLatch();
-
- final OneInputStreamTask<String, String> task = new OneInputStreamTask<>();
-
- final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
-
- StreamConfig streamConfig = testHarness.getStreamConfig();
-
- streamConfig.setStreamOperator(new AsyncCheckpointOperator());
-
- StreamMockEnvironment mockEnv = new StreamMockEnvironment(
- testHarness.jobConfig,
- testHarness.taskConfig,
- testHarness.memorySize,
- new MockInputSplitProvider(),
- testHarness.bufferSize) {
-
- @Override
- public ExecutionConfig getExecutionConfig() {
- return testHarness.executionConfig;
- }
-
- @Override
- public void acknowledgeCheckpoint(long checkpointId) {
- super.acknowledgeCheckpoint(checkpointId);
- }
-
- @Override
- public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) {
- super.acknowledgeCheckpoint(checkpointId, state);
-
- // block on the latch, to verify that triggerCheckpoint returns below,
- // even though the async checkpoint would not finish
- try {
- delayCheckpointLatch.await();
- } catch (InterruptedException e) {
- e.printStackTrace();
- }
-
- assertTrue(state instanceof StreamTaskStateList);
- StreamTaskStateList stateList = (StreamTaskStateList) state;
-
- // should be only one state
- StreamTaskState taskState = stateList.getState(this.getUserClassLoader())[0];
- StateHandle<?> operatorState = taskState.getOperatorState();
- assertTrue("It must be a TestStateHandle", operatorState instanceof TestStateHandle);
- TestStateHandle testState = (TestStateHandle) operatorState;
- assertEquals(42, testState.checkpointId);
- assertEquals(17, testState.timestamp);
-
- // we now know that the checkpoint went through
- ensureCheckpointLatch.trigger();
- }
- };
-
- testHarness.invoke(mockEnv);
-
- // wait for the task to be running
- for (Field field: StreamTask.class.getDeclaredFields()) {
- if (field.getName().equals("isRunning")) {
- field.setAccessible(true);
- while (!field.getBoolean(task)) {
- Thread.sleep(10);
- }
-
- }
- }
-
- task.triggerCheckpoint(42, 17);
-
- // now we allow the checkpoint
- delayCheckpointLatch.trigger();
-
- // wait for the checkpoint to go through
- ensureCheckpointLatch.await();
-
- testHarness.endInput();
- testHarness.waitForTaskCompletion();
- }
-
-
- // ------------------------------------------------------------------------
-
- public static class AsyncCheckpointOperator
- extends AbstractStreamOperator<String>
- implements OneInputStreamOperator<String, String> {
- @Override
- public void processElement(StreamRecord<String> element) throws Exception {
- // we also don't care
- }
-
- @Override
- public void processWatermark(Watermark mark) throws Exception {
- // not interested
- }
-
-
- @Override
- public StreamTaskState snapshotOperatorState(final long checkpointId, final long timestamp) throws Exception {
- StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
-
- AsynchronousStateHandle<String> asyncState =
- new DataInputViewAsynchronousStateHandle(checkpointId, timestamp);
-
- taskState.setOperatorState(asyncState);
-
- return taskState;
- }
-
- @Override
- public void restoreState(StreamTaskState taskState) throws Exception {
- super.restoreState(taskState);
- }
- }
-
- private static class DataInputViewAsynchronousStateHandle extends AsynchronousStateHandle<String> {
-
- private final long checkpointId;
- private final long timestamp;
-
- public DataInputViewAsynchronousStateHandle(long checkpointId, long timestamp) {
- this.checkpointId = checkpointId;
- this.timestamp = timestamp;
- }
-
- @Override
- public StateHandle<String> materialize() throws Exception {
- return new TestStateHandle(checkpointId, timestamp);
- }
-
- @Override
- public long getStateSize() {
- return 0;
- }
-
- @Override
- public void close() throws IOException {}
- }
-
- private static class TestStateHandle implements StateHandle<String> {
-
- public final long checkpointId;
- public final long timestamp;
-
- public TestStateHandle(long checkpointId, long timestamp) {
- this.checkpointId = checkpointId;
- this.timestamp = timestamp;
- }
-
- @Override
- public String getState(ClassLoader userCodeClassLoader) throws Exception {
- return null;
- }
-
- @Override
- public void discardState() throws Exception {}
-
- @Override
- public long getStateSize() {
- return 0;
- }
-
- @Override
- public void close() throws IOException {}
- }
-
- public static class DummyMapFunction<T> implements MapFunction<T, T> {
- @Override
- public T map(T value) { return value; }
- }
-}
+///*
+// * 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.functions.MapFunction;
+//import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+//import org.apache.flink.core.testutils.OneShotLatch;
+//import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
+//import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
+//import org.apache.flink.streaming.api.graph.StreamConfig;
+//import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+//import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+//import org.apache.flink.streaming.api.watermark.Watermark;
+//import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+//import org.junit.Test;
+//import org.junit.runner.RunWith;
+//import org.powermock.core.classloader.annotations.PowerMockIgnore;
+//import org.powermock.core.classloader.annotations.PrepareForTest;
+//import org.powermock.modules.junit4.PowerMockRunner;
+//
+//import java.io.IOException;
+//import java.lang.reflect.Field;
+//
+//import static org.junit.Assert.assertEquals;
+//import static org.junit.Assert.assertTrue;
+//
+///**
+// * Tests for asynchronous checkpoints.
+// */
+//@RunWith(PowerMockRunner.class)
+//@PrepareForTest(ResultPartitionWriter.class)
+//@PowerMockIgnore({"javax.management.*", "com.sun.jndi.*"})
+//@SuppressWarnings("serial")
+//public class StreamTaskAsyncCheckpointTest {
+//
+// /**
+// * This ensures that asynchronous state handles are actually materialized asynchonously.
+// *
+// * <p>We use latches to block at various stages and see if the code still continues through
+// * the parts that are not asynchronous. If the checkpoint is not done asynchronously the
+// * test will simply lock forever.
+// * @throws Exception
+// */
+// @Test
+// public void testAsyncCheckpoints() throws Exception {
+// final OneShotLatch delayCheckpointLatch = new OneShotLatch();
+// final OneShotLatch ensureCheckpointLatch = new OneShotLatch();
+//
+// final OneInputStreamTask<String, String> task = new OneInputStreamTask<>();
+//
+// final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(task, BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
+//
+// StreamConfig streamConfig = testHarness.getStreamConfig();
+//
+// streamConfig.setStreamOperator(new AsyncCheckpointOperator());
+//
+// StreamMockEnvironment mockEnv = new StreamMockEnvironment(
+// testHarness.jobConfig,
+// testHarness.taskConfig,
+// testHarness.memorySize,
+// new MockInputSplitProvider(),
+// testHarness.bufferSize) {
+//
+// @Override
+// public ExecutionConfig getExecutionConfig() {
+// return testHarness.executionConfig;
+// }
+//
+// @Override
+// public void acknowledgeCheckpoint(long checkpointId) {
+// super.acknowledgeCheckpoint(checkpointId);
+// }
+//
+// @Override
+// public void acknowledgeCheckpoint(long checkpointId, StateHandle<?> state) {
+// super.acknowledgeCheckpoint(checkpointId, state);
+//
+// // block on the latch, to verify that triggerCheckpoint returns below,
+// // even though the async checkpoint would not finish
+// try {
+// delayCheckpointLatch.await();
+// } catch (InterruptedException e) {
+// e.printStackTrace();
+// }
+//
+// assertTrue(state instanceof StreamTaskStateList);
+// StreamTaskStateList stateList = (StreamTaskStateList) state;
+//
+// // should be only one state
+// StreamTaskState taskState = stateList.getState(this.getUserClassLoader())[0];
+// StateHandle<?> operatorState = taskState.getOperatorState();
+// assertTrue("It must be a TestStateHandle", operatorState instanceof TestStateHandle);
+// TestStateHandle testState = (TestStateHandle) operatorState;
+// assertEquals(42, testState.checkpointId);
+// assertEquals(17, testState.timestamp);
+//
+// // we now know that the checkpoint went through
+// ensureCheckpointLatch.trigger();
+// }
+// };
+//
+// testHarness.invoke(mockEnv);
+//
+// // wait for the task to be running
+// for (Field field: StreamTask.class.getDeclaredFields()) {
+// if (field.getName().equals("isRunning")) {
+// field.setAccessible(true);
+// while (!field.getBoolean(task)) {
+// Thread.sleep(10);
+// }
+//
+// }
+// }
+//
+// task.triggerCheckpoint(42, 17);
+//
+// // now we allow the checkpoint
+// delayCheckpointLatch.trigger();
+//
+// // wait for the checkpoint to go through
+// ensureCheckpointLatch.await();
+//
+// testHarness.endInput();
+// testHarness.waitForTaskCompletion();
+// }
+//
+//
+// // ------------------------------------------------------------------------
+//
+// public static class AsyncCheckpointOperator
+// extends AbstractStreamOperator<String>
+// implements OneInputStreamOperator<String, String> {
+// @Override
+// public void processElement(StreamRecord<String> element) throws Exception {
+// // we also don't care
+// }
+//
+// @Override
+// public void processWatermark(Watermark mark) throws Exception {
+// // not interested
+// }
+//
+//
+// @Override
+// public StreamTaskState snapshotOperatorState(final long checkpointId, final long timestamp) throws Exception {
+// StreamTaskState taskState = super.snapshotOperatorState(checkpointId, timestamp);
+//
+// AsynchronousStateHandle<String> asyncState =
+// new DataInputViewAsynchronousStateHandle(checkpointId, timestamp);
+//
+// taskState.setOperatorState(asyncState);
+//
+// return taskState;
+// }
+//
+// @Override
+// public void restoreState(StreamTaskState taskState) throws Exception {
+// super.restoreState(taskState);
+// }
+// }
+//
+// private static class DataInputViewAsynchronousStateHandle extends AsynchronousStateHandle<String> {
+//
+// private final long checkpointId;
+// private final long timestamp;
+//
+// public DataInputViewAsynchronousStateHandle(long checkpointId, long timestamp) {
+// this.checkpointId = checkpointId;
+// this.timestamp = timestamp;
+// }
+//
+// @Override
+// public StateHandle<String> materialize() throws Exception {
+// return new TestStateHandle(checkpointId, timestamp);
+// }
+//
+// @Override
+// public long getStateSize() {
+// return 0;
+// }
+//
+// @Override
+// public void close() throws IOException {}
+// }
+//
+// private static class TestStateHandle implements StateHandle<String> {
+//
+// public final long checkpointId;
+// public final long timestamp;
+//
+// public TestStateHandle(long checkpointId, long timestamp) {
+// this.checkpointId = checkpointId;
+// this.timestamp = timestamp;
+// }
+//
+// @Override
+// public String getState(ClassLoader userCodeClassLoader) throws Exception {
+// return null;
+// }
+//
+// @Override
+// public void discardState() throws Exception {}
+//
+// @Override
+// public long getStateSize() {
+// return 0;
+// }
+//
+// @Override
+// public void close() throws IOException {}
+// }
+//
+// public static class DummyMapFunction<T> implements MapFunction<T, T> {
+// @Override
+// public T map(T value) { return value; }
+// }
+//}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java
index 675e7b6..bc255ff 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/MockContext.java
@@ -22,24 +22,18 @@ import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.Callable;
-import java.util.concurrent.Executors;
+import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.accumulators.Accumulator;
import org.apache.flink.api.common.typeinfo.TypeInformation;
-import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.functions.KeySelector;
-import org.apache.flink.configuration.Configuration;
-import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
import org.apache.flink.runtime.operators.testutils.MockEnvironment;
import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
-import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.Output;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
@@ -75,47 +69,41 @@ public class MockContext<IN, OUT> {
return output;
}
- public static <IN, OUT> List<OUT> createAndExecute(OneInputStreamOperator<IN, OUT> operator, List<IN> inputs) {
+ public static <IN, OUT> List<OUT> createAndExecute(OneInputStreamOperator<IN, OUT> operator, List<IN> inputs) throws Exception {
return createAndExecuteForKeyedStream(operator, inputs, null, null);
}
public static <IN, OUT, KEY> List<OUT> createAndExecuteForKeyedStream(
OneInputStreamOperator<IN, OUT> operator, List<IN> inputs,
- KeySelector<IN, KEY> keySelector, TypeInformation<KEY> keyType) {
+ KeySelector<IN, KEY> keySelector, TypeInformation<KEY> keyType) throws Exception {
+
+ OneInputStreamOperatorTestHarness<IN, OUT> testHarness =
+ new OneInputStreamOperatorTestHarness<>(operator);
+
+ testHarness.configureForKeyedStream(keySelector, keyType);
+
+ testHarness.setup();
+ testHarness.open();
- MockContext<IN, OUT> mockContext = new MockContext<IN, OUT>(inputs);
+ operator.open();
- StreamConfig config = new StreamConfig(new Configuration());
- if (keySelector != null && keyType != null) {
- config.setStateKeySerializer(keyType.createSerializer(new ExecutionConfig()));
- config.setStatePartitioner(0, keySelector);
+ for (IN in: inputs) {
+ testHarness.processElement(new StreamRecord<>(in));
}
-
- final ScheduledExecutorService timerService = Executors.newSingleThreadScheduledExecutor();
- final Object lock = new Object();
- final StreamTask<?, ?> mockTask = createMockTaskWithTimer(timerService, lock);
-
- operator.setup(mockTask, config, mockContext.output);
- try {
- operator.open();
-
- StreamRecord<IN> record = new StreamRecord<IN>(null);
- for (IN in: inputs) {
- record = record.replace(in);
- synchronized (lock) {
- operator.setKeyContextElement1(record);
- operator.processElement(record);
- }
- }
- operator.close();
- } catch (Exception e) {
- throw new RuntimeException("Cannot invoke operator.", e);
- } finally {
- timerService.shutdownNow();
+ testHarness.close();
+
+ ConcurrentLinkedQueue<Object> output = testHarness.getOutput();
+
+ List<OUT> result = new ArrayList<>();
+
+ for (Object o : output) {
+ if (o instanceof StreamRecord) {
+ result.add((OUT) ((StreamRecord) o).getValue());
+ }
}
- return mockContext.getOutputs();
+ return result;
}
private static StreamTask<?, ?> createMockTaskWithTimer(
@@ -149,22 +137,6 @@ public class MockContext<IN, OUT> {
}
}).when(task).registerTimer(anyLong(), any(Triggerable.class));
-
- try {
- doAnswer(new Answer<AbstractStateBackend>() {
- @Override
- public AbstractStateBackend answer(InvocationOnMock invocationOnMock) throws Throwable {
- final String operatorIdentifier = (String) invocationOnMock.getArguments()[0];
- final TypeSerializer<?> keySerializer = (TypeSerializer<?>) invocationOnMock.getArguments()[1];
- MemoryStateBackend backend = MemoryStateBackend.create();
- backend.initializeForJob(new DummyEnvironment("dummty", 1, 0), operatorIdentifier, keySerializer);
- return backend;
- }
- }).when(task).createStateBackend(any(String.class), any(TypeSerializer.class));
- } catch (Exception e) {
- e.printStackTrace();
- }
-
return task;
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
index 6855989..6cb46d6 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
@@ -27,28 +27,23 @@ import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.operators.testutils.MockEnvironment;
import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
-import org.apache.flink.runtime.state.AsynchronousKvStateSnapshot;
-import org.apache.flink.runtime.state.AsynchronousStateHandle;
-import org.apache.flink.runtime.state.KvStateSnapshot;
+import org.apache.flink.runtime.state.AbstractStateBackend;
+import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.streaming.api.graph.StreamConfig;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.Output;
-import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.operators.Triggerable;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.runtime.tasks.DefaultTimeServiceProvider;
import org.apache.flink.streaming.runtime.tasks.StreamTask;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import org.apache.flink.streaming.runtime.tasks.TimeServiceProvider;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
-import java.io.Serializable;
import java.util.Collection;
-import java.util.HashMap;
-import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Executors;
@@ -73,9 +68,9 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
final ConcurrentLinkedQueue<Object> outputList;
final StreamConfig config;
-
+
final ExecutionConfig executionConfig;
-
+
final Object checkpointLock;
final TimeServiceProvider timeServiceProvider;
@@ -89,8 +84,8 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
* Whether setup() was called on the operator. This is reset when calling close().
*/
private boolean setupCalled = false;
-
-
+
+
public OneInputStreamOperatorTestHarness(OneInputStreamOperator<IN, OUT> operator) {
this(operator, new ExecutionConfig());
}
@@ -107,17 +102,20 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
TimeServiceProvider testTimeProvider) {
this.operator = operator;
this.outputList = new ConcurrentLinkedQueue<Object>();
- this.config = new StreamConfig(new Configuration());
+ Configuration underlyingConfig = new Configuration();
+ this.config = new StreamConfig(underlyingConfig);
+ this.config.setCheckpointingEnabled(true);
this.executionConfig = executionConfig;
this.checkpointLock = new Object();
- final Environment env = new MockEnvironment("MockTwoInputTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024);
+ final Environment env = new MockEnvironment("MockTwoInputTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024, underlyingConfig);
mockTask = mock(StreamTask.class);
timeServiceProvider = testTimeProvider;
when(mockTask.getName()).thenReturn("Mock Task");
when(mockTask.getCheckpointLock()).thenReturn(checkpointLock);
when(mockTask.getConfiguration()).thenReturn(config);
+ when(mockTask.getTaskConfiguration()).thenReturn(underlyingConfig);
when(mockTask.getEnvironment()).thenReturn(env);
when(mockTask.getExecutionConfig()).thenReturn(executionConfig);
when(mockTask.getUserCodeClassLoader()).thenReturn(this.getClass().getClassLoader());
@@ -173,8 +171,9 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
ClosureCleaner.clean(keySelector, false);
config.setStatePartitioner(0, keySelector);
config.setStateKeySerializer(keyType.createSerializer(executionConfig));
+ config.setKeyGroupAssigner(new HashKeyGroupAssigner<K>(10));
}
-
+
/**
* Get all the output from the task. This contains StreamRecords and Events interleaved. Use
* {@link org.apache.flink.streaming.util.TestHarnessUtil#getStreamRecordsFromOutput(java.util.List)}
@@ -206,47 +205,30 @@ public class OneInputStreamOperatorTestHarness<IN, OUT> {
}
/**
- * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#snapshotOperatorState(long, long)}
+ * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#snapshotState(org.apache.flink.core.fs.FSDataOutputStream, long, long)} ()}
*/
- public StreamTaskState snapshot(long checkpointId, long timestamp) throws Exception {
- StreamTaskState snapshot = operator.snapshotOperatorState(checkpointId, timestamp);
- // materialize asynchronous state handles
- if (snapshot != null) {
- if (snapshot.getFunctionState() instanceof AsynchronousStateHandle) {
- AsynchronousStateHandle<Serializable> asyncState = (AsynchronousStateHandle<Serializable>) snapshot.getFunctionState();
- snapshot.setFunctionState(asyncState.materialize());
- }
- if (snapshot.getOperatorState() instanceof AsynchronousStateHandle) {
- AsynchronousStateHandle<?> asyncState = (AsynchronousStateHandle<?>) snapshot.getOperatorState();
- snapshot.setOperatorState(asyncState.materialize());
- }
- if (snapshot.getKvStates() != null) {
- Set<String> keys = snapshot.getKvStates().keySet();
- HashMap<String, KvStateSnapshot<?, ?, ?, ?, ?>> kvStates = snapshot.getKvStates();
- for (String key: keys) {
- if (kvStates.get(key) instanceof AsynchronousKvStateSnapshot) {
- AsynchronousKvStateSnapshot<?, ?, ?, ?, ?> asyncHandle = (AsynchronousKvStateSnapshot<?, ?, ?, ?, ?>) kvStates.get(key);
- kvStates.put(key, asyncHandle.materialize());
- }
- }
- }
-
- }
- return snapshot;
+ public StreamStateHandle snapshot(long checkpointId, long timestamp) throws Exception {
+ // simply use an in-memory handle
+ MemoryStateBackend backend = new MemoryStateBackend();
+ AbstractStateBackend.CheckpointStateOutputStream outStream =
+ backend.createCheckpointStateOutputStream(checkpointId, timestamp);
+ operator.snapshotState(outStream, checkpointId, timestamp);
+ return outStream.closeAndGetHandle();
}
/**
- * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#notifyOfCompletedCheckpoint(long)}
+ * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#notifyOfCompletedCheckpoint(long)} ()}
*/
public void notifyOfCompletedCheckpoint(long checkpointId) throws Exception {
operator.notifyOfCompletedCheckpoint(checkpointId);
}
+
/**
- * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#restoreState(StreamTaskState)}
+ * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#restoreState(org.apache.flink.core.fs.FSDataInputStream)} ()}
*/
- public void restore(StreamTaskState snapshot, long recoveryTimestamp) throws Exception {
- operator.restoreState(snapshot);
+ public void restore(StreamStateHandle snapshot) throws Exception {
+ operator.restoreState(snapshot.openInputStream());
}
/**
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java
----------------------------------------------------------------------
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java
index faf38a3..779436a 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/WindowingTestHarness.java
@@ -22,6 +22,9 @@ import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.core.fs.FSDataOutputStream;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.streaming.api.functions.windowing.WindowFunction;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.api.windowing.assigners.WindowAssigner;
@@ -30,7 +33,6 @@ import org.apache.flink.streaming.api.windowing.windows.Window;
import org.apache.flink.streaming.runtime.operators.windowing.WindowOperator;
import org.apache.flink.streaming.runtime.operators.windowing.functions.InternalIterableWindowFunction;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
import org.apache.flink.streaming.runtime.tasks.TestTimeServiceProvider;
import org.apache.flink.util.Collector;
import org.apache.flink.util.Preconditions;
@@ -167,20 +169,20 @@ public class WindowingTestHarness<K, IN, W extends Window> {
/**
* Takes a snapshot of the current state of the operator. This can be used to test fault-tolerance.
*/
- public StreamTaskState snapshot(long checkpointId, long timestamp) throws Exception {
+ public StreamStateHandle snapshot(long checkpointId, long timestamp) throws Exception {
return testHarness.snapshot(checkpointId, timestamp);
}
/**
- * Resumes execution from a provided {@link StreamTaskState}. This is used to test recovery after a failure.
+ * Resumes execution from a provided {@link StreamStateHandle}. This is used to test recovery after a failure.
*/
- public void restore(StreamTaskState snapshot, long recoveryTime) throws Exception {
+ public void restore(StreamStateHandle stateHandle) throws Exception {
Preconditions.checkArgument(!isOpen,
"You are trying to restore() while the operator is still open. " +
"Please call close() first.");
testHarness.setup();
- testHarness.restore(snapshot, recoveryTime);
+ testHarness.restore(stateHandle);
openOperator();
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
index 199a6af..97c8339 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/EventTimeWindowCheckpointingITCase.java
@@ -73,6 +73,8 @@ import static org.junit.Assert.*;
@RunWith(Parameterized.class)
public class EventTimeWindowCheckpointingITCase extends TestLogger {
+
+ private static final int MAX_MEM_STATE_SIZE = 10 * 1024 * 1024;
private static final int PARALLELISM = 4;
private static ForkableFlinkMiniCluster cluster;
@@ -109,7 +111,7 @@ public class EventTimeWindowCheckpointingITCase extends TestLogger {
public void initStateBackend() throws IOException {
switch (stateBackendEnum) {
case MEM:
- this.stateBackend = new MemoryStateBackend();
+ this.stateBackend = new MemoryStateBackend(MAX_MEM_STATE_SIZE);
break;
case FILE: {
String backups = tempFolder.newFolder().getAbsolutePath();
@@ -119,7 +121,7 @@ public class EventTimeWindowCheckpointingITCase extends TestLogger {
case ROCKSDB: {
String rocksDb = tempFolder.newFolder().getAbsolutePath();
String rocksDbBackups = tempFolder.newFolder().toURI().toString();
- RocksDBStateBackend rdb = new RocksDBStateBackend(rocksDbBackups, new MemoryStateBackend());
+ RocksDBStateBackend rdb = new RocksDBStateBackend(rocksDbBackups, new MemoryStateBackend(MAX_MEM_STATE_SIZE));
rdb.setDbStoragePath(rocksDb);
this.stateBackend = rdb;
break;
@@ -127,7 +129,7 @@ public class EventTimeWindowCheckpointingITCase extends TestLogger {
case ROCKSDB_FULLY_ASYNC: {
String rocksDb = tempFolder.newFolder().getAbsolutePath();
String rocksDbBackups = tempFolder.newFolder().toURI().toString();
- RocksDBStateBackend rdb = new RocksDBStateBackend(rocksDbBackups, new MemoryStateBackend());
+ RocksDBStateBackend rdb = new RocksDBStateBackend(rocksDbBackups, new MemoryStateBackend(MAX_MEM_STATE_SIZE));
rdb.setDbStoragePath(rocksDb);
rdb.enableFullyAsyncSnapshots();
this.stateBackend = rdb;
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
index 0de2a75..8d1baeb 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/RescalingITCase.java
@@ -35,6 +35,7 @@ import org.apache.flink.runtime.instance.ActorGateway;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.messages.JobManagerMessages;
import org.apache.flink.runtime.state.HashKeyGroupAssigner;
+import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.filesystem.FsStateBackendFactory;
import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages;
import org.apache.flink.streaming.api.checkpoint.Checkpointed;
@@ -145,7 +146,7 @@ public class RescalingITCase extends TestLogger {
for (int key = 0; key < numberKeys; key++) {
int keyGroupIndex = keyGroupAssigner.getKeyGroupIndex(key);
- expectedResult.add(Tuple2.of(keyGroupIndex % parallelism, numberElements * key));
+ expectedResult.add(Tuple2.of(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, keyGroupIndex), numberElements * key));
}
assertEquals(expectedResult, actualResult);
@@ -188,7 +189,7 @@ public class RescalingITCase extends TestLogger {
for (int key = 0; key < numberKeys; key++) {
int keyGroupIndex = keyGroupAssigner2.getKeyGroupIndex(key);
- expectedResult2.add(Tuple2.of(keyGroupIndex % parallelism2, key * (numberElements + numberElements2)));
+ expectedResult2.add(Tuple2.of(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism2, keyGroupIndex), key * (numberElements + numberElements2)));
}
assertEquals(expectedResult2, actualResult2);
@@ -351,7 +352,8 @@ public class RescalingITCase extends TestLogger {
for (int key = 0; key < numberKeys; key++) {
int keyGroupIndex = keyGroupAssigner.getKeyGroupIndex(key);
- expectedResult.add(Tuple2.of(keyGroupIndex % parallelism, numberElements * key));
+// expectedResult.add(Tuple2.of(keyGroupIndex % parallelism, numberElements * key));
+ expectedResult.add(Tuple2.of(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism, keyGroupIndex) , numberElements * key));
}
assertEquals(expectedResult, actualResult);
@@ -401,7 +403,7 @@ public class RescalingITCase extends TestLogger {
for (int key = 0; key < numberKeys; key++) {
int keyGroupIndex = keyGroupAssigner2.getKeyGroupIndex(key);
- expectedResult2.add(Tuple2.of(keyGroupIndex % parallelism2, key * (numberElements + numberElements2)));
+ expectedResult2.add(Tuple2.of(KeyGroupRange.computeOperatorIndexForKeyGroup(maxParallelism, parallelism2, keyGroupIndex), key * (numberElements + numberElements2)));
}
assertEquals(expectedResult2, actualResult2);
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-tests/src/test/java/org/apache/flink/test/checkpointing/SavepointITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/SavepointITCase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/SavepointITCase.java
index 4fc310c..550ba75 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/SavepointITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/SavepointITCase.java
@@ -30,11 +30,10 @@ import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.akka.AkkaUtils;
-import org.apache.flink.runtime.blob.BlobKey;
import org.apache.flink.runtime.checkpoint.SubtaskState;
import org.apache.flink.runtime.checkpoint.TaskState;
import org.apache.flink.runtime.checkpoint.savepoint.SavepointStoreFactory;
-import org.apache.flink.runtime.checkpoint.savepoint.SavepointV0;
+import org.apache.flink.runtime.checkpoint.savepoint.SavepointV1;
import org.apache.flink.runtime.deployment.TaskDeploymentDescriptor;
import org.apache.flink.runtime.execution.SuppressRestartsException;
import org.apache.flink.runtime.instance.ActorGateway;
@@ -46,8 +45,10 @@ import org.apache.flink.runtime.messages.JobManagerMessages.DisposeSavepoint;
import org.apache.flink.runtime.messages.JobManagerMessages.TriggerSavepoint;
import org.apache.flink.runtime.messages.JobManagerMessages.TriggerSavepointFailure;
import org.apache.flink.runtime.messages.JobManagerMessages.TriggerSavepointSuccess;
+import org.apache.flink.runtime.state.ChainedStateHandle;
import org.apache.flink.runtime.state.CheckpointListener;
-import org.apache.flink.runtime.state.filesystem.AbstractFileStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.filesystem.FileStateHandle;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
import org.apache.flink.runtime.state.filesystem.FsStateBackendFactory;
import org.apache.flink.runtime.testingUtils.TestingJobManagerMessages.NotifyWhenJobRemoved;
@@ -61,8 +62,6 @@ import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.DiscardingSink;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskState;
-import org.apache.flink.streaming.runtime.tasks.StreamTaskStateList;
import org.apache.flink.test.util.ForkableFlinkMiniCluster;
import org.apache.flink.testutils.junit.RetryOnFailure;
import org.apache.flink.testutils.junit.RetryRule;
@@ -71,7 +70,6 @@ import org.junit.Rule;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import scala.Option;
import scala.concurrent.Await;
import scala.concurrent.Future;
import scala.concurrent.duration.Deadline;
@@ -219,7 +217,7 @@ public class SavepointITCase extends TestLogger {
new RequestSavepoint(savepointPath),
deadline.timeLeft());
- SavepointV0 savepoint = (SavepointV0) ((ResponseSavepoint) Await.result(
+ SavepointV1 savepoint = (SavepointV1) ((ResponseSavepoint) Await.result(
savepointFuture, deadline.timeLeft())).savepoint();
LOG.info("Retrieved savepoint: " + savepointPath + ".");
@@ -334,7 +332,7 @@ public class SavepointITCase extends TestLogger {
assertNotNull(subtaskState);
errMsg = "Initial operator state mismatch.";
- assertEquals(errMsg, subtaskState.getState(), tdd.getOperatorState());
+ assertEquals(errMsg, subtaskState.getChainedStateHandle(), tdd.getOperatorState());
}
}
@@ -345,7 +343,7 @@ public class SavepointITCase extends TestLogger {
LOG.info("Disposing savepoint " + savepointPath + ".");
Future<Object> disposeFuture = jobManager.ask(
- new DisposeSavepoint(savepointPath, Option.<List<BlobKey>>empty()),
+ new DisposeSavepoint(savepointPath),
deadline.timeLeft());
errMsg = "Failed to dispose savepoint " + savepointPath + ".";
@@ -360,14 +358,13 @@ public class SavepointITCase extends TestLogger {
for (TaskState stateForTaskGroup : savepoint.getTaskStates()) {
for (SubtaskState subtaskState : stateForTaskGroup.getStates()) {
- StreamTaskStateList taskStateList = (StreamTaskStateList) subtaskState.getState()
- .deserializeValue(ClassLoader.getSystemClassLoader());
+ ChainedStateHandle<StreamStateHandle> streamTaskState = subtaskState.getChainedStateHandle();
- for (StreamTaskState taskState : taskStateList.getState(
- ClassLoader.getSystemClassLoader())) {
-
- AbstractFileStateHandle fsState = (AbstractFileStateHandle) taskState.getFunctionState();
- checkpointFiles.add(new File(fsState.getFilePath().toUri()));
+ for (int i = 0; i < streamTaskState.getLength(); i++) {
+ if (streamTaskState.get(i) != null) {
+ FileStateHandle fileStateHandle = (FileStateHandle) streamTaskState.get(i);
+ checkpointFiles.add(new File(fileStateHandle.getFilePath().toUri()));
+ }
}
}
}
@@ -499,21 +496,20 @@ public class SavepointITCase extends TestLogger {
new RequestSavepoint(savepointPath),
deadline.timeLeft());
- SavepointV0 savepoint = (SavepointV0) ((ResponseSavepoint) Await.result(
+ SavepointV1 savepoint = (SavepointV1) ((ResponseSavepoint) Await.result(
savepointFuture, deadline.timeLeft())).savepoint();
LOG.info("Retrieved savepoint: " + savepointPath + ".");
// Check that all checkpoint files have been removed
for (TaskState stateForTaskGroup : savepoint.getTaskStates()) {
for (SubtaskState subtaskState : stateForTaskGroup.getStates()) {
- StreamTaskStateList taskStateList = (StreamTaskStateList) subtaskState.getState()
- .deserializeValue(ClassLoader.getSystemClassLoader());
+ ChainedStateHandle<StreamStateHandle> streamTaskState = subtaskState.getChainedStateHandle();
- for (StreamTaskState taskState : taskStateList.getState(
- ClassLoader.getSystemClassLoader())) {
-
- AbstractFileStateHandle fsState = (AbstractFileStateHandle) taskState.getFunctionState();
- checkpointFiles.add(new File(fsState.getFilePath().toUri()));
+ for (int i = 0; i < streamTaskState.getLength(); i++) {
+ if (streamTaskState.get(i) != null) {
+ FileStateHandle fileStateHandle = (FileStateHandle) streamTaskState.get(i);
+ checkpointFiles.add(new File(fileStateHandle.getFilePath().toUri()));
+ }
}
}
}
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java b/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java
index cdc7a80..3bc3cf5 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/classloading/ClassLoaderITCase.java
@@ -303,7 +303,7 @@ public class ClassLoaderITCase extends TestLogger {
// Dispose savepoint
LOG.info("Disposing savepoint at " + savepointPath);
- Future<Object> disposeFuture = jm.ask(new DisposeSavepoint(savepointPath, Option.apply(blobKeys)), deadline.timeLeft());
+ Future<Object> disposeFuture = jm.ask(new DisposeSavepoint(savepointPath), deadline.timeLeft());
Object disposeResponse = Await.result(disposeFuture, deadline.timeLeft());
if (disposeResponse.getClass() == JobManagerMessages.getDisposeSavepointSuccess().getClass()) {
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java b/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java
index c7d9a42..b3e2137 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/state/StateHandleSerializationTest.java
@@ -19,8 +19,8 @@
package org.apache.flink.test.state;
import org.apache.flink.runtime.state.KvStateSnapshot;
-import org.apache.flink.runtime.state.StateHandle;
+import org.apache.flink.runtime.state.StateObject;
import org.junit.Test;
import org.reflections.Reflections;
@@ -34,7 +34,7 @@ import static org.junit.Assert.*;
public class StateHandleSerializationTest {
/**
- * This test validates that all subclasses of {@link StateHandle} have a proper
+ * This test validates that all subclasses of {@link StateObject} have a proper
* serial version UID.
*/
@Test
@@ -46,7 +46,7 @@ public class StateHandleSerializationTest {
@SuppressWarnings("unchecked")
Set<Class<?>> stateHandleImplementations = (Set<Class<?>>) (Set<?>)
- reflections.getSubTypesOf(StateHandle.class);
+ reflections.getSubTypesOf(StateObject.class);
for (Class<?> clazz : stateHandleImplementations) {
validataSerialVersionUID(clazz);
@@ -73,7 +73,7 @@ public class StateHandleSerializationTest {
private static void validataSerialVersionUID(Class<?> clazz) {
// all non-interface types must have a serial version UID
if (!clazz.isInterface()) {
- assertFalse("Anonymous state handle classes have problematic serialization behavior",
+ assertFalse("Anonymous state handle classes have problematic serialization behavior: " + clazz,
clazz.isAnonymousClass());
try {
http://git-wip-us.apache.org/repos/asf/flink/blob/847ead01/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
----------------------------------------------------------------------
diff --git a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
index 6288946..41455cf 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/streaming/runtime/StateBackendITCase.java
@@ -34,13 +34,10 @@ import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.client.JobExecutionException;
import org.apache.flink.runtime.execution.Environment;
import org.apache.flink.runtime.state.AbstractStateBackend;
-import org.apache.flink.runtime.state.StateHandle;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.util.StreamingMultipleProgramsTestBase;
import org.junit.Test;
-import java.io.Serializable;
-
import static org.junit.Assert.fail;
public class StateBackendITCase extends StreamingMultipleProgramsTestBase {
@@ -132,13 +129,6 @@ public class StateBackendITCase extends StreamingMultipleProgramsTestBase {
long timestamp) throws Exception {
return null;
}
-
- @Override
- public <S extends Serializable> StateHandle<S> checkpointStateSerializable(S state,
- long checkpointID,
- long timestamp) throws Exception {
- return null;
- }
}
static final class SuccessException extends Exception {