You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by pn...@apache.org on 2020/04/10 09:05:28 UTC
[flink] 04/07: [FLINK-16744][task] Implement channel state reading
and writing for unaligned checkpoints
This is an automated email from the ASF dual-hosted git repository.
pnowojski pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 53d9f36ac1ae5a55f0fad9132117d1468e6eaeb1
Author: Roman Khachatryan <kh...@gmail.com>
AuthorDate: Sun Mar 29 22:58:36 2020 +0200
[FLINK-16744][task] Implement channel state reading and writing for unaligned checkpoints
---
.../channel/ChannelStateCheckpointWriter.java | 203 +++++++++++++
.../checkpoint/channel/ChannelStateReaderImpl.java | 118 ++++++++
.../checkpoint/channel/ChannelStateSerializer.java | 177 +++++++++++
.../channel/ChannelStateStreamReader.java | 111 +++++++
.../channel/ChannelStateWriteRequest.java | 168 +++++++++++
.../ChannelStateWriteRequestDispatcher.java | 36 +++
.../ChannelStateWriteRequestDispatcherImpl.java | 95 ++++++
.../channel/ChannelStateWriteRequestExecutor.java | 49 +++
.../ChannelStateWriteRequestExecutorImpl.java | 157 ++++++++++
.../checkpoint/channel/ChannelStateWriter.java | 97 ++++--
.../checkpoint/channel/ChannelStateWriterImpl.java | 187 ++++++++++++
.../channel/RefCountingFSDataInputStream.java | 129 ++++++++
.../runtime/io/network/buffer/BufferBuilder.java | 5 +
.../channel/ChannelStateCheckpointWriterTest.java | 144 +++++++++
.../channel/ChannelStateReaderImplTest.java | 205 +++++++++++++
.../channel/ChannelStateSerializerImplTest.java | 149 +++++++++
.../ChannelStateWriteRequestDispatcherTest.java | 118 ++++++++
.../ChannelStateWriteRequestExecutorImplTest.java | 215 +++++++++++++
.../channel/ChannelStateWriterImplTest.java | 332 +++++++++++++++++++++
.../channel/CheckpointInProgressRequestTest.java | 82 +++++
.../buffer/BufferBuilderAndConsumerTest.java | 19 ++
.../runtime/state/ChannelPersistenceITCase.java | 179 +++++++++++
22 files changed, 2954 insertions(+), 21 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriter.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriter.java
new file mode 100644
index 0000000..628fcd5
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriter.java
@@ -0,0 +1,203 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter.ChannelStateWriteResult;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.runtime.state.CheckpointStreamFactory.CheckpointStateOutputStream;
+import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.util.Preconditions;
+import org.apache.flink.util.function.RunnableWithException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.BiFunction;
+
+import static org.apache.flink.runtime.state.CheckpointedStateScope.EXCLUSIVE;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * Writes channel state for a specific checkpoint-subtask-attempt triple.
+ */
+@NotThreadSafe
+class ChannelStateCheckpointWriter {
+ private static final Logger LOG = LoggerFactory.getLogger(ChannelStateCheckpointWriter.class);
+
+ private final DataOutputStream dataStream;
+ private final CheckpointStateOutputStream checkpointStream;
+ private final ChannelStateWriteResult result;
+ private final Map<InputChannelInfo, List<Long>> inputChannelOffsets = new HashMap<>();
+ private final Map<ResultSubpartitionInfo, List<Long>> resultSubpartitionOffsets = new HashMap<>();
+ private final ChannelStateSerializer serializer;
+ private final long checkpointId;
+ private boolean allInputsReceived = false;
+ private boolean allOutputsReceived = false;
+ private final RunnableWithException onComplete;
+
+ ChannelStateCheckpointWriter(
+ CheckpointStartRequest startCheckpointItem,
+ CheckpointStreamFactory streamFactory,
+ ChannelStateSerializer serializer,
+ RunnableWithException onComplete) throws Exception {
+ this(
+ startCheckpointItem.getCheckpointId(),
+ startCheckpointItem.getTargetResult(),
+ streamFactory.createCheckpointStateOutputStream(EXCLUSIVE),
+ serializer,
+ onComplete);
+ }
+
+ ChannelStateCheckpointWriter(
+ long checkpointId,
+ ChannelStateWriteResult result,
+ CheckpointStateOutputStream stream,
+ ChannelStateSerializer serializer,
+ RunnableWithException onComplete) throws Exception {
+ this(checkpointId, result, serializer, onComplete, stream, new DataOutputStream(stream));
+ }
+
+ ChannelStateCheckpointWriter(
+ long checkpointId,
+ ChannelStateWriteResult result,
+ ChannelStateSerializer serializer,
+ RunnableWithException onComplete,
+ CheckpointStateOutputStream checkpointStateOutputStream,
+ DataOutputStream dataStream) throws Exception {
+ this.checkpointId = checkpointId;
+ this.result = checkNotNull(result);
+ this.checkpointStream = checkNotNull(checkpointStateOutputStream);
+ this.serializer = checkNotNull(serializer);
+ this.dataStream = checkNotNull(dataStream);
+ this.onComplete = checkNotNull(onComplete);
+ runWithChecks(() -> serializer.writeHeader(dataStream));
+ }
+
+ void writeInput(InputChannelInfo info, Buffer... flinkBuffers) throws Exception {
+ write(inputChannelOffsets, info, flinkBuffers, !allInputsReceived);
+ }
+
+ void writeOutput(ResultSubpartitionInfo info, Buffer... flinkBuffers) throws Exception {
+ write(resultSubpartitionOffsets, info, flinkBuffers, !allOutputsReceived);
+ }
+
+ private <K> void write(Map<K, List<Long>> offsets, K key, Buffer[] flinkBuffers, boolean precondition) throws Exception {
+ try {
+ if (result.isDone()) {
+ return;
+ }
+ runWithChecks(() -> {
+ checkState(precondition);
+ offsets
+ .computeIfAbsent(key, unused -> new ArrayList<>())
+ .add(checkpointStream.getPos());
+ serializer.writeData(dataStream, flinkBuffers);
+ });
+ } finally {
+ for (Buffer flinkBuffer : flinkBuffers) {
+ flinkBuffer.recycleBuffer();
+ }
+ }
+ }
+
+ void completeInput() throws Exception {
+ LOG.debug("complete input, output completed: {}", allOutputsReceived);
+ complete(!allInputsReceived, () -> allInputsReceived = true);
+ }
+
+ void completeOutput() throws Exception {
+ LOG.debug("complete output, input completed: {}", allInputsReceived);
+ complete(!allOutputsReceived, () -> allOutputsReceived = true);
+ }
+
+ private void complete(boolean precondition, RunnableWithException complete) throws Exception {
+ if (result.isDone()) {
+ // likely after abort - only need to set the flag run onComplete callback
+ doComplete(precondition, complete, onComplete);
+ } else {
+ runWithChecks(() -> doComplete(precondition, complete, onComplete, this::finishWriteAndResult));
+ }
+ }
+
+ private void finishWriteAndResult() throws IOException {
+ dataStream.flush();
+ StreamStateHandle underlying = checkpointStream.closeAndGetHandle();
+ complete(
+ result.inputChannelStateHandles,
+ inputChannelOffsets,
+ (chan, offsets) -> new InputChannelStateHandle(chan, underlying, offsets));
+ complete(
+ result.resultSubpartitionStateHandles,
+ resultSubpartitionOffsets,
+ (chan, offsets) -> new ResultSubpartitionStateHandle(chan, underlying, offsets));
+ }
+
+ private void doComplete(boolean precondition, RunnableWithException complete, RunnableWithException... callbacks) throws Exception {
+ Preconditions.checkArgument(precondition);
+ complete.run();
+ if (allInputsReceived && allOutputsReceived) {
+ for (RunnableWithException callback : callbacks) {
+ callback.run();
+ }
+ }
+ }
+
+ private <I, H extends AbstractChannelStateHandle<I>> void complete(
+ CompletableFuture<Collection<H>> future,
+ Map<I, List<Long>> offsets,
+ BiFunction<I, List<Long>, H> buildHandle) {
+ final Collection<H> handles = new ArrayList<>();
+ for (Map.Entry<I, List<Long>> e : offsets.entrySet()) {
+ handles.add(buildHandle.apply(e.getKey(), e.getValue()));
+ }
+ future.complete(handles);
+ LOG.debug("channel state write completed, checkpointId: {}, handles: {}", checkpointId, handles);
+ }
+
+ private void runWithChecks(RunnableWithException r) throws Exception {
+ try {
+ checkState(!result.isDone(), "result is already completed", result);
+ r.run();
+ } catch (Exception e) {
+ fail(e);
+ throw e;
+ }
+ }
+
+ public void fail(Throwable e) throws Exception {
+ result.fail(e);
+ checkpointStream.close();
+ dataStream.close();
+ }
+
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateReaderImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateReaderImpl.java
new file mode 100644
index 0000000..99b95a0
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateReaderImpl.java
@@ -0,0 +1,118 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
+import org.apache.flink.runtime.checkpoint.channel.RefCountingFSDataInputStream.RefCountingFSDataInputStreamFactory;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
+import org.apache.flink.runtime.jobgraph.OperatorID;
+import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.util.Preconditions;
+
+import org.apache.flink.shaded.guava18.com.google.common.io.Closer;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+
+import static java.util.Arrays.asList;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * {@link ChannelStateReader} implementation. Usage considerations:
+ * <ol>
+ * <li>state of a channel can be read once per instance of this class; once done it returns
+ * {@link org.apache.flink.runtime.checkpoint.channel.ChannelStateReader.ReadResult#NO_MORE_DATA NO_MORE_DATA}</li>
+ * <li>reader/writer indices of the passed buffer are respected and updated</li>
+ * <li>buffers must be prepared (cleared) before passing to reader</li>
+ * <li>buffers must be released after use</li>
+ * </ol>
+ */
+@NotThreadSafe
+@Internal
+public class ChannelStateReaderImpl implements ChannelStateReader {
+ private static final Logger log = LoggerFactory.getLogger(ChannelStateReaderImpl.class);
+
+ private final Map<InputChannelInfo, ChannelStateStreamReader> inputChannelHandleReaders;
+ private final Map<ResultSubpartitionInfo, ChannelStateStreamReader> resultSubpartitionHandleReaders;
+
+ public ChannelStateReaderImpl(TaskStateSnapshot snapshot) {
+ this(snapshot, new ChannelStateSerializerImpl());
+ }
+
+ ChannelStateReaderImpl(TaskStateSnapshot snapshot, ChannelStateDeserializer serializer) {
+ RefCountingFSDataInputStreamFactory streamFactory = new RefCountingFSDataInputStreamFactory(serializer);
+ final HashMap<InputChannelInfo, ChannelStateStreamReader> inputChannelHandleReadersTmp = new HashMap<>();
+ final HashMap<ResultSubpartitionInfo, ChannelStateStreamReader> resultSubpartitionHandleReadersTmp = new HashMap<>();
+ for (Map.Entry<OperatorID, OperatorSubtaskState> e : snapshot.getSubtaskStateMappings()) {
+ addReaders(inputChannelHandleReadersTmp, e.getValue().getInputChannelState(), streamFactory);
+ addReaders(resultSubpartitionHandleReadersTmp, e.getValue().getResultSubpartitionState(), streamFactory);
+ }
+ inputChannelHandleReaders = inputChannelHandleReadersTmp; // memory barrier to allow another thread call clear()
+ resultSubpartitionHandleReaders = resultSubpartitionHandleReadersTmp; // memory barrier to allow another thread call clear()
+ }
+
+ private <T> void addReaders(
+ Map<T, ChannelStateStreamReader> readerMap,
+ Collection<? extends AbstractChannelStateHandle<T>> handles,
+ RefCountingFSDataInputStreamFactory streamFactory) {
+ for (AbstractChannelStateHandle<T> handle : handles) {
+ checkState(!readerMap.containsKey(handle.getInfo()), "multiple states exist for channel: " + handle.getInfo());
+ readerMap.put(handle.getInfo(), new ChannelStateStreamReader(handle, streamFactory));
+ }
+ }
+
+ @Override
+ public ReadResult readInputData(InputChannelInfo info, Buffer buffer) throws IOException {
+ log.debug("readInputData, resultSubpartitionInfo: {} , buffer {}", info, buffer);
+ return getReader(info, inputChannelHandleReaders).readInto(buffer);
+ }
+
+ @Override
+ public ReadResult readOutputData(ResultSubpartitionInfo info, BufferBuilder bufferBuilder) throws IOException {
+ log.debug("readOutputData, resultSubpartitionInfo: {} , bufferBuilder {}", info, bufferBuilder);
+ return getReader(info, resultSubpartitionHandleReaders).readInto(bufferBuilder);
+ }
+
+ private <K> ChannelStateStreamReader getReader(K info, Map<K, ChannelStateStreamReader> readerMap) {
+ Preconditions.checkArgument(readerMap.containsKey(info), String.format("unknown channel %s. Known channels: %s", info, readerMap.keySet()));
+ return readerMap.get(info);
+ }
+
+ @Override
+ public void close() throws Exception {
+ try (Closer closer = Closer.create()) {
+ for (Map<?, ChannelStateStreamReader> map : asList(inputChannelHandleReaders, resultSubpartitionHandleReaders)) {
+ for (ChannelStateStreamReader reader : map.values()) {
+ closer.register(reader);
+ }
+ map.clear();
+ }
+ }
+ }
+
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializer.java
new file mode 100644
index 0000000..84e13d1
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializer.java
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
+import org.apache.flink.util.Preconditions;
+
+import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+
+import static java.lang.Math.addExact;
+import static java.lang.Math.min;
+
+interface ChannelStateSerializer {
+
+ void writeHeader(DataOutputStream dataStream) throws IOException;
+
+ void writeData(DataOutputStream stream, Buffer... flinkBuffers) throws IOException;
+}
+
+interface ChannelStateDeserializer {
+
+ void readHeader(InputStream stream) throws IOException;
+
+ int readLength(InputStream stream) throws IOException;
+
+ int readData(InputStream stream, ChannelStateByteBuffer buffer, int bytes) throws IOException;
+}
+
+/**
+ * Wrapper around various buffers to receive channel state data.
+ */
+@Internal
+@NotThreadSafe
+interface ChannelStateByteBuffer {
+
+ boolean isWritable();
+
+ /**
+ * Read up to <code>bytesToRead</code> bytes into this buffer from the given {@link InputStream}.
+ * @return the total number of bytes read into this buffer.
+ */
+ int writeBytes(InputStream input, int bytesToRead) throws IOException;
+
+ static ChannelStateByteBuffer wrap(Buffer buffer) {
+ return new ChannelStateByteBuffer() {
+
+ private final ByteBuf byteBuf = buffer.asByteBuf();
+
+ @Override
+ public boolean isWritable() {
+ return byteBuf.isWritable();
+ }
+
+ @Override
+ public int writeBytes(InputStream input, int bytesToRead) throws IOException {
+ return byteBuf.writeBytes(input, Math.min(bytesToRead, byteBuf.writableBytes()));
+ }
+ };
+ }
+
+ static ChannelStateByteBuffer wrap(BufferBuilder bufferBuilder) {
+ final byte[] buf = new byte[1024];
+ return new ChannelStateByteBuffer() {
+ @Override
+ public boolean isWritable() {
+ return !bufferBuilder.isFull();
+ }
+
+ @Override
+ public int writeBytes(InputStream input, int bytesToRead) throws IOException {
+ int left = bytesToRead;
+ for (int toRead = getToRead(left); toRead > 0; toRead = getToRead(left)) {
+ int read = input.read(buf, 0, toRead);
+ int copied = bufferBuilder.append(java.nio.ByteBuffer.wrap(buf, 0, read));
+ Preconditions.checkState(copied == read);
+ left -= read;
+ }
+ bufferBuilder.commit();
+ return bytesToRead - left;
+ }
+
+ private int getToRead(int bytesToRead) {
+ return min(bytesToRead, min(buf.length, bufferBuilder.getWritableBytes()));
+ }
+ };
+ }
+
+ static ChannelStateByteBuffer wrap(byte[] bytes) {
+ return new ChannelStateByteBuffer() {
+ private int written = 0;
+
+ @Override
+ public boolean isWritable() {
+ return written < bytes.length;
+ }
+
+ @Override
+ public int writeBytes(InputStream input, int bytesToRead) throws IOException {
+ final int bytesRead = input.read(bytes, written, bytes.length - written);
+ written += bytesRead;
+ return bytesRead;
+ }
+ };
+ }
+}
+
+class ChannelStateSerializerImpl implements ChannelStateSerializer, ChannelStateDeserializer {
+ private static final int SERIALIZATION_VERSION = 0;
+
+ @Override
+ public void writeHeader(DataOutputStream dataStream) throws IOException {
+ dataStream.writeInt(SERIALIZATION_VERSION);
+ }
+
+ @Override
+ public void writeData(DataOutputStream stream, Buffer... flinkBuffers) throws IOException {
+ stream.writeInt(getSize(flinkBuffers));
+ for (Buffer buffer : flinkBuffers) {
+ ByteBuf nettyByteBuf = buffer.asByteBuf();
+ nettyByteBuf.getBytes(nettyByteBuf.readerIndex(), stream, nettyByteBuf.readableBytes());
+ }
+ }
+
+ private int getSize(Buffer[] buffers) {
+ int len = 0;
+ for (Buffer buffer : buffers) {
+ len = addExact(len, buffer.readableBytes());
+ }
+ return len;
+ }
+
+ @Override
+ public void readHeader(InputStream stream) throws IOException {
+ int version = readInt(stream);
+ Preconditions.checkArgument(version == SERIALIZATION_VERSION, "unsupported version: " + version);
+ }
+
+ @Override
+ public int readLength(InputStream stream) throws IOException {
+ int len = readInt(stream);
+ Preconditions.checkArgument(len >= 0, "negative state size");
+ return len;
+ }
+
+ @Override
+ public int readData(InputStream stream, ChannelStateByteBuffer buffer, int bytes) throws IOException {
+ return buffer.writeBytes(stream, bytes);
+ }
+
+ private static int readInt(InputStream stream) throws IOException {
+ return new DataInputStream(stream).readInt();
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateStreamReader.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateStreamReader.java
new file mode 100644
index 0000000..6747750
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateStreamReader.java
@@ -0,0 +1,111 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateReader.ReadResult;
+import org.apache.flink.runtime.checkpoint.channel.RefCountingFSDataInputStream.RefCountingFSDataInputStreamFactory;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
+import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.util.Preconditions;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Queue;
+
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateByteBuffer.wrap;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateReader.ReadResult.HAS_MORE_DATA;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateReader.ReadResult.NO_MORE_DATA;
+
+/**
+ * Reads the state of a single channel pointed by {@link org.apache.flink.runtime.state.AbstractChannelStateHandle AbstractChannelStateHandle}.
+ * Once all data is read, this class can't be used anymore.
+ * Uses {@link RefCountingFSDataInputStream} internally.
+ */
+@NotThreadSafe
+class ChannelStateStreamReader implements Closeable {
+
+ private final RefCountingFSDataInputStream stream;
+ private final ChannelStateDeserializer serializer;
+ private final Queue<Long> offsets;
+ private int remainingBytes = -1;
+ private boolean closed = false;
+
+ ChannelStateStreamReader(AbstractChannelStateHandle<?> handle, RefCountingFSDataInputStreamFactory streamFactory) {
+ this(streamFactory.getOrCreate(handle), handle.getOffsets(), streamFactory.getSerializer());
+ }
+
+ private ChannelStateStreamReader(RefCountingFSDataInputStream stream, List<Long> offsets, ChannelStateDeserializer serializer) {
+ this.stream = stream;
+ this.stream.incRef();
+ this.serializer = serializer;
+ this.offsets = new LinkedList<>(offsets);
+ }
+
+ ReadResult readInto(Buffer buffer) throws IOException {
+ return readInto(wrap(buffer));
+ }
+
+ ReadResult readInto(BufferBuilder bufferBuilder) throws IOException {
+ return readInto(wrap(bufferBuilder));
+ }
+
+ private ReadResult readInto(ChannelStateByteBuffer buffer) throws IOException {
+ Preconditions.checkState(!closed, "reader is closed");
+ readWhilePossible(buffer);
+ if (haveMoreData()) {
+ return HAS_MORE_DATA;
+ } else {
+ closed = true;
+ stream.decRef();
+ return NO_MORE_DATA;
+ }
+ }
+
+ private void readWhilePossible(ChannelStateByteBuffer buffer) throws IOException {
+ while (haveMoreData() && buffer.isWritable()) {
+ if (remainingBytes <= 0) {
+ advanceOffset();
+ }
+ int bytesRead = serializer.readData(stream, buffer, remainingBytes);
+ remainingBytes -= bytesRead;
+ }
+ }
+
+ private boolean haveMoreData() {
+ return remainingBytes > 0 || !offsets.isEmpty();
+ }
+
+ @SuppressWarnings("ConstantConditions")
+ private void advanceOffset() throws IOException {
+ stream.seek(offsets.poll());
+ remainingBytes = serializer.readLength(stream);
+ }
+
+ @Override
+ public void close() throws IOException {
+ if (!closed) {
+ closed = true;
+ stream.close();
+ }
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequest.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequest.java
new file mode 100644
index 0000000..bd6c7ba
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequest.java
@@ -0,0 +1,168 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter.ChannelStateWriteResult;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.state.CheckpointStorageLocationReference;
+import org.apache.flink.util.Preconditions;
+import org.apache.flink.util.function.ThrowingConsumer;
+
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+
+import static org.apache.flink.runtime.checkpoint.channel.CheckpointInProgressRequestState.CANCELLED;
+import static org.apache.flink.runtime.checkpoint.channel.CheckpointInProgressRequestState.COMPLETED;
+import static org.apache.flink.runtime.checkpoint.channel.CheckpointInProgressRequestState.EXECUTING;
+import static org.apache.flink.runtime.checkpoint.channel.CheckpointInProgressRequestState.FAILED;
+import static org.apache.flink.runtime.checkpoint.channel.CheckpointInProgressRequestState.NEW;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+interface ChannelStateWriteRequest {
+ long getCheckpointId();
+
+ void cancel(Throwable cause);
+
+ static CheckpointInProgressRequest completeInput(long checkpointId) {
+ return new CheckpointInProgressRequest("completeInput", checkpointId, ChannelStateCheckpointWriter::completeInput, false);
+ }
+
+ static CheckpointInProgressRequest completeOutput(long checkpointId) {
+ return new CheckpointInProgressRequest("completeOutput", checkpointId, ChannelStateCheckpointWriter::completeOutput, false);
+ }
+
+ static ChannelStateWriteRequest write(long checkpointId, InputChannelInfo info, Buffer... flinkBuffers) {
+ return new CheckpointInProgressRequest("writeInput", checkpointId, writer -> writer.writeInput(info, flinkBuffers), recycle(flinkBuffers), false);
+ }
+
+ static ChannelStateWriteRequest write(long checkpointId, ResultSubpartitionInfo info, Buffer... flinkBuffers) {
+ return new CheckpointInProgressRequest("writeOutput", checkpointId, writer -> writer.writeOutput(info, flinkBuffers), recycle(flinkBuffers), false);
+ }
+
+ static ChannelStateWriteRequest start(long checkpointId, ChannelStateWriteResult targetResult, CheckpointStorageLocationReference locationReference) {
+ return new CheckpointStartRequest(checkpointId, targetResult, locationReference);
+ }
+
+ static ChannelStateWriteRequest abort(long checkpointId, Throwable cause) {
+ return new CheckpointInProgressRequest("abort", checkpointId, writer -> writer.fail(cause), true);
+ }
+
+ static Consumer<Throwable> recycle(Buffer[] flinkBuffers) {
+ return unused -> {
+ for (Buffer b : flinkBuffers) {
+ b.recycleBuffer();
+ }
+ };
+ }
+}
+
+final class CheckpointStartRequest implements ChannelStateWriteRequest {
+ private final ChannelStateWriteResult targetResult;
+ private final CheckpointStorageLocationReference locationReference;
+ private final long checkpointId;
+
+ CheckpointStartRequest(long checkpointId, ChannelStateWriteResult targetResult, CheckpointStorageLocationReference locationReference) {
+ this.checkpointId = checkpointId;
+ this.targetResult = checkNotNull(targetResult);
+ this.locationReference = checkNotNull(locationReference);
+ }
+
+ @Override
+ public long getCheckpointId() {
+ return checkpointId;
+ }
+
+ ChannelStateWriteResult getTargetResult() {
+ return targetResult;
+ }
+
+ public CheckpointStorageLocationReference getLocationReference() {
+ return locationReference;
+ }
+
+ @Override
+ public void cancel(Throwable cause) {
+ targetResult.fail(cause);
+ }
+
+ @Override
+ public String toString() {
+ return "start " + checkpointId;
+ }
+}
+
+enum CheckpointInProgressRequestState {
+ NEW, EXECUTING, COMPLETED, FAILED, CANCELLED
+}
+
+final class CheckpointInProgressRequest implements ChannelStateWriteRequest {
+ private final ThrowingConsumer<ChannelStateCheckpointWriter, Exception> action;
+ private final Consumer<Throwable> discardAction;
+ private final long checkpointId;
+ private final String name;
+ private final boolean ignoreMissingWriter;
+ private final AtomicReference<CheckpointInProgressRequestState> state = new AtomicReference<>(NEW);
+
+ CheckpointInProgressRequest(String name, long checkpointId, ThrowingConsumer<ChannelStateCheckpointWriter, Exception> action, boolean ignoreMissingWriter) {
+ this(name, checkpointId, action, unused -> {
+ }, ignoreMissingWriter);
+ }
+
+ CheckpointInProgressRequest(String name, long checkpointId, ThrowingConsumer<ChannelStateCheckpointWriter, Exception> action, Consumer<Throwable> discardAction, boolean ignoreMissingWriter) {
+ this.checkpointId = checkpointId;
+ this.action = checkNotNull(action);
+ this.discardAction = checkNotNull(discardAction);
+ this.name = checkNotNull(name);
+ this.ignoreMissingWriter = ignoreMissingWriter;
+ }
+
+ @Override
+ public long getCheckpointId() {
+ return checkpointId;
+ }
+
+ @Override
+ public void cancel(Throwable cause) {
+ if (state.compareAndSet(NEW, CANCELLED) || state.compareAndSet(FAILED, CANCELLED)) {
+ discardAction.accept(cause);
+ }
+ }
+
+ void execute(ChannelStateCheckpointWriter channelStateCheckpointWriter) throws Exception {
+ Preconditions.checkState(state.compareAndSet(NEW, EXECUTING));
+ try {
+ action.accept(channelStateCheckpointWriter);
+ state.set(COMPLETED);
+ } catch (Exception e) {
+ state.set(FAILED);
+ throw e;
+ }
+ }
+
+ void onWriterMissing() {
+ if (!ignoreMissingWriter) {
+ throw new IllegalArgumentException("writer not found while processing request: " + toString());
+ }
+ }
+
+ @Override
+ public String toString() {
+ return name + " " + checkpointId;
+ }
+
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestDispatcher.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestDispatcher.java
new file mode 100644
index 0000000..461acc4
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestDispatcher.java
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+interface ChannelStateWriteRequestDispatcher {
+
+ void dispatch(ChannelStateWriteRequest request) throws Exception;
+
+ void fail(Throwable cause);
+
+ ChannelStateWriteRequestDispatcher NO_OP = new ChannelStateWriteRequestDispatcher() {
+ @Override
+ public void dispatch(ChannelStateWriteRequest request) {
+ }
+
+ @Override
+ public void fail(Throwable cause) {
+ }
+ };
+}
+
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestDispatcherImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestDispatcherImpl.java
new file mode 100644
index 0000000..843663e
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestDispatcherImpl.java
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.runtime.state.CheckpointStorageWorkerView;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * Maintains a set of {@link ChannelStateCheckpointWriter writers} per checkpoint and translates incoming
+ * {@link ChannelStateWriteRequest requests} to their corresponding methods.
+ */
+final class ChannelStateWriteRequestDispatcherImpl implements ChannelStateWriteRequestDispatcher {
+ private static final Logger LOG = LoggerFactory.getLogger(ChannelStateWriteRequestDispatcherImpl.class);
+
+ private final Map<Long, ChannelStateCheckpointWriter> writers; // limited indirectly by results max size
+ private final CheckpointStorageWorkerView streamFactoryResolver;
+ private final ChannelStateSerializer serializer;
+
+ ChannelStateWriteRequestDispatcherImpl(CheckpointStorageWorkerView streamFactoryResolver, ChannelStateSerializer serializer) {
+ this.writers = new HashMap<>();
+ this.streamFactoryResolver = checkNotNull(streamFactoryResolver);
+ this.serializer = checkNotNull(serializer);
+ }
+
+ @Override
+ public void dispatch(ChannelStateWriteRequest request) throws Exception {
+ LOG.debug("process {}", request);
+ try {
+ dispatchInternal(request);
+ } catch (Exception e) {
+ request.cancel(e);
+ throw e;
+ }
+ }
+
+ private void dispatchInternal(ChannelStateWriteRequest request) throws Exception {
+ if (request instanceof CheckpointStartRequest) {
+ checkState(!writers.containsKey(request.getCheckpointId()), "writer not found for request " + request);
+ writers.put(request.getCheckpointId(), buildWriter((CheckpointStartRequest) request));
+ } else if (request instanceof CheckpointInProgressRequest) {
+ ChannelStateCheckpointWriter writer = writers.get(request.getCheckpointId());
+ CheckpointInProgressRequest req = (CheckpointInProgressRequest) request;
+ if (writer == null) {
+ req.onWriterMissing();
+ } else {
+ req.execute(writer);
+ }
+ } else {
+ throw new IllegalArgumentException("unknown request type: " + request);
+ }
+ }
+
+ private ChannelStateCheckpointWriter buildWriter(CheckpointStartRequest request) throws Exception {
+ return new ChannelStateCheckpointWriter(
+ request,
+ streamFactoryResolver.resolveCheckpointStorageLocation(request.getCheckpointId(), request.getLocationReference()),
+ serializer,
+ () -> writers.remove(request.getCheckpointId()));
+ }
+
+ @Override
+ public void fail(Throwable cause) {
+ for (ChannelStateCheckpointWriter writer : writers.values()) {
+ try {
+ writer.fail(cause);
+ } catch (Exception ex) {
+ LOG.warn("unable to fail write channel state writer", cause);
+ }
+ }
+ writers.clear();
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestExecutor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestExecutor.java
new file mode 100644
index 0000000..e908766
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestExecutor.java
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import java.io.Closeable;
+
+/**
+ * Executes {@link ChannelStateWriteRequest}s potentially asynchronously. An exception thrown during the execution
+ * should be re-thrown on any next call.
+ */
+interface ChannelStateWriteRequestExecutor extends Closeable {
+
+ /**
+ * @throws IllegalStateException if called more than once or after {@link #close()}
+ */
+ void start() throws IllegalStateException;
+
+ /**
+ * Send {@link ChannelStateWriteRequest} to this worker. If this method throws an exception then client must
+ * {@link ChannelStateWriteRequest#cancel cancel} it.
+ * @throws IllegalStateException if worker is not running
+ * @throws Exception if any exception occurred during processing this or other items previously
+ */
+ void submit(ChannelStateWriteRequest r) throws Exception;
+
+ /**
+ * Send {@link ChannelStateWriteRequest} to this worker to be processed first. If this method throws an exception then client must
+ * {@link ChannelStateWriteRequest#cancel cancel} it.
+ * @throws IllegalStateException if worker is not running
+ * @throws Exception if any exception occurred during processing this or other items previously
+ */
+ void submitPriority(ChannelStateWriteRequest r) throws Exception;
+
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestExecutorImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestExecutorImpl.java
new file mode 100644
index 0000000..cbcc3f7
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestExecutorImpl.java
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.function.RunnableWithException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.BlockingDeque;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.LinkedBlockingDeque;
+
+/**
+ * Executes {@link ChannelStateWriteRequest}s in a separate thread. Any exception occurred during execution causes this
+ * thread to stop and the exception to be re-thrown on any subsequent call.
+ */
+@ThreadSafe
+class ChannelStateWriteRequestExecutorImpl implements ChannelStateWriteRequestExecutor {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ChannelStateWriteRequestExecutorImpl.class);
+ private static final int DEFAULT_HANDOVER_CAPACITY = 10_000;
+
+ private final ChannelStateWriteRequestDispatcher dispatcher;
+ private final BlockingDeque<ChannelStateWriteRequest> deque;
+ private final Thread thread;
+ private volatile Exception thrown = null;
+ private volatile boolean wasClosed = false;
+
+ ChannelStateWriteRequestExecutorImpl(ChannelStateWriteRequestDispatcher dispatcher) {
+ this(dispatcher, new LinkedBlockingDeque<>(DEFAULT_HANDOVER_CAPACITY));
+ }
+
+ ChannelStateWriteRequestExecutorImpl(ChannelStateWriteRequestDispatcher dispatcher, BlockingDeque<ChannelStateWriteRequest> deque) {
+ this.dispatcher = dispatcher;
+ this.deque = deque;
+ this.thread = new Thread(this::run);
+ this.thread.setDaemon(true);
+ }
+
+ @VisibleForTesting
+ void run() {
+ try {
+ loop();
+ } catch (Exception ex) {
+ thrown = ex;
+ } finally {
+ cleanupRequests();
+ dispatcher.fail(thrown == null ? new CancellationException() : thrown);
+ }
+ LOG.debug("loop terminated");
+ }
+
+ private void loop() throws Exception {
+ while (!wasClosed) {
+ try {
+ dispatcher.dispatch(deque.take());
+ } catch (InterruptedException e) {
+ if (!wasClosed) {
+ LOG.debug("interrupted while waiting for a request (continue waiting)", e);
+ } else {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+ }
+
+ private void cleanupRequests() {
+ Throwable cause = thrown == null ? new CancellationException() : thrown;
+ List<ChannelStateWriteRequest> drained = new ArrayList<>();
+ deque.drainTo(drained);
+ LOG.info("discarding {} drained requests", drained.size());
+ for (ChannelStateWriteRequest request : drained) {
+ request.cancel(cause);
+ }
+ }
+
+ @Override
+ public void start() throws IllegalStateException {
+ this.thread.start();
+ }
+
+ @Override
+ public void submit(ChannelStateWriteRequest request) throws Exception {
+ submitInternal(request, () -> deque.add(request));
+ }
+
+ @Override
+ public void submitPriority(ChannelStateWriteRequest request) throws Exception {
+ submitInternal(request, () -> deque.addFirst(request));
+ }
+
+ private void submitInternal(ChannelStateWriteRequest request, RunnableWithException action) throws Exception {
+ try {
+ action.run();
+ } catch (Exception ex) {
+ request.cancel(ex);
+ throw ex;
+ }
+ ensureRunning();
+ }
+
+ private void ensureRunning() throws Exception {
+ // this check should be performed *at least after* enqueuing a request
+ // checking before is not enough because (check + enqueue) is not atomic
+ if (wasClosed || !thread.isAlive()) {
+ cleanupRequests();
+ throw ExceptionUtils.firstOrSuppressed(new IllegalStateException("not running"), thrown);
+ }
+ }
+
+ @Override
+ public void close() throws IOException {
+ wasClosed = true;
+ while (thread.isAlive()) {
+ thread.interrupt();
+ try {
+ thread.join();
+ } catch (InterruptedException e) {
+ if (!thread.isAlive()) {
+ Thread.currentThread().interrupt();
+ }
+ LOG.debug("interrupted while waiting for the writer thread to die", e);
+ }
+ }
+ if (thrown != null) {
+ throw new IOException(thrown);
+ }
+ }
+
+ @VisibleForTesting
+ Thread getThread() {
+ return thread;
+ }
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriter.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriter.java
index 7c6a57c..3de2631 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriter.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriter.java
@@ -18,19 +18,62 @@
package org.apache.flink.runtime.checkpoint.channel;
import org.apache.flink.annotation.Internal;
+import org.apache.flink.runtime.checkpoint.CheckpointOptions;
import org.apache.flink.runtime.io.network.buffer.Buffer;
-import org.apache.flink.runtime.state.StateObject;
+import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.ResultSubpartitionStateHandle;
+import java.io.Closeable;
import java.util.Collection;
import java.util.Collections;
import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.Future;
/**
* Writes channel state during checkpoint/savepoint.
*/
@Internal
-public interface ChannelStateWriter extends AutoCloseable {
+public interface ChannelStateWriter extends Closeable {
+
+ /**
+ * Channel state write result.
+ */
+ class ChannelStateWriteResult {
+ final CompletableFuture<Collection<InputChannelStateHandle>> inputChannelStateHandles;
+ final CompletableFuture<Collection<ResultSubpartitionStateHandle>> resultSubpartitionStateHandles;
+
+ ChannelStateWriteResult() {
+ this(new CompletableFuture<>(), new CompletableFuture<>());
+ }
+
+ ChannelStateWriteResult(
+ CompletableFuture<Collection<InputChannelStateHandle>> inputChannelStateHandles,
+ CompletableFuture<Collection<ResultSubpartitionStateHandle>> resultSubpartitionStateHandles) {
+ this.inputChannelStateHandles = inputChannelStateHandles;
+ this.resultSubpartitionStateHandles = resultSubpartitionStateHandles;
+ }
+
+ public CompletableFuture<Collection<InputChannelStateHandle>> getInputChannelStateHandles() {
+ return inputChannelStateHandles;
+ }
+
+ public CompletableFuture<Collection<ResultSubpartitionStateHandle>> getResultSubpartitionStateHandles() {
+ return resultSubpartitionStateHandles;
+ }
+
+ public static final ChannelStateWriteResult EMPTY = new ChannelStateWriteResult(
+ CompletableFuture.completedFuture(Collections.emptyList()),
+ CompletableFuture.completedFuture(Collections.emptyList())
+ );
+
+ public void fail(Throwable e) {
+ inputChannelStateHandles.completeExceptionally(e);
+ resultSubpartitionStateHandles.completeExceptionally(e);
+ }
+
+ boolean isDone() {
+ return inputChannelStateHandles.isDone() && resultSubpartitionStateHandles.isDone();
+ }
+ }
/**
* Sequence number for the buffers that were saved during the previous execution attempt; then restored; and now are
@@ -46,60 +89,66 @@ public interface ChannelStateWriter extends AutoCloseable {
/**
* Initiate write of channel state for the given checkpoint id.
*/
- void start(long checkpointId);
+ void start(long checkpointId, CheckpointOptions checkpointOptions);
/**
* Add in-flight buffers from the {@link org.apache.flink.runtime.io.network.partition.consumer.InputChannel InputChannel}.
- * Must be called after {@link #start(long)} and before {@link #finish(long)}.
+ * Must be called after {@link #start} (long)} and before {@link #finishInput(long)}.
+ * Buffers are recycled after they are written or exception occurs.
* @param startSeqNum sequence number of the 1st passed buffer.
* It is intended to use for incremental snapshots.
* If no data is passed it is ignored.
- * @param data zero or more buffers ordered by their sequence numbers
+ * @param data zero or more <b>data</b> buffers ordered by their sequence numbers
+ * @throws IllegalArgumentException if one or more passed buffers {@link Buffer#isBuffer() isn't a buffer}
* @see org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter#SEQUENCE_NUMBER_RESTORED
* @see org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter#SEQUENCE_NUMBER_UNKNOWN
*/
- void addInputData(long checkpointId, InputChannelInfo info, int startSeqNum, Buffer... data);
+ void addInputData(long checkpointId, InputChannelInfo info, int startSeqNum, Buffer... data) throws IllegalArgumentException;
/**
* Add in-flight buffers from the {@link org.apache.flink.runtime.io.network.partition.ResultSubpartition ResultSubpartition}.
- * Must be called after {@link #start(long)} and before {@link #finish(long)}.
+ * Must be called after {@link #start} and before {@link #finishOutput(long)}.
+ * Buffers are recycled after they are written or exception occurs.
* @param startSeqNum sequence number of the 1st passed buffer.
* It is intended to use for incremental snapshots.
* If no data is passed it is ignored.
- * @param data zero or more buffers ordered by their sequence numbers
+ * @param data zero or more <b>data</b> buffers ordered by their sequence numbers
+ * @throws IllegalArgumentException if one or more passed buffers {@link Buffer#isBuffer() isn't a buffer}
* @see org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter#SEQUENCE_NUMBER_RESTORED
* @see org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter#SEQUENCE_NUMBER_UNKNOWN
*/
- void addOutputData(long checkpointId, ResultSubpartitionInfo info, int startSeqNum, Buffer... data);
+ void addOutputData(long checkpointId, ResultSubpartitionInfo info, int startSeqNum, Buffer... data) throws IllegalArgumentException;
/**
* Finalize write of channel state data for the given checkpoint id.
- * Must be called after {@link #start(long)} and all of the input data of the given checkpoint added.
+ * Must be called after {@link #start(long, CheckpointOptions)} and all of the input data of the given checkpoint added.
* When both {@link #finishInput} and {@link #finishOutput} were called the results can be (eventually) obtained
- * using {@link #getWriteCompletionFuture}
+ * using {@link #getWriteResult}
*/
void finishInput(long checkpointId);
/**
* Finalize write of channel state data for the given checkpoint id.
- * Must be called after {@link #start(long)} and all of the output data of the given checkpoint added.
+ * Must be called after {@link #start(long, CheckpointOptions)} and all of the output data of the given checkpoint added.
* When both {@link #finishInput} and {@link #finishOutput} were called the results can be (eventually) obtained
- * using {@link #getWriteCompletionFuture}
+ * using {@link #getWriteResult}
*/
void finishOutput(long checkpointId);
/**
- * Must be called after {@link #start(long)}.
+ * Aborts the checkpoint and fails pending result for this checkpoint.
*/
- Future<Collection<StateObject>> getWriteCompletionFuture(long checkpointId);
+ void abort(long checkpointId, Throwable cause);
- @Override
- void close() throws Exception;
+ /**
+ * Must be called after {@link #start(long, CheckpointOptions)}.
+ */
+ ChannelStateWriteResult getWriteResult(long checkpointId);
ChannelStateWriter NO_OP = new ChannelStateWriter() {
@Override
- public void start(long checkpointId) {
+ public void start(long checkpointId, CheckpointOptions checkpointOptions) {
}
@Override
@@ -119,8 +168,14 @@ public interface ChannelStateWriter extends AutoCloseable {
}
@Override
- public Future<Collection<StateObject>> getWriteCompletionFuture(long checkpointId) {
- return CompletableFuture.completedFuture(Collections.emptyList());
+ public void abort(long checkpointId, Throwable cause) {
+ }
+
+ @Override
+ public ChannelStateWriteResult getWriteResult(long checkpointId) {
+ return new ChannelStateWriteResult(
+ CompletableFuture.completedFuture(Collections.emptyList()),
+ CompletableFuture.completedFuture(Collections.emptyList()));
}
@Override
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImpl.java
new file mode 100644
index 0000000..6aa2c75
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImpl.java
@@ -0,0 +1,187 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.runtime.checkpoint.CheckpointOptions;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.state.CheckpointStorageWorkerView;
+import org.apache.flink.runtime.state.CheckpointStreamFactory;
+import org.apache.flink.util.Preconditions;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.ThreadSafe;
+
+import java.io.IOException;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.completeInput;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.completeOutput;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.write;
+
+/**
+ * {@link ChannelStateWriter} implemented using
+ * {@link CheckpointStreamFactory.CheckpointStateOutputStream CheckpointStateOutputStreams}. Internally, it has by default
+ * <ul>
+ * <li>one stream per checkpoint; having multiple streams would mean more files written and more connections opened
+ * (and more latency on restore)</li>
+ * <li>one thread; having multiple threads means more connections, couples with the implementation and increases complexity</li>
+ * </ul>
+ * Thread-safety: this class is thread-safe when used with a thread-safe {@link ChannelStateWriteRequestExecutor executor}
+ * (e.g. default {@link ChannelStateWriteRequestExecutorImpl}.
+ */
+@Internal
+@ThreadSafe
+public class ChannelStateWriterImpl implements ChannelStateWriter {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ChannelStateWriterImpl.class);
+ private static final int DEFAULT_MAX_CHECKPOINTS = 5; // currently, only single in-flight checkpoint is supported
+
+ private final ChannelStateWriteRequestExecutor executor;
+ private final ConcurrentMap<Long, ChannelStateWriteResult> results;
+ private final int maxCheckpoints;
+
+ /**
+ * Creates a {@link ChannelStateWriterImpl} with {@link #DEFAULT_MAX_CHECKPOINTS} as {@link #maxCheckpoints}.
+ */
+ public ChannelStateWriterImpl(CheckpointStorageWorkerView streamFactoryResolver) {
+ this(streamFactoryResolver, DEFAULT_MAX_CHECKPOINTS);
+ }
+
+ /**
+ * Creates a {@link ChannelStateWriterImpl} with {@link ChannelStateSerializerImpl default} {@link ChannelStateSerializer},
+ * and a {@link ChannelStateWriteRequestExecutorImpl}.
+ *
+ * @param maxCheckpoints maximum number of checkpoints to be written currently or finished but not taken yet.
+ * @param streamFactoryResolver a factory to obtain output stream factory for a given checkpoint
+ */
+ ChannelStateWriterImpl(CheckpointStorageWorkerView streamFactoryResolver, int maxCheckpoints) {
+ this(
+ new ConcurrentHashMap<>(maxCheckpoints),
+ new ChannelStateWriteRequestExecutorImpl(new ChannelStateWriteRequestDispatcherImpl(streamFactoryResolver, new ChannelStateSerializerImpl())),
+ maxCheckpoints
+ );
+ }
+
+ ChannelStateWriterImpl(ConcurrentMap<Long, ChannelStateWriteResult> results, ChannelStateWriteRequestExecutor executor, int maxCheckpoints) {
+ this.results = results;
+ this.maxCheckpoints = maxCheckpoints;
+ this.executor = executor;
+ }
+
+ @Override
+ public void start(long checkpointId, CheckpointOptions checkpointOptions) {
+ LOG.debug("start checkpoint {} ({})", checkpointId, checkpointOptions);
+ ChannelStateWriteResult result = new ChannelStateWriteResult();
+ ChannelStateWriteResult put = results.computeIfAbsent(checkpointId, id -> {
+ Preconditions.checkState(results.size() < maxCheckpoints, "results.size() > maxCheckpoints", results.size(), maxCheckpoints);
+ enqueue(new CheckpointStartRequest(checkpointId, result, checkpointOptions.getTargetLocation()), false);
+ return result;
+ });
+ Preconditions.checkArgument(put == result, "result future already present for checkpoint id: " + checkpointId);
+ }
+
+ @Override
+ public void addInputData(long checkpointId, InputChannelInfo info, int startSeqNum, Buffer... data) {
+ LOG.debug("add input data, checkpoint id: {}, channel: {}, startSeqNum: {}, num buffers: {}",
+ checkpointId, info, startSeqNum, data == null ? 0 : data.length);
+ enqueue(write(checkpointId, info, checkBufferType(data)), false);
+ }
+
+ @Override
+ public void addOutputData(long checkpointId, ResultSubpartitionInfo info, int startSeqNum, Buffer... data) {
+ LOG.debug("add output data, checkpoint id: {}, channel: {}, startSeqNum: {}, num buffers: {}",
+ checkpointId, info, startSeqNum, data == null ? 0 : data.length);
+ enqueue(write(checkpointId, info, checkBufferType(data)), false);
+ }
+
+ @Override
+ public void finishInput(long checkpointId) {
+ LOG.debug("finish input data, checkpoint id: {}", checkpointId);
+ enqueue(completeInput(checkpointId), false);
+ }
+
+ @Override
+ public void finishOutput(long checkpointId) {
+ LOG.debug("finish output data, checkpoint id: {}", checkpointId);
+ enqueue(completeOutput(checkpointId), false);
+ }
+
+ @Override
+ public void abort(long checkpointId, Throwable cause) {
+ LOG.debug("abort, checkpoint id: {}", checkpointId);
+ enqueue(ChannelStateWriteRequest.abort(checkpointId, cause), true); // abort already started
+ enqueue(ChannelStateWriteRequest.abort(checkpointId, cause), false); // abort enqueued but not started
+ results.remove(checkpointId);
+ }
+
+ @Override
+ public ChannelStateWriteResult getWriteResult(long checkpointId) {
+ LOG.debug("requested write result, checkpoint id: {}", checkpointId);
+ ChannelStateWriteResult result = results.remove(checkpointId);
+ Preconditions.checkArgument(result != null, "channel state write result not found for checkpoint id " + checkpointId);
+ return result;
+ }
+
+ public void open() {
+ executor.start();
+ }
+
+ @Override
+ public void close() throws IOException {
+ results.clear();
+ executor.close();
+ }
+
+ private void enqueue(ChannelStateWriteRequest request, boolean atTheFront) {
+ // state check and previous errors check are performed inside the worker
+ try {
+ if (atTheFront) {
+ executor.submitPriority(request);
+ } else {
+ executor.submit(request);
+ }
+ } catch (Exception e) {
+ request.cancel(e);
+ throw new RuntimeException("unable to send request to worker", e);
+ }
+ }
+
+ private static Buffer[] checkBufferType(Buffer... data) {
+ if (data == null) {
+ return new Buffer[0];
+ }
+ try {
+ for (Buffer buffer : data) {
+ Preconditions.checkArgument(buffer.isBuffer());
+ }
+ } catch (Exception e) {
+ for (Buffer buffer : data) {
+ if (buffer.isBuffer()) {
+ buffer.recycleBuffer();
+ }
+ }
+ throw e;
+ }
+ return data;
+ }
+
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RefCountingFSDataInputStream.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RefCountingFSDataInputStream.java
new file mode 100644
index 0000000..86181a2
--- /dev/null
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RefCountingFSDataInputStream.java
@@ -0,0 +1,129 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.core.fs.FSDataInputStream;
+import org.apache.flink.runtime.state.AbstractChannelStateHandle;
+import org.apache.flink.runtime.state.StreamStateHandle;
+import org.apache.flink.util.Preconditions;
+import org.apache.flink.util.function.SupplierWithException;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+@NotThreadSafe
+class RefCountingFSDataInputStream extends FSDataInputStream {
+
+ private enum State {NEW, OPENED, CLOSED}
+
+ private final SupplierWithException<FSDataInputStream, IOException> streamSupplier;
+ private FSDataInputStream stream;
+ private final ChannelStateDeserializer serializer;
+ private int refCount = 0;
+ private State state = State.NEW;
+
+ private RefCountingFSDataInputStream(
+ SupplierWithException<FSDataInputStream, IOException> streamSupplier,
+ ChannelStateDeserializer serializer) {
+ this.streamSupplier = checkNotNull(streamSupplier);
+ this.serializer = checkNotNull(serializer);
+ }
+
+ void incRef() {
+ checkNotClosed();
+ refCount++;
+ }
+
+ void decRef() throws IOException {
+ checkNotClosed();
+ refCount--;
+ if (refCount == 0) {
+ close();
+ }
+ }
+
+ @Override
+ public int read() throws IOException {
+ ensureOpen();
+ return stream.read();
+ }
+
+ @Override
+ public void seek(long pos) throws IOException {
+ ensureOpen();
+ stream.seek(pos);
+ }
+
+ @Override
+ public long getPos() throws IOException {
+ ensureOpen();
+ return stream.getPos();
+ }
+
+ public void close() throws IOException {
+ state = State.CLOSED;
+ if (stream != null) {
+ stream.close();
+ stream = null;
+ }
+ }
+
+ private void ensureOpen() throws IOException {
+ checkNotClosed();
+ if (state == State.NEW) {
+ stream = Preconditions.checkNotNull(streamSupplier.get());
+ serializer.readHeader(stream);
+ state = State.OPENED;
+ }
+ }
+
+ private void checkNotClosed() {
+ checkState(state != State.CLOSED, "stream is closed");
+ }
+
+ @NotThreadSafe
+ static class RefCountingFSDataInputStreamFactory {
+ private final Map<StreamStateHandle, RefCountingFSDataInputStream> streams = new HashMap<>(); // not clearing: expecting short life
+ private final ChannelStateDeserializer serializer;
+
+ RefCountingFSDataInputStreamFactory(ChannelStateDeserializer serializer) {
+ this.serializer = checkNotNull(serializer);
+ }
+
+ <T> RefCountingFSDataInputStream getOrCreate(AbstractChannelStateHandle<T> handle) {
+ StreamStateHandle streamStateHandle = handle.getDelegate();
+ RefCountingFSDataInputStream stream = streams.get(streamStateHandle);
+ if (stream == null) {
+ stream = new RefCountingFSDataInputStream(streamStateHandle::openInputStream, serializer);
+ streams.put(streamStateHandle, stream);
+ }
+ return stream;
+ }
+
+ ChannelStateDeserializer getSerializer() {
+ return serializer;
+ }
+ }
+
+}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferBuilder.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferBuilder.java
index bcd42d2..77d3526 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferBuilder.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferBuilder.java
@@ -117,6 +117,11 @@ public class BufferBuilder {
return positionMarker.getCached() == getMaxCapacity();
}
+ public int getWritableBytes() {
+ checkState(positionMarker.getCached() <= getMaxCapacity());
+ return getMaxCapacity() - positionMarker.getCached();
+ }
+
public int getMaxCapacity() {
return memorySegment.size();
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriterTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriterTest.java
new file mode 100644
index 0000000..196c948
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateCheckpointWriterTest.java
@@ -0,0 +1,144 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.core.memory.HeapMemorySegment;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter.ChannelStateWriteResult;
+import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.memory.MemCheckpointStreamFactory.MemoryCheckpointOutputStream;
+import org.apache.flink.util.function.RunnableWithException;
+
+import org.junit.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * {@link ChannelStateCheckpointWriter} test.
+ */
+public class ChannelStateCheckpointWriterTest {
+ private static final RunnableWithException NO_OP_RUNNABLE = () -> {
+ };
+ private final Random random = new Random();
+
+ @Test
+ public void testRecyclingBuffers() throws Exception {
+ ChannelStateCheckpointWriter writer = createWriter(new ChannelStateWriteResult());
+ NetworkBuffer buffer = new NetworkBuffer(HeapMemorySegment.FACTORY.allocateUnpooledSegment(10, null), FreeingBufferRecycler.INSTANCE);
+ writer.writeInput(new InputChannelInfo(1, 2), buffer);
+ assertTrue(buffer.isRecycled());
+ }
+
+ @Test
+ public void testFlush() throws Exception {
+ class FlushRecorder extends DataOutputStream {
+ private boolean flushed = false;
+
+ private FlushRecorder() {
+ super(new ByteArrayOutputStream());
+ }
+
+ @Override
+ public void flush() throws IOException {
+ flushed = true;
+ super.flush();
+ }
+ }
+
+ FlushRecorder dataStream = new FlushRecorder();
+ final ChannelStateCheckpointWriter writer = new ChannelStateCheckpointWriter(
+ 1L,
+ new ChannelStateWriteResult(),
+ new ChannelStateSerializerImpl(),
+ NO_OP_RUNNABLE,
+ new MemoryCheckpointOutputStream(42),
+ dataStream
+ );
+
+ writer.completeInput();
+ writer.completeOutput();
+
+ assertTrue(dataStream.flushed);
+ }
+
+ @Test
+ public void testResultCompletion() throws Exception {
+ ChannelStateWriteResult result = new ChannelStateWriteResult();
+ ChannelStateCheckpointWriter writer = createWriter(result);
+ writer.completeInput();
+ assertFalse(result.isDone());
+ writer.completeOutput();
+ assertTrue(result.isDone());
+ }
+
+ @Test
+ public void testRecordingOffsets() throws Exception {
+ Map<InputChannelInfo, Integer> offsetCounts = new HashMap<>();
+ offsetCounts.put(new InputChannelInfo(1, 1), 1);
+ offsetCounts.put(new InputChannelInfo(1, 2), 2);
+ offsetCounts.put(new InputChannelInfo(1, 3), 99);
+
+ ChannelStateWriteResult result = new ChannelStateWriteResult();
+ ChannelStateCheckpointWriter writer = createWriter(result);
+ for (Map.Entry<InputChannelInfo, Integer> e : offsetCounts.entrySet()) {
+ for (int i = 0; i < e.getValue(); i++) {
+ write(writer, e.getKey(), getData(100));
+ }
+ }
+ writer.completeInput();
+ writer.completeOutput();
+
+ for (InputChannelStateHandle handle : result.inputChannelStateHandles.get()) {
+ assertEquals(handle.getOffsets().size(), (int) offsetCounts.remove(handle.getInfo()));
+ }
+ assertTrue(offsetCounts.isEmpty());
+ }
+
+ private byte[] getData(int len) {
+ byte[] bytes = new byte[len];
+ random.nextBytes(bytes);
+ return bytes;
+ }
+
+ private void write(ChannelStateCheckpointWriter writer, InputChannelInfo channelInfo, byte[] data) throws Exception {
+ NetworkBuffer buffer = new NetworkBuffer(HeapMemorySegment.FACTORY.allocateUnpooledSegment(data.length, null), FreeingBufferRecycler.INSTANCE);
+ buffer.setBytes(0, data);
+ writer.writeInput(channelInfo, buffer);
+ }
+
+ private ChannelStateCheckpointWriter createWriter(ChannelStateWriteResult result) throws Exception {
+ return new ChannelStateCheckpointWriter(
+ 1L,
+ result,
+ new MemoryCheckpointOutputStream(1000),
+ new ChannelStateSerializerImpl(),
+ NO_OP_RUNNABLE
+ );
+ }
+
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateReaderImplTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateReaderImplTest.java
new file mode 100644
index 0000000..8c780f3
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateReaderImplTest.java
@@ -0,0 +1,205 @@
+package org.apache.flink.runtime.checkpoint.channel;
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import org.apache.flink.core.memory.HeapMemorySegment;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.runtime.checkpoint.StateObjectCollection;
+import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
+import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.jobgraph.OperatorID;
+import org.apache.flink.runtime.state.InputChannelStateHandle;
+import org.apache.flink.runtime.state.memory.ByteStreamStateHandle;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+
+import static java.util.Collections.singletonList;
+import static java.util.stream.Collectors.toMap;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateReader.ReadResult.HAS_MORE_DATA;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateReader.ReadResult.NO_MORE_DATA;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+
+/**
+ * {@link ChannelStateReaderImpl} test.
+ */
+public class ChannelStateReaderImplTest {
+
+ private static final InputChannelInfo CHANNEL = new InputChannelInfo(1, 2);
+ private static final byte[] DATA = generateData(10);
+ private ChannelStateReaderImpl reader;
+
+ @Before
+ public void init() {
+ reader = getReader(CHANNEL, DATA);
+ }
+
+ @After
+ public void tearDown() throws Exception {
+ reader.close();
+ }
+
+ @Test
+ public void testDifferentBufferSizes() throws Exception {
+ for (int bufferSize = 1; bufferSize < 2 * DATA.length; bufferSize++) {
+ try (ChannelStateReaderImpl reader = getReader(CHANNEL, DATA)) { // re-create reader to re-read the same channel
+ readAndVerify(bufferSize, CHANNEL, DATA, reader);
+ }
+ }
+ }
+
+ @Test
+ public void testWithOffsets() throws IOException {
+ Map<InputChannelStateHandle, byte[]> handlesAndBytes = generateHandlesWithBytes(10, 20);
+ ChannelStateReader reader = new ChannelStateReaderImpl(taskStateSnapshot(handlesAndBytes.keySet()), new ChannelStateSerializerImpl());
+ for (Map.Entry<InputChannelStateHandle, byte[]> e : handlesAndBytes.entrySet()) {
+ readAndVerify(42, e.getKey().getInfo(), e.getValue(), reader);
+ }
+ }
+
+ @Test(expected = Exception.class)
+ public void testReadOnlyOnce() throws IOException {
+ reader.readInputData(CHANNEL, getBuffer(DATA.length));
+ reader.readInputData(CHANNEL, getBuffer(DATA.length));
+ }
+
+ @Test(expected = Exception.class)
+ public void testReadClosed() throws Exception {
+ reader.close();
+ reader.readInputData(CHANNEL, getBuffer(DATA.length));
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testReadWrongChannelState() throws IOException {
+ InputChannelInfo wrongChannel = new InputChannelInfo(CHANNEL.getGateIdx() + 1, CHANNEL.getInputChannelIdx() + 1);
+ reader.readInputData(wrongChannel, getBuffer(DATA.length));
+ }
+
+ private TaskStateSnapshot taskStateSnapshot(Collection<InputChannelStateHandle> inputChannelStateHandles) {
+ return new TaskStateSnapshot(Collections.singletonMap(
+ new OperatorID(),
+ new OperatorSubtaskState(
+ StateObjectCollection.empty(),
+ StateObjectCollection.empty(),
+ StateObjectCollection.empty(),
+ StateObjectCollection.empty(),
+ new StateObjectCollection<>(inputChannelStateHandles),
+ StateObjectCollection.empty()
+ )));
+ }
+
+ static byte[] generateData(int len) {
+ byte[] bytes = new byte[len];
+ new Random().nextBytes(bytes);
+ return bytes;
+ }
+
+ private byte[] toBytes(NetworkBuffer buffer) {
+ byte[] buf = new byte[buffer.readableBytes()];
+ buffer.readBytes(buf);
+ return buf;
+ }
+
+ private ChannelStateDeserializer deserializer(byte[] data) {
+ return new ChannelStateDeserializer() {
+ @Override
+ public void readHeader(InputStream stream) {
+ }
+
+ @Override
+ public int readLength(InputStream stream) {
+ return data.length;
+ }
+
+ @Override
+ public int readData(InputStream stream, ChannelStateByteBuffer buffer, int bytes) throws IOException {
+ return buffer.writeBytes(stream, bytes);
+ }
+ };
+ }
+
+ private ChannelStateReaderImpl getReader(InputChannelInfo channel, byte[] data) {
+ return new ChannelStateReaderImpl(
+ taskStateSnapshot(singletonList(new InputChannelStateHandle(channel, new ByteStreamStateHandle("", data), singletonList(0L)))),
+ deserializer(DATA));
+ }
+
+ private void readAndVerify(int bufferSize, InputChannelInfo channelInfo, byte[] data, ChannelStateReader reader) throws IOException {
+ int dataSize = data.length;
+ int iterations = dataSize / bufferSize + (-(dataSize % bufferSize) >>> 31);
+ NetworkBuffer buffer = getBuffer(bufferSize);
+ try {
+ for (int i = 0; i < iterations; i++) {
+ String hint = String.format("dataSize=%d, bufferSize=%d, iteration=%d/%d", dataSize, bufferSize, i + 1, iterations);
+ boolean isLast = i == iterations - 1;
+ assertEquals(hint, isLast ? NO_MORE_DATA : HAS_MORE_DATA, reader.readInputData(channelInfo, buffer));
+ assertEquals(hint, isLast ? dataSize - bufferSize * i : bufferSize, buffer.readableBytes());
+ assertArrayEquals(hint, Arrays.copyOfRange(data, i * bufferSize, Math.min(dataSize, (i + 1) * bufferSize)), toBytes(buffer));
+ buffer.resetReaderIndex();
+ buffer.resetWriterIndex();
+ }
+ } finally {
+ buffer.release();
+ }
+ }
+
+ private NetworkBuffer getBuffer(int len) {
+ return new NetworkBuffer(HeapMemorySegment.FACTORY.allocateUnpooledSegment(len, null), FreeingBufferRecycler.INSTANCE);
+ }
+
+ private Map<InputChannelStateHandle, byte[]> generateHandlesWithBytes(int numHandles, int handleDataSize) throws IOException {
+ Map<Integer, byte[]> offsetsAndBytes = new HashMap<>();
+ ByteArrayOutputStream baos = new ByteArrayOutputStream(100);
+ DataOutputStream out = new DataOutputStream(baos);
+ ChannelStateSerializerImpl serializer = new ChannelStateSerializerImpl();
+ serializer.writeHeader(out);
+ for (int i = 0; i < numHandles; i++) {
+ offsetsAndBytes.put(baos.size(), writeSomeBytes(handleDataSize, out, serializer));
+ }
+ ByteStreamStateHandle sharedUnderlyingHandle = new ByteStreamStateHandle("", baos.toByteArray());
+ return offsetsAndBytes.entrySet().stream().collect(toMap(
+ e -> new InputChannelStateHandle(new InputChannelInfo(e.getKey(), e.getKey()), sharedUnderlyingHandle, singletonList((long) e.getKey())),
+ Map.Entry::getValue));
+ }
+
+ private byte[] writeSomeBytes(int bytesCount, DataOutputStream out, ChannelStateSerializer serializer) throws IOException {
+ byte[] bytes = generateData(bytesCount);
+ NetworkBuffer buf = getBuffer(bytesCount);
+ try {
+ buf.writeBytes(bytes);
+ serializer.writeData(out, buf);
+ return bytes;
+ } finally {
+ buf.release();
+ }
+ }
+
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializerImplTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializerImplTest.java
new file mode 100644
index 0000000..e3a79cb
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateSerializerImplTest.java
@@ -0,0 +1,149 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.core.memory.HeapMemorySegment;
+import org.apache.flink.core.memory.MemorySegmentFactory;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
+import org.apache.flink.runtime.io.network.buffer.BufferConsumer;
+import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+
+import org.junit.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.Arrays;
+import java.util.Random;
+
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateByteBuffer.wrap;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateReaderImplTest.generateData;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+
+/**
+ * {@link ChannelStateSerializerImpl} test.
+ */
+public class ChannelStateSerializerImplTest {
+
+ private final Random random = new Random();
+
+ @Test
+ public void testReadWrite() throws IOException {
+ byte[] data = generateData(123);
+ ChannelStateSerializerImpl serializer = new ChannelStateSerializerImpl();
+ try (ByteArrayOutputStream baos = new ByteArrayOutputStream(data.length)) {
+ write(data, serializer, baos);
+ readAndCheck(data, serializer, new ByteArrayInputStream(baos.toByteArray()));
+ }
+ }
+
+ @Test
+ public void testReadWriteWithMultipleBuffers() throws IOException {
+ int bufSize = 10;
+ int[] numBuffersToWriteAtOnce = {0, 1, 2, 3};
+ byte[] data = generateData(bufSize);
+ ChannelStateSerializer s = new ChannelStateSerializerImpl();
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ DataOutputStream out = new DataOutputStream(baos);
+ s.writeHeader(out);
+ for (int count : numBuffersToWriteAtOnce) {
+ Buffer[] buffers = new Buffer[count];
+ Arrays.fill(buffers, getBuffer(data));
+ s.writeData(out, buffers);
+ }
+ out.close();
+
+ ChannelStateDeserializer d = new ChannelStateSerializerImpl();
+ ByteArrayInputStream is = new ByteArrayInputStream(baos.toByteArray());
+ d.readHeader(is);
+ for (int count : numBuffersToWriteAtOnce) {
+ int expected = bufSize * count;
+ assertEquals(expected, d.readLength(is));
+ byte[] readBuf = new byte[expected];
+ assertEquals(expected, d.readData(is, wrap(readBuf), Integer.MAX_VALUE));
+ for (int i = 0; i < count; i++) {
+ assertArrayEquals(data, Arrays.copyOfRange(readBuf, i * bufSize, (i + 1) * bufSize));
+ }
+ }
+ }
+
+ @Test
+ public void testReadToBufferBuilder() throws IOException {
+ byte[] data = generateData(100);
+ BufferBuilder bufferBuilder = new BufferBuilder(HeapMemorySegment.FACTORY.allocateUnpooledSegment(data.length, null), FreeingBufferRecycler.INSTANCE);
+ BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
+
+ new ChannelStateSerializerImpl().readData(new ByteArrayInputStream(data), wrap(bufferBuilder), Integer.MAX_VALUE);
+
+ assertFalse(bufferBuilder.isFinished());
+
+ bufferBuilder.finish();
+ Buffer buffer = bufferConsumer.build();
+
+ assertEquals(data.length, buffer.readableBytes());
+ byte[] actual = new byte[buffer.readableBytes()];
+ buffer.asByteBuf().readBytes(actual);
+ assertArrayEquals(data, actual);
+ }
+
+ private NetworkBuffer getBuffer(byte[] data) {
+ NetworkBuffer buffer = new NetworkBuffer(HeapMemorySegment.FACTORY.allocateUnpooledSegment(data.length, null), FreeingBufferRecycler.INSTANCE);
+ buffer.writeBytes(data);
+ return buffer;
+ }
+
+ private byte[] readBytes(NetworkBuffer buffer) {
+ byte[] tmp = new byte[buffer.readableBytes()];
+ buffer.readBytes(tmp);
+ return tmp;
+ }
+
+ private void write(byte[] data, ChannelStateSerializerImpl serializer, OutputStream baos) throws IOException {
+ DataOutputStream out = new DataOutputStream(baos);
+ serializer.writeHeader(out);
+ NetworkBuffer buffer = new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(data.length), FreeingBufferRecycler.INSTANCE);
+ try {
+ buffer.writeBytes(data);
+ serializer.writeData(out, buffer);
+ out.flush();
+ } finally {
+ buffer.release();
+ }
+ }
+
+ private void readAndCheck(byte[] data, ChannelStateSerializerImpl serializer, ByteArrayInputStream is) throws IOException {
+ serializer.readHeader(is);
+ int size = serializer.readLength(is);
+ assertEquals(data.length, size);
+ NetworkBuffer buffer = new NetworkBuffer(MemorySegmentFactory.allocateUnpooledSegment(data.length), FreeingBufferRecycler.INSTANCE);
+ try {
+ int read = serializer.readData(is, wrap(buffer), size);
+ assertEquals(size, read);
+ assertArrayEquals(data, readBytes(buffer));
+ } finally {
+ buffer.release();
+ }
+ }
+
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestDispatcherTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestDispatcherTest.java
new file mode 100644
index 0000000..f953c22
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestDispatcherTest.java
@@ -0,0 +1,118 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter.ChannelStateWriteResult;
+import org.apache.flink.runtime.state.CheckpointStorageLocationReference;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
+
+import java.util.List;
+import java.util.Optional;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singletonList;
+import static java.util.Optional.empty;
+import static java.util.Optional.of;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.completeInput;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.completeOutput;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequest.write;
+import static org.apache.flink.runtime.state.ChannelPersistenceITCase.getStreamFactoryFactory;
+import static org.junit.Assert.fail;
+
+/**
+ * {@link ChannelStateWriteRequestDispatcherImpl} tests.
+ */
+@RunWith(Parameterized.class)
+public class ChannelStateWriteRequestDispatcherTest {
+
+ private final List<ChannelStateWriteRequest> requests;
+ private final Optional<Class<?>> expectedException;
+ public static final long CHECKPOINT_ID = 42L;
+
+ @Parameters
+ public static Object[][] data() {
+
+ return new Object[][]{
+ // valid calls
+ new Object[]{empty(), asList(start(), completeIn(), completeOut())},
+ new Object[]{empty(), asList(start(), writeIn(), completeIn())},
+ new Object[]{empty(), asList(start(), writeOut(), completeOut())},
+ new Object[]{empty(), asList(start(), completeIn(), writeOut())},
+ new Object[]{empty(), asList(start(), completeOut(), writeIn())},
+ // invalid without start
+ new Object[]{of(IllegalArgumentException.class), singletonList(writeIn())},
+ new Object[]{of(IllegalArgumentException.class), singletonList(writeOut())},
+ new Object[]{of(IllegalArgumentException.class), singletonList(completeIn())},
+ new Object[]{of(IllegalArgumentException.class), singletonList(completeOut())},
+ // invalid double complete
+ new Object[]{of(IllegalArgumentException.class), asList(start(), completeIn(), completeIn())},
+ new Object[]{of(IllegalArgumentException.class), asList(start(), completeOut(), completeOut())},
+ // invalid write after complete
+ new Object[]{of(IllegalStateException.class), asList(start(), completeIn(), writeIn())},
+ new Object[]{of(IllegalStateException.class), asList(start(), completeOut(), writeOut())},
+ // invalid double start
+ new Object[]{of(IllegalStateException.class), asList(start(), start())}
+ };
+ }
+
+ private static CheckpointInProgressRequest completeOut() {
+ return completeOutput(CHECKPOINT_ID);
+ }
+
+ private static CheckpointInProgressRequest completeIn() {
+ return completeInput(CHECKPOINT_ID);
+ }
+
+ private static ChannelStateWriteRequest writeIn() {
+ return write(CHECKPOINT_ID, new InputChannelInfo(1, 1));
+ }
+
+ private static ChannelStateWriteRequest writeOut() {
+ return write(CHECKPOINT_ID, new ResultSubpartitionInfo(1, 1));
+ }
+
+ private static CheckpointStartRequest start() {
+ return new CheckpointStartRequest(CHECKPOINT_ID, new ChannelStateWriteResult(), new CheckpointStorageLocationReference(new byte[]{1}));
+ }
+
+ public ChannelStateWriteRequestDispatcherTest(Optional<Class<?>> expectedException, List<ChannelStateWriteRequest> requests) {
+ this.requests = requests;
+ this.expectedException = expectedException;
+ }
+
+ @Test
+ public void doRun() {
+ ChannelStateWriteRequestDispatcher processor = new ChannelStateWriteRequestDispatcherImpl(getStreamFactoryFactory(), new ChannelStateSerializerImpl());
+ try {
+ for (ChannelStateWriteRequest request : requests) {
+ processor.dispatch(request);
+ }
+ } catch (Throwable t) {
+ if (expectedException.filter(e -> e.isInstance(t)).isPresent()) {
+ return;
+ }
+ throw new RuntimeException("unexpected exception", t);
+ }
+ expectedException.ifPresent(e -> fail("expected exception " + e));
+ }
+
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestExecutorImplTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestExecutorImplTest.java
new file mode 100644
index 0000000..5aad9c6
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriteRequestExecutorImplTest.java
@@ -0,0 +1,215 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.function.BiConsumerWithException;
+
+import org.junit.Test;
+
+import javax.annotation.Nonnull;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.concurrent.LinkedBlockingDeque;
+
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriteRequestDispatcher.NO_OP;
+import static org.apache.flink.util.ExceptionUtils.findThrowable;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+/**
+ * {@link ChannelStateWriteRequestExecutorImpl} test.
+ */
+public class ChannelStateWriteRequestExecutorImplTest {
+
+ @Test(expected = IllegalStateException.class)
+ public void testCloseAfterSubmit() throws Exception {
+ testCloseAfterSubmit(ChannelStateWriteRequestExecutor::submit);
+ }
+
+ @Test(expected = IllegalStateException.class)
+ public void testCloseAfterSubmitPriority() throws Exception {
+ testCloseAfterSubmit(ChannelStateWriteRequestExecutor::submitPriority);
+ }
+
+ @Test
+ public void testSubmitFailure() throws Exception {
+ testSubmitFailure(ChannelStateWriteRequestExecutor::submit);
+ }
+
+ @Test
+ public void testSubmitPriorityFailure() throws Exception {
+ testSubmitFailure(ChannelStateWriteRequestExecutor::submitPriority);
+ }
+
+ private void testCloseAfterSubmit(BiConsumerWithException<ChannelStateWriteRequestExecutor, ChannelStateWriteRequest, Exception> requestFun) throws Exception {
+ WorkerClosingDeque closingDeque = new WorkerClosingDeque();
+ ChannelStateWriteRequestExecutorImpl worker = new ChannelStateWriteRequestExecutorImpl(NO_OP, closingDeque);
+ closingDeque.setWorker(worker);
+ TestWriteRequest request = new TestWriteRequest();
+ requestFun.accept(worker, request);
+ assertTrue(closingDeque.isEmpty());
+ assertFalse(request.isCancelled());
+ }
+
+ private void testSubmitFailure(BiConsumerWithException<ChannelStateWriteRequestExecutor, ChannelStateWriteRequest, Exception> submitAction) throws Exception {
+ TestWriteRequest request = new TestWriteRequest();
+ LinkedBlockingDeque<ChannelStateWriteRequest> deque = new LinkedBlockingDeque<>();
+ try {
+ submitAction.accept(new ChannelStateWriteRequestExecutorImpl(NO_OP, deque), request);
+ } catch (IllegalStateException e) {
+ // expected: executor not started;
+ return;
+ } finally {
+ assertTrue(request.cancelled);
+ assertTrue(deque.isEmpty());
+ }
+ throw new RuntimeException("expected exception not thrown");
+ }
+
+ @Test
+ @SuppressWarnings("CallToThreadRun")
+ public void testCleanup() throws IOException {
+ TestWriteRequest request = new TestWriteRequest();
+ LinkedBlockingDeque<ChannelStateWriteRequest> deque = new LinkedBlockingDeque<>();
+ deque.add(request);
+ TestRequestDispatcher requestProcessor = new TestRequestDispatcher();
+ ChannelStateWriteRequestExecutorImpl worker = new ChannelStateWriteRequestExecutorImpl(requestProcessor, deque);
+
+ worker.close();
+ worker.run();
+
+ assertTrue(requestProcessor.isStopped());
+ assertTrue(deque.isEmpty());
+ assertTrue(request.isCancelled());
+ }
+
+ @Test
+ public void testIgnoresInterruptsWhileRunning() throws Exception {
+ TestRequestDispatcher requestProcessor = new TestRequestDispatcher();
+ LinkedBlockingDeque<ChannelStateWriteRequest> deque = new LinkedBlockingDeque<>();
+ try (ChannelStateWriteRequestExecutorImpl worker = new ChannelStateWriteRequestExecutorImpl(requestProcessor, deque)) {
+ worker.start();
+ worker.getThread().interrupt();
+ worker.submit(new TestWriteRequest());
+ worker.getThread().interrupt();
+ while (!deque.isEmpty()) {
+ Thread.sleep(100);
+ }
+ }
+ }
+
+ @Test
+ public void testCanBeClosed() throws IOException {
+ TestRequestDispatcher requestProcessor = new TestRequestDispatcher();
+ try (ChannelStateWriteRequestExecutorImpl worker = new ChannelStateWriteRequestExecutorImpl(requestProcessor)) {
+ worker.start();
+ }
+ }
+
+ @Test
+ public void testRecordsException() throws IOException {
+ TestException testException = new TestException();
+ TestRequestDispatcher throwingRequestProcessor = new TestRequestDispatcher() {
+ @Override
+ public void dispatch(ChannelStateWriteRequest request) {
+ throw testException;
+ }
+ };
+ LinkedBlockingDeque<ChannelStateWriteRequest> deque = new LinkedBlockingDeque<>(Arrays.asList(new TestWriteRequest()));
+ ChannelStateWriteRequestExecutorImpl worker = new ChannelStateWriteRequestExecutorImpl(throwingRequestProcessor, deque);
+ worker.run();
+ try {
+ worker.close();
+ } catch (IOException e) {
+ if (findThrowable(e, TestException.class).filter(found -> found == testException).isPresent()) {
+ return;
+ } else {
+ throw e;
+ }
+ }
+ fail("exception not thrown");
+ }
+
+ private static class TestWriteRequest implements ChannelStateWriteRequest {
+ private boolean cancelled = false;
+
+ @Override
+ public long getCheckpointId() {
+ return 0;
+ }
+
+ @Override
+ public void cancel(Throwable cause) {
+ cancelled = true;
+ }
+
+ public boolean isCancelled() {
+ return cancelled;
+ }
+ }
+
+ private static class WorkerClosingDeque extends LinkedBlockingDeque<ChannelStateWriteRequest> {
+ private ChannelStateWriteRequestExecutor worker;
+
+ @Override
+ public void put(@Nonnull ChannelStateWriteRequest request) throws InterruptedException {
+ super.putFirst(request);
+ try {
+ worker.close();
+ } catch (IOException e) {
+ ExceptionUtils.rethrow(e);
+ }
+ }
+
+ @Override
+ public void putFirst(@Nonnull ChannelStateWriteRequest request) throws InterruptedException {
+ super.putFirst(request);
+ try {
+ worker.close();
+ } catch (IOException e) {
+ ExceptionUtils.rethrow(e);
+ }
+ }
+
+ public void setWorker(ChannelStateWriteRequestExecutor worker) {
+ this.worker = worker;
+ }
+ }
+
+ private static class TestRequestDispatcher implements ChannelStateWriteRequestDispatcher {
+ private boolean isStopped;
+
+ @Override
+ public void dispatch(ChannelStateWriteRequest request) {
+ }
+
+ @Override
+ public void fail(Throwable cause) {
+ isStopped = true;
+ }
+
+ public boolean isStopped() {
+ return isStopped;
+ }
+
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImplTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImplTest.java
new file mode 100644
index 0000000..0c30660
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateWriterImplTest.java
@@ -0,0 +1,332 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.apache.flink.core.memory.HeapMemorySegment;
+import org.apache.flink.runtime.checkpoint.CheckpointOptions;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter.ChannelStateWriteResult;
+import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.util.function.BiConsumerWithException;
+import org.apache.flink.util.function.RunnableWithException;
+
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Consumer;
+
+import static org.apache.flink.runtime.state.ChannelPersistenceITCase.getStreamFactoryFactory;
+import static org.apache.flink.util.ExceptionUtils.findThrowable;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * {@link ChannelStateWriterImpl} lifecycle tests.
+ */
+public class ChannelStateWriterImplTest {
+ private static final long CHECKPOINT_ID = 42L;
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testAddEventBuffer() {
+ NetworkBuffer dataBuf = getBuffer();
+ NetworkBuffer eventBuf = getBuffer();
+ eventBuf.tagAsEvent();
+ ChannelStateWriterImpl writer = openWriter();
+ callStart(writer);
+ try {
+ writer.addInputData(CHECKPOINT_ID, new InputChannelInfo(1, 1), 1, eventBuf, dataBuf);
+ } finally {
+ assertTrue(dataBuf.isRecycled());
+ }
+ }
+
+ @Test
+ public void testResultCompletion() throws IOException {
+ ChannelStateWriteResult result;
+ try (ChannelStateWriterImpl writer = openWriter()) {
+ callStart(writer);
+ result = writer.getWriteResult(CHECKPOINT_ID);
+ assertFalse(result.resultSubpartitionStateHandles.isDone());
+ assertFalse(result.inputChannelStateHandles.isDone());
+ }
+ assertTrue(result.inputChannelStateHandles.isDone());
+ assertTrue(result.resultSubpartitionStateHandles.isDone());
+ }
+
+ @Test
+ public void testAbort() throws Exception {
+ NetworkBuffer buffer = getBuffer();
+ runWithSyncWorker((writer, worker) -> {
+ callStart(writer);
+ ChannelStateWriteResult result = writer.getWriteResult(CHECKPOINT_ID);
+ callAddInputData(writer, buffer);
+ callAbort(writer);
+ worker.processAllRequests();
+ assertTrue(result.isDone());
+ assertTrue(buffer.isRecycled());
+ });
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testAbortClearsResults() throws Exception {
+ NetworkBuffer buffer = getBuffer();
+ runWithSyncWorker((writer, worker) -> {
+ callStart(writer);
+ callAbort(writer);
+ worker.processAllRequests();
+ writer.getWriteResult(CHECKPOINT_ID);
+ });
+ }
+
+ @Test
+ public void testAbortIgnoresMissing() throws Exception {
+ runWithSyncWorker(this::callAbort);
+ }
+
+ @Test(expected = TestException.class)
+ public void testBuffersRecycledOnError() throws Exception {
+ unwrappingError(TestException.class, () -> {
+ NetworkBuffer buffer = getBuffer();
+ try (ChannelStateWriterImpl writer = new ChannelStateWriterImpl(new ConcurrentHashMap<>(), failingWorker(), 5)) {
+ writer.open();
+ callAddInputData(writer, buffer);
+ } finally {
+ assertTrue(buffer.isRecycled());
+ }
+ });
+ }
+
+ @Test
+ public void testBuffersRecycledOnClose() throws Exception {
+ NetworkBuffer buffer = getBuffer();
+ runWithSyncWorker(writer -> {
+ callStart(writer);
+ callAddInputData(writer, buffer);
+ assertFalse(buffer.isRecycled());
+ });
+ assertTrue(buffer.isRecycled());
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testNoAddDataAfterFinished() throws Exception {
+ unwrappingError(IllegalArgumentException.class, () -> runWithSyncWorker(
+ writer -> {
+ callStart(writer);
+ callFinish(writer);
+ callAddInputData(writer);
+ }
+ ));
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testAddDataNotStarted() throws Exception {
+ unwrappingError(IllegalArgumentException.class, () -> runWithSyncWorker(writer -> callAddInputData(writer)));
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testFinishNotStarted() throws Exception {
+ unwrappingError(IllegalArgumentException.class, () -> runWithSyncWorker(this::callFinish));
+ }
+
+ @Test(expected = IllegalArgumentException.class)
+ public void testRethrowOnClose() throws Exception {
+ unwrappingError(IllegalArgumentException.class, () -> runWithSyncWorker(
+ writer -> {
+ try {
+ callFinish(writer);
+ } catch (IllegalArgumentException e) {
+ // ignore here - should rethrow in close
+ }
+ }
+ ));
+ }
+
+ @Test(expected = TestException.class)
+ public void testRethrowOnNextCall() throws Exception {
+ SyncChannelStateWriteRequestExecutor worker = new SyncChannelStateWriteRequestExecutor();
+ ChannelStateWriterImpl writer = new ChannelStateWriterImpl(new ConcurrentHashMap<>(), worker, 5);
+ writer.open();
+ worker.setThrown(new TestException());
+ unwrappingError(TestException.class, () -> callStart(writer));
+ }
+
+ @Test(expected = IllegalStateException.class)
+ public void testLimit() throws IOException {
+ int maxCheckpoints = 3;
+ try (ChannelStateWriterImpl writer = new ChannelStateWriterImpl(getStreamFactoryFactory(), maxCheckpoints)) {
+ writer.open();
+ for (int i = 0; i < maxCheckpoints; i++) {
+ writer.start(i, CheckpointOptions.forCheckpointWithDefaultLocation());
+ }
+ writer.start(maxCheckpoints, CheckpointOptions.forCheckpointWithDefaultLocation());
+ }
+ }
+
+ @Test(expected = IllegalStateException.class)
+ public void testStartNotOpened() throws Exception {
+ unwrappingError(IllegalStateException.class, () -> {
+ try (ChannelStateWriterImpl writer = new ChannelStateWriterImpl(getStreamFactoryFactory())) {
+ callStart(writer);
+ }
+ });
+ }
+
+ @Test(expected = IllegalStateException.class)
+ public void testNoStartAfterClose() throws Exception {
+ unwrappingError(IllegalStateException.class, () -> {
+ ChannelStateWriterImpl writer = openWriter();
+ writer.close();
+ writer.start(42, CheckpointOptions.forCheckpointWithDefaultLocation());
+ });
+ }
+
+ @Test(expected = IllegalStateException.class)
+ public void testNoAddDataAfterClose() throws Exception {
+ unwrappingError(IllegalStateException.class, () -> {
+ ChannelStateWriterImpl writer = openWriter();
+ callStart(writer);
+ writer.close();
+ callAddInputData(writer);
+ });
+ }
+
+ private static <T extends Throwable> void unwrappingError(Class<T> clazz, RunnableWithException r) throws Exception {
+ try {
+ r.run();
+ } catch (Exception e) {
+ throw findThrowable(e, clazz).map(te -> (Exception) te).orElse(e);
+ }
+ }
+
+ private NetworkBuffer getBuffer() {
+ return new NetworkBuffer(HeapMemorySegment.FACTORY.allocateUnpooledSegment(123, null), FreeingBufferRecycler.INSTANCE);
+ }
+
+ private ChannelStateWriteRequestExecutor failingWorker() {
+ return new ChannelStateWriteRequestExecutor() {
+ @Override
+ public void close() {
+ }
+
+ @Override
+ public void submit(ChannelStateWriteRequest e) {
+ throw new TestException();
+ }
+
+ @Override
+ public void submitPriority(ChannelStateWriteRequest e) {
+ throw new TestException();
+ }
+
+ @Override
+ public void start() throws IllegalStateException {
+ }
+ };
+ }
+
+ private void runWithSyncWorker(Consumer<ChannelStateWriter> writerConsumer) throws Exception {
+ runWithSyncWorker((channelStateWriter, syncChannelStateWriterWorker) -> writerConsumer.accept(channelStateWriter));
+ }
+
+ private void runWithSyncWorker(BiConsumerWithException<ChannelStateWriter, SyncChannelStateWriteRequestExecutor, Exception> testFn) throws Exception {
+ try (
+ SyncChannelStateWriteRequestExecutor worker = new SyncChannelStateWriteRequestExecutor();
+ ChannelStateWriterImpl writer = new ChannelStateWriterImpl(new ConcurrentHashMap<>(), worker, 5)
+ ) {
+ writer.open();
+ testFn.accept(writer, worker);
+ worker.processAllRequests();
+ }
+ }
+
+ private ChannelStateWriterImpl openWriter() {
+ ChannelStateWriterImpl writer = new ChannelStateWriterImpl(getStreamFactoryFactory());
+ writer.open();
+ return writer;
+ }
+
+ private void callStart(ChannelStateWriter writer) {
+ writer.start(CHECKPOINT_ID, CheckpointOptions.forCheckpointWithDefaultLocation());
+ }
+
+ private void callAddInputData(ChannelStateWriter writer, NetworkBuffer... buffer) {
+ writer.addInputData(CHECKPOINT_ID, new InputChannelInfo(1, 1), 1, buffer);
+ }
+
+ private void callAbort(ChannelStateWriter writer) {
+ writer.abort(CHECKPOINT_ID, new TestException());
+ }
+
+ private void callFinish(ChannelStateWriter writer) {
+ writer.finishInput(CHECKPOINT_ID);
+ writer.finishOutput(CHECKPOINT_ID);
+ }
+
+}
+
+class TestException extends RuntimeException {
+}
+
+class SyncChannelStateWriteRequestExecutor implements ChannelStateWriteRequestExecutor {
+ private final ChannelStateWriteRequestDispatcher requestProcessor;
+ private final Deque<ChannelStateWriteRequest> deque;
+ private Exception thrown;
+
+ SyncChannelStateWriteRequestExecutor() {
+ deque = new ArrayDeque<>();
+ requestProcessor = new ChannelStateWriteRequestDispatcherImpl(getStreamFactoryFactory(), new ChannelStateSerializerImpl());
+ }
+
+ @Override
+ public void submit(ChannelStateWriteRequest e) throws Exception {
+ deque.offer(e);
+ if (thrown != null) {
+ throw thrown;
+ }
+ }
+
+ @Override
+ public void submitPriority(ChannelStateWriteRequest e) throws Exception {
+ deque.offerFirst(e);
+ if (thrown != null) {
+ throw thrown;
+ }
+ }
+
+ @Override
+ public void start() throws IllegalStateException {
+ }
+
+ @Override
+ public void close() {
+ }
+
+ void processAllRequests() throws Exception {
+ while (!deque.isEmpty()) {
+ requestProcessor.dispatch(deque.poll());
+ }
+ }
+
+ public void setThrown(Exception thrown) {
+ this.thrown = thrown;
+ }
+}
+
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/CheckpointInProgressRequestTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/CheckpointInProgressRequestTest.java
new file mode 100644
index 0000000..3617b8f
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/CheckpointInProgressRequestTest.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.checkpoint.channel;
+
+import org.junit.Test;
+
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * {@link CheckpointInProgressRequest} test.
+ */
+public class CheckpointInProgressRequestTest {
+
+ /**
+ * Tests that a request can only be cancelled once. This is important for requests to write data to prevent double
+ * recycling of their buffers.
+ */
+ @Test
+ public void testNoCancelTwice() throws Exception {
+ AtomicInteger counter = new AtomicInteger();
+ CyclicBarrier barrier = new CyclicBarrier(10);
+ CheckpointInProgressRequest request = cancelCountingRequest(counter, barrier);
+ Thread[] threads = new Thread[barrier.getParties()];
+ for (int i = 0; i < barrier.getParties(); i++) {
+ threads[i] = new Thread(() -> {
+ request.cancel(new RuntimeException("test"));
+ await(barrier);
+ });
+ }
+ for (int i = 0; i < barrier.getParties(); i++) {
+ threads[i].start();
+ threads[i].join();
+ }
+
+ assertEquals(1, counter.get());
+ }
+
+ private CheckpointInProgressRequest cancelCountingRequest(AtomicInteger cancelCounter, CyclicBarrier cb) {
+ return new CheckpointInProgressRequest(
+ "test",
+ 1L,
+ unused -> {
+ },
+ unused -> {
+ cancelCounter.incrementAndGet();
+ await(cb);
+ },
+ false
+ );
+ }
+
+ private void await(CyclicBarrier cb) {
+ if (cb.getNumberWaiting() == 0) {
+ // skip waiting if waited inside cancel
+ return;
+ }
+ try {
+ cb.await();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/BufferBuilderAndConsumerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/BufferBuilderAndConsumerTest.java
index 1033c5e..c7dfe3e 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/BufferBuilderAndConsumerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/BufferBuilderAndConsumerTest.java
@@ -206,6 +206,25 @@ public class BufferBuilderAndConsumerTest {
testIsFinished(BUFFER_INT_SIZE);
}
+ @Test
+ public void testWritableBytes() {
+ BufferBuilder bufferBuilder = createBufferBuilder();
+ assertEquals(bufferBuilder.getMaxCapacity(), bufferBuilder.getWritableBytes());
+
+ ByteBuffer byteBuffer = toByteBuffer(1, 2, 3);
+ bufferBuilder.append(byteBuffer);
+ assertEquals(bufferBuilder.getMaxCapacity() - byteBuffer.position(), bufferBuilder.getWritableBytes());
+
+ assertEquals(bufferBuilder.getMaxCapacity() - byteBuffer.position(), bufferBuilder.getWritableBytes());
+ }
+
+ @Test
+ public void testWritableBytesWhenFull() {
+ BufferBuilder bufferBuilder = createBufferBuilder();
+ bufferBuilder.append(toByteBuffer(new int[bufferBuilder.getMaxCapacity()]));
+ assertEquals(0, bufferBuilder.getWritableBytes());
+ }
+
private static void testIsFinished(int writes) {
BufferBuilder bufferBuilder = createBufferBuilder();
BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/state/ChannelPersistenceITCase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ChannelPersistenceITCase.java
new file mode 100644
index 0000000..3f5e2cc
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/state/ChannelPersistenceITCase.java
@@ -0,0 +1,179 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.state;
+
+import org.apache.flink.core.memory.HeapMemorySegment;
+import org.apache.flink.core.memory.MemorySegment;
+import org.apache.flink.runtime.checkpoint.CheckpointOptions;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.runtime.checkpoint.StateObjectCollection;
+import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateReader;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateReader.ReadResult;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateReaderImpl;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter.ChannelStateWriteResult;
+import org.apache.flink.runtime.checkpoint.channel.ChannelStateWriterImpl;
+import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
+import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
+import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.jobgraph.OperatorID;
+import org.apache.flink.runtime.state.memory.NonPersistentMetadataCheckpointStorageLocation;
+import org.apache.flink.util.function.BiFunctionWithException;
+
+import org.junit.Test;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.stream.Collectors;
+
+import static java.util.Collections.singletonMap;
+import static org.apache.flink.runtime.checkpoint.CheckpointType.CHECKPOINT;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateReader.ReadResult.NO_MORE_DATA;
+import static org.apache.flink.runtime.checkpoint.channel.ChannelStateWriter.SEQUENCE_NUMBER_UNKNOWN;
+import static org.apache.flink.util.Preconditions.checkState;
+import static org.junit.Assert.assertArrayEquals;
+
+/**
+ * ChannelPersistenceITCase.
+ */
+public class ChannelPersistenceITCase {
+ private static final Random RANDOM = new Random(System.currentTimeMillis());
+
+ @Test
+ public void testReadWritten() throws Exception {
+ long checkpointId = 1L;
+
+ InputChannelInfo inputChannelInfo = new InputChannelInfo(2, 3);
+ byte[] inputChannelInfoData = randomBytes(1024);
+
+ ResultSubpartitionInfo resultSubpartitionInfo = new ResultSubpartitionInfo(4, 5);
+ byte[] resultSubpartitionInfoData = randomBytes(1024);
+
+ ChannelStateWriteResult handles = write(
+ checkpointId,
+ singletonMap(inputChannelInfo, inputChannelInfoData),
+ singletonMap(resultSubpartitionInfo, resultSubpartitionInfoData)
+ );
+
+ assertArrayEquals(inputChannelInfoData, read(
+ toTaskStateSnapshot(handles),
+ inputChannelInfoData.length,
+ (reader, mem) -> reader.readInputData(inputChannelInfo, new NetworkBuffer(mem, FreeingBufferRecycler.INSTANCE))
+ ));
+
+ assertArrayEquals(resultSubpartitionInfoData, read(
+ toTaskStateSnapshot(handles),
+ resultSubpartitionInfoData.length,
+ (reader, mem) -> reader.readOutputData(resultSubpartitionInfo, new BufferBuilder(mem, FreeingBufferRecycler.INSTANCE))
+ ));
+ }
+
+ private byte[] randomBytes(int size) {
+ byte[] bytes = new byte[size];
+ RANDOM.nextBytes(bytes);
+ return bytes;
+ }
+
+ private ChannelStateWriteResult write(long checkpointId, Map<InputChannelInfo, byte[]> icMap, Map<ResultSubpartitionInfo, byte[]> rsMap) throws Exception {
+ int maxStateSize = sizeOfBytes(icMap) + sizeOfBytes(rsMap) + Long.BYTES * 2;
+ Map<InputChannelInfo, Buffer> icBuffers = wrapWithBuffers(icMap);
+ Map<ResultSubpartitionInfo, Buffer> rsBuffers = wrapWithBuffers(rsMap);
+ try (ChannelStateWriterImpl writer = new ChannelStateWriterImpl(getStreamFactoryFactory(maxStateSize))) {
+ writer.open();
+ writer.start(checkpointId, new CheckpointOptions(CHECKPOINT, new CheckpointStorageLocationReference("poly".getBytes())));
+ for (Map.Entry<InputChannelInfo, Buffer> e : icBuffers.entrySet()) {
+ writer.addInputData(checkpointId, e.getKey(), SEQUENCE_NUMBER_UNKNOWN, e.getValue());
+ }
+ writer.finishInput(checkpointId);
+ for (Map.Entry<ResultSubpartitionInfo, Buffer> e : rsBuffers.entrySet()) {
+ writer.addOutputData(checkpointId, e.getKey(), SEQUENCE_NUMBER_UNKNOWN, e.getValue());
+ }
+ writer.finishOutput(checkpointId);
+ ChannelStateWriteResult result = writer.getWriteResult(checkpointId);
+ result.getResultSubpartitionStateHandles().join(); // prevent abnormal complete in close
+ return result;
+ }
+ }
+
+ public static CheckpointStorageWorkerView getStreamFactoryFactory() {
+ return getStreamFactoryFactory(42);
+ }
+
+ public static CheckpointStorageWorkerView getStreamFactoryFactory(int maxStateSize) {
+ return new CheckpointStorageWorkerView() {
+ @Override
+ public CheckpointStreamFactory resolveCheckpointStorageLocation(long checkpointId, CheckpointStorageLocationReference reference) {
+ return new NonPersistentMetadataCheckpointStorageLocation(maxStateSize);
+ }
+
+ @Override
+ public CheckpointStreamFactory.CheckpointStateOutputStream createTaskOwnedStateStream() {
+ throw new UnsupportedOperationException();
+ }
+ };
+ }
+
+ private byte[] read(TaskStateSnapshot taskStateSnapshot, int size, BiFunctionWithException<ChannelStateReader, MemorySegment, ReadResult, Exception> readFn) throws Exception {
+ byte[] dst = new byte[size];
+ HeapMemorySegment mem = HeapMemorySegment.FACTORY.wrap(dst);
+ try {
+ checkState(NO_MORE_DATA == readFn.apply(new ChannelStateReaderImpl(taskStateSnapshot), mem));
+ } finally {
+ mem.free();
+ }
+ return dst;
+ }
+
+ private TaskStateSnapshot toTaskStateSnapshot(ChannelStateWriteResult t) throws Exception {
+ return new TaskStateSnapshot(singletonMap(new OperatorID(),
+ new OperatorSubtaskState(
+ StateObjectCollection.empty(),
+ StateObjectCollection.empty(),
+ StateObjectCollection.empty(),
+ StateObjectCollection.empty(),
+ new StateObjectCollection<>(t.getInputChannelStateHandles().get()),
+ new StateObjectCollection<>(t.getResultSubpartitionStateHandles().get())
+ )
+ ));
+ }
+
+ @SuppressWarnings("unchecked")
+ private <C> List<C> collect(Collection<StateObject> handles, Class<C> clazz) {
+ return handles.stream().filter(clazz::isInstance).map(h -> (C) h).collect(Collectors.toList());
+ }
+
+ private static int sizeOfBytes(Map<?, byte[]> map) {
+ return map.values().stream().mapToInt(d -> d.length).sum();
+ }
+
+ private <K> Map<K, Buffer> wrapWithBuffers(Map<K, byte[]> icMap) {
+ return icMap.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> wrapWithBuffer(e.getValue())));
+ }
+
+ private static Buffer wrapWithBuffer(byte[] data) {
+ NetworkBuffer buffer = new NetworkBuffer(HeapMemorySegment.FACTORY.allocateUnpooledSegment(data.length, null), FreeingBufferRecycler.INSTANCE);
+ buffer.writeBytes(data);
+ return buffer;
+ }
+
+}