armour_rpc/

client.rs

use std::ops::Bound;
use std::path::Path;

use compio::io::{
    AsyncRead, AsyncWrite,
    framed::{Framed, frame::LengthDelimited},
};
use compio::net::{OwnedReadHalf, OwnedWriteHalf, TcpStream, ToSocketAddrsAsync, UnixStream};
use futures_util::{SinkExt, StreamExt};

use crate::codec::ClientCodec;
use crate::error::{Result, RpcError};
use crate::protocol::*;

pub struct RpcClient<R, W> {
    hashname: u64,
    framed: Framed<R, W, ClientCodec, LengthDelimited, Request, Response>,
}

fn hash_name(name: &str) -> u64 {
    xxhash_rust::xxh3::xxh3_64(name.as_bytes())
}

impl RpcClient<OwnedReadHalf<TcpStream>, OwnedWriteHalf<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<OwnedReadHalf<UnixStream>, OwnedWriteHalf<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<R: AsyncRead + Unpin + 'static, W: AsyncWrite + Unpin + 'static> RpcClient<R, W> {
    fn new(reader: R, writer: W, hashname: u64) -> Self {
        let framed = Framed::new::<Request, Response>(ClientCodec, LengthDelimited::new())
            .with_reader(reader)
            .with_writer(writer);
        Self { hashname, framed }
    }

    /// Send a request and return the raw Ok payload bytes, or an error.
    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<Option<u32>> {
        let data = self
            .call(OpCode::Contains, RequestPayload::Key(key.to_vec()))
            .await?;
        decode_optional_len(&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>>,
    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let data = self
            .call(OpCode::Range, RequestPayload::Range { start, end })
            .await?;
        decode_key_values(&data)
    }

    pub async fn range_keys(
        &mut self,
        start: Bound<Vec<u8>>,
        end: Bound<Vec<u8>>,
    ) -> Result<Vec<Vec<u8>>> {
        let data = self
            .call(OpCode::RangeKeys, RequestPayload::Range { start, end })
            .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 apply_batch(&mut self, items: Vec<(Vec<u8>, Option<Vec<u8>>)>) -> Result<()> {
        self.call(OpCode::ApplyBatch, RequestPayload::Batch(items))
            .await?;
        Ok(())
    }
}

// --- Response payload decoders ---

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;
    let mut pairs = Vec::with_capacity(count);
    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;
    let mut keys = Vec::with_capacity(count);
    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)
}