[#1727] feat(server): Introduce local allocation buffer to store blocks in memory #2492

xianjingfeng · 2025-05-28T10:16:59Z

What changes were proposed in this pull request?

Introduce local allocation buffer to store blocks in memory.

Why are the changes needed?

Fix: #1727

Does this PR introduce any user-facing change?

set rss.server.buffer.lab.enable to true in server.conf

How was this patch tested?

CI and verify in production environment

xianjingfeng · 2025-05-28T10:31:44Z

Test result of terasort(data size:1T, parallelism:8000)

github-actions · 2025-05-28T10:54:06Z

Test Results

3 049 files + 30 3 049 suites +30 6h 47m 55s ⏱️ -21s
1 186 tests + 8 1 185 ✅ + 8 1 💤 ±0 0 ❌ ±0
15 042 runs +120 15 027 ✅ +120 15 💤 ±0 0 ❌ ±0

Results for commit 497c524. ± Comparison against base commit b45e986.

This pull request removes 4 and adds 12 tests. Note that renamed tests count towards both.

org.apache.uniffle.common.ShufflePartitionedBlockTest ‑ testNotEquals{int, long, long, int}[1]
org.apache.uniffle.common.ShufflePartitionedBlockTest ‑ testNotEquals{int, long, long, int}[2]
org.apache.uniffle.common.ShufflePartitionedBlockTest ‑ testNotEquals{int, long, long, int}[3]
org.apache.uniffle.common.ShufflePartitionedBlockTest ‑ testNotEquals{int, long, long, int}[4]

org.apache.uniffle.server.buffer.LABShuffleBufferWithLinkedListTest ‑ appendMultiBlocksTest
org.apache.uniffle.server.buffer.LABShuffleBufferWithLinkedListTest ‑ appendRepeatBlockTest
org.apache.uniffle.server.buffer.LABShuffleBufferWithLinkedListTest ‑ appendTest
org.apache.uniffle.server.buffer.LABShuffleBufferWithLinkedListTest ‑ getShuffleDataTest
org.apache.uniffle.server.buffer.LABShuffleBufferWithLinkedListTest ‑ getShuffleDataWithExpectedTaskIdsFilterTest
org.apache.uniffle.server.buffer.LABShuffleBufferWithLinkedListTest ‑ getShuffleDataWithLocalOrderTest
org.apache.uniffle.server.buffer.LABShuffleBufferWithLinkedListTest ‑ toFlushEventTest
org.apache.uniffle.server.buffer.LABShuffleBufferWithSkipListTest ‑ appendMultiBlocksTest
org.apache.uniffle.server.buffer.LABShuffleBufferWithSkipListTest ‑ appendRepeatBlockTest
org.apache.uniffle.server.buffer.LABShuffleBufferWithSkipListTest ‑ appendTest
…

♻️ This comment has been updated with latest results.

jerqi · 2025-05-28T11:05:00Z

cc @frankliee

jerqi · 2025-05-29T03:04:29Z

@rickyma Could you help me review this pull request?

server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java

server/src/main/java/org/apache/uniffle/server/buffer/AbstractShuffleBuffer.java

jerqi · 2025-05-29T03:52:11Z

server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java

@@ -226,9 +238,9 @@ public StatusCode registerBuffer(
      ShuffleServerMetrics.gaugeTotalPartitionNum.inc();
      ShuffleBuffer shuffleBuffer;
      if (shuffleBufferType == ShuffleBufferType.SKIP_LIST) {
-        shuffleBuffer = new ShuffleBufferWithSkipList();


We can add the class SkipListLabShuffleBuffer and LinkedListLabShuffleBuffer, and add ShuffleBufferFactory.

xianjingfeng · 2025-05-30T01:38:25Z

Do you mean to do it like below? I tried this solution and found that many methods need to be implemented. And the code looks inelegant.@jerqi

And I also tried the solution like below. And i think it's better.

jerqi · 2025-05-30T02:23:42Z

Do you mean to do it like below? I tried this solution and found that many methods need to be implemented. And the code looks inelegant.@jerqi And I also tried the solution like below. And i think it's better.

The second way is ok for me, too. You can extract an interface SupportsLAB for second way.

zuston · 2025-06-03T02:25:37Z

Good to see this PR for the performance improvement. I will take a look in the later of this day

zuston

Looks good for this abstrction

common/src/main/java/org/apache/uniffle/common/util/ByteBufUtils.java

jerqi

LGTM, cc @frankliee , Could you take a look? I will wait util tomorrow.

server/src/main/java/org/apache/uniffle/server/buffer/lab/LAB.java

advancedxy · 2025-06-05T11:12:35Z

I am quite busy these days. But I will take a look at this later today.

advancedxy

Thanks for this, the comparison result is promising.
Left some minor comments, others lgtm.

advancedxy · 2025-06-05T11:32:07Z

server/src/main/java/org/apache/uniffle/server/buffer/lab/Chunk.java

+
+  abstract void allocateDataBuffer();
+
+  public int alloc(int size) {


the implementation and the method name doesn't seem aligned? Could you add some java doc and/or change the method name here?

server/src/main/java/org/apache/uniffle/server/buffer/lab/ChunkCreator.java

advancedxy · 2025-06-05T11:38:33Z

server/src/main/java/org/apache/uniffle/server/buffer/lab/ChunkCreator.java

+  private Chunk createChunk(boolean pool, int size) {
+    Chunk chunk;
+    int id = chunkID.getAndIncrement();
+    assert id > 0;


this seem redundant?

Maybe chunk id will be less than 0
https://github.com/apache/hbase/blob/d84824b350377854a5f60759e775ab077847477f/hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/ChunkCreator.java#L240

I don't think so .. chunkID should always greater than zero?

Even if it's possible to overflow, we should use Preconditions.check instead of assert, which will be no-op if assertion if not enabled.

It's possible to overflow, let's use Preconditions.check.

server/src/main/java/org/apache/uniffle/server/buffer/lab/ChunkCreator.java

advancedxy · 2025-06-05T11:48:29Z

common/src/main/java/org/apache/uniffle/common/ShufflePartitionedBlock.java

-    return dataLength == that.dataLength
-        && crc == that.crc
-        && blockId == that.blockId
-        && data.equals(that.data);


Why this change? I think at least we should keep dataLength and crc check?

common/src/main/java/org/apache/uniffle/common/ShufflePartitionedBlock.java

server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java

server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java

advancedxy · 2025-06-05T11:52:34Z

server/src/main/java/org/apache/uniffle/server/buffer/lab/LAB.java

+  private Chunk currChunk;
+
+  List<Integer> chunks = new LinkedList<>();
+  private final int maxAlloc;


nit: this field name is a bit of vague, could this be called as capacity or something similar?

It means when a block's size is too big (bigger than maxAlloc), it will not be allocated on LAB.

zuston · 2025-06-10T07:47:29Z

Gentle ping @advancedxy

advancedxy

One minor comment, others lgtm.

…apacityRatio (#2554) ### What changes were proposed in this pull request? Change the default value of chunkPoolCapacityRatio ### Why are the changes needed? The proportion and the frequency of small blocks is not high. If this value is set too high, it may cause off-heap memory overflow. Fix: #2492 ### Does this PR introduce any user-facing change? No. ### How was this patch tested? Verify in production environment.

xianjingfeng added 5 commits May 28, 2025 15:59

add lab and ChunkCreator

7223603

optimize

4ca2871

support LAB

55e9dc7

fix code style

c1bd89f

nit

c5214cf

xianjingfeng requested review from maobaolong, zuston, advancedxy, jerqi and rickyma and removed request for maobaolong and zuston May 28, 2025 10:19

fix ut

17281a1

jerqi reviewed May 29, 2025

View reviewed changes

server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java Outdated Show resolved Hide resolved

jerqi reviewed May 29, 2025

View reviewed changes

server/src/main/java/org/apache/uniffle/server/buffer/AbstractShuffleBuffer.java Outdated Show resolved Hide resolved

jerqi reviewed May 29, 2025

View reviewed changes

fix ut

786eddb

xianjingfeng changed the title ~~[#1727] feat(server): Introduce local allocation buffer to store blocks in memory~~ [WIP][#1727] feat(server): Introduce local allocation buffer to store blocks in memory May 29, 2025

xianjingfeng added 3 commits May 29, 2025 17:32

fix ut

3d9c7c3

nit

1772d1a

revert ReassignAndStageRetryTest

8b900c3

optimize

3959c47

optimize

9c556f9

zuston reviewed Jun 4, 2025

View reviewed changes

common/src/main/java/org/apache/uniffle/common/util/ByteBufUtils.java Outdated Show resolved Hide resolved

zuston previously approved these changes Jun 4, 2025

View reviewed changes

rename createByteBufFrom to copy

5f98abe

xianjingfeng dismissed zuston’s stale review via 5f98abe June 5, 2025 02:05

optimize

9c117c0

jerqi previously approved these changes Jun 5, 2025

View reviewed changes

frankliee reviewed Jun 5, 2025

View reviewed changes

server/src/main/java/org/apache/uniffle/server/buffer/lab/LAB.java Show resolved Hide resolved

xianjingfeng added 2 commits June 5, 2025 11:17

add comments for LAB

b28b627

set rss.server.buffer.lab.chunkSize to 100K

f981e6e

xianjingfeng dismissed jerqi’s stale review via f981e6e June 5, 2025 03:21

xianjingfeng requested a review from frankliee June 5, 2025 03:23

zuston previously approved these changes Jun 5, 2025

View reviewed changes

jerqi previously approved these changes Jun 5, 2025

View reviewed changes

advancedxy reviewed Jun 5, 2025

View reviewed changes

nit

9cd9949

xianjingfeng dismissed stale reviews from jerqi and zuston via 9cd9949 June 6, 2025 09:13

xianjingfeng requested a review from advancedxy June 6, 2025 09:20

advancedxy reviewed Jun 10, 2025

View reviewed changes

use Preconditions.check instead of assert

497c524

jerqi requested a review from advancedxy June 11, 2025 03:42

advancedxy approved these changes Jun 11, 2025

View reviewed changes

jerqi merged commit c83a1b5 into apache:master Jun 12, 2025
41 checks passed

xianjingfeng mentioned this pull request Jun 27, 2025

[Improvement][LAB] Add some metrics for LAB #2525

Closed

3 tasks

xianjingfeng mentioned this pull request Jul 24, 2025

[#2492][FOLLOWUP] improvement: change the default value of chunkPoolCapacityRatio #2554

Merged


		abstract void allocateDataBuffer();

		public int alloc(int size) {

[#1727] feat(server): Introduce local allocation buffer to store blocks in memory #2492

[#1727] feat(server): Introduce local allocation buffer to store blocks in memory #2492

Uh oh!

Conversation

xianjingfeng commented May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Uh oh!

xianjingfeng commented May 28, 2025

Uh oh!

github-actions bot commented May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Test Results

Uh oh!

jerqi commented May 28, 2025

Uh oh!

jerqi commented May 29, 2025

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xianjingfeng commented May 30, 2025

Uh oh!

jerqi commented May 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

zuston commented Jun 3, 2025

Uh oh!

zuston left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jerqi left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

advancedxy commented Jun 5, 2025

Uh oh!

advancedxy left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zuston commented Jun 10, 2025

Uh oh!

advancedxy left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

xianjingfeng commented May 28, 2025 •

edited

Loading

github-actions bot commented May 28, 2025 •

edited

Loading

jerqi commented May 30, 2025 •

edited

Loading