127 lines
3.2 KiB
Rust
127 lines
3.2 KiB
Rust
use std::{
|
|
io::{Stdout, Write},
|
|
sync::{
|
|
atomic::{AtomicU64, AtomicUsize, Ordering},
|
|
Arc, Mutex, OnceLock,
|
|
},
|
|
};
|
|
|
|
pub mod protocol;
|
|
pub use protocol::{Body, Message};
|
|
|
|
pub type DynNode = Arc<Mutex<dyn Node>>;
|
|
|
|
pub trait Node {
|
|
fn handle(&mut self, runner: &Runner, msg: &Message);
|
|
}
|
|
|
|
pub struct Runner {
|
|
msg_id: AtomicU64,
|
|
node: DynNode,
|
|
node_id: OnceLock<String>,
|
|
nodes: OnceLock<Vec<String>>,
|
|
steps: AtomicUsize,
|
|
output: Arc<Mutex<Stdout>>,
|
|
}
|
|
|
|
impl Runner {
|
|
pub fn new(node: DynNode, output: Stdout) -> Self {
|
|
Runner {
|
|
node,
|
|
msg_id: AtomicU64::new(1),
|
|
nodes: OnceLock::new(),
|
|
node_id: OnceLock::new(),
|
|
steps: AtomicUsize::new(0),
|
|
output: Arc::new(Mutex::new(output)),
|
|
}
|
|
}
|
|
|
|
pub fn run(&self, input: std::sync::mpsc::Receiver<Message>) {
|
|
for msg in input.iter() {
|
|
self.steps.fetch_add(1, Ordering::SeqCst);
|
|
let typ = &msg.body.typ;
|
|
if let "init" = typ.as_str() {
|
|
self.init(&msg);
|
|
let body = Body::from_type("init_ok");
|
|
self.reply(&msg, body);
|
|
} else {
|
|
let mut n = self.node.lock().unwrap();
|
|
n.handle(self, &msg);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn node_id(&self) -> String {
|
|
self.node_id.get().cloned().unwrap_or_default()
|
|
}
|
|
|
|
pub fn next_msg_id(&self) -> u64 {
|
|
self.msg_id.fetch_add(1, Ordering::SeqCst)
|
|
}
|
|
|
|
pub fn cur_msg_id(&self) -> u64 {
|
|
self.msg_id.load(Ordering::SeqCst)
|
|
}
|
|
|
|
pub fn nodes(&self) -> &[String] {
|
|
self.nodes.get().unwrap()
|
|
}
|
|
|
|
pub fn steps(&self) -> usize {
|
|
self.steps.load(Ordering::SeqCst)
|
|
}
|
|
|
|
pub fn init(&self, msg: &Message) {
|
|
let node_id = msg
|
|
.body
|
|
.payload
|
|
.get("node_id")
|
|
.unwrap()
|
|
.as_str()
|
|
.unwrap()
|
|
.to_owned();
|
|
|
|
let nodes = msg
|
|
.body
|
|
.payload
|
|
.get("node_ids")
|
|
.unwrap()
|
|
.as_array()
|
|
.unwrap()
|
|
.iter()
|
|
.map(|s| s.as_str().unwrap().to_string())
|
|
.collect::<Vec<_>>();
|
|
|
|
let _ = self.node_id.get_or_init(|| node_id.to_owned());
|
|
let _ = self.nodes.get_or_init(|| nodes.to_vec());
|
|
}
|
|
|
|
pub fn reply(&self, req: &Message, body: Body) {
|
|
let mut body = body;
|
|
let dest = req.src.clone();
|
|
let in_reply_to = req.body.msg_id;
|
|
body.in_reply_to = in_reply_to;
|
|
self.send(&dest, body);
|
|
}
|
|
|
|
pub fn send(&self, dest: &str, body: Body) {
|
|
let mut body = body;
|
|
if body.msg_id == 0 {
|
|
body.msg_id = self.next_msg_id();
|
|
}
|
|
let msg = Message {
|
|
src: self.node_id().to_string(),
|
|
dest: dest.to_string(),
|
|
body,
|
|
};
|
|
let msg = serde_json::to_string(&msg).unwrap();
|
|
self.writeln(&msg);
|
|
}
|
|
|
|
fn writeln(&self, msg: &str) {
|
|
let mut out = self.output.lock().unwrap();
|
|
let msg = format!("{msg}\n");
|
|
out.write_all(msg.as_bytes()).unwrap();
|
|
out.flush().unwrap();
|
|
}
|
|
}
|