From 05722008a27fd80f71797a712bb77aeca4ade43c Mon Sep 17 00:00:00 2001 From: Oleksandr Deundiak Date: Mon, 30 Dec 2024 15:22:26 +0100 Subject: [PATCH] feat(rust): improve ockam_node --- Cargo.lock | 14 +- examples/rust/get_started/examples/bob.rs | 2 +- examples/rust/get_started/src/echoer.rs | 2 +- examples/rust/get_started/src/hop.rs | 4 +- examples/rust/get_started/src/logger.rs | 10 +- examples/rust/get_started/src/relay.rs | 2 +- .../src/tcp_interceptor/transport/listener.rs | 4 +- .../src/tcp_interceptor/workers/listener.rs | 6 +- .../src/tcp_interceptor/workers/processor.rs | 20 +- .../ockam/ockam/src/relay_service/relay.rs | 2 +- .../ockam/src/relay_service/relay_service.rs | 2 + .../rust/ockam/ockam/src/remote/lifecycle.rs | 6 +- .../rust/ockam/ockam_abac/src/abac/abac.rs | 11 +- .../ockam/ockam_abac/src/abac/outgoing.rs | 2 +- .../ockam/ockam_abac/src/policy/outgoing.rs | 2 +- .../rust/ockam/ockam_api/src/hop.rs | 7 +- .../influxdb/gateway/token_lease_refresher.rs | 2 +- .../src/influxdb/lease_issuer/node_service.rs | 6 +- .../src/kafka/tests/integration_test.rs | 2 +- .../src/kafka/tests/interceptor_test.rs | 14 +- .../ockam_api/src/nodes/connection/mod.rs | 4 +- .../nodes/service/background_node_client.rs | 9 +- .../src/nodes/service/in_memory_node.rs | 2 +- .../src/nodes/service/kafka_services.rs | 6 +- .../src/nodes/service/node_services.rs | 4 +- .../ockam_api/src/nodes/service/relay.rs | 2 +- .../src/nodes/service/secure_channel.rs | 4 +- .../service/tcp_inlets/session_replacer.rs | 26 +- .../src/nodes/service/tcp_outlets.rs | 1 - .../ockam_api/src/nodes/service/transport.rs | 12 +- .../ockam_api/src/nodes/service/worker.rs | 6 +- .../ockam_api/src/nodes/service/workers.rs | 2 +- .../ockam_api/src/rendezvous_healthcheck.rs | 2 +- .../ockam/ockam_api/src/session/session.rs | 7 +- .../ockam/ockam_api/src/test_utils/mod.rs | 10 +- .../rust/ockam/ockam_api/tests/authority.rs | 6 +- .../rust/ockam/ockam_api/tests/common/mod.rs | 2 - .../ockam/ockam_api/tests/common/session.rs | 4 +- .../ockam_api/tests/common/trace_code.rs | 73 -- .../rust/ockam/ockam_api/tests/latency.rs | 2 +- .../rust/ockam/ockam_api/tests/session.rs | 4 +- .../tcp_outlet/invitation_access_control.rs | 2 +- .../rust/ockam/ockam_app_lib/src/state/mod.rs | 4 +- .../rust/ockam/ockam_core/Cargo.toml | 2 +- .../rust/ockam/ockam_core/src/api.rs | 2 +- .../rust/ockam/ockam_core/src/compat.rs | 15 +- .../src/flow_control/flow_control_id.rs | 2 +- .../flow_controls/flow_controls.rs | 10 +- .../src/identity/secure_channel_metadata.rs | 22 +- .../ockam/ockam_core/src/routing/address.rs | 7 + .../ockam_core/src/routing/address_meta.rs | 12 +- .../ockam/ockam_core/src/routing/mailbox.rs | 68 +- .../remote_retriever_creator.rs | 2 +- .../src/secure_channel/decryptor.rs | 6 +- .../src/secure_channel/encryptor_worker.rs | 10 +- .../handshake/handshake_worker.rs | 23 +- .../src/secure_channel/listener.rs | 2 +- .../src/secure_channels/secure_channels.rs | 4 +- .../src/secure_channels/secure_client.rs | 10 +- .../ockam/ockam_identity/tests/channel.rs | 74 +- .../tests/common/message_flow_auth.rs | 2 +- .../ockam_identity/tests/update_route.rs | 14 +- .../rust/ockam/ockam_node/src/async_drop.rs | 48 -- .../ockam/ockam_node/src/context/context.rs | 86 ++- .../src/context/context_lifecycle.rs | 143 ++-- .../rust/ockam/ockam_node/src/context/mod.rs | 1 - .../ockam_node/src/context/receive_message.rs | 2 +- .../ockam_node/src/context/register_router.rs | 2 +- .../ockam_node/src/context/send_message.rs | 21 +- .../ockam/ockam_node/src/context/stop_env.rs | 2 +- .../ockam_node/src/context/transports.rs | 2 +- .../src/context/worker_lifecycle.rs | 64 +- .../rust/ockam/ockam_node/src/debugger.rs | 27 +- .../rust/ockam/ockam_node/src/delayed.rs | 28 +- .../rust/ockam/ockam_node/src/executor.rs | 116 +--- .../rust/ockam/ockam_node/src/lib.rs | 5 - .../rust/ockam/ockam_node/src/metrics.rs | 21 +- .../rust/ockam/ockam_node/src/node.rs | 24 +- .../ockam/ockam_node/src/processor_builder.rs | 190 ++---- .../ockam_node/src/relay/processor_relay.rs | 78 ++- .../ockam_node/src/relay/worker_relay.rs | 99 +-- .../rust/ockam/ockam_node/src/router/mod.rs | 89 +-- .../ockam/ockam_node/src/router/processor.rs | 53 +- .../ockam/ockam_node/src/router/record.rs | 646 ++++++++---------- .../ockam/ockam_node/src/router/shutdown.rs | 84 +-- .../rust/ockam/ockam_node/src/router/state.rs | 47 +- .../ockam/ockam_node/src/router/worker.rs | 60 +- .../rust/ockam/ockam_node/src/runtime.rs | 16 - .../rust/ockam/ockam_node/src/shutdown.rs | 2 + .../ockam/ockam_node/src/worker_builder.rs | 184 ++--- .../ockam/ockam_node/src/workers/echoer.rs | 2 +- .../tests/{router.rs => metadata.rs} | 169 ++--- .../rust/ockam/ockam_node/tests/tests.rs | 66 +- .../ockam_transport_ble/src/router/mod.rs | 10 +- .../src/workers/listener.rs | 2 +- .../src/workers/receiver.rs | 2 +- .../ockam_transport_ble/src/workers/sender.rs | 6 +- .../ockam_transport_core/src/transport.rs | 2 +- .../src/portal/inlet_listener.rs | 18 +- .../src/portal/inlet_shared_state.rs | 12 +- .../src/portal/interceptor.rs | 8 +- .../src/portal/outlet_listener.rs | 6 +- .../src/portal/portal_receiver.rs | 5 +- .../src/portal/portal_worker.rs | 28 +- .../privileged_portal/privileged_portals.rs | 11 +- .../src/privileged_portal/registry/inlet.rs | 8 +- .../workers/internal_processor.rs | 6 +- .../src/transport/connection.rs | 8 +- .../src/transport/lifecycle.rs | 8 +- .../src/transport/listener.rs | 4 +- .../src/transport/portals.rs | 45 +- .../src/workers/listener.rs | 7 +- .../src/workers/receiver.rs | 11 +- .../ockam_transport_tcp/src/workers/sender.rs | 29 +- .../ockam_transport_tcp/tests/lifecycle.rs | 8 +- .../ockam/ockam_transport_tcp/tests/portal.rs | 6 +- .../src/puncture/negotiation/listener.rs | 2 +- .../src/puncture/negotiation/negotiation.rs | 16 +- .../src/puncture/puncture/puncture.rs | 5 +- .../src/puncture/puncture/receiver.rs | 20 +- .../ockam_transport_udp/src/transport/bind.rs | 4 +- .../src/transport/lifecycle.rs | 8 +- .../src/transport/puncture.rs | 4 +- .../src/workers/receiver.rs | 2 +- .../ockam_transport_udp/src/workers/sender.rs | 4 +- .../ockam_transport_uds/src/router/handle.rs | 1 + .../src/router/uds_router.rs | 19 +- .../src/workers/listener.rs | 4 +- .../src/workers/receiver.rs | 2 +- .../ockam_transport_uds/src/workers/sender.rs | 12 +- .../src/router/mod.rs | 5 +- .../src/workers/listener.rs | 2 +- .../src/workers/receiver.rs | 2 +- .../src/workers/sender.rs | 18 +- 134 files changed, 1506 insertions(+), 1852 deletions(-) delete mode 100644 implementations/rust/ockam/ockam_api/tests/common/trace_code.rs delete mode 100644 implementations/rust/ockam/ockam_node/src/async_drop.rs delete mode 100644 implementations/rust/ockam/ockam_node/src/runtime.rs rename implementations/rust/ockam/ockam_node/tests/{router.rs => metadata.rs} (56%) diff --git a/Cargo.lock b/Cargo.lock index 20230e6fb2e..80d8ca87c62 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -983,7 +983,7 @@ checksum = "c51b96c5a8ed8705b40d655273bc4212cbbf38d4e3be2788f36306f154523ec7" dependencies = [ "bytes 1.8.0", "core-error", - "hashbrown 0.15.1", + "hashbrown 0.15.2", "log", "object", "thiserror 1.0.69", @@ -3052,18 +3052,18 @@ checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" dependencies = [ "ahash", "allocator-api2", - "serde", ] [[package]] name = "hashbrown" -version = "0.15.1" +version = "0.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a9bfc1af68b1726ea47d3d5109de126281def866b33970e10fbab11b5dafab3" +checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289" dependencies = [ "allocator-api2", "equivalent", "foldhash", + "serde", ] [[package]] @@ -3569,7 +3569,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "707907fe3c25f5424cce2cb7e1cbcafee6bdbe735ca90ef77c29e84591e5b9da" dependencies = [ "equivalent", - "hashbrown 0.15.1", + "hashbrown 0.15.2", ] [[package]] @@ -4637,7 +4637,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aedf0a2d09c573ed1d8d85b30c119153926a2b36dce0ab28322c09a117a4683e" dependencies = [ "crc32fast", - "hashbrown 0.15.1", + "hashbrown 0.15.2", "indexmap 2.6.0", "memchr", ] @@ -4880,7 +4880,7 @@ dependencies = [ "cfg-if", "core2", "futures-util", - "hashbrown 0.14.5", + "hashbrown 0.15.2", "hex", "miette", "minicbor", diff --git a/examples/rust/get_started/examples/bob.rs b/examples/rust/get_started/examples/bob.rs index d45a0faf64a..e8fdd6b8f1a 100644 --- a/examples/rust/get_started/examples/bob.rs +++ b/examples/rust/get_started/examples/bob.rs @@ -14,7 +14,7 @@ impl Worker for Echoer { type Message = String; async fn handle_message(&mut self, ctx: &mut Context, msg: Routed) -> Result<()> { - println!("\n[✓] Address: {}, Received: {:?}", ctx.address(), msg); + println!("\n[✓] Address: {}, Received: {:?}", ctx.primary_address(), msg); // Echo the message body back on its return_route. ctx.send(msg.return_route().clone(), msg.into_body()?).await diff --git a/examples/rust/get_started/src/echoer.rs b/examples/rust/get_started/src/echoer.rs index 578ef57c2c5..6dac38aa74f 100644 --- a/examples/rust/get_started/src/echoer.rs +++ b/examples/rust/get_started/src/echoer.rs @@ -8,7 +8,7 @@ impl Worker for Echoer { type Message = String; async fn handle_message(&mut self, ctx: &mut Context, msg: Routed) -> Result<()> { - println!("Address: {}, Received: {:?}", ctx.address(), msg); + println!("Address: {}, Received: {:?}", ctx.primary_address(), msg); // Echo the message body back on its return_route. ctx.send(msg.return_route().clone(), msg.into_body()?).await diff --git a/examples/rust/get_started/src/hop.rs b/examples/rust/get_started/src/hop.rs index 7f2bdac44f5..b0218c71870 100644 --- a/examples/rust/get_started/src/hop.rs +++ b/examples/rust/get_started/src/hop.rs @@ -10,10 +10,10 @@ impl Worker for Hop { /// This handle function takes any incoming message and forwards /// it to the next hop in it's onward route async fn handle_message(&mut self, ctx: &mut Context, msg: Routed) -> Result<()> { - println!("Address: {}, Received: {:?}", ctx.address(), msg); + println!("Address: {}, Received: {:?}", ctx.primary_address(), msg); // Send the message to the next worker on its onward_route - ctx.forward(msg.into_local_message().step_forward(&ctx.address())?) + ctx.forward(msg.into_local_message().step_forward(ctx.primary_address())?) .await } } diff --git a/examples/rust/get_started/src/logger.rs b/examples/rust/get_started/src/logger.rs index 96817597bc0..57de1eae3d4 100644 --- a/examples/rust/get_started/src/logger.rs +++ b/examples/rust/get_started/src/logger.rs @@ -15,11 +15,15 @@ impl Worker for Logger { let payload = local_msg.payload_ref(); if let Ok(str) = String::from_utf8(payload.to_vec()) { - println!("Address: {}, Received string: {}", ctx.address(), str); + println!("Address: {}, Received string: {}", ctx.primary_address(), str); } else { - println!("Address: {}, Received binary: {}", ctx.address(), hex::encode(payload)); + println!( + "Address: {}, Received binary: {}", + ctx.primary_address(), + hex::encode(payload) + ); } - ctx.forward(local_msg.step_forward(&ctx.address())?).await + ctx.forward(local_msg.step_forward(ctx.primary_address())?).await } } diff --git a/examples/rust/get_started/src/relay.rs b/examples/rust/get_started/src/relay.rs index 3f3a32cf0f7..f4416126b21 100644 --- a/examples/rust/get_started/src/relay.rs +++ b/examples/rust/get_started/src/relay.rs @@ -24,7 +24,7 @@ impl Worker for Relay { /// This handle function takes any incoming message and forwards /// it to the next hop in it's onward route async fn handle_message(&mut self, ctx: &mut Context, msg: Routed) -> Result<()> { - println!("Address: {}, Received: {:?}", ctx.address(), msg); + println!("Address: {}, Received: {:?}", ctx.primary_address(), msg); let next_on_route = self.route.next()?.clone(); diff --git a/examples/rust/mitm_node/src/tcp_interceptor/transport/listener.rs b/examples/rust/mitm_node/src/tcp_interceptor/transport/listener.rs index 4e0b6eb35fa..2b510f8281a 100644 --- a/examples/rust/mitm_node/src/tcp_interceptor/transport/listener.rs +++ b/examples/rust/mitm_node/src/tcp_interceptor/transport/listener.rs @@ -25,7 +25,7 @@ impl TcpMitmTransport { } /// Interrupt an active TCP listener given its `Address` - pub async fn stop_listener(&self, address: &Address) -> Result<()> { - self.ctx.stop_processor(address.clone()).await + pub fn stop_listener(&self, address: &Address) -> Result<()> { + self.ctx.stop_address(address.clone()) } } diff --git a/examples/rust/mitm_node/src/tcp_interceptor/workers/listener.rs b/examples/rust/mitm_node/src/tcp_interceptor/workers/listener.rs index 245785a65c9..53ff6d6c46d 100644 --- a/examples/rust/mitm_node/src/tcp_interceptor/workers/listener.rs +++ b/examples/rust/mitm_node/src/tcp_interceptor/workers/listener.rs @@ -45,15 +45,15 @@ impl Processor for TcpMitmListenProcessor { type Context = Context; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(CLUSTER_NAME).await?; + ctx.set_cluster(CLUSTER_NAME)?; - self.registry.add_listener(&ctx.address()); + self.registry.add_listener(ctx.primary_address()); Ok(()) } async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { - self.registry.remove_listener(&ctx.address()); + self.registry.remove_listener(ctx.primary_address()); Ok(()) } diff --git a/examples/rust/mitm_node/src/tcp_interceptor/workers/processor.rs b/examples/rust/mitm_node/src/tcp_interceptor/workers/processor.rs index a558b294098..e11a71e88c4 100644 --- a/examples/rust/mitm_node/src/tcp_interceptor/workers/processor.rs +++ b/examples/rust/mitm_node/src/tcp_interceptor/workers/processor.rs @@ -59,20 +59,20 @@ impl Processor for TcpMitmProcessor { type Context = Context; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(CLUSTER_NAME).await?; + ctx.set_cluster(CLUSTER_NAME)?; self.registry - .add_processor(&ctx.address(), self.role, self.write_half.clone()); + .add_processor(ctx.primary_address(), self.role, self.write_half.clone()); - debug!("Initialize {}", ctx.address()); + debug!("Initialize {}", ctx.primary_address()); Ok(()) } async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { - self.registry.remove_processor(&ctx.address()); + self.registry.remove_processor(ctx.primary_address()); - debug!("Shutdown {}", ctx.address()); + debug!("Shutdown {}", ctx.primary_address()); Ok(()) } @@ -83,9 +83,9 @@ impl Processor for TcpMitmProcessor { let len = match self.read_half.read(&mut buf).await { Ok(l) if l != 0 => l, _ => { - info!("Connection was closed; dropping stream {}", ctx.address()); + info!("Connection was closed; dropping stream {}", ctx.primary_address()); - let _ = ctx.stop_processor(self.address_of_other_processor.clone()).await; + let _ = ctx.stop_address(self.address_of_other_processor.clone()); return Ok(false); } @@ -93,12 +93,12 @@ impl Processor for TcpMitmProcessor { match self.write_half.lock().await.write_all(&buf[..len]).await { Ok(_) => { - debug!("Forwarded {} bytes from {}", len, ctx.address()); + debug!("Forwarded {} bytes from {}", len, ctx.primary_address()); } _ => { - debug!("Connection was closed; dropping stream {}", ctx.address()); + debug!("Connection was closed; dropping stream {}", ctx.primary_address()); - let _ = ctx.stop_processor(self.address_of_other_processor.clone()).await; + let _ = ctx.stop_address(self.address_of_other_processor.clone()); return Ok(false); } diff --git a/implementations/rust/ockam/ockam/src/relay_service/relay.rs b/implementations/rust/ockam/ockam/src/relay_service/relay.rs index 5a5bd4a71ae..819490b1668 100644 --- a/implementations/rust/ockam/ockam/src/relay_service/relay.rs +++ b/implementations/rust/ockam/ockam/src/relay_service/relay.rs @@ -65,7 +65,7 @@ impl Worker for Relay { ctx.forward( LocalMessage::new() .with_onward_route(self.forward_route.clone()) - .with_return_route(route![ctx.address()]) + .with_return_route(route![ctx.primary_address()]) .with_payload(payload), ) .await?; diff --git a/implementations/rust/ockam/ockam/src/relay_service/relay_service.rs b/implementations/rust/ockam/ockam/src/relay_service/relay_service.rs index 9a2b1049d45..ed28c717570 100644 --- a/implementations/rust/ockam/ockam/src/relay_service/relay_service.rs +++ b/implementations/rust/ockam/ockam/src/relay_service/relay_service.rs @@ -33,6 +33,7 @@ impl RelayService { options.setup_flow_control_for_relay_service(ctx.flow_controls(), alias); additional_mailboxes.push(Mailbox::new( alias.clone(), + None, options.service_incoming_access_control.clone(), Arc::new(DenyAll), )); @@ -45,6 +46,7 @@ impl RelayService { .with_mailboxes(Mailboxes::new( Mailbox::new( address.clone(), + None, service_incoming_access_control, Arc::new(DenyAll), ), diff --git a/implementations/rust/ockam/ockam/src/remote/lifecycle.rs b/implementations/rust/ockam/ockam/src/remote/lifecycle.rs index ac0768553ac..151c278dfd2 100644 --- a/implementations/rust/ockam/ockam/src/remote/lifecycle.rs +++ b/implementations/rust/ockam/ockam/src/remote/lifecycle.rs @@ -32,12 +32,14 @@ impl RemoteRelay { ) -> Mailboxes { let main_internal = Mailbox::new( addresses.main_internal, + None, Arc::new(DenyAll), outgoing_access_control, ); let main_remote = Mailbox::new( addresses.main_remote, + None, Arc::new(AllowAll), Arc::new(AllowAll), ); @@ -72,7 +74,7 @@ impl RemoteRelay { let addresses = Addresses::generate(RelayType::Static); let mut callback_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( addresses.completion_callback.clone(), Arc::new(AllowSourceAddress(addresses.main_remote.clone())), Arc::new(DenyAll), @@ -117,7 +119,7 @@ impl RemoteRelay { let addresses = Addresses::generate(RelayType::Ephemeral); let mut callback_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( addresses.completion_callback.clone(), Arc::new(AllowSourceAddress(addresses.main_remote.clone())), Arc::new(DenyAll), diff --git a/implementations/rust/ockam/ockam_abac/src/abac/abac.rs b/implementations/rust/ockam/ockam_abac/src/abac/abac.rs index 0a9b7aa6976..4a129c1c854 100644 --- a/implementations/rust/ockam/ockam_abac/src/abac/abac.rs +++ b/implementations/rust/ockam/ockam_abac/src/abac/abac.rs @@ -59,20 +59,19 @@ impl Abac { } impl Abac { - pub async fn get_outgoing_identifier( + pub fn get_outgoing_identifier( ctx: &Context, relay_msg: &RelayMessage, ) -> Result> { - let terminal = if let Some(terminal) = ctx - .find_terminal_address(relay_msg.onward_route().clone()) - .await? + let metadata = if let Some((_address, metadata)) = + ctx.find_terminal_address(relay_msg.onward_route().iter())? { - terminal + metadata } else { return Ok(None); }; - if let Ok(metadata) = SecureChannelMetadata::from_terminal_address(&terminal) { + if let Ok(metadata) = SecureChannelMetadata::from_terminal_address_metadata(&metadata) { Ok(Some(metadata.their_identifier().into())) } else { Ok(None) diff --git a/implementations/rust/ockam/ockam_abac/src/abac/outgoing.rs b/implementations/rust/ockam/ockam_abac/src/abac/outgoing.rs index dabc04632c6..168d5add66d 100644 --- a/implementations/rust/ockam/ockam_abac/src/abac/outgoing.rs +++ b/implementations/rust/ockam/ockam_abac/src/abac/outgoing.rs @@ -82,7 +82,7 @@ impl OutgoingAbac { /// Returns true if the sender of the message is validated by the expression stored in AbacAccessControl pub async fn is_authorized_impl(&self, relay_msg: &RelayMessage) -> Result { - let identifier = match Abac::get_outgoing_identifier(&self.ctx, relay_msg).await? { + let identifier = match Abac::get_outgoing_identifier(&self.ctx, relay_msg)? { Some(identifier) => identifier, None => { debug! { diff --git a/implementations/rust/ockam/ockam_abac/src/policy/outgoing.rs b/implementations/rust/ockam/ockam_abac/src/policy/outgoing.rs index 7c7f97defc4..2311a3a3eae 100644 --- a/implementations/rust/ockam/ockam_abac/src/policy/outgoing.rs +++ b/implementations/rust/ockam/ockam_abac/src/policy/outgoing.rs @@ -45,7 +45,7 @@ impl OutgoingAccessControl for OutgoingPolicyAccessControl { return Ok(false); }; - let identifier = match Abac::get_outgoing_identifier(&self.ctx, relay_msg).await? { + let identifier = match Abac::get_outgoing_identifier(&self.ctx, relay_msg)? { Some(identifier) => identifier, None => { debug! { diff --git a/implementations/rust/ockam/ockam_api/src/hop.rs b/implementations/rust/ockam/ockam_api/src/hop.rs index b07a067427c..856f8a10a65 100644 --- a/implementations/rust/ockam/ockam_api/src/hop.rs +++ b/implementations/rust/ockam/ockam_api/src/hop.rs @@ -12,7 +12,10 @@ impl Worker for Hop { /// it to the next hop in it's onward route async fn handle_message(&mut self, ctx: &mut Context, msg: Routed) -> Result<()> { // Send the message on its onward_route - ctx.forward(msg.into_local_message().step_forward(&ctx.address())?) - .await + ctx.forward( + msg.into_local_message() + .step_forward(ctx.primary_address())?, + ) + .await } } diff --git a/implementations/rust/ockam/ockam_api/src/influxdb/gateway/token_lease_refresher.rs b/implementations/rust/ockam/ockam_api/src/influxdb/gateway/token_lease_refresher.rs index 92c54dbfe8b..ffc6a3f0a17 100644 --- a/implementations/rust/ockam/ockam_api/src/influxdb/gateway/token_lease_refresher.rs +++ b/implementations/rust/ockam/ockam_api/src/influxdb/gateway/token_lease_refresher.rs @@ -25,7 +25,7 @@ impl TokenLeaseRefresher { lease_issuer_route: MultiAddr, ) -> Result { let token = Arc::new(RwLock::new(None)); - let mailboxes = Mailboxes::main( + let mailboxes = Mailboxes::primary( Address::random_tagged("LeaseRetriever"), Arc::new(DenyAll), Arc::new(AllowAll), diff --git a/implementations/rust/ockam/ockam_api/src/influxdb/lease_issuer/node_service.rs b/implementations/rust/ockam/ockam_api/src/influxdb/lease_issuer/node_service.rs index 71ff1b0c9e3..d32e5f35cf9 100644 --- a/implementations/rust/ockam/ockam_api/src/influxdb/lease_issuer/node_service.rs +++ b/implementations/rust/ockam/ockam_api/src/influxdb/lease_issuer/node_service.rs @@ -127,10 +127,8 @@ impl InMemoryNode { match self.registry.influxdb_services.get(&address).await { None => Ok(None), Some(_) => { - context.stop_worker(address.clone()).await?; - context - .stop_processor(format!("{address}-processor")) - .await?; + context.stop_address(address.clone())?; + context.stop_address(format!("{address}-processor"))?; self.registry.influxdb_services.remove(&address).await; Ok(Some(())) } diff --git a/implementations/rust/ockam/ockam_api/src/kafka/tests/integration_test.rs b/implementations/rust/ockam/ockam_api/src/kafka/tests/integration_test.rs index 016fc7f062d..9255b6dad9c 100644 --- a/implementations/rust/ockam/ockam_api/src/kafka/tests/integration_test.rs +++ b/implementations/rust/ockam/ockam_api/src/kafka/tests/integration_test.rs @@ -170,7 +170,7 @@ async fn producer__flow_with_mock_kafka__content_encryption_and_decryption( .await; drop(consumer_mock_kafka); // drop the outlet and re-create it when we need it later - context.stop_worker("kafka_consumer_outlet").await?; + context.stop_address("kafka_consumer_outlet")?; } let mut producer_mock_kafka = TcpServerSimulator::start("127.0.0.1:0").await; diff --git a/implementations/rust/ockam/ockam_api/src/kafka/tests/interceptor_test.rs b/implementations/rust/ockam/ockam_api/src/kafka/tests/interceptor_test.rs index bea6e701729..a6a6d57ebd6 100644 --- a/implementations/rust/ockam/ockam_api/src/kafka/tests/interceptor_test.rs +++ b/implementations/rust/ockam/ockam_api/src/kafka/tests/interceptor_test.rs @@ -82,13 +82,13 @@ async fn kafka_portal_worker__pieces_of_kafka_message__message_assembled( // send 2 distinct pieces and see if the kafka message is re-assembled back context .send( - route![portal_inlet_address.clone(), context.address()], + route![portal_inlet_address.clone(), context.primary_address()], PortalMessage::Payload(first_piece_of_payload, None).to_neutral_message()?, ) .await?; context .send( - route![portal_inlet_address, context.address()], + route![portal_inlet_address, context.primary_address()], PortalMessage::Payload(second_piece_of_payload, None).to_neutral_message()?, ) .await?; @@ -127,7 +127,7 @@ async fn kafka_portal_worker__double_kafka_message__message_assembled( let double_payload = request_buffer.as_ref(); context .send( - route![portal_inlet_address.clone(), context.address()], + route![portal_inlet_address.clone(), context.primary_address()], PortalMessage::Payload(double_payload, None).to_neutral_message()?, ) .await?; @@ -172,7 +172,7 @@ async fn kafka_portal_worker__bigger_than_limit_kafka_message__error( for chunk in huge_payload.chunks(MAX_PAYLOAD_SIZE) { let _error = context .send( - route![portal_inlet_address.clone(), context.address()], + route![portal_inlet_address.clone(), context.primary_address()], PortalMessage::Payload(chunk, None).to_neutral_message()?, ) .await; @@ -314,7 +314,7 @@ async fn setup_only_worker(context: &mut Context, handle: &NodeManagerHandle) -> PortalInterceptorWorker::create_inlet_interceptor( context, None, - route![context.address()], + route![context.primary_address()], Arc::new(AllowAll), Arc::new(AllowAll), Arc::new(KafkaMessageInterceptorWrapper::new( @@ -412,7 +412,7 @@ async fn kafka_portal_worker__metadata_exchange__response_changed( let portal_inlet_address = PortalInterceptorWorker::create_inlet_interceptor( context, None, - route![context.address()], + route![context.primary_address()], Arc::new(AllowAll), Arc::new(AllowAll), Arc::new(KafkaMessageInterceptorWrapper::new( @@ -438,7 +438,7 @@ async fn kafka_portal_worker__metadata_exchange__response_changed( context .send( - route![portal_inlet_address, context.address()], + route![portal_inlet_address, context.primary_address()], PortalMessage::Payload(&request_buffer, None).to_neutral_message()?, ) .await?; diff --git a/implementations/rust/ockam/ockam_api/src/nodes/connection/mod.rs b/implementations/rust/ockam/ockam_api/src/nodes/connection/mod.rs index 4b118410436..17a6aa76006 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/connection/mod.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/connection/mod.rs @@ -91,7 +91,7 @@ impl Connection { if let Some(tcp_connection) = self.tcp_connection.as_ref() { let address = tcp_connection.sender_address().clone(); - if let Err(error) = node_manager.tcp_transport.disconnect(address.clone()).await { + if let Err(error) = node_manager.tcp_transport.disconnect(address.clone()) { match error.code().kind { Kind::NotFound => { debug!("cannot find and disconnect tcp worker `{tcp_connection}`"); @@ -306,7 +306,7 @@ impl ConnectionBuilder { // last piece only if it's a terminal (a service connecting to another node) if last_pass && is_last { let is_terminal = ctx - .get_metadata(address.clone()) + .get_metadata(address.clone())? .map(|m| m.is_terminal) .unwrap_or(false); diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/background_node_client.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/background_node_client.rs index 3c37dc8e94c..5762082cbd5 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/background_node_client.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/background_node_client.rs @@ -140,7 +140,7 @@ impl BackgroundNodeClient { .success() .into_diagnostic(); - let _ = tcp_connection.stop(ctx).await; + let _ = tcp_connection.stop(ctx); res } @@ -158,7 +158,7 @@ impl BackgroundNodeClient { let (tcp_connection, client) = self.make_client().await?; let res = client.ask(ctx, req).await.into_diagnostic(); - let _ = tcp_connection.stop(ctx).await; + let _ = tcp_connection.stop(ctx); res } @@ -175,7 +175,7 @@ impl BackgroundNodeClient { .success() .into_diagnostic(); - let _ = tcp_connection.stop(ctx).await; + let _ = tcp_connection.stop(ctx); res } @@ -191,7 +191,7 @@ impl BackgroundNodeClient { let (tcp_connection, client) = self.make_client().await?; let res = client.tell(ctx, req).await.into_diagnostic(); - let _ = tcp_connection.stop(ctx).await; + let _ = tcp_connection.stop(ctx); res } @@ -200,7 +200,6 @@ impl BackgroundNodeClient { self.create_tcp_connection() .await? .stop(ctx) - .await .into_diagnostic() } diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/in_memory_node.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/in_memory_node.rs index a626603f779..485e77cc69f 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/in_memory_node.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/in_memory_node.rs @@ -192,7 +192,7 @@ impl InMemoryNode { } for addr in DefaultAddress::iter() { - let result = ctx.stop_worker(addr).await; + let result = ctx.stop_address(addr); // when stopping we can safely ignore missing services if let Err(err) = result { if err.code().kind == Kind::NotFound { diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/kafka_services.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/kafka_services.rs index 05ed74d73d0..9543c085e4b 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/kafka_services.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/kafka_services.rs @@ -351,11 +351,11 @@ impl InMemoryNode { if kind.eq(e.kind()) { match e.kind() { KafkaServiceKind::Inlet => { - ctx.stop_worker(address.clone()).await?; + ctx.stop_address(address.clone())?; } KafkaServiceKind::Outlet => { - ctx.stop_worker(KAFKA_OUTLET_INTERCEPTOR_ADDRESS).await?; - ctx.stop_worker(KAFKA_OUTLET_BOOTSTRAP_ADDRESS).await?; + ctx.stop_address(KAFKA_OUTLET_INTERCEPTOR_ADDRESS)?; + ctx.stop_address(KAFKA_OUTLET_BOOTSTRAP_ADDRESS)?; } } self.registry.kafka_services.remove(&address).await; diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/node_services.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/node_services.rs index 27a11c0ad6e..ac70d31e21d 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/node_services.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/node_services.rs @@ -202,8 +202,8 @@ impl NodeManager { ))); } - if ctx.is_worker_registered_at(&addr) { - ctx.stop_worker(addr.clone()).await? + if ctx.is_worker_registered_at(&addr)? { + ctx.stop_address(addr.clone())? }; let (incoming_ac, outgoing_ac) = self diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/relay.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/relay.rs index af7b3429da1..c9e44c5fac4 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/relay.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/relay.rs @@ -361,7 +361,7 @@ impl SessionReplacer for RelaySessionReplacer { } if let Some(relay_address) = self.relay_worker_address.take() { - match self.context.stop_worker(relay_address.clone()).await { + match self.context.stop_address(relay_address.clone()) { Ok(_) => { debug!(%relay_address, "Successfully stopped relay"); } diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/secure_channel.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/secure_channel.rs index 916fc710ce6..1348f26d406 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/secure_channel.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/secure_channel.rs @@ -271,7 +271,7 @@ impl NodeManager { format!("Secure channel with address, {}, not found", addr), )); } - self.secure_channels.stop_secure_channel(ctx, addr).await?; + self.secure_channels.stop_secure_channel(ctx, addr)?; self.registry.secure_channels.remove_by_addr(addr).await; Ok(()) } @@ -421,7 +421,7 @@ impl NodeManager { addr: &Address, ) -> Result { debug!("deleting secure channel listener: {addr}"); - ctx.stop_worker(addr.clone()).await?; + ctx.stop_address(addr.clone())?; self.registry .secure_channel_listeners .remove(addr) diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/tcp_inlets/session_replacer.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/tcp_inlets/session_replacer.rs index 9fb91fe0052..d4285afd2d7 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/tcp_inlets/session_replacer.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/tcp_inlets/session_replacer.rs @@ -127,7 +127,7 @@ impl InletSessionReplacer { } async fn create_impl(&mut self, node_manager: &NodeManager) -> Result { - self.pause_inlet().await; + self.pause_inlet(); self.close_connection(node_manager).await; let connection = node_manager @@ -158,9 +158,7 @@ impl InletSessionReplacer { // Finally, attempt to create/update inlet using the new route let inlet_address = match self.inlet.clone() { Some(inlet) => { - inlet - .unpause(&self.context, normalized_stripped_route.clone()) - .await?; + inlet.unpause(&self.context, normalized_stripped_route.clone())?; inlet.processor_address().cloned() } @@ -213,16 +211,16 @@ impl InletSessionReplacer { }) } - async fn pause_inlet(&mut self) { + fn pause_inlet(&mut self) { if let Some(inlet) = self.inlet.as_mut() { - inlet.pause().await; + inlet.pause(); } } - async fn close_inlet(&mut self) { + fn close_inlet(&mut self) { if let Some(inlet) = self.inlet.take() { // The previous inlet worker needs to be stopped: - let result = inlet.stop(&self.context).await; + let result = inlet.stop(&self.context); if let Err(err) = result { error!( @@ -290,7 +288,7 @@ impl SessionReplacer for InletSessionReplacer { return; }; - self.close_inlet().await; + self.close_inlet(); self.close_connection(&node_manager).await; } } @@ -388,7 +386,7 @@ impl AdditionalSessionReplacer for InletSessionReplacer { additional_sc.update_remote_node_route(route![puncture.sender_address()])?; let new_route = route![additional_sc.clone()]; - inlet.unpause(&self.context, new_route.clone()).await?; + inlet.unpause(&self.context, new_route.clone())?; self.additional_route = Some(new_route.clone()); @@ -402,20 +400,20 @@ impl AdditionalSessionReplacer for InletSessionReplacer { match self.main_route.as_ref() { Some(main_route) if enable_fallback => { // Switch Inlet to the main route - let res = inlet.unpause(&self.context, main_route.clone()).await; + let res = inlet.unpause(&self.context, main_route.clone()); if let Some(err) = res.err() { error!("Error switching Inlet to the main route {}", err); } } _ => { - inlet.pause().await; + inlet.pause(); } } } if let Some(secure_channel) = self.additional_secure_channel.take() { - let res = self.context.stop_worker(secure_channel).await; + let res = self.context.stop_address(secure_channel); if let Some(err) = res.err() { error!("Error closing secure channel {}", err); @@ -423,7 +421,7 @@ impl AdditionalSessionReplacer for InletSessionReplacer { } if let Some(puncture) = self.udp_puncture.take() { - let res = puncture.stop(&self.context).await; + let res = puncture.stop(&self.context); if let Some(err) = res.err() { error!("Error stopping puncture {}", err); diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/tcp_outlets.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/tcp_outlets.rs index f3f56d44ed1..1bf3b32271d 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/tcp_outlets.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/tcp_outlets.rs @@ -226,7 +226,6 @@ impl NodeManager { if let Err(e) = self .tcp_transport .stop_outlet(deleted_outlet.worker_addr.clone()) - .await { warn!(%worker_addr, %e, "Failed to stop outlet worker"); } diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/transport.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/transport.rs index d1976925510..dfc953de8a8 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/transport.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/transport.rs @@ -63,7 +63,7 @@ impl NodeManager { Ok(listener.into()) } - async fn delete_tcp_connection(&self, address: String) -> Result<(), String> { + fn delete_tcp_connection(&self, address: String) -> Result<(), String> { let sender_address = match address.parse::() { Ok(socket_address) => self .tcp_transport() @@ -77,11 +77,10 @@ impl NodeManager { self.tcp_transport .disconnect(sender_address.clone()) - .await .map_err(|err| format!("Unable to disconnect from {sender_address}: {err}")) } - async fn delete_tcp_listener(&self, address: String) -> Result<(), String> { + fn delete_tcp_listener(&self, address: String) -> Result<(), String> { let listener_address = match address.parse::() { Ok(socket_address) => self .tcp_transport() @@ -95,7 +94,6 @@ impl NodeManager { self.tcp_transport .stop_listener(&listener_address) - .await .map_err(|err| format!("Unable to stop listener {listener_address}: {err}")) } } @@ -178,7 +176,7 @@ impl NodeManagerWorker { }) } - pub(super) async fn delete_tcp_connection( + pub(super) fn delete_tcp_connection( &self, delete: DeleteTransport, ) -> Result, Response> { @@ -186,12 +184,11 @@ impl NodeManagerWorker { self.node_manager .delete_tcp_connection(delete.address) - .await .map(|status| Response::ok().body(status)) .map_err(|msg| Response::bad_request_no_request(&msg)) } - pub(super) async fn delete_tcp_listener( + pub(super) fn delete_tcp_listener( &self, delete: DeleteTransport, ) -> Result, Response> { @@ -199,7 +196,6 @@ impl NodeManagerWorker { self.node_manager .delete_tcp_listener(delete.address) - .await .map(|status| Response::ok().body(status)) .map_err(|msg| Response::bad_request_no_request(&msg)) } diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/worker.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/worker.rs index 2580e6bcd63..15ff4093a12 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/worker.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/worker.rs @@ -23,7 +23,7 @@ impl NodeManagerWorker { // TODO: This is never called. pub async fn stop(&self, ctx: &Context) -> Result<()> { self.node_manager.stop(ctx).await?; - ctx.stop_worker(NODEMANAGER_ADDR).await?; + ctx.stop_address(NODEMANAGER_ADDR)?; Ok(()) } } @@ -69,7 +69,7 @@ impl NodeManagerWorker { encode_response(req, self.create_tcp_connection(ctx, dec.decode()?).await)? } (Delete, ["node", "tcp", "connection"]) => { - encode_response(req, self.delete_tcp_connection(dec.decode()?).await)? + encode_response(req, self.delete_tcp_connection(dec.decode()?))? } // ==*== Tcp Listeners ==*== @@ -81,7 +81,7 @@ impl NodeManagerWorker { encode_response(req, self.create_tcp_listener(dec.decode()?).await)? } (Delete, ["node", "tcp", "listener"]) => { - encode_response(req, self.delete_tcp_listener(dec.decode()?).await)? + encode_response(req, self.delete_tcp_listener(dec.decode()?))? } // ==*== Secure channels ==*== diff --git a/implementations/rust/ockam/ockam_api/src/nodes/service/workers.rs b/implementations/rust/ockam/ockam_api/src/nodes/service/workers.rs index 8a350ed4141..bf8bc00145b 100644 --- a/implementations/rust/ockam/ockam_api/src/nodes/service/workers.rs +++ b/implementations/rust/ockam/ockam_api/src/nodes/service/workers.rs @@ -11,7 +11,7 @@ impl NodeManagerWorker { ctx: &Context, ) -> Result, Response> { let list = ctx - .list_workers() + .list_workers()? .into_iter() .map(|addr| WorkerStatus::new(addr.address())) .collect(); diff --git a/implementations/rust/ockam/ockam_api/src/rendezvous_healthcheck.rs b/implementations/rust/ockam/ockam_api/src/rendezvous_healthcheck.rs index 182395dce07..b64eea44393 100644 --- a/implementations/rust/ockam/ockam_api/src/rendezvous_healthcheck.rs +++ b/implementations/rust/ockam/ockam_api/src/rendezvous_healthcheck.rs @@ -147,7 +147,7 @@ impl RendezvousHealthcheckTask { ) }); - self.udp.unbind(bind).await?; + self.udp.unbind(bind)?; res } diff --git a/implementations/rust/ockam/ockam_api/src/session/session.rs b/implementations/rust/ockam/ockam_api/src/session/session.rs index 0beebe7f158..66d573f7ed8 100644 --- a/implementations/rust/ockam/ockam_api/src/session/session.rs +++ b/implementations/rust/ockam/ockam_api/src/session/session.rs @@ -375,10 +375,7 @@ impl Session { .await; additional_state.shared_state.status.set_down(); - _ = self - .ctx - .stop_worker(additional_state.collector_address) - .await; + _ = self.ctx.stop_address(additional_state.collector_address); } } @@ -393,7 +390,7 @@ impl Session { // ping_receiver_handle task will shut down itself when Collector Worker drops the sender - _ = self.ctx.stop_worker(self.collector_address.clone()).await; + _ = self.ctx.stop_address(self.collector_address.clone()); } /// Stop everything diff --git a/implementations/rust/ockam/ockam_api/src/test_utils/mod.rs b/implementations/rust/ockam/ockam_api/src/test_utils/mod.rs index e2ec9d23c72..e1944d6c5fe 100644 --- a/implementations/rust/ockam/ockam_api/src/test_utils/mod.rs +++ b/implementations/rust/ockam/ockam_api/src/test_utils/mod.rs @@ -2,7 +2,7 @@ use crate::config::lookup::InternetAddress; use crate::nodes::service::{NodeManagerCredentialRetrieverOptions, NodeManagerTrustOptions}; -use ockam_node::{Context, NodeBuilder}; +use ockam_node::{Context, Executor, NodeBuilder}; use sqlx::__rt::timeout; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::ops::Deref; @@ -200,6 +200,7 @@ pub async fn start_tcp_echo_server() -> EchoServerHandle { } pub struct TestNode { + pub executor: Executor, pub context: Context, pub node_manager_handle: NodeManagerHandle, } @@ -209,16 +210,14 @@ impl TestNode { /// needs be cleaned-up before a test is executed pub async fn clean() -> Result<()> { if let Some(configuration) = DatabaseConfiguration::postgres()? { - let db = SqlxDatabase::create_no_migration(&configuration) - .await - .unwrap(); + let db = SqlxDatabase::create_no_migration(&configuration).await?; db.drop_all_postgres_tables().await?; }; Ok(()) } pub async fn create(runtime: Arc, listen_addr: Option<&str>) -> Self { - let (mut context, _executor) = NodeBuilder::new().with_runtime(runtime.clone()).build(); + let (mut context, executor) = NodeBuilder::new().with_runtime(runtime).build(); let node_manager_handle = start_manager_for_tests( &mut context, listen_addr, @@ -233,6 +232,7 @@ impl TestNode { .expect("cannot start node manager"); Self { + executor, context, node_manager_handle, } diff --git a/implementations/rust/ockam/ockam_api/tests/authority.rs b/implementations/rust/ockam/ockam_api/tests/authority.rs index 7ec768ffc0f..094d42862cc 100644 --- a/implementations/rust/ockam/ockam_api/tests/authority.rs +++ b/implementations/rust/ockam/ockam_api/tests/authority.rs @@ -10,7 +10,7 @@ async fn authority_starts_with_default_configuration(ctx: &mut Context) -> Resul let configuration = default_configuration().await?; start_authority_node(ctx, &configuration).await?; - let workers = ctx.list_workers(); + let workers = ctx.list_workers()?; assert!(!workers.contains(&Address::from(DefaultAddress::DIRECT_AUTHENTICATOR))); assert!(!workers.contains(&Address::from(DefaultAddress::ENROLLMENT_TOKEN_ACCEPTOR))); @@ -28,7 +28,7 @@ async fn authority_starts_direct_authenticator(ctx: &mut Context) -> Result<()> configuration.no_direct_authentication = false; start_authority_node(ctx, &configuration).await?; - let workers = ctx.list_workers(); + let workers = ctx.list_workers()?; assert!(workers.contains(&Address::from(DefaultAddress::DIRECT_AUTHENTICATOR))); assert!(!workers.contains(&Address::from(DefaultAddress::ENROLLMENT_TOKEN_ACCEPTOR))); @@ -46,7 +46,7 @@ async fn authority_starts_enrollment_token(ctx: &mut Context) -> Result<()> { configuration.no_token_enrollment = false; start_authority_node(ctx, &configuration).await?; - let workers = ctx.list_workers(); + let workers = ctx.list_workers()?; assert!(!workers.contains(&Address::from(DefaultAddress::DIRECT_AUTHENTICATOR))); assert!(workers.contains(&Address::from(DefaultAddress::ENROLLMENT_TOKEN_ACCEPTOR))); diff --git a/implementations/rust/ockam/ockam_api/tests/common/mod.rs b/implementations/rust/ockam/ockam_api/tests/common/mod.rs index 9a3370c152a..eae79b75ff0 100644 --- a/implementations/rust/ockam/ockam_api/tests/common/mod.rs +++ b/implementations/rust/ockam/ockam_api/tests/common/mod.rs @@ -4,5 +4,3 @@ pub mod common; pub mod session; #[allow(dead_code)] pub mod test_spans; -#[allow(dead_code)] -pub mod trace_code; diff --git a/implementations/rust/ockam/ockam_api/tests/common/session.rs b/implementations/rust/ockam/ockam_api/tests/common/session.rs index 86e2723d5be..285da5285ea 100644 --- a/implementations/rust/ockam/ockam_api/tests/common/session.rs +++ b/implementations/rust/ockam/ockam_api/tests/common/session.rs @@ -81,11 +81,11 @@ impl Worker for MockHop { async fn handle_message(&mut self, ctx: &mut Context, msg: Routed) -> Result<()> { if !self.responsive.load(Ordering::Relaxed) { - info!("Drop Hop message, {}", ctx.address()); + info!("Drop Hop message, {}", ctx.primary_address()); return Ok(()); } - info!("Forward Hop message {}", ctx.address()); + info!("Forward Hop message {}", ctx.primary_address()); let msg = msg.into_local_message(); let msg = msg.pop_front_onward_route()?; diff --git a/implementations/rust/ockam/ockam_api/tests/common/trace_code.rs b/implementations/rust/ockam/ockam_api/tests/common/trace_code.rs deleted file mode 100644 index 148dbac9cd7..00000000000 --- a/implementations/rust/ockam/ockam_api/tests/common/trace_code.rs +++ /dev/null @@ -1,73 +0,0 @@ -use crate::common::test_spans::Trace; -use ockam_api::logs::{ExportingConfiguration, LoggingConfiguration, LoggingTracing}; -use ockam_core::{AsyncTryClone, OpenTelemetryContext}; -use ockam_node::{Context, NodeBuilder}; -use opentelemetry::global; -use opentelemetry::trace::{FutureExt, Tracer}; -use opentelemetry_sdk::export::trace::SpanData; -use opentelemetry_sdk::testing::logs::InMemoryLogsExporter; -use opentelemetry_sdk::testing::trace::InMemorySpanExporter; -use std::future::Future; - -/// Run an async function using a tracer and return: -/// -/// - the return value of the function -/// - all the exported spans -pub fn trace_code(f: impl Fn(Context) -> F + Send + Sync + 'static) -> (F::Output, Vec) -where - F: Future + Send + 'static, - F::Output: Send + 'static, -{ - let spans_exporter = InMemorySpanExporter::default(); - let guard = LoggingTracing::setup_with_exporters( - spans_exporter.clone(), - InMemoryLogsExporter::default(), - &LoggingConfiguration::off() - .unwrap() - .set_all_crates() - .set_log_level(tracing_core::metadata::Level::TRACE), - &ExportingConfiguration::foreground().unwrap(), - "test", - None, - ); - - let (mut ctx, mut executor) = NodeBuilder::new().build(); - - let tracer = global::tracer("ockam-test"); - let result = executor - .execute_no_abort(async move { - let result = tracer - .in_span("root", |_| { - ctx.set_tracing_context(OpenTelemetryContext::current()); - async { f(ctx.async_try_clone().await.unwrap()).await }.with_current_context() - }) - .await; - let _ = ctx.stop().await; - result - }) - .unwrap(); - - // get the exported spans - guard.force_flush(); - let spans = spans_exporter.get_finished_spans().unwrap(); - (result, spans) -} - -/// Return a string displaying the traces for all the given spans -pub fn display_traces(spans: Vec) -> String { - let mut traces = Trace::from_span_data(spans); - traces.sort_by_key(|t| t.to_string()); - - // remove some uninteresting traces based on their root name - traces.retain(|t| { - !["start", "shutdown", "initialize", "process"] - .iter() - .any(|v| t.0.to_string().split('\n').collect::>()[0].ends_with(v)) - }); - - traces - .iter() - .map(|t| t.to_string()) - .collect::>() - .join("\n") -} diff --git a/implementations/rust/ockam/ockam_api/tests/latency.rs b/implementations/rust/ockam/ockam_api/tests/latency.rs index 4eb7bf974ee..fe8789c6567 100644 --- a/implementations/rust/ockam/ockam_api/tests/latency.rs +++ b/implementations/rust/ockam/ockam_api/tests/latency.rs @@ -64,7 +64,7 @@ pub fn measure_message_latency_two_nodes() { first_node .context .flow_controls() - .add_consumer(first_node.context.address(), &flow_control_id); + .add_consumer(first_node.context.primary_address(), &flow_control_id); } let payload = NeutralMessage::from(vec![1, 2, 3, 4]); diff --git a/implementations/rust/ockam/ockam_api/tests/session.rs b/implementations/rust/ockam/ockam_api/tests/session.rs index 1fbb2d0f3d8..e052b9f5710 100644 --- a/implementations/rust/ockam/ockam_api/tests/session.rs +++ b/implementations/rust/ockam/ockam_api/tests/session.rs @@ -114,12 +114,12 @@ async fn start_monitoring__available__should_be_up_fast(ctx: &mut Context) -> Re ctx.start_worker(Address::from_string("echo"), MockEchoer::new()) .await?; - assert!(!ctx.is_worker_registered_at(session.collector_address())); + assert!(!ctx.is_worker_registered_at(session.collector_address())?); // Start the Session in a separate task session.start_monitoring().await?; - assert!(ctx.is_worker_registered_at(session.collector_address())); + assert!(ctx.is_worker_registered_at(session.collector_address())?); let mut time_to_restore = 0; diff --git a/implementations/rust/ockam/ockam_app_lib/src/shared_service/tcp_outlet/invitation_access_control.rs b/implementations/rust/ockam/ockam_app_lib/src/shared_service/tcp_outlet/invitation_access_control.rs index d33d81c5c31..ed02f8fa48b 100644 --- a/implementations/rust/ockam/ockam_app_lib/src/shared_service/tcp_outlet/invitation_access_control.rs +++ b/implementations/rust/ockam/ockam_app_lib/src/shared_service/tcp_outlet/invitation_access_control.rs @@ -190,7 +190,7 @@ pub struct InvitationOutgoingAccessControl { #[async_trait] impl OutgoingAccessControl for InvitationOutgoingAccessControl { async fn is_authorized(&self, relay_message: &RelayMessage) -> ockam_core::Result { - let identifier = match Abac::get_outgoing_identifier(&self.ctx, relay_message).await? { + let identifier = match Abac::get_outgoing_identifier(&self.ctx, relay_message)? { Some(identifier) => identifier, None => { debug!("identity identifier not found; access denied"); diff --git a/implementations/rust/ockam/ockam_app_lib/src/state/mod.rs b/implementations/rust/ockam/ockam_app_lib/src/state/mod.rs index b1c89ddb383..bc70f1a8894 100644 --- a/implementations/rust/ockam/ockam_app_lib/src/state/mod.rs +++ b/implementations/rust/ockam/ockam_app_lib/src/state/mod.rs @@ -319,8 +319,8 @@ impl AppState { info!("stopped the old node manager"); - for w in self.context.list_workers() { - let _ = self.context.stop_worker(w.address()).await; + for w in self.context.list_workers()? { + let _ = self.context.stop_address(w.address()); } info!("stopped all the ctx workers"); diff --git a/implementations/rust/ockam/ockam_core/Cargo.toml b/implementations/rust/ockam/ockam_core/Cargo.toml index 17d11801368..f88a124faef 100644 --- a/implementations/rust/ockam/ockam_core/Cargo.toml +++ b/implementations/rust/ockam/ockam_core/Cargo.toml @@ -76,7 +76,7 @@ backtrace = { version = "0.3", default-features = false, features = ["std", "ser cfg-if = "1.0" core2 = { version = "0.4.0", default-features = false, optional = true } futures-util = { version = "0.3.30", default-features = false, features = ["alloc", "async-await-macro", "sink"] } -hashbrown = { version = "0.14", default-features = false, features = ["ahash", "serde"] } +hashbrown = { version = "0.15.2", features = ["serde", "default-hasher"] } hex = { version = "0.4", default-features = false, optional = true } miette = { version = "7.2.0", features = ["fancy-no-backtrace"], optional = true } minicbor = { version = "0.25.1", default-features = false, features = ["derive"] } diff --git a/implementations/rust/ockam/ockam_core/src/api.rs b/implementations/rust/ockam/ockam_core/src/api.rs index 63c144ebc17..7409c3e188d 100644 --- a/implementations/rust/ockam/ockam_core/src/api.rs +++ b/implementations/rust/ockam/ockam_core/src/api.rs @@ -125,7 +125,7 @@ impl Serialize for Reply { match self { Reply::Successful(t) => t.serialize(serializer), Reply::Failed(e, Some(s)) => { - let mut map = HashMap::new(); + let mut map: HashMap<&str, String> = Default::default(); map.insert("error", e.to_string()); map.insert("status", s.to_string()); serializer.collect_map(map) diff --git a/implementations/rust/ockam/ockam_core/src/compat.rs b/implementations/rust/ockam/ockam_core/src/compat.rs index 1fba92283d5..dc8d008cf17 100644 --- a/implementations/rust/ockam/ockam_core/src/compat.rs +++ b/implementations/rust/ockam/ockam_core/src/compat.rs @@ -30,6 +30,17 @@ pub mod collections { pub use alloc::collections::{BTreeMap, BTreeSet, BinaryHeap, LinkedList, VecDeque}; pub use hashbrown::{HashMap, HashSet}; + + /// hash map + pub mod hash_map { + pub use hashbrown::hash_map::{Entry, EntryRef}; + } + + /// btree map + #[cfg(feature = "alloc")] + pub mod btree_map { + pub use alloc::collections::btree_map::Entry; + } } /// Provides a `std::error::Error` trait. @@ -224,7 +235,7 @@ pub mod str { pub mod sync { use core::convert::Infallible; - pub use alloc::sync::Arc; + pub use alloc::sync::{Arc, Weak}; /// Wrap `spin::RwLock` as it does not return LockResult like `std::sync::Mutex`. #[derive(Debug)] @@ -293,7 +304,7 @@ pub mod sync { /// Provides `std::sync` for `std` targets. #[cfg(feature = "std")] pub mod sync { - pub use std::sync::Arc; + pub use std::sync::{Arc, Weak}; pub use std::sync::{Mutex, RwLock, RwLockWriteGuard}; } diff --git a/implementations/rust/ockam/ockam_core/src/flow_control/flow_control_id.rs b/implementations/rust/ockam/ockam_core/src/flow_control/flow_control_id.rs index 4769b7b4e77..5c657fa7c96 100644 --- a/implementations/rust/ockam/ockam_core/src/flow_control/flow_control_id.rs +++ b/implementations/rust/ockam/ockam_core/src/flow_control/flow_control_id.rs @@ -7,7 +7,7 @@ use minicbor::{CborLen, Decode, Encode}; use serde::{Deserialize, Serialize}; /// Unique random identifier of a Flow Control -#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, CborLen)] +#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Encode, Decode, CborLen)] #[rustfmt::skip] #[cbor(map)] pub struct FlowControlId { diff --git a/implementations/rust/ockam/ockam_core/src/flow_control/flow_controls/flow_controls.rs b/implementations/rust/ockam/ockam_core/src/flow_control/flow_controls/flow_controls.rs index 156bf092490..ceacba18192 100644 --- a/implementations/rust/ockam/ockam_core/src/flow_control/flow_controls/flow_controls.rs +++ b/implementations/rust/ockam/ockam_core/src/flow_control/flow_controls/flow_controls.rs @@ -1,4 +1,4 @@ -use crate::compat::collections::BTreeMap; +use crate::compat::collections::HashMap; use crate::compat::sync::{Arc, RwLock}; use crate::flow_control::{ConsumersInfo, FlowControlId, ProducerInfo}; use crate::Address; @@ -7,12 +7,12 @@ use crate::Address; #[derive(Clone, Debug)] pub struct FlowControls { // All known consumers - pub(super) consumers: Arc>>, + pub(super) consumers: Arc>>, // All known producers - pub(super) producers: Arc>>, + pub(super) producers: Arc>>, // Allows to find producer by having its additional Address, // e.g. Decryptor by its Encryptor Address or TCP Receiver by its TCP Sender Address - pub(super) producers_additional_addresses: Arc>>, + pub(super) producers_additional_addresses: Arc>>, // All known spawners - pub(super) spawners: Arc>>, + pub(super) spawners: Arc>>, } diff --git a/implementations/rust/ockam/ockam_core/src/identity/secure_channel_metadata.rs b/implementations/rust/ockam/ockam_core/src/identity/secure_channel_metadata.rs index d03f3b83ad6..12b8958205a 100644 --- a/implementations/rust/ockam/ockam_core/src/identity/secure_channel_metadata.rs +++ b/implementations/rust/ockam/ockam_core/src/identity/secure_channel_metadata.rs @@ -1,6 +1,6 @@ use crate::compat::string::{String, ToString}; use crate::errcode::{Kind, Origin}; -use crate::{AddressAndMetadata, Error, LocalInfoIdentifier, Result, SECURE_CHANNEL_IDENTIFIER}; +use crate::{AddressMetadata, Error, LocalInfoIdentifier, Result, SECURE_CHANNEL_IDENTIFIER}; /// SecureChannel Metadata used for Terminal Address pub struct SecureChannelMetadata { @@ -34,19 +34,15 @@ impl SecureChannelMetadata { } /// Get the Identifier of the other side of the Secure Channel - pub fn from_terminal_address(terminal: &AddressAndMetadata) -> Result { + pub fn from_terminal_address_metadata(terminal: &AddressMetadata) -> Result { let identifier = if let Some(identifier) = - terminal - .metadata - .attributes - .iter() - .find_map(|(key, value)| { - if key == SECURE_CHANNEL_IDENTIFIER { - Some(value.clone()) - } else { - None - } - }) { + terminal.attributes.iter().find_map(|(key, value)| { + if key == SECURE_CHANNEL_IDENTIFIER { + Some(value.clone()) + } else { + None + } + }) { identifier } else { return Err(Self::error_type_id()); diff --git a/implementations/rust/ockam/ockam_core/src/routing/address.rs b/implementations/rust/ockam/ockam_core/src/routing/address.rs index 03cf25b8061..e3f7e925a9d 100644 --- a/implementations/rust/ockam/ockam_core/src/routing/address.rs +++ b/implementations/rust/ockam/ockam_core/src/routing/address.rs @@ -40,6 +40,13 @@ pub struct Address { #[n(1)] inner: Vec, } +// FIXME +impl From<&Address> for Address { + fn from(value: &Address) -> Self { + value.clone() + } +} + impl Address { /// Creates a new address from separate transport type and data parts. /// diff --git a/implementations/rust/ockam/ockam_core/src/routing/address_meta.rs b/implementations/rust/ockam/ockam_core/src/routing/address_meta.rs index 843816166ec..9d5ca25683a 100644 --- a/implementations/rust/ockam/ockam_core/src/routing/address_meta.rs +++ b/implementations/rust/ockam/ockam_core/src/routing/address_meta.rs @@ -1,9 +1,8 @@ use crate::compat::string::String; use crate::compat::vec::Vec; -use crate::Address; /// Additional metadata for address -#[derive(Debug, Clone, Eq, PartialEq, Default)] +#[derive(Debug, Clone, Eq, PartialEq)] pub struct AddressMetadata { /// Indicates that this Address will forward message to another route, therefore the next /// hop after this one belongs to another node @@ -11,12 +10,3 @@ pub struct AddressMetadata { /// Arbitrary set of attributes pub attributes: Vec<(String, String)>, } - -/// A set of metadata for a particular address -#[derive(Debug, Clone, Eq, PartialEq)] -pub struct AddressAndMetadata { - /// Address - pub address: Address, - /// Metadata - pub metadata: AddressMetadata, -} diff --git a/implementations/rust/ockam/ockam_core/src/routing/mailbox.rs b/implementations/rust/ockam/ockam_core/src/routing/mailbox.rs index 098a59f02a8..86ce9b883e3 100644 --- a/implementations/rust/ockam/ockam_core/src/routing/mailbox.rs +++ b/implementations/rust/ockam/ockam_core/src/routing/mailbox.rs @@ -1,6 +1,8 @@ use crate::access_control::IncomingAccessControl; use crate::compat::{sync::Arc, vec::Vec}; -use crate::{debugger, Address, DenyAll, OutgoingAccessControl, RelayMessage, Result}; +use crate::{ + debugger, Address, AddressMetadata, DenyAll, OutgoingAccessControl, RelayMessage, Result, +}; use core::cmp::Ordering; use core::fmt::{self, Debug}; @@ -10,6 +12,7 @@ use core::fmt::{self, Debug}; #[derive(Clone)] pub struct Mailbox { address: Address, + metadata: Option, incoming: Arc, outgoing: Arc, } @@ -48,11 +51,13 @@ impl Mailbox { /// Create a new `Mailbox` with the given [`Address`], [`IncomingAccessControl`] and [`OutgoingAccessControl`] pub fn new( address: impl Into
, + metadata: Option, incoming: Arc, outgoing: Arc, ) -> Self { Self { address: address.into(), + metadata, incoming, outgoing, } @@ -62,6 +67,7 @@ impl Mailbox { pub fn deny_all(address: impl Into
) -> Self { Self { address: address.into(), + metadata: Default::default(), incoming: Arc::new(DenyAll), outgoing: Arc::new(DenyAll), } @@ -81,12 +87,17 @@ impl Mailbox { pub fn outgoing_access_control(&self) -> &Arc { &self.outgoing } + + /// Metadata + pub fn metadata(&self) -> &Option { + &self.metadata + } } /// A collection of [`Mailbox`]es for a specific [`Worker`](crate::Worker), [`Processor`](crate::Processor) or `Context` #[derive(Clone)] pub struct Mailboxes { - main_mailbox: Mailbox, + primary_mailbox: Mailbox, additional_mailboxes: Vec, } @@ -95,30 +106,31 @@ impl Debug for Mailboxes { write!( f, "{:?} + {:?}", - self.main_mailbox, self.additional_mailboxes + self.primary_mailbox, self.additional_mailboxes ) } } impl Mailboxes { - /// Create [`Mailboxes`] given main [`Mailbox`] and collection of additional [`Mailbox`]es - pub fn new(main_mailbox: Mailbox, additional_mailboxes: Vec) -> Self { + /// Create [`Mailboxes`] given primary [`Mailbox`] and collection of additional [`Mailbox`]es + pub fn new(primary_mailbox: Mailbox, additional_mailboxes: Vec) -> Self { Self { - main_mailbox, + primary_mailbox, additional_mailboxes, } } - /// Create [`Mailboxes`] with only main [`Mailbox`] for the given + /// Create [`Mailboxes`] with only primary [`Mailbox`] for the given /// [`Address`] with [`IncomingAccessControl`] and [`OutgoingAccessControl`] - pub fn main( + pub fn primary( address: impl Into
, incoming_access_control: Arc, outgoing_access_control: Arc, ) -> Self { Self { - main_mailbox: Mailbox::new( + primary_mailbox: Mailbox::new( address.into(), + None, incoming_access_control, outgoing_access_control, ), @@ -127,26 +139,18 @@ impl Mailboxes { } /// Return all additional [`Address`]es represented by these [`Mailboxes`] - pub fn additional_addresses(&self) -> Vec
{ - self.additional_mailboxes - .iter() - .map(|x| x.address.clone()) - .collect() - } - - /// Return the main [`Address`] of this [`Mailboxes`] - pub fn main_address(&self) -> Address { - self.main_mailbox.address.clone() + pub fn additional_addresses(&self) -> impl Iterator { + self.additional_mailboxes.iter().map(|x| &x.address) } - /// Return the main [`Address`] of this [`Mailboxes`] - pub fn main_address_ref(&self) -> &Address { - &self.main_mailbox.address + /// Return the primary [`Address`] of this [`Mailboxes`] + pub fn primary_address(&self) -> &Address { + &self.primary_mailbox.address } /// Return `true` if the given [`Address`] is included in this [`Mailboxes`] pub fn contains(&self, msg_addr: &Address) -> bool { - if &self.main_mailbox.address == msg_addr { + if &self.primary_mailbox.address == msg_addr { true } else { self.additional_mailboxes @@ -157,8 +161,8 @@ impl Mailboxes { /// Return a reference to the [`Mailbox`] with the given [`Address`] pub fn find_mailbox(&self, msg_addr: &Address) -> Option<&Mailbox> { - if &self.main_mailbox.address == msg_addr { - Some(&self.main_mailbox) + if &self.primary_mailbox.address == msg_addr { + Some(&self.primary_mailbox) } else { self.additional_mailboxes .iter() @@ -200,17 +204,9 @@ impl Mailboxes { } } - /// Return all (mail + additional) [`Address`]es represented by this [`Mailboxes`] - pub fn addresses(&self) -> Vec
{ - let mut addresses = Vec::with_capacity(self.additional_mailboxes.len() + 1); - addresses.push(self.main_mailbox.address.clone()); - addresses.append(&mut self.additional_addresses()); - addresses - } - - /// Return a reference to the main [`Mailbox`] for this [`Mailboxes`] - pub fn main_mailbox(&self) -> &Mailbox { - &self.main_mailbox + /// Return a reference to the primary [`Mailbox`] for this [`Mailboxes`] + pub fn primary_mailbox(&self) -> &Mailbox { + &self.primary_mailbox } /// Return a reference to the additional [`Mailbox`]es for this [`Mailboxes`] diff --git a/implementations/rust/ockam/ockam_identity/src/credentials/retriever/remote_retriever/remote_retriever_creator.rs b/implementations/rust/ockam/ockam_identity/src/credentials/retriever/remote_retriever/remote_retriever_creator.rs index c90c2a7b7ad..e0e18d5685e 100644 --- a/implementations/rust/ockam/ockam_identity/src/credentials/retriever/remote_retriever/remote_retriever_creator.rs +++ b/implementations/rust/ockam/ockam_identity/src/credentials/retriever/remote_retriever/remote_retriever_creator.rs @@ -99,7 +99,7 @@ impl CredentialRetrieverCreator for RemoteCredentialRetrieverCreator { "Creating new RemoteCredentialRetriever for: {}, authority: {}", subject, self.info.issuer ); - let mailboxes = Mailboxes::main( + let mailboxes = Mailboxes::primary( Address::random_tagged("RemoteCredentialRetriever"), Arc::new(DenyAll), Arc::new(AllowAll), diff --git a/implementations/rust/ockam/ockam_identity/src/secure_channel/decryptor.rs b/implementations/rust/ockam/ockam_identity/src/secure_channel/decryptor.rs index d37ce85fbac..476824e8685 100644 --- a/implementations/rust/ockam/ockam_identity/src/secure_channel/decryptor.rs +++ b/implementations/rust/ockam/ockam_identity/src/secure_channel/decryptor.rs @@ -147,13 +147,13 @@ impl DecryptorHandler { } } - async fn handle_close(&mut self, ctx: &mut Context) -> Result<()> { + fn handle_close(&mut self, ctx: &mut Context) -> Result<()> { // Prevent sending another Close message self.shared_state .should_send_close .store(false, Ordering::Relaxed); // Should be enough to stop the encryptor, since it will stop the decryptor - ctx.stop_worker(self.addresses.encryptor.clone()).await + ctx.stop_address(self.addresses.encryptor.clone()) } async fn handle_refresh_credentials( @@ -215,7 +215,7 @@ impl DecryptorHandler { SecureChannelMessage::RefreshCredentials(decrypted_msg) => { self.handle_refresh_credentials(ctx, decrypted_msg).await? } - SecureChannelMessage::Close => self.handle_close(ctx).await?, + SecureChannelMessage::Close => self.handle_close(ctx)?, }; Ok(()) diff --git a/implementations/rust/ockam/ockam_identity/src/secure_channel/encryptor_worker.rs b/implementations/rust/ockam/ockam_identity/src/secure_channel/encryptor_worker.rs index b958fdac970..108bcc6afc5 100644 --- a/implementations/rust/ockam/ockam_identity/src/secure_channel/encryptor_worker.rs +++ b/implementations/rust/ockam/ockam_identity/src/secure_channel/encryptor_worker.rs @@ -110,7 +110,7 @@ impl EncryptorWorker { Err(err) => { let address = self.addresses.encryptor.clone(); error!("Error while encrypting: {err} at: {address}"); - ctx.stop_worker(address).await?; + ctx.stop_address(address)?; Err(err) } } @@ -162,7 +162,7 @@ impl EncryptorWorker { .await?; if should_stop { - ctx.stop_worker(self.addresses.encryptor.clone()).await?; + ctx.stop_address(self.addresses.encryptor.clone())?; } Ok(()) @@ -334,7 +334,7 @@ impl Worker for EncryptorWorker { Ok(()) } - #[instrument(skip_all, name = "EncryptorWorker::handle_message", fields(worker = % ctx.address()))] + #[instrument(skip_all, name = "EncryptorWorker::handle_message", fields(worker = % ctx.primary_address()))] async fn handle_message( &mut self, ctx: &mut Self::Context, @@ -367,9 +367,7 @@ impl Worker for EncryptorWorker { credential_retriever.unsubscribe(&self.addresses.encryptor_internal)?; } - let _ = context - .stop_worker(self.addresses.decryptor_internal.clone()) - .await; + let _ = context.stop_address(self.addresses.decryptor_internal.clone()); if self.shared_state.should_send_close.load(Ordering::Relaxed) { let _ = self.send_close_channel(context).await; } diff --git a/implementations/rust/ockam/ockam_identity/src/secure_channel/handshake/handshake_worker.rs b/implementations/rust/ockam/ockam_identity/src/secure_channel/handshake/handshake_worker.rs index 40c8b3383a2..ee492b8bafe 100644 --- a/implementations/rust/ockam/ockam_identity/src/secure_channel/handshake/handshake_worker.rs +++ b/implementations/rust/ockam/ockam_identity/src/secure_channel/handshake/handshake_worker.rs @@ -4,8 +4,8 @@ use ockam_core::compat::boxed::Box; use ockam_core::compat::sync::{Arc, RwLock}; use ockam_core::errcode::{Kind, Origin}; use ockam_core::{ - AllowAll, Any, DenyAll, Error, Mailbox, Mailboxes, NeutralMessage, OutgoingAccessControl, - Route, Routed, SecureChannelMetadata, + AddressMetadata, AllowAll, Any, DenyAll, Error, Mailbox, Mailboxes, NeutralMessage, + OutgoingAccessControl, Route, Routed, SecureChannelMetadata, }; use ockam_core::{Result, Worker}; use ockam_node::callback::CallbackSender; @@ -116,7 +116,7 @@ impl Worker for HandshakeWorker { } async fn shutdown(&mut self, context: &mut Self::Context) -> Result<()> { - let _ = context.stop_worker(self.addresses.encryptor.clone()).await; + let _ = context.stop_address(self.addresses.encryptor.clone()); self.secure_channels .secure_channel_registry .unregister_channel(&self.addresses.encryptor); @@ -351,6 +351,7 @@ impl HandshakeWorker { ) -> Mailboxes { let remote_mailbox = Mailbox::new( addresses.decryptor_remote.clone(), + None, // Doesn't matter since we check incoming messages cryptographically, // but this may be reduced to allowing only from the transport connection that was used // to create this channel initially @@ -360,11 +361,13 @@ impl HandshakeWorker { ); let internal_mailbox = Mailbox::new( addresses.decryptor_internal.clone(), + None, Arc::new(DenyAll), decryptor_outgoing_access_control, ); let api_mailbox = Mailbox::new( addresses.decryptor_api.clone(), + None, Arc::new(AllowAll), Arc::new(AllowAll), ); @@ -424,16 +427,24 @@ impl HandshakeWorker { let main_mailbox = Mailbox::new( self.addresses.encryptor.clone(), + Some(AddressMetadata { + is_terminal: true, + attributes: vec![SecureChannelMetadata::attribute( + their_identifier.clone().into(), + )], + }), Arc::new(AllowAll), Arc::new(AllowAll), ); let api_mailbox = Mailbox::new( self.addresses.encryptor_api.clone(), + None, Arc::new(AllowAll), Arc::new(AllowAll), ); let internal_mailbox = Mailbox::new( self.addresses.encryptor_internal.clone(), + None, Arc::new(AllowAll), Arc::new(DenyAll), ); @@ -443,12 +454,6 @@ impl HandshakeWorker { main_mailbox, vec![api_mailbox, internal_mailbox], )) - .terminal_with_attributes( - self.addresses.encryptor.clone(), - vec![SecureChannelMetadata::attribute( - their_identifier.clone().into(), - )], - ) .start(context) .await?; } diff --git a/implementations/rust/ockam/ockam_identity/src/secure_channel/listener.rs b/implementations/rust/ockam/ockam_identity/src/secure_channel/listener.rs index 31a9838360e..3ab347d2430 100644 --- a/implementations/rust/ockam/ockam_identity/src/secure_channel/listener.rs +++ b/implementations/rust/ockam/ockam_identity/src/secure_channel/listener.rs @@ -70,7 +70,7 @@ impl Worker for SecureChannelListenerWorker { let addresses = Addresses::generate(Role::Responder); let flow_control_id = self.options.setup_flow_control_for_channel( ctx.flow_controls(), - ctx.address_ref(), + ctx.primary_address(), &addresses, ); let decryptor_outgoing_access_control = self diff --git a/implementations/rust/ockam/ockam_identity/src/secure_channels/secure_channels.rs b/implementations/rust/ockam/ockam_identity/src/secure_channels/secure_channels.rs index 9546970d2b2..e70687cd16a 100644 --- a/implementations/rust/ockam/ockam_identity/src/secure_channels/secure_channels.rs +++ b/implementations/rust/ockam/ockam_identity/src/secure_channels/secure_channels.rs @@ -292,7 +292,7 @@ impl SecureChannels { } /// Stop a SecureChannel given an encryptor address - pub async fn stop_secure_channel(&self, ctx: &Context, channel: &Address) -> Result<()> { - ctx.stop_worker(channel.clone()).await + pub fn stop_secure_channel(&self, ctx: &Context, channel: &Address) -> Result<()> { + ctx.stop_address(channel.clone()) } } diff --git a/implementations/rust/ockam/ockam_identity/src/secure_channels/secure_client.rs b/implementations/rust/ockam/ockam_identity/src/secure_channels/secure_client.rs index 3516bff31dd..bb7ba8dca9d 100644 --- a/implementations/rust/ockam/ockam_identity/src/secure_channels/secure_client.rs +++ b/implementations/rust/ockam/ockam_identity/src/secure_channels/secure_client.rs @@ -229,10 +229,9 @@ impl SecureClient { let response = client.request(ctx, req).await; let _ = self .secure_channels - .stop_secure_channel(ctx, secure_channel.encryptor_address()) - .await; + .stop_secure_channel(ctx, secure_channel.encryptor_address()); if let Some(transport_address) = transport_address { - let _ = self.transport.disconnect(transport_address).await; + let _ = self.transport.disconnect(transport_address); } // we delay the unwrapping of the response to make sure that the secure channel is // properly stopped first @@ -274,10 +273,9 @@ impl SecureClient { let (secure_channel, transport_address) = self.create_secure_channel(ctx).await?; let _ = self .secure_channels - .stop_secure_channel(ctx, secure_channel.encryptor_address()) - .await; + .stop_secure_channel(ctx, secure_channel.encryptor_address()); if let Some(transport_address) = transport_address { - let _ = self.transport.disconnect(transport_address).await; + let _ = self.transport.disconnect(transport_address); } Ok(()) diff --git a/implementations/rust/ockam/ockam_identity/tests/channel.rs b/implementations/rust/ockam/ockam_identity/tests/channel.rs index cc08e8af09f..a82c0335e6e 100644 --- a/implementations/rust/ockam/ockam_identity/tests/channel.rs +++ b/implementations/rust/ockam/ockam_identity/tests/channel.rs @@ -41,7 +41,7 @@ async fn test_channel(ctx: &mut Context) -> Result<()> { .await?; let mut child_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "child", Arc::new(AllowAll), Arc::new(AllowAll), @@ -53,7 +53,7 @@ async fn test_channel(ctx: &mut Context) -> Result<()> { child_ctx .send( - route![alice_channel.clone(), child_ctx.address()], + route![alice_channel.clone(), child_ctx.primary_address()], "Hello, Bob!".to_string(), ) .await?; @@ -246,7 +246,7 @@ async fn test_channel_rejected_trust_policy(ctx: &mut Context) -> Result<()> { .await?; let mut child_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "child", Arc::new(AllowAll), Arc::new(AllowAll), @@ -255,7 +255,7 @@ async fn test_channel_rejected_trust_policy(ctx: &mut Context) -> Result<()> { child_ctx .send( - route![alice_channel, child_ctx.address()], + route![alice_channel, child_ctx.primary_address()], "Hello, Bob!".to_string(), ) .await?; @@ -295,7 +295,7 @@ async fn test_channel_send_multiple_messages_both_directions(ctx: &mut Context) .await?; let mut child_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "child", Arc::new(AllowAll), Arc::new(AllowAll), @@ -305,11 +305,11 @@ async fn test_channel_send_multiple_messages_both_directions(ctx: &mut Context) for n in 0..50 { child_ctx .flow_controls() - .add_consumer(child_ctx.address(), &sc_listener_flow_control_id); + .add_consumer(child_ctx.primary_address(), &sc_listener_flow_control_id); let payload = format!("Hello, Bob! {}", n); child_ctx .send( - route![alice_channel.clone(), child_ctx.address()], + route![alice_channel.clone(), child_ctx.primary_address()], payload.clone(), ) .await?; @@ -320,7 +320,7 @@ async fn test_channel_send_multiple_messages_both_directions(ctx: &mut Context) child_ctx .flow_controls() - .add_consumer(child_ctx.address(), &sc_flow_control_id); + .add_consumer(child_ctx.primary_address(), &sc_flow_control_id); let payload = format!("Hello, Alice! {}", n); child_ctx.send(return_route, payload.clone()).await?; @@ -366,7 +366,7 @@ async fn test_channel_registry(ctx: &mut Context) -> Result<()> { assert_eq!(alice_channel_data.their_id(), &bob); let mut bob_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "bob", Arc::new(AllowAll), Arc::new(AllowAll), @@ -428,7 +428,7 @@ async fn test_channel_api(ctx: &mut Context) -> Result<()> { .await?; let mut bob_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "bob", Arc::new(AllowAll), Arc::new(AllowAll), @@ -555,7 +555,7 @@ async fn test_tunneled_secure_channel_works(ctx: &mut Context) -> Result<()> { .await?; let mut child_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "child", Arc::new(AllowAll), Arc::new(AllowAll), @@ -567,7 +567,7 @@ async fn test_tunneled_secure_channel_works(ctx: &mut Context) -> Result<()> { child_ctx .send( - route![alice_another_channel.clone(), child_ctx.address()], + route![alice_another_channel.clone(), child_ctx.primary_address()], "Hello, Bob!".to_string(), ) .await?; @@ -649,7 +649,7 @@ async fn test_double_tunneled_secure_channel_works(ctx: &mut Context) -> Result< .await?; let mut child_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "child", Arc::new(AllowAll), Arc::new(AllowAll), @@ -661,7 +661,10 @@ async fn test_double_tunneled_secure_channel_works(ctx: &mut Context) -> Result< child_ctx .send( - route![alice_yet_another_channel.clone(), child_ctx.address()], + route![ + alice_yet_another_channel.clone(), + child_ctx.primary_address() + ], "Hello, Bob!".to_string(), ) .await?; @@ -725,7 +728,7 @@ async fn test_many_times_tunneled_secure_channel_works(ctx: &mut Context) -> Res } let mut child_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "child", Arc::new(AllowAll), Arc::new(AllowAll), @@ -737,7 +740,10 @@ async fn test_many_times_tunneled_secure_channel_works(ctx: &mut Context) -> Res child_ctx .send( - route![channels.last().unwrap().clone(), child_ctx.address()], + route![ + channels.last().unwrap().clone(), + child_ctx.primary_address() + ], "Hello, Bob!".to_string(), ) .await?; @@ -1013,7 +1019,23 @@ async fn test_channel_delete_ephemeral_keys(ctx: &mut Context) -> Result<()> { assert_eq!(bob_sc_vault.number_of_ephemeral_x25519_secrets(), 0); assert_eq!(bob_sc_vault.number_of_static_x25519_secrets().await?, 1); - ctx.stop().await?; + let alice_sc = secure_channels_alice + .secure_channel_registry() + .get_channel_list() + .first() + .unwrap() + .clone(); + secure_channels_alice.stop_secure_channel(ctx, alice_sc.encryptor_messaging_address())?; + + let bob_sc = secure_channels_bob + .secure_channel_registry() + .get_channel_list() + .first() + .unwrap() + .clone(); + secure_channels_bob.stop_secure_channel(ctx, bob_sc.encryptor_messaging_address())?; + + ctx.sleep(Duration::from_millis(250)).await; // when the channel is closed only purpose key should be left assert_eq!(alice_identity_vault.number_of_keys().await?, 1); @@ -1069,9 +1091,7 @@ async fn should_stop_encryptor__and__decryptor__in__secure_channel( let channel2 = sc_list[1].clone(); // This will stop both ends of the channel - secure_channels - .stop_secure_channel(ctx, channel1.encryptor_messaging_address()) - .await?; + secure_channels.stop_secure_channel(ctx, channel1.encryptor_messaging_address())?; ctx.sleep(Duration::from_millis(250)).await; @@ -1083,7 +1103,7 @@ async fn should_stop_encryptor__and__decryptor__in__secure_channel( 0 ); - let workers = ctx.list_workers(); + let workers = ctx.list_workers()?; assert!(!workers.contains(channel1.decryptor_messaging_address())); assert!(!workers.contains(channel1.encryptor_messaging_address())); assert!(!workers.contains(channel2.decryptor_messaging_address())); @@ -1119,17 +1139,15 @@ async fn address_metadata__encryptor__should_be_terminal(ctx: &mut Context) -> R ) .await?; - let meta = ctx - .find_terminal_address(route!["app", sc.clone(), "test"]) - .await? - .unwrap(); + let route = route!["app", sc.clone(), "test"]; + let (address, meta) = ctx.find_terminal_address(route.iter())?.unwrap(); - assert_eq!(meta.address, sc.into()); + assert_eq!(address, &sc.into()); assert_eq!( - meta.metadata.attributes, + meta.attributes, vec![(SECURE_CHANNEL_IDENTIFIER.to_string(), hex::encode(bob.0))] ); - assert!(meta.metadata.is_terminal); + assert!(meta.is_terminal); Ok(()) } diff --git a/implementations/rust/ockam/ockam_identity/tests/common/message_flow_auth.rs b/implementations/rust/ockam/ockam_identity/tests/common/message_flow_auth.rs index 12bb645d774..117bf24b56b 100644 --- a/implementations/rust/ockam/ockam_identity/tests/common/message_flow_auth.rs +++ b/implementations/rust/ockam/ockam_identity/tests/common/message_flow_auth.rs @@ -66,7 +66,7 @@ async fn check_message_flow_with_ctx( let msg: [u8; 4] = random(); let msg = hex::encode(msg); ctx.send( - route![address.clone(), receiving_ctx.address()], + route![address.clone(), receiving_ctx.primary_address()], msg.clone(), ) .await?; diff --git a/implementations/rust/ockam/ockam_identity/tests/update_route.rs b/implementations/rust/ockam/ockam_identity/tests/update_route.rs index d74aabbd36c..e2fcdd4a024 100644 --- a/implementations/rust/ockam/ockam_identity/tests/update_route.rs +++ b/implementations/rust/ockam/ockam_identity/tests/update_route.rs @@ -21,7 +21,7 @@ async fn test_update_decryptor_route(ctx: &mut Context) -> Result<()> { .await?; let mut child_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "child", Arc::new(AllowAll), Arc::new(AllowAll), @@ -35,7 +35,7 @@ async fn test_update_decryptor_route(ctx: &mut Context) -> Result<()> { child_ctx .send( - route![alice_channel.clone(), child_ctx.address()], + route![alice_channel.clone(), child_ctx.primary_address()], "Hello, Bob!".to_string(), ) .await?; @@ -54,7 +54,7 @@ async fn test_update_decryptor_route(ctx: &mut Context) -> Result<()> { child_ctx .send( - route![alice_channel.clone(), child_ctx.address()], + route![alice_channel.clone(), child_ctx.primary_address()], "Hello, Bob!".to_string(), ) .await?; @@ -113,7 +113,7 @@ async fn test_update_decryptor_route_tcp(ctx: &mut Context) -> Result<()> { .await?; let mut child_ctx = ctx - .new_detached_with_mailboxes(Mailboxes::main( + .new_detached_with_mailboxes(Mailboxes::primary( "child", Arc::new(AllowAll), Arc::new(AllowAll), @@ -127,7 +127,7 @@ async fn test_update_decryptor_route_tcp(ctx: &mut Context) -> Result<()> { child_ctx .send( - route![alice_channel.clone(), child_ctx.address()], + route![alice_channel.clone(), child_ctx.primary_address()], "Hello, Bob!".to_string(), ) .await?; @@ -142,13 +142,13 @@ async fn test_update_decryptor_route_tcp(ctx: &mut Context) -> Result<()> { assert_eq!("Hello, Alice!", msg.into_body()?); - tcp_connection1.stop(ctx).await?; + tcp_connection1.stop(ctx)?; alice_channel.update_remote_node_route(route![tcp_connection2])?; child_ctx .send( - route![alice_channel.clone(), child_ctx.address()], + route![alice_channel.clone(), child_ctx.primary_address()], "Hello, Bob!".to_string(), ) .await?; diff --git a/implementations/rust/ockam/ockam_node/src/async_drop.rs b/implementations/rust/ockam/ockam_node/src/async_drop.rs deleted file mode 100644 index 4d90bd62de1..00000000000 --- a/implementations/rust/ockam/ockam_node/src/async_drop.rs +++ /dev/null @@ -1,48 +0,0 @@ -use crate::router::Router; -use crate::tokio::sync::oneshot::{self, Receiver, Sender}; -use alloc::sync::Arc; -use ockam_core::Address; - -/// A helper to implement Drop mechanisms, but async -/// -/// This mechanism uses a Oneshot channel, which doesn't require an -/// async context to send into it (i.e. we can send a single message -/// from a `Drop` handler without needing to block on a future!) -/// -/// The receiver is then tasked to de-allocate the specified resource. -/// -/// This is not a very generic interface, i.e. it will only generate -/// stop_worker messages. If we want to reuse this mechanism, we may -/// also want to extend the API so that other resources can specify -/// additional metadata to generate messages. -pub struct AsyncDrop { - rx: Receiver
, - router: Arc, -} - -impl AsyncDrop { - /// Create a new AsyncDrop and AsyncDrop sender - /// - /// The `sender` parameter can simply be cloned from the parent - /// Context that creates this hook, while the `address` field must - /// refer to the address of the context that will be deallocated - /// this way. - pub fn new(router: Arc) -> (Self, Sender
) { - let (tx, rx) = oneshot::channel(); - (Self { rx, router }, tx) - } - - /// Wait for the cancellation of the channel and then send a - /// message to the router - /// - /// Because this code is run detached from its original context, - /// we can't handle any errors. - pub async fn run(self) { - if let Ok(addr) = self.rx.await { - debug!("Received AsyncDrop request for address: {}", addr); - if let Err(e) = self.router.stop_worker(&addr, true).await { - debug!("Failed sending AsyncDrop request to router: {}", e); - } - } - } -} diff --git a/implementations/rust/ockam/ockam_node/src/context/context.rs b/implementations/rust/ockam/ockam_node/src/context/context.rs index da1dd88a74d..6a1ed6af223 100644 --- a/implementations/rust/ockam/ockam_node/src/context/context.rs +++ b/implementations/rust/ockam/ockam_node/src/context/context.rs @@ -1,6 +1,5 @@ use crate::channel_types::SmallReceiver; use crate::tokio::runtime::Handle; -use crate::AsyncDropSender; use core::sync::atomic::AtomicUsize; use ockam_core::compat::collections::HashMap; use ockam_core::compat::sync::{Arc, RwLock}; @@ -10,30 +9,39 @@ use ockam_core::flow_control::FlowControls; #[cfg(feature = "std")] use ockam_core::OpenTelemetryContext; use ockam_core::{ - async_trait, Address, AddressAndMetadata, AddressMetadata, Error, Mailboxes, RelayMessage, - Result, TransportType, + async_trait, Address, AddressMetadata, Error, Mailboxes, RelayMessage, Result, TransportType, }; use crate::router::Router; #[cfg(feature = "std")] use core::fmt::{Debug, Formatter}; +use ockam_core::compat::sync::Weak; use ockam_core::errcode::{Kind, Origin}; use ockam_transport_core::Transport; /// A default timeout in seconds pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30); +/// Context mode depending on the fact if it's attached to a Worker or a Processor +#[derive(Clone, Copy, Debug)] +pub enum ContextMode { + /// Without a Worker or a Processor + Detached, + /// With a Worker or a Processor + Attached, +} + /// Context contains Node state and references to the runtime. pub struct Context { pub(super) mailboxes: Mailboxes, - pub(super) router: Arc, - pub(super) rt: Handle, + pub(super) router: Weak, + pub(super) runtime_handle: Handle, pub(super) receiver: SmallReceiver, - pub(super) async_drop_sender: Option, pub(super) mailbox_count: Arc, /// List of transports used to resolve external addresses to local workers in routes pub(super) transports: Arc>>>, pub(super) flow_controls: FlowControls, + pub(super) mode: ContextMode, #[cfg(feature = "std")] pub(super) tracing_context: OpenTelemetryContext, } @@ -50,7 +58,8 @@ impl Debug for Context { fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { f.debug_struct("Context") .field("mailboxes", &self.mailboxes) - .field("runtime", &self.rt) + .field("runtime", &self.runtime_handle) + .field("mode", &self.mode) .finish() } } @@ -58,7 +67,7 @@ impl Debug for Context { impl Context { /// Return runtime clone pub fn runtime(&self) -> &Handle { - &self.rt + &self.runtime_handle } /// Return mailbox_count clone @@ -67,23 +76,25 @@ impl Context { } /// Return a reference to sender - pub(crate) fn router(&self) -> Arc { - self.router.clone() + pub(crate) fn router(&self) -> Result> { + // FIXME + self.router + .upgrade() + .ok_or_else(|| Error::new(Origin::Node, Kind::Shutdown, "Failed to upgrade router")) } - /// Return the primary address of the current worker - pub fn address(&self) -> Address { - self.mailboxes.main_address() + pub(crate) fn router_weak(&self) -> Weak { + self.router.clone() } /// Return the primary address of the current worker - pub fn address_ref(&self) -> &Address { - self.mailboxes.main_address_ref() + pub fn primary_address(&self) -> &Address { + self.mailboxes.primary_address() } - /// Return all addresses of the current worker - pub fn addresses(&self) -> Vec
{ - self.mailboxes.addresses() + /// Return additional addresses of the current worker + pub fn additional_addresses(&self) -> impl Iterator { + self.mailboxes.additional_addresses() } /// Return a reference to the mailboxes of this context @@ -125,44 +136,31 @@ impl Context { /// /// Clusters are de-allocated in reverse order of their /// initialisation when the node is stopped. - pub async fn set_cluster>(&self, label: S) -> Result<()> { - self.router.set_cluster(self.address(), label.into()) + pub fn set_cluster>(&self, label: S) -> Result<()> { + self.router()? + .set_cluster(self.primary_address(), label.into()) } /// Return a list of all available worker addresses on a node - pub fn list_workers(&self) -> Vec
{ - self.router.list_workers() + pub fn list_workers(&self) -> Result> { + Ok(self.router()?.list_workers()) } /// Return true if a worker is already registered at this address - pub fn is_worker_registered_at(&self, address: &Address) -> bool { - self.router.is_worker_registered_at(address) - } - - /// Send a shutdown acknowledgement to the router - pub(crate) async fn stop_ack(&self) -> Result<()> { - self.router.stop_ack(self.address()).await + pub fn is_worker_registered_at(&self, address: &Address) -> Result { + Ok(self.router()?.is_worker_registered_at(address)) } /// Finds the terminal address of a route, if present - pub async fn find_terminal_address( + pub fn find_terminal_address<'a>( &self, - route: impl Into>, - ) -> Result> { - let addresses = route.into(); - if addresses.iter().any(|a| !a.transport_type().is_local()) { - return Err(Error::new( - Origin::Node, - Kind::Invalid, - "Only local addresses are allowed while looking for a terminal address", - )); - } - - Ok(self.router.find_terminal_address(addresses)) + addresses: impl Iterator, + ) -> Result> { + Ok(self.router()?.find_terminal_address(addresses)) } /// Read metadata for the provided address - pub fn get_metadata(&self, address: impl Into
) -> Option { - self.router.get_address_metadata(&address.into()) + pub fn get_metadata(&self, address: impl Into
) -> Result> { + Ok(self.router()?.get_address_metadata(&address.into())) } } diff --git a/implementations/rust/ockam/ockam_node/src/context/context_lifecycle.rs b/implementations/rust/ockam/ockam_node/src/context/context_lifecycle.rs index ef694245ebd..9a46a0af64a 100644 --- a/implementations/rust/ockam/ockam_node/src/context/context_lifecycle.rs +++ b/implementations/rust/ockam/ockam_node/src/context/context_lifecycle.rs @@ -1,40 +1,49 @@ -use core::time::Duration; - #[cfg(not(feature = "std"))] use crate::tokio; +use core::time::Duration; use ockam_core::compat::collections::HashMap; +use ockam_core::compat::sync::Weak; use ockam_core::compat::time::now; use ockam_core::compat::{boxed::Box, sync::Arc, sync::RwLock}; use ockam_core::flow_control::FlowControls; #[cfg(feature = "std")] use ockam_core::OpenTelemetryContext; use ockam_core::{ - Address, AsyncTryClone, DenyAll, IncomingAccessControl, Mailboxes, OutgoingAccessControl, - Result, TransportType, + Address, AllowAll, AsyncTryClone, DenyAll, IncomingAccessControl, Mailbox, Mailboxes, + OutgoingAccessControl, Result, TransportType, }; use ockam_transport_core::Transport; -use tokio::runtime::Handle; - -use crate::async_drop::AsyncDrop; -use crate::channel_types::{message_channel, small_channel, SmallReceiver}; +use crate::channel_types::{message_channel, oneshot_channel, OneshotReceiver}; use crate::router::Router; -use crate::{debugger, Context}; +use crate::{debugger, Context, ContextMode}; use crate::{relay::CtrlSignal, router::SenderPair}; - -/// A special type of `Context` that has no worker relay and inherits -/// the parent `Context`'s access control -pub type DetachedContext = Context; - -/// A special sender type that connects a type to an AsyncDrop handler -pub type AsyncDropSender = tokio::sync::oneshot::Sender
; +use tokio::runtime::Handle; impl Drop for Context { fn drop(&mut self) { - if let Some(sender) = self.async_drop_sender.take() { - trace!("De-allocated detached context {}", self.address()); - if let Err(e) = sender.send(self.address()) { - warn!("Encountered error while dropping detached context: {}", e); + if let ContextMode::Detached = self.mode { + let router = match self.router() { + Ok(router) => router, + Err(_) => { + debug!( + "Can't upgrade router inside Context::drop: {}", + self.primary_address() + ); + return; + } + }; + + match router.stop_address(self.primary_address(), true) { + Ok(_) => {} + Err(e) => { + // That is OK during node stop + debug!( + "Encountered error while dropping detached context: {}. Err={}", + self.primary_address(), + e + ); + } } } } @@ -62,24 +71,24 @@ impl Context { /// `async_drop_sender` must be provided when creating a detached /// Context type (i.e. not backed by a worker relay). #[allow(clippy::too_many_arguments)] - pub(crate) fn new( - rt: Handle, - router: Arc, + fn new( + runtime_handle: Handle, + router: Weak, mailboxes: Mailboxes, - async_drop_sender: Option, + mode: ContextMode, transports: Arc>>>, flow_controls: &FlowControls, #[cfg(feature = "std")] tracing_context: OpenTelemetryContext, - ) -> (Self, SenderPair, SmallReceiver) { + ) -> (Self, SenderPair, OneshotReceiver) { let (mailbox_tx, receiver) = message_channel(); - let (ctrl_tx, ctrl_rx) = small_channel(); + let (ctrl_tx, ctrl_rx) = oneshot_channel(); ( Self { - rt, + runtime_handle, router, mailboxes, + mode, receiver, - async_drop_sender, mailbox_count: Arc::new(0.into()), transports, flow_controls: flow_controls.clone(), @@ -94,32 +103,40 @@ impl Context { ) } - pub(crate) fn copy_with_mailboxes( - &self, - mailboxes: Mailboxes, - ) -> (Context, SenderPair, SmallReceiver) { - Context::new( - self.runtime().clone(), - self.router(), + pub(crate) fn create_app_context( + runtime_handle: Handle, + router: Weak, + flow_controls: &FlowControls, + #[cfg(feature = "std")] tracing_context: OpenTelemetryContext, + ) -> (Self, SenderPair, OneshotReceiver) { + let addr: Address = "app".into(); + let mailboxes = Mailboxes::new( + Mailbox::new(addr, None, Arc::new(AllowAll), Arc::new(AllowAll)), + vec![], + ); + + Self::new( + runtime_handle, + router, mailboxes, - None, - self.transports.clone(), - &self.flow_controls, + ContextMode::Detached, + Default::default(), + flow_controls, #[cfg(feature = "std")] - self.tracing_context(), + tracing_context, ) } - pub(crate) fn copy_with_mailboxes_detached( + pub(crate) fn new_with_mailboxes( &self, mailboxes: Mailboxes, - drop_sender: AsyncDropSender, - ) -> (Context, SenderPair, SmallReceiver) { - Context::new( + mode: ContextMode, + ) -> (Context, SenderPair, OneshotReceiver) { + Self::new( self.runtime().clone(), - self.router(), + self.router_weak(), mailboxes, - Some(drop_sender), + mode, self.transports.clone(), &self.flow_controls, #[cfg(feature = "std")] @@ -164,10 +181,7 @@ impl Context { } /// TODO basically we can just rename `Self::new_detached_impl()` - pub async fn new_detached_with_mailboxes( - &self, - mailboxes: Mailboxes, - ) -> Result { + pub async fn new_detached_with_mailboxes(&self, mailboxes: Mailboxes) -> Result { let ctx = self.new_detached_impl(mailboxes).await?; debugger::log_inherit_context("DETACHED_WITH_MB", self, &ctx); @@ -210,8 +224,8 @@ impl Context { address: impl Into
, incoming: impl IncomingAccessControl, outgoing: impl OutgoingAccessControl, - ) -> Result { - let mailboxes = Mailboxes::main(address.into(), Arc::new(incoming), Arc::new(outgoing)); + ) -> Result { + let mailboxes = Mailboxes::primary(address.into(), Arc::new(incoming), Arc::new(outgoing)); let ctx = self.new_detached_impl(mailboxes).await?; debugger::log_inherit_context("DETACHED", self, &ctx); @@ -219,24 +233,12 @@ impl Context { Ok(ctx) } - async fn new_detached_impl(&self, mailboxes: Mailboxes) -> Result { - // A detached Context exists without a worker relay, which - // requires special shutdown handling. To allow the Drop - // handler to interact with the Node runtime, we use an - // AsyncDrop handler. - // - // This handler is spawned and listens for an event from the - // Drop handler, and then forwards a message to the Node - // router. - let (async_drop, drop_sender) = AsyncDrop::new(self.router.clone()); - self.rt.spawn(async_drop.run()); - + async fn new_detached_impl(&self, mailboxes: Mailboxes) -> Result { // Create a new context and get access to the mailbox senders - let addresses = mailboxes.addresses(); - let (ctx, sender, _) = self.copy_with_mailboxes_detached(mailboxes, drop_sender); + let (ctx, sender, _) = self.new_with_mailboxes(mailboxes, ContextMode::Detached); - self.router - .start_worker(addresses, sender, true, vec![], self.mailbox_count.clone())?; + self.router()? + .add_worker(ctx.mailboxes(), sender, true, self.mailbox_count.clone())?; Ok(ctx) } @@ -255,12 +257,11 @@ mod tests { // after a copy with new mailboxes the list of transports should be intact let mailboxes = Mailboxes::new(Mailbox::deny_all("address"), vec![]); - let (copy, _, _) = ctx.copy_with_mailboxes(mailboxes.clone()); + let (copy, _, _) = ctx.new_with_mailboxes(mailboxes.clone(), ContextMode::Attached); assert!(copy.is_transport_registered(transport.transport_type())); // after a detached copy with new mailboxes the list of transports should be intact - let (_, drop_sender) = AsyncDrop::new(ctx.router.clone()); - let (copy, _, _) = ctx.copy_with_mailboxes_detached(mailboxes, drop_sender); + let (copy, _, _) = ctx.new_with_mailboxes(mailboxes, ContextMode::Attached); assert!(copy.is_transport_registered(transport.transport_type())); Ok(()) } @@ -277,7 +278,7 @@ mod tests { Ok(address) } - async fn disconnect(&self, _address: Address) -> Result<()> { + fn disconnect(&self, _address: Address) -> Result<()> { Ok(()) } } diff --git a/implementations/rust/ockam/ockam_node/src/context/mod.rs b/implementations/rust/ockam/ockam_node/src/context/mod.rs index d37581d40d3..9356fabe43c 100644 --- a/implementations/rust/ockam/ockam_node/src/context/mod.rs +++ b/implementations/rust/ockam/ockam_node/src/context/mod.rs @@ -9,6 +9,5 @@ mod transports; mod worker_lifecycle; pub use context::*; -pub use context_lifecycle::*; pub use receive_message::*; pub use send_message::*; diff --git a/implementations/rust/ockam/ockam_node/src/context/receive_message.rs b/implementations/rust/ockam/ockam_node/src/context/receive_message.rs index 8334d0b32c9..95c7c05dac5 100644 --- a/implementations/rust/ockam/ockam_node/src/context/receive_message.rs +++ b/implementations/rust/ockam/ockam_node/src/context/receive_message.rs @@ -61,7 +61,7 @@ impl Context { pub(crate) async fn receiver_next(&mut self) -> Result> { loop { let relay_msg = if let Some(msg) = self.receiver.recv().await.map(|msg| { - trace!("{}: received new message!", self.address()); + trace!("{}: received new message!", self.primary_address()); // First we update the mailbox fill metrics self.mailbox_count.fetch_sub(1, Ordering::Acquire); diff --git a/implementations/rust/ockam/ockam_node/src/context/register_router.rs b/implementations/rust/ockam/ockam_node/src/context/register_router.rs index 65e16c61999..0d9e8c3e09d 100644 --- a/implementations/rust/ockam/ockam_node/src/context/register_router.rs +++ b/implementations/rust/ockam/ockam_node/src/context/register_router.rs @@ -5,6 +5,6 @@ impl Context { // TODO: This method should be deprecated /// Register a router for a specific address type pub fn register>(&self, type_: TransportType, addr: A) -> Result<()> { - self.router.register_router(type_, addr.into()) + self.router()?.register_router(type_, addr.into()) } } diff --git a/implementations/rust/ockam/ockam_node/src/context/send_message.rs b/implementations/rust/ockam/ockam_node/src/context/send_message.rs index 91dd00ce763..51f3d7614fe 100644 --- a/implementations/rust/ockam/ockam_node/src/context/send_message.rs +++ b/implementations/rust/ockam/ockam_node/src/context/send_message.rs @@ -88,6 +88,7 @@ impl Context { let mailboxes = Mailboxes::new( Mailbox::new( address.clone(), + None, Arc::new(AllowAll), Arc::new(AllowOnwardAddress(next.clone())), ), @@ -168,7 +169,7 @@ impl Context { R: Into, M: Message + Send + 'static, { - self.send_from_address(route.into(), msg, self.address()) + self.send_from_address(route.into(), msg, self.primary_address().clone()) .await } @@ -184,8 +185,13 @@ impl Context { R: Into, M: Message + Send + 'static, { - self.send_from_address_impl(route.into(), msg, self.address(), local_info) - .await + self.send_from_address_impl( + route.into(), + msg, + self.primary_address().clone(), + local_info, + ) + .await } /// Send a message to an address or via a fully-qualified route @@ -240,7 +246,7 @@ impl Context { } }; - let sender = self.router.resolve(&addr)?; + let sender = self.router()?.resolve(&addr)?; // Pack the payload into a TransportMessage let payload = msg.encode().map_err(|_| NodeError::Data.internal())?; @@ -265,7 +271,7 @@ impl Context { } // Pack local message into a RelayMessage wrapper - let relay_msg = RelayMessage::new(sending_address.clone(), addr, local_msg); + let relay_msg = RelayMessage::new(sending_address, addr, local_msg); debugger::log_outgoing_message(self, &relay_msg); @@ -300,7 +306,8 @@ impl Context { /// [`Context::send`]: crate::Context::send /// [`LocalMessage`]: ockam_core::LocalMessage pub async fn forward(&self, local_msg: LocalMessage) -> Result<()> { - self.forward_from_address(local_msg, self.address()).await + self.forward_from_address(local_msg, self.primary_address().clone()) + .await } /// Forward a transport message to its next routing destination @@ -337,7 +344,7 @@ impl Context { return Err(err); } }; - let sender = self.router.resolve(&addr)?; + let sender = self.router()?.resolve(&addr)?; // Pack the transport message into a RelayMessage wrapper let relay_msg = RelayMessage::new(sending_address, addr, local_msg); diff --git a/implementations/rust/ockam/ockam_node/src/context/stop_env.rs b/implementations/rust/ockam/ockam_node/src/context/stop_env.rs index e6ce046b5ba..a3d2234170d 100644 --- a/implementations/rust/ockam/ockam_node/src/context/stop_env.rs +++ b/implementations/rust/ockam/ockam_node/src/context/stop_env.rs @@ -17,6 +17,6 @@ impl Context { /// This call will hang until a safe shutdown has been completed /// or the desired timeout has been reached. pub async fn stop_timeout(&self, seconds: u8) -> Result<()> { - self.router.clone().stop_graceful(seconds).await + self.router()?.stop_graceful(seconds).await } } diff --git a/implementations/rust/ockam/ockam_node/src/context/transports.rs b/implementations/rust/ockam/ockam_node/src/context/transports.rs index 607490c9a2b..3bdb2f0a444 100644 --- a/implementations/rust/ockam/ockam_node/src/context/transports.rs +++ b/implementations/rust/ockam/ockam_node/src/context/transports.rs @@ -145,7 +145,7 @@ mod tests { Ok(Address::new(LOCAL, address.inner())) } - async fn disconnect(&self, _address: Address) -> Result<()> { + fn disconnect(&self, _address: Address) -> Result<()> { Ok(()) } } diff --git a/implementations/rust/ockam/ockam_node/src/context/worker_lifecycle.rs b/implementations/rust/ockam/ockam_node/src/context/worker_lifecycle.rs index c8cae5aa7da..150418a0887 100644 --- a/implementations/rust/ockam/ockam_node/src/context/worker_lifecycle.rs +++ b/implementations/rust/ockam/ockam_node/src/context/worker_lifecycle.rs @@ -4,20 +4,6 @@ use ockam_core::{ Address, IncomingAccessControl, OutgoingAccessControl, Processor, Result, Worker, }; -enum AddressType { - Worker, - Processor, -} - -impl AddressType { - fn str(&self) -> &'static str { - match self { - AddressType::Worker => "worker", - AddressType::Processor => "processor", - } - } -} - impl Context { /// Start a new worker instance at the given address. Default AccessControl is AllowAll /// @@ -190,51 +176,13 @@ impl Context { Ok(()) } - /// Shut down a local worker by its primary address - /// - /// Approximate flow of stopping a worker: - /// - /// 1. StopWorker message -> Router - /// 2. Get AddressRecord - /// 3. Drop sender - /// 4. WorkerRelay calls Worker::shutdown - /// 5. StopAck message -> Router (from main_address) - /// 6. router.map.free_address(main_address) is called (given Router state is running): - /// remote main_address from router.map.stopping (it's not their anyway, unless in was a cluster and node was shutting down) - /// Remove AddressRecord from router.map.address_records_map (return error if not found) - /// Remove all alias in router.map.alias_map - /// Remote all meta from router.map.address_metadata - pub async fn stop_worker>(&self, addr: A) -> Result<()> { - self.stop_address(addr.into(), AddressType::Worker).await - } - - /// Shut down a local processor by its address - /// - /// Approximate flow of stopping a processor: - /// - /// 1. StopProcessor message -> Router - /// 2. Get AddressRecord - /// 3. Call AddressRecord::stop: - /// Send CtrlSignal::InterruptStop to AddressRecord::ctrl_tx - /// Set AddressRecord::state = AddressState::Stopping - /// 4. ProcessorRelay calls Processor::shutdown - /// 5. StopAck message -> Router (from main_address) - /// 6. router.map.free_address(main_address) is called (given Router state is running): - /// remote main_address from router.map.stopping (it's not their anyways unless in was a cluster and node was shutting down) - /// Remove AddressRecord from router.map.address_records_map (return error if not found) - /// Remove all alias in router.map.alias_map - /// Remote all meta from router.map.address_metadata - pub async fn stop_processor>(&self, addr: A) -> Result<()> { - self.stop_address(addr.into(), AddressType::Processor).await + /// Stop a Worker or a Processor running on given Address + pub fn stop_address(&self, addr: impl Into
) -> Result<()> { + self.stop_address_impl(&addr.into()) } - async fn stop_address(&self, addr: Address, t: AddressType) -> Result<()> { - debug!("Shutting down {} {}", t.str(), addr); - - // Send the stop request - match t { - AddressType::Worker => self.router.stop_worker(&addr, false).await, - AddressType::Processor => self.router.stop_processor(&addr).await, - } + /// Stop a Worker or a Processor running on given Address + pub fn stop_address_impl(&self, address: &Address) -> Result<()> { + self.router()?.stop_address(address, false) } } diff --git a/implementations/rust/ockam/ockam_node/src/debugger.rs b/implementations/rust/ockam/ockam_node/src/debugger.rs index 65315fde17f..21936be9d41 100644 --- a/implementations/rust/ockam/ockam_node/src/debugger.rs +++ b/implementations/rust/ockam/ockam_node/src/debugger.rs @@ -8,7 +8,6 @@ use ockam_core::{Address, Mailbox, Mailboxes}; #[cfg(feature = "debugger")] use ockam_core::compat::{ - collections::BTreeMap, sync::{Arc, RwLock}, vec::Vec, }; @@ -23,13 +22,13 @@ use core::{ #[derive(Default)] struct Debugger { /// Map context inheritance from parent main `Mailbox` to child [`Mailboxes`] - inherited_mb: Arc>>>, + inherited_mb: Arc>>>, /// Map message destination to source - incoming: Arc>>>, + incoming: Arc>>>, /// Map message destination `Mailbox` to source [`Mailbox`] - incoming_mb: Arc>>>, + incoming_mb: Arc>>>, /// Map message source to destinations - outgoing: Arc>>>, + outgoing: Arc>>>, } /// Return a mutable reference to the global debugger instance @@ -88,9 +87,9 @@ pub fn log_incoming_message(_receiving_ctx: &Context, _relay_msg: &RelayMessage) tracing::trace!( "log_incoming_message #{:03}: {} -> {} ({})", COUNTER.fetch_add(1, Ordering::Relaxed), - _relay_msg.source(), // sending address - _relay_msg.destination(), // receiving address - _receiving_ctx.address(), // actual receiving context address + _relay_msg.source(), // sending address + _relay_msg.destination(), // receiving address + _receiving_ctx.primary_address(), // actual receiving context address ); match instance().incoming.write() { @@ -137,9 +136,9 @@ pub fn log_outgoing_message(_sending_ctx: &Context, _relay_msg: &RelayMessage) { tracing::trace!( "log_outgoing_message #{:03}: {} ({}) -> {}", COUNTER.fetch_add(1, Ordering::Relaxed), - _relay_msg.source(), // sending address - _sending_ctx.address(), // actual sending context address - _relay_msg.destination(), // receiving address + _relay_msg.source(), // sending address + _sending_ctx.primary_address(), // actual sending context address + _relay_msg.destination(), // receiving address ); match instance().outgoing.write() { @@ -187,7 +186,7 @@ pub fn log_inherit_context(_tag: &str, _parent: &Context, _child: &Context) { match instance().inherited_mb.write() { Ok(mut inherited_mb) => { - let parent = _parent.mailboxes().main_mailbox().clone(); + let parent = _parent.mailboxes().primary_mailbox().clone(); let children = _child.mailboxes().clone(); inherited_mb .entry(parent) @@ -250,8 +249,8 @@ pub fn generate_graphs(w: &mut BufWriter) -> io::Result<()> { } // generate mailboxes set - use ockam_core::compat::collections::BTreeSet; - let mut mailboxes = BTreeSet::new(); + use ockam_core::compat::collections::HashSet; + let mut mailboxes = HashSet::new(); if let Ok(inherited_mb) = instance().inherited_mb.read() { for (parent, children) in inherited_mb.iter() { for child in children.iter() { diff --git a/implementations/rust/ockam/ockam_node/src/delayed.rs b/implementations/rust/ockam/ockam_node/src/delayed.rs index 9899d9024a8..01bece87a98 100644 --- a/implementations/rust/ockam/ockam_node/src/delayed.rs +++ b/implementations/rust/ockam/ockam_node/src/delayed.rs @@ -28,7 +28,7 @@ impl DelayedEvent { msg: M, ) -> Result { let destination_addr = destination_addr.into(); - let mailboxes = Mailboxes::main( + let mailboxes = Mailboxes::primary( Address::random_tagged("DelayedEvent.create"), Arc::new(DenyAll), Arc::new(AllowOnwardAddress(destination_addr.clone())), @@ -46,8 +46,8 @@ impl DelayedEvent { } /// Address used to send messages to destination address - pub fn address(&self) -> Address { - self.ctx.address() + pub fn address(&self) -> &Address { + self.ctx.primary_address() } } @@ -60,7 +60,7 @@ impl DelayedEvent { } /// Schedule heartbeat. Cancels already scheduled heartbeat if there is such heartbeat - pub async fn schedule(&mut self, duration: Duration) -> Result<()> { + pub fn schedule(&mut self, duration: Duration) -> Result<()> { self.cancel(); let destination_addr = self.destination_addr.clone(); @@ -136,11 +136,11 @@ mod tests { ctx.start_worker("counting_worker", worker).await?; - heartbeat.schedule(Duration::from_millis(100)).await?; + heartbeat.schedule(Duration::from_millis(100))?; sleep(Duration::from_millis(150)).await; - heartbeat.schedule(Duration::from_millis(100)).await?; + heartbeat.schedule(Duration::from_millis(100))?; sleep(Duration::from_millis(150)).await; - heartbeat.schedule(Duration::from_millis(100)).await?; + heartbeat.schedule(Duration::from_millis(100))?; sleep(Duration::from_millis(150)).await; assert_eq!(3, msgs_count.load(Ordering::Relaxed)); @@ -160,9 +160,9 @@ mod tests { ctx.start_worker("counting_worker", worker).await?; - heartbeat.schedule(Duration::from_millis(100)).await?; - heartbeat.schedule(Duration::from_millis(100)).await?; - heartbeat.schedule(Duration::from_millis(100)).await?; + heartbeat.schedule(Duration::from_millis(100))?; + heartbeat.schedule(Duration::from_millis(100))?; + heartbeat.schedule(Duration::from_millis(100))?; sleep(Duration::from_millis(150)).await; assert_eq!(1, msgs_count.load(Ordering::Relaxed)); @@ -182,9 +182,9 @@ mod tests { ctx.start_worker("counting_worker", worker).await?; - heartbeat.schedule(Duration::from_millis(100)).await?; + heartbeat.schedule(Duration::from_millis(100))?; sleep(Duration::from_millis(150)).await; - heartbeat.schedule(Duration::from_millis(200)).await?; + heartbeat.schedule(Duration::from_millis(200))?; sleep(Duration::from_millis(100)).await; heartbeat.cancel(); sleep(Duration::from_millis(300)).await; @@ -206,9 +206,9 @@ mod tests { ctx.start_worker("counting_worker", worker).await?; - heartbeat.schedule(Duration::from_millis(100)).await?; + heartbeat.schedule(Duration::from_millis(100))?; sleep(Duration::from_millis(150)).await; - heartbeat.schedule(Duration::from_millis(200)).await?; + heartbeat.schedule(Duration::from_millis(200))?; sleep(Duration::from_millis(100)).await; drop(heartbeat); sleep(Duration::from_millis(300)).await; diff --git a/implementations/rust/ockam/ockam_node/src/executor.rs b/implementations/rust/ockam/ockam_node/src/executor.rs index 510a1221cc9..d83ed5b8f64 100644 --- a/implementations/rust/ockam/ockam_node/src/executor.rs +++ b/implementations/rust/ockam/ockam_node/src/executor.rs @@ -1,11 +1,9 @@ -#[cfg(feature = "std")] -use crate::runtime; -use crate::{ - router::{Router, SenderPair}, - tokio::runtime::Runtime, -}; +use crate::{router::Router, tokio::runtime::Runtime}; use core::future::Future; -use ockam_core::{compat::sync::Arc, Address, Result}; +use ockam_core::{ + compat::sync::{Arc, Weak}, + Result, +}; #[cfg(feature = "metrics")] use crate::metrics::Metrics; @@ -25,6 +23,13 @@ use ockam_core::{ Error, }; +// TODO: Either make this the only one place we create runtime, or vice verse: remove it from here +#[cfg(feature = "std")] +pub(crate) static RUNTIME: once_cell::sync::Lazy>> = + once_cell::sync::Lazy::new(|| { + ockam_core::compat::sync::Mutex::new(Some(Runtime::new().unwrap())) + }); + /// Underlying Ockam node executor /// /// This type is a small wrapper around an inner async runtime (`tokio` by @@ -32,7 +37,7 @@ use ockam_core::{ /// `ockam::node` function annotation instead! pub struct Executor { /// Reference to the runtime needed to spawn tasks - rt: Arc, + runtime: Arc, /// Application router router: Arc, /// Metrics collection endpoint @@ -42,38 +47,25 @@ pub struct Executor { impl Executor { /// Create a new Ockam node [`Executor`] instance - pub fn new(rt: Arc, flow_controls: &FlowControls) -> Self { - let router = Arc::new(Router::new(flow_controls)); + pub fn new(runtime: Arc, flow_controls: &FlowControls) -> Self { + let router = Arc::new(Router::new(runtime.handle().clone(), flow_controls)); #[cfg(feature = "metrics")] - let metrics = Metrics::new(&rt, router.get_metrics_readout()); + let metrics = Metrics::new(runtime.handle().clone(), router.get_metrics_readout()); Self { - rt, + runtime, router, #[cfg(feature = "metrics")] metrics, } } - /// Get access to the internal message sender - pub(crate) fn router(&self) -> Arc { - self.router.clone() - } - - /// Initialize the root application worker - pub(crate) fn initialize_system>( - &self, - address: S, - senders: SenderPair, - ) -> Result<()> { - trace!("Initializing node executor"); - self.router.init(address.into(), senders) + /// Get access to the Router + pub(crate) fn router(&self) -> Weak { + Arc::downgrade(&self.router) } /// Initialise and run the Ockam node executor context /// - /// In this background this launches async execution of the Ockam - /// router, while blocking execution on the provided future. - /// /// Any errors encountered by the router or provided application /// code will be returned from this function. #[cfg(feature = "std")] @@ -87,16 +79,11 @@ impl Executor { #[cfg(feature = "metrics")] let alive = Arc::new(AtomicBool::from(true)); #[cfg(feature = "metrics")] - self.rt.spawn( - self.metrics - .clone() - .run(alive.clone()) - .with_current_context(), - ); + self.metrics.clone().spawn(alive.clone()); // Spawn user code second let future = Executor::wrapper(self.router.clone(), future); - let join_body = self.rt.spawn(future.with_current_context()); + let join_body = self.runtime.spawn(future.with_current_context()); // Shut down metrics collector #[cfg(feature = "metrics")] @@ -104,51 +91,12 @@ impl Executor { // Last join user code let res = self - .rt + .runtime .block_on(join_body) .map_err(|e| Error::new(Origin::Executor, Kind::Unknown, e))?; - Ok(res) - } - - /// Initialise and run the Ockam node executor context - /// - /// In this background this launches async execution of the Ockam - /// router, while blocking execution on the provided future. - /// - /// Any errors encountered by the router or provided application - /// code will be returned from this function. - /// - /// Don't abort the router in case of a failure - #[cfg(feature = "std")] - pub fn execute_no_abort(&mut self, future: F) -> Result - where - F: Future + Send + 'static, - T: Send + 'static, - { - // Spawn the metrics collector first - #[cfg(feature = "metrics")] - let alive = Arc::new(AtomicBool::from(true)); - #[cfg(feature = "metrics")] - self.rt.spawn( - self.metrics - .clone() - .run(alive.clone()) - .with_current_context(), - ); - - // Spawn user code second - let join_body = self.rt.spawn(future.with_current_context()); - - // Shut down metrics collector - #[cfg(feature = "metrics")] - alive.fetch_or(true, Ordering::Acquire); - - // Last join user code - let res = self - .rt - .block_on(join_body) - .map_err(|e| Error::new(Origin::Executor, Kind::Unknown, e))?; + // TODO: Shutdown Runtime if we exclusively own it. Which should be always except when we + // run multiple nodes inside the same process Ok(res) } @@ -161,17 +109,11 @@ impl Executor { { match future.await { Ok(val) => { + // TODO: Add timeout here router.wait_termination().await; Ok(val) } Err(e) => { - // We earlier sent the AbortNode message to the router here. - // It failed because the router state was not set to `Stopping` - // But sending Graceful shutdown message works because, it internally does that. - // - // I think way AbortNode is implemented right now, it is more of an - // internal/private message not meant to be directly used, without changing the - // router state. if let Err(error) = router.stop_graceful(1).await { error!("Failed to stop gracefully: {}", error); } @@ -189,7 +131,7 @@ impl Executor { F: Future + Send + 'static, F::Output: Send + 'static, { - let lock = runtime::RUNTIME.lock().unwrap(); + let lock = RUNTIME.lock().unwrap(); let rt = lock.as_ref().expect("Runtime was consumed"); let join_body = rt.spawn(future.with_current_context()); rt.block_on(join_body.with_current_context()) @@ -210,12 +152,12 @@ impl Executor { F: Future + Send + 'static, F::Output: Send + 'static, { - let _join = self.rt.spawn(future); + let _join = self.runtime.spawn(future); let router = self.router.clone(); // Block this task executing the primary message router, // returning any critical failures that it encounters. - crate::tokio::runtime::execute(&self.rt, async move { + crate::tokio::runtime::execute(&self.runtime, async move { router.wait_termination().await; }); Ok(()) diff --git a/implementations/rust/ockam/ockam_node/src/lib.rs b/implementations/rust/ockam/ockam_node/src/lib.rs index 06dcfe71497..c76261909e4 100644 --- a/implementations/rust/ockam/ockam_node/src/lib.rs +++ b/implementations/rust/ockam/ockam_node/src/lib.rs @@ -52,7 +52,6 @@ pub mod callback; /// Helper workers pub mod workers; -mod async_drop; mod context; mod delayed; mod error; @@ -68,10 +67,6 @@ pub mod storage; mod worker_builder; -/// Singleton for the runtime executor -#[cfg(feature = "std")] -pub mod runtime; - #[cfg(feature = "watchdog")] mod watchdog; diff --git a/implementations/rust/ockam/ockam_node/src/metrics.rs b/implementations/rust/ockam/ockam_node/src/metrics.rs index 0d8d1a1e4e6..00e2d4a520d 100644 --- a/implementations/rust/ockam/ockam_node/src/metrics.rs +++ b/implementations/rust/ockam/ockam_node/src/metrics.rs @@ -3,27 +3,34 @@ use core::{ sync::atomic::{AtomicBool, AtomicUsize, Ordering}, time::Duration, }; -use ockam_core::compat::{collections::BTreeMap, sync::Arc}; +use ockam_core::compat::{collections::HashMap, sync::Arc}; use ockam_core::env::get_env; +use opentelemetry::trace::FutureExt; use std::{fs::OpenOptions, io::Write}; +use tokio::runtime::Handle; pub struct Metrics { - rt: Arc, + runtime_handle: Handle, router: (Arc, Arc), } impl Metrics { /// Create a new Metrics collector with access to the runtime pub(crate) fn new( - rt: &Arc, + runtime_handle: Handle, router: (Arc, Arc), ) -> Arc { Arc::new(Self { - rt: Arc::clone(rt), + runtime_handle, router, }) } + pub(crate) fn spawn(self: Arc, alive: Arc) { + self.runtime_handle + .spawn(self.run(alive).with_current_context()); + } + /// Spawned by the Executor to periodically collect metrics pub(crate) async fn run(self: Arc, alive: Arc) { let path = match get_env::("OCKAM_METRICS_PATH") { @@ -66,13 +73,13 @@ impl Metrics { freq: u64, acc: &mut MetricsReport, ) -> MetricsReport { - let m = self.rt.metrics(); + let m = self.runtime_handle.metrics(); let tokio_workers = m.num_workers(); let router_addr_count = self.router.0.load(Ordering::Acquire); let router_cluster_count = self.router.1.load(Ordering::Acquire); - let mut tokio_busy_ms = BTreeMap::new(); + let mut tokio_busy_ms = HashMap::new(); for wid in 0..tokio_workers { // Get the previously accumulated let acc_ms = acc.tokio_busy_ms.get(&wid).unwrap_or(&0); @@ -96,7 +103,7 @@ impl Metrics { #[derive(Default)] #[allow(unused)] pub struct MetricsReport { - tokio_busy_ms: BTreeMap, + tokio_busy_ms: HashMap, router_addr_count: usize, router_cluster_count: usize, } diff --git a/implementations/rust/ockam/ockam_node/src/node.rs b/implementations/rust/ockam/ockam_node/src/node.rs index 83ac54103e4..48548fdf86f 100644 --- a/implementations/rust/ockam/ockam_node/src/node.rs +++ b/implementations/rust/ockam/ockam_node/src/node.rs @@ -4,7 +4,6 @@ use ockam_core::compat::sync::Arc; use ockam_core::flow_control::FlowControls; #[cfg(feature = "std")] use ockam_core::OpenTelemetryContext; -use ockam_core::{Address, AllowAll, Mailbox, Mailboxes}; /// A minimal worker implementation that does nothing pub struct NullWorker; @@ -116,8 +115,6 @@ impl NodeBuilder { #[cfg(not(feature = "std"))] Arc::new(Runtime::new().expect("cannot initialize the tokio runtime")) }); - let exe = Executor::new(rt.clone(), &flow_controls); - let addr: Address = "app".into(); #[cfg(feature = "watchdog")] { @@ -125,17 +122,17 @@ impl NodeBuilder { watchdog.start_watchdog_loop(&rt); } + let handle = rt.handle().clone(); + let exe = Executor::new(rt, &flow_controls); + + let router = exe.router().upgrade().unwrap(); + // The root application worker needs a mailbox and relay to accept // messages from workers, and to buffer incoming transcoded data. - let (ctx, sender, _) = Context::new( - rt.handle().clone(), - exe.router(), - Mailboxes::new( - Mailbox::new(addr, Arc::new(AllowAll), Arc::new(AllowAll)), - vec![], - ), - None, - Default::default(), + + let (ctx, sender, _) = Context::create_app_context( + handle.clone(), + Arc::downgrade(&router), &flow_controls, #[cfg(feature = "std")] OpenTelemetryContext::current(), @@ -144,7 +141,8 @@ impl NodeBuilder { debugger::log_inherit_context("NODE", &ctx, &ctx); // Register this mailbox handle with the executor - exe.initialize_system("app", sender) + router + .add_worker(ctx.mailboxes(), sender, true, ctx.mailbox_count()) .expect("router initialization failed"); // Then return the root context and executor diff --git a/implementations/rust/ockam/ockam_node/src/processor_builder.rs b/implementations/rust/ockam/ockam_node/src/processor_builder.rs index 4edc3232cc4..5d1198946e5 100644 --- a/implementations/rust/ockam/ockam_node/src/processor_builder.rs +++ b/implementations/rust/ockam/ockam_node/src/processor_builder.rs @@ -1,9 +1,9 @@ -use crate::debugger; +use crate::{debugger, ContextMode}; use crate::{relay::ProcessorRelay, Context}; -use alloc::string::String; -use ockam_core::compat::{sync::Arc, vec::Vec}; +use ockam_core::compat::string::String; +use ockam_core::compat::sync::Arc; use ockam_core::{ - Address, AddressAndMetadata, AddressMetadata, DenyAll, IncomingAccessControl, Mailboxes, + Address, AddressMetadata, DenyAll, IncomingAccessControl, Mailbox, Mailboxes, OutgoingAccessControl, Processor, Result, }; @@ -32,14 +32,46 @@ impl

ProcessorBuilder

where P: Processor, { - /// Worker with only one [`Address`] + /// Processor with only one [`Address`] pub fn with_address(self, address: impl Into

) -> ProcessorBuilderOneAddress

{ + self.with_address_and_metadata_impl(address, None) + } + + /// Processor with single terminal [`Address`] + pub fn with_terminal_address( + self, + address: impl Into

, + ) -> ProcessorBuilderOneAddress

{ + self.with_address_and_metadata( + address, + AddressMetadata { + is_terminal: true, + attributes: vec![], + }, + ) + } + + /// Processor with single terminal [`Address`] and metadata + pub fn with_address_and_metadata( + self, + address: impl Into

, + metadata: AddressMetadata, + ) -> ProcessorBuilderOneAddress

{ + self.with_address_and_metadata_impl(address, Some(metadata)) + } + + /// Processor with single terminal [`Address`] and metadata + pub fn with_address_and_metadata_impl( + self, + address: impl Into

, + metadata: Option, + ) -> ProcessorBuilderOneAddress

{ ProcessorBuilderOneAddress { incoming_ac: Arc::new(DenyAll), outgoing_ac: Arc::new(DenyAll), processor: self.processor, address: address.into(), - metadata: None, + metadata, } } @@ -48,7 +80,6 @@ where ProcessorBuilderMultipleAddresses { mailboxes, processor: self.processor, - metadata_list: vec![], } } } @@ -59,72 +90,15 @@ where { mailboxes: Mailboxes, processor: P, - metadata_list: Vec, } impl

ProcessorBuilderMultipleAddresses

where P: Processor, { - /// Mark the provided address as terminal - pub fn terminal(self, address: impl Into

) -> Self { - self.terminal_with_attributes(address.into(), vec![]) - } - - /// Mark the provided address as terminal - pub fn terminal_with_attributes( - mut self, - address: impl Into
, - attributes: Vec<(String, String)>, - ) -> Self { - let address = address.into(); - let metadata = self.metadata_list.iter_mut().find(|m| m.address == address); - - if let Some(metadata) = metadata { - metadata.metadata.is_terminal = true; - metadata.metadata.attributes = attributes; - } else { - self.metadata_list.push(AddressAndMetadata { - address, - metadata: AddressMetadata { - is_terminal: true, - attributes, - }, - }); - } - self - } - - /// Adds metadata attribute for the provided address - pub fn with_metadata_attribute( - mut self, - address: impl Into
, - key: impl Into, - value: impl Into, - ) -> Self { - let address = address.into(); - let metadata = self.metadata_list.iter_mut().find(|m| m.address == address); - - if let Some(metadata) = metadata { - metadata - .metadata - .attributes - .push((key.into(), value.into())); - } else { - self.metadata_list.push(AddressAndMetadata { - address, - metadata: AddressMetadata { - is_terminal: false, - attributes: vec![(key.into(), value.into())], - }, - }); - } - self - } - /// Consume this builder and start a new Ockam [`Processor`] from the given context pub async fn start(self, context: &Context) -> Result<()> { - start(context, self.mailboxes, self.processor, self.metadata_list).await + start(context, self.mailboxes, self.processor).await } } @@ -136,55 +110,38 @@ where outgoing_ac: Arc, address: Address, processor: P, - metadata: Option, + metadata: Option, } impl

ProcessorBuilderOneAddress

where P: Processor, { - /// Mark the address as terminal - pub fn terminal(self) -> Self { - self.terminal_with_attributes(vec![]) - } - - /// Mark the address as terminal - pub fn terminal_with_attributes(mut self, attributes: Vec<(String, String)>) -> Self { - if let Some(metadata) = self.metadata.as_mut() { - metadata.metadata.is_terminal = true; - metadata.metadata.attributes = attributes; - } else { - self.metadata = Some(AddressAndMetadata { - address: self.address.clone(), - metadata: AddressMetadata { - is_terminal: true, - attributes, - }, - }); - } + /// Mark the provided address as terminal + pub fn terminal(mut self) -> Self { + self.metadata + .get_or_insert(AddressMetadata { + is_terminal: false, + attributes: vec![], + }) + .is_terminal = true; self } - /// Adds metadata attribute + /// Adds metadata attribute for the provided address pub fn with_metadata_attribute( mut self, key: impl Into, value: impl Into, ) -> Self { - if let Some(metadata) = self.metadata.as_mut() { - metadata - .metadata - .attributes - .push((key.into(), value.into())); - } else { - self.metadata = Some(AddressAndMetadata { - address: self.address.clone(), - metadata: AddressMetadata { - is_terminal: false, - attributes: vec![(key.into(), value.into())], - }, - }); - } + self.metadata + .get_or_insert(AddressMetadata { + is_terminal: false, + attributes: vec![], + }) + .attributes + .push((key.into(), value.into())); + self } @@ -192,9 +149,16 @@ where pub async fn start(self, context: &Context) -> Result<()> { start( context, - Mailboxes::main(self.address, self.incoming_ac, self.outgoing_ac), + Mailboxes::new( + Mailbox::new( + self.address, + self.metadata, + self.incoming_ac, + self.outgoing_ac, + ), + vec![], + ), self.processor, - self.metadata.map(|m| vec![m]).unwrap_or_default(), ) .await } @@ -242,32 +206,24 @@ where } /// Consume this builder and start a new Ockam [`Processor`] from the given context - -pub async fn start

( - context: &Context, - mailboxes: Mailboxes, - processor: P, - metadata: Vec, -) -> Result<()> +pub async fn start

(context: &Context, mailboxes: Mailboxes, processor: P) -> Result<()> where P: Processor, { debug!( "Initializing ockam processor '{}' with access control in:{:?} out:{:?}", - mailboxes.main_address(), - mailboxes.main_mailbox().incoming_access_control(), - mailboxes.main_mailbox().outgoing_access_control(), + mailboxes.primary_address(), + mailboxes.primary_mailbox().incoming_access_control(), + mailboxes.primary_mailbox().outgoing_access_control(), ); - let addresses = mailboxes.addresses(); - // Pass it to the context - let (ctx, sender, ctrl_rx) = context.copy_with_mailboxes(mailboxes); + let (ctx, sender, ctrl_rx) = context.new_with_mailboxes(mailboxes, ContextMode::Attached); debugger::log_inherit_context("PROCESSOR", context, &ctx); - let router = context.router(); - router.start_processor(addresses, sender, metadata)?; + let router = context.router()?; + router.add_processor(ctx.mailboxes(), sender)?; // Then initialise the processor message relay ProcessorRelay::

::init(context.runtime(), processor, ctx, ctrl_rx); diff --git a/implementations/rust/ockam/ockam_node/src/relay/processor_relay.rs b/implementations/rust/ockam/ockam_node/src/relay/processor_relay.rs index 1a4e200a753..8169af1046c 100644 --- a/implementations/rust/ockam/ockam_node/src/relay/processor_relay.rs +++ b/implementations/rust/ockam/ockam_node/src/relay/processor_relay.rs @@ -1,4 +1,4 @@ -use crate::channel_types::SmallReceiver; +use crate::channel_types::OneshotReceiver; use crate::{relay::CtrlSignal, tokio::runtime::Handle, Context}; use ockam_core::{Processor, Result}; @@ -20,20 +20,19 @@ where #[cfg_attr(not(feature = "std"), allow(unused_mut))] #[cfg_attr(not(feature = "std"), allow(unused_variables))] - async fn run(self, mut ctrl_rx: SmallReceiver) { + async fn run(self, ctrl_rx: OneshotReceiver) { let mut ctx = self.ctx; let mut processor = self.processor; - let ctx_addr = ctx.address(); match processor.initialize(&mut ctx).await { Ok(()) => {} Err(e) => { error!( "Failure during '{}' processor initialisation: {}", - ctx.address(), + ctx.primary_address(), e ); - shutdown_and_stop_ack(&mut processor, &mut ctx, &ctx_addr).await; + shutdown_and_stop_ack(&mut processor, &mut ctx, false).await; return; } } @@ -51,10 +50,15 @@ where #[cfg(feature = "debugger")] error!( "Error encountered during '{}' processing: {:?}", - ctx_addr, e + ctx.primary_address(), + e ); #[cfg(not(feature = "debugger"))] - error!("Error encountered during '{}' processing: {}", ctx_addr, e); + error!( + "Error encountered during '{}' processing: {}", + ctx.primary_address(), + e + ); } } } @@ -62,15 +66,15 @@ where Result::<()>::Ok(()) }; + let mut stopped_from_router = false; #[cfg(feature = "std")] { - // This future resolves when a stop control signal is received - let shutdown_signal = async { ctrl_rx.recv().await }; - - // Then select over the two futures + // Select over the two futures tokio::select! { - _ = shutdown_signal => { - debug!("Shutting down processor {} due to shutdown signal", ctx_addr); + // This future resolves when a stop control signal is received + _ = ctrl_rx => { + debug!("Shutting down processor {} due to shutdown signal", ctx.primary_address()); + stopped_from_router = true; }, _ = run_loop => {} }; @@ -79,12 +83,12 @@ where // TODO wait on run_loop until we have a no_std select! implementation #[cfg(not(feature = "std"))] match run_loop.await { - Ok(_) => trace!("Processor shut down cleanly {}", ctx_addr), + Ok(_) => trace!("Processor shut down cleanly {}", ctx.primary_address()), Err(err) => error!("processor run loop aborted with error: {:?}", err), }; // If we reach this point the router has signaled us to shut down - shutdown_and_stop_ack(&mut processor, &mut ctx, &ctx_addr).await; + shutdown_and_stop_ack(&mut processor, &mut ctx, stopped_from_router).await; } /// Create a processor relay with two node contexts @@ -92,30 +96,56 @@ where rt: &Handle, processor: P, ctx: Context, - ctrl_rx: SmallReceiver, + ctrl_rx: OneshotReceiver, ) { let relay = ProcessorRelay::

::new(processor, ctx); rt.spawn(relay.run(ctrl_rx)); } } -async fn shutdown_and_stop_ack

( - processor: &mut P, - ctx: &mut Context, - ctx_addr: &ockam_core::Address, -) where +async fn shutdown_and_stop_ack

(processor: &mut P, ctx: &mut Context, stopped_from_router: bool) +where P: Processor, { match processor.shutdown(ctx).await { Ok(()) => {} Err(e) => { - error!("Failure during '{}' processor shutdown: {}", ctx_addr, e); + error!( + "Failure during '{}' processor shutdown: {}", + ctx.primary_address(), + e + ); + } + } + + let router = match ctx.router() { + Ok(router) => router, + Err(_) => { + error!( + "Failure during '{}' processor shutdown. Can't get router", + ctx.primary_address() + ); + return; + } + }; + + if !stopped_from_router { + if let Err(e) = router.stop_address(ctx.primary_address(), !stopped_from_router) { + error!( + "Failure during '{}' processor shutdown: {}", + ctx.primary_address(), + e + ); } } // Finally send the router a stop ACK -- log errors trace!("Sending shutdown ACK"); - ctx.stop_ack().await.unwrap_or_else(|e| { - error!("Failed to send stop ACK for '{}': {}", ctx_addr, e); + router.stop_ack(ctx.primary_address()).unwrap_or_else(|e| { + error!( + "Failed to send stop ACK for '{}': {}", + ctx.primary_address(), + e + ); }); } diff --git a/implementations/rust/ockam/ockam_node/src/relay/worker_relay.rs b/implementations/rust/ockam/ockam_node/src/relay/worker_relay.rs index 0d2aaa3adbb..7bdda8d3b71 100644 --- a/implementations/rust/ockam/ockam_node/src/relay/worker_relay.rs +++ b/implementations/rust/ockam/ockam_node/src/relay/worker_relay.rs @@ -1,4 +1,4 @@ -use crate::channel_types::SmallReceiver; +use crate::channel_types::OneshotReceiver; use crate::relay::CtrlSignal; use crate::tokio::runtime::Handle; use crate::Context; @@ -60,7 +60,7 @@ where let relay_msg = match self.ctx.receiver_next().await? { Some(msg) => msg, None => { - trace!("No more messages for worker {}", self.ctx.address()); + trace!("No more messages for worker {}", self.ctx.primary_address()); return Ok(false); } }; @@ -95,22 +95,20 @@ where #[cfg_attr(not(feature = "std"), allow(unused_mut))] #[cfg_attr(not(feature = "std"), allow(unused_variables))] - async fn run(mut self, mut ctrl_rx: SmallReceiver) { + async fn run(mut self, mut ctrl_rx: OneshotReceiver) { match self.worker.initialize(&mut self.ctx).await { Ok(()) => {} Err(e) => { error!( "Failure during '{}' worker initialisation: {}", - self.ctx.address(), + self.ctx.primary_address(), e ); - self.shutdown_and_stop_ack().await; + shutdown_and_stop_ack(&mut self.worker, &mut self.ctx, false).await; return; } } - let address = self.ctx.address(); - #[cfg(feature = "std")] loop { crate::tokio::select! { @@ -125,17 +123,15 @@ where // An error occurred -- log and continue Err(e) => { #[cfg(feature = "debugger")] - error!("Error encountered during '{}' message handling: {:?}", address, e); + error!("Error encountered during '{}' message handling: {:?}", self.ctx.primary_address(), e); #[cfg(not(feature = "debugger"))] - error!("Error encountered during '{}' message handling: {}", address, e); + error!("Error encountered during '{}' message handling: {}", self.ctx.primary_address(), e); } } }, - result = ctrl_rx.recv() => { - if result.is_some() { - debug!("Relay received shutdown signal, terminating!"); - break; - } + _ = &mut ctrl_rx => { + debug!(primary_address=%self.ctx.primary_address(), "Relay received shutdown signal, terminating!"); + break; // We are stopping } @@ -153,41 +149,66 @@ where // An error occurred -- log and continue Err(e) => error!( "Error encountered during '{}' message handling: {}", - address, e + self.ctx.primary_address(), + e ), } } - self.shutdown_and_stop_ack().await; + shutdown_and_stop_ack(&mut self.worker, &mut self.ctx, true).await; } - async fn shutdown_and_stop_ack(&mut self) { - // Run the shutdown hook for this worker - match self.worker.shutdown(&mut self.ctx).await { - Ok(()) => {} - Err(e) => { - error!( - "Failure during '{}' worker shutdown: {}", - self.ctx.address(), - e - ); - } - } + /// Build and spawn a new worker relay, returning a send handle to it + pub(crate) fn init(rt: &Handle, worker: W, ctx: Context, ctrl_rx: OneshotReceiver) { + let relay = WorkerRelay::new(worker, ctx); + rt.spawn(relay.run(ctrl_rx)); + } +} - // Finally send the router a stop ACK -- log errors - trace!("Sending shutdown ACK"); - self.ctx.stop_ack().await.unwrap_or_else(|e| { +async fn shutdown_and_stop_ack(worker: &mut W, ctx: &mut Context, stopped_from_router: bool) +where + W: Worker, +{ + // Run the shutdown hook for this worker + match worker.shutdown(ctx).await { + Ok(()) => {} + Err(e) => { error!( - "Failed to send stop ACK for worker '{}': {}", - self.ctx.address(), + "Failure during '{}' worker shutdown: {}", + ctx.primary_address(), e - ) - }); + ); + } } - /// Build and spawn a new worker relay, returning a send handle to it - pub(crate) fn init(rt: &Handle, worker: W, ctx: Context, ctrl_rx: SmallReceiver) { - let relay = WorkerRelay::new(worker, ctx); - rt.spawn(relay.run(ctrl_rx)); + let router = match ctx.router() { + Ok(router) => router, + Err(_) => { + error!( + "Failure during '{}' worker shutdown. Can't get router", + ctx.primary_address() + ); + return; + } + }; + + if !stopped_from_router { + if let Err(e) = router.stop_address(ctx.primary_address(), !stopped_from_router) { + error!( + "Failure during '{}' worker shutdown: {}", + ctx.primary_address(), + e + ); + } } + + // Finally send the router a stop ACK -- log errors + trace!("Sending shutdown ACK"); + router.stop_ack(ctx.primary_address()).unwrap_or_else(|e| { + error!( + "Failed to send stop ACK for worker '{}': {}", + ctx.primary_address(), + e + ) + }); } diff --git a/implementations/rust/ockam/ockam_node/src/router/mod.rs b/implementations/rust/ockam/ockam_node/src/router/mod.rs index c04fe5220bc..d1d3bc5b0e3 100644 --- a/implementations/rust/ockam/ockam_node/src/router/mod.rs +++ b/implementations/rust/ockam/ockam_node/src/router/mod.rs @@ -7,19 +7,18 @@ pub mod worker; #[cfg(feature = "metrics")] use core::sync::atomic::AtomicUsize; -use crate::channel_types::{MessageSender, SmallSender}; +use crate::channel_types::{MessageSender, OneshotSender}; use crate::relay::CtrlSignal; +use crate::tokio::runtime::Handle; use crate::{NodeError, NodeReason}; use alloc::string::String; -use alloc::sync::Arc; use alloc::vec::Vec; -use ockam_core::compat::collections::BTreeMap; +use ockam_core::compat::collections::hash_map::Entry; +use ockam_core::compat::collections::HashMap; use ockam_core::compat::sync::RwLock as SyncRwLock; use ockam_core::errcode::{Kind, Origin}; use ockam_core::flow_control::FlowControls; -use ockam_core::{ - Address, AddressAndMetadata, AddressMetadata, Error, RelayMessage, Result, TransportType, -}; +use ockam_core::{Address, AddressMetadata, Error, RelayMessage, Result, TransportType}; use record::{AddressRecord, InternalMap, WorkerMeta}; use state::{NodeState, RouterState}; @@ -27,7 +26,7 @@ use state::{NodeState, RouterState}; #[derive(Debug)] pub struct SenderPair { pub msgs: MessageSender, - pub ctrl: SmallSender, + pub ctrl: OneshotSender, } /// A combined address type and local worker router @@ -40,12 +39,13 @@ pub struct SenderPair { /// registers itself with this router. Only one router can be /// registered per address type. pub struct Router { + runtime_handle: Handle, /// Keep track of some additional router state information state: RouterState, /// Internal address state map: InternalMap, /// Externally registered router components - external: SyncRwLock>, + external: SyncRwLock>, } enum RouteType { @@ -62,33 +62,17 @@ fn determine_type(next: &Address) -> RouteType { } impl Router { - pub fn new(flow_controls: &FlowControls) -> Self { + pub fn new(runtime_handle: Handle, flow_controls: &FlowControls) -> Self { Self { + runtime_handle, state: RouterState::new(), map: InternalMap::new(flow_controls), external: Default::default(), } } - pub fn init(&self, addr: Address, senders: SenderPair) -> Result<()> { - self.map.insert_address_record( - addr.clone(), - AddressRecord::new( - vec![addr.clone()], - senders.msgs, - senders.ctrl, - Arc::new(0.into()), // don't track for app worker (yet?) - WorkerMeta { - processor: false, - detached: true, - }, - ), - vec![], - ) - } - - pub fn set_cluster(&self, addr: Address, label: String) -> Result<()> { - self.map.set_cluster(label, addr) + pub fn set_cluster(&self, addr: &Address, label: String) -> Result<()> { + self.map.set_cluster(addr, label) } pub fn list_workers(&self) -> Vec

{ @@ -99,30 +83,23 @@ impl Router { self.map.is_worker_registered_at(address) } - pub async fn stop_ack(&self, addr: Address) -> Result<()> { - let running = self.state.running(); - debug!(%running, "Handling shutdown ACK for {}", addr); - self.map.free_address(&addr); - - if !running { - // The router is shutting down - if !self.map.cluster_done() { - // We are not done yet. - // The last worker should call another `stop_ack` - return Ok(()); - } - - // Check if there is a next cluster - let finished = self.stop_next_cluster().await?; - if finished { - self.state.terminate().await; - } + pub fn stop_ack(&self, primary_address: &Address) -> Result<()> { + let running = self.state.is_running(); + debug!(%running, "Handling shutdown ACK for {}", primary_address); + + let empty = self.map.stop_ack(primary_address); + if !running && empty { + self.state.set_to_stopped(); } + Ok(()) } - pub fn find_terminal_address(&self, addresses: Vec
) -> Option { - self.map.find_terminal_address(&addresses) + pub fn find_terminal_address<'a>( + &self, + addresses: impl Iterator, + ) -> Option<(&'a Address, AddressMetadata)> { + self.map.find_terminal_address(addresses) } pub fn get_address_metadata(&self, address: &Address) -> Option { @@ -130,8 +107,7 @@ impl Router { } pub fn register_router(&self, tt: TransportType, addr: Address) -> Result<()> { - let mut guard = self.external.write().unwrap(); - if let alloc::collections::btree_map::Entry::Vacant(e) = guard.entry(tt) { + if let Entry::Vacant(e) = self.external.write().unwrap().entry(tt) { e.insert(addr); Ok(()) } else { @@ -161,8 +137,17 @@ impl Router { .ok_or_else(|| NodeError::NodeState(NodeReason::Unknown).internal()) } - pub async fn wait_termination(self: Arc) { - self.state.wait_termination().await; + pub async fn wait_termination(&self) { + self.state.wait_until_stopped().await; + } + + /// Stop the worker + pub fn stop_address(&self, addr: &Address, skip_sending_stop_signal: bool) -> Result<()> { + debug!("Stopping address '{}'", addr); + + self.map.stop(addr, skip_sending_stop_signal)?; + + Ok(()) } #[cfg(feature = "metrics")] diff --git a/implementations/rust/ockam/ockam_node/src/router/processor.rs b/implementations/rust/ockam/ockam_node/src/router/processor.rs index a06a72321d9..10fc6edce5f 100644 --- a/implementations/rust/ockam/ockam_node/src/router/processor.rs +++ b/implementations/rust/ockam/ockam_node/src/router/processor.rs @@ -1,19 +1,13 @@ use super::{AddressRecord, NodeState, Router, SenderPair, WorkerMeta}; -use crate::{error::NodeError, RouterReason}; -use ockam_core::compat::{sync::Arc, vec::Vec}; +use ockam_core::compat::sync::Arc; use ockam_core::errcode::{Kind, Origin}; -use ockam_core::{Address, AddressAndMetadata, Error, Result}; +use ockam_core::{Error, Mailboxes, Result}; impl Router { /// Start a processor - pub(crate) fn start_processor( - &self, - addrs: Vec
, - senders: SenderPair, - addresses_metadata: Vec, - ) -> Result<()> { - if self.state.running() { - self.start_processor_impl(addrs, senders, addresses_metadata) + pub(crate) fn add_processor(&self, mailboxes: &Mailboxes, senders: SenderPair) -> Result<()> { + if self.state.is_running() { + self.add_processor_impl(mailboxes, senders) } else { match self.state.node_state() { NodeState::Stopping => Err(Error::new( @@ -22,26 +16,18 @@ impl Router { "The node is shutting down", ))?, NodeState::Running => unreachable!(), - NodeState::Terminated => unreachable!(), + NodeState::Stopped => unreachable!(), } } } - fn start_processor_impl( - &self, - addrs: Vec
, - senders: SenderPair, - addresses_metadata: Vec, - ) -> Result<()> { - let primary_addr = addrs - .first() - .ok_or_else(|| NodeError::RouterState(RouterReason::EmptyAddressSet).internal())?; - - debug!("Starting new processor '{}'", &primary_addr); + fn add_processor_impl(&self, mailboxes: &Mailboxes, senders: SenderPair) -> Result<()> { + debug!("Starting new processor '{}'", mailboxes.primary_address()); let SenderPair { msgs, ctrl } = senders; let record = AddressRecord::new( - addrs.clone(), + mailboxes.primary_address().clone(), + mailboxes.additional_addresses().cloned().collect(), msgs, ctrl, // We don't keep track of the mailbox count for processors @@ -56,25 +42,8 @@ impl Router { }, ); - self.map - .insert_address_record(primary_addr.clone(), record, addresses_metadata) - } - - /// Stop the processor - pub(crate) async fn stop_processor(&self, addr: &Address) -> Result<()> { - trace!("Stopping processor '{}'", addr); - - // Resolve any secondary address to the primary address - let primary_address = match self.map.get_primary_address(addr) { - Some(p) => p.clone(), - None => { - return Err(Error::new(Origin::Node, Kind::NotFound, "No such address") - .context("Address", addr.clone())) - } - }; + self.map.insert_address_record(record, mailboxes)?; - // Then send processor shutdown signal - self.map.stop(&primary_address).await?; Ok(()) } } diff --git a/implementations/rust/ockam/ockam_node/src/router/record.rs b/implementations/rust/ockam/ockam_node/src/router/record.rs index 6aa197cf93b..1732195c1e4 100644 --- a/implementations/rust/ockam/ockam_node/src/router/record.rs +++ b/implementations/rust/ockam/ockam_node/src/router/record.rs @@ -1,55 +1,67 @@ -use crate::channel_types::{MessageSender, SmallSender}; +use crate::channel_types::{MessageSender, OneshotSender}; use crate::error::{NodeError, NodeReason}; use crate::relay::CtrlSignal; use alloc::string::String; use core::default::Default; use core::fmt::Debug; use core::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; -use ockam_core::compat::sync::Mutex as SyncMutex; +use ockam_core::compat::collections::hash_map::{Entry, EntryRef}; use ockam_core::compat::sync::RwLock as SyncRwLock; -use ockam_core::compat::sync::RwLockWriteGuard as SyncRwLockWriteGuard; use ockam_core::errcode::{Kind, Origin}; use ockam_core::{ compat::{ - collections::{BTreeMap, BTreeSet}, + collections::{HashMap, HashSet}, sync::Arc, vec::Vec, }, flow_control::FlowControls, - Address, AddressAndMetadata, AddressMetadata, Error, RelayMessage, Result, + Address, AddressMetadata, Error, Mailbox, Mailboxes, RelayMessage, Result, }; #[derive(Default)] struct AddressMaps { + // NOTE: It's crucial that if more that one of these structures is needed to perform an + // operation, we should always acquire locks in the order they're declared here. Otherwise, it + // can cause a deadlock. /// Registry of primary address to worker address record state - address_records_map: BTreeMap, + records: SyncRwLock>, /// Alias-registry to map arbitrary address to primary addresses - alias_map: BTreeMap, + aliases: SyncRwLock>, /// Registry of arbitrary metadata for each address, lazily populated - address_metadata_map: BTreeMap, + metadata: SyncRwLock>, +} + +#[derive(Default)] +struct ClusterMaps { + /// The order in which clusters are allocated and de-allocated + order: Vec, + /// Key is Label + map: HashMap>, // Only primary addresses here } /// Address states and associated logic pub struct InternalMap { - address_maps: SyncRwLock, - /// The order in which clusters are allocated and de-allocated - cluster_order: SyncMutex>, - /// Cluster data records - clusters: SyncMutex>>, - /// Track stop information for Clusters - stopping: SyncMutex>, + // NOTE: It's crucial that if more that one of these structures is needed to perform an + // operation, we should always acquire locks in the order they're declared here. Otherwise, it + // can cause a deadlock. + address_maps: AddressMaps, + cluster_maps: SyncRwLock, + /// Track addresses that are being stopped atm + stopping: SyncRwLock>, /// Access to [`FlowControls`] to clean resources flow_controls: FlowControls, /// Metrics collection and sharing #[cfg(feature = "metrics")] metrics: (Arc, Arc), } + impl InternalMap { pub(crate) fn resolve(&self, addr: &Address) -> Result> { - let guard = self.address_maps.read().unwrap(); + let records = self.address_maps.records.read().unwrap(); + let aliases = self.address_maps.aliases.read().unwrap(); - let address_record = if let Some(primary_address) = guard.alias_map.get(addr) { - guard.address_records_map.get(primary_address) + let address_record = if let Some(primary_address) = aliases.get(addr) { + records.get(primary_address) } else { trace!("Resolving worker address '{addr}'... FAILED; no such worker"); return Err(Error::new(Origin::Node, Kind::NotFound, "No such address") @@ -58,17 +70,9 @@ impl InternalMap { match address_record { Some(address_record) => { - if let Some(sender) = address_record.sender() { - trace!("Resolving worker address '{addr}'... OK"); - address_record.increment_msg_count(); - Ok(sender) - } else { - trace!("Resolving worker address '{addr}'... REJECTED; worker shutting down"); - Err( - Error::new(Origin::Node, Kind::Shutdown, "Worker shutting down") - .context("Address", addr.clone()), - ) - } + trace!("Resolving worker address '{addr}'... OK"); + address_record.increment_msg_count(); + Ok(address_record.sender.clone()) } None => { trace!("Resolving worker address '{addr}'... FAILED; no such worker"); @@ -83,9 +87,8 @@ impl InternalMap { pub(super) fn new(flow_controls: &FlowControls) -> Self { Self { address_maps: Default::default(), - cluster_order: SyncMutex::new(Default::default()), - clusters: SyncMutex::new(Default::default()), - stopping: SyncMutex::new(Default::default()), + cluster_maps: Default::default(), + stopping: Default::default(), flow_controls: flow_controls.clone(), #[cfg(feature = "metrics")] metrics: Default::default(), @@ -94,57 +97,72 @@ impl InternalMap { } impl InternalMap { - #[cfg_attr(not(feature = "std"), allow(dead_code))] - pub(super) fn clear_address_records_map(&self) { - self.address_maps - .write() - .unwrap() - .address_records_map - .clear() - } + pub(super) fn stop(&self, address: &Address, skip_sending_stop_signal: bool) -> Result<()> { + // To guarantee consistency we'll first acquire lock on all the maps we need to touch + // and only then start modifications + let mut records = self.address_maps.records.write().unwrap(); + let mut aliases = self.address_maps.aliases.write().unwrap(); + let mut metadata = self.address_maps.metadata.write().unwrap(); + let mut stopping = self.stopping.write().unwrap(); + + let primary_address = aliases + .get(address) + .ok_or_else(|| { + Error::new(Origin::Node, Kind::NotFound, "No such address") + .context("Address", address.clone()) + })? + .clone(); - pub(super) async fn stop(&self, primary_address: &Address) -> Result<()> { - { - let mut guard = self.stopping.lock().unwrap(); - if guard.contains(primary_address) { - return Ok(()); - } else { - guard.insert(primary_address.clone()); - } - } + self.flow_controls.cleanup_address(&primary_address); - let send_signal_future = { - let guard = self.address_maps.read().unwrap(); - if let Some(record) = guard.address_records_map.get(primary_address) { - record.stop() - } else { - return Err(Error::new(Origin::Node, Kind::NotFound, "No such address") - .context("Address", primary_address.clone())); - } + let record = if let Some(record) = records.remove(&primary_address) { + record + } else { + return Err(Error::new(Origin::Node, Kind::NotFound, "No such address") + .context("Address", address.clone())); }; - // we can't call `.await` while holding the lock - if let Some(send_signal_future) = send_signal_future { - send_signal_future.await - } else { - Ok(()) + for address in &record.additional_addresses { + metadata.remove(address); + aliases.remove(address); + } + + metadata.remove(&primary_address); + aliases.remove(&primary_address); + + // Detached doesn't need any stop confirmation, since they don't have a Relay = don't have + // an async task running in a background that should be stopped. + if !record.meta.detached { + stopping.insert(primary_address); } + + record.stop(skip_sending_stop_signal)?; + + Ok(()) + } + + pub(super) fn stop_ack(&self, primary_address: &Address) -> bool { + // FIXME: Detached workers + let mut stopping = self.stopping.write().unwrap(); + stopping.remove(primary_address); + + stopping.is_empty() } pub(super) fn is_worker_registered_at(&self, primary_address: &Address) -> bool { self.address_maps + .records .read() .unwrap() - .address_records_map .contains_key(primary_address) // TODO: we should also check aliases } pub(super) fn list_workers(&self) -> Vec
{ self.address_maps + .records .read() .unwrap() - .address_records_map .keys() .cloned() .collect() @@ -152,103 +170,132 @@ impl InternalMap { pub(super) fn insert_address_record( &self, - primary_address: Address, record: AddressRecord, - addresses_metadata: Vec, + mailboxes: &Mailboxes, ) -> Result<()> { - let mut guard = self.address_maps.write().unwrap(); + let mut records = self.address_maps.records.write().unwrap(); - // check if the address already exists - if let Some(record) = guard.address_records_map.get(&primary_address) { - if record.check_integrity() { - let node = NodeError::Address(primary_address.clone()); + let entry = records.entry(record.primary_address.clone()); + + let entry = match entry { + Entry::Occupied(_) => { + let node = NodeError::Address(record.primary_address); return Err(node.already_exists()); - } else { - self.free_address_impl(&mut guard, &primary_address); } - } + Entry::Vacant(entry) => entry, + }; - // check if aliases are already in use - for addr in record.address_set() { - if guard.alias_map.contains_key(addr) { - let node = NodeError::Address(primary_address.clone()); - return Err(node.already_exists()); + // It may fail, so we don't insert record before that + Self::insert_aliases(&mut self.address_maps.aliases.write().unwrap(), &record)?; + Self::insert_all_metadata(&mut self.address_maps.metadata.write().unwrap(), mailboxes); + + entry.insert(record); + + Ok(()) + } + + fn insert_aliases( + aliases: &mut HashMap, + record: &AddressRecord, + ) -> Result<()> { + Self::insert_alias(aliases, &record.primary_address, &record.primary_address)?; + + for i in 0..record.additional_addresses.len() { + match Self::insert_alias( + aliases, + &record.primary_address, + &record.additional_addresses[i], + ) { + Ok(_) => {} + Err(err) => { + // Rollback + for j in 0..i { + aliases.remove(&record.additional_addresses[j]); + } + + return Err(err); + } } } - record.address_set.iter().for_each(|addr| { - guard - .alias_map - .insert(addr.clone(), primary_address.clone()); - }); + Ok(()) + } - for metadata in addresses_metadata { - if !record.address_set().contains(&metadata.address) { - warn!( - "Address {} is not in the set of addresses", - metadata.address - ); - continue; + fn insert_alias( + aliases: &mut HashMap, + primary_address: &Address, + alias: &Address, + ) -> Result<()> { + match aliases.insert(alias.clone(), primary_address.clone()) { + None => Ok(()), + Some(old_value) => { + // Rollback + aliases.insert(alias.clone(), old_value); + + // FIXME: better error + let node = NodeError::Address(primary_address.clone()); + Err(node.already_exists()) } + } + } + + fn insert_all_metadata( + metadata: &mut HashMap, + mailboxes: &Mailboxes, + ) { + Self::insert_mailbox_metadata(metadata, mailboxes.primary_mailbox()); - let entry = guard - .address_metadata_map - .entry(metadata.address) - .or_default(); - *entry = metadata.metadata; + for mailbox in mailboxes.additional_mailboxes() { + Self::insert_mailbox_metadata(metadata, mailbox); } + } - _ = guard.address_records_map.insert(primary_address, record); - Ok(()) + fn insert_mailbox_metadata( + metadata: &mut HashMap, + mailbox: &Mailbox, + ) { + if let Some(meta) = mailbox.metadata().clone() { + metadata.insert(mailbox.address().clone(), meta.clone()); + } } - pub(super) fn find_terminal_address( + pub(super) fn find_terminal_address<'a>( &self, - addresses: &[Address], - ) -> Option { - addresses.iter().find_map(|address| { - self.address_maps - .read() - .unwrap() - .address_metadata_map - .get(address) - .filter(|&meta| meta.is_terminal) - .map(|meta| AddressAndMetadata { - address: address.clone(), - metadata: meta.clone(), - }) - }) + addresses: impl Iterator, + ) -> Option<(&'a Address, AddressMetadata)> { + let metadata = self.address_maps.metadata.read().unwrap(); + for address in addresses { + if let Some(metadata) = metadata.get(address) { + if metadata.is_terminal { + return Some((address, metadata.clone())); + } + } + } + + None } pub(super) fn get_address_metadata(&self, address: &Address) -> Option { self.address_maps + .metadata .read() .unwrap() - .address_metadata_map .get(address) .cloned() } - - pub(super) fn get_primary_address(&self, alias_address: &Address) -> Option
{ - self.address_maps - .read() - .unwrap() - .alias_map - .get(alias_address) - .cloned() - } } impl InternalMap { #[cfg(feature = "metrics")] pub(super) fn update_metrics(&self) { self.metrics.0.store( - self.address_maps.read().unwrap().address_records_map.len(), + self.address_maps.records.read().unwrap().len(), + Ordering::Release, + ); + self.metrics.1.store( + self.cluster_maps.read().unwrap().map.len(), Ordering::Release, ); - self.metrics - .1 - .store(self.clusters.lock().unwrap().len(), Ordering::Release); } #[cfg(feature = "metrics")] @@ -262,159 +309,132 @@ impl InternalMap { } /// Add an address to a particular cluster - pub(super) fn set_cluster(&self, label: String, primary: Address) -> Result<()> { - let address_records_guard = self.address_maps.read().unwrap(); - let rec = address_records_guard - .address_records_map - .get(&primary) - .ok_or_else(|| NodeError::Address(primary).not_found())?; - - let mut clusters_guard = self.clusters.lock().unwrap(); - - // If this is the first time we see this cluster ID - if !clusters_guard.contains_key(&label) { - clusters_guard.insert(label.clone(), BTreeSet::new()); - self.cluster_order.lock().unwrap().push(label.clone()); + pub(super) fn set_cluster(&self, primary: &Address, label: String) -> Result<()> { + if !self + .address_maps + .records + .read() + .unwrap() + .contains_key(primary) + { + return Err(NodeError::Address(primary.clone()).not_found())?; } - // Add all addresses to the cluster set - for addr in rec.address_set() { - clusters_guard - .get_mut(&label) - .expect("No such cluster??") - .insert(addr.clone()); + // If this is the first time we see this cluster ID + let mut cluster_maps = self.cluster_maps.write().unwrap(); + match cluster_maps.map.entry_ref(&label) { + EntryRef::Occupied(mut occupied_entry) => { + occupied_entry.get_mut().insert(primary.clone()); + } + EntryRef::Vacant(vacant_entry) => { + vacant_entry.insert_entry(HashSet::from([primary.clone()])); + cluster_maps.order.push(label.clone()); + } } Ok(()) } - /// Retrieve the next cluster in reverse-initialisation order - /// Return None if there is no next cluster or if the cluster - /// contained no more active addresses - pub(super) fn next_cluster(&self) -> Option> { - // loop until either: - // - there are no more clusters - // - we found a non-empty list of active addresses in a cluster - loop { - let name = self.cluster_order.lock().unwrap().pop()?; - let addrs = self.clusters.lock().unwrap().remove(&name)?; - let active_addresses: Vec
= self - .address_maps - .read() - .unwrap() - .address_records_map - .iter() - .filter_map(|(primary, _)| { - if addrs.contains(primary) { - Some(primary.clone()) - } else { - None - } - }) - .collect(); - if active_addresses.is_empty() { - continue; - } else { - return Some(active_addresses); + /// Stop all workers not in a cluster, returns their primary addresses + pub(super) fn stop_all_non_cluster_workers(&self) -> bool { + // All clustered addresses + let clustered = self.cluster_maps.read().unwrap().map.iter().fold( + HashSet::new(), + |mut acc, (_, set)| { + acc.extend(set.iter().cloned()); + acc + }, + ); + + let mut records = self.address_maps.records.write().unwrap(); + let mut stopping = self.stopping.write().unwrap(); + + let records_to_stop = records.extract_if(|addr, _record| { + // Filter all clustered workers + !clustered.contains(addr) + }); + + let mut was_empty = true; + + for (_, record) in records_to_stop { + // Detached doesn't need any stop confirmation, since they don't have a Relay = don't have + // an async task running in a background that should be stopped. + if !record.meta.detached { + was_empty = false; + stopping.insert(record.primary_address.clone()); } - } - } - /// Mark this address as "having started to stop" - pub(super) fn init_stop(&self, addr: Address) { - self.stopping.lock().unwrap().insert(addr); - } + if let Err(err) = record.stop(false) { + // FIXME: Insert into stopping + error!("Error stopping address. Err={}", err); + } + } - /// Check whether the current cluster of addresses was stopped - pub(super) fn cluster_done(&self) -> bool { - self.stopping.lock().unwrap().is_empty() + was_empty } - /// Stop all workers not in a cluster, returns their primary addresses - pub(super) async fn stop_all_non_cluster_workers(&self) -> Vec
{ - let mut futures = Vec::new(); - let mut addresses = Vec::new(); - { - let clustered = - self.clusters - .lock() - .unwrap() - .iter() - .fold(BTreeSet::new(), |mut acc, (_, set)| { - acc.append(&mut set.clone()); - acc - }); - - let guard = self.address_maps.read().unwrap(); - let records: Vec<&AddressRecord> = guard - .address_records_map - .iter() - .filter_map(|(addr, rec)| { - if clustered.contains(addr) { - None - } else { - Some(rec) - } - }) - // Filter all detached workers because they don't matter :( - .filter(|rec| !rec.meta.detached) - .collect(); - - for record in records { - if let Some(first_address) = record.address_set.first() { - debug!("Stopping address {}", first_address); - addresses.push(first_address.clone()); - let send_stop_signal_future = record.stop(); - if let Some(send_stop_signal_future) = send_stop_signal_future { - futures.push((first_address.clone(), send_stop_signal_future)); + pub(super) fn stop_all_cluster_workers(&self) -> bool { + let mut records = self.address_maps.records.write().unwrap(); + let mut cluster_maps = self.cluster_maps.write().unwrap(); + let mut stopping = self.stopping.write().unwrap(); + + let mut was_empty = true; + while let Some(label) = cluster_maps.order.pop() { + if let Some(addrs) = cluster_maps.map.remove(&label) { + for addr in addrs { + match records.remove(&addr) { + Some(record) => { + was_empty = false; + // Detached doesn't need any stop confirmation, since they don't have a Relay = don't have + // an async task running in a background that should be stopped. + if !record.meta.detached { + stopping.insert(record.primary_address.clone()); + } + match record.stop(false) { + Ok(_) => {} + Err(err) => { + error!( + "Error stopping address {} from cluster {}. Err={}", + addr, label, err + ); + } + } + } + None => { + warn!( + "Stopping address {} from cluster {} but it doesn't exist", + addr, label + ); + } } - } else { - error!("Empty Address Set during graceful shutdown") } } } - // We can only call `.await` outside the lock - for (first_address, send_stop_signal_future) in futures { - send_stop_signal_future.await.unwrap_or_else(|e| { - error!("Failed to stop address {}: {}", first_address, e); - }); - } - - addresses + was_empty } - /// Permanently free all remaining resources associated to a particular address - pub(super) fn free_address(&self, primary: &Address) { - let mut guard = self.address_maps.write().unwrap(); - self.free_address_impl(&mut guard, primary); - } - - fn free_address_impl(&self, guard: &mut SyncRwLockWriteGuard, primary: &Address) { - self.stopping.lock().unwrap().remove(primary); - self.flow_controls.cleanup_address(primary); + pub(super) fn force_clear_records(&self) -> Vec
{ + let mut records = self.address_maps.records.write().unwrap(); - let removed = guard.address_records_map.remove(primary); - if let Some(record) = removed { - for alias_address in record.address_set { - guard.address_metadata_map.remove(&alias_address); - guard.alias_map.remove(&alias_address); - } - } + records.drain().map(|(address, _record)| address).collect() } } /// Additional metadata for worker records #[derive(Debug)] pub struct WorkerMeta { + // FIXME + #[allow(dead_code)] pub processor: bool, pub detached: bool, } pub struct AddressRecord { - address_set: Vec
, - sender: SyncMutex>>, - ctrl_tx: SmallSender, // Unused for not-detached workers + primary_address: Address, + additional_addresses: Vec
, + sender: MessageSender, + ctrl_tx: OneshotSender, state: AtomicU8, meta: WorkerMeta, msg_count: Arc, @@ -423,8 +443,9 @@ pub struct AddressRecord { impl Debug for AddressRecord { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { f.debug_struct("AddressRecord") - .field("address_set", &self.address_set) - .field("sender", &self.sender.lock().unwrap().is_some()) + .field("primary_address", &self.primary_address) + .field("additional_addresses", &self.additional_addresses) + .field("sender", &self.sender) .field("ctrl_tx", &self.ctrl_tx) .field("state", &self.state) .field("meta", &self.meta) @@ -434,28 +455,18 @@ impl Debug for AddressRecord { } impl AddressRecord { - pub fn address_set(&self) -> &[Address] { - &self.address_set - } - - pub fn sender(&self) -> Option> { - self.sender.lock().unwrap().clone() - } - - pub fn drop_sender(&self) { - *self.sender.lock().unwrap() = None; - } - pub fn new( - address_set: Vec
, + primary_address: Address, + additional_addresses: Vec
, sender: MessageSender, - ctrl_tx: SmallSender, + ctrl_tx: OneshotSender, msg_count: Arc, meta: WorkerMeta, ) -> Self { AddressRecord { - address_set, - sender: SyncMutex::new(Some(sender)), + primary_address, + additional_addresses, + sender, ctrl_tx, state: AtomicU8::new(AddressState::Running as u8), msg_count, @@ -469,12 +480,14 @@ impl AddressRecord { } /// Signal this worker to stop -- it will no longer be able to receive messages - /// Since the worker is held behind a lock, we cannot hold the synchronous lock - /// and call an `.await` to send the signal. - /// To avoid this shortcoming, we return a future that can be awaited on instead. - pub fn stop(&self) -> Option>> { + pub fn stop(self, skip_sending_stop_signal: bool) -> Result<()> { + trace!("AddressRecord::stop called for {:?}", self.primary_address); if self.state.load(Ordering::Relaxed) != AddressState::Running as u8 { - return None; + trace!( + "AddressRecord::stop called for {:?}. Already stopping", + self.primary_address + ); + return Ok(()); } // the stop can be triggered only once @@ -485,29 +498,17 @@ impl AddressRecord { Ordering::Relaxed, ); if result.is_err() { - return None; + return Ok(()); } - if self.meta.processor { - let ctrl_tx = self.ctrl_tx.clone(); - Some(async move { - ctrl_tx - .send(CtrlSignal::InterruptStop) - .await - .map_err(|_| NodeError::NodeState(NodeReason::Unknown).internal())?; - Ok(()) - }) - } else { - self.drop_sender(); - None + // TODO: Recheck + if !self.meta.detached && !skip_sending_stop_signal { + self.ctrl_tx + .send(CtrlSignal::InterruptStop) + .map_err(|_| NodeError::NodeState(NodeReason::Unknown).internal())?; } - } - /// Check the integrity of this record - #[inline] - pub fn check_integrity(&self) -> bool { - self.state.load(Ordering::Relaxed) == AddressState::Running as u8 - && self.sender.lock().unwrap().is_some() + Ok(()) } } @@ -517,70 +518,9 @@ impl AddressRecord { pub enum AddressState { /// The runner is looping in its main body (either handling messages or a manual run-loop) Running, - /// The runner was signaled to shut-down (running `shutdown()`) + /// The runner was signaled to shut-down (running `stop()`) Stopping, /// The runner has experienced an error and is waiting for supervisor intervention #[allow(unused)] Faulty, } - -#[cfg(test)] -mod test { - use super::*; - use crate::channel_types::small_channel; - use crate::router::record::InternalMap; - - #[test] - fn test_next_cluster() { - let map = InternalMap::new(&FlowControls::new()); - - // create 3 clusters - // cluster 1 has one active address - // cluster 2 has no active address - // cluster 3 has one active address - map.address_maps - .write() - .unwrap() - .address_records_map - .insert("address1".into(), create_address_record("address1")); - let _ = map.set_cluster("CLUSTER1".into(), "address1".into()); - - map.address_maps - .write() - .unwrap() - .address_records_map - .insert("address2".into(), create_address_record("address2")); - let _ = map.set_cluster("CLUSTER2".into(), "address2".into()); - map.free_address(&"address2".into()); - - map.address_maps - .write() - .unwrap() - .address_records_map - .insert("address3".into(), create_address_record("address3")); - let _ = map.set_cluster("CLUSTER3".into(), "address3".into()); - - // get the active addresses for cluster 3, clusters are popped in reverse order - assert_eq!(map.next_cluster(), Some(vec!["address3".into()])); - // get the active addresses for cluster 1, cluster 2 is skipped - assert_eq!(map.next_cluster(), Some(vec!["address1".into()])); - // finally there are no more clusters with active addresses - assert_eq!(map.next_cluster(), None); - } - - /// HELPERS - fn create_address_record(primary: &str) -> AddressRecord { - let (tx1, _) = small_channel(); - let (tx2, _) = small_channel(); - AddressRecord::new( - vec![primary.into()], - tx1, - tx2, - Arc::new(AtomicUsize::new(1)), - WorkerMeta { - processor: false, - detached: false, - }, - ) - } -} diff --git a/implementations/rust/ockam/ockam_node/src/router/shutdown.rs b/implementations/rust/ockam/ockam_node/src/router/shutdown.rs index bccf02eba45..a91da3168e1 100644 --- a/implementations/rust/ockam/ockam_node/src/router/shutdown.rs +++ b/implementations/rust/ockam/ockam_node/src/router/shutdown.rs @@ -1,7 +1,8 @@ use super::Router; -use alloc::sync::Arc; -use alloc::vec::Vec; -use ockam_core::{Address, Result}; +use crate::tokio::time; +use core::time::Duration; +use ockam_core::compat::sync::Arc; +use ockam_core::Result; /// Register a stop ACK /// @@ -9,14 +10,14 @@ use ockam_core::{Address, Result}; /// If not, we do nothing. If so, we trigger the next cluster to stop. impl Router { - /// Implement the graceful shutdown strategy + /// Implement the graceful stop strategy #[cfg_attr(not(feature = "std"), allow(unused_variables))] pub async fn stop_graceful(self: Arc, seconds: u8) -> Result<()> { // Mark the router as shutting down to prevent spawning info!("Initiate graceful node shutdown"); // This changes the router state to `Stopping` - let receiver = self.state.shutdown(); - let mut receiver = if let Some(receiver) = receiver { + let receiver = self.state.set_to_stopping(); + let receiver = if let Some(receiver) = receiver { receiver } else { debug!("Node is already terminated"); @@ -24,58 +25,43 @@ impl Router { }; // Start by shutting down clusterless workers - let cluster = self.map.stop_all_non_cluster_workers().await; + let no_non_cluster_workers = self.map.stop_all_non_cluster_workers(); - // If there _are_ no clusterless workers we go to the next cluster - if cluster.is_empty() { - let result = self.stop_next_cluster().await; - if let Ok(true) = result { - self.state.terminate().await; - return Ok(()); - } - } + // Not stop cluster addresses + let no_cluster_workers = self.map.stop_all_cluster_workers(); - // Otherwise: keep track of addresses we are stopping - cluster - .into_iter() - .for_each(|addr| self.map.init_stop(addr)); + if no_non_cluster_workers && no_cluster_workers { + // No stop ack will arrive because we didn't stop anything + self.state.set_to_stopped(); + return Ok(()); + } // Start a timeout task to interrupt us... - #[cfg(feature = "std")] - { - use core::time::Duration; - use tokio::{task, time}; - - let dur = Duration::from_secs(seconds as u64); - task::spawn(async move { - time::sleep(dur).await; - warn!("Shutdown timeout reached; aborting node!"); - self.map.clear_address_records_map(); - self.state.terminate().await; - }); - } + let dur = Duration::from_secs(seconds as u64); - receiver.recv().await; - Ok(()) - } + let r = self.clone(); + let timeout_handle = self.runtime_handle.spawn(async move { + time::sleep(dur).await; - pub(crate) async fn stop_next_cluster(&self) -> Result { - let next_cluster_addresses = self.map.next_cluster(); + warn!("Shutdown timeout reached; aborting node!"); + let uncleared_addresses = r.map.force_clear_records(); - match next_cluster_addresses { - Some(vec) => { - self.stop_cluster_addresses(vec).await?; - Ok(false) + if !uncleared_addresses.is_empty() { + error!( + "Router internal inconsistency detected.\ + Records map is not empty after stopping all workers. Addresses: {:?}", + uncleared_addresses + ); } - // If not, we are done! - None => Ok(true), - } - } - async fn stop_cluster_addresses(&self, addresses: Vec
) -> Result<()> { - for address in addresses.iter() { - self.map.stop(address).await?; - } + r.state.set_to_stopped(); + }); + + receiver.await.unwrap(); //FIXME + + #[cfg(feature = "std")] + timeout_handle.abort(); + Ok(()) } } diff --git a/implementations/rust/ockam/ockam_node/src/router/state.rs b/implementations/rust/ockam/ockam_node/src/router/state.rs index c36637f6b20..1e8e46188a6 100644 --- a/implementations/rust/ockam/ockam_node/src/router/state.rs +++ b/implementations/rust/ockam/ockam_node/src/router/state.rs @@ -1,21 +1,22 @@ //! Router run state utilities -use crate::channel_types::{SmallReceiver, SmallSender}; +use crate::channel_types::{OneshotReceiver, OneshotSender}; use alloc::vec::Vec; use core::ops::{Deref, DerefMut}; use core::sync::atomic::AtomicBool; use ockam_core::compat::sync::Mutex as SyncMutex; +// TODO: Merge RouterState and NodeState. /// Node state #[derive(Clone)] pub enum NodeState { Running, Stopping, - Terminated, + Stopped, } pub struct RouterState { - node_state: SyncMutex, - termination_senders: SyncMutex>>, + node_state: SyncMutex, // TODO: Use AtomicU8 instead and remove is_running field + termination_senders: SyncMutex>>, is_running: AtomicBool, } @@ -29,35 +30,35 @@ impl RouterState { } /// Set the router state to `Stopping` and return a receiver - /// to wait for the shutdown to complete. + /// to wait for the stop to complete. /// When `None` is returned, the router is already terminated. - pub(super) fn shutdown(&self) -> Option> { - let mut guard = self.node_state.lock().unwrap(); - match guard.deref_mut() { + pub(super) fn set_to_stopping(&self) -> Option> { + let mut node_state = self.node_state.lock().unwrap(); + match node_state.deref_mut() { NodeState::Running => { - let (sender, receiver) = crate::channel_types::small_channel(); - *guard = NodeState::Stopping; + let (sender, receiver) = crate::channel_types::oneshot_channel(); + *node_state = NodeState::Stopping; self.is_running .store(false, core::sync::atomic::Ordering::Relaxed); self.termination_senders.lock().unwrap().push(sender); Some(receiver) } NodeState::Stopping => { - let (sender, receiver) = crate::channel_types::small_channel(); + let (sender, receiver) = crate::channel_types::oneshot_channel(); self.termination_senders.lock().unwrap().push(sender); Some(receiver) } - NodeState::Terminated => None, + NodeState::Stopped => None, } } - /// Set the router to `Terminated` state and notify all tasks waiting for shutdown - pub(super) async fn terminate(&self) { + /// Set the router to `Stopped` state and notify all tasks waiting for shutdown + pub(super) fn set_to_stopped(&self) { self.is_running .store(false, core::sync::atomic::Ordering::Relaxed); let previous = { let mut guard = self.node_state.lock().unwrap(); - core::mem::replace(guard.deref_mut(), NodeState::Terminated) + core::mem::replace(guard.deref_mut(), NodeState::Stopped) }; match previous { @@ -68,31 +69,31 @@ impl RouterState { core::mem::take(guard.deref_mut()) }; for sender in senders { - let _ = sender.send(()).await; + let _ = sender.send(()); } } - NodeState::Terminated => {} + NodeState::Stopped => {} } } - pub(super) async fn wait_termination(&self) { - let mut receiver = { + pub(super) async fn wait_until_stopped(&self) { + let receiver = { let guard = self.node_state.lock().unwrap(); match guard.deref() { NodeState::Running | NodeState::Stopping => { - let (sender, receiver) = crate::channel_types::small_channel(); + let (sender, receiver) = crate::channel_types::oneshot_channel(); self.termination_senders.lock().unwrap().push(sender); receiver } - NodeState::Terminated => { + NodeState::Stopped => { return; } } }; - receiver.recv().await; + receiver.await.unwrap() // FIXME } - pub(super) fn running(&self) -> bool { + pub(super) fn is_running(&self) -> bool { self.is_running.load(core::sync::atomic::Ordering::Relaxed) } diff --git a/implementations/rust/ockam/ockam_node/src/router/worker.rs b/implementations/rust/ockam/ockam_node/src/router/worker.rs index 502b8531917..ca8ea9a012a 100644 --- a/implementations/rust/ockam/ockam_node/src/router/worker.rs +++ b/implementations/rust/ockam/ockam_node/src/router/worker.rs @@ -1,23 +1,18 @@ use crate::router::{AddressRecord, NodeState, Router, SenderPair, WorkerMeta}; -use crate::{error::NodeError, RouterReason}; use core::sync::atomic::AtomicUsize; use ockam_core::errcode::{Kind, Origin}; -use ockam_core::{ - compat::{sync::Arc, vec::Vec}, - Address, AddressAndMetadata, Error, Result, -}; +use ockam_core::{compat::sync::Arc, Error, Mailboxes, Result}; impl Router { /// Start a new worker - pub fn start_worker( + pub fn add_worker( &self, - addrs: Vec
, + mailboxes: &Mailboxes, senders: SenderPair, detached: bool, - addresses_metadata: Vec, metrics: Arc, ) -> Result<()> { - if !self.state.running() { + if !self.state.is_running() { match self.state.node_state() { NodeState::Stopping => Err(Error::new( Origin::Node, @@ -25,31 +20,27 @@ impl Router { "The node is shutting down", ))?, NodeState::Running => unreachable!(), - NodeState::Terminated => unreachable!(), + NodeState::Stopped => unreachable!(), } } else { - self.start_worker_impl(addrs, senders, detached, addresses_metadata, metrics) + self.add_worker_impl(mailboxes, senders, detached, metrics) } } - fn start_worker_impl( + fn add_worker_impl( &self, - addrs: Vec
, + mailboxes: &Mailboxes, senders: SenderPair, detached: bool, - addresses_metadata: Vec, metrics: Arc, ) -> Result<()> { - let primary_addr = addrs - .first() - .ok_or_else(|| NodeError::RouterState(RouterReason::EmptyAddressSet).internal())?; - - debug!("Starting new worker '{}'", primary_addr); + debug!("Starting new worker '{}'", mailboxes.primary_address()); let SenderPair { msgs, ctrl } = senders; // Create an address record and insert it into the internal map let address_record = AddressRecord::new( - addrs.clone(), + mailboxes.primary_address().clone(), + mailboxes.additional_addresses().cloned().collect(), msgs, ctrl, metrics, @@ -59,34 +50,7 @@ impl Router { }, ); - self.map - .insert_address_record(primary_addr.clone(), address_record, addresses_metadata) - } - - /// Stop the worker - pub async fn stop_worker(&self, addr: &Address, detached: bool) -> Result<()> { - debug!("Stopping worker '{}'", addr); - - // Resolve any secondary address to the primary address - let primary_address = match self.map.get_primary_address(addr) { - Some(p) => p, - None => { - return Err(Error::new(Origin::Node, Kind::NotFound, "No such address") - .context("Address", addr)) - } - }; - - // If we are dropping a real worker, then we simply close the - // mailbox channel to trigger a graceful worker self-shutdown. - // The worker shutdown will call `free_address()`. - // - // For detached workers (i.e. Context's without a mailbox relay - // running) we simply free the address. - if detached { - self.map.free_address(&primary_address); - } else { - self.map.stop(&primary_address).await?; - } + self.map.insert_address_record(address_record, mailboxes)?; Ok(()) } diff --git a/implementations/rust/ockam/ockam_node/src/runtime.rs b/implementations/rust/ockam/ockam_node/src/runtime.rs deleted file mode 100644 index c354d06f214..00000000000 --- a/implementations/rust/ockam/ockam_node/src/runtime.rs +++ /dev/null @@ -1,16 +0,0 @@ -use once_cell::sync::Lazy; -use std::sync::Mutex; -use tokio::runtime::Runtime; - -pub(crate) static RUNTIME: Lazy>> = - Lazy::new(|| Mutex::new(Some(Runtime::new().unwrap()))); - -/// Return the Runtime singleton -/// This function can only be accessed once -pub fn take() -> Runtime { - RUNTIME - .lock() - .unwrap() - .take() - .expect("Runtime was consumed") -} diff --git a/implementations/rust/ockam/ockam_node/src/shutdown.rs b/implementations/rust/ockam/ockam_node/src/shutdown.rs index a75b1f8281b..8582064378b 100644 --- a/implementations/rust/ockam/ockam_node/src/shutdown.rs +++ b/implementations/rust/ockam/ockam_node/src/shutdown.rs @@ -1,3 +1,5 @@ +// TODO: Use either stop or shutdown everywhere + /// Specify the type of node shutdown /// /// For most users `ShutdownType::Graceful()` is recommended. The diff --git a/implementations/rust/ockam/ockam_node/src/worker_builder.rs b/implementations/rust/ockam/ockam_node/src/worker_builder.rs index 3d980757551..7eb3328ce6b 100644 --- a/implementations/rust/ockam/ockam_node/src/worker_builder.rs +++ b/implementations/rust/ockam/ockam_node/src/worker_builder.rs @@ -1,9 +1,9 @@ -use crate::debugger; +use crate::{debugger, ContextMode}; use crate::{relay::WorkerRelay, Context}; -use alloc::string::String; -use ockam_core::compat::{sync::Arc, vec::Vec}; +use ockam_core::compat::string::String; +use ockam_core::compat::sync::Arc; use ockam_core::{ - Address, AddressAndMetadata, AddressMetadata, AllowAll, IncomingAccessControl, Mailboxes, + Address, AddressMetadata, AllowAll, IncomingAccessControl, Mailbox, Mailboxes, OutgoingAccessControl, Result, Worker, }; @@ -34,12 +34,41 @@ where { /// Worker with only one [`Address`] pub fn with_address(self, address: impl Into
) -> WorkerBuilderOneAddress { + self.with_address_and_metadata_impl(address, None) + } + + /// Worker with single terminal [`Address`] + pub fn with_terminal_address(self, address: impl Into
) -> WorkerBuilderOneAddress { + self.with_address_and_metadata( + address, + AddressMetadata { + is_terminal: true, + attributes: vec![], + }, + ) + } + + /// Worker with single terminal [`Address`] and metadata + pub fn with_address_and_metadata( + self, + address: impl Into
, + metadata: AddressMetadata, + ) -> WorkerBuilderOneAddress { + self.with_address_and_metadata_impl(address, Some(metadata)) + } + + /// Worker with single terminal [`Address`] and metadata + pub fn with_address_and_metadata_impl( + self, + address: impl Into
, + metadata: Option, + ) -> WorkerBuilderOneAddress { WorkerBuilderOneAddress { incoming_ac: Arc::new(AllowAll), outgoing_ac: Arc::new(AllowAll), worker: self.worker, address: address.into(), - metadata: None, + metadata, } } @@ -48,7 +77,6 @@ where WorkerBuilderMultipleAddresses { mailboxes, worker: self.worker, - metadata_list: vec![], } } } @@ -59,72 +87,15 @@ where { mailboxes: Mailboxes, worker: W, - metadata_list: Vec, } impl WorkerBuilderMultipleAddresses where W: Worker, { - /// Mark the provided address as terminal - pub fn terminal(self, address: impl Into
) -> Self { - self.terminal_with_attributes(address.into(), vec![]) - } - - /// Mark the provided address as terminal - pub fn terminal_with_attributes( - mut self, - address: impl Into
, - attributes: Vec<(String, String)>, - ) -> Self { - let address = address.into(); - let metadata = self.metadata_list.iter_mut().find(|m| m.address == address); - - if let Some(metadata) = metadata { - metadata.metadata.attributes = attributes; - metadata.metadata.is_terminal = true; - } else { - self.metadata_list.push(AddressAndMetadata { - address, - metadata: AddressMetadata { - is_terminal: true, - attributes, - }, - }); - } - self - } - - /// Adds metadata attribute for the provided address - pub fn with_metadata_attribute( - mut self, - address: impl Into
, - key: impl Into, - value: impl Into, - ) -> Self { - let address = address.into(); - let metadata = self.metadata_list.iter_mut().find(|m| m.address == address); - - if let Some(metadata) = metadata { - metadata - .metadata - .attributes - .push((key.into(), value.into())); - } else { - self.metadata_list.push(AddressAndMetadata { - address, - metadata: AddressMetadata { - is_terminal: false, - attributes: vec![(key.into(), value.into())], - }, - }); - } - self - } - /// Consume this builder and start a new Ockam [`Worker`] from the given context pub async fn start(self, context: &Context) -> Result<()> { - start(context, self.mailboxes, self.worker, self.metadata_list).await + start(context, self.mailboxes, self.worker).await } } @@ -136,55 +107,38 @@ where outgoing_ac: Arc, address: Address, worker: W, - metadata: Option, + metadata: Option, } impl WorkerBuilderOneAddress where W: Worker, { - /// Mark the address as terminal - pub fn terminal(self) -> Self { - self.terminal_with_attributes(vec![]) - } - - /// Mark the address as terminal - pub fn terminal_with_attributes(mut self, attributes: Vec<(String, String)>) -> Self { - if let Some(metadata) = self.metadata.as_mut() { - metadata.metadata.is_terminal = true; - metadata.metadata.attributes = attributes; - } else { - self.metadata = Some(AddressAndMetadata { - address: self.address.clone(), - metadata: AddressMetadata { - is_terminal: true, - attributes, - }, - }); - } + /// Mark the provided address as terminal + pub fn terminal(mut self) -> Self { + self.metadata + .get_or_insert(AddressMetadata { + is_terminal: false, + attributes: vec![], + }) + .is_terminal = true; self } - /// Adds metadata attribute + /// Adds metadata attribute for the provided address pub fn with_metadata_attribute( mut self, key: impl Into, value: impl Into, ) -> Self { - if let Some(metadata) = self.metadata.as_mut() { - metadata - .metadata - .attributes - .push((key.into(), value.into())); - } else { - self.metadata = Some(AddressAndMetadata { - address: self.address.clone(), - metadata: AddressMetadata { - is_terminal: false, - attributes: vec![(key.into(), value.into())], - }, - }); - } + self.metadata + .get_or_insert(AddressMetadata { + is_terminal: false, + attributes: vec![], + }) + .attributes + .push((key.into(), value.into())); + self } @@ -192,9 +146,16 @@ where pub async fn start(self, context: &Context) -> Result<()> { start( context, - Mailboxes::main(self.address, self.incoming_ac, self.outgoing_ac), + Mailboxes::new( + Mailbox::new( + self.address, + self.metadata, + self.incoming_ac, + self.outgoing_ac, + ), + vec![], + ), self.worker, - self.metadata.map(|m| vec![m]).unwrap_or_default(), ) .await } @@ -242,31 +203,24 @@ where } /// Consume this builder and start a new Ockam [`Worker`] from the given context -async fn start( - context: &Context, - mailboxes: Mailboxes, - worker: W, - metadata: Vec, -) -> Result<()> +async fn start(context: &Context, mailboxes: Mailboxes, worker: W) -> Result<()> where W: Worker, { debug!( "Initializing ockam worker '{}' with access control in:{:?} out:{:?}", - mailboxes.main_address(), - mailboxes.main_mailbox().incoming_access_control(), - mailboxes.main_mailbox().outgoing_access_control(), + mailboxes.primary_address(), + mailboxes.primary_mailbox().incoming_access_control(), + mailboxes.primary_mailbox().outgoing_access_control(), ); - let addresses = mailboxes.addresses(); - // Pass it to the context - let (ctx, sender, ctrl_rx) = context.copy_with_mailboxes(mailboxes); + let (ctx, sender, ctrl_rx) = context.new_with_mailboxes(mailboxes, ContextMode::Attached); debugger::log_inherit_context("WORKER", context, &ctx); - let router = context.router(); - router.start_worker(addresses, sender, false, metadata, context.mailbox_count())?; + let router = context.router()?; + router.add_worker(ctx.mailboxes(), sender, false, context.mailbox_count())?; // Then initialise the worker message relay WorkerRelay::init(context.runtime(), worker, ctx, ctrl_rx); diff --git a/implementations/rust/ockam/ockam_node/src/workers/echoer.rs b/implementations/rust/ockam/ockam_node/src/workers/echoer.rs index b8e3c0faa1c..13e2862d83a 100644 --- a/implementations/rust/ockam/ockam_node/src/workers/echoer.rs +++ b/implementations/rust/ockam/ockam_node/src/workers/echoer.rs @@ -15,7 +15,7 @@ impl Worker for Echoer { type Message = String; async fn handle_message(&mut self, ctx: &mut Context, msg: Routed) -> Result<()> { - debug!("Address: {}, Received: {:?}", ctx.address(), msg); + debug!("Address: {}, Received: {:?}", ctx.primary_address(), msg); // Echo the message body back on its return_route. ctx.send(msg.return_route().clone(), msg.into_body()?).await diff --git a/implementations/rust/ockam/ockam_node/tests/router.rs b/implementations/rust/ockam/ockam_node/tests/metadata.rs similarity index 56% rename from implementations/rust/ockam/ockam_node/tests/router.rs rename to implementations/rust/ockam/ockam_node/tests/metadata.rs index 4ea18dba91a..ffe3c1f41d8 100644 --- a/implementations/rust/ockam/ockam_node/tests/router.rs +++ b/implementations/rust/ockam/ockam_node/tests/metadata.rs @@ -1,7 +1,8 @@ use ockam_core::{ - async_trait, route, Address, DenyAll, Mailbox, Mailboxes, Processor, Result, TransportType, + async_trait, route, AddressMetadata, DenyAll, Mailbox, Mailboxes, Processor, Result, }; use ockam_node::{Context, NullWorker, ProcessorBuilder, WorkerBuilder}; +use std::string::ToString; use std::sync::Arc; struct NullProcessor; @@ -24,65 +25,64 @@ async fn find_terminal_for_processor(context: &mut Context) -> Result<()> { .await?; assert!(context - .find_terminal_address(route!["simple_processor", "non-existing"]) - .await? + .find_terminal_address(route!["simple_processor", "non-existing"].iter())? .is_none()); ProcessorBuilder::new(NullProcessor {}) - .with_address("terminal_processor") - .terminal() + .with_terminal_address("terminal_processor") .start(context) .await?; assert_eq!( context - .find_terminal_address(route![ - "simple_worker", - "terminal_processor", - "non-existing" - ]) - .await? + .find_terminal_address( + route!["simple_worker", "terminal_processor", "non-existing"].iter() + )? .unwrap() - .address, - "terminal_processor".into() + .0, + &"terminal_processor".into() ); - context.stop().await + Ok(()) } #[ockam_macros::test] async fn find_terminal_for_processor_alias(context: &mut Context) -> Result<()> { ProcessorBuilder::new(NullProcessor {}) .with_mailboxes(Mailboxes::new( - Mailbox::new("main", Arc::new(DenyAll), Arc::new(DenyAll)), - vec![Mailbox::new("alias", Arc::new(DenyAll), Arc::new(DenyAll))], + Mailbox::new("main", None, Arc::new(DenyAll), Arc::new(DenyAll)), + vec![Mailbox::new( + "alias", + Some(AddressMetadata { + is_terminal: true, + attributes: vec![], + }), + Arc::new(DenyAll), + Arc::new(DenyAll), + )], )) - .terminal("alias") .start(context) .await?; assert!(context - .find_terminal_address(route!["main", "non-existing"]) - .await? + .find_terminal_address(route!["main", "non-existing"].iter())? .is_none()); assert_eq!( context - .find_terminal_address(route!["main", "alias", "other"]) - .await? + .find_terminal_address(route!["main", "alias", "other"].iter())? .unwrap() - .address, - "alias".into() + .0, + &"alias".into() ); - context.stop_processor("main").await?; + context.stop_address("main")?; ockam_node::compat::tokio::time::sleep(std::time::Duration::from_millis(10)).await; assert!(context - .find_terminal_address(route!["main", "alias", "other"]) - .await? + .find_terminal_address(route!["main", "alias", "other"].iter())? .is_none()); - context.stop().await + Ok(()) } #[ockam_macros::test] @@ -94,7 +94,7 @@ async fn provide_and_read_processor_address_metadata(context: &mut Context) -> R .start(context) .await?; - let meta = context.get_metadata("processor_address").unwrap(); + let meta = context.get_metadata("processor_address")?.unwrap(); assert!(!meta.is_terminal); @@ -106,13 +106,13 @@ async fn provide_and_read_processor_address_metadata(context: &mut Context) -> R ] ); - assert_eq!(context.get_metadata("non-existing-worker"), None); + assert_eq!(context.get_metadata("non-existing-worker")?, None); - context.stop_processor("processor_address").await?; + context.stop_address("processor_address")?; ockam_node::compat::tokio::time::sleep(std::time::Duration::from_millis(10)).await; - assert_eq!(context.get_metadata("processor_address"), None); + assert_eq!(context.get_metadata("processor_address")?, None); - context.stop().await + Ok(()) } #[ockam_macros::test] @@ -123,81 +123,70 @@ async fn find_terminal_for_worker(context: &mut Context) -> Result<()> { .await?; assert!(context - .find_terminal_address(route!["simple_worker", "non-existing"]) - .await? + .find_terminal_address(route!["simple_worker", "non-existing"].iter())? .is_none()); WorkerBuilder::new(NullWorker {}) - .with_address("terminal_worker") - .terminal() + .with_terminal_address("terminal_worker") .start(context) .await?; assert_eq!( context - .find_terminal_address(route!["simple_worker", "terminal_worker", "non-existing"]) - .await? + .find_terminal_address( + route!["simple_worker", "terminal_worker", "non-existing"].iter() + )? .unwrap() - .address, - "terminal_worker".into() + .0, + &"terminal_worker".into() ); - let remote = Address::new_with_string(TransportType::new(1), "127.0.0.1"); - assert!(context - .find_terminal_address(route![ - "simple_worker", - remote, - "terminal_worker", - "non-existing" - ]) - .await - .is_err()); - - context.stop_worker("terminal_worker").await?; - ockam_node::compat::tokio::time::sleep(std::time::Duration::from_millis(10)).await; + context.stop_address("terminal_worker")?; assert_eq!( - context - .find_terminal_address(route!["terminal_worker"]) - .await?, + context.find_terminal_address(route!["terminal_worker"].iter())?, None ); - context.stop().await + Ok(()) } #[ockam_macros::test] async fn find_terminal_for_worker_alias(context: &mut Context) -> Result<()> { WorkerBuilder::new(NullWorker {}) .with_mailboxes(Mailboxes::new( - Mailbox::new("main", Arc::new(DenyAll), Arc::new(DenyAll)), - vec![Mailbox::new("alias", Arc::new(DenyAll), Arc::new(DenyAll))], + Mailbox::new("main", None, Arc::new(DenyAll), Arc::new(DenyAll)), + vec![Mailbox::new( + "alias", + Some(AddressMetadata { + is_terminal: true, + attributes: vec![], + }), + Arc::new(DenyAll), + Arc::new(DenyAll), + )], )) - .terminal("alias") .start(context) .await?; assert!(context - .find_terminal_address(route!["main", "non-existing"]) - .await? + .find_terminal_address(route!["main", "non-existing"].iter())? .is_none()); assert_eq!( context - .find_terminal_address(route!["main", "alias", "other"]) - .await? + .find_terminal_address(route!["main", "alias", "other"].iter())? .unwrap() - .address, - "alias".into() + .0, + &"alias".into() ); - context.stop_worker("main").await?; + context.stop_address("main")?; ockam_node::compat::tokio::time::sleep(std::time::Duration::from_millis(10)).await; assert!(context - .find_terminal_address(route!["main", "alias", "other"]) - .await? + .find_terminal_address(route!["main", "alias", "other"].iter())? .is_none()); - context.stop().await + Ok(()) } #[ockam_macros::test] @@ -209,7 +198,7 @@ async fn provide_and_read_address_metadata(context: &mut Context) -> Result<()> .start(context) .await?; - let meta = context.get_metadata("worker_address").unwrap(); + let meta = context.get_metadata("worker_address")?.unwrap(); assert!(!meta.is_terminal); @@ -221,28 +210,42 @@ async fn provide_and_read_address_metadata(context: &mut Context) -> Result<()> ] ); - assert_eq!(context.get_metadata("non-existing-worker"), None); + assert_eq!(context.get_metadata("non-existing-worker")?, None); - context.stop_worker("worker_address").await?; + context.stop_address("worker_address")?; ockam_node::compat::tokio::time::sleep(std::time::Duration::from_millis(10)).await; - assert_eq!(context.get_metadata("worker_address"), None); + assert_eq!(context.get_metadata("worker_address")?, None); - context.stop().await + Ok(()) } #[ockam_macros::test] async fn provide_and_read_address_metadata_worker_alias(context: &mut Context) -> Result<()> { WorkerBuilder::new(NullWorker {}) .with_mailboxes(Mailboxes::new( - Mailbox::new("main", Arc::new(DenyAll), Arc::new(DenyAll)), - vec![Mailbox::new("alias", Arc::new(DenyAll), Arc::new(DenyAll))], + Mailbox::new( + "main", + Some(AddressMetadata { + is_terminal: false, + attributes: vec![("TEST_KEY".to_string(), "TEST_VALUE".to_string())], + }), + Arc::new(DenyAll), + Arc::new(DenyAll), + ), + vec![Mailbox::new( + "alias", + Some(AddressMetadata { + is_terminal: false, + attributes: vec![("TEST_KEY_2".to_string(), "TEST_VALUE_2".to_string())], + }), + Arc::new(DenyAll), + Arc::new(DenyAll), + )], )) - .with_metadata_attribute("main", "TEST_KEY", "TEST_VALUE") - .with_metadata_attribute("alias", "TEST_KEY_2", "TEST_VALUE_2") .start(context) .await?; - let meta = context.get_metadata("alias").unwrap(); + let meta = context.get_metadata("alias")?.unwrap(); assert!(!meta.is_terminal); @@ -251,9 +254,9 @@ async fn provide_and_read_address_metadata_worker_alias(context: &mut Context) - vec![("TEST_KEY_2".to_string(), "TEST_VALUE_2".to_string())] ); - context.stop_worker("main").await?; + context.stop_address("main")?; ockam_node::compat::tokio::time::sleep(std::time::Duration::from_millis(10)).await; - assert_eq!(context.get_metadata("alias"), None); + assert_eq!(context.get_metadata("alias")?, None); - context.stop().await + Ok(()) } diff --git a/implementations/rust/ockam/ockam_node/tests/tests.rs b/implementations/rust/ockam/ockam_node/tests/tests.rs index 2bb8f7f98a9..e0a1a9ec30f 100644 --- a/implementations/rust/ockam/ockam_node/tests/tests.rs +++ b/implementations/rust/ockam/ockam_node/tests/tests.rs @@ -14,6 +14,7 @@ use serde::{Deserialize, Serialize}; use std::sync::atomic::{AtomicI8, AtomicU32}; use std::time::{SystemTime, UNIX_EPOCH}; use tokio::time::sleep; +use tracing::debug; #[allow(non_snake_case)] #[ockam_macros::test] @@ -38,7 +39,7 @@ async fn receive_timeout__1_sec__should_return_from_call(ctx: &mut Context) -> R #[allow(non_snake_case)] #[test] fn start_and_shutdown_node__many_iterations__should_not_fail() { - for _ in 0..100 { + for _ in 0..1 { let (ctx, mut executor) = NodeBuilder::new().build(); executor .execute(async move { @@ -65,6 +66,7 @@ fn start_and_shutdown_node__many_iterations__should_not_fail() { .unwrap() } } + struct SimpleWorker { initialize_was_called: Arc, shutdown_was_called: Arc, @@ -79,6 +81,8 @@ impl Worker for SimpleWorker { self.initialize_was_called.store(true, Ordering::Relaxed); assert!(self.initialize_was_called.load(Ordering::Relaxed)); + debug!("INITIALIZE"); + Ok(()) } @@ -87,6 +91,8 @@ impl Worker for SimpleWorker { assert!(self.initialize_was_called.load(Ordering::Relaxed)); assert!(self.shutdown_was_called.load(Ordering::Relaxed)); + debug!("SHUTDOWN"); + Ok(()) } @@ -99,6 +105,30 @@ impl Worker for SimpleWorker { } } +#[allow(non_snake_case)] +#[ockam_macros::test] +async fn simple_worker__run_node_lifecycle__should_not_fail(ctx: &mut Context) -> Result<()> { + let initialize_was_called = Arc::new(AtomicBool::new(false)); + let shutdown_was_called = Arc::new(AtomicBool::new(false)); + + let initialize_was_called_clone = initialize_was_called.clone(); + let shutdown_was_called_clone = shutdown_was_called.clone(); + + let worker = SimpleWorker { + initialize_was_called: initialize_was_called_clone, + shutdown_was_called: shutdown_was_called_clone, + }; + + ctx.start_worker("simple_worker", worker).await?; + + let msg: String = ctx + .send_and_receive(route!["simple_worker"], "Hello".to_string()) + .await?; + assert_eq!(msg, "Hello"); + + Ok(()) +} + #[allow(non_snake_case)] #[ockam_macros::test] async fn simple_worker__run_node_lifecycle__worker_lifecycle_should_be_full( @@ -124,11 +154,11 @@ async fn simple_worker__run_node_lifecycle__worker_lifecycle_should_be_full( ctx.stop().await?; // Wait till tokio Runtime is shut down - // std::thread::sleep(Duration::new(1, 0)); sleep(Duration::new(1, 0)).await; assert!(initialize_was_called.load(Ordering::Relaxed)); assert!(shutdown_was_called.load(Ordering::Relaxed)); + Ok(()) } @@ -172,7 +202,7 @@ async fn worker_initialize_fail_should_shutdown(ctx: &mut Context) -> Result<()> sleep(Duration::new(1, 0)).await; assert!(shutdown_was_called.load(Ordering::Relaxed)); - assert!(!ctx.list_workers().contains(&address)); + assert!(!ctx.list_workers()?.contains(&address)); Ok(()) } @@ -207,7 +237,7 @@ async fn processor_initialize_fail_should_shutdown(ctx: &mut Context) -> Result< assert!(res.is_ok()); sleep(Duration::new(1, 0)).await; assert!(shutdown_was_called.load(Ordering::Relaxed)); - assert!(!ctx.list_workers().contains(&address)); + assert!(!ctx.list_workers()?.contains(&address)); Ok(()) } @@ -336,7 +366,7 @@ async fn waiting_processor__shutdown__should_be_interrupted(ctx: &mut Context) - ctx.start_processor("waiting_processor", processor).await?; sleep(Duration::from_secs(1)).await; - ctx.stop_processor("waiting_processor").await?; + ctx.stop_address("waiting_processor")?; sleep(Duration::from_secs(1)).await; assert!(initialize_was_called.load(Ordering::Relaxed)); @@ -403,13 +433,15 @@ async fn waiting_processor__messaging__should_work(ctx: &mut Context) -> Result< ctx.start_processor_with_access_control("messaging_processor", processor, AllowAll, AllowAll) .await?; - sleep(Duration::new(1, 0)).await; + sleep(Duration::from_millis(250)).await; let msg: String = ctx .send_and_receive(route!["messaging_processor"], "Keep working".to_string()) .await?; assert_eq!("OK", msg); + sleep(Duration::from_millis(250)).await; + assert!(initialize_was_called.load(Ordering::Relaxed)); assert!(!shutdown_was_called.load(Ordering::Relaxed)); @@ -418,6 +450,8 @@ async fn waiting_processor__messaging__should_work(ctx: &mut Context) -> Result< .await?; assert_eq!("I go home", msg); + sleep(Duration::from_millis(250)).await; + assert!(initialize_was_called.load(Ordering::Relaxed)); assert!(shutdown_was_called.load(Ordering::Relaxed)); @@ -462,7 +496,7 @@ impl Worker for StopFromHandleMessageWorker { type Context = Context; async fn handle_message(&mut self, ctx: &mut Context, _msg: Routed) -> Result<()> { self.counter_a.fetch_add(1, Ordering::Relaxed); - ctx.stop_worker(ctx.address()).await?; + ctx.stop_address(ctx.primary_address())?; self.counter_b.fetch_add(1, Ordering::Relaxed); Ok(()) } @@ -614,7 +648,7 @@ impl Worker for CountingErrorWorker { ) -> Result<()> { let _ = self.counter.fetch_add(1, Ordering::Relaxed); - Err(ockam_core::Error::new(Origin::Core, Kind::Misuse, "")) + Err(ockam_core::Error::new(Origin::Core, Kind::Misuse, "test")) } } @@ -646,3 +680,19 @@ async fn message_handle__error_during_handling__keep_worker_running( Ok(()) } + +#[test] +fn test1() { + let (ctx, mut executor) = NodeBuilder::new().build(); + + executor + .execute::<_, (), ockam_core::Error>(async move { + ctx.sleep(Duration::from_secs(1)).await; + + ctx.stop().await?; + + Ok(()) + }) + .unwrap() + .unwrap(); +} diff --git a/implementations/rust/ockam/ockam_transport_ble/src/router/mod.rs b/implementations/rust/ockam/ockam_transport_ble/src/router/mod.rs index 4a90bac6cf9..05cefc8cca7 100644 --- a/implementations/rust/ockam/ockam_transport_ble/src/router/mod.rs +++ b/implementations/rust/ockam/ockam_transport_ble/src/router/mod.rs @@ -110,7 +110,7 @@ impl Worker for BleRouter { type Message = Any; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await?; + ctx.set_cluster(crate::CLUSTER_NAME)?; Ok(()) } @@ -167,9 +167,15 @@ impl BleRouter { // TODO: @ac let mailboxes = Mailboxes::new( - Mailbox::new(main_addr.clone(), Arc::new(AllowAll), Arc::new(AllowAll)), + Mailbox::new( + main_addr.clone(), + None, + Arc::new(AllowAll), + Arc::new(AllowAll), + ), vec![Mailbox::new( api_addr, + None, Arc::new(AllowAll), Arc::new(AllowAll), )], diff --git a/implementations/rust/ockam/ockam_transport_ble/src/workers/listener.rs b/implementations/rust/ockam/ockam_transport_ble/src/workers/listener.rs index 697c471f774..5696bac7740 100644 --- a/implementations/rust/ockam/ockam_transport_ble/src/workers/listener.rs +++ b/implementations/rust/ockam/ockam_transport_ble/src/workers/listener.rs @@ -56,7 +56,7 @@ where type Context = Context; async fn initialize(&mut self, ctx: &mut Self::Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await + ctx.set_cluster(crate::CLUSTER_NAME) } async fn process(&mut self, ctx: &mut Self::Context) -> Result { diff --git a/implementations/rust/ockam/ockam_transport_ble/src/workers/receiver.rs b/implementations/rust/ockam/ockam_transport_ble/src/workers/receiver.rs index 8c8fe2bc998..0d762c4163f 100644 --- a/implementations/rust/ockam/ockam_transport_ble/src/workers/receiver.rs +++ b/implementations/rust/ockam/ockam_transport_ble/src/workers/receiver.rs @@ -39,7 +39,7 @@ where type Context = Context; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await + ctx.set_cluster(crate::CLUSTER_NAME) } async fn process(&mut self, ctx: &mut Context) -> Result { diff --git a/implementations/rust/ockam/ockam_transport_ble/src/workers/sender.rs b/implementations/rust/ockam/ockam_transport_ble/src/workers/sender.rs index fdc1a94a855..ca59918f006 100644 --- a/implementations/rust/ockam/ockam_transport_ble/src/workers/sender.rs +++ b/implementations/rust/ockam/ockam_transport_ble/src/workers/sender.rs @@ -83,7 +83,7 @@ where type Message = TransportMessage; async fn initialize(&mut self, ctx: &mut Self::Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await?; + ctx.set_cluster(crate::CLUSTER_NAME)?; debug!("initialize for peer: {:?}", self.peer); @@ -135,7 +135,7 @@ where Ok(_) => (), Err(e) => { error!("Failed to send fragment to peer {}: {:?}", self.peer, e); - ctx.stop_worker(ctx.address()).await?; + ctx.stop_address(ctx.primary_address())?; } } @@ -147,7 +147,7 @@ where Ok(_) => (), Err(e) => { error!("Failed to send fragment to peer {}: {:?}", self.peer, e); - ctx.stop_worker(ctx.address()).await?; + ctx.stop_address(ctx.primary_address())?; } } diff --git a/implementations/rust/ockam/ockam_transport_core/src/transport.rs b/implementations/rust/ockam/ockam_transport_core/src/transport.rs index 51556c07225..4af80cb4137 100644 --- a/implementations/rust/ockam/ockam_transport_core/src/transport.rs +++ b/implementations/rust/ockam/ockam_transport_core/src/transport.rs @@ -19,7 +19,7 @@ pub trait Transport: Send + Sync + 'static { async fn resolve_address(&self, address: Address) -> Result
; /// Stop all workers and free all resources associated with the connection - async fn disconnect(&self, address: Address) -> Result<()>; + fn disconnect(&self, address: Address) -> Result<()>; } /// Helper that creates a length-prefixed buffer containing the given diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/portal/inlet_listener.rs b/implementations/rust/ockam/ockam_transport_tcp/src/portal/inlet_listener.rs index 07e2941b9e6..57fb9f2bd38 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/portal/inlet_listener.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/portal/inlet_listener.rs @@ -4,11 +4,10 @@ use crate::portal::{InletSharedState, ReadHalfMaybeTls, WriteHalfMaybeTls}; use crate::{portal::TcpPortalWorker, TcpInlet, TcpInletOptions, TcpRegistry}; use log::warn; use ockam_core::compat::net::SocketAddr; -use ockam_core::compat::sync::Arc; +use ockam_core::compat::sync::{Arc, RwLock as SyncRwLock}; use ockam_core::errcode::{Kind, Origin}; use ockam_core::{async_trait, compat::boxed::Box, Result}; use ockam_core::{Address, Processor, Route}; -use ockam_node::compat::asynchronous::RwLock; use ockam_node::Context; use ockam_transport_core::{HostnamePort, TransportError}; use rustls::pki_types::CertificateDer; @@ -27,7 +26,7 @@ use tracing::{debug, error, instrument}; pub(crate) struct TcpInletListenProcessor { registry: TcpRegistry, inner: TcpListener, - inlet_shared_state: Arc>, + inlet_shared_state: Arc>, options: TcpInletOptions, } @@ -35,7 +34,7 @@ impl TcpInletListenProcessor { pub fn new( registry: TcpRegistry, inner: TcpListener, - inlet_shared_state: Arc>, + inlet_shared_state: Arc>, options: TcpInletOptions, ) -> Self { Self { @@ -67,8 +66,8 @@ impl TcpInletListenProcessor { }; let socket_addr = inner.local_addr().map_err(TransportError::from)?; let inlet_shared_state = - InletSharedState::create(ctx, outlet_listener_route, options.is_paused).await?; - let inlet_shared_state = Arc::new(RwLock::new(inlet_shared_state)); + InletSharedState::create(ctx, outlet_listener_route, options.is_paused)?; + let inlet_shared_state = Arc::new(SyncRwLock::new(inlet_shared_state)); let processor = Self::new(registry, inner, inlet_shared_state.clone(), options); ctx.start_processor(processor_address.clone(), processor) @@ -162,7 +161,8 @@ impl Processor for TcpInletListenProcessor { #[instrument(skip_all, name = "TcpInletListenProcessor::initialize")] async fn initialize(&mut self, ctx: &mut Self::Context) -> Result<()> { - self.registry.add_inlet_listener_processor(&ctx.address()); + self.registry + .add_inlet_listener_processor(ctx.primary_address()); Ok(()) } @@ -170,7 +170,7 @@ impl Processor for TcpInletListenProcessor { #[instrument(skip_all, name = "TcpInletListenProcessor::shutdown")] async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { self.registry - .remove_inlet_listener_processor(&ctx.address()); + .remove_inlet_listener_processor(ctx.primary_address()); Ok(()) } @@ -181,7 +181,7 @@ impl Processor for TcpInletListenProcessor { let addresses = Addresses::generate(PortalType::Inlet); - let inlet_shared_state = self.inlet_shared_state.read().await.clone(); + let inlet_shared_state = self.inlet_shared_state.read().unwrap().clone(); if inlet_shared_state.is_paused() { // Just drop the stream diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/portal/inlet_shared_state.rs b/implementations/rust/ockam/ockam_transport_tcp/src/portal/inlet_shared_state.rs index 4322b3458b8..799533de155 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/portal/inlet_shared_state.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/portal/inlet_shared_state.rs @@ -19,10 +19,10 @@ pub struct InletSharedState { } impl InletSharedState { - pub async fn create(ctx: &Context, route: Route, is_paused: bool) -> Result { + pub fn create(ctx: &Context, route: Route, is_paused: bool) -> Result { let their_identifier = - if let Some(terminal) = ctx.find_terminal_address(route.clone()).await? { - SecureChannelMetadata::from_terminal_address(&terminal) + if let Some((_address, metadata)) = ctx.find_terminal_address(route.iter())? { + SecureChannelMetadata::from_terminal_address_metadata(&metadata) .map(|m| m.their_identifier()) .ok() } else { @@ -53,10 +53,10 @@ impl InletSharedState { self.route_index } - pub async fn update_route(&mut self, ctx: &Context, new_route: Route) -> Result<()> { + pub fn update_route(&mut self, ctx: &Context, new_route: Route) -> Result<()> { let their_identifier = - if let Some(terminal) = ctx.find_terminal_address(new_route.clone()).await? { - SecureChannelMetadata::from_terminal_address(&terminal) + if let Some((_address, metadata)) = ctx.find_terminal_address(new_route.iter())? { + SecureChannelMetadata::from_terminal_address_metadata(&metadata) .map(|m| m.their_identifier()) .ok() } else { diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/portal/interceptor.rs b/implementations/rust/ockam/ockam_transport_tcp/src/portal/interceptor.rs index cb44928531b..30cf3437993 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/portal/interceptor.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/portal/interceptor.rs @@ -295,13 +295,11 @@ impl Worker for PortalInterceptorWorker { if !disconnect_received { debug!( "{:?} received disconnect event from {:?}", - context.address(), + context.primary_address(), return_route ); - context - .stop_worker(self.other_worker_address.clone()) - .await?; - context.stop_worker(context.address()).await?; + context.stop_address(self.other_worker_address.clone())?; + context.stop_address(context.primary_address())?; } } PortalMessage::Ping => self.forward(context, routed_message).await?, diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/portal/outlet_listener.rs b/implementations/rust/ockam/ockam_transport_tcp/src/portal/outlet_listener.rs index 52d2666dabe..95e5fec2d13 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/portal/outlet_listener.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/portal/outlet_listener.rs @@ -59,14 +59,16 @@ impl Worker for TcpOutletListenWorker { #[instrument(skip_all, name = "TcpOutletListenWorker::initialize")] async fn initialize(&mut self, ctx: &mut Self::Context) -> Result<()> { - self.registry.add_outlet_listener_worker(&ctx.address()); + self.registry + .add_outlet_listener_worker(ctx.primary_address()); Ok(()) } #[instrument(skip_all, name = "TcpOutletListenWorker::shutdown")] async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { - self.registry.remove_outlet_listener_worker(&ctx.address()); + self.registry + .remove_outlet_listener_worker(ctx.primary_address()); Ok(()) } diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/portal/portal_receiver.rs b/implementations/rust/ockam/ockam_transport_tcp/src/portal/portal_receiver.rs index 6210a023038..27da7989f4c 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/portal/portal_receiver.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/portal/portal_receiver.rs @@ -52,7 +52,8 @@ impl Processor for TcpPortalRecvPr #[instrument(skip_all, name = "TcpPortalRecvProcessor::initialize")] async fn initialize(&mut self, ctx: &mut Self::Context) -> Result<()> { - self.registry.add_portal_receiver_processor(&ctx.address()); + self.registry + .add_portal_receiver_processor(ctx.primary_address()); Ok(()) } @@ -60,7 +61,7 @@ impl Processor for TcpPortalRecvPr #[instrument(skip_all, name = "TcpPortalRecvProcessor::shutdown")] async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { self.registry - .remove_portal_receiver_processor(&ctx.address()); + .remove_portal_receiver_processor(ctx.primary_address()); Ok(()) } diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/portal/portal_worker.rs b/implementations/rust/ockam/ockam_transport_tcp/src/portal/portal_worker.rs index ef4c9f5049a..9f6e2cddd4d 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/portal/portal_worker.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/portal/portal_worker.rs @@ -177,12 +177,14 @@ impl TcpPortalWorker { let internal_mailbox = Mailbox::new( addresses.sender_internal, + None, Arc::new(AllowSourceAddress(addresses.receiver_internal)), Arc::new(DenyAll), ); let remote_mailbox = Mailbox::new( addresses.sender_remote, + None, incoming_access_control, outgoing_access_control, ); @@ -238,12 +240,14 @@ impl TcpPortalWorker { let remote = Mailbox::new( self.addresses.receiver_remote.clone(), + None, Arc::new(DenyAll), self.outgoing_access_control.clone(), ); let internal = Mailbox::new( self.addresses.receiver_internal.clone(), + None, Arc::new(DenyAll), Arc::new(AllowOnwardAddress(self.addresses.sender_internal.clone())), ); @@ -278,10 +282,9 @@ impl TcpPortalWorker { } #[instrument(skip_all)] - async fn stop_receiver(&self, ctx: &Context) -> Result<()> { + fn stop_receiver(&self, ctx: &Context) -> Result<()> { if ctx - .stop_processor(self.addresses.receiver_remote.clone()) - .await + .stop_address(self.addresses.receiver_remote.clone()) .is_ok() { debug!( @@ -295,9 +298,8 @@ impl TcpPortalWorker { } #[instrument(skip_all)] - async fn stop_sender(&self, ctx: &Context) -> Result<()> { - ctx.stop_worker(self.addresses.sender_internal.clone()) - .await + fn stop_sender(&self, ctx: &Context) -> Result<()> { + ctx.stop_address(self.addresses.sender_internal.clone()) } /// Start the portal disconnection process @@ -314,20 +316,20 @@ impl TcpPortalWorker { // connection and shut down both processor and worker DisconnectionReason::FailedTx => { self.notify_remote_about_disconnection(ctx).await?; - self.stop_receiver(ctx).await?; + self.stop_receiver(ctx)?; // Sleep, so that if connection is dropped on both sides at the same time, the other // side had time to notify us about the closure. Otherwise, the message won't be // delivered which can lead to a warning message from a secure channel (or whatever // is used to deliver the message). Can be removed though ctx.sleep(Duration::from_secs(2)).await; - self.stop_sender(ctx).await?; + self.stop_sender(ctx)?; } // Packets were dropped while traveling to us, let's notify the other end about dropped // connection and DisconnectionReason::InvalidCounter => { self.notify_remote_about_disconnection(ctx).await?; - self.stop_receiver(ctx).await?; - self.stop_sender(ctx).await?; + self.stop_receiver(ctx)?; + self.stop_sender(ctx)?; } // We couldn't read data from the tcp connection // Receiver should have already notified the other end and should shut down itself @@ -337,13 +339,13 @@ impl TcpPortalWorker { // delivered which can lead to a warning message from a secure channel (or whatever // is used to deliver the message). Can be removed though ctx.sleep(Duration::from_secs(2)).await; - self.stop_sender(ctx).await?; + self.stop_sender(ctx)?; } // Other end notifies us that the tcp connection is dropped // Let's shut down both processor and worker DisconnectionReason::Remote => { - self.stop_receiver(ctx).await?; - self.stop_sender(ctx).await?; + self.stop_receiver(ctx)?; + self.stop_sender(ctx)?; } } diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/privileged_portals.rs b/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/privileged_portals.rs index 93b73a65326..ca4487e69d9 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/privileged_portals.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/privileged_portals.rs @@ -6,12 +6,12 @@ use caps::{CapSet, Capability}; use core::fmt::Debug; use log::{debug, error}; use nix::unistd::Uid; +use ockam_core::compat::sync::{Arc, RwLock as SyncRwLock}; use ockam_core::{Address, DenyAll, Result, Route}; -use ockam_node::compat::asynchronous::{resolve_peer, RwLock}; +use ockam_node::compat::asynchronous::resolve_peer; use ockam_node::{ProcessorBuilder, WorkerBuilder}; use ockam_transport_core::{HostnamePort, TransportError}; use std::net::IpAddr; -use std::sync::Arc; use tokio::net::TcpListener; use tokio::sync::mpsc::channel; use tracing::instrument; @@ -90,9 +90,8 @@ impl TcpTransport { let tcp_packet_writer = self.start_raw_socket_processor_if_needed().await?; - let inlet_shared_state = - InletSharedState::create(self.ctx(), outlet_route.clone(), false).await?; - let inlet_shared_state = Arc::new(RwLock::new(inlet_shared_state)); + let inlet_shared_state = InletSharedState::create(self.ctx(), outlet_route.clone(), false)?; + let inlet_shared_state = Arc::new(SyncRwLock::new(inlet_shared_state)); let remote_worker_address = Address::random_tagged("Ebpf.RemoteWorker.Inlet"); let internal_worker_address = Address::random_tagged("Ebpf.InternalWorker.Inlet"); @@ -229,7 +228,7 @@ impl TcpTransport { &self, addr: impl Into
+ Clone + Debug, ) -> Result<()> { - self.ctx().stop_worker(addr).await?; + self.ctx().stop_address(addr)?; // TODO: eBPF Remove from the registry // self.ebpf_support.outlet_registry diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/registry/inlet.rs b/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/registry/inlet.rs index 628c4edfbfe..9fbc71e4d44 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/registry/inlet.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/registry/inlet.rs @@ -1,10 +1,8 @@ use crate::portal::InletSharedState; use crate::privileged_portal::packet::RawSocketReadResult; use crate::privileged_portal::{ConnectionIdentifier, Port}; -use ockam_core::compat::sync::Arc; -use ockam_core::compat::sync::RwLock as SyncRwLock; +use ockam_core::compat::sync::{Arc, RwLock as SyncRwLock}; use ockam_core::{Address, LocalInfoIdentifier}; -use ockam_node::compat::asynchronous::RwLock as AsyncRwLock; use std::collections::HashMap; use std::net::Ipv4Addr; use tokio::net::TcpListener; @@ -32,7 +30,7 @@ impl InletRegistry { sender: Sender, port: Port, tcp_listener: TcpListener, - inlet_shared_state: Arc>, + inlet_shared_state: Arc>, ) -> Inlet { let mut inlets = self.inlets.write().unwrap(); @@ -72,7 +70,7 @@ pub struct Inlet { /// Port pub port: Port, /// Route to the corresponding Outlet - pub inlet_shared_state: Arc>, + pub inlet_shared_state: Arc>, /// Hold to mark the port as taken pub _tcp_listener: Arc, /// Same map with different key diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/workers/internal_processor.rs b/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/workers/internal_processor.rs index 70528406221..bcd15dcc3cb 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/workers/internal_processor.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/privileged_portal/workers/internal_processor.rs @@ -72,7 +72,7 @@ impl Processor for InternalProcessor { match &self.mode { // Client -> Inlet packet PortalMode::Inlet { inlet } => { - let inlet_shared_state = inlet.inlet_shared_state.read().await.clone(); + let inlet_shared_state = inlet.inlet_shared_state.read().unwrap().clone(); if inlet_shared_state.is_paused() { return Ok(true); @@ -133,7 +133,7 @@ impl Processor for InternalProcessor { .with_onward_route(inlet_shared_state.route().clone()) .with_return_route(route![inlet.remote_worker_address.clone()]) .with_payload(cbor_encode_preallocate(&portal_packet)?), - ctx.address(), + ctx.primary_address().clone(), ) .await?; } @@ -174,7 +174,7 @@ impl Processor for InternalProcessor { .with_onward_route(return_route) .with_return_route(route![outlet.remote_worker_address.clone()]) .with_payload(cbor_encode_preallocate(&portal_packet)?), - ctx.address(), + ctx.primary_address().clone(), ) .await?; } diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/transport/connection.rs b/implementations/rust/ockam/ockam_transport_tcp/src/transport/connection.rs index 12ea8729c50..33909e72f6c 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/transport/connection.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/transport/connection.rs @@ -58,8 +58,8 @@ impl TcpConnection { /// Stops the [`TcpConnection`], this method must be called to avoid /// leakage of the connection. /// Simply dropping this object won't close the connection - pub async fn stop(&self, context: &Context) -> Result<()> { - context.stop_worker(self.sender_address.clone()).await + pub fn stop(&self, context: &Context) -> Result<()> { + context.stop_address(self.sender_address.clone()) } /// Corresponding [`TcpSendWorker`](super::workers::TcpSendWorker) [`Address`] that can be used /// in a route to send messages to the other side of the TCP connection @@ -151,7 +151,7 @@ impl TcpTransport { } /// Interrupt an active TCP connection given its Sender `Address` - pub async fn disconnect(&self, address: impl Into
) -> Result<()> { - self.ctx.stop_worker(address.into()).await + pub fn disconnect(&self, address: impl Into
) -> Result<()> { + self.ctx.stop_address(address.into()) } } diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/transport/lifecycle.rs b/implementations/rust/ockam/ockam_transport_tcp/src/transport/lifecycle.rs index 96e555c1019..bd4fb2d9122 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/transport/lifecycle.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/transport/lifecycle.rs @@ -129,8 +129,8 @@ impl Transport for TcpTransport { } } - async fn disconnect(&self, address: Address) -> Result<()> { - self.disconnect(address).await + fn disconnect(&self, address: Address) -> Result<()> { + self.disconnect(address) } } @@ -145,7 +145,7 @@ mod tests { async fn test_resolve_address(ctx: &mut Context) -> Result<()> { let tcp = TcpTransport::create(ctx).await?; let tcp_address = "127.0.0.1:0"; - let initial_workers = ctx.list_workers(); + let initial_workers = ctx.list_workers()?; let listener = TcpListener::bind(tcp_address) .await .map_err(TransportError::from)?; @@ -163,7 +163,7 @@ mod tests { .await?; // there are 2 additional workers - let mut additional_workers = ctx.list_workers(); + let mut additional_workers = ctx.list_workers()?; additional_workers.retain(|w| !initial_workers.contains(w)); assert_eq!(additional_workers.len(), 2); diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/transport/listener.rs b/implementations/rust/ockam/ockam_transport_tcp/src/transport/listener.rs index 1f2b7d8647d..841fee2616e 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/transport/listener.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/transport/listener.rs @@ -87,7 +87,7 @@ impl TcpTransport { } /// Interrupt an active TCP listener given its `Address` - pub async fn stop_listener(&self, address: &Address) -> Result<()> { - self.ctx.stop_processor(address.clone()).await + pub fn stop_listener(&self, address: &Address) -> Result<()> { + self.ctx.stop_address(address.clone()) } } diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/transport/portals.rs b/implementations/rust/ockam/ockam_transport_tcp/src/transport/portals.rs index 9909ff6549d..e0e0a2206b0 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/transport/portals.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/transport/portals.rs @@ -3,10 +3,9 @@ use crate::{portal::TcpOutletListenWorker, TcpInletOptions, TcpOutletOptions, Tc use core::fmt; use core::fmt::{Debug, Formatter}; use ockam_core::compat::net::SocketAddr; -use ockam_core::compat::sync::Arc; +use ockam_core::compat::sync::{Arc, RwLock as SyncRwLock}; use ockam_core::flow_control::FlowControls; use ockam_core::{route, Address, Result, Route}; -use ockam_node::compat::asynchronous::RwLock; use ockam_node::Context; use ockam_transport_core::{parse_socket_addr, HostnamePort}; use tracing::instrument; @@ -26,7 +25,7 @@ impl TcpTransport { /// /// let tcp = TcpTransport::create(&ctx).await?; /// tcp.create_inlet("inlet", route_path, TcpInletOptions::new()).await?; - /// # tcp.stop_inlet("inlet").await?; + /// # tcp.stop_inlet("inlet")?; /// # Ok(()) } /// ``` #[instrument(skip(self), fields(address = ? bind_addr.clone().into(), outlet_route = ? outlet_route.clone()))] @@ -58,12 +57,12 @@ impl TcpTransport { /// /// let tcp = TcpTransport::create(&ctx).await?; /// tcp.create_inlet("inlet", route, TcpInletOptions::new()).await?; - /// tcp.stop_inlet("inlet").await?; + /// tcp.stop_inlet("inlet")?; /// # Ok(()) } /// ``` #[instrument(skip(self), fields(address = ? addr.clone().into()))] - pub async fn stop_inlet(&self, addr: impl Into
+ Clone + Debug) -> Result<()> { - self.ctx.stop_processor(addr).await?; + pub fn stop_inlet(&self, addr: impl Into
+ Clone + Debug) -> Result<()> { + self.ctx.stop_address(addr)?; Ok(()) } @@ -84,7 +83,7 @@ impl TcpTransport { /// /// let tcp = TcpTransport::create(&ctx).await?; /// tcp.create_outlet("outlet", HostnamePort::new("localhost", 9000), TcpOutletOptions::new()).await?; - /// # tcp.stop_outlet("outlet").await?; + /// # tcp.stop_outlet("outlet")?; /// # Ok(()) } /// ``` #[instrument(skip(self), fields(address = ? address.clone().into(), peer=peer.clone().to_string()))] @@ -117,12 +116,12 @@ impl TcpTransport { /// /// let tcp = TcpTransport::create(&ctx).await?; /// tcp.create_outlet("outlet", HostnamePort::new("127.0.0.1", 5000), TcpOutletOptions::new()).await?; - /// tcp.stop_outlet("outlet").await?; + /// tcp.stop_outlet("outlet")?; /// # Ok(()) } /// ``` #[instrument(skip(self), fields(address = % addr.clone().into()))] - pub async fn stop_outlet(&self, addr: impl Into
+ Clone + Debug) -> Result<()> { - self.ctx.stop_worker(addr).await?; + pub fn stop_outlet(&self, addr: impl Into
+ Clone + Debug) -> Result<()> { + self.ctx.stop_address(addr)?; Ok(()) } } @@ -131,7 +130,7 @@ impl TcpTransport { #[derive(Clone, Debug)] pub struct TcpInlet { socket_address: SocketAddr, - inlet_shared_state: Arc>, + inlet_shared_state: Arc>, state: TcpInletState, } @@ -169,7 +168,7 @@ impl TcpInlet { pub fn new_regular( socket_address: SocketAddr, processor_address: Address, - inlet_shared_state: Arc>, + inlet_shared_state: Arc>, ) -> Self { Self { socket_address, @@ -182,7 +181,7 @@ impl TcpInlet { pub fn new_privileged( socket_address: SocketAddr, portal_worker_address: Address, - inlet_shared_state: Arc>, + inlet_shared_state: Arc>, ) -> Self { Self { socket_address, @@ -223,12 +222,12 @@ impl TcpInlet { /// only newly accepted connections will use the new route. /// For privileged Portals old connections can continue work in case the Identifier of the /// Outlet node didn't change - pub async fn update_outlet_node_route(&self, ctx: &Context, new_route: Route) -> Result<()> { - let mut inlet_shared_state = self.inlet_shared_state.write().await; + pub fn update_outlet_node_route(&self, ctx: &Context, new_route: Route) -> Result<()> { + let mut inlet_shared_state = self.inlet_shared_state.write().unwrap(); let new_route = Self::build_new_full_route(new_route, inlet_shared_state.route())?; let next = new_route.next()?.clone(); - inlet_shared_state.update_route(ctx, new_route).await?; + inlet_shared_state.update_route(ctx, new_route)?; self.update_flow_controls(ctx.flow_controls(), next); @@ -236,8 +235,8 @@ impl TcpInlet { } /// Pause TCP Inlet, all incoming TCP streams will be dropped. - pub async fn pause(&self) { - let mut inlet_shared_state = self.inlet_shared_state.write().await; + pub fn pause(&self) { + let mut inlet_shared_state = self.inlet_shared_state.write().unwrap(); inlet_shared_state.set_is_paused(true); } @@ -258,13 +257,13 @@ impl TcpInlet { } /// Unpause TCP Inlet and update the outlet route. - pub async fn unpause(&self, ctx: &Context, new_route: Route) -> Result<()> { - let mut inlet_shared_state = self.inlet_shared_state.write().await; + pub fn unpause(&self, ctx: &Context, new_route: Route) -> Result<()> { + let mut inlet_shared_state = self.inlet_shared_state.write().unwrap(); let new_route = Self::build_new_full_route(new_route, inlet_shared_state.route())?; let next = new_route.next()?.clone(); - inlet_shared_state.update_route(ctx, new_route).await?; + inlet_shared_state.update_route(ctx, new_route)?; inlet_shared_state.set_is_paused(false); self.update_flow_controls(ctx.flow_controls(), next); @@ -273,13 +272,13 @@ impl TcpInlet { } /// Stop the Inlet - pub async fn stop(&self, ctx: &Context) -> Result<()> { + pub fn stop(&self, ctx: &Context) -> Result<()> { match &self.state { TcpInletState::Privileged { .. } => { // TODO: eBPF } TcpInletState::Regular { processor_address } => { - ctx.stop_processor(processor_address.clone()).await?; + ctx.stop_address(processor_address.clone())?; } } diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/workers/listener.rs b/implementations/rust/ockam/ockam_transport_tcp/src/workers/listener.rs index 39034526dfd..563569c0d92 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/workers/listener.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/workers/listener.rs @@ -55,10 +55,10 @@ impl Processor for TcpListenProcessor { #[instrument(skip_all, name = "TcpListenProcessor::initialize")] async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await?; + ctx.set_cluster(crate::CLUSTER_NAME)?; self.registry.add_listener_processor(TcpListenerInfo::new( - ctx.address(), + ctx.primary_address().clone(), self.socket_address, self.options.flow_control_id.clone(), )); @@ -68,7 +68,8 @@ impl Processor for TcpListenProcessor { #[instrument(skip_all, name = "TcpListenProcessor::shutdown")] async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { - self.registry.remove_listener_processor(&ctx.address()); + self.registry + .remove_listener_processor(ctx.primary_address()); Ok(()) } diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/workers/receiver.rs b/implementations/rust/ockam/ockam_transport_tcp/src/workers/receiver.rs index 13b8e969319..914458a908b 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/workers/receiver.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/workers/receiver.rs @@ -80,11 +80,13 @@ impl TcpRecvProcessor { let mailbox = Mailbox::new( addresses.receiver_address().clone(), + None, Arc::new(DenyAll), receiver_outgoing_access_control, ); let internal = Mailbox::new( addresses.receiver_internal_address().clone(), + None, Arc::new(DenyAll), Arc::new(AllowOnwardAddress( addresses.sender_internal_address().clone(), @@ -119,10 +121,10 @@ impl Processor for TcpRecvProcessor { #[instrument(skip_all, name = "TcpRecvProcessor::initialize")] async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await?; + ctx.set_cluster(crate::CLUSTER_NAME)?; self.registry.add_receiver_processor(TcpReceiverInfo::new( - ctx.address(), + ctx.primary_address().clone(), self.addresses.sender_address().clone(), self.socket_address, self.mode, @@ -158,7 +160,8 @@ impl Processor for TcpRecvProcessor { #[instrument(skip_all, name = "TcpRecvProcessor::shutdown")] async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { - self.registry.remove_receiver_processor(&ctx.address()); + self.registry + .remove_receiver_processor(ctx.primary_address()); Ok(()) } @@ -174,7 +177,7 @@ impl Processor for TcpRecvProcessor { /// Context to avoid spawning a zombie task. /// 3. We must also stop the TcpReceive loop when the worker gets /// killed by the user or node. - #[instrument(skip_all, name = "TcpRecvProcessor::process", fields(worker = %ctx.address()))] + #[instrument(skip_all, name = "TcpRecvProcessor::process", fields(worker = %ctx.primary_address()))] async fn process(&mut self, ctx: &mut Context) -> Result { // Read the message length let len = match self.read_half.read_u32().await { diff --git a/implementations/rust/ockam/ockam_transport_tcp/src/workers/sender.rs b/implementations/rust/ockam/ockam_transport_tcp/src/workers/sender.rs index 9f5cbb0e9a1..4c7b98dc56f 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/src/workers/sender.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/src/workers/sender.rs @@ -4,7 +4,7 @@ use ockam_core::flow_control::FlowControlId; use ockam_core::{ async_trait, compat::{net::SocketAddr, sync::Arc}, - AllowAll, AllowSourceAddress, DenyAll, LocalMessage, + AddressMetadata, AllowAll, AllowSourceAddress, DenyAll, LocalMessage, }; use ockam_core::{Any, Decodable, Mailbox, Mailboxes, Message, Result, Routed, Worker}; use ockam_node::{Context, WorkerBuilder}; @@ -89,12 +89,17 @@ impl TcpSendWorker { let main_mailbox = Mailbox::new( addresses.sender_address().clone(), + Some(AddressMetadata { + is_terminal: true, + attributes: vec![], + }), Arc::new(AllowAll), Arc::new(DenyAll), ); let internal_mailbox = Mailbox::new( addresses.sender_internal_address().clone(), + None, Arc::new(AllowSourceAddress( addresses.receiver_internal_address().clone(), )), @@ -103,7 +108,6 @@ impl TcpSendWorker { WorkerBuilder::new(sender_worker) .with_mailboxes(Mailboxes::new(main_mailbox.clone(), vec![internal_mailbox])) - .terminal(addresses.sender_address().clone()) .start(ctx) .await?; @@ -111,9 +115,8 @@ impl TcpSendWorker { } #[instrument(skip_all, name = "TcpSendWorker::stop")] - async fn stop(&self, ctx: &Context) -> Result<()> { - ctx.stop_worker(self.addresses.sender_address().clone()) - .await?; + fn stop(&self, ctx: &Context) -> Result<()> { + ctx.stop_address(self.addresses.sender_address().clone())?; Ok(()) } @@ -166,7 +169,7 @@ impl Worker for TcpSendWorker { #[instrument(skip_all, name = "TcpSendWorker::initialize")] async fn initialize(&mut self, ctx: &mut Self::Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await?; + ctx.set_cluster(crate::CLUSTER_NAME)?; self.registry.add_sender_worker(TcpSenderInfo::new( self.addresses.sender_address().clone(), @@ -187,7 +190,7 @@ impl Worker for TcpSendWorker { "Failed to send protocol version to peer {}", self.socket_address ); - self.stop(ctx).await?; + self.stop(ctx)?; return Ok(()); } @@ -201,9 +204,7 @@ impl Worker for TcpSendWorker { .remove_sender_worker(self.addresses.sender_address()); if self.rx_should_be_stopped { - let _ = ctx - .stop_processor(self.addresses.receiver_address().clone()) - .await; + let _ = ctx.stop_address(self.addresses.receiver_address().clone()); } Ok(()) @@ -211,7 +212,7 @@ impl Worker for TcpSendWorker { // TcpSendWorker will receive messages from the TcpRouter to send // across the TcpStream to our friend - #[instrument(skip_all, name = "TcpSendWorker::handle_message", fields(worker = %ctx.address()))] + #[instrument(skip_all, name = "TcpSendWorker::handle_message", fields(worker = %ctx.primary_address()))] async fn handle_message( &mut self, ctx: &mut Context, @@ -230,7 +231,7 @@ impl Worker for TcpSendWorker { // No need to stop Receiver as it notified us about connection drop and will // stop itself self.rx_should_be_stopped = false; - self.stop(ctx).await?; + self.stop(ctx)?; return Ok(()); } @@ -243,14 +244,14 @@ impl Worker for TcpSendWorker { if let Err(err) = self.serialize_message(local_message) { // Close the stream - self.stop(ctx).await?; + self.stop(ctx)?; return Err(err); }; if self.write_half.write_all(&self.buffer).await.is_err() { warn!("Failed to send message to peer {}", self.socket_address); - self.stop(ctx).await?; + self.stop(ctx)?; return Ok(()); } diff --git a/implementations/rust/ockam/ockam_transport_tcp/tests/lifecycle.rs b/implementations/rust/ockam/ockam_transport_tcp/tests/lifecycle.rs index 93d6d2da98b..3ad2caa88d6 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/tests/lifecycle.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/tests/lifecycle.rs @@ -91,7 +91,7 @@ async fn tcp_lifecycle__disconnect__should_stop_worker(ctx: &mut Context) -> Res .await?; assert_eq!(reply2, msg2, "Should receive the same message"); - transport.disconnect(connection1.clone()).await?; + transport.disconnect(connection1.clone())?; let res = ctx .send(route![connection1.clone(), "echoer"], msg1.clone()) .await; @@ -102,7 +102,7 @@ async fn tcp_lifecycle__disconnect__should_stop_worker(ctx: &mut Context) -> Res .await?; assert_eq!(reply3, msg3, "Should receive the same message"); - transport.disconnect(connection2.clone()).await?; + transport.disconnect(connection2.clone())?; let res = ctx .send(route![connection2.clone(), "echoer"], msg3.clone()) .await; @@ -144,9 +144,7 @@ async fn tcp_lifecycle__stop_listener__should_stop_accepting_connections( .await?; assert_eq!(reply1, msg1, "Should receive the same message"); - transport - .stop_listener(listener.processor_address()) - .await?; + transport.stop_listener(listener.processor_address())?; ctx.sleep(Duration::from_millis(10)).await; let res = transport diff --git a/implementations/rust/ockam/ockam_transport_tcp/tests/portal.rs b/implementations/rust/ockam/ockam_transport_tcp/tests/portal.rs index 30f509eaf01..43a89c1defb 100644 --- a/implementations/rust/ockam/ockam_transport_tcp/tests/portal.rs +++ b/implementations/rust/ockam/ockam_transport_tcp/tests/portal.rs @@ -298,11 +298,9 @@ async fn portal__update_route__should_succeed(ctx: &mut Context) -> Result<()> { write_binary(&mut stream, payload1).await; read_assert_binary(&mut stream, payload2).await; - node_connection1.stop(ctx).await?; + node_connection1.stop(ctx)?; - inlet - .update_outlet_node_route(ctx, route![node_connection2]) - .await?; + inlet.update_outlet_node_route(ctx, route![node_connection2])?; let mut stream = TcpStream::connect(inlet.socket_address()).await.unwrap(); write_binary(&mut stream, payload1).await; diff --git a/implementations/rust/ockam/ockam_transport_udp/src/puncture/negotiation/listener.rs b/implementations/rust/ockam/ockam_transport_udp/src/puncture/negotiation/listener.rs index 84791d2cb8d..da0259730c5 100644 --- a/implementations/rust/ockam/ockam_transport_udp/src/puncture/negotiation/listener.rs +++ b/implementations/rust/ockam/ockam_transport_udp/src/puncture/negotiation/listener.rs @@ -75,7 +75,7 @@ impl UdpPunctureNegotiationListener { "Error getting UDP public address for the responder: {}", err ); - udp.unbind(udp_bind.sender_address().clone()).await?; + udp.unbind(udp_bind.sender_address().clone())?; return Err(err); } }; diff --git a/implementations/rust/ockam/ockam_transport_udp/src/puncture/negotiation/negotiation.rs b/implementations/rust/ockam/ockam_transport_udp/src/puncture/negotiation/negotiation.rs index 331b9c41761..00ef9c502eb 100644 --- a/implementations/rust/ockam/ockam_transport_udp/src/puncture/negotiation/negotiation.rs +++ b/implementations/rust/ockam/ockam_transport_udp/src/puncture/negotiation/negotiation.rs @@ -34,7 +34,7 @@ impl UdpPunctureNegotiation { { // To be able to receive the response ctx.flow_controls() - .add_consumer(child_ctx.address(), &flow_control_id); + .add_consumer(child_ctx.primary_address(), &flow_control_id); } // We create a new bind for each puncture. Ownership will be transferred to the @@ -49,7 +49,7 @@ impl UdpPunctureNegotiation { debug!( "Initializing UdpPunctureNegotiation Initiator at {}", - child_ctx.address_ref() + child_ctx.primary_address() ); let client = RendezvousClient::new(&udp_bind, rendezvous_route); let my_udp_public_address = match client.get_my_address(ctx).await { @@ -59,14 +59,14 @@ impl UdpPunctureNegotiation { "Error getting UDP public address for the initiator: {}", err ); - udp.unbind(udp_bind.sender_address().clone()).await?; + udp.unbind(udp_bind.sender_address().clone())?; return Err(err); } }; info!( "UdpPunctureNegotiation Initiator {} got its public address: {}", - child_ctx.address_ref(), + child_ctx.primary_address(), my_udp_public_address ); @@ -94,11 +94,11 @@ impl UdpPunctureNegotiation { Err(err) => { error!( "Error receiving response for Udp Puncture at: {}. {}", - child_ctx.address_ref(), + child_ctx.primary_address(), err ); - udp.unbind(udp_bind.sender_address().clone()).await?; + udp.unbind(udp_bind.sender_address().clone())?; return Err(err); } }; @@ -108,11 +108,11 @@ impl UdpPunctureNegotiation { Err(err) => { error!( "Invalid response for Udp Puncture at: {}. {}", - child_ctx.address_ref(), + child_ctx.primary_address(), err ); - udp.unbind(udp_bind.sender_address().clone()).await?; + udp.unbind(udp_bind.sender_address().clone())?; return Err(err); } }; diff --git a/implementations/rust/ockam/ockam_transport_udp/src/puncture/puncture/puncture.rs b/implementations/rust/ockam/ockam_transport_udp/src/puncture/puncture/puncture.rs index f0262092d95..9a8c9fe73b9 100644 --- a/implementations/rust/ockam/ockam_transport_udp/src/puncture/puncture/puncture.rs +++ b/implementations/rust/ockam/ockam_transport_udp/src/puncture/puncture/puncture.rs @@ -76,9 +76,8 @@ impl UdpPuncture { } /// Stop the receiver (which will shut down everything else as well) - pub async fn stop(&self, ctx: &Context) -> Result<()> { - ctx.stop_worker(self.addresses.receiver_address().clone()) - .await + pub fn stop(&self, ctx: &Context) -> Result<()> { + ctx.stop_address(self.addresses.receiver_address().clone()) } /// Flow Control Id diff --git a/implementations/rust/ockam/ockam_transport_udp/src/puncture/puncture/receiver.rs b/implementations/rust/ockam/ockam_transport_udp/src/puncture/puncture/receiver.rs index a49cd0bf8ac..34489add449 100644 --- a/implementations/rust/ockam/ockam_transport_udp/src/puncture/puncture/receiver.rs +++ b/implementations/rust/ockam/ockam_transport_udp/src/puncture/puncture/receiver.rs @@ -69,6 +69,7 @@ impl UdpPunctureReceiverWorker { let remote_mailbox = Mailbox::new( addresses.remote_address().clone(), + None, Arc::new(AllowAll), Arc::new(AllowAll), ); @@ -77,13 +78,15 @@ impl UdpPunctureReceiverWorker { let receiver_mailbox = Mailbox::new( addresses.receiver_address().clone(), + None, Arc::new(DenyAll), options.create_receiver_outgoing_access_control(ctx.flow_controls()), ); let heartbeat_mailbox = Mailbox::new( addresses.heartbeat_address().clone(), - Arc::new(AllowSourceAddress(heartbeat.address())), + None, + Arc::new(AllowSourceAddress(heartbeat.address().clone())), Arc::new(DenyAll), ); @@ -216,8 +219,7 @@ impl UdpPunctureReceiverWorker { .send(UdpPunctureNotification::Closed); // Shut down itself - ctx.stop_worker(self.addresses.remote_address().clone()) - .await?; + ctx.stop_address(self.addresses.remote_address().clone())?; return Ok(()); } @@ -256,7 +258,7 @@ impl UdpPunctureReceiverWorker { let res = self.handle_heartbeat_impl(ctx).await; // Schedule next heartbeat here in case something errors - self.heartbeat.schedule(HEARTBEAT_INTERVAL).await?; + self.heartbeat.schedule(HEARTBEAT_INTERVAL)?; res } @@ -268,19 +270,15 @@ impl Worker for UdpPunctureReceiverWorker { type Context = Context; async fn initialize(&mut self, _context: &mut Self::Context) -> Result<()> { - self.heartbeat.schedule(Duration::ZERO).await?; - - Ok(()) + self.heartbeat.schedule(Duration::ZERO) } async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { self.heartbeat.cancel(); - _ = ctx - .stop_worker(self.addresses.sender_address().clone()) - .await; + _ = ctx.stop_address(self.addresses.sender_address().clone()); - _ = ctx.stop_worker(self.bind.sender_address().clone()).await; + _ = ctx.stop_address(self.bind.sender_address().clone()); Ok(()) } diff --git a/implementations/rust/ockam/ockam_transport_udp/src/transport/bind.rs b/implementations/rust/ockam/ockam_transport_udp/src/transport/bind.rs index 74cb2bd14f7..c558f76b21e 100644 --- a/implementations/rust/ockam/ockam_transport_udp/src/transport/bind.rs +++ b/implementations/rust/ockam/ockam_transport_udp/src/transport/bind.rs @@ -142,8 +142,8 @@ impl UdpTransport { } /// Interrupt an active TCP connection given its Sender `Address` - pub async fn unbind(&self, address: impl Into
) -> Result<()> { - self.ctx.stop_worker(address.into()).await + pub fn unbind(&self, address: impl Into
) -> Result<()> { + self.ctx.stop_address(address.into()) } } diff --git a/implementations/rust/ockam/ockam_transport_udp/src/transport/lifecycle.rs b/implementations/rust/ockam/ockam_transport_udp/src/transport/lifecycle.rs index ae94e0d3e91..8ab7f260d41 100644 --- a/implementations/rust/ockam/ockam_transport_udp/src/transport/lifecycle.rs +++ b/implementations/rust/ockam/ockam_transport_udp/src/transport/lifecycle.rs @@ -68,8 +68,8 @@ impl Transport for UdpTransport { } } - async fn disconnect(&self, address: Address) -> Result<()> { - self.unbind(address).await + fn disconnect(&self, address: Address) -> Result<()> { + self.unbind(address) } } @@ -84,7 +84,7 @@ mod tests { async fn test_resolve_address(ctx: &mut Context) -> Result<()> { let udp = UdpTransport::create(ctx).await?; let udp_address = "127.0.0.1:0"; - let initial_workers = ctx.list_workers(); + let initial_workers = ctx.list_workers()?; let socket = UdpSocket::bind(udp_address) .await .map_err(TransportError::from)?; @@ -95,7 +95,7 @@ mod tests { .await?; // there are 2 additional workers - let mut additional_workers = ctx.list_workers(); + let mut additional_workers = ctx.list_workers()?; additional_workers.retain(|w| !initial_workers.contains(w)); assert_eq!(additional_workers.len(), 2); diff --git a/implementations/rust/ockam/ockam_transport_udp/src/transport/puncture.rs b/implementations/rust/ockam/ockam_transport_udp/src/transport/puncture.rs index 5b1674b1777..a1354b9d3cd 100644 --- a/implementations/rust/ockam/ockam_transport_udp/src/transport/puncture.rs +++ b/implementations/rust/ockam/ockam_transport_udp/src/transport/puncture.rs @@ -25,7 +25,7 @@ impl UdpTransport { } /// Stop a puncture - pub async fn stop_puncture(&self, puncture: UdpPuncture) -> Result<()> { - puncture.stop(&self.ctx).await + pub fn stop_puncture(&self, puncture: UdpPuncture) -> Result<()> { + puncture.stop(&self.ctx) } } diff --git a/implementations/rust/ockam/ockam_transport_udp/src/workers/receiver.rs b/implementations/rust/ockam/ockam_transport_udp/src/workers/receiver.rs index 490e92e207d..94ae1e876cc 100644 --- a/implementations/rust/ockam/ockam_transport_udp/src/workers/receiver.rs +++ b/implementations/rust/ockam/ockam_transport_udp/src/workers/receiver.rs @@ -145,7 +145,7 @@ impl Processor for UdpReceiverProcessor { type Context = Context; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await + ctx.set_cluster(crate::CLUSTER_NAME) } async fn process(&mut self, ctx: &mut Self::Context) -> Result { diff --git a/implementations/rust/ockam/ockam_transport_udp/src/workers/sender.rs b/implementations/rust/ockam/ockam_transport_udp/src/workers/sender.rs index ad696e2e5ee..8a0ec66fa1f 100644 --- a/implementations/rust/ockam/ockam_transport_udp/src/workers/sender.rs +++ b/implementations/rust/ockam/ockam_transport_udp/src/workers/sender.rs @@ -47,9 +47,7 @@ impl Worker for UdpSenderWorker { type Context = Context; async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { - let _ = ctx - .stop_processor(self.addresses.receiver_address().clone()) - .await; + let _ = ctx.stop_address(self.addresses.receiver_address().clone()); Ok(()) } diff --git a/implementations/rust/ockam/ockam_transport_uds/src/router/handle.rs b/implementations/rust/ockam/ockam_transport_uds/src/router/handle.rs index fff7d4a2d35..0e017526a9b 100644 --- a/implementations/rust/ockam/ockam_transport_uds/src/router/handle.rs +++ b/implementations/rust/ockam/ockam_transport_uds/src/router/handle.rs @@ -29,6 +29,7 @@ impl AsyncTryClone for UdsRouterHandle { let mailboxes = Mailboxes::new( Mailbox::new( Address::random_tagged("UdsRouterHandle.async_try_clone.detached"), + None, Arc::new(DenyAll), Arc::new(DenyAll), ), diff --git a/implementations/rust/ockam/ockam_transport_uds/src/router/uds_router.rs b/implementations/rust/ockam/ockam_transport_uds/src/router/uds_router.rs index af5d98aba75..c40c79f32b8 100644 --- a/implementations/rust/ockam/ockam_transport_uds/src/router/uds_router.rs +++ b/implementations/rust/ockam/ockam_transport_uds/src/router/uds_router.rs @@ -51,8 +51,18 @@ impl UdsRouter { }; let handle = router.create_self_handle().await?; - let main_mailbox = Mailbox::new(main_addr.clone(), Arc::new(AllowAll), Arc::new(AllowAll)); - let api_mailbox = Mailbox::new(api_addr.clone(), Arc::new(AllowAll), Arc::new(AllowAll)); + let main_mailbox = Mailbox::new( + main_addr.clone(), + None, + Arc::new(AllowAll), + Arc::new(AllowAll), + ); + let api_mailbox = Mailbox::new( + api_addr.clone(), + None, + Arc::new(AllowAll), + Arc::new(AllowAll), + ); WorkerBuilder::new(router) .with_mailboxes(Mailboxes::new(main_mailbox, vec![api_mailbox])) @@ -132,7 +142,7 @@ impl UdsRouter { self.handle_unregister(self_address.clone()).await?; - self.ctx.stop_worker(self_address).await?; + self.ctx.stop_address(self_address)?; Ok(()) } @@ -259,8 +269,7 @@ impl Worker for UdsRouter { type Message = Any; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await?; - Ok(()) + ctx.set_cluster(crate::CLUSTER_NAME) } async fn handle_message(&mut self, ctx: &mut Context, msg: Routed) -> Result<()> { diff --git a/implementations/rust/ockam/ockam_transport_uds/src/workers/listener.rs b/implementations/rust/ockam/ockam_transport_uds/src/workers/listener.rs index c4138cdb4c0..37e4d283c25 100644 --- a/implementations/rust/ockam/ockam_transport_uds/src/workers/listener.rs +++ b/implementations/rust/ockam/ockam_transport_uds/src/workers/listener.rs @@ -61,7 +61,7 @@ impl Processor for UdsListenProcessor { type Context = Context; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await + ctx.set_cluster(crate::CLUSTER_NAME) } /// Listen for and accept incoming UDS connections. @@ -92,12 +92,14 @@ impl Processor for UdsListenProcessor { let tx_mailbox = Mailbox::new( pair.tx_addr(), + None, Arc::new(AllowSourceAddress(self.router_handle.main_addr().clone())), Arc::new(DenyAll), ); let internal_mailbox = Mailbox::new( send_worker.internal_addr().clone(), + None, Arc::new(AllowSourceAddress(send_worker.rx_addr().clone())), Arc::new(DenyAll), ); diff --git a/implementations/rust/ockam/ockam_transport_uds/src/workers/receiver.rs b/implementations/rust/ockam/ockam_transport_uds/src/workers/receiver.rs index bee1fd79e97..4e0c95e282c 100644 --- a/implementations/rust/ockam/ockam_transport_uds/src/workers/receiver.rs +++ b/implementations/rust/ockam/ockam_transport_uds/src/workers/receiver.rs @@ -37,7 +37,7 @@ impl Processor for UdsRecvProcessor { type Context = Context; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await + ctx.set_cluster(crate::CLUSTER_NAME) } /// Get the next message from the connection if there are any diff --git a/implementations/rust/ockam/ockam_transport_uds/src/workers/sender.rs b/implementations/rust/ockam/ockam_transport_uds/src/workers/sender.rs index 86ee9e74fa6..0c2e55965a7 100644 --- a/implementations/rust/ockam/ockam_transport_uds/src/workers/sender.rs +++ b/implementations/rust/ockam/ockam_transport_uds/src/workers/sender.rs @@ -140,12 +140,14 @@ impl UdsSendWorker { let tx_mailbox = Mailbox::new( pair.tx_addr(), + None, Arc::new(ockam_core::AllowSourceAddress(udsrouter_main_addr)), Arc::new(ockam_core::DenyAll), ); let internal_mailbox = Mailbox::new( worker.internal_addr().clone(), + None, Arc::new(ockam_core::AllowSourceAddress(worker.rx_addr().clone())), Arc::new(ockam_core::DenyAll), ); @@ -159,9 +161,11 @@ impl UdsSendWorker { } async fn stop_and_unregister(&self, ctx: &Context) -> Result<()> { - self.router_handle.unregister(ctx.address()).await?; + self.router_handle + .unregister(ctx.primary_address().clone()) + .await?; - ctx.stop_worker(ctx.address()).await?; + ctx.stop_address(ctx.primary_address())?; Ok(()) } @@ -176,7 +180,7 @@ impl Worker for UdsSendWorker { /// /// Spawn a UDS Recceiver worker to processes incoming UDS messages async fn initialize(&mut self, ctx: &mut Self::Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await?; + ctx.set_cluster(crate::CLUSTER_NAME)?; let path = match self.peer.as_pathname() { Some(p) => p, @@ -236,7 +240,7 @@ impl Worker for UdsSendWorker { async fn shutdown(&mut self, ctx: &mut Self::Context) -> Result<()> { if self.rx_should_be_stopped { - let _ = ctx.stop_processor(self.rx_addr().clone()).await; + let _ = ctx.stop_address(self.rx_addr().clone()); } Ok(()) diff --git a/implementations/rust/ockam/ockam_transport_websocket/src/router/mod.rs b/implementations/rust/ockam/ockam_transport_websocket/src/router/mod.rs index 049c9234ec0..59acefd049a 100644 --- a/implementations/rust/ockam/ockam_transport_websocket/src/router/mod.rs +++ b/implementations/rust/ockam/ockam_transport_websocket/src/router/mod.rs @@ -81,11 +81,13 @@ impl WebSocketRouter { let mailboxes = Mailboxes::new( Mailbox::new( main_addr.clone(), + None, Arc::new(AllowAll), // FIXME: @ac Arc::new(AllowAll), // FIXME: @ac ), vec![Mailbox::new( api_addr, + None, Arc::new(AllowAll), // FIXME: @ac Arc::new(AllowAll), // FIXME: @ac )], @@ -117,8 +119,7 @@ impl Worker for WebSocketRouter { type Context = Context; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await?; - Ok(()) + ctx.set_cluster(crate::CLUSTER_NAME) } async fn handle_message(&mut self, ctx: &mut Context, msg: Routed) -> Result<()> { diff --git a/implementations/rust/ockam/ockam_transport_websocket/src/workers/listener.rs b/implementations/rust/ockam/ockam_transport_websocket/src/workers/listener.rs index ff4d47cbf29..13fdb338446 100644 --- a/implementations/rust/ockam/ockam_transport_websocket/src/workers/listener.rs +++ b/implementations/rust/ockam/ockam_transport_websocket/src/workers/listener.rs @@ -49,7 +49,7 @@ impl Processor for WebSocketListenProcessor { type Context = Context; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await + ctx.set_cluster(crate::CLUSTER_NAME) } async fn process(&mut self, ctx: &mut Self::Context) -> Result { diff --git a/implementations/rust/ockam/ockam_transport_websocket/src/workers/receiver.rs b/implementations/rust/ockam/ockam_transport_websocket/src/workers/receiver.rs index 5c7e1b06d6e..6de472569dc 100644 --- a/implementations/rust/ockam/ockam_transport_websocket/src/workers/receiver.rs +++ b/implementations/rust/ockam/ockam_transport_websocket/src/workers/receiver.rs @@ -46,7 +46,7 @@ where type Context = Context; async fn initialize(&mut self, ctx: &mut Context) -> Result<()> { - ctx.set_cluster(crate::CLUSTER_NAME).await + ctx.set_cluster(crate::CLUSTER_NAME) } /// Get next message from the WebSocket stream if there is diff --git a/implementations/rust/ockam/ockam_transport_websocket/src/workers/sender.rs b/implementations/rust/ockam/ockam_transport_websocket/src/workers/sender.rs index 0449da48039..4005c145d84 100644 --- a/implementations/rust/ockam/ockam_transport_websocket/src/workers/sender.rs +++ b/implementations/rust/ockam/ockam_transport_websocket/src/workers/sender.rs @@ -61,11 +61,13 @@ impl WorkerPair { let mailboxes = Mailboxes::new( Mailbox::new( tx_addr.clone(), + None, Arc::new(AllowAll), // FIXME: @ac Arc::new(AllowAll), // FIXME: @ac ), vec![Mailbox::new( internal_addr, + None, Arc::new(AllowAll), // FIXME: @ac Arc::new(AllowAll), // FIXME: @ac )], @@ -105,11 +107,13 @@ impl WorkerPair { let mailboxes = Mailboxes::new( Mailbox::new( tx_addr.clone(), + None, Arc::new(AllowAll), // FIXME: @ac Arc::new(AllowAll), // FIXME: @ac ), vec![Mailbox::new( internal_addr, + None, Arc::new(AllowAll), // FIXME: @ac Arc::new(AllowAll), // FIXME: @ac )], @@ -164,18 +168,18 @@ where return Err(TransportError::GenericIo)?; } - ctx.set_cluster(crate::CLUSTER_NAME).await?; - self.schedule_heartbeat().await?; + ctx.set_cluster(crate::CLUSTER_NAME)?; + self.schedule_heartbeat()?; Ok(()) } - async fn schedule_heartbeat(&mut self) -> Result<()> { + fn schedule_heartbeat(&mut self) -> Result<()> { let heartbeat_interval = match &self.heartbeat_interval { Some(hi) => *hi, None => return Ok(()), }; - self.heartbeat.schedule(heartbeat_interval).await + self.heartbeat.schedule(heartbeat_interval) } /// Receive messages from the `WebSocketRouter` to send @@ -199,7 +203,7 @@ where .is_err() { warn!("Failed to send heartbeat to peer {}", self.peer); - ctx.stop_worker(ctx.address()).await?; + ctx.stop_address(ctx.primary_address())?; return Ok(()); } @@ -214,13 +218,13 @@ where let msg = WebSocketMessage::from(msg.into_transport_message().encode()?); if ws_sink.send(msg).await.is_err() { warn!("Failed to send message to peer {}", self.peer); - ctx.stop_worker(ctx.address()).await?; + ctx.stop_address(ctx.primary_address())?; return Ok(()); } debug!("Sent message to peer {}", self.peer); } - self.schedule_heartbeat().await?; + self.schedule_heartbeat()?; Ok(()) }