diff --git a/src/client.rs b/src/client.rs index 8e58307..c6520b1 100644 --- a/src/client.rs +++ b/src/client.rs @@ -144,6 +144,32 @@ impl Client { Ok(r.try_get_u8()? == 1) } + pub async fn ttl(&mut self, key: &str) -> Result> { + let mut bytes = BytesMut::new(); + bytes.put_u16(3); + bytes.put_slice(b"ttl"); + + let key_length: u16 = key + .len() + .try_into() + .map_err(|_| AppError::KeyLength(key.len()))?; + + bytes.put_u16(key_length); + bytes.put_slice(key.as_bytes()); + + self.connection.write(bytes.into()).await?; + + let mut r = self.get_response().await?; + + let ttl = match r.try_get_u8()? { + 1 => Some(r.try_get_u64()?), + 0 => None, + _ => return Err(AppError::InvalidCommandResponse), + }; + + Ok(ttl) + } + async fn get_response(&mut self) -> Result { self.connection .read_bytes() diff --git a/src/commands/mod.rs b/src/commands/mod.rs index b79ab75..a33d14f 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -1,7 +1,8 @@ -pub mod delete; +mod delete; mod get; -pub mod has; -pub mod set; +mod has; +mod set; +mod ttl; use std::io::Cursor; @@ -10,6 +11,7 @@ use delete::Delete; use get::Get; use has::Has; use set::Set; +use ttl::Ttl; use crate::{Result, connection::Connection, database::Database, errors::AppError}; @@ -19,6 +21,7 @@ pub enum Command { Set(Set), Delete(Delete), Has(Has), + Ttl(Ttl), } impl Command { @@ -28,6 +31,7 @@ impl Command { Command::Set(set) => set.execute(db, connection).await, Command::Delete(delete) => delete.execute(db, connection).await, Command::Has(has) => has.execute(db, connection).await, + Command::Ttl(ttl) => ttl.execute(db, connection).await, } } @@ -51,6 +55,7 @@ impl Command { "set" => Self::Set(Set::parse(bytes)?), "delete" => Self::Delete(Delete::parse(bytes)?), "has" => Self::Has(Has::parse(bytes)?), + "ttl" => Self::Ttl(Ttl::parse(bytes)?), _ => return Err(AppError::UnknownCommand(command_name)), }; diff --git a/src/commands/ttl.rs b/src/commands/ttl.rs new file mode 100644 index 0000000..1ef2e44 --- /dev/null +++ b/src/commands/ttl.rs @@ -0,0 +1,42 @@ +use std::io::Cursor; + +use bytes::{Buf as _, BufMut, Bytes, BytesMut}; + +use crate::{Result, connection::Connection, database::Database, errors::AppError}; + +#[derive(Debug, Clone)] +pub struct Ttl { + key: String, +} + +impl Ttl { + pub async fn execute(self, db: &Database, connection: &mut Connection) -> Result<()> { + let ttl = db.ttl(&self.key).await; + + let Some(ttl) = ttl else { + connection.write(Bytes::from_static(&[0])).await?; + return Ok(()); + }; + + let mut buf = BytesMut::new(); + + buf.put_u8(1); + buf.put_u64(ttl); + + connection.write(buf.into()).await?; + + Ok(()) + } + + pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result { + let key_length = bytes.try_get_u16()? as usize; + + if bytes.remaining() < key_length { + return Err(AppError::IncompleteCommandBuffer); + } + + let key = String::from_utf8(bytes.copy_to_bytes(key_length).to_vec())?; + + Ok(Self { key }) + } +} diff --git a/src/database.rs b/src/database.rs index e9623d7..84922aa 100644 --- a/src/database.rs +++ b/src/database.rs @@ -128,6 +128,16 @@ impl Database { state.entries.contains_key(key) } + pub async fn ttl(&self, key: &str) -> Option { + self.state + .lock() + .await + .entries + .get(key) + .map(|v| v.expiration.map(|e| (e - Instant::now()).as_secs())) + .flatten() + } + pub async fn shutdown(&mut self) { self.state.lock().await.shutdown = true; self.notify.notify_one(); diff --git a/src/tests.rs b/src/tests.rs index 1602be9..9e1f8cd 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -15,12 +15,18 @@ async fn expiration() -> Result<(), Box> { let mut client = client("127.0.0.1:6171").await?; client - .set("test-key", "test-value".as_bytes(), Some(2)) + .set("test-key", "test-value".as_bytes(), Some(3)) .await .unwrap(); + assert!(client.has("test-key").await.unwrap()); + assert_eq!(client.ttl("test-key").await.unwrap(), Some(2)); + tokio::time::sleep(Duration::from_secs(1)).await; + assert!(client.has("test-key").await.unwrap()); + assert_eq!(client.ttl("test-key").await.unwrap(), Some(1)); + tokio::time::sleep(Duration::from_secs(2)).await; assert!(!client.has("test-key").await.unwrap());