You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by ro...@apache.org on 2023/02/13 10:02:55 UTC
[flink] branch master updated: [FLINK-28440][state/changelog] record reference count of StreamStateHandle in TaskChangelogRegistry
This is an automated email from the ASF dual-hosted git repository.
roman pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new c0aa73df4df [FLINK-28440][state/changelog] record reference count of StreamStateHandle in TaskChangelogRegistry
c0aa73df4df is described below
commit c0aa73df4df4e39c138f2cddaeb8efad6c831d03
Author: wangfeifan <zo...@163.com>
AuthorDate: Wed Feb 1 00:04:50 2023 +0800
[FLINK-28440][state/changelog] record reference count of StreamStateHandle in TaskChangelogRegistry
Co-authored-by: Yanfei Lei <18...@users.noreply.github.com>
---
.../fs/AbstractStateChangeFsUploader.java | 19 +--
.../flink/changelog/fs/FsStateChangelogWriter.java | 8 +-
.../flink/changelog/fs/TaskChangelogRegistry.java | 28 ++---
.../changelog/fs/TaskChangelogRegistryImpl.java | 46 ++++---
.../fs/DiscardRecordableStateChangeUploader.java | 73 +++++++++++
.../changelog/fs/FsStateChangelogWriterTest.java | 140 +++++++++++++++++++++
.../fs/TaskChangelogRegistryImplTest.java | 22 ++--
.../fs/TestingBatchingUploadScheduler.java | 52 ++++++++
.../state/changelog/ChangelogStateDiscardTest.java | 39 +++---
9 files changed, 352 insertions(+), 75 deletions(-)
diff --git a/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/AbstractStateChangeFsUploader.java b/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/AbstractStateChangeFsUploader.java
index 3dd3dfcc70f..55bdfeac179 100644
--- a/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/AbstractStateChangeFsUploader.java
+++ b/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/AbstractStateChangeFsUploader.java
@@ -33,7 +33,6 @@ import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.function.BiFunction;
-import java.util.stream.Collectors;
/** Base implementation of StateChangeUploader. */
public abstract class AbstractStateChangeFsUploader implements StateChangeUploader {
@@ -80,22 +79,16 @@ public abstract class AbstractStateChangeFsUploader implements StateChangeUpload
for (UploadTask task : tasks) {
tasksOffsets.put(task, format.write(stream, task.changeSets));
}
+
+ long numOfChangeSets = tasks.stream().flatMap(t -> t.getChangeSets().stream()).count();
+
StreamStateHandle handle = stream.getHandle(handleFactory);
- changelogRegistry.startTracking(
- handle,
- tasks.stream()
- .flatMap(t -> t.getChangeSets().stream())
- .map(StateChangeSet::getLogId)
- .collect(Collectors.toSet()));
+ changelogRegistry.startTracking(handle, numOfChangeSets);
+
if (stream instanceof DuplicatingOutputStreamWithPos) {
StreamStateHandle localHandle =
((DuplicatingOutputStreamWithPos) stream).getSecondaryHandle(handleFactory);
- changelogRegistry.startTracking(
- localHandle,
- tasks.stream()
- .flatMap(t -> t.getChangeSets().stream())
- .map(StateChangeSet::getLogId)
- .collect(Collectors.toSet()));
+ changelogRegistry.startTracking(localHandle, numOfChangeSets);
return new UploadTasksResult(tasksOffsets, handle, localHandle);
}
// WARN: streams have to be closed before returning the results
diff --git a/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/FsStateChangelogWriter.java b/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/FsStateChangelogWriter.java
index f72c50a9d88..83a4b2d5861 100644
--- a/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/FsStateChangelogWriter.java
+++ b/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/FsStateChangelogWriter.java
@@ -288,9 +288,9 @@ class FsStateChangelogWriter implements StateChangelogWriter<ChangelogStateHandl
} else {
// uploaded already truncated, i.e. materialized state changes,
// or closed
- changelogRegistry.notUsed(result.streamStateHandle, logId);
+ changelogRegistry.release(result.streamStateHandle);
if (result.localStreamHandle != null) {
- changelogRegistry.notUsed(result.localStreamHandle, logId);
+ changelogRegistry.release(result.localStreamHandle);
}
}
}
@@ -517,9 +517,9 @@ class FsStateChangelogWriter implements StateChangelogWriter<ChangelogStateHandl
private void notifyStateNotUsed(Map<SequenceNumber, UploadResult> notUsedState) {
LOG.trace("Uploaded state to discard: {}", notUsedState);
for (UploadResult result : notUsedState.values()) {
- changelogRegistry.notUsed(result.streamStateHandle, logId);
+ changelogRegistry.release(result.streamStateHandle);
if (result.localStreamHandle != null) {
- changelogRegistry.notUsed(result.localStreamHandle, logId);
+ changelogRegistry.release(result.localStreamHandle);
}
}
}
diff --git a/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/TaskChangelogRegistry.java b/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/TaskChangelogRegistry.java
index ff089fecb82..e93c925b6fb 100644
--- a/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/TaskChangelogRegistry.java
+++ b/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/TaskChangelogRegistry.java
@@ -21,8 +21,6 @@ import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.runtime.state.StreamStateHandle;
-import java.util.Set;
-import java.util.UUID;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
@@ -32,43 +30,45 @@ import java.util.concurrent.Executors;
* org.apache.flink.runtime.state.changelog.StateChangelogWriter StateChangelogWriters} of a {@link
* org.apache.flink.runtime.state.changelog.StateChangelogStorage StateChangelogStorage}.
*
- * <p>Initially, when {@link #startTracking(StreamStateHandle, Set) starting the tracking}, the
+ * <p>Initially, when {@link #startTracking(StreamStateHandle, long) starting the tracking}, the
* ownership of a changelog segments is not clear, and it is assumed that JM <strong>might</strong>
- * be the owner. Once the backends are not using the segments, JM can not become an owner anymore.
- * the state is discarded.
+ * be the owner. The refCount of the StateObject refers to the number of changelog segments contains
+ * in the StateObject. {@link #release(StreamStateHandle)} should be called when every changelog
+ * segment become not used, and it will count down the refCount by one. Once the refCount reaches
+ * zero, JM can not become an owner anymore, the state is discarded.
*
* <p>However, if at any point it becomes known that JM is the owner, tracking is {@link
* #stopTracking(StreamStateHandle) stopped} and the state will not be discarded.
*
- * <p>It is the client responsibility to make sure that JM can not become an owner when calling
- * {@link #notUsed(StreamStateHandle, UUID)}.
+ * <p>It is the client responsibility to call {@link #release(StreamStateHandle)} when every
+ * corresponding changelog segment becomes not used.
*/
@Internal
public interface TaskChangelogRegistry {
- /** Start tracking the state uploaded for the given backends. */
- void startTracking(StreamStateHandle handle, Set<UUID> backendIDs);
+ /** Start tracking the state uploaded. The refCount is the number of StateChangeSets. */
+ void startTracking(StreamStateHandle handle, long refCount);
/** Stop tracking the state, so that it's not tracked (some other component is doing that). */
void stopTracking(StreamStateHandle handle);
/**
- * Mark the state as unused by the given backend, e.g. if it was pre-emptively uploaded and
- * materialized. Once no backend is using the state, it is discarded (unless it was {@link
+ * Decrease the reference count of the state by one, e.g. if it was pre-emptively uploaded and
+ * materialized. Once the reference count reaches zero, it is discarded (unless it was {@link
* #stopTracking(StreamStateHandle) unregistered} earlier).
*/
- void notUsed(StreamStateHandle handle, UUID backendId);
+ void release(StreamStateHandle handle);
TaskChangelogRegistry NO_OP =
new TaskChangelogRegistry() {
@Override
- public void startTracking(StreamStateHandle handle, Set<UUID> backendIDs) {}
+ public void startTracking(StreamStateHandle handle, long refCount) {}
@Override
public void stopTracking(StreamStateHandle handle) {}
@Override
- public void notUsed(StreamStateHandle handle, UUID backendId) {}
+ public void release(StreamStateHandle handle) {}
};
static TaskChangelogRegistry defaultChangelogRegistry(int numAsyncDiscardThreads) {
diff --git a/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/TaskChangelogRegistryImpl.java b/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/TaskChangelogRegistryImpl.java
index 55376148363..79e164626e0 100644
--- a/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/TaskChangelogRegistryImpl.java
+++ b/flink-dstl/flink-dstl-dfs/src/main/java/org/apache/flink/changelog/fs/TaskChangelogRegistryImpl.java
@@ -20,6 +20,7 @@ package org.apache.flink.changelog.fs;
import org.apache.flink.annotation.Internal;
import org.apache.flink.runtime.state.PhysicalStateHandleID;
import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.util.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -27,10 +28,7 @@ import org.slf4j.LoggerFactory;
import javax.annotation.concurrent.ThreadSafe;
import java.util.Map;
-import java.util.Set;
-import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.Executor;
@Internal
@@ -38,7 +36,7 @@ import java.util.concurrent.Executor;
class TaskChangelogRegistryImpl implements TaskChangelogRegistry {
private static final Logger LOG = LoggerFactory.getLogger(TaskChangelogRegistryImpl.class);
- private final Map<PhysicalStateHandleID, Set<UUID>> entries = new ConcurrentHashMap<>();
+ private final Map<PhysicalStateHandleID, Long> entries = new ConcurrentHashMap<>();
private final Executor executor;
public TaskChangelogRegistryImpl(Executor executor) {
@@ -46,12 +44,13 @@ class TaskChangelogRegistryImpl implements TaskChangelogRegistry {
}
@Override
- public void startTracking(StreamStateHandle handle, Set<UUID> backendIDs) {
+ public void startTracking(StreamStateHandle handle, long refCount) {
+ Preconditions.checkState(refCount > 0, "Initial refCount of state must larger than zero");
LOG.debug(
"start tracking state, key: {}, state: {}",
handle.getStreamStateHandleID(),
handle);
- entries.put(handle.getStreamStateHandleID(), new CopyOnWriteArraySet<>(backendIDs));
+ entries.put(handle.getStreamStateHandleID(), refCount);
}
@Override
@@ -62,19 +61,30 @@ class TaskChangelogRegistryImpl implements TaskChangelogRegistry {
}
@Override
- public void notUsed(StreamStateHandle handle, UUID backendId) {
+ public void release(StreamStateHandle handle) {
PhysicalStateHandleID key = handle.getStreamStateHandleID();
- LOG.debug("backend {} not using state, key: {}, state: {}", backendId, key, handle);
- Set<UUID> backends = entries.get(key);
- if (backends == null) {
- LOG.warn("backend {} was not using state, key: {}, state: {}", backendId, key, handle);
- return;
- }
- backends.remove(backendId);
- if (backends.isEmpty() && entries.remove(key) != null) {
- LOG.debug("state is not used by any backend, schedule discard: {}/{}", key, handle);
- scheduleDiscard(handle);
- }
+ LOG.debug("state reference count decreased by one, key: {}, state: {}", key, handle);
+
+ entries.compute(
+ key,
+ (handleID, refCount) -> {
+ if (refCount == null) {
+ LOG.warn("state is not in tracking, key: {}, state: {}", key, handle);
+ return null;
+ }
+
+ long newRefCount = refCount - 1;
+ if (newRefCount == 0) {
+ LOG.debug(
+ "state is not used by any backend, schedule discard: {}/{}",
+ key,
+ handle);
+ scheduleDiscard(handle);
+ return null;
+ } else {
+ return newRefCount;
+ }
+ });
}
private void scheduleDiscard(StreamStateHandle handle) {
diff --git a/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/DiscardRecordableStateChangeUploader.java b/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/DiscardRecordableStateChangeUploader.java
new file mode 100644
index 00000000000..90199f304f5
--- /dev/null
+++ b/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/DiscardRecordableStateChangeUploader.java
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.changelog.fs;
+
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.changelog.fs.StateChangeUploadScheduler.UploadTask;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.runtime.state.TestingStreamStateHandle;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * {@link StateChangeUploader} implementation which can verify that the returned StreamStateHandle
+ * was deleted.
+ */
+public class DiscardRecordableStateChangeUploader implements StateChangeUploader {
+ private final TaskChangelogRegistry changelogRegistry;
+
+ public DiscardRecordableStateChangeUploader(TaskChangelogRegistry changelogRegistry) {
+ this.changelogRegistry = changelogRegistry;
+ }
+
+ @Override
+ public UploadTasksResult upload(Collection<UploadTask> tasks) throws IOException {
+ Map<UploadTask, Map<StateChangeSet, Tuple2<Long, Long>>> tasksOffsets = new HashMap<>();
+
+ for (UploadTask task : tasks) {
+ Map<StateChangeSet, Tuple2<Long, Long>> offsets = new HashMap<>();
+ for (StateChangeSet changeSet : task.getChangeSets()) {
+ offsets.put(changeSet, Tuple2.of(0L, 0L)); // fake offsets
+ }
+ tasksOffsets.put(task, offsets);
+ }
+
+ long numOfChangeSets = tasks.stream().flatMap(t -> t.getChangeSets().stream()).count();
+
+ // fake StreamStateHandle without data, just for discarding records
+ StreamStateHandle handle = new TestingStreamStateHandle();
+ changelogRegistry.startTracking(handle, numOfChangeSets);
+ return new UploadTasksResult(tasksOffsets, handle);
+ }
+
+ public boolean isDiscarded(StreamStateHandle handle) {
+ if (handle instanceof TestingStreamStateHandle) {
+ return ((TestingStreamStateHandle) handle).isDisposed();
+ } else {
+ throw new IllegalStateException(
+ "only accept StreamStateHandle created by DiscardRecordableStateChangeUploader");
+ }
+ }
+
+ @Override
+ public void close() throws Exception {}
+}
diff --git a/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/FsStateChangelogWriterTest.java b/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/FsStateChangelogWriterTest.java
index d2a3aa61f79..176da94efb3 100644
--- a/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/FsStateChangelogWriterTest.java
+++ b/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/FsStateChangelogWriterTest.java
@@ -17,13 +17,16 @@
package org.apache.flink.changelog.fs;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.runtime.mailbox.SyncMailboxExecutor;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.state.SnapshotResult;
+import org.apache.flink.runtime.state.StreamStateHandle;
import org.apache.flink.runtime.state.TestLocalRecoveryConfig;
import org.apache.flink.runtime.state.changelog.ChangelogStateHandleStreamImpl;
import org.apache.flink.runtime.state.changelog.LocalChangelogRegistry;
import org.apache.flink.runtime.state.changelog.SequenceNumber;
+import org.apache.flink.util.concurrent.Executors;
import org.apache.flink.util.function.BiConsumerWithException;
import org.junit.jupiter.api.Test;
@@ -106,6 +109,143 @@ class FsStateChangelogWriterTest {
});
}
+ @Test
+ void testFileAvailableAfterPreUpload() throws Exception {
+ long appendPersistThreshold = 100;
+
+ TaskChangelogRegistry taskChangelogRegistry =
+ new TaskChangelogRegistryImpl(Executors.directExecutor());
+
+ try (DiscardRecordableStateChangeUploader uploader =
+ new DiscardRecordableStateChangeUploader(taskChangelogRegistry);
+ TestingBatchingUploadScheduler uploadScheduler =
+ new TestingBatchingUploadScheduler(uploader);
+ FsStateChangelogWriter writer =
+ new FsStateChangelogWriter(
+ UUID.randomUUID(),
+ KeyGroupRange.of(KEY_GROUP, KEY_GROUP),
+ uploadScheduler,
+ appendPersistThreshold,
+ new SyncMailboxExecutor(),
+ taskChangelogRegistry,
+ TestLocalRecoveryConfig.disabled(),
+ LocalChangelogRegistry.NO_OP)) {
+ SequenceNumber initialSqn = writer.initialSequenceNumber();
+
+ writer.append(KEY_GROUP, getBytes(10)); // sqn: 0
+
+ // checkpoint 1 trigger
+ SequenceNumber checkpoint1sqn = writer.nextSequenceNumber();
+ writer.persist(initialSqn);
+ uploadScheduler.scheduleAll(); // checkpoint 1 completed
+ writer.confirm(initialSqn, checkpoint1sqn, 1);
+
+ writer.append(KEY_GROUP, getBytes(10)); // sqn: 1
+
+ // materialization 1 trigger
+ SequenceNumber materializationSqn = writer.nextSequenceNumber();
+
+ writer.append(KEY_GROUP, getBytes(10)); // sqn: 2
+
+ // materialization 1 completed
+ // checkpoint 2 trigger
+ SequenceNumber checkpoint2sqn = writer.nextSequenceNumber();
+ writer.persist(materializationSqn);
+ uploadScheduler.scheduleAll(); // checkpoint 2 completed
+ writer.confirm(materializationSqn, checkpoint2sqn, 2);
+
+ // checkpoint 1 subsumed
+ writer.truncate(
+ materializationSqn.compareTo(checkpoint1sqn) < 0
+ ? materializationSqn
+ : checkpoint1sqn);
+
+ writer.append(KEY_GROUP, getBytes(10)); // sqn: 3
+
+ // checkpoint 3 trigger
+ SequenceNumber checkpoint3sqn = writer.nextSequenceNumber();
+ writer.persist(materializationSqn);
+ uploadScheduler.scheduleAll(); // checkpoint 3 completed
+ writer.confirm(materializationSqn, checkpoint3sqn, 3);
+
+ // trigger pre-emptive upload
+ writer.append(KEY_GROUP, getBytes(100)); // sqn: 4
+ uploadScheduler.scheduleAll();
+
+ // checkpoint 2 subsumed
+ writer.truncate(
+ materializationSqn.compareTo(checkpoint2sqn) < 0
+ ? materializationSqn
+ : checkpoint2sqn);
+
+ // checkpoint 4 trigger
+ SequenceNumber checkpoint4sqn = writer.nextSequenceNumber();
+ CompletableFuture<SnapshotResult<ChangelogStateHandleStreamImpl>> future =
+ writer.persist(materializationSqn);
+ uploadScheduler.scheduleAll(); // checkpoint 4 completed
+ writer.confirm(materializationSqn, checkpoint4sqn, 4);
+
+ SnapshotResult<ChangelogStateHandleStreamImpl> result = future.get();
+ ChangelogStateHandleStreamImpl resultHandle = result.getJobManagerOwnedSnapshot();
+
+ for (Tuple2<StreamStateHandle, Long> handleAndOffset :
+ resultHandle.getHandlesAndOffsets()) {
+ assertThat(uploader.isDiscarded(handleAndOffset.f0))
+ .isFalse(); // all handles should not be discarded
+ }
+ }
+ }
+
+ @Test
+ void testFileAvailableAfterClose() throws Exception {
+ long appendPersistThreshold = 100;
+
+ TaskChangelogRegistry taskChangelogRegistry =
+ new TaskChangelogRegistryImpl(Executors.directExecutor());
+
+ try (DiscardRecordableStateChangeUploader uploader =
+ new DiscardRecordableStateChangeUploader(taskChangelogRegistry);
+ TestingBatchingUploadScheduler uploadScheduler =
+ new TestingBatchingUploadScheduler(uploader)) {
+ FsStateChangelogWriter writer =
+ new FsStateChangelogWriter(
+ UUID.randomUUID(),
+ KeyGroupRange.of(KEY_GROUP, KEY_GROUP),
+ uploadScheduler,
+ appendPersistThreshold,
+ new SyncMailboxExecutor(),
+ taskChangelogRegistry,
+ TestLocalRecoveryConfig.disabled(),
+ LocalChangelogRegistry.NO_OP);
+
+ SequenceNumber initialSqn = writer.initialSequenceNumber();
+
+ writer.append(KEY_GROUP, getBytes(10)); // sqn: 0
+
+ // checkpoint 1 trigger
+ SequenceNumber checkpoint1sqn = writer.nextSequenceNumber();
+ CompletableFuture<SnapshotResult<ChangelogStateHandleStreamImpl>> future =
+ writer.persist(initialSqn);
+
+ // trigger pre-emptive upload
+ writer.append(KEY_GROUP, getBytes(100)); // sqn: 1
+
+ uploadScheduler.scheduleAll(); // checkpoint 1 completed
+
+ // close task before confirm checkpoint 1
+ writer.truncateAndClose(checkpoint1sqn);
+
+ SnapshotResult<ChangelogStateHandleStreamImpl> result = future.get();
+ ChangelogStateHandleStreamImpl resultHandle = result.getJobManagerOwnedSnapshot();
+
+ for (Tuple2<StreamStateHandle, Long> handleAndOffset :
+ resultHandle.getHandlesAndOffsets()) {
+ assertThat(uploader.isDiscarded(handleAndOffset.f0))
+ .isFalse(); // all handles should not be discarded
+ }
+ }
+ }
+
@Test
void testNoReUploadBeforeCompletion() throws Exception {
withWriter(
diff --git a/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/TaskChangelogRegistryImplTest.java b/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/TaskChangelogRegistryImplTest.java
index e5dc73ad7be..0d64cd654b4 100644
--- a/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/TaskChangelogRegistryImplTest.java
+++ b/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/TaskChangelogRegistryImplTest.java
@@ -21,11 +21,6 @@ import org.apache.flink.runtime.state.TestingStreamStateHandle;
import org.junit.Test;
-import java.util.Arrays;
-import java.util.HashSet;
-import java.util.List;
-import java.util.UUID;
-
import static org.apache.flink.util.concurrent.Executors.directExecutor;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
@@ -37,11 +32,11 @@ public class TaskChangelogRegistryImplTest {
public void testDiscardedWhenNotUsed() {
TaskChangelogRegistry registry = new TaskChangelogRegistryImpl(directExecutor());
TestingStreamStateHandle handle = new TestingStreamStateHandle();
- List<UUID> backends = Arrays.asList(UUID.randomUUID(), UUID.randomUUID());
- registry.startTracking(handle, new HashSet<>(backends));
- for (UUID backend : backends) {
+ long refCount = 2;
+ registry.startTracking(handle, refCount);
+ for (int i = 0; i < refCount; i++) {
assertFalse(handle.isDisposed());
- registry.notUsed(handle, backend);
+ registry.release(handle);
}
assertTrue(handle.isDisposed());
}
@@ -50,10 +45,13 @@ public class TaskChangelogRegistryImplTest {
public void testNotDiscardedIfStoppedTracking() {
TaskChangelogRegistry registry = new TaskChangelogRegistryImpl(directExecutor());
TestingStreamStateHandle handle = new TestingStreamStateHandle();
- List<UUID> backends = Arrays.asList(UUID.randomUUID(), UUID.randomUUID());
- registry.startTracking(handle, new HashSet<>(backends));
+ long refCount = 2;
+ registry.startTracking(handle, refCount);
registry.stopTracking(handle);
- backends.forEach(id -> registry.notUsed(handle, id));
+ for (int i = 0; i < refCount; i++) {
+ assertFalse(handle.isDisposed());
+ registry.release(handle);
+ }
assertFalse(handle.isDisposed());
}
}
diff --git a/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/TestingBatchingUploadScheduler.java b/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/TestingBatchingUploadScheduler.java
new file mode 100644
index 00000000000..d5bc2c6b44c
--- /dev/null
+++ b/flink-dstl/flink-dstl-dfs/src/test/java/org/apache/flink/changelog/fs/TestingBatchingUploadScheduler.java
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.changelog.fs;
+
+import java.io.IOException;
+import java.util.LinkedList;
+import java.util.Queue;
+
+/** Implementation class for {@link StateChangeUploadScheduler} to test. */
+class TestingBatchingUploadScheduler implements StateChangeUploadScheduler {
+ private final Queue<UploadTask> tasks;
+ private final StateChangeUploader uploader;
+
+ public TestingBatchingUploadScheduler(StateChangeUploader uploader) {
+ this.tasks = new LinkedList<>();
+ this.uploader = uploader;
+ }
+
+ @Override
+ public void upload(UploadTask uploadTask) throws IOException {
+ tasks.add(uploadTask);
+ }
+
+ public void scheduleAll() throws IOException {
+ if (tasks.size() > 0) {
+ uploader.upload(tasks).complete();
+ }
+ tasks.clear();
+ }
+
+ @Override
+ public void close() throws Exception {
+ tasks.clear();
+ uploader.close();
+ }
+}
diff --git a/flink-state-backends/flink-statebackend-changelog/src/test/java/org/apache/flink/state/changelog/ChangelogStateDiscardTest.java b/flink-state-backends/flink-statebackend-changelog/src/test/java/org/apache/flink/state/changelog/ChangelogStateDiscardTest.java
index 838742a133b..1f3b3f2b855 100644
--- a/flink-state-backends/flink-statebackend-changelog/src/test/java/org/apache/flink/state/changelog/ChangelogStateDiscardTest.java
+++ b/flink-state-backends/flink-statebackend-changelog/src/test/java/org/apache/flink/state/changelog/ChangelogStateDiscardTest.java
@@ -69,9 +69,10 @@ import java.util.stream.Collectors;
import static java.util.Collections.emptyList;
import static java.util.Collections.singletonList;
import static java.util.function.Function.identity;
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.summingLong;
import static java.util.stream.Collectors.toList;
import static java.util.stream.Collectors.toMap;
-import static java.util.stream.Collectors.toSet;
import static org.apache.flink.changelog.fs.StateChangeUploadScheduler.directScheduler;
import static org.apache.flink.runtime.checkpoint.CheckpointType.CHECKPOINT;
import static org.apache.flink.runtime.state.SnapshotResult.empty;
@@ -376,7 +377,10 @@ public class ChangelogStateDiscardTest {
// todo: make the contract more explicit or extract common code
Map<UploadTask, Map<StateChangeSet, Tuple2<Long, Long>>> taskOffsets =
tasks.stream().collect(toMap(identity(), this::mapOffsets));
- tasks.forEach(task -> startTracking(registry, handle, task));
+
+ long refCount = tasks.stream().flatMap(t -> t.getChangeSets().stream()).count();
+ registry.startTracking(handle, refCount);
+
return new UploadTasksResult(taskOffsets, handle);
}
@@ -420,17 +424,33 @@ public class ChangelogStateDiscardTest {
*/
public List<UploadResult> completeUploads(
Function<UploadTask, List<UploadResult>> resultsProvider) {
+
List<UploadResult> allResults = new ArrayList<>();
+ List<Tuple2<UploadTask, List<UploadResult>>> taskResults = new ArrayList<>();
uploads.forEach(
task -> {
List<UploadResult> results = resultsProvider.apply(task);
- for (UploadResult result : results) {
- startTracking(registry, result.getStreamStateHandle(), task);
- }
+ taskResults.add(Tuple2.of(task, results));
allResults.addAll(results);
+ });
+
+ Map<StreamStateHandle, Long> stateHandleAndRefCounts =
+ allResults.stream()
+ .collect(
+ groupingBy(
+ UploadResult::getStreamStateHandle,
+ summingLong(x -> 1L)));
+ stateHandleAndRefCounts.forEach(
+ (handle, refCount) -> registry.startTracking(handle, refCount));
+
+ taskResults.forEach(
+ taskResult -> {
+ UploadTask task = taskResult.f0;
+ List<UploadResult> results = taskResult.f1;
task.complete(results);
checkState(task.isFinished());
});
+
uploads.clear();
return allResults;
}
@@ -439,15 +459,6 @@ public class ChangelogStateDiscardTest {
public void close() {}
}
- private static void startTracking(
- TaskChangelogRegistry registry,
- StreamStateHandle streamStateHandle,
- UploadTask upload) {
- registry.startTracking(
- streamStateHandle,
- upload.getChangeSets().stream().map(StateChangeSet::getLogId).collect(toSet()));
- }
-
private static void materialize(
ChangelogKeyedStateBackend<String> backend, StateChangelogWriter<?> writer) {
backend.handleMaterializationResult(empty(), 0L, writer.nextSequenceNumber());