diff --git a/crates/shared/src/models/task.rs b/crates/shared/src/models/task.rs index 2ef1597c..934595dd 100644 --- a/crates/shared/src/models/task.rs +++ b/crates/shared/src/models/task.rs @@ -87,6 +87,8 @@ pub struct Task { pub scheduling_config: Option, #[serde(default)] pub storage_config: Option, + #[serde(default)] + pub restart_count: u32, } impl Default for Task { @@ -103,6 +105,7 @@ impl Default for Task { updated_at: None, scheduling_config: None, storage_config: None, + restart_count: 0, } } } @@ -158,6 +161,7 @@ impl TryFrom for Task { updated_at: None, scheduling_config: request.scheduling_config, storage_config: request.storage_config, + restart_count: 0, }) } } diff --git a/crates/worker/src/docker/service.rs b/crates/worker/src/docker/service.rs index 7fd7e8a2..f919cc42 100644 --- a/crates/worker/src/docker/service.rs +++ b/crates/worker/src/docker/service.rs @@ -340,9 +340,12 @@ impl DockerService { pub async fn restart_task(&self) -> Result<(), Box> { let current_task = self.state.get_current_task().await; match current_task { - Some(task) => { + Some(mut task) => { let container_id = format!("{}-{}", TASK_PREFIX, task.id); self.docker_manager.restart_container(&container_id).await?; + task.restart_count = task.restart_count.saturating_add(1); + task.state = TaskState::RESTARTING; + self.state.set_current_task(Some(task.clone())).await; Ok(()) } None => Ok(()), @@ -387,6 +390,10 @@ mod tests { .state .set_current_task(Some(task_clone)) .await; + docker_service.restart_task().await.unwrap(); + let updated_task = docker_service.state.get_current_task().await.unwrap(); + assert_eq!(updated_task.restart_count, 1); + assert_eq!(updated_task.state, TaskState::RESTARTING); assert_eq!( docker_service.state.get_current_task().await.unwrap().name, diff --git a/crates/worker/src/docker/state.rs b/crates/worker/src/docker/state.rs index 7d98ed2b..5cca66bb 100644 --- a/crates/worker/src/docker/state.rs +++ b/crates/worker/src/docker/state.rs @@ -63,4 +63,12 @@ impl DockerState { let mut is_running_guard = self.is_running.lock().await; *is_running_guard = is_running; } + + pub async fn increment_restart_count(&self) { + let mut current_task = self.current_task.lock().await; + if let Some(task) = current_task.as_mut() { + task.restart_count += 1; + } + } } +