Skip to main content

librt_rs/
lib.rs

1use tokio::runtime::{Builder, Runtime};
2use tracing::{info, warn};
3
4pub mod signal;
5
6#[cfg(feature = "multicore")]
7pub fn build(thread_name: &str) -> Runtime {
8    // The proxy creates an additional admin thread, but it would be wasteful to
9    // allocate a whole core to it; so we let the main runtime consume all
10    // available cores. The number of available cores is determined by checking
11    // the environment or by inspecting the host or cgroups.
12    //
13    // The basic scheduler is used when the threaded scheduler would provide no
14    // benefit.
15    let mut cores = std::env::var("RT_CORES")
16        .ok()
17        .and_then(|v| {
18            let opt = v.parse::<usize>().ok().filter(|n| *n > 0);
19            if opt.is_none() {
20                warn!(RT_CORES = %v, "Ignoring invalid configuration");
21            }
22            opt
23        })
24        .unwrap_or(0);
25
26    let cpus = num_cpus::get();
27    debug_assert!(cpus > 0, "At least one CPU must be available");
28    if cores > cpus {
29        warn!(
30            cpus,
31            RT_CORES = cores,
32            "Ignoring configuration due to insufficient resources"
33        );
34        cores = cpus;
35    }
36
37    match cores {
38        // `0` is unexpected, but it's a wild world out there.
39        0 | 1 => {
40            info!("Using single-threaded runtime");
41            Builder::new_current_thread()
42                .enable_all()
43                .thread_name(thread_name)
44                .build()
45                .expect("failed to build basic runtime!")
46        }
47        num_cpus => {
48            info!(%cores, "Using multi-threaded runtime");
49            Builder::new_multi_thread()
50                .enable_all()
51                .thread_name(thread_name)
52                .worker_threads(num_cpus)
53                .max_blocking_threads(num_cpus)
54                .build()
55                .expect("failed to build threaded runtime!")
56        }
57    }
58}
59
60#[cfg(not(feature = "multicore"))]
61pub fn build(thread_name: &str) -> Runtime {
62    Builder::new_current_thread()
63        .enable_all()
64        .thread_name(thread_name)
65        .build()
66        .expect("failed to build basic runtime!")
67}