【Rust投稿】从零实现消息中间件(4)-SERVER.CLIENT

原创

Rust语言中文社区 2022-06-27 13:49:44 ©著作权

文章标签 sed github 负载均衡 文章分类 Java 后端开发

©著作权归作者所有：来自51CTO博客作者Rust语言中文社区的原创作品，请联系作者获取转载授权，否则将追究法律责任

这部分主要说的是服务器端对于来自client连接的数据的处理. 主要功能包括

接收消息
收到sub消息,就记录到全局列表中
收到pub消息,就发送给相关订阅的client
出错,删除订阅,关闭连接

数据结构定义

Client中除了cid以外,其他两项都使用了Mutex进行保护,上一篇讲到过,凡是多线程读写的都需要Arc<Mutex>保护.

srv: 主要还是pub sub的时候都需要访问全局的sublist.
msg_sender: 之所以用Mutex保护是因为除了client自己要发送消息,当其他client pub消息的时候也要通过这个ClientMessageSender发送消息
ClientMessageSender在我们这个版本中则非常简单,就是一个TcpStream的writer.rust #[derive(Debug)] pub struct Client<T: SubListTrait> { pub srv: Arc<Mutex<ServerState<T>>>, pub cid: u64, pub msg_sender: Arc<Mutex<ClientMessageSender>>, } #[derive(Debug)] pub struct ClientMessageSender { writer: WriteHalf<TcpStream>, }

代码实现

process_connection

创建Client以及可以共享使用的ClientMessageSender
启动client_task

impl<T: SubListTrait + Send + 'static> Client<T> {

pub fn process_connection(
    cid: u64,
    srv: Arc<Mutex<ServerState<T>>>,
    conn: TcpStream,
) -> Arc<Mutex<ClientMessageSender>> {
let (reader, writer) = tokio::io::split(conn);
let msg_sender = Arc::new(Mutex::new(ClientMessageSender::new(writer)));
let c = Client {
        srv: srv,
        cid,
        msg_sender: msg_sender.clone(),
};
    tokio::spawn(Client::client_task(c, reader));
    msg_sender
}

...
}

client_task

主要功能:

读取,解析消息
分发消息给相应的处理函数

process_error
process_sub
process_pub

这个其实就是一个tcp连接的主循环,说到这里我想把tokio::spawn 和 go语言中的go关键字做一个类比.
在go中TcpServer接收到一个连接以后,紧接着就是单独起一个goroutine来处理.类似于go client.processConnection(),而到了Rust中基本上可以等价为

tokio::spawn(async move{
 Client::process_connection();
});

当然Rust重要复杂很多,涉及到所有权,生命周期等一系列问题.

async fn client_task(self, mut reader: ReadHalf<TcpStream>) {
let mut buf = [0; 1024];
let mut parser = Parser::new();
let mut subs = HashMap::new();
loop {
let r = reader.read(&mut buf[..]).await;
if r.is_err() {
let e = r.unwrap_err();
self.process_error(e, subs).await;
return;
}
let r = r.unwrap();
let n = r;
if n == 0 {
self.process_error(NError::new(ERROR_CONNECTION_CLOSED), subs)
.await;
return;
}
let mut buf = &buf[0..n];
loop {
let r = parser.parse(&buf[..]);
if r.is_err() {
self.process_error(r.unwrap_err(), subs).await;
return;
}
let (result, left) = r.unwrap();

match result {
                    ParseResult::NoMsg => {
break;
}
                    ParseResult::Sub(ref sub) => {
if let Err(e) = self.process_sub(sub, &mut subs).await {
self.process_error(e, subs).await;
return;
}
}
                    ParseResult::Pub(ref pub_arg) => {
if let Err(e) = self.process_pub(pub_arg).await {
self.process_error(e, subs).await;
return;
}
}
}
if left == buf.len() {
break;
}
                buf = &buf[left..];
}
}
}

从整个代码中也可以看出client_task的主要工作就是接受消息,并处理.

process_error

删除所有订阅
关闭连接rust async fn process_error<E: Error>(&self, err: E, subs: HashMap<String, ArcSubscription>) { println!("client {} process err {:?}", self.cid, err); { let mut sublist = &mut self.srv.lock().await.sublist; for (_, sub) in subs { sublist.remove(sub); } } let r = self.msg_sender.lock().await.writer.shutdown().await; if r.is_err() { println!("shutdown err {:?}", r.unwrap_err()); } }

process_sub

对于收到的sub则是

全局订阅列表中保存一份
本地连接保存一份,这样连接断开的时候好删除
为了避免内存分配,我们的SubArg里面使用的还是Parer缓冲区中的内存,当我们需要在连接之外访问这些信息的时候,我们就必须单独保存一份了,这里我们用的是sub.subject.to_string()来分配一个新的内存.rust async fn process_sub( &self, sub: &SubArg<'_>, subs: &mut HashMap<String, ArcSubscription>, ) -> crate::error::Result<()> { let sub = Subscription { subject: sub.subject.to_string(), queue: sub.queue.map(|q| q.to_string()), sid: sub.sid.to_string(), msg_sender: self.msg_sender.clone(), }; let sub = Arc::new(sub); subs.insert(sub.subject.clone(), sub.clone()); let sublist = &mut self.srv.lock().await.sublist; sublist.insert(sub); Ok(()) }

process_pub

收到pub消息,

查找所有的订阅
将消息逐一转发给他们
转发的过程中要稍微麻烦一点,因为考虑到设计中的负载均衡问题,qsubs则是从同一个queue中随机选择一个来推送消息.rust async fn process_pub(&self, pub_arg: &PubArg<'_>) -> crate::error::Result<()> { let sub_result = { let sub_list = &mut self.srv.lock().await.sublist; sub_list.match_subject(pub_arg.subject)? }; for sub in sub_result.subs.iter() { self.send_message(sub.as_ref(), pub_arg) .await .map_err(|e| NError::new(ERROR_CONNECTION_CLOSED))?; } //qsubs 要考虑负载均衡问题 let mut rng = rand::rngs::StdRng::from_entropy(); for qsubs in sub_result.qsubs.iter() { let n = rng.next_u32(); let n = n as usize % qsubs.len(); let sub = qsubs.get(n).unwrap(); self.send_message(sub.as_ref(), pub_arg) .await .map_err(|e| NError::new(ERROR_CONNECTION_CLOSED))?; } Ok(()) }

send_message

就是拼装消息格式
因为是第一个版本,也是展示关键api的使用,里面用到了大量的await,实际上没有必要.
实际项目中,肯定会使用缓冲区来做.

///消息格式
///```
/// MSG <subject> <sid> <size>\r\n
/// <message>\r\n
/// ```
async fn send_message(&self, sub: &Subscription, pub_arg: &PubArg<'_>) -> std::io::Result<()> {
let writer = &mut sub.msg_sender.lock().await.writer;
    writer.write("MSG ".as_bytes()).await?;
    writer.write(sub.subject.as_bytes()).await?;
    writer.write(" ".as_bytes()).await?;
    writer.write(sub.sid.as_bytes()).await?;
    writer.write(" ".as_bytes()).await?;
    writer.write(pub_arg.size_buf.as_bytes()).await?;
    writer.write("\r\n".as_bytes()).await?;
    writer.write(pub_arg.msg).await?;
    writer.write("\r\n".as_bytes()).await?;
Ok(())
}