feat: using dashmap to store sessions

Author: kallebysantos

Cargo.lock

@@ -17,6 +17,7 @@ ext_ai_v8_utilities.workspace = true
 anyhow.workspace = true
 clap = { workspace = true, features = ["derive"] }
 ctor.workspace = true
+dashmap.workspace = true
 faster-hex.workspace = true
 futures.workspace = true
 futures-util = { workspace = true, features = ["io"] }

ext/ai/Cargo.toml

@@ -1,12 +1,11 @@
+use dashmap::DashMap;
 use deno_core::error::AnyError;
 use futures::io::AllowStdIo;
 use once_cell::sync::Lazy;
 use reqwest::Url;
-use std::collections::HashMap;
 use std::hash::Hasher;
 use std::sync::Arc;
 use std::sync::Mutex;
-use tokio::sync::Mutex as AsyncMutex;
 use tokio_util::compat::FuturesAsyncWriteCompatExt;
 use tracing::debug;
 use tracing::instrument;
@@ -26,8 +25,8 @@ use ort::session::Session;
 
 use crate::onnx::ensure_onnx_env_init;
 
-static SESSIONS: Lazy<AsyncMutex<HashMap<String, Arc<Mutex<Session>>>>> =
-  Lazy::new(|| AsyncMutex::new(HashMap::new()));
+static SESSIONS: Lazy<DashMap<String, Arc<Mutex<Session>>>> =
+  Lazy::new(DashMap::new);
 
 #[derive(Debug)]
 pub struct SessionWithId {
@@ -136,16 +135,14 @@ pub(crate) async fn load_session_from_bytes(
     faster_hex::hex_string(&hasher.finish().to_be_bytes())
   };
 
-  let mut sessions = SESSIONS.lock().await;
-
-  if let Some(session) = sessions.get(&session_id) {
+  if let Some(session) = SESSIONS.get(&session_id) {
     return Ok((session_id, session.clone()).into());
   }
 
   trace!(session_id, "new session");
   let session = create_session(model_bytes)?;
 
-  sessions.insert(session_id.clone(), session.clone());
+  SESSIONS.insert(session_id.clone(), session.clone());
 
   Ok((session_id, session).into())
 }
@@ -156,9 +153,7 @@ pub(crate) async fn load_session_from_url(
 ) -> Result<SessionWithId, Error> {
   let session_id = fxhash::hash(model_url.as_str()).to_string();
 
-  let mut sessions = SESSIONS.lock().await;
-
-  if let Some(session) = sessions.get(&session_id) {
+  if let Some(session) = SESSIONS.get(&session_id) {
     debug!(session_id, "use existing session");
     return Ok((session_id, session.clone()).into());
   }
@@ -174,22 +169,23 @@ pub(crate) async fn load_session_from_url(
   let session = create_session(model_bytes.as_slice())?;
 
   debug!(session_id, "new session");
-  sessions.insert(session_id.clone(), session.clone());
+  SESSIONS.insert(session_id.clone(), session.clone());
 
   Ok((session_id, session).into())
 }
 
 pub(crate) async fn get_session(id: &str) -> Option<Arc<Mutex<Session>>> {
-  SESSIONS.lock().await.get(id).cloned()
+  SESSIONS.get(id).map(|value| value.pair().1.clone())
 }
 
 pub async fn cleanup() -> Result<usize, AnyError> {
   let mut remove_counter = 0;
   {
-    let mut guard = SESSIONS.lock().await;
+    //let mut guard = SESSIONS.lock().await;
     let mut to_be_removed = vec![];
 
-    for (key, session) in &mut *guard {
+    for v in SESSIONS.iter() {
+      let (key, session) = v.pair();
       // Since we're currently referencing the session at this point
       // It also will increments the counter, so we need to check: counter > 1
       if Arc::strong_count(session) > 1 {
@@ -200,7 +196,7 @@ pub async fn cleanup() -> Result<usize, AnyError> {
     }
 
     for key in to_be_removed {
-      let old_store = guard.remove(&key);
+      let old_store = SESSIONS.remove(&key);
       debug_assert!(old_store.is_some());
 
       remove_counter += 1;