Skip to content

Commit 14c70d1

Browse files
committed
Java SDK: Clear the XComs of a task instance at the start of each try
Signed-off-by: PoAn Yang <payang@apache.org>
1 parent 4083290 commit 14c70d1

7 files changed

Lines changed: 380 additions & 39 deletions

File tree

‎java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ package org.apache.airflow.sdk.execution
2222
import kotlinx.coroutines.runBlocking
2323
import org.apache.airflow.sdk.execution.comm.ConnectionResult
2424
import org.apache.airflow.sdk.execution.comm.DeleteVariable
25+
import org.apache.airflow.sdk.execution.comm.DeleteXCom
2526
import org.apache.airflow.sdk.execution.comm.GetConnection
2627
import org.apache.airflow.sdk.execution.comm.GetVariable
2728
import org.apache.airflow.sdk.execution.comm.GetXCom
@@ -74,6 +75,14 @@ interface Client {
7475
runId: String,
7576
mapIndex: Int,
7677
)
78+
79+
fun deleteXCom(
80+
key: String,
81+
dagId: String,
82+
taskId: String,
83+
runId: String,
84+
mapIndex: Int,
85+
)
7786
}
7887

7988
/**
@@ -157,4 +166,22 @@ class CoordinatorClient(
157166
}
158167
return runBlocking { exec.communicate<XComResult>(message) }
159168
}
169+
170+
override fun deleteXCom(
171+
key: String,
172+
dagId: String,
173+
taskId: String,
174+
runId: String,
175+
mapIndex: Int,
176+
) {
177+
val message =
178+
DeleteXCom().also {
179+
it.key = key
180+
it.dagId = dagId
181+
it.taskId = taskId
182+
it.runId = runId
183+
it.mapIndex = mapIndex
184+
}
185+
runBlocking { exec.communicate<Unit>(message) }
186+
}
160187
}

‎java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,27 @@ internal object TaskRunner {
7777
val taskDef =
7878
bundle.taskDef(request.ti.dagId, request.ti.taskId)
7979
?: return TaskResult.of(TaskState.State.REMOVED)
80+
// When a try starts, the runtime has to delete the XComs in xcom_keys_to_clear, because Airflow
81+
// does not delete them. The Python task runner deletes them before it renders templates. Here
82+
// they are deleted before the task class is instantiated, so they are already gone if the
83+
// constructor throws during this try.
84+
for (key in request.tiContext.xcomKeysToClear.orEmpty()) {
85+
try {
86+
client.impl.deleteXCom(
87+
key = key,
88+
dagId = request.ti.dagId,
89+
taskId = request.ti.taskId,
90+
runId = request.ti.runId,
91+
mapIndex = request.ti.mapIndex ?: -1,
92+
)
93+
} catch (e: Throwable) {
94+
logger.error(
95+
"Error clearing XCom",
96+
mapOf("ti" to request.ti, "key" to key, "error" to e, "trace" to e.stackTraceToString()),
97+
)
98+
return TaskResult.failure(request.tiContext.shouldRetry)
99+
}
100+
}
80101
val definition = taskDef.definition
81102
val instance =
82103
try {

‎java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,14 @@ internal class FakeXComTransport(
6767
runId: String,
6868
mapIndex: Int,
6969
) = throw NotImplementedError()
70+
71+
override fun deleteXCom(
72+
key: String,
73+
dagId: String,
74+
taskId: String,
75+
runId: String,
76+
mapIndex: Int,
77+
) = throw NotImplementedError()
7078
}
7179

7280
internal fun startupDetails(argBindings: List<Map<String, Any?>>?): StartupDetails =

‎java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,14 @@ private class FakeTransport(
5959
runId: String,
6060
mapIndex: Int,
6161
) = throw NotImplementedError()
62+
63+
override fun deleteXCom(
64+
key: String,
65+
dagId: String,
66+
taskId: String,
67+
runId: String,
68+
mapIndex: Int,
69+
) = throw NotImplementedError()
6270
}
6371

6472
class ClientTest {

‎java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import org.junit.jupiter.api.DisplayName
3333
import org.junit.jupiter.api.Test
3434
import org.junit.jupiter.api.Timeout
3535
import org.msgpack.core.MessagePack
36+
import org.msgpack.core.buffer.ArrayBufferInput
3637
import java.io.ByteArrayOutputStream
3738
import java.util.concurrent.ArrayBlockingQueue
3839
import java.util.concurrent.TimeUnit
@@ -60,6 +61,111 @@ class ServerTest {
6061
writeByteArray(payload)
6162
}
6263

64+
private fun errorResponseFrame(id: Int): ByteArray {
65+
val out = ByteArrayOutputStream()
66+
MessagePack.newDefaultPacker(out).use { packer ->
67+
packer.packArrayHeader(3)
68+
packer.packInt(id)
69+
packer.packNil()
70+
packer.packMapHeader(3)
71+
packer.packString("type")
72+
packer.packString("ErrorResponse")
73+
packer.packString("error")
74+
packer.packString("API_SERVER_ERROR")
75+
packer.packString("detail")
76+
packer.packMapHeader(1)
77+
packer.packString("status_code")
78+
packer.packInt(500)
79+
}
80+
return out.toByteArray()
81+
}
82+
83+
private fun startupFrame(
84+
mapIndex: Int?,
85+
tiContext: Map<String, Any?>,
86+
): ByteArray {
87+
val body =
88+
mapOf(
89+
"type" to "StartupDetails",
90+
"ti" to
91+
mapOf(
92+
"id" to "4d828a62-a417-4936-a7a6-2b3fabacecab",
93+
"task_id" to "extract",
94+
"dag_id" to "dag1",
95+
"run_id" to "run1",
96+
"try_number" to 2,
97+
"map_index" to mapIndex,
98+
"dag_version_id" to "4d828a62-a417-4936-a7a6-2b3fabacecab",
99+
),
100+
"ti_context" to mapOf("dag_run" to mapOf("dag_id" to "dag1", "run_id" to "run1"), "max_tries" to 1) + tiContext,
101+
"dag_rel_path" to "/dev/null",
102+
"bundle_info" to mapOf("name" to "any-name", "version" to "any-version"),
103+
"start_date" to "2024-12-01T01:00:00Z",
104+
"sentry_integration" to "",
105+
)
106+
return Frame.encodeRequest(0, body).fold(ByteArray(0)) { acc, buffer -> acc + buffer.toByteArray() }
107+
}
108+
109+
/**
110+
* Runs [taskClass] as task "extract" of Dag "dag1" through [Server.dispatchTask]. The fake
111+
* supervisor first sends StartupDetails, with [mapIndex] as the map_index of the task instance and
112+
* [tiContext] added to ti_context. Until the task reports its final state, [answer] builds the
113+
* reply to each request from the request ID. Returns the bodies of those requests in arrival
114+
* order, and the body of the final state.
115+
*/
116+
private fun serveTask(
117+
taskClass: Class<out Task>,
118+
mapIndex: Int?,
119+
tiContext: Map<String, Any?>,
120+
answer: (Int) -> ByteArray = ::ackFrame,
121+
): Pair<List<Map<*, *>>, Map<*, *>> {
122+
val toServer = ByteChannel(autoFlush = true)
123+
val fromServer = ByteChannel(autoFlush = true)
124+
val comm = CoordinatorComm(toServer, fromServer)
125+
val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0))
126+
val bundle = Bundle(listOf(DagDef("dag1").addTask(TaskDef("extract", taskClass))))
127+
128+
val requests = mutableListOf<Map<*, *>>()
129+
var terminal: Map<*, *>? = null
130+
val supervisor =
131+
Thread {
132+
runBlocking {
133+
toServer.writeFrame(startupFrame(mapIndex, tiContext))
134+
while (terminal == null) {
135+
val prefix = fromServer.readByteArray(4)
136+
val payload = fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt())
137+
val request = Frame.decodeRaw(ArrayBufferInput(payload))
138+
val body = request.rawBody as Map<*, *>
139+
if (body["type"] in setOf("SucceedTask", "TaskState", "RetryTask")) {
140+
terminal = body
141+
toServer.writeFrame(ackFrame(request.id))
142+
} else {
143+
requests += body
144+
toServer.writeFrame(answer(request.id))
145+
}
146+
}
147+
}
148+
}
149+
supervisor.start()
150+
151+
runBlocking { server.dispatchTask(bundle, comm) }
152+
supervisor.join()
153+
comm.close()
154+
return requests to checkNotNull(terminal)
155+
}
156+
157+
private fun deleteXComRequest(
158+
key: String,
159+
mapIndex: Int,
160+
) = mapOf(
161+
"type" to "DeleteXCom",
162+
"key" to key,
163+
"dag_id" to "dag1",
164+
"run_id" to "run1",
165+
"task_id" to "extract",
166+
"map_index" to mapIndex.toLong(),
167+
)
168+
63169
@Test
64170
@DisplayName("Should run the task and report the result as a normal awaited request")
65171
@Timeout(value = 30, unit = TimeUnit.SECONDS)
@@ -112,6 +218,64 @@ class ServerTest {
112218
comm.close()
113219
}
114220

221+
@Test
222+
@DisplayName("Should delete each listed XCom before the task sends anything")
223+
@Timeout(value = 30, unit = TimeUnit.SECONDS)
224+
fun deletesListedXComsBeforeTheTaskSendsAnything() {
225+
val cases =
226+
listOf(
227+
Triple(-1, listOf("return_value", "summary"), -1),
228+
Triple(0, listOf("return_value", "summary"), 0),
229+
Triple(2, listOf("return_value", "summary"), 2),
230+
Triple(null, listOf("return_value"), -1),
231+
Triple(-1, emptyList(), -1),
232+
Triple(-1, null, -1),
233+
)
234+
for ((mapIndex, keys, deletedMapIndex) in cases) {
235+
val (requests, terminal) = serveTask(PushingTask::class.java, mapIndex, mapOf("xcom_keys_to_clear" to keys))
236+
237+
Assertions.assertEquals(keys.orEmpty().map { deleteXComRequest(it, deletedMapIndex) }, requests.dropLast(1)) {
238+
"unexpected requests for map_index $mapIndex and keys $keys: $requests"
239+
}
240+
Assertions.assertEquals("SetXCom", requests.last()["type"])
241+
Assertions.assertEquals("SucceedTask", terminal["type"])
242+
}
243+
}
244+
245+
@Test
246+
@DisplayName("Should fail without running the task when an XCom cannot be deleted")
247+
@Timeout(value = 30, unit = TimeUnit.SECONDS)
248+
fun failsWithoutRunningTheTaskWhenAnXComCannotBeDeleted() {
249+
for (shouldRetry in listOf(true, false)) {
250+
val (requests, terminal) =
251+
serveTask(
252+
PushingTask::class.java,
253+
-1,
254+
mapOf("xcom_keys_to_clear" to listOf("return_value", "summary"), "should_retry" to shouldRetry),
255+
::errorResponseFrame,
256+
)
257+
258+
Assertions.assertEquals(listOf(deleteXComRequest("return_value", -1)), requests) {
259+
"the runtime should stop at the first XCom it cannot delete and not run the task"
260+
}
261+
if (shouldRetry) {
262+
Assertions.assertEquals("RetryTask", terminal["type"])
263+
} else {
264+
Assertions.assertEquals("TaskState", terminal["type"])
265+
Assertions.assertEquals("failed", terminal["state"])
266+
}
267+
}
268+
}
269+
270+
class PushingTask : Task {
271+
override fun execute(
272+
context: Context,
273+
client: Client,
274+
) {
275+
client.setXCom(value = "rows")
276+
}
277+
}
278+
115279
private companion object {
116280
// [2, msg, null] with msg coming from
117281
// https://github.com/astronomer/airflow/blob/f39c8da8/task-sdk/tests/task_sdk/execution_time/test_comms.py#L73-L108

0 commit comments

Comments
 (0)