zng_env/
process.rs

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

zng_env/
process.rs
