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

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

1. 接收消息
   

2. 收到sub消息,就记录到全局列表中
   

3. 收到pub消息,就发送给相关订阅的client
   

4. 出错,删除订阅,关闭连接
   
   

数据结构定义

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

1. 删除所有订阅
   

2. 关闭连接
   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则是

1. 全局订阅列表中保存一份
   

2. 本地连接保存一份,这样连接断开的时候好删除
   为了避免内存分配,我们的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消息,

1. 查找所有的订阅
   

2. 将消息逐一转发给他们
   转发的过程中要稍微麻烦一点,因为考虑到设计中的负载均衡问题,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(())
}

其他

相关代码都在我的github rnats 欢迎围观

https://github.com/nkbai/learnrustbynats