ckb_async_runtime/
native.rs

1use ckb_spawn::Spawn;
2use core::future::Future;
3use std::sync::atomic::{AtomicU32, Ordering};
4use std::thread::available_parallelism;
5use tokio::runtime::{Builder, Handle as TokioHandle, Runtime};
6
7use tokio::sync::mpsc::{Receiver, Sender};
8use tokio::task::JoinHandle;
9
10// Handle is a newtype wrap and unwrap tokio::Handle, it is workaround with Rust Orphan Rules.
11// We need `Handle` impl ckb spawn trait decouple tokio dependence
12
13/// Handle to the runtime.
14#[derive(Debug, Clone)]
15pub struct Handle {
16    pub(crate) inner: TokioHandle,
17    guard: Option<Sender<()>>,
18}
19
20impl Handle {
21    /// Create a new Handle
22    pub fn new(inner: TokioHandle, guard: Option<Sender<()>>) -> Self {
23        Self { inner, guard }
24    }
25
26    /// Drop the guard
27    pub fn drop_guard(&mut self) {
28        let _ = self.guard.take();
29    }
30}
31
32impl Handle {
33    /// Enter the runtime context. This allows you to construct types that must
34    /// have an executor available on creation such as [`tokio::time::Sleep`] or [`tokio::net::TcpStream`].
35    /// It will also allow you to call methods such as [`tokio::spawn`].
36    pub fn enter<F, R>(&self, f: F) -> R
37    where
38        F: FnOnce() -> R,
39    {
40        let _enter = self.inner.enter();
41        f()
42    }
43
44    /// Spawns a future onto the runtime.
45    ///
46    /// This spawns the given future onto the runtime's executor
47    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
48    where
49        F: Future + Send + 'static,
50        F::Output: Send + 'static,
51    {
52        let tokio_task_guard = self.guard.clone();
53
54        self.inner.spawn(async move {
55            // move tokio_task_guard into the spawned future
56            // so that it will be dropped when the future is finished
57            let _guard = tokio_task_guard;
58            future.await
59        })
60    }
61
62    /// Run a future to completion on the Tokio runtime from a synchronous context.
63    pub fn block_on<F: Future>(&self, future: F) -> F::Output {
64        self.inner.block_on(future)
65    }
66
67    /// Spawns a future onto the runtime blocking pool.
68    ///
69    /// This spawns the given future onto the runtime's blocking executor
70    pub fn spawn_blocking<F, R>(&self, f: F) -> JoinHandle<R>
71    where
72        F: FnOnce() -> R + Send + 'static,
73        R: Send + 'static,
74    {
75        self.inner.spawn_blocking(f)
76    }
77
78    /// Transform to inner tokio handler
79    pub fn into_inner(self) -> TokioHandle {
80        self.inner
81    }
82}
83
84/// Create a new runtime with unique name.
85fn new_runtime(worker_num: Option<usize>) -> Runtime {
86    Builder::new_multi_thread()
87        .enable_all()
88        .worker_threads(worker_num.unwrap_or_else(|| available_parallelism().unwrap().into()))
89        .thread_name_fn(|| {
90            static ATOMIC_ID: AtomicU32 = AtomicU32::new(0);
91            let id = ATOMIC_ID
92                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| {
93                    // A long thread name will cut to 15 characters in debug tools.
94                    // Such as "top", "htop", "gdb" and so on.
95                    // It's a kernel limit.
96                    //
97                    // So if we want to see the whole name in debug tools,
98                    // this number should have 6 digits at most,
99                    // since the prefix uses 9 characters in below code.
100                    //
101                    // There still has a issue:
102                    // When id wraps around, we couldn't know whether the old id
103                    // is released or not.
104                    // But we can ignore this, because it's almost impossible.
105                    if n >= 999_999 { Some(0) } else { Some(n + 1) }
106                })
107                .expect("impossible since the above closure must return Some(number)");
108            format!("GlobalRt-{id}")
109        })
110        .build()
111        .expect("ckb runtime initialized")
112}
113
114/// Create new threaded_scheduler tokio Runtime, return `Runtime`
115pub fn new_global_runtime(worker_num: Option<usize>) -> (Handle, Receiver<()>, Runtime) {
116    let runtime = new_runtime(worker_num);
117    let handle = runtime.handle().clone();
118    let (guard, handle_stop_rx): (Sender<()>, Receiver<()>) = tokio::sync::mpsc::channel::<()>(1);
119
120    (Handle::new(handle, Some(guard)), handle_stop_rx, runtime)
121}
122
123/// Create new threaded_scheduler tokio Runtime, return `Handle` and background thread join handle,
124/// NOTICE: This is only used in testing
125pub fn new_background_runtime() -> Handle {
126    let runtime = new_runtime(None);
127    let handle = runtime.handle().clone();
128
129    let (guard, mut handle_stop_rx): (Sender<()>, Receiver<()>) =
130        tokio::sync::mpsc::channel::<()>(1);
131    let _thread = std::thread::Builder::new()
132        .name("GlobalRtBuilder".to_string())
133        .spawn(move || {
134            let ret = runtime.block_on(async move { handle_stop_rx.recv().await });
135            ckb_logger::debug!("Global runtime finished {:?}", ret);
136        })
137        .expect("tokio runtime started");
138
139    Handle::new(handle, Some(guard))
140}
141
142impl Spawn for Handle {
143    fn spawn_task<F>(&self, future: F)
144    where
145        F: Future<Output = ()> + Send + 'static,
146    {
147        self.spawn(future);
148    }
149}
ckb_async_runtime/native.rs

ckb_async_runtime/
native.rs
