Add backend status stream endpoint

Author: paulgb

plane/src/admin.rs

@@ -1,18 +1,19 @@
-use std::path::PathBuf;
-
 use crate::{
-    client::{PlaneClient, PlaneClientError},
+    client::{sse::SseStream, PlaneClient, PlaneClientError},
     names::{BackendName, DroneName, Name, ProxyName},
     protocol::{CertManagerRequest, CertManagerResponse, MessageFromProxy, MessageToProxy},
     types::{
-        BackendStatus, ClusterName, ClusterState, ConnectRequest, DockerExecutorConfig,
-        DronePoolName, KeyConfig, Mount, NodeState, SpawnConfig, Subdomain,
+        backend_state::BackendStatusStreamEntry, BackendStatus, ClusterName, ClusterState,
+        ConnectRequest, DockerExecutorConfig, DronePoolName, KeyConfig, Mount, NodeState,
+        SpawnConfig, Subdomain,
     },
     PLANE_GIT_HASH, PLANE_VERSION,
 };
 use chrono::Duration;
 use clap::{Parser, Subcommand};
 use colored::Colorize;
+use futures_util::{Stream, StreamExt};
+use std::path::PathBuf;
 use url::Url;
 
 fn show_error(error: &PlaneClientError) {
@@ -156,6 +157,9 @@ pub enum AdminCommand {
     ClusterState {
         cluster: ClusterName,
     },
+    BackendStatus {
+        backend: BackendName,
+    },
 }
 
 pub async fn run_admin_command(opts: AdminOpts) {
@@ -165,6 +169,23 @@ pub async fn run_admin_command(opts: AdminOpts) {
     }
 }
 
+pub async fn print_status_stream(
+    mut stream: SseStream<BackendStatusStreamEntry>,
+    until: BackendStatus,
+) {
+    while let Some(status) = stream.next().await {
+        println!(
+            "Status: {} at {}",
+            status.status.to_string().magenta(),
+            status.time.0.to_string().bright_cyan()
+        );
+
+        if status.status >= until {
+            break;
+        }
+    }
+}
+
 pub async fn run_admin_command_inner(opts: AdminOpts) -> Result<(), PlaneClientError> {
     let client = PlaneClient::new(opts.controller);
 
@@ -227,19 +248,8 @@ pub async fn run_admin_command_inner(opts: AdminOpts) -> Result<(), PlaneClientE
             }
 
             if !immediate {
-                let mut stream = client.backend_status_stream(&response.backend_id).await?;
-
-                while let Some(status) = stream.next().await {
-                    println!(
-                        "Status: {} at {}",
-                        status.status.to_string().magenta(),
-                        status.time.0.to_string().bright_cyan()
-                    );
-
-                    if status.status >= BackendStatus::Ready {
-                        break;
-                    }
-                }
+                let stream = client.backend_status_stream(&response.backend_id).await?;
+                print_status_stream(stream, BackendStatus::Ready).await;
             }
         }
         AdminCommand::Terminate {
@@ -259,19 +269,8 @@ pub async fn run_admin_command_inner(opts: AdminOpts) -> Result<(), PlaneClientE
             );
 
             if !immediate {
-                let mut stream = client.backend_status_stream(&backend).await?;
-
-                while let Some(status) = stream.next().await {
-                    println!(
-                        "Status: {} at {}",
-                        status.status.to_string().magenta(),
-                        status.time.0.to_string().bright_cyan()
-                    );
-
-                    if status.status >= BackendStatus::Terminated {
-                        break;
-                    }
-                }
+                let stream = client.backend_status_stream(&backend).await?;
+                print_status_stream(stream, BackendStatus::Terminated).await;
             }
         }
         AdminCommand::Drain { cluster, drone } => {
@@ -366,6 +365,10 @@ pub async fn run_admin_command_inner(opts: AdminOpts) -> Result<(), PlaneClientE
             let cluster_state = client.cluster_state(&cluster).await?;
             show_cluster_state(&cluster_state);
         }
+        AdminCommand::BackendStatus { backend } => {
+            let stream = client.backend_status_stream(&backend).await?;
+            print_status_stream(stream, BackendStatus::Terminated).await;
+        }
     };
 
     Ok(())

plane/src/client/mod.rs

@@ -13,7 +13,7 @@ use reqwest::{Response, StatusCode};
 use serde::de::DeserializeOwned;
 use url::{form_urlencoded, Url};
 pub mod controller_address;
-mod sse;
+pub mod sse;
 
 #[derive(thiserror::Error, Debug)]
 pub enum PlaneClientError {