Merge pull request #93 from k82cn/retry_task

Retry tasks.

common/src/apis.rs

@@ -89,6 +89,10 @@ pub struct Task {
 }
 
 impl Task {
+    pub fn is_completed(&self) -> bool {
+        self.state == TaskState::Succeed || self.state == TaskState::Failed
+    }
+
     pub fn gid(&self) -> TaskGID {
         TaskGID {
             ssn_id: self.ssn_id,
@@ -155,12 +159,6 @@ pub struct SessionContext {
     pub common_data: Option<CommonData>,
 }
 
-impl Task {
-    pub fn is_completed(&self) -> bool {
-        self.state == TaskState::Succeed || self.state == TaskState::Failed
-    }
-}
-
 impl Session {
     pub fn is_closed(&self) -> bool {
         self.status.state == SessionState::Closed

session_manager/src/storage/engine/mod.rs

@@ -41,6 +41,7 @@ pub trait Engine: Send + Sync + 'static {
         task_input: Option<TaskInput>,
     ) -> Result<Task, FlameError>;
     async fn get_task(&self, gid: TaskGID) -> Result<Task, FlameError>;
+    async fn retry_task(&self, gid: TaskGID) -> Result<Task, FlameError>;
     async fn delete_task(&self, gid: TaskGID) -> Result<Task, FlameError>;
     async fn update_task_state(&self, gid: TaskGID, state: TaskState) -> Result<Task, FlameError>;
     async fn find_tasks(&self, ssn_id: SessionID) -> Result<Vec<Task>, FlameError>;

session_manager/src/storage/engine/sqlite.rs

@@ -298,6 +298,30 @@ impl Engine for SqliteEngine {
 
         task.try_into()
     }
+
+    async fn retry_task(&self, gid: TaskGID) -> Result<Task, FlameError> {
+        let mut tx = self
+            .pool
+            .begin()
+            .await
+            .map_err(|e| FlameError::Storage(e.to_string()))?;
+
+        let sql = r#"UPDATE tasks SET state=? WHERE id=? AND ssn_id=? RETURNING *"#;
+        let task: TaskDao = sqlx::query_as(sql)
+            .bind(TaskState::Pending as i32)
+            .bind(gid.task_id)
+            .bind(gid.ssn_id)
+            .fetch_one(&mut *tx)
+            .await
+            .map_err(|e| FlameError::Storage(e.to_string()))?;
+
+        tx.commit()
+            .await
+            .map_err(|e| FlameError::Storage(e.to_string()))?;
+
+        task.try_into()
+    }
+
     async fn update_task_state(&self, gid: TaskGID, state: TaskState) -> Result<Task, FlameError> {
         let mut tx = self
             .pool

session_manager/src/storage/mod.rs

@@ -90,6 +90,11 @@ impl Storage {
             let task_list = self.engine.find_tasks(ssn.id).await?;
             let mut ssn = ssn.clone();
             for task in task_list {
+                let task = match task.state {
+                    TaskState::Running => self.engine.retry_task(task.gid()).await?,
+                    _ => task,
+                };
+
                 ssn.update_task(&task);
             }