Add exit code control and hook lifecycle features

5 files changed, 75 insertions, 12 deletions

diff --git a/mingling_core/src/lib.rs b/mingling_core/src/lib.rs
index f4f7d18..83edda8 100644
--- a/mingling_core/src/lib.rs
+++ b/mingling_core/src/lib.rs
@@ -69,3 +69,6 @@ pub mod comp;
 
 #[cfg(feature = "comp")]
 pub use crate::comp::*;
+
+pub use crate::setup::exit_code_control::current_exit_code;
+pub use crate::setup::exit_code_control::update_exit_code;
diff --git a/mingling_core/src/program.rs b/mingling_core/src/program.rs
index 0d23770..b8a409e 100644
--- a/mingling_core/src/program.rs
+++ b/mingling_core/src/program.rs
@@ -207,7 +207,7 @@ where
     }
 
     /// Run the command line program
-    pub async fn exec(self)
+    pub async fn exec(self) -> i32
     where
         C: 'static + Send + Sync,
     {
@@ -217,29 +217,44 @@ where
             Err(e) => match e {
                 ProgramExecuteError::DispatcherNotFound => {
                     eprintln!("Dispatcher not found");
-                    return;
+                    return 1;
                 }
                 ProgramExecuteError::RendererNotFound(renderer_name) => {
                     eprintln!("Renderer `{}` not found", renderer_name);
-                    return;
+                    return 1;
                 }
                 ProgramExecuteError::Other(e) => {
                     eprintln!("{}", e);
-                    return;
+                    return 1;
                 }
             },
         };
 
         // Render result
         if stdout_setting.render_output && !result.is_empty() {
+            let exit_code = result.exit_code;
             print!("{}", result);
+
             if let Err(e) = std::io::Write::flush(&mut std::io::stdout())
                 && stdout_setting.error_output
             {
                 eprintln!("{}", e);
+                1
+            } else {
+                exit_code
             }
+        } else {
+            0
         }
     }
+
+    /// Run the command line program, then exit
+    pub async fn exec_and_exit(self)
+    where
+        C: 'static + Send + Sync,
+    {
+        std::process::exit(self.exec().await)
+    }
 }
 
 // Sync program
diff --git a/mingling_core/src/program/hook.rs b/mingling_core/src/program/hook.rs
index d422e12..edbf996 100644
--- a/mingling_core/src/program/hook.rs
+++ b/mingling_core/src/program/hook.rs
@@ -56,7 +56,18 @@ pub struct ProgramAnonymousHook {
     pub post_render: Option<fn(result: &RenderResult)>,
 
     /// Executes before the program ends
-    pub finish: Option<fn()>,
+    pub finish: Option<fn() -> i32>,
+}
+
+impl<C> Program<C>
+where
+    C: ProgramCollect<Enum = C>,
+{
+    /// Adds an anonymous hook to the program. The hook will be called at the appropriate
+    /// lifecycle events, but receives string representations instead of typed references.
+    pub fn with_hook_anonymous(&mut self, hook: ProgramAnonymousHook) {
+        self.anonymous_hooks.push(hook);
+    }
 }
 
 impl<C> Program<C>
@@ -69,12 +80,6 @@ where
         self.hooks.push(hook);
     }
 
-    /// Adds an anonymous hook to the program. The hook will be called at the appropriate
-    /// lifecycle events, but receives string representations instead of typed references.
-    pub fn with_hook_anonymous(&mut self, hook: ProgramAnonymousHook) {
-        self.anonymous_hooks.push(hook);
-    }
-
     pub(crate) fn run_hook_on_begin(&self) {
         if !self.user_context.run_hook {
             return;
@@ -342,7 +347,7 @@ impl ProgramAnonymousHook {
     }
 
     /// Sets the handler for the `finish` event.
-    pub fn on_finish(mut self, handler: fn()) -> Self {
+    pub fn on_finish(mut self, handler: fn() -> i32) -> Self {
         let _ = self.finish.insert(handler);
         self
     }
diff --git a/mingling_core/src/program/setup.rs b/mingling_core/src/program/setup.rs
index 28aa49b..289dfee 100644
--- a/mingling_core/src/program/setup.rs
+++ b/mingling_core/src/program/setup.rs
@@ -3,6 +3,10 @@ use crate::{ProgramCollect, program::Program};
 mod basic;
 pub use basic::*;
 
+#[doc(hidden)]
+pub mod exit_code_control;
+pub use exit_code_control::ExitCodeSetup;
+
 #[cfg(feature = "general_renderer")]
 mod general_renderer;
 
diff --git a/mingling_core/src/program/setup/exit_code_control.rs b/mingling_core/src/program/setup/exit_code_control.rs
new file mode 100644
index 0000000..0844a24
--- /dev/null
+++ b/mingling_core/src/program/setup/exit_code_control.rs
@@ -0,0 +1,36 @@
+use std::sync::atomic::{AtomicI32, Ordering};
+
+use crate::{ProgramCollect, hook::ProgramAnonymousHook, setup::ProgramSetup};
+
+static EXIT_CODE: AtomicI32 = AtomicI32::new(0);
+
+/// Provides the ability to control the program's exit code, which is returned when the program ends.
+///
+/// - Use `mingling::update_exit_code` to update the exit code.
+/// - Use `mingling::current_exit_code` to query the current exit code.
+pub struct ExitCodeSetup;
+
+impl<C> ProgramSetup<C> for ExitCodeSetup
+where
+    C: ProgramCollect<Enum = C>,
+{
+    fn setup(&mut self, program: &mut crate::Program<C>) {
+        program
+            .with_hook_anonymous(ProgramAnonymousHook::empty().on_finish(|| current_exit_code()));
+    }
+}
+
+/// Updates the program's exit code.
+///
+/// This function sets the value that will be returned when the program exits.
+/// The new code will take effect immediately and be used when the program finishes.
+pub fn update_exit_code(code: i32) {
+    EXIT_CODE.store(code, Ordering::SeqCst);
+}
+
+/// Returns the current exit code.
+///
+/// This function queries the value that will be returned when the program exits.
+pub fn current_exit_code() -> i32 {
+    EXIT_CODE.load(Ordering::SeqCst)
+}