feat(api): add batch cancellation for tasks and workers

BobAnkh · BobAnkh · commit fcc1ca1f1269 · 2025-10-25T13:59:57.000Z
diff --git a/netmito/src/api/tasks.rs b/netmito/src/api/tasks.rs
@@ -32,6 +32,7 @@ pub fn tasks_router(st: InfraPool) -> Router<InfraPool> {
         .route("/{uuid}/artifacts/{content_type}", delete(delete_artifact))
         .route("/{uuid}/artifacts", post(upload_artifact))
         .route("/query", post(query_tasks))
+        .route("/cancel", post(cancel_tasks))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
             user_auth_middleware,
@@ -149,6 +150,24 @@ pub async fn query_tasks(
     Ok(Json(tasks))
 }
 
+pub async fn cancel_tasks(
+    Extension(u): Extension<AuthUser>,
+    State(pool): State<InfraPool>,
+    Json(req): Json<TasksCancelByFilterReq>,
+) -> Result<Json<TasksCancelByFilterResp>, ApiError> {
+    let resp = service::task::cancel_tasks_by_filter(u.id, &pool, req)
+        .await
+        .map_err(|e| match e {
+            crate::error::Error::AuthError(err) => ApiError::AuthError(err),
+            crate::error::Error::ApiError(e) => e,
+            _ => {
+                tracing::error!("{}", e);
+                ApiError::InternalServerError
+            }
+        })?;
+    Ok(Json(resp))
+}
+
 pub async fn upload_artifact(
     Extension(u): Extension<AuthUser>,
     State(pool): State<InfraPool>,
diff --git a/netmito/src/api/workers.rs b/netmito/src/api/workers.rs
@@ -54,6 +54,7 @@ pub fn workers_router(st: InfraPool) -> Router<InfraPool> {
         )
         .route("/", post(user_register_worker))
         .route("/query", post(user_query_workers))
+        .route("/shutdown", post(user_shutdown_workers))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
             user_auth_middleware,
@@ -364,3 +365,21 @@ pub async fn user_query_workers(
         })?;
     Ok(Json(resp))
 }
+
+pub async fn user_shutdown_workers(
+    Extension(u): Extension<AuthUser>,
+    State(pool): State<InfraPool>,
+    Json(req): Json<WorkersShutdownByFilterReq>,
+) -> Result<Json<WorkersShutdownByFilterResp>, ApiError> {
+    let resp = service::worker::shutdown_workers_by_filter(u.id, &pool, req)
+        .await
+        .map_err(|e| match e {
+            crate::error::Error::AuthError(err) => ApiError::AuthError(err),
+            crate::error::Error::ApiError(e) => e,
+            _ => {
+                tracing::error!("{}", e);
+                ApiError::InternalServerError
+            }
+        })?;
+    Ok(Json(resp))
+}
diff --git a/netmito/src/client/http.rs b/netmito/src/client/http.rs
@@ -737,6 +737,30 @@ impl MitoHttpClient {
         }
     }
 
+    pub async fn shutdown_workers_by_filter(
+        &mut self,
+        req: WorkersShutdownByFilterReq,
+    ) -> crate::error::Result<WorkersShutdownByFilterResp> {
+        self.url.set_path("workers/shutdown");
+        let resp = self
+            .http_client
+            .post(self.url.as_str())
+            .bearer_auth(&self.credential)
+            .json(&req)
+            .send()
+            .await
+            .map_err(map_reqwest_err)?;
+        if resp.status().is_success() {
+            let resp = resp
+                .json::<WorkersShutdownByFilterResp>()
+                .await
+                .map_err(RequestError::from)?;
+            Ok(resp)
+        } else {
+            Err(get_error_from_resp(resp).await.into())
+        }
+    }
+
     pub async fn replace_worker_tags(
         &mut self,
         uuid: Uuid,
@@ -842,6 +866,30 @@ impl MitoHttpClient {
         }
     }
 
+    pub async fn cancel_tasks_by_filter(
+        &mut self,
+        req: TasksCancelByFilterReq,
+    ) -> crate::error::Result<TasksCancelByFilterResp> {
+        self.url.set_path("tasks/cancel");
+        let resp = self
+            .http_client
+            .post(self.url.as_str())
+            .bearer_auth(&self.credential)
+            .json(&req)
+            .send()
+            .await
+            .map_err(map_reqwest_err)?;
+        if resp.status().is_success() {
+            let resp = resp
+                .json::<TasksCancelByFilterResp>()
+                .await
+                .map_err(RequestError::from)?;
+            Ok(resp)
+        } else {
+            Err(get_error_from_resp(resp).await.into())
+        }
+    }
+
     pub async fn update_task_labels(
         &mut self,
         uuid: Uuid,
diff --git a/netmito/src/client/mod.rs b/netmito/src/client/mod.rs
@@ -736,6 +736,15 @@ impl MitoClient {
             .await
     }
 
+    pub async fn workers_batch_cancel(
+        &mut self,
+        args: CancelWorkersArgs,
+    ) -> crate::error::Result<WorkersShutdownByFilterResp> {
+        self.http_client
+            .shutdown_workers_by_filter(args.into())
+            .await
+    }
+
     pub async fn workers_update_tags(
         &mut self,
         args: WorkerUpdateTagsArgs,
@@ -776,6 +785,13 @@ impl MitoClient {
         self.http_client.cancel_task_by_uuid(uuid).await
     }
 
+    pub async fn tasks_batch_cancel(
+        &mut self,
+        args: CancelTasksArgs,
+    ) -> crate::error::Result<TasksCancelByFilterResp> {
+        self.http_client.cancel_tasks_by_filter(args.into()).await
+    }
+
     pub async fn tasks_update_labels(
         &mut self,
         args: UpdateTaskLabelsArgs,
@@ -1012,6 +1028,18 @@ impl MitoClient {
                         tracing::error!("{}", e);
                     }
                 },
+                WorkersCommands::CancelMany(args) => match self.workers_batch_cancel(args).await {
+                    Ok(resp) => {
+                        tracing::info!(
+                            "Shutdown {} workers in group {}",
+                            resp.shutdown_count,
+                            resp.group_name
+                        );
+                    }
+                    Err(e) => {
+                        tracing::error!("{}", e);
+                    }
+                },
                 WorkersCommands::UpdateTags(args) => match self.workers_update_tags(args).await {
                     Ok(_) => {
                         tracing::info!("Worker tags updated successfully");
@@ -1308,6 +1336,18 @@ impl MitoClient {
                         tracing::error!("{}", e);
                     }
                 },
+                TasksCommands::CancelMany(args) => match self.tasks_batch_cancel(args).await {
+                    Ok(resp) => {
+                        tracing::info!(
+                            "Cancelled {} tasks in group {}",
+                            resp.cancelled_count,
+                            resp.group_name
+                        );
+                    }
+                    Err(e) => {
+                        tracing::error!("{}", e);
+                    }
+                },
                 TasksCommands::UpdateLabels(args) => match self.tasks_update_labels(args).await {
                     Ok(_) => {
                         tracing::info!("Task labels updated successfully");
diff --git a/netmito/src/config/client/tasks.rs b/netmito/src/config/client/tasks.rs
@@ -5,7 +5,10 @@ use uuid::Uuid;
 use crate::{
     config::client::parse_resources,
     entity::state::{TaskExecState, TaskState},
-    schema::{ChangeTaskReq, RemoteResourceDownload, TaskSpec, TasksQueryReq, UpdateTaskLabelsReq},
+    schema::{
+        ChangeTaskReq, RemoteResourceDownload, TaskSpec, TasksCancelByFilterReq, TasksQueryReq,
+        UpdateTaskLabelsReq,
+    },
 };
 
 use super::{parse_key_val, parse_watch_task, ArtifactsArgs};
@@ -26,6 +29,8 @@ pub enum TasksCommands {
     Query(QueryTasksArgs),
     /// Cancel a task
     Cancel(CancelTaskArgs),
+    /// Cancel multiple tasks subject to the filter
+    CancelMany(CancelTasksArgs),
     /// Replace labels of a task
     UpdateLabels(UpdateTaskLabelsArgs),
     /// Update the spec of a task
@@ -117,6 +122,31 @@ pub struct CancelTaskArgs {
     pub uuid: Uuid,
 }
 
+#[derive(Serialize, Debug, Deserialize, Args, Clone)]
+pub struct CancelTasksArgs {
+    /// The username of the creator who submitted the tasks
+    #[arg(short, long, num_args = 0.., value_delimiter = ',')]
+    pub creators: Vec<String>,
+    /// The name of the group the tasks belong to
+    #[arg(short, long)]
+    pub group: Option<String>,
+    /// The tags of the tasks
+    #[arg(short, long, num_args = 0.., value_delimiter = ',')]
+    pub tags: Vec<String>,
+    /// The labels of the tasks
+    #[arg(short, long, num_args = 0.., value_delimiter = ',')]
+    pub labels: Vec<String>,
+    /// The state of the tasks
+    #[arg(short, long, num_args = 0.., value_delimiter = ',')]
+    pub state: Vec<TaskState>,
+    /// The exit status of the tasks, support operators like `=`(default), `!=`, `<`, `<=`, `>`, `>=`
+    #[arg(short, long)]
+    pub exit_status: Option<String>,
+    /// The priority of the tasks, support operators like `=`(default), `!=`, `<`, `<=`, `>`, `>=`
+    #[arg(short, long)]
+    pub priority: Option<String>,
+}
+
 #[derive(Serialize, Debug, Deserialize, Args, Clone)]
 pub struct UpdateTaskLabelsArgs {
     /// The UUID of the task
@@ -222,3 +252,33 @@ impl From<UpdateTaskLabelsArgs> for UpdateTaskLabelsReq {
         }
     }
 }
+
+impl From<CancelTasksArgs> for TasksCancelByFilterReq {
+    fn from(args: CancelTasksArgs) -> Self {
+        Self {
+            creator_usernames: if args.creators.is_empty() {
+                None
+            } else {
+                Some(args.creators.into_iter().collect())
+            },
+            group_name: args.group,
+            tags: if args.tags.is_empty() {
+                None
+            } else {
+                Some(args.tags.into_iter().collect())
+            },
+            labels: if args.labels.is_empty() {
+                None
+            } else {
+                Some(args.labels.into_iter().collect())
+            },
+            states: if args.state.is_empty() {
+                None
+            } else {
+                Some(args.state.into_iter().collect())
+            },
+            exit_status: args.exit_status,
+            priority: args.priority,
+        }
+    }
+}
diff --git a/netmito/src/config/client/workers.rs b/netmito/src/config/client/workers.rs
@@ -6,7 +6,7 @@ use crate::{
     entity::role::GroupWorkerRole,
     schema::{
         RemoveGroupWorkerRoleReq, ReplaceWorkerLabelsReq, ReplaceWorkerTagsReq,
-        UpdateGroupWorkerRoleReq, WorkersQueryReq,
+        UpdateGroupWorkerRoleReq, WorkerShutdownOp, WorkersQueryReq, WorkersShutdownByFilterReq,
     },
 };
 
@@ -22,6 +22,8 @@ pub struct WorkersArgs {
 pub enum WorkersCommands {
     /// Cancel a worker
     Cancel(CancelWorkerArgs),
+    /// Cancel multiple workers subject to the filter
+    CancelMany(CancelWorkersArgs),
     /// Replace tags of a worker
     UpdateTags(WorkerUpdateTagsArgs),
     /// Replace labels of a worker
@@ -46,6 +48,29 @@ pub struct CancelWorkerArgs {
     pub force: bool,
 }
 
+#[derive(Serialize, Debug, Deserialize, Args, Clone)]
+pub struct CancelWorkersArgs {
+    /// The name of the group has access to the workers
+    #[arg(short, long)]
+    pub group: Option<String>,
+    /// The role of the group on the workers
+    #[arg(short, long, num_args = 0.., value_delimiter = ',')]
+    pub role: Vec<GroupWorkerRole>,
+    /// The tags of the workers
+    #[arg(short, long, num_args = 0.., value_delimiter = ',')]
+    pub tags: Vec<String>,
+    /// The labels of the workers
+    #[arg(short, long, num_args = 0.., value_delimiter = ',')]
+    pub labels: Vec<String>,
+    /// The username of the creator
+    #[arg(long)]
+    pub creator: Option<String>,
+    /// Whether to force the workers to shutdown.
+    /// If not specified, the workers will be try to shutdown gracefully
+    #[arg(short, long)]
+    pub force: bool,
+}
+
 #[derive(Serialize, Debug, Deserialize, Args, Clone)]
 pub struct WorkerUpdateTagsArgs {
     /// The UUID of the worker
@@ -169,3 +194,32 @@ impl From<RemoveWorkerGroupArgs> for RemoveGroupWorkerRoleReq {
         }
     }
 }
+
+impl From<CancelWorkersArgs> for WorkersShutdownByFilterReq {
+    fn from(args: CancelWorkersArgs) -> Self {
+        Self {
+            group_name: args.group,
+            role: if args.role.is_empty() {
+                None
+            } else {
+                Some(args.role.into_iter().collect())
+            },
+            tags: if args.tags.is_empty() {
+                None
+            } else {
+                Some(args.tags.into_iter().collect())
+            },
+            labels: if args.labels.is_empty() {
+                None
+            } else {
+                Some(args.labels.into_iter().collect())
+            },
+            creator_username: args.creator,
+            op: if args.force {
+                WorkerShutdownOp::Force
+            } else {
+                WorkerShutdownOp::Graceful
+            },
+        }
+    }
+}
diff --git a/netmito/src/entity/role.rs b/netmito/src/entity/role.rs
@@ -6,7 +6,18 @@ use serde::{Deserialize, Serialize};
 
 /// The role of a user to a group.
 #[derive(
-    EnumIter, DeriveActiveEnum, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, ValueEnum,
+    EnumIter,
+    DeriveActiveEnum,
+    Clone,
+    Debug,
+    PartialEq,
+    Eq,
+    Hash,
+    Serialize,
+    Deserialize,
+    ValueEnum,
+    Ord,
+    PartialOrd,
 )]
 #[sea_orm(rs_type = "i32", db_type = "Integer")]
 pub enum UserGroupRole {
diff --git a/netmito/src/schema.rs b/netmito/src/schema.rs
diff --git a/netmito/src/service/task.rs b/netmito/src/service/task.rs
diff --git a/netmito/src/service/worker/mod.rs b/netmito/src/service/worker/mod.rs
