use crate::{Clock, Error, Handle};
use bytes::Bytes;
use futures::{
channel::mpsc,
task::{waker_ref, ArcWake},
SinkExt, StreamExt,
};
use governor::clock::{Clock as GClock, ReasonablyRealtime};
use prometheus_client::{
encoding::EncodeLabelSet,
metrics::{counter::Counter, family::Family, gauge::Gauge},
registry::Registry,
};
use rand::{prelude::SliceRandom, rngs::StdRng, CryptoRng, RngCore, SeedableRng};
use sha2::{Digest, Sha256};
use std::{
collections::{BinaryHeap, HashMap},
future::Future,
mem::replace,
net::{IpAddr, Ipv4Addr, SocketAddr},
ops::Range,
pin::Pin,
sync::{Arc, Mutex},
task::{self, Poll, Waker},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use tracing::trace;
const EPHEMERAL_PORT_RANGE: Range<u16> = 32768..61000;
const ROOT_TASK: &str = "root";
#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
struct Work {
label: String,
}
#[derive(Debug)]
struct Metrics {
tasks_spawned: Family<Work, Counter>,
tasks_running: Family<Work, Gauge>,
task_polls: Family<Work, Counter>,
bandwidth: Counter,
}
impl Metrics {
pub fn init(registry: Arc<Mutex<Registry>>) -> Self {
let metrics = Self {
task_polls: Family::default(),
tasks_running: Family::default(),
tasks_spawned: Family::default(),
bandwidth: Counter::default(),
};
{
let mut registry = registry.lock().unwrap();
registry.register(
"tasks_spawned",
"Total number of tasks spawned",
metrics.tasks_spawned.clone(),
);
registry.register(
"tasks_running",
"Number of tasks currently running",
metrics.tasks_running.clone(),
);
registry.register(
"task_polls",
"Total number of task polls",
metrics.task_polls.clone(),
);
registry.register(
"bandwidth",
"Total amount of data sent over network",
metrics.bandwidth.clone(),
);
}
metrics
}
}
pub struct Auditor {
hash: Mutex<String>,
}
impl Auditor {
fn new() -> Self {
Self {
hash: Mutex::new(String::new()),
}
}
fn process_task(&self, task: u128, label: &str) {
let mut hash = self.hash.lock().unwrap();
let mut hasher = Sha256::new();
hasher.update(hash.as_bytes());
hasher.update(b"process_task");
hasher.update(task.to_be_bytes());
hasher.update(label.as_bytes());
*hash = format!("{:x}", hasher.finalize());
}
fn bind(&self, address: SocketAddr) {
let mut hash = self.hash.lock().unwrap();
let mut hasher = Sha256::new();
hasher.update(hash.as_bytes());
hasher.update(b"bind");
hasher.update(address.to_string().as_bytes());
*hash = format!("{:x}", hasher.finalize());
}
fn dial(&self, dialer: SocketAddr, dialee: SocketAddr) {
let mut hash = self.hash.lock().unwrap();
let mut hasher = Sha256::new();
hasher.update(hash.as_bytes());
hasher.update(b"dial");
hasher.update(dialer.to_string().as_bytes());
hasher.update(dialee.to_string().as_bytes());
*hash = format!("{:x}", hasher.finalize());
}
fn accept(&self, dialee: SocketAddr, dialer: SocketAddr) {
let mut hash = self.hash.lock().unwrap();
let mut hasher = Sha256::new();
hasher.update(hash.as_bytes());
hasher.update(b"accept");
hasher.update(dialee.to_string().as_bytes());
hasher.update(dialer.to_string().as_bytes());
*hash = format!("{:x}", hasher.finalize());
}
fn send(&self, sender: SocketAddr, receiver: SocketAddr, message: Bytes) {
let mut hash = self.hash.lock().unwrap();
let mut hasher = Sha256::new();
hasher.update(hash.as_bytes());
hasher.update(b"send");
hasher.update(sender.to_string().as_bytes());
hasher.update(receiver.to_string().as_bytes());
hasher.update(&message);
*hash = format!("{:x}", hasher.finalize());
}
fn recv(&self, receiver: SocketAddr, sender: SocketAddr, message: Bytes) {
let mut hash = self.hash.lock().unwrap();
let mut hasher = Sha256::new();
hasher.update(hash.as_bytes());
hasher.update(b"recv");
hasher.update(receiver.to_string().as_bytes());
hasher.update(sender.to_string().as_bytes());
hasher.update(&message);
*hash = format!("{:x}", hasher.finalize());
}
fn rand(&self, method: String) {
let mut hash = self.hash.lock().unwrap();
let mut hasher = Sha256::new();
hasher.update(hash.as_bytes());
hasher.update(b"rand");
hasher.update(method.as_bytes());
*hash = format!("{:x}", hasher.finalize());
}
pub fn state(&self) -> String {
self.hash.lock().unwrap().clone()
}
}
struct Task {
id: u128,
label: String,
tasks: Arc<Tasks>,
root: bool,
future: Mutex<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
completed: Mutex<bool>,
}
impl ArcWake for Task {
fn wake_by_ref(arc_self: &Arc<Self>) {
arc_self.tasks.enqueue(arc_self.clone());
}
}
struct Tasks {
counter: Mutex<u128>,
queue: Mutex<Vec<Arc<Task>>>,
}
impl Tasks {
fn register(
arc_self: &Arc<Self>,
label: &str,
root: bool,
future: Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
) {
let mut queue = arc_self.queue.lock().unwrap();
let id = {
let mut l = arc_self.counter.lock().unwrap();
let old = *l;
*l = l.checked_add(1).expect("task counter overflow");
old
};
queue.push(Arc::new(Task {
id,
label: label.to_string(),
root,
future: Mutex::new(future),
tasks: arc_self.clone(),
completed: Mutex::new(false),
}));
}
fn enqueue(&self, task: Arc<Task>) {
let mut queue = self.queue.lock().unwrap();
queue.push(task);
}
fn drain(&self) -> Vec<Arc<Task>> {
let mut queue = self.queue.lock().unwrap();
let len = queue.len();
replace(&mut *queue, Vec::with_capacity(len))
}
fn len(&self) -> usize {
self.queue.lock().unwrap().len()
}
}
#[derive(Clone)]
pub struct Config {
pub registry: Arc<Mutex<Registry>>,
pub seed: u64,
pub cycle: Duration,
}
impl Default for Config {
fn default() -> Self {
Self {
registry: Arc::new(Mutex::new(Registry::default())),
seed: 42,
cycle: Duration::from_millis(1),
}
}
}
pub struct Executor {
cycle: Duration,
metrics: Arc<Metrics>,
auditor: Arc<Auditor>,
rng: Mutex<StdRng>,
prefix: Mutex<String>,
time: Mutex<SystemTime>,
tasks: Arc<Tasks>,
sleeping: Mutex<BinaryHeap<Alarm>>,
}
impl Executor {
pub fn init(cfg: Config) -> (Runner, Context, Arc<Auditor>) {
let metrics = Arc::new(Metrics::init(cfg.registry));
let auditor = Arc::new(Auditor::new());
let executor = Arc::new(Self {
cycle: cfg.cycle,
metrics: metrics.clone(),
auditor: auditor.clone(),
rng: Mutex::new(StdRng::seed_from_u64(cfg.seed)),
prefix: Mutex::new(String::new()),
time: Mutex::new(UNIX_EPOCH),
tasks: Arc::new(Tasks {
queue: Mutex::new(Vec::new()),
counter: Mutex::new(0),
}),
sleeping: Mutex::new(BinaryHeap::new()),
});
(
Runner {
executor: executor.clone(),
},
Context {
executor,
networking: Arc::new(Networking::new(metrics, auditor.clone())),
},
auditor,
)
}
pub fn seeded(seed: u64) -> (Runner, Context, Arc<Auditor>) {
let cfg = Config {
seed,
..Config::default()
};
Self::init(cfg)
}
#[allow(clippy::should_implement_trait)]
pub fn default() -> (Runner, Context, Arc<Auditor>) {
Self::init(Config::default())
}
}
pub struct Runner {
executor: Arc<Executor>,
}
impl crate::Runner for Runner {
fn start<F>(self, f: F) -> F::Output
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let output = Arc::new(Mutex::new(None));
Tasks::register(
&self.executor.tasks,
ROOT_TASK,
true,
Box::pin({
let output = output.clone();
async move {
*output.lock().unwrap() = Some(f.await);
}
}),
);
let mut iter = 0;
loop {
let mut tasks = self.executor.tasks.drain();
{
let mut rng = self.executor.rng.lock().unwrap();
tasks.shuffle(&mut *rng);
}
trace!(iter, tasks = tasks.len(), "starting loop");
for task in tasks {
{
let mut prefix = self.executor.prefix.lock().unwrap();
prefix.clone_from(&task.label);
}
self.executor.auditor.process_task(task.id, &task.label);
if *task.completed.lock().unwrap() {
trace!(id = task.id, "skipping already completed task");
continue;
}
trace!(id = task.id, "processing task");
let waker = waker_ref(&task);
let mut context = task::Context::from_waker(&waker);
let mut future = task.future.lock().unwrap();
self.executor
.metrics
.task_polls
.get_or_create(&Work {
label: task.label.clone(),
})
.inc();
let pending = future.as_mut().poll(&mut context).is_pending();
if pending {
trace!(id = task.id, "task is still pending");
continue;
}
*task.completed.lock().unwrap() = true;
trace!(id = task.id, "task is complete");
if task.root {
return output.lock().unwrap().take().unwrap();
}
}
let mut current;
{
let mut time = self.executor.time.lock().unwrap();
*time = time
.checked_add(self.executor.cycle)
.expect("executor time overflow");
current = *time;
}
trace!(
now = current.duration_since(UNIX_EPOCH).unwrap().as_millis(),
"time advanced",
);
if self.executor.tasks.len() == 0 {
let mut skip = None;
{
let sleeping = self.executor.sleeping.lock().unwrap();
if let Some(next) = sleeping.peek() {
if next.time > current {
skip = Some(next.time);
}
}
}
if skip.is_some() {
{
let mut time = self.executor.time.lock().unwrap();
*time = skip.unwrap();
current = *time;
}
trace!(
now = current.duration_since(UNIX_EPOCH).unwrap().as_millis(),
"time skipped",
);
}
}
let mut to_wake = Vec::new();
let mut remaining;
{
let mut sleeping = self.executor.sleeping.lock().unwrap();
while let Some(next) = sleeping.peek() {
if next.time <= current {
let sleeper = sleeping.pop().unwrap();
to_wake.push(sleeper.waker);
} else {
break;
}
}
remaining = sleeping.len();
}
for waker in to_wake {
waker.wake();
}
remaining += self.executor.tasks.len();
if remaining == 0 {
panic!("runtime stalled");
}
iter += 1;
}
}
}
#[derive(Clone)]
pub struct Context {
executor: Arc<Executor>,
networking: Arc<Networking>,
}
impl crate::Spawner for Context {
fn spawn<F, T>(&self, label: &str, f: F) -> Handle<T>
where
F: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
let label = {
let prefix = self.executor.prefix.lock().unwrap();
if prefix.is_empty() || *prefix == ROOT_TASK {
label.to_string()
} else {
format!("{}_{}", *prefix, label)
}
};
if label == ROOT_TASK {
panic!("root task cannot be spawned");
}
let work = Work {
label: label.clone(),
};
self.executor
.metrics
.tasks_spawned
.get_or_create(&work)
.inc();
let gauge = self
.executor
.metrics
.tasks_running
.get_or_create(&work)
.clone();
let (f, handle) = Handle::init(f, gauge, false);
Tasks::register(&self.executor.tasks, &label, false, Box::pin(f));
handle
}
}
struct Sleeper {
executor: Arc<Executor>,
time: SystemTime,
registered: bool,
}
struct Alarm {
time: SystemTime,
waker: Waker,
}
impl PartialEq for Alarm {
fn eq(&self, other: &Self) -> bool {
self.time.eq(&other.time)
}
}
impl Eq for Alarm {}
impl PartialOrd for Alarm {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for Alarm {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
other.time.cmp(&self.time)
}
}
impl Future for Sleeper {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
{
let current_time = *self.executor.time.lock().unwrap();
if current_time >= self.time {
return Poll::Ready(());
}
}
if !self.registered {
self.registered = true;
self.executor.sleeping.lock().unwrap().push(Alarm {
time: self.time,
waker: cx.waker().clone(),
});
}
Poll::Pending
}
}
impl Clock for Context {
fn current(&self) -> SystemTime {
*self.executor.time.lock().unwrap()
}
fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + 'static {
let time = self
.current()
.checked_add(duration)
.expect("overflow when setting wake time");
Sleeper {
executor: self.executor.clone(),
time,
registered: false,
}
}
fn sleep_until(&self, deadline: SystemTime) -> impl Future<Output = ()> + Send + 'static {
Sleeper {
executor: self.executor.clone(),
time: deadline,
registered: false,
}
}
}
impl GClock for Context {
type Instant = SystemTime;
fn now(&self) -> Self::Instant {
self.current()
}
}
impl ReasonablyRealtime for Context {}
type Dialable = mpsc::UnboundedSender<(
SocketAddr,
mpsc::UnboundedSender<Bytes>, mpsc::UnboundedReceiver<Bytes>, )>;
struct Networking {
metrics: Arc<Metrics>,
auditor: Arc<Auditor>,
ephemeral: Mutex<u16>,
listeners: Mutex<HashMap<SocketAddr, Dialable>>,
}
impl Networking {
fn new(metrics: Arc<Metrics>, auditor: Arc<Auditor>) -> Self {
Self {
metrics,
auditor,
ephemeral: Mutex::new(EPHEMERAL_PORT_RANGE.start),
listeners: Mutex::new(HashMap::new()),
}
}
fn bind(&self, socket: SocketAddr) -> Result<Listener, Error> {
self.auditor.bind(socket);
if EPHEMERAL_PORT_RANGE.contains(&socket.port()) {
return Err(Error::BindFailed);
}
let mut listeners = self.listeners.lock().unwrap();
if listeners.contains_key(&socket) {
return Err(Error::BindFailed);
}
let (sender, receiver) = mpsc::unbounded();
listeners.insert(socket, sender.clone());
Ok(Listener {
auditor: self.auditor.clone(),
address: socket,
listener: receiver,
metrics: self.metrics.clone(),
})
}
async fn dial(&self, socket: SocketAddr) -> Result<(Sink, Stream), Error> {
let dialer = {
let mut ephemeral = self.ephemeral.lock().unwrap();
let dialer = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), *ephemeral);
*ephemeral = ephemeral
.checked_add(1)
.expect("ephemeral port range exhausted");
dialer
};
self.auditor.dial(dialer, socket);
let mut sender = {
let listeners = self.listeners.lock().unwrap();
let sender = listeners.get(&socket).ok_or(Error::ConnectionFailed)?;
sender.clone()
};
let (dialer_sender, dialer_receiver) = mpsc::unbounded();
let (dialee_sender, dialee_receiver) = mpsc::unbounded();
sender
.send((dialer, dialer_sender, dialee_receiver))
.await
.map_err(|_| Error::ConnectionFailed)?;
Ok((
Sink {
metrics: self.metrics.clone(),
auditor: self.auditor.clone(),
me: dialer,
peer: socket,
sender: dialee_sender,
},
Stream {
auditor: self.auditor.clone(),
me: dialer,
peer: socket,
receiver: dialer_receiver,
},
))
}
}
impl crate::Network<Listener, Sink, Stream> for Context {
async fn bind(&self, socket: SocketAddr) -> Result<Listener, Error> {
self.networking.bind(socket)
}
async fn dial(&self, socket: SocketAddr) -> Result<(Sink, Stream), Error> {
self.networking.dial(socket).await
}
}
pub struct Listener {
metrics: Arc<Metrics>,
auditor: Arc<Auditor>,
address: SocketAddr,
listener: mpsc::UnboundedReceiver<(
SocketAddr,
mpsc::UnboundedSender<Bytes>,
mpsc::UnboundedReceiver<Bytes>,
)>,
}
impl crate::Listener<Sink, Stream> for Listener {
async fn accept(&mut self) -> Result<(SocketAddr, Sink, Stream), Error> {
let (socket, sender, receiver) = self.listener.next().await.ok_or(Error::ReadFailed)?;
self.auditor.accept(self.address, socket);
Ok((
socket,
Sink {
metrics: self.metrics.clone(),
auditor: self.auditor.clone(),
me: self.address,
peer: socket,
sender,
},
Stream {
auditor: self.auditor.clone(),
me: self.address,
peer: socket,
receiver,
},
))
}
}
pub struct Sink {
metrics: Arc<Metrics>,
auditor: Arc<Auditor>,
me: SocketAddr,
peer: SocketAddr,
sender: mpsc::UnboundedSender<Bytes>,
}
impl crate::Sink for Sink {
async fn send(&mut self, msg: Bytes) -> Result<(), Error> {
let len = msg.len();
self.auditor.send(self.me, self.peer, msg.clone());
self.sender
.send(msg)
.await
.map_err(|_| Error::WriteFailed)?;
self.metrics.bandwidth.inc_by(len as u64);
Ok(())
}
}
pub struct Stream {
auditor: Arc<Auditor>,
me: SocketAddr,
peer: SocketAddr,
receiver: mpsc::UnboundedReceiver<Bytes>,
}
impl crate::Stream for Stream {
async fn recv(&mut self) -> Result<Bytes, Error> {
let msg = self.receiver.next().await.ok_or(Error::ReadFailed)?;
self.auditor.recv(self.me, self.peer, msg.clone());
Ok(msg)
}
}
impl RngCore for Context {
fn next_u32(&mut self) -> u32 {
self.executor.auditor.rand("next_u32".to_string());
self.executor.rng.lock().unwrap().next_u32()
}
fn next_u64(&mut self) -> u64 {
self.executor.auditor.rand("next_u64".to_string());
self.executor.rng.lock().unwrap().next_u64()
}
fn fill_bytes(&mut self, dest: &mut [u8]) {
self.executor.auditor.rand("fill_bytes".to_string());
self.executor.rng.lock().unwrap().fill_bytes(dest)
}
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
self.executor.auditor.rand("try_fill_bytes".to_string());
self.executor.rng.lock().unwrap().try_fill_bytes(dest)
}
}
impl CryptoRng for Context {}
#[cfg(test)]
mod tests {
use super::*;
use crate::utils::run_tasks;
use futures::task::noop_waker;
fn run_with_seed(seed: u64) -> (String, Vec<usize>) {
let (executor, runtime, auditor) = Executor::seeded(seed);
let messages = run_tasks(5, executor, runtime);
(auditor.state(), messages)
}
#[test]
fn test_same_seed_same_order() {
let mut outputs = Vec::new();
for seed in 0..1000 {
let output = run_with_seed(seed);
outputs.push(output);
}
for seed in 0..1000 {
let output = run_with_seed(seed);
assert_eq!(output, outputs[seed as usize]);
}
}
#[test]
fn test_different_seeds_different_order() {
let output1 = run_with_seed(12345);
let output2 = run_with_seed(54321);
assert_ne!(output1, output2);
}
#[test]
fn test_alarm_min_heap() {
let now = SystemTime::now();
let alarms = vec![
Alarm {
time: now + Duration::new(10, 0),
waker: noop_waker(),
},
Alarm {
time: now + Duration::new(5, 0),
waker: noop_waker(),
},
Alarm {
time: now + Duration::new(15, 0),
waker: noop_waker(),
},
Alarm {
time: now + Duration::new(5, 0),
waker: noop_waker(),
},
];
let mut heap = BinaryHeap::new();
for alarm in alarms {
heap.push(alarm);
}
let mut sorted_times = Vec::new();
while let Some(alarm) = heap.pop() {
sorted_times.push(alarm.time);
}
assert_eq!(
sorted_times,
vec![
now + Duration::new(5, 0),
now + Duration::new(5, 0),
now + Duration::new(10, 0),
now + Duration::new(15, 0),
]
);
}
}