Implement Recoverable connection for AMQP operations. (Azure#2657)

This PR does three different things:

It renames the ConnectionManager type to RecoverableConnection to better reflect what is going on.
It splits out the authorization functionality from ConnectionManager into a new struct named Authorizer to simplify the logic in connection_manager.rs
It consolidates all the AMQP active protocol interactions to RecoverableConnection. This was done to (a) simplify the AMQP interactions for each of the specific protocol elements and (b) lay the foundation for connection, session, and link recovery. See the rather large comment on the RecoverableConnection for more details on how this all works and fits together.
There are a few bonus changes in this PR as well:

Adds a send_ref function to the AmqpSender struct which allows sending a reference to a message (saving a message clone).
AmqpOrderedMap::iter() passes references to the inner elements to be consistent with the standard collection iter() method.

Author: LarryOsterman

sdk/core/azure_core_amqp/CHANGELOG.md

@@ -6,6 +6,10 @@
 
 ### Breaking Changes
 
+- `AmqpClaimsBasedSecurity` now takes ownership of the associated session rather than simply referencing the associated session. This means that all CBS authentication operations should be performed on dedicated AmqpSession objects.
+
+- `AmqpOrderedMap::iter` now iterates over references to key and value, not clones of the key and value, thus eliminating unnecessary clones.
+
 ### Bugs Fixed
 
 ### Other Changes

sdk/core/azure_core_amqp/src/cbs.rs

@@ -1,18 +1,17 @@
 // Copyright (c) Microsoft Corporation. All Rights reserved
 // Licensed under the MIT license.
 
-use async_trait::async_trait;
-use azure_core::error::Result;
-
 use super::session::AmqpSession;
+use azure_core::{credentials::Secret, error::Result};
 
 #[cfg(all(feature = "fe2o3_amqp", not(target_arch = "wasm32")))]
-type CbsImplementation<'a> = super::fe2o3::cbs::Fe2o3ClaimsBasedSecurity<'a>;
+type CbsImplementation = super::fe2o3::cbs::Fe2o3ClaimsBasedSecurity;
 
 #[cfg(any(not(any(feature = "fe2o3_amqp")), target_arch = "wasm32"))]
-type CbsImplementation<'a> = super::noop::NoopAmqpClaimsBasedSecurity<'a>;
+type CbsImplementation = super::noop::NoopAmqpClaimsBasedSecurity;
 
-#[async_trait]
+#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
 pub trait AmqpClaimsBasedSecurityApis {
     /// Asynchronously attaches the Claims-Based Security (CBS) node to the AMQP session.
     ///
@@ -52,30 +51,31 @@ pub trait AmqpClaimsBasedSecurityApis {
         &self,
         path: String,
         token_type: Option<String>,
-        secret: String,
+        secret: &Secret,
         expires_on: time::OffsetDateTime,
     ) -> Result<()>;
 }
 
-pub struct AmqpClaimsBasedSecurity<'a> {
-    implementation: CbsImplementation<'a>,
+pub struct AmqpClaimsBasedSecurity {
+    implementation: CbsImplementation,
 }
 
-impl<'a> AmqpClaimsBasedSecurity<'a> {
-    pub fn new(session: &'a AmqpSession) -> Result<Self> {
+impl AmqpClaimsBasedSecurity {
+    pub fn new(session: AmqpSession) -> Result<Self> {
         Ok(Self {
             implementation: CbsImplementation::new(session)?,
         })
     }
 }
 
-#[async_trait]
-impl AmqpClaimsBasedSecurityApis for AmqpClaimsBasedSecurity<'_> {
+#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
+impl AmqpClaimsBasedSecurityApis for AmqpClaimsBasedSecurity {
     async fn authorize_path(
         &self,
         path: String,
         token_type: Option<String>,
-        secret: String,
+        secret: &Secret,
         expires_on: time::OffsetDateTime,
     ) -> Result<()> {
         self.implementation

sdk/core/azure_core_amqp/src/connection.rs

@@ -2,7 +2,6 @@
 // Licensed under the MIT license.
 
 use super::value::{AmqpOrderedMap, AmqpSymbol, AmqpValue};
-use async_trait::async_trait;
 use azure_core::{error::Result, http::Url};
 use std::fmt::Debug;
 use time::Duration;
@@ -40,7 +39,8 @@ pub struct AmqpConnectionOptions {
 
 impl AmqpConnectionOptions {}
 
-#[async_trait]
+#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
 pub trait AmqpConnectionApis {
     async fn open(
         &self,
@@ -62,7 +62,8 @@ pub struct AmqpConnection {
     pub(crate) implementation: ConnectionImplementation,
 }
 
-#[async_trait]
+#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
 impl AmqpConnectionApis for AmqpConnection {
     async fn open(
         &self,

sdk/core/azure_core_amqp/src/fe2o3/cbs.rs

@@ -3,22 +3,21 @@
 // cspell:: words amqp servicebus sastoken
 
 use crate::{cbs::AmqpClaimsBasedSecurityApis, session::AmqpSession, AmqpError};
-use async_trait::async_trait;
-use azure_core::error::Result;
+use azure_core::{credentials::Secret, error::Result};
 use fe2o3_amqp_cbs::token::CbsToken;
 use fe2o3_amqp_types::primitives::Timestamp;
 use std::borrow::BorrowMut;
 use std::sync::OnceLock;
 use tokio::sync::Mutex;
 use tracing::{debug, trace};
 
-pub(crate) struct Fe2o3ClaimsBasedSecurity<'a> {
+pub(crate) struct Fe2o3ClaimsBasedSecurity {
     cbs: OnceLock<Mutex<fe2o3_amqp_cbs::client::CbsClient>>,
-    session: &'a AmqpSession,
+    session: AmqpSession,
 }
 
-impl<'a> Fe2o3ClaimsBasedSecurity<'a> {
-    pub fn new(session: &'a AmqpSession) -> Result<Self> {
+impl Fe2o3ClaimsBasedSecurity {
+    pub fn new(session: AmqpSession) -> Result<Self> {
         Ok(Self {
             cbs: OnceLock::new(),
             session,
@@ -45,16 +44,17 @@ impl<'a> Fe2o3ClaimsBasedSecurity<'a> {
     }
 }
 
-impl Fe2o3ClaimsBasedSecurity<'_> {}
+impl Fe2o3ClaimsBasedSecurity {}
 
-impl Drop for Fe2o3ClaimsBasedSecurity<'_> {
+impl Drop for Fe2o3ClaimsBasedSecurity {
     fn drop(&mut self) {
         debug!("Dropping Fe2o3ClaimsBasedSecurity.");
     }
 }
 
-#[async_trait]
-impl AmqpClaimsBasedSecurityApis for Fe2o3ClaimsBasedSecurity<'_> {
+#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
+impl AmqpClaimsBasedSecurityApis for Fe2o3ClaimsBasedSecurity {
     async fn attach(&self) -> Result<()> {
         let session = self.session.implementation.get()?;
         let mut session = session.lock().await;
@@ -80,7 +80,7 @@ impl AmqpClaimsBasedSecurityApis for Fe2o3ClaimsBasedSecurity<'_> {
         &self,
         path: String,
         token_type: Option<String>,
-        secret: String,
+        secret: &Secret,
         expires_at: time::OffsetDateTime,
     ) -> Result<()> {
         trace!(
@@ -89,7 +89,7 @@ impl AmqpClaimsBasedSecurityApis for Fe2o3ClaimsBasedSecurity<'_> {
             expires_at
         );
         let cbs_token = CbsToken::new(
-            secret,
+            secret.secret(),
             token_type.unwrap_or("jwt".to_string()),
             Some(Timestamp::from(
                 expires_at

sdk/core/azure_core_amqp/src/fe2o3/connection.rs

@@ -5,7 +5,6 @@ use super::error::{Fe2o3ConnectionError, Fe2o3ConnectionOpenError, Fe2o3Transpor
 use crate::connection::{AmqpConnectionApis, AmqpConnectionOptions};
 use crate::error::AmqpErrorKind;
 use crate::value::{AmqpOrderedMap, AmqpSymbol, AmqpValue};
-use async_trait::async_trait;
 use azure_core::{http::Url, Result};
 use fe2o3_amqp::connection::ConnectionHandle;
 use std::{borrow::BorrowMut, sync::OnceLock};
@@ -45,7 +44,8 @@ impl Drop for Fe2o3AmqpConnection {
     }
 }
 
-#[async_trait]
+#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
 impl AmqpConnectionApis for Fe2o3AmqpConnection {
     async fn open(
         &self,

sdk/core/azure_core_amqp/src/fe2o3/management.rs

@@ -11,10 +11,8 @@ use crate::{
     value::{AmqpOrderedMap, AmqpValue},
     AmqpError,
 };
-use async_trait::async_trait;
 use azure_core::{credentials::AccessToken, Result};
 use fe2o3_amqp_management::operations::ReadResponse;
-use fe2o3_amqp_types::{messaging::ApplicationProperties, primitives::SimpleValue};
 use std::sync::{Arc, OnceLock};
 use tokio::sync::Mutex;
 use tracing::debug;
@@ -64,7 +62,8 @@ impl Fe2o3AmqpManagement {
     }
 }
 
-#[async_trait]
+#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
 impl AmqpManagementApis for Fe2o3AmqpManagement {
     async fn attach(&self) -> Result<()> {
         let management = fe2o3_amqp_management::client::MgmtClient::builder()
@@ -185,21 +184,24 @@ impl fe2o3_amqp_management::Request for WithApplicationPropertiesRequest<'_> {
     fn encode_application_properties(
         &mut self,
     ) -> Option<fe2o3_amqp_types::messaging::ApplicationProperties> {
-        let builder = ApplicationProperties::builder();
-        let builder = self
+        let mut fe2o3_application_properties = self
             .application_properties
             .iter()
-            .fold(builder, |builder, (key, value)| {
-                builder.insert(
+            .map(|(key, value)| {
+                (
                     key.clone(),
                     Into::<fe2o3_amqp_types::primitives::SimpleValue>::into(value),
                 )
             })
-            .insert(
-                "security_token",
-                Into::<SimpleValue>::into(self.access_token.token.secret()),
-            );
-        Some(builder.build())
+            .collect::<fe2o3_amqp_types::primitives::OrderedMap<_, _>>();
+        fe2o3_application_properties.insert(
+            "security_token".to_string(),
+            fe2o3_amqp_types::primitives::SimpleValue::from(self.access_token.token.secret()),
+        );
+
+        Some(fe2o3_amqp_types::messaging::ApplicationProperties(
+            fe2o3_application_properties,
+        ))
     }
     fn encode_body(self) -> Self::Body {}
 }

sdk/core/azure_core_amqp/src/fe2o3/messaging/message_fields.rs

@@ -253,6 +253,19 @@ impl From<AmqpAnnotationKey> for fe2o3_amqp_types::messaging::annotations::Owned
     }
 }
 
+impl From<&AmqpAnnotationKey> for fe2o3_amqp_types::messaging::annotations::OwnedKey {
+    fn from(key: &AmqpAnnotationKey) -> Self {
+        match key {
+            AmqpAnnotationKey::Ulong(key) => {
+                fe2o3_amqp_types::messaging::annotations::OwnedKey::Ulong(*key)
+            }
+            AmqpAnnotationKey::Symbol(key) => {
+                fe2o3_amqp_types::messaging::annotations::OwnedKey::Symbol(key.into())
+            }
+        }
+    }
+}
+
 impl From<&fe2o3_amqp_types::messaging::annotations::OwnedKey>
     for crate::messaging::AmqpAnnotationKey
 {

sdk/core/azure_core_amqp/src/fe2o3/messaging/message_source.rs

@@ -32,7 +32,7 @@ impl From<AmqpSource> for fe2o3_amqp_types::messaging::Source {
             builder = builder.dynamic_node_properties(
                 dynamic_node_properties
                     .iter()
-                    .map(|(k, v)| (k.0.into(), v.into()))
+                    .map(|(k, v)| (k.into(), v.into()))
                     .collect::<fe2o3_amqp_types::definitions::Fields>(),
             );
         }
@@ -43,7 +43,7 @@ impl From<AmqpSource> for fe2o3_amqp_types::messaging::Source {
             builder = builder.filter(
                 filter
                     .iter()
-                    .map(|(k, v)| (k.0.into(), v.into()))
+                    .map(|(k, v)| (k.into(), v.into()))
                     .collect::<fe2o3_amqp_types::messaging::FilterSet>(),
             );
         }

sdk/core/azure_core_amqp/src/fe2o3/messaging/messaging_types.rs

@@ -3,15 +3,19 @@
 
 use std::sync::OnceLock;
 
+use typespec_macros::SafeDebug;
+
 use crate::messaging::{
     AmqpDelivery, AmqpDeliveryApis, AmqpMessage, AmqpOutcome, DeliveryNumber, DeliveryTag,
     DistributionMode, TerminusDurability, TerminusExpiryPolicy,
 };
-
+#[derive(SafeDebug)]
+#[safe(true)]
 pub(crate) struct Fe2o3AmqpDelivery {
     pub(crate) delivery: fe2o3_amqp::link::delivery::Delivery<
         fe2o3_amqp_types::messaging::Body<fe2o3_amqp_types::primitives::Value>,
     >,
+    #[safe(false)]
     message: OnceLock<AmqpMessage>,
     delivery_tag: OnceLock<DeliveryTag>,
 }

sdk/core/azure_core_amqp/src/fe2o3/receiver.rs

@@ -8,7 +8,6 @@ use crate::{
     session::AmqpSession,
     AmqpError,
 };
-use async_trait::async_trait;
 use azure_core::error::Result;
 use std::borrow::BorrowMut;
 use std::sync::OnceLock;
@@ -44,7 +43,8 @@ impl From<&fe2o3_amqp::link::receiver::CreditMode> for ReceiverCreditMode {
     }
 }
 
-#[async_trait]
+#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
 impl AmqpReceiverApis for Fe2o3AmqpReceiver {
     async fn attach(
         &self,