Skip to content

[AURON #2352] Lock spills before new spiller #2350

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

Open
wForget wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[AURON #2352] Lock spills before new spiller #2350

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3335d71
fix: lock spills before new spiller
wForget Jun 22, 2026
90f5392
Potential fix for pull request finding
wForget Jun 23, 2026
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
5 changes: 3 additions & 2 deletions native-engine/datafusion-ext-plans/src/shuffle/sort_repartitioner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,15 +98,16 @@ impl MemConsumer for SortShuffleRepartitioner {
async fn spill(&self) -> Result<()> {
let data = self.data.lock().await.drain();
let spill_metrics = self.exec_ctx.spill_metrics().clone();
let mut spills = self.spills.lock().await;

@weiqingy weiqingy Jun 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

In this file the on-heap spill_id comes from newSpill() (auron-memmgr/src/spill.rs:182) and try_new_spill takes no index, so spills.len() doesn't feed a spill id here the way it does in agg_table.rs (try_into_spill(spill, spill_idx)). Is the race from the issue actually reachable in sort_repartitioner, or is this more defensive? Either's fine — just want to make sure I'm reading it right.

@wForget wForget Jun 24, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I found this issue while reviewing the code, and didn't encounter any actual failures. However, this race condition is possible; although sort_repartitioner doesn't take spill_id, it still affects the spill index on the JVM side. If I understand correctly, we should always ensure that the native-side spills length is consistent with the JVM-side length, so we need to lock the spills first.

auron/spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/SparkOnHeapSpillManager.scala

Lines 106 to 107 in ca3bece

val spill = OnHeapSpill(this, spills.length)
spills.append(Some(spill))

let spill = tokio::task::spawn_blocking(move || {
let mut spill = try_new_spill(&spill_metrics)?;
let offsets = data.write(spill.get_buf_writer())?;
Ok::<_, DataFusionError>(Offsetted::new(offsets, spill))
})
.await
.expect("tokio spawn_blocking error")?;
Comment on lines +101 to 108

self.spills.lock().await.push(spill);
spills.push(spill);
drop(spills);
self.update_mem_used(0).await?;
Ok(())
}
Expand Down
Loading