Skip to content

[SPARK-24334] Fix race condition in ArrowPythonRunner causes unclean shutdown of Arrow memory allocator #21397

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
icexelloss wants to merge 3 commits into from
Closed

[SPARK-24334] Fix race condition in ArrowPythonRunner causes unclean shutdown of Arrow memory allocator #21397

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
435ccff
Remove the Arrow allocation cleanup routine in TaskCompletionListen t…
icexelloss May 22, 2018
69c9104
Add unit test
icexelloss May 22, 2018
756a73a
Remove tests. Add comments
icexelloss May 27, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
Remove the Arrow allocation cleanup routine in TaskCompletionListen t… 
…o avoid race condition
  • Loading branch information
@icexelloss
icexelloss committed May 22, 2018
commit 435ccfff44995ca5ad487e77128b2cae4ff1cfd5
26 changes: 15 additions & 11 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,19 +70,13 @@ class ArrowPythonRunner(
val arrowSchema = ArrowUtils.toArrowSchema(schema, timeZoneId)
val allocator = ArrowUtils.rootAllocator.newChildAllocator(
s"stdout writer for $pythonExec", 0, Long.MaxValue)

val root = VectorSchemaRoot.create(arrowSchema, allocator)
val arrowWriter = ArrowWriter.create(root)

context.addTaskCompletionListener { _ =>
root.close()
allocator.close()
}

val writer = new ArrowStreamWriter(root, null, dataOut)
writer.start()

Utils.tryWithSafeFinally {
val arrowWriter = ArrowWriter.create(root)
val writer = new ArrowStreamWriter(root, null, dataOut)
writer.start()

while (inputIterator.hasNext) {
val nextBatch = inputIterator.next()

Expand All @@ -94,8 +88,18 @@ class ArrowPythonRunner(
writer.writeBatch()
arrowWriter.reset()
}
} {
writer.end()
Copy link
Contributor Author

@icexelloss icexelloss May 22, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is moved out of the finally block because this function will throw exception if the thread is interrupted. writer.end() also doesn't do any clean up - it just writes some final bits to the output channel so we don't need to call it in clean up.

Copy link
Contributor

@advancedxy advancedxy May 22, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So basically your change makes sure that root and allocator is safely closed in the final block and hence we don't need the safe belt in TaskCompletionListener?

Copy link
Contributor Author

@icexelloss icexelloss May 22, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The routine in TaskCompletionListener is not really a safe belt, it actually causes the race condition. But you are right, I believe putting the cleanup routine in finally block is enough.

Copy link
Member

@felixcheung felixcheung May 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe add a comment to explain the same - that this should be outside of the finally?

} {
// If we close root and allocator in TaskCompletionListener, there could be a race
// condition where the writer thread keeps writing to the VectorSchemaRoot while
// it's being closed by the TaskCompletion listener.
// Closing root and allocator here is cleaner because root and allocator is owned
// by the writer thread and is only visible to the writer thread.
//
// If the writer thread is interrupted by TaskCompletionListener, it should either
// (1) in the try block, in which case it will get an InterruptedException when
// performing io, and goes into the finally block or (2) in the finally block,
// in which case it will ignore the interruption and close the resources.
root.close()
allocator.close()
}
Expand Down