You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@flink.apache.org by "TanYuxin-tyx (via GitHub)" <gi...@apache.org> on 2023/06/20 12:37:23 UTC
[GitHub] [flink] TanYuxin-tyx opened a new pull request, #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
TanYuxin-tyx opened a new pull request, #22833:
URL: https://github.com/apache/flink/pull/22833
<!--
*Thank you very much for contributing to Apache Flink - we are happy that you want to help us improve Flink. To help the community review your contribution in the best possible way, please go through the checklist below, which will get the contribution into a shape in which it can be best reviewed.*
*Please understand that we do not do this to make contributions to Flink a hassle. In order to uphold a high standard of quality for code contributions, while at the same time managing a large number of contributions, we need contributors to prepare the contributions well, and give reviewers enough contextual information for the review. Please also understand that contributions that do not follow this guide will take longer to review and thus typically be picked up with lower priority by the community.*
## Contribution Checklist
- Make sure that the pull request corresponds to a [JIRA issue](https://issues.apache.org/jira/projects/FLINK/issues). Exceptions are made for typos in JavaDoc or documentation files, which need no JIRA issue.
- Name the pull request in the form "[FLINK-XXXX] [component] Title of the pull request", where *FLINK-XXXX* should be replaced by the actual issue number. Skip *component* if you are unsure about which is the best component.
Typo fixes that have no associated JIRA issue should be named following this pattern: `[hotfix] [docs] Fix typo in event time introduction` or `[hotfix] [javadocs] Expand JavaDoc for PuncuatedWatermarkGenerator`.
- Fill out the template below to describe the changes contributed by the pull request. That will give reviewers the context they need to do the review.
- Make sure that the change passes the automated tests, i.e., `mvn clean verify` passes. You can set up Azure Pipelines CI to do that following [this guide](https://cwiki.apache.org/confluence/display/FLINK/Azure+Pipelines#AzurePipelines-Tutorial:SettingupAzurePipelinesforaforkoftheFlinkrepository).
- Each pull request should address only one issue, not mix up code from multiple issues.
- Each commit in the pull request has a meaningful commit message (including the JIRA id)
- Once all items of the checklist are addressed, remove the above text and this checklist, leaving only the filled out template below.
**(The sections below can be removed for hotfixes of typos)**
-->
## What is the purpose of the change
*Implement the SortBufferAccumulator for the tiered storage*
## Brief change log
- *Implement the SortBufferAccumulator for the tiered storage*
## Verifying this change
This change will add tests.
## Does this pull request potentially affect one of the following parts:
- Dependencies (does it add or upgrade a dependency): (no)
- The public API, i.e., is any changed class annotated with `@Public(Evolving)`: (no)
- The serializers: (no)
- The runtime per-record code paths (performance sensitive): (no)
- Anything that affects deployment or recovery: JobManager (and its components), Checkpointing, Kubernetes/Yarn, ZooKeeper: (no)
- The S3 file system connector: (no)
## Documentation
- Does this pull request introduce a new feature? (no)
- If yes, how is the feature documented? (not documented)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1245156676
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
+ numExpectBuffers += currentBufferWrittenBytes == 0 ? 0 : 1;
+ assertThat(numReceivedFinishedBuffer.get()).isEqualTo(numExpectBuffers);
+ }
+
+ @Test
+ void testWriteLargeRecord() throws IOException {
+ int numBuffers = 15;
+ Random random = new Random();
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ AtomicInteger numReceivedBuffers = new AtomicInteger(0);
+ bufferAccumulator.setup(
+ (subpartitionIndex, buffers) -> {
+ numReceivedBuffers.set(numReceivedBuffers.get() + buffers.size());
+ buffers.forEach(Buffer::recycleBuffer);
+ });
+ ByteBuffer largeRecord = generateRandomData(BUFFER_SIZE_BYTES * numBuffers, random);
+ bufferAccumulator.receive(
+ largeRecord,
+ new TieredStorageSubpartitionId(0),
+ Buffer.DataType.DATA_BUFFER,
+ false);
+ assertThat(numReceivedBuffers.get()).isEqualTo(numBuffers);
Review Comment:
Fixed.
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
+ numExpectBuffers += currentBufferWrittenBytes == 0 ? 0 : 1;
+ assertThat(numReceivedFinishedBuffer.get()).isEqualTo(numExpectBuffers);
+ }
+
+ @Test
+ void testWriteLargeRecord() throws IOException {
+ int numBuffers = 15;
+ Random random = new Random();
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ AtomicInteger numReceivedBuffers = new AtomicInteger(0);
+ bufferAccumulator.setup(
+ (subpartitionIndex, buffers) -> {
+ numReceivedBuffers.set(numReceivedBuffers.get() + buffers.size());
Review Comment:
Fixed.
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
+ numExpectBuffers += currentBufferWrittenBytes == 0 ? 0 : 1;
+ assertThat(numReceivedFinishedBuffer.get()).isEqualTo(numExpectBuffers);
Review Comment:
Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] xintongsong commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1239483891
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/BufferAccumulator.java:
##########
@@ -46,7 +46,10 @@ public interface BufferAccumulator extends AutoCloseable {
* transformed into finished buffers.
*/
void receive(
- ByteBuffer record, TieredStorageSubpartitionId subpartitionId, Buffer.DataType dataType)
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
Review Comment:
When `isBroadcast` is true, what `subpartitionId` are we expecting? Will it be ignored and thus can be an arbitrary id, or even `null`? Or it is required to be some special value?
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
Review Comment:
It took me a while to figure out why `sortBufferContainer.hasRemaining()` being `false` should result in writing a large record. It would be easier if we add a simple comment explaining: the container is empty, yet we failed to write the record into it, which suggests the record is larger than the container can hold.
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferContainer.java:
##########
@@ -0,0 +1,454 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+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;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedList;
+
+import static org.apache.flink.util.Preconditions.checkArgument;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The buffer container for accumulating the records into {@link Buffer}s. After accumulating, the
+ * {@link SortBufferAccumulator} will read the sorted buffers.
+ */
+public class SortBufferContainer {
+
+ /**
+ * Size of an index entry: 4 bytes for record length, 4 bytes for data type and 8 bytes for
+ * pointer to next entry.
+ */
+ private static final int INDEX_ENTRY_SIZE = 4 + 4 + 8;
+
+ /** A list of {@link MemorySegment}s used to store data in memory. */
+ private final LinkedList<MemorySegment> freeSegments;
+
+ /** A segment list as a joint buffer which stores all records and index entries. */
+ private final ArrayList<MemorySegment> dataSegments;
+
+ /** {@link BufferRecycler} used to recycle {@link #freeSegments}. */
+ private final BufferRecycler bufferRecycler;
+
+ /** Addresses of the first record's index entry for each subpartition. */
+ private final long[] subpartitionFirstBufferIndexEntries;
+
+ /** Addresses of the last record's index entry for each subpartition. */
+ private final long[] subpartitionLastBufferIndexEntries;
+
+ /** Size of buffers requested from buffer pool. All buffers must be of the same size. */
+ private final int bufferSizeBytes;
+
+ /** Number of guaranteed buffers can be allocated from the buffer pool for data sort. */
+ private final int numBuffersForSort;
+
+ // ------------------------------------------------------------------------
+ // The statistics and states
+ // ------------------------------------------------------------------------
+
+ /** Total number of bytes already appended to this sort buffer. */
+ private long numTotalBytes;
+
+ /** Total number of bytes already read from this sort buffer. */
+ private long numTotalBytesRead;
+
+ /** Whether this sort buffer is finished. One can only read a finished sort buffer. */
+ private boolean isFinished;
+
+ /** Whether this sort buffer is released. A released sort buffer can not be used. */
+ private boolean isReleased;
+
+ // ------------------------------------------------------------------------
+ // For writing
+ // ------------------------------------------------------------------------
+
+ /** Array index in the segment list of the current available buffer for writing. */
+ private int writeBufferIndex;
+
+ /** Next position in the current available buffer for writing. */
+ private int writeOffsetInCurrentBuffer;
+
+ // ------------------------------------------------------------------------
+ // For reading
+ // ------------------------------------------------------------------------
+
+ /** Index entry address of the current record or event to be read. */
+ private long readBufferIndexEntry;
+
+ /**
+ * Record the bytes remaining after the last read, which must be initialized before reading a
+ * new record.
+ */
+ private int recordRemainingBytesToRead;
+
+ /** The subpartition that is reading data from. */
+ private int readingSubpartitionId = -1;
+
+ SortBufferContainer(
+ LinkedList<MemorySegment> freeSegments,
+ BufferRecycler bufferRecycler,
+ int numSubpartitions,
+ int bufferSizeBytes,
+ int numBuffersForSort) {
+ checkArgument(bufferSizeBytes > INDEX_ENTRY_SIZE, "Buffer size is too small.");
+ checkArgument(numBuffersForSort > 0, "No guaranteed buffers for sort.");
+ checkState(numBuffersForSort <= freeSegments.size(), "Wrong number of free segments.");
+
+ this.freeSegments = checkNotNull(freeSegments);
+ this.bufferRecycler = checkNotNull(bufferRecycler);
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffersForSort = numBuffersForSort;
+ this.dataSegments = new ArrayList<>();
+ this.subpartitionFirstBufferIndexEntries = new long[numSubpartitions];
+ this.subpartitionLastBufferIndexEntries = new long[numSubpartitions];
+
+ Arrays.fill(subpartitionFirstBufferIndexEntries, -1L);
+ Arrays.fill(subpartitionLastBufferIndexEntries, -1L);
+ }
+
+ // ------------------------------------------------------------------------
+ // Called by SortBufferAccumulator
+ // ------------------------------------------------------------------------
+
+ /**
+ * Note that no partial records will be written to this {@link SortBufferContainer}, which means
+ * that either all data of target record will be written or nothing will be written.
+ *
+ * @param record the record to be written
+ * @param subpartitionId the subpartition id
+ * @param dataType the data type of the record
+ * @return true if the {@link SortBufferContainer} is full, or return false if the contianer is
+ * not full
Review Comment:
This is against conventions. Usually, a write-like method returns true means the write operation has succeed.
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+
+ flushDataBuffer(sortBufferContainer);
+ sortBufferContainer.release();
+ if (record.hasRemaining()) {
+ receive(record, subpartitionId, dataType, isBroadcast);
+ }
+ }
+
+ @Override
+ public void close() {
+ flushUnicastDataBuffer();
+ flushBroadcastDataBuffer();
+ releaseFreeBuffers();
+ releaseDataBuffer(unicastDataBuffer);
+ releaseDataBuffer(broadcastDataBuffer);
+ }
+
+ // ------------------------------------------------------------------------
+ // Internal Methods
+ // ------------------------------------------------------------------------
+
+ private SortBufferContainer getUnicastDataBuffer() {
+ flushBroadcastDataBuffer();
+
+ if (unicastDataBuffer != null
+ && !unicastDataBuffer.isFinished()
+ && !unicastDataBuffer.isReleased()) {
+ return unicastDataBuffer;
+ }
+
+ unicastDataBuffer = createNewDataBuffer();
+ return unicastDataBuffer;
+ }
+
+ private SortBufferContainer getBroadcastDataBuffer() {
+ flushUnicastDataBuffer();
+
+ if (broadcastDataBuffer != null
+ && !broadcastDataBuffer.isFinished()
+ && !broadcastDataBuffer.isReleased()) {
+ return broadcastDataBuffer;
+ }
+
+ broadcastDataBuffer = createNewDataBuffer();
+ return broadcastDataBuffer;
+ }
+
+ private SortBufferContainer createNewDataBuffer() {
+ requestNetworkBuffers();
+
+ return new SortBufferContainer(
+ freeSegments,
+ this::recycleBuffer,
+ numSubpartitions,
+ bufferSizeBytes,
+ numBuffersForSort);
+ }
+
+ private void requestGuaranteedBuffers() {
+ int effectiveRequiredBuffers = effectiveNumRequestedBuffers();
+
+ while (freeSegments.size() < effectiveRequiredBuffers) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegments.add(checkNotNull(buffer).getMemorySegment());
+ if (bufferRecycler == null) {
+ bufferRecycler = buffer.getRecycler();
+ }
+ }
+ }
+
+ private void requestNetworkBuffers() {
+ requestGuaranteedBuffers();
+
+ // Use the half of the buffers for writing, and the other half for reading
+ numBuffersForSort = freeSegments.size() / 2;
+ }
+
+ private void flushDataBuffer(SortBufferContainer sortBufferContainer) {
+ if (sortBufferContainer == null
+ || sortBufferContainer.isReleased()
+ || !sortBufferContainer.hasRemaining()) {
+ return;
+ }
+ sortBufferContainer.finish();
+
+ do {
+ MemorySegment freeSegment = getFreeSegment();
+ Pair<Integer, Buffer> bufferAndSubpartitionId =
+ sortBufferContainer.readBuffer(freeSegment);
+ if (bufferAndSubpartitionId == null) {
+ if (freeSegment != null) {
+ recycleBuffer(freeSegment);
+ }
+ break;
+ }
+ addFinishedBuffer(bufferAndSubpartitionId);
+ } while (true);
+
+ releaseFreeBuffers();
+ sortBufferContainer.release();
+ }
+
+ private void flushBroadcastDataBuffer() {
+ if (broadcastDataBuffer != null) {
+ flushDataBuffer(broadcastDataBuffer);
+ broadcastDataBuffer.release();
+ broadcastDataBuffer = null;
+ }
+ }
+
+ private void flushUnicastDataBuffer() {
+ if (unicastDataBuffer != null) {
+ flushDataBuffer(unicastDataBuffer);
+ unicastDataBuffer.release();
+ unicastDataBuffer = null;
+ }
+ }
+
+ private void flushContainerWhenEndOfPartition(
+ boolean isEndOfPartition, SortBufferContainer sortBufferContainer) {
+ if (isEndOfPartition) {
+ flushDataBuffer(sortBufferContainer);
+ }
+ }
+
+ private void writeLargeRecord(
+ ByteBuffer record, int targetSubpartition, Buffer.DataType dataType) {
+
+ checkState(dataType != Buffer.DataType.EVENT_BUFFER);
+ while (record.hasRemaining()) {
+ int toCopy = Math.min(record.remaining(), bufferSizeBytes);
+ MemorySegment writeBuffer = checkNotNull(getFreeSegment());
+ writeBuffer.put(0, record, toCopy);
+
+ addFinishedBuffer(
+ Pair.of(
+ targetSubpartition,
+ new NetworkBuffer(
+ writeBuffer, checkNotNull(bufferRecycler), dataType, toCopy)));
+ }
+
+ releaseFreeBuffers();
+ }
+
+ private MemorySegment getFreeSegment() {
+ MemorySegment freeSegment = freeSegments.poll();
+ if (freeSegment == null) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegment = buffer.getMemorySegment();
+ }
+ return freeSegment;
+ }
+
+ private int effectiveNumRequestedBuffers() {
+ return Math.min(numSubpartitions + 1, numBuffers);
Review Comment:
Why `numSubpartitions + 1`?
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferContainer.java:
##########
@@ -0,0 +1,454 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+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;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedList;
+
+import static org.apache.flink.util.Preconditions.checkArgument;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The buffer container for accumulating the records into {@link Buffer}s. After accumulating, the
+ * {@link SortBufferAccumulator} will read the sorted buffers.
+ */
+public class SortBufferContainer {
Review Comment:
1. The name is confusing. I'd suggest simply `SortBuffer` or `BufferSorter`.
2. I think the important characteristic of this component is that, you can write data into it in any order but read data from it in the order grouped by subpartitions. This is not clear enough from the JavaDoc.
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+
+ flushDataBuffer(sortBufferContainer);
+ sortBufferContainer.release();
+ if (record.hasRemaining()) {
Review Comment:
Shouldn't this always be `true`?
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
Review Comment:
It is implicit that getting one of the two data buffers would result in the other one being flushed.
I'd suggest the following:
- Have only one field `currentDataBuffer`, instead of two nullable `broadcastDataBuffer` and `unicastDataBuffer` and only one of them can have non-null value at anytime.
- Instead of calling `getBroadcast/UnicastDataBuffer`, we can have a method `switchCurrentDataBufferIfNeeded(isBroadcast)`, in which we can update the value of `currentDataBuffer`. This would make the flushing less surprising.
- We might need another field to remember whether the current data buffer is broadcast or unicast.
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+
+ flushDataBuffer(sortBufferContainer);
+ sortBufferContainer.release();
+ if (record.hasRemaining()) {
+ receive(record, subpartitionId, dataType, isBroadcast);
+ }
+ }
+
+ @Override
+ public void close() {
+ flushUnicastDataBuffer();
+ flushBroadcastDataBuffer();
+ releaseFreeBuffers();
+ releaseDataBuffer(unicastDataBuffer);
+ releaseDataBuffer(broadcastDataBuffer);
+ }
+
+ // ------------------------------------------------------------------------
+ // Internal Methods
+ // ------------------------------------------------------------------------
+
+ private SortBufferContainer getUnicastDataBuffer() {
+ flushBroadcastDataBuffer();
+
+ if (unicastDataBuffer != null
+ && !unicastDataBuffer.isFinished()
+ && !unicastDataBuffer.isReleased()) {
+ return unicastDataBuffer;
+ }
+
+ unicastDataBuffer = createNewDataBuffer();
+ return unicastDataBuffer;
+ }
+
+ private SortBufferContainer getBroadcastDataBuffer() {
+ flushUnicastDataBuffer();
+
+ if (broadcastDataBuffer != null
+ && !broadcastDataBuffer.isFinished()
+ && !broadcastDataBuffer.isReleased()) {
+ return broadcastDataBuffer;
+ }
+
+ broadcastDataBuffer = createNewDataBuffer();
+ return broadcastDataBuffer;
+ }
+
+ private SortBufferContainer createNewDataBuffer() {
+ requestNetworkBuffers();
+
+ return new SortBufferContainer(
+ freeSegments,
+ this::recycleBuffer,
+ numSubpartitions,
+ bufferSizeBytes,
+ numBuffersForSort);
+ }
+
+ private void requestGuaranteedBuffers() {
+ int effectiveRequiredBuffers = effectiveNumRequestedBuffers();
+
+ while (freeSegments.size() < effectiveRequiredBuffers) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegments.add(checkNotNull(buffer).getMemorySegment());
+ if (bufferRecycler == null) {
+ bufferRecycler = buffer.getRecycler();
+ }
+ }
+ }
+
+ private void requestNetworkBuffers() {
+ requestGuaranteedBuffers();
+
+ // Use the half of the buffers for writing, and the other half for reading
+ numBuffersForSort = freeSegments.size() / 2;
+ }
+
+ private void flushDataBuffer(SortBufferContainer sortBufferContainer) {
+ if (sortBufferContainer == null
+ || sortBufferContainer.isReleased()
+ || !sortBufferContainer.hasRemaining()) {
+ return;
+ }
+ sortBufferContainer.finish();
+
+ do {
+ MemorySegment freeSegment = getFreeSegment();
+ Pair<Integer, Buffer> bufferAndSubpartitionId =
+ sortBufferContainer.readBuffer(freeSegment);
+ if (bufferAndSubpartitionId == null) {
+ if (freeSegment != null) {
+ recycleBuffer(freeSegment);
+ }
+ break;
+ }
+ addFinishedBuffer(bufferAndSubpartitionId);
+ } while (true);
+
+ releaseFreeBuffers();
+ sortBufferContainer.release();
+ }
+
+ private void flushBroadcastDataBuffer() {
+ if (broadcastDataBuffer != null) {
+ flushDataBuffer(broadcastDataBuffer);
+ broadcastDataBuffer.release();
+ broadcastDataBuffer = null;
+ }
+ }
+
+ private void flushUnicastDataBuffer() {
+ if (unicastDataBuffer != null) {
+ flushDataBuffer(unicastDataBuffer);
+ unicastDataBuffer.release();
+ unicastDataBuffer = null;
+ }
+ }
+
+ private void flushContainerWhenEndOfPartition(
+ boolean isEndOfPartition, SortBufferContainer sortBufferContainer) {
+ if (isEndOfPartition) {
+ flushDataBuffer(sortBufferContainer);
+ }
+ }
+
+ private void writeLargeRecord(
+ ByteBuffer record, int targetSubpartition, Buffer.DataType dataType) {
+
+ checkState(dataType != Buffer.DataType.EVENT_BUFFER);
+ while (record.hasRemaining()) {
+ int toCopy = Math.min(record.remaining(), bufferSizeBytes);
+ MemorySegment writeBuffer = checkNotNull(getFreeSegment());
+ writeBuffer.put(0, record, toCopy);
+
+ addFinishedBuffer(
+ Pair.of(
+ targetSubpartition,
+ new NetworkBuffer(
+ writeBuffer, checkNotNull(bufferRecycler), dataType, toCopy)));
+ }
+
+ releaseFreeBuffers();
+ }
+
+ private MemorySegment getFreeSegment() {
+ MemorySegment freeSegment = freeSegments.poll();
+ if (freeSegment == null) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegment = buffer.getMemorySegment();
+ }
+ return freeSegment;
+ }
+
+ private int effectiveNumRequestedBuffers() {
+ return Math.min(numSubpartitions + 1, numBuffers);
+ }
+
+ private void releaseDataBuffer(SortBufferContainer sortBufferContainer) {
+ if (sortBufferContainer != null) {
+ sortBufferContainer.release();
+ }
+ }
+
+ private void addFinishedBuffer(Pair<Integer, Buffer> bufferAndSubpartitionId) {
Review Comment:
```suggestion
private void flushBuffer(Pair<Integer, Buffer> bufferAndSubpartitionId) {
```
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+
+ flushDataBuffer(sortBufferContainer);
+ sortBufferContainer.release();
+ if (record.hasRemaining()) {
+ receive(record, subpartitionId, dataType, isBroadcast);
+ }
+ }
+
+ @Override
+ public void close() {
+ flushUnicastDataBuffer();
+ flushBroadcastDataBuffer();
+ releaseFreeBuffers();
+ releaseDataBuffer(unicastDataBuffer);
+ releaseDataBuffer(broadcastDataBuffer);
+ }
+
+ // ------------------------------------------------------------------------
+ // Internal Methods
+ // ------------------------------------------------------------------------
+
+ private SortBufferContainer getUnicastDataBuffer() {
+ flushBroadcastDataBuffer();
+
+ if (unicastDataBuffer != null
+ && !unicastDataBuffer.isFinished()
+ && !unicastDataBuffer.isReleased()) {
+ return unicastDataBuffer;
+ }
+
+ unicastDataBuffer = createNewDataBuffer();
+ return unicastDataBuffer;
+ }
+
+ private SortBufferContainer getBroadcastDataBuffer() {
+ flushUnicastDataBuffer();
+
+ if (broadcastDataBuffer != null
+ && !broadcastDataBuffer.isFinished()
+ && !broadcastDataBuffer.isReleased()) {
+ return broadcastDataBuffer;
+ }
+
+ broadcastDataBuffer = createNewDataBuffer();
+ return broadcastDataBuffer;
+ }
+
+ private SortBufferContainer createNewDataBuffer() {
+ requestNetworkBuffers();
+
+ return new SortBufferContainer(
+ freeSegments,
+ this::recycleBuffer,
+ numSubpartitions,
+ bufferSizeBytes,
+ numBuffersForSort);
+ }
+
+ private void requestGuaranteedBuffers() {
+ int effectiveRequiredBuffers = effectiveNumRequestedBuffers();
+
+ while (freeSegments.size() < effectiveRequiredBuffers) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegments.add(checkNotNull(buffer).getMemorySegment());
+ if (bufferRecycler == null) {
+ bufferRecycler = buffer.getRecycler();
+ }
+ }
+ }
+
+ private void requestNetworkBuffers() {
+ requestGuaranteedBuffers();
+
+ // Use the half of the buffers for writing, and the other half for reading
+ numBuffersForSort = freeSegments.size() / 2;
+ }
+
+ private void flushDataBuffer(SortBufferContainer sortBufferContainer) {
+ if (sortBufferContainer == null
+ || sortBufferContainer.isReleased()
+ || !sortBufferContainer.hasRemaining()) {
+ return;
+ }
+ sortBufferContainer.finish();
+
+ do {
+ MemorySegment freeSegment = getFreeSegment();
+ Pair<Integer, Buffer> bufferAndSubpartitionId =
+ sortBufferContainer.readBuffer(freeSegment);
+ if (bufferAndSubpartitionId == null) {
+ if (freeSegment != null) {
+ recycleBuffer(freeSegment);
+ }
+ break;
+ }
+ addFinishedBuffer(bufferAndSubpartitionId);
+ } while (true);
+
+ releaseFreeBuffers();
+ sortBufferContainer.release();
+ }
+
+ private void flushBroadcastDataBuffer() {
+ if (broadcastDataBuffer != null) {
+ flushDataBuffer(broadcastDataBuffer);
+ broadcastDataBuffer.release();
+ broadcastDataBuffer = null;
+ }
+ }
+
+ private void flushUnicastDataBuffer() {
+ if (unicastDataBuffer != null) {
+ flushDataBuffer(unicastDataBuffer);
+ unicastDataBuffer.release();
+ unicastDataBuffer = null;
+ }
+ }
+
+ private void flushContainerWhenEndOfPartition(
+ boolean isEndOfPartition, SortBufferContainer sortBufferContainer) {
+ if (isEndOfPartition) {
+ flushDataBuffer(sortBufferContainer);
+ }
+ }
Review Comment:
This is unused.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1244830366
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/BufferAccumulator.java:
##########
@@ -44,9 +44,21 @@ public interface BufferAccumulator extends AutoCloseable {
/**
* Receives the records from tiered store producer, these records will be accumulated and
* transformed into finished buffers.
+ *
+ * <p>Note that when isBroadcast is true, for a broadcast-only partition, the subpartitionId
+ * value will always be 0. Conversely, for a non-broadcast-only partition (such as an all-to-all
Review Comment:
yes, Removed the description of `all-to-all` partition.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1241083798
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferContainer.java:
##########
@@ -0,0 +1,454 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+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;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedList;
+
+import static org.apache.flink.util.Preconditions.checkArgument;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The buffer container for accumulating the records into {@link Buffer}s. After accumulating, the
+ * {@link SortBufferAccumulator} will read the sorted buffers.
+ */
+public class SortBufferContainer {
+
+ /**
+ * Size of an index entry: 4 bytes for record length, 4 bytes for data type and 8 bytes for
+ * pointer to next entry.
+ */
+ private static final int INDEX_ENTRY_SIZE = 4 + 4 + 8;
+
+ /** A list of {@link MemorySegment}s used to store data in memory. */
+ private final LinkedList<MemorySegment> freeSegments;
+
+ /** A segment list as a joint buffer which stores all records and index entries. */
+ private final ArrayList<MemorySegment> dataSegments;
+
+ /** {@link BufferRecycler} used to recycle {@link #freeSegments}. */
+ private final BufferRecycler bufferRecycler;
+
+ /** Addresses of the first record's index entry for each subpartition. */
+ private final long[] subpartitionFirstBufferIndexEntries;
+
+ /** Addresses of the last record's index entry for each subpartition. */
+ private final long[] subpartitionLastBufferIndexEntries;
+
+ /** Size of buffers requested from buffer pool. All buffers must be of the same size. */
+ private final int bufferSizeBytes;
+
+ /** Number of guaranteed buffers can be allocated from the buffer pool for data sort. */
+ private final int numBuffersForSort;
+
+ // ------------------------------------------------------------------------
+ // The statistics and states
+ // ------------------------------------------------------------------------
+
+ /** Total number of bytes already appended to this sort buffer. */
+ private long numTotalBytes;
+
+ /** Total number of bytes already read from this sort buffer. */
+ private long numTotalBytesRead;
+
+ /** Whether this sort buffer is finished. One can only read a finished sort buffer. */
+ private boolean isFinished;
+
+ /** Whether this sort buffer is released. A released sort buffer can not be used. */
+ private boolean isReleased;
+
+ // ------------------------------------------------------------------------
+ // For writing
+ // ------------------------------------------------------------------------
+
+ /** Array index in the segment list of the current available buffer for writing. */
+ private int writeBufferIndex;
+
+ /** Next position in the current available buffer for writing. */
+ private int writeOffsetInCurrentBuffer;
+
+ // ------------------------------------------------------------------------
+ // For reading
+ // ------------------------------------------------------------------------
+
+ /** Index entry address of the current record or event to be read. */
+ private long readBufferIndexEntry;
+
+ /**
+ * Record the bytes remaining after the last read, which must be initialized before reading a
+ * new record.
+ */
+ private int recordRemainingBytesToRead;
+
+ /** The subpartition that is reading data from. */
+ private int readingSubpartitionId = -1;
+
+ SortBufferContainer(
+ LinkedList<MemorySegment> freeSegments,
+ BufferRecycler bufferRecycler,
+ int numSubpartitions,
+ int bufferSizeBytes,
+ int numBuffersForSort) {
+ checkArgument(bufferSizeBytes > INDEX_ENTRY_SIZE, "Buffer size is too small.");
+ checkArgument(numBuffersForSort > 0, "No guaranteed buffers for sort.");
+ checkState(numBuffersForSort <= freeSegments.size(), "Wrong number of free segments.");
+
+ this.freeSegments = checkNotNull(freeSegments);
+ this.bufferRecycler = checkNotNull(bufferRecycler);
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffersForSort = numBuffersForSort;
+ this.dataSegments = new ArrayList<>();
+ this.subpartitionFirstBufferIndexEntries = new long[numSubpartitions];
+ this.subpartitionLastBufferIndexEntries = new long[numSubpartitions];
+
+ Arrays.fill(subpartitionFirstBufferIndexEntries, -1L);
+ Arrays.fill(subpartitionLastBufferIndexEntries, -1L);
+ }
+
+ // ------------------------------------------------------------------------
+ // Called by SortBufferAccumulator
+ // ------------------------------------------------------------------------
+
+ /**
+ * Note that no partial records will be written to this {@link SortBufferContainer}, which means
+ * that either all data of target record will be written or nothing will be written.
+ *
+ * @param record the record to be written
+ * @param subpartitionId the subpartition id
+ * @param dataType the data type of the record
+ * @return true if the {@link SortBufferContainer} is full, or return false if the contianer is
+ * not full
Review Comment:
Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] reswqa merged pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "reswqa (via GitHub)" <gi...@apache.org>.
reswqa merged PR #22833:
URL: https://github.com/apache/flink/pull/22833
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1244833111
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager tieredStorageMemoryManager =
+ createStorageMemoryManager(numBuffers);
+
+ int numExpectBytes = 0;
+ int numExpectBuffers = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, tieredStorageMemoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ int currentBufferWriteBytes = 0;
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (numExpectBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ numExpectBytes = 0;
Review Comment:
Yes, renamed it to `currentBufferWrittenBytes`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] reswqa commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "reswqa (via GitHub)" <gi...@apache.org>.
reswqa commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1244913594
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
+ numExpectBuffers += currentBufferWrittenBytes == 0 ? 0 : 1;
+ assertThat(numReceivedFinishedBuffer.get()).isEqualTo(numExpectBuffers);
+ }
+
+ @Test
+ void testWriteLargeRecord() throws IOException {
+ int numBuffers = 15;
+ Random random = new Random();
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ AtomicInteger numReceivedBuffers = new AtomicInteger(0);
+ bufferAccumulator.setup(
+ (subpartitionIndex, buffers) -> {
+ numReceivedBuffers.set(numReceivedBuffers.get() + buffers.size());
+ buffers.forEach(Buffer::recycleBuffer);
+ });
+ ByteBuffer largeRecord = generateRandomData(BUFFER_SIZE_BYTES * numBuffers, random);
+ bufferAccumulator.receive(
+ largeRecord,
+ new TieredStorageSubpartitionId(0),
+ Buffer.DataType.DATA_BUFFER,
+ false);
+ assertThat(numReceivedBuffers.get()).isEqualTo(numBuffers);
Review Comment:
```suggestion
assertThat(numReceivedBuffers).hasValue(numBuffers);
```
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
Review Comment:
```suggestion
assertThat(currentBufferWrittenBytes).isLessThan(BUFFER_SIZE_BYTES);
```
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
+ numExpectBuffers += currentBufferWrittenBytes == 0 ? 0 : 1;
+ assertThat(numReceivedFinishedBuffer.get()).isEqualTo(numExpectBuffers);
Review Comment:
```suggestion
assertThat(numReceivedFinishedBuffer).hasValue(numExpectBuffers);
```
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
+ numExpectBuffers += currentBufferWrittenBytes == 0 ? 0 : 1;
+ assertThat(numReceivedFinishedBuffer.get()).isEqualTo(numExpectBuffers);
+ }
+
+ @Test
+ void testWriteLargeRecord() throws IOException {
+ int numBuffers = 15;
+ Random random = new Random();
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ AtomicInteger numReceivedBuffers = new AtomicInteger(0);
+ bufferAccumulator.setup(
+ (subpartitionIndex, buffers) -> {
+ numReceivedBuffers.set(numReceivedBuffers.get() + buffers.size());
Review Comment:
```suggestion
numReceivedBuffers.getAndAdd(buffers.size());
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1244831973
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager tieredStorageMemoryManager =
+ createStorageMemoryManager(numBuffers);
+
+ int numExpectBytes = 0;
+ int numExpectBuffers = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, tieredStorageMemoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ int currentBufferWriteBytes = 0;
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (numExpectBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ numExpectBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ numExpectBytes += numBytes + indexEntrySize;
+ currentBufferWriteBytes += numBytes + indexEntrySize;
Review Comment:
Nice catch, Removed. It's a wrongly useless legacy variable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1244832264
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,269 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link SortBufferAccumulator} can help use less buffers to accumulate data, which
+ * decouples the buffer usage with the number of parallelism. The number of buffers used by the
+ * {@link SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBuffer} is full, or
+ * switching from broadcast to non-broadcast(or vice versa), the buffer in the sort buffer will be
+ * flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
Review Comment:
Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1244834064
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager tieredStorageMemoryManager =
+ createStorageMemoryManager(numBuffers);
+
+ int numExpectBytes = 0;
+ int numExpectBuffers = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, tieredStorageMemoryManager)) {
Review Comment:
Fixed with an added comment.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1241087268
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+
+ flushDataBuffer(sortBufferContainer);
+ sortBufferContainer.release();
+ if (record.hasRemaining()) {
+ receive(record, subpartitionId, dataType, isBroadcast);
+ }
+ }
+
+ @Override
+ public void close() {
+ flushUnicastDataBuffer();
+ flushBroadcastDataBuffer();
+ releaseFreeBuffers();
+ releaseDataBuffer(unicastDataBuffer);
+ releaseDataBuffer(broadcastDataBuffer);
+ }
+
+ // ------------------------------------------------------------------------
+ // Internal Methods
+ // ------------------------------------------------------------------------
+
+ private SortBufferContainer getUnicastDataBuffer() {
+ flushBroadcastDataBuffer();
+
+ if (unicastDataBuffer != null
+ && !unicastDataBuffer.isFinished()
+ && !unicastDataBuffer.isReleased()) {
+ return unicastDataBuffer;
+ }
+
+ unicastDataBuffer = createNewDataBuffer();
+ return unicastDataBuffer;
+ }
+
+ private SortBufferContainer getBroadcastDataBuffer() {
+ flushUnicastDataBuffer();
+
+ if (broadcastDataBuffer != null
+ && !broadcastDataBuffer.isFinished()
+ && !broadcastDataBuffer.isReleased()) {
+ return broadcastDataBuffer;
+ }
+
+ broadcastDataBuffer = createNewDataBuffer();
+ return broadcastDataBuffer;
+ }
+
+ private SortBufferContainer createNewDataBuffer() {
+ requestNetworkBuffers();
+
+ return new SortBufferContainer(
+ freeSegments,
+ this::recycleBuffer,
+ numSubpartitions,
+ bufferSizeBytes,
+ numBuffersForSort);
+ }
+
+ private void requestGuaranteedBuffers() {
+ int effectiveRequiredBuffers = effectiveNumRequestedBuffers();
+
+ while (freeSegments.size() < effectiveRequiredBuffers) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegments.add(checkNotNull(buffer).getMemorySegment());
+ if (bufferRecycler == null) {
+ bufferRecycler = buffer.getRecycler();
+ }
+ }
+ }
+
+ private void requestNetworkBuffers() {
+ requestGuaranteedBuffers();
+
+ // Use the half of the buffers for writing, and the other half for reading
+ numBuffersForSort = freeSegments.size() / 2;
+ }
+
+ private void flushDataBuffer(SortBufferContainer sortBufferContainer) {
+ if (sortBufferContainer == null
+ || sortBufferContainer.isReleased()
+ || !sortBufferContainer.hasRemaining()) {
+ return;
+ }
+ sortBufferContainer.finish();
+
+ do {
+ MemorySegment freeSegment = getFreeSegment();
+ Pair<Integer, Buffer> bufferAndSubpartitionId =
+ sortBufferContainer.readBuffer(freeSegment);
+ if (bufferAndSubpartitionId == null) {
+ if (freeSegment != null) {
+ recycleBuffer(freeSegment);
+ }
+ break;
+ }
+ addFinishedBuffer(bufferAndSubpartitionId);
+ } while (true);
+
+ releaseFreeBuffers();
+ sortBufferContainer.release();
+ }
+
+ private void flushBroadcastDataBuffer() {
+ if (broadcastDataBuffer != null) {
+ flushDataBuffer(broadcastDataBuffer);
+ broadcastDataBuffer.release();
+ broadcastDataBuffer = null;
+ }
+ }
+
+ private void flushUnicastDataBuffer() {
+ if (unicastDataBuffer != null) {
+ flushDataBuffer(unicastDataBuffer);
+ unicastDataBuffer.release();
+ unicastDataBuffer = null;
+ }
+ }
+
+ private void flushContainerWhenEndOfPartition(
+ boolean isEndOfPartition, SortBufferContainer sortBufferContainer) {
+ if (isEndOfPartition) {
+ flushDataBuffer(sortBufferContainer);
+ }
+ }
+
+ private void writeLargeRecord(
+ ByteBuffer record, int targetSubpartition, Buffer.DataType dataType) {
+
+ checkState(dataType != Buffer.DataType.EVENT_BUFFER);
+ while (record.hasRemaining()) {
+ int toCopy = Math.min(record.remaining(), bufferSizeBytes);
+ MemorySegment writeBuffer = checkNotNull(getFreeSegment());
+ writeBuffer.put(0, record, toCopy);
+
+ addFinishedBuffer(
+ Pair.of(
+ targetSubpartition,
+ new NetworkBuffer(
+ writeBuffer, checkNotNull(bufferRecycler), dataType, toCopy)));
+ }
+
+ releaseFreeBuffers();
+ }
+
+ private MemorySegment getFreeSegment() {
+ MemorySegment freeSegment = freeSegments.poll();
+ if (freeSegment == null) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegment = buffer.getMemorySegment();
+ }
+ return freeSegment;
+ }
+
+ private int effectiveNumRequestedBuffers() {
+ return Math.min(numSubpartitions + 1, numBuffers);
Review Comment:
Removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1245157313
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
Review Comment:
Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1245157313
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
Review Comment:
Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on PR #22833:
URL: https://github.com/apache/flink/pull/22833#issuecomment-1606798278
@xintongsong Thanks for helping review, I addressed the comments with a fixup commit.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1245156676
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
+ numExpectBuffers += currentBufferWrittenBytes == 0 ? 0 : 1;
+ assertThat(numReceivedFinishedBuffer.get()).isEqualTo(numExpectBuffers);
+ }
+
+ @Test
+ void testWriteLargeRecord() throws IOException {
+ int numBuffers = 15;
+ Random random = new Random();
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ AtomicInteger numReceivedBuffers = new AtomicInteger(0);
+ bufferAccumulator.setup(
+ (subpartitionIndex, buffers) -> {
+ numReceivedBuffers.set(numReceivedBuffers.get() + buffers.size());
+ buffers.forEach(Buffer::recycleBuffer);
+ });
+ ByteBuffer largeRecord = generateRandomData(BUFFER_SIZE_BYTES * numBuffers, random);
+ bufferAccumulator.receive(
+ largeRecord,
+ new TieredStorageSubpartitionId(0),
+ Buffer.DataType.DATA_BUFFER,
+ false);
+ assertThat(numReceivedBuffers.get()).isEqualTo(numBuffers);
Review Comment:
Fixed.
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
+ numExpectBuffers += currentBufferWrittenBytes == 0 ? 0 : 1;
+ assertThat(numReceivedFinishedBuffer.get()).isEqualTo(numExpectBuffers);
+ }
+
+ @Test
+ void testWriteLargeRecord() throws IOException {
+ int numBuffers = 15;
+ Random random = new Random();
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ AtomicInteger numReceivedBuffers = new AtomicInteger(0);
+ bufferAccumulator.setup(
+ (subpartitionIndex, buffers) -> {
+ numReceivedBuffers.set(numReceivedBuffers.get() + buffers.size());
Review Comment:
Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1244830703
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,269 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link SortBufferAccumulator} can help use less buffers to accumulate data, which
+ * decouples the buffer usage with the number of parallelism. The number of buffers used by the
+ * {@link SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBuffer} is full, or
+ * switching from broadcast to non-broadcast(or vice versa), the buffer in the sort buffer will be
+ * flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBuffer}. */
+ private int numBuffersForSort;
Review Comment:
Indeed, fixed.
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager tieredStorageMemoryManager =
+ createStorageMemoryManager(numBuffers);
+
+ int numExpectBytes = 0;
+ int numExpectBuffers = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, tieredStorageMemoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ int currentBufferWriteBytes = 0;
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (numExpectBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ numExpectBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ numExpectBytes += numBytes + indexEntrySize;
+ currentBufferWriteBytes += numBytes + indexEntrySize;
+ if (currentBufferWriteBytes >= BUFFER_SIZE_BYTES) {
+ currentBufferWriteBytes = currentBufferWriteBytes % BUFFER_SIZE_BYTES;
+ }
+ }
+ }
+
+ assertThat(numExpectBytes < BUFFER_SIZE_BYTES).isTrue();
+ assertThat(numReceivedFinishedBuffer.get())
+ .isEqualTo(
+ numExpectBuffers
+ + numExpectBytes / BUFFER_SIZE_BYTES
+ + (numExpectBytes % BUFFER_SIZE_BYTES == 0 ? 0 : 1));
+ }
+
+ @Test
+ void testNoBuffersForSort() throws IOException {
+ int numBuffers = 10;
+ int bufferSize = 1024;
+ Random random = new Random(1111);
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ try (SortBufferAccumulator sortBufferAccumulator =
+ new SortBufferAccumulator(1, 1, bufferSize, memoryManager)) {
+ sortBufferAccumulator.setup((subpartitionIndex, buffers) -> {});
+ assertThatThrownBy(
+ () -> {
Review Comment:
Fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] reswqa commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "reswqa (via GitHub)" <gi...@apache.org>.
reswqa commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1244602826
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/BufferAccumulator.java:
##########
@@ -44,9 +44,21 @@ public interface BufferAccumulator extends AutoCloseable {
/**
* Receives the records from tiered store producer, these records will be accumulated and
* transformed into finished buffers.
+ *
+ * <p>Note that when isBroadcast is true, for a broadcast-only partition, the subpartitionId
+ * value will always be 0. Conversely, for a non-broadcast-only partition (such as an all-to-all
Review Comment:
> for a non-broadcast-only partition (such as an all-to-all partition).
I'm a bit confused: Isn't broadcast only partition is also a `all-to-all` partition?
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager tieredStorageMemoryManager =
+ createStorageMemoryManager(numBuffers);
+
+ int numExpectBytes = 0;
+ int numExpectBuffers = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, tieredStorageMemoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ int currentBufferWriteBytes = 0;
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (numExpectBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ numExpectBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ numExpectBytes += numBytes + indexEntrySize;
+ currentBufferWriteBytes += numBytes + indexEntrySize;
Review Comment:
What is the purpose of this `currentBufferWriteBytes `?
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager tieredStorageMemoryManager =
+ createStorageMemoryManager(numBuffers);
+
+ int numExpectBytes = 0;
+ int numExpectBuffers = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, tieredStorageMemoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ int currentBufferWriteBytes = 0;
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (numExpectBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ numExpectBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ numExpectBytes += numBytes + indexEntrySize;
+ currentBufferWriteBytes += numBytes + indexEntrySize;
+ if (currentBufferWriteBytes >= BUFFER_SIZE_BYTES) {
+ currentBufferWriteBytes = currentBufferWriteBytes % BUFFER_SIZE_BYTES;
+ }
+ }
+ }
+
+ assertThat(numExpectBytes < BUFFER_SIZE_BYTES).isTrue();
+ assertThat(numReceivedFinishedBuffer.get())
+ .isEqualTo(
+ numExpectBuffers
+ + numExpectBytes / BUFFER_SIZE_BYTES
+ + (numExpectBytes % BUFFER_SIZE_BYTES == 0 ? 0 : 1));
+ }
+
+ @Test
+ void testNoBuffersForSort() throws IOException {
+ int numBuffers = 10;
+ int bufferSize = 1024;
+ Random random = new Random(1111);
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ try (SortBufferAccumulator sortBufferAccumulator =
+ new SortBufferAccumulator(1, 1, bufferSize, memoryManager)) {
+ sortBufferAccumulator.setup((subpartitionIndex, buffers) -> {});
+ assertThatThrownBy(
+ () -> {
Review Comment:
Can be replaced with expression lambda.
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,269 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link SortBufferAccumulator} can help use less buffers to accumulate data, which
+ * decouples the buffer usage with the number of parallelism. The number of buffers used by the
+ * {@link SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBuffer} is full, or
+ * switching from broadcast to non-broadcast(or vice versa), the buffer in the sort buffer will be
+ * flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
Review Comment:
```suggestion
private final TieredStorageMemoryManager memoryManager;
```
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager tieredStorageMemoryManager =
+ createStorageMemoryManager(numBuffers);
+
+ int numExpectBytes = 0;
+ int numExpectBuffers = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, tieredStorageMemoryManager)) {
Review Comment:
This seems to potentially rely on an assumption: there is only one buffer for sort, and when it is full, it will flush. It is best to add some comments to explain, otherwise the following code is not very intuitive.
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBuffer.java:
##########
@@ -0,0 +1,455 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+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;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedList;
+
+import static org.apache.flink.util.Preconditions.checkArgument;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The {@link SortBuffer} is used to accumulate the records into {@link Buffer}s. The {@link
+ * SortBuffer} allows for writing data in arbitrary subpartition orders but supports reading of data
+ * in the order grouped by subpartitions. Note that the {@link SortBuffer} only supports reading
+ * after the write process finished, and can not support reading while writing.
+ */
+public class SortBuffer {
Review Comment:
There are too many codes that duplicated with `SortBasedDataBuffer`, I wonder can we reuse it to some extent?
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,269 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link SortBufferAccumulator} can help use less buffers to accumulate data, which
+ * decouples the buffer usage with the number of parallelism. The number of buffers used by the
+ * {@link SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBuffer} is full, or
+ * switching from broadcast to non-broadcast(or vice versa), the buffer in the sort buffer will be
+ * flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBuffer}. */
+ private int numBuffersForSort;
Review Comment:
Can this field be a local variable?
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
Review Comment:
IIUC, This test was not covered the situation of largeRecord.
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager tieredStorageMemoryManager =
+ createStorageMemoryManager(numBuffers);
+
+ int numExpectBytes = 0;
+ int numExpectBuffers = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, tieredStorageMemoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ int currentBufferWriteBytes = 0;
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (numExpectBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ numExpectBytes = 0;
Review Comment:
This name is a bit strange, it feels like all the bytes expected, but in reality, it's not.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on PR #22833:
URL: https://github.com/apache/flink/pull/22833#issuecomment-1612392096
@xintongsong @reswqa Thanks for the careful review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on PR #22833:
URL: https://github.com/apache/flink/pull/22833#issuecomment-1612313953
@flinkbot run azure
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1245159077
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBuffer.java:
##########
@@ -0,0 +1,455 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+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;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedList;
+
+import static org.apache.flink.util.Preconditions.checkArgument;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The {@link SortBuffer} is used to accumulate the records into {@link Buffer}s. The {@link
+ * SortBuffer} allows for writing data in arbitrary subpartition orders but supports reading of data
+ * in the order grouped by subpartitions. Note that the {@link SortBuffer} only supports reading
+ * after the write process finished, and can not support reading while writing.
+ */
+public class SortBuffer {
Review Comment:
Introduced `SortBuffer` and `TieredStorageSortBuffer` to reuse the code.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] flinkbot commented on pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "flinkbot (via GitHub)" <gi...@apache.org>.
flinkbot commented on PR #22833:
URL: https://github.com/apache/flink/pull/22833#issuecomment-1598705770
<!--
Meta data
{
"version" : 1,
"metaDataEntries" : [ {
"hash" : "6bae94bd7afb6d8fad6634fa4f4e4ebbf45919cb",
"status" : "UNKNOWN",
"url" : "TBD",
"triggerID" : "6bae94bd7afb6d8fad6634fa4f4e4ebbf45919cb",
"triggerType" : "PUSH"
} ]
}-->
## CI report:
* 6bae94bd7afb6d8fad6634fa4f4e4ebbf45919cb UNKNOWN
<details>
<summary>Bot commands</summary>
The @flinkbot bot supports the following commands:
- `@flinkbot run azure` re-run the last Azure build
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1241087268
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+
+ flushDataBuffer(sortBufferContainer);
+ sortBufferContainer.release();
+ if (record.hasRemaining()) {
+ receive(record, subpartitionId, dataType, isBroadcast);
+ }
+ }
+
+ @Override
+ public void close() {
+ flushUnicastDataBuffer();
+ flushBroadcastDataBuffer();
+ releaseFreeBuffers();
+ releaseDataBuffer(unicastDataBuffer);
+ releaseDataBuffer(broadcastDataBuffer);
+ }
+
+ // ------------------------------------------------------------------------
+ // Internal Methods
+ // ------------------------------------------------------------------------
+
+ private SortBufferContainer getUnicastDataBuffer() {
+ flushBroadcastDataBuffer();
+
+ if (unicastDataBuffer != null
+ && !unicastDataBuffer.isFinished()
+ && !unicastDataBuffer.isReleased()) {
+ return unicastDataBuffer;
+ }
+
+ unicastDataBuffer = createNewDataBuffer();
+ return unicastDataBuffer;
+ }
+
+ private SortBufferContainer getBroadcastDataBuffer() {
+ flushUnicastDataBuffer();
+
+ if (broadcastDataBuffer != null
+ && !broadcastDataBuffer.isFinished()
+ && !broadcastDataBuffer.isReleased()) {
+ return broadcastDataBuffer;
+ }
+
+ broadcastDataBuffer = createNewDataBuffer();
+ return broadcastDataBuffer;
+ }
+
+ private SortBufferContainer createNewDataBuffer() {
+ requestNetworkBuffers();
+
+ return new SortBufferContainer(
+ freeSegments,
+ this::recycleBuffer,
+ numSubpartitions,
+ bufferSizeBytes,
+ numBuffersForSort);
+ }
+
+ private void requestGuaranteedBuffers() {
+ int effectiveRequiredBuffers = effectiveNumRequestedBuffers();
+
+ while (freeSegments.size() < effectiveRequiredBuffers) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegments.add(checkNotNull(buffer).getMemorySegment());
+ if (bufferRecycler == null) {
+ bufferRecycler = buffer.getRecycler();
+ }
+ }
+ }
+
+ private void requestNetworkBuffers() {
+ requestGuaranteedBuffers();
+
+ // Use the half of the buffers for writing, and the other half for reading
+ numBuffersForSort = freeSegments.size() / 2;
+ }
+
+ private void flushDataBuffer(SortBufferContainer sortBufferContainer) {
+ if (sortBufferContainer == null
+ || sortBufferContainer.isReleased()
+ || !sortBufferContainer.hasRemaining()) {
+ return;
+ }
+ sortBufferContainer.finish();
+
+ do {
+ MemorySegment freeSegment = getFreeSegment();
+ Pair<Integer, Buffer> bufferAndSubpartitionId =
+ sortBufferContainer.readBuffer(freeSegment);
+ if (bufferAndSubpartitionId == null) {
+ if (freeSegment != null) {
+ recycleBuffer(freeSegment);
+ }
+ break;
+ }
+ addFinishedBuffer(bufferAndSubpartitionId);
+ } while (true);
+
+ releaseFreeBuffers();
+ sortBufferContainer.release();
+ }
+
+ private void flushBroadcastDataBuffer() {
+ if (broadcastDataBuffer != null) {
+ flushDataBuffer(broadcastDataBuffer);
+ broadcastDataBuffer.release();
+ broadcastDataBuffer = null;
+ }
+ }
+
+ private void flushUnicastDataBuffer() {
+ if (unicastDataBuffer != null) {
+ flushDataBuffer(unicastDataBuffer);
+ unicastDataBuffer.release();
+ unicastDataBuffer = null;
+ }
+ }
+
+ private void flushContainerWhenEndOfPartition(
+ boolean isEndOfPartition, SortBufferContainer sortBufferContainer) {
+ if (isEndOfPartition) {
+ flushDataBuffer(sortBufferContainer);
+ }
+ }
+
+ private void writeLargeRecord(
+ ByteBuffer record, int targetSubpartition, Buffer.DataType dataType) {
+
+ checkState(dataType != Buffer.DataType.EVENT_BUFFER);
+ while (record.hasRemaining()) {
+ int toCopy = Math.min(record.remaining(), bufferSizeBytes);
+ MemorySegment writeBuffer = checkNotNull(getFreeSegment());
+ writeBuffer.put(0, record, toCopy);
+
+ addFinishedBuffer(
+ Pair.of(
+ targetSubpartition,
+ new NetworkBuffer(
+ writeBuffer, checkNotNull(bufferRecycler), dataType, toCopy)));
+ }
+
+ releaseFreeBuffers();
+ }
+
+ private MemorySegment getFreeSegment() {
+ MemorySegment freeSegment = freeSegments.poll();
+ if (freeSegment == null) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegment = buffer.getMemorySegment();
+ }
+ return freeSegment;
+ }
+
+ private int effectiveNumRequestedBuffers() {
+ return Math.min(numSubpartitions + 1, numBuffers);
Review Comment:
Actually, I think `numSubpartitions` is enough. Changed it to `numSubpartitions`.
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+
+ flushDataBuffer(sortBufferContainer);
+ sortBufferContainer.release();
+ if (record.hasRemaining()) {
+ receive(record, subpartitionId, dataType, isBroadcast);
+ }
+ }
+
+ @Override
+ public void close() {
+ flushUnicastDataBuffer();
+ flushBroadcastDataBuffer();
+ releaseFreeBuffers();
+ releaseDataBuffer(unicastDataBuffer);
+ releaseDataBuffer(broadcastDataBuffer);
+ }
+
+ // ------------------------------------------------------------------------
+ // Internal Methods
+ // ------------------------------------------------------------------------
+
+ private SortBufferContainer getUnicastDataBuffer() {
+ flushBroadcastDataBuffer();
+
+ if (unicastDataBuffer != null
+ && !unicastDataBuffer.isFinished()
+ && !unicastDataBuffer.isReleased()) {
+ return unicastDataBuffer;
+ }
+
+ unicastDataBuffer = createNewDataBuffer();
+ return unicastDataBuffer;
+ }
+
+ private SortBufferContainer getBroadcastDataBuffer() {
+ flushUnicastDataBuffer();
+
+ if (broadcastDataBuffer != null
+ && !broadcastDataBuffer.isFinished()
+ && !broadcastDataBuffer.isReleased()) {
+ return broadcastDataBuffer;
+ }
+
+ broadcastDataBuffer = createNewDataBuffer();
+ return broadcastDataBuffer;
+ }
+
+ private SortBufferContainer createNewDataBuffer() {
+ requestNetworkBuffers();
+
+ return new SortBufferContainer(
+ freeSegments,
+ this::recycleBuffer,
+ numSubpartitions,
+ bufferSizeBytes,
+ numBuffersForSort);
+ }
+
+ private void requestGuaranteedBuffers() {
+ int effectiveRequiredBuffers = effectiveNumRequestedBuffers();
+
+ while (freeSegments.size() < effectiveRequiredBuffers) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegments.add(checkNotNull(buffer).getMemorySegment());
+ if (bufferRecycler == null) {
+ bufferRecycler = buffer.getRecycler();
+ }
+ }
+ }
+
+ private void requestNetworkBuffers() {
+ requestGuaranteedBuffers();
+
+ // Use the half of the buffers for writing, and the other half for reading
+ numBuffersForSort = freeSegments.size() / 2;
+ }
+
+ private void flushDataBuffer(SortBufferContainer sortBufferContainer) {
+ if (sortBufferContainer == null
+ || sortBufferContainer.isReleased()
+ || !sortBufferContainer.hasRemaining()) {
+ return;
+ }
+ sortBufferContainer.finish();
+
+ do {
+ MemorySegment freeSegment = getFreeSegment();
+ Pair<Integer, Buffer> bufferAndSubpartitionId =
+ sortBufferContainer.readBuffer(freeSegment);
+ if (bufferAndSubpartitionId == null) {
+ if (freeSegment != null) {
+ recycleBuffer(freeSegment);
+ }
+ break;
+ }
+ addFinishedBuffer(bufferAndSubpartitionId);
+ } while (true);
+
+ releaseFreeBuffers();
+ sortBufferContainer.release();
+ }
+
+ private void flushBroadcastDataBuffer() {
+ if (broadcastDataBuffer != null) {
+ flushDataBuffer(broadcastDataBuffer);
+ broadcastDataBuffer.release();
+ broadcastDataBuffer = null;
+ }
+ }
+
+ private void flushUnicastDataBuffer() {
+ if (unicastDataBuffer != null) {
+ flushDataBuffer(unicastDataBuffer);
+ unicastDataBuffer.release();
+ unicastDataBuffer = null;
+ }
+ }
+
+ private void flushContainerWhenEndOfPartition(
+ boolean isEndOfPartition, SortBufferContainer sortBufferContainer) {
+ if (isEndOfPartition) {
+ flushDataBuffer(sortBufferContainer);
+ }
+ }
+
+ private void writeLargeRecord(
+ ByteBuffer record, int targetSubpartition, Buffer.DataType dataType) {
+
+ checkState(dataType != Buffer.DataType.EVENT_BUFFER);
+ while (record.hasRemaining()) {
+ int toCopy = Math.min(record.remaining(), bufferSizeBytes);
+ MemorySegment writeBuffer = checkNotNull(getFreeSegment());
+ writeBuffer.put(0, record, toCopy);
+
+ addFinishedBuffer(
+ Pair.of(
+ targetSubpartition,
+ new NetworkBuffer(
+ writeBuffer, checkNotNull(bufferRecycler), dataType, toCopy)));
+ }
+
+ releaseFreeBuffers();
+ }
+
+ private MemorySegment getFreeSegment() {
+ MemorySegment freeSegment = freeSegments.poll();
+ if (freeSegment == null) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegment = buffer.getMemorySegment();
+ }
+ return freeSegment;
+ }
+
+ private int effectiveNumRequestedBuffers() {
+ return Math.min(numSubpartitions + 1, numBuffers);
+ }
+
+ private void releaseDataBuffer(SortBufferContainer sortBufferContainer) {
+ if (sortBufferContainer != null) {
+ sortBufferContainer.release();
+ }
+ }
+
+ private void addFinishedBuffer(Pair<Integer, Buffer> bufferAndSubpartitionId) {
Review Comment:
Fixed.
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+
+ flushDataBuffer(sortBufferContainer);
+ sortBufferContainer.release();
+ if (record.hasRemaining()) {
+ receive(record, subpartitionId, dataType, isBroadcast);
+ }
+ }
+
+ @Override
+ public void close() {
+ flushUnicastDataBuffer();
+ flushBroadcastDataBuffer();
+ releaseFreeBuffers();
+ releaseDataBuffer(unicastDataBuffer);
+ releaseDataBuffer(broadcastDataBuffer);
+ }
+
+ // ------------------------------------------------------------------------
+ // Internal Methods
+ // ------------------------------------------------------------------------
+
+ private SortBufferContainer getUnicastDataBuffer() {
+ flushBroadcastDataBuffer();
+
+ if (unicastDataBuffer != null
+ && !unicastDataBuffer.isFinished()
+ && !unicastDataBuffer.isReleased()) {
+ return unicastDataBuffer;
+ }
+
+ unicastDataBuffer = createNewDataBuffer();
+ return unicastDataBuffer;
+ }
+
+ private SortBufferContainer getBroadcastDataBuffer() {
+ flushUnicastDataBuffer();
+
+ if (broadcastDataBuffer != null
+ && !broadcastDataBuffer.isFinished()
+ && !broadcastDataBuffer.isReleased()) {
+ return broadcastDataBuffer;
+ }
+
+ broadcastDataBuffer = createNewDataBuffer();
+ return broadcastDataBuffer;
+ }
+
+ private SortBufferContainer createNewDataBuffer() {
+ requestNetworkBuffers();
+
+ return new SortBufferContainer(
+ freeSegments,
+ this::recycleBuffer,
+ numSubpartitions,
+ bufferSizeBytes,
+ numBuffersForSort);
+ }
+
+ private void requestGuaranteedBuffers() {
+ int effectiveRequiredBuffers = effectiveNumRequestedBuffers();
+
+ while (freeSegments.size() < effectiveRequiredBuffers) {
+ BufferBuilder bufferBuilder = storeMemoryManager.requestBufferBlocking(this);
+ Buffer buffer = bufferBuilder.createBufferConsumerFromBeginning().build();
+ freeSegments.add(checkNotNull(buffer).getMemorySegment());
+ if (bufferRecycler == null) {
+ bufferRecycler = buffer.getRecycler();
+ }
+ }
+ }
+
+ private void requestNetworkBuffers() {
+ requestGuaranteedBuffers();
+
+ // Use the half of the buffers for writing, and the other half for reading
+ numBuffersForSort = freeSegments.size() / 2;
+ }
+
+ private void flushDataBuffer(SortBufferContainer sortBufferContainer) {
+ if (sortBufferContainer == null
+ || sortBufferContainer.isReleased()
+ || !sortBufferContainer.hasRemaining()) {
+ return;
+ }
+ sortBufferContainer.finish();
+
+ do {
+ MemorySegment freeSegment = getFreeSegment();
+ Pair<Integer, Buffer> bufferAndSubpartitionId =
+ sortBufferContainer.readBuffer(freeSegment);
+ if (bufferAndSubpartitionId == null) {
+ if (freeSegment != null) {
+ recycleBuffer(freeSegment);
+ }
+ break;
+ }
+ addFinishedBuffer(bufferAndSubpartitionId);
+ } while (true);
+
+ releaseFreeBuffers();
+ sortBufferContainer.release();
+ }
+
+ private void flushBroadcastDataBuffer() {
+ if (broadcastDataBuffer != null) {
+ flushDataBuffer(broadcastDataBuffer);
+ broadcastDataBuffer.release();
+ broadcastDataBuffer = null;
+ }
+ }
+
+ private void flushUnicastDataBuffer() {
+ if (unicastDataBuffer != null) {
+ flushDataBuffer(unicastDataBuffer);
+ unicastDataBuffer.release();
+ unicastDataBuffer = null;
+ }
+ }
+
+ private void flushContainerWhenEndOfPartition(
+ boolean isEndOfPartition, SortBufferContainer sortBufferContainer) {
+ if (isEndOfPartition) {
+ flushDataBuffer(sortBufferContainer);
+ }
+ }
Review Comment:
Removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1241087057
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+
+ flushDataBuffer(sortBufferContainer);
+ sortBufferContainer.release();
+ if (record.hasRemaining()) {
Review Comment:
yes, replace this with a `checkState`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1241086989
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
Review Comment:
Fixed with a `currentDataBuffer `, a new `switchCurrentDataBufferIfNeeded(isBroadcast)`, a field of `isBroadcastDataBuffer`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1241087452
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferContainer.java:
##########
@@ -0,0 +1,454 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+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;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedList;
+
+import static org.apache.flink.util.Preconditions.checkArgument;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The buffer container for accumulating the records into {@link Buffer}s. After accumulating, the
+ * {@link SortBufferAccumulator} will read the sorted buffers.
+ */
+public class SortBufferContainer {
Review Comment:
1. Renamed it to `SortBuffer`.
2. Fixed the doc of `SortBuffer`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1241083885
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.core.memory.MemorySegment;
+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.BufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * The sort-based implementation of the {@link BufferAccumulator}. The {@link BufferAccumulator}
+ * receives the records from {@link TieredStorageProducerClient} and the records will accumulate and
+ * transform to finished buffers. The accumulated buffers will be transferred to the corresponding
+ * tier dynamically.
+ *
+ * <p>The {@link BufferAccumulator} can help use less buffers to accumulate data, which decouples
+ * the buffer usage with the number of parallelism. The number of buffers used by the {@link
+ * SortBufferAccumulator} will be numBuffers at most. Once the {@link SortBufferContainer} is full,
+ * or receiving a different type of buffer, or receiving the end-of-partition event, the buffer in
+ * the sort buffer container will be flushed to the tiers.
+ *
+ * <p>Note that this class need not be thread-safe, because it should only be accessed from the main
+ * thread.
+ */
+public class SortBufferAccumulator implements BufferAccumulator {
+
+ /** The number of the subpartitions. */
+ private final int numSubpartitions;
+
+ /** The total number of the buffers used by the {@link SortBufferAccumulator}. */
+ private final int numBuffers;
+
+ /** The byte size of one single buffer. */
+ private final int bufferSizeBytes;
+
+ /** The empty buffers without storing data. */
+ private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
+
+ /** The memory manager of the tiered storage. */
+ private final TieredStorageMemoryManager storeMemoryManager;
+
+ /** The number of buffers for sorting used in the {@link SortBufferContainer}. */
+ private int numBuffersForSort;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating broadcast data. Note that this can be null
+ * before using it to store records, and this buffer container will be released once flushed.
+ */
+ @Nullable private SortBufferContainer broadcastDataBuffer;
+
+ /**
+ * The {@link SortBufferContainer} for accumulating non-broadcast data. Note that this can be
+ * null before using it to store records, and this buffer container will be released once
+ * flushed.
+ */
+ @Nullable private SortBufferContainer unicastDataBuffer;
+
+ /**
+ * The buffer recycler. Note that this can be null before requesting buffers from the memory
+ * manager.
+ */
+ @Nullable private BufferRecycler bufferRecycler;
+
+ /**
+ * The {@link SortBufferAccumulator}'s accumulated buffer flusher is not prepared during
+ * construction, requiring the field to be initialized during setup. Therefore, it is necessary
+ * to verify whether this field is null before using it.
+ */
+ @Nullable
+ private BiConsumer<TieredStorageSubpartitionId, List<Buffer>> accumulatedBufferFlusher;
+
+ public SortBufferAccumulator(
+ int numSubpartitions,
+ int numBuffers,
+ int bufferSizeBytes,
+ TieredStorageMemoryManager storeMemoryManager) {
+ this.numSubpartitions = numSubpartitions;
+ this.bufferSizeBytes = bufferSizeBytes;
+ this.numBuffers = numBuffers;
+ this.storeMemoryManager = storeMemoryManager;
+ }
+
+ @Override
+ public void setup(BiConsumer<TieredStorageSubpartitionId, List<Buffer>> bufferFlusher) {
+ this.accumulatedBufferFlusher = bufferFlusher;
+ }
+
+ @Override
+ public void receive(
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
+ throws IOException {
+ int targetSubpartition = subpartitionId.getSubpartitionId();
+ SortBufferContainer sortBufferContainer =
+ isBroadcast ? getBroadcastDataBuffer() : getUnicastDataBuffer();
+ if (!sortBufferContainer.writeRecord(record, targetSubpartition, dataType)) {
+ return;
+ }
+
+ if (!sortBufferContainer.hasRemaining()) {
+ sortBufferContainer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
Review Comment:
Added docs to describe this situation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1241083625
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/BufferAccumulator.java:
##########
@@ -46,7 +46,10 @@ public interface BufferAccumulator extends AutoCloseable {
* transformed into finished buffers.
*/
void receive(
- ByteBuffer record, TieredStorageSubpartitionId subpartitionId, Buffer.DataType dataType)
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
Review Comment:
For a broadcast-only partition, When isBroadcast is true, the subpartitionId is 0.
For a non-broadcast-only partition (for exampe, a all-to-all edge), this will be all the subpartitionId range from 0 to numSupartitions(This is from a for loop).
##########
flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/BufferAccumulator.java:
##########
@@ -46,7 +46,10 @@ public interface BufferAccumulator extends AutoCloseable {
* transformed into finished buffers.
*/
void receive(
- ByteBuffer record, TieredStorageSubpartitionId subpartitionId, Buffer.DataType dataType)
+ ByteBuffer record,
+ TieredStorageSubpartitionId subpartitionId,
+ Buffer.DataType dataType,
+ boolean isBroadcast)
Review Comment:
For a broadcast-only partition, When isBroadcast is true, the subpartitionId is 0.
For a non-broadcast-only partition (for example, a all-to-all edge), this will be all the subpartitionId range from 0 to numSupartitions(This is from a for loop).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1244833430
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,148 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
Review Comment:
Added a test for large record.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [flink] TanYuxin-tyx commented on a diff in pull request #22833: [FLINK-32283][network] Implement the SortBufferAccumulator for the tiered storage
Posted by "TanYuxin-tyx (via GitHub)" <gi...@apache.org>.
TanYuxin-tyx commented on code in PR #22833:
URL: https://github.com/apache/flink/pull/22833#discussion_r1245157005
##########
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java:
##########
@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+
+import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.runtime.io.network.partition.hybrid.tiered.TieredStorageTestUtils.generateRandomData;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link SortBufferAccumulator}. */
+class SortBufferAccumulatorTest {
+
+ private static final int NUM_TOTAL_BUFFERS = 1000;
+
+ private static final int BUFFER_SIZE_BYTES = 1024;
+
+ private static final float NUM_BUFFERS_TRIGGER_FLUSH_RATIO = 0.6f;
+
+ private NetworkBufferPool globalPool;
+
+ @BeforeEach
+ void before() {
+ globalPool = new NetworkBufferPool(NUM_TOTAL_BUFFERS, BUFFER_SIZE_BYTES);
+ }
+
+ @AfterEach
+ void after() {
+ globalPool.destroy();
+ }
+
+ @Test
+ void testAccumulateRecordsAndGenerateBuffers() throws IOException {
+ int numBuffers = 10;
+ int numRecords = 1000;
+ int indexEntrySize = 16;
+ TieredStorageSubpartitionId subpartitionId = new TieredStorageSubpartitionId(0);
+ Random random = new Random(1234);
+ TieredStorageMemoryManager memoryManager = createStorageMemoryManager(numBuffers);
+
+ int numExpectBuffers = 0;
+ int currentBufferWrittenBytes = 0;
+ AtomicInteger numReceivedFinishedBuffer = new AtomicInteger(0);
+
+ // The test use only one buffer for sort, and when it is full, the sort buffer will be
+ // flushed.
+ try (SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(1, 2, BUFFER_SIZE_BYTES, memoryManager)) {
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
+ buffers.forEach(
+ buffer -> {
+ numReceivedFinishedBuffer.incrementAndGet();
+ buffer.recycleBuffer();
+ })));
+ boolean isBroadcastForPreviousRecord = false;
+ for (int i = 0; i < numRecords; i++) {
+ int numBytes = random.nextInt(BUFFER_SIZE_BYTES) + 1;
+ ByteBuffer record = generateRandomData(numBytes, random);
+ boolean isBroadcast = random.nextBoolean();
+ bufferAccumulator.receive(
+ record, subpartitionId, Buffer.DataType.DATA_BUFFER, isBroadcast);
+ if (currentBufferWrittenBytes + numBytes + indexEntrySize > BUFFER_SIZE_BYTES
+ || i > 0 && isBroadcastForPreviousRecord != isBroadcast) {
+ numExpectBuffers++;
+ currentBufferWrittenBytes = 0;
+ }
+
+ isBroadcastForPreviousRecord = isBroadcast;
+ currentBufferWrittenBytes += numBytes + indexEntrySize;
+ }
+ }
+
+ assertThat(currentBufferWrittenBytes < BUFFER_SIZE_BYTES).isTrue();
+ numExpectBuffers += currentBufferWrittenBytes == 0 ? 0 : 1;
+ assertThat(numReceivedFinishedBuffer.get()).isEqualTo(numExpectBuffers);
Review Comment:
Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org