zng_env/
process.rs

1use std::{
2    mem,
3    sync::atomic::{AtomicU8, 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 supposed 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_init` 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 = "0.2"
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        const _: () = {
96            #[$crate::__linkme::distributed_slice($crate::ZNG_ENV_ON_PROCESS_START)]
97            #[linkme(crate = $crate::__linkme)]
98            #[doc(hidden)]
99            static _ON_PROCESS_START: fn(&$crate::ProcessStartArgs) = _on_process_start;
100            #[doc(hidden)]
101            fn _on_process_start(args: &$crate::ProcessStartArgs) {
102                fn on_process_start(args: &$crate::ProcessStartArgs, handler: impl FnOnce(&$crate::ProcessStartArgs)) {
103                    handler(args)
104                }
105                on_process_start(args, $closure)
106            }
107        };
108    };
109}
110
111#[cfg(target_arch = "wasm32")]
112#[doc(hidden)]
113#[macro_export]
114macro_rules! __on_process_start {
115    ($closure:expr) => {
116        $crate::wasm_process_start! {$crate,$closure}
117    };
118}
119
120#[doc(hidden)]
121#[cfg(target_arch = "wasm32")]
122pub use wasm_bindgen::prelude::wasm_bindgen;
123
124#[doc(hidden)]
125#[cfg(target_arch = "wasm32")]
126pub use zng_env_proc_macros::wasm_process_start;
127use zng_txt::Txt;
128
129#[cfg(target_arch = "wasm32")]
130std::thread_local! {
131    #[doc(hidden)]
132    pub static WASM_INIT: std::cell::RefCell<Vec<fn(&ProcessStartArgs)>> = const { std::cell::RefCell::new(vec![]) };
133}
134
135#[cfg(not(target_arch = "wasm32"))]
136#[doc(hidden)]
137#[linkme::distributed_slice]
138pub static ZNG_ENV_ON_PROCESS_START: [fn(&ProcessStartArgs)];
139
140#[cfg(not(target_arch = "wasm32"))]
141pub(crate) fn process_init() -> impl Drop {
142    process_init_impl(&ZNG_ENV_ON_PROCESS_START)
143}
144
145fn process_init_impl(handlers: &[fn(&ProcessStartArgs)]) -> MainExitHandler {
146    let process_state = std::mem::replace(
147        &mut *zng_unique_id::hot_static_ref!(PROCESS_LIFETIME_STATE).lock(),
148        ProcessLifetimeState::Inited,
149    );
150    assert_eq!(process_state, ProcessLifetimeState::BeforeInit, "init!() already called");
151
152    let mut yielded = vec![];
153    // TODO(breaking) remove this
154    let mut yield_until_app = vec![];
155    let mut next_handlers_count = handlers.len();
156    for h in handlers {
157        next_handlers_count -= 1;
158        let args = ProcessStartArgs {
159            next_handlers_count,
160            yield_count: 0,
161            yield_requested: AtomicU8::new(0),
162        };
163        h(&args);
164        match args.yield_requested.load(Ordering::Relaxed) {
165            ProcessStartArgs::YIELD_ONCE => {
166                yielded.push(h);
167                next_handlers_count += 1;
168            }
169            ProcessStartArgs::YIELD_UNTIL_APP => {
170                yield_until_app.push(h);
171            }
172            _ => {}
173        }
174    }
175
176    let mut yield_count = 0;
177    while !yielded.is_empty() {
178        yield_count += 1;
179        if yield_count > ProcessStartArgs::MAX_YIELD_COUNT {
180            eprintln!("start handlers requested `yield_start` more them 32 times");
181            break;
182        }
183
184        next_handlers_count = yielded.len();
185        for h in mem::take(&mut yielded) {
186            next_handlers_count -= 1;
187            let args = ProcessStartArgs {
188                next_handlers_count,
189                yield_count,
190                yield_requested: AtomicU8::new(0),
191            };
192            h(&args);
193            match args.yield_requested.load(Ordering::Relaxed) {
194                ProcessStartArgs::YIELD_ONCE => {
195                    yielded.push(h);
196                    next_handlers_count += 1;
197                }
198                ProcessStartArgs::YIELD_UNTIL_APP => {
199                    yield_until_app.push(h);
200                }
201                _ => {}
202            }
203        }
204    }
205
206    for h in yield_until_app {
207        let args = ProcessStartArgs {
208            next_handlers_count: 0,
209            yield_count: ProcessStartArgs::MAX_YIELD_COUNT,
210            yield_requested: AtomicU8::new(0),
211        };
212        h(&args);
213        if args.yield_requested.load(Ordering::Relaxed) != 0 {
214            eprintln!("handler requested `yield_until_app` and then yielded again")
215        }
216    }
217    MainExitHandler
218}
219
220#[cfg(target_arch = "wasm32")]
221pub(crate) fn process_init() -> impl Drop {
222    std::panic::set_hook(Box::new(console_error_panic_hook::hook));
223
224    let window = web_sys::window().expect("cannot 'init!', no window object");
225    let module = js_sys::Reflect::get(&window, &"__zng_env_init_module".into())
226        .expect("cannot 'init!', missing module in 'window.__zng_env_init_module'");
227
228    if module == wasm_bindgen::JsValue::undefined() || module == wasm_bindgen::JsValue::null() {
229        panic!("cannot 'init!', missing module in 'window.__zng_env_init_module'");
230    }
231
232    let module: js_sys::Object = module.into();
233
234    for entry in js_sys::Object::entries(&module) {
235        let entry: js_sys::Array = entry.into();
236        let ident = entry.get(0).as_string().expect("expected ident at entry[0]");
237
238        if ident.starts_with("__zng_env_start_") {
239            let func: js_sys::Function = entry.get(1).into();
240            if let Err(e) = func.call0(&wasm_bindgen::JsValue::NULL) {
241                panic!("'init!' function error, {e:?}");
242            }
243        }
244    }
245
246    process_init_impl(&WASM_INIT.with_borrow_mut(std::mem::take))
247}
248
249/// Arguments for [`on_process_start`] handlers.
250///
251/// Empty in this release.
252pub struct ProcessStartArgs {
253    /// Number of start handlers yet to run.
254    pub next_handlers_count: usize,
255
256    /// Number of times this handler has yielded.
257    ///
258    /// If this exceeds 32 times the handler is ignored.
259    pub yield_count: u16,
260
261    yield_requested: AtomicU8,
262}
263impl ProcessStartArgs {
264    /// Yield requests after this are ignored.
265    pub const MAX_YIELD_COUNT: u16 = 32;
266
267    const YIELD_ONCE: u8 = 1;
268    const YIELD_UNTIL_APP: u8 = 2;
269
270    /// Let other process start handlers run first.
271    ///
272    /// The handler must call this if it takes over the process and it cannot determinate if it should from the environment.
273    ///
274    /// ```
275    /// # macro_rules! on_process_start { ($($tt:tt)*) => { } }
276    /// fn run_foo_process() {}
277    /// on_process_start!(|args| {
278    ///     if args.yield_count == 0 {
279    ///         return args.yield_once();
280    ///     }
281    ///
282    ///     // yielded once, handlers that affect all processes (loggers, tracers) are inited now
283    ///     if std::env::var("IS_FOO").is_ok() {
284    ///         // take over as "foo" process
285    ///         run_foo_process();
286    ///         zng_env::exit(0);
287    ///     }
288    /// });
289    /// ```
290    pub fn yield_once(&self) {
291        self.yield_requested.store(Self::YIELD_ONCE, Ordering::Relaxed);
292    }
293
294    /// Deprecated.
295    ///
296    /// There is actually no way to know if the process is an app-process until `APP.on_init` runs
297    /// because any process can start an app.
298    #[deprecated = "use `APP.on_init` to register a closure that runs if the process becomes an app process"]
299    pub fn yield_until_app(&self) -> bool {
300        if self.next_handlers_count > 0 && self.yield_count < Self::MAX_YIELD_COUNT {
301            // yield until we are the last handler, this ensures we are running in the app-process
302            self.yield_requested.store(Self::YIELD_UNTIL_APP, Ordering::Relaxed);
303            return true;
304        }
305
306        // init name early, since there is interest
307        crate::init_process_name("app-process");
308
309        false
310    }
311}
312
313struct MainExitHandler;
314impl Drop for MainExitHandler {
315    fn drop(&mut self) {
316        run_exit_handlers(if std::thread::panicking() { 101 } else { 0 })
317    }
318}
319
320type ExitHandler = Box<dyn FnOnce(&ProcessExitArgs) + Send + 'static>;
321
322zng_unique_id::hot_static! {
323    static ON_PROCESS_EXIT: Mutex<Vec<ExitHandler>> = Mutex::new(vec![]);
324}
325
326/// Terminates the current process with the specified exit code.
327///
328/// This function must be used instead of `std::process::exit` as it runs the [`on_process_exit`].
329pub fn exit(code: i32) -> ! {
330    run_exit_handlers(code);
331    std::process::exit(code)
332}
333
334fn run_exit_handlers(code: i32) {
335    *zng_unique_id::hot_static_ref!(PROCESS_LIFETIME_STATE).lock() = ProcessLifetimeState::Exiting;
336
337    let on_exit = mem::take(&mut *zng_unique_id::hot_static_ref!(ON_PROCESS_EXIT).lock());
338    let args = ProcessExitArgs { code };
339    for h in on_exit {
340        h(&args);
341    }
342}
343
344/// Arguments for [`on_process_exit`] handlers.
345#[non_exhaustive]
346pub struct ProcessExitArgs {
347    /// Exit code that will be used.
348    pub code: i32,
349}
350
351/// Register a `handler` to run once when the current process exits.
352///
353/// Note that the handler is only called if the process is terminated by [`exit`], or by the executable main
354/// function returning if [`init!`] is called on it.
355///
356/// [`init!`]: crate::init!
357pub fn on_process_exit(handler: impl FnOnce(&ProcessExitArgs) + Send + 'static) {
358    zng_unique_id::hot_static_ref!(ON_PROCESS_EXIT).lock().push(Box::new(handler))
359}
360
361/// Defines the state of the current process instance.
362///
363/// Use [`process_lifetime_state()`] to get.
364#[derive(Debug, Clone, Copy, PartialEq, Eq)]
365pub enum ProcessLifetimeState {
366    /// Init not called yet.
367    BeforeInit,
368    /// Init called and the function where it is called has not returned yet.
369    Inited,
370    /// Init called and the function where it is called is returning.
371    Exiting,
372}
373
374zng_unique_id::hot_static! {
375    static PROCESS_LIFETIME_STATE: Mutex<ProcessLifetimeState> = Mutex::new(ProcessLifetimeState::BeforeInit);
376}
377zng_unique_id::hot_static! {
378    static PROCESS_NAME: Mutex<Txt> = Mutex::new(Txt::from_static(""));
379}
380
381/// Get the state of the current process instance.
382pub fn process_lifetime_state() -> ProcessLifetimeState {
383    *zng_unique_id::hot_static_ref!(PROCESS_LIFETIME_STATE).lock()
384}
385
386/// Gets a process runtime name.
387///
388/// The primary use of this name is to identify the process in logs, see [`set_process_name`] for details about the logged name.
389/// On set or init the name is logged as an info message "pid: {pid}, name: {name}".
390///
391/// # Common Names
392///
393/// All Zng provided process handlers name the process.
394///
395/// * `"app-process"` - Set by `APP` if no other name was set before the app starts building.
396/// * `"view-process"` - Set by the view-process implementer when running in multi process mode.
397/// * `"crash-handler-process"` - Set by the crash-handler when running with crash handling.
398/// * `"crash-dialog-process"` - Set by the crash-handler on the crash dialog process.
399/// * `"worker-process ({worker_name}, {pid})"` - Set by task worker processes if no name was set before the task runner server starts.
400pub fn process_name() -> Txt {
401    zng_unique_id::hot_static_ref!(PROCESS_NAME).lock().clone()
402}
403
404/// Changes the process runtime name.
405///
406/// This sets [`process_name`] and traces an info message "pid: {pid}, name: {name}". If the same PID is named multiple times
407/// the last name should be used when presenting the process in trace viewers.
408///
409/// The process name ideally should be set only by the [`on_process_start!`] "process takeover" handlers. You can use [`init_process_name`]
410/// to only set the name if it has not been set yet.
411pub fn set_process_name(name: impl Into<Txt>) {
412    set_process_name_impl(name.into(), true);
413}
414
415/// Set the process runtime name if it has not been named yet.
416///
417/// See [`set_process_name`] for more details.
418///
419/// Returns `true` if the name was set.
420pub fn init_process_name(name: impl Into<Txt>) -> bool {
421    set_process_name_impl(name.into(), false)
422}
423
424fn set_process_name_impl(new_name: Txt, replace: bool) -> bool {
425    let mut name = zng_unique_id::hot_static_ref!(PROCESS_NAME).lock();
426    if replace || name.is_empty() {
427        *name = new_name;
428        drop(name);
429        tracing::info!("pid: {}, name: {}", std::process::id(), process_name());
430        true
431    } else {
432        false
433    }
434}
435
436/// Panics with an standard message if `zng::env::init!()` was not called or was not called correctly.
437pub fn assert_inited() {
438    match process_lifetime_state() {
439        ProcessLifetimeState::BeforeInit => panic!("env not inited, please call `zng::env::init!()` in main"),
440        ProcessLifetimeState::Inited => {}
441        ProcessLifetimeState::Exiting => {
442            panic!("env not inited correctly, please call `zng::env::init!()` at the beginning of the actual main function")
443        }
444    }
445}
zng_env/process.rs

zng_env/
process.rs
