Skip to content

Commit 0ffd4d9

Browse files
committed
fix: wait for scheduler tasks shutdown in parallel
1 parent 416131b commit 0ffd4d9

File tree

1 file changed

+19
-8
lines changed

1 file changed

+19
-8
lines changed

src/scheduler.rs

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use rand::Rng;
1111
use tokio::sync::{oneshot, RwLock, RwLockWriteGuard};
1212
use tokio::task;
1313
use tokio_util::sync::CancellationToken;
14+
use tokio_util::task::TaskTracker;
1415

1516
use self::connectivity::ConnectivityStore;
1617
use crate::config::{self, Config};
@@ -989,16 +990,26 @@ impl Scheduler {
989990

990991
// Actually shutdown tasks.
991992
let timeout_duration = std::time::Duration::from_secs(30);
993+
994+
let tracker = TaskTracker::new();
992995
for b in once(self.inbox).chain(self.oboxes) {
993-
tokio::time::timeout(timeout_duration, b.handle)
994-
.await
995-
.log_err(context)
996-
.ok();
996+
let context = context.clone();
997+
tracker.spawn(async move {
998+
tokio::time::timeout(timeout_duration, b.handle)
999+
.await
1000+
.log_err(&context)
1001+
});
9971002
}
998-
tokio::time::timeout(timeout_duration, self.smtp_handle)
999-
.await
1000-
.log_err(context)
1001-
.ok();
1003+
{
1004+
let context = context.clone();
1005+
tracker.spawn(async move {
1006+
tokio::time::timeout(timeout_duration, self.smtp_handle)
1007+
.await
1008+
.log_err(&context)
1009+
});
1010+
}
1011+
tracker.close();
1012+
tracker.wait().await;
10021013

10031014
// Abort tasks, then await them to ensure the `Future` is dropped.
10041015
// Just aborting the task may keep resources such as `Context` clone

0 commit comments

Comments
 (0)