add db cli command for marking a drones backends as lost (#848)

Author: rolyatmax

plane/.sqlx/query-5b6d66015eb731e4e65c9d28795c707b6d93b58a6b814fc1babb40e61fa8d748.json

@@ -49,8 +49,18 @@ enum Command {
         cleanup_batch_size: Option<i32>,
     },
     MarkBackendLost {
-        #[arg(required = true)]
+        #[arg(required = false)]
         backends: Vec<BackendName>,
+
+        /// If provided, all alive backends on the drone will be marked as lost.
+        /// Must also provide a cluster.
+        #[clap(long)]
+        drone: Option<DroneName>,
+
+        /// If provided, all alive backends on the drone will be marked as lost.
+        /// Must also provide a drone.
+        #[clap(long)]
+        cluster: Option<ClusterName>,
     },
 }
 
@@ -138,10 +148,30 @@ async fn main_inner(opts: Opts) -> anyhow::Result<()> {
                 );
             }
         }
-        Command::MarkBackendLost { backends } => {
+        Command::MarkBackendLost {
+            backends,
+            drone,
+            cluster,
+        } => {
             let stdin = std::io::stdin();
 
-            for backend in backends {
+            let backends_to_mark = match (drone, cluster, backends.is_empty()) {
+                (Some(drone), Some(cluster), true) => db
+                    .backend()
+                    .list_alive_backends_for_drone(&cluster, &drone)
+                    .await?
+                    .into_iter()
+                    .map(|b| b.id)
+                    .collect(),
+                (None, None, false) => backends,
+                _ => {
+                    return Err(anyhow::anyhow!(
+                        "Must either provide a list of backends, or a drone and cluster"
+                    ));
+                }
+            };
+
+            for backend in backends_to_mark {
                 let Some(backend) = db.backend().backend(&backend).await? else {
                     println!("Could not find backend: {}, skipping...", backend);
                     continue;

plane/src/bin/db-cli.rs

@@ -4,7 +4,7 @@ use super::{
 };
 use crate::{
     log_types::BackendAddr,
-    names::{BackendActionName, BackendName},
+    names::{BackendActionName, BackendName, DroneName},
     protocol::{BackendAction, RouteInfo},
     types::{
         backend_state::BackendStatusStreamEntry, BackendState, BackendStatus, BearerToken,
@@ -290,6 +290,56 @@ impl<'a> BackendDatabase<'a> {
         Ok(result)
     }
 
+    pub async fn list_alive_backends_for_drone(
+        &self,
+        cluster: &ClusterName,
+        drone: &DroneName,
+    ) -> sqlx::Result<Vec<BackendRow>> {
+        let query_result = sqlx::query!(
+            r#"
+            select
+                id,
+                cluster,
+                last_status,
+                last_status_time,
+                state,
+                drone_id,
+                expiration_time,
+                allowed_idle_seconds,
+                last_keepalive,
+                now() as "as_of!"
+            from backend
+            where
+                drone_id = (select id from node where name = $1 and cluster = $2) and
+                last_status != 'terminated'
+            "#,
+            drone.to_string(),
+            cluster.to_string(),
+        )
+        .fetch_all(&self.db.pool)
+        .await?;
+
+        let mut result = Vec::new();
+
+        for row in query_result {
+            result.push(BackendRow {
+                id: BackendName::try_from(row.id)
+                    .map_err(|_| sqlx::Error::Decode("Failed to decode backend name.".into()))?,
+                cluster: row.cluster,
+                last_status_time: row.last_status_time,
+                state: serde_json::from_value(row.state)
+                    .map_err(|_| sqlx::Error::Decode("Failed to decode backend state.".into()))?,
+                last_keepalive: row.last_keepalive,
+                drone_id: NodeId::from(row.drone_id),
+                expiration_time: row.expiration_time,
+                allowed_idle_seconds: row.allowed_idle_seconds,
+                as_of: row.as_of,
+            });
+        }
+
+        Ok(result)
+    }
+
     pub async fn route_info_for_static_token(
         &self,
         token: &BearerToken,

plane/src/database/backend.rs

@@ -259,11 +259,7 @@ impl BackendManager {
         }
     }
 
-    pub async fn terminate(
-        self: &Arc<Self>,
-        kind: TerminationKind,
-        reason: TerminationReason,
-    ) -> Result<()> {
+    pub async fn terminate(self: &Arc<Self>, kind: TerminationKind, reason: TerminationReason) {
         let state = self
             .state
             .lock()
@@ -276,8 +272,6 @@ impl BackendManager {
             TerminationKind::Hard => state.to_hard_terminating(reason),
         };
         self.set_state(new_state);
-
-        Ok(())
     }
 
     pub fn mark_terminated(self: &Arc<Self>, exit_code: Option<i32>) -> Result<()> {

plane/src/drone/backend_manager.rs

@@ -232,7 +232,7 @@ impl Executor {
                     manager.clone()
                 };
 
-                manager.terminate(*kind, *reason).await?;
+                manager.terminate(*kind, *reason).await;
             }
         }