use std::time::Duration;
use yo_common::{Code, Error, Result};
use yo_kv::{Expire, SetOptions, Str};
use crate::db::Handle;
#[derive(Clone)]
pub struct Strings {
pub(crate) db: Handle,
}
impl core::fmt::Debug for Strings {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Strings")
.field("keys", &self.len().ok())
.finish()
}
}
impl Strings {
pub fn get(&self, key: impl AsRef<[u8]>) -> Result<Option<Vec<u8>>> {
self.with(key, |v| v.to_vec())
}
pub fn with<R>(
&self,
key: impl AsRef<[u8]>,
f: impl FnOnce(Str<'_>) -> R,
) -> Result<Option<R>> {
self.db
.run(|inner| Ok(inner.strings.get(key.as_ref())?.map(f)))
}
pub fn exists(&self, key: impl AsRef<[u8]>) -> Result<bool> {
self.db.run(|inner| Ok(inner.strings.exists(key.as_ref())))
}
pub fn len_of(&self, key: impl AsRef<[u8]>) -> Result<usize> {
self.db.run(|inner| inner.strings.strlen(key.as_ref()))
}
pub fn set(&self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) -> Result<()> {
self.db
.run(|inner| inner.strings.set_plain(key.as_ref(), value.as_ref()))
}
pub fn set_if_missing(&self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) -> Result<bool> {
self.db
.run(|inner| inner.strings.setnx(key.as_ref(), value.as_ref()))
}
pub fn set_for(
&self,
key: impl AsRef<[u8]>,
value: impl AsRef<[u8]>,
ttl: Duration,
) -> Result<()> {
let ms = u64::try_from(ttl.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)?;
inner
.strings
.set(
key.as_ref(),
value.as_ref(),
SetOptions::PLAIN.expiring(Expire::At(at)),
)
.map(|_| ())
})
}
pub fn ttl(&self, key: impl AsRef<[u8]>) -> Result<Option<Duration>> {
self.db.run(|inner| {
let now = inner.strings.clock().now_ms();
Ok(inner
.strings
.expire_at(key.as_ref())
.map(|at| Duration::from_millis(at.saturating_sub(now))))
})
}
pub fn replace(
&self,
key: impl AsRef<[u8]>,
value: impl AsRef<[u8]>,
) -> Result<Option<Vec<u8>>> {
self.db
.run(|inner| inner.strings.getset(key.as_ref(), value.as_ref()))
}
pub fn take(&self, key: impl AsRef<[u8]>) -> Result<Option<Vec<u8>>> {
self.db.run(|inner| inner.strings.getdel(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 set_many<K: AsRef<[u8]>, V: AsRef<[u8]>>(&self, pairs: &[(K, V)]) -> Result<()> {
self.db.run(|inner| {
inner
.strings
.mset(pairs.iter().map(|(k, v)| (k.as_ref(), v.as_ref())))
})
}
pub fn get_many<K: AsRef<[u8]>>(&self, keys: &[K]) -> Result<Vec<Option<Vec<u8>>>> {
self.db.run(|inner| {
Ok(keys
.iter()
.map(|k| inner.strings.mget_one(k.as_ref()).map(|v| v.to_vec()))
.collect())
})
}
pub fn incr(&self, key: impl AsRef<[u8]>) -> Result<i64> {
self.db.run(|inner| inner.strings.incr(key.as_ref()))
}
pub fn incr_by(&self, key: impl AsRef<[u8]>, by: i64) -> Result<i64> {
self.db.run(|inner| inner.strings.incrby(key.as_ref(), by))
}
pub fn decr(&self, key: impl AsRef<[u8]>) -> Result<i64> {
self.db.run(|inner| inner.strings.decr(key.as_ref()))
}
pub fn incr_by_float(&self, key: impl AsRef<[u8]>, by: f64) -> Result<f64> {
self.db
.run(|inner| inner.strings.incrbyfloat(key.as_ref(), by))
}
pub fn append(&self, key: impl AsRef<[u8]>, tail: impl AsRef<[u8]>) -> Result<usize> {
self.db
.run(|inner| inner.strings.append(key.as_ref(), tail.as_ref()))
}
pub fn len(&self) -> Result<usize> {
self.db.run(|inner| Ok(inner.strings.len()))
}
pub fn is_empty(&self) -> Result<bool> {
self.db.run(|inner| Ok(inner.strings.is_empty()))
}
pub fn expired_keys(&self) -> Result<u64> {
self.db.run(|inner| Ok(inner.strings.expired_keys()))
}
}
fn too_far() -> Error {
Error::new(
Code::Invalid,
"that deadline is further away than a millisecond timestamp reaches",
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{MEMORY, open};
#[test]
fn the_string_commands_are_the_ones_a_redis_client_would_send() {
let db = open(MEMORY).unwrap();
let keys = db.strings();
keys.set("greeting", "hello").unwrap();
assert_eq!(
keys.get("greeting").unwrap().as_deref(),
Some(&b"hello"[..])
);
assert_eq!(keys.len_of("greeting").unwrap(), 5);
assert_eq!(keys.append("greeting", " there").unwrap(), 11);
assert!(keys.exists("greeting").unwrap());
assert!(!keys.set_if_missing("greeting", "other").unwrap());
assert!(keys.set_if_missing("fresh", "yes").unwrap());
assert_eq!(
keys.replace("greeting", "hi").unwrap().as_deref(),
Some(&b"hello there"[..])
);
assert_eq!(keys.take("greeting").unwrap().as_deref(), Some(&b"hi"[..]));
assert!(!keys.exists("greeting").unwrap());
assert!(!keys.del("greeting").unwrap());
}
#[test]
fn a_counter_starts_at_zero_without_being_created() {
let db = open(MEMORY).unwrap();
let keys = db.strings();
assert_eq!(keys.incr("hits").unwrap(), 1);
assert_eq!(keys.incr_by("hits", 9).unwrap(), 10);
assert_eq!(keys.decr("hits").unwrap(), 9);
assert_eq!(keys.incr_by_float("ratio", 1.5).unwrap(), 1.5);
keys.set("word", "nope").unwrap();
let e = keys.incr("word").expect_err("that is not a number");
assert_eq!(e.code(), Code::Invalid);
assert_eq!(e.message(), "value is not an integer or out of range");
}
#[test]
fn many_keys_at_once_go_in_and_come_out_together() {
let db = open(MEMORY).unwrap();
let keys = db.strings();
keys.set_many(&[("a", "1"), ("b", "2"), ("c", "3")])
.unwrap();
let got = keys.get_many(&["a", "c", "missing"]).unwrap();
assert_eq!(got[0].as_deref(), Some(&b"1"[..]));
assert_eq!(got[1].as_deref(), Some(&b"3"[..]));
assert_eq!(got[2], None);
assert_eq!(keys.len().unwrap(), 3);
assert!(!keys.is_empty().unwrap());
}
#[test]
fn a_deadline_is_what_makes_time_start_moving() {
let db = open(MEMORY).unwrap();
let keys = db.strings();
keys.set("plain", "v").unwrap();
assert_eq!(keys.ttl("plain").unwrap(), None);
assert!(!db.reads_the_clock());
keys.set_for("short", "v", Duration::from_millis(50))
.unwrap();
assert!(db.reads_the_clock());
assert!(keys.ttl("short").unwrap().unwrap() <= Duration::from_millis(50));
std::thread::sleep(Duration::from_millis(60));
assert_eq!(keys.get("short").unwrap(), None);
assert_eq!(keys.expired_keys().unwrap(), 1);
assert_eq!(keys.get("plain").unwrap().as_deref(), Some(&b"v"[..]));
}
#[test]
fn a_key_is_bytes_and_not_only_text() {
let db = open(MEMORY).unwrap();
let keys = db.strings();
keys.set(b"\x00\xff", vec![1u8, 2, 3]).unwrap();
assert_eq!(
keys.get(b"\x00\xff").unwrap().as_deref(),
Some(&[1u8, 2, 3][..])
);
assert_eq!(keys.with(b"\x00\xff", |v| v.len()).unwrap(), Some(3));
assert!(format!("{keys:?}").contains("keys"));
}
}