async: Refactor message encoding/deocding.

Make message encoding/decoding uniform.

Signed-off-by: wllenyj <[email protected]>

Author: wllenyj

src/asynchronous/client.rs

@@ -9,23 +9,16 @@ use std::os::unix::io::RawFd;
 use std::sync::{Arc, Mutex};
 
 use nix::unistd::close;
-use tokio::{
-    self,
-    io::split,
-    sync::mpsc::{channel, Receiver, Sender},
-    sync::Notify,
-};
+use tokio::{self, io::split, sync::mpsc, sync::Notify};
 
 use crate::common::client_connect;
 use crate::error::{Error, Result};
 use crate::proto::{Code, Codec, GenMessage, Message, Request, Response, MESSAGE_TYPE_RESPONSE};
+use crate::r#async::stream::{ResultReceiver, ResultSender};
 use crate::r#async::utils;
 
-type RequestSender = Sender<(GenMessage, Sender<Result<Vec<u8>>>)>;
-type RequestReceiver = Receiver<(GenMessage, Sender<Result<Vec<u8>>>)>;
-
-type ResponseSender = Sender<Result<Vec<u8>>>;
-type ResponseReceiver = Receiver<Result<Vec<u8>>>;
+type RequestSender = mpsc::Sender<(GenMessage, ResultSender)>;
+type RequestReceiver = mpsc::Receiver<(GenMessage, ResultSender)>;
 
 /// A ttrpc Client (async).
 #[derive(Clone)]
@@ -44,7 +37,7 @@ impl Client {
         let stream = utils::new_unix_stream_from_raw_fd(fd);
 
         let (mut reader, mut writer) = split(stream);
-        let (req_tx, mut rx): (RequestSender, RequestReceiver) = channel(100);
+        let (req_tx, mut rx): (RequestSender, RequestReceiver) = mpsc::channel(100);
 
         let req_map = Arc::new(Mutex::new(HashMap::new()));
         let req_map2 = req_map.clone();
@@ -131,7 +124,7 @@ impl Client {
                                         return;
                                     }
 
-                                    resp_tx2.send(Ok(msg.payload)).await.unwrap_or_else(|_e| error!("The request has returned"));
+                                    resp_tx2.send(Ok(msg)).await.unwrap_or_else(|_e| error!("The request has returned"));
                                 });
                             }
                             Err(e) => {
@@ -170,7 +163,7 @@ impl Client {
             .try_into()
             .map_err(|e: protobuf::error::ProtobufError| Error::Others(e.to_string()))?;
 
-        let (tx, mut rx): (ResponseSender, ResponseReceiver) = channel(100);
+        let (tx, mut rx): (ResultSender, ResultReceiver) = mpsc::channel(100);
         self.req_tx
             .send((msg, tx))
             .await
@@ -190,9 +183,9 @@ impl Client {
             .ok_or_else(|| Error::Others("Receive packet from receiver error".to_string()))?
         };
 
-        let buf = result?;
-        let res =
-            Response::decode(&buf).map_err(err_to_others_err!(e, "Unpack response error "))?;
+        let msg = result?;
+        let res = Response::decode(&msg.payload)
+            .map_err(err_to_others_err!(e, "Unpack response error "))?;
 
         let status = res.get_status();
         if status.get_code() != Code::OK {

src/asynchronous/server.rs

@@ -17,22 +17,23 @@ use futures::StreamExt as _;
 use nix::unistd;
 use tokio::{
     self,
-    io::{split, AsyncRead, AsyncWrite, AsyncWriteExt},
+    io::{split, AsyncRead, AsyncWrite},
     net::UnixListener,
     select, spawn,
-    sync::mpsc::{channel, Receiver, Sender},
+    sync::mpsc::{channel, Sender},
     time::timeout,
 };
 #[cfg(target_os = "linux")]
 use tokio_vsock::VsockListener;
 
-use crate::asynchronous::stream::{receive, respond, respond_with_status};
+use crate::asynchronous::stream::{respond, respond_with_status};
 use crate::asynchronous::unix_incoming::UnixIncoming;
 use crate::common::{self, Domain};
 use crate::context;
 use crate::error::{get_status, Error, Result};
-use crate::proto::{Code, MessageHeader, Status, MESSAGE_TYPE_REQUEST};
+use crate::proto::{Code, GenMessage, MessageHeader, Response, Status, MESSAGE_TYPE_REQUEST};
 use crate::r#async::shutdown;
+use crate::r#async::stream::{MessageReceiver, MessageSender};
 use crate::r#async::utils;
 use crate::r#async::{MethodHandler, TtrpcContext};
 
@@ -244,15 +245,15 @@ async fn spawn_connection_handler<S>(
 {
     spawn(async move {
         let (mut reader, mut writer) = split(stream);
-        let (tx, mut rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel(100);
+        let (tx, mut rx): (MessageSender, MessageReceiver) = channel(100);
 
         let server_shutdown = shutdown_waiter.clone();
         let (disconnect_notifier, disconnect_waiter) =
             shutdown::with_timeout(DEFAULT_CONN_SHUTDOWN_TIMEOUT);
 
         spawn(async move {
-            while let Some(buf) = rx.recv().await {
-                if let Err(e) = writer.write_all(&buf).await {
+            while let Some(msg) = rx.recv().await {
+                if let Err(e) = msg.write_to(&mut writer).await {
                     error!("write_message got error: {:?}", e);
                 }
             }
@@ -264,8 +265,8 @@ async fn spawn_connection_handler<S>(
             let handler_shutdown_waiter = disconnect_waiter.clone();
 
             select! {
-                resp = receive(&mut reader) => {
-                    match resp {
+                res = GenMessage::read_from(&mut reader) => {
+                    match res {
                         Ok(message) => {
                             spawn(async move {
                                 select! {
@@ -304,7 +305,7 @@ async fn do_handle_request(
     methods: Arc<HashMap<String, Box<dyn MethodHandler + Send + Sync>>>,
     header: MessageHeader,
     body: &[u8],
-) -> StdResult<(u32, Vec<u8>), Status> {
+) -> StdResult<Option<Response>, Status> {
     let req = utils::body_to_request(body)?;
     let path = utils::get_path(&req.service, &req.method);
     let method = methods
@@ -328,6 +329,7 @@ async fn do_handle_request(
             .handler(ctx, req)
             .await
             .map_err(get_unknown_status_and_log_err)
+            .map(Some)
     } else {
         timeout(
             Duration::from_nanos(req.timeout_nano as u64),
@@ -343,32 +345,39 @@ async fn do_handle_request(
             // Handler finished
             r.map_err(get_unknown_status_and_log_err)
         })
+        .map(Some)
     }
 }
 
 async fn handle_request(
-    tx: Sender<Vec<u8>>,
+    tx: MessageSender,
     fd: RawFd,
     methods: Arc<HashMap<String, Box<dyn MethodHandler + Send + Sync>>>,
-    message: (MessageHeader, Vec<u8>),
+    message: GenMessage,
 ) {
-    let (header, body) = message;
+    let GenMessage {
+        header,
+        payload: body,
+    } = message;
     let stream_id = header.stream_id;
 
     if header.type_ != MESSAGE_TYPE_REQUEST {
         return;
     }
 
     match do_handle_request(fd, methods, header, &body).await {
-        Ok((stream_id, resp_body)) => {
-            if let Err(x) = respond(tx.clone(), stream_id, resp_body).await {
-                error!("respond got error {:?}", x);
+        Ok(opt_msg) => match opt_msg {
+            Some(msg) => {
+                if let Err(x) = respond(tx.clone(), stream_id, msg).await {
+                    error!("respond got error {:?}", x);
+                }
             }
-        }
-        Err(status) => {
-            if let Err(x) = respond_with_status(tx.clone(), stream_id, status).await {
-                error!("respond got error {:?}", x);
+            None => {
+                unimplemented!();
             }
+        },
+        Err(status) => {
+            respond_with_status(tx.clone(), stream_id, status).await;
         }
     }
 }

src/asynchronous/stream.rs

@@ -3,129 +3,40 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use protobuf::Message;
-use tokio::io::AsyncReadExt;
+use tokio::sync::mpsc;
 
-use crate::error::{get_rpc_status, sock_error_msg, Error, Result};
-use crate::proto::{
-    Code, Response, Status, MESSAGE_HEADER_LENGTH, MESSAGE_LENGTH_MAX, MESSAGE_TYPE_RESPONSE,
-};
-use crate::r#async::utils;
+use crate::error::{Error, Result};
+use crate::proto::{Codec, GenMessage, Response, Status};
 use crate::MessageHeader;
 
-async fn receive_count<T>(reader: &mut T, count: usize) -> Result<Vec<u8>>
-where
-    T: AsyncReadExt + std::marker::Unpin,
-{
-    let mut content = vec![0u8; count];
-    if let Err(e) = reader.read_exact(&mut content).await {
-        return Err(Error::Socket(e.to_string()));
-    }
+pub type MessageSender = mpsc::Sender<GenMessage>;
+pub type MessageReceiver = mpsc::Receiver<GenMessage>;
 
-    Ok(content)
-}
-
-async fn receive_header<T>(reader: &mut T) -> Result<MessageHeader>
-where
-    T: AsyncReadExt + std::marker::Unpin,
-{
-    let buf = receive_count(reader, MESSAGE_HEADER_LENGTH).await?;
-    let size = buf.len();
-    if size != MESSAGE_HEADER_LENGTH {
-        return Err(sock_error_msg(
-            size,
-            format!("Message header length {} is too small", size),
-        ));
-    }
-
-    let mh = MessageHeader::from(&buf);
-
-    Ok(mh)
-}
-
-pub(crate) async fn receive<T>(reader: &mut T) -> Result<(MessageHeader, Vec<u8>)>
-where
-    T: AsyncReadExt + std::marker::Unpin,
-{
-    let mh = receive_header(reader).await?;
-    trace!("Got Message header {:?}", mh);
-
-    if mh.length > MESSAGE_LENGTH_MAX as u32 {
-        return Err(get_rpc_status(
-            Code::INVALID_ARGUMENT,
-            format!(
-                "message length {} exceed maximum message size of {}",
-                mh.length, MESSAGE_LENGTH_MAX
-            ),
-        ));
-    }
-
-    let buf = receive_count(reader, mh.length as usize).await?;
-    let size = buf.len();
-    if size != mh.length as usize {
-        return Err(sock_error_msg(
-            size,
-            format!("Message length {} is not {}", size, mh.length),
-        ));
-    }
-    trace!("Got Message body {:?}", buf);
-
-    Ok((mh, buf))
-}
-
-fn header_to_buf(mh: MessageHeader) -> Vec<u8> {
-    mh.into()
-}
+pub type ResultSender = mpsc::Sender<Result<GenMessage>>;
+pub type ResultReceiver = mpsc::Receiver<Result<GenMessage>>;
 
-pub(crate) fn to_res_buf(stream_id: u32, mut body: Vec<u8>) -> Vec<u8> {
-    let header = utils::get_response_header_from_body(stream_id, &body);
-    let mut buf = header_to_buf(header);
-    buf.append(&mut body);
-
-    buf
-}
-
-fn get_response_body(res: &Response) -> Result<Vec<u8>> {
-    let mut buf = Vec::with_capacity(res.compute_size() as usize);
-    let mut s = protobuf::CodedOutputStream::vec(&mut buf);
-    res.write_to(&mut s).map_err(err_to_others_err!(e, ""))?;
-    s.flush().map_err(err_to_others_err!(e, ""))?;
-
-    Ok(buf)
-}
-
-pub(crate) async fn respond(
-    tx: tokio::sync::mpsc::Sender<Vec<u8>>,
-    stream_id: u32,
-    body: Vec<u8>,
-) -> Result<()> {
-    let buf = to_res_buf(stream_id, body);
+pub(crate) async fn respond(tx: MessageSender, stream_id: u32, resp: Response) -> Result<()> {
+    let payload = resp
+        .encode()
+        .map_err(err_to_others_err!(e, "Encode Response failed."))?;
+    let msg = GenMessage {
+        header: MessageHeader::new_response(stream_id, payload.len() as u32),
+        payload,
+    };
 
-    tx.send(buf)
+    tx.send(msg)
         .await
         .map_err(err_to_others_err!(e, "Send packet to sender error "))
 }
 
-pub(crate) async fn respond_with_status(
-    tx: tokio::sync::mpsc::Sender<Vec<u8>>,
-    stream_id: u32,
-    status: Status,
-) -> Result<()> {
+pub(crate) async fn respond_with_status(tx: MessageSender, stream_id: u32, status: Status) {
     let mut res = Response::new();
     res.set_status(status);
-    let mut body = get_response_body(&res)?;
-
-    let mh = MessageHeader {
-        length: body.len() as u32,
-        stream_id,
-        type_: MESSAGE_TYPE_RESPONSE,
-        flags: 0,
-    };
 
-    let mut buf = header_to_buf(mh);
-    buf.append(&mut body);
-
-    tx.send(buf)
+    respond(tx, stream_id, res)
         .await
-        .map_err(err_to_others_err!(e, "Send packet to sender error "))
+        .map_err(|e| {
+            error!("respond with status got error {:?}", e);
+        })
+        .ok();
 }

src/asynchronous/utils.rs

@@ -12,7 +12,7 @@ use protobuf::{CodedInputStream, Message};
 use tokio::net::UnixStream;
 
 use crate::error::{get_status, Result};
-use crate::proto::{Code, MessageHeader, Request, Status, MESSAGE_TYPE_RESPONSE};
+use crate::proto::{Code, MessageHeader, Request, Response, Status};
 
 /// Handle request in async mode.
 #[macro_export]
@@ -48,12 +48,7 @@ macro_rules! async_request_handler {
             },
         }
 
-        let mut buf = Vec::with_capacity(res.compute_size() as usize);
-        let mut s = protobuf::CodedOutputStream::vec(&mut buf);
-        res.write_to(&mut s).map_err(ttrpc::err_to_others!(e, ""))?;
-        s.flush().map_err(ttrpc::err_to_others!(e, ""))?;
-
-        return Ok(($ctx.mh.stream_id, buf));
+        return Ok(res);
     };
 }
 
@@ -88,7 +83,7 @@ macro_rules! async_client_request {
 /// Trait that implements handler which is a proxy to the desired method (async).
 #[async_trait]
 pub trait MethodHandler {
-    async fn handler(&self, ctx: TtrpcContext, req: Request) -> Result<(u32, Vec<u8>)>;
+    async fn handler(&self, ctx: TtrpcContext, req: Request) -> Result<Response>;
 }
 
 /// The context of ttrpc (async).
@@ -100,15 +95,6 @@ pub struct TtrpcContext {
     pub timeout_nano: i64,
 }
 
-pub(crate) fn get_response_header_from_body(stream_id: u32, body: &[u8]) -> MessageHeader {
-    MessageHeader {
-        length: body.len() as u32,
-        stream_id,
-        type_: MESSAGE_TYPE_RESPONSE,
-        flags: 0,
-    }
-}
-
 pub(crate) fn new_unix_stream_from_raw_fd(fd: RawFd) -> UnixStream {
     let std_stream: std::os::unix::net::UnixStream;
     unsafe {