Add concurrency CI failure checks

Signed-off-by: lucarlig <luca.carlig@ibm.com>

Author: lucarlig

.github/workflows/ci.yml

@@ -55,6 +55,25 @@ jobs:
           tool: cargo-nextest
       - run: cargo nextest run --locked --workspace
 
+  loom:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6.0.2
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2.9.1
+      - run: cargo test --locked -p contextforge-gateway-rs-lib --features loom-tests,internal-test-hooks --test concurrency_loom
+
+  miri:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6.0.2
+      - uses: dtolnay/rust-toolchain@nightly
+        with:
+          components: miri
+      - uses: Swatinem/rust-cache@v2.9.1
+      - run: cargo miri setup
+      - run: cargo miri test --locked -p contextforge-gateway-rs-lib --features miri-tests,internal-test-hooks --test miri_checks
+
   build:
     runs-on: ubuntu-latest
     steps:

Cargo.lock

@@ -55,6 +55,9 @@ typed-builder.workspace = true
 [features]
 default = []
 with_tools = []
+loom-tests = []
+miri-tests = []
+internal-test-hooks = []
 
 
 [dev-dependencies]
@@ -64,6 +67,7 @@ axum-test = "20.0.0"
 test-log = "0.2.20"
 axum-server = { version = "0.8.0", features = ["tls-rustls"] }
 futures.workspace = true
+loom = "0.7.2"
 
 [lints]
 workspace = true

crates/contextforge-gateway-rs-lib/Cargo.toml

@@ -29,7 +29,7 @@ use crate::{
     SessionId,
     gateway::{
         mcp_call_validator::InitializeCallValidator,
-        session_manager::SessionManager,
+        session_manager::{SessionManager, return_transport_entry},
         session_store::{UserSession, UserSessionStore},
     },
 };
@@ -252,7 +252,8 @@ where
 
         let mut transports = self.transports.lock().await;
         for (name, svc) in backend_services {
-            transports.entry(BackendTransportKey::from((&name, session_id))).and_modify(|e| e.service = svc);
+            let key = BackendTransportKey::from((&name, session_id));
+            return_transport_entry(&mut transports, &key, svc, |entry| &mut entry.service);
         }
         drop(transports);
 
@@ -399,7 +400,8 @@ where
 
         let mut transports = self.transports.lock().await;
         for (name, svc) in backend_services {
-            transports.entry(BackendTransportKey::from((&name, session_id))).and_modify(|e| e.service = svc);
+            let key = BackendTransportKey::from((&name, session_id));
+            return_transport_entry(&mut transports, &key, svc, |entry| &mut entry.service);
         }
         drop(transports);
 
@@ -781,6 +783,17 @@ fn split_resource_name<'a, T: AsRef<str>, N: AsRef<str>>(
     None
 }
 
+#[cfg(feature = "internal-test-hooks")]
+pub fn test_split_tool_name_owned(tool_name: &str, backend_names: &[&str]) -> Option<(String, String)> {
+    split_tool_name(&tool_name, backend_names).map(|pair| (pair.backend_name.to_owned(), pair.tool_name.to_owned()))
+}
+
+#[cfg(feature = "internal-test-hooks")]
+pub fn test_split_resource_name_owned(resource_uri: &str, backend_names: &[&str]) -> Option<(String, String)> {
+    split_resource_name(&resource_uri, backend_names)
+        .map(|pair| (pair.backend_name.to_owned(), pair.resource_uri.to_owned()))
+}
+
 #[cfg(test)]
 mod tests {
     // Note this useful idiom: importing names from outer (for mod tests) scope.

crates/contextforge-gateway-rs-lib/src/gateway/mcp_gateway.rs

@@ -4,3 +4,13 @@ mod session_manager;
 mod session_store;
 
 pub use mcp_gateway::{LocalUserSessionStore, McpService};
+
+#[cfg(feature = "internal-test-hooks")]
+#[doc(hidden)]
+pub mod test_support {
+    pub use super::mcp_gateway::{test_split_resource_name_owned, test_split_tool_name_owned};
+    pub use super::session_manager::{borrow_transport_entry, return_transport_entry};
+    pub use super::session_store::test_support::{
+        test_session_mapping_msgpack_round_trip, test_user_session_msgpack_round_trip,
+    };
+}

crates/contextforge-gateway-rs-lib/src/gateway/mod.rs

@@ -1,4 +1,8 @@
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::HashMap,
+    hash::{BuildHasher, Hash},
+    sync::Arc,
+};
 
 use contextforge_gateway_rs_apis::user_store::VirtualHost;
 use tokio::sync::Mutex;
@@ -7,6 +11,32 @@ use tracing::{debug, info};
 use super::mcp_gateway::{BackendTransportKey, BackendTransportService, ServiceHolder};
 use crate::layers::session_id::SessionId;
 
+pub fn borrow_transport_entry<K, V, Service, Hasher>(
+    transports: &mut HashMap<K, V, Hasher>,
+    key: &K,
+    service_slot: impl FnOnce(&mut V) -> &mut Option<Service>,
+) -> Option<Option<Service>>
+where
+    K: Eq + Hash,
+    Hasher: BuildHasher,
+{
+    transports.get_mut(key).map(|entry| service_slot(entry).take())
+}
+
+pub fn return_transport_entry<K, V, Service, Hasher>(
+    transports: &mut HashMap<K, V, Hasher>,
+    key: &K,
+    running_service: Option<Service>,
+    service_slot: impl FnOnce(&mut V) -> &mut Option<Service>,
+) where
+    K: Eq + Hash,
+    Hasher: BuildHasher,
+{
+    if let Some(entry) = transports.get_mut(key) {
+        *service_slot(entry) = running_service;
+    }
+}
+
 pub struct SessionManager<'a> {
     virtual_host: &'a VirtualHost,
     session_id: &'a SessionId,
@@ -32,9 +62,9 @@ impl<'a> SessionManager<'a> {
         names
             .into_iter()
             .filter_map(|name| {
-                transports
-                    .get_mut(&BackendTransportKey::from((&name, self.session_id)))
-                    .map(|b| ServiceHolder::new(name, b.service.take()))
+                let key = BackendTransportKey::from((&name, self.session_id));
+                borrow_transport_entry(&mut transports, &key, |entry| &mut entry.service)
+                    .map(|service| ServiceHolder::new(name, service))
             })
             .collect()
     }
@@ -44,9 +74,8 @@ impl<'a> SessionManager<'a> {
         info!("Returning transports {:?} {backend_transports:?}", self.session_id);
         let mut transports = self.transports.lock().await;
         for svc_holder in backend_transports {
-            transports
-                .entry(BackendTransportKey::from((&svc_holder.name, self.session_id)))
-                .and_modify(|e| e.service = svc_holder.running_service);
+            let key = BackendTransportKey::from((&svc_holder.name, self.session_id));
+            return_transport_entry(&mut transports, &key, svc_holder.running_service, |entry| &mut entry.service);
         }
     }
 

crates/contextforge-gateway-rs-lib/src/gateway/session_manager.rs

@@ -69,6 +69,54 @@ impl UserSession {
     }
 }
 
+#[cfg(feature = "internal-test-hooks")]
+pub(super) mod test_support {
+    use std::sync::Arc;
+
+    use serde::Deserialize;
+
+    use super::{SessionMapping, UserSession};
+
+    #[derive(Debug, Deserialize, PartialEq, Eq)]
+    struct OwnedUserSession {
+        name: String,
+        principal: String,
+        downstream_session_id: Arc<str>,
+    }
+
+    pub fn test_session_mapping_msgpack_round_trip(entries: &[(&str, Option<&str>)]) -> Vec<(String, Option<String>)> {
+        let mut mapping = SessionMapping::new();
+        for (backend_name, upstream_session_id) in entries {
+            let upstream_session_id = upstream_session_id.map(Arc::<str>::from);
+            mapping.push((*backend_name).to_owned(), upstream_session_id.as_ref());
+        }
+
+        let encoded = rmp_serde::encode::to_vec(&mapping).expect("session mapping should encode");
+        let decoded: SessionMapping = rmp_serde::decode::from_slice(&encoded).expect("session mapping should decode");
+
+        entries
+            .iter()
+            .map(|(backend_name, _)| {
+                let upstream_session_id =
+                    decoded.get(backend_name).and_then(super::SessionMap::session).map(|id| id.to_string());
+                ((*backend_name).to_owned(), upstream_session_id)
+            })
+            .collect()
+    }
+
+    pub fn test_user_session_msgpack_round_trip(principal: &str, downstream_session_id: &str) -> bool {
+        let session = UserSession::new(principal.to_owned(), Arc::<str>::from(downstream_session_id));
+        let encoded = rmp_serde::encode::to_vec(&session).expect("user session should encode");
+        let decoded: OwnedUserSession = rmp_serde::decode::from_slice(&encoded).expect("user session should decode");
+        decoded
+            == OwnedUserSession {
+                name: "UserSession".to_owned(),
+                principal: principal.to_owned(),
+                downstream_session_id: Arc::<str>::from(downstream_session_id),
+            }
+    }
+}
+
 #[async_trait]
 pub trait UserSessionStore: Send + Sync {
     async fn get_session<'a>(&self, key: &'a UserSession) -> Result<Option<SessionMapping>, SessionStoreError>;

crates/contextforge-gateway-rs-lib/src/gateway/session_store/mod.rs

@@ -29,6 +29,9 @@ use transports::{DownstreamTls, Tcp};
 use typed_builder::TypedBuilder;
 
 pub use crate::common::{Config, LogRotation};
+#[cfg(feature = "internal-test-hooks")]
+#[doc(hidden)]
+pub use crate::gateway::test_support;
 
 pub type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 pub type Result<T> = std::result::Result<T, Error>;