net/
async_runtime.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
4use std::sync::OnceLock;
5use std::sync::atomic::{AtomicUsize, Ordering};
6use std::thread;
7use std::time::Duration;
8
9use futures::Future;
10use net_traits::AsyncRuntime;
11use tokio::runtime::{Builder, Handle, Runtime};
12
13/// The actual runtime,
14/// to be used as part of shut-down.
15pub struct AsyncRuntimeHolder {
16    runtime: Option<Runtime>,
17}
18
19impl AsyncRuntimeHolder {
20    pub(crate) fn new(runtime: Runtime) -> Self {
21        Self {
22            runtime: Some(runtime),
23        }
24    }
25}
26
27impl AsyncRuntime for AsyncRuntimeHolder {
28    fn shutdown(&mut self) {
29        self.runtime
30            .take()
31            .expect("Runtime should have been initialized on start-up.")
32            .shutdown_timeout(Duration::from_millis(100))
33    }
34}
35
36/// A shared handle to the runtime,
37/// to be initialized on start-up.
38static ASYNC_RUNTIME_HANDLE: OnceLock<Handle> = OnceLock::new();
39
40pub fn init_async_runtime() -> Box<dyn AsyncRuntime> {
41    // Initialize a tokio runtime.
42    let runtime = Builder::new_multi_thread()
43        .thread_name_fn(|| {
44            static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
45            let id = ATOMIC_ID.fetch_add(1, Ordering::Relaxed);
46            format!("tokio-runtime-{}", id)
47        })
48        .worker_threads(
49            thread::available_parallelism()
50                .map(|i| i.get())
51                .unwrap_or(servo_config::pref!(threadpools_fallback_worker_num) as usize)
52                .min(servo_config::pref!(threadpools_async_runtime_workers_max).max(1) as usize),
53        )
54        .enable_io()
55        .enable_time()
56        .build()
57        .expect("Unable to build tokio-runtime runtime");
58
59    // Make the runtime available to users inside this crate.
60    ASYNC_RUNTIME_HANDLE
61        .set(runtime.handle().clone())
62        .expect("Runtime handle should be initialized once on start-up");
63
64    // Return an async runtime for use in shutdown.
65    Box::new(AsyncRuntimeHolder::new(runtime))
66}
67
68/// Spawn a task using the handle to the runtime.
69pub fn spawn_task<F>(task: F)
70where
71    F: Future + 'static + std::marker::Send,
72    F::Output: Send + 'static,
73{
74    ASYNC_RUNTIME_HANDLE
75        .get()
76        .expect("Runtime handle should be initialized on start-up")
77        .spawn(task);
78}
79
80/// Spawn a blocking task using the handle to the runtime.
81pub fn spawn_blocking_task<F, R>(task: F) -> F::Output
82where
83    F: Future,
84{
85    ASYNC_RUNTIME_HANDLE
86        .get()
87        .expect("Runtime handle should be initialized on start-up")
88        .block_on(task)
89}