FrankYang0529 commented on code in PR #69254:
URL: https://github.com/apache/airflow/pull/69254#discussion_r3651650345
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt:
##########
@@ -57,77 +61,54 @@ class CoordinatorComm(
}
private val nextId = AtomicInt(0)
- private var shutDownRequested = false
- private val commMutex = Mutex()
-
- suspend fun startProcessing() {
- while (!shutDownRequested) {
- processOnce(::handleIncoming)
- }
- logger.debug("Goodbye")
- }
-
- private suspend fun processOnce(handle: suspend (IncomingFrame) -> Unit) {
- val prefix = reader.readByteArray(4) // First 4 bytes as length.
- if (prefix.size != 4) { // Something is terribly wrong. Let's bail.
- logger.error("Need 4 prefix bytes", mapOf("actual" to prefix.size))
- shutDownRequested = true
- return
- }
-
- val payloadLength = Frame.parseLengthPrefix(prefix)
- val payload = reader.readByteArray(payloadLength)
- if (payload.size != payloadLength) { // Something is terribly wrong. Let's
bail.
- logger.error(
- "Payload length not right",
- mapOf("expect" to payloadLength, "receive" to payload.size),
- )
- shutDownRequested = true
- return
+ private val writeMutex = Mutex()
+ private val stateMutex = Mutex()
+ private val pending = mutableMapOf<Int, CompletableDeferred<IncomingFrame>>()
+ private var dispatcherStarted = false
+ private var readError: ApiError? = null
+
+ private val dispatcherScope = CoroutineScope(Dispatchers.IO +
SupervisorJob())
+
+ suspend fun readMessage(): IncomingFrame =
+ stateMutex.withLock {
+ check(!dispatcherStarted) { "readMessage cannot be used after the
dispatcher has started" }
+ readFrame()
}
- val frame = decode(payload)
- logger.debug("Handling", mapOf("id" to frame.id))
- handle(frame)
- }
private suspend fun sendMessage(
id: Int,
body: Any,
) {
val data = encode(OutgoingFrame(id, body))
logger.debug("Sending", mapOf("id" to id, "body" to body))
- writer.writeByteArray(Frame.lengthPrefix(data.size))
- writer.writeByteArray(data)
- }
-
- suspend fun handleIncoming(frame: IncomingFrame) {
- when (val request = frame.body) {
- null -> {}
- is ErrorResponse -> throw ApiError("[${request.error}]
${request.detail}")
- is StartupDetails -> {
- communicate<Unit>(runTask(bundle, request, this))
- shutDownRequested = true
- }
+ writeMutex.withLock {
+ writer.writeByteArray(Frame.lengthPrefix(data.size))
+ writer.writeByteArray(data)
}
Review Comment:
Updated it.
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt:
##########
@@ -164,4 +168,25 @@ class Server(
}
}
}
+
+ internal suspend fun dispatchTask(
+ bundle: Bundle,
+ coordinator: CoordinatorComm,
+ ) {
+ val frame = coordinator.readMessage()
+ when (val body = frame.body) {
+ is StartupDetails -> runTaskAndReport(bundle, body, coordinator)
+ is ErrorResponse -> throw ApiError("[${body.error}] ${body.detail}")
+ else -> logger.debug("Ignoring unexpected initial frame", mapOf("id" to
frame.id))
Review Comment:
Updated it.
--
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]
