Merge pull request #206 from k82cn/fix_ptr

fix: enhance ptr struct.

Author: k82cn

common/src/apis.rs

@@ -20,7 +20,7 @@ use rustix::system;
 use ::rpc::flame::ApplicationSpec;
 use rpc::flame as rpc;
 
-use crate::ptr::MutexPtr;
+use crate::MutexPtr;
 use crate::FlameError;
 
 pub const DEFAULT_MAX_INSTANCES: u32 = 1_000_000;

common/src/lib.rs

@@ -13,11 +13,11 @@ limitations under the License.
 
 pub mod apis;
 pub mod ctx;
-pub mod ptr;
 pub mod trace;
 
 use serde_json::json;
 use std::collections::HashMap;
+use std::sync::Arc;
 use thiserror::Error;
 use time::macros::format_description;
 use tonic::Status;
@@ -78,21 +78,17 @@ impl From<FromEnvError> for FlameError {
     }
 }
 
-#[macro_export]
-macro_rules! lock_ptr {
-    ( $mutex_arc:expr ) => {
-        $mutex_arc
-            .lock()
-            .map_err(|_| FlameError::Internal("mutex ptr".to_string()))
-    };
+pub type MutexPtr<T> = Arc<std::sync::Mutex<T>>;
+
+pub fn new_ptr<T>(t: T) -> MutexPtr<T> {
+    Arc::new(std::sync::Mutex::new(t))
 }
 
 #[macro_export]
-macro_rules! lock_async_ptr {
+macro_rules! lock_ptr {
     ( $mutex_arc:expr ) => {
         $mutex_arc
             .lock()
-            .await
             .map_err(|_| FlameError::Internal("mutex ptr".to_string()))
     };
 }

common/src/ptr.rs

@@ -22,7 +22,7 @@ use common::apis::{
     Application, ExecutorID, ExecutorState, Node, NodeState, ResourceRequirement, Session,
     SessionID, SessionState, Task, TaskID, TaskState,
 };
-use common::ptr::MutexPtr;
+use common::MutexPtr;
 use common::{lock_ptr, FlameError};
 
 pub type SessionInfoPtr = Arc<SessionInfo>;

session_manager/src/model/mod.rs

@@ -21,7 +21,7 @@ use crate::model::{ExecutorInfoPtr, NodeInfo, NodeInfoPtr, SessionInfo, SessionI
 use crate::scheduler::plugins::fairshare::FairShare;
 use crate::scheduler::Context;
 
-use common::ptr::{self, MutexPtr};
+use common::{self, MutexPtr, new_ptr};
 use common::{lock_ptr, FlameError};
 
 mod fairshare;
@@ -83,7 +83,7 @@ impl PluginManager {
         }
 
         Ok(Arc::new(PluginManager {
-            plugins: ptr::new_ptr(plugins),
+            plugins: new_ptr(plugins),
         }))
     }
 

session_manager/src/scheduler/plugins/mod.rs

@@ -22,7 +22,7 @@ use common::apis::{
     ExecutorState, Node, NodePtr, ResourceRequirement, Session, SessionID, SessionPtr,
     SessionState, Task, TaskGID, TaskID, TaskInput, TaskOutput, TaskPtr, TaskResult, TaskState,
 };
-use common::ptr::{self, MutexPtr};
+use common::{self, MutexPtr};
 use common::{ctx::FlameContext, lock_ptr, trace::TraceFn, trace_fn, FlameError};
 
 use crate::model::{
@@ -49,10 +49,10 @@ pub async fn new_ptr(config: &FlameContext) -> Result<StoragePtr, FlameError> {
     Ok(Arc::new(Storage {
         context: config.clone(),
         engine: engine::connect(&config.storage).await?,
-        sessions: ptr::new_ptr(HashMap::new()),
-        executors: ptr::new_ptr(HashMap::new()),
-        nodes: ptr::new_ptr(HashMap::new()),
-        applications: ptr::new_ptr(HashMap::new()),
+        sessions: common::new_ptr(HashMap::new()),
+        executors: common::new_ptr(HashMap::new()),
+        nodes: common::new_ptr(HashMap::new()),
+        applications: common::new_ptr(HashMap::new()),
     }))
 }
 
@@ -133,7 +133,7 @@ impl Storage {
     pub async fn register_node(&self, node: &Node) -> Result<(), FlameError> {
         trace_fn!("Storage::register_node");
         let mut node_map = lock_ptr!(self.nodes)?;
-        node_map.insert(node.name.clone(), ptr::new_ptr(node.clone()));
+        node_map.insert(node.name.clone(), common::new_ptr(node.clone()));
         Ok(())
     }
 
@@ -145,7 +145,7 @@ impl Storage {
         // trace_fn!("Storage::sync_node");
 
         let mut node_map = lock_ptr!(self.nodes)?;
-        node_map.insert(node.name.clone(), ptr::new_ptr(node.clone()));
+        node_map.insert(node.name.clone(), common::new_ptr(node.clone()));
 
         let mut res = vec![];
 
@@ -310,7 +310,7 @@ impl Storage {
         let app = self.engine.register_application(name, attr).await?;
 
         let mut app_map = lock_ptr!(self.applications)?;
-        app_map.insert(app.name.clone(), ptr::new_ptr(app.clone()));
+        app_map.insert(app.name.clone(), common::new_ptr(app.clone()));
 
         Ok(())
     }
@@ -345,7 +345,7 @@ impl Storage {
         let app = self.engine.update_application(name.clone(), attr).await?;
 
         let mut app_map = lock_ptr!(self.applications)?;
-        app_map.insert(name.clone(), ptr::new_ptr(app.clone()));
+        app_map.insert(name.clone(), common::new_ptr(app.clone()));
 
         Ok(())
     }