Merge pull request #66 from minghuaw/avoid_copy_partial_payload

Avoid copy partial payload for multi-frame delivery

Author: minghuaw

examples/protocol_test/src/bin/connector.rs

@@ -4,8 +4,8 @@ use fe2o3_amqp::{types::primitives::Value, Connection, Delivery, Receiver, Sende
 use tracing::{instrument, Level};
 use tracing_subscriber::FmtSubscriber;
 
-// const SASL_PLAIN: &str = "";
-const SASL_PLAIN: &str = "guest:guest";
+const SASL_PLAIN: &str = "";
+// const SASL_PLAIN: &str = "guest:guest";
 const BASE_ADDR: &str = "localhost:5672";
 
 #[instrument]

examples/protocol_test/src/bin/listener.rs

@@ -34,7 +34,7 @@ async fn session_main(mut session: ListenerSessionHandle) {
                 let handle = tokio::spawn(async move {
                     tracing::info!("Incoming link is connected (remote: sender, local: receiver");
                     let delivery = recver.recv::<Value>().await.unwrap();
-                    tracing::info!(message = ?delivery.message());
+                    tracing::info!(id = ?delivery.delivery_id());
                     recver.accept(&delivery).await.unwrap();
                     if let Err(e) = recver.close().await {
                         // The remote may close the session
@@ -72,7 +72,7 @@ async fn main() {
     // let connection_acceptor = ConnectionAcceptor::new("test_conn_listener");
     let connection_acceptor = ConnectionAcceptor::builder()
         .container_id("example_connection_acceptor")
-        .sasl_acceptor(SaslPlainMechanism::new("guest", "guest"))
+        // .sasl_acceptor(SaslPlainMechanism::new("guest", "guest"))
         .build();
 
     while let Ok((stream, addr)) = tcp_listener.accept().await {

examples/protocol_test/src/bin/receiver.rs

@@ -70,10 +70,7 @@ async fn main() {
     // let body = delivery.into_body();
     println!("{:?}", delivery.delivery_id());
 
-    if let Err(err) = receiver.close().await {
-        println!("{}", err);
-    }
-
+    receiver.close().await.unwrap();
     session.end().await.unwrap();
     connection.close().await.unwrap();
 }

examples/protocol_test/src/bin/send_large_content.rs

@@ -206,7 +206,7 @@ async fn main() -> Result<()> {
 
     // let args: Vec<String> = env::args().collect();
     // let test_sender = TestSender::try_from(args)?;
-    let message_iter = create_message_sizes("string", "[10, 10, 10, 10]")?;
+    let message_iter = create_message_sizes("string", "[10, 10]")?;
     let test_sender = TestSender {
         broker_addr: "localhost:5672".to_string(),
         target_addr: "q1".to_string(),

examples/protocol_test/src/bin/sender.rs

@@ -22,40 +22,41 @@ async fn main() {
     let subscriber = FmtSubscriber::builder()
         // all spans/events with a level higher than TRACE (e.g, debug, info, warn, etc.)
         // will be written to stdout.
-        .with_max_level(Level::TRACE)
+        .with_max_level(Level::DEBUG)
         // .with_max_level(Level::DEBUG)
         // completes the builder.
         .finish();
 
     tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed");
 
-    let addr = "localhost:5671";
-    let domain = "localhost";
-    let stream = TcpStream::connect(addr).await.unwrap();
-    let connector = native_tls::TlsConnector::builder()
-        .danger_accept_invalid_certs(true)
-        .build()
-        .unwrap();
-    let connector = tokio_native_tls::TlsConnector::from(connector);
-    let tls_stream = connector.connect(domain, stream).await.unwrap();
-
-    // let mut connection = Connection::open("connection-1", "amqp://guest:guest@localhost:5671")
+    // let addr = "localhost:5671";
+    // let domain = "localhost";
+    // let stream = TcpStream::connect(addr).await.unwrap();
+    // let connector = native_tls::TlsConnector::builder()
+    //     .danger_accept_invalid_certs(true)
+    //     .build()
+    //     .unwrap();
+    // let connector = tokio_native_tls::TlsConnector::from(connector);
+    // let tls_stream = connector.connect(domain, stream).await.unwrap();
+
+    // let mut connection = Connection::builder()
+    //     .container_id("connection-1")
+    //     .scheme("amqp")
+    //     .max_frame_size(1000)
+    //     .channel_max(9)
+    //     .idle_time_out(50_000 as u32)
+    //     // .sasl_profile(SaslProfile::Plain {
+    //     //     username: "guest".into(),
+    //     //     password: "guest".into(),
+    //     // })
+    //     .open_with_stream(tls_stream)
+    //     // .open("amqp://localhost:5672")
+    //     // .open("amqp://guest:guest@localhost:5672")
+    //     // .open("localhost:5672")
     //     .await
     //     .unwrap();
-    let mut connection = Connection::builder()
-        .container_id("connection-1")
-        .scheme("amqp")
-        .max_frame_size(1000)
-        .channel_max(9)
-        .idle_time_out(50_000 as u32)
-        // .sasl_profile(SaslProfile::Plain {
-        //     username: "guest".into(),
-        //     password: "guest".into(),
-        // })
-        .open_with_stream(tls_stream)
-        // .open("amqp://localhost:5672")
-        // .open("amqp://guest:guest@localhost:5672")
-        // .open("localhost:5672")
+
+    let mut connection = Connection::open("connection-1", "amqp://localhost:5672")
         .await
         .unwrap();
 

fe2o3-amqp/src/control.rs

@@ -87,7 +87,7 @@ pub(crate) enum SessionControl {
 impl std::fmt::Display for SessionControl {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            SessionControl::End(_) => write!(f, "End"),
+            SessionControl::End(err) => write!(f, "End({:?})", err),
             SessionControl::AllocateLink {
                 link_name: _,
                 link_relay: _,

fe2o3-amqp/src/endpoint/link.rs

@@ -14,7 +14,7 @@ use tokio::sync::mpsc;
 use crate::{
     control::SessionControl,
     link::{delivery::Delivery, state::LinkState, LinkFrame},
-    Payload,
+    Payload, util::{AsByteIterator, IntoReader},
 };
 
 use super::{OutputHandle, Settlement};
@@ -163,10 +163,10 @@ pub(crate) trait ReceiverLink: Link + LinkExt {
 
     // More than one transfer frames should be hanlded by the
     // `Receiver`
-    async fn on_complete_transfer<'a, T>(
+    async fn on_complete_transfer<'a, T, P>(
         &'a mut self,
         transfer: Transfer,
-        payload: Payload,
+        payload: P,
     ) -> Result<
         (
             Delivery<T>,
@@ -175,7 +175,8 @@ pub(crate) trait ReceiverLink: Link + LinkExt {
         Self::TransferError,
     >
     where
-        T: DecodeIntoMessage + Send;
+        T: DecodeIntoMessage + Send,
+        for<'b> P: IntoReader + AsByteIterator<'b> + Send + 'a;
 
     async fn dispose(
         &mut self,

fe2o3-amqp/src/link/mod.rs

@@ -5,7 +5,6 @@ mod source;
 use std::{collections::BTreeMap, marker::PhantomData, sync::Arc};
 
 use async_trait::async_trait;
-use bytes::Buf;
 use fe2o3_amqp_types::{
     definitions::{
         self, AmqpError, DeliveryNumber, DeliveryTag, MessageFormat, ReceiverSettleMode, Role,

fe2o3-amqp/src/link/receiver.rs

@@ -3,7 +3,6 @@
 use std::time::Duration;
 
 use async_trait::async_trait;
-use bytes::BytesMut;
 use fe2o3_amqp_types::{
     definitions::{self, DeliveryNumber, DeliveryTag, SequenceNo},
     messaging::{
@@ -21,7 +20,7 @@ use crate::{
     control::SessionControl,
     endpoint::{self, LinkAttach, LinkDetach, LinkExt},
     session::SessionHandle,
-    Payload,
+    Payload, 
 };
 
 use super::{
@@ -63,20 +62,20 @@ macro_rules! or_assign {
 #[derive(Debug)]
 pub(crate) struct IncompleteTransfer {
     pub performative: Transfer,
-    pub buffer: BytesMut,
+    pub buffer: Vec<Payload>,
     pub section_number: u32,
     pub section_offset: u64,
 }
 
 impl IncompleteTransfer {
     pub fn new(transfer: Transfer, partial_payload: Payload) -> Self {
         let (number, offset) = section_number_and_offset(partial_payload.as_ref());
-        let mut buffer = BytesMut::new();
-        // TODO: anyway to make this not copying the bytes?
-        buffer.extend(partial_payload);
+        // let mut buffer = BytesMut::new();
+        // // TODO: anyway to make this not copying the bytes?
+        // buffer.extend(partial_payload);
         Self {
             performative: transfer,
-            buffer,
+            buffer: vec![partial_payload],
             section_number: number,
             section_offset: offset,
         }
@@ -140,7 +139,7 @@ impl IncompleteTransfer {
             self.section_offset = offset;
         }
 
-        self.buffer.extend(other);
+        self.buffer.push(other);
     }
 }
 
@@ -654,9 +653,10 @@ where
             match self.incomplete_transfer.take() {
                 Some(mut incomplete) => {
                     incomplete.or_assign(transfer)?;
-                    incomplete.buffer.extend(payload);
+                    incomplete.buffer.push(payload);
+
                     self.link
-                        .on_complete_transfer(incomplete.performative, incomplete.buffer.freeze())
+                        .on_complete_transfer(incomplete.performative, incomplete.buffer)
                         .await?
                 }
                 None => {

fe2o3-amqp/src/link/receiver_link.rs

@@ -1,6 +1,8 @@
 use fe2o3_amqp_types::messaging::message::DecodeIntoMessage;
 use serde_amqp::format_code::EncodingCodes;
 
+use crate::util::{AsByteIterator, IntoReader};
+
 use super::*;
 
 const DESCRIBED_TYPE: u8 = EncodingCodes::DescribedType as u8;
@@ -146,10 +148,10 @@ where
         }
     }
 
-    async fn on_complete_transfer<'a, T>(
+    async fn on_complete_transfer<'a, T, P>(
         &'a mut self,
         transfer: Transfer,
-        payload: Payload,
+        payload: P,
     ) -> Result<
         (
             Delivery<T>,
@@ -159,6 +161,7 @@ where
     >
     where
         T: DecodeIntoMessage + Send,
+        for<'b> P: IntoReader + AsByteIterator<'b> + Send + 'a,
     {
         // ReceiverFlowState will not wait until link credit is available.
         // Will return with an error if there is not enough link credit.
@@ -182,7 +185,7 @@ where
         let (message, delivery_state) = if settled_by_sender {
             // If the message is pre-settled, there is no need to
             // add to the unsettled map and no need to reply to the Sender
-            let message = T::decode_into_message(payload.reader())
+            let message = T::decode_into_message(payload.into_reader())
                 .map_err(|_| Self::TransferError::MessageDecodeError)?;
             (message, None)
         } else {
@@ -208,7 +211,7 @@ where
                 // once it has arrived without waiting for the sender to settle first.
                 ReceiverSettleMode::First => {
                     // Spontaneously settle the message with an Accept
-                    let message = T::decode_into_message(payload.reader())
+                    let message = T::decode_into_message(payload.into_reader())
                         .map_err(|_| Self::TransferError::MessageDecodeError)?;
 
                     (message, Some(DeliveryState::Accepted(Accepted {})))
@@ -218,9 +221,9 @@ where
                 // disposition from the sender.
                 ReceiverSettleMode::Second => {
                     // Add to unsettled map
-                    let section_offset = rfind_offset_of_complete_message(payload.as_ref())
+                    let section_offset = rfind_offset_of_complete_message(&payload)
                         .ok_or(Self::TransferError::MessageDecodeError)?;
-                    let message = T::decode_into_message(payload.reader())
+                    let message = T::decode_into_message(payload.into_reader())
                         .map_err(|_| Self::TransferError::MessageDecodeError)?;
                     let section_number = message.sections();
 
@@ -320,13 +323,16 @@ where
 }
 
 /// Finds offset of a complete message
-fn rfind_offset_of_complete_message(bytes: &[u8]) -> Option<u64> {
+fn rfind_offset_of_complete_message<'a, B>(bytes: &'a B) -> Option<u64> 
+where
+    B: AsByteIterator<'a>,
+{
     // For a complete message, the only need is to check Footer or Body
-
-    let len = bytes.len();
-    let mut iter = bytes
-        .iter()
-        .zip(bytes.iter().skip(1).zip(bytes.iter().skip(2)));
+    let b0 = bytes.as_byte_iterator();
+    let b1 = bytes.as_byte_iterator().skip(1);
+    let b2 = bytes.as_byte_iterator().skip(2);
+    let len = b0.len();
+    let mut iter = b0.zip(b1.zip(b2));
 
     iter.rposition(|(&b0, (&b1, &b2))| {
         matches!(