From e458c04ec83c7f2801082587523fb41a1d763f22 Mon Sep 17 00:00:00 2001
From: Luca BRUNO <luca.bruno@coreos.com>
Date: Fri, 5 Nov 2021 11:02:29 +0000
Subject: [PATCH] system: run and return exit code on stop

This adds an additional running method to `System`, which allows
consumers to retrieve the stopping code.
Additionally, the existing `run()` is reworked to use that
internally, ensuring both codepaths are tested and aligned in
behavior.
---
 actix-rt/CHANGES.md    |  3 +++
 actix-rt/src/system.rs | 29 +++++++++++++++--------------
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/actix-rt/CHANGES.md b/actix-rt/CHANGES.md
index 4a0c5cfb..70af5a25 100644
--- a/actix-rt/CHANGES.md
+++ b/actix-rt/CHANGES.md
@@ -2,6 +2,9 @@
 
 ## Unreleased - 2021-xx-xx
 
+* Add `System::run_until_stop` 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..d841c95c 100644
--- a/actix-rt/src/system.rs
+++ b/actix-rt/src/system.rs
@@ -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<()> {
+        // run loop
+        let code = self.run_until_stop()?;
+        match 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_until_stop(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(|e| io::Error::new(io::ErrorKind::Other, e))
     }
 
     /// Runs the provided future, blocking the current thread until the future completes.