(celeborn) branch main updated: [CELEBORN-2360] Fix ReviseLostShuffles RPC deserialization

[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 8903a31d8 [CELEBORN-2360] Fix ReviseLostShuffles RPC deserialization
8903a31d8 is described below

commit 8903a31d81dedf3b9bd52e5773cbb98220d6b14d
Author: afterincomparableyum 
<224495379+afterincomparable...@users.noreply.github.com>
AuthorDate: Tue Jun 23 10:56:50 2026 +0800

    [CELEBORN-2360] Fix ReviseLostShuffles RPC deserialization
    
    ### What changes were proposed in this pull request?
    
    CELEBORN-1601 added ReviseLostShuffles support but only wired up the encode 
side (toTransportMessage), the master handler, and the client sender. The 
corresponding decode cases in ControlMessages.fromTransportMessage were never 
added. Since that match has no default case, receiving a REVISE_LOST_SHUFFLES 
(89) request on the master or a REVISE_LOST_SHUFFLES_RESPONSE (90) on the 
client throws a MatchError, so the feature is broken over RPC. The existing 
test exercises handleReviseLostSh [...]
    
    To fix, I added the missing REVISE_LOST_SHUFFLES_VALUE and 
REVISE_LOST_SHUFFLES_RESPONSE_VALUE cases so the messages deserialize into 
PbReviseLostShuffles / PbReviseLostShufflesResponse.
    
    I also added a round trip test in UtilsSuite covering both messages.
    
    ### Why are the changes needed?
    
    To fix an existing bug.
    
    ### Does this PR resolve a correctness bug?
    
    - [x] Yes
    
    ### Does this PR introduce _any_ user-facing change?
    
    - [ ] Yes
    
    ### How was this patch tested?
    
    CI/CD
    
    Closes #3735 from afterincomparableyum/CELEBORN-2360.
    
    Authored-by: afterincomparableyum 
<224495379+afterincomparable...@users.noreply.github.com>
    Signed-off-by: Nicholas Jiang <programg...@163.com>
---
 .../common/protocol/message/ControlMessages.scala     |  6 ++++++
 .../org/apache/celeborn/common/util/UtilsSuite.scala  | 19 +++++++++++++++++--
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git 
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
 
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
index e12d4b697..f1e34aa54 100644
--- 
a/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
+++ 
b/common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala
@@ -1331,6 +1331,12 @@ object ControlMessages extends Logging {
       case BATCH_UNREGISTER_SHUFFLES_VALUE =>
         PbBatchUnregisterShuffles.parseFrom(message.getPayload)
 
+      case REVISE_LOST_SHUFFLES_VALUE =>
+        PbReviseLostShuffles.parseFrom(message.getPayload)
+
+      case REVISE_LOST_SHUFFLES_RESPONSE_VALUE =>
+        PbReviseLostShufflesResponse.parseFrom(message.getPayload)
+
       case UNREGISTER_SHUFFLE_RESPONSE_VALUE =>
         PbUnregisterShuffleResponse.parseFrom(message.getPayload)
 
diff --git 
a/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala 
b/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
index 8be472b64..c83a83b95 100644
--- a/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
+++ b/common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala
@@ -29,8 +29,8 @@ import 
org.apache.celeborn.common.client.{MasterEndpointResolver, StaticMasterEn
 import org.apache.celeborn.common.exception.CelebornException
 import org.apache.celeborn.common.identity.DefaultIdentityProvider
 import org.apache.celeborn.common.network.protocol.SerdeVersion
-import org.apache.celeborn.common.protocol.{PartitionLocation, 
TransportModuleConstants}
-import 
org.apache.celeborn.common.protocol.message.ControlMessages.{GetReducerFileGroupResponse,
 MapperEnd}
+import org.apache.celeborn.common.protocol.{PartitionLocation, 
PbReviseLostShuffles, PbReviseLostShufflesResponse, TransportModuleConstants}
+import 
org.apache.celeborn.common.protocol.message.ControlMessages.{GetReducerFileGroupResponse,
 MapperEnd, ReviseLostShuffles, ReviseLostShufflesResponse}
 import org.apache.celeborn.common.protocol.message.StatusCode
 
 class UtilsSuite extends CelebornFunSuite {
@@ -174,6 +174,21 @@ class UtilsSuite extends CelebornFunSuite {
     mapperEnd.bytesWrittenPerPartition.array should contain 
theSameElementsInOrderAs mapperEndTrans.bytesWrittenPerPartition
   }
 
+  test("ReviseLostShuffles class convert with pb") {
+    val req = ReviseLostShuffles("app-1", util.Arrays.asList[Integer](1, 2, 
3), "req-1")
+    val reqTrans = Utils.fromTransportMessage(Utils.toTransportMessage(req))
+      .asInstanceOf[PbReviseLostShuffles]
+    assert(req.getAppId == reqTrans.getAppId)
+    assert(req.getLostShufflesList == reqTrans.getLostShufflesList)
+    assert(req.getRequestId == reqTrans.getRequestId)
+
+    val resp = ReviseLostShufflesResponse(true, "ok")
+    val respTrans = Utils.fromTransportMessage(Utils.toTransportMessage(resp))
+      .asInstanceOf[PbReviseLostShufflesResponse]
+    assert(resp.getSuccess == respTrans.getSuccess)
+    assert(resp.getMessage == respTrans.getMessage)
+  }
+
   test("validate HDFS compatible fs path") {
     val hdfsPath = "hdfs://xxx:9000/xxxx/xx-xx/x-x-x"
     val simpleHdfsPath = "hdfs:///xxxx/xx-xx/x-x-x"