tokio-ring.rs

syndicate-macros/examples/tokio-ring.rs

@@ -0,0 +1,175 @@
+use std::env;
+use std::sync::Arc;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
+
+use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
+
+type Ref<T> = UnboundedSender<T>;
+
+#[derive(Debug)]
+enum Instruction {
+    SetPeer(Arc<Ref<Instruction>>),
+    HandleMessage(u64),
+}
+
+struct Forwarder {
+    hop_limit: u64,
+    supervisor: Arc<Ref<Instruction>>,
+    peer: Option<Arc<Ref<Instruction>>>,
+}
+
+impl Drop for Forwarder {
+    fn drop(&mut self) {
+        let r = self.peer.take();
+        let _ = tokio::spawn(async move {
+            drop(r);
+        });
+    }
+}
+
+enum Action { Continue, Stop }
+
+trait Actor<T> {
+    fn message(&mut self, message: T) -> Action;
+}
+
+fn send<T: std::marker::Send + 'static>(ch: &Arc<Ref<T>>, message: T) -> () {
+    match ch.send(message) {
+        Ok(()) => (),
+        Err(v) => panic!("Aiee! Could not send {:?}", v),
+    }
+}
+
+fn spawn<T: std::marker::Send + 'static, R: Actor<T> + std::marker::Send + 'static>(rt: Option<Arc<AtomicU64>>, mut ac: R) -> Arc<Ref<T>> {
+    let (tx, mut rx) = unbounded_channel();
+    if let Some(ref c) = rt {
+        c.fetch_add(1, Ordering::SeqCst);
+    }
+    tokio::spawn(async move {
+        loop {
+            match rx.recv().await {
+                None => break,
+                Some(message) => {
+                    match ac.message(message) {
+                        Action::Continue => continue,
+                        Action::Stop => break,
+                    }
+                }
+            }
+        }
+        if let Some(c) = rt {
+            c.fetch_sub(1, Ordering::SeqCst);
+        }
+    });
+    Arc::new(tx)
+}
+
+impl Actor<Instruction> for Forwarder {
+    fn message(&mut self, message: Instruction) -> Action {
+        match message {
+            Instruction::SetPeer(r) => {
+                tracing::info!("Setting peer {:?}", r);
+                self.peer = Some(r);
+            }
+            Instruction::HandleMessage(n) => {
+                let target = if n >= self.hop_limit { &self.supervisor } else { self.peer.as_ref().expect("peer") };
+                send(target, Instruction::HandleMessage(n + 1));
+            }
+        }
+        Action::Continue
+    }
+}
+
+struct Supervisor {
+    latency_mode: bool,
+    total_transfers: u64,
+    remaining_to_receive: u32,
+    start_time: Option<std::time::Instant>,
+}
+
+impl Actor<Instruction> for Supervisor {
+    fn message(&mut self, message: Instruction) -> Action {
+        match message {
+            Instruction::SetPeer(_) => {
+                tracing::info!("Start");
+                self.start_time = Some(std::time::Instant::now());
+            },
+            Instruction::HandleMessage(_n) => {
+                self.remaining_to_receive -= 1;
+                if self.remaining_to_receive == 0 {
+                    let stop_time = std::time::Instant::now();
+                    let duration = stop_time - self.start_time.unwrap();
+                    tracing::info!("Stop after {:?}; {:?} messages, so {:?} Hz ({} mode)",
+                                   duration,
+                                   self.total_transfers,
+                                   (1000.0 * self.total_transfers as f64) / duration.as_millis() as f64,
+                                   if self.latency_mode { "latency" } else { "throughput" });
+                    return Action::Stop;
+                }
+            },
+        }
+        Action::Continue
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error + std::marker::Send + std::marker::Sync>> {
+    syndicate::convenient_logging()?;
+
+    let args: Vec<String> = env::args().collect();
+    let n_actors: u32 = args.get(1).unwrap_or(&"1000000".to_string()).parse()?;
+    let n_rounds: u32 = args.get(2).unwrap_or(&"200".to_string()).parse()?;
+    let latency_mode: bool = match args.get(3).unwrap_or(&"throughput".to_string()).as_str() {
+        "latency" => true,
+        "throughput" => false,
+        _other => return Err("Invalid throughput/latency mode".into()),
+    };
+    tracing::info!("Will run {:?} actors for {:?} rounds", n_actors, n_rounds);
+
+    let count = Arc::new(AtomicU64::new(0));
+
+    let total_transfers: u64 = n_actors as u64 * n_rounds as u64;
+    let (hop_limit, injection_count) = if latency_mode {
+        (total_transfers, 1)
+    } else {
+        (n_rounds as u64, n_actors)
+    };
+
+    let me = spawn(Some(count.clone()), Supervisor {
+        latency_mode,
+        total_transfers,
+        remaining_to_receive: injection_count,
+        start_time: None,
+    });
+
+    let mut forwarders: Vec<Arc<Ref<Instruction>>> = Vec::new();
+    for _i in 0 .. n_actors {
+        if _i % 10000 == 0 { tracing::info!("Actor {:?}", _i); }
+        forwarders.push(spawn(None, Forwarder {
+            hop_limit,
+            supervisor: me.clone(),
+            peer: forwarders.last().cloned(),
+        }));
+    }
+    send(&forwarders[0], Instruction::SetPeer(forwarders.last().expect("an entity").clone()));
+    send(&me, Instruction::SetPeer(me.clone()));
+
+    let mut injected: u32 = 0;
+    for f in forwarders.into_iter() {
+        if injected >= injection_count {
+            break;
+        }
+        send(&f, Instruction::HandleMessage(0));
+        injected += 1;
+    }
+
+    loop {
+        if count.load(Ordering::SeqCst) == 0 {
+            break;
+        }
+        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+    }
+
+    Ok(())
+}