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 2021/06/07 09:01:53 UTC
[flink] branch release-1.13 updated (b73f6c4 -> 6a2f806)
This is an automated email from the ASF dual-hosted git repository.
pnowojski pushed a change to branch release-1.13
in repository https://gitbox.apache.org/repos/asf/flink.git.
from b73f6c4 [FLINK-22855][docs-zh] Translate the 'Overview of Python API' page into Chinese.
new 8511ecb [FLINK-22376][runtime] BufferProvider is able to provide pure MemorySegment in order to avoid extracting it from the bufferBuilder
new 0bf1fa2 [FLINK-22376][runtime] Buffer is used inside of BufferBuilder for references counting
new 79a7b6d [FLINK-22376][runtime] RecoveredChannelStateHandler requires the ownership of the input Buffer
new 6a2f806 [hotfix] Clarification ownerships of the buffer for Remote InputChannel#onBuffer in javadoc
The 4 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../InflightDataRescalingDescriptor.java | 3 +-
.../checkpoint/channel/ChannelStateSerializer.java | 13 +--
.../channel/RecoveredChannelStateHandler.java | 83 ++++++++-------
.../channel/SequentialChannelStateReaderImpl.java | 18 ++--
.../network/api/serialization/EventSerializer.java | 4 +-
.../runtime/io/network/buffer/BufferBuilder.java | 27 ++---
.../runtime/io/network/buffer/BufferConsumer.java | 32 +-----
.../runtime/io/network/buffer/BufferProvider.java | 18 ++++
.../runtime/io/network/buffer/LocalBufferPool.java | 11 +-
.../partition/BufferWritingResultPartition.java | 2 +
.../network/partition/PartitionSortedBuffer.java | 28 +++---
.../partition/SortMergeResultPartition.java | 25 +++--
.../partition/consumer/RemoteInputChannel.java | 4 +
.../channel/ChannelStateChunkReaderTest.java | 23 ++++-
.../InputChannelRecoveredStateHandlerTest.java | 111 +++++++++++++++++++++
.../channel/RecoveredChannelStateHandlerTest.java | 17 ++--
...esultSubpartitionRecoveredStateHandlerTest.java | 107 ++++++++++++++++++++
.../SpanningRecordSerializationTest.java | 17 ++--
.../io/network/api/writer/RecordWriterTest.java | 13 +--
.../buffer/BufferBuilderAndConsumerTest.java | 95 ++++++++++++++----
.../io/network/buffer/BufferBuilderTestUtils.java | 23 +++--
.../io/network/buffer/LocalBufferPoolTest.java | 89 ++++++++---------
.../io/network/buffer/NetworkBufferPoolTest.java | 2 +-
.../runtime/io/network/buffer/NoOpBufferPool.java | 10 ++
.../io/network/buffer/UnpooledBufferPool.java | 8 +-
.../BoundedBlockingSubpartitionWriteReadTest.java | 5 +-
.../partition/consumer/LocalInputChannelTest.java | 11 +-
.../partition/consumer/RemoteInputChannelTest.java | 9 +-
.../runtime/io/network/util/TestBufferFactory.java | 11 +-
.../io/network/util/TestPooledBufferProvider.java | 82 +++++++--------
.../io/network/util/TestSubpartitionProducer.java | 5 +-
.../consumer/StreamTestSingleInputGate.java | 1 +
.../runtime/io/StreamTaskNetworkInputTest.java | 12 ++-
.../DemultiplexingRecordDeserializerTest.java | 64 ++++++------
34 files changed, 671 insertions(+), 312 deletions(-)
create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/InputChannelRecoveredStateHandlerTest.java
copy flink-table/flink-table-common/src/test/java/org/apache/flink/table/module/CoreModuleTest.java => flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandlerTest.java (65%)
create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ResultSubpartitionRecoveredStateHandlerTest.java
[flink] 03/04: [FLINK-22376][runtime] RecoveredChannelStateHandler
requires the ownership of the input Buffer
Posted by pn...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
pnowojski pushed a commit to branch release-1.13
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 79a7b6db86e68be4aa265e52210046d3265149ad
Author: Anton Kalashnikov <ka...@yandex.ru>
AuthorDate: Tue Jun 1 17:08:47 2021 +0300
[FLINK-22376][runtime] RecoveredChannelStateHandler requires the ownership of the input Buffer
(cherry picked from commit 2293362f0a144f47e182039c26c0d7047f87232a)
---
.../InflightDataRescalingDescriptor.java | 3 +-
.../channel/RecoveredChannelStateHandler.java | 83 ++++++++-------
.../channel/SequentialChannelStateReaderImpl.java | 18 ++--
.../channel/ChannelStateChunkReaderTest.java | 23 ++++-
.../InputChannelRecoveredStateHandlerTest.java | 111 +++++++++++++++++++++
.../channel/RecoveredChannelStateHandlerTest.java | 34 +++++++
...esultSubpartitionRecoveredStateHandlerTest.java | 107 ++++++++++++++++++++
7 files changed, 330 insertions(+), 49 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java
index 3a1ab27..53d4792 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java
@@ -128,7 +128,8 @@ public class InflightDataRescalingDescriptor implements Serializable {
private final MappingType mappingType;
- enum MappingType {
+ /** Type of mapping which should be used for this in-flight data. */
+ public enum MappingType {
IDENTITY,
RESCALING
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java
index 86894c9..4a724a7 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandler.java
@@ -17,7 +17,6 @@
package org.apache.flink.runtime.checkpoint.channel;
-import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptor;
import org.apache.flink.runtime.checkpoint.RescaleMappings;
import org.apache.flink.runtime.io.network.api.SubtaskConnectionDescriptor;
@@ -51,11 +50,21 @@ interface RecoveredChannelStateHandler<Info, Context> extends AutoCloseable {
this.buffer = buffer;
this.context = context;
}
+
+ public void close() {
+ buffer.close();
+ }
}
BufferWithContext<Context> getBuffer(Info info) throws IOException, InterruptedException;
- void recover(Info info, int oldSubtaskIndex, Context context) throws IOException;
+ /**
+ * Recover the data from buffer. This method is taking over the ownership of the
+ * bufferWithContext and is fully responsible for cleaning it up both on the happy path and in
+ * case of an error.
+ */
+ void recover(Info info, int oldSubtaskIndex, BufferWithContext<Context> bufferWithContext)
+ throws IOException;
}
class InputChannelRecoveredStateHandler
@@ -84,8 +93,12 @@ class InputChannelRecoveredStateHandler
}
@Override
- public void recover(InputChannelInfo channelInfo, int oldSubtaskIndex, Buffer buffer)
+ public void recover(
+ InputChannelInfo channelInfo,
+ int oldSubtaskIndex,
+ BufferWithContext<Buffer> bufferWithContext)
throws IOException {
+ Buffer buffer = bufferWithContext.context;
try {
if (buffer.readableBytes() > 0) {
for (final RecoveredInputChannel channel : getMappedChannels(channelInfo)) {
@@ -143,8 +156,7 @@ class InputChannelRecoveredStateHandler
}
class ResultSubpartitionRecoveredStateHandler
- implements RecoveredChannelStateHandler<
- ResultSubpartitionInfo, Tuple2<BufferBuilder, BufferConsumer>> {
+ implements RecoveredChannelStateHandler<ResultSubpartitionInfo, BufferBuilder> {
private final ResultPartitionWriter[] writers;
private final boolean notifyAndBlockOnCompletion;
@@ -165,51 +177,48 @@ class ResultSubpartitionRecoveredStateHandler
}
@Override
- public BufferWithContext<Tuple2<BufferBuilder, BufferConsumer>> getBuffer(
- ResultSubpartitionInfo subpartitionInfo) throws IOException, InterruptedException {
+ public BufferWithContext<BufferBuilder> getBuffer(ResultSubpartitionInfo subpartitionInfo)
+ throws IOException, InterruptedException {
// request the buffer from any mapped subpartition as they all will receive the same buffer
final List<CheckpointedResultSubpartition> channels = getMappedChannels(subpartitionInfo);
BufferBuilder bufferBuilder = channels.get(0).requestBufferBuilderBlocking();
- return new BufferWithContext<>(
- wrap(bufferBuilder),
- Tuple2.of(bufferBuilder, bufferBuilder.createBufferConsumer()));
+ return new BufferWithContext<>(wrap(bufferBuilder), bufferBuilder);
}
@Override
public void recover(
ResultSubpartitionInfo subpartitionInfo,
int oldSubtaskIndex,
- Tuple2<BufferBuilder, BufferConsumer> bufferBuilderAndConsumer)
+ BufferWithContext<BufferBuilder> bufferWithContext)
throws IOException {
- try {
- bufferBuilderAndConsumer.f0.finish();
- if (bufferBuilderAndConsumer.f1.isDataAvailable()) {
- NetworkActionsLogger.traceRecover(
- "ResultSubpartitionRecoveredStateHandler#recover",
- bufferBuilderAndConsumer.f1,
- subpartitionInfo);
- final List<CheckpointedResultSubpartition> channels =
- getMappedChannels(subpartitionInfo);
- for (final CheckpointedResultSubpartition channel : channels) {
- // channel selector is created from the downstream's point of view: the subtask
- // of
- // downstream = subpartition index of recovered buffer
- final SubtaskConnectionDescriptor channelSelector =
- new SubtaskConnectionDescriptor(
- subpartitionInfo.getSubPartitionIdx(), oldSubtaskIndex);
- channel.add(
- EventSerializer.toBufferConsumer(channelSelector, false),
- Integer.MIN_VALUE);
- boolean added =
- channel.add(bufferBuilderAndConsumer.f1.copy(), Integer.MIN_VALUE);
- if (!added) {
- throw new IOException(
- "Buffer consumer couldn't be added to ResultSubpartition");
+ try (BufferBuilder bufferBuilder = bufferWithContext.context) {
+ try (BufferConsumer bufferConsumer =
+ bufferBuilder.createBufferConsumerFromBeginning()) {
+ bufferBuilder.finish();
+ if (bufferConsumer.isDataAvailable()) {
+ NetworkActionsLogger.traceRecover(
+ "ResultSubpartitionRecoveredStateHandler#recover",
+ bufferConsumer,
+ subpartitionInfo);
+ final List<CheckpointedResultSubpartition> channels =
+ getMappedChannels(subpartitionInfo);
+ for (final CheckpointedResultSubpartition channel : channels) {
+ // channel selector is created from the downstream's point of view: the
+ // subtask of downstream = subpartition index of recovered buffer
+ final SubtaskConnectionDescriptor channelSelector =
+ new SubtaskConnectionDescriptor(
+ subpartitionInfo.getSubPartitionIdx(), oldSubtaskIndex);
+ channel.add(
+ EventSerializer.toBufferConsumer(channelSelector, false),
+ Integer.MIN_VALUE);
+ boolean added = channel.add(bufferConsumer.copy(), Integer.MIN_VALUE);
+ if (!added) {
+ throw new IOException(
+ "Buffer consumer couldn't be added to ResultSubpartition");
+ }
}
}
}
- } finally {
- bufferBuilderAndConsumer.f1.close();
}
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java
index 8329c71..88b23c3 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/channel/SequentialChannelStateReaderImpl.java
@@ -196,19 +196,19 @@ class ChannelStateChunkReader {
while (length > 0) {
RecoveredChannelStateHandler.BufferWithContext<Context> bufferWithContext =
stateHandler.getBuffer(channelInfo);
- try {
- try (Closeable ignored =
- NetworkActionsLogger.measureIO(
- "ChannelStateChunkReader#readChunk", bufferWithContext.buffer)) {
- while (length > 0 && bufferWithContext.buffer.isWritable()) {
- length -= serializer.readData(source, bufferWithContext.buffer, length);
- }
+ try (Closeable ignored =
+ NetworkActionsLogger.measureIO(
+ "ChannelStateChunkReader#readChunk", bufferWithContext.buffer)) {
+ while (length > 0 && bufferWithContext.buffer.isWritable()) {
+ length -= serializer.readData(source, bufferWithContext.buffer, length);
}
- stateHandler.recover(channelInfo, oldSubtaskIndex, bufferWithContext.context);
} catch (Exception e) {
- bufferWithContext.buffer.recycle();
+ bufferWithContext.close();
throw e;
}
+
+ // Passing the ownership of buffer to inside.
+ stateHandler.recover(channelInfo, oldSubtaskIndex, bufferWithContext);
}
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateChunkReaderTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateChunkReaderTest.java
index c163d3d..0e88489 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateChunkReaderTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ChannelStateChunkReaderTest.java
@@ -59,6 +59,22 @@ public class ChannelStateChunkReaderTest {
}
@Test
+ public void testBufferRecycledOnSuccess() throws IOException, InterruptedException {
+ ChannelStateSerializer serializer = new ChannelStateSerializerImpl();
+ TestRecoveredChannelStateHandler handler = new TestRecoveredChannelStateHandler();
+
+ try (FSDataInputStream stream = getStream(serializer, 10)) {
+ new ChannelStateChunkReader(serializer)
+ .readChunk(stream, serializer.getHeaderLength(), handler, "channelInfo", 0);
+ } finally {
+ checkState(!handler.requestedBuffers.isEmpty());
+ assertTrue(
+ handler.requestedBuffers.stream()
+ .allMatch(TestChannelStateByteBuffer::isRecycled));
+ }
+ }
+
+ @Test
public void testBuffersNotRequestedForEmptyStream() throws IOException, InterruptedException {
ChannelStateSerializer serializer = new ChannelStateSerializerImpl();
TestRecoveredChannelStateHandler handler = new TestRecoveredChannelStateHandler();
@@ -109,7 +125,10 @@ public class ChannelStateChunkReaderTest {
}
@Override
- public void recover(Object o, int oldSubtaskIndex, Object o2) {}
+ public void recover(
+ Object o, int oldSubtaskIndex, BufferWithContext<Object> bufferWithContext) {
+ bufferWithContext.close();
+ }
@Override
public void close() throws Exception {}
@@ -151,7 +170,7 @@ public class ChannelStateChunkReaderTest {
}
@Override
- public void recycle() {
+ public void close() {
checkArgument(!recycled);
recycled = true;
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/InputChannelRecoveredStateHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/InputChannelRecoveredStateHandlerTest.java
new file mode 100644
index 0000000..d499b90
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/InputChannelRecoveredStateHandlerTest.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.InflightDataRescalingDescriptor;
+import org.apache.flink.runtime.checkpoint.RescaleMappings;
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.consumer.InputChannelBuilder;
+import org.apache.flink.runtime.io.network.partition.consumer.InputGate;
+import org.apache.flink.runtime.io.network.partition.consumer.SingleInputGate;
+import org.apache.flink.runtime.io.network.partition.consumer.SingleInputGateBuilder;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.HashSet;
+
+import static org.junit.Assert.assertEquals;
+
+/** Test of different implementation of {@link InputChannelRecoveredStateHandler}. */
+public class InputChannelRecoveredStateHandlerTest extends RecoveredChannelStateHandlerTest {
+ private static final int preAllocatedSegments = 3;
+ private NetworkBufferPool networkBufferPool;
+ private SingleInputGate inputGate;
+ private InputChannelRecoveredStateHandler icsHander;
+ private InputChannelInfo channelInfo;
+
+ @Before
+ public void setUp() {
+ // given: Segment provider with defined number of allocated segments.
+ networkBufferPool = new NetworkBufferPool(preAllocatedSegments, 1024);
+
+ // and: Configured input gate with recovered channel.
+ inputGate =
+ new SingleInputGateBuilder()
+ .setChannelFactory(InputChannelBuilder::buildLocalRecoveredChannel)
+ .setSegmentProvider(networkBufferPool)
+ .build();
+
+ icsHander = buildInputChannelStateHandler(inputGate);
+
+ channelInfo = new InputChannelInfo(0, 0);
+ }
+
+ private InputChannelRecoveredStateHandler buildInputChannelStateHandler(
+ SingleInputGate inputGate) {
+ return new InputChannelRecoveredStateHandler(
+ new InputGate[] {inputGate},
+ new InflightDataRescalingDescriptor(
+ new InflightDataRescalingDescriptor
+ .InflightDataGateOrPartitionRescalingDescriptor[] {
+ new InflightDataRescalingDescriptor
+ .InflightDataGateOrPartitionRescalingDescriptor(
+ new int[] {1},
+ RescaleMappings.identity(1, 1),
+ new HashSet<>(),
+ InflightDataRescalingDescriptor
+ .InflightDataGateOrPartitionRescalingDescriptor
+ .MappingType.IDENTITY)
+ }));
+ }
+
+ @Test
+ public void testRecycleBufferBeforeRecoverWasCalled() throws Exception {
+ // when: Request the buffer.
+ RecoveredChannelStateHandler.BufferWithContext<Buffer> bufferWithContext =
+ icsHander.getBuffer(channelInfo);
+
+ // and: Recycle buffer outside.
+ bufferWithContext.buffer.close();
+
+ // Close the gate for flushing the cached recycled buffers to the segment provider.
+ inputGate.close();
+
+ // then: All pre-allocated segments should be successfully recycled.
+ assertEquals(preAllocatedSegments, networkBufferPool.getNumberOfAvailableMemorySegments());
+ }
+
+ @Test
+ public void testRecycleBufferAfterRecoverWasCalled() throws Exception {
+ // when: Request the buffer.
+ RecoveredChannelStateHandler.BufferWithContext<Buffer> bufferWithContext =
+ icsHander.getBuffer(channelInfo);
+
+ // and: Recycle buffer outside.
+ icsHander.recover(channelInfo, 0, bufferWithContext);
+
+ // Close the gate for flushing the cached recycled buffers to the segment provider.
+ inputGate.close();
+
+ // then: All pre-allocated segments should be successfully recycled.
+ assertEquals(preAllocatedSegments, networkBufferPool.getNumberOfAvailableMemorySegments());
+ }
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandlerTest.java
new file mode 100644
index 0000000..a74c128
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/RecoveredChannelStateHandlerTest.java
@@ -0,0 +1,34 @@
+/*
+ * 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;
+
+/**
+ * Base class which contains all tests which should be implemented for every implementation of
+ * {@link InputChannelRecoveredStateHandler}.
+ */
+public abstract class RecoveredChannelStateHandlerTest {
+
+ @Test
+ public abstract void testRecycleBufferBeforeRecoverWasCalled() throws Exception;
+
+ @Test
+ public abstract void testRecycleBufferAfterRecoverWasCalled() throws Exception;
+}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ResultSubpartitionRecoveredStateHandlerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ResultSubpartitionRecoveredStateHandlerTest.java
new file mode 100644
index 0000000..b188927
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/channel/ResultSubpartitionRecoveredStateHandlerTest.java
@@ -0,0 +1,107 @@
+/*
+ * 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.InflightDataRescalingDescriptor;
+import org.apache.flink.runtime.checkpoint.RescaleMappings;
+import org.apache.flink.runtime.io.network.api.writer.ResultPartitionWriter;
+import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.ResultPartition;
+import org.apache.flink.runtime.io.network.partition.ResultPartitionBuilder;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.HashSet;
+
+import static org.junit.Assert.assertEquals;
+
+/** Test of different implementation of {@link ResultSubpartitionRecoveredStateHandler}. */
+public class ResultSubpartitionRecoveredStateHandlerTest extends RecoveredChannelStateHandlerTest {
+ private static final int preAllocatedSegments = 3;
+ private NetworkBufferPool networkBufferPool;
+ private ResultPartition partition;
+ private ResultSubpartitionRecoveredStateHandler rstHandler;
+ private ResultSubpartitionInfo channelInfo;
+
+ @Before
+ public void setUp() throws IOException {
+ // given: Segment provider with defined number of allocated segments.
+ channelInfo = new ResultSubpartitionInfo(0, 0);
+
+ networkBufferPool = new NetworkBufferPool(preAllocatedSegments, 1024);
+ partition = new ResultPartitionBuilder().setNetworkBufferPool(networkBufferPool).build();
+ partition.setup();
+
+ rstHandler = buildResultStateHandler(partition);
+ }
+
+ private ResultSubpartitionRecoveredStateHandler buildResultStateHandler(
+ ResultPartition partition) {
+ return new ResultSubpartitionRecoveredStateHandler(
+ new ResultPartitionWriter[] {partition},
+ false,
+ new InflightDataRescalingDescriptor(
+ new InflightDataRescalingDescriptor
+ .InflightDataGateOrPartitionRescalingDescriptor[] {
+ new InflightDataRescalingDescriptor
+ .InflightDataGateOrPartitionRescalingDescriptor(
+ new int[] {1},
+ RescaleMappings.identity(1, 1),
+ new HashSet<>(),
+ InflightDataRescalingDescriptor
+ .InflightDataGateOrPartitionRescalingDescriptor
+ .MappingType.IDENTITY)
+ }));
+ }
+
+ @Test
+ public void testRecycleBufferBeforeRecoverWasCalled() throws Exception {
+ // when: Request the buffer.
+ RecoveredChannelStateHandler.BufferWithContext<BufferBuilder> bufferWithContext =
+ rstHandler.getBuffer(new ResultSubpartitionInfo(0, 0));
+
+ // and: Recycle buffer outside.
+ bufferWithContext.buffer.close();
+
+ // Close the partition for flushing the cached recycled buffers to the segment provider.
+ partition.close();
+
+ // then: All pre-allocated segments should be successfully recycled.
+ assertEquals(preAllocatedSegments, networkBufferPool.getNumberOfAvailableMemorySegments());
+ }
+
+ @Test
+ public void testRecycleBufferAfterRecoverWasCalled() throws Exception {
+ // when: Request the buffer.
+ RecoveredChannelStateHandler.BufferWithContext<BufferBuilder> bufferWithContext =
+ rstHandler.getBuffer(channelInfo);
+
+ // and: Pass the buffer to recovery.
+ rstHandler.recover(channelInfo, 0, bufferWithContext);
+
+ // Close the partition for flushing the cached recycled buffers to the segment provider.
+ partition.close();
+
+ // then: All pre-allocated segments should be successfully recycled.
+ assertEquals(preAllocatedSegments, networkBufferPool.getNumberOfAvailableMemorySegments());
+ }
+}
[flink] 04/04: [hotfix] Clarification ownerships of the buffer for
Remote InputChannel#onBuffer in javadoc
Posted by pn...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
pnowojski pushed a commit to branch release-1.13
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 6a2f80605b1cc2545ac1400ae6a131fb7b5093b4
Author: Anton Kalashnikov <ka...@yandex.ru>
AuthorDate: Wed Jun 2 19:39:04 2021 +0300
[hotfix] Clarification ownerships of the buffer for Remote InputChannel#onBuffer in javadoc
(cherry picked from commit bc8c739320210432128bbb486259b70a7250de9c)
---
.../runtime/io/network/partition/consumer/RemoteInputChannel.java | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java
index 128fd18..0cd918f 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannel.java
@@ -443,6 +443,10 @@ public class RemoteInputChannel extends InputChannel {
}
}
+ /**
+ * Handles the input buffer. This method is taking over the ownership of the buffer and is fully
+ * responsible for cleaning it up both on the happy path and in case of an error.
+ */
public void onBuffer(Buffer buffer, int sequenceNumber, int backlog) throws IOException {
boolean recycleBuffer = true;
[flink] 02/04: [FLINK-22376][runtime] Buffer is used inside of
BufferBuilder for references counting
Posted by pn...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
pnowojski pushed a commit to branch release-1.13
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 0bf1fa2cab4e352011966ff191fbabdcef3daaed
Author: Anton Kalashnikov <ka...@yandex.ru>
AuthorDate: Fri May 28 11:50:12 2021 +0200
[FLINK-22376][runtime] Buffer is used inside of BufferBuilder for references counting
(cherry picked from commit 1051ea9c108d4c56623013ddfd2dd866df69eaa7)
---
.../checkpoint/channel/ChannelStateSerializer.java | 13 +--
.../network/api/serialization/EventSerializer.java | 4 +-
.../runtime/io/network/buffer/BufferBuilder.java | 21 ++---
.../runtime/io/network/buffer/BufferConsumer.java | 32 +-------
.../partition/BufferWritingResultPartition.java | 2 +
.../SpanningRecordSerializationTest.java | 17 ++--
.../io/network/api/writer/RecordWriterTest.java | 13 +--
.../buffer/BufferBuilderAndConsumerTest.java | 95 +++++++++++++++++-----
.../io/network/buffer/BufferBuilderTestUtils.java | 23 ++++--
.../io/network/buffer/LocalBufferPoolTest.java | 88 ++++++++++----------
.../io/network/buffer/NetworkBufferPoolTest.java | 2 +-
.../BoundedBlockingSubpartitionWriteReadTest.java | 5 +-
.../partition/consumer/LocalInputChannelTest.java | 11 ++-
.../partition/consumer/RemoteInputChannelTest.java | 9 +-
.../io/network/util/TestSubpartitionProducer.java | 5 +-
.../consumer/StreamTestSingleInputGate.java | 1 +
.../runtime/io/StreamTaskNetworkInputTest.java | 12 +--
.../DemultiplexingRecordDeserializerTest.java | 64 ++++++++-------
18 files changed, 242 insertions(+), 175 deletions(-)
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
index 2d21b6c..252d25c 100644
--- 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
@@ -57,11 +57,12 @@ interface ChannelStateSerializer {
/** Wrapper around various buffers to receive channel state data. */
@Internal
@NotThreadSafe
-interface ChannelStateByteBuffer {
+interface ChannelStateByteBuffer extends AutoCloseable {
boolean isWritable();
- void recycle();
+ @Override
+ void close();
/**
* Read up to <code>bytesToRead</code> bytes into this buffer from the given {@link
@@ -82,7 +83,7 @@ interface ChannelStateByteBuffer {
}
@Override
- public void recycle() {
+ public void close() {
buffer.recycleBuffer();
}
@@ -102,8 +103,8 @@ interface ChannelStateByteBuffer {
}
@Override
- public void recycle() {
- bufferBuilder.recycle();
+ public void close() {
+ bufferBuilder.close();
}
@Override
@@ -135,7 +136,7 @@ interface ChannelStateByteBuffer {
}
@Override
- public void recycle() {}
+ public void close() {}
@Override
public int writeBytes(InputStream input, int bytesToRead) throws IOException {
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/api/serialization/EventSerializer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/api/serialization/EventSerializer.java
index fd6b1c7..54f4a9e 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/api/serialization/EventSerializer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/api/serialization/EventSerializer.java
@@ -306,7 +306,9 @@ public class EventSerializer {
MemorySegment data = MemorySegmentFactory.wrap(serializedEvent.array());
return new BufferConsumer(
- data, FreeingBufferRecycler.INSTANCE, getDataType(event, hasPriority));
+ new NetworkBuffer(
+ data, FreeingBufferRecycler.INSTANCE, getDataType(event, hasPriority)),
+ data.size());
}
public static AbstractEvent fromBuffer(Buffer buffer, ClassLoader classLoader)
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 a5ae7ce..36d21e6 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
@@ -34,18 +34,17 @@ import static org.apache.flink.util.Preconditions.checkState;
* written data.
*/
@NotThreadSafe
-public class BufferBuilder {
+public class BufferBuilder implements AutoCloseable {
+ private final Buffer buffer;
private final MemorySegment memorySegment;
- private final BufferRecycler recycler;
-
private final SettablePositionMarker positionMarker = new SettablePositionMarker();
private boolean bufferConsumerCreated = false;
public BufferBuilder(MemorySegment memorySegment, BufferRecycler recycler) {
this.memorySegment = checkNotNull(memorySegment);
- this.recycler = checkNotNull(recycler);
+ this.buffer = new NetworkBuffer(memorySegment, recycler);
}
/**
@@ -73,7 +72,7 @@ public class BufferBuilder {
checkState(
!bufferConsumerCreated, "Two BufferConsumer shouldn't exist for one BufferBuilder");
bufferConsumerCreated = true;
- return new BufferConsumer(memorySegment, recycler, positionMarker, currentReaderPosition);
+ return new BufferConsumer(buffer.retainBuffer(), positionMarker, currentReaderPosition);
}
/** Same as {@link #append(ByteBuffer)} but additionally {@link #commit()} the appending. */
@@ -143,16 +142,12 @@ public class BufferBuilder {
}
public int getMaxCapacity() {
- return memorySegment.size();
- }
-
- @VisibleForTesting
- public BufferRecycler getRecycler() {
- return recycler;
+ return buffer.getMaxCapacity();
}
- public void recycle() {
- recycler.recycle(memorySegment);
+ @Override
+ public void close() {
+ buffer.recycleBuffer();
}
/**
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferConsumer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferConsumer.java
index 191d899..084f0ed 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferConsumer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferConsumer.java
@@ -46,41 +46,15 @@ public class BufferConsumer implements Closeable {
private int currentReaderPosition;
- /** Constructs {@link BufferConsumer} instance with the initial reader position. */
- public BufferConsumer(
- MemorySegment memorySegment,
- BufferRecycler recycler,
- PositionMarker currentWriterPosition,
- int currentReaderPosition) {
- this(
- new NetworkBuffer(checkNotNull(memorySegment), checkNotNull(recycler)),
- currentWriterPosition,
- currentReaderPosition);
- }
-
- /** Constructs {@link BufferConsumer} instance with static content. */
- public BufferConsumer(
- MemorySegment memorySegment, BufferRecycler recycler, Buffer.DataType dataType) {
- this(memorySegment, recycler, memorySegment.size(), dataType);
- }
-
/** Constructs {@link BufferConsumer} instance with static content of a certain size. */
- public BufferConsumer(
- MemorySegment memorySegment,
- BufferRecycler recycler,
- int size,
- Buffer.DataType dataType) {
- this(
- new NetworkBuffer(checkNotNull(memorySegment), checkNotNull(recycler), dataType),
- () -> -size,
- 0);
- checkState(memorySegment.size() > 0);
+ public BufferConsumer(Buffer buffer, int size) {
+ this(buffer, () -> -size, 0);
checkState(
isFinished(),
"BufferConsumer with static size must be finished after construction!");
}
- private BufferConsumer(
+ public BufferConsumer(
Buffer buffer,
BufferBuilder.PositionMarker currentWriterPosition,
int currentReaderPosition) {
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/BufferWritingResultPartition.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/BufferWritingResultPartition.java
index bc3278d..abb7e1d 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/BufferWritingResultPartition.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/BufferWritingResultPartition.java
@@ -349,6 +349,7 @@ public abstract class BufferWritingResultPartition extends ResultPartition {
numBytesOut.inc(bufferBuilder.finish());
numBuffersOut.inc();
unicastBufferBuilders[targetSubpartition] = null;
+ bufferBuilder.close();
}
}
@@ -362,6 +363,7 @@ public abstract class BufferWritingResultPartition extends ResultPartition {
if (broadcastBufferBuilder != null) {
numBytesOut.inc(broadcastBufferBuilder.finish() * numSubpartitions);
numBuffersOut.inc(numSubpartitions);
+ broadcastBufferBuilder.close();
broadcastBufferBuilder = null;
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/api/serialization/SpanningRecordSerializationTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/api/serialization/SpanningRecordSerializationTest.java
index c1aa03a..6924a88 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/api/serialization/SpanningRecordSerializationTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/api/serialization/SpanningRecordSerializationTest.java
@@ -292,15 +292,16 @@ public class SpanningRecordSerializationTest extends TestLogger {
}
private static Buffer appendLeftOverBytes(Buffer buffer, byte[] leftOverBytes) {
- BufferBuilder bufferBuilder =
+ try (BufferBuilder bufferBuilder =
new BufferBuilder(
MemorySegmentFactory.allocateUnpooledSegment(
buffer.readableBytes() + leftOverBytes.length),
- FreeingBufferRecycler.INSTANCE);
- try (BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer()) {
- bufferBuilder.append(buffer.getNioBufferReadable());
- bufferBuilder.appendAndCommit(ByteBuffer.wrap(leftOverBytes));
- return bufferConsumer.build();
+ FreeingBufferRecycler.INSTANCE)) {
+ try (BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer()) {
+ bufferBuilder.append(buffer.getNioBufferReadable());
+ bufferBuilder.appendAndCommit(ByteBuffer.wrap(leftOverBytes));
+ return bufferConsumer.build();
+ }
}
}
@@ -332,6 +333,10 @@ public class SpanningRecordSerializationTest extends TestLogger {
BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
bufferConsumer.build().recycleBuffer();
+ // Closing the BufferBuilder here just allow to be sure that Buffer will be recovered when
+ // BufferConsumer will be closed.
+ bufferBuilder.close();
+
bufferBuilder.appendAndCommit(serializedRecord);
return new BufferAndSerializerResult(
bufferBuilder,
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/api/writer/RecordWriterTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/api/writer/RecordWriterTest.java
index aacfd54..ec78fe2 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/api/writer/RecordWriterTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/api/writer/RecordWriterTest.java
@@ -312,13 +312,14 @@ public class RecordWriterTest {
assertTrue(recordWriter.getAvailableFuture().isDone());
// request one buffer from the local pool to make it unavailable afterwards
- final BufferBuilder bufferBuilder = localPool.requestBufferBuilder(0);
- assertNotNull(bufferBuilder);
- assertFalse(recordWriter.getAvailableFuture().isDone());
+ try (BufferBuilder bufferBuilder = localPool.requestBufferBuilder(0)) {
+ assertNotNull(bufferBuilder);
+ assertFalse(recordWriter.getAvailableFuture().isDone());
- // recycle the buffer to make the local pool available again
- final Buffer buffer = BufferBuilderTestUtils.buildSingleBuffer(bufferBuilder);
- buffer.recycleBuffer();
+ // recycle the buffer to make the local pool available again
+ final Buffer buffer = BufferBuilderTestUtils.buildSingleBuffer(bufferBuilder);
+ buffer.recycleBuffer();
+ }
assertTrue(recordWriter.getAvailableFuture().isDone());
assertEquals(recordWriter.AVAILABLE, recordWriter.getAvailableFuture());
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 c6db226..029e2da 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
@@ -18,7 +18,7 @@
package org.apache.flink.runtime.io.network.buffer;
-import org.apache.flink.core.memory.MemorySegmentFactory;
+import org.apache.flink.core.memory.MemorySegment;
import org.junit.Test;
@@ -28,6 +28,7 @@ import java.nio.ByteBuffer;
import java.nio.IntBuffer;
import java.util.ArrayList;
+import static org.apache.flink.core.memory.MemorySegmentFactory.allocateUnpooledSegment;
import static org.apache.flink.runtime.io.network.buffer.BufferBuilderTestUtils.buildSingleBuffer;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
@@ -46,6 +47,8 @@ public class BufferBuilderAndConsumerTest {
assertEquals(3 * Integer.BYTES, bufferBuilder.appendAndCommit(toByteBuffer(1, 2, 3)));
+ bufferBuilder.close();
+
Buffer buffer = bufferConsumer.build();
assertFalse(buffer.isRecycled());
buffer.recycleBuffer();
@@ -152,33 +155,37 @@ public class BufferBuilderAndConsumerTest {
@Test
public void buildEmptyBuffer() {
- Buffer buffer = buildSingleBuffer(createBufferBuilder());
- assertEquals(0, buffer.getSize());
- assertContent(buffer, FreeingBufferRecycler.INSTANCE);
+ try (BufferBuilder bufferBuilder = createBufferBuilder()) {
+ Buffer buffer = buildSingleBuffer(bufferBuilder);
+ assertEquals(0, buffer.getSize());
+ assertContent(buffer, FreeingBufferRecycler.INSTANCE);
+ }
}
@Test
public void buildingBufferMultipleTimes() {
- BufferBuilder bufferBuilder = createBufferBuilder();
- try (BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer()) {
- bufferBuilder.appendAndCommit(toByteBuffer(0, 1));
- bufferBuilder.appendAndCommit(toByteBuffer(2));
+ try (BufferBuilder bufferBuilder = createBufferBuilder()) {
+ try (BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer()) {
+ bufferBuilder.appendAndCommit(toByteBuffer(0, 1));
+ bufferBuilder.appendAndCommit(toByteBuffer(2));
- assertContent(bufferConsumer, 0, 1, 2);
+ assertContent(bufferConsumer, 0, 1, 2);
- bufferBuilder.appendAndCommit(toByteBuffer(3, 42));
- bufferBuilder.appendAndCommit(toByteBuffer(44));
+ bufferBuilder.appendAndCommit(toByteBuffer(3, 42));
+ bufferBuilder.appendAndCommit(toByteBuffer(44));
- assertContent(bufferConsumer, 3, 42, 44);
+ assertContent(bufferConsumer, 3, 42, 44);
- ArrayList<Integer> originalValues = new ArrayList<>();
- while (!bufferBuilder.isFull()) {
- bufferBuilder.appendAndCommit(toByteBuffer(1337));
- originalValues.add(1337);
- }
+ ArrayList<Integer> originalValues = new ArrayList<>();
+ while (!bufferBuilder.isFull()) {
+ bufferBuilder.appendAndCommit(toByteBuffer(1337));
+ originalValues.add(1337);
+ }
- assertContent(
- bufferConsumer, originalValues.stream().mapToInt(Integer::intValue).toArray());
+ assertContent(
+ bufferConsumer,
+ originalValues.stream().mapToInt(Integer::intValue).toArray());
+ }
}
}
@@ -220,6 +227,43 @@ public class BufferBuilderAndConsumerTest {
assertEquals(0, bufferBuilder.getWritableBytes());
}
+ @Test
+ public void recycleWithoutConsumer() {
+ // given: Recycler with the counter of recycle invocation.
+ CountedRecycler recycler = new CountedRecycler();
+ BufferBuilder bufferBuilder =
+ new BufferBuilder(allocateUnpooledSegment(BUFFER_SIZE), recycler);
+
+ // when: Invoke the recycle.
+ bufferBuilder.close();
+
+ // then: Recycling successfully finished.
+ assertEquals(1, recycler.recycleInvocationCounter);
+ }
+
+ @Test
+ public void recycleConsumerAndBufferBuilder() {
+ // given: Recycler with the counter of recycling invocation.
+ CountedRecycler recycler = new CountedRecycler();
+ BufferBuilder bufferBuilder =
+ new BufferBuilder(allocateUnpooledSegment(BUFFER_SIZE), recycler);
+
+ // and: One buffer consumer.
+ BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
+
+ // when: Invoke the recycle of BufferBuilder.
+ bufferBuilder.close();
+
+ // then: Nothing happened because BufferBuilder has already consumer.
+ assertEquals(0, recycler.recycleInvocationCounter);
+
+ // when: Close the consumer.
+ bufferConsumer.close();
+
+ // then: Recycling successfully finished.
+ assertEquals(1, recycler.recycleInvocationCounter);
+ }
+
private static void testIsFinished(int writes) {
BufferBuilder bufferBuilder = createBufferBuilder();
BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
@@ -283,7 +327,16 @@ public class BufferBuilderAndConsumerTest {
private static BufferBuilder createBufferBuilder() {
return new BufferBuilder(
- MemorySegmentFactory.allocateUnpooledSegment(BUFFER_SIZE),
- FreeingBufferRecycler.INSTANCE);
+ allocateUnpooledSegment(BUFFER_SIZE), FreeingBufferRecycler.INSTANCE);
+ }
+
+ private static class CountedRecycler implements BufferRecycler {
+ int recycleInvocationCounter;
+
+ @Override
+ public void recycle(MemorySegment memorySegment) {
+ recycleInvocationCounter++;
+ memorySegment.free();
+ }
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/BufferBuilderTestUtils.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/BufferBuilderTestUtils.java
index f08e128..77eb28d 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/BufferBuilderTestUtils.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/BufferBuilderTestUtils.java
@@ -39,12 +39,20 @@ public class BufferBuilderTestUtils {
return createFilledBufferBuilder(size, 0);
}
+ public static BufferBuilder createBufferBuilder(MemorySegment memorySegment) {
+ return createFilledBufferBuilder(memorySegment, 0);
+ }
+
public static BufferBuilder createFilledBufferBuilder(int size, int dataSize) {
checkArgument(size >= dataSize);
+ return createFilledBufferBuilder(
+ MemorySegmentFactory.allocateUnpooledSegment(size), dataSize);
+ }
+
+ public static BufferBuilder createFilledBufferBuilder(
+ MemorySegment memorySegment, int dataSize) {
BufferBuilder bufferBuilder =
- new BufferBuilder(
- MemorySegmentFactory.allocateUnpooledSegment(size),
- FreeingBufferRecycler.INSTANCE);
+ new BufferBuilder(memorySegment, FreeingBufferRecycler.INSTANCE);
return fillBufferBuilder(bufferBuilder, dataSize);
}
@@ -83,6 +91,7 @@ public class BufferBuilderTestUtils {
if (isFinished) {
bufferBuilder.finish();
+ bufferBuilder.close();
}
return bufferConsumer;
@@ -90,9 +99,11 @@ public class BufferBuilderTestUtils {
public static BufferConsumer createEventBufferConsumer(int size, Buffer.DataType dataType) {
return new BufferConsumer(
- MemorySegmentFactory.allocateUnpooledSegment(size),
- FreeingBufferRecycler.INSTANCE,
- dataType);
+ new NetworkBuffer(
+ MemorySegmentFactory.allocateUnpooledSegment(size),
+ FreeingBufferRecycler.INSTANCE,
+ dataType),
+ size);
}
public static Buffer buildBufferWithAscendingInts(int bufferSize, int numInts, int nextValue) {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
index 545d37a..d713a3c 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
@@ -449,15 +449,15 @@ public class LocalBufferPoolTest extends TestLogger {
assertFalse(localBufferPool.getAvailableFuture().isDone());
// recycle segments
- bufferBuilder11.getRecycler().recycle(bufferBuilder11.getMemorySegment());
+ bufferBuilder11.close();
assertFalse(localBufferPool.getAvailableFuture().isDone());
- bufferBuilder21.getRecycler().recycle(bufferBuilder21.getMemorySegment());
+ bufferBuilder21.close();
assertFalse(localBufferPool.getAvailableFuture().isDone());
- bufferBuilder02.getRecycler().recycle(bufferBuilder02.getMemorySegment());
+ bufferBuilder02.close();
assertTrue(localBufferPool.getAvailableFuture().isDone());
- bufferBuilder01.getRecycler().recycle(bufferBuilder01.getMemorySegment());
+ bufferBuilder01.close();
assertTrue(localBufferPool.getAvailableFuture().isDone());
- bufferBuilder22.getRecycler().recycle(bufferBuilder22.getMemorySegment());
+ bufferBuilder22.close();
assertTrue(localBufferPool.getAvailableFuture().isDone());
}
@@ -468,55 +468,53 @@ public class LocalBufferPoolTest extends TestLogger {
assertTrue(localBufferPool.isAvailable());
// request one buffer
- final BufferBuilder bufferBuilder =
- checkNotNull(localBufferPool.requestBufferBuilderBlocking());
- CompletableFuture<?> availableFuture = localBufferPool.getAvailableFuture();
- assertFalse(availableFuture.isDone());
-
- // set the pool size
- final int numLocalBuffers = 5;
- localBufferPool.setNumBuffers(numLocalBuffers);
- assertTrue(availableFuture.isDone());
- assertTrue(localBufferPool.isAvailable());
-
- // drain the local buffer pool
- final Deque<Buffer> buffers = new ArrayDeque<>(LocalBufferPoolTest.numBuffers);
- for (int i = 0; i < numLocalBuffers - 1; i++) {
+ try (BufferBuilder bufferBuilder =
+ checkNotNull(localBufferPool.requestBufferBuilderBlocking())) {
+ CompletableFuture<?> availableFuture = localBufferPool.getAvailableFuture();
+ assertFalse(availableFuture.isDone());
+
+ // set the pool size
+ final int numLocalBuffers = 5;
+ localBufferPool.setNumBuffers(numLocalBuffers);
+ assertTrue(availableFuture.isDone());
assertTrue(localBufferPool.isAvailable());
- buffers.add(checkNotNull(localBufferPool.requestBuffer()));
- }
- assertFalse(localBufferPool.isAvailable());
- buffers.pop().recycleBuffer();
- assertTrue(localBufferPool.isAvailable());
+ // drain the local buffer pool
+ final Deque<Buffer> buffers = new ArrayDeque<>(LocalBufferPoolTest.numBuffers);
+ for (int i = 0; i < numLocalBuffers - 1; i++) {
+ assertTrue(localBufferPool.isAvailable());
+ buffers.add(checkNotNull(localBufferPool.requestBuffer()));
+ }
+ assertFalse(localBufferPool.isAvailable());
- // recycle the requested segments to global buffer pool
- for (final Buffer buffer : buffers) {
- buffer.recycleBuffer();
- }
- assertTrue(localBufferPool.isAvailable());
+ buffers.pop().recycleBuffer();
+ assertTrue(localBufferPool.isAvailable());
- // scale down (first buffer still taken), but there should still be one segment locally
- // available
- localBufferPool.setNumBuffers(2);
- assertTrue(localBufferPool.isAvailable());
+ // recycle the requested segments to global buffer pool
+ for (final Buffer buffer : buffers) {
+ buffer.recycleBuffer();
+ }
+ assertTrue(localBufferPool.isAvailable());
- final Buffer buffer2 = checkNotNull(localBufferPool.requestBuffer());
- assertFalse(localBufferPool.isAvailable());
+ // scale down (first buffer still taken), but there should still be one segment locally
+ // available
+ localBufferPool.setNumBuffers(2);
+ assertTrue(localBufferPool.isAvailable());
- buffer2.recycleBuffer();
- assertTrue(localBufferPool.isAvailable());
+ final Buffer buffer2 = checkNotNull(localBufferPool.requestBuffer());
+ assertFalse(localBufferPool.isAvailable());
- // reset the pool size
- localBufferPool.setNumBuffers(1);
- CompletableFuture<?> availableFuture2 = localBufferPool.getAvailableFuture();
- assertFalse(availableFuture2.isDone());
+ buffer2.recycleBuffer();
+ assertTrue(localBufferPool.isAvailable());
+
+ // reset the pool size
+ localBufferPool.setNumBuffers(1);
+ assertFalse(localBufferPool.getAvailableFuture().isDone());
+ // recycle the requested buffer
+ }
- // recycle the requested buffer
- bufferBuilder.createBufferConsumer().close();
- bufferBuilder.recycle();
assertTrue(localBufferPool.isAvailable());
- assertTrue(availableFuture2.isDone());
+ assertTrue(localBufferPool.getAvailableFuture().isDone());
}
/** For FLINK-20547: https://issues.apache.org/jira/browse/FLINK-20547. */
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NetworkBufferPoolTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NetworkBufferPoolTest.java
index 917c61a..235ca3c 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NetworkBufferPoolTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NetworkBufferPoolTest.java
@@ -695,7 +695,7 @@ public class NetworkBufferPoolTest extends TestLogger {
// recycle all the requested buffers
for (BufferBuilder bufferBuilder : segmentsRequested) {
- bufferBuilder.createBufferConsumer().close();
+ bufferBuilder.close();
}
} finally {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/BoundedBlockingSubpartitionWriteReadTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/BoundedBlockingSubpartitionWriteReadTest.java
index 37dad72..64afe99 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/BoundedBlockingSubpartitionWriteReadTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/BoundedBlockingSubpartitionWriteReadTest.java
@@ -26,6 +26,7 @@ import org.apache.flink.runtime.io.disk.FileChannelManagerImpl;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.BufferConsumer;
import org.apache.flink.runtime.io.network.buffer.BufferDecompressor;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
import org.apache.flink.runtime.io.network.partition.ResultSubpartition.BufferAndBacklog;
import org.apache.flink.runtime.util.EnvironmentInformation;
@@ -232,7 +233,9 @@ public class BoundedBlockingSubpartitionWriteReadTest {
}
partition.add(
- new BufferConsumer(memory, (ignored) -> {}, pos, Buffer.DataType.DATA_BUFFER));
+ new BufferConsumer(
+ new NetworkBuffer(memory, (ignored) -> {}, Buffer.DataType.DATA_BUFFER),
+ pos));
// we need to flush after every buffer as long as the add() contract is that
// buffer are immediately added and can be filled further after that (for low latency
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannelTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannelTest.java
index 650dc4b..ee48c4a 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannelTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/LocalInputChannelTest.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.io.network.partition.consumer;
+import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.checkpoint.CheckpointOptions;
import org.apache.flink.runtime.checkpoint.channel.RecordingChannelStateWriter;
import org.apache.flink.runtime.concurrent.FutureUtils;
@@ -32,6 +33,7 @@ import org.apache.flink.runtime.io.network.buffer.BufferConsumer;
import org.apache.flink.runtime.io.network.buffer.BufferPool;
import org.apache.flink.runtime.io.network.buffer.BufferProvider;
import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
import org.apache.flink.runtime.io.network.partition.BufferAvailabilityListener;
import org.apache.flink.runtime.io.network.partition.BufferWritingResultPartition;
@@ -100,11 +102,14 @@ public class LocalInputChannelTest {
CheckpointBarrier barrier =
new CheckpointBarrier(
1L, 0L, CheckpointOptions.alignedWithTimeout(getDefault(), 123L));
+ MemorySegment memorySegment = EventSerializer.toBuffer(barrier, false).getMemorySegment();
BufferConsumer barrierHolder =
new BufferConsumer(
- EventSerializer.toBuffer(barrier, false).getMemorySegment(),
- FreeingBufferRecycler.INSTANCE,
- Buffer.DataType.EVENT_BUFFER);
+ new NetworkBuffer(
+ memorySegment,
+ FreeingBufferRecycler.INSTANCE,
+ Buffer.DataType.EVENT_BUFFER),
+ memorySegment.size());
BufferConsumer data = BufferBuilderTestUtils.createFilledFinishedBufferConsumer(1);
RecordingChannelStateWriter stateWriter = new RecordingChannelStateWriter();
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannelTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannelTest.java
index 5182e20..1828147 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannelTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/RemoteInputChannelTest.java
@@ -35,6 +35,7 @@ import org.apache.flink.runtime.io.network.TestingPartitionRequestClient;
import org.apache.flink.runtime.io.network.api.CheckpointBarrier;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.Buffer.DataType;
+import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
import org.apache.flink.runtime.io.network.buffer.BufferListener.NotificationResult;
import org.apache.flink.runtime.io.network.buffer.BufferPool;
import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
@@ -1237,9 +1238,11 @@ public class RemoteInputChannelTest {
final Callable<Void> bufferPoolInteractionsTask =
() -> {
for (int i = 0; i < retries; ++i) {
- Buffer buffer =
- buildSingleBuffer(bufferPool.requestBufferBuilderBlocking());
- buffer.recycleBuffer();
+ try (BufferBuilder bufferBuilder =
+ bufferPool.requestBufferBuilderBlocking()) {
+ Buffer buffer = buildSingleBuffer(bufferBuilder);
+ buffer.recycleBuffer();
+ }
}
return null;
};
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestSubpartitionProducer.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestSubpartitionProducer.java
index 2b8808a..973e68f 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestSubpartitionProducer.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestSubpartitionProducer.java
@@ -22,6 +22,7 @@ import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.core.memory.MemorySegmentFactory;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.BufferConsumer;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
import org.apache.flink.runtime.io.network.partition.ResultSubpartition;
import org.apache.flink.runtime.io.network.util.TestProducerSource.BufferAndChannel;
@@ -77,7 +78,9 @@ public class TestSubpartitionProducer implements Callable<Boolean> {
MemorySegment segment = MemorySegmentFactory.wrap(bufferAndChannel.getBuffer());
subpartition.add(
new BufferConsumer(
- segment, MemorySegment::free, Buffer.DataType.DATA_BUFFER));
+ new NetworkBuffer(
+ segment, MemorySegment::free, Buffer.DataType.DATA_BUFFER),
+ segment.size()));
// Check for interrupted flag after adding data to prevent resource leaks
if (Thread.interrupted()) {
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/StreamTestSingleInputGate.java b/flink-streaming-java/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/StreamTestSingleInputGate.java
index 47ced2c..4b45fe2 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/StreamTestSingleInputGate.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/runtime/io/network/partition/consumer/StreamTestSingleInputGate.java
@@ -121,6 +121,7 @@ public class StreamTestSingleInputGate<T> extends TestSingleInputGate {
BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
bufferBuilder.appendAndCommit(serializedRecord);
bufferBuilder.finish();
+ bufferBuilder.close();
// Call getCurrentBuffer to ensure size is set
return Optional.of(
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/StreamTaskNetworkInputTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/StreamTaskNetworkInputTest.java
index 969867b..6be7815 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/StreamTaskNetworkInputTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/StreamTaskNetworkInputTest.java
@@ -236,12 +236,14 @@ public class StreamTaskNetworkInputTest {
}
private BufferOrEvent createDataBuffer() throws IOException {
- BufferBuilder bufferBuilder = BufferBuilderTestUtils.createEmptyBufferBuilder(PAGE_SIZE);
- BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
- serializeRecord(42L, bufferBuilder);
- serializeRecord(44L, bufferBuilder);
+ try (BufferBuilder bufferBuilder =
+ BufferBuilderTestUtils.createEmptyBufferBuilder(PAGE_SIZE)) {
+ BufferConsumer bufferConsumer = bufferBuilder.createBufferConsumer();
+ serializeRecord(42L, bufferBuilder);
+ serializeRecord(44L, bufferBuilder);
- return new BufferOrEvent(bufferConsumer.build(), new InputChannelInfo(0, 0));
+ return new BufferOrEvent(bufferConsumer.build(), new InputChannelInfo(0, 0));
+ }
}
private StreamTaskNetworkInput<Long> createStreamTaskNetworkInput(List<BufferOrEvent> buffers) {
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/recovery/DemultiplexingRecordDeserializerTest.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/recovery/DemultiplexingRecordDeserializerTest.java
index 447c2be..0986468 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/recovery/DemultiplexingRecordDeserializerTest.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/runtime/io/recovery/DemultiplexingRecordDeserializerTest.java
@@ -19,6 +19,7 @@ package org.apache.flink.streaming.runtime.io.recovery;
import org.apache.flink.api.common.typeutils.base.LongSerializer;
import org.apache.flink.core.memory.DataOutputSerializer;
+import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo;
import org.apache.flink.runtime.io.disk.iomanager.IOManager;
import org.apache.flink.runtime.io.disk.iomanager.IOManagerAsync;
@@ -56,6 +57,7 @@ import java.util.concurrent.ThreadLocalRandom;
import java.util.stream.Collectors;
import static java.util.Collections.emptySet;
+import static org.apache.flink.core.memory.MemorySegmentFactory.allocateUnpooledSegment;
import static org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptorUtil.array;
import static org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptorUtil.mappings;
import static org.apache.flink.runtime.checkpoint.InflightDataRescalingDescriptorUtil.rescalingDescriptor;
@@ -115,15 +117,17 @@ public class DemultiplexingRecordDeserializerTest {
long start = selector.getInputSubtaskIndex() << 4 | selector.getOutputSubtaskIndex();
- final BufferBuilder bufferBuilder = createBufferBuilder(128);
- Buffer buffer = writeLongs(bufferBuilder, start + 1L, start + 2L, start + 3L);
+ MemorySegment memorySegment = allocateUnpooledSegment(128);
+ try (BufferBuilder bufferBuilder = createBufferBuilder(memorySegment)) {
+ Buffer buffer = writeLongs(bufferBuilder, start + 1L, start + 2L, start + 3L);
- deserializer.select(selector);
- deserializer.setNextBuffer(buffer);
+ deserializer.select(selector);
+ deserializer.setNextBuffer(buffer);
+ }
assertEquals(
Arrays.asList(start + 1L, start + 2L, start + 3L), readLongs(deserializer));
- assertTrue(bufferBuilder.getMemorySegment().isFreed());
+ assertTrue(memorySegment.isFreed());
}
}
@@ -151,23 +155,25 @@ public class DemultiplexingRecordDeserializerTest {
deserializer.getVirtualChannelSelectors());
for (int i = 0; i < 100; i++) {
- final BufferBuilder bufferBuilder = createBufferBuilder(128);
- // add one even and one odd number
- Buffer buffer = writeLongs(bufferBuilder, i, i + 1L);
-
- SubtaskConnectionDescriptor selector =
- Iterables.get(deserializer.getVirtualChannelSelectors(), i / 10 % 2);
- deserializer.select(selector);
- deserializer.setNextBuffer(buffer);
-
- if (selector.getInputSubtaskIndex() == 41) {
- assertEquals(Arrays.asList((long) i, i + 1L), readLongs(deserializer));
- } else {
- // only odd should occur in output
- assertEquals(Arrays.asList(i / 2 * 2 + 1L), readLongs(deserializer));
+ MemorySegment memorySegment = allocateUnpooledSegment(128);
+ try (BufferBuilder bufferBuilder = createBufferBuilder(memorySegment)) {
+ // add one even and one odd number
+ Buffer buffer = writeLongs(bufferBuilder, i, i + 1L);
+
+ SubtaskConnectionDescriptor selector =
+ Iterables.get(deserializer.getVirtualChannelSelectors(), i / 10 % 2);
+ deserializer.select(selector);
+ deserializer.setNextBuffer(buffer);
+
+ if (selector.getInputSubtaskIndex() == 41) {
+ assertEquals(Arrays.asList((long) i, i + 1L), readLongs(deserializer));
+ } else {
+ // only odd should occur in output
+ assertEquals(Arrays.asList(i / 2 * 2 + 1L), readLongs(deserializer));
+ }
}
- assertTrue(bufferBuilder.getMemorySegment().isFreed());
+ assertTrue(memorySegment.isFreed());
}
}
@@ -190,13 +196,15 @@ public class DemultiplexingRecordDeserializerTest {
deserializer.getVirtualChannelSelectors().iterator();
iterator.hasNext(); ) {
SubtaskConnectionDescriptor selector = iterator.next();
- final BufferBuilder bufferBuilder = createBufferBuilder(128);
- final long ts =
- 42L + selector.getInputSubtaskIndex() + selector.getOutputSubtaskIndex();
- Buffer buffer = write(bufferBuilder, new Watermark(ts));
-
- deserializer.select(selector);
- deserializer.setNextBuffer(buffer);
+ MemorySegment memorySegment = allocateUnpooledSegment(128);
+ try (BufferBuilder bufferBuilder = createBufferBuilder(memorySegment)) {
+ final long ts =
+ 42L + selector.getInputSubtaskIndex() + selector.getOutputSubtaskIndex();
+ Buffer buffer = write(bufferBuilder, new Watermark(ts));
+
+ deserializer.select(selector);
+ deserializer.setNextBuffer(buffer);
+ }
if (iterator.hasNext()) {
assertEquals(Collections.emptyList(), read(deserializer));
@@ -205,7 +213,7 @@ public class DemultiplexingRecordDeserializerTest {
assertEquals(Arrays.asList(new Watermark(42)), read(deserializer));
}
- assertTrue(bufferBuilder.getMemorySegment().isFreed());
+ assertTrue(memorySegment.isFreed());
}
}
[flink] 01/04: [FLINK-22376][runtime] BufferProvider is able to
provide pure MemorySegment in order to avoid extracting it from the
bufferBuilder
Posted by pn...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
pnowojski pushed a commit to branch release-1.13
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 8511ecbc33e613d8554911ba8b943bf7a98ca77c
Author: Anton Kalashnikov <ka...@yandex.ru>
AuthorDate: Fri May 28 11:51:37 2021 +0200
[FLINK-22376][runtime] BufferProvider is able to provide pure MemorySegment in order to avoid extracting it from the bufferBuilder
(cherry picked from commit 397468706943aeed4ec2cf72f846e4457e3e1cd4)
---
.../runtime/io/network/buffer/BufferBuilder.java | 6 --
.../runtime/io/network/buffer/BufferProvider.java | 18 +++++
.../runtime/io/network/buffer/LocalBufferPool.java | 11 ++-
.../network/partition/PartitionSortedBuffer.java | 28 ++++----
.../partition/SortMergeResultPartition.java | 25 ++++---
.../io/network/buffer/LocalBufferPoolTest.java | 3 +-
.../runtime/io/network/buffer/NoOpBufferPool.java | 10 +++
.../io/network/buffer/UnpooledBufferPool.java | 8 ++-
.../runtime/io/network/util/TestBufferFactory.java | 11 ++-
.../io/network/util/TestPooledBufferProvider.java | 82 +++++++++++-----------
10 files changed, 120 insertions(+), 82 deletions(-)
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 c09827f..a5ae7ce 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
@@ -18,7 +18,6 @@
package org.apache.flink.runtime.io.network.buffer;
-import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.core.memory.MemorySegment;
import javax.annotation.concurrent.NotThreadSafe;
@@ -156,11 +155,6 @@ public class BufferBuilder {
recycler.recycle(memorySegment);
}
- @VisibleForTesting
- public MemorySegment getMemorySegment() {
- return memorySegment;
- }
-
/**
* Holds a reference to the current writer position. Negative values indicate that writer
* ({@link BufferBuilder} has finished. Value {@code Integer.MIN_VALUE} represents finished
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferProvider.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferProvider.java
index 8606f61..e5d38df 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferProvider.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferProvider.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.io.network.buffer;
+import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.io.AvailabilityProvider;
import javax.annotation.Nullable;
@@ -85,4 +86,21 @@ public interface BufferProvider extends AvailabilityProvider {
/** Returns whether the buffer provider has been destroyed. */
boolean isDestroyed();
+
+ /**
+ * Returns a {@link MemorySegment} instance from the buffer provider.
+ *
+ * @return {@code null} if no memory segment is available or the buffer provider has been
+ * destroyed.
+ */
+ @Nullable
+ MemorySegment requestMemorySegment();
+
+ /**
+ * Returns a {@link MemorySegment} instance from the buffer provider.
+ *
+ * <p>If there is no memory segment available, the call will block until one becomes available
+ * again or the buffer provider has been destroyed.
+ */
+ MemorySegment requestMemorySegmentBlocking() throws InterruptedException;
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
index 6748504..6c14824 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
@@ -281,7 +281,12 @@ class LocalBufferPool implements BufferPool {
@Override
public BufferBuilder requestBufferBuilderBlocking() throws InterruptedException {
- return toBufferBuilder(requestMemorySegmentBlocking(UNKNOWN_CHANNEL), UNKNOWN_CHANNEL);
+ return toBufferBuilder(requestMemorySegmentBlocking(), UNKNOWN_CHANNEL);
+ }
+
+ @Override
+ public MemorySegment requestMemorySegmentBlocking() throws InterruptedException {
+ return requestMemorySegmentBlocking(UNKNOWN_CHANNEL);
}
@Override
@@ -359,8 +364,8 @@ class LocalBufferPool implements BufferPool {
return segment;
}
- @Nullable
- private MemorySegment requestMemorySegment() {
+ @Override
+ public MemorySegment requestMemorySegment() {
return requestMemorySegment(UNKNOWN_CHANNEL);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PartitionSortedBuffer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PartitionSortedBuffer.java
index 81a2330..b432d30 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PartitionSortedBuffer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/PartitionSortedBuffer.java
@@ -21,7 +21,6 @@ package org.apache.flink.runtime.io.network.partition;
import org.apache.flink.core.memory.MemorySegment;
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.BufferPool;
import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
@@ -67,7 +66,7 @@ public class PartitionSortedBuffer implements SortBuffer {
/** A segment list as a joint buffer which stores all records and index entries. */
@GuardedBy("lock")
- private final ArrayList<MemorySegment> buffers = new ArrayList<>();
+ private final ArrayList<MemorySegment> segments = new ArrayList<>();
/** Addresses of the first record's index entry for each subpartition. */
private final long[] firstIndexEntryAddresses;
@@ -184,7 +183,7 @@ public class PartitionSortedBuffer implements SortBuffer {
}
private void writeIndex(int channelIndex, int numRecordBytes, Buffer.DataType dataType) {
- MemorySegment segment = buffers.get(writeSegmentIndex);
+ MemorySegment segment = segments.get(writeSegmentIndex);
// record length takes the high 32 bits and data type takes the low 32 bits
segment.putLong(writeSegmentOffset, ((long) numRecordBytes << 32) | dataType.ordinal());
@@ -197,7 +196,7 @@ public class PartitionSortedBuffer implements SortBuffer {
if (lastIndexEntryAddress >= 0) {
// link the previous index entry of the given channel to the new index entry
- segment = buffers.get(getSegmentIndexFromPointer(lastIndexEntryAddress));
+ segment = segments.get(getSegmentIndexFromPointer(lastIndexEntryAddress));
segment.putLong(
getSegmentOffsetFromPointer(lastIndexEntryAddress) + 8, indexEntryAddress);
} else {
@@ -210,7 +209,7 @@ public class PartitionSortedBuffer implements SortBuffer {
private void writeRecord(ByteBuffer source) {
while (source.hasRemaining()) {
- MemorySegment segment = buffers.get(writeSegmentIndex);
+ MemorySegment segment = segments.get(writeSegmentIndex);
int toCopy = Math.min(bufferSize - writeSegmentOffset, source.remaining());
segment.put(writeSegmentOffset, source, toCopy);
@@ -222,7 +221,7 @@ public class PartitionSortedBuffer implements SortBuffer {
private boolean allocateBuffersForRecord(int numRecordBytes) throws IOException {
int numBytesRequired = INDEX_ENTRY_SIZE + numRecordBytes;
int availableBytes =
- writeSegmentIndex == buffers.size() ? 0 : bufferSize - writeSegmentOffset;
+ writeSegmentIndex == segments.size() ? 0 : bufferSize - writeSegmentOffset;
// return directly if current available bytes is adequate
if (availableBytes >= numBytesRequired) {
@@ -262,22 +261,21 @@ public class PartitionSortedBuffer implements SortBuffer {
throw new IllegalStateException("Sort buffer is already released.");
}
- buffers.add(segment);
+ segments.add(segment);
}
}
private MemorySegment requestBufferFromPool() throws IOException {
try {
// blocking request buffers if there is still guaranteed memory
- if (buffers.size() < numGuaranteedBuffers) {
- return bufferPool.requestBufferBuilderBlocking().getMemorySegment();
+ if (segments.size() < numGuaranteedBuffers) {
+ return bufferPool.requestMemorySegmentBlocking();
}
} catch (InterruptedException e) {
throw new IOException("Interrupted while requesting buffer.");
}
- BufferBuilder buffer = bufferPool.requestBufferBuilder();
- return buffer != null ? buffer.getMemorySegment() : null;
+ return bufferPool.requestMemorySegment();
}
private void updateWriteSegmentIndexAndOffset(int numBytes) {
@@ -303,7 +301,7 @@ public class PartitionSortedBuffer implements SortBuffer {
do {
int sourceSegmentIndex = getSegmentIndexFromPointer(readIndexEntryAddress);
int sourceSegmentOffset = getSegmentOffsetFromPointer(readIndexEntryAddress);
- MemorySegment sourceSegment = buffers.get(sourceSegmentIndex);
+ MemorySegment sourceSegment = segments.get(sourceSegmentIndex);
long lengthAndDataType = sourceSegment.getLong(sourceSegmentOffset);
int length = getSegmentIndexFromPointer(lengthAndDataType);
@@ -376,7 +374,7 @@ public class PartitionSortedBuffer implements SortBuffer {
int sourceRemainingBytes =
Math.min(bufferSize - sourceSegmentOffset, recordRemainingBytes);
int numBytes = Math.min(targetSegmentSize - targetSegmentOffset, sourceRemainingBytes);
- MemorySegment sourceSegment = buffers.get(sourceSegmentIndex);
+ MemorySegment sourceSegment = segments.get(sourceSegmentIndex);
sourceSegment.copyTo(sourceSegmentOffset, targetSegment, targetSegmentOffset, numBytes);
recordRemainingBytes -= numBytes;
@@ -445,10 +443,10 @@ public class PartitionSortedBuffer implements SortBuffer {
isReleased = true;
- for (MemorySegment segment : buffers) {
+ for (MemorySegment segment : segments) {
bufferPool.recycle(segment);
}
- buffers.clear();
+ segments.clear();
numTotalBytes = 0;
numTotalRecords = 0;
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/SortMergeResultPartition.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/SortMergeResultPartition.java
index 4442ca1..9382503 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/SortMergeResultPartition.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/SortMergeResultPartition.java
@@ -78,7 +78,7 @@ public class SortMergeResultPartition extends ResultPartition {
/** Buffers cut from the network buffer pool for data writing. */
@GuardedBy("lock")
- private final List<MemorySegment> writeBuffers = new ArrayList<>();
+ private final List<MemorySegment> writeSegments = new ArrayList<>();
/** Size of network buffer and write buffer. */
private final int networkBufferSize;
@@ -172,9 +172,8 @@ public class SortMergeResultPartition extends ResultPartition {
synchronized (lock) {
try {
for (int i = 0; i < numWriteBuffers; ++i) {
- MemorySegment segment =
- bufferPool.requestBufferBuilderBlocking().getMemorySegment();
- writeBuffers.add(segment);
+ MemorySegment segment = bufferPool.requestMemorySegmentBlocking();
+ writeSegments.add(segment);
}
} catch (InterruptedException exception) {
// the setup method does not allow InterruptedException
@@ -310,13 +309,13 @@ public class SortMergeResultPartition extends ResultPartition {
fileWriter.startNewRegion(isBroadcast);
List<BufferWithChannel> toWrite = new ArrayList<>();
- Queue<MemorySegment> segments = getWriteBuffers();
+ Queue<MemorySegment> segments = getWriteSegments();
while (sortBuffer.hasRemaining()) {
if (segments.isEmpty()) {
fileWriter.writeBuffers(toWrite);
toWrite.clear();
- segments = getWriteBuffers();
+ segments = getWriteSegments();
}
BufferWithChannel bufferWithChannel =
@@ -339,10 +338,10 @@ public class SortMergeResultPartition extends ResultPartition {
flushSortBuffer(unicastSortBuffer, false);
}
- private Queue<MemorySegment> getWriteBuffers() {
+ private Queue<MemorySegment> getWriteSegments() {
synchronized (lock) {
- checkState(!writeBuffers.isEmpty(), "Task has been canceled.");
- return new ArrayDeque<>(writeBuffers);
+ checkState(!writeSegments.isEmpty(), "Task has been canceled.");
+ return new ArrayDeque<>(writeSegments);
}
}
@@ -371,13 +370,13 @@ public class SortMergeResultPartition extends ResultPartition {
fileWriter.startNewRegion(isBroadcast);
List<BufferWithChannel> toWrite = new ArrayList<>();
- Queue<MemorySegment> segments = getWriteBuffers();
+ Queue<MemorySegment> segments = getWriteSegments();
while (record.hasRemaining()) {
if (segments.isEmpty()) {
fileWriter.writeBuffers(toWrite);
toWrite.clear();
- segments = getWriteBuffers();
+ segments = getWriteSegments();
}
int toCopy = Math.min(record.remaining(), networkBufferSize);
@@ -414,10 +413,10 @@ public class SortMergeResultPartition extends ResultPartition {
private void releaseWriteBuffers() {
synchronized (lock) {
if (bufferPool != null) {
- for (MemorySegment segment : writeBuffers) {
+ for (MemorySegment segment : writeSegments) {
bufferPool.recycle(segment);
}
- writeBuffers.clear();
+ writeSegments.clear();
}
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
index c5115a6..545d37a 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
@@ -514,6 +514,7 @@ public class LocalBufferPoolTest extends TestLogger {
// recycle the requested buffer
bufferBuilder.createBufferConsumer().close();
+ bufferBuilder.recycle();
assertTrue(localBufferPool.isAvailable());
assertTrue(availableFuture2.isDone());
}
@@ -524,7 +525,7 @@ public class LocalBufferPoolTest extends TestLogger {
NetworkBufferPool globalPool = new TestNetworkBufferPool(numBuffers, memorySegmentSize);
try {
BufferPool localPool = new LocalBufferPool(globalPool, 1);
- MemorySegment segment = localPool.requestBufferBuilderBlocking().getMemorySegment();
+ MemorySegment segment = localPool.requestMemorySegmentBlocking();
localPool.setNumBuffers(2);
localPool.recycle(segment);
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NoOpBufferPool.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NoOpBufferPool.java
index b93de38..a7cc030 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NoOpBufferPool.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NoOpBufferPool.java
@@ -67,6 +67,16 @@ public class NoOpBufferPool implements BufferPool {
}
@Override
+ public MemorySegment requestMemorySegment() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public MemorySegment requestMemorySegmentBlocking() throws InterruptedException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
public int getNumberOfRequiredMemorySegments() {
throw new UnsupportedOperationException();
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/UnpooledBufferPool.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/UnpooledBufferPool.java
index d2598c5..73b0804 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/UnpooledBufferPool.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/UnpooledBufferPool.java
@@ -41,11 +41,17 @@ public class UnpooledBufferPool implements BufferPool {
return new BufferBuilder(requestMemorySegment(), this);
}
- private MemorySegment requestMemorySegment() {
+ @Override
+ public MemorySegment requestMemorySegment() {
return MemorySegmentFactory.allocateUnpooledOffHeapMemory(SEGMENT_SIZE, null);
}
@Override
+ public MemorySegment requestMemorySegmentBlocking() throws InterruptedException {
+ return requestMemorySegment();
+ }
+
+ @Override
public BufferBuilder requestBufferBuilderBlocking() throws InterruptedException {
return requestBufferBuilder();
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestBufferFactory.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestBufferFactory.java
index e1877bf..22e72b5 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestBufferFactory.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestBufferFactory.java
@@ -18,6 +18,7 @@
package org.apache.flink.runtime.io.network.util;
+import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.core.memory.MemorySegmentFactory;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.BufferRecycler;
@@ -51,14 +52,18 @@ public class TestBufferFactory {
this.bufferRecycler = checkNotNull(bufferRecycler);
}
- public synchronized Buffer create() {
+ public synchronized MemorySegment createMemorySegment() {
if (numberOfCreatedBuffers >= poolSize) {
return null;
}
numberOfCreatedBuffers++;
- return new NetworkBuffer(
- MemorySegmentFactory.allocateUnpooledSegment(bufferSize), bufferRecycler);
+ return MemorySegmentFactory.allocateUnpooledSegment(bufferSize);
+ }
+
+ public synchronized Buffer create() {
+ MemorySegment memorySegment = createMemorySegment();
+ return memorySegment == null ? null : new NetworkBuffer(memorySegment, bufferRecycler);
}
public synchronized int getNumberOfCreatedBuffers() {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestPooledBufferProvider.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestPooledBufferProvider.java
index f92bc98..da6c4f0 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestPooledBufferProvider.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/util/TestPooledBufferProvider.java
@@ -36,9 +36,10 @@ import java.util.concurrent.LinkedBlockingDeque;
import static org.apache.flink.util.Preconditions.checkArgument;
+/** */
public class TestPooledBufferProvider implements BufferProvider {
- private final BlockingQueue<Buffer> buffers = new LinkedBlockingDeque<>();
+ private final BlockingQueue<MemorySegment> segments = new LinkedBlockingDeque<>();
private final TestBufferFactory bufferFactory;
@@ -51,25 +52,22 @@ public class TestPooledBufferProvider implements BufferProvider {
public TestPooledBufferProvider(int poolSize, int bufferSize) {
checkArgument(poolSize > 0);
- this.bufferRecycler = new PooledBufferProviderRecycler(buffers);
+ this.bufferRecycler = new PooledBufferProviderRecycler(segments);
this.bufferFactory = new TestBufferFactory(poolSize, bufferSize, bufferRecycler);
}
@Override
public Buffer requestBuffer() {
- final Buffer buffer = buffers.poll();
- if (buffer != null) {
- return buffer;
- }
+ MemorySegment memorySegment = requestMemorySegment();
- return bufferFactory.create();
+ return memorySegment == null ? null : new NetworkBuffer(memorySegment, bufferRecycler);
}
@Override
public BufferBuilder requestBufferBuilder() {
- Buffer buffer = requestBuffer();
- if (buffer != null) {
- return new BufferBuilder(buffer.getMemorySegment(), buffer.getRecycler());
+ MemorySegment memorySegment = requestMemorySegment();
+ if (memorySegment != null) {
+ return new BufferBuilder(memorySegment, bufferRecycler);
}
return null;
}
@@ -79,24 +77,9 @@ public class TestPooledBufferProvider implements BufferProvider {
return requestBufferBuilder();
}
- private Buffer requestBufferBlocking() throws InterruptedException {
- Buffer buffer = buffers.poll();
- if (buffer != null) {
- return buffer;
- }
-
- buffer = bufferFactory.create();
- if (buffer != null) {
- return buffer;
- }
-
- return buffers.take();
- }
-
@Override
public BufferBuilder requestBufferBuilderBlocking() throws InterruptedException {
- Buffer buffer = requestBufferBlocking();
- return new BufferBuilder(buffer.getMemorySegment(), buffer.getRecycler());
+ return new BufferBuilder(requestMemorySegmentBlocking(), bufferRecycler);
}
@Override
@@ -116,49 +99,68 @@ public class TestPooledBufferProvider implements BufferProvider {
}
@Override
- public CompletableFuture<?> getAvailableFuture() {
- return AVAILABLE;
+ public MemorySegment requestMemorySegment() {
+ final MemorySegment buffer = segments.poll();
+ if (buffer != null) {
+ return buffer;
+ }
+
+ return bufferFactory.createMemorySegment();
}
- public int getNumberOfAvailableBuffers() {
- return buffers.size();
+ @Override
+ public MemorySegment requestMemorySegmentBlocking() throws InterruptedException {
+ MemorySegment buffer = segments.poll();
+ if (buffer != null) {
+ return buffer;
+ }
+
+ buffer = bufferFactory.createMemorySegment();
+ if (buffer != null) {
+ return buffer;
+ }
+
+ return segments.take();
+ }
+
+ @Override
+ public CompletableFuture<?> getAvailableFuture() {
+ return AVAILABLE;
}
- public int getNumberOfCreatedBuffers() {
- return bufferFactory.getNumberOfCreatedBuffers();
+ public int getNumberOfAvailableSegments() {
+ return segments.size();
}
private static class PooledBufferProviderRecycler implements BufferRecycler {
private final Object listenerRegistrationLock = new Object();
- private final Queue<Buffer> buffers;
+ private final Queue<MemorySegment> segments;
private final ConcurrentLinkedQueue<BufferListener> registeredListeners =
Queues.newConcurrentLinkedQueue();
- public PooledBufferProviderRecycler(Queue<Buffer> buffers) {
- this.buffers = buffers;
+ public PooledBufferProviderRecycler(Queue<MemorySegment> segments) {
+ this.segments = segments;
}
@Override
public void recycle(MemorySegment segment) {
synchronized (listenerRegistrationLock) {
- final Buffer buffer = new NetworkBuffer(segment, this);
-
BufferListener listener = registeredListeners.poll();
if (listener == null) {
- buffers.add(buffer);
+ segments.add(segment);
} else {
- listener.notifyBufferAvailable(buffer);
+ listener.notifyBufferAvailable(new NetworkBuffer(segment, this));
}
}
}
boolean registerListener(BufferListener listener) {
synchronized (listenerRegistrationLock) {
- if (buffers.isEmpty()) {
+ if (segments.isEmpty()) {
registeredListeners.add(listener);
return true;