Skip to content

Commit a29ddff

Browse files
committed
fix: attempt each resource cleanup independently in CometExecIterator.close
One input's close failure no longer skips the remaining batch, native util, or shuffle inputs; failures accumulate as suppressed exceptions and the native plan is still released exactly once. Addresses review feedback on #5494.
1 parent 2ba46fb commit a29ddff

2 files changed

Lines changed: 32 additions & 15 deletions

File tree

‎spark/src/main/scala/org/apache/comet/CometExecIterator.scala‎

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -229,27 +229,33 @@ class CometExecIterator(
229229
if (!closed) {
230230
closed = true
231231

232+
// Attempt every resource's cleanup independently, so that one failure does not skip the
233+
// remaining resources: this close() is the only chance to release them, since `closed` is
234+
// already set and the task-completion retry is a no-op. The first failure is rethrown with
235+
// any later ones attached as suppressed exceptions.
232236
var failure: Throwable = null
233-
try {
237+
def attempt(cleanup: => Unit): Unit = {
238+
try {
239+
cleanup
240+
} catch {
241+
case t: Throwable =>
242+
if (failure == null) failure = t else failure.addSuppressed(t)
243+
}
244+
}
245+
246+
attempt {
234247
if (currentBatch != null) {
235248
currentBatch.close()
236249
currentBatch = null
237250
}
238-
nativeUtil.close()
239-
shuffleBlockIterators.values.foreach(_.close())
240-
} catch {
241-
case t: Throwable => failure = t
242251
}
252+
attempt(nativeUtil.close())
253+
shuffleBlockIterators.values.foreach(it => attempt(it.close()))
243254

244-
// Release the native plan even if the teardown above failed. Since this is the only chance
245-
// to do so, skipping it would strand the native execution context, which owns this plan's
246-
// task-shared memory pool reference and several JNI global refs.
247-
try {
248-
nativeLib.releasePlan(plan)
249-
} catch {
250-
case t: Throwable =>
251-
if (failure == null) failure = t else failure.addSuppressed(t)
252-
}
255+
// Released last and exactly once, even if the teardown above failed: dropping the native
256+
// execution context frees this plan's task-shared memory pool reference and several JNI
257+
// global refs.
258+
attempt(nativeLib.releasePlan(plan))
253259

254260
if (failure != null) {
255261
throw failure

‎spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,14 @@ class CometExecIteratorLifecycleSuite extends CometTestBase {
141141
new CometShuffleBlockIterator(new ByteArrayInputStream(Array.emptyByteArray)) {
142142
override def close(): Unit = throw boom
143143
}
144+
@volatile var laterInputClosed = false
145+
val trackingBlockIter =
146+
new CometShuffleBlockIterator(new ByteArrayInputStream(Array.emptyByteArray)) {
147+
override def close(): Unit = {
148+
laterInputClosed = true
149+
super.close()
150+
}
151+
}
144152
val limitOp =
145153
CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get
146154
val iter = new CometExecIterator(
@@ -151,10 +159,13 @@ class CometExecIteratorLifecycleSuite extends CometTestBase {
151159
nativeMetrics = CometMetricNode(Map.empty),
152160
numParts = 1,
153161
partitionIndex = 0,
154-
shuffleBlockIterators = Map(0 -> throwingBlockIter))
162+
shuffleBlockIterators = Map(0 -> throwingBlockIter, 1 -> trackingBlockIter))
155163

156164
val thrown = intercept[java.io.IOException](iter.close())
157165
assert(thrown eq boom)
166+
// One input's close failure must not skip the remaining resources: this close() is the only
167+
// chance to release them, since the task-completion retry is a no-op once `closed` is set.
168+
assert(laterInputClosed, "a later shuffle input was not closed after an earlier one threw")
158169
// The first close() must have marked the iterator closed and released the plan despite the
159170
// teardown failure: a second close() re-running releasePlan would free the native
160171
// execution context twice, and skipping the release would strand it.

0 commit comments

Comments
 (0)