test: verify Celeborn reflection compatibility

.github/workflows/pr_build_linux.yml

@@ -176,8 +176,51 @@
       - name: Compile (skip tests)
         run: ./mvnw -B install -DskipTests -Dmaven.test.skip=true -Pspark-4.1
 
+  celeborn-reflection-compatibility:
+    needs: lint
+    name: Celeborn ${{ matrix.celeborn_version }} reflection compatibility
+    runs-on: ubuntu-24.04
+    container:
+      image: amd64/rust
+      env:
+        JAVA_TOOL_OPTIONS: --add-exports=java.base/sun.nio.ch=ALL-UNNAMED --add-exports=java.base/sun.util.calendar=ALL-UNNAMED --add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.lang=ALL-UNNAMED
+    strategy:
+      fail-fast: false
+      matrix:
+        celeborn_version: ["0.6.0", "0.7.0"]
+    steps:
+      - uses: actions/checkout@v7
+
+      - name: Setup Rust & Java toolchain
+        uses: ./.github/actions/setup-builder
+        with:
+          rust-version: ${{ env.RUST_VERSION }}
+          jdk-version: 17
+
+      - name: Cache Maven dependencies
+        uses: actions/cache@v6
+        with:
+          path: |
+            ~/.m2/repository
+            /root/.m2/repository
+          key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }}-celeborn-${{ matrix.celeborn_version }}
+          restore-keys: |
+            ${{ runner.os }}-java-maven-
+
+      - name: Verify reflected Celeborn internals
+        env:
+          SPARK_LOCAL_HOSTNAME: localhost
+          SPARK_LOCAL_IP: 127.0.0.1
+        run: |
+          SPARK_HOME="$GITHUB_WORKSPACE" ./mvnw -B clean test \
+            -Pspark-3.5,scala-2.12,celeborn-reflection-compatibility \
+            -Dceleborn.version="${{ matrix.celeborn_version }}" \
+            -Dtest=none \
+            -Dsuites=org.apache.comet.shuffle.CelebornReflectionCompatibilitySuite \
+            -DfailIfNoTests=false
+
   # Build native library once and share with all test jobs
   build-native:
     needs: lint
     name: Build Native Library
     runs-on: ubuntu-24.04

dev/ci/check-suites.py

@@ -36,6 +36,7 @@ def file_to_class_name(path: Path) -> str | None:
         "org.apache.comet.parquet.ParquetReadFromS3Suite", # manual test suite
         "org.apache.comet.IcebergReadFromS3Suite", # manual test suite
         "org.apache.comet.cloud.s3.CometS3CredentialBridgeSuite", # manual test suite
+        "org.apache.comet.shuffle.CelebornReflectionCompatibilitySuite", # dedicated version matrix
         "org.apache.spark.sql.comet.CometPlanStabilitySuite", # abstract
         "org.apache.spark.sql.comet.ParquetDatetimeRebaseSuite", # abstract
         "org.apache.comet.exec.CometColumnarShuffleSuite" # abstract

pom.xml

@@ -52,6 +52,8 @@ under the License.
       property without an error.
     -->
     <delta.version>4.1.0</delta.version>
+    <!-- Used only to verify reflected Celeborn internals; never bundled with Comet. -->
+    <celeborn.version>0.7.0</celeborn.version>
     <maven.compiler.source>${java.version}</maven.compiler.source>
     <maven.compiler.target>${java.version}</maven.compiler.target>
     <maven-compiler-plugin.version>3.11.0</maven-compiler-plugin.version>
@@ -1046,6 +1048,7 @@ under the License.
               <includes>
                 <include>src/main/scala/**/*.scala</include>
                 <include>src/test/scala/**/*.scala</include>
+                <include>src/test/celeborn-reflection-compatibility/**/*.scala</include>
                 <!-- Include spark shim sources -->
                 <include>src/main/spark-*/**/*.scala</include>
                 <include>src/test/spark-*/**/*.scala</include>

spark/pom.xml

@@ -222,6 +222,41 @@ under the License.
   </dependencies>
 
   <profiles>
+    <profile>
+      <id>celeborn-reflection-compatibility</id>
+      <dependencies>
+        <dependency>
+          <!-- Verify reflection against a released client without adding a production dependency. -->
+          <groupId>org.apache.celeborn</groupId>
+          <artifactId>celeborn-client-spark-3-shaded_${scala.binary.version}</artifactId>
+          <version>${celeborn.version}</version>
+          <scope>test</scope>
+        </dependency>
+      </dependencies>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>org.codehaus.mojo</groupId>
+            <artifactId>build-helper-maven-plugin</artifactId>
+            <executions>
+              <execution>
+                <id>add-celeborn-reflection-compatibility-test-source</id>
+                <phase>generate-test-sources</phase>
+                <goals>
+                  <goal>add-test-source</goal>
+                </goals>
+                <configuration>
+                  <sources>
+                    <source>src/test/celeborn-reflection-compatibility</source>
+                  </sources>
+                </configuration>
+              </execution>
+            </executions>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
+
     <!-- Iceberg dependencies vary by Spark version (Iceberg 1.8.1 for Spark 3.x, 1.10.0 for Spark 4.0, 1.11.0 for Spark 4.1) -->
     <profile>
       <id>spark-3.4</id>

spark/src/test/celeborn-reflection-compatibility/org/apache/comet/shuffle/CelebornReflectionCompatibilitySuite.scala

@@ -0,0 +1,221 @@
+/*
+ * 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.comet.shuffle
+
+import java.lang.reflect.{Field, Method, Modifier}
+import java.util.{List => JList, Map => JMap, Optional}
+import java.util.concurrent.{ConcurrentHashMap, ExecutorService}
+import java.util.concurrent.atomic.{AtomicReference, LongAdder}
+
+import org.scalatest.funsuite.AnyFunSuite
+
+class CelebornReflectionCompatibilitySuite extends AnyFunSuite {
+
+  private val classLoader = getClass.getClassLoader
+
+  private lazy val shuffleClient = load("org.apache.celeborn.client.ShuffleClientImpl")
+  private lazy val pushState = load("org.apache.celeborn.common.write.PushState")
+  private lazy val inFlightRequestTracker =
+    load("org.apache.celeborn.common.write.InFlightRequestTracker")
+  private lazy val transportClientFactory =
+    load("org.apache.celeborn.common.network.client.TransportClientFactory")
+  private lazy val transportClient =
+    load("org.apache.celeborn.common.network.client.TransportClient")
+  private lazy val transportResponseHandler =
+    load("org.apache.celeborn.common.network.client.TransportResponseHandler")
+  private lazy val pushRequestInfo = load("org.apache.celeborn.common.write.PushRequestInfo")
+
+  private lazy val celebornVersion = Option(shuffleClient.getPackage)
+    .flatMap(pkg => Option(pkg.getImplementationVersion))
+    .getOrElse(fail("Celeborn client JAR does not declare its implementation version"))
+
+  private lazy val celebornReleaseLine = celebornVersion.split("\\.").take(2).mkString(".")
+
+  test("load a released Celeborn 0.6 or 0.7 client") {
+    assert(shuffleClient.getPackage.getImplementationTitle == "celeborn-client-spark-3-shaded")
+    assert(
+      Set("0.6", "0.7").contains(celebornReleaseLine),
+      s"unsupported Celeborn compatibility-test version $celebornVersion")
+  }
+
+  test("resolve the Celeborn shuffle client members used by the partition pusher") {
+    instanceMethod(
+      shuffleClient,
+      "pushOrMergeData",
+      java.lang.Integer.TYPE,
+      java.lang.Integer.TYPE,
+      java.lang.Integer.TYPE,
+      java.lang.Integer.TYPE,
+      java.lang.Integer.TYPE,
+      classOf[Array[Byte]],
+      java.lang.Integer.TYPE,
+      java.lang.Integer.TYPE,
+      java.lang.Integer.TYPE,
+      java.lang.Integer.TYPE,
+      java.lang.Boolean.TYPE,
+      java.lang.Boolean.TYPE)
+    instanceMethod(
+      shuffleClient,
+      "cleanup",
+      java.lang.Void.TYPE,
+      java.lang.Integer.TYPE,
+      java.lang.Integer.TYPE,
+      java.lang.Integer.TYPE)
+    instanceMethod(shuffleClient, "getPushState", pushState, classOf[String])
+
+    instanceField(shuffleClient, "pushStates", classOf[JMap[_, _]])
+    instanceField(shuffleClient, "pushDataRetryPool", classOf[ExecutorService])
+
+    celebornReleaseLine match {
+      case "0.6" =>
+        instanceMethod(
+          shuffleClient,
+          "mapperEnd",
+          java.lang.Void.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE)
+        assert(!shuffleClient.getMethods.exists(_.getName == "computeBatchCRC"))
+        intercept[NoSuchFieldException](declaredField(shuffleClient, "cryptoHandler"))
+
+      case "0.7" =>
+        instanceMethod(
+          shuffleClient,
+          "mapperEnd",
+          java.lang.Void.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE)
+        instanceMethod(
+          shuffleClient,
+          "computeBatchCRC",
+          java.lang.Void.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE,
+          classOf[Array[Byte]],
+          java.lang.Integer.TYPE,
+          java.lang.Integer.TYPE)
+        instanceField(shuffleClient, "cryptoHandler", classOf[Optional[_]])
+
+      case other => fail(s"unsupported Celeborn release line $other")
+    }
+  }
+
+  test("resolve the Celeborn push-state members used for admission") {
+    exactInstanceField(pushState, "exception", classOf[AtomicReference[_]])
+    instanceField(pushState, "inFlightRequestTracker", inFlightRequestTracker)
+    exactInstanceField(inFlightRequestTracker, "totalInflightReqs", classOf[LongAdder])
+  }
+
+  test("resolve the Celeborn transport members used to track callback ownership") {
+    instanceMethod(shuffleClient, "getDataClientFactory", transportClientFactory)
+
+    val transportClientBootstrap =
+      load("org.apache.celeborn.common.network.client.TransportClientBootstrap")
+    assert(transportClientBootstrap.isInterface)
+    instanceMethod(transportClientBootstrap, "doBootstrap", java.lang.Void.TYPE, transportClient)
+
+    val bootstraps = instanceField(transportClientFactory, "clientBootstraps", classOf[JList[_]])
+    assert(bootstraps.getGenericType.getTypeName.contains(transportClientBootstrap.getName))
+    instanceField(transportClientFactory, "connectionPool", classOf[JMap[_, _]])
+
+    val clientPool =
+      load("org.apache.celeborn.common.network.client.TransportClientFactory$ClientPool")
+    val clients = instanceField(clientPool, "clients", classOf[Array[_]])
+    assert(clients.getType.getComponentType == transportClient)
+    val locks = instanceField(clientPool, "locks", classOf[Array[_]])
+    assert(locks.getType.getComponentType == classOf[Object])
+
+    val channel = declaredField(transportClient, "channel")
+    assert(!Modifier.isStatic(channel.getModifiers))
+    assert(channel.getType.isInterface)
+    instanceMethod(transportClient, "getHandler", transportResponseHandler)
+    instanceField(transportResponseHandler, "outstandingPushes", classOf[ConcurrentHashMap[_, _]])
+
+    val callback = declaredField(pushRequestInfo, "callback")
+    assert(!Modifier.isStatic(callback.getModifiers))
+    assert(!Modifier.isFinal(callback.getModifiers))
+    assert(callback.getType.isInterface)
+
+    val writeFuture = channel.getType.getMethod("writeAndFlush", classOf[Object]).getReturnType
+    assert(writeFuture.isInterface)
+    instanceMethod(writeFuture, "isDone", java.lang.Boolean.TYPE)
+    assert(writeFuture.getMethods.exists { method =>
+      method.getName == "addListener" &&
+      method.getParameterCount == 1 &&
+      method.getParameterTypes.head.isInterface
+    })
+  }
+
+  private def load(name: String): Class[_] = Class.forName(name, false, classLoader)
+
+  private def instanceField(owner: Class[_], name: String, expectedType: Class[_]): Field = {
+    val result = declaredField(owner, name)
+    assert(
+      !Modifier.isStatic(result.getModifiers),
+      s"${owner.getName}.$name must be an instance field")
+    assert(
+      expectedType.isAssignableFrom(result.getType),
+      s"${owner.getName}.$name has type ${result.getType.getName}, expected ${expectedType.getName}")
+    result
+  }
+
+  private def exactInstanceField(owner: Class[_], name: String, expectedType: Class[_]): Field = {
+    val result = instanceField(owner, name, expectedType)
+    assert(
+      result.getType == expectedType,
+      s"${owner.getName}.$name has type ${result.getType.getName}, expected ${expectedType.getName}")
+    result
+  }
+
+  private def declaredField(owner: Class[_], name: String): Field = {
+    var current = owner
+    while (current != null) {
+      try {
+        val result = current.getDeclaredField(name)
+        result.setAccessible(true)
+        return result
+      } catch {
+        case _: NoSuchFieldException => current = current.getSuperclass
+      }
+    }
+    throw new NoSuchFieldException(s"${owner.getName}.$name")
+  }
+
+  private def instanceMethod(
+      owner: Class[_],
+      name: String,
+      returnType: Class[_],
+      parameterTypes: Class[_]*): Method = {
+    val result = owner.getMethod(name, parameterTypes: _*)
+    assert(
+      !Modifier.isStatic(result.getModifiers),
+      s"${owner.getName}.$name must be an instance method")
+    assert(
+      result.getReturnType == returnType,
+      s"${owner.getName}.$name returns ${result.getReturnType.getName}, expected ${returnType.getName}")
+    result
+  }
+}