Add proper delivery handling

CHANGES.md

@@ -1,5 +1,9 @@
 # Changes
 
+## [1.1.0] - 2024-03-04
+
+* Add proper delivery handling
+
 ## [codec-0.9.2] - 2024-02-01
 
 * Add more buffer length checks

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ntex-amqp"
-version = "1.0.2"
+version = "1.1.0"
 authors = ["ntex contributors <[email protected]>"]
 description = "AMQP 1.0 Client/Server framework"
 documentation = "https://docs.rs/ntex-amqp"

src/delivery.rs

@@ -0,0 +1,294 @@
+use std::cell::Cell as StdCell;
+
+use ntex::{channel::pool, util::Bytes};
+use ntex_amqp_codec::protocol::{
+    DeliveryNumber, DeliveryState, Disposition, DispositionInner, Error, ErrorCondition, Rejected,
+    Role, TransferBody,
+};
+use ntex_amqp_codec::types::{Str, Symbol};
+
+use crate::session::Session;
+use crate::{cell::Cell, error::AmqpProtocolError, sndlink::SenderLinkInner};
+
+bitflags::bitflags! {
+    #[derive(Copy, Clone, Debug)]
+    struct Flags: u8 {
+        const SENDER         = 0b0000_0001;
+        const LOCAL_SETTLED  = 0b0000_0100;
+        const REMOTE_SETTLED = 0b0000_1000;
+    }
+}
+
+pub struct Delivery {
+    id: DeliveryNumber,
+    session: Session,
+    flags: StdCell<Flags>,
+}
+
+#[derive(Default, Debug)]
+pub(crate) struct DeliveryInner {
+    settled: bool,
+    state: Option<DeliveryState>,
+    error: Option<AmqpProtocolError>,
+    tx: Option<pool::Sender<()>>,
+}
+
+impl Delivery {
+    pub fn remote_state(&self) -> Option<DeliveryState> {
+        if let Some(inner) = self
+            .session
+            .inner
+            .get_mut()
+            .unsettled_deliveries
+            .get_mut(&self.id)
+        {
+            inner.state.clone()
+        } else {
+            None
+        }
+    }
+
+    pub fn is_remote_settled(&self) -> bool {
+        self.is_set(Flags::REMOTE_SETTLED)
+    }
+
+    pub fn settle(&mut self, state: DeliveryState) {
+        // remote side is settled, not need to send disposition
+        if self.is_set(Flags::REMOTE_SETTLED) {
+            return;
+        }
+
+        if !self.is_set(Flags::LOCAL_SETTLED) {
+            self.set_flag(Flags::LOCAL_SETTLED);
+
+            let disp = Disposition(Box::new(DispositionInner {
+                role: if self.is_set(Flags::SENDER) {
+                    Role::Sender
+                } else {
+                    Role::Receiver
+                },
+                first: self.id,
+                last: None,
+                settled: true,
+                state: Some(state),
+                batchable: false,
+            }));
+            self.session.inner.get_mut().post_frame(disp.into());
+        }
+    }
+
+    pub fn update_state(&mut self, state: DeliveryState) {
+        // remote side is settled, not need to send disposition
+        if self.is_set(Flags::REMOTE_SETTLED) || self.is_set(Flags::LOCAL_SETTLED) {
+            return;
+        }
+
+        let disp = Disposition(Box::new(DispositionInner {
+            role: if self.is_set(Flags::SENDER) {
+                Role::Sender
+            } else {
+                Role::Receiver
+            },
+            first: self.id,
+            last: None,
+            settled: false,
+            state: Some(state),
+            batchable: false,
+        }));
+        self.session.inner.get_mut().post_frame(disp.into());
+    }
+
+    fn is_set(&self, flag: Flags) -> bool {
+        self.flags.get().contains(flag)
+    }
+
+    fn set_flag(&self, flag: Flags) {
+        let mut flags = self.flags.get();
+        flags.insert(flag);
+        self.flags.set(flags);
+    }
+
+    pub async fn wait(&self) -> Result<Option<DeliveryState>, AmqpProtocolError> {
+        let rx = if let Some(inner) = self
+            .session
+            .inner
+            .get_mut()
+            .unsettled_deliveries
+            .get_mut(&self.id)
+        {
+            if let Some(st) = self.check_inner(inner) {
+                return st;
+            }
+
+            let (tx, rx) = self.session.inner.get_ref().pool_notify.channel();
+            inner.tx = Some(tx);
+            rx
+        } else {
+            return Ok(None);
+        };
+        if rx.await.is_err() {
+            return Err(AmqpProtocolError::ConnectionDropped);
+        }
+
+        if let Some(inner) = self
+            .session
+            .inner
+            .get_mut()
+            .unsettled_deliveries
+            .get_mut(&self.id)
+        {
+            if inner.settled {
+                self.set_flag(Flags::REMOTE_SETTLED);
+            }
+            if let Some(st) = self.check_inner(inner) {
+                return st;
+            }
+        }
+        Ok(None)
+    }
+
+    fn check_inner(
+        &self,
+        inner: &mut DeliveryInner,
+    ) -> Option<Result<Option<DeliveryState>, AmqpProtocolError>> {
+        if let Some(ref st) = inner.state {
+            if matches!(st, DeliveryState::Modified(..)) {
+                // non terminal state
+                Some(Ok(Some(inner.state.take().unwrap())))
+            } else {
+                // return clone of terminal state
+                Some(Ok(Some(st.clone())))
+            }
+        } else if let Some(ref err) = inner.error {
+            Some(Err(err.clone()))
+        } else {
+            None
+        }
+    }
+}
+
+impl Drop for Delivery {
+    fn drop(&mut self) {
+        let inner = self.session.inner.get_mut();
+
+        if inner.unsettled_deliveries.contains_key(&self.id) {
+            if !self.is_set(Flags::REMOTE_SETTLED) && !self.is_set(Flags::LOCAL_SETTLED) {
+                let err = Error::build()
+                    .condition(ErrorCondition::Custom(Symbol(Str::Static(
+                        "Internal error",
+                    ))))
+                    .finish();
+
+                let disp = Disposition(Box::new(DispositionInner {
+                    role: if self.is_set(Flags::SENDER) {
+                        Role::Sender
+                    } else {
+                        Role::Receiver
+                    },
+                    first: self.id,
+                    last: None,
+                    settled: true,
+                    state: Some(DeliveryState::Rejected(Rejected { error: Some(err) })),
+                    batchable: false,
+                }));
+                inner.post_frame(disp.into());
+            }
+
+            inner.unsettled_deliveries.remove(&self.id);
+        }
+    }
+}
+
+impl DeliveryInner {
+    pub(crate) fn new() -> Self {
+        Self {
+            tx: None,
+            state: None,
+            error: None,
+            settled: false,
+        }
+    }
+
+    pub(crate) fn set_error(&mut self, error: AmqpProtocolError) {
+        self.error = Some(error);
+        if let Some(tx) = self.tx.take() {
+            let _ = tx.send(());
+        }
+    }
+
+    pub(crate) fn handle_disposition(&mut self, disp: Disposition) {
+        if disp.settled() {
+            self.settled = true;
+        }
+        if let Some(state) = disp.state() {
+            self.state = Some(state.clone());
+        }
+        if let Some(tx) = self.tx.take() {
+            let _ = tx.send(());
+        }
+    }
+}
+
+impl Drop for DeliveryInner {
+    fn drop(&mut self) {
+        if let Some(tx) = self.tx.take() {
+            let _ = tx.send(());
+        }
+    }
+}
+
+pub struct DeliveryBuilder {
+    tag: Option<Bytes>,
+    settled: bool,
+    data: TransferBody,
+    sender: Cell<SenderLinkInner>,
+}
+
+impl DeliveryBuilder {
+    pub(crate) fn new(data: TransferBody, sender: Cell<SenderLinkInner>) -> Self {
+        Self {
+            tag: None,
+            settled: false,
+            data,
+            sender,
+        }
+    }
+
+    pub fn tag(mut self, tag: Bytes) -> Self {
+        self.tag = Some(tag);
+        self
+    }
+
+    pub fn settled(mut self) -> Self {
+        self.settled = true;
+        self
+    }
+
+    pub async fn send(self) -> Result<Delivery, AmqpProtocolError> {
+        let inner = self.sender.get_ref();
+
+        if let Some(ref err) = inner.error {
+            Err(err.clone())
+        } else {
+            if inner
+                .max_message_size
+                .map(|l| self.data.len() > l as usize)
+                .unwrap_or_default()
+            {
+                Err(AmqpProtocolError::BodyTooLarge)
+            } else {
+                let id = self.sender.get_mut().send(self.data, self.tag).await?;
+
+                Ok(Delivery {
+                    id,
+                    session: self.sender.get_ref().session.clone(),
+                    flags: StdCell::new(if self.settled {
+                        Flags::SENDER | Flags::LOCAL_SETTLED
+                    } else {
+                        Flags::SENDER
+                    }),
+                })
+            }
+        }
+    }
+}

src/lib.rs

@@ -1,4 +1,4 @@
-#![deny(rust_2018_idioms, warnings, unreachable_pub)]
+//#![deny(rust_2018_idioms, warnings, unreachable_pub)]
 #![allow(clippy::type_complexity, clippy::let_underscore_future)]
 
 #[macro_use]
@@ -13,6 +13,7 @@ pub mod client;
 mod connection;
 mod control;
 mod default;
+mod delivery;
 mod dispatcher;
 pub mod error;
 pub mod error_code;
@@ -26,6 +27,7 @@ pub mod types;
 
 pub use self::connection::{Connection, ConnectionRef};
 pub use self::control::{ControlFrame, ControlFrameKind};
+pub use self::delivery::{Delivery, DeliveryBuilder};
 pub use self::rcvlink::{ReceiverLink, ReceiverLinkBuilder};
 pub use self::session::Session;
 pub use self::sndlink::{SenderLink, SenderLinkBuilder};

src/rcvlink.rs

@@ -6,9 +6,8 @@ use std::{
 use ntex::util::{ByteString, BytesMut, PoolRef, Stream};
 use ntex::{channel::oneshot, task::LocalWaker};
 use ntex_amqp_codec::protocol::{
-    self as codec, Attach, DeliveryNumber, Disposition, Error, Handle, LinkError,
-    ReceiverSettleMode, Role, SenderSettleMode, Source, TerminusDurability, TerminusExpiryPolicy,
-    Transfer, TransferBody,
+    self as codec, Attach, Disposition, Error, Handle, LinkError, ReceiverSettleMode, Role,
+    SenderSettleMode, Source, TerminusDurability, TerminusExpiryPolicy, Transfer, TransferBody,
 };
 use ntex_amqp_codec::types::{Symbol, Variant};
 use ntex_amqp_codec::Encode;
@@ -133,14 +132,6 @@ impl ReceiverLink {
             .post_frame(disp.into());
     }
 
-    /// Wait for disposition with specified number
-    pub fn wait_disposition(
-        &self,
-        id: DeliveryNumber,
-    ) -> impl Future<Output = Result<Disposition, AmqpProtocolError>> {
-        self.inner.get_mut().session.wait_disposition(id)
-    }
-
     pub fn close(&self) -> impl Future<Output = Result<(), AmqpProtocolError>> {
         self.inner.get_mut().close(None)
     }