feat: build queue (#532)

Author: chesedo

Cargo.lock

@@ -24,6 +24,10 @@ pub enum Command {
 
     /// Manage project names
     ProjectNames,
+
+    /// Viewing and managing stats
+    #[command(subcommand)]
+    Stats(StatsCommand),
 }
 
 #[derive(Subcommand, Debug)]
@@ -55,3 +59,13 @@ pub enum AcmeCommand {
         credentials: PathBuf,
     },
 }
+
+#[derive(Subcommand, Debug)]
+pub enum StatsCommand {
+    /// View load stats
+    Load {
+        /// Clear the loads counter
+        #[arg(long)]
+        clear: bool,
+    },
+}

admin/src/args.rs

@@ -1,7 +1,7 @@
 use anyhow::{Context, Result};
 use serde::{de::DeserializeOwned, Serialize};
 use shuttle_common::{
-    models::{project, ToJson},
+    models::{project, stats, ToJson},
     project::ProjectName,
 };
 use tracing::trace;
@@ -43,6 +43,15 @@ impl Client {
         self.get("/admin/projects").await
     }
 
+    pub async fn get_load(&self) -> Result<stats::LoadResponse> {
+        self.get("/admin/stats/load").await
+    }
+
+    pub async fn clear_load(&self) -> Result<stats::LoadResponse> {
+        self.delete("/admin/stats/load", Option::<String>::None)
+            .await
+    }
+
     async fn post<T: Serialize, R: DeserializeOwned>(
         &self,
         path: &str,
@@ -67,6 +76,30 @@ impl Client {
             .context("failed to extract json body from post response")
     }
 
+    async fn delete<T: Serialize, R: DeserializeOwned>(
+        &self,
+        path: &str,
+        body: Option<T>,
+    ) -> Result<R> {
+        trace!(self.api_key, "using api key");
+
+        let mut builder = reqwest::Client::new()
+            .delete(format!("{}{}", self.api_url, path))
+            .bearer_auth(&self.api_key);
+
+        if let Some(body) = body {
+            builder = builder.json(&body);
+        }
+
+        builder
+            .send()
+            .await
+            .context("failed to make delete request")?
+            .to_json()
+            .await
+            .context("failed to extract json body from delete response")
+    }
+
     async fn get<R: DeserializeOwned>(&self, path: &str) -> Result<R> {
         reqwest::Client::new()
             .get(format!("{}{}", self.api_url, path))

admin/src/client.rs

@@ -1,6 +1,6 @@
 use clap::Parser;
 use shuttle_admin::{
-    args::{AcmeCommand, Args, Command},
+    args::{AcmeCommand, Args, Command, StatsCommand},
     client::Client,
     config::get_api_key,
 };
@@ -141,6 +141,20 @@ async fn main() {
 
             res
         }
+        Command::Stats(StatsCommand::Load { clear }) => {
+            let resp = if clear {
+                client.clear_load().await.expect("to delete load stats")
+            } else {
+                client.get_load().await.expect("to get load stats")
+            };
+
+            let has_capacity = if resp.has_capacity { "a" } else { "no" };
+
+            format!(
+                "Currently {} builds are running and there is {} capacity for new builds",
+                resp.builds_count, has_capacity
+            )
+        }
     };
 
     println!("{res}");

admin/src/main.rs

@@ -4,6 +4,7 @@ pub mod project;
 pub mod resource;
 pub mod secret;
 pub mod service;
+pub mod stats;
 pub mod user;
 
 use anyhow::{Context, Result};

common/src/models/mod.rs

@@ -0,0 +1,13 @@
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+#[derive(Deserialize, Serialize)]
+pub struct LoadRequest {
+    pub id: Uuid,
+}
+
+#[derive(Deserialize, Serialize)]
+pub struct LoadResponse {
+    pub builds_count: usize,
+    pub has_capacity: bool,
+}

common/src/models/stats.rs

@@ -2,6 +2,7 @@ use std::{net::SocketAddr, path::PathBuf};
 
 use clap::Parser;
 use fqdn::FQDN;
+use hyper::Uri;
 use shuttle_common::{project::ProjectName, Port};
 
 /// Program to handle the deploys for a single project
@@ -33,6 +34,10 @@ pub struct Args {
     #[clap(long, default_value = "0.0.0.0:8000")]
     pub proxy_address: SocketAddr,
 
+    /// Address to reach gateway's control plane at
+    #[clap(long, default_value = "http://gateway:8001")]
+    pub gateway_uri: Uri,
+
     /// Project being served by this deployer
     #[clap(long)]
     pub project: ProjectName,

deployer/src/args.rs

@@ -362,9 +362,9 @@ mod tests {
 
     use crate::{
         deployment::{
-            deploy_layer::LogType, provisioner_factory, runtime_logger,
-            storage_manager::StorageManager, ActiveDeploymentsGetter, Built, DeploymentManager,
-            Queued,
+            deploy_layer::LogType, gateway_client::BuildQueueClient, provisioner_factory,
+            runtime_logger, storage_manager::StorageManager, ActiveDeploymentsGetter, Built,
+            DeploymentManager, Queued,
         },
         persistence::{SecretRecorder, State},
     };
@@ -529,16 +529,29 @@ mod tests {
         }
     }
 
+    #[derive(Clone)]
+    struct StubBuildQueueClient;
+
+    #[async_trait::async_trait]
+    impl BuildQueueClient for StubBuildQueueClient {
+        async fn get_slot(
+            &self,
+            _id: Uuid,
+        ) -> Result<bool, crate::deployment::gateway_client::Error> {
+            Ok(true)
+        }
+
+        async fn release_slot(
+            &self,
+            _id: Uuid,
+        ) -> Result<(), crate::deployment::gateway_client::Error> {
+            Ok(())
+        }
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn deployment_to_be_queued() {
-        let deployment_manager = DeploymentManager::new(
-            StubAbstractProvisionerFactory,
-            StubRuntimeLoggerFactory,
-            RECORDER.clone(),
-            RECORDER.clone(),
-            StubActiveDeploymentGetter,
-            PathBuf::from("/tmp"),
-        );
+        let deployment_manager = get_deployment_manager();
 
         let queued = get_queue("sleep-async");
         let id = queued.id;
@@ -650,14 +663,7 @@ mod tests {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn deployment_self_stop() {
-        let deployment_manager = DeploymentManager::new(
-            StubAbstractProvisionerFactory,
-            StubRuntimeLoggerFactory,
-            RECORDER.clone(),
-            RECORDER.clone(),
-            StubActiveDeploymentGetter,
-            PathBuf::from("/tmp"),
-        );
+        let deployment_manager = get_deployment_manager();
 
         let queued = get_queue("self-stop");
         let id = queued.id;
@@ -730,14 +736,7 @@ mod tests {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn deployment_bind_panic() {
-        let deployment_manager = DeploymentManager::new(
-            StubAbstractProvisionerFactory,
-            StubRuntimeLoggerFactory,
-            RECORDER.clone(),
-            RECORDER.clone(),
-            StubActiveDeploymentGetter,
-            PathBuf::from("/tmp"),
-        );
+        let deployment_manager = get_deployment_manager();
 
         let queued = get_queue("bind-panic");
         let id = queued.id;
@@ -810,14 +809,7 @@ mod tests {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn deployment_main_panic() {
-        let deployment_manager = DeploymentManager::new(
-            StubAbstractProvisionerFactory,
-            StubRuntimeLoggerFactory,
-            RECORDER.clone(),
-            RECORDER.clone(),
-            StubActiveDeploymentGetter,
-            PathBuf::from("/tmp"),
-        );
+        let deployment_manager = get_deployment_manager();
 
         let queued = get_queue("main-panic");
         let id = queued.id;
@@ -885,14 +877,7 @@ mod tests {
 
     #[tokio::test]
     async fn deployment_from_run() {
-        let deployment_manager = DeploymentManager::new(
-            StubAbstractProvisionerFactory,
-            StubRuntimeLoggerFactory,
-            RECORDER.clone(),
-            RECORDER.clone(),
-            StubActiveDeploymentGetter,
-            PathBuf::from("/tmp"),
-        );
+        let deployment_manager = get_deployment_manager();
 
         let id = Uuid::new_v4();
         deployment_manager
@@ -940,14 +925,7 @@ mod tests {
 
     #[tokio::test]
     async fn scope_with_nil_id() {
-        let deployment_manager = DeploymentManager::new(
-            StubAbstractProvisionerFactory,
-            StubRuntimeLoggerFactory,
-            RECORDER.clone(),
-            RECORDER.clone(),
-            StubActiveDeploymentGetter,
-            PathBuf::from("/tmp"),
-        );
+        let deployment_manager = get_deployment_manager();
 
         let id = Uuid::nil();
         deployment_manager
@@ -973,6 +951,18 @@ mod tests {
         );
     }
 
+    fn get_deployment_manager() -> DeploymentManager {
+        DeploymentManager::builder()
+            .abstract_factory(StubAbstractProvisionerFactory)
+            .runtime_logger_factory(StubRuntimeLoggerFactory)
+            .build_log_recorder(RECORDER.clone())
+            .secret_recorder(RECORDER.clone())
+            .active_deployment_getter(StubActiveDeploymentGetter)
+            .artifacts_path(PathBuf::from("/tmp"))
+            .queue_client(StubBuildQueueClient)
+            .build()
+    }
+
     fn get_queue(name: &str) -> Queued {
         let enc = GzEncoder::new(Vec::new(), Compression::fast());
         let mut tar = tar::Builder::new(enc);