use std::collections::{BTreeMap, HashMap};
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use async_trait::async_trait;
use rand::{RngCore, SeedableRng};
use rand_chacha::ChaCha8Rng;
use tokio::sync::Mutex as AsyncMutex;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use crate::receptionist::Receptionist;
use crate::runtime::Runtime;
use crate::sim::SimRuntime;
use super::super::config::{NodeClass, NodeIdentity};
use super::super::error::ClusterError;
use super::super::framing::ControlMessage;
use super::{
Connection, ConnectionEvent, IncomingConnection, Net, NodeId, PeerAddr, RecvHalf, SendHalf,
run_control_stream_writer,
};
#[derive(Clone)]
struct LatencyConfig {
runtime: SimRuntime,
base: Duration,
jitter: Duration,
rng: Arc<Mutex<ChaCha8Rng>>,
}
impl LatencyConfig {
fn sample_delay(&self) -> Duration {
if self.jitter.is_zero() {
return self.base;
}
let jitter_ns = self.jitter.as_nanos() as u64;
let n = self.rng.lock().unwrap().next_u64() % (jitter_ns + 1);
self.base + Duration::from_nanos(n)
}
}
type Stream = (Box<dyn SendHalf>, Box<dyn RecvHalf>);
type Nodes = Arc<Mutex<BTreeMap<String, mpsc::UnboundedSender<Stream>>>>;
type Peers = Arc<Mutex<BTreeMap<String, SimDialTarget>>>;
struct SimSend {
tx: mpsc::UnboundedSender<Vec<u8>>,
severed: CancellationToken,
latency: Option<LatencyConfig>,
last_delivery: Duration,
}
#[async_trait]
impl SendHalf for SimSend {
async fn write_all(&mut self, buf: &[u8]) -> Result<(), ClusterError> {
if self.severed.is_cancelled() {
return Err(ClusterError::Transport("sim link severed".into()));
}
if let Some(cfg) = &self.latency {
let now = cfg.runtime.now();
let delay = cfg.sample_delay();
let deliver_at = self.last_delivery.max(now) + delay;
self.last_delivery = deliver_at;
let sleep_for = deliver_at.saturating_sub(now);
let tx = self.tx.clone();
let severed = self.severed.clone();
let rt = cfg.runtime.clone();
let chunk = buf.to_vec();
cfg.runtime.spawn(Box::pin(async move {
rt.sleep(sleep_for).await;
if !severed.is_cancelled() {
let _ = tx.send(chunk);
}
}));
Ok(())
} else {
self.tx
.send(buf.to_vec())
.map_err(|_| ClusterError::Transport("sim stream: peer closed".into()))
}
}
fn finish(&mut self) -> Result<(), ClusterError> {
Ok(())
}
}
struct SimRecv {
rx: mpsc::UnboundedReceiver<Vec<u8>>,
leftover: Vec<u8>,
severed: CancellationToken,
}
impl SimRecv {
fn new(rx: mpsc::UnboundedReceiver<Vec<u8>>, severed: CancellationToken) -> Self {
Self {
rx,
leftover: Vec::new(),
severed,
}
}
}
#[async_trait]
impl RecvHalf for SimRecv {
async fn read(&mut self, buf: &mut [u8]) -> Result<Option<usize>, ClusterError> {
if self.leftover.is_empty() {
tokio::select! {
biased;
chunk = self.rx.recv() => match chunk {
Some(chunk) => self.leftover = chunk,
None => return Ok(None), },
_ = self.severed.cancelled() => return Ok(None),
}
}
let n = buf.len().min(self.leftover.len());
buf[..n].copy_from_slice(&self.leftover[..n]);
self.leftover.drain(..n);
Ok(Some(n))
}
}
fn stream_pair(severed: CancellationToken, latency: Option<LatencyConfig>) -> (Stream, Stream) {
let (a_tx, b_rx) = mpsc::unbounded_channel::<Vec<u8>>();
let (b_tx, a_rx) = mpsc::unbounded_channel::<Vec<u8>>();
let near: Stream = (
Box::new(SimSend {
tx: a_tx,
severed: severed.clone(),
latency: latency.clone(),
last_delivery: Duration::ZERO,
}),
Box::new(SimRecv::new(a_rx, severed.clone())),
);
let far: Stream = (
Box::new(SimSend {
tx: b_tx,
severed: severed.clone(),
latency,
last_delivery: Duration::ZERO,
}),
Box::new(SimRecv::new(b_rx, severed)),
);
(near, far)
}
struct SimConnection {
remote_id: NodeId,
peer_accept_tx: mpsc::UnboundedSender<Stream>,
my_accept_rx: AsyncMutex<mpsc::UnboundedReceiver<Stream>>,
closed: CancellationToken,
latency: Option<LatencyConfig>,
}
#[async_trait]
impl Connection for SimConnection {
fn remote_id(&self) -> NodeId {
self.remote_id.clone()
}
async fn open_bi(&self) -> Result<Stream, ClusterError> {
let (near, far) = stream_pair(self.closed.clone(), self.latency.clone());
self.peer_accept_tx
.send(far)
.map_err(|_| ClusterError::Connection("sim peer not accepting streams".into()))?;
Ok(near)
}
async fn accept_bi(&self) -> Option<Stream> {
let mut rx = self.my_accept_rx.lock().await;
tokio::select! {
biased;
s = rx.recv() => s,
_ = self.closed.cancelled() => None,
}
}
fn close(&self, _code: u32, _reason: &[u8]) {
self.closed.cancel();
}
}
#[derive(Clone)]
struct SimDialTarget {
identity: NodeIdentity,
node_class: NodeClass,
node_metadata: HashMap<String, String>,
is_edge_client: bool,
raw_inbound_tx: mpsc::UnboundedSender<RawInbound>,
}
struct RawInbound {
remote_identity: NodeIdentity,
remote_class: NodeClass,
remote_metadata: HashMap<String, String>,
remote_is_edge: bool,
connection: SimConnection,
control_send: Box<dyn SendHalf>,
control_recv: Box<dyn RecvHalf>,
peer_accept_tx: mpsc::UnboundedSender<Stream>,
link: CancellationToken,
}
struct SimPeerConn {
control_tx: mpsc::UnboundedSender<ControlMessage>,
peer_accept_tx: mpsc::UnboundedSender<Stream>,
severed: CancellationToken,
}
#[derive(Clone)]
pub struct SimFabric {
runtime: SimRuntime,
nodes: Nodes,
peers: Peers,
latency: Option<LatencyConfig>,
}
impl SimFabric {
pub fn new(runtime: SimRuntime) -> Self {
Self {
runtime,
nodes: Arc::new(Mutex::new(BTreeMap::new())),
peers: Arc::new(Mutex::new(BTreeMap::new())),
latency: None,
}
}
pub fn runtime(&self) -> &SimRuntime {
&self.runtime
}
pub fn set_latency(&mut self, base: Duration, jitter: Duration) {
let seed = self.runtime.derive_seed("net-faults");
self.latency = Some(LatencyConfig {
runtime: self.runtime.clone(),
base,
jitter,
rng: Arc::new(Mutex::new(ChaCha8Rng::seed_from_u64(seed))),
});
}
pub fn node(&self, key: impl Into<String>) -> (SimNet, mpsc::UnboundedReceiver<Stream>) {
let key = key.into();
let (tx, rx) = mpsc::unbounded_channel();
self.nodes.lock().unwrap().insert(key.clone(), tx);
let identity = NodeIdentity::new_seeded(key.clone(), NodeId(key.clone()), "0.0.0.0", 0, 0);
let (conn_events_tx, _conn_events_rx) = mpsc::unbounded_channel();
let net = SimNet {
node_id: NodeId(key),
identity,
node_class: NodeClass::Worker,
node_metadata: HashMap::new(),
nodes: Arc::clone(&self.nodes),
peers: Arc::clone(&self.peers),
connections: Arc::new(Mutex::new(BTreeMap::new())),
conn_events_tx,
runtime: self.runtime.clone(),
shutdown: CancellationToken::new(),
latency: self.latency.clone(),
};
(net, rx)
}
pub fn bind(
&self,
identity: NodeIdentity,
node_class: NodeClass,
node_metadata: HashMap<String, String>,
shutdown: CancellationToken,
) -> (
Arc<SimNet>,
mpsc::UnboundedReceiver<IncomingConnection>,
mpsc::UnboundedReceiver<ConnectionEvent>,
) {
let (incoming_tx, incoming_rx) = mpsc::unbounded_channel();
let (conn_events_tx, conn_events_rx) = mpsc::unbounded_channel();
let (raw_inbound_tx, raw_inbound_rx) = mpsc::unbounded_channel::<RawInbound>();
self.peers.lock().unwrap().insert(
identity.endpoint_id.0.clone(),
SimDialTarget {
identity: identity.clone(),
node_class: node_class.clone(),
node_metadata: node_metadata.clone(),
is_edge_client: false,
raw_inbound_tx,
},
);
let net = Arc::new(SimNet {
node_id: identity.endpoint_id.clone(),
identity,
node_class,
node_metadata,
nodes: Arc::clone(&self.nodes),
peers: Arc::clone(&self.peers),
connections: Arc::new(Mutex::new(BTreeMap::new())),
conn_events_tx,
runtime: self.runtime.clone(),
shutdown,
latency: self.latency.clone(),
});
let accept_net = Arc::clone(&net);
self.runtime.spawn(Box::pin(run_accept_loop(
accept_net,
raw_inbound_rx,
incoming_tx,
)));
(net, incoming_rx, conn_events_rx)
}
}
async fn run_accept_loop(
net: Arc<SimNet>,
mut raw_inbound_rx: mpsc::UnboundedReceiver<RawInbound>,
incoming_tx: mpsc::UnboundedSender<IncomingConnection>,
) {
while let Some(raw) = raw_inbound_rx.recv().await {
let RawInbound {
remote_identity,
remote_class,
remote_metadata,
remote_is_edge,
connection,
control_send,
control_recv,
peer_accept_tx,
link,
} = raw;
let node_key = remote_identity.node_id_string();
let (control_out_tx, control_out_rx) = mpsc::unbounded_channel();
net.runtime.spawn(Box::pin(run_control_stream_writer(
control_send,
control_out_rx,
net.shutdown.clone(),
)));
net.connections.lock().unwrap().insert(
node_key.clone(),
SimPeerConn {
control_tx: control_out_tx.clone(),
peer_accept_tx,
severed: link,
},
);
let _ = net
.conn_events_tx
.send(ConnectionEvent::Connected(node_key.clone()));
let ic = IncomingConnection {
remote_identity,
connection: Box::new(connection),
control_tx: control_out_tx,
control_recv,
node_class: remote_class,
node_metadata: remote_metadata,
is_edge_client: remote_is_edge,
};
if incoming_tx.send(ic).is_err() {
break; }
}
}
pub struct SimNet {
node_id: NodeId,
identity: NodeIdentity,
node_class: NodeClass,
node_metadata: HashMap<String, String>,
nodes: Nodes,
peers: Peers,
connections: Arc<Mutex<BTreeMap<String, SimPeerConn>>>,
conn_events_tx: mpsc::UnboundedSender<ConnectionEvent>,
runtime: SimRuntime,
shutdown: CancellationToken,
latency: Option<LatencyConfig>,
}
#[async_trait]
impl Net for SimNet {
fn node_id(&self) -> NodeId {
self.node_id.clone()
}
fn local_addr(&self) -> SocketAddr {
self.identity.socket_addr()
}
async fn connect(&self, addr: PeerAddr) -> Result<IncomingConnection, ClusterError> {
let target = self
.peers
.lock()
.unwrap()
.get(&addr.id.0)
.cloned()
.ok_or_else(|| ClusterError::NodeNotFound(addr.id.0.clone()))?;
let link = CancellationToken::new();
let (a_to_b_tx, a_to_b_rx) = mpsc::unbounded_channel::<Stream>();
let (b_to_a_tx, b_to_a_rx) = mpsc::unbounded_channel::<Stream>();
let our_conn = SimConnection {
remote_id: target.identity.endpoint_id.clone(),
peer_accept_tx: a_to_b_tx.clone(),
my_accept_rx: AsyncMutex::new(b_to_a_rx),
closed: link.clone(),
latency: self.latency.clone(),
};
let their_conn = SimConnection {
remote_id: self.identity.endpoint_id.clone(),
peer_accept_tx: b_to_a_tx.clone(),
my_accept_rx: AsyncMutex::new(a_to_b_rx),
closed: link.clone(),
latency: self.latency.clone(),
};
let (near, far) = stream_pair(link.clone(), self.latency.clone());
let (our_ctrl_send, our_ctrl_recv) = near;
let (their_ctrl_send, their_ctrl_recv) = far;
let (our_ctrl_tx, our_ctrl_rx) = mpsc::unbounded_channel();
self.runtime.spawn(Box::pin(run_control_stream_writer(
our_ctrl_send,
our_ctrl_rx,
self.shutdown.clone(),
)));
let target_key = target.identity.node_id_string();
self.connections.lock().unwrap().insert(
target_key.clone(),
SimPeerConn {
control_tx: our_ctrl_tx.clone(),
peer_accept_tx: a_to_b_tx,
severed: link.clone(),
},
);
let _ = self
.conn_events_tx
.send(ConnectionEvent::Connected(target_key));
let raw = RawInbound {
remote_identity: self.identity.clone(),
remote_class: self.node_class.clone(),
remote_metadata: self.node_metadata.clone(),
remote_is_edge: false,
connection: their_conn,
control_send: their_ctrl_send,
control_recv: their_ctrl_recv,
peer_accept_tx: b_to_a_tx,
link,
};
target
.raw_inbound_tx
.send(raw)
.map_err(|_| ClusterError::Transport(format!("node {} not accepting", addr.id.0)))?;
Ok(IncomingConnection {
remote_identity: target.identity,
connection: Box::new(our_conn),
control_tx: our_ctrl_tx,
control_recv: our_ctrl_recv,
node_class: target.node_class,
node_metadata: target.node_metadata,
is_edge_client: target.is_edge_client,
})
}
async fn send_control(&self, node_id: &str, msg: ControlMessage) -> Result<(), ClusterError> {
let tx = self
.connections
.lock()
.unwrap()
.get(node_id)
.map(|c| c.control_tx.clone())
.ok_or_else(|| ClusterError::NodeNotFound(node_id.to_string()))?;
tx.send(msg)
.map_err(|_| ClusterError::Transport("sim control channel closed".into()))
}
async fn broadcast_control(&self, msg: &ControlMessage) {
let txs: Vec<_> = self
.connections
.lock()
.unwrap()
.values()
.map(|c| c.control_tx.clone())
.collect();
for tx in txs {
let _ = tx.send(msg.clone());
}
}
async fn open_actor_stream(&self, node_id: &str) -> Result<Stream, ClusterError> {
let conn = self
.connections
.lock()
.unwrap()
.get(node_id)
.map(|c| (c.peer_accept_tx.clone(), c.severed.clone()));
if let Some((peer_accept_tx, severed)) = conn {
let (near, far) = stream_pair(severed, self.latency.clone());
peer_accept_tx.send(far).map_err(|_| {
ClusterError::Transport(format!("node {node_id} not accepting streams"))
})?;
return Ok(near);
}
let target = self
.nodes
.lock()
.unwrap()
.get(node_id)
.cloned()
.ok_or_else(|| ClusterError::NodeNotFound(node_id.to_string()))?;
let (near, far) = stream_pair(CancellationToken::new(), self.latency.clone());
target.send(far).map_err(|_| {
ClusterError::Transport(format!("node {node_id} not accepting streams"))
})?;
Ok(near)
}
async fn connected_nodes(&self) -> Vec<String> {
self.connections.lock().unwrap().keys().cloned().collect()
}
async fn remove_connection(&self, node_id: &str) {
if self.connections.lock().unwrap().remove(node_id).is_some() {
let _ = self
.conn_events_tx
.send(ConnectionEvent::Disconnected(node_id.to_string()));
}
}
}
impl SimNet {
pub fn partition(&self, peer_node_key: &str) -> bool {
match self.connections.lock().unwrap().get(peer_node_key) {
Some(conn) => {
conn.severed.cancel();
true
}
None => false,
}
}
}
pub fn serve(
runtime: &SimRuntime,
receptionist: Receptionist,
peer_key: String,
mut accept_rx: mpsc::UnboundedReceiver<Stream>,
) {
let rt = runtime.clone();
runtime.spawn(Box::pin(async move {
let (ctrl_tx, _ctrl_rx) = mpsc::unbounded_channel::<(String, ControlMessage)>();
while let Some((send, recv)) = accept_rx.recv().await {
let receptionist = receptionist.clone();
let ctrl_tx = ctrl_tx.clone();
let peer_key = peer_key.clone();
rt.spawn(Box::pin(async move {
crate::cluster::handle_incoming_stream(receptionist, send, recv, ctrl_tx, peer_key)
.await;
}));
}
}));
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ResponseRegistry;
use crate::actor::ActorContext;
use crate::cluster::remote::run_actor_stream_writer;
use crate::endpoint::Endpoint;
use crate::prelude::*;
use crate::sim::SimWorld;
use serde::{Deserialize, Serialize};
struct Counter;
#[derive(Default)]
struct CounterState {
count: i64,
}
impl Actor for Counter {
type State = CounterState;
}
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = i64, remote = "netsim::Increment")]
struct Increment {
amount: i64,
}
#[handlers]
impl Counter {
#[handler]
fn increment(
&self,
_ctx: &ActorContext<Self>,
state: &mut CounterState,
msg: Increment,
) -> i64 {
state.count += msg.amount;
state.count
}
}
fn two_node_roundtrip(seed: u64) -> i64 {
let mut world = SimWorld::new(seed);
let rt = world.runtime().clone();
let _counter = world
.system()
.start("counter/0", Counter, CounterState::default());
let recept_a = world.system().receptionist().clone();
let fabric = SimFabric::new(rt.clone());
let (_net_a, accept_a) = fabric.node("node-a");
let (net_b, _accept_b) = fabric.node("node-b");
serve(&rt, recept_a, "node-b".to_string(), accept_a);
let (wire_tx, wire_rx) = mpsc::unbounded_channel();
let response_registry = ResponseRegistry::new();
let ep: Endpoint<Counter> =
Endpoint::remote("counter/0".to_string(), wire_tx, response_registry.clone());
let net_b: Arc<dyn Net> = Arc::new(net_b);
let runtime: Arc<dyn Runtime> = Arc::new(rt.clone());
rt.spawn(Box::pin(run_actor_stream_writer(
net_b,
runtime,
"node-a".to_string(),
"counter/0".to_string(),
wire_rx,
response_registry,
)));
world.block_on(async move { ep.send(Increment { amount: 5 }).await.unwrap() })
}
#[test]
fn cross_node_message_and_reply_over_sim_net() {
assert_eq!(two_node_roundtrip(0xC0FFEE), 5);
assert_eq!(two_node_roundtrip(1), 5);
}
#[test]
fn latency_in_stream_fifo_order_preserved() {
use crate::sim::SimWorld;
let mut world = SimWorld::new(0xABCD);
let rt = world.runtime().clone();
let cfg = LatencyConfig {
runtime: rt.clone(),
base: Duration::from_millis(10),
jitter: Duration::from_millis(40),
rng: Arc::new(Mutex::new(ChaCha8Rng::seed_from_u64(
rt.derive_seed("fifo-test"),
))),
};
let severed = CancellationToken::new();
let (near, far) = stream_pair(severed, Some(cfg));
let (mut send, _) = near;
let (_, mut recv) = far;
world.block_on(async move {
send.write_all(&[1, 2, 3]).await.unwrap();
send.write_all(&[4, 5, 6]).await.unwrap();
});
world.advance(Duration::from_millis(100));
let got: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
let got2 = got.clone();
world.block_on(async move {
let mut buf = [0u8; 8];
let n = recv.read(&mut buf).await.unwrap().expect("first chunk");
got2.lock().unwrap().extend_from_slice(&buf[..n]);
let n = recv.read(&mut buf).await.unwrap().expect("second chunk");
got2.lock().unwrap().extend_from_slice(&buf[..n]);
});
let got = got.lock().unwrap();
assert_eq!(&got[..3], &[1, 2, 3], "first chunk arrived first (FIFO)");
assert_eq!(&got[3..6], &[4, 5, 6], "second chunk arrived second (FIFO)");
}
#[test]
fn latency_in_flight_chunk_dropped_on_sever() {
use crate::sim::SimWorld;
let mut world = SimWorld::new(0xBEEF);
let rt = world.runtime().clone();
let cfg = LatencyConfig {
runtime: rt.clone(),
base: Duration::from_millis(200),
jitter: Duration::ZERO,
rng: Arc::new(Mutex::new(ChaCha8Rng::seed_from_u64(
rt.derive_seed("inflight-test"),
))),
};
let severed = CancellationToken::new();
let (near, far) = stream_pair(severed.clone(), Some(cfg));
let (mut send, _) = near;
let (_, mut recv) = far;
world.block_on(async move {
send.write_all(&[9, 8, 7]).await.unwrap();
});
severed.cancel();
world.advance(Duration::from_millis(300));
let result: Arc<Mutex<Option<Option<usize>>>> = Arc::new(Mutex::new(None));
let result2 = result.clone();
world.block_on(async move {
let mut buf = [0u8; 8];
let r = recv.read(&mut buf).await.unwrap();
*result2.lock().unwrap() = Some(r);
});
match *result.lock().unwrap() {
Some(None) => {} Some(Some(n)) => {
panic!("in-flight chunk should have been dropped on sever, but got {n} bytes")
}
None => panic!("block_on did not complete"),
}
}
}