use std::ops::Bound;
use std::path::Path;
use compio::io::{
AsyncRead, AsyncWrite,
framed::{Framed, frame::LengthDelimited},
};
use compio::net::{TcpStream, ToSocketAddrsAsync, UnixStream};
use futures_util::{SinkExt, StreamExt};
use crate::codec::ClientCodec;
use crate::error::{Result, RpcError};
use crate::protocol::*;
pub struct RpcClient<S> {
hashname: u64,
framed: Framed<S, S, ClientCodec, LengthDelimited, Request, Response>,
}
fn hash_name(name: &str) -> u64 {
xxhash_rust::xxh3::xxh3_64(name.as_bytes())
}
impl RpcClient<TcpStream> {
pub async fn connect_tcp(addr: impl ToSocketAddrsAsync, hashname: u64) -> Result<Self> {
let stream = TcpStream::connect(addr).await?;
let (reader, writer) = stream.into_split();
Ok(Self::new(reader, writer, hashname))
}
pub async fn connect_tcp_by_name(addr: impl ToSocketAddrsAsync, name: &str) -> Result<Self> {
Self::connect_tcp(addr, hash_name(name)).await
}
}
impl RpcClient<UnixStream> {
pub async fn connect_uds(path: impl AsRef<Path>, hashname: u64) -> Result<Self> {
let stream = UnixStream::connect(path).await?;
let (reader, writer) = stream.into_split();
Ok(Self::new(reader, writer, hashname))
}
pub async fn connect_uds_by_name(path: impl AsRef<Path>, name: &str) -> Result<Self> {
Self::connect_uds(path, hash_name(name)).await
}
}
impl<S: AsyncRead + AsyncWrite + Unpin + 'static> RpcClient<S> {
fn new(reader: S, writer: S, hashname: u64) -> Self {
let framed = Framed::new::<Request, Response>(ClientCodec, LengthDelimited::new())
.with_reader(reader)
.with_writer(writer);
Self { hashname, framed }
}
async fn call(&mut self, op: OpCode, payload: RequestPayload) -> Result<Vec<u8>> {
let request = Request {
op,
hashname: self.hashname,
payload,
};
self.framed.send(request).await?;
let response = self.framed.next().await.ok_or(RpcError::UnexpectedEof)??;
match response {
Response::Ok(data) => Ok(data),
Response::Err { code, message } => Err(RpcError::Server { code, message }),
}
}
pub async fn get(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>> {
let data = self
.call(OpCode::Get, RequestPayload::Key(key.to_vec()))
.await?;
decode_optional_data(&data)
}
pub async fn contains(&mut self, key: &[u8]) -> Result<bool> {
let data = self
.call(OpCode::Contains, RequestPayload::Key(key.to_vec()))
.await?;
decode_bool(&data)
}
pub async fn first(&mut self) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
let data = self.call(OpCode::First, RequestPayload::Empty).await?;
decode_optional_kv(&data)
}
pub async fn last(&mut self) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
let data = self.call(OpCode::Last, RequestPayload::Empty).await?;
decode_optional_kv(&data)
}
pub async fn range(
&mut self,
start: Bound<Vec<u8>>,
end: Bound<Vec<u8>>,
limit: u32,
) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
let data = self
.call(OpCode::Range, RequestPayload::Range { start, end, limit })
.await?;
decode_key_values(&data)
}
pub async fn range_keys(
&mut self,
start: Bound<Vec<u8>>,
end: Bound<Vec<u8>>,
limit: u32,
) -> Result<Vec<Vec<u8>>> {
let data = self
.call(
OpCode::RangeKeys,
RequestPayload::Range { start, end, limit },
)
.await?;
decode_keys(&data)
}
pub async fn upsert(
&mut self,
key: UpsertKey,
flag: Option<bool>,
value: Vec<u8>,
) -> Result<Vec<u8>> {
let data = self
.call(OpCode::Upsert, RequestPayload::Upsert { key, flag, value })
.await?;
decode_key(&data)
}
pub async fn remove(&mut self, key: &[u8], soft: bool) -> Result<()> {
self.call(
OpCode::Remove,
RequestPayload::Remove {
key: key.to_vec(),
soft,
},
)
.await?;
Ok(())
}
pub async fn take(&mut self, key: &[u8], soft: bool) -> Result<Option<Vec<u8>>> {
let data = self
.call(
OpCode::Take,
RequestPayload::Take {
key: key.to_vec(),
soft,
},
)
.await?;
decode_optional_data(&data)
}
pub async fn count(&mut self, exact: bool) -> Result<u64> {
let data = self
.call(OpCode::Count, RequestPayload::Count { exact })
.await?;
decode_count(&data)
}
pub async fn entry_len(&mut self, key: &[u8]) -> Result<Option<u32>> {
let data = self
.call(
OpCode::EntryLen,
RequestPayload::EntryLen { key: key.to_vec() },
)
.await?;
decode_optional_len(&data)
}
pub async fn apply_batch(&mut self, items: Vec<(Vec<u8>, Option<Vec<u8>>)>) -> Result<()> {
self.call(OpCode::ApplyBatch, RequestPayload::Batch(items))
.await?;
Ok(())
}
#[cfg(feature = "schema")]
pub async fn get_schema(&mut self) -> Result<crate::protocol::SchemaResponse> {
let data = self.call(OpCode::GetSchema, RequestPayload::Empty).await?;
serde_json::from_slice(&data).map_err(|e| RpcError::Protocol(e.to_string()))
}
}
fn decode_optional_data(buf: &[u8]) -> Result<Option<Vec<u8>>> {
let mut pos = 0;
let flag = read_u8(buf, &mut pos)?;
if flag == 0 {
return Ok(None);
}
Ok(Some(read_bytes(buf, &mut pos)?))
}
fn decode_optional_len(buf: &[u8]) -> Result<Option<u32>> {
let mut pos = 0;
let flag = read_u8(buf, &mut pos)?;
if flag == 0 {
return Ok(None);
}
Ok(Some(read_u32_be(buf, &mut pos)?))
}
fn decode_optional_kv(buf: &[u8]) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
let mut pos = 0;
let flag = read_u8(buf, &mut pos)?;
if flag == 0 {
return Ok(None);
}
let key = read_bytes(buf, &mut pos)?;
let val = read_bytes(buf, &mut pos)?;
Ok(Some((key, val)))
}
fn decode_key_values(buf: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
let mut pos = 0;
let count = read_u32_be(buf, &mut pos)? as usize;
const MAX_PREALLOC: usize = 64 * 1024;
let mut pairs = Vec::with_capacity(count.min(MAX_PREALLOC));
for _ in 0..count {
let key = read_bytes(buf, &mut pos)?;
let val = read_bytes(buf, &mut pos)?;
pairs.push((key, val));
}
Ok(pairs)
}
fn decode_keys(buf: &[u8]) -> Result<Vec<Vec<u8>>> {
let mut pos = 0;
let count = read_u32_be(buf, &mut pos)? as usize;
const MAX_PREALLOC: usize = 64 * 1024;
let mut keys = Vec::with_capacity(count.min(MAX_PREALLOC));
for _ in 0..count {
keys.push(read_bytes(buf, &mut pos)?);
}
Ok(keys)
}
fn decode_count(buf: &[u8]) -> Result<u64> {
let mut pos = 0;
read_u64_be(buf, &mut pos)
}
fn decode_key(buf: &[u8]) -> Result<Vec<u8>> {
let mut pos = 0;
read_bytes(buf, &mut pos)
}
pub(crate) fn decode_bool(buf: &[u8]) -> Result<bool> {
let mut pos = 0;
let v = read_u8(buf, &mut pos)?;
Ok(v != 0)
}