armour_rpc/
client.rs

1use std::ops::Bound;
2use std::path::Path;
3
4use compio::io::{
5    AsyncRead, AsyncWrite,
6    framed::{Framed, frame::LengthDelimited},
7};
8use compio::net::{TcpStream, ToSocketAddrsAsync, UnixStream};
9use futures_util::{SinkExt, StreamExt};
10
11use crate::codec::ClientCodec;
12use crate::error::{Result, RpcError};
13use crate::protocol::*;
14
15pub struct RpcClient<S> {
16    hashname: u64,
17    framed: Framed<S, S, ClientCodec, LengthDelimited, Request, Response>,
18}
19
20fn hash_name(name: &str) -> u64 {
21    xxhash_rust::xxh3::xxh3_64(name.as_bytes())
22}
23
24impl RpcClient<TcpStream> {
25    pub async fn connect_tcp(addr: impl ToSocketAddrsAsync, hashname: u64) -> Result<Self> {
26        let stream = TcpStream::connect(addr).await?;
27        let (reader, writer) = stream.into_split();
28        Ok(Self::new(reader, writer, hashname))
29    }
30
31    pub async fn connect_tcp_by_name(addr: impl ToSocketAddrsAsync, name: &str) -> Result<Self> {
32        Self::connect_tcp(addr, hash_name(name)).await
33    }
34}
35
36impl RpcClient<UnixStream> {
37    pub async fn connect_uds(path: impl AsRef<Path>, hashname: u64) -> Result<Self> {
38        let stream = UnixStream::connect(path).await?;
39        let (reader, writer) = stream.into_split();
40        Ok(Self::new(reader, writer, hashname))
41    }
42
43    pub async fn connect_uds_by_name(path: impl AsRef<Path>, name: &str) -> Result<Self> {
44        Self::connect_uds(path, hash_name(name)).await
45    }
46}
47
48impl<S: AsyncRead + AsyncWrite + Unpin + 'static> RpcClient<S> {
49    fn new(reader: S, writer: S, hashname: u64) -> Self {
50        let framed = Framed::new::<Request, Response>(ClientCodec, LengthDelimited::new())
51            .with_reader(reader)
52            .with_writer(writer);
53        Self { hashname, framed }
54    }
55
56    /// Send a request and return the raw Ok payload bytes, or an error.
57    async fn call(&mut self, op: OpCode, payload: RequestPayload) -> Result<Vec<u8>> {
58        let request = Request {
59            op,
60            hashname: self.hashname,
61            payload,
62        };
63        self.framed.send(request).await?;
64        let response = self.framed.next().await.ok_or(RpcError::UnexpectedEof)??;
65        match response {
66            Response::Ok(data) => Ok(data),
67            Response::Err { code, message } => Err(RpcError::Server { code, message }),
68        }
69    }
70
71    pub async fn get(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>> {
72        let data = self
73            .call(OpCode::Get, RequestPayload::Key(key.to_vec()))
74            .await?;
75        decode_optional_data(&data)
76    }
77
78    pub async fn contains(&mut self, key: &[u8]) -> Result<bool> {
79        let data = self
80            .call(OpCode::Contains, RequestPayload::Key(key.to_vec()))
81            .await?;
82        decode_bool(&data)
83    }
84
85    pub async fn first(&mut self) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
86        let data = self.call(OpCode::First, RequestPayload::Empty).await?;
87        decode_optional_kv(&data)
88    }
89
90    pub async fn last(&mut self) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
91        let data = self.call(OpCode::Last, RequestPayload::Empty).await?;
92        decode_optional_kv(&data)
93    }
94
95    pub async fn range(
96        &mut self,
97        start: Bound<Vec<u8>>,
98        end: Bound<Vec<u8>>,
99        limit: u32,
100    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
101        let data = self
102            .call(OpCode::Range, RequestPayload::Range { start, end, limit })
103            .await?;
104        decode_key_values(&data)
105    }
106
107    pub async fn range_keys(
108        &mut self,
109        start: Bound<Vec<u8>>,
110        end: Bound<Vec<u8>>,
111        limit: u32,
112    ) -> Result<Vec<Vec<u8>>> {
113        let data = self
114            .call(
115                OpCode::RangeKeys,
116                RequestPayload::Range { start, end, limit },
117            )
118            .await?;
119        decode_keys(&data)
120    }
121
122    pub async fn upsert(
123        &mut self,
124        key: UpsertKey,
125        flag: Option<bool>,
126        value: Vec<u8>,
127    ) -> Result<Vec<u8>> {
128        let data = self
129            .call(OpCode::Upsert, RequestPayload::Upsert { key, flag, value })
130            .await?;
131        decode_key(&data)
132    }
133
134    pub async fn remove(&mut self, key: &[u8], soft: bool) -> Result<()> {
135        self.call(
136            OpCode::Remove,
137            RequestPayload::Remove {
138                key: key.to_vec(),
139                soft,
140            },
141        )
142        .await?;
143        Ok(())
144    }
145
146    pub async fn take(&mut self, key: &[u8], soft: bool) -> Result<Option<Vec<u8>>> {
147        let data = self
148            .call(
149                OpCode::Take,
150                RequestPayload::Take {
151                    key: key.to_vec(),
152                    soft,
153                },
154            )
155            .await?;
156        decode_optional_data(&data)
157    }
158
159    pub async fn count(&mut self, exact: bool) -> Result<u64> {
160        let data = self
161            .call(OpCode::Count, RequestPayload::Count { exact })
162            .await?;
163        decode_count(&data)
164    }
165
166    pub async fn entry_len(&mut self, key: &[u8]) -> Result<Option<u32>> {
167        let data = self
168            .call(
169                OpCode::EntryLen,
170                RequestPayload::EntryLen { key: key.to_vec() },
171            )
172            .await?;
173        decode_optional_len(&data)
174    }
175
176    pub async fn apply_batch(&mut self, items: Vec<(Vec<u8>, Option<Vec<u8>>)>) -> Result<()> {
177        self.call(OpCode::ApplyBatch, RequestPayload::Batch(items))
178            .await?;
179        Ok(())
180    }
181
182    /// Fetch the schema (`Typ` + `KeyScheme` + name/version) of the bound collection.
183    ///
184    /// The server returns a JSON-serialised [`SchemaResponse`](crate::protocol::SchemaResponse).
185    /// The client's `hashname` (set at connect time) selects the collection.
186    #[cfg(feature = "schema")]
187    pub async fn get_schema(&mut self) -> Result<crate::protocol::SchemaResponse> {
188        let data = self.call(OpCode::GetSchema, RequestPayload::Empty).await?;
189        serde_json::from_slice(&data).map_err(|e| RpcError::Protocol(e.to_string()))
190    }
191}
192
193// --- Response payload decoders ---
194
195fn decode_optional_data(buf: &[u8]) -> Result<Option<Vec<u8>>> {
196    let mut pos = 0;
197    let flag = read_u8(buf, &mut pos)?;
198    if flag == 0 {
199        return Ok(None);
200    }
201    Ok(Some(read_bytes(buf, &mut pos)?))
202}
203
204fn decode_optional_len(buf: &[u8]) -> Result<Option<u32>> {
205    let mut pos = 0;
206    let flag = read_u8(buf, &mut pos)?;
207    if flag == 0 {
208        return Ok(None);
209    }
210    Ok(Some(read_u32_be(buf, &mut pos)?))
211}
212
213fn decode_optional_kv(buf: &[u8]) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
214    let mut pos = 0;
215    let flag = read_u8(buf, &mut pos)?;
216    if flag == 0 {
217        return Ok(None);
218    }
219    let key = read_bytes(buf, &mut pos)?;
220    let val = read_bytes(buf, &mut pos)?;
221    Ok(Some((key, val)))
222}
223
224fn decode_key_values(buf: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
225    let mut pos = 0;
226    let count = read_u32_be(buf, &mut pos)? as usize;
227    const MAX_PREALLOC: usize = 64 * 1024;
228    let mut pairs = Vec::with_capacity(count.min(MAX_PREALLOC));
229    for _ in 0..count {
230        let key = read_bytes(buf, &mut pos)?;
231        let val = read_bytes(buf, &mut pos)?;
232        pairs.push((key, val));
233    }
234    Ok(pairs)
235}
236
237fn decode_keys(buf: &[u8]) -> Result<Vec<Vec<u8>>> {
238    let mut pos = 0;
239    let count = read_u32_be(buf, &mut pos)? as usize;
240    const MAX_PREALLOC: usize = 64 * 1024;
241    let mut keys = Vec::with_capacity(count.min(MAX_PREALLOC));
242    for _ in 0..count {
243        keys.push(read_bytes(buf, &mut pos)?);
244    }
245    Ok(keys)
246}
247
248fn decode_count(buf: &[u8]) -> Result<u64> {
249    let mut pos = 0;
250    read_u64_be(buf, &mut pos)
251}
252
253fn decode_key(buf: &[u8]) -> Result<Vec<u8>> {
254    let mut pos = 0;
255    read_bytes(buf, &mut pos)
256}
257
258pub(crate) fn decode_bool(buf: &[u8]) -> Result<bool> {
259    let mut pos = 0;
260    let v = read_u8(buf, &mut pos)?;
261    Ok(v != 0)
262}
armour_rpc/client.rs

armour_rpc/
client.rs
