subspace_process/
lib.rs

1use futures::channel::oneshot;
2use futures::channel::oneshot::Canceled;
3use futures::future::{Either, FusedFuture};
4use std::fmt::Display;
5use std::future::Future;
6use std::ops::Deref;
7use std::pin::{Pin, pin};
8use std::process::exit;
9use std::task::{Context, Poll};
10use std::{io, panic, thread};
11use tokio::runtime::Handle;
12use tokio::{signal, task};
13use tracing::level_filters::LevelFilter;
14use tracing::{debug, info, warn};
15use tracing_subscriber::layer::SubscriberExt;
16use tracing_subscriber::util::SubscriberInitExt;
17use tracing_subscriber::{EnvFilter, Layer, fmt};
18
19#[cfg(test)]
20mod tests;
21
22pub fn init_logger() {
23    // sc_informant and other substrate crates colorize log content via `console`;
24    // tracing-subscriber 0.3.20+ escapes those embedded ANSI codes into literal `\x1b[`
25    // text in the message. Turn console color off at the source so nothing is emitted.
26    //
27    // The workspace uses console 0.16.x, but substrate still links against console
28    // 0.15.x, which keeps its own global color state. Disable colors in both so the
29    // workaround remains effective for substrate's messages.
30    console::set_colors_enabled(false);
31    console::set_colors_enabled_stderr(false);
32    console_0_15::set_colors_enabled(false);
33    console_0_15::set_colors_enabled_stderr(false);
34
35    // TODO: Workaround for https://github.com/tokio-rs/tracing/issues/2214, also on
36    //  Windows terminal doesn't support the same colors as bash does
37    let enable_color = if cfg!(windows) {
38        false
39    } else {
40        supports_color::on(supports_color::Stream::Stderr).is_some()
41    };
42
43    let res = tracing_subscriber::registry()
44        .with(
45            fmt::layer().with_ansi(enable_color).with_filter(
46                EnvFilter::builder()
47                    .with_default_directive(LevelFilter::INFO.into())
48                    .from_env_lossy(),
49            ),
50        )
51        .try_init();
52
53    if let Err(e) = res {
54        // In production, this might be a bug in the logging setup.
55        // In some tests, it is expected.
56        eprintln!(
57            "Failed to initialize logger: {e}. \
58            This is expected when running nexttest test functions under `cargo test`."
59        );
60    }
61}
62
63/// Joins async join handle on drop.
64/// This future is fused, and will return `Poll::Pending` if polled after completion.
65#[derive(Debug)]
66pub struct AsyncJoinOnDrop<T> {
67    handle: Option<task::JoinHandle<T>>,
68    abort_on_drop: bool,
69}
70
71impl<T> Drop for AsyncJoinOnDrop<T> {
72    #[inline]
73    fn drop(&mut self) {
74        if let Some(handle) = self.handle.take() {
75            if self.abort_on_drop {
76                handle.abort();
77            }
78
79            if !handle.is_finished() {
80                task::block_in_place(move || {
81                    let _ = Handle::current().block_on(handle);
82                });
83            }
84        }
85    }
86}
87
88impl<T> AsyncJoinOnDrop<T> {
89    /// Create new instance.
90    #[inline]
91    pub fn new(handle: task::JoinHandle<T>, abort_on_drop: bool) -> Self {
92        Self {
93            handle: Some(handle),
94            abort_on_drop,
95        }
96    }
97}
98
99impl<T> FusedFuture for AsyncJoinOnDrop<T> {
100    fn is_terminated(&self) -> bool {
101        self.handle.is_none()
102    }
103}
104
105impl<T> Future for AsyncJoinOnDrop<T> {
106    type Output = Result<T, task::JoinError>;
107
108    #[inline]
109    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
110        if let Some(handle) = self.handle.as_mut() {
111            let result = Pin::new(handle).poll(cx);
112            if result.is_ready() {
113                // Drop the handle, because if we poll it again, it will panic.
114                self.handle.take();
115            }
116            result
117        } else {
118            Poll::Pending
119        }
120    }
121}
122
123/// Joins synchronous join handle on drop
124pub(crate) struct JoinOnDrop(Option<thread::JoinHandle<()>>);
125
126impl Drop for JoinOnDrop {
127    #[inline]
128    fn drop(&mut self) {
129        self.0
130            .take()
131            .expect("Always called exactly once; qed")
132            .join()
133            .expect("Panic if background thread panicked");
134    }
135}
136
137impl JoinOnDrop {
138    // Create new instance
139    #[inline]
140    pub(crate) fn new(handle: thread::JoinHandle<()>) -> Self {
141        Self(Some(handle))
142    }
143}
144
145impl Deref for JoinOnDrop {
146    type Target = thread::JoinHandle<()>;
147
148    #[inline]
149    fn deref(&self) -> &Self::Target {
150        self.0.as_ref().expect("Only dropped in Drop impl; qed")
151    }
152}
153
154/// Runs future on a dedicated thread with the specified name. Will block on drop until background
155/// thread with future is stopped, ensuring nothing is left in memory.
156///
157/// Some OSes (like Linux) truncate thread names at 15 characters due to kernel limits.
158pub fn run_future_in_dedicated_thread<CreateFut, Fut, T>(
159    create_future: CreateFut,
160    thread_name: String,
161) -> io::Result<impl Future<Output = Result<T, Canceled>> + Send>
162where
163    CreateFut: (FnOnce() -> Fut) + Send + 'static,
164    Fut: Future<Output = T> + 'static,
165    T: Send + 'static,
166{
167    let (drop_tx, drop_rx) = oneshot::channel::<()>();
168    let (result_tx, result_rx) = oneshot::channel();
169    let handle = Handle::current();
170    let join_handle = thread::Builder::new().name(thread_name).spawn(move || {
171        let _tokio_handle_guard = handle.enter();
172
173        let future = pin!(create_future());
174
175        let result = match handle.block_on(futures::future::select(future, drop_rx)) {
176            Either::Left((result, _)) => result,
177            Either::Right(_) => {
178                // Outer future was dropped, nothing left to do
179                return;
180            }
181        };
182        if let Err(_error) = result_tx.send(result) {
183            debug!(
184                thread_name = ?thread::current().name(),
185                "Future finished, but receiver was already dropped",
186            );
187        }
188    })?;
189    // Ensure thread will not be left hanging forever
190    let join_on_drop = JoinOnDrop::new(join_handle);
191
192    Ok(async move {
193        let result = result_rx.await;
194        drop(drop_tx);
195        drop(join_on_drop);
196        result
197    })
198}
199
200/// Wait for the process to receive a shutdown signal, and log the supplied process kind.
201#[cfg(unix)]
202pub async fn shutdown_signal(process_kind: impl Display) {
203    use futures::FutureExt;
204    use std::pin::pin;
205
206    let mut sigint = signal::unix::signal(signal::unix::SignalKind::interrupt())
207        .expect("Setting signal handlers must never fail");
208    let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())
209        .expect("Setting signal handlers must never fail");
210
211    futures::future::select(
212        pin!(sigint.recv().map(|_| {
213            info!("Received SIGINT, shutting down {process_kind}...");
214        }),),
215        pin!(sigterm.recv().map(|_| {
216            info!("Received SIGTERM, shutting down {process_kind}...");
217        }),),
218    )
219    .await;
220}
221
222/// Wait for the process to receive a shutdown signal, and log the supplied process kind.
223#[cfg(not(unix))]
224pub async fn shutdown_signal(process_kind: impl Display) {
225    signal::ctrl_c()
226        .await
227        .expect("Setting signal handlers must never fail");
228
229    info!("Received Ctrl+C, shutting down {process_kind}...");
230}
231
232/// Raise the file descriptor limit for the process to the maximum possible value.
233pub fn raise_fd_limit() {
234    match fdlimit::raise_fd_limit() {
235        Ok(fdlimit::Outcome::LimitRaised { from, to }) => {
236            debug!(
237                "Increased file descriptor limit from previous (most likely soft) limit {} to \
238                new (most likely hard) limit {}",
239                from, to
240            );
241        }
242        Ok(fdlimit::Outcome::Unsupported) => {
243            // Unsupported platform (a platform other than Linux or macOS)
244        }
245        Err(error) => {
246            warn!(
247                "Failed to increase file descriptor limit for the process due to an error: {}.",
248                error
249            );
250        }
251    }
252}
253
254/// Install a panic handler which exits on panics, rather than unwinding. Unwinding can hang the
255/// tokio runtime waiting for stuck tasks or threads.
256pub fn set_exit_on_panic() {
257    let default_panic_hook = panic::take_hook();
258    panic::set_hook(Box::new(move |panic_info| {
259        default_panic_hook(panic_info);
260        exit(1);
261    }));
262}
subspace_process/lib.rs

subspace_process/
lib.rs
