From ae758a6ec0f820c6ff5e1bd35f68f7a4e86b880b Mon Sep 17 00:00:00 2001
From: Dominic Burkart <dominic.burkart@datadoghq.com>
Date: Tue, 20 Jun 2023 10:09:13 +0200
Subject: [PATCH] refactor internal of ExitCode to ExitStatus to support the
 windows service api

---
 src/app.rs            | 20 ++++++++++++--------
 src/main.rs           |  7 +++++--
 src/vector_windows.rs |  9 +++++----
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/src/app.rs b/src/app.rs
index eeb7821b88198d..6dca25555f3c75 100644
--- a/src/app.rs
+++ b/src/app.rs
@@ -1,7 +1,6 @@
 #![allow(missing_docs)]
 use std::{
-    collections::HashMap, num::NonZeroUsize, path::PathBuf, process::ExitCode as Exit,
-    time::Duration,
+    collections::HashMap, num::NonZeroUsize, path::PathBuf, process::ExitStatus, time::Duration,
 };
 
 use exitcode::ExitCode;
@@ -35,6 +34,11 @@ use crate::{
     trace,
 };
 
+#[cfg(unix)]
+use std::os::unix::process::ExitStatusExt;
+#[cfg(windows)]
+use std::os::windows::process::ExitStatusExt;
+
 pub static WORKER_THREADS: OnceNonZeroUsize = OnceNonZeroUsize::new();
 
 use crate::internal_events::{VectorQuit, VectorStarted, VectorStopped};
@@ -148,7 +152,7 @@ impl ApplicationConfig {
 }
 
 impl Application {
-    pub fn run() -> Exit {
+    pub fn run() -> ExitStatus {
         let (runtime, app) = Self::prepare_start().unwrap_or_else(|code| std::process::exit(code));
 
         runtime.block_on(app.run())
@@ -245,7 +249,7 @@ pub struct StartedApplication {
 }
 
 impl StartedApplication {
-    pub async fn run(self) -> Exit {
+    pub async fn run(self) -> ExitStatus {
         self.main().await.shutdown().await
     }
 
@@ -320,7 +324,7 @@ pub struct FinishedApplication {
 }
 
 impl FinishedApplication {
-    pub async fn shutdown(self) -> Exit {
+    pub async fn shutdown(self) -> ExitStatus {
         let FinishedApplication {
             signal,
             mut signal_rx,
@@ -338,12 +342,12 @@ impl FinishedApplication {
             SignalTo::Shutdown => {
                 emit!(VectorStopped);
                 tokio::select! {
-                    _ = topology_controller.stop() => Exit::SUCCESS, // Graceful shutdown finished
+                    _ = topology_controller.stop() => ExitStatus::from_raw(exitcode::OK as u32), // Graceful shutdown finished
                     _ = signal_rx.recv() => {
                         // It is highly unlikely that this event will exit from topology.
                         emit!(VectorQuit);
                         // Dropping the shutdown future will immediately shut the server down
-                        Exit::FAILURE
+                        ExitStatus::from_raw(exitcode::UNAVAILABLE as u32)
                     }
                 }
             }
@@ -351,7 +355,7 @@ impl FinishedApplication {
                 // It is highly unlikely that this event will exit from topology.
                 emit!(VectorQuit);
                 drop(topology_controller);
-                Exit::FAILURE
+                ExitStatus::from_raw(exitcode::UNAVAILABLE as u32)
             }
             _ => unreachable!(),
         }
diff --git a/src/main.rs b/src/main.rs
index 7efb8017dac326..66818155ab9eae 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -37,7 +37,8 @@ fn main() -> ExitCode {
         }
     }
 
-    Application::run()
+    let exit_code = Application::run().code().unwrap_or(exitcode::UNAVAILABLE) as u8;
+    ExitCode::from(exit_code)
 }
 
 #[cfg(windows)]
@@ -46,5 +47,7 @@ pub fn main() -> ExitCode {
     // to run vector as a service. If we fail, we consider that we are in
     // interactive mode and then fallback to console mode.  See
     // https://docs.microsoft.com/en-us/dotnet/api/system.environment.userinteractive?redirectedfrom=MSDN&view=netcore-3.1#System_Environment_UserInteractive
-    vector::vector_windows::run().unwrap_or_else(|_| Application::run())
+    let exit_code = vector::vector_windows::run()
+        .unwrap_or_else(|_| Application::run().code().unwrap_or(exitcode::UNAVAILABLE));
+    ExitCode::from(exit_code as u8)
 }
diff --git a/src/vector_windows.rs b/src/vector_windows.rs
index 4410f63600076f..5ba3548ce5a254 100644
--- a/src/vector_windows.rs
+++ b/src/vector_windows.rs
@@ -362,8 +362,9 @@ fn win_main(arguments: Vec<OsString>) {
     if let Err(_e) = run_service(arguments) {}
 }
 
-pub fn run() -> Result<()> {
-    service_dispatcher::start(SERVICE_NAME, ffi_service_main)
+pub fn run() -> Result<i32> {
+    service_dispatcher::start(SERVICE_NAME, ffi_service_main).map(|()| 0_i32)
+    // Always returns 0 exit code as errors are handled by the service dispatcher.
 }
 
 fn run_service(_arguments: Vec<OsString>) -> Result<()> {
@@ -407,10 +408,10 @@ fn run_service(_arguments: Vec<OsString>) -> Result<()> {
                 current_state: ServiceState::Stopped,
                 controls_accepted: ServiceControlAccept::empty(),
                 exit_code: {
-                    if program_completion_status.id() == 0 {
+                    if program_completion_status.success() {
                         ServiceExitCode::Win32(NO_ERROR)
                     } else {
-                        // we could not gracefully shut down in time, likely due to timeout.
+                        // we didn't gracefully shutdown within grace period.
                         ServiceExitCode::Win32(ERROR)
                     }
                 },