From ce8ec15eaae000f7cb00a1a08b2081727b1aa02a Mon Sep 17 00:00:00 2001
From: Luca Bruno <lucab@redhat.com>
Date: Mon, 15 Nov 2021 18:49:02 +0000
Subject: [PATCH] system: run and return exit code on stop (#411)

Co-authored-by: Rob Ede <robjtede@icloud.com>
---
 actix-rt/CHANGES.md     |  2 ++
 actix-rt/src/system.rs  | 42 ++++++++++++++++++++++++-----------------
 actix-rt/tests/tests.rs | 13 +++++++++++--
 3 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/actix-rt/CHANGES.md b/actix-rt/CHANGES.md
index efb78d53..bee05bdf 100644
--- a/actix-rt/CHANGES.md
+++ b/actix-rt/CHANGES.md
@@ -1,7 +1,9 @@
 # Changes
 
 ## Unreleased - 2021-xx-xx
+* Add `System::run_with_code` to allow retrieving the exit code on stop. [#411]
 
+[#411]: https://github.com/actix/actix-net/pull/411
 
 ## 2.4.0 - 2021-11-05
 * Add `Arbiter::try_current` for situations where thread may or may not have Arbiter context. [#408]
diff --git a/actix-rt/src/system.rs b/actix-rt/src/system.rs
index e32d6209..395c4f96 100644
--- a/actix-rt/src/system.rs
+++ b/actix-rt/src/system.rs
@@ -175,8 +175,8 @@ impl System {
     }
 }
 
-#[cfg(not(feature = "io-uring"))]
 /// Runner that keeps a [System]'s event loop alive until stop message is received.
+#[cfg(not(feature = "io-uring"))]
 #[must_use = "A SystemRunner does nothing unless `run` is called."]
 #[derive(Debug)]
 pub struct SystemRunner {
@@ -190,23 +190,24 @@ pub struct SystemRunner {
 impl SystemRunner {
     /// Starts event loop and will return once [System] is [stopped](System::stop).
     pub fn run(self) -> io::Result<()> {
+        let exit_code = self.run_with_code()?;
+
+        match exit_code {
+            0 => Ok(()),
+            nonzero => Err(io::Error::new(
+                io::ErrorKind::Other,
+                format!("Non-zero exit code: {}", nonzero),
+            )),
+        }
+    }
+
+    /// Runs the event loop until [stopped](System::stop_with_code), returning the exit code.
+    pub fn run_with_code(self) -> io::Result<i32> {
         let SystemRunner { rt, stop_rx, .. } = self;
 
         // run loop
-        match rt.block_on(stop_rx) {
-            Ok(code) => {
-                if code != 0 {
-                    Err(io::Error::new(
-                        io::ErrorKind::Other,
-                        format!("Non-zero exit code: {}", code),
-                    ))
-                } else {
-                    Ok(())
-                }
-            }
-
-            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
-        }
+        rt.block_on(stop_rx)
+            .map_err(|err| io::Error::new(io::ErrorKind::Other, err))
     }
 
     /// Runs the provided future, blocking the current thread until the future completes.
@@ -216,8 +217,8 @@ impl SystemRunner {
     }
 }
 
-#[cfg(feature = "io-uring")]
 /// Runner that keeps a [System]'s event loop alive until stop message is received.
+#[cfg(feature = "io-uring")]
 #[must_use = "A SystemRunner does nothing unless `run` is called."]
 #[derive(Debug)]
 pub struct SystemRunner;
@@ -226,7 +227,14 @@ pub struct SystemRunner;
 impl SystemRunner {
     /// Starts event loop and will return once [System] is [stopped](System::stop).
     pub fn run(self) -> io::Result<()> {
-        unimplemented!("SystemRunner::run is not implemented yet")
+        unimplemented!("SystemRunner::run is not implemented for io-uring feature yet");
+    }
+
+    /// Runs the event loop until [stopped](System::stop_with_code), returning the exit code.
+    pub fn run_with_code(self) -> io::Result<i32> {
+        unimplemented!(
+            "SystemRunner::run_with_code is not implemented for io-uring feature yet"
+        );
     }
 
     /// Runs the provided future, blocking the current thread until the future completes.
diff --git a/actix-rt/tests/tests.rs b/actix-rt/tests/tests.rs
index 83950221..551a395d 100644
--- a/actix-rt/tests/tests.rs
+++ b/actix-rt/tests/tests.rs
@@ -24,6 +24,15 @@ fn await_for_timer() {
     );
 }
 
+#[cfg(not(feature = "io-uring"))]
+#[test]
+fn run_with_code() {
+    let sys = System::new();
+    System::current().stop_with_code(42);
+    let exit_code = sys.run_with_code().expect("system stop should not error");
+    assert_eq!(exit_code, 42);
+}
+
 #[test]
 fn join_another_arbiter() {
     let time = Duration::from_secs(1);
@@ -99,8 +108,8 @@ fn wait_for_spawns() {
 
     let handle = rt.spawn(async {
         println!("running on the runtime");
-        // assertion panic is caught at task boundary
-        assert_eq!(1, 2);
+        // panic is caught at task boundary
+        panic!("intentional test panic");
     });
 
     assert!(rt.block_on(handle).is_err());