chore: merge shuttle-ecs-common to main

runtime/src/alpha.rs

@@ -29,20 +29,65 @@ use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::sync::CancellationToken;
 use tonic::{transport::Server, Request, Response, Status};
 
-use crate::args::args;
 use crate::print_version;
 
-// uses custom macro instead of clap to reduce dependency weight
-args! {
-    pub struct Args {
-        // The port to open the gRPC control layer on.
-        // The address to expose for the service is given in the StartRequest.
-        "--port" => pub port: u16,
+#[derive(Default)]
+struct Args {
+    /// Enable compatibility with beta platform
+    beta: bool,
+    /// Alpha (required): Port to open gRPC server on
+    port: Option<u16>,
+    /// Beta (required): Address to bind the gRPC server to
+    address: Option<SocketAddr>,
+}
+
+impl Args {
+    // uses simple arg parsing logic instead of clap to reduce dependency weight
+    fn parse() -> anyhow::Result<Self> {
+        let mut args = Self::default();
+
+        // The first argument is the path of the executable
+        let mut args_iter = std::env::args().skip(1);
+
+        while let Some(arg) = args_iter.next() {
+            match arg.as_str() {
+                "--beta" => {
+                    args.beta = true;
+                }
+                "--port" => {
+                    let port = args_iter
+                        .next()
+                        .context("missing port value")?
+                        .parse()
+                        .context("invalid port value")?;
+                    args.port = Some(port);
+                }
+                "--address" => {
+                    let address = args_iter
+                        .next()
+                        .context("missing address value")?
+                        .parse()
+                        .context("invalid address value")?;
+                    args.address = Some(address);
+                }
+                _ => {}
+            }
+        }
+
+        if args.beta {
+            if args.address.is_none() {
+                return Err(anyhow::anyhow!("--address is required with --beta"));
+            }
+        } else if args.port.is_none() {
+            return Err(anyhow::anyhow!("--port is required"));
+        }
+
+        Ok(args)
     }
 }
 
 pub async fn start(loader: impl Loader + Send + 'static, runner: impl Runner + Send + 'static) {
-    // `--version` overrides any other arguments.
+    // `--version` overrides any other arguments. Used by cargo-shuttle to check compatibility on local runs.
     if std::env::args().any(|arg| arg == "--version") {
         print_version();
         return;
@@ -88,7 +133,11 @@ pub async fn start(loader: impl Loader + Send + 'static, runner: impl Runner + S
     }
 
     // where to serve the gRPC control layer
-    let addr = SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), args.port);
+    let addr = if args.beta {
+        args.address.unwrap()
+    } else {
+        SocketAddr::new(Ipv4Addr::LOCALHOST.into(), args.port.unwrap())
+    };
 
     let mut server_builder = Server::builder()
         .http2_keepalive_interval(Some(Duration::from_secs(60)))
@@ -99,7 +148,7 @@ pub async fn start(loader: impl Loader + Send + 'static, runner: impl Runner + S
     let cloned_token = token.clone();
 
     let router = {
-        let alpha = Alpha::new(loader, runner, token);
+        let alpha = Alpha::new(args.beta, loader, runner, token);
 
         let svc = RuntimeServer::new(alpha);
         server_builder.add_service(svc)
@@ -125,6 +174,8 @@ pub enum State {
 }
 
 pub struct Alpha<L, R> {
+    /// alter behaviour to interact with the new platform
+    beta: bool,
     // Mutexes are for interior mutability
     stopped_tx: Sender<(StopReason, String)>,
     kill_tx: Mutex<Option<oneshot::Sender<String>>>,
@@ -138,10 +189,11 @@ pub struct Alpha<L, R> {
 }
 
 impl<L, R> Alpha<L, R> {
-    pub fn new(loader: L, runner: R, cancellation_token: CancellationToken) -> Self {
+    pub fn new(beta: bool, loader: L, runner: R, cancellation_token: CancellationToken) -> Self {
         let (stopped_tx, _stopped_rx) = broadcast::channel(10);
 
         Self {
+            beta,
             stopped_tx,
             kill_tx: Mutex::new(None),
             loader: Mutex::new(Some(loader)),
@@ -251,24 +303,26 @@ where
             }
         };
 
-        println!("setting current state to healthy");
         *self.state.lock().unwrap() = State::Loading;
 
         let state = self.state.clone();
         let cancellation_token = self.cancellation_token.clone();
 
-        // Ensure that the runtime is set to unhealthy if it doesn't reach the running state after
-        // it has sent a load response, so that the ECS task will fail.
-        tokio::spawn(async move {
-            // Note: The timeout is quite low as we are not actually provisioning resources after
-            // sending the load response.
-            tokio::time::sleep(Duration::from_secs(180)).await;
-            if !matches!(state.lock().unwrap().deref(), State::Running) {
-                println!("the runtime failed to enter the running state before timing out");
-
-                cancellation_token.cancel();
-            }
-        });
+        // State and cancellation is not used in alpha
+        if self.beta {
+            // Ensure that the runtime is set to unhealthy if it doesn't reach the running state after
+            // it has sent a load response, so that the ECS task will fail.
+            tokio::spawn(async move {
+                // Note: The timeout is quite low as we are not actually provisioning resources after
+                // sending the load response.
+                tokio::time::sleep(Duration::from_secs(180)).await;
+                if !matches!(state.lock().unwrap().deref(), State::Running) {
+                    println!("the runtime failed to enter the running state before timing out");
+
+                    cancellation_token.cancel();
+                }
+            });
+        }
 
         Ok(Response::new(LoadResponse {
             success: true,

runtime/src/lib.rs

@@ -16,7 +16,6 @@ pub use async_trait::async_trait;
 pub use tokio;
 
 mod alpha;
-mod args;
 
 const NAME: &str = env!("CARGO_PKG_NAME");
 const VERSION: &str = env!("CARGO_PKG_VERSION");