From 7bfc42368bcb7688e75cb4c3bf0ccdf18e621d38 Mon Sep 17 00:00:00 2001
From: Erin Power <xampprocky@gmail.com>
Date: Wed, 26 Oct 2022 17:06:52 +0200
Subject: [PATCH] Try to recover when the main process fails

---
 src/cli.rs                        | 36 ++++++++++++++++++++++---------
 src/cli/generate_config_schema.rs |  2 +-
 src/cli/manage.rs                 |  4 ++--
 src/cli/run.rs                    |  2 +-
 4 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/src/cli.rs b/src/cli.rs
index 51f8ab60db..b2828c59d4 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -57,7 +57,7 @@ pub struct Cli {
 }
 
 /// The various Quilkin commands.
-#[derive(clap::Subcommand)]
+#[derive(Clone, clap::Subcommand)]
 pub enum Commands {
     Run(Run),
     GenerateConfigSchema(GenerateConfigSchema),
@@ -137,16 +137,32 @@ impl Cli {
             shutdown_tx.send(()).ok();
         });
 
-        let fut = match self.command {
-            Commands::Run(runner) => tokio::spawn({
-                let shutdown_rx = shutdown_rx.clone();
-                async move { runner.run(config, shutdown_rx).await }
-            }),
-            Commands::Manage(manager) => tokio::spawn(async move { manager.manage(config).await }),
-            Commands::GenerateConfigSchema(generator) => {
-                tokio::spawn(std::future::ready(generator.generate_config_schema()))
+        let fut = tryhard::retry_fn({
+            let shutdown_rx = shutdown_rx.clone();
+            move || match self.command.clone() {
+                Commands::Run(runner) => {
+                    let config = config.clone();
+                    let shutdown_rx = shutdown_rx.clone();
+                    tokio::spawn(
+                        async move { runner.run(config.clone(), shutdown_rx.clone()).await },
+                    )
+                }
+                Commands::Manage(manager) => {
+                    let config = config.clone();
+                    tokio::spawn(async move { manager.manage(config.clone()).await })
+                }
+                Commands::GenerateConfigSchema(generator) => {
+                    tokio::spawn(std::future::ready(generator.generate_config_schema()))
+                }
             }
-        };
+        })
+        .retries(3)
+        .on_retry(|_, _, error| {
+            let error = error.to_string();
+            async move {
+                tracing::warn!(%error, "error would have caused fatal crash");
+            }
+        });
 
         tokio::select! {
             result = fut => result?,
diff --git a/src/cli/generate_config_schema.rs b/src/cli/generate_config_schema.rs
index 4d2a194fe8..a3c45e662d 100644
--- a/src/cli/generate_config_schema.rs
+++ b/src/cli/generate_config_schema.rs
@@ -15,7 +15,7 @@
  */
 
 /// Generates JSON schema files for known filters.
-#[derive(clap::Args)]
+#[derive(clap::Args, Clone)]
 pub struct GenerateConfigSchema {
     /// The directory to write configuration files.
     #[clap(short, long, default_value = ".")]
diff --git a/src/cli/manage.rs b/src/cli/manage.rs
index 9290f1d1f6..80c4215f41 100644
--- a/src/cli/manage.rs
+++ b/src/cli/manage.rs
@@ -16,7 +16,7 @@
 
 /// Runs Quilkin as a xDS management server, using `provider` as
 /// a configuration source.
-#[derive(clap::Args)]
+#[derive(clap::Args, Clone)]
 pub struct Manage {
     #[clap(short, long, env = "QUILKIN_PORT")]
     port: Option<u16>,
@@ -26,7 +26,7 @@ pub struct Manage {
 }
 
 /// The available xDS source providers.
-#[derive(clap::Subcommand)]
+#[derive(Clone, clap::Subcommand)]
 pub enum Providers {
     /// Watches Agones' game server CRDs for `Allocated` game server endpoints,
     /// and for a `ConfigMap` that specifies the filter configuration.
diff --git a/src/cli/run.rs b/src/cli/run.rs
index 09ed32f17e..bd1e22183a 100644
--- a/src/cli/run.rs
+++ b/src/cli/run.rs
@@ -20,7 +20,7 @@ use std::net::SocketAddr;
 use crate::filters::FilterFactory;
 
 /// Run Quilkin as a UDP reverse proxy.
-#[derive(clap::Args)]
+#[derive(clap::Args, Clone)]
 #[non_exhaustive]
 pub struct Run {
     /// One or more `quilkin manage` endpoints to listen to for config changes