fix: change size signature to &mut

linyihai · Dec 2, 2022 · 54adcd7 · 54adcd7
1 parent 80bf9eb
commit 54adcd7
Show file tree

Hide file tree

Showing 11 changed files with 43 additions and 35 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/examples/src/hello/grpc_client.rs b/examples/src/hello/grpc_client.rs
@@ -16,7 +16,7 @@ lazy_static! {
 #[volo::main]
 async fn main() {
     let req = volo_gen::proto_gen::hello::HelloRequest {
-        name: "Volo".to_string(),
+        name: Some("Volo".to_string()),
     };
     let resp = CLIENT.clone().hello(req).await;
     match resp {

diff --git a/examples/src/hello/grpc_server.rs b/examples/src/hello/grpc_server.rs
@@ -12,7 +12,7 @@ impl volo_gen::proto_gen::hello::HelloService for S {
     ) -> Result<volo_grpc::Response<volo_gen::proto_gen::hello::HelloResponse>, volo_grpc::Status>
     {
         let resp = volo_gen::proto_gen::hello::HelloResponse {
-            message: format!("Hello, {}!", req.get_ref().name),
+            message: Some(format!("Hello, {}!", req.get_ref().name.as_ref().unwrap())),
         };
         Ok(volo_grpc::Response::new(resp))
     }

diff --git a/volo-build/src/thrift_backend.rs b/volo-build/src/thrift_backend.rs
@@ -86,7 +86,7 @@ impl VoloThriftBackend {
                             #decode_async
                         }
 
-                    fn size<T: ::pilota::thrift::TLengthProtocol>(&self, protocol: &T) -> usize {
+                    fn size<T: ::pilota::thrift::TLengthProtocol>(&self, protocol: &mut T) -> usize {
                         match self {
                             #(Self::#variant_names(value) => {
                                 ::volo_thrift::Message::size(value, protocol)
@@ -118,7 +118,7 @@ impl VoloThriftBackend {
                             #decode_async
                         }
 
-                    fn size<T: ::pilota::thrift::TLengthProtocol>(&self, protocol: &T) -> usize {
+                    fn size<T: ::pilota::thrift::TLengthProtocol>(&self, protocol: &mut T) -> usize {
                         match self {
                             #(Self::#variant_names(value) => {
                                 ::volo_thrift::Message::size(value, protocol)
@@ -175,7 +175,7 @@ impl VoloThriftBackend {
                             #decode_async
                         }
 
-                    fn size<T: ::pilota::thrift::TLengthProtocol>(&self, protocol: &T) -> usize {
+                    fn size<T: ::pilota::thrift::TLengthProtocol>(&self, protocol: &mut T) -> usize {
                         match self {
                             #(
                                 Self::#variant_names(value) => {

diff --git a/volo-thrift/src/codec/default/framed.rs b/volo-thrift/src/codec/default/framed.rs
@@ -196,9 +196,9 @@ where
         &mut self,
         cx: &mut Cx,
         msg: &ThriftMessage<Msg>,
-    ) -> crate::Result<usize> {
-        self.inner_size = self.inner.size(cx, msg)? as i32;
-        check_framed_size(self.inner_size, self.max_frame_size)?;
+    ) -> crate::Result<(usize, usize)> {
+        let (real_size, malloc_size) = self.inner.size(cx, msg)?;
+        self.inner_size = real_size as i32;
         // only calc framed size if role is client or server has detected framed in decode
         if cx.rpc_info().role() == Role::Client
             || cx
@@ -207,9 +207,10 @@ where
                 .unwrap_or(&HasFramed(false))
                 .0
         {
-            Ok(self.inner_size as usize + FRAMED_HEADER_SIZE)
+            check_framed_size(self.inner_size, self.max_frame_size)?;
+            Ok((real_size + FRAMED_HEADER_SIZE, malloc_size + FRAMED_HEADER_SIZE))
         } else {
-            Ok(self.inner_size as usize)
+            Ok((real_size, malloc_size))
         }
     }
 }

diff --git a/volo-thrift/src/codec/default/mod.rs b/volo-thrift/src/codec/default/mod.rs
@@ -60,11 +60,13 @@ pub trait ZeroCopyEncoder: Send + Sync + 'static {
     ///
     /// To avoid the overhead of calculating the size again in the [`encode`] method, the
     /// implementation can cache the size in the struct.
+    ///
+    /// The returned value is (real_size, recommended_malloc_size).
     fn size<Msg: Send + EntryMessage, Cx: ThriftContext>(
         &mut self,
         cx: &mut Cx,
         msg: &ThriftMessage<Msg>,
-    ) -> Result<usize>;
+    ) -> Result<(usize, usize)>;
 }
 
 /// [`ZeroCopyDecoder`] tries to decode a message without copying large data, so the [`BytesMut`] in
@@ -119,10 +121,10 @@ impl<E: ZeroCopyEncoder, W: AsyncWrite + Unpin + Send + Sync + 'static> Encoder
         msg: ThriftMessage<Req>,
     ) -> Result<()> {
         // first, we need to get the size of the message
-        let size = self.encoder.size(cx, &msg)?;
-        trace!("[VOLO] codec encode message size: {}", size);
+        let (real_size, malloc_size) = self.encoder.size(cx, &msg)?;
+        trace!("[VOLO] codec encode message real size: {}, malloc size: {}", real_size, malloc_size);
         // then we reserve the size of the message in the linked bytes
-        self.linked_bytes.reserve(size);
+        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)?;
         self.linked_bytes

diff --git a/volo-thrift/src/codec/default/thrift.rs b/volo-thrift/src/codec/default/thrift.rs
@@ -1,6 +1,8 @@
 use bytes::BytesMut;
 use linkedbytes::LinkedBytes;
-use pilota::thrift::{binary::TBinaryProtocol, ProtocolErrorKind, TAsyncBinaryProtocol};
+use pilota::thrift::{
+    binary::TBinaryProtocol, ProtocolErrorKind, TAsyncBinaryProtocol, TLengthProtocol,
+};
 use tokio::io::AsyncRead;
 use volo::util::buf_reader::BufReader;
 
@@ -209,14 +211,15 @@ impl ZeroCopyEncoder for ThriftCodec {
         &mut self,
         cx: &mut Cx,
         msg: &ThriftMessage<Msg>,
-    ) -> crate::Result<usize> {
+    ) -> crate::Result<(usize, usize)> {
         // for the client side, the match expression will always be `&self.protocol`
         // TODO: use the protocol in TTHeader?
         match cx.extensions().get::<Protocol>().unwrap_or(&self.protocol) {
             Protocol::Binary => {
-                let p = TBinaryProtocol::new((), true);
-                let size = msg.size(&p);
-                Ok(size)
+                let mut p = TBinaryProtocol::new((), true);
+                let real_size = msg.size(&mut p);
+                let malloc_size = real_size - p.zero_copy_len();
+                Ok((real_size, malloc_size))
             }
             p => Err(crate::Error::Pilota(
                 pilota::thrift::error::new_protocol_error(

diff --git a/volo-thrift/src/codec/default/ttheader.rs b/volo-thrift/src/codec/default/ttheader.rs
@@ -181,8 +181,9 @@ where
         &mut self,
         cx: &mut Cx,
         msg: &ThriftMessage<Msg>,
-    ) -> crate::Result<usize> {
-        self.inner_size = self.inner.size(cx, msg)?;
+    ) -> crate::Result<(usize, usize)> {
+        let (real_size, malloc_size) = self.inner.size(cx, msg)?;
+        self.inner_size = real_size;
         // only calc ttheader size if role is client or server has detected ttheader in decode
         if cx.rpc_info().role() == Role::Client
             || cx
@@ -191,9 +192,10 @@ where
                 .unwrap_or(&HasTTHeader(false))
                 .0
         {
-            Ok(self.inner_size + encode_size(cx)?)
+            let size = encode_size(cx)?;
+            Ok((real_size + size, malloc_size + size))
         } else {
-            Ok(self.inner_size)
+            Ok((real_size, malloc_size))
         }
     }
 }

diff --git a/volo-thrift/src/error.rs b/volo-thrift/src/error.rs
@@ -248,7 +248,7 @@ impl Message for ApplicationError {
         Ok(ApplicationError { kind, message })
     }
 
-    fn size<T: TLengthProtocol>(&self, protocol: &T) -> usize {
+    fn size<T: TLengthProtocol>(&self, protocol: &mut T) -> usize {
         protocol.write_struct_begin_len(&TAPPLICATION_EXCEPTION)
             + protocol.write_field_begin_len(&ERROR_MESSAGE_FIELD)
             + protocol.write_string_len(&self.message)

diff --git a/volo-thrift/src/message.rs b/volo-thrift/src/message.rs
@@ -26,7 +26,7 @@ pub trait EntryMessage: Sized + Send {
     where
         R: AsyncRead + Unpin + Send;
 
-    fn size<T: TLengthProtocol>(&self, protocol: &T) -> usize;
+    fn size<T: TLengthProtocol>(&self, protocol: &mut T) -> usize;
 }
 
 #[async_trait::async_trait]
@@ -57,7 +57,7 @@ where
             .map(Arc::new)
     }
 
-    fn size<T: TLengthProtocol>(&self, protocol: &T) -> usize {
+    fn size<T: TLengthProtocol>(&self, protocol: &mut T) -> usize {
         (**self).size(protocol)
     }
 }
diff --git a/volo-thrift/src/message_wrapper.rs b/volo-thrift/src/message_wrapper.rs
@@ -51,7 +51,7 @@ impl EntryMessage for DummyMessage {
         unreachable!()
     }
 
-    fn size<T: TLengthProtocol>(&self, _protocol: &T) -> usize {
+    fn size<T: TLengthProtocol>(&self, _protocol: &mut T) -> usize {
         unreachable!()
     }
 }
@@ -91,7 +91,7 @@ impl<U> ThriftMessage<U>
 where
     U: EntryMessage,
 {
-    pub(crate) fn size<T: TLengthProtocol>(&self, protocol: &T) -> usize {
+    pub(crate) fn size<T: TLengthProtocol>(&self, protocol: &mut T) -> usize {
         let ident = TMessageIdentifier::new(
             self.meta.method.clone(),
             self.meta.msg_type,