use std::collections::VecDeque;
use std::error::Error;
use std::sync::mpsc::{self, Receiver, TryRecvError};
use std::sync::{Arc, Mutex};
use beamr::atom::{Atom, AtomTable};
use beamr::ets::OwnedTerm;
use beamr::module::ModuleRegistry;
use beamr::native::BifRegistryImpl;
use beamr::process::ExitReason;
use beamr::scheduler::WasmScheduler;
use beamr::term::Term;
use super::handle::{CommandQueue, ShardCommand, ShardCommandKind, ShardError};
use super::native::{self, ShardNativeHandler};
type TestResult = Result<(), Box<dyn Error>>;
fn cooperative_scheduler() -> WasmScheduler {
let atom_table = Arc::new(AtomTable::with_common_atoms());
let modules = Arc::new(ModuleRegistry::new());
let bifs = Arc::new(BifRegistryImpl::new());
WasmScheduler::new(atom_table, modules, bifs)
}
fn wake() -> OwnedTerm {
OwnedTerm::immediate(Term::atom(Atom::OK))
}
const MAX_TURNS_PER_REPLY: usize = 16;
fn enqueue_wake_and_pump<T>(
scheduler: &mut WasmScheduler,
pid: u64,
commands: &CommandQueue,
command: ShardCommand,
reply_rx: &Receiver<T>,
) -> Result<T, Box<dyn Error>> {
native::lock_queue(commands).push_back(command);
scheduler.send_owned(pid, &wake())?;
for _ in 0..MAX_TURNS_PER_REPLY {
let _exited = scheduler.run_native_until_idle();
match reply_rx.try_recv() {
Ok(value) => return Ok(value),
Err(TryRecvError::Empty) => {}
Err(TryRecvError::Disconnected) => {
return Err("shard reply channel disconnected before a reply".into());
}
}
}
Err("shard produced no reply within the cooperative turn budget".into())
}
#[test]
fn shard_actor_put_get_round_trip_on_cooperative_scheduler() -> TestResult {
let mut scheduler = cooperative_scheduler();
let dir = tempfile::tempdir()?;
let store_dir = dir.path().join("store");
let wal_path = dir.path().join("shard.wal");
let commands: CommandQueue = Arc::new(Mutex::new(VecDeque::new()));
let factory = ShardNativeHandler::make_factory(store_dir, wal_path, Arc::clone(&commands));
let pid = scheduler.spawn_native_root(factory);
let exited = scheduler.run_native_until_idle();
assert!(
!exited.contains(&pid),
"the freshly booted shard must park, not exit, before any command"
);
assert_eq!(
scheduler.native_exit_reason(pid),
None,
"the shard is live (no native exit) after booting"
);
let key = b"wr-9a/key".to_vec();
let value = b"cooperative-shard-value".to_vec();
let (put_tx, put_rx) = mpsc::sync_channel(1);
let put_command = ShardCommand {
id: 1,
kind: ShardCommandKind::Put {
key: key.clone(),
value: value.clone(),
ttl: None,
reply: put_tx,
},
};
let put_result: Result<(), ShardError> =
enqueue_wake_and_pump(&mut scheduler, pid, &commands, put_command, &put_rx)?;
put_result?;
assert_eq!(
scheduler.native_exit_reason(pid),
None,
"the shard stays live after the put"
);
let (get_tx, get_rx) = mpsc::sync_channel(1);
let get_command = ShardCommand {
id: 2,
kind: ShardCommandKind::Get { key, reply: get_tx },
};
let read_back: Result<Option<Vec<u8>>, ShardError> =
enqueue_wake_and_pump(&mut scheduler, pid, &commands, get_command, &get_rx)?;
assert_eq!(
read_back?,
Some(value),
"the value put through the cooperative shard reads back equal"
);
let (stop_tx, stop_rx) = mpsc::sync_channel(1);
let stop_command = ShardCommand {
id: 3,
kind: ShardCommandKind::Shutdown { reply: stop_tx },
};
let stop_result: Result<(), ShardError> =
enqueue_wake_and_pump(&mut scheduler, pid, &commands, stop_command, &stop_rx)?;
stop_result?;
for _ in 0..MAX_TURNS_PER_REPLY {
if scheduler.native_exit_reason(pid) == Some(ExitReason::Normal) {
break;
}
let _exited = scheduler.run_native_until_idle();
}
assert_eq!(
scheduler.native_exit_reason(pid),
Some(ExitReason::Normal),
"the shard stopped cleanly (Normal) after the shutdown command"
);
Ok(())
}