krpc/

server.rs

extern crate rmp_serde as rmps;
use std::collections::HashMap;
use std::future::Future;

use super::msg::*;
#[cfg(unix)]
use std::path::Path;
use std::pin::Pin;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
#[cfg(unix)]
use tokio::net::UnixListener;
use tokio::sync::mpsc;

include!("publish.rs");
include!("session.rs");

pub type AsyncCallback = Pin<Box<dyn Future<Output = Bytes> + Send>>;
pub type Callback = Box<dyn Fn(Msg) -> AsyncCallback + Send + Sync>;

/// 用以声明回调函数,主要用于服务端绑定lambda函数(支持同步和异步)
///
/// 支持如下格式:
/// - 前缀`krpc::clone!(a,...)` `async` `move`都是可选的,如果lambda函数中使用的变量需要在函数体外被使用,需用`clone!`标注这些变量,例如`clone!(tx,var)`
/// - ||{...},无参数也无返回值
/// - |id:i32|{...},有参数无返回值
/// - ||->String {...},无参数有返回值
/// - |id:i32,s:String|->String {...},有参数且有返回值
#[macro_export]
macro_rules! callback {
	($($(krpc::)?clone!($($var:ident),*),)? $(async)? $(move)? || $b:block) => {
        //无参数也无返回值
        {
            $($(let $var = $var.clone();)*)?
            Box::new(move |msg: $crate::msg::Msg| -> $crate::AsyncCallback {
                $($(let $var = $var.clone();)*)?
                Box::pin(async move {
                    $b;
                    msg.encode_without_body($crate::msg::Mode::Respond)
                })
            })
        }
    };
    ($($(krpc::)?clone!($($var:ident),*),)? $(async)? $(move)? |$($a:ident :$t:ty),+| $b:block) => {
        //有参数无返回值
        {
            $($(let $var = $var.clone();)*)?
            Box::new(move |msg: $crate::msg::Msg| -> $crate::AsyncCallback {
                $($(let $var = $var.clone();)*)?
                Box::pin(async move {
                    let body = msg.parse::<($($t,)+)>();
                    match body {
                        Ok(($($a,)+)) => {
                            $b;
                            msg.encode_without_body($crate::msg::Mode::Respond)
                        }
                        Err(_) => msg.encode_without_body($crate::msg::Mode::NotMatch),
                    }
                })
            })
        }
    };
    ($($(krpc::)?clone!($($var:ident),*),)? $(async)? $(move)? || -> $r:ty $b:block) => {
        //无参数有返回值
        {
            $($(let $var = $var.clone();)*)?
            Box::new(move |msg: $crate::msg::Msg| -> $crate::AsyncCallback {
                $($(let $var = $var.clone();)*)?
                Box::pin(async move {
                    let ret:$r = $b;
                    msg.encode($crate::msg::Mode::Respond, &ret)
                })
            })
        }
    };
    ($($(krpc::)?clone!($($var:ident),*),)? $(async)? $(move)? |$($a:ident :$t:ty),+| -> $r:ty $b:block) => {
        //有参数且有返回值
        {
            $($(let $var = $var.clone();)*)?
            Box::new(move |msg: $crate::msg::Msg| -> $crate::AsyncCallback {
                $($(let $var = $var.clone();)*)?
                Box::pin(async move {
                    let body = msg.parse::<($($t,)+)>();
                    match body {
                        Ok(($($a,)+)) => {
                            let ret:$r = $b;
                            msg.encode($crate::msg::Mode::Respond, &ret)
                        }
                        Err(_) => msg.encode_without_body($crate::msg::Mode::NotMatch),
                    }
                })
            })
        }
    };
    (@with_return $($(krpc::)?clone!($($var:ident),*),)? $(async)? $(move)? || $b:block) => {
        //无参数也无返回值且会提前return
        {
            $($(let $var = $var.clone();)*)?
            Box::new(move |msg: $crate::msg::Msg| -> $crate::AsyncCallback {
                $($(let $var = $var.clone();)*)?
                Box::pin(async move {
                    let fu = async move {$b};
					fu.await;
                    msg.encode_without_body($crate::msg::Mode::Respond)
                })
            })
        }
    };
    (@with_return $($(krpc::)?clone!($($var:ident),*),)? $(async)? $(move)? |$($a:ident :$t:ty),+| $b:block) => {
        //有参数无返回值且会提前return
        {
            $($(let $var = $var.clone();)*)?
            Box::new(move |msg: $crate::msg::Msg| -> $crate::AsyncCallback {
                $($(let $var = $var.clone();)*)?
                Box::pin(async move {
                    let body = msg.parse::<($($t,)+)>();
                    match body {
                        Ok(($($a,)+)) => {
                            let fu = async move {$b};
							fu.await;
							//$b;
                            msg.encode_without_body($crate::msg::Mode::Respond)
                        }
                        Err(_) => msg.encode_without_body($crate::msg::Mode::NotMatch),
                    }
                })
            })
        }
    };
    (@with_return $($(krpc::)?clone!($($var:ident),*),)? $(async)? $(move)? || -> $r:ty $b:block) => {
        //无参数有返回值且会提前return
        {
            $($(let $var = $var.clone();)*)?
            Box::new(move |msg: $crate::msg::Msg| -> $crate::AsyncCallback {
                $($(let $var = $var.clone();)*)?
                Box::pin(async move {
					let fu = async move {$b};
                    let ret:$r = fu.await;
                    msg.encode($crate::msg::Mode::Respond, &ret)
                })
            })
        }
    };
    (@with_return $($(krpc::)?clone!($($var:ident),*),)? $(async)? $(move)? |$($a:ident :$t:ty),+| -> $r:ty $b:block) => {
        //有参数且有返回值且会提前return
        {
            $($(let $var = $var.clone();)*)?
            Box::new(move |msg: $crate::msg::Msg| -> $crate::AsyncCallback {
                $($(let $var = $var.clone();)*)?
                Box::pin(async move {
                    let body = msg.parse::<($($t,)+)>();
                    match body {
                        Ok(($($a,)+)) => {
							let fu = async move {$b};
		                    let ret:$r = fu.await;
                            msg.encode($crate::msg::Mode::Respond, &ret)
                        }
                        Err(_) => msg.encode_without_body($crate::msg::Mode::NotMatch),
                    }
                })
            })
        }
    };
}

/// RPC服务端
pub struct Server {
	funcs: HashMap<&'static [u8], Callback>,
	sub_tx: mpsc::Sender<Content>,
	sub_rx: mpsc::Receiver<Content>,
	call_tx: Option<mpsc::Sender<Pair>>,
	call_rx: Option<mpsc::Receiver<Pair>>,
}

/// 用以监听Server
macro_rules! listen {
	($self:ident, $listener:ident, $onSubcribe:ident) => {
		if let Some(mut rx) = $self.call_rx.take() {
			//用于RPC的调用请求并回复
			let funcs = $self.funcs;
			tokio::spawn(async move {
				loop {
					let Some(Pair { msg, tx }) = rx.recv().await else {break;};
					if let Some(f) = funcs.get(msg.name()) {
						let _ = tx.send(f(msg).await).await;
					} else {
						let _ = tx.send(msg.encode_without_body(Mode::NotFound)).await;
					}
				}
			});
		}
		//用于接收RPC的订阅请求以及发布信息请求
		let mut rx = $self.sub_rx;
		tokio::spawn(async move {
			let mut ps = PubSub { subs: HashMap::new() };
			loop {
				let Some(c) = rx.recv().await else {break;};
				match c {
					Content::Sub(id, pair) => ps.add(id, pair, &mut $onSubcribe).await,
					Content::Pub(name, data) => ps.publish(name, data).await,
					Content::Del(id) => ps.del(id),
				}
			}
		});
		//用于接收RPC客户端连接
		let mut id = 0u32;
		loop {
			let (stream, _) = $listener.accept().await?;
			let (readstream, writestream) = stream.into_split();
			let (tx, rx) = mpsc::channel(3); //创建用于发送的通道
			id += 1;
			{
				let sub_tx = $self.sub_tx.clone();
				let call_tx = $self.call_tx.clone();
				tokio::spawn(async move {
					let _sub_tx = sub_tx.clone();
					let _ = recv(readstream, sub_tx, call_tx, tx, id).await;
					let _ = _sub_tx.send(Content::Del(id)).await;
				});
			}
			tokio::spawn(send(writestream, rx));
		}
	};
}

impl Server {
	/// 创建一个RPC服务
	#[inline]
	pub fn new() -> Self {
		let (stx, srx) = mpsc::channel(3);
		Server { funcs: HashMap::new(), sub_tx: stx, sub_rx: srx, call_tx: None, call_rx: None }
	}
	/// 绑定RPC回调函数，请使用callback!宏必要时结合clone!宏进行bind操作
	#[inline]
	pub fn bind(&mut self, name: &'static str, f: Callback) {
		if self.call_rx.is_none() {
			let (ctx, crx) = mpsc::channel(3);
			self.call_tx = Some(ctx);
			self.call_rx = Some(crx);
		}
		self.funcs.insert(name.as_bytes(), f);
	}
	/// 获取发布器用于发布消息
	#[inline]
	pub fn publisher(&self) -> Publisher {
		Publisher { tx: self.sub_tx.clone() }
	}
	/// 通过unix域套接字异步运行服务端监听指定的路径
	#[cfg(unix)]
	pub async fn run_by_unix(
		mut self, path: &str, mut on_subcribe: impl FnMut(&str) -> Option<Bytes> + Send + Sync + 'static,
	) -> std::io::Result<()> {
		let path = Path::new(path);
		if path.exists() {
			std::fs::remove_file(path)?;
		} else {
			match path.parent() {
				Some(dir) if !dir.exists() => std::fs::create_dir_all(dir)?,
				_ => (),
			}
		}
		let listener = UnixListener::bind(path)?;
		listen!(self, listener, on_subcribe);
	}
	/// 通过tcp套接字异步运行服务端监听指定的地址
	pub async fn run_by_tcp(
		mut self, addr: &str, mut on_subcribe: impl FnMut(&str) -> Option<Bytes> + Send + Sync + 'static,
	) -> std::io::Result<()> {
		let listener = TcpListener::bind(addr).await?;
		listen!(self, listener, on_subcribe);
	}
	/// 异步运行服务端监听指定的地址，如果是类unix系统则使用unix域套接字通信，如果是其他系统则使用tcp通信
	#[cfg(unix)]
	pub async fn run(self, addr: &str) -> std::io::Result<()> {
		self.run_by_unix(addr, |_| None).await
	}
	#[cfg(unix)]
	pub async fn run_with(
		self, addr: &str, on_subcribe: impl FnMut(&str) -> Option<Bytes> + Send + Sync + 'static,
	) -> std::io::Result<()> {
		self.run_by_unix(addr, on_subcribe).await
	}
	#[cfg(not(unix))]
	pub async fn run(self, addr: &str) -> std::io::Result<()> {
		self.run_by_tcp(addr, |_| None).await
	}
	#[cfg(not(unix))]
	pub async fn run_with(
		self, addr: &str, on_subcribe: impl FnMut(&str) -> Option<Bytes> + Send + Sync + 'static,
	) -> std::io::Result<()> {
		self.run_by_tcp(addr, on_subcribe).await
	}
}