[multiplex] merge flush and use sharded hashmap #321

volo-thrift/Cargo.toml

@@ -48,6 +48,7 @@ tokio = { workspace = true, features = [
     "parking_lot",
 ] }
 tracing.workspace = true
+tokio-condvar = "0.1.0"
 
 [features]
 default = []

volo-thrift/src/client/mod.rs

@@ -463,9 +463,10 @@ impl<IL, OL, C, Req, Resp, MkT, MkC, LB> ClientBuilder<IL, OL, C, Req, Resp, MkT
 }
 
 #[derive(Clone)]
-pub struct MessageService<Resp, MkT, MkC>
+pub struct MessageService<Req, Resp, MkT, MkC>
 where
-    Resp: EntryMessage + Send + 'static,
+    Req: EntryMessage + Send + 'static + Sync,
+    Resp: EntryMessage + Send + 'static + Sync,
     MkT: MakeTransport,
     MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf> + Sync,
 {
@@ -474,14 +475,14 @@ where
     #[cfg(feature = "multiplex")]
     inner: motore::utils::Either<
         pingpong::Client<Resp, MkT, MkC>,
-        crate::transport::multiplex::Client<Resp, MkT, MkC>,
+        crate::transport::multiplex::Client<Req, Resp, MkT, MkC>,
     >,
     read_biz_error: bool,
 }
 
-impl<Req, Resp, MkT, MkC> Service<ClientContext, Req> for MessageService<Resp, MkT, MkC>
+impl<Req, Resp, MkT, MkC> Service<ClientContext, Req> for MessageService<Req, Resp, MkT, MkC>
 where
-    Req: EntryMessage + 'static + Send,
+    Req: Send + 'static + EntryMessage + Sync,
     Resp: Send + 'static + EntryMessage + Sync,
     MkT: MakeTransport,
     MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf> + Sync,
@@ -531,8 +532,8 @@ where
         + Clone
         + Sync,
     Req: EntryMessage + Send + 'static + Sync + Clone,
-    Resp: EntryMessage + Send + 'static,
-    IL: Layer<MessageService<Resp, MkT, MkC>>,
+    Resp: EntryMessage + Send + 'static + Sync,
+    IL: Layer<MessageService<Req, Resp, MkT, MkC>>,
     IL::Service:
         Service<ClientContext, Req, Response = Option<Resp>> + Sync + Clone + Send + 'static,
     <IL::Service as Service<ClientContext, Req>>::Error: Send + Into<ClientError>,

volo-thrift/src/codec/default/mod.rs

@@ -116,8 +116,7 @@ pub struct DefaultEncoder<E, W> {
 impl<E: ZeroCopyEncoder, W: AsyncWrite + Unpin + Send + Sync + 'static> Encoder
     for DefaultEncoder<E, W>
 {
-    #[inline]
-    async fn encode<Req: Send + EntryMessage, Cx: ThriftContext>(
+    async fn send<Req: Send + EntryMessage, Cx: ThriftContext>(
         &mut self,
         cx: &mut Cx,
         msg: ThriftMessage<Req>,
@@ -179,6 +178,52 @@ impl<E: ZeroCopyEncoder, W: AsyncWrite + Unpin + Send + Sync + 'static> Encoder
         }
         // write_result
     }
+
+    #[inline]
+    async fn encode<Req: Send + EntryMessage, Cx: ThriftContext>(
+        &mut self,
+        cx: &mut Cx,
+        msg: ThriftMessage<Req>,
+    ) -> Result<(), ThriftException> {
+        cx.stats_mut().record_encode_start_at();
+
+        // first, we need to get the size of the message
+        let (real_size, malloc_size) = self.encoder.size(cx, &msg)?;
+        trace!(
+            "[VOLO] codec encode message real size: {}, malloc size: {}",
+            real_size,
+            malloc_size
+        );
+        cx.stats_mut().set_write_size(real_size);
+
+        // then we reserve the size of the message in the linked bytes
+        self.linked_bytes.reserve(malloc_size);
+        // after that, we encode the message into the linked bytes
+        self.encoder
+            .encode(cx, &mut self.linked_bytes, msg)
+            .map_err(|e| {
+                // record the error time
+                cx.stats_mut().record_encode_end_at();
+                e
+            })?;
+
+        cx.stats_mut().record_encode_end_at();
+        Ok(())
+    }
+
+    async fn flush(&mut self) -> Result<(), ThriftException> {
+        let write_result: Result<(), ThriftException> = self
+            .linked_bytes
+            .write_all_vectored(&mut self.writer)
+            .await
+            .map_err(|e| e.into());
+        write_result?;
+        self.writer.flush().await.map_err(Into::into)
+    }
+
+    async fn reset(&mut self) {
+        self.linked_bytes.reset();
+    }
 }
 
 pub struct DefaultDecoder<D, R> {

volo-thrift/src/codec/mod.rs

@@ -26,11 +26,19 @@ pub trait Decoder: Send + 'static {
 ///
 /// Note: [`Encoder`] should be designed to be ready for reuse.
 pub trait Encoder: Send + 'static {
+    fn reset(&mut self) -> impl Future<Output = ()> + Send;
+    fn send<Req: Send + EntryMessage, Cx: ThriftContext>(
+        &mut self,
+        cx: &mut Cx,
+        msg: ThriftMessage<Req>,
+    ) -> impl Future<Output = Result<(), ThriftException>> + Send;
     fn encode<Req: Send + EntryMessage, Cx: ThriftContext>(
         &mut self,
         cx: &mut Cx,
         msg: ThriftMessage<Req>,
     ) -> impl Future<Output = Result<(), ThriftException>> + Send;
+
+    fn flush(&mut self) -> impl Future<Output = Result<(), ThriftException>> + Send;
 }
 
 /// [`MakeCodec`] receives an [`AsyncRead`] and an [`AsyncWrite`] and returns a

volo-thrift/src/transport/mod.rs

@@ -1,5 +1,4 @@
 pub(crate) mod incoming;
-#[cfg(feature = "multiplex")]
 pub mod multiplex;
 pub mod pingpong;
 pub mod pool;

volo-thrift/src/transport/multiplex/client.rs

@@ -14,18 +14,18 @@ use crate::{
     ClientError, EntryMessage, ThriftMessage,
 };
 
-pub struct MakeClientTransport<MkT, MkC, Resp>
+pub struct MakeClientTransport<MkT, MkC, Req, Resp>
 where
     MkT: MakeTransport,
     MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf>,
 {
     make_transport: MkT,
     make_codec: MkC,
-    _phantom: PhantomData<fn() -> Resp>,
+    _phantom: PhantomData<(fn() -> Resp, fn() -> Req)>,
 }
 
-impl<MkT: MakeTransport, MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf>, Resp> Clone
-    for MakeClientTransport<MkT, MkC, Resp>
+impl<MkT: MakeTransport, MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf>, Req, Resp> Clone
+    for MakeClientTransport<MkT, MkC, Req, Resp>
 {
     fn clone(&self) -> Self {
         Self {
@@ -36,7 +36,7 @@ impl<MkT: MakeTransport, MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf>, Resp> Cl
     }
 }
 
-impl<MkT, MkC, Resp> MakeClientTransport<MkT, MkC, Resp>
+impl<MkT, MkC, Req, Resp> MakeClientTransport<MkT, MkC, Req, Resp>
 where
     MkT: MakeTransport,
     MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf>,
@@ -51,13 +51,14 @@ where
     }
 }
 
-impl<MkT, MkC, Resp> UnaryService<Address> for MakeClientTransport<MkT, MkC, Resp>
+impl<MkT, MkC, Req, Resp> UnaryService<Address> for MakeClientTransport<MkT, MkC, Req, Resp>
 where
     MkT: MakeTransport,
     MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf> + Sync,
-    Resp: EntryMessage + Send + 'static,
+    Resp: EntryMessage + Send + 'static + Sync,
+    Req: EntryMessage + Send + 'static + Sync,
 {
-    type Response = ThriftTransport<MkC::Encoder, Resp>;
+    type Response = ThriftTransport<MkC::Encoder, Req, Resp>;
     type Error = io::Error;
 
     async fn call(&self, target: Address) -> Result<Self::Response, Self::Error> {
@@ -72,22 +73,24 @@ where
     }
 }
 
-pub struct Client<Resp, MkT, MkC>
+pub struct Client<Req, Resp, MkT, MkC>
 where
     MkT: MakeTransport,
     MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf> + Sync,
-    Resp: EntryMessage + Send + 'static,
+    Resp: EntryMessage + Send + 'static + Sync,
+    Req: EntryMessage + Send + 'static + Sync,
 {
     #[allow(clippy::type_complexity)]
-    make_transport: PooledMakeTransport<MakeClientTransport<MkT, MkC, Resp>, Address>,
+    make_transport: PooledMakeTransport<MakeClientTransport<MkT, MkC, Req, Resp>, Address>,
     _marker: PhantomData<Resp>,
 }
 
-impl<Resp, MkT, MkC> Clone for Client<Resp, MkT, MkC>
+impl<Req, Resp, MkT, MkC> Clone for Client<Req, Resp, MkT, MkC>
 where
     MkT: MakeTransport,
     MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf> + Sync,
-    Resp: EntryMessage + Send + 'static,
+    Resp: EntryMessage + Send + 'static + Sync,
+    Req: EntryMessage + Send + 'static + Sync,
 {
     fn clone(&self) -> Self {
         Self {
@@ -97,11 +100,12 @@ where
     }
 }
 
-impl<Resp, MkT, MkC> Client<Resp, MkT, MkC>
+impl<Req, Resp, MkT, MkC> Client<Req, Resp, MkT, MkC>
 where
     MkT: MakeTransport,
     MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf> + Sync,
-    Resp: EntryMessage + Send + 'static,
+    Resp: EntryMessage + Send + 'static + Sync,
+    Req: EntryMessage + Send + 'static + Sync,
 {
     pub fn new(make_transport: MkT, pool_cfg: Option<Config>, make_codec: MkC) -> Self {
         let make_transport = MakeClientTransport::new(make_transport, make_codec);
@@ -113,9 +117,9 @@ where
     }
 }
 
-impl<Req, Resp, MkT, MkC> Service<ClientContext, ThriftMessage<Req>> for Client<Resp, MkT, MkC>
+impl<Req, Resp, MkT, MkC> Service<ClientContext, ThriftMessage<Req>> for Client<Req, Resp, MkT, MkC>
 where
-    Req: Send + 'static + EntryMessage,
+    Req: Send + 'static + EntryMessage + Sync,
     Resp: EntryMessage + Send + 'static + Sync,
     MkT: MakeTransport,
     MkC: MakeCodec<MkT::ReadHalf, MkT::WriteHalf> + Sync,

volo-thrift/src/transport/multiplex/mod.rs

@@ -1,6 +1,7 @@
 mod client;
 mod server;
 mod thrift_transport;
+pub mod utils;
 
 pub use client::Client;
 pub use server::serve;

volo-thrift/src/transport/multiplex/server.rs

@@ -56,7 +56,7 @@ pub async fn serve<Svc, Req, Resp, E, D>(
                                         if let Err(e) = metainfo::METAINFO
                                             .scope(
                                                 RefCell::new(mi),
-                                                encoder.encode::<Resp, ServerContext>(&mut cx, msg),
+                                                encoder.send::<Resp, ServerContext>(&mut cx, msg),
                                             )
                                             .await
                                         {