1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
use std::future::Future;
use tokio::task::JoinHandle;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RuntimeMode {
/// Spawns tasks on the current runtime that has invoked [`ServiceBuilder::build`] or
/// [`ServiceBuilder::build_mutable`].
///
/// This is the default.
#[default]
Current,
/// Spawns a single threaded [`tokio::runtime::Runtime`] for each shard.
///
/// This is useful when you are running into limitations with having a single runtime for all the
/// shards, because currently, a single tokio runtime becomes in-efficient when it spans a large number of worker
/// threads across multiple NUMA nodes.
///
/// See [this issue](https://github.com/tokio-rs/tokio/issues/5076) for more information.
RuntimePerShard,
}
impl RuntimeMode {
pub(super) fn spawn_shard(
&self,
shard_id: u8,
shard_future: impl Future<Output = ()> + Send + 'static,
) -> JoinHandle<()> {
match self {
Self::Current => tokio::task::spawn(shard_future),
Self::RuntimePerShard => {
let runtime = tokio::runtime::Builder::new_current_thread()
.thread_name_fn(move || format!("shard-{}-runtime-thread", shard_id))
.enable_all()
.build()
.expect("failed to build multi thread runtime");
let (shard_finished_tx, shard_finished_rx) = tokio::sync::oneshot::channel();
let join_handle = runtime.spawn(async move {
shard_future.await;
shard_finished_tx.send(()).ok();
});
std::thread::Builder::new()
.name(format!("shard-{}-runtime-main-thread", shard_id))
.spawn(move || runtime.block_on(shard_finished_rx))
.unwrap_or_else(|_| {
panic!("failed to spawn runtime main thread for shard={}", shard_id)
});
join_handle
}
}
}
}