use crate::key::check_prefix;
use crate::{clip, Bytes, Cond, Error, Key, Meta, Object, ObjectStore, Result};
use std::collections::{BTreeMap, VecDeque};
use std::future::{ready, Ready};
use std::ops::Range;
use std::sync::Mutex;
#[derive(Debug)]
#[non_exhaustive]
pub enum Fault {
Pass,
Fail(Error),
LostResponse(Error),
}
#[derive(Debug, Default)]
struct Inner {
objects: BTreeMap<String, Object>,
next_generation: u64,
faults: VecDeque<Fault>,
ops: u64,
}
#[derive(Debug, Default)]
pub struct Mem {
inner: Mutex<Inner>,
}
impl Mem {
pub fn inject(&self, fault: Fault) {
self.lock().faults.push_back(fault);
}
pub fn inject_after(&self, skip: u64, fault: Fault) {
let mut g = self.lock();
for _ in 0..skip {
g.faults.push_back(Fault::Pass);
}
g.faults.push_back(fault);
}
pub fn clear_faults(&self) {
self.lock().faults.clear();
}
#[must_use]
pub fn ops(&self) -> u64 {
self.lock().ops
}
#[must_use]
pub fn keys(&self) -> Vec<String> {
self.lock().objects.keys().cloned().collect()
}
fn lock(&self) -> std::sync::MutexGuard<'_, Inner> {
self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
}
fn run<T>(&self, op: impl FnOnce(&mut Inner) -> Result<T>) -> Result<T> {
let mut g = self.lock();
g.ops += 1;
match g.faults.pop_front() {
Some(Fault::Fail(e)) => Err(e),
Some(Fault::LostResponse(e)) => {
let _ = op(&mut g);
Err(e)
}
Some(Fault::Pass) | None => op(&mut g),
}
}
}
impl ObjectStore for Mem {
type Get<'a> = Ready<Result<Option<Object>>>;
type GetRange<'a> = Ready<Result<Bytes>>;
type Head<'a> = Ready<Result<Option<Meta>>>;
type Put<'a> = Ready<Result<u64>>;
type List<'a> = Ready<Result<Vec<Meta>>>;
type Delete<'a> = Ready<Result<()>>;
fn get<'a>(&'a self, key: &'a str) -> Self::Get<'a> {
ready(Key::parse(key).and_then(|k| self.run(|g| Ok(g.objects.get(k.as_str()).cloned()))))
}
fn get_range<'a>(&'a self, key: &'a str, range: Range<u64>) -> Self::GetRange<'a> {
ready(Key::parse(key).and_then(|k| {
self.run(|g| {
let o = g.objects.get(k.as_str()).ok_or(Error::NotFound)?;
Ok(o.bytes.slice(clip(&range, o.bytes.len() as u64)))
})
}))
}
fn head<'a>(&'a self, key: &'a str) -> Self::Head<'a> {
ready(Key::parse(key).and_then(|k| {
self.run(|g| {
Ok(g.objects.get(k.as_str()).map(|o| Meta {
key: k.as_str().to_owned(),
size: o.bytes.len() as u64,
generation: o.generation,
}))
})
}))
}
fn put<'a>(&'a self, key: &'a str, body: Bytes, cond: Cond) -> Self::Put<'a> {
ready(Key::parse(key).and_then(|k| {
self.run(|g| {
let current = g.objects.get(k.as_str()).map(|o| o.generation);
match (cond, current) {
(Cond::CreateOnly, Some(_)) => return Err(Error::Conflict),
(Cond::IfGeneration(want), cur) if cur != Some(want) => {
return Err(Error::PreconditionFailed)
}
_ => {}
}
g.next_generation += 1;
let generation = g.next_generation;
g.objects.insert(k.as_str().to_owned(), Object { bytes: body, generation });
Ok(generation)
})
}))
}
fn list<'a>(&'a self, prefix: &'a str) -> Self::List<'a> {
ready(check_prefix(prefix).and_then(|()| {
self.run(|g| {
Ok(g.objects
.range(prefix.to_owned()..)
.take_while(|(k, _)| k.starts_with(prefix))
.map(|(k, o)| Meta {
key: k.clone(),
size: o.bytes.len() as u64,
generation: o.generation,
})
.collect())
})
}))
}
fn delete<'a>(&'a self, key: &'a str) -> Self::Delete<'a> {
ready(Key::parse(key).and_then(|k| {
self.run(|g| {
g.objects.remove(k.as_str());
Ok(())
})
}))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn inject_after_targets_the_nth_operation() {
let m = Mem::default();
m.inject_after(2, Fault::Fail(Error::unavailable("boom", true)));
let put = |k| futures::executor::block_on(m.put(k, Bytes::new(), Cond::None));
assert!(put("a").is_ok());
assert!(put("b").is_ok());
assert!(put("c").is_err());
assert!(put("d").is_ok());
}
use futures::executor::block_on;
#[test]
fn faults_apply_in_order() {
let m = Mem::default();
m.inject(Fault::LostResponse(Error::unavailable("lost", true)));
m.inject(Fault::Fail(Error::unavailable("down", true)));
assert!(block_on(m.put("a", Bytes::from_static(b"1"), Cond::CreateOnly)).is_err());
assert_eq!(block_on(m.get("a")).unwrap_err().to_string(), "store unavailable: down");
assert!(block_on(m.get("a")).unwrap().is_some(), "the lost-response write was applied");
assert!(matches!(
block_on(m.put("a", Bytes::new(), Cond::CreateOnly)),
Err(Error::Conflict)
));
assert_eq!(m.ops(), 4);
}
}