zng_env/
process.rs

1use std::{
2    mem,
3    sync::atomic::{AtomicBool, Ordering},
4};
5
6use parking_lot::Mutex;
7
8#[doc(hidden)]
9#[cfg(not(target_arch = "wasm32"))]
10pub use linkme as __linkme;
11
12/// Register a `FnOnce(&ProcessStartArgs)` closure to be called on [`init!`].
13///
14/// Components that spawn special process instances implemented on the same executable
15/// can use this macro to inject their own "main" without needing to ask the user to plug an init
16/// function on the executable main. The component can spawn an instance of the current executable
17/// with marker environment variables that identify the component's process.
18///
19/// [`init!`]: crate::init!
20///
21/// # Examples
22///
23/// The example below declares a "main" for a foo component and a function that spawns it.
24///
25/// ```
26/// zng_env::on_process_start!(|args| {
27///     if args.yield_count == 0 {
28///         return args.yield_once();
29///     }
30///
31///     if std::env::var("FOO_MARKER").is_ok() {
32///         println!("Spawned as foo!");
33///         zng_env::exit(0);
34///     }
35/// });
36///
37/// fn main() {
38///     zng_env::init!(); // foo_main OR
39///     // normal main
40/// }
41///
42/// pub fn spawn_foo() -> std::io::Result<()> {
43///     std::process::Command::new(std::env::current_exe()?).env("FOO_MARKER", "").spawn()?;
44///     Ok(())
45/// }
46/// ```
47///
48/// Note that the handler yields once, this gives a chance for all handlers to run first before the handler is called again
49/// and takes over the process. It is good practice to yield at least once to ensure handlers that are supported to affect all
50/// processes actually init, as an example, the trace recorder may never start for the process if it does not yield.
51///
52/// Also note the use of custom [`exit`], it is important to call it to collaborate with [`on_process_exit`] handlers.
53///
54/// # App Context
55///
56/// This event happens on the executable process context, before any `APP` context starts, you can use
57/// `zng::app::on_app_start` here to register a handler to be called in the app context, if and when it starts.
58///
59/// # Web Assembly
60///
61/// Crates that declare `on_process_start` must have the [`wasm_bindgen`] dependency to compile for the `wasm32` target.
62///
63/// In `Cargo.toml` add this dependency:
64///
65/// ```toml
66/// [target.'cfg(target_arch = "wasm32")'.dependencies]
67/// wasm-bindgen = "*"
68/// ```
69///
70/// Try to match the version used by `zng-env`.
71///
72/// # Linker Optimizer Issues
73///
74/// The macOS system linker can "optimize" away crates that are only referenced via this macro, that is, a crate dependency
75/// that is not otherwise directly addressed by code. To workaround this issue you can add a bogus reference to the crate code, something
76/// that is not trivial to optimize away. Unfortunately this code must be added on the dependent crate, or on an intermediary dependency,
77/// if your crate is at risk of being used this way please document this issue.
78///
79/// See [`zng#437`] for an example of how to fix this issue.
80///
81/// [`wasm_bindgen`]: https://crates.io/crates/wasm-bindgen
82/// [`zng#437`]: https://github.com/zng-ui/zng/pull/437
83#[macro_export]
84macro_rules! on_process_start {
85    ($closure:expr) => {
86        $crate::__on_process_start! {$closure}
87    };
88}
89
90#[cfg(not(target_arch = "wasm32"))]
91#[doc(hidden)]
92#[macro_export]
93macro_rules! __on_process_start {
94    ($closure:expr) => {
95        #[$crate::__linkme::distributed_slice($crate::ZNG_ENV_ON_PROCESS_START)]
96        #[linkme(crate = $crate::__linkme)]
97        #[doc(hidden)]
98        static _ON_PROCESS_START: fn(&$crate::ProcessStartArgs) = _on_process_start;
99        #[doc(hidden)]
100        fn _on_process_start(args: &$crate::ProcessStartArgs) {
101            fn on_process_start(args: &$crate::ProcessStartArgs, handler: impl FnOnce(&$crate::ProcessStartArgs)) {
102                handler(args)
103            }
104            on_process_start(args, $closure)
105        }
106    };
107}
108
109#[cfg(target_arch = "wasm32")]
110#[doc(hidden)]
111#[macro_export]
112macro_rules! __on_process_start {
113    ($closure:expr) => {
114        $crate::wasm_process_start! {$crate,$closure}
115    };
116}
117
118#[doc(hidden)]
119#[cfg(target_arch = "wasm32")]
120pub use wasm_bindgen::prelude::wasm_bindgen;
121
122#[doc(hidden)]
123#[cfg(target_arch = "wasm32")]
124pub use zng_env_proc_macros::wasm_process_start;
125use zng_txt::Txt;
126
127#[cfg(target_arch = "wasm32")]
128std::thread_local! {
129    #[doc(hidden)]
130    pub static WASM_INIT: std::cell::RefCell<Vec<fn(&ProcessStartArgs)>> = const { std::cell::RefCell::new(vec![]) };
131}
132
133#[cfg(not(target_arch = "wasm32"))]
134#[doc(hidden)]
135#[linkme::distributed_slice]
136pub static ZNG_ENV_ON_PROCESS_START: [fn(&ProcessStartArgs)];
137
138#[cfg(not(target_arch = "wasm32"))]
139pub(crate) fn process_init() -> impl Drop {
140    process_init_impl(&ZNG_ENV_ON_PROCESS_START)
141}
142
143fn process_init_impl(handlers: &[fn(&ProcessStartArgs)]) -> MainExitHandler {
144    let process_state = std::mem::replace(
145        &mut *zng_unique_id::hot_static_ref!(PROCESS_LIFETIME_STATE).lock(),
146        ProcessLifetimeState::Inited,
147    );
148    assert_eq!(process_state, ProcessLifetimeState::BeforeInit, "init!() already called");
149
150    let mut yielded = vec![];
151    let mut next_handlers_count = handlers.len();
152    for h in handlers {
153        next_handlers_count -= 1;
154        let args = ProcessStartArgs {
155            next_handlers_count,
156            yield_count: 0,
157            yield_requested: AtomicBool::new(false),
158        };
159        h(&args);
160        if args.yield_requested.load(Ordering::Relaxed) {
161            yielded.push(h);
162            next_handlers_count += 1;
163        }
164    }
165
166    let mut yield_count = 0;
167    while !yielded.is_empty() {
168        yield_count += 1;
169        if yield_count > ProcessStartArgs::MAX_YIELD_COUNT {
170            eprintln!("start handlers requested `yield_start` more them 32 times");
171            break;
172        }
173
174        next_handlers_count = yielded.len();
175        for h in mem::take(&mut yielded) {
176            next_handlers_count -= 1;
177            let args = ProcessStartArgs {
178                next_handlers_count,
179                yield_count,
180                yield_requested: AtomicBool::new(false),
181            };
182            h(&args);
183            if args.yield_requested.load(Ordering::Relaxed) {
184                yielded.push(h);
185                next_handlers_count += 1;
186            }
187        }
188    }
189    MainExitHandler
190}
191
192#[cfg(target_arch = "wasm32")]
193pub(crate) fn process_init() -> impl Drop {
194    std::panic::set_hook(Box::new(console_error_panic_hook::hook));
195
196    let window = web_sys::window().expect("cannot 'init!', no window object");
197    let module = js_sys::Reflect::get(&window, &"__zng_env_init_module".into())
198        .expect("cannot 'init!', missing module in 'window.__zng_env_init_module'");
199
200    if module == wasm_bindgen::JsValue::undefined() || module == wasm_bindgen::JsValue::null() {
201        panic!("cannot 'init!', missing module in 'window.__zng_env_init_module'");
202    }
203
204    let module: js_sys::Object = module.into();
205
206    for entry in js_sys::Object::entries(&module) {
207        let entry: js_sys::Array = entry.into();
208        let ident = entry.get(0).as_string().expect("expected ident at entry[0]");
209
210        if ident.starts_with("__zng_env_start_") {
211            let func: js_sys::Function = entry.get(1).into();
212            if let Err(e) = func.call0(&wasm_bindgen::JsValue::NULL) {
213                panic!("'init!' function error, {e:?}");
214            }
215        }
216    }
217
218    process_init_impl(&WASM_INIT.with_borrow_mut(std::mem::take))
219}
220
221/// Arguments for [`on_process_start`] handlers.
222///
223/// Empty in this release.
224pub struct ProcessStartArgs {
225    /// Number of start handlers yet to run.
226    pub next_handlers_count: usize,
227
228    /// Number of times this handler has yielded.
229    ///
230    /// If this exceeds 32 times the handler is ignored.
231    pub yield_count: u16,
232
233    yield_requested: AtomicBool,
234}
235impl ProcessStartArgs {
236    /// Yield requests after this are ignored.
237    pub const MAX_YIELD_COUNT: u16 = 32;
238
239    /// Let other process start handlers run first.
240    ///
241    /// The handler must call this if it takes over the process and it cannot determinate if it should from the environment.
242    pub fn yield_once(&self) {
243        self.yield_requested.store(true, Ordering::Relaxed);
244    }
245}
246
247struct MainExitHandler;
248impl Drop for MainExitHandler {
249    fn drop(&mut self) {
250        run_exit_handlers(if std::thread::panicking() { 101 } else { 0 })
251    }
252}
253
254type ExitHandler = Box<dyn FnOnce(&ProcessExitArgs) + Send + 'static>;
255
256zng_unique_id::hot_static! {
257    static ON_PROCESS_EXIT: Mutex<Vec<ExitHandler>> = Mutex::new(vec![]);
258}
259
260/// Terminates the current process with the specified exit code.
261///
262/// This function must be used instead of `std::process::exit` as it runs the [`on_process_exit`].
263pub fn exit(code: i32) -> ! {
264    run_exit_handlers(code);
265    std::process::exit(code)
266}
267
268fn run_exit_handlers(code: i32) {
269    *zng_unique_id::hot_static_ref!(PROCESS_LIFETIME_STATE).lock() = ProcessLifetimeState::Exiting;
270
271    let on_exit = mem::take(&mut *zng_unique_id::hot_static_ref!(ON_PROCESS_EXIT).lock());
272    let args = ProcessExitArgs { code };
273    for h in on_exit {
274        h(&args);
275    }
276}
277
278/// Arguments for [`on_process_exit`] handlers.
279#[non_exhaustive]
280pub struct ProcessExitArgs {
281    /// Exit code that will be used.
282    pub code: i32,
283}
284
285/// Register a `handler` to run once when the current process exits.
286///
287/// Note that the handler is only called if the process is terminated by [`exit`], or by the executable main
288/// function returning if [`init!`] is called on it.
289///
290/// [`init!`]: crate::init!
291pub fn on_process_exit(handler: impl FnOnce(&ProcessExitArgs) + Send + 'static) {
292    zng_unique_id::hot_static_ref!(ON_PROCESS_EXIT).lock().push(Box::new(handler))
293}
294
295/// Defines the state of the current process instance.
296///
297/// Use [`process_lifetime_state()`] to get.
298#[derive(Debug, Clone, Copy, PartialEq, Eq)]
299pub enum ProcessLifetimeState {
300    /// Init not called yet.
301    BeforeInit,
302    /// Init called and the function where it is called has not returned yet.
303    Inited,
304    /// Init called and the function where it is called is returning.
305    Exiting,
306}
307
308zng_unique_id::hot_static! {
309    static PROCESS_LIFETIME_STATE: Mutex<ProcessLifetimeState> = Mutex::new(ProcessLifetimeState::BeforeInit);
310}
311zng_unique_id::hot_static! {
312    static PROCESS_NAME: Mutex<Txt> = Mutex::new(Txt::from_static(""));
313}
314
315/// Get the state of the current process instance.
316pub fn process_lifetime_state() -> ProcessLifetimeState {
317    *zng_unique_id::hot_static_ref!(PROCESS_LIFETIME_STATE).lock()
318}
319
320/// Gets a process runtime name.
321///
322/// The primary use of this name is to identify the process in logs, see [`set_process_name`] for details about the logged name.
323/// On set or init the name is logged as an info message "pid: {pid}, name: {name}".
324///
325/// # Common Names
326///
327/// All Zng provided process handlers name the process.
328///
329/// * `"app-process"` - Set by `APP` if no other name was set before the app starts building.
330/// * `"view-process"` - Set by the view-process implementer when running in multi process mode.
331/// * `"crash-handler-process"` - Set by the crash-handler when running with crash handling.
332/// * `"crash-dialog-process"` - Set by the crash-handler on the crash dialog process.
333/// * `"worker-process ({worker_name}, {pid})"` - Set by task worker processes if no name was set before the task runner server starts.
334pub fn process_name() -> Txt {
335    zng_unique_id::hot_static_ref!(PROCESS_NAME).lock().clone()
336}
337
338/// Changes the process runtime name.
339///
340/// This sets [`process_name`] and traces an info message "pid: {pid}, name: {name}". If the same PID is named multiple times
341/// the last name should be used when presenting the process in trace viewers.
342///
343/// The process name ideally should be set only by the [`on_process_start!`] "process takeover" handlers. You can use [`init_process_name`]
344/// to only set the name if it has not been set yet.
345pub fn set_process_name(name: impl Into<Txt>) {
346    set_process_name_impl(name.into(), true);
347}
348
349/// Set the process runtime name if it has not been named yet.
350///
351/// See [`set_process_name`] for more details.
352///
353/// Returns `true` if the name was set.
354pub fn init_process_name(name: impl Into<Txt>) -> bool {
355    set_process_name_impl(name.into(), false)
356}
357
358fn set_process_name_impl(new_name: Txt, replace: bool) -> bool {
359    let mut name = zng_unique_id::hot_static_ref!(PROCESS_NAME).lock();
360    if replace || name.is_empty() {
361        *name = new_name;
362        drop(name);
363        tracing::info!("pid: {}, name: {}", std::process::id(), process_name());
364        true
365    } else {
366        false
367    }
368}
369
370/// Panics with an standard message if `zng::env::init!()` was not called or was not called correctly.
371pub fn assert_inited() {
372    match process_lifetime_state() {
373        ProcessLifetimeState::BeforeInit => panic!("env not inited, please call `zng::env::init!()` in main"),
374        ProcessLifetimeState::Inited => {}
375        ProcessLifetimeState::Exiting => {
376            panic!("env not inited correctly, please call `zng::env::init!()` at the beginning of the actual main function")
377        }
378    }
379}