(celeborn) branch main updated: [CELEBORN-2321] Avoid locking disk writers during memory split checks

[nicholasjiang]

This is an automated email from the ASF dual-hosted git repository.

SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git


The following commit(s) were added to refs/heads/main by this push:
     new 70ff956d0 [CELEBORN-2321] Avoid locking disk writers during memory 
split checks
70ff956d0 is described below

commit 70ff956d049ab071d6ff93ff55acc4fea0a46635
Author: Chao Sun <c...@openai.com>
AuthorDate: Sun May 10 12:35:50 2026 +0800

    [CELEBORN-2321] Avoid locking disk writers during memory split checks
    
    ### Why are the changes needed?
    
    `needHardSplitForMemoryShuffleStorage()` runs on the push path. Disk-backed 
writers can never require this memory-only split check, but the method 
currently acquires the writer lock before returning `false`. For the common 
disk-backed case, that adds avoidable contention with writes and evictions on a 
hot path.
    
    ### What changes were proposed in this PR?
    
    This PR adds an unlocked fast path for non-memory writers so they return 
immediately without taking the `PartitionDataWriter` monitor. For memory-backed 
writers, it rechecks `currentTierWriter` after entering the synchronized block 
before evaluating the existing hard-split conditions, which preserves the 
original behavior if the writer tier changes concurrently.
    
    ### Does this PR resolve a correctness bug?
    
    No.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    - Attempted `build/mvn -pl worker -am -DskipTests compile` on current 
`main`.
    - The Maven reactor fails before reaching `worker` because 
`celeborn-master_2.12` cannot resolve snapshot test-jar artifacts for 
`celeborn-common_2.12` and `celeborn-service_2.12`; that failure is unrelated 
to this change.
    
    Closes #3680 from 
sunchao/dev/chao/codex/celeborn-fast-memory-split-check-oss-main.
    
    Authored-by: Chao Sun <c...@openai.com>
    Signed-off-by: SteNicholas <programg...@163.com>
---
 .../deploy/worker/storage/PartitionDataWriter.java   | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git 
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
 
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
index 5fbb6e976..cf7ee4e5f 100644
--- 
a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
+++ 
b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
@@ -122,13 +122,23 @@ public class PartitionDataWriter implements 
DeviceObserver {
     currentTierWriter.flush(false, false);
   }
 
-  public synchronized boolean needHardSplitForMemoryShuffleStorage() {
-    if (!(currentTierWriter instanceof MemoryTierWriter)) {
+  public boolean needHardSplitForMemoryShuffleStorage() {
+    // Disk-backed writers never need this memory-only split check. Avoid 
contending with writes and
+    // evictions on the hot push path for the common case.
+    TierWriterBase tierWriter = currentTierWriter;
+    if (!(tierWriter instanceof MemoryTierWriter)) {
       return false;
     }
-    return !storageManager.localOrDfsStorageAvailable()
-        && (currentTierWriter.fileInfo().getFileLength() > 
memoryFileStorageMaxFileSize
-            || !MemoryManager.instance().memoryFileStorageAvailable());
+
+    synchronized (this) {
+      tierWriter = currentTierWriter;
+      if (!(tierWriter instanceof MemoryTierWriter)) {
+        return false;
+      }
+      return !storageManager.localOrDfsStorageAvailable()
+          && (tierWriter.fileInfo().getFileLength() > 
memoryFileStorageMaxFileSize
+              || !MemoryManager.instance().memoryFileStorageAvailable());
+    }
   }
 
   public synchronized void write(ByteBuf data) throws IOException {