Compare commits
4 Commits
51fdaa4120
...
c527bb0072
| Author | SHA1 | Date | |
|---|---|---|---|
| c527bb0072 | |||
| cff5c37a40 | |||
| 87eb32eb5d | |||
| 02aaef1560 |
@@ -39,6 +39,13 @@ enum Commands {
|
||||
Persist {
|
||||
key: String,
|
||||
},
|
||||
|
||||
#[command(name = "mget")]
|
||||
MGet {
|
||||
#[arg(num_args = 1..)]
|
||||
keys: Vec<String>,
|
||||
},
|
||||
|
||||
#[command(aliases = &["exit", "q"])]
|
||||
Quit,
|
||||
}
|
||||
@@ -108,6 +115,13 @@ async fn main() -> Result<()> {
|
||||
let value = client.persist(&key).await?;
|
||||
println!("{value:?}");
|
||||
}
|
||||
|
||||
Commands::MGet { keys } => {
|
||||
let value = client.m_get(keys).await?;
|
||||
|
||||
println!("{value:?}");
|
||||
}
|
||||
|
||||
Commands::Quit => break,
|
||||
}
|
||||
}
|
||||
|
||||
164
src/buffer.rs
Normal file
164
src/buffer.rs
Normal file
@@ -0,0 +1,164 @@
|
||||
use std::num::TryFromIntError;
|
||||
|
||||
use bytes::{Buf, BufMut, Bytes};
|
||||
|
||||
use crate::{Result, errors::AppError};
|
||||
|
||||
pub trait ArchiveBuf<B: Buf> {
|
||||
fn try_get_string(&mut self) -> Result<String>;
|
||||
fn try_get_short_string(&mut self) -> Result<String>;
|
||||
fn try_get_bytes(&mut self) -> Result<Bytes>;
|
||||
fn try_get_bool(&mut self) -> Result<bool>;
|
||||
fn try_get_option<T, F, E>(&mut self, f: F) -> Result<Option<T>>
|
||||
where
|
||||
T: Sized,
|
||||
F: FnOnce(&mut B) -> std::result::Result<T, E>,
|
||||
AppError: From<E>;
|
||||
fn try_get_vec<T, F, E>(&mut self, f: F) -> Result<Vec<T>>
|
||||
where
|
||||
T: Sized,
|
||||
F: FnMut(&mut B) -> std::result::Result<T, E>,
|
||||
AppError: From<E>;
|
||||
}
|
||||
|
||||
pub trait ArchiveBufMut<B: Buf> {
|
||||
fn put_string(&mut self, s: &str) -> Result<()>;
|
||||
fn put_short_string(&mut self, s: &str) -> Result<()>;
|
||||
fn put_bytes_with_length<T: AsRef<[u8]>>(&mut self, bytes: T);
|
||||
fn put_option<T, F>(&mut self, value: Option<T>, f: F)
|
||||
where
|
||||
T: Sized,
|
||||
F: FnOnce(&mut B, T);
|
||||
fn try_put_vec<T, F, V, E>(&mut self, value: V, f: F) -> Result<()>
|
||||
where
|
||||
T: Sized,
|
||||
F: FnMut(&T, &mut B) -> std::result::Result<(), E>,
|
||||
V: AsRef<[T]>,
|
||||
AppError: From<E> + From<TryFromIntError>;
|
||||
}
|
||||
|
||||
impl<B: Buf> ArchiveBuf<B> for B {
|
||||
fn try_get_string(&mut self) -> Result<String> {
|
||||
let len = self.try_get_u32()? as usize;
|
||||
|
||||
if self.remaining() < len {
|
||||
return Err(AppError::IncompleteBuffer);
|
||||
}
|
||||
|
||||
Ok(String::from_utf8(self.copy_to_bytes(len).to_vec())?)
|
||||
}
|
||||
|
||||
fn try_get_short_string(&mut self) -> Result<String> {
|
||||
let len = self.try_get_u16()? as usize;
|
||||
|
||||
if self.remaining() < len {
|
||||
return Err(AppError::IncompleteBuffer);
|
||||
}
|
||||
|
||||
Ok(String::from_utf8(self.copy_to_bytes(len).to_vec())?)
|
||||
}
|
||||
|
||||
fn try_get_bytes(&mut self) -> Result<Bytes> {
|
||||
let len = self.try_get_u32()? as usize;
|
||||
|
||||
if self.remaining() < len {
|
||||
return Err(AppError::IncompleteBuffer);
|
||||
}
|
||||
|
||||
let data = self.copy_to_bytes(len);
|
||||
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
fn try_get_bool(&mut self) -> Result<bool> {
|
||||
Ok(match self.try_get_u8()? {
|
||||
1 => true,
|
||||
0 => false,
|
||||
_ => return Err(AppError::UnexpectedData),
|
||||
})
|
||||
}
|
||||
|
||||
fn try_get_option<T, F, E>(&mut self, f: F) -> Result<Option<T>>
|
||||
where
|
||||
T: Sized,
|
||||
F: FnOnce(&mut B) -> std::result::Result<T, E>,
|
||||
AppError: From<E>,
|
||||
{
|
||||
if !self.try_get_bool()? {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(f(self)?))
|
||||
}
|
||||
|
||||
fn try_get_vec<T, F, E>(&mut self, mut f: F) -> Result<Vec<T>>
|
||||
where
|
||||
T: Sized,
|
||||
F: FnMut(&mut B) -> std::result::Result<T, E>,
|
||||
AppError: From<E>,
|
||||
{
|
||||
let len = self.try_get_u16()?;
|
||||
|
||||
let mut vec = Vec::with_capacity(len.into());
|
||||
|
||||
for _ in 0..len {
|
||||
vec.push(f(self)?);
|
||||
}
|
||||
|
||||
Ok(vec)
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: Buf + BufMut> ArchiveBufMut<B> for B {
|
||||
fn put_string(&mut self, s: &str) -> Result<()> {
|
||||
self.put_u32(s.len().try_into()?);
|
||||
self.put_slice(s.as_bytes());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn put_short_string(&mut self, s: &str) -> Result<()> {
|
||||
self.put_u16(s.len().try_into()?);
|
||||
self.put_slice(s.as_bytes());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn put_bytes_with_length<T: AsRef<[u8]>>(&mut self, bytes: T) {
|
||||
let bytes = bytes.as_ref();
|
||||
|
||||
self.put_u32(bytes.len() as u32);
|
||||
self.put_slice(bytes);
|
||||
}
|
||||
|
||||
fn put_option<T, F>(&mut self, value: Option<T>, f: F)
|
||||
where
|
||||
T: Sized,
|
||||
F: FnOnce(&mut B, T),
|
||||
{
|
||||
let Some(value) = value else {
|
||||
self.put_u8(0);
|
||||
return;
|
||||
};
|
||||
|
||||
self.put_u8(1);
|
||||
f(self, value);
|
||||
}
|
||||
|
||||
fn try_put_vec<T, F, V, E>(&mut self, vec: V, mut f: F) -> Result<()>
|
||||
where
|
||||
T: Sized,
|
||||
F: FnMut(&T, &mut B) -> std::result::Result<(), E>,
|
||||
V: AsRef<[T]>,
|
||||
AppError: From<E> + From<TryFromIntError>,
|
||||
{
|
||||
let vec = vec.as_ref();
|
||||
self.put_u16(vec.len().try_into()?);
|
||||
|
||||
for element in vec {
|
||||
f(element, self)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
217
src/client.rs
217
src/client.rs
@@ -1,7 +1,16 @@
|
||||
use bytes::{Buf, BufMut as _, Bytes, BytesMut};
|
||||
use bytes::{Buf, Bytes, BytesMut};
|
||||
use tokio::net::{TcpStream, ToSocketAddrs};
|
||||
|
||||
use crate::{Result, connection::Connection, errors::AppError};
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::ArchiveBuf,
|
||||
commands::{
|
||||
delete::Delete, expire::Expire, get::Get, has::Has, m_get::MGet, m_set::MSet,
|
||||
persist::Persist, set::Set, ttl::Ttl,
|
||||
},
|
||||
connection::Connection,
|
||||
errors::AppError,
|
||||
};
|
||||
|
||||
pub struct Client {
|
||||
connection: Connection,
|
||||
@@ -19,35 +28,16 @@ impl Client {
|
||||
pub async fn get(&mut self, key: &str) -> Result<Option<Bytes>> {
|
||||
let mut bytes = BytesMut::new();
|
||||
|
||||
bytes.put_u16(3);
|
||||
bytes.put_slice(b"get");
|
||||
|
||||
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());
|
||||
let cmd = Get::new(key.to_owned());
|
||||
cmd.put(&mut bytes)?;
|
||||
|
||||
self.connection.write(bytes.into()).await?;
|
||||
|
||||
let mut r = self.get_response().await?;
|
||||
let response = match r.try_get_u8()? {
|
||||
0 => None,
|
||||
1 => {
|
||||
let len = r.try_get_u32()? as usize;
|
||||
|
||||
if r.remaining() < len {
|
||||
return Err(AppError::InvalidCommandResponse);
|
||||
}
|
||||
let value = r.try_get_option(ArchiveBuf::try_get_bytes)?;
|
||||
|
||||
Some(r.copy_to_bytes(len))
|
||||
}
|
||||
_ => return Err(AppError::InvalidCommandResponse),
|
||||
};
|
||||
|
||||
Ok(response)
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
pub async fn set(
|
||||
@@ -56,172 +46,137 @@ impl Client {
|
||||
data: &[u8],
|
||||
expiration_secs: Option<u64>,
|
||||
) -> Result<()> {
|
||||
let mut bytes = BytesMut::new();
|
||||
let mut buf = BytesMut::new();
|
||||
|
||||
bytes.put_u16(3);
|
||||
bytes.put_slice(b"set");
|
||||
let cmd = Set::new(key.to_owned(), data.into(), expiration_secs);
|
||||
cmd.put(&mut buf)?;
|
||||
|
||||
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());
|
||||
|
||||
bytes.put_u32(data.len() as u32);
|
||||
bytes.put_slice(data);
|
||||
|
||||
match expiration_secs {
|
||||
Some(seconds) => {
|
||||
bytes.put_u8(1);
|
||||
bytes.put_u64(seconds);
|
||||
}
|
||||
None => {
|
||||
bytes.put_u8(0);
|
||||
}
|
||||
}
|
||||
|
||||
self.connection.write(bytes.into()).await?;
|
||||
self.connection.write(buf.into()).await?;
|
||||
|
||||
let mut r = self.get_response().await?;
|
||||
|
||||
match r.try_get_u8()? {
|
||||
1 => return Ok(()),
|
||||
_ => return Err(AppError::InvalidCommandResponse),
|
||||
if !r.try_get_bool()? {
|
||||
return Err(AppError::InvalidCommandResponse);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn delete(&mut self, key: &str) -> Result<Option<Bytes>> {
|
||||
let mut bytes = BytesMut::new();
|
||||
|
||||
bytes.put_u16(6);
|
||||
bytes.put_slice(b"delete");
|
||||
|
||||
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());
|
||||
let cmd = Delete::new(key.to_owned());
|
||||
cmd.put(&mut bytes)?;
|
||||
|
||||
self.connection.write(bytes.into()).await?;
|
||||
|
||||
let mut r = self.get_response().await?;
|
||||
|
||||
let response = match r.try_get_u8()? {
|
||||
1 => {
|
||||
let len = r.try_get_u32()?;
|
||||
let bytes = r.copy_to_bytes(len as usize);
|
||||
let value = r.try_get_option(ArchiveBuf::try_get_bytes)?;
|
||||
|
||||
Some(bytes)
|
||||
}
|
||||
0 => None,
|
||||
_ => return Err(AppError::InvalidCommandResponse),
|
||||
};
|
||||
|
||||
Ok(response)
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
pub async fn has(&mut self, key: &str) -> Result<bool> {
|
||||
let mut bytes = BytesMut::new();
|
||||
bytes.put_u16(3);
|
||||
bytes.put_slice(b"has");
|
||||
|
||||
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());
|
||||
let cmd = Has::new(key.to_owned());
|
||||
cmd.put(&mut bytes)?;
|
||||
|
||||
self.connection.write(bytes.into()).await?;
|
||||
|
||||
let mut r = self.get_response().await?;
|
||||
|
||||
Ok(r.try_get_u8()? == 1)
|
||||
let has = r.try_get_bool()?;
|
||||
|
||||
Ok(has)
|
||||
}
|
||||
|
||||
pub async fn ttl(&mut self, key: &str) -> Result<Option<u64>> {
|
||||
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());
|
||||
let cmd = Ttl::new(key.to_owned());
|
||||
cmd.put(&mut 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),
|
||||
};
|
||||
let ttl = r.try_get_option(Bytes::try_get_u64)?;
|
||||
|
||||
Ok(ttl)
|
||||
}
|
||||
|
||||
pub async fn expire(&mut self, key: &str, seconds: u64) -> Result<bool> {
|
||||
let mut bytes = BytesMut::new();
|
||||
bytes.put_u16(6);
|
||||
bytes.put_slice(b"expire");
|
||||
|
||||
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());
|
||||
|
||||
bytes.put_u64(seconds);
|
||||
let cmd = Expire::new(key.to_owned(), seconds);
|
||||
cmd.put(&mut bytes)?;
|
||||
|
||||
self.connection.write(bytes.into()).await?;
|
||||
|
||||
let mut r = self.get_response().await?;
|
||||
|
||||
let success = match r.try_get_u8()? {
|
||||
1 => true,
|
||||
0 => false,
|
||||
_ => return Err(AppError::InvalidCommandResponse),
|
||||
};
|
||||
|
||||
Ok(success)
|
||||
Ok(r.try_get_bool()?)
|
||||
}
|
||||
|
||||
pub async fn persist(&mut self, key: &str) -> Result<bool> {
|
||||
let mut bytes = BytesMut::new();
|
||||
bytes.put_u16(7);
|
||||
bytes.put_slice(b"persist");
|
||||
|
||||
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());
|
||||
let cmd = Persist::new(key.to_owned());
|
||||
cmd.put(&mut bytes)?;
|
||||
|
||||
self.connection.write(bytes.into()).await?;
|
||||
|
||||
let mut r = self.get_response().await?;
|
||||
|
||||
let success = match r.try_get_u8()? {
|
||||
1 => true,
|
||||
0 => false,
|
||||
_ => return Err(AppError::InvalidCommandResponse),
|
||||
};
|
||||
Ok(r.try_get_bool()?)
|
||||
}
|
||||
|
||||
Ok(success)
|
||||
pub async fn m_set(
|
||||
&mut self,
|
||||
keys: Vec<&str>,
|
||||
data: Vec<&[u8]>,
|
||||
expirations: Vec<Option<u64>>,
|
||||
) -> Result<()> {
|
||||
let mut bytes = BytesMut::new();
|
||||
|
||||
let len = keys.len().min(data.len()).min(expirations.len());
|
||||
|
||||
let mut sets = Vec::with_capacity(len);
|
||||
|
||||
for i in 0..len {
|
||||
sets.push(Set::new(keys[i].to_owned(), data[i].into(), expirations[i]));
|
||||
}
|
||||
|
||||
let cmd = MSet::new(sets);
|
||||
cmd.put(&mut bytes)?;
|
||||
|
||||
self.connection.write(bytes.into()).await?;
|
||||
|
||||
let mut r = self.get_response().await?;
|
||||
|
||||
if !r.try_get_bool()? {
|
||||
return Err(AppError::InvalidCommandResponse);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn m_get(&mut self, keys: Vec<String>) -> Result<Vec<Option<Bytes>>> {
|
||||
let mut bytes = BytesMut::new();
|
||||
|
||||
let gets: Vec<Get> = keys.into_iter().map(Get::new).collect();
|
||||
let cmd = MGet::new(gets);
|
||||
cmd.put(&mut bytes)?;
|
||||
|
||||
self.connection.write(bytes.into()).await?;
|
||||
|
||||
let mut r = self.get_response().await?;
|
||||
|
||||
let values = r.try_get_vec(|b| b.try_get_option(ArchiveBuf::try_get_bytes))?;
|
||||
|
||||
Ok(values)
|
||||
}
|
||||
|
||||
async fn get_response(&mut self) -> Result<Bytes> {
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use bytes::{Buf as _, BufMut as _, BytesMut};
|
||||
use bytes::BytesMut;
|
||||
|
||||
use crate::{Result, connection::Connection, database::Database, errors::AppError};
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::{ArchiveBuf as _, ArchiveBufMut},
|
||||
connection::Connection,
|
||||
database::Database,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Delete {
|
||||
@@ -10,33 +15,33 @@ pub struct Delete {
|
||||
}
|
||||
|
||||
impl Delete {
|
||||
pub fn new(key: String) -> Self {
|
||||
Self { key }
|
||||
}
|
||||
|
||||
pub async fn execute(self, db: &Database, connection: &mut Connection) -> Result<()> {
|
||||
let value = db.delete(&self.key).await;
|
||||
|
||||
let mut buf = BytesMut::new();
|
||||
match value {
|
||||
Some(v) => {
|
||||
buf.put_u8(1);
|
||||
buf.put_u32(v.len() as u32);
|
||||
buf.put_slice(&v);
|
||||
}
|
||||
None => buf.put_u8(0),
|
||||
}
|
||||
buf.put_option(value, ArchiveBufMut::put_bytes_with_length);
|
||||
|
||||
connection.write(buf.into()).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
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())?;
|
||||
pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let key = buf.try_get_short_string()?;
|
||||
|
||||
Ok(Self { key })
|
||||
}
|
||||
|
||||
pub fn put(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string("delete")?;
|
||||
self.put_without_cmd_name(buf)
|
||||
}
|
||||
|
||||
pub fn put_without_cmd_name(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string(&self.key)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use bytes::{Buf as _, Bytes};
|
||||
use bytes::{Buf as _, BufMut as _, Bytes, BytesMut};
|
||||
|
||||
use crate::{Result, connection::Connection, database::Database, errors::AppError};
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::{ArchiveBuf as _, ArchiveBufMut as _},
|
||||
connection::Connection,
|
||||
database::Database,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Expire {
|
||||
@@ -11,27 +16,38 @@ pub struct Expire {
|
||||
}
|
||||
|
||||
impl Expire {
|
||||
pub fn new(key: String, seconds: u64) -> Self {
|
||||
Self { key, seconds }
|
||||
}
|
||||
|
||||
pub async fn execute(self, db: &Database, connection: &mut Connection) -> Result<()> {
|
||||
let value = db.expire(&self.key, self.seconds).await?;
|
||||
let success = db.expire(&self.key, self.seconds).await?;
|
||||
|
||||
connection
|
||||
.write(Bytes::from_static(if value { &[1] } else { &[0] }))
|
||||
.write(Bytes::from_static(if success { &[1] } else { &[0] }))
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
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())?;
|
||||
|
||||
let seconds = bytes.try_get_u64()?;
|
||||
pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let key = buf.try_get_short_string()?;
|
||||
let seconds = buf.try_get_u64()?;
|
||||
|
||||
Ok(Self { key, seconds })
|
||||
}
|
||||
|
||||
pub fn put(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string("expire")?;
|
||||
|
||||
self.put_without_cmd_name(buf)
|
||||
}
|
||||
|
||||
pub fn put_without_cmd_name(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string(&self.key)?;
|
||||
|
||||
buf.put_u64(self.seconds);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,44 +1,47 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use bytes::{Buf as _, BufMut as _, BytesMut};
|
||||
use bytes::BytesMut;
|
||||
|
||||
use crate::{Result, connection::Connection, database::Database, errors::AppError};
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::{ArchiveBuf as _, ArchiveBufMut},
|
||||
connection::Connection,
|
||||
database::Database,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Get {
|
||||
key: String,
|
||||
pub(super) key: String,
|
||||
}
|
||||
|
||||
impl Get {
|
||||
pub fn new(key: String) -> Self {
|
||||
Self { key }
|
||||
}
|
||||
|
||||
pub async fn execute(self, db: &Database, connection: &mut Connection) -> Result<()> {
|
||||
let value = db.get(&self.key).await;
|
||||
|
||||
let mut buf = BytesMut::new();
|
||||
match value {
|
||||
Some(v) => {
|
||||
buf.put_u8(1);
|
||||
buf.put_u32(v.len() as u32);
|
||||
buf.put_slice(&v);
|
||||
}
|
||||
None => {
|
||||
buf.put_u8(0);
|
||||
}
|
||||
}
|
||||
buf.put_option(value, ArchiveBufMut::put_bytes_with_length);
|
||||
|
||||
connection.write(buf.into()).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
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())?;
|
||||
pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let key = buf.try_get_short_string()?;
|
||||
|
||||
Ok(Self { key })
|
||||
}
|
||||
|
||||
pub fn put(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string("get")?;
|
||||
self.put_without_cmd_name(buf)
|
||||
}
|
||||
|
||||
pub fn put_without_cmd_name(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string(&self.key)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use bytes::{Buf as _, Bytes};
|
||||
use bytes::{Bytes, BytesMut};
|
||||
|
||||
use crate::{Result, connection::Connection, database::Database, errors::AppError};
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::{ArchiveBuf as _, ArchiveBufMut as _},
|
||||
connection::Connection,
|
||||
database::Database,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Has {
|
||||
@@ -10,6 +15,10 @@ pub struct Has {
|
||||
}
|
||||
|
||||
impl Has {
|
||||
pub fn new(key: String) -> Self {
|
||||
Self { key }
|
||||
}
|
||||
|
||||
pub async fn execute(self, db: &Database, connection: &mut Connection) -> Result<()> {
|
||||
let value = db.has(&self.key).await;
|
||||
|
||||
@@ -20,15 +29,18 @@ impl Has {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
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())?;
|
||||
pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let key = buf.try_get_short_string()?;
|
||||
|
||||
Ok(Self { key })
|
||||
}
|
||||
|
||||
pub fn put(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string("has")?;
|
||||
self.put_without_cmd_name(buf)
|
||||
}
|
||||
|
||||
pub fn put_without_cmd_name(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string(&self.key)
|
||||
}
|
||||
}
|
||||
|
||||
60
src/commands/m_get.rs
Normal file
60
src/commands/m_get.rs
Normal file
@@ -0,0 +1,60 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::{ArchiveBuf as _, ArchiveBufMut},
|
||||
connection::Connection,
|
||||
database::Database,
|
||||
errors::AppError,
|
||||
};
|
||||
use bytes::BytesMut;
|
||||
|
||||
use super::get::Get;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MGet {
|
||||
gets: Vec<Get>,
|
||||
}
|
||||
|
||||
impl MGet {
|
||||
pub fn new(gets: Vec<Get>) -> Self {
|
||||
Self { gets }
|
||||
}
|
||||
|
||||
pub async fn execute(self, db: &Database, connection: &mut Connection) -> Result<()> {
|
||||
let mut values = Vec::with_capacity(self.gets.len());
|
||||
|
||||
for get in self.gets {
|
||||
values.push(db.get(&get.key).await);
|
||||
}
|
||||
|
||||
let mut buf = BytesMut::new();
|
||||
|
||||
buf.try_put_vec(values, |data, buf| {
|
||||
buf.put_option(data.as_deref(), ArchiveBufMut::put_bytes_with_length);
|
||||
|
||||
Ok::<(), AppError>(())
|
||||
})?;
|
||||
|
||||
connection.write(buf.into()).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let gets = buf.try_get_vec(Get::parse)?;
|
||||
|
||||
Ok(Self { gets })
|
||||
}
|
||||
|
||||
pub fn put(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string("mget")?;
|
||||
self.put_without_cmd_name(buf)
|
||||
}
|
||||
|
||||
pub fn put_without_cmd_name(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.try_put_vec(&self.gets, Get::put_without_cmd_name)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
49
src/commands/m_set.rs
Normal file
49
src/commands/m_set.rs
Normal file
@@ -0,0 +1,49 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::{ArchiveBuf as _, ArchiveBufMut as _},
|
||||
connection::Connection,
|
||||
database::Database,
|
||||
};
|
||||
use bytes::{Bytes, BytesMut};
|
||||
|
||||
use super::set::Set;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MSet {
|
||||
sets: Vec<Set>,
|
||||
}
|
||||
|
||||
impl MSet {
|
||||
pub fn new(sets: Vec<Set>) -> Self {
|
||||
Self { sets }
|
||||
}
|
||||
|
||||
pub async fn execute(self, db: &Database, connection: &mut Connection) -> Result<()> {
|
||||
for set in self.sets {
|
||||
db.set(set.key, set.data, set.expiration).await?;
|
||||
}
|
||||
|
||||
connection.write(Bytes::from_static(&[1])).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let sets = buf.try_get_vec(Set::parse)?;
|
||||
|
||||
Ok(Self { sets })
|
||||
}
|
||||
|
||||
pub fn put(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string("mset")?;
|
||||
self.put_without_cmd_name(buf)
|
||||
}
|
||||
|
||||
pub fn put_without_cmd_name(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.try_put_vec(&self.sets, Set::put_without_cmd_name)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,23 +1,29 @@
|
||||
mod delete;
|
||||
mod expire;
|
||||
mod get;
|
||||
mod has;
|
||||
mod persist;
|
||||
mod set;
|
||||
mod ttl;
|
||||
pub mod delete;
|
||||
pub mod expire;
|
||||
pub mod get;
|
||||
pub mod has;
|
||||
pub mod m_get;
|
||||
pub mod m_set;
|
||||
pub mod persist;
|
||||
pub mod set;
|
||||
pub mod ttl;
|
||||
|
||||
use std::io::Cursor;
|
||||
|
||||
use bytes::{Buf, BytesMut};
|
||||
use bytes::BytesMut;
|
||||
use delete::Delete;
|
||||
use expire::Expire;
|
||||
use get::Get;
|
||||
use has::Has;
|
||||
use m_get::MGet;
|
||||
use m_set::MSet;
|
||||
use persist::Persist;
|
||||
use set::Set;
|
||||
use ttl::Ttl;
|
||||
|
||||
use crate::{Result, connection::Connection, database::Database, errors::AppError};
|
||||
use crate::{
|
||||
Result, buffer::ArchiveBuf as _, connection::Connection, database::Database, errors::AppError,
|
||||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Command {
|
||||
@@ -28,6 +34,9 @@ pub enum Command {
|
||||
Ttl(Ttl),
|
||||
Expire(Expire),
|
||||
Persist(Persist),
|
||||
|
||||
MSet(MSet),
|
||||
MGet(MGet),
|
||||
}
|
||||
|
||||
impl Command {
|
||||
@@ -40,21 +49,18 @@ impl Command {
|
||||
Command::Ttl(ttl) => ttl.execute(db, connection).await,
|
||||
Command::Expire(expire) => expire.execute(db, connection).await,
|
||||
Command::Persist(persist) => persist.execute(db, connection).await,
|
||||
|
||||
Command::MSet(m_set) => m_set.execute(db, connection).await,
|
||||
Command::MGet(m_get) => m_get.execute(db, connection).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse(bytes: &BytesMut) -> Result<(Self, u64)> {
|
||||
let mut buffer = Cursor::new(&bytes[..]);
|
||||
let mut buf = Cursor::new(&bytes[..]);
|
||||
|
||||
let name_length = buffer.try_get_u16()? as usize;
|
||||
let name = buf.try_get_short_string()?;
|
||||
|
||||
if buffer.remaining() < name_length {
|
||||
return Err(AppError::IncompleteCommandBuffer);
|
||||
}
|
||||
|
||||
let name = String::from_utf8(buffer.copy_to_bytes(name_length).to_vec())?;
|
||||
|
||||
Self::parse_inner(name, &mut buffer)
|
||||
Self::parse_inner(name, &mut buf)
|
||||
}
|
||||
|
||||
fn parse_inner(command_name: String, bytes: &mut Cursor<&[u8]>) -> Result<(Self, u64)> {
|
||||
@@ -66,6 +72,8 @@ impl Command {
|
||||
"ttl" => Self::Ttl(Ttl::parse(bytes)?),
|
||||
"expire" => Self::Expire(Expire::parse(bytes)?),
|
||||
"persist" => Self::Persist(Persist::parse(bytes)?),
|
||||
"mset" => Self::MSet(MSet::parse(bytes)?),
|
||||
"mget" => Self::MGet(MGet::parse(bytes)?),
|
||||
_ => return Err(AppError::UnknownCommand(command_name)),
|
||||
};
|
||||
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use bytes::{Buf as _, Bytes};
|
||||
use bytes::{Bytes, BytesMut};
|
||||
|
||||
use crate::{Result, connection::Connection, database::Database, errors::AppError};
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::{ArchiveBuf as _, ArchiveBufMut as _},
|
||||
connection::Connection,
|
||||
database::Database,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Persist {
|
||||
@@ -10,6 +15,10 @@ pub struct Persist {
|
||||
}
|
||||
|
||||
impl Persist {
|
||||
pub fn new(key: String) -> Self {
|
||||
Self { key }
|
||||
}
|
||||
|
||||
pub async fn execute(self, db: &Database, connection: &mut Connection) -> Result<()> {
|
||||
let value = db.persist(&self.key).await?;
|
||||
|
||||
@@ -20,15 +29,18 @@ impl Persist {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
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())?;
|
||||
pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let key = buf.try_get_short_string()?;
|
||||
|
||||
Ok(Self { key })
|
||||
}
|
||||
|
||||
pub fn put(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string("persist")?;
|
||||
self.put_without_cmd_name(buf)
|
||||
}
|
||||
|
||||
pub fn put_without_cmd_name(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string(&self.key)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,16 +1,29 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use crate::{Result, connection::Connection, database::Database, errors::AppError};
|
||||
use bytes::{Buf as _, Bytes};
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::{ArchiveBuf as _, ArchiveBufMut as _},
|
||||
connection::Connection,
|
||||
database::Database,
|
||||
};
|
||||
use bytes::{Buf as _, BufMut as _, Bytes, BytesMut};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Set {
|
||||
key: String,
|
||||
data: Box<[u8]>,
|
||||
expiration: Option<u64>,
|
||||
pub(super) key: String,
|
||||
pub(super) data: Box<[u8]>,
|
||||
pub(super) expiration: Option<u64>,
|
||||
}
|
||||
|
||||
impl Set {
|
||||
pub fn new(key: String, data: Box<[u8]>, expiration: Option<u64>) -> Self {
|
||||
Self {
|
||||
key,
|
||||
data,
|
||||
expiration,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute(self, db: &Database, connection: &mut Connection) -> Result<()> {
|
||||
db.set(self.key, self.data, self.expiration).await?;
|
||||
|
||||
@@ -19,28 +32,12 @@ impl Set {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let key_length = bytes.try_get_u16()? as usize;
|
||||
pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let key = buf.try_get_short_string()?;
|
||||
|
||||
if bytes.remaining() < key_length {
|
||||
return Err(AppError::IncompleteCommandBuffer);
|
||||
}
|
||||
let data = buf.try_get_bytes()?;
|
||||
|
||||
let key = String::from_utf8(bytes.copy_to_bytes(key_length).to_vec())?;
|
||||
|
||||
let value_length = bytes.try_get_u32()? as usize;
|
||||
|
||||
if bytes.remaining() < value_length {
|
||||
return Err(AppError::IncompleteCommandBuffer);
|
||||
}
|
||||
|
||||
let data = bytes.copy_to_bytes(value_length);
|
||||
|
||||
let expiration: Option<u64> = match bytes.try_get_u8()? {
|
||||
1 => Some(bytes.try_get_u64()?),
|
||||
0 => None,
|
||||
_ => return Err(AppError::UnexpectedCommandData),
|
||||
};
|
||||
let expiration = buf.try_get_option(Cursor::try_get_u64)?;
|
||||
|
||||
Ok(Self {
|
||||
key,
|
||||
@@ -48,4 +45,20 @@ impl Set {
|
||||
expiration,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn put(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string("set")?;
|
||||
|
||||
self.put_without_cmd_name(buf)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn put_without_cmd_name(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string(&self.key)?;
|
||||
buf.put_bytes_with_length(&self.data);
|
||||
buf.put_option(self.expiration, BytesMut::put_u64);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use bytes::{Buf as _, BufMut, Bytes, BytesMut};
|
||||
use bytes::{BufMut, BytesMut};
|
||||
|
||||
use crate::{Result, connection::Connection, database::Database, errors::AppError};
|
||||
use crate::{
|
||||
Result,
|
||||
buffer::{ArchiveBuf as _, ArchiveBufMut as _},
|
||||
connection::Connection,
|
||||
database::Database,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Ttl {
|
||||
@@ -10,33 +15,36 @@ pub struct Ttl {
|
||||
}
|
||||
|
||||
impl Ttl {
|
||||
pub fn new(key: String) -> Self {
|
||||
Self { key }
|
||||
}
|
||||
|
||||
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);
|
||||
buf.put_option(ttl, BytesMut::put_u64);
|
||||
|
||||
connection.write(buf.into()).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
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())?;
|
||||
pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
|
||||
let key = buf.try_get_short_string()?;
|
||||
|
||||
Ok(Self { key })
|
||||
}
|
||||
|
||||
pub fn put(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string("ttl")?;
|
||||
|
||||
self.put_without_cmd_name(buf)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn put_without_cmd_name(&self, buf: &mut BytesMut) -> Result<()> {
|
||||
buf.put_short_string(&self.key)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,8 +6,8 @@ pub enum AppError {
|
||||
Io(#[from] std::io::Error),
|
||||
#[error("A TryGetError occurred")]
|
||||
TryGet(#[from] bytes::TryGetError),
|
||||
#[error("The buffer is missing data for a complete command")]
|
||||
IncompleteCommandBuffer,
|
||||
#[error("The buffer is missing data")]
|
||||
IncompleteBuffer,
|
||||
#[error("A Utf8Error occurred")]
|
||||
FromUtf8(#[from] std::string::FromUtf8Error),
|
||||
#[error("The command {0} was not recognized")]
|
||||
@@ -18,6 +18,8 @@ pub enum AppError {
|
||||
NoResponse,
|
||||
#[error("Expected a different response for the executed command")]
|
||||
InvalidCommandResponse,
|
||||
#[error("The binary command data is not structured correctly")]
|
||||
UnexpectedCommandData,
|
||||
#[error("The binary data is not structured correctly")]
|
||||
UnexpectedData,
|
||||
#[error("Failed to convert integer")]
|
||||
TryFromInt(#[from] std::num::TryFromIntError),
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ use errors::AppError;
|
||||
#[cfg(test)]
|
||||
pub mod tests;
|
||||
|
||||
pub mod buffer;
|
||||
pub mod client;
|
||||
pub mod commands;
|
||||
pub mod config;
|
||||
|
||||
48
src/tests.rs
48
src/tests.rs
@@ -50,6 +50,54 @@ async fn expiration() -> Result<(), Box<dyn std::error::Error>> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn m_set_m_get() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let config = ServerConfig::builder()
|
||||
.host("127.0.0.1".into())
|
||||
.port(6172)
|
||||
.build();
|
||||
let mut server = Server::new(&config).await?;
|
||||
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel();
|
||||
let server_handle = tokio::spawn(async move { server.run(shutdown_rx).await });
|
||||
|
||||
let mut client = client("127.0.0.1:6172").await?;
|
||||
|
||||
client
|
||||
.m_set(
|
||||
vec!["key-0", "key-1", "key-2"],
|
||||
vec![b"value-0", b"value-1", b"value-2"],
|
||||
vec![None, Some(2), None],
|
||||
)
|
||||
.await?;
|
||||
|
||||
assert_eq!(
|
||||
client
|
||||
.m_get(vec!["key-0".into(), "key-1".into(), "key-2".into()])
|
||||
.await?,
|
||||
vec![
|
||||
Some("value-0".into()),
|
||||
Some("value-1".into()),
|
||||
Some("value-2".into())
|
||||
]
|
||||
);
|
||||
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
|
||||
assert_eq!(
|
||||
client
|
||||
.m_get(vec!["key-0".into(), "key-1".into(), "key-2".into()])
|
||||
.await?,
|
||||
vec![Some("value-0".into()), None, Some("value-2".into())]
|
||||
);
|
||||
|
||||
shutdown_tx.send(()).unwrap();
|
||||
|
||||
server_handle.await??;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn client<A: tokio::net::ToSocketAddrs>(
|
||||
addr: A,
|
||||
) -> Result<Client, Box<dyn std::error::Error>> {
|
||||
|
||||
Reference in New Issue
Block a user