ranflarion commented on code in PR #25763:
URL: https://github.com/apache/datafusion/pull/25763#discussion_r4363493455
##########
datafusion/common-runtime/src/join_set.rs:
##########
@@ -142,34 +218,271 @@ impl<T: 'static> JoinSet<T> {
&mut self,
cx: &mut Context<'_>,
) -> Poll<Option<Result<T, JoinError>>> {
- self.inner.poll_join_next(cx)
+ let joined = ready!(self.inner.poll_join_next_with_id(cx));
+ self.forget(joined.as_ref());
+ Poll::Ready(joined.map(|res| res.map(|(_, output)| output)))
}
/// [JoinSet::join_next_with_id](tokio::task::JoinSet::join_next_with_id)
- Await the next completed task with its ID.
pub async fn join_next_with_id(&mut self) -> Option<Result<(Id, T),
JoinError>> {
- self.inner.join_next_with_id().await
+ let joined = self.inner.join_next_with_id().await;
+ self.forget(joined.as_ref());
+ joined
}
///
[JoinSet::try_join_next_with_id](tokio::task::JoinSet::try_join_next_with_id) -
Try to join the next completed task with its ID.
pub fn try_join_next_with_id(&mut self) -> Option<Result<(Id, T),
JoinError>> {
- self.inner.try_join_next_with_id()
+ let joined = self.inner.try_join_next_with_id();
+ self.forget(joined.as_ref());
+ joined
}
///
[JoinSet::poll_join_next_with_id](tokio::task::JoinSet::poll_join_next_with_id)
- Poll for the next completed task with its ID.
pub fn poll_join_next_with_id(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Option<Result<(Id, T), JoinError>>> {
- self.inner.poll_join_next_with_id(cx)
+ let joined = ready!(self.inner.poll_join_next_with_id(cx));
+ self.forget(joined.as_ref());
+ Poll::Ready(joined)
}
/// [JoinSet::shutdown](tokio::task::JoinSet::shutdown) - Abort all tasks
and wait for shutdown.
pub async fn shutdown(&mut self) {
+ self.reclaim_all();
self.inner.shutdown().await
}
/// [JoinSet::join_all](tokio::task::JoinSet::join_all) - Await all tasks.
- pub async fn join_all(self) -> Vec<T> {
- self.inner.join_all().await
+ pub async fn join_all(mut self) -> Vec<T> {
+ // Tokio's loop, run on this set so that dropping the future reclaims
before aborting
+ let mut output = Vec::with_capacity(self.len());
+ while let Some(res) = self.join_next().await {
+ match res {
+ Ok(t) => output.push(t),
+ Err(err) if err.is_panic() =>
std::panic::resume_unwind(err.into_panic()),
+ Err(err) => panic!("{err}"),
+ }
+ }
+ output
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::future::pending;
+
+ use tokio::sync::oneshot;
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn drop_releases_idle_tasks_in_place() {
+ let owned = Arc::new(());
+ let released = Arc::downgrade(&owned);
+ let mut set = JoinSet::new();
+ let mut started = vec![];
+ for _ in 0..4 {
+ let owned = Arc::clone(&owned);
+ let (started_tx, started_rx) = oneshot::channel();
+ started.push(started_rx);
+ set.spawn_reclaimable(async move {
+ started_tx.send(()).unwrap();
+ pending::<()>().await;
+ drop(owned);
+ });
+ }
+ drop(owned);
+ for started_rx in started {
+ started_rx.await.unwrap();
+ }
+ set.reclaim
+ .values()
+ .for_each(ReclaimHandle::wait_until_idle);
+
+ drop(set);
+ assert_eq!(released.strong_count(), 0);
+ }
+
+ #[tokio::test]
+ async fn joined_tasks_are_forgotten() {
+ let mut set = JoinSet::new();
+ for i in 0..3 {
+ set.spawn_reclaimable(async move { i });
+ }
+ let mut outputs = vec![];
+ while let Some(output) = set.join_next().await {
+ outputs.push(output.unwrap());
+ }
+ outputs.sort_unstable();
+ assert_eq!(outputs, vec![0, 1, 2]);
+ assert!(set.reclaim.is_empty());
+ }
+
+ #[tokio::test]
+ async fn detached_tasks_keep_running() {
+ let mut set = JoinSet::new();
+ let (resume_tx, resume_rx) = oneshot::channel::<()>();
+ let (done_tx, done_rx) = oneshot::channel();
+ set.spawn_reclaimable(async move {
+ resume_rx.await.unwrap();
+ done_tx.send(()).unwrap();
+ });
+ set.detach_all();
+ drop(set);
+
+ resume_tx.send(()).unwrap();
+ done_rx.await.unwrap();
+ }
+
+ #[tokio::test]
+ async fn abort_all_reports_a_panic_from_a_reclaimed_destructor() {
+ struct PanicOnDrop;
+ impl Drop for PanicOnDrop {
+ fn drop(&mut self) {
+ panic!("destructor panic");
+ }
+ }
+ let mut set = JoinSet::new();
+ set.spawn_reclaimable(async move {
+ let _panics = PanicOnDrop;
+ pending::<()>().await;
+ });
+ // Lets the task start and park
+ tokio::task::yield_now().await;
+
+ set.abort_all();
+ assert!(set.join_next().await.unwrap().unwrap_err().is_panic());
+ }
+
+ #[test]
+ fn runtime_shutdown_does_not_wait_for_a_reclaiming_destructor() {
+ use std::sync::mpsc::{Receiver, Sender, channel};
+ use std::time::Duration;
+
+ struct BlockInDrop {
+ entered: Sender<()>,
+ resume: Receiver<()>,
+ }
+ impl Drop for BlockInDrop {
+ fn drop(&mut self) {
+ self.entered.send(()).unwrap();
+ let _ = self.resume.recv();
+ }
+ }
+
+ let runtime = tokio::runtime::Builder::new_multi_thread()
+ .worker_threads(1)
+ .build()
+ .unwrap();
+ let (entered_tx, entered_rx) = channel();
+ let (resume_tx, resume_rx) = channel();
+ let (started_tx, started_rx) = channel();
+ let guard = BlockInDrop {
+ entered: entered_tx,
+ resume: resume_rx,
+ };
+ let mut set = JoinSet::new();
+ set.spawn_reclaimable_on(
+ async move {
+ started_tx.send(()).unwrap();
+ pending::<()>().await;
+ drop(guard);
+ },
+ runtime.handle(),
+ );
+ started_rx.recv().unwrap();
+ set.reclaim
+ .values()
+ .for_each(ReclaimHandle::wait_until_idle);
+
+ // The set's drop runs the destructor on this thread until the test
resumes it
+ let dropper = std::thread::spawn(move || drop(set));
+ entered_rx.recv().unwrap();
+ let (shut_down_tx, shut_down_rx) = channel();
+ std::thread::spawn(move || {
+ runtime.shutdown_timeout(Duration::from_secs(10));
+ shut_down_tx.send(()).unwrap();
+ });
+ let shut_down =
shut_down_rx.recv_timeout(Duration::from_secs(5)).is_ok();
+ resume_tx.send(()).unwrap();
+ dropper.join().unwrap();
+ assert!(
+ shut_down,
+ "shutdown waited for a destructor on another thread"
+ );
+ }
+
+ #[test]
+ fn dropping_join_all_reclaims_before_aborting() {
+ use std::sync::mpsc::{Receiver, Sender, channel};
+ use std::time::Duration;
+
+ struct BlockInDrop {
Review Comment:
moved `PanicOnDrop`, `BlockInDrop` and the one-worker runtime into a
`#[cfg(test)]` `test_utils` module
##########
datafusion/common-runtime/src/join_set.rs:
##########
@@ -56,6 +68,27 @@ impl<T> JoinSet<T> {
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
+
+ fn reclaim_all(&mut self) {
+ for (_, task) in self.reclaim.drain() {
+ task.reclaim();
+ }
+ }
+
+ fn forget<R>(&mut self, joined: Option<&Result<(Id, R), JoinError>>) {
Review Comment:
It does hash, hashbrown's `remove_entry` calls `make_hash` before probing
even when the table is empty. Added the early return, so sets without
reclaimable tasks skip it.
--
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]
