You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by sr...@apache.org on 2018/08/02 12:34:46 UTC
[flink] 01/02: [FLINK-9981][state] Testing and performance tuning
for RocksDB-based priority queue
This is an automated email from the ASF dual-hosted git repository.
srichter pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit a20fd1d51ba9449f05b96b00e0580f937329bd47
Author: Stefan Richter <s....@data-artisans.com>
AuthorDate: Wed Jul 18 12:54:49 2018 +0200
[FLINK-9981][state] Testing and performance tuning for RocksDB-based priority queue
This closes #6438.
---
.../generated/rocks_db_configuration.html | 4 +-
.../core/memory/ByteArrayInputStreamWithPos.java | 18 +-
.../memory/ByteArrayInputStreamWithPosTest.java | 14 +
.../test-scripts/test_resume_savepoint.sh | 2 +-
flink-runtime/pom.xml | 6 -
.../flink/runtime/state/InternalPriorityQueue.java | 12 -
.../state/TieBreakingPriorityComparator.java | 117 -----
.../state/heap/AbstractHeapPriorityQueue.java | 222 ++++++++
.../heap/CachingInternalPriorityQueueSet.java | 356 -------------
.../runtime/state/heap/HeapPriorityQueue.java | 204 +-------
.../runtime/state/heap/HeapPriorityQueueSet.java | 6 +-
.../HeapPriorityQueueSnapshotRestoreWrapper.java | 10 -
.../heap/KeyGroupPartitionedPriorityQueue.java | 48 +-
.../runtime/state/heap/TreeOrderedSetCache.java | 135 -----
.../state/InternalPriorityQueueTestBase.java | 99 +++-
.../CachingInternalPriorityQueueSetTestBase.java | 43 --
.../heap/KeyGroupPartitionedPriorityQueueTest.java | 49 +-
.../state/heap/OrderedSetCacheTestBase.java | 102 ----
.../SimpleCachingInternalPriorityQueueSetTest.java | 35 --
.../flink/runtime/state/heap/TestOrderedStore.java | 60 ---
.../state/heap/TreeOrderedSetCacheTest.java | 30 --
.../state/RocksDBCachingPriorityQueueSet.java | 560 +++++++++++++++++++++
.../streaming/state/RocksDBKeyedStateBackend.java | 107 ++--
.../contrib/streaming/state/RocksDBOptions.java | 9 +-
.../streaming/state/RocksDBOrderedSetStore.java | 274 ----------
.../streaming/state/RocksDBStateBackend.java | 4 +-
.../streaming/state/TreeOrderedSetCache.java | 104 ++++
...ternalPriorityQueueSetWithRocksDBStoreTest.java | 65 ---
...rtitionedPriorityQueueWithRocksDBStoreTest.java | 56 ++-
.../streaming/state/RocksDBAsyncSnapshotTest.java | 2 +-
.../state/RocksDBOrderedSetStoreTest.java | 136 -----
.../state/RocksDBStateBackendConfigTest.java | 42 ++
.../api/operators/InternalTimeServiceManager.java | 20 +-
...rService.java => InternalTimerServiceImpl.java} | 48 +-
.../InternalTimerServiceSerializationProxy.java | 10 +-
...Test.java => InternalTimerServiceImplTest.java} | 50 +-
.../EventTimeWindowCheckpointingITCase.java | 39 +-
37 files changed, 1279 insertions(+), 1819 deletions(-)
diff --git a/docs/_includes/generated/rocks_db_configuration.html b/docs/_includes/generated/rocks_db_configuration.html
index 57b9511..8983f8b 100644
--- a/docs/_includes/generated/rocks_db_configuration.html
+++ b/docs/_includes/generated/rocks_db_configuration.html
@@ -13,9 +13,9 @@
<td>The local directory (on the TaskManager) where RocksDB puts its files.</td>
</tr>
<tr>
- <td><h5>state.backend.rocksdb.timer-service.impl</h5></td>
+ <td><h5>state.backend.rocksdb.timer-service.factory</h5></td>
<td style="word-wrap: break-word;">"HEAP"</td>
- <td>This determines the timer service implementation. Options are either HEAP (heap-based, default) or ROCKSDB for an implementation based on RocksDB.</td>
+ <td>This determines the factory for timer service state implementation. Options are either HEAP (heap-based, default) or ROCKSDB for an implementation based on RocksDB .</td>
</tr>
</tbody>
</table>
diff --git a/flink-core/src/main/java/org/apache/flink/core/memory/ByteArrayInputStreamWithPos.java b/flink-core/src/main/java/org/apache/flink/core/memory/ByteArrayInputStreamWithPos.java
index 1447e96..bc81593 100644
--- a/flink-core/src/main/java/org/apache/flink/core/memory/ByteArrayInputStreamWithPos.java
+++ b/flink-core/src/main/java/org/apache/flink/core/memory/ByteArrayInputStreamWithPos.java
@@ -30,20 +30,23 @@ import java.io.InputStream;
@Internal
public class ByteArrayInputStreamWithPos extends InputStream {
+ private static final byte[] EMPTY = new byte[0];
+
protected byte[] buffer;
protected int position;
protected int count;
protected int mark = 0;
+ public ByteArrayInputStreamWithPos() {
+ this(EMPTY);
+ }
+
public ByteArrayInputStreamWithPos(byte[] buffer) {
this(buffer, 0, buffer.length);
}
public ByteArrayInputStreamWithPos(byte[] buffer, int offset, int length) {
- this.position = offset;
- this.buffer = buffer;
- this.mark = offset;
- this.count = Math.min(buffer.length, offset + length);
+ setBuffer(buffer, offset, length);
}
@Override
@@ -122,4 +125,11 @@ public class ByteArrayInputStreamWithPos extends InputStream {
Preconditions.checkArgument(pos >= 0 && pos <= count, "Position out of bounds.");
this.position = pos;
}
+
+ public void setBuffer(byte[] buffer, int offset, int length) {
+ this.count = Math.min(buffer.length, offset + length);
+ setPosition(offset);
+ this.buffer = buffer;
+ this.mark = offset;
+ }
}
diff --git a/flink-core/src/test/java/org/apache/flink/core/memory/ByteArrayInputStreamWithPosTest.java b/flink-core/src/test/java/org/apache/flink/core/memory/ByteArrayInputStreamWithPosTest.java
index 4797c0a..62ad911 100644
--- a/flink-core/src/test/java/org/apache/flink/core/memory/ByteArrayInputStreamWithPosTest.java
+++ b/flink-core/src/test/java/org/apache/flink/core/memory/ByteArrayInputStreamWithPosTest.java
@@ -81,4 +81,18 @@ public class ByteArrayInputStreamWithPosTest {
thrown.expectMessage("Position out of bounds.");
stream.setPosition(-1);
}
+
+ @Test
+ public void testSetBuffer() {
+ ByteArrayInputStreamWithPos in = new ByteArrayInputStreamWithPos();
+ Assert.assertEquals(-1, in.read());
+ byte[] testData = new byte[]{0x42, 0x43, 0x44, 0x45};
+ int off = 1;
+ int len = 2;
+ in.setBuffer(testData, off, len);
+ for (int i = 0; i < len; ++i) {
+ Assert.assertEquals(testData[i + off], in.read());
+ }
+ Assert.assertEquals(-1, in.read());
+ }
}
diff --git a/flink-end-to-end-tests/test-scripts/test_resume_savepoint.sh b/flink-end-to-end-tests/test-scripts/test_resume_savepoint.sh
index 26a7e85..fd0bf91 100755
--- a/flink-end-to-end-tests/test-scripts/test_resume_savepoint.sh
+++ b/flink-end-to-end-tests/test-scripts/test_resume_savepoint.sh
@@ -59,7 +59,7 @@ backup_config
change_conf "taskmanager.numberOfTaskSlots" "1" "${NUM_SLOTS}"
if [ $STATE_BACKEND_ROCKS_TIMER_SERVICE_TYPE == 'rocks' ]; then
- set_conf "state.backend.rocksdb.timer-service.impl" "rocksdb"
+ set_conf "state.backend.rocksdb.timer-service.factory" "rocksdb"
fi
setup_flink_slf4j_metric_reporter
diff --git a/flink-runtime/pom.xml b/flink-runtime/pom.xml
index 271582f..e163881 100644
--- a/flink-runtime/pom.xml
+++ b/flink-runtime/pom.xml
@@ -116,12 +116,6 @@ under the License.
</dependency>
<dependency>
- <groupId>it.unimi.dsi</groupId>
- <artifactId>fastutil</artifactId>
- <version>8.2.1</version>
- </dependency>
-
- <dependency>
<groupId>org.scala-lang</groupId>
<artifactId>scala-library</artifactId>
</dependency>
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/InternalPriorityQueue.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/InternalPriorityQueue.java
index dc46c8a..fb3ee82 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/InternalPriorityQueue.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/InternalPriorityQueue.java
@@ -26,8 +26,6 @@ import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.Collection;
-import java.util.function.Consumer;
-import java.util.function.Predicate;
/**
* Interface for collection that gives in order access to elements w.r.t their priority.
@@ -38,16 +36,6 @@ import java.util.function.Predicate;
public interface InternalPriorityQueue<T> {
/**
- * Polls from the top of the queue as long as the the queue is not empty and passes the elements to
- * {@link Consumer} until a {@link Predicate} rejects an offered element. The rejected element is not
- * removed from the queue and becomes the new head.
- *
- * @param canConsume bulk polling ends once this returns false. The rejected element is nor removed and not consumed.
- * @param consumer consumer function for elements accepted by canConsume.
- */
- void bulkPoll(@Nonnull Predicate<T> canConsume, @Nonnull Consumer<T> consumer);
-
- /**
* Retrieves and removes the first element (w.r.t. the order) of this set,
* or returns {@code null} if this set is empty.
*
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/TieBreakingPriorityComparator.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/TieBreakingPriorityComparator.java
deleted file mode 100644
index 94aed45..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/TieBreakingPriorityComparator.java
+++ /dev/null
@@ -1,117 +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.typeutils.TypeSerializer;
-import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
-import org.apache.flink.util.FlinkRuntimeException;
-
-import javax.annotation.Nonnull;
-
-import java.io.IOException;
-import java.util.Comparator;
-
-/**
- * This class is an adapter between {@link PriorityComparator} and a full {@link Comparator} that respects the
- * contract between {@link Comparator#compare(Object, Object)} and {@link Object#equals(Object)}. This is currently
- * needed for implementations of
- * {@link org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet.OrderedSetCache} that are implemented
- * on top of a data structure that relies on the this contract, e.g. a tree set. We should replace this in the near
- * future.
- *
- * @param <T> type of the compared elements.
- */
-public class TieBreakingPriorityComparator<T> implements Comparator<T>, PriorityComparator<T> {
-
- /** The {@link PriorityComparator} to which we delegate in a first step. */
- @Nonnull
- private final PriorityComparator<T> priorityComparator;
-
- /** Serializer for instances of the compared objects. */
- @Nonnull
- private final TypeSerializer<T> serializer;
-
- /** Stream that we use in serialization. */
- @Nonnull
- private final ByteArrayOutputStreamWithPos outStream;
-
- /** {@link org.apache.flink.core.memory.DataOutputView} around outStream. */
- @Nonnull
- private final DataOutputViewStreamWrapper outView;
-
- public TieBreakingPriorityComparator(
- @Nonnull PriorityComparator<T> priorityComparator,
- @Nonnull TypeSerializer<T> serializer,
- @Nonnull ByteArrayOutputStreamWithPos outStream,
- @Nonnull DataOutputViewStreamWrapper outView) {
-
- this.priorityComparator = priorityComparator;
- this.serializer = serializer;
- this.outStream = outStream;
- this.outView = outView;
- }
-
- @SuppressWarnings("unchecked")
- @Override
- public int compare(T o1, T o2) {
-
- // first we compare priority, this should be the most commonly hit case
- int cmp = priorityComparator.comparePriority(o1, o2);
-
- if (cmp != 0) {
- return cmp;
- }
-
- // here we start tie breaking and do our best to comply with the compareTo/equals contract, first we try
- // to simply find an existing way to fully compare.
- if (o1 instanceof Comparable && o1.getClass().equals(o2.getClass())) {
- return ((Comparable<T>) o1).compareTo(o2);
- }
-
- // if objects are not equal, their serialized form should somehow differ as well. this can be costly, and...
- // TODO we should have an alternative approach in the future, e.g. a cache that does not rely on compare to check equality.
- try {
- outStream.reset();
- serializer.serialize(o1, outView);
- int leftLen = outStream.getPosition();
- serializer.serialize(o2, outView);
- int rightLen = outStream.getPosition() - leftLen;
- return compareBytes(outStream.getBuf(), 0, leftLen, leftLen, rightLen);
- } catch (IOException ex) {
- throw new FlinkRuntimeException("Serializer problem in comparator.", ex);
- }
- }
-
- @Override
- public int comparePriority(T left, T right) {
- return priorityComparator.comparePriority(left, right);
- }
-
- public static int compareBytes(byte[] bytes, int offLeft, int leftLen, int offRight, int rightLen) {
- int maxLen = Math.min(leftLen, rightLen);
- for (int i = 0; i < maxLen; ++i) {
- int cmp = bytes[offLeft + i] - bytes[offRight + i];
- if (cmp != 0) {
- return cmp;
- }
- }
- return leftLen - rightLen;
- }
-}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapPriorityQueue.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapPriorityQueue.java
new file mode 100644
index 0000000..f2981bd
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/AbstractHeapPriorityQueue.java
@@ -0,0 +1,222 @@
+/*
+ * 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.runtime.state.InternalPriorityQueue;
+import org.apache.flink.util.CloseableIterator;
+
+import javax.annotation.Nonnegative;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+
+import static org.apache.flink.util.CollectionUtil.MAX_ARRAY_SIZE;
+
+/**
+ * Abstract base class for heap (object array) based implementations of priority queues, with support for fast deletes
+ * via {@link HeapPriorityQueueElement}.
+ *
+ * @param <T> type of the elements contained in the priority queue.
+ */
+public abstract class AbstractHeapPriorityQueue<T extends HeapPriorityQueueElement>
+ implements InternalPriorityQueue<T> {
+
+ /** The array that represents the heap-organized priority queue. */
+ @Nonnull
+ protected T[] queue;
+
+ /** The current size of the priority queue. */
+ @Nonnegative
+ protected int size;
+
+ @SuppressWarnings("unchecked")
+ public AbstractHeapPriorityQueue(@Nonnegative int minimumCapacity) {
+ this.queue = (T[]) new HeapPriorityQueueElement[getHeadElementIndex() + minimumCapacity];
+ this.size = 0;
+ }
+
+ @Override
+ @Nullable
+ public T poll() {
+ return size() > 0 ? removeInternal(getHeadElementIndex()) : null;
+ }
+
+ @Override
+ @Nullable
+ public T peek() {
+ // References to removed elements are expected to become set to null.
+ return queue[getHeadElementIndex()];
+ }
+
+ @Override
+ public boolean add(@Nonnull T toAdd) {
+ addInternal(toAdd);
+ return toAdd.getInternalIndex() == getHeadElementIndex();
+ }
+
+ @Override
+ public boolean remove(@Nonnull T toRemove) {
+ final int elementIndex = toRemove.getInternalIndex();
+ removeInternal(elementIndex);
+ return elementIndex == getHeadElementIndex();
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return size() == 0;
+ }
+
+ @Override
+ public int size() {
+ return size;
+ }
+
+ @Override
+ public void addAll(@Nullable Collection<? extends T> toAdd) {
+
+ if (toAdd == null) {
+ return;
+ }
+
+ resizeForBulkLoad(toAdd.size());
+
+ for (T element : toAdd) {
+ add(element);
+ }
+ }
+
+ @SuppressWarnings({"unchecked"})
+ @Nonnull
+ public <O> O[] toArray(O[] out) {
+ final int heapArrayOffset = getHeadElementIndex();
+ if (out.length < size) {
+ return (O[]) Arrays.copyOfRange(queue, heapArrayOffset, heapArrayOffset + size, out.getClass());
+ } else {
+ System.arraycopy(queue, heapArrayOffset, out, 0, size);
+ if (out.length > size) {
+ out[size] = null;
+ }
+ return out;
+ }
+ }
+
+ /**
+ * Returns an iterator over the elements in this queue. The iterator
+ * does not return the elements in any particular order.
+ *
+ * @return an iterator over the elements in this queue.
+ */
+ @Nonnull
+ @Override
+ public CloseableIterator<T> iterator() {
+ return new HeapIterator();
+ }
+
+ /**
+ * Clears the queue.
+ */
+ public void clear() {
+ final int arrayOffset = getHeadElementIndex();
+ Arrays.fill(queue, arrayOffset, arrayOffset + size, null);
+ size = 0;
+ }
+
+ protected void resizeForBulkLoad(int totalSize) {
+ if (totalSize > queue.length) {
+ int desiredSize = totalSize + (totalSize >>> 3);
+ resizeQueueArray(desiredSize, totalSize);
+ }
+ }
+
+ protected void resizeQueueArray(int desiredSize, int minRequiredSize) {
+ if (isValidArraySize(desiredSize)) {
+ queue = Arrays.copyOf(queue, desiredSize);
+ } else if (isValidArraySize(minRequiredSize)) {
+ queue = Arrays.copyOf(queue, MAX_ARRAY_SIZE);
+ } else {
+ throw new OutOfMemoryError("Required minimum heap size " + minRequiredSize +
+ " exceeds maximum size of " + MAX_ARRAY_SIZE + ".");
+ }
+ }
+
+ protected void moveElementToIdx(T element, int idx) {
+ queue[idx] = element;
+ element.setInternalIndex(idx);
+ }
+
+ /**
+ * Implements how to remove the element at the given index from the queue.
+ *
+ * @param elementIndex the index to remove.
+ * @return the removed element.
+ */
+ protected abstract T removeInternal(@Nonnegative int elementIndex);
+
+ /**
+ * Implements how to add an element to the queue.
+ *
+ * @param toAdd the element to add.
+ */
+ protected abstract void addInternal(@Nonnull T toAdd);
+
+ /**
+ * Returns the start index of the queue elements in the array.
+ */
+ protected abstract int getHeadElementIndex();
+
+ private static boolean isValidArraySize(int size) {
+ return size >= 0 && size <= MAX_ARRAY_SIZE;
+ }
+
+ /**
+ * {@link Iterator} implementation for {@link HeapPriorityQueue}.
+ * {@link Iterator#remove()} is not supported.
+ */
+ private final class HeapIterator implements CloseableIterator<T> {
+
+ private int runningIdx;
+ private final int endIdx;
+
+ HeapIterator() {
+ this.runningIdx = getHeadElementIndex();
+ this.endIdx = runningIdx + size;
+ }
+
+ @Override
+ public boolean hasNext() {
+ return runningIdx < endIdx;
+ }
+
+ @Override
+ public T next() {
+ if (runningIdx >= endIdx) {
+ throw new NoSuchElementException("Iterator has no next element.");
+ }
+ return queue[runningIdx++];
+ }
+
+ @Override
+ public void close() {
+ }
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CachingInternalPriorityQueueSet.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CachingInternalPriorityQueueSet.java
deleted file mode 100644
index b1ad0df..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/CachingInternalPriorityQueueSet.java
+++ /dev/null
@@ -1,356 +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.heap;
-
-import org.apache.flink.runtime.state.InternalPriorityQueue;
-import org.apache.flink.util.CloseableIterator;
-import org.apache.flink.util.FlinkRuntimeException;
-
-import javax.annotation.Nonnegative;
-import javax.annotation.Nonnull;
-import javax.annotation.Nullable;
-
-import java.util.Collection;
-import java.util.function.Consumer;
-import java.util.function.Predicate;
-
-import static org.apache.flink.runtime.state.heap.KeyGroupPartitionedPriorityQueue.ENABLE_RELAXED_FIRING_ORDER_OPTIMIZATION;
-
-/**
- * This class is an implementation of a {@link InternalPriorityQueue} with set semantics that internally consists of
- * two different storage types. The first storage is a (potentially slow) ordered set store manages the ground truth
- * about the elements in this queue. The second storage is a (fast) ordered set cache, typically with some limited
- * capacity. The cache is used to improve performance of accesses to the underlying store and contains an ordered
- * (partial) view on the top elements in the ordered store. We are currently applying simple write-through to keep cache
- * and store in sync on updates and refill the cache from the store when it is empty and we expect more elements
- * contained in the store.
- *
- * @param <E> the type if the managed elements.
- */
-public class CachingInternalPriorityQueueSet<E> implements InternalPriorityQueue<E>, HeapPriorityQueueElement {
-
- /** A ordered set cache that contains a (partial) view on the top elements in the ordered store. */
- @Nonnull
- private final OrderedSetCache<E> orderedCache;
-
- /** A store with ordered set semantics that contains the ground truth of all inserted elements. */
- @Nonnull
- private final OrderedSetStore<E> orderedStore;
-
- /** This flag is true if there could be elements in the backend that are not in the cache (false positives ok). */
- private boolean storeOnlyElements;
-
- /** Management data for the {@link HeapPriorityQueueElement} trait. */
- private int pqManagedIndex;
-
- @SuppressWarnings("unchecked")
- public CachingInternalPriorityQueueSet(
- @Nonnull OrderedSetCache<E> orderedCache,
- @Nonnull OrderedSetStore<E> orderedStore) {
- this.orderedCache = orderedCache;
- this.orderedStore = orderedStore;
- // We are careful and set this to true. Could be set to false if we know for sure that the store is empty, but
- // checking this could be an expensive operation.
- this.storeOnlyElements = true;
- this.pqManagedIndex = HeapPriorityQueueElement.NOT_CONTAINED;
- }
-
- @Nullable
- @Override
- public E peek() {
-
- checkRefillCacheFromStore();
-
- return orderedCache.peekFirst();
- }
-
- @Override
- public void bulkPoll(@Nonnull Predicate<E> canConsume, @Nonnull Consumer<E> consumer) {
- if (ENABLE_RELAXED_FIRING_ORDER_OPTIMIZATION) {
- bulkPollRelaxedOrder(canConsume, consumer);
- } else {
- bulkPollStrictOrder(canConsume, consumer);
- }
- }
-
- private void bulkPollRelaxedOrder(@Nonnull Predicate<E> canConsume, @Nonnull Consumer<E> consumer) {
- if (orderedCache.isEmpty()) {
- bulkPollStore(canConsume, consumer);
- } else {
- while (!orderedCache.isEmpty() && canConsume.test(orderedCache.peekFirst())) {
- final E next = orderedCache.removeFirst();
- orderedStore.remove(next);
- consumer.accept(next);
- }
-
- if (orderedCache.isEmpty()) {
- bulkPollStore(canConsume, consumer);
- }
- }
- }
-
- private void bulkPollStrictOrder(@Nonnull Predicate<E> canConsume, @Nonnull Consumer<E> consumer) {
- E element;
- while ((element = peek()) != null && canConsume.test(element)) {
- poll();
- consumer.accept(element);
- }
- }
-
- private void bulkPollStore(@Nonnull Predicate<E> canConsume, @Nonnull Consumer<E> consumer) {
- try (CloseableIterator<E> iterator = orderedStore.orderedIterator()) {
- while (iterator.hasNext()) {
- final E next = iterator.next();
- if (canConsume.test(next)) {
- orderedStore.remove(next);
- consumer.accept(next);
- } else {
- orderedCache.add(next);
- while (iterator.hasNext() && !orderedCache.isFull()) {
- orderedCache.add(iterator.next());
- }
- break;
- }
- }
- } catch (Exception e) {
- throw new FlinkRuntimeException("Exception while bulk polling store.", e);
- }
- }
-
- @Nullable
- @Override
- public E poll() {
-
- checkRefillCacheFromStore();
-
- final E first = orderedCache.removeFirst();
-
- if (first != null) {
- // write-through sync
- orderedStore.remove(first);
- }
-
- return first;
- }
-
- @Override
- public boolean add(@Nonnull E toAdd) {
-
- checkRefillCacheFromStore();
-
- // write-through sync
- orderedStore.add(toAdd);
-
- final boolean cacheFull = orderedCache.isFull();
-
- if ((!cacheFull && !storeOnlyElements) || orderedCache.isInLowerBound(toAdd)) {
-
- if (cacheFull) {
- // we drop the element with lowest priority from the cache
- orderedCache.removeLast();
- // the dropped element is now only in the store
- storeOnlyElements = true;
- }
-
- orderedCache.add(toAdd);
- return toAdd.equals(orderedCache.peekFirst());
- } else {
- // we only added to the store
- storeOnlyElements = true;
- return false;
- }
- }
-
- @Override
- public boolean remove(@Nonnull E toRemove) {
-
- checkRefillCacheFromStore();
-
- boolean newHead = toRemove.equals(orderedCache.peekFirst());
- // write-through sync
- orderedStore.remove(toRemove);
- orderedCache.remove(toRemove);
- return newHead;
- }
-
- @Override
- public void addAll(@Nullable Collection<? extends E> toAdd) {
-
- if (toAdd == null) {
- return;
- }
-
- for (E element : toAdd) {
- add(element);
- }
- }
-
- @Override
- public int size() {
- return orderedStore.size();
- }
-
- @Override
- public boolean isEmpty() {
- checkRefillCacheFromStore();
- return orderedCache.isEmpty();
- }
-
- @Nonnull
- @Override
- public CloseableIterator<E> iterator() {
- if (storeOnlyElements) {
- return orderedStore.orderedIterator();
- } else {
- return orderedCache.orderedIterator();
- }
- }
-
- @Override
- public int getInternalIndex() {
- return pqManagedIndex;
- }
-
- @Override
- public void setInternalIndex(int updateIndex) {
- this.pqManagedIndex = updateIndex;
- }
-
- /**
- * Refills the cache from the store when the cache is empty and we expect more elements in the store.
- *
- * TODO: We can think about exploiting the property that the store is already sorted when bulk-filling the cache.
- */
- private void checkRefillCacheFromStore() {
- if (storeOnlyElements && orderedCache.isEmpty()) {
- try (final CloseableIterator<E> iterator = orderedStore.orderedIterator()) {
- while (iterator.hasNext() && !orderedCache.isFull()) {
- orderedCache.add(iterator.next());
- }
- storeOnlyElements = iterator.hasNext();
- } catch (Exception e) {
- throw new FlinkRuntimeException("Exception while refilling store from iterator.", e);
- }
- }
- }
-
- /**
- * Interface for an ordered cache with set semantics of elements to be used in
- * {@link CachingInternalPriorityQueueSet}. Cache implementations typically have a limited capacity as indicated
- * via the {@link #isFull()} method.
- *
- * @param <E> the type of contained elements.
- */
- public interface OrderedSetCache<E> {
-
- /**
- * Adds the given element to the cache (if not yet contained). This method should only be called if the cache
- * is not full.
- * @param element element to add to the cache.
- */
- void add(@Nonnull E element);
-
- /**
- * Removes the given element from the cache (if contained).
- * @param element element to remove from the cache.
- */
- void remove(@Nonnull E element);
-
- /**
- * Returns <code>true</code> if the cache is full and no more elements can be added.
- */
- boolean isFull();
-
- /**
- * Returns <code>true</code> if the cache is empty, i.e. contains ne elements.
- */
- boolean isEmpty();
-
- /**
- * Returns true, if the element is compares smaller than the currently largest element in the cache.
- */
- boolean isInLowerBound(@Nonnull E toCheck);
-
- /**
- * Removes and returns the first (smallest) element from the cache.
- */
- @Nullable
- E removeFirst();
-
- /**
- * Removes and returns the last (larges) element from the cache.
- */
- @Nullable
- E removeLast();
-
- /**
- * Returns the first (smallest) element from the cache (without removing it).
- */
- @Nullable
- E peekFirst();
-
- /**
- * Returns the last (larges) element from the cache (without removing it).
- */
- @Nullable
- E peekLast();
-
- /**
- * Returns an iterator over the store that returns element in order. The iterator must be closed by the client
- * after usage.
- */
- @Nonnull
- CloseableIterator<E> orderedIterator();
- }
-
- /**
- * Interface for an ordered store with set semantics of elements to be used in
- * {@link CachingInternalPriorityQueueSet}. Stores are assumed to have (practically) unlimited capacity, but their
- * operations could all be expensive.
- *
- * @param <E> the type of contained elements.
- */
- public interface OrderedSetStore<E> {
-
- /**
- * Adds the given element to the store (if not yet contained).
- * @param element element to add to the store.
- */
- void add(@Nonnull E element);
-
- /**
- * Removed the given element from the cache (if contained).
- * @param element element to remove from the cache.
- */
- void remove(@Nonnull E element);
-
- /**
- * Returns the number of elements in the store.
- */
- @Nonnegative
- int size();
-
- /**
- * Returns an iterator over the store that returns element in order. The iterator must be closed by the client
- * after usage.
- */
- @Nonnull
- CloseableIterator<E> orderedIterator();
- }
-}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueue.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueue.java
index 22b2419..2dd5701 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueue.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueue.java
@@ -18,22 +18,10 @@
package org.apache.flink.runtime.state.heap;
-import org.apache.flink.runtime.state.InternalPriorityQueue;
import org.apache.flink.runtime.state.PriorityComparator;
-import org.apache.flink.util.CloseableIterator;
import javax.annotation.Nonnegative;
import javax.annotation.Nonnull;
-import javax.annotation.Nullable;
-
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Iterator;
-import java.util.NoSuchElementException;
-import java.util.function.Consumer;
-import java.util.function.Predicate;
-
-import static org.apache.flink.util.CollectionUtil.MAX_ARRAY_SIZE;
/**
* Basic heap-based priority queue for {@link HeapPriorityQueueElement} objects. This heap supports fast deletes
@@ -51,7 +39,7 @@ import static org.apache.flink.util.CollectionUtil.MAX_ARRAY_SIZE;
* @param <T> type of the contained elements.
*/
public class HeapPriorityQueue<T extends HeapPriorityQueueElement>
- implements InternalPriorityQueue<T> {
+ extends AbstractHeapPriorityQueue<T> {
/**
* The index of the head element in the array that represents the heap.
@@ -61,17 +49,8 @@ public class HeapPriorityQueue<T extends HeapPriorityQueueElement>
/**
* Comparator for the priority of contained elements.
*/
- private final PriorityComparator<T> elementPriorityComparator;
-
- /**
- * The array that represents the heap-organized priority queue.
- */
- private T[] queue;
-
- /**
- * The current size of the priority queue.
- */
- private int size;
+ @Nonnull
+ protected final PriorityComparator<T> elementPriorityComparator;
/**
* Creates an empty {@link HeapPriorityQueue} with the requested initial capacity.
@@ -83,124 +62,31 @@ public class HeapPriorityQueue<T extends HeapPriorityQueueElement>
public HeapPriorityQueue(
@Nonnull PriorityComparator<T> elementPriorityComparator,
@Nonnegative int minimumCapacity) {
-
+ super(minimumCapacity);
this.elementPriorityComparator = elementPriorityComparator;
- this.queue = (T[]) new HeapPriorityQueueElement[QUEUE_HEAD_INDEX + minimumCapacity];
}
- @Override
- public void bulkPoll(@Nonnull Predicate<T> canConsume, @Nonnull Consumer<T> consumer) {
- T element;
- while ((element = peek()) != null && canConsume.test(element)) {
- poll();
- consumer.accept(element);
+ public void adjustModifiedElement(@Nonnull T element) {
+ final int elementIndex = element.getInternalIndex();
+ if (element == queue[elementIndex]) {
+ adjustElementAtIndex(element, elementIndex);
}
}
@Override
- @Nullable
- public T poll() {
- return size() > 0 ? removeElementAtIndex(QUEUE_HEAD_INDEX) : null;
- }
-
- @Override
- @Nullable
- public T peek() {
- return size() > 0 ? queue[QUEUE_HEAD_INDEX] : null;
- }
-
- /**
- * Adds the element to add to the heap. This element should not be managed by any other {@link HeapPriorityQueue}.
- *
- * @return <code>true</code> if the operation changed the head element or if is it unclear if the head element changed.
- * Only returns <code>false</code> iff the head element was not changed by this operation.
- */
- @Override
- public boolean add(@Nonnull T toAdd) {
- return addInternal(toAdd);
- }
-
- /**
- * This remove is based on object identity, not the result of equals. We use the objects managed index to find
- * the instance in the queue array.
- *
- * @return <code>true</code> if the operation changed the head element or if is it unclear if the head element changed.
- * Only returns <code>false</code> iff the head element was not changed by this operation.
- */
- @Override
- public boolean remove(@Nonnull T toRemove) {
- return removeInternal(toRemove);
- }
-
- @Override
- public boolean isEmpty() {
- return size() == 0;
- }
-
- @Override
- @Nonnegative
- public int size() {
- return size;
- }
-
- public void clear() {
- size = 0;
- Arrays.fill(queue, null);
- }
-
- @SuppressWarnings({"unchecked"})
- @Nonnull
- public <O> O[] toArray(O[] out) {
- if (out.length < size) {
- return (O[]) Arrays.copyOfRange(queue, QUEUE_HEAD_INDEX, QUEUE_HEAD_INDEX + size, out.getClass());
- } else {
- System.arraycopy(queue, QUEUE_HEAD_INDEX, out, 0, size);
- if (out.length > size) {
- out[size] = null;
- }
- return out;
- }
- }
-
- /**
- * Returns an iterator over the elements in this queue. The iterator
- * does not return the elements in any particular order.
- *
- * @return an iterator over the elements in this queue.
- */
- @Nonnull
- public CloseableIterator<T> iterator() {
- return new HeapIterator();
+ protected int getHeadElementIndex() {
+ return QUEUE_HEAD_INDEX;
}
@Override
- public void addAll(@Nullable Collection<? extends T> restoredElements) {
-
- if (restoredElements == null) {
- return;
- }
-
- resizeForBulkLoad(restoredElements.size());
-
- for (T element : restoredElements) {
- add(element);
- }
- }
-
- private boolean addInternal(@Nonnull T element) {
+ protected void addInternal(@Nonnull T element) {
final int newSize = increaseSizeByOne();
moveElementToIdx(element, newSize);
siftUp(newSize);
- return element.getInternalIndex() == QUEUE_HEAD_INDEX;
}
- private boolean removeInternal(@Nonnull T elementToRemove) {
- final int elementIndex = elementToRemove.getInternalIndex();
- removeElementAtIndex(elementIndex);
- return elementIndex == QUEUE_HEAD_INDEX;
- }
-
- private T removeElementAtIndex(int removeIdx) {
+ @Override
+ protected T removeInternal(int removeIdx) {
T[] heap = this.queue;
T removedValue = heap[removeIdx];
@@ -220,13 +106,6 @@ public class HeapPriorityQueue<T extends HeapPriorityQueueElement>
return removedValue;
}
- public void adjustModifiedElement(@Nonnull T element) {
- final int elementIndex = element.getInternalIndex();
- if (element == queue[elementIndex]) {
- adjustElementAtIndex(element, elementIndex);
- }
- }
-
private void adjustElementAtIndex(T element, int index) {
siftDown(index);
if (queue[index] == element) {
@@ -285,11 +164,6 @@ public class HeapPriorityQueue<T extends HeapPriorityQueueElement>
return elementPriorityComparator.comparePriority(a, b) < 0;
}
- private void moveElementToIdx(T element, int idx) {
- queue[idx] = element;
- element.setInternalIndex(idx);
- }
-
private int increaseSizeByOne() {
final int oldArraySize = queue.length;
final int minRequiredNewSize = ++size;
@@ -300,56 +174,4 @@ public class HeapPriorityQueue<T extends HeapPriorityQueueElement>
// TODO implement shrinking as well?
return minRequiredNewSize;
}
-
- private void resizeForBulkLoad(int totalSize) {
- if (totalSize > queue.length) {
- int desiredSize = totalSize + (totalSize >>> 3);
- resizeQueueArray(desiredSize, totalSize);
- }
- }
-
- private void resizeQueueArray(int desiredSize, int minRequiredSize) {
- if (isValidArraySize(desiredSize)) {
- queue = Arrays.copyOf(queue, desiredSize);
- } else if (isValidArraySize(minRequiredSize)) {
- queue = Arrays.copyOf(queue, MAX_ARRAY_SIZE);
- } else {
- throw new OutOfMemoryError("Required minimum heap size " + minRequiredSize +
- " exceeds maximum size of " + MAX_ARRAY_SIZE + ".");
- }
- }
-
- private static boolean isValidArraySize(int size) {
- return size >= 0 && size <= MAX_ARRAY_SIZE;
- }
-
- /**
- * {@link Iterator} implementation for {@link HeapPriorityQueue}.
- * {@link Iterator#remove()} is not supported.
- */
- private class HeapIterator implements CloseableIterator<T> {
-
- private int iterationIdx;
-
- HeapIterator() {
- this.iterationIdx = QUEUE_HEAD_INDEX - 1;
- }
-
- @Override
- public boolean hasNext() {
- return iterationIdx < size;
- }
-
- @Override
- public T next() {
- if (iterationIdx >= size) {
- throw new NoSuchElementException("Iterator has no next element.");
- }
- return queue[++iterationIdx];
- }
-
- @Override
- public void close() {
- }
- }
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueueSet.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueueSet.java
index 79f319c..c0215ac 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueueSet.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueueSet.java
@@ -106,11 +106,7 @@ public class HeapPriorityQueueSet<T extends HeapPriorityQueueElement>
@Nullable
public T poll() {
final T toRemove = super.poll();
- if (toRemove != null) {
- return getDedupMapForElement(toRemove).remove(toRemove);
- } else {
- return null;
- }
+ return toRemove != null ? getDedupMapForElement(toRemove).remove(toRemove) : null;
}
/**
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueueSnapshotRestoreWrapper.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueueSnapshotRestoreWrapper.java
index 5fd67f0..b2b2843 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueueSnapshotRestoreWrapper.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueueSnapshotRestoreWrapper.java
@@ -68,16 +68,6 @@ public class HeapPriorityQueueSnapshotRestoreWrapper<T extends HeapPriorityQueue
@Override
public StateSnapshot stateSnapshot() {
final T[] queueDump = (T[]) priorityQueue.toArray(new HeapPriorityQueueElement[priorityQueue.size()]);
-
- final TypeSerializer<T> elementSerializer = metaInfo.getElementSerializer();
-
- // turn the flat copy into a deep copy if required.
- if (!elementSerializer.isImmutableType()) {
- for (int i = 0; i < queueDump.length; ++i) {
- queueDump[i] = elementSerializer.copy(queueDump[i]);
- }
- }
-
return new HeapPriorityQueueStateSnapshot<>(
queueDump,
keyExtractorFunction,
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/KeyGroupPartitionedPriorityQueue.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/KeyGroupPartitionedPriorityQueue.java
index d8b0a5a..7978764 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/KeyGroupPartitionedPriorityQueue.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/KeyGroupPartitionedPriorityQueue.java
@@ -35,8 +35,6 @@ import javax.annotation.Nullable;
import java.util.Collection;
import java.util.HashSet;
import java.util.Set;
-import java.util.function.Consumer;
-import java.util.function.Predicate;
/**
* This implementation of {@link InternalPriorityQueue} is internally partitioned into sub-queues per key-group and
@@ -49,11 +47,9 @@ import java.util.function.Predicate;
public class KeyGroupPartitionedPriorityQueue<T, PQ extends InternalPriorityQueue<T> & HeapPriorityQueueElement>
implements InternalPriorityQueue<T>, KeyGroupedInternalPriorityQueue<T> {
- static final boolean ENABLE_RELAXED_FIRING_ORDER_OPTIMIZATION = false;
-
/** A heap of heap sets. Each sub-heap represents the partition for a key-group.*/
@Nonnull
- private final HeapPriorityQueue<PQ> heapOfkeyGroupedHeaps;
+ private final HeapPriorityQueue<PQ> heapOfKeyGroupedHeaps;
/** All elements from keyGroupHeap, indexed by their key-group id, relative to firstKeyGroup. */
@Nonnull
@@ -83,55 +79,30 @@ public class KeyGroupPartitionedPriorityQueue<T, PQ extends InternalPriorityQueu
this.totalKeyGroups = totalKeyGroups;
this.firstKeyGroup = keyGroupRange.getStartKeyGroup();
this.keyGroupedHeaps = (PQ[]) new InternalPriorityQueue[keyGroupRange.getNumberOfKeyGroups()];
- this.heapOfkeyGroupedHeaps = new HeapPriorityQueue<>(
+ this.heapOfKeyGroupedHeaps = new HeapPriorityQueue<>(
new InternalPriorityQueueComparator<>(elementPriorityComparator),
keyGroupRange.getNumberOfKeyGroups());
for (int i = 0; i < keyGroupedHeaps.length; i++) {
final PQ keyGroupSubHeap =
- orderedCacheFactory.create(firstKeyGroup + i, totalKeyGroups, elementPriorityComparator);
+ orderedCacheFactory.create(firstKeyGroup + i, totalKeyGroups, keyExtractor, elementPriorityComparator);
keyGroupedHeaps[i] = keyGroupSubHeap;
- heapOfkeyGroupedHeaps.add(keyGroupSubHeap);
- }
- }
-
- @Override
- public void bulkPoll(@Nonnull Predicate<T> canConsume, @Nonnull Consumer<T> consumer) {
- if (ENABLE_RELAXED_FIRING_ORDER_OPTIMIZATION) {
- bulkPollRelaxedOrder(canConsume, consumer);
- } else {
- bulkPollStrictOrder(canConsume, consumer);
- }
- }
-
- private void bulkPollRelaxedOrder(@Nonnull Predicate<T> canConsume, @Nonnull Consumer<T> consumer) {
- PQ headList = heapOfkeyGroupedHeaps.peek();
- while (headList.peek() != null && canConsume.test(headList.peek())) {
- headList.bulkPoll(canConsume, consumer);
- heapOfkeyGroupedHeaps.adjustModifiedElement(headList);
- }
- }
-
- private void bulkPollStrictOrder(@Nonnull Predicate<T> canConsume, @Nonnull Consumer<T> consumer) {
- T element;
- while ((element = peek()) != null && canConsume.test(element)) {
- poll();
- consumer.accept(element);
+ heapOfKeyGroupedHeaps.add(keyGroupSubHeap);
}
}
@Nullable
@Override
public T poll() {
- final PQ headList = heapOfkeyGroupedHeaps.peek();
+ final PQ headList = heapOfKeyGroupedHeaps.peek();
final T head = headList.poll();
- heapOfkeyGroupedHeaps.adjustModifiedElement(headList);
+ heapOfKeyGroupedHeaps.adjustModifiedElement(headList);
return head;
}
@Nullable
@Override
public T peek() {
- return heapOfkeyGroupedHeaps.peek().peek();
+ return heapOfKeyGroupedHeaps.peek().peek();
}
@Override
@@ -140,7 +111,7 @@ public class KeyGroupPartitionedPriorityQueue<T, PQ extends InternalPriorityQueu
// the branch checks if the head element has (potentially) changed.
if (list.add(toAdd)) {
- heapOfkeyGroupedHeaps.adjustModifiedElement(list);
+ heapOfKeyGroupedHeaps.adjustModifiedElement(list);
// could we have a new head?
return toAdd.equals(peek());
} else {
@@ -157,7 +128,7 @@ public class KeyGroupPartitionedPriorityQueue<T, PQ extends InternalPriorityQueu
// the branch checks if the head element has (potentially) changed.
if (list.remove(toRemove)) {
- heapOfkeyGroupedHeaps.adjustModifiedElement(list);
+ heapOfKeyGroupedHeaps.adjustModifiedElement(list);
// could we have a new head?
return toRemove.equals(oldHead);
} else {
@@ -330,6 +301,7 @@ public class KeyGroupPartitionedPriorityQueue<T, PQ extends InternalPriorityQueu
PQS create(
@Nonnegative int keyGroupId,
@Nonnegative int numKeyGroups,
+ @Nonnull KeyExtractorFunction<T> keyExtractorFunction,
@Nonnull PriorityComparator<T> elementPriorityComparator);
}
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/TreeOrderedSetCache.java b/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/TreeOrderedSetCache.java
deleted file mode 100644
index 14c281e..0000000
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/TreeOrderedSetCache.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.heap;
-
-import org.apache.flink.util.CloseableIterator;
-import org.apache.flink.util.Preconditions;
-
-import it.unimi.dsi.fastutil.objects.ObjectAVLTreeSet;
-
-import javax.annotation.Nonnegative;
-import javax.annotation.Nonnull;
-import javax.annotation.Nullable;
-
-import java.util.Comparator;
-
-/**
- * Implementation of {@link org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet.OrderedSetCache} based
- * on an AVL-Tree. We chose the implementation from fastutil over JDK for performance reasons.
- *
- * <p>Maintainer notes: We can consider the following potential performance improvements. First, we could introduce a
- * bulk-load method to OrderedSetCache to exploit the fact that adding from an OrderedSetStore is already happening in
- * sorted order, e.g. there are more efficient ways to construct search trees from sorted elements. Second, we could
- * replace the internal AVL-Tree with an extended variant of {@link HeapPriorityQueueSet} that is organized as a
- * Min-Max-Heap.
- *
- * @param <E> type of the contained elements.
- */
-public class TreeOrderedSetCache<E> implements CachingInternalPriorityQueueSet.OrderedSetCache<E> {
-
- /** The tree is used to store cached elements. */
- @Nonnull
- private final ObjectAVLTreeSet<E> avlTree;
-
- /** The element comparator. */
- @Nonnull
- private final Comparator<E> elementComparator;
-
- /** The maximum capacity of the cache. */
- @Nonnegative
- private final int capacity;
-
- /**
- * Creates a new {@link TreeOrderedSetCache} with the given capacity and element comparator. Capacity must be > 0.
- * @param elementComparator comparator for the cached elements.
- * @param capacity the capacity of the cache. Must be > 0.
- */
- public TreeOrderedSetCache(@Nonnull Comparator<E> elementComparator, @Nonnegative int capacity) {
- Preconditions.checkArgument(capacity > 0, "Cache capacity must be greater than 0.");
- this.avlTree = new ObjectAVLTreeSet<>(elementComparator);
- this.elementComparator = elementComparator;
- this.capacity = capacity;
- }
-
- @Override
- public void add(@Nonnull E element) {
- assert !isFull();
- avlTree.add(element);
- }
-
- @Override
- public void remove(@Nonnull E element) {
- avlTree.remove(element);
- }
-
- @Override
- public boolean isFull() {
- return avlTree.size() == capacity;
- }
-
- @Override
- public boolean isEmpty() {
- return avlTree.isEmpty();
- }
-
- @Override
- public boolean isInLowerBound(@Nonnull E toCheck) {
- return avlTree.isEmpty() || elementComparator.compare(peekLast(), toCheck) > 0;
- }
-
- @Nullable
- @Override
- public E removeFirst() {
- if (avlTree.isEmpty()) {
- return null;
- }
- final E first = avlTree.first();
- avlTree.remove(first);
- return first;
- }
-
- @Nullable
- @Override
- public E removeLast() {
- if (avlTree.isEmpty()) {
- return null;
- }
- final E last = avlTree.last();
- avlTree.remove(last);
- return last;
- }
-
- @Nullable
- @Override
- public E peekFirst() {
- return !avlTree.isEmpty() ? avlTree.first() : null;
- }
-
- @Nullable
- @Override
- public E peekLast() {
- return !avlTree.isEmpty() ? avlTree.last() : null;
- }
-
- @Nonnull
- @Override
- public CloseableIterator<E> orderedIterator() {
- return CloseableIterator.adapterForIterator(avlTree.iterator());
- }
-}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/InternalPriorityQueueTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/InternalPriorityQueueTestBase.java
index 0cd551c..510d277 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/InternalPriorityQueueTestBase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/InternalPriorityQueueTestBase.java
@@ -21,13 +21,17 @@ package org.apache.flink.runtime.state;
import org.apache.flink.api.common.typeutils.CompatibilityResult;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.common.typeutils.TypeSerializerConfigSnapshot;
+import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
import org.apache.flink.core.memory.DataInputView;
import org.apache.flink.core.memory.DataOutputView;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
import org.apache.flink.runtime.state.heap.HeapPriorityQueueElement;
import org.apache.flink.util.CloseableIterator;
import org.apache.flink.util.MathUtils;
import org.apache.flink.util.TestLogger;
+import org.apache.flink.shaded.guava18.com.google.common.primitives.UnsignedBytes;
+
import org.junit.Assert;
import org.junit.Test;
@@ -35,6 +39,7 @@ import javax.annotation.Nonnull;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Comparator;
import java.util.HashSet;
import java.util.Iterator;
@@ -53,14 +58,15 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
protected static final KeyExtractorFunction<TestElement> KEY_EXTRACTOR_FUNCTION = TestElement::getKey;
protected static final PriorityComparator<TestElement> TEST_ELEMENT_PRIORITY_COMPARATOR =
(left, right) -> Long.compare(left.getPriority(), right.getPriority());
- protected static final Comparator<TestElement> TEST_ELEMENT_COMPARATOR = (o1, o2) -> {
- int priorityCmp = TEST_ELEMENT_PRIORITY_COMPARATOR.comparePriority(o1, o2);
- if (priorityCmp != 0) {
- return priorityCmp;
- }
- // to fully comply with compareTo/equals contract.
- return Long.compare(o1.getKey(), o2.getKey());
- };
+ protected static final Comparator<TestElement> TEST_ELEMENT_COMPARATOR = new TestElementComparator();
+
+ protected Comparator<Long> getTestElementPriorityComparator() {
+ return Long::compareTo;
+ }
+
+ private long getHighestPriorityValueForComparator() {
+ return getTestElementPriorityComparator().compare(-1L , 1L) > 0 ? Long.MAX_VALUE : Long.MIN_VALUE;
+ }
protected static void insertRandomElements(
@Nonnull InternalPriorityQueue<TestElement> priorityQueue,
@@ -107,13 +113,14 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
public void testPeekPollOrder() {
final int initialCapacity = 4;
final int testSize = 1000;
+ final Comparator<Long> comparator = getTestElementPriorityComparator();
InternalPriorityQueue<TestElement> priorityQueue =
newPriorityQueue(initialCapacity);
HashSet<TestElement> checkSet = new HashSet<>(testSize);
insertRandomElements(priorityQueue, checkSet, testSize);
- long lastPriorityValue = Long.MIN_VALUE;
+ long lastPriorityValue = getHighestPriorityValueForComparator();
int lastSize = priorityQueue.size();
Assert.assertEquals(testSize, lastSize);
TestElement testElement;
@@ -122,7 +129,7 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
Assert.assertEquals(lastSize, priorityQueue.size());
Assert.assertEquals(testElement, priorityQueue.poll());
Assert.assertTrue(checkSet.remove(testElement));
- Assert.assertTrue(testElement.getPriority() >= lastPriorityValue);
+ Assert.assertTrue(comparator.compare(testElement.getPriority(), lastPriorityValue) >= 0);
lastPriorityValue = testElement.getPriority();
--lastSize;
}
@@ -136,7 +143,7 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
public void testRemoveInsertMixKeepsOrder() {
InternalPriorityQueue<TestElement> priorityQueue = newPriorityQueue(3);
-
+ final Comparator<Long> comparator = getTestElementPriorityComparator();
final ThreadLocalRandom random = ThreadLocalRandom.current();
final int testSize = 300;
final int addCounterMax = testSize / 4;
@@ -148,6 +155,8 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
// check that the whole set is still in order
while (!checkSet.isEmpty()) {
+ final long highestPrioValue = getHighestPriorityValueForComparator();
+
Iterator<TestElement> iterator = checkSet.iterator();
TestElement element = iterator.next();
iterator.remove();
@@ -160,16 +169,23 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
priorityQueue.remove(element);
}
- long lastPriorityValue = removesHead ? element.getPriority() : Long.MIN_VALUE;
+ long currentPriorityWatermark;
+
+ // test some bulk polling from time to time
+ if (removesHead) {
+ currentPriorityWatermark = element.getPriority();
+ } else {
+ currentPriorityWatermark = highestPrioValue;
+ }
while ((element = priorityQueue.poll()) != null) {
- Assert.assertTrue(element.getPriority() >= lastPriorityValue);
- lastPriorityValue = element.getPriority();
+ Assert.assertTrue(comparator.compare(element.getPriority(), currentPriorityWatermark) >= 0);
+ currentPriorityWatermark = element.getPriority();
if (--iterationsTillNextAdds == 0) {
// some random adds
iterationsTillNextAdds = random.nextInt(addCounterMax);
insertRandomElements(priorityQueue, new HashSet<>(checkSet), 1 + random.nextInt(3));
- lastPriorityValue = priorityQueue.peek().getPriority();
+ currentPriorityWatermark = priorityQueue.peek().getPriority();
}
}
@@ -182,6 +198,7 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
@Test
public void testPoll() {
InternalPriorityQueue<TestElement> priorityQueue = newPriorityQueue(3);
+ final Comparator<Long> comparator = getTestElementPriorityComparator();
Assert.assertNull(priorityQueue.poll());
@@ -189,12 +206,12 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
HashSet<TestElement> checkSet = new HashSet<>(testSize);
insertRandomElements(priorityQueue, checkSet, testSize);
- long lastPriorityValue = Long.MIN_VALUE;
+ long lastPriorityValue = getHighestPriorityValueForComparator();
while (!priorityQueue.isEmpty()) {
TestElement removed = priorityQueue.poll();
Assert.assertNotNull(removed);
Assert.assertTrue(checkSet.remove(removed));
- Assert.assertTrue(removed.getPriority() >= lastPriorityValue);
+ Assert.assertTrue(comparator.compare(removed.getPriority(), lastPriorityValue) >= 0);
lastPriorityValue = removed.getPriority();
}
Assert.assertTrue(checkSet.isEmpty());
@@ -286,18 +303,21 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
InternalPriorityQueue<TestElement> priorityQueue =
newPriorityQueue(1);
- TestElement lowPrioElement = new TestElement(4711L, 42L);
- TestElement highPrioElement = new TestElement(815L, 23L);
- Assert.assertTrue(priorityQueue.add(lowPrioElement));
+ final List<TestElement> testElements =
+ Arrays.asList(new TestElement(4711L, 42L), new TestElement(815L, 23L));
+
+ testElements.sort((l, r) -> getTestElementPriorityComparator().compare(r.priority, l.priority));
+
+ Assert.assertTrue(priorityQueue.add(testElements.get(0)));
if (testSetSemanticsAgainstDuplicateElements()) {
- priorityQueue.add(lowPrioElement.deepCopy());
+ priorityQueue.add(testElements.get(0).deepCopy());
}
Assert.assertEquals(1, priorityQueue.size());
- Assert.assertTrue(priorityQueue.add(highPrioElement));
+ Assert.assertTrue(priorityQueue.add(testElements.get(1)));
Assert.assertEquals(2, priorityQueue.size());
- Assert.assertEquals(highPrioElement, priorityQueue.poll());
+ Assert.assertEquals(testElements.get(1), priorityQueue.poll());
Assert.assertEquals(1, priorityQueue.size());
- Assert.assertEquals(lowPrioElement, priorityQueue.poll());
+ Assert.assertEquals(testElements.get(0), priorityQueue.poll());
Assert.assertEquals(0, priorityQueue.size());
}
@@ -378,6 +398,14 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
public TestElement deepCopy() {
return new TestElement(key, priority);
}
+
+ @Override
+ public String toString() {
+ return "TestElement{" +
+ "key=" + key +
+ ", priority=" + priority +
+ '}';
+ }
}
/**
@@ -471,4 +499,27 @@ public abstract class InternalPriorityQueueTestBase extends TestLogger {
throw new UnsupportedOperationException();
}
}
+
+ /**
+ * Comparator for test elements, operating on the serialized bytes of the elements.
+ */
+ protected static class TestElementComparator implements Comparator<TestElement> {
+
+ @Override
+ public int compare(TestElement o1, TestElement o2) {
+
+ ByteArrayOutputStreamWithPos os = new ByteArrayOutputStreamWithPos();
+ DataOutputViewStreamWrapper ow = new DataOutputViewStreamWrapper(os);
+ try {
+ TestElementSerializer.INSTANCE.serialize(o1, ow);
+ byte[] a1 = os.toByteArray();
+ os.reset();
+ TestElementSerializer.INSTANCE.serialize(o2, ow);
+ byte[] a2 = os.toByteArray();
+ return UnsignedBytes.lexicographicalComparator().compare(a1, a2);
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+ }
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/CachingInternalPriorityQueueSetTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/CachingInternalPriorityQueueSetTestBase.java
deleted file mode 100644
index 79fd556..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/CachingInternalPriorityQueueSetTestBase.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.heap;
-
-import org.apache.flink.runtime.state.InternalPriorityQueue;
-import org.apache.flink.runtime.state.InternalPriorityQueueTestBase;
-
-/**
- * Test for {@link CachingInternalPriorityQueueSet}.
- */
-public abstract class CachingInternalPriorityQueueSetTestBase extends InternalPriorityQueueTestBase {
-
- @Override
- protected InternalPriorityQueue<TestElement> newPriorityQueue(int initialCapacity) {
- final CachingInternalPriorityQueueSet.OrderedSetCache<TestElement> cache = createOrderedSetCache();
- final CachingInternalPriorityQueueSet.OrderedSetStore<TestElement> store = createOrderedSetStore();
- return new CachingInternalPriorityQueueSet<>(cache, store);
- }
-
- @Override
- protected boolean testSetSemanticsAgainstDuplicateElements() {
- return true;
- }
-
- protected abstract CachingInternalPriorityQueueSet.OrderedSetStore<TestElement> createOrderedSetStore();
-
- protected abstract CachingInternalPriorityQueueSet.OrderedSetCache<TestElement> createOrderedSetCache();
-}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/KeyGroupPartitionedPriorityQueueTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/KeyGroupPartitionedPriorityQueueTest.java
index d348e10..175c615 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/KeyGroupPartitionedPriorityQueueTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/KeyGroupPartitionedPriorityQueueTest.java
@@ -19,6 +19,11 @@ package org.apache.flink.runtime.state.heap;
import org.apache.flink.runtime.state.InternalPriorityQueue;
import org.apache.flink.runtime.state.InternalPriorityQueueTestBase;
+import org.apache.flink.runtime.state.KeyExtractorFunction;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.PriorityComparator;
+
+import javax.annotation.Nonnull;
/**
* Test for {@link KeyGroupPartitionedPriorityQueue}.
@@ -34,20 +39,46 @@ public class KeyGroupPartitionedPriorityQueueTest extends InternalPriorityQueueT
KEY_GROUP_RANGE, KEY_GROUP_RANGE.getNumberOfKeyGroups());
}
- protected KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<
- TestElement, CachingInternalPriorityQueueSet<TestElement>> newFactory(int initialCapacity) {
+ private KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<
+ TestElement, KeyGroupHeapPQSet<TestElement>> newFactory(int initialCapacity) {
- return (keyGroupId, numKeyGroups, elementComparator) -> {
- CachingInternalPriorityQueueSet.OrderedSetCache<TestElement> cache =
- new TreeOrderedSetCache<>(TEST_ELEMENT_COMPARATOR, 32);
- CachingInternalPriorityQueueSet.OrderedSetStore<TestElement> store =
- new TestOrderedStore<>(TEST_ELEMENT_COMPARATOR);
- return new CachingInternalPriorityQueueSet<>(cache, store);
- };
+ return (keyGroupId, numKeyGroups, keyExtractorFunction, elementComparator) ->
+ new KeyGroupHeapPQSet<>(
+ elementComparator,
+ keyExtractorFunction,
+ initialCapacity,
+ KeyGroupRange.of(keyGroupId, keyGroupId),
+ numKeyGroups);
}
@Override
protected boolean testSetSemanticsAgainstDuplicateElements() {
return true;
}
+
+ private static class KeyGroupHeapPQSet<T extends HeapPriorityQueueElement> extends HeapPriorityQueueSet<T>
+ implements HeapPriorityQueueElement {
+
+ private int internalIndex;
+
+ public KeyGroupHeapPQSet(
+ @Nonnull PriorityComparator<T> elementPriorityComparator,
+ @Nonnull KeyExtractorFunction<T> keyExtractor,
+ int minimumCapacity,
+ @Nonnull KeyGroupRange keyGroupRange,
+ int totalNumberOfKeyGroups) {
+ super(elementPriorityComparator, keyExtractor, minimumCapacity, keyGroupRange, totalNumberOfKeyGroups);
+ this.internalIndex = HeapPriorityQueueElement.NOT_CONTAINED;
+ }
+
+ @Override
+ public int getInternalIndex() {
+ return internalIndex;
+ }
+
+ @Override
+ public void setInternalIndex(int newIndex) {
+ internalIndex = newIndex;
+ }
+ }
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/OrderedSetCacheTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/OrderedSetCacheTestBase.java
deleted file mode 100644
index 1d627ed..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/OrderedSetCacheTestBase.java
+++ /dev/null
@@ -1,102 +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.heap;
-
-import org.junit.Assert;
-import org.junit.Test;
-
-import java.util.Random;
-import java.util.TreeSet;
-
-/**
- * Test base for instances of
- * {@link org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet.OrderedSetCache}.
- */
-public abstract class OrderedSetCacheTestBase {
-
- @Test
- public void testOrderedSetCacheContract() {
- final Random random = new Random(0x42);
- final int capacity = 5000;
- final int keySpaceUpperBound = 100 * capacity;
- final TreeSet<Integer> checkSet = new TreeSet<>(Integer::compareTo);
- final CachingInternalPriorityQueueSet.OrderedSetCache<Integer> testInstance = createInstance(capacity);
-
- Assert.assertTrue(testInstance.isEmpty());
-
- while (checkSet.size() < capacity) {
- Assert.assertEquals(checkSet.size() >= capacity, testInstance.isFull());
- if (!checkSet.isEmpty() && random.nextInt(10) == 0) {
- final int toDelete = pickContainedRandomElement(checkSet, random);
- Assert.assertTrue(checkSet.remove(toDelete));
- testInstance.remove(toDelete);
- } else {
- final int randomValue = random.nextInt(keySpaceUpperBound);
- checkSet.add(randomValue);
- testInstance.add(randomValue);
- }
- Assert.assertEquals(checkSet.isEmpty(), testInstance.isEmpty());
-
- Assert.assertEquals(checkSet.first(), testInstance.peekFirst());
- Assert.assertEquals(checkSet.last(), testInstance.peekLast());
-
- Assert.assertFalse(testInstance.isInLowerBound(checkSet.last()));
- Assert.assertTrue(testInstance.isInLowerBound(checkSet.last() - 1));
- }
-
- Assert.assertTrue(testInstance.isFull());
- Assert.assertFalse(testInstance.isInLowerBound(checkSet.last()));
- Assert.assertTrue(testInstance.isInLowerBound(checkSet.last() - 1));
-
- testInstance.remove(pickNotContainedRandomElement(checkSet, random, keySpaceUpperBound));
- Assert.assertTrue(testInstance.isFull());
-
- int containedKey = pickContainedRandomElement(checkSet, random);
-
- Assert.assertTrue(checkSet.remove(containedKey));
- testInstance.remove(containedKey);
-
- Assert.assertFalse(testInstance.isFull());
-
- for (int i = 0; i < capacity; ++i) {
- if (random.nextInt(1) == 0) {
- Assert.assertEquals(checkSet.pollFirst(), testInstance.removeFirst());
- } else {
- Assert.assertEquals(checkSet.pollLast(), testInstance.removeLast());
- }
- }
-
- Assert.assertFalse(testInstance.isFull());
- Assert.assertTrue(testInstance.isEmpty());
- }
-
- private int pickNotContainedRandomElement(TreeSet<Integer> checkSet, Random random, int upperBound) {
- int notContainedKey;
- do {
- notContainedKey = random.nextInt(upperBound);
- } while (checkSet.contains(notContainedKey));
- return notContainedKey;
- }
-
- private int pickContainedRandomElement(TreeSet<Integer> checkSet, Random random) {
- assert !checkSet.isEmpty();
- return checkSet.ceiling(1 + random.nextInt(checkSet.last()));
- }
-
- protected abstract CachingInternalPriorityQueueSet.OrderedSetCache<Integer> createInstance(int capacity);
-}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/SimpleCachingInternalPriorityQueueSetTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/SimpleCachingInternalPriorityQueueSetTest.java
deleted file mode 100644
index 4c4eb46..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/SimpleCachingInternalPriorityQueueSetTest.java
+++ /dev/null
@@ -1,35 +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.heap;
-
-/**
- * Test for {@link CachingInternalPriorityQueueSet}.
- */
-public class SimpleCachingInternalPriorityQueueSetTest extends CachingInternalPriorityQueueSetTestBase {
-
- @Override
- protected CachingInternalPriorityQueueSet.OrderedSetStore<TestElement> createOrderedSetStore() {
- return new TestOrderedStore<>(TEST_ELEMENT_COMPARATOR);
- }
-
- @Override
- protected CachingInternalPriorityQueueSet.OrderedSetCache<TestElement> createOrderedSetCache() {
- return new TreeOrderedSetCache<>(TEST_ELEMENT_COMPARATOR, 3);
- }
-}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/TestOrderedStore.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/TestOrderedStore.java
deleted file mode 100644
index 36a3341..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/TestOrderedStore.java
+++ /dev/null
@@ -1,60 +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.heap;
-
-import org.apache.flink.util.CloseableIterator;
-
-import javax.annotation.Nonnull;
-
-import java.util.Comparator;
-import java.util.TreeSet;
-
-/**
- * Simple implementation of {@link org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet.OrderedSetStore}
- * for tests.
- */
-public class TestOrderedStore<T> implements CachingInternalPriorityQueueSet.OrderedSetStore<T> {
-
- private final TreeSet<T> treeSet;
-
- public TestOrderedStore(Comparator<T> comparator) {
- this.treeSet = new TreeSet<>(comparator);
- }
-
- @Override
- public void add(@Nonnull T element) {
- treeSet.add(element);
- }
-
- @Override
- public void remove(@Nonnull T element) {
- treeSet.remove(element);
- }
-
- @Override
- public int size() {
- return treeSet.size();
- }
-
- @Nonnull
- @Override
- public CloseableIterator<T> orderedIterator() {
- return CloseableIterator.adapterForIterator(treeSet.iterator());
- }
-}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/TreeOrderedSetCacheTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/TreeOrderedSetCacheTest.java
deleted file mode 100644
index cfe823e..0000000
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/state/heap/TreeOrderedSetCacheTest.java
+++ /dev/null
@@ -1,30 +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.heap;
-
-/**
- * Test for {@link TreeOrderedSetCache}.
- */
-public class TreeOrderedSetCacheTest extends OrderedSetCacheTestBase {
-
- @Override
- protected CachingInternalPriorityQueueSet.OrderedSetCache<Integer> createInstance(int capacity) {
- return new TreeOrderedSetCache<>(Integer::compareTo, capacity);
- }
-}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBCachingPriorityQueueSet.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBCachingPriorityQueueSet.java
new file mode 100644
index 0000000..51251c0
--- /dev/null
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBCachingPriorityQueueSet.java
@@ -0,0 +1,560 @@
+/*
+ * 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.typeutils.TypeSerializer;
+import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
+import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
+import org.apache.flink.runtime.state.InternalPriorityQueue;
+import org.apache.flink.runtime.state.heap.HeapPriorityQueueElement;
+import org.apache.flink.util.CloseableIterator;
+import org.apache.flink.util.FlinkRuntimeException;
+
+import org.apache.flink.shaded.guava18.com.google.common.primitives.UnsignedBytes;
+
+import org.rocksdb.ColumnFamilyHandle;
+import org.rocksdb.RocksDB;
+import org.rocksdb.RocksDBException;
+
+import javax.annotation.Nonnegative;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+import java.util.TreeSet;
+
+import static org.apache.flink.contrib.streaming.state.RocksDBCachingPriorityQueueSet.OrderedByteArraySetCache.LEXICOGRAPHIC_BYTE_COMPARATOR;
+
+/**
+ * A priority queue with set semantics, implemented on top of RocksDB. This uses a {@link TreeSet} to cache the bytes
+ * of up to the first n elements from RocksDB in memory to reduce interaction with RocksDB, in particular seek
+ * operations. Cache uses a simple write-through policy.
+ *
+ * @param <E> the type of the contained elements in the queue.
+ */
+public class RocksDBCachingPriorityQueueSet<E extends HeapPriorityQueueElement>
+ implements InternalPriorityQueue<E>, HeapPriorityQueueElement {
+
+ /** Serialized empty value to insert into RocksDB. */
+ private static final byte[] DUMMY_BYTES = new byte[] {};
+
+ /** The RocksDB instance that serves as store. */
+ @Nonnull
+ private final RocksDB db;
+
+ /** Handle to the column family of the RocksDB instance in which the elements are stored. */
+ @Nonnull
+ private final ColumnFamilyHandle columnFamilyHandle;
+
+ /**
+ * Serializer for the contained elements. The lexicographical order of the bytes of serialized objects must be
+ * aligned with their logical order.
+ */
+ @Nonnull
+ private final TypeSerializer<E> byteOrderProducingSerializer;
+
+ /** Wrapper to batch all writes to RocksDB. */
+ @Nonnull
+ private final RocksDBWriteBatchWrapper batchWrapper;
+
+ /** The key-group id in serialized form. */
+ @Nonnull
+ private final byte[] groupPrefixBytes;
+
+ /** Output stream that helps to serialize elements. */
+ @Nonnull
+ private final ByteArrayOutputStreamWithPos outputStream;
+
+ /** Output view that helps to serialize elements, must wrap the output stream. */
+ @Nonnull
+ private final DataOutputViewStreamWrapper outputView;
+
+ @Nonnull
+ private final ByteArrayInputStreamWithPos inputStream;
+
+ @Nonnull
+ private final DataInputViewStreamWrapper inputView;
+
+ /** In memory cache that holds a head-subset of the elements stored in RocksDB. */
+ @Nonnull
+ private final OrderedByteArraySetCache orderedCache;
+
+ /** This holds the key that we use to seek to the first element in RocksDB, to improve seek/iterator performance. */
+ @Nonnull
+ private byte[] seekHint;
+
+ /** Cache for the head element in de-serialized form. */
+ @Nullable
+ private E peekCache;
+
+ /** This flag is true iff all elements in RocksDB are also contained in the cache. */
+ private boolean allElementsInCache;
+
+ /** Index for management as a {@link HeapPriorityQueueElement}. */
+ private int internalIndex;
+
+ RocksDBCachingPriorityQueueSet(
+ @Nonnegative int keyGroupId,
+ @Nonnegative int keyGroupPrefixBytes,
+ @Nonnull RocksDB db,
+ @Nonnull ColumnFamilyHandle columnFamilyHandle,
+ @Nonnull TypeSerializer<E> byteOrderProducingSerializer,
+ @Nonnull ByteArrayOutputStreamWithPos outputStream,
+ @Nonnull ByteArrayInputStreamWithPos inputStream,
+ @Nonnull RocksDBWriteBatchWrapper batchWrapper,
+ @Nonnull OrderedByteArraySetCache orderedByteArraySetCache) {
+ this.db = db;
+ this.columnFamilyHandle = columnFamilyHandle;
+ this.byteOrderProducingSerializer = byteOrderProducingSerializer;
+ this.outputStream = outputStream;
+ this.inputStream = inputStream;
+ this.batchWrapper = batchWrapper;
+ this.allElementsInCache = false;
+ this.outputView = new DataOutputViewStreamWrapper(outputStream);
+ this.inputView = new DataInputViewStreamWrapper(inputStream);
+ this.orderedCache = orderedByteArraySetCache;
+ this.groupPrefixBytes = createKeyGroupBytes(keyGroupId, keyGroupPrefixBytes);
+ this.seekHint = groupPrefixBytes;
+ this.internalIndex = HeapPriorityQueueElement.NOT_CONTAINED;
+ }
+
+ @Nullable
+ @Override
+ public E peek() {
+
+ checkRefillCacheFromStore();
+
+ if (peekCache != null) {
+ return peekCache;
+ }
+
+ byte[] firstBytes = orderedCache.peekFirst();
+ if (firstBytes != null) {
+ peekCache = deserializeElement(firstBytes);
+ return peekCache;
+ } else {
+ return null;
+ }
+ }
+
+ @Nullable
+ @Override
+ public E poll() {
+
+ checkRefillCacheFromStore();
+
+ final byte[] firstBytes = orderedCache.pollFirst();
+
+ if (firstBytes == null) {
+ return null;
+ }
+
+ // write-through sync
+ removeFromRocksDB(firstBytes);
+
+ if (orderedCache.isEmpty()) {
+ seekHint = firstBytes;
+ }
+
+ if (peekCache != null) {
+ E fromCache = peekCache;
+ peekCache = null;
+ return fromCache;
+ } else {
+ return deserializeElement(firstBytes);
+ }
+ }
+
+ @Override
+ public boolean add(@Nonnull E toAdd) {
+
+ checkRefillCacheFromStore();
+
+ final byte[] toAddBytes = serializeElement(toAdd);
+
+ final boolean cacheFull = orderedCache.isFull();
+
+ if ((!cacheFull && allElementsInCache) ||
+ LEXICOGRAPHIC_BYTE_COMPARATOR.compare(toAddBytes, orderedCache.peekLast()) < 0) {
+
+ if (cacheFull) {
+ // we drop the element with lowest priority from the cache
+ orderedCache.pollLast();
+ // the dropped element is now only in the store
+ allElementsInCache = false;
+ }
+
+ if (orderedCache.add(toAddBytes)) {
+ // write-through sync
+ addToRocksDB(toAddBytes);
+ if (toAddBytes == orderedCache.peekFirst()) {
+ peekCache = null;
+ return true;
+ }
+ }
+ } else {
+ // we only added to the store
+ addToRocksDB(toAddBytes);
+ allElementsInCache = false;
+ }
+ return false;
+ }
+
+ @Override
+ public boolean remove(@Nonnull E toRemove) {
+
+ checkRefillCacheFromStore();
+
+ final byte[] oldHead = orderedCache.peekFirst();
+
+ if (oldHead == null) {
+ return false;
+ }
+
+ final byte[] toRemoveBytes = serializeElement(toRemove);
+
+ // write-through sync
+ removeFromRocksDB(toRemoveBytes);
+ orderedCache.remove(toRemoveBytes);
+
+ if (orderedCache.isEmpty()) {
+ seekHint = toRemoveBytes;
+ peekCache = null;
+ return true;
+ }
+
+ if (oldHead != orderedCache.peekFirst()) {
+ peekCache = null;
+ return true;
+ }
+
+ return false;
+ }
+
+ @Override
+ public void addAll(@Nullable Collection<? extends E> toAdd) {
+
+ if (toAdd == null) {
+ return;
+ }
+
+ for (E element : toAdd) {
+ add(element);
+ }
+ }
+
+ @Override
+ public boolean isEmpty() {
+ checkRefillCacheFromStore();
+ return orderedCache.isEmpty();
+ }
+
+ @Nonnull
+ @Override
+ public CloseableIterator<E> iterator() {
+ return new DeserializingIteratorWrapper(orderedBytesIterator());
+ }
+
+ /**
+ * This implementation comes at a relatively high cost per invocation. It should not be called repeatedly when it is
+ * clear that the value did not change. Currently this is only truly used to realize certain higher-level tests.
+ */
+ @Override
+ public int size() {
+
+ if (allElementsInCache) {
+ return orderedCache.size();
+ } else {
+ int count = 0;
+ try (final RocksBytesIterator iterator = orderedBytesIterator()) {
+ while (iterator.hasNext()) {
+ iterator.next();
+ ++count;
+ }
+ }
+ return count;
+ }
+ }
+
+ @Override
+ public int getInternalIndex() {
+ return internalIndex;
+ }
+
+ @Override
+ public void setInternalIndex(int newIndex) {
+ this.internalIndex = newIndex;
+ }
+
+ @Nonnull
+ private RocksBytesIterator orderedBytesIterator() {
+ flushWriteBatch();
+ return new RocksBytesIterator(
+ new RocksIteratorWrapper(
+ db.newIterator(columnFamilyHandle)));
+ }
+
+ /**
+ * Ensures that recent writes are flushed and reflect in the RocksDB instance.
+ */
+ private void flushWriteBatch() {
+ try {
+ batchWrapper.flush();
+ } catch (RocksDBException e) {
+ throw new FlinkRuntimeException(e);
+ }
+ }
+
+ private void addToRocksDB(@Nonnull byte[] toAddBytes) {
+ try {
+ batchWrapper.put(columnFamilyHandle, toAddBytes, DUMMY_BYTES);
+ } catch (RocksDBException e) {
+ throw new FlinkRuntimeException(e);
+ }
+ }
+
+ private void removeFromRocksDB(@Nonnull byte[] toRemoveBytes) {
+ try {
+ batchWrapper.remove(columnFamilyHandle, toRemoveBytes);
+ } catch (RocksDBException e) {
+ throw new FlinkRuntimeException(e);
+ }
+ }
+
+ private void checkRefillCacheFromStore() {
+ if (!allElementsInCache && orderedCache.isEmpty()) {
+ try (final RocksBytesIterator iterator = orderedBytesIterator()) {
+ orderedCache.bulkLoadFromOrderedIterator(iterator);
+ allElementsInCache = !iterator.hasNext();
+ } catch (Exception e) {
+ throw new FlinkRuntimeException("Exception while refilling store from iterator.", e);
+ }
+ }
+ }
+
+ private static boolean isPrefixWith(byte[] bytes, byte[] prefixBytes) {
+ for (int i = 0; i < prefixBytes.length; ++i) {
+ if (bytes[i] != prefixBytes[i]) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ @Nonnull
+ private byte[] createKeyGroupBytes(int keyGroupId, int numPrefixBytes) {
+
+ outputStream.reset();
+
+ try {
+ RocksDBKeySerializationUtils.writeKeyGroup(keyGroupId, numPrefixBytes, outputView);
+ } catch (IOException e) {
+ throw new FlinkRuntimeException("Could not write key-group bytes.", e);
+ }
+
+ return outputStream.toByteArray();
+ }
+
+ @Nonnull
+ private byte[] serializeElement(@Nonnull E element) {
+ try {
+ outputStream.reset();
+ outputView.write(groupPrefixBytes);
+ byteOrderProducingSerializer.serialize(element, outputView);
+ return outputStream.toByteArray();
+ } catch (IOException e) {
+ throw new FlinkRuntimeException("Error while serializing the element.", e);
+ }
+ }
+
+ @Nonnull
+ private E deserializeElement(@Nonnull byte[] bytes) {
+ try {
+ inputStream.setBuffer(bytes, groupPrefixBytes.length, bytes.length);
+ return byteOrderProducingSerializer.deserialize(inputView);
+ } catch (IOException e) {
+ throw new FlinkRuntimeException("Error while deserializing the element.", e);
+ }
+ }
+
+ /**
+ * Wraps an iterator over byte-arrays with deserialization logic, so that it iterates over deserialized objects.
+ */
+ private class DeserializingIteratorWrapper implements CloseableIterator<E> {
+
+ /** The iterator over byte-arrays with the serialized objects. */
+ @Nonnull
+ private final CloseableIterator<byte[]> bytesIterator;
+
+ private DeserializingIteratorWrapper(@Nonnull CloseableIterator<byte[]> bytesIterator) {
+ this.bytesIterator = bytesIterator;
+ }
+
+ @Override
+ public void close() throws Exception {
+ bytesIterator.close();
+ }
+
+ @Override
+ public boolean hasNext() {
+ return bytesIterator.hasNext();
+ }
+
+ @Override
+ public E next() {
+ return deserializeElement(bytesIterator.next());
+ }
+ }
+
+ /**
+ * Adapter between RocksDB iterator and Java iterator. This is also closeable to release the native resources after
+ * use.
+ */
+ private class RocksBytesIterator implements CloseableIterator<byte[]> {
+
+ /** The RocksDb iterator to which we forward ops. */
+ @Nonnull
+ private final RocksIteratorWrapper iterator;
+
+ /** Cache for the current element of the iteration. */
+ @Nullable
+ private byte[] currentElement;
+
+ private RocksBytesIterator(@Nonnull RocksIteratorWrapper iterator) {
+ this.iterator = iterator;
+ try {
+ // We use our knowledge about the lower bound to issue a seek that is as close to the first element in
+ // the key-group as possible, i.e. we generate the next possible key after seekHint by appending one
+ // zero-byte.
+ iterator.seek(Arrays.copyOf(seekHint, seekHint.length + 1));
+ currentElement = nextElementIfAvailable();
+ } catch (Exception ex) {
+ // ensure resource cleanup also in the face of (runtime) exceptions in the constructor.
+ iterator.close();
+ throw new FlinkRuntimeException("Could not initialize ordered iterator.", ex);
+ }
+ }
+
+ @Override
+ public void close() {
+ iterator.close();
+ }
+
+ @Override
+ public boolean hasNext() {
+ return currentElement != null;
+ }
+
+ @Override
+ public byte[] next() {
+ final byte[] returnElement = this.currentElement;
+ if (returnElement == null) {
+ throw new NoSuchElementException("Iterator has no more elements!");
+ }
+ iterator.next();
+ currentElement = nextElementIfAvailable();
+ return returnElement;
+ }
+
+ private byte[] nextElementIfAvailable() {
+ final byte[] elementBytes;
+ return iterator.isValid()
+ && isPrefixWith((elementBytes = iterator.key()), groupPrefixBytes) ? elementBytes : null;
+ }
+ }
+
+ /**
+ * Cache that is organized as an ordered set for byte-arrays. The byte-arrays are sorted in lexicographic order
+ * of their content. Caches typically have a bounded size.
+ */
+ public interface OrderedByteArraySetCache {
+
+ /** Comparator for byte arrays. */
+ Comparator<byte[]> LEXICOGRAPHIC_BYTE_COMPARATOR = UnsignedBytes.lexicographicalComparator();
+
+ /**
+ * Returns the number of contained elements.
+ */
+ int size();
+
+ /**
+ * Returns the maximum number of elements that can be stored in the cache.
+ */
+ int maxSize();
+
+ /**
+ * Returns <code>size() == 0</code>.
+ */
+ boolean isEmpty();
+
+ /**
+ * Returns <code>size() == maxSize()</code>.
+ */
+ boolean isFull();
+
+ /**
+ * Adds the given element, if it was not already contained. Returns <code>true</code> iff the cache was modified.
+ */
+ boolean add(@Nonnull byte[] toAdd);
+
+ /**
+ * Removes the given element, if it is contained. Returns <code>true</code> iff the cache was modified.
+ */
+ boolean remove(@Nonnull byte[] toRemove);
+
+ /**
+ * Returns the first element or <code>null</code> if empty.
+ */
+ @Nullable
+ byte[] peekFirst();
+
+ /**
+ * Returns the last element or <code>null</code> if empty.
+ */
+ @Nullable
+ byte[] peekLast();
+
+ /**
+ * Returns and removes the first element or returns <code>null</code> if empty.
+ */
+ @Nullable
+ byte[] pollFirst();
+
+ /**
+ * Returns and removes the last element or returns <code>null</code> if empty.
+ */
+ @Nullable
+ byte[] pollLast();
+
+ /**
+ * Clears the cache and adds up to <code>maxSize()</code> elements from the iterator to the cache.
+ * Iterator must be ordered in the same order as this cache.
+ *
+ * @param orderedIterator iterator with elements in-order.
+ */
+ void bulkLoadFromOrderedIterator(@Nonnull Iterator<byte[]> orderedIterator);
+ }
+}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
index 72aa4fb..2b4e34c 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
@@ -84,13 +84,10 @@ import org.apache.flink.runtime.state.StateObject;
import org.apache.flink.runtime.state.StateUtil;
import org.apache.flink.runtime.state.StreamCompressionDecorator;
import org.apache.flink.runtime.state.StreamStateHandle;
-import org.apache.flink.runtime.state.TieBreakingPriorityComparator;
import org.apache.flink.runtime.state.UncompressedStreamCompressionDecorator;
-import org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet;
import org.apache.flink.runtime.state.heap.HeapPriorityQueueElement;
import org.apache.flink.runtime.state.heap.HeapPriorityQueueSetFactory;
import org.apache.flink.runtime.state.heap.KeyGroupPartitionedPriorityQueue;
-import org.apache.flink.runtime.state.heap.TreeOrderedSetCache;
import org.apache.flink.runtime.state.metainfo.StateMetaInfoSnapshot;
import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
import org.apache.flink.util.ExceptionUtils;
@@ -262,6 +259,7 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
/** Factory for priority queue state. */
private final PriorityQueueSetFactory priorityQueueFactory;
+ /** Shared wrapper for batch writes to the RocksDB instance. */
private RocksDBWriteBatchWrapper writeBatchWrapper;
public RocksDBKeyedStateBackend(
@@ -458,6 +456,11 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
return keyGroupPrefixBytes;
}
+ @VisibleForTesting
+ PriorityQueueSetFactory getPriorityQueueFactory() {
+ return priorityQueueFactory;
+ }
+
public WriteOptions getWriteOptions() {
return writeOptions;
}
@@ -481,6 +484,9 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
final CheckpointStreamFactory streamFactory,
CheckpointOptions checkpointOptions) throws Exception {
+ // flush everything into db before taking a snapshot
+ writeBatchWrapper.flush();
+
return snapshotStrategy.performSnapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
}
@@ -1775,9 +1781,6 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
final CloseableRegistry snapshotCloseableRegistry = new CloseableRegistry();
- // flush everything into db before taking a snapshot
- writeBatchWrapper.flush();
-
final RocksDBFullSnapshotOperation<K> snapshotOperation =
new RocksDBFullSnapshotOperation<>(
RocksDBKeyedStateBackend.this,
@@ -2643,79 +2646,54 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
class RocksDBPriorityQueueSetFactory implements PriorityQueueSetFactory {
/** Default cache size per key-group. */
- private static final int DEFAULT_CACHES_SIZE = 1024;
+ private static final int DEFAULT_CACHES_SIZE = 128; //TODO make this configurable
/** A shared buffer to serialize elements for the priority queue. */
@Nonnull
private final ByteArrayOutputStreamWithPos elementSerializationOutStream;
- /** A shared adapter wrapper around elementSerializationOutStream to become a {@link DataOutputView}. */
+ /** A shared buffer to de-serialize elements for the priority queue. */
@Nonnull
- private final DataOutputViewStreamWrapper elementSerializationOutView;
+ private final ByteArrayInputStreamWithPos elementSerializationInStream;
RocksDBPriorityQueueSetFactory() {
this.elementSerializationOutStream = new ByteArrayOutputStreamWithPos();
- this.elementSerializationOutView = new DataOutputViewStreamWrapper(elementSerializationOutStream);
+ this.elementSerializationInStream = new ByteArrayInputStreamWithPos();
}
@Nonnull
@Override
public <T extends HeapPriorityQueueElement & PriorityComparable & Keyed> KeyGroupedInternalPriorityQueue<T>
- create(
- @Nonnull String stateName,
- @Nonnull TypeSerializer<T> byteOrderedElementSerializer) {
-
- final PriorityComparator<T> priorityComparator =
- PriorityComparator.forPriorityComparableObjects();
+ create(@Nonnull String stateName, @Nonnull TypeSerializer<T> byteOrderedElementSerializer) {
- Tuple2<ColumnFamilyHandle, RegisteredStateMetaInfoBase> entry =
- kvStateInformation.get(stateName);
-
- if (entry == null) {
- RegisteredPriorityQueueStateBackendMetaInfo<T> metaInfo =
- new RegisteredPriorityQueueStateBackendMetaInfo<>(stateName, byteOrderedElementSerializer);
-
- final ColumnFamilyHandle columnFamilyHandle = createColumnFamily(stateName);
-
- entry = new Tuple2<>(columnFamilyHandle, metaInfo);
- kvStateInformation.put(stateName, entry);
- }
+ final Tuple2<ColumnFamilyHandle, RegisteredStateMetaInfoBase> entry =
+ tryRegisterPriorityQueueMetaInfo(stateName, byteOrderedElementSerializer);
final ColumnFamilyHandle columnFamilyHandle = entry.f0;
- @Nonnull
- TieBreakingPriorityComparator<T> tieBreakingComparator =
- new TieBreakingPriorityComparator<>(
- priorityComparator,
- byteOrderedElementSerializer,
- elementSerializationOutStream,
- elementSerializationOutView);
-
return new KeyGroupPartitionedPriorityQueue<>(
KeyExtractorFunction.forKeyedObjects(),
- priorityComparator,
- new KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<T, CachingInternalPriorityQueueSet<T>>() {
+ PriorityComparator.forPriorityComparableObjects(),
+ new KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<T, RocksDBCachingPriorityQueueSet<T>>() {
@Nonnull
@Override
- public CachingInternalPriorityQueueSet<T> create(
+ public RocksDBCachingPriorityQueueSet<T> create(
int keyGroupId,
int numKeyGroups,
+ @Nonnull KeyExtractorFunction<T> keyExtractor,
@Nonnull PriorityComparator<T> elementPriorityComparator) {
-
- CachingInternalPriorityQueueSet.OrderedSetCache<T> cache =
- new TreeOrderedSetCache<>(tieBreakingComparator, DEFAULT_CACHES_SIZE);
- CachingInternalPriorityQueueSet.OrderedSetStore<T> store =
- new RocksDBOrderedSetStore<>(
- keyGroupId,
- keyGroupPrefixBytes,
- db,
- columnFamilyHandle,
- byteOrderedElementSerializer,
- elementSerializationOutStream,
- elementSerializationOutView,
- writeBatchWrapper);
-
- return new CachingInternalPriorityQueueSet<>(cache, store);
+ TreeOrderedSetCache orderedSetCache = new TreeOrderedSetCache(DEFAULT_CACHES_SIZE);
+ return new RocksDBCachingPriorityQueueSet<>(
+ keyGroupId,
+ keyGroupPrefixBytes,
+ db,
+ columnFamilyHandle,
+ byteOrderedElementSerializer,
+ elementSerializationOutStream,
+ elementSerializationInStream,
+ writeBatchWrapper,
+ orderedSetCache
+ );
}
},
keyGroupRange,
@@ -2723,6 +2701,27 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
}
}
+ @Nonnull
+ private <T> Tuple2<ColumnFamilyHandle, RegisteredStateMetaInfoBase> tryRegisterPriorityQueueMetaInfo(
+ @Nonnull String stateName,
+ @Nonnull TypeSerializer<T> byteOrderedElementSerializer) {
+
+ Tuple2<ColumnFamilyHandle, RegisteredStateMetaInfoBase> entry =
+ kvStateInformation.get(stateName);
+
+ if (entry == null) {
+ RegisteredPriorityQueueStateBackendMetaInfo<T> metaInfo =
+ new RegisteredPriorityQueueStateBackendMetaInfo<>(stateName, byteOrderedElementSerializer);
+
+ final ColumnFamilyHandle columnFamilyHandle = createColumnFamily(stateName);
+
+ entry = new Tuple2<>(columnFamilyHandle, metaInfo);
+ kvStateInformation.put(stateName, entry);
+ }
+
+ return entry;
+ }
+
@Override
public boolean requiresLegacySynchronousTimerSnapshots() {
return priorityQueueFactory instanceof HeapPriorityQueueSetFactory;
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOptions.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOptions.java
index 18f9ec9..c85a7b2 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOptions.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOptions.java
@@ -39,9 +39,10 @@ public class RocksDBOptions {
/**
* Choice of timer service implementation.
*/
- public static final ConfigOption<String> TIMER_SERVICE_IMPL = ConfigOptions
- .key("state.backend.rocksdb.timer-service.impl")
+ public static final ConfigOption<String> TIMER_SERVICE_FACTORY = ConfigOptions
+ .key("state.backend.rocksdb.timer-service.factory")
.defaultValue(HEAP.name())
- .withDescription(String.format("This determines the timer service implementation. Options are either %s " +
- "(heap-based, default) or %s for an implementation based on RocksDB.", HEAP.name(), ROCKSDB.name()));
+ .withDescription(String.format("This determines the factory for timer service state implementation. Options " +
+ "are either %s (heap-based, default) or %s for an implementation based on RocksDB .",
+ HEAP.name(), ROCKSDB.name()));
}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOrderedSetStore.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOrderedSetStore.java
deleted file mode 100644
index 4068c50..0000000
--- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOrderedSetStore.java
+++ /dev/null
@@ -1,274 +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.typeutils.TypeSerializer;
-import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
-import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
-import org.apache.flink.core.memory.DataInputViewStreamWrapper;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
-import org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet;
-import org.apache.flink.util.CloseableIterator;
-import org.apache.flink.util.FlinkRuntimeException;
-
-import org.rocksdb.ColumnFamilyHandle;
-import org.rocksdb.RocksDB;
-import org.rocksdb.RocksDBException;
-
-import javax.annotation.Nonnegative;
-import javax.annotation.Nonnull;
-import javax.annotation.Nullable;
-
-import java.io.IOException;
-import java.util.NoSuchElementException;
-
-/**
- * Implementation of {@link org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet.OrderedSetStore}
- * based on RocksDB.
- *
- * <p>IMPORTANT: The store is ordered and the order is determined by the lexicographic order of the byte sequences
- * produced by the provided serializer for the elements!
- *
- * @param <T> the type of stored elements.
- */
-public class RocksDBOrderedSetStore<T> implements CachingInternalPriorityQueueSet.OrderedSetStore<T> {
-
- /** Serialized empty value to insert into RocksDB. */
- private static final byte[] DUMMY_BYTES = new byte[] {0};
-
- /** The RocksDB instance that serves as store. */
- @Nonnull
- private final RocksDB db;
-
- /** Handle to the column family of the RocksDB instance in which the elements are stored. */
- @Nonnull
- private final ColumnFamilyHandle columnFamilyHandle;
-
- /**
- * Serializer for the contained elements. The lexicographical order of the bytes of serialized objects must be
- * aligned with their logical order.
- */
- @Nonnull
- private final TypeSerializer<T> byteOrderProducingSerializer;
-
- /** Wrapper to batch all writes to RocksDB. */
- @Nonnull
- private final RocksDBWriteBatchWrapper batchWrapper;
-
- /** The key-group id in serialized form. */
- @Nonnull
- private final byte[] groupPrefixBytes;
-
- /** Output stream that helps to serialize elements. */
- @Nonnull
- private final ByteArrayOutputStreamWithPos outputStream;
-
- /** Output view that helps to serialize elements, must wrap the output stream. */
- @Nonnull
- private final DataOutputViewStreamWrapper outputView;
-
- public RocksDBOrderedSetStore(
- @Nonnegative int keyGroupId,
- @Nonnegative int keyGroupPrefixBytes,
- @Nonnull RocksDB db,
- @Nonnull ColumnFamilyHandle columnFamilyHandle,
- @Nonnull TypeSerializer<T> byteOrderProducingSerializer,
- @Nonnull ByteArrayOutputStreamWithPos outputStream,
- @Nonnull DataOutputViewStreamWrapper outputView,
- @Nonnull RocksDBWriteBatchWrapper batchWrapper) {
- this.db = db;
- this.columnFamilyHandle = columnFamilyHandle;
- this.byteOrderProducingSerializer = byteOrderProducingSerializer;
- this.outputStream = outputStream;
- this.outputView = outputView;
- this.batchWrapper = batchWrapper;
- this.groupPrefixBytes = createKeyGroupBytes(keyGroupId, keyGroupPrefixBytes);
- }
-
- private byte[] createKeyGroupBytes(int keyGroupId, int numPrefixBytes) {
-
- outputStream.reset();
-
- try {
- RocksDBKeySerializationUtils.writeKeyGroup(keyGroupId, numPrefixBytes, outputView);
- } catch (IOException e) {
- throw new FlinkRuntimeException("Could not write key-group bytes.", e);
- }
-
- return outputStream.toByteArray();
- }
-
- @Override
- public void add(@Nonnull T element) {
- byte[] elementBytes = serializeElement(element);
- try {
- batchWrapper.put(columnFamilyHandle, elementBytes, DUMMY_BYTES);
- } catch (RocksDBException e) {
- throw new FlinkRuntimeException("Error while getting element from RocksDB.", e);
- }
- }
-
- @Override
- public void remove(@Nonnull T element) {
- byte[] elementBytes = serializeElement(element);
- try {
- batchWrapper.remove(columnFamilyHandle, elementBytes);
- } catch (RocksDBException e) {
- throw new FlinkRuntimeException("Error while removing element from RocksDB.", e);
- }
- }
-
- /**
- * This implementation comes at a relatively high cost per invocation. It should not be called repeatedly when it is
- * clear that the value did not change. Currently this is only truly used to realize certain higher-level tests.
- *
- * @see org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet.OrderedSetStore
- */
- @Override
- public int size() {
-
- int count = 0;
- try (final RocksToJavaIteratorAdapter iterator = orderedIterator()) {
- while (iterator.hasNext()) {
- iterator.next();
- ++count;
- }
- }
-
- return count;
- }
-
- @Nonnull
- @Override
- public RocksToJavaIteratorAdapter orderedIterator() {
-
- flushWriteBatch();
- return new RocksToJavaIteratorAdapter(
- new RocksIteratorWrapper(
- db.newIterator(columnFamilyHandle)));
- }
-
- /**
- * Ensures that recent writes are flushed and reflect in the RocksDB instance.
- */
- private void flushWriteBatch() {
- try {
- batchWrapper.flush();
- } catch (RocksDBException e) {
- throw new FlinkRuntimeException(e);
- }
- }
-
- private static boolean isPrefixWith(byte[] bytes, byte[] prefixBytes) {
- for (int i = 0; i < prefixBytes.length; ++i) {
- if (bytes[i] != prefixBytes[i]) {
- return false;
- }
- }
- return true;
- }
-
- private byte[] serializeElement(T element) {
- try {
- outputStream.reset();
- outputView.write(groupPrefixBytes);
- byteOrderProducingSerializer.serialize(element, outputView);
- return outputStream.toByteArray();
- } catch (IOException e) {
- throw new FlinkRuntimeException("Error while serializing the element.", e);
- }
- }
-
- private T deserializeElement(byte[] bytes) {
- try {
- // TODO introduce a stream in which we can change the internal byte[] to avoid creating instances per call
- ByteArrayInputStreamWithPos inputStream = new ByteArrayInputStreamWithPos(bytes);
- DataInputViewStreamWrapper inputView = new DataInputViewStreamWrapper(inputStream);
- inputView.skipBytes(groupPrefixBytes.length);
- return byteOrderProducingSerializer.deserialize(inputView);
- } catch (IOException e) {
- throw new FlinkRuntimeException("Error while deserializing the element.", e);
- }
- }
-
- /**
- * Adapter between RocksDB iterator and Java iterator. This is also closeable to release the native resources after
- * use.
- */
- private class RocksToJavaIteratorAdapter implements CloseableIterator<T> {
-
- /** The RocksDb iterator to which we forward ops. */
- @Nonnull
- private final RocksIteratorWrapper iterator;
-
- /** Cache for the current element of the iteration. */
- @Nullable
- private T currentElement;
-
- private RocksToJavaIteratorAdapter(@Nonnull RocksIteratorWrapper iterator) {
- this.iterator = iterator;
- try {
- // TODO we could check if it is more efficient to make the seek more specific, e.g. with a provided hint
- // that is lexicographically closer the first expected element in the key-group. I wonder if this could
- // help to improve the seek if there are many tombstones for elements at the beginning of the key-group
- // (like for elements that have been removed in previous polling, before they are compacted away).
- iterator.seek(groupPrefixBytes);
- deserializeNextElementIfAvailable();
- } catch (Exception ex) {
- // ensure resource cleanup also in the face of (runtime) exceptions in the constructor.
- iterator.close();
- throw new FlinkRuntimeException("Could not initialize ordered iterator.", ex);
- }
- }
-
- @Override
- public void close() {
- iterator.close();
- }
-
- @Override
- public boolean hasNext() {
- return currentElement != null;
- }
-
- @Override
- public T next() {
- final T returnElement = this.currentElement;
- if (returnElement == null) {
- throw new NoSuchElementException("Iterator has no more elements!");
- }
- iterator.next();
- deserializeNextElementIfAvailable();
- return returnElement;
- }
-
- private void deserializeNextElementIfAvailable() {
- if (iterator.isValid()) {
- final byte[] elementBytes = iterator.key();
- if (isPrefixWith(elementBytes, groupPrefixBytes)) {
- this.currentElement = deserializeElement(elementBytes);
- } else {
- this.currentElement = null;
- }
- } else {
- this.currentElement = null;
- }
- }
- }
-}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
index 0564849..b8bd73c 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
@@ -60,7 +60,7 @@ import java.util.List;
import java.util.Random;
import java.util.UUID;
-import static org.apache.flink.contrib.streaming.state.RocksDBOptions.TIMER_SERVICE_IMPL;
+import static org.apache.flink.contrib.streaming.state.RocksDBOptions.TIMER_SERVICE_FACTORY;
import static org.apache.flink.util.Preconditions.checkNotNull;
/**
@@ -271,7 +271,7 @@ public class RocksDBStateBackend extends AbstractStateBackend implements Configu
this.enableIncrementalCheckpointing = original.enableIncrementalCheckpointing.resolveUndefined(
config.getBoolean(CheckpointingOptions.INCREMENTAL_CHECKPOINTS));
- final String priorityQueueTypeString = config.getString(TIMER_SERVICE_IMPL);
+ final String priorityQueueTypeString = config.getString(TIMER_SERVICE_FACTORY);
this.priorityQueueStateType = priorityQueueTypeString.length() > 0 ?
PriorityQueueStateType.valueOf(priorityQueueTypeString.toUpperCase()) : original.priorityQueueStateType;
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/TreeOrderedSetCache.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/TreeOrderedSetCache.java
new file mode 100644
index 0000000..c48cc74
--- /dev/null
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/TreeOrderedSetCache.java
@@ -0,0 +1,104 @@
+/*
+ * 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 javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+import java.util.Iterator;
+import java.util.TreeSet;
+
+/**
+ * Implementation of a lexicographically ordered set of byte-arrays, based on a {@link TreeSet}.
+ */
+public class TreeOrderedSetCache implements RocksDBCachingPriorityQueueSet.OrderedByteArraySetCache {
+
+ /** Maximum capacity. */
+ private final int maxSize;
+
+ @Nonnull
+ private final TreeSet<byte[]> treeSet;
+
+ TreeOrderedSetCache(int maxSize) {
+ this.maxSize = maxSize;
+ this.treeSet = new TreeSet<>(LEXICOGRAPHIC_BYTE_COMPARATOR);
+ }
+
+ @Override
+ public int size() {
+ return treeSet.size();
+ }
+
+ @Override
+ public int maxSize() {
+ return maxSize;
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return treeSet.isEmpty();
+ }
+
+ @Override
+ public boolean isFull() {
+ return treeSet.size() >= maxSize;
+ }
+
+ @Override
+ public boolean add(@Nonnull byte[] toAdd) {
+ return treeSet.add(toAdd);
+ }
+
+ @Override
+ public boolean remove(@Nonnull byte[] toRemove) {
+ return treeSet.remove(toRemove);
+ }
+
+ @Nullable
+ @Override
+ public byte[] peekFirst() {
+ return !isEmpty() ? treeSet.first() : null;
+ }
+
+ @Nullable
+ @Override
+ public byte[] peekLast() {
+ return !isEmpty() ? treeSet.last() : null;
+ }
+
+ @Nullable
+ @Override
+ public byte[] pollFirst() {
+ return !isEmpty() ? treeSet.pollFirst() : null;
+ }
+
+ @Nullable
+ @Override
+ public byte[] pollLast() {
+ return !isEmpty() ? treeSet.pollLast() : null;
+ }
+
+ @Override
+ public void bulkLoadFromOrderedIterator(@Nonnull Iterator<byte[]> orderedIterator) {
+ treeSet.clear();
+ for (int i = maxSize; --i >= 0 && orderedIterator.hasNext(); ) {
+ treeSet.add(orderedIterator.next());
+ }
+ }
+}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/CachingInternalPriorityQueueSetWithRocksDBStoreTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/CachingInternalPriorityQueueSetWithRocksDBStoreTest.java
deleted file mode 100644
index 5f26835..0000000
--- a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/CachingInternalPriorityQueueSetWithRocksDBStoreTest.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.flink.core.memory.ByteArrayOutputStreamWithPos;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
-import org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet;
-import org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSetTestBase;
-import org.apache.flink.runtime.state.heap.TreeOrderedSetCache;
-
-import org.junit.Rule;
-
-/**
- * Test for {@link CachingInternalPriorityQueueSet} with a
- * {@link org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet.OrderedSetStore} based on RocksDB.
- */
-public class CachingInternalPriorityQueueSetWithRocksDBStoreTest extends CachingInternalPriorityQueueSetTestBase {
-
- @Rule
- public final RocksDBResource rocksDBResource = new RocksDBResource();
-
- @Override
- protected CachingInternalPriorityQueueSet.OrderedSetStore<TestElement> createOrderedSetStore() {
- return createRocksDBStore(0, 1, rocksDBResource);
- }
-
- @Override
- protected CachingInternalPriorityQueueSet.OrderedSetCache<TestElement> createOrderedSetCache() {
- return new TreeOrderedSetCache<>(TEST_ELEMENT_COMPARATOR, 32);
- }
-
- public static CachingInternalPriorityQueueSet.OrderedSetStore<TestElement> createRocksDBStore(
- int keyGroupId,
- int totalKeyGroups,
- RocksDBResource rocksDBResource) {
- ByteArrayOutputStreamWithPos outputStream = new ByteArrayOutputStreamWithPos(16);
- DataOutputViewStreamWrapper outputView = new DataOutputViewStreamWrapper(outputStream);
- int prefixBytes = RocksDBKeySerializationUtils.computeRequiredBytesInKeyGroupPrefix(totalKeyGroups);
- return new RocksDBOrderedSetStore<>(
- keyGroupId,
- prefixBytes,
- rocksDBResource.getRocksDB(),
- rocksDBResource.getDefaultColumnFamily(),
- TestElementSerializer.INSTANCE,
- outputStream,
- outputView,
- rocksDBResource.getBatchWrapper());
- }
-}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/KeyGroupPartitionedPriorityQueueWithRocksDBStoreTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/KeyGroupPartitionedPriorityQueueWithRocksDBStoreTest.java
index 42782a4..3cdfdbe 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/KeyGroupPartitionedPriorityQueueWithRocksDBStoreTest.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/KeyGroupPartitionedPriorityQueueWithRocksDBStoreTest.java
@@ -18,36 +18,54 @@
package org.apache.flink.contrib.streaming.state;
-import org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet;
+import org.apache.flink.core.memory.ByteArrayInputStreamWithPos;
+import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
+import org.apache.flink.runtime.state.InternalPriorityQueue;
+import org.apache.flink.runtime.state.InternalPriorityQueueTestBase;
import org.apache.flink.runtime.state.heap.KeyGroupPartitionedPriorityQueue;
-import org.apache.flink.runtime.state.heap.KeyGroupPartitionedPriorityQueueTest;
-import org.apache.flink.runtime.state.heap.TreeOrderedSetCache;
import org.junit.Rule;
/**
- * Test of {@link KeyGroupPartitionedPriorityQueue} powered by a {@link RocksDBOrderedSetStore}.
+ * Test of {@link KeyGroupPartitionedPriorityQueue} powered by a {@link RocksDBCachingPriorityQueueSet}.
*/
-public class KeyGroupPartitionedPriorityQueueWithRocksDBStoreTest extends KeyGroupPartitionedPriorityQueueTest {
+public class KeyGroupPartitionedPriorityQueueWithRocksDBStoreTest extends InternalPriorityQueueTestBase {
@Rule
public final RocksDBResource rocksDBResource = new RocksDBResource();
@Override
- protected KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<
- TestElement, CachingInternalPriorityQueueSet<TestElement>> newFactory(
- int initialCapacity) {
-
- return (keyGroupId, numKeyGroups, elementComparator) -> {
- CachingInternalPriorityQueueSet.OrderedSetCache<TestElement> cache =
- new TreeOrderedSetCache<>(TEST_ELEMENT_COMPARATOR, 32);
- CachingInternalPriorityQueueSet.OrderedSetStore<TestElement> store =
- RocksDBOrderedSetStoreTest.createRocksDBOrderedStore(
- rocksDBResource,
- TestElementSerializer.INSTANCE,
- keyGroupId,
- numKeyGroups);
- return new CachingInternalPriorityQueueSet<>(cache, store);
+ protected InternalPriorityQueue<TestElement> newPriorityQueue(int initialCapacity) {
+ return new KeyGroupPartitionedPriorityQueue<>(
+ KEY_EXTRACTOR_FUNCTION,
+ TEST_ELEMENT_PRIORITY_COMPARATOR,
+ newFactory(),
+ KEY_GROUP_RANGE, KEY_GROUP_RANGE.getNumberOfKeyGroups());
+ }
+
+ @Override
+ protected boolean testSetSemanticsAgainstDuplicateElements() {
+ return true;
+ }
+
+ private KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<
+ TestElement, RocksDBCachingPriorityQueueSet<TestElement>> newFactory() {
+
+ return (keyGroupId, numKeyGroups, keyExtractorFunction, elementComparator) -> {
+ ByteArrayOutputStreamWithPos outputStreamWithPos = new ByteArrayOutputStreamWithPos(32);
+ ByteArrayInputStreamWithPos inputStreamWithPos = new ByteArrayInputStreamWithPos();
+ int keyGroupPrefixBytes = RocksDBKeySerializationUtils.computeRequiredBytesInKeyGroupPrefix(numKeyGroups);
+ TreeOrderedSetCache orderedSetCache = new TreeOrderedSetCache(32);
+ return new RocksDBCachingPriorityQueueSet<>(
+ keyGroupId,
+ keyGroupPrefixBytes,
+ rocksDBResource.getRocksDB(),
+ rocksDBResource.getDefaultColumnFamily(),
+ TestElementSerializer.INSTANCE,
+ outputStreamWithPos,
+ inputStreamWithPos,
+ rocksDBResource.getBatchWrapper(),
+ orderedSetCache);
};
}
}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
index a8e5a36..e344638 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java
@@ -256,7 +256,7 @@ public class RocksDBAsyncSnapshotTest extends TestLogger {
File dbDir = temporaryFolder.newFolder();
- final RocksDBStateBackend.PriorityQueueStateType timerServicePriorityQueueType = RocksDBStateBackend.PriorityQueueStateType.valueOf(RocksDBOptions.TIMER_SERVICE_IMPL.defaultValue());
+ final RocksDBStateBackend.PriorityQueueStateType timerServicePriorityQueueType = RocksDBStateBackend.PriorityQueueStateType.valueOf(RocksDBOptions.TIMER_SERVICE_FACTORY.defaultValue());
final int skipStreams;
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBOrderedSetStoreTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBOrderedSetStoreTest.java
deleted file mode 100644
index 0b1d07b..0000000
--- a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBOrderedSetStoreTest.java
+++ /dev/null
@@ -1,136 +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.typeutils.TypeSerializer;
-import org.apache.flink.api.common.typeutils.base.IntSerializer;
-import org.apache.flink.core.memory.ByteArrayOutputStreamWithPos;
-import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
-import org.apache.flink.runtime.state.heap.CachingInternalPriorityQueueSet;
-import org.apache.flink.util.CloseableIterator;
-
-import org.junit.Assert;
-import org.junit.Rule;
-import org.junit.Test;
-
-import javax.annotation.Nonnegative;
-import javax.annotation.Nonnull;
-
-import java.util.NoSuchElementException;
-
-/**
- * Test for RocksDBOrderedStore.
- */
-public class RocksDBOrderedSetStoreTest {
-
- @Rule
- public final RocksDBResource rocksDBResource = new RocksDBResource();
-
- @Test
- public void testOrderedIterator() throws Exception {
- CachingInternalPriorityQueueSet.OrderedSetStore<Integer> store = createRocksDBOrderedStore();
-
- //test empty iterator
- try (final CloseableIterator<Integer> emptyIterator = store.orderedIterator()) {
- Assert.assertFalse(emptyIterator.hasNext());
- try {
- emptyIterator.next();
- Assert.fail();
- } catch (NoSuchElementException expected) {
- }
- }
-
- store.add(43);
- store.add(42);
- store.add(41);
- store.add(41);
- store.remove(42);
-
- // test in-order iteration
- try (final CloseableIterator<Integer> iterator = store.orderedIterator()) {
- Assert.assertTrue(iterator.hasNext());
- Assert.assertEquals(Integer.valueOf(41), iterator.next());
- Assert.assertTrue(iterator.hasNext());
- Assert.assertEquals(Integer.valueOf(43), iterator.next());
- Assert.assertFalse(iterator.hasNext());
- try {
- iterator.next();
- Assert.fail();
- } catch (NoSuchElementException expected) {
- }
- }
- }
-
- @Test
- public void testAddRemoveSize() {
-
- CachingInternalPriorityQueueSet.OrderedSetStore<Integer> store = createRocksDBOrderedStore();
-
- // test empty size
- Assert.assertEquals(0, store.size());
-
- // test add uniques
- store.remove(41);
- Assert.assertEquals(0, store.size());
- store.add(41);
- Assert.assertEquals(1, store.size());
- store.add(42);
- Assert.assertEquals(2, store.size());
- store.add(43);
- Assert.assertEquals(3, store.size());
- store.add(44);
- Assert.assertEquals(4, store.size());
- store.add(45);
- Assert.assertEquals(5, store.size());
-
- // test remove
- store.remove(41);
- Assert.assertEquals(4, store.size());
- store.remove(41);
- Assert.assertEquals(4, store.size());
-
- // test set semantics by attempt to insert duplicate
- store.add(42);
- Assert.assertEquals(4, store.size());
- }
-
- public static <E> RocksDBOrderedSetStore<E> createRocksDBOrderedStore(
- @Nonnull RocksDBResource rocksDBResource,
- @Nonnull TypeSerializer<E> byteOrderSerializer,
- @Nonnegative int keyGroupId,
- @Nonnegative int totalKeyGroups) {
-
- ByteArrayOutputStreamWithPos outputStreamWithPos = new ByteArrayOutputStreamWithPos(32);
- DataOutputViewStreamWrapper outputView = new DataOutputViewStreamWrapper(outputStreamWithPos);
- int keyGroupPrefixBytes = RocksDBKeySerializationUtils.computeRequiredBytesInKeyGroupPrefix(totalKeyGroups);
- return new RocksDBOrderedSetStore<>(
- keyGroupId,
- keyGroupPrefixBytes,
- rocksDBResource.getRocksDB(),
- rocksDBResource.getDefaultColumnFamily(),
- byteOrderSerializer,
- outputStreamWithPos,
- outputView,
- rocksDBResource.getBatchWrapper());
- }
-
- protected RocksDBOrderedSetStore<Integer> createRocksDBOrderedStore() {
- return createRocksDBOrderedStore(rocksDBResource, IntSerializer.INSTANCE, 0, 1);
- }
-}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
index d6c3144..b7d05b3 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
@@ -35,11 +35,13 @@ import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.StateBackend;
import org.apache.flink.runtime.state.TestTaskStateManager;
import org.apache.flink.runtime.state.filesystem.FsStateBackend;
+import org.apache.flink.runtime.state.heap.HeapPriorityQueueSetFactory;
import org.apache.flink.runtime.state.memory.MemoryStateBackend;
import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
import org.apache.flink.runtime.util.TestingTaskManagerRuntimeInfo;
import org.apache.flink.util.IOUtils;
+import org.junit.Assert;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
@@ -48,6 +50,7 @@ import org.rocksdb.CompactionStyle;
import org.rocksdb.DBOptions;
import java.io.File;
+import java.util.Collections;
import static org.hamcrest.CoreMatchers.anyOf;
import static org.hamcrest.MatcherAssert.assertThat;
@@ -126,6 +129,45 @@ public class RocksDBStateBackendConfigTest {
}
@Test
+ public void testConfigureTimerService() throws Exception {
+
+ final Environment env = getMockEnvironment(tempFolder.newFolder());
+
+ // Fix the option key string
+ Assert.assertEquals("state.backend.rocksdb.timer-service.factory", RocksDBOptions.TIMER_SERVICE_FACTORY.key());
+
+ // Fix the option value string and ensure all are covered
+ Assert.assertEquals(2, RocksDBStateBackend.PriorityQueueStateType.values().length);
+ Assert.assertEquals("ROCKSDB", RocksDBStateBackend.PriorityQueueStateType.ROCKSDB.toString());
+ Assert.assertEquals("HEAP", RocksDBStateBackend.PriorityQueueStateType.HEAP.toString());
+
+ // Fix the default
+ Assert.assertEquals(
+ RocksDBStateBackend.PriorityQueueStateType.HEAP.toString(),
+ RocksDBOptions.TIMER_SERVICE_FACTORY.defaultValue());
+
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());
+
+ RocksDBKeyedStateBackend<Integer> keyedBackend = createKeyedStateBackend(rocksDbBackend, env);
+ keyedBackend.restore(Collections.emptyList());
+ Assert.assertEquals(HeapPriorityQueueSetFactory.class, keyedBackend.getPriorityQueueFactory().getClass());
+ keyedBackend.dispose();
+
+ Configuration conf = new Configuration();
+ conf.setString(
+ RocksDBOptions.TIMER_SERVICE_FACTORY,
+ RocksDBStateBackend.PriorityQueueStateType.ROCKSDB.toString());
+
+ rocksDbBackend = rocksDbBackend.configure(conf);
+ keyedBackend = createKeyedStateBackend(rocksDbBackend, env);
+ keyedBackend.restore(Collections.emptyList());
+ Assert.assertEquals(
+ RocksDBKeyedStateBackend.RocksDBPriorityQueueSetFactory.class,
+ keyedBackend.getPriorityQueueFactory().getClass());
+ keyedBackend.dispose();
+ }
+
+ @Test
public void testStoragePathWithFilePrefix() throws Exception {
final File folder = tempFolder.newFolder();
final String dbStoragePath = new Path(folder.toURI().toString()).toString();
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimeServiceManager.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimeServiceManager.java
index 29b89c2..b54a1a9 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimeServiceManager.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimeServiceManager.java
@@ -37,7 +37,7 @@ import java.util.Map;
/**
* An entity keeping all the time-related services available to all operators extending the
* {@link AbstractStreamOperator}. Right now, this is only a
- * {@link HeapInternalTimerService timer services}.
+ * {@link InternalTimerServiceImpl timer services}.
*
* <b>NOTE:</b> These services are only available to keyed operators.
*
@@ -58,7 +58,7 @@ public class InternalTimeServiceManager<K> {
private final PriorityQueueSetFactory priorityQueueSetFactory;
private final ProcessingTimeService processingTimeService;
- private final Map<String, HeapInternalTimerService<K, ?>> timerServices;
+ private final Map<String, InternalTimerServiceImpl<K, ?>> timerServices;
private final boolean useLegacySynchronousSnapshots;
@@ -86,7 +86,7 @@ public class InternalTimeServiceManager<K> {
TimerSerializer<K, N> timerSerializer,
Triggerable<K, N> triggerable) {
- HeapInternalTimerService<K, N> timerService = registerOrGetTimerService(name, timerSerializer);
+ InternalTimerServiceImpl<K, N> timerService = registerOrGetTimerService(name, timerSerializer);
timerService.startTimerService(
timerSerializer.getKeySerializer(),
@@ -97,11 +97,11 @@ public class InternalTimeServiceManager<K> {
}
@SuppressWarnings("unchecked")
- <N> HeapInternalTimerService<K, N> registerOrGetTimerService(String name, TimerSerializer<K, N> timerSerializer) {
- HeapInternalTimerService<K, N> timerService = (HeapInternalTimerService<K, N>) timerServices.get(name);
+ <N> InternalTimerServiceImpl<K, N> registerOrGetTimerService(String name, TimerSerializer<K, N> timerSerializer) {
+ InternalTimerServiceImpl<K, N> timerService = (InternalTimerServiceImpl<K, N>) timerServices.get(name);
if (timerService == null) {
- timerService = new HeapInternalTimerService<>(
+ timerService = new InternalTimerServiceImpl<>(
localKeyGroupRange,
keyContext,
processingTimeService,
@@ -113,7 +113,7 @@ public class InternalTimeServiceManager<K> {
return timerService;
}
- Map<String, HeapInternalTimerService<K, ?>> getRegisteredTimerServices() {
+ Map<String, InternalTimerServiceImpl<K, ?>> getRegisteredTimerServices() {
return Collections.unmodifiableMap(timerServices);
}
@@ -126,7 +126,7 @@ public class InternalTimeServiceManager<K> {
}
public void advanceWatermark(Watermark watermark) throws Exception {
- for (HeapInternalTimerService<?, ?> service : timerServices.values()) {
+ for (InternalTimerServiceImpl<?, ?> service : timerServices.values()) {
service.advanceWatermark(watermark.getTimestamp());
}
}
@@ -164,7 +164,7 @@ public class InternalTimeServiceManager<K> {
@VisibleForTesting
public int numProcessingTimeTimers() {
int count = 0;
- for (HeapInternalTimerService<?, ?> timerService : timerServices.values()) {
+ for (InternalTimerServiceImpl<?, ?> timerService : timerServices.values()) {
count += timerService.numProcessingTimeTimers();
}
return count;
@@ -173,7 +173,7 @@ public class InternalTimeServiceManager<K> {
@VisibleForTesting
public int numEventTimeTimers() {
int count = 0;
- for (HeapInternalTimerService<?, ?> timerService : timerServices.values()) {
+ for (InternalTimerServiceImpl<?, ?> timerService : timerServices.values()) {
count += timerService.numEventTimeTimers();
}
return count;
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/HeapInternalTimerService.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimerServiceImpl.java
similarity index 92%
rename from flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/HeapInternalTimerService.java
rename to flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimerServiceImpl.java
index 6c1b188..3bded503 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/HeapInternalTimerService.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimerServiceImpl.java
@@ -43,7 +43,7 @@ import static org.apache.flink.util.Preconditions.checkNotNull;
/**
* {@link InternalTimerService} that stores timers on the Java heap.
*/
-public class HeapInternalTimerService<K, N> implements InternalTimerService<N>, ProcessingTimeCallback {
+public class InternalTimerServiceImpl<K, N> implements InternalTimerService<N>, ProcessingTimeCallback {
private final ProcessingTimeService processingTimeService;
@@ -95,7 +95,7 @@ public class HeapInternalTimerService<K, N> implements InternalTimerService<N>,
/** The restored timers snapshot, if any. */
private InternalTimersSnapshot<K, N> restoredTimersSnapshot;
- HeapInternalTimerService(
+ InternalTimerServiceImpl(
KeyGroupRange localKeyGroupRange,
KeyContext keyContext,
ProcessingTimeService processingTimeService,
@@ -117,7 +117,7 @@ public class HeapInternalTimerService<K, N> implements InternalTimerService<N>,
}
/**
- * Starts the local {@link HeapInternalTimerService} by:
+ * Starts the local {@link InternalTimerServiceImpl} by:
* <ol>
* <li>Setting the {@code keySerialized} and {@code namespaceSerializer} for the timers it will contain.</li>
* <li>Setting the {@code triggerTarget} which contains the action to be performed when a timer fires.</li>
@@ -227,37 +227,29 @@ public class HeapInternalTimerService<K, N> implements InternalTimerService<N>,
// inside the callback.
nextTimer = null;
- processingTimeTimersQueue.bulkPoll(
- (timer) -> (timer.getTimestamp() <= time),
- (timer) -> {
- keyContext.setCurrentKey(timer.getKey());
- try {
- triggerTarget.onProcessingTime(timer);
- } catch (Exception e) {
- throw new FlinkRuntimeException("Problem in trigger target.", e);
- }
- });
+ InternalTimer<K, N> timer;
- if (nextTimer == null) {
- final TimerHeapInternalTimer<K, N> timer = processingTimeTimersQueue.peek();
- if (timer != null) {
- nextTimer = processingTimeService.registerTimer(timer.getTimestamp(), this);
- }
+ while ((timer = processingTimeTimersQueue.peek()) != null && timer.getTimestamp() <= time) {
+ processingTimeTimersQueue.poll();
+ keyContext.setCurrentKey(timer.getKey());
+ triggerTarget.onProcessingTime(timer);
+ }
+
+ if (timer != null && nextTimer == null) {
+ nextTimer = processingTimeService.registerTimer(timer.getTimestamp(), this);
}
}
public void advanceWatermark(long time) throws Exception {
currentWatermark = time;
- eventTimeTimersQueue.bulkPoll(
- (timer) -> (timer.getTimestamp() <= time),
- (timer) -> {
- keyContext.setCurrentKey(timer.getKey());
- try {
- triggerTarget.onEventTime(timer);
- } catch (Exception e) {
- throw new FlinkRuntimeException("Problem in trigger target.", e);
- }
- });
+
+ InternalTimer<K, N> timer;
+
+ while ((timer = eventTimeTimersQueue.peek()) != null && timer.getTimestamp() <= time) {
+ eventTimeTimersQueue.poll();
+ keyContext.setCurrentKey(timer.getKey());
+ triggerTarget.onEventTime(timer);
+ }
}
/**
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimerServiceSerializationProxy.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimerServiceSerializationProxy.java
index ce490b5..dea17f9 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimerServiceSerializationProxy.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/InternalTimerServiceSerializationProxy.java
@@ -77,13 +77,13 @@ public class InternalTimerServiceSerializationProxy<K> extends PostVersionedIORe
@Override
public void write(DataOutputView out) throws IOException {
super.write(out);
- final Map<String, HeapInternalTimerService<K, ?>> registeredTimerServices =
+ final Map<String, InternalTimerServiceImpl<K, ?>> registeredTimerServices =
timerServicesManager.getRegisteredTimerServices();
out.writeInt(registeredTimerServices.size());
- for (Map.Entry<String, HeapInternalTimerService<K, ?>> entry : registeredTimerServices.entrySet()) {
+ for (Map.Entry<String, InternalTimerServiceImpl<K, ?>> entry : registeredTimerServices.entrySet()) {
String serviceName = entry.getKey();
- HeapInternalTimerService<K, ?> timerService = entry.getValue();
+ InternalTimerServiceImpl<K, ?> timerService = entry.getValue();
out.writeUTF(serviceName);
InternalTimersSnapshotReaderWriters
@@ -104,7 +104,7 @@ public class InternalTimerServiceSerializationProxy<K> extends PostVersionedIORe
.getReaderForVersion(readerVersion, userCodeClassLoader)
.readTimersSnapshot(in);
- HeapInternalTimerService<K, ?> timerService = registerOrGetTimerService(
+ InternalTimerServiceImpl<K, ?> timerService = registerOrGetTimerService(
serviceName,
restoredTimersSnapshot);
@@ -113,7 +113,7 @@ public class InternalTimerServiceSerializationProxy<K> extends PostVersionedIORe
}
@SuppressWarnings("unchecked")
- private <N> HeapInternalTimerService<K, N> registerOrGetTimerService(
+ private <N> InternalTimerServiceImpl<K, N> registerOrGetTimerService(
String serviceName, InternalTimersSnapshot<?, ?> restoredTimersSnapshot) {
final TypeSerializer<K> keySerializer = (TypeSerializer<K>) restoredTimersSnapshot.getKeySerializer();
final TypeSerializer<N> namespaceSerializer = (TypeSerializer<N>) restoredTimersSnapshot.getNamespaceSerializer();
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/HeapInternalTimerServiceTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/InternalTimerServiceImplTest.java
similarity index 95%
rename from flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/HeapInternalTimerServiceTest.java
rename to flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/InternalTimerServiceImplTest.java
index 957a535..f2da6da 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/HeapInternalTimerServiceTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/InternalTimerServiceImplTest.java
@@ -61,10 +61,10 @@ import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
/**
- * Tests for {@link HeapInternalTimerService}.
+ * Tests for {@link InternalTimerServiceImpl}.
*/
@RunWith(Parameterized.class)
-public class HeapInternalTimerServiceTest {
+public class InternalTimerServiceImplTest {
private final int maxParallelism;
private final KeyGroupRange testKeyGroupRange;
@@ -73,7 +73,7 @@ public class HeapInternalTimerServiceTest {
return any();
}
- public HeapInternalTimerServiceTest(int startKeyGroup, int endKeyGroup, int maxParallelism) {
+ public InternalTimerServiceImplTest(int startKeyGroup, int endKeyGroup, int maxParallelism) {
this.testKeyGroupRange = new KeyGroupRange(startKeyGroup, endKeyGroup);
this.maxParallelism = maxParallelism;
}
@@ -89,7 +89,7 @@ public class HeapInternalTimerServiceTest {
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> service = createInternalTimerService(
+ InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
testKeyGroupList,
keyContext,
processingTimeService,
@@ -117,7 +117,7 @@ public class HeapInternalTimerServiceTest {
final PriorityQueueSetFactory priorityQueueSetFactory =
createQueueFactory(keyGroupRange, totalNoOfKeyGroups);
- HeapInternalTimerService<Integer, String> timerService = createInternalTimerService(
+ InternalTimerServiceImpl<Integer, String> timerService = createInternalTimerService(
keyGroupRange,
keyContext,
new TestProcessingTimeService(),
@@ -182,7 +182,7 @@ public class HeapInternalTimerServiceTest {
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
PriorityQueueSetFactory priorityQueueSetFactory =
new HeapPriorityQueueSetFactory(testKeyGroupRange, maxParallelism, 128);
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, priorityQueueSetFactory);
int key = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
@@ -244,7 +244,7 @@ public class HeapInternalTimerServiceTest {
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());
int key = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
@@ -277,7 +277,7 @@ public class HeapInternalTimerServiceTest {
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- final HeapInternalTimerService<Integer, String> timerService =
+ final InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());
int key = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
@@ -328,7 +328,7 @@ public class HeapInternalTimerServiceTest {
TestKeyContext keyContext = new TestKeyContext();
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());
processingTimeService.setCurrentTime(17L);
@@ -346,7 +346,7 @@ public class HeapInternalTimerServiceTest {
TestKeyContext keyContext = new TestKeyContext();
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());
timerService.advanceWatermark(17);
@@ -366,7 +366,7 @@ public class HeapInternalTimerServiceTest {
TestKeyContext keyContext = new TestKeyContext();
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());
// get two different keys
@@ -411,7 +411,7 @@ public class HeapInternalTimerServiceTest {
TestKeyContext keyContext = new TestKeyContext();
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());
// get two different keys
@@ -458,7 +458,7 @@ public class HeapInternalTimerServiceTest {
TestKeyContext keyContext = new TestKeyContext();
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());
// get two different keys
@@ -515,7 +515,7 @@ public class HeapInternalTimerServiceTest {
TestKeyContext keyContext = new TestKeyContext();
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());
// get two different keys
@@ -590,7 +590,7 @@ public class HeapInternalTimerServiceTest {
TestKeyContext keyContext = new TestKeyContext();
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());
// get two different keys
@@ -665,7 +665,7 @@ public class HeapInternalTimerServiceTest {
TestKeyContext keyContext = new TestKeyContext();
TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
final PriorityQueueSetFactory queueFactory = createQueueFactory();
- HeapInternalTimerService<Integer, String> timerService =
+ InternalTimerServiceImpl<Integer, String> timerService =
createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, queueFactory);
int midpoint = testKeyGroupRange.getStartKeyGroup() +
@@ -730,7 +730,7 @@ public class HeapInternalTimerServiceTest {
TestProcessingTimeService processingTimeService1 = new TestProcessingTimeService();
TestProcessingTimeService processingTimeService2 = new TestProcessingTimeService();
- HeapInternalTimerService<Integer, String> timerService1 = restoreTimerService(
+ InternalTimerServiceImpl<Integer, String> timerService1 = restoreTimerService(
snapshot1,
snapshotVersion,
mockTriggerable1,
@@ -739,7 +739,7 @@ public class HeapInternalTimerServiceTest {
subKeyGroupRange1,
queueFactory);
- HeapInternalTimerService<Integer, String> timerService2 = restoreTimerService(
+ InternalTimerServiceImpl<Integer, String> timerService2 = restoreTimerService(
snapshot2,
snapshotVersion,
mockTriggerable2,
@@ -806,13 +806,13 @@ public class HeapInternalTimerServiceTest {
return result;
}
- private static HeapInternalTimerService<Integer, String> createAndStartInternalTimerService(
+ private static InternalTimerServiceImpl<Integer, String> createAndStartInternalTimerService(
Triggerable<Integer, String> triggerable,
KeyContext keyContext,
ProcessingTimeService processingTimeService,
KeyGroupRange keyGroupList,
PriorityQueueSetFactory priorityQueueSetFactory) {
- HeapInternalTimerService<Integer, String> service = createInternalTimerService(
+ InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
keyGroupList,
keyContext,
processingTimeService,
@@ -824,7 +824,7 @@ public class HeapInternalTimerServiceTest {
return service;
}
- private static HeapInternalTimerService<Integer, String> restoreTimerService(
+ private static InternalTimerServiceImpl<Integer, String> restoreTimerService(
Map<Integer, byte[]> state,
int snapshotVersion,
Triggerable<Integer, String> triggerable,
@@ -834,7 +834,7 @@ public class HeapInternalTimerServiceTest {
PriorityQueueSetFactory priorityQueueSetFactory) throws Exception {
// create an empty service
- HeapInternalTimerService<Integer, String> service = createInternalTimerService(
+ InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
keyGroupsList,
keyContext,
processingTimeService,
@@ -848,7 +848,7 @@ public class HeapInternalTimerServiceTest {
try (ByteArrayInputStream inputStream = new ByteArrayInputStream(state.get(keyGroupIndex))) {
InternalTimersSnapshot<?, ?> restoredTimersSnapshot =
InternalTimersSnapshotReaderWriters
- .getReaderForVersion(snapshotVersion, HeapInternalTimerServiceTest.class.getClassLoader())
+ .getReaderForVersion(snapshotVersion, InternalTimerServiceImplTest.class.getClassLoader())
.readTimersSnapshot(new DataInputViewStreamWrapper(inputStream));
service.restoreTimersForKeyGroup(restoredTimersSnapshot, keyGroupIndex);
@@ -886,7 +886,7 @@ public class HeapInternalTimerServiceTest {
});
}
- private static <K, N> HeapInternalTimerService<K, N> createInternalTimerService(
+ private static <K, N> InternalTimerServiceImpl<K, N> createInternalTimerService(
KeyGroupRange keyGroupsList,
KeyContext keyContext,
ProcessingTimeService processingTimeService,
@@ -896,7 +896,7 @@ public class HeapInternalTimerServiceTest {
TimerSerializer<K, N> timerSerializer = new TimerSerializer<>(keySerializer, namespaceSerializer);
- return new HeapInternalTimerService<>(
+ return new InternalTimerServiceImpl<>(
keyGroupsList,
keyContext,
processingTimeService,
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 e9a2e45..4fa9020 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
@@ -29,6 +29,7 @@ import org.apache.flink.configuration.AkkaOptions;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.HighAvailabilityOptions;
import org.apache.flink.configuration.TaskManagerOptions;
+import org.apache.flink.contrib.streaming.state.RocksDBOptions;
import org.apache.flink.contrib.streaming.state.RocksDBStateBackend;
import org.apache.flink.core.fs.Path;
import org.apache.flink.runtime.state.AbstractStateBackend;
@@ -167,26 +168,18 @@ public class EventTimeWindowCheckpointingITCase extends TestLogger {
break;
}
case ROCKSDB_FULLY_ASYNC: {
- String rocksDb = tempFolder.newFolder().getAbsolutePath();
- String backups = tempFolder.newFolder().getAbsolutePath();
- RocksDBStateBackend rdb = new RocksDBStateBackend(new FsStateBackend("file://" + backups));
- rdb.setDbStoragePath(rocksDb);
- this.stateBackend = rdb;
+ setupRocksDB(-1, false);
break;
}
case ROCKSDB_INCREMENTAL:
+ // Test RocksDB based timer service as well
+ config.setString(
+ RocksDBOptions.TIMER_SERVICE_FACTORY,
+ RocksDBStateBackend.PriorityQueueStateType.ROCKSDB.toString());
+ setupRocksDB(16, true);
+ break;
case ROCKSDB_INCREMENTAL_ZK: {
- String rocksDb = tempFolder.newFolder().getAbsolutePath();
- String backups = tempFolder.newFolder().getAbsolutePath();
- // we use the fs backend with small threshold here to test the behaviour with file
- // references, not self contained byte handles
- RocksDBStateBackend rdb =
- new RocksDBStateBackend(
- new FsStateBackend(
- new Path("file://" + backups).toUri(), 16),
- true);
- rdb.setDbStoragePath(rocksDb);
- this.stateBackend = rdb;
+ setupRocksDB(16, true);
break;
}
default:
@@ -195,6 +188,20 @@ public class EventTimeWindowCheckpointingITCase extends TestLogger {
return config;
}
+ private void setupRocksDB(int fileSizeThreshold, boolean incrementalCheckpoints) throws IOException {
+ String rocksDb = tempFolder.newFolder().getAbsolutePath();
+ String backups = tempFolder.newFolder().getAbsolutePath();
+ // we use the fs backend with small threshold here to test the behaviour with file
+ // references, not self contained byte handles
+ RocksDBStateBackend rdb =
+ new RocksDBStateBackend(
+ new FsStateBackend(
+ new Path("file://" + backups).toUri(), fileSizeThreshold),
+ incrementalCheckpoints);
+ rdb.setDbStoragePath(rocksDb);
+ this.stateBackend = rdb;
+ }
+
protected Configuration createClusterConfig() throws IOException {
TemporaryFolder temporaryFolder = new TemporaryFolder();
temporaryFolder.create();