sunchao commented on code in PR #5513:
URL: https://github.com/apache/datafusion-comet/pull/5513#discussion_r3876647694
##########
spark/src/test/scala/org/apache/comet/shuffle/CelebornShufflePartitionPusherSuite.scala:
##########
@@ -405,6 +1144,374 @@ class RecordingCelebornPushClient {
cryptoHandler.fold(length)(handler => handler.encrypt(bytes, offset,
length).length)
acceptedBytes.getOrElse(transportPayloadLength + 16)
}
+
+ @throws[IOException]
+ def mapperEnd(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ numMappers: Int,
+ numPartitions: Int): Unit = {
+ mapperEndCalls.incrementAndGet()
+ lastMapperEnd = (shuffleId, mapId, attemptId, numMappers, numPartitions)
+ if (mapperEndFailure != null) {
+ throw mapperEndFailure
+ }
+ }
+
+ @throws[IOException]
+ def cleanup(shuffleId: Int, mapId: Int, attemptId: Int): Unit = {
+ cleanupCalls.incrementAndGet()
+ lastCleanup = (shuffleId, mapId, attemptId)
+ if (cleanupFailure != null) {
+ throw cleanupFailure
+ }
+ }
+}
+
+/** Mirrors stock Celeborn's private request tracker without requiring its
optional dependency. */
+final class RecordingCelebornInFlightTracker {
+ val totalInflightReqs: LongAdder = new LongAdder()
+}
+
+/** Mirrors the public PushState failure slot and its stock private tracker
member. */
+final class RecordingCelebornPushState {
+ val inFlightRequestTracker: RecordingCelebornInFlightTracker =
+ new RecordingCelebornInFlightTracker()
+ val exception: AtomicReference[IOException] = new
AtomicReference[IOException]()
+}
+
+/** Exposes the same lifecycle and completion state as the public Apache
Celeborn client. */
+class AsyncRecordingCelebornPushClient extends RecordingCelebornPushClient {
+ val pushStates: ConcurrentHashMap[String, RecordingCelebornPushState] =
+ new ConcurrentHashMap[String, RecordingCelebornPushState]()
+
+ def getPushState(mapKey: String): RecordingCelebornPushState =
+ pushStates.computeIfAbsent(mapKey, _ => new RecordingCelebornPushState())
+
+ def currentState(shuffleId: Int, mapId: Int, attemptId: Int):
RecordingCelebornPushState =
+ pushStates.get(s"$shuffleId-$mapId-$attemptId")
+
+ def complete(state: RecordingCelebornPushState): Unit = state.synchronized {
+ state.inFlightRequestTracker.totalInflightReqs.decrement()
+ state.notifyAll()
+ }
+
+ def failWithoutRemovingRequest(state: RecordingCelebornPushState, failure:
IOException): Unit =
+ state.synchronized {
+ state.exception.compareAndSet(null, failure)
+ state.notifyAll()
+ }
+
+ @throws[IOException]
+ override def pushOrMergeData(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ partitionId: Int,
+ bytes: Array[Byte],
+ offset: Int,
+ length: Int,
+ numMappers: Int,
+ numPartitions: Int,
+ doPush: Boolean,
+ skipCompress: Boolean): Int = {
+ val state = getPushState(s"$shuffleId-$mapId-$attemptId")
+ state.inFlightRequestTracker.totalInflightReqs.increment()
+ super.pushOrMergeData(
+ shuffleId,
+ mapId,
+ attemptId,
+ partitionId,
+ bytes,
+ offset,
+ length,
+ numMappers,
+ numPartitions,
+ doPush,
+ skipCompress)
+ }
+
+ @throws[IOException]
+ override def mapperEnd(
+ shuffleId: Int,
+ mapId: Int,
+ attemptId: Int,
+ numMappers: Int,
+ numPartitions: Int): Unit = {
+ super.mapperEnd(shuffleId, mapId, attemptId, numMappers, numPartitions)
+ val key = s"$shuffleId-$mapId-$attemptId"
+ val state = pushStates.get(key)
+ if (state != null) {
+ state.synchronized {
+ while (state.exception.get() == null &&
+ state.inFlightRequestTracker.totalInflightReqs.sum() > 0) {
+ state.wait(25)
+ }
+ val failure = state.exception.get()
+ if (failure != null) {
+ throw failure
+ }
+ }
+ pushStates.remove(key, state)
+ }
+ }
+
+ @throws[IOException]
+ override def cleanup(shuffleId: Int, mapId: Int, attemptId: Int): Unit = {
+ super.cleanup(shuffleId, mapId, attemptId)
+ val removed = pushStates.remove(s"$shuffleId-$mapId-$attemptId")
+ if (removed != null) {
+ removed.synchronized {
+ removed.exception.compareAndSet(null, new IOException("Cleaned Up"))
+ removed.notifyAll()
+ }
+ }
+ }
+}
+
+/** Reproduces stock Celeborn's client-factory, handler, and callback
completion boundaries. */
+final class TransportRecordingCelebornPushClient extends
AsyncRecordingCelebornPushClient {
+ val dataClientFactory: RecordingCelebornTransportClientFactory =
+ new RecordingCelebornTransportClientFactory
+ val retryExecutor = new RecordingCelebornRetryExecutor
+ val pushDataRetryPool: ExecutorService = retryExecutor
+
+ def getDataClientFactory: RecordingCelebornTransportClientFactory =
dataClientFactory
+
+ var openConnectionBeforePush: Boolean = false
+ var beforePushBegins: () => Unit = () => ()
+ var beforePushReturns: RecordingCelebornPushState => Unit = _ => ()
+ var retriesBeforeFailure: Int = 0
+ var retryCallback: RecordingCelebornTransportCallbackApi => Unit =
+ _.onFailure(new IOException("revive failed"))
Review Comment:
[P2] Keep the callback initializers valid under Scala 2.12
At exact head `520ca35f5dbaff1a9a47ad02bb7af027d93910d1`, these two new
mutable callback initializers prevent the default Spark 3.4 and 3.5 builds from
compiling their test sources. Scala 2.12 parses the bare `_` after a `var`
declaration as its default initializer, then rejects the remaining `=>` at line
1284 and `.` at line 1287. Root Maven `test` with `-Pspark-3.5 -Pjdk17` and
`clean test` with `-Pspark-3.4 -Pjdk17` both stop in `scala:testCompile` before
any selected suite can execute. The current [Spark 3.5 CI
job](https://github.com/apache/datafusion-comet/actions/runs/33125216230/job/98702084424)
and [Spark 3.4 CI
job](https://github.com/apache/datafusion-comet/actions/runs/33125216230/job/98702084393)
report the same two errors.
For the causal control, this suite's exact base
`eabb5d4773091b983d8fce713f0e34b1cf93f877` and previous
`16d71677d5975bb2aaa0d835a77f781d487c7262` sources both pass the parser with
Scala 2.12.17 and 2.12.18; the current source fails. Parenthesizing only these
initializer expressions in a scratch copy restores parsing. Spark 4.0/Scala
2.13 passes all 63 focused tests, so that profile does not cover this
compatibility regression. Please express both lambdas in a form Scala 2.12
accepts so normal Spark 3.4/3.5 test and package builds can complete.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
