(celeborn) branch main updated: [CELEBORN-2033] updateProduceBytes should be called even if updateProduceBytes throws exception

[fe...@apache.org]

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

feiwang 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 bbd3bb481 [CELEBORN-2033] updateProduceBytes should be called even if updateProduceBytes throws exception
bbd3bb481 is described below

commit bbd3bb4814a7505d6b7499ab1029997958564705
Author: Xianming Lei <xi...@shopee.com>
AuthorDate: Wed Jun 11 10:54:24 2025 -0700

    [CELEBORN-2033] updateProduceBytes should be called even if updateProduceBytes throws exception
    
    ### What changes were proposed in this pull request?
    updateProduceBytes should be called even if updateProduceBytes throws exception
    
    ### Why are the changes needed?
    To make UserProduceSpeed ​​metrics more accurate
    
    ### Does this PR introduce _any_ user-facing change?
    No.
    
    ### How was this patch tested?
    Existing UTs.
    
    Closes #3322 from leixm/CELEBORN-2033.
    
    Authored-by: Xianming Lei <xi...@shopee.com>
    Signed-off-by: Wang, Fei <fw...@ebay.com>
---
 .../apache/celeborn/service/deploy/worker/storage/TierWriter.scala  | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala
index 278b6f945..00dd90bd5 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala
@@ -412,6 +412,8 @@ class LocalTierWriter(
 
   override def writeInternal(buf: ByteBuf): Unit = {
     val numBytes = buf.readableBytes()
+    if (userCongestionControlContext != null)
+      userCongestionControlContext.updateProduceBytes(numBytes)
     val flushBufferReadableBytes = flushBuffer.readableBytes
     if (flushBufferReadableBytes != 0 && flushBufferReadableBytes + numBytes >= flusherBufferSize) {
       flush(false)
@@ -422,13 +424,9 @@ class LocalTierWriter(
     } catch {
       case oom: OutOfMemoryError =>
         MemoryManager.instance.incrementDiskBuffer(numBytes)
-        if (userCongestionControlContext != null)
-          userCongestionControlContext.updateProduceBytes(numBytes)
         throw oom
     }
     MemoryManager.instance.incrementDiskBuffer(numBytes)
-    if (userCongestionControlContext != null)
-      userCongestionControlContext.updateProduceBytes(numBytes)
   }
 
   override def evict(file: TierWriterBase): Unit = ???