Add restart count increment on task restarts

crates/shared/src/models/task.rs

@@ -87,6 +87,8 @@ pub struct Task {
     pub scheduling_config: Option<SchedulingConfig>,
     #[serde(default)]
     pub storage_config: Option<StorageConfig>,
+    #[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<TaskRequest> for Task {
             updated_at: None,
             scheduling_config: request.scheduling_config,
             storage_config: request.storage_config,
+            restart_count: 0,
         })
     }
 }

crates/worker/src/docker/service.rs

@@ -340,9 +340,12 @@ impl DockerService {
     pub async fn restart_task(&self) -> Result<(), Box<dyn std::error::Error>> {
         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,

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;
+        }
+    }
 }
+