Skip to content

Resolve lost wakeup in SpillPoolReader with multiple concurrent SpillPoolWriters #23522

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
pepijnve wants to merge 1 commit into
base: main
Choose a base branch
from
+62 −63
Open

Resolve lost wakeup in SpillPoolReader with multiple concurrent SpillPoolWriters #23522

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
125 changes: 62 additions & 63 deletions datafusion/physical-plan/src/spill/spill_pool.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,14 +55,12 @@ struct SpillPoolShared {
spill_manager: Arc<SpillManager>,
/// Pool-level waker to notify when new files are available (single reader)
waker: Option<Waker>,
/// Whether the writer has been dropped (no more files will be added)
writer_dropped: bool,
/// Writer's reference to the current file (shared by all cloned writers).
/// Has its own lock to allow I/O without blocking queue access.
current_write_file: Option<Arc<Mutex<ActiveSpillFileShared>>>,
/// Number of active writer clones. Only when this reaches zero should
/// `writer_dropped` be set to true. This prevents premature EOF signaling
/// when one writer clone is dropped while others are still active.
/// FIFO queue of open write files. The queue may contain multiple items when multiple
/// writers concurrently write to the pool.
/// Each write file has its own lock to allow I/O without blocking queue access.
current_write_files: VecDeque<Arc<Mutex<ActiveSpillFileShared>>>,
/// Number of active writer clones. As long as this value is greater than zero, readers should
/// assume batches may still be pushed. This prevents premature EOF signaling.
active_writer_count: usize,
}

Expand All @@ -73,8 +71,7 @@ impl SpillPoolShared {
files: VecDeque::new(),
spill_manager,
waker: None,
writer_dropped: false,
current_write_file: None,
current_write_files: VecDeque::new(),
active_writer_count: 1,
}
}
Expand Down Expand Up @@ -171,7 +168,9 @@ impl SpillPoolWriter {
let mut shared = self.shared.lock();

// Create new file if we don't have one yet
if shared.current_write_file.is_none() {
let current_write_file = if !shared.current_write_files.is_empty() {
shared.current_write_files.pop_front().unwrap()
} else {
let spill_manager = Arc::clone(&shared.spill_manager);
// Release shared lock before disk I/O (fine-grained locking)
drop(shared);
Expand All @@ -194,52 +193,48 @@ impl SpillPoolWriter {
// Re-acquire lock and push to shared queue
shared = self.shared.lock();
shared.files.push_back(Arc::clone(&file_shared));
shared.current_write_file = Some(file_shared);
shared.wake(); // Wake readers waiting for new files
}
file_shared
};

let current_write_file = shared.current_write_file.take();
// Release shared lock before file I/O (fine-grained locking)
// This allows readers to access the queue while we do disk I/O
drop(shared);

// Write batch to current file - lock only the specific file
if let Some(current_file) = current_write_file {
// Now lock just this file for I/O (separate from shared lock)
let mut file_shared = current_file.lock();

// Append the batch
if let Some(ref mut writer) = file_shared.writer {
writer.append_batch(batch)?;
// make sure we flush the writer for readers
writer.flush()?;
file_shared.batches_written += 1;
file_shared.estimated_size += batch_size;
}
let mut file_shared = current_write_file.lock();

// Append the batch
if let Some(ref mut writer) = file_shared.writer {
writer.append_batch(batch)?;
// make sure we flush the writer for readers
writer.flush()?;
file_shared.batches_written += 1;
file_shared.estimated_size += batch_size;
}

// Wake reader waiting on this specific file
file_shared.wake();
// Wake reader waiting on this specific file
file_shared.wake();

// Check if we need to rotate
let needs_rotation = file_shared.estimated_size > self.max_file_size_bytes;
// Check if we need to rotate
let needs_rotation = file_shared.estimated_size > self.max_file_size_bytes;

if needs_rotation {
// Finish the IPC writer
if let Some(mut writer) = file_shared.writer.take() {
writer.finish()?;
}
// Mark as finished so readers know not to wait for more data
file_shared.writer_finished = true;
// Wake reader waiting on this file (it's now finished)
file_shared.wake();
// Don't put back current_write_file - let it rotate
} else {
// Release file lock
drop(file_shared);
// Put back the current file for further writing
let mut shared = self.shared.lock();
shared.current_write_file = Some(current_file);
if needs_rotation {
// Finish the IPC writer
if let Some(mut writer) = file_shared.writer.take() {
writer.finish()?;
}
// Mark as finished so readers know not to wait for more data
file_shared.writer_finished = true;
// Wake reader waiting on this file (it's now finished)
file_shared.wake();
// Don't put back current_write_file - let it rotate
} else {
// Release file lock
drop(file_shared);
// Put back the current file for further writing
let mut shared = self.shared.lock();
shared.current_write_files.push_back(current_write_file);
}

Ok(())
Expand All @@ -259,31 +254,35 @@ impl Drop for SpillPoolWriter {
return;
}

// Finalize the current file when the last writer is dropped
if let Some(current_file) = shared.current_write_file.take() {
// Release shared lock before locking file
// Finalize any spill files that were not finished yet

if !shared.current_write_files.is_empty() {
// Copy and clear `current_write_files` so we can release shared lock before locking files
let files = shared.current_write_files.clone();
shared.current_write_files.clear();
drop(shared);

let mut file_shared = current_file.lock();
for file in files {
let mut file_shared = file.lock();

// Finish the current writer if it exists
if let Some(mut writer) = file_shared.writer.take() {
// Ignore errors on drop - we're in destructor
let _ = writer.finish();
}
// Finish the current writer if it exists
if let Some(mut writer) = file_shared.writer.take() {
// Ignore errors on drop - we're in destructor
let _ = writer.finish();
}

// Mark as finished so readers know not to wait for more data
file_shared.writer_finished = true;
// Mark as finished so readers know not to wait for more data
file_shared.writer_finished = true;

// Wake reader waiting on this file (it's now finished)
file_shared.wake();
// Wake reader waiting on this file (it's now finished)
file_shared.wake();
drop(file_shared);
}

drop(file_shared);
shared = self.shared.lock();
}

// Mark writer as dropped and wake pool-level readers
shared.writer_dropped = true;
// Wake pool-level readers
shared.wake();
}
}
Expand Down Expand Up @@ -723,7 +722,7 @@ impl Stream for SpillPoolReader {
}

// No files in queue - check if writer is done
if shared.writer_dropped {
if shared.active_writer_count == 0 {
// Writer is done and no more files will be added - EOF
return Poll::Ready(None);
}
Expand Down
Loading