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#[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
221pub struct ProcessStartArgs {
225 pub next_handlers_count: usize,
227
228 pub yield_count: u16,
232
233 yield_requested: AtomicBool,
234}
235impl ProcessStartArgs {
236 pub const MAX_YIELD_COUNT: u16 = 32;
238
239 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
260pub 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#[non_exhaustive]
280pub struct ProcessExitArgs {
281 pub code: i32,
283}
284
285pub 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
299pub enum ProcessLifetimeState {
300 BeforeInit,
302 Inited,
304 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
315pub fn process_lifetime_state() -> ProcessLifetimeState {
317 *zng_unique_id::hot_static_ref!(PROCESS_LIFETIME_STATE).lock()
318}
319
320pub fn process_name() -> Txt {
335 zng_unique_id::hot_static_ref!(PROCESS_NAME).lock().clone()
336}
337
338pub fn set_process_name(name: impl Into<Txt>) {
346 set_process_name_impl(name.into(), true);
347}
348
349pub 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
370pub 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}