use super::msg::*;
use std::collections::HashMap;
use std::future::Future;
pub use std::ops::Deref;
use std::pin::Pin;
use std::sync::atomic::{AtomicU32, Ordering};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
#[cfg(unix)]
use tokio::net::UnixStream;
use tokio::sync::{mpsc, oneshot};
type Return<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
pub trait Call {
fn call<'a>(&'a self, id: u32, data: Bytes) -> Return<'a, Result<Msg, Error>>;
}
pub trait SubcribeCallback<T> {
fn callback(&mut self, data: T) -> bool;
}
pub struct MyStream {
tx: mpsc::Sender<(u32, oneshot::Sender<Result<Msg, Error>>, Bytes)>,
}
impl MyStream {
fn new<Stream: AsyncReadExt + AsyncWriteExt + std::marker::Unpin + Send + 'static>(mut stream: Stream) -> Self {
let (tx, mut rx) = mpsc::channel::<(u32, oneshot::Sender<Result<Msg, Error>>, Bytes)>(1);
tokio::spawn(async move {
let mut header = [0u8; RPC_HEADER_LEN];
let mut callers = HashMap::<u32, oneshot::Sender<Result<Msg, Error>>>::new();
let error = loop {
tokio::select! {
Some((id, otx, data)) = rx.recv() => {
callers.insert(id, otx);
if let Err(err) = stream.write_all(&data[..]).await {
break Error::from(err);
}
}
ret = stream.read_exact(&mut header[..]) => {
match ret {
Ok(0) => break Error::new("对端已关闭读取数据长度为0"),
Ok(_) => {
if let Ok(mut msg) = Msg::decode(&header[..]) {
match msg.mode() {
Mode::Respond | Mode::Publish => {
if let Some(buf) = msg.body() {
if let Err(err) = stream.read_exact(buf).await {
break Error::from(err);
}
}
}
_ => (),
}
if let Some(otx) = callers.remove(&msg.id()) {
let _ = otx.send(Ok(msg));
}
}
}
Err(err) => break Error::from(err),
}
}
}
};
callers.into_iter().for_each(|(_, otx)| {
let _ = otx.send(Err(error.clone()));
});
});
Self { tx }
}
}
impl Call for MyStream {
fn call<'a>(&'a self, id: u32, data: Bytes) -> Return<'a, Result<Msg, Error>> {
Box::pin(async move {
let (tx, rx) = oneshot::channel();
let _ = self.tx.send((id, tx, data)).await;
rx.await.unwrap()
})
}
}
pub struct Client<Stream: Call + Send> {
stream: Stream,
id: AtomicU32,
}
const HEARTBEAT: &'static str = "heartbeat";
impl<Stream: Call + Send> Client<Stream> {
pub fn new(stream: Stream) -> Self {
Self { stream, id: AtomicU32::new(0) }
}
fn encode_heartbeat(&self) -> (u32, Bytes) {
let id = self.id.fetch_add(1, Ordering::SeqCst);
let msg = Msg::new(id, HEARTBEAT);
let data = msg.encode_without_body(Mode::HeartBeat);
(id, data)
}
fn encode_without_arg(&self, name: &str) -> (u32, Bytes) {
let id = self.id.fetch_add(1, Ordering::SeqCst);
let msg = Msg::new(id, name);
let data = msg.encode_without_body(Mode::Request);
(id, data)
}
fn encode_with_arg<Args>(&self, name: &str, args: Args) -> (u32, Bytes)
where
Args: serde::ser::Serialize,
{
let id = self.id.fetch_add(1, Ordering::SeqCst);
let msg = Msg::new(id, name);
let data = msg.encode(Mode::Request, &args);
(id, data)
}
fn decode_heartbeat(msg: Msg) -> Result<(), Error> {
if msg.name() != HEARTBEAT.as_bytes() {
return Err(Error::new("心跳回复函数名称不匹配"));
}
if !msg.headeronly() {
return Err(Error::new("心跳不应当返回消息体"));
}
match msg.mode() {
Mode::HeartBeat => Ok(()),
_ => Err(Error::new("返回消息模式不正确")),
}
}
fn decode_without_ret(msg: Msg, name: &str) -> Result<(), Error> {
if msg.name() != name.as_bytes() {
return Err(Error::new("回复函数名称不匹配"));
}
if !msg.headeronly() {
return Err(Error::new("不应当返回消息体"));
}
match msg.mode() {
Mode::Respond => Ok(()),
Mode::NotFound => Err(Error::new("没有找到相应的函数")),
Mode::NotMatch => Err(Error::new("函数参数不匹配")),
Mode::NoAccess => Err(Error::new("没有权限")),
_ => Err(Error::new("返回消息模式不正确")),
}
}
fn decode_with_ret<Ret>(msg: Msg, name: &str) -> Result<Ret, Error>
where
Ret: for<'a> serde::de::Deserialize<'a>,
{
if msg.name() != name.as_bytes() {
return Err(Error::new("回复函数名称不匹配"));
}
match msg.mode() {
Mode::Respond => {
if msg.headeronly() {
Err(Error::new("没有消息体"))
} else {
msg.parse()
}
}
Mode::NotFound => Err(Error::new("没有找到相应的函数")),
Mode::NotMatch => Err(Error::new("函数参数不匹配")),
Mode::NoAccess => Err(Error::new("没有权限")),
_ => Err(Error::new("返回消息模式不正确")),
}
}
#[inline]
pub async fn heartbeat(&self) -> Result<(), Error> {
let (id, data) = self.encode_heartbeat();
let msg = self.stream.call(id, data).await?;
Self::decode_heartbeat(msg)
}
#[inline]
pub async fn call_without_arg_ret(&self, name: &str) -> Result<(), Error> {
let (id, data) = self.encode_without_arg(name);
let msg = self.stream.call(id, data).await?;
Self::decode_without_ret(msg, name)
}
#[inline]
pub async fn call_with_ret<Ret>(&self, name: &str) -> Result<Ret, Error>
where
Ret: for<'a> serde::de::Deserialize<'a>,
{
let (id, data) = self.encode_without_arg(name);
let msg = self.stream.call(id, data).await?;
Self::decode_with_ret(msg, name)
}
#[inline]
pub async fn call_with_arg<Args>(&self, name: &str, args: Args) -> Result<(), Error>
where
Args: serde::ser::Serialize,
{
let (id, data) = self.encode_with_arg(name, args);
let msg = self.stream.call(id, data).await?;
Self::decode_without_ret(msg, name)
}
#[inline]
pub async fn call_with_arg_ret<Args, Ret>(&self, name: &str, args: Args) -> Result<Ret, Error>
where
Args: serde::ser::Serialize,
Ret: for<'a> serde::de::Deserialize<'a>,
{
let (id, data) = self.encode_with_arg(name, args);
let msg = self.stream.call(id, data).await?;
Self::decode_with_ret(msg, name)
}
pub async fn subcribe_with_lambda<Ret, F>(&self, topic: &str, mut f: F) -> Result<(), Error>
where
Ret: for<'a> serde::de::Deserialize<'a>,
F: FnMut(Ret),
{
loop {
let id = self.id.fetch_add(1, Ordering::SeqCst);
let msg = Msg::new(id, topic);
let data = msg.encode_without_body(Mode::Subcribe);
let msg = self.stream.call(id, data).await?;
if msg.name() != topic.as_bytes() {
break Err(Error::new("订阅主题名称不匹配"));
}
match msg.mode() {
Mode::Publish => {
if msg.headeronly() {
break Err(Error::new("没有订阅到消息体"));
} else {
f(msg.parse()?);
continue;
}
}
Mode::NotFound => break Err(Error::new("没有找到相应的函数")),
Mode::NotMatch => break Err(Error::new("函数参数不匹配")),
Mode::NoAccess => break Err(Error::new("没有权限")),
_ => break Err(Error::new("返回消息模式不正确")),
}
}
}
pub async fn subcribe_with_trait<Ret, T>(&self, topic: &str, t: &mut T) -> Result<(), Error>
where
Ret: for<'a> serde::de::Deserialize<'a>,
T: SubcribeCallback<Ret>,
{
loop {
let id = self.id.fetch_add(1, Ordering::SeqCst);
let msg = Msg::new(id, topic);
let data = msg.encode_without_body(Mode::Subcribe);
let msg = self.stream.call(id, data).await?;
if msg.name() != topic.as_bytes() {
break Err(Error::new("订阅主题名称不匹配"));
}
match msg.mode() {
Mode::Publish => {
if msg.headeronly() {
break Err(Error::new("没有订阅到消息体"));
} else {
if t.callback(msg.parse()?) {
continue; } else {
break Ok(()); }
}
}
Mode::NotFound => break Err(Error::new("没有找到相应的函数")),
Mode::NotMatch => break Err(Error::new("函数参数不匹配")),
Mode::NoAccess => break Err(Error::new("没有权限")),
_ => break Err(Error::new("返回消息模式不正确")),
}
}
}
}
pub type TCPClient = Client<MyStream>;
#[cfg(unix)]
pub type UnixClient = Client<MyStream>;
#[inline]
pub async fn new_tcp_client(addr: &str) -> std::io::Result<TCPClient> {
Ok(Client::new(MyStream::new(TcpStream::connect(addr).await?)))
}
#[inline]
#[cfg(unix)]
pub async fn new_unix_client(path: &str) -> std::io::Result<UnixClient> {
Ok(Client::new(MyStream::new(UnixStream::connect(path).await?)))
}
#[macro_export]
macro_rules! call {
(@call $connect:ident, $addr:expr) => {
async move {
match $crate::$connect($addr).await {
Ok(client) => client.heartbeat().await,
Err(e) => Err($crate::msg::Error::from(e)),
}
}
};
(@call $connect:ident, $addr:expr, $func:ident()) => {
async move {
match $crate::$connect($addr).await {
Ok(client) => client.call_without_arg_ret(stringify!($func)).await,
Err(e) => Err($crate::msg::Error::from(e)),
}
}
};
(@call $connect:ident, $addr:expr, $func:ident() -> $ret:ty) => {
async move {
match $crate::$connect($addr).await {
Ok(client) => {
let result: $ret = client.call_with_ret(stringify!($func)).await?;
Ok(result)
}
Err(e) => Err($crate::msg::Error::from(e)),
}
}
};
(@call $connect:ident, $addr:expr, $func:ident($($arg:expr),+)) => {
async move {
match $crate::$connect($addr).await {
Ok(client) => client.call_with_arg(stringify!($func), ($($arg,)+)).await,
Err(e) => Err($crate::msg::Error::from(e)),
}
}
};
(@call $connect:ident, $addr:expr, $func:ident($($arg:expr),+) -> $ret:ty) => {
async move {
match $crate::$connect($addr).await {
Ok(client) => {
let result: $ret = client.call_with_arg_ret(stringify!($func), ($($arg,)+)).await?;
Ok(result)
}
Err(e) => Err($crate::msg::Error::from(e)),
}
}
};
(tcp, $($var:tt)+) => {
async move {
call!(@call new_tcp_client, $($var)+).await
}
};
(unix, $($var:tt)+) => {
async move {
call!(@call new_unix_client, $($var)+).await
}
};
}
#[macro_export]
macro_rules! subcribe {
(@sub $connect:ident, $addr:expr, $topic:expr, |$($arg:ident:$argType:ty),+|$body:block) => {
async move {
match $crate::$connect($addr).await {
Ok(client) => client.subcribe_with_lambda($topic, |($($arg,)+):($($argType,)+)|$body).await,
Err(e) => Err($crate::msg::Error::from(e)),
}
}
};
(@sub $connect:ident, $addr:expr, $topic:expr, $var:expr) => {
async move {
match $crate::$connect($addr).await {
Ok(client) => client.subcribe_with_trait($topic, $var).await,
Err(e) => Err($crate::msg::Error::from(e)),
}
}
};
(tcp, $($var:tt)+) => {
async move {
subcribe!(@sub new_tcp_client, $($var)+).await
}
};
(unix, $($var:tt)+) => {
async move {
subcribe!(@sub new_unix_client, $($var)+).await
}
};
}
#[macro_export]
macro_rules! define_new_type {
(@method fn $name:ident$(<$generic:tt>)?()) => {
pub async fn $name$(<$generic>)?(&self) -> Result<(), $crate::msg::Error> {
self.0.call_without_arg_ret(stringify!($name)).await
}
};
(@method fn $name:ident$(<$generic:tt>)?($($arg:ident:$argType:ty,)+)) => {
pub async fn $name$(<$generic>)?(&self, $($arg:$argType,)+) -> Result<(), $crate::msg::Error>
where
$($argType: serde::ser::Serialize,)+
{
self.0.call_with_arg(stringify!($name), ($($arg,)+)).await
}
};
(@method fn $name:ident$(<$generic:tt>)?()->$ret:ty) => {
pub async fn $name$(<$generic>)?(&self) -> Result<$ret, $crate::msg::Error>
where
$ret: for<'a> serde::de::Deserialize<'a>,
{
self.0.call_with_ret(stringify!($name)).await
}
};
(@method fn $name:ident$(<$generic:tt>)?($($arg:ident:$argType:ty,)+)->$ret:ty) => {
pub async fn $name$(<$generic>)?(&self, $($arg:$argType,)+) -> Result<$ret, $crate::msg::Error>
where
$($argType: serde::ser::Serialize,)+
$ret: for<'a> serde::de::Deserialize<'a>,
{
self.0.call_with_arg_ret(stringify!($name), ($($arg,)+)).await
}
};
(@method sub $name:ident($topic:ident:$topicType:ty, $f:ident:$(impl)? FnMut($ArgType:ty))) => {
pub async fn $name<F>(&self, $topic:$topicType, $f:F) -> Result<(), $crate::msg::Error>
where
F: FnMut($ArgType),
$topicType: Deref<Target = str>,
$ArgType: for<'a> serde::de::Deserialize<'a>,
{
self.0.subcribe_with_lambda(&$topic, $f).await
}
};
(@method sub $name:ident($topic:ident:$topicType:ty, $var:ident:&mut $ArgType:ty)) => {
pub async fn $name<Ret>(&self, $topic:&str, $var:&mut $ArgType) -> Result<(), $crate::msg::Error>
where
Ret: for<'a> serde::de::Deserialize<'a>,
$topicType: Deref<Target = str>,
$ArgType: SubcribeCallback<Ret>,
{
self.0.subcribe_with_trait(&$topic, $var).await
}
};
($f:ident, $t:ident, $StructName:ident, $(fn $name:ident$(<$generic:tt>)?($($arg:ident:$argType:ty),*)$(->$ret:ty)?),+) => {
struct $StructName($t);
impl $StructName {
pub async fn new(path:&str) -> std::io::Result<Self> {
Ok(Self($f(path).await?))
}
pub async fn heartbeat(&self) -> Result<(), $crate::msg::Error> {
self.0.heartbeat().await
}
$(define_new_type!(@method fn $name$(<$generic>)?($($arg:$argType,)*)$(->$ret)?);)+
}
};
($f:ident, $t:ident, $StructName:ident, $(sub $name:ident($topic:ident:$topicType:ty, $arg:ident:$($argType:tt)+)),+) => {
struct $StructName($t);
impl $StructName {
pub async fn new(path:&str) -> std::io::Result<Self> {
Ok(Self($f(path).await?))
}
pub async fn heartbeat(&self) -> Result<(), $crate::msg::Error> {
self.0.heartbeat().await
}
$(define_new_type!(@method sub $name($topic:$topicType,$arg:$($argType)+));)+
}
};
($f:ident, $t:ident, $StructName:ident, $(fn $name:ident($($arg:ident:$argType:ty),*)$(->$ret:ty)?),+, $(sub $name2:ident($topic:ident:$topicType:ty, $arg2:ident:$($argType2:tt)+)),+) => {
struct $StructName($t);
impl $StructName {
pub async fn new(path:&str) -> std::io::Result<Self> {
Ok(Self($f(path).await?))
}
pub async fn heartbeat(&self) -> Result<(), $crate::msg::Error> {
self.0.heartbeat().await
}
$(define_new_type!(@method fn $name($($arg:$argType,)*)$(->$ret)?);)+
$(define_new_type!(@method sub $name2($topic:$topicType,$arg2:$($argType2)+));)+
}
};
}
#[macro_export]
macro_rules! define {
(tcp, $($var:tt)+) => {
define_new_type!(new_tcp_client, TCPClient, $($var)+);
};
(unix, $($var:tt)+) => {
define_new_type!(new_unix_client, UnixClient, $($var)+);
};
}