From d4ed288aa9944c74d8662f3508d3e3786be1f457 Mon Sep 17 00:00:00 2001
From: Klaus Ma <klausm@nvidia.com>
Date: Thu, 11 Jan 2024 12:44:38 +0000
Subject: [PATCH] Update task.

Signed-off-by: Klaus Ma <klausm@nvidia.com>
---
 common/src/apis.rs                 | 41 ++----------------------------
 session_manager/src/storage/mod.rs |  6 ++---
 2 files changed, 5 insertions(+), 42 deletions(-)

diff --git a/common/src/apis.rs b/common/src/apis.rs
index 8d2c712..e233a60 100644
--- a/common/src/apis.rs
+++ b/common/src/apis.rs
@@ -166,7 +166,7 @@ impl Session {
         self.status.state == SessionState::Closed
     }
 
-    pub fn add_task(&mut self, task: &Task) {
+    pub fn update_task(&mut self, task: &Task) {
         let task_ptr = TaskPtr::new(task.clone().into());
 
         self.tasks.insert(task.id, task_ptr.clone());
@@ -185,43 +185,6 @@ impl Session {
 
         None
     }
-
-    // pub fn update_task_state(
-    //     &mut self,
-    //     task_ptr: TaskPtr,
-    //     state: TaskState,
-    // ) -> Result<(), FlameError> {
-    //     let mut task = lock_ptr!(task_ptr)?;
-    //     match self.tasks_index.get_mut(&task.state) {
-    //         None => {
-    //             log::error!(
-    //                 "Failed to find task <{}> in state map <{}>.",
-    //                 task.id,
-    //                 task.state.to_string()
-    //             );
-
-    //             return Err(FlameError::NotFound(format!(
-    //                 "task <{}> in state map <{}>",
-    //                 task.id, task.state
-    //             )));
-    //         }
-
-    //         Some(index) => {
-    //             index.remove(&task.id);
-    //         }
-    //     }
-
-    //     self.tasks.remove(&task.id);
-
-    //     task.state = state;
-    //     // Also set completion time.
-    //     if state == TaskState::Succeed || state == TaskState::Failed {
-    //         task.completion_time = Some(Utc::now());
-    //     }
-    //     self.add_task(&task);
-
-    //     Ok(())
-    // }
 }
 
 impl Clone for Session {
@@ -241,7 +204,7 @@ impl Clone for Session {
         for (id, t) in &self.tasks {
             match t.lock() {
                 Ok(t) => {
-                    ssn.add_task(&t);
+                    ssn.update_task(&t);
                 }
                 Err(_) => {
                     log::error!("Failed to lock task: <{}>, ignore it during clone.", id);
diff --git a/session_manager/src/storage/mod.rs b/session_manager/src/storage/mod.rs
index 2467e02..8bc10a1 100644
--- a/session_manager/src/storage/mod.rs
+++ b/session_manager/src/storage/mod.rs
@@ -90,7 +90,7 @@ impl Storage {
             let task_list = self.engine.find_tasks(ssn.id).await?;
             let mut ssn = ssn.clone();
             for task in task_list {
-                ssn.add_task(&task);
+                ssn.update_task(&task);
             }
 
             let mut ssn_map = lock_ptr!(self.sessions)?;
@@ -184,7 +184,7 @@ impl Storage {
 
         let ssn = self.get_session_ptr(ssn_id)?;
         let mut ssn = lock_ptr!(ssn)?;
-        ssn.add_task(&task);
+        ssn.update_task(&task);
 
         Ok(task)
     }
@@ -225,7 +225,7 @@ impl Storage {
         let task = self.engine.update_task_state(gid, state).await?;
 
         let mut ssn_ptr = lock_ptr!(ssn)?;
-        ssn_ptr.add_task(&task);
+        ssn_ptr.update_task(&task);
 
         Ok(())
     }