-
Notifications
You must be signed in to change notification settings - Fork 142
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
perf: use separate dropper thread for
current_thread
flavor (#666)
`tokio::test` run with `current_thread` flavor by default and it's typical use-case. However, we need to perform some async operations on drop, and it used to spawn a new thread for each drop. It's inefficient and was a temporary solution. Now, it uses: - for `multi-thread` runtime flavor: `tokio::task::block_in_place` - for `current-thread` runtime flavor: lazily spawns a separate thread for dropping tasks, thus it will be only 1 thread regardless of the number of containers, networks, etc.
- Loading branch information
Showing
5 changed files
with
63 additions
and
38 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,54 @@ | ||
use std::sync::OnceLock; | ||
|
||
use futures::future::BoxFuture; | ||
|
||
static DROP_TASK_SENDER: OnceLock<tokio::sync::mpsc::UnboundedSender<BoxFuture<'static, ()>>> = | ||
OnceLock::new(); | ||
|
||
/// A helper to perform async operations in `Drop` implementation. | ||
/// | ||
/// The behavior depends on the runtime flavor used in the test: | ||
/// - `multi-threaded` runtime: it will use `tokio::task::block_in_place` to run the provided future | ||
/// - `current-thread` runtime: it spawns a separate tokio runtime in a dedicated thread to run the provided futures. | ||
/// * Only 1 drop-worker for the process, regardless of number of containers and drops. | ||
// We can consider creating `AsyncDrop` trait + `AsyncDropGuard<T: AsyncDrop>` wrapper to make it more ergonomic. | ||
// However, we have a only couple of places where we need this functionality. | ||
pub(crate) fn async_drop(future: impl std::future::Future<Output = ()> + Send + 'static) { | ||
let handle = tokio::runtime::Handle::current(); | ||
match handle.runtime_flavor() { | ||
tokio::runtime::RuntimeFlavor::CurrentThread => { | ||
let (tx, rx) = std::sync::mpsc::sync_channel(1); | ||
dropper_task_sender() | ||
.send(Box::pin(async move { | ||
future.await; | ||
let _ = tx.send(()); | ||
})) | ||
.expect("drop-worker must be running: failed to send drop task"); | ||
let _ = rx.recv(); | ||
} | ||
tokio::runtime::RuntimeFlavor::MultiThread => { | ||
tokio::task::block_in_place(move || handle.block_on(future)) | ||
} | ||
_ => unreachable!("unsupported runtime flavor"), | ||
} | ||
} | ||
|
||
fn dropper_task_sender() -> &'static tokio::sync::mpsc::UnboundedSender<BoxFuture<'static, ()>> { | ||
DROP_TASK_SENDER.get_or_init(|| { | ||
let (dropper_tx, mut dropper_rx) = tokio::sync::mpsc::unbounded_channel(); | ||
std::thread::spawn(move || { | ||
tokio::runtime::Builder::new_current_thread() | ||
.thread_name("testcontainers-drop-worker") | ||
.enable_all() | ||
.build() | ||
.expect("failed to create dropper runtime") | ||
.block_on(async move { | ||
while let Some(future) = dropper_rx.recv().await { | ||
future.await; | ||
} | ||
}); | ||
}); | ||
|
||
dropper_tx | ||
}) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters