use std::{cell::RefCell, future::Future, rc::Rc, time::Duration};
use futures::{channel::oneshot, future::LocalBoxFuture, task::SpawnError};
use lenso_kernel::{DeterministicDriver, DriverTask, RuntimeDriver};
use crate::FaultInjector;
#[derive(Clone, Debug, Default)]
pub struct TestSimulator {
driver: DeterministicDriver,
faults: FaultInjector,
}
impl TestSimulator {
pub fn new() -> Self {
Self::default()
}
pub fn gate(&self, name: impl Into<String>) -> SimulatorGate {
SimulatorGate {
name: name.into(),
state: Rc::new(RefCell::new(GateState::default())),
}
}
pub fn resource(&self, name: impl Into<String>) -> SimulatorResource {
SimulatorResource {
name: name.into(),
state: Rc::new(RefCell::new(ResourceState::default())),
}
}
pub fn faults(&self) -> FaultInjector {
self.faults.clone()
}
pub fn run<F: Future>(&self, future: F) -> F::Output {
self.driver.run(future)
}
pub fn pump(&self) {
self.run(self.driver.yield_now());
}
pub fn spawn<F>(&self, future: F) -> Result<DriverTask, SpawnError>
where
F: Future<Output = ()> + 'static,
{
self.driver.spawn_local(Box::pin(future))
}
pub fn advance(&self, duration: Duration) {
self.driver.advance(duration);
}
pub fn now(&self) -> Duration {
self.driver.now()
}
pub fn sleep_until(&self, deadline: Duration) -> LocalBoxFuture<'static, ()> {
self.driver.sleep_until(deadline)
}
pub fn driver(&self) -> DeterministicDriver {
self.driver.clone()
}
}
#[derive(Debug, Default)]
struct ResourceState {
acquisitions: usize,
frozen: bool,
thaw_notifiers: Vec<oneshot::Sender<()>>,
}
#[derive(Clone, Debug)]
pub struct SimulatorResource {
name: String,
state: Rc<RefCell<ResourceState>>,
}
impl SimulatorResource {
pub fn name(&self) -> &str {
&self.name
}
pub async fn acquire(&self) {
loop {
let thaw = {
let mut state = self.state.borrow_mut();
if state.frozen {
let (sender, receiver) = oneshot::channel();
state.thaw_notifiers.push(sender);
Some(receiver)
} else {
state.acquisitions += 1;
None
}
};
let Some(thaw) = thaw else {
return;
};
let _ = thaw.await;
}
}
#[must_use]
pub fn freeze(&self) -> bool {
let mut state = self.state.borrow_mut();
if state.frozen {
return false;
}
state.frozen = true;
true
}
#[must_use]
pub fn thaw(&self) -> bool {
let notifiers = {
let mut state = self.state.borrow_mut();
if !state.frozen {
return false;
}
state.frozen = false;
std::mem::take(&mut state.thaw_notifiers)
};
for notifier in notifiers {
let _ = notifier.send(());
}
true
}
pub fn acquisition_count(&self) -> usize {
self.state.borrow().acquisitions
}
pub fn is_frozen(&self) -> bool {
self.state.borrow().frozen
}
}
#[derive(Debug, Default)]
struct GateState {
reached: usize,
released: bool,
reached_notifiers: Vec<oneshot::Sender<()>>,
release_notifiers: Vec<oneshot::Sender<()>>,
}
#[derive(Clone, Debug)]
pub struct SimulatorGate {
name: String,
state: Rc<RefCell<GateState>>,
}
impl SimulatorGate {
pub fn name(&self) -> &str {
&self.name
}
pub async fn wait(&self) {
let release = {
let mut state = self.state.borrow_mut();
state.reached += 1;
for notifier in state.reached_notifiers.drain(..) {
let _ = notifier.send(());
}
if state.released {
None
} else {
let (sender, receiver) = oneshot::channel();
state.release_notifiers.push(sender);
Some(receiver)
}
};
if let Some(release) = release {
let _ = release.await;
}
}
pub async fn reached(&self) {
let notification = {
let mut state = self.state.borrow_mut();
if state.reached > 0 {
None
} else {
let (sender, receiver) = oneshot::channel();
state.reached_notifiers.push(sender);
Some(receiver)
}
};
if let Some(notification) = notification {
let _ = notification.await;
}
}
#[must_use]
pub fn release(&self) -> bool {
let notifiers = {
let mut state = self.state.borrow_mut();
if state.released {
return false;
}
state.released = true;
std::mem::take(&mut state.release_notifiers)
};
for notifier in notifiers {
let _ = notifier.send(());
}
true
}
pub fn reached_count(&self) -> usize {
self.state.borrow().reached
}
pub fn is_released(&self) -> bool {
self.state.borrow().released
}
}
#[cfg(test)]
mod tests {
use std::{cell::RefCell, rc::Rc};
use futures::FutureExt;
use lenso_kernel::TaskOutcome;
use super::*;
#[test]
fn gate_blocks_work_until_the_test_releases_it() {
let simulator = TestSimulator::new();
let gate = simulator.gate("provider.before-response");
let task = simulator
.spawn({
let gate = gate.clone();
async move {
gate.wait().await;
}
})
.unwrap();
simulator.run(gate.reached());
assert_eq!(gate.name(), "provider.before-response");
assert_eq!(gate.reached_count(), 1);
assert!(!gate.is_released());
assert!(gate.release());
assert!(!gate.release());
assert_eq!(simulator.run(task), TaskOutcome::Completed);
}
#[test]
fn gates_allow_test_chosen_completion_order() {
let simulator = TestSimulator::new();
let first = simulator.gate("operation.first");
let second = simulator.gate("operation.second");
let completed = Rc::new(RefCell::new(Vec::new()));
let first_task = simulator
.spawn({
let first = first.clone();
let completed = completed.clone();
async move {
first.wait().await;
completed.borrow_mut().push("first");
}
})
.unwrap();
let second_task = simulator
.spawn({
let second = second.clone();
let completed = completed.clone();
async move {
second.wait().await;
completed.borrow_mut().push("second");
}
})
.unwrap();
simulator.run(async {
first.reached().await;
second.reached().await;
});
assert!(second.release());
assert_eq!(simulator.run(second_task), TaskOutcome::Completed);
assert_eq!(&*completed.borrow(), &["second"]);
assert!(first.release());
assert_eq!(simulator.run(first_task), TaskOutcome::Completed);
assert_eq!(&*completed.borrow(), &["second", "first"]);
}
#[test]
fn cancelling_one_blocked_operation_does_not_release_its_peer() {
let simulator = TestSimulator::new();
let cancelled = simulator.gate("operation.cancelled");
let peer = simulator.gate("operation.peer");
let cancelled_task = simulator
.spawn({
let cancelled = cancelled.clone();
async move {
cancelled.wait().await;
}
})
.unwrap();
let peer_task = simulator
.spawn({
let peer = peer.clone();
async move {
peer.wait().await;
}
})
.unwrap();
simulator.run(async {
cancelled.reached().await;
peer.reached().await;
});
cancelled_task.cancel();
assert_eq!(simulator.run(cancelled_task), TaskOutcome::Cancelled);
assert!(!peer.is_released());
assert!(peer.release());
assert_eq!(simulator.run(peer_task), TaskOutcome::Completed);
}
#[test]
fn advancing_virtual_time_wakes_elapsed_work_when_the_test_pumps() {
let simulator = TestSimulator::new();
let completed = Rc::new(RefCell::new(false));
let sleep = simulator.sleep_until(Duration::from_millis(5));
let task = simulator
.spawn({
let completed = completed.clone();
async move {
sleep.await;
*completed.borrow_mut() = true;
}
})
.unwrap();
simulator.pump();
assert!(!*completed.borrow());
simulator.advance(Duration::from_millis(4));
simulator.pump();
assert!(!*completed.borrow());
simulator.advance(Duration::from_millis(1));
assert_eq!(simulator.now(), Duration::from_millis(5));
assert!(!*completed.borrow());
simulator.pump();
assert!(*completed.borrow());
assert_eq!(simulator.run(task), TaskOutcome::Completed);
}
#[test]
fn sleeping_at_or_before_now_completes_without_advancing_time() {
let simulator = TestSimulator::new();
simulator.advance(Duration::from_millis(5));
for deadline in [Duration::from_millis(4), simulator.now()] {
assert_eq!(simulator.sleep_until(deadline).now_or_never(), Some(()));
}
assert_eq!(simulator.now(), Duration::from_millis(5));
}
#[test]
fn dropping_and_cancelling_sleepers_does_not_complete_peer_work() {
let simulator = TestSimulator::new();
let completed = Rc::new(RefCell::new(Vec::new()));
let deadline = Duration::from_millis(5);
drop(simulator.sleep_until(deadline));
let cancelled = simulator
.spawn({
let sleep = simulator.sleep_until(deadline);
let completed = completed.clone();
async move {
sleep.await;
completed.borrow_mut().push("cancelled");
}
})
.unwrap();
let peer = simulator
.spawn({
let sleep = simulator.sleep_until(deadline);
let completed = completed.clone();
async move {
sleep.await;
completed.borrow_mut().push("peer");
}
})
.unwrap();
simulator.pump();
cancelled.cancel();
assert_eq!(simulator.run(cancelled), TaskOutcome::Cancelled);
simulator.advance(Duration::from_millis(4));
simulator.pump();
assert!(completed.borrow().is_empty());
simulator.advance(Duration::from_millis(1));
assert_eq!(simulator.run(peer), TaskOutcome::Completed);
assert_eq!(&*completed.borrow(), &["peer"]);
}
#[test]
fn freezing_one_resource_owner_does_not_block_another() {
let simulator = TestSimulator::new();
let auth_store = simulator.resource("auth-store");
let object_store = simulator.resource("object-store");
assert!(auth_store.freeze());
assert!(!auth_store.freeze());
let blocked = simulator
.spawn({
let auth_store = auth_store.clone();
async move {
auth_store.acquire().await;
}
})
.unwrap();
let independent = simulator
.spawn({
let object_store = object_store.clone();
async move {
object_store.acquire().await;
}
})
.unwrap();
simulator.pump();
assert_eq!(auth_store.acquisition_count(), 0);
assert_eq!(object_store.acquisition_count(), 1);
assert_eq!(simulator.run(independent), TaskOutcome::Completed);
assert!(auth_store.is_frozen());
assert!(auth_store.thaw());
assert!(!auth_store.thaw());
assert_eq!(simulator.run(blocked), TaskOutcome::Completed);
assert_eq!(auth_store.acquisition_count(), 1);
}
}