feat(rust): rework Sessions

Author: SanjoDeundiak

implementations/rust/ockam/ockam_api/src/lib.rs

@@ -33,6 +33,7 @@ pub mod minicbor_url;
 pub mod nodes;
 pub mod okta;
 pub mod port_range;
+pub mod session;
 pub mod uppercase;
 mod version;
 
@@ -41,7 +42,6 @@ mod influxdb_token_lease;
 
 pub mod logs;
 mod schema;
-mod session;
 
 mod date;
 mod rendezvous_healthcheck;
@@ -54,7 +54,7 @@ pub use error::*;
 pub use influxdb_token_lease::*;
 pub use nodes::service::default_address::*;
 pub use rendezvous_healthcheck::*;
-pub use session::sessions::ConnectionStatus;
+pub use session::connection_status::ConnectionStatus;
 pub use ui::*;
 pub use util::*;
 pub use version::*;

implementations/rust/ockam/ockam_api/src/nodes/connection/mod.rs

@@ -19,7 +19,6 @@ pub(crate) use plain_tcp::PlainTcpInstantiator;
 pub(crate) use project::ProjectInstantiator;
 pub(crate) use secure::SecureChannelInstantiator;
 use std::fmt::{Debug, Formatter};
-use std::sync::Arc;
 
 #[derive(Clone)]
 pub struct Connection {
@@ -43,18 +42,18 @@ pub struct Connection {
 
 impl Connection {
     /// Shorthand to add the address as consumer to the flow control
-    pub fn add_consumer(&self, context: Arc<Context>, address: &Address) {
+    pub fn add_consumer(&self, context: &Context, address: &Address) {
         if let Some(flow_control_id) = &self.flow_control_id {
             context
                 .flow_controls()
                 .add_consumer(address.clone(), flow_control_id);
         }
     }
 
-    pub fn add_default_consumers(&self, ctx: Arc<Context>) {
-        self.add_consumer(ctx.clone(), &DefaultAddress::KEY_EXCHANGER_LISTENER.into());
-        self.add_consumer(ctx.clone(), &DefaultAddress::SECURE_CHANNEL_LISTENER.into());
-        self.add_consumer(ctx.clone(), &DefaultAddress::UPPERCASE_SERVICE.into());
+    pub fn add_default_consumers(&self, ctx: &Context) {
+        self.add_consumer(ctx, &DefaultAddress::KEY_EXCHANGER_LISTENER.into());
+        self.add_consumer(ctx, &DefaultAddress::SECURE_CHANNEL_LISTENER.into());
+        self.add_consumer(ctx, &DefaultAddress::UPPERCASE_SERVICE.into());
         self.add_consumer(ctx, &DefaultAddress::ECHO_SERVICE.into());
     }
 
@@ -182,7 +181,7 @@ pub trait Instantiator: Send + Sync + 'static {
     /// The returned [`Changes`] will be used to update the builder state.
     async fn instantiate(
         &self,
-        ctx: Arc<Context>,
+        ctx: &Context,
         node_manager: &NodeManager,
         transport_route: Route,
         extracted: (MultiAddr, MultiAddr, MultiAddr),
@@ -217,7 +216,7 @@ impl ConnectionBuilder {
     /// user make sure higher protocol abstraction are called before lower level ones
     pub async fn instantiate(
         mut self,
-        ctx: Arc<Context>,
+        ctx: &Context,
         node_manager: &NodeManager,
         instantiator: impl Instantiator,
     ) -> Result<Self, ockam_core::Error> {
@@ -233,14 +232,14 @@ impl ConnectionBuilder {
                     // the transport route should include only the pieces before the match
                     self.transport_route = self
                         .recalculate_transport_route(
-                            &ctx,
+                            ctx,
                             self.current_multiaddr.split(start).0,
                             false,
                         )
                         .await?;
                     let mut changes = instantiator
                         .instantiate(
-                            ctx.clone(),
+                            ctx,
                             node_manager,
                             self.transport_route.clone(),
                             self.extract(start, instantiator.matches().len()),
@@ -271,7 +270,7 @@ impl ConnectionBuilder {
         }
 
         self.transport_route = self
-            .recalculate_transport_route(&ctx, self.current_multiaddr.clone(), true)
+            .recalculate_transport_route(ctx, self.current_multiaddr.clone(), true)
             .await?;
 
         Ok(Self {

implementations/rust/ockam/ockam_api/src/nodes/connection/plain_tcp.rs

@@ -1,7 +1,6 @@
 use crate::error::ApiError;
 use crate::nodes::connection::{Changes, ConnectionBuilder, Instantiator};
 use crate::{multiaddr_to_route, route_to_multiaddr};
-use std::sync::Arc;
 
 use crate::nodes::NodeManager;
 use ockam_core::{async_trait, Error, Route};
@@ -30,7 +29,7 @@ impl Instantiator for PlainTcpInstantiator {
 
     async fn instantiate(
         &self,
-        _ctx: Arc<Context>,
+        _ctx: &Context,
         node_manager: &NodeManager,
         _transport_route: Route,
         extracted: (MultiAddr, MultiAddr, MultiAddr),

implementations/rust/ockam/ockam_api/src/nodes/connection/project.rs

@@ -2,7 +2,6 @@ use crate::error::ApiError;
 use crate::nodes::connection::{Changes, Instantiator};
 use crate::nodes::NodeManager;
 use crate::{multiaddr_to_route, try_address_to_multiaddr};
-use std::sync::Arc;
 
 use ockam_core::{async_trait, Error, Route};
 use ockam_multiaddr::proto::Project;
@@ -36,7 +35,7 @@ impl Instantiator for ProjectInstantiator {
 
     async fn instantiate(
         &self,
-        ctx: Arc<Context>,
+        ctx: &Context,
         node_manager: &NodeManager,
         _transport_route: Route,
         extracted: (MultiAddr, MultiAddr, MultiAddr),
@@ -66,7 +65,7 @@ impl Instantiator for ProjectInstantiator {
         debug!("create a secure channel to the project {project_identifier}");
         let sc = node_manager
             .create_secure_channel_internal(
-                &ctx,
+                ctx,
                 tcp.route,
                 &self.identifier.clone(),
                 Some(vec![project_identifier]),

implementations/rust/ockam/ockam_api/src/nodes/connection/secure.rs

@@ -1,4 +1,3 @@
-use std::sync::Arc;
 use std::time::Duration;
 
 use crate::nodes::connection::{Changes, Instantiator};
@@ -41,7 +40,7 @@ impl Instantiator for SecureChannelInstantiator {
 
     async fn instantiate(
         &self,
-        ctx: Arc<Context>,
+        ctx: &Context,
         node_manager: &NodeManager,
         transport_route: Route,
         extracted: (MultiAddr, MultiAddr, MultiAddr),

implementations/rust/ockam/ockam_api/src/nodes/models/portal.rs

@@ -16,7 +16,7 @@ use crate::colors::color_primary;
 use crate::error::ApiError;
 
 use crate::output::Output;
-use crate::session::sessions::ConnectionStatus;
+use crate::session::connection_status::ConnectionStatus;
 use crate::{route_to_multiaddr, try_address_to_multiaddr};
 
 /// Request body to create an inlet

implementations/rust/ockam/ockam_api/src/nodes/models/relay.rs

@@ -10,6 +10,8 @@ use ockam_multiaddr::MultiAddr;
 use crate::colors::OckamColor;
 use crate::error::ApiError;
 use crate::output::{colorize_connection_status, Output};
+use crate::session::replacer::ReplacerOutputKind;
+use crate::session::session::Session;
 use crate::{route_to_multiaddr, ConnectionStatus};
 
 /// Request body when instructing a node to create a relay
@@ -94,6 +96,20 @@ impl RelayInfo {
         }
     }
 
+    pub fn from_session(session: &Session, destination_address: MultiAddr, alias: String) -> Self {
+        let relay_info = Self::new(destination_address, alias, session.connection_status());
+        if let Some(outcome) = session.last_outcome() {
+            match outcome {
+                ReplacerOutputKind::Relay(info) => relay_info.with(info),
+                ReplacerOutputKind::Inlet(_) => {
+                    panic!("InletInfo should not be in the registry")
+                }
+            }
+        } else {
+            relay_info
+        }
+    }
+
     pub fn with(self, remote_relay_info: RemoteRelayInfo) -> Self {
         Self {
             forwarding_route: Some(remote_relay_info.forwarding_route().to_string()),

implementations/rust/ockam/ockam_api/src/nodes/registry.rs

@@ -1,16 +1,18 @@
 use crate::cli_state::random_name;
-use crate::nodes::models::relay::RelayInfo;
-use crate::session::sessions::{ReplacerOutputKind, Session};
 use crate::DefaultAddress;
+
 use ockam::identity::Identifier;
 use ockam::identity::{SecureChannel, SecureChannelListener};
 use ockam_core::compat::collections::BTreeMap;
 use ockam_core::{Address, Route};
 use ockam_multiaddr::MultiAddr;
-use ockam_node::compat::asynchronous::RwLock;
+use ockam_node::compat::asynchronous::{Mutex, RwLock};
 use ockam_transport_core::HostnamePort;
+
+use crate::session::session::Session;
 use std::borrow::Borrow;
 use std::fmt::Display;
+use std::sync::Arc;
 
 #[derive(Default)]
 pub(crate) struct SecureChannelRegistry {
@@ -127,15 +129,15 @@ impl KafkaServiceInfo {
 pub(crate) struct InletInfo {
     pub(crate) bind_addr: String,
     pub(crate) outlet_addr: MultiAddr,
-    pub(crate) session: Session,
+    pub(crate) session: Arc<Mutex<Session>>,
 }
 
 impl InletInfo {
     pub(crate) fn new(bind_addr: &str, outlet_addr: MultiAddr, session: Session) -> Self {
         Self {
             bind_addr: bind_addr.to_owned(),
             outlet_addr,
-            session,
+            session: Arc::new(Mutex::new(session)),
         }
     }
 }
@@ -160,34 +162,7 @@ impl OutletInfo {
 pub struct RegistryRelayInfo {
     pub(crate) destination_address: MultiAddr,
     pub(crate) alias: String,
-    pub(crate) session: Session,
-}
-
-impl From<RegistryRelayInfo> for RelayInfo {
-    fn from(registry_relay_info: RegistryRelayInfo) -> Self {
-        let relay_info = RelayInfo::new(
-            registry_relay_info.destination_address.clone(),
-            registry_relay_info.alias.clone(),
-            registry_relay_info.session.connection_status(),
-        );
-
-        let current_relay_status =
-            registry_relay_info
-                .session
-                .status()
-                .map(|info| match info.kind {
-                    ReplacerOutputKind::Inlet(_) => {
-                        panic!("InletInfo should not be in the registry")
-                    }
-                    ReplacerOutputKind::Relay(info) => info,
-                });
-
-        if let Some(current_relay_status) = current_relay_status {
-            relay_info.with(current_relay_status)
-        } else {
-            relay_info
-        }
-    }
+    pub(crate) session: Arc<Mutex<Session>>,
 }
 
 #[derive(Default)]

implementations/rust/ockam/ockam_api/src/nodes/service/in_memory_node.rs

@@ -179,7 +179,14 @@ impl InMemoryNode {
     }
 
     pub async fn stop(&self, ctx: &Context) -> Result<()> {
-        self.medic_handle.stop_medic(ctx).await?;
+        for session in self.registry.inlets.values().await {
+            session.session.lock().await.stop().await;
+        }
+
+        for session in self.registry.relays.values().await {
+            session.session.lock().await.stop().await;
+        }
+
         for addr in DefaultAddress::iter() {
             let result = ctx.stop_worker(addr).await;
             // when stopping we can safely ignore missing services
@@ -191,6 +198,7 @@ impl InMemoryNode {
                 }
             }
         }
+
         Ok(())
     }