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::*;
#[derive(Debug, Clone)]
pub struct CollectionInfo {
pub name: String,
pub partition_name: String,
pub hashname: u64,
pub typ_hash: u64,
pub version: u16,
pub count: u64,
}
pub async fn list_collections_tcp(addr: impl ToSocketAddrsAsync) -> Result<Vec<CollectionInfo>> {
let stream = TcpStream::connect(addr).await?;
let (reader, writer) = stream.into_split();
list_collections_call(reader, writer).await
}
pub async fn list_collections_uds(path: impl AsRef<Path>) -> Result<Vec<CollectionInfo>> {
let stream = UnixStream::connect(path).await?;
let (reader, writer) = stream.into_split();
list_collections_call(reader, writer).await
}
async fn list_collections_call<R, W>(reader: R, writer: W) -> Result<Vec<CollectionInfo>>
where
R: AsyncRead + Unpin + 'static,
W: AsyncWrite + Unpin + 'static,
{
let mut framed = Framed::new::<Request, Response>(ClientCodec, LengthDelimited::new())
.with_reader(reader)
.with_writer(writer);
let request = Request {
op: OpCode::ListCollections,
hashname: 0,
payload: RequestPayload::ListCollections,
};
framed.send(request).await?;
let response = framed.next().await.ok_or(RpcError::UnexpectedEof)??;
match response {
Response::Ok(data) => decode_collections(&data),
Response::Err { code, message } => Err(RpcError::Server { code, message }),
}
}
fn decode_collections(buf: &[u8]) -> Result<Vec<CollectionInfo>> {
let mut pos = 0;
let count = read_u32_be(buf, &mut pos)? as usize;
const MAX_PREALLOC: usize = 64 * 1024;
let mut collections = Vec::with_capacity(count.min(MAX_PREALLOC));
for _ in 0..count {
let name_bytes = read_bytes(buf, &mut pos)?;
let name = String::from_utf8(name_bytes).map_err(|e| RpcError::Protocol(e.to_string()))?;
let partition_bytes = read_bytes(buf, &mut pos)?;
let partition_name =
String::from_utf8(partition_bytes).map_err(|e| RpcError::Protocol(e.to_string()))?;
let hashname = read_u64_be(buf, &mut pos)?;
let typ_hash = read_u64_be(buf, &mut pos)?;
let version = read_u16_be(buf, &mut pos)?;
let count = read_u64_be(buf, &mut pos)?;
collections.push(CollectionInfo {
name,
partition_name,
hashname,
typ_hash,
version,
count,
});
}
Ok(collections)
}