Skip to content

Commit

Permalink
Add SwarmCallback to replace message CallbackFn and Validator
Browse files Browse the repository at this point in the history
  • Loading branch information
@Ma233
Ma233 committed Oct 8, 2023
1 parent 55967c0 commit 8e6f255
Show file tree
Hide file tree
Showing 16 changed files with 152 additions and 198 deletions.
144 changes: 25 additions & 119 deletions core/src/message/handlers/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,11 @@ use std::sync::Arc;
use async_recursion::async_recursion;
use async_trait::async_trait;

use super::CustomMessage;
use super::Message;
use super::MessagePayload;
use crate::dht::vnode::VirtualNode;
use crate::dht::Did;
use crate::dht::PeerRing;
use crate::error::Error;
use crate::error::Result;
use crate::message::ConnectNodeReport;
use crate::message::ConnectNodeSend;
Expand All @@ -40,44 +38,6 @@ pub mod subring;
/// Type alias for message payload.
pub type Payload = MessagePayload<Message>;

/// Trait of message callback.
#[cfg_attr(feature = "wasm", async_trait(?Send))]
#[cfg_attr(not(feature = "wasm"), async_trait)]
pub trait MessageCallback {
/// Message handler for custom message
async fn custom_message(
&self,
ctx: &MessagePayload<Message>,
msg: &CustomMessage,
) -> Vec<MessageHandlerEvent>;
/// Message handler for builtin message
async fn builtin_message(&self, ctx: &MessagePayload<Message>) -> Vec<MessageHandlerEvent>;
}

/// Trait of message validator.
#[cfg_attr(feature = "wasm", async_trait(?Send))]
#[cfg_attr(not(feature = "wasm"), async_trait)]
pub trait MessageValidator {
/// Externality validator
async fn validate(&self, ctx: &MessagePayload<Message>) -> Option<String>;
}

/// Boxed Callback, for non-wasm, it should be Sized, Send and Sync.
#[cfg(not(feature = "wasm"))]
pub type CallbackFn = Box<dyn MessageCallback + Send + Sync>;

/// Boxed Callback
#[cfg(feature = "wasm")]
pub type CallbackFn = Box<dyn MessageCallback>;

/// Boxed Validator
#[cfg(not(feature = "wasm"))]
pub type ValidatorFn = Box<dyn MessageValidator + Send + Sync>;

/// Boxed Validator, for non-wasm, it should be Sized, Send and Sync.
#[cfg(feature = "wasm")]
pub type ValidatorFn = Box<dyn MessageValidator>;

type NextHop = Did;

/// MessageHandlerEvent that will be handled by Swarm.
Expand Down Expand Up @@ -127,10 +87,6 @@ pub enum MessageHandlerEvent {
#[derive(Clone)]
pub struct MessageHandler {
dht: Arc<PeerRing>,
/// CallbackFn implement `customMessage` and `builtin_message`.
callback: Arc<Option<CallbackFn>>,
/// A specific validator implement ValidatorFn.
validator: Arc<Option<ValidatorFn>>,
}

/// Generic trait for handle message ,inspired by Actor-Model.
Expand All @@ -147,47 +103,8 @@ pub trait HandleMsg<T> {

impl MessageHandler {
/// Create a new MessageHandler Instance.
pub fn new(
dht: Arc<PeerRing>,
callback: Option<CallbackFn>,
validator: Option<ValidatorFn>,
) -> Self {
Self {
dht,
callback: Arc::new(callback),
validator: Arc::new(validator),
}
}

/// Invoke callback, which will be call after builtin handler.
async fn invoke_callback(&self, payload: &MessagePayload<Message>) -> Vec<MessageHandlerEvent> {
if let Some(ref cb) = *self.callback {
match payload.data {
Message::CustomMessage(ref msg) => {
if self.dht.did == payload.relay.destination {
tracing::debug!("INVOKE CUSTOM MESSAGE CALLBACK {}", &payload.tx_id);
return cb.custom_message(payload, msg).await;
}
}
_ => return cb.builtin_message(payload).await,
};
} else if let Message::CustomMessage(ref msg) = payload.data {
if self.dht.did == payload.relay.destination {
tracing::warn!("No callback registered, skip invoke_callback of {:?}", msg);
}
}
vec![]
}

/// Validate message.
async fn validate(&self, payload: &MessagePayload<Message>) -> Result<()> {
if let Some(ref v) = *self.validator {
v.validate(payload)
.await
.map(|info| Err(Error::InvalidMessage(info)))
.unwrap_or(Ok(()))?;
};
Ok(())
pub fn new(dht: Arc<PeerRing>) -> Self {
Self { dht }
}

/// Handle builtin message.
Expand All @@ -203,9 +120,7 @@ impl MessageHandler {
}
tracing::debug!("START HANDLE MESSAGE: {} {}", &payload.tx_id, &payload.data);

self.validate(payload).await?;

let mut events = match &payload.data {
let events = match &payload.data {
Message::JoinDHT(ref msg) => self.handle(payload, msg).await,
Message::LeaveDHT(ref msg) => self.handle(payload, msg).await,
Message::ConnectNodeSend(ref msg) => self.handle(payload, msg).await,
Expand All @@ -223,10 +138,6 @@ impl MessageHandler {
Message::QueryForTopoInfoReport(ref msg) => self.handle(payload, msg).await,
}?;

tracing::debug!("INVOKE CALLBACK {}", &payload.tx_id);

events.extend(self.invoke_callback(payload).await);

tracing::debug!("FINISH HANDLE MESSAGE {}", &payload.tx_id);
Ok(events)
}
Expand All @@ -243,12 +154,13 @@ pub mod tests {
use crate::dht::Did;
use crate::ecc::SecretKey;
use crate::message::PayloadSender;
use crate::swarm::callback::SwarmCallback;
use crate::swarm::Swarm;
use crate::tests::default::prepare_node_with_callback;
use crate::tests::manually_establish_connection;

#[derive(Clone)]
struct MessageCallbackInstance {
struct CallbackInstance {
#[allow(clippy::type_complexity)]
handler_messages: Arc<Mutex<Vec<(Did, Vec<u8>)>>>,
}
Expand All @@ -259,40 +171,34 @@ pub mod tests {
let key2 = SecretKey::random();

#[async_trait]
impl MessageCallback for MessageCallbackInstance {
async fn custom_message(
impl SwarmCallback for CallbackInstance {
async fn on_payload(
&self,
ctx: &MessagePayload<Message>,
msg: &CustomMessage,
) -> Vec<MessageHandlerEvent> {
self.handler_messages
.lock()
.await
.push((ctx.addr, msg.0.clone()));
println!("{:?}, {:?}, {:?}", ctx, ctx.addr, msg);
vec![]
}
payload: &MessagePayload<Message>,
) -> std::result::Result<(), Box<dyn std::error::Error>> {
println!("{:?}, {:?}, {:?}", payload, payload.addr, payload.data);

if let Message::CustomMessage(ref msg) = payload.data {
self.handler_messages
.lock()
.await
.push((payload.addr, msg.0.clone()))
}

async fn builtin_message(
&self,
ctx: &MessagePayload<Message>,
) -> Vec<MessageHandlerEvent> {
println!("{:?}, {:?}", ctx, ctx.addr);
vec![]
Ok(())
}
}

let msg_callback1 = MessageCallbackInstance {
let cb1 = CallbackInstance {
handler_messages: Arc::new(Mutex::new(vec![])),
};
let msg_callback2 = MessageCallbackInstance {

let cb2 = CallbackInstance {
handler_messages: Arc::new(Mutex::new(vec![])),
};
let cb1: CallbackFn = Box::new(msg_callback1.clone());
let cb2: CallbackFn = Box::new(msg_callback2.clone());

let (node1, _path1) = prepare_node_with_callback(key1, Some(cb1)).await;
let (node2, _path2) = prepare_node_with_callback(key2, Some(cb2)).await;
let (node1, _path1) = prepare_node_with_callback(key1, Some(cb1.clone().boxed())).await;
let (node2, _path2) = prepare_node_with_callback(key2, Some(cb2.clone().boxed())).await;

manually_establish_connection(&node1, &node2).await;

Expand Down Expand Up @@ -343,12 +249,12 @@ pub mod tests {

sleep(Duration::from_secs(5)).await;

assert_eq!(msg_callback1.handler_messages.lock().await.as_slice(), &[
assert_eq!(cb1.handler_messages.lock().await.as_slice(), &[
(node2.did(), "Hello world 2 to 1 - 1".as_bytes().to_vec()),
(node2.did(), "Hello world 2 to 1 - 2".as_bytes().to_vec())
]);

assert_eq!(msg_callback2.handler_messages.lock().await.as_slice(), &[
assert_eq!(cb2.handler_messages.lock().await.as_slice(), &[
(node1.did(), "Hello world 1 to 2 - 1".as_bytes().to_vec()),
(node1.did(), "Hello world 1 to 2 - 2".as_bytes().to_vec()),
(node1.did(), "Hello world 1 to 2 - 3".as_bytes().to_vec())
Expand Down
3 changes: 0 additions & 3 deletions core/src/message/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,9 @@ pub mod handlers;
pub use handlers::storage::ChordStorageInterface;
pub use handlers::storage::ChordStorageInterfaceCacheChecker;
pub use handlers::subring::SubringInterface;
pub use handlers::CallbackFn;
pub use handlers::HandleMsg;
pub use handlers::MessageCallback;
pub use handlers::MessageHandler;
pub use handlers::MessageHandlerEvent;
pub use handlers::ValidatorFn;

mod protocols;
pub use protocols::MessageRelay;
38 changes: 18 additions & 20 deletions core/src/swarm/builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,21 +4,24 @@

use std::sync::Arc;

use rings_transport::core::callback::Callback;
use rings_transport::core::callback::TransportCallback;

use crate::channels::Channel;
use crate::dht::PeerRing;
use crate::message::CallbackFn;
use crate::message::MessageHandler;
use crate::message::ValidatorFn;
use crate::session::SessionSk;
use crate::storage::PersistenceStorage;
use crate::swarm::callback::BoxedSwarmCallback;
use crate::swarm::callback::InnerSwarmCallback;
use crate::swarm::callback::SwarmCallback;
use crate::swarm::MeasureImpl;
use crate::swarm::Swarm;
use crate::types::channel::Channel as ChannelTrait;
use crate::types::Transport;

struct DefaultCallback;
impl SwarmCallback for DefaultCallback {}

/// Creates a SwarmBuilder to configure a Swarm.
pub struct SwarmBuilder {
ice_servers: String,
Expand All @@ -28,8 +31,7 @@ pub struct SwarmBuilder {
session_sk: SessionSk,
session_ttl: Option<usize>,
measure: Option<MeasureImpl>,
message_callback: Option<CallbackFn>,
message_validator: Option<ValidatorFn>,
callback: Option<BoxedSwarmCallback>,
}

impl SwarmBuilder {
Expand All @@ -43,8 +45,7 @@ impl SwarmBuilder {
session_sk,
session_ttl: None,
measure: None,
message_callback: None,
message_validator: None,
callback: None,
}
}

Expand Down Expand Up @@ -73,15 +74,9 @@ impl SwarmBuilder {
self
}

/// Bind message callback function for Swarm.
pub fn message_callback(mut self, callback: CallbackFn) -> Self {
self.message_callback = Some(callback);
self
}

/// Bind message vilidator function implementation for Swarm.
pub fn message_validator(mut self, validator: ValidatorFn) -> Self {
self.message_validator = Some(validator);
/// Bind callback for Swarm.
pub fn callback(mut self, callback: BoxedSwarmCallback) -> Self {
self.callback = Some(callback);
self
}

Expand All @@ -95,12 +90,15 @@ impl SwarmBuilder {
self.dht_storage,
));

let message_handler =
MessageHandler::new(dht.clone(), self.message_callback, self.message_validator);
let message_handler = MessageHandler::new(dht.clone());
let swarm_callback = self.callback.unwrap_or_else(|| DefaultCallback {}.boxed());

let transport_event_channel = Channel::new();

let transport = Box::new(Transport::new(&self.ice_servers, self.external_address));
let callback = Arc::new(SwarmCallback::new(transport_event_channel.sender()).boxed());
let transport_callback = Arc::new(
InnerSwarmCallback::new(transport_event_channel.sender(), swarm_callback).boxed(),
);

Swarm {
transport_event_channel,
Expand All @@ -109,7 +107,7 @@ impl SwarmBuilder {
session_sk: self.session_sk,
message_handler,
transport,
callback,
transport_callback,
}
}
}
Loading

0 comments on commit 8e6f255

Please sign in to comment.