use std::time::{Duration, SystemTime, UNIX_EPOCH};
use yo_common::{Code, Error, Result};
use yo_kv::{Applied, Ask, Cond, Kind, MAX_AT, Moved};
use crate::db::Handle;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Ttl {
Missing,
Forever,
In(Duration),
}
impl Ttl {
#[must_use]
pub fn left(self) -> Option<Duration> {
match self {
Ttl::In(left) => Some(left),
_ => None,
}
}
#[must_use]
pub fn found(self) -> bool {
!matches!(self, Ttl::Missing)
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum When {
#[default]
Always,
Unset,
AlreadySet,
Later,
Earlier,
EarlierAndAlreadySet,
}
impl From<When> for Cond {
fn from(when: When) -> Cond {
match when {
When::Always => Cond::Always,
When::Unset => Cond::NotSet,
When::AlreadySet => Cond::AlreadySet,
When::Later => Cond::Greater,
When::Earlier => Cond::Less,
When::EarlierAndAlreadySet => Cond::LessAndSet,
}
}
}
#[derive(Clone)]
pub struct Keys {
pub(crate) db: Handle,
}
impl core::fmt::Debug for Keys {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Keys").finish_non_exhaustive()
}
}
impl Keys {
pub fn exists(&self, key: impl AsRef<[u8]>) -> Result<bool> {
self.db.run(|inner| Ok(inner.strings.exists(key.as_ref())))
}
pub fn count<K: AsRef<[u8]>>(&self, keys: &[K]) -> Result<usize> {
self.db.run(|inner| {
Ok(keys
.iter()
.filter(|key| inner.strings.exists(key.as_ref()))
.count())
})
}
pub fn kind(&self, key: impl AsRef<[u8]>) -> Result<Option<Kind>> {
self.db.run(|inner| Ok(inner.strings.kind_of(key.as_ref())))
}
pub fn del(&self, key: impl AsRef<[u8]>) -> Result<bool> {
self.db.run(|inner| Ok(inner.strings.del(key.as_ref())))
}
pub fn del_many<K: AsRef<[u8]>>(&self, keys: &[K]) -> Result<usize> {
self.db.run(|inner| {
Ok(keys
.iter()
.filter(|key| inner.strings.del(key.as_ref()))
.count())
})
}
pub fn expire_in(&self, key: impl AsRef<[u8]>, after: Duration) -> Result<bool> {
self.expire_in_when(key, after, When::Always)
}
pub fn expire_in_when(
&self,
key: impl AsRef<[u8]>,
after: Duration,
when: When,
) -> Result<bool> {
let ms = u64::try_from(after.as_millis()).map_err(|_| too_far())?;
self.db.deadlines(|inner| {
let at = inner
.strings
.clock()
.now_ms()
.checked_add(ms)
.ok_or_else(too_far)?;
apply(
inner
.strings
.expire(key.as_ref(), reachable(at)?, when.into()),
)
})
}
pub fn expire_at(&self, key: impl AsRef<[u8]>, at: SystemTime) -> Result<bool> {
self.expire_at_when(key, at, When::Always)
}
pub fn expire_at_when(
&self,
key: impl AsRef<[u8]>,
at: SystemTime,
when: When,
) -> Result<bool> {
let ms = moment(at)?;
self.db
.deadlines(|inner| apply(inner.strings.expire(key.as_ref(), ms, when.into())))
}
pub fn ttl(&self, key: impl AsRef<[u8]>) -> Result<Ttl> {
self.db.run(|inner| {
let now = inner.strings.clock().now_ms();
Ok(match inner.strings.deadline_of(key.as_ref()) {
Ask::Missing => Ttl::Missing,
Ask::NoDeadline => Ttl::Forever,
Ask::At(at) => Ttl::In(Duration::from_millis(at.saturating_sub(now))),
})
})
}
pub fn deadline(&self, key: impl AsRef<[u8]>) -> Result<Option<SystemTime>> {
self.db.run(|inner| {
Ok(match inner.strings.deadline_of(key.as_ref()) {
Ask::At(at) => Some(UNIX_EPOCH + Duration::from_millis(at)),
Ask::Missing | Ask::NoDeadline => None,
})
})
}
pub fn persist(&self, key: impl AsRef<[u8]>) -> Result<bool> {
self.db.run(|inner| Ok(inner.strings.persist(key.as_ref())))
}
pub fn rename(&self, src: impl AsRef<[u8]>, dst: impl AsRef<[u8]>) -> Result<Moved> {
self.db
.run(|inner| Ok(inner.strings.rename(src.as_ref(), dst.as_ref(), false)))
}
pub fn rename_if_new(&self, src: impl AsRef<[u8]>, dst: impl AsRef<[u8]>) -> Result<Moved> {
self.db
.run(|inner| Ok(inner.strings.rename(src.as_ref(), dst.as_ref(), true)))
}
pub fn copy(&self, src: impl AsRef<[u8]>, dst: impl AsRef<[u8]>) -> Result<Moved> {
self.db
.run(|inner| Ok(inner.strings.copy(src.as_ref(), dst.as_ref(), false)))
}
pub fn copy_over(&self, src: impl AsRef<[u8]>, dst: impl AsRef<[u8]>) -> Result<Moved> {
self.db
.run(|inner| Ok(inner.strings.copy(src.as_ref(), dst.as_ref(), true)))
}
pub fn each(&self, mut f: impl FnMut(&[u8])) -> Result<()> {
self.db.run(|inner| {
inner.strings.keys(&mut f);
Ok(())
})
}
pub fn all(&self) -> Result<Vec<Vec<u8>>> {
let mut out = Vec::new();
self.each(|key| out.push(key.to_vec()))?;
Ok(out)
}
pub fn matching(&self, pattern: impl AsRef<[u8]>) -> Result<Vec<Vec<u8>>> {
let pattern = pattern.as_ref();
let mut out = Vec::new();
self.each(|key| {
if yo_common::glob_matches(pattern, key) {
out.push(key.to_vec());
}
})?;
Ok(out)
}
pub fn random(&self) -> Result<Option<Vec<u8>>> {
self.db
.run(|inner| Ok(inner.strings.random_key().map(<[u8]>::to_vec)))
}
}
fn apply(done: Applied) -> Result<bool> {
Ok(match done {
Applied::Ok | Applied::Deleted => true,
Applied::Missing | Applied::NotMet => false,
})
}
fn moment(at: SystemTime) -> Result<u64> {
let ms = at
.duration_since(UNIX_EPOCH)
.map_or(0, |since| since.as_millis());
reachable(u64::try_from(ms).map_err(|_| too_far())?)
}
fn reachable(at: u64) -> Result<u64> {
if at > MAX_AT {
return Err(too_far());
}
Ok(at)
}
fn too_far() -> Error {
Error::new(
Code::Invalid,
"that deadline is further away than a millisecond timestamp reaches, which is the year 4199",
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{MEMORY, open};
#[test]
fn a_key_is_there_until_it_is_not() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
assert!(!keys.exists("k").unwrap());
assert_eq!(keys.kind("k").unwrap(), None);
assert!(!keys.del("k").unwrap());
db.strings().set("k", "v").unwrap();
assert!(keys.exists("k").unwrap());
assert_eq!(keys.kind("k").unwrap(), Some(Kind::String));
assert!(keys.del("k").unwrap());
assert!(!keys.exists("k").unwrap());
}
#[test]
fn a_rename_carries_the_deadline_and_a_copy_is_a_second_value() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("a", "v1").unwrap();
keys.expire_in("a", Duration::from_secs(100)).unwrap();
db.strings().set("b", "v2").unwrap();
assert_eq!(keys.rename("a", "b").unwrap(), Moved::Ok);
assert_eq!(db.strings().get("b").unwrap().as_deref(), Some(&b"v1"[..]));
assert!(keys.ttl("b").unwrap().left().is_some(), "a's and not b's");
assert!(!keys.exists("a").unwrap());
db.set("s").add("m1").unwrap();
assert_eq!(keys.copy("s", "t").unwrap(), Moved::Ok);
db.set("t").add("m2").unwrap();
assert_eq!(db.set("s").len().unwrap(), 1, "the original is intact");
assert_eq!(db.set("t").len().unwrap(), 2);
}
#[test]
fn the_three_answers_a_move_can_give_are_three_and_not_two() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("a", "v1").unwrap();
db.strings().set("b", "v2").unwrap();
assert_eq!(keys.rename_if_new("nosuch", "z").unwrap(), Moved::Missing);
assert_eq!(keys.rename_if_new("a", "b").unwrap(), Moved::Taken);
assert_eq!(keys.copy("a", "b").unwrap(), Moved::Taken);
assert_eq!(db.strings().get("b").unwrap().as_deref(), Some(&b"v2"[..]));
assert_eq!(keys.copy_over("a", "b").unwrap(), Moved::Ok);
assert_eq!(db.strings().get("b").unwrap().as_deref(), Some(&b"v1"[..]));
assert_eq!(keys.rename("a", "a").unwrap(), Moved::Ok);
assert_eq!(keys.rename_if_new("a", "a").unwrap(), Moved::Taken);
}
#[test]
fn several_keys_at_once_count_the_way_redis_counts_them() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("a", "1").unwrap();
db.strings().set("b", "2").unwrap();
assert_eq!(keys.count(&["a", "b", "missing"]).unwrap(), 2);
assert_eq!(keys.count(&["a", "a"]).unwrap(), 2, "the same key twice");
assert_eq!(keys.del_many(&["a", "b", "missing"]).unwrap(), 2);
assert_eq!(keys.count(&["a", "b"]).unwrap(), 0);
}
#[test]
fn a_deadline_lands_on_a_key_whatever_the_key_holds() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("s", "v").unwrap();
db.set("t").add("member").unwrap();
for key in ["s", "t"] {
assert!(keys.expire_in(key, Duration::from_secs(600)).unwrap());
let left = keys.ttl(key).unwrap().left().expect("a deadline");
assert!(left <= Duration::from_secs(600) && left > Duration::from_secs(590));
}
assert_eq!(keys.kind("t").unwrap(), Some(Kind::Set), "still a set");
assert_eq!(db.set("t").len().unwrap(), 1, "with its member");
}
#[test]
fn the_three_answers_are_three_and_not_two() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
assert_eq!(keys.ttl("nothing").unwrap(), Ttl::Missing);
assert!(!keys.ttl("nothing").unwrap().found());
db.strings().set("k", "v").unwrap();
assert_eq!(keys.ttl("k").unwrap(), Ttl::Forever);
assert!(keys.ttl("k").unwrap().found());
assert_eq!(keys.ttl("k").unwrap().left(), None, "forever has none left");
keys.expire_in("k", Duration::from_secs(60)).unwrap();
assert!(matches!(keys.ttl("k").unwrap(), Ttl::In(_)));
}
#[test]
fn a_moment_that_has_gone_removes_the_key_now() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("k", "v").unwrap();
assert!(keys.expire_at("k", UNIX_EPOCH).unwrap(), "it was applied");
assert!(!keys.exists("k").unwrap(), "and applying it took the key");
}
#[test]
fn a_deadline_comes_back_as_the_moment_it_was_set_to() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("k", "v").unwrap();
assert_eq!(keys.deadline("k").unwrap(), None, "no deadline yet");
let at = UNIX_EPOCH + Duration::from_millis(4_000_000_000_000);
assert!(keys.expire_at("k", at).unwrap());
assert_eq!(keys.deadline("k").unwrap(), Some(at));
assert!(keys.persist("k").unwrap());
assert_eq!(keys.deadline("k").unwrap(), None);
assert!(
!keys.persist("k").unwrap(),
"there was nothing left to take"
);
assert!(keys.exists("k").unwrap(), "and the key is still here");
}
#[test]
fn a_condition_decides_whether_the_deadline_moves() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("k", "v").unwrap();
let hour = Duration::from_secs(3600);
let day = Duration::from_secs(86400);
assert!(keys.expire_in_when("k", hour, When::Unset).unwrap());
assert!(
!keys.expire_in_when("k", day, When::Unset).unwrap(),
"taken"
);
assert!(keys.expire_in_when("k", day, When::AlreadySet).unwrap());
assert!(!keys.expire_in_when("k", hour, When::Later).unwrap(), "in");
assert!(keys.expire_in_when("k", hour, When::Earlier).unwrap());
assert!(keys.expire_in_when("k", day, When::Later).unwrap());
keys.persist("k").unwrap();
assert!(
!keys.expire_in_when("k", hour, When::Later).unwrap(),
"no deadline is infinitely far out, so nothing is further"
);
assert!(
keys.expire_in_when("k", hour, When::Earlier).unwrap(),
"and by the same reading everything is nearer"
);
keys.persist("k").unwrap();
assert!(
!keys
.expire_in_when("k", hour, When::EarlierAndAlreadySet)
.unwrap(),
"unless XX takes that reading away"
);
}
#[test]
fn a_condition_that_says_no_leaves_a_key_that_would_have_gone() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("k", "v").unwrap();
keys.expire_in("k", Duration::from_secs(60)).unwrap();
assert!(
!keys.expire_at_when("k", UNIX_EPOCH, When::Unset).unwrap(),
"the condition is checked before the moment is"
);
assert!(keys.exists("k").unwrap());
}
#[test]
fn a_deadline_past_the_year_4199_is_refused_rather_than_clamped() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("k", "v").unwrap();
let err = keys
.expire_in("k", Duration::from_secs(u64::MAX))
.unwrap_err();
assert_eq!(err.code(), Code::Invalid);
let far = UNIX_EPOCH + Duration::from_millis(MAX_AT + 1);
assert_eq!(keys.expire_at("k", far).unwrap_err().code(), Code::Invalid);
assert_eq!(keys.ttl("k").unwrap(), Ttl::Forever, "and nothing moved");
}
#[test]
fn nothing_reads_the_clock_until_a_deadline_exists_to_read_it_for() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("k", "v").unwrap();
keys.exists("k").unwrap();
keys.kind("k").unwrap();
keys.ttl("k").unwrap();
assert!(!db.reads_the_clock(), "asking is not creating");
keys.expire_in("k", Duration::from_secs(60)).unwrap();
assert!(db.reads_the_clock());
}
#[test]
fn a_walk_sees_every_key_whatever_it_holds() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
assert!(keys.all().unwrap().is_empty());
assert_eq!(keys.random().unwrap(), None);
db.strings().set("a", "v").unwrap();
db.set("s").add("m").unwrap();
for i in 0..1_000 {
db.strings().set(format!("n:{i}"), "v").unwrap();
}
let mut all = keys.all().unwrap();
all.sort();
assert_eq!(all.len(), 1_002);
assert_eq!(all[0], b"a");
assert_eq!(keys.matching("n:*").unwrap().len(), 1_000);
assert_eq!(keys.matching("s").unwrap(), vec![b"s".to_vec()]);
assert!(keys.matching("nothing").unwrap().is_empty());
let mut picked = std::collections::HashSet::new();
for _ in 0..100 {
picked.insert(keys.random().unwrap().expect("the database is not empty"));
}
assert!(
picked.len() > 5,
"randomkey is stuck on {} keys",
picked.len()
);
assert!(picked.iter().all(|k| all.contains(k)));
}
#[test]
fn a_walk_does_not_hand_out_a_key_that_has_expired() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("alive", "v").unwrap();
db.strings().set("dead", "v").unwrap();
keys.expire_at("dead", UNIX_EPOCH + Duration::from_secs(1))
.unwrap();
assert_eq!(keys.all().unwrap(), vec![b"alive".to_vec()]);
assert_eq!(keys.random().unwrap(), Some(b"alive".to_vec()));
}
#[test]
fn a_walk_cannot_be_reentered() {
let db = open(MEMORY).unwrap();
let keys = db.keys();
db.strings().set("k", "v").unwrap();
let mut inner = Ok(true);
keys.each(|_| inner = keys.exists("k")).unwrap();
assert_eq!(inner.unwrap_err().code(), Code::Invalid);
}
}