kosiew commented on code in PR #24891:
URL: https://github.com/apache/datafusion/pull/24891#discussion_r3932497267
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -1645,4 +1696,556 @@ mod tests {
Ok(())
}
+
+ type WriteHook = Arc<dyn Fn() + Send + Sync>;
+ /// What a read of a spill file does. The two failure modes reach the two
+ /// different error paths of [`SpillPoolFile::poll_file`]: one fails while
it
+ /// builds the stream, the other fails while it polls a stream that is
+ /// already built.
+ enum ReadBehavior {
+ /// Read the real bytes, after this delay.
+ Delay(Duration),
+ /// Fail to open the file, so the stream is never built.
+ FailOpen(DataFusionError),
+ /// Open the file, then give this error as the first item of the
stream.
+ FailFirstItem(DataFusionError),
+ }
+
+ type ReadHook = Arc<dyn Fn() -> ReadBehavior + Send + Sync>;
+
+ /// Test double for a spill file that runs `write_hook` before every disk
+ /// write or flush, and consults `read_hook` before every read. Writers
write
+ /// while holding their file's lock, so a hook that blocks pauses a
+ /// `push_batch` at exactly the point where it has a file checked out for
+ /// writing; a read delay makes an older file's bytes arrive after a newer
+ /// file's, and a read error fails the read. Data still goes to real
+ /// temporary files.
+ struct IoHookFactory {
+ inner: Arc<DiskManager>,
+ write_hook: WriteHook,
+ read_hook: ReadHook,
+ }
+
+ impl TempFileFactory for IoHookFactory {
+ fn create_temp_file(&self, description: &str) -> Result<Arc<dyn
SpillFile>> {
+ Ok(Arc::new(IoHookFile {
+ inner: self.inner.create_tmp_file(description)?,
+ write_hook: Arc::clone(&self.write_hook),
+ read_hook: Arc::clone(&self.read_hook),
+ }))
+ }
+ }
+
+ struct IoHookFile {
+ inner: Arc<dyn SpillFile>,
+ write_hook: WriteHook,
+ read_hook: ReadHook,
+ }
+
+ impl SpillFile for IoHookFile {
+ fn path(&self) -> Option<&std::path::Path> {
+ self.inner.path()
+ }
+
+ fn size(&self) -> Option<u64> {
+ self.inner.size()
+ }
+
+ fn read_stream(
+ &self,
+ ) -> Result<Pin<Box<dyn Stream<Item = Result<bytes::Bytes>> + Send>>> {
+ let delay = match (self.read_hook)() {
+ ReadBehavior::Delay(delay) => delay,
+ ReadBehavior::FailOpen(e) => return Err(e),
+ ReadBehavior::FailFirstItem(e) => {
+ return Ok(Box::pin(futures::stream::once(async move {
Err(e) })));
+ }
+ };
+ let mut inner = Some(self.inner.read_stream()?);
+ Ok(Box::pin(
+ futures::stream::once(tokio::time::sleep(delay))
+ .flat_map(move |_| inner.take().expect("polled once")),
+ ))
+ }
+
+ fn open_writer(&self) -> Result<Box<dyn SpillWriter>> {
+ Ok(Box::new(IoHookWriter {
+ inner: self.inner.open_writer()?,
+ hook: Arc::clone(&self.write_hook),
+ }))
+ }
+ }
+
+ struct IoHookWriter {
+ inner: Box<dyn SpillWriter>,
+ hook: WriteHook,
+ }
+
+ impl std::io::Write for IoHookWriter {
+ fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+ (self.hook)();
+ self.inner.write(buf)
+ }
+
+ fn flush(&mut self) -> std::io::Result<()> {
+ (self.hook)();
+ self.inner.flush()
+ }
+ }
+
+ impl SpillWriter for IoHookWriter {
+ fn finish(&mut self) -> Result<()> {
+ self.inner.finish()
+ }
+ }
+
+ /// A `SpillManager` whose files run `write_hook` before every disk write
and
+ /// consult `read_hook` before every read, plus the `DiskManager` that owns
+ /// the files so tests can check disk usage.
+ fn spill_manager_with_io_hooks(
+ write_hook: WriteHook,
+ read_hook: ReadHook,
+ ) -> Result<(Arc<SpillManager>, Arc<DiskManager>)> {
+ let disk_manager = Arc::new(DiskManagerBuilder::default().build()?);
+ let runtime = RuntimeEnvBuilder::new()
+ .with_disk_manager_builder(DiskManagerBuilder::default().with_mode(
+ DiskManagerMode::Custom(Arc::new(IoHookFactory {
+ inner: Arc::clone(&disk_manager),
+ write_hook,
+ read_hook,
+ })),
+ ))
+ .build_arc()?;
+ let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+ let spill_manager =
+ Arc::new(SpillManager::new(runtime, metrics,
create_test_schema()));
+ Ok((spill_manager, disk_manager))
+ }
+
+ /// Holds the first disk write on the pool until the test releases it.
+ struct WriteGate {
+ taken: AtomicBool,
+ entered: Barrier,
+ release: Barrier,
+ }
+
+ impl WriteGate {
+ fn hold_if_first(&self) {
+ if !self.taken.swap(true, Ordering::SeqCst) {
+ self.entered.wait();
+ self.release.wait();
+ }
+ }
+ }
+
+ /// Regression test for
<https://github.com/apache/datafusion/issues/24883>.
+ ///
+ /// Two writers push concurrently. The first writer is held inside its
first
+ /// write to disk (holding its file's lock) while the second writer pushes.
+ /// On the unfixed pool that leaves two open files with one batch each, and
+ /// the reader, which only ever drained the oldest file, parked on that
file
+ /// once it had read its single batch even though the second batch was on
+ /// disk. `RepartitionExec` blocks on this stream once per pushed batch
and,
+ /// with its channel gate closed, cannot push anything else to wake it.
+ ///
+ /// The reader must yield both batches while both writers are still alive.
+ #[tokio::test]
+ async fn test_reader_does_not_wait_on_drained_file_while_another_has_data()
+ -> Result<()> {
+ let gate = Arc::new(WriteGate {
+ taken: AtomicBool::new(false),
+ entered: Barrier::new(2),
+ release: Barrier::new(2),
+ });
+ let write_hook: WriteHook = {
+ let gate = Arc::clone(&gate);
+ Arc::new(move || gate.hold_if_first())
+ };
+ let (spill_manager, _disk_manager) = spill_manager_with_io_hooks(
+ write_hook,
+ Arc::new(|| ReadBehavior::Delay(Duration::ZERO)),
+ )?;
+
+ let (writer1, mut reader) = mpsc_channel(1024 * 1024, spill_manager);
+ let writer2 = writer1.clone();
+
+ // Writer 1 pushes and is held inside its first disk write, holding its
+ // file's lock. The writers run on plain threads because `push_batch`
+ // blocks.
+ let writer1 = std::thread::spawn(move || {
+ writer1.push_batch(&create_test_batch(0, 10)).unwrap();
+ writer1
+ });
+ gate.entered.wait();
+
+ // Writer 2 pushes while writer 1 is held.
+ let writer2 = std::thread::spawn(move || {
+ writer2.push_batch(&create_test_batch(10, 10)).unwrap();
+ writer2
+ });
+ // Let writer 2 either complete (into a second file) or block on writer
+ // 1's file lock, depending on the pool implementation.
+ std::thread::sleep(Duration::from_millis(200));
Review Comment:
Could we replace the fixed 200 ms sleep with a synchronization point that
confirms writer 2 has acquired or created its file before releasing writer 1?
On a heavily loaded runner, writer 2 might not get scheduled until after writer
1 is released. In that case it could reuse writer 1's returned file, leaving
the test with only one file. The test could then pass on the old implementation
without actually exercising the two-open-file case. This is just test
hardening, since the fuzzer gives us additional coverage as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
