[EventHubs] Retry a CBS authentication failure once on a fresh connection

sdk/core/azure-core-amqp/inc/azure/core/amqp/internal/claims_based_security.hpp

@@ -7,8 +7,10 @@
 
 #include <azure/core/context.hpp>
 
+#include <exception>
 #include <stdexcept>
 #include <string>
+#include <utility>
 
 namespace Azure { namespace Core { namespace Amqp { namespace _detail {
   class ClaimsBasedSecurityImpl;
@@ -53,6 +55,23 @@ namespace Azure { namespace Core { namespace Amqp { namespace _detail {
     CbsOpenResult Result;
   };
 
+#if ENABLE_UAMQP
+  /** @brief Identifies a failed uAMQP CBS put-token operation. */
+  class CbsPutTokenFailedException final : public std::runtime_error {
+  public:
+    CbsPutTokenFailedException(std::exception_ptr original, std::string const& what)
+        : std::runtime_error(what), m_original{std::move(original)}
+    {
+    }
+
+    std::exception_ptr GetOriginal() const { return m_original; }
+    [[noreturn]] void RethrowOriginal() const { std::rethrow_exception(m_original); }
+
+  private:
+    std::exception_ptr m_original;
+  };
+#endif
+
   enum class CbsTokenType
   {
     Invalid,

sdk/core/azure-core-amqp/src/amqp/connection.cpp

@@ -195,9 +195,14 @@ namespace Azure { namespace Core { namespace Amqp { namespace _detail {
             = claimsBasedSecurity->PutToken(tokenType, audienceUrl, token, expiresOn, context);
         if (std::get<0>(result) != CbsOperationResult::Ok)
         {
-          throw Azure::Core::Credentials::AuthenticationException(
+          auto failure = Azure::Core::Credentials::AuthenticationException(
               "Could not authenticate client. Error Status: " + std::to_string(std::get<1>(result))
               + " reason: " + std::get<2>(result));
+#if ENABLE_UAMQP
+          throw CbsPutTokenFailedException(std::make_exception_ptr(failure), failure.what());
+#else
+          throw failure;
+#endif
         }
         Log::Stream(Logger::Level::Verbose) << "Close CBS object";
         claimsBasedSecurity->Close(context);

sdk/core/azure-core-amqp/test/ut/mock_amqp_server.hpp

@@ -17,6 +17,7 @@
 #include <azure/core/amqp/internal/session.hpp>
 
 #include <memory>
+#include <utility>
 
 #include <gtest/gtest.h>
 
@@ -62,9 +63,20 @@ namespace Azure { namespace Core { namespace Amqp { namespace Tests {
       {
       }
 
+      virtual ~MockServiceEndpoint() = default;
+
       const std::string& GetName() const { return m_name; }
 
-      bool OnLinkAttached(
+      void DetachLink(
+          Azure::Core::Amqp::_internal::Session const& session,
+          Azure::Core::Amqp::_internal::LinkEndpoint const& linkEndpoint,
+          bool closeLink,
+          Models::_internal::AmqpError const& error) const
+      {
+        session.SendDetach(linkEndpoint, closeLink, error);
+      }
+
+      virtual bool OnLinkAttached(
           Azure::Core::Amqp::_internal::Session const& session,
           std::string const& linkName,
           Azure::Core::Amqp::_internal::LinkEndpoint& linkEndpoint,
@@ -544,24 +556,27 @@ namespace Azure { namespace Core { namespace Amqp { namespace Tests {
     public:
       AmqpServerMock(
           std::string name = testing::UnitTest::GetInstance()->current_test_info()->name())
-          : m_connectionId{"Mock Server for " + name}, m_testPort{FindAvailableSocket()}
+          : AmqpServerMock(FindAvailableSocket(), std::move(name), true)
       {
-        // Every server mock has CBS endpoint support
-        MockServiceEndpointOptions options;
-        options.EnableTrace = m_enableTrace;
-        options.ListenerContext = m_listenerContext;
-        AddServiceEndpoint(std::make_shared<AmqpClaimBasedSecurity>(options));
       }
       AmqpServerMock(
           uint16_t listeningPort,
           std::string name = testing::UnitTest::GetInstance()->current_test_info()->name())
+          : AmqpServerMock(listeningPort, std::move(name), true)
+      {
+      }
+
+      AmqpServerMock(uint16_t listeningPort, std::string name, bool addCbsEndpoint)
           : m_connectionId{"Mock Server for " + name}, m_testPort{listeningPort}
       {
-        // Every server mock has CBS endpoint support
-        MockServiceEndpointOptions options;
-        options.EnableTrace = m_enableTrace;
-        options.ListenerContext = m_listenerContext;
-        AddServiceEndpoint(std::make_shared<AmqpClaimBasedSecurity>(options));
+        if (addCbsEndpoint)
+        {
+          // Every server mock has CBS endpoint support
+          MockServiceEndpointOptions options;
+          options.EnableTrace = m_enableTrace;
+          options.ListenerContext = m_listenerContext;
+          AddServiceEndpoint(std::make_shared<AmqpClaimBasedSecurity>(options));
+        }
       }
 
       virtual ~AmqpServerMock()
@@ -579,6 +594,7 @@ namespace Azure { namespace Core { namespace Amqp { namespace Tests {
       }
 
       uint16_t GetPort() const { return m_testPort; }
+      std::size_t GetConnectionCount() const { return m_connections.size(); }
       Azure::Core::Context& GetListenerContext() { return m_listenerContext; }
 
       void StartListening()

sdk/eventhubs/azure-messaging-eventhubs/CHANGELOG.md

@@ -11,6 +11,7 @@
 
 ### Bugs Fixed
 
+- The uAMQP backend now retries a CBS PutToken failure or exact unauthorized send or receive error once on a fresh connection, subject to `RetryOptions`, while preserving the final error.
 - [[#7254]](https://github.com/Azure/azure-sdk-for-cpp/issues/7254) A teardown of the cached sender no longer runs while another thread sends on that sender. `ProducerClient::Send` gives each attempt a copy of the sender, and a failed attempt on one thread closed the object that a second thread was using. On the Rust AMQP backend that close frees the sender, so the race was a use after free. Each partition now has a guard that lets sends run at the same time and makes a teardown wait for the sends in flight. `ProducerClient::Close` uses the same guard, and it now logs a failed close and continues instead of leaving the other objects open.
 - [[#7254]](https://github.com/Azure/azure-sdk-for-cpp/issues/7254) `ProducerClient::CreateBatch` now builds a new sender when it cannot read the maximum message size. The client caches a sender for each partition, and a cached sender holds a link that the service detaches after 30 idle minutes. The size of a batch comes from the attached link, so this call was the first one to touch the dead link, and it threw. The `Send(EventData)` overloads go through this call, so the whole producer failed after an idle period even though `Send` builds a new sender on each attempt.
 - [[#7254]](https://github.com/Azure/azure-sdk-for-cpp/issues/7254) Updated producer retries to honor `EventHubsException::IsTransient`, treat empty AMQP error conditions as transient, stop immediately for unknown and known non-transient failures, preserve bounded retries for AMQP runtime failures, and make backoff cancellable through `Azure::Core::Context`. Retry accounting now always performs the initial attempt and treats `MaxRetries` as additional retry attempts.

sdk/eventhubs/azure-messaging-eventhubs/inc/azure/messaging/eventhubs/consumer_client.hpp

@@ -15,6 +15,9 @@
 #include <azure/core/diagnostics/logger.hpp>
 #include <azure/core/http/policies/policy.hpp>
 #include <azure/core/internal/diagnostics/log.hpp>
+
+#include <memory>
+#include <vector>
 namespace Azure { namespace Messaging { namespace EventHubs {
   namespace _detail {
     class EventHubsPropertiesClient;
@@ -199,6 +202,12 @@ namespace Azure { namespace Messaging { namespace EventHubs {
         Core::Context const& context = {});
 
   private:
+#if ENABLE_UAMQP
+    std::mutex m_partitionClientStatesLock;
+    std::vector<std::shared_ptr<_detail::PartitionClientState>> m_partitionClientStates;
+    bool m_partitionClientStatesClosing{false};
+#endif
+
     /// The connection string for the Event Hubs namespace
     std::string m_connectionString;
 

sdk/eventhubs/azure-messaging-eventhubs/inc/azure/messaging/eventhubs/partition_client.hpp

@@ -11,10 +11,15 @@
 #include <azure/core/http/policies/policy.hpp>
 #include <azure/core/nullable.hpp>
 
+#include <memory>
+
 namespace Azure { namespace Messaging { namespace EventHubs {
   namespace _detail {
     class PartitionClientFactory;
-  }
+    struct PartitionClientState;
+  } // namespace _detail
+
+  class ConsumerClient;
   /**brief PartitionClientOptions provides options for the ConsumerClient::CreatePartitionClient
    * function.
    */
@@ -89,6 +94,15 @@ namespace Azure { namespace Messaging { namespace EventHubs {
 
   private:
     friend class _detail::PartitionClientFactory;
+    friend class ConsumerClient;
+
+#if ENABLE_UAMQP
+    std::shared_ptr<_detail::PartitionClientState> m_state;
+    explicit PartitionClient(std::shared_ptr<_detail::PartitionClientState> state);
+    std::shared_ptr<_detail::PartitionClientState> GetState() const { return m_state; }
+#endif
+
+#if ENABLE_RUST_AMQP
     /// The message receiver used to receive events from the partition.
     Azure::Core::Amqp::_internal::MessageReceiver m_receiver;
 
@@ -115,7 +129,9 @@ namespace Azure { namespace Messaging { namespace EventHubs {
      * response to being throttled or encountering a transient error.
      */
     Azure::Core::Http::Policies::RetryOptions m_retryOptions{};
+#endif
 
+#if ENABLE_RUST_AMQP
     /** Creates a new PartitionClient
      *
      * @param messageReceiver Message Receiver for the partition client.
@@ -133,10 +149,15 @@ namespace Azure { namespace Messaging { namespace EventHubs {
         std::string receiverName,
         PartitionClientOptions options,
         Core::Http::Policies::RetryOptions retryOptions);
+#endif
 
+#if ENABLE_RUST_AMQP || ENABLE_UAMQP
     /// Closes the faulted receiver and attaches a new one starting after the last offset.
     void RebuildReceiver(Core::Context const& context);
+#endif
 
+#if ENABLE_RUST_AMQP
     std::string GetStartExpression(Models::StartPosition const& startPosition);
+#endif
   };
 }}} // namespace Azure::Messaging::EventHubs

sdk/eventhubs/azure-messaging-eventhubs/inc/azure/messaging/eventhubs/producer_client.hpp

@@ -195,6 +195,8 @@ namespace Azure { namespace Messaging { namespace EventHubs {
     }
 
   private:
+    struct ProducerCallState;
+
     /// The connection string for the Event Hubs namespace
     std::string m_connectionString;
 
@@ -259,6 +261,16 @@ namespace Azure { namespace Messaging { namespace EventHubs {
     // Ensure that a message sender for the specified partition has been created.
     void EnsureSender(std::string const& partitionId, Azure::Core::Context const& context);
 
+    EventDataBatch CreateBatch(
+        EventDataBatchOptions const& options,
+        Azure::Core::Context const& context,
+        ProducerCallState& callState);
+
+    void Send(
+        EventDataBatch const& eventDataBatch,
+        Core::Context const& context,
+        ProducerCallState& callState);
+
     // Calls EnsureSender, and discards a failed attach unless the context is cancelled.
     void EnsureSenderOrInvalidate(
         std::string const& partitionId,
@@ -268,7 +280,8 @@ namespace Azure { namespace Messaging { namespace EventHubs {
     // security open reported CbsOpenResult::Error.
     void EstablishSenderWithRetry(
         std::string const& partitionId,
-        Azure::Core::Context const& context);
+        Azure::Core::Context const& context,
+        ProducerCallState& callState);
 
     // Discards the sender, session, and connection for the partition. A null generation,
     // as Close passes, removes whatever is present regardless of generation.

sdk/eventhubs/azure-messaging-eventhubs/src/consumer_client.cpp

@@ -67,6 +67,18 @@ namespace Azure { namespace Messaging { namespace EventHubs {
   void ConsumerClient::Close(Azure::Core::Context const& context)
   {
     Log::Stream(Logger::Level::Verbose) << "Close consumer client.";
+#if ENABLE_UAMQP
+    std::vector<std::shared_ptr<_detail::PartitionClientState>> partitionClientStates;
+    {
+      std::lock_guard<std::mutex> lock(m_partitionClientStatesLock);
+      if (m_partitionClientStatesClosing)
+      {
+        return;
+      }
+      m_partitionClientStatesClosing = true;
+      partitionClientStates = std::move(m_partitionClientStates);
+    }
+#endif
     {
       std::unique_lock<std::mutex> lock(m_propertiesClientLock);
       if (m_propertiesClient)
@@ -83,6 +95,12 @@ namespace Azure { namespace Messaging { namespace EventHubs {
         m_propertiesClient.reset();
       }
     }
+#if ENABLE_UAMQP
+    for (auto const& state : partitionClientStates)
+    {
+      _detail::ClosePartitionClientState(state, context);
+    }
+#endif
     Log::Stream(Logger::Level::Verbose) << "Closing message receivers.";
     // Tear down the sessions and then the connections, in that order.
     _detail::ForEachBestEffort(
@@ -218,6 +236,38 @@ namespace Azure { namespace Messaging { namespace EventHubs {
     std::string suffix = !partitionId.empty() ? "/Partitions/" + partitionId : "";
     std::string hostUrl = m_hostUrl + suffix;
 
+#if ENABLE_UAMQP
+    auto partition = _detail::PartitionClientFactory::CreatePartitionClient(
+        m_fullyQualifiedNamespace,
+        m_credential,
+        m_targetPort,
+        m_consumerClientOptions.ApplicationID,
+        m_consumerClientOptions.CppStandardVersion,
+        "Consumer for " + m_consumerClientOptions.ApplicationID + " on " + partitionId,
+        std::move(hostUrl),
+        m_consumerClientOptions.Name,
+        options,
+        m_consumerClientOptions.RetryOptions,
+        context);
+    bool closeLatePartition = false;
+    {
+      std::lock_guard<std::mutex> lock(m_partitionClientStatesLock);
+      if (m_partitionClientStatesClosing)
+      {
+        closeLatePartition = true;
+      }
+      else
+      {
+        m_partitionClientStates.push_back(partition.GetState());
+      }
+    }
+    if (closeLatePartition)
+    {
+      _detail::ClosePartitionClientState(partition.GetState(), context);
+      throw Azure::Core::OperationCancelledException("Consumer client is closed.");
+    }
+    return partition;
+#elif ENABLE_RUST_AMQP
     EnsureSession(partitionId, context);
 
     return _detail::PartitionClientFactory::CreatePartitionClient(
@@ -227,6 +277,7 @@ namespace Azure { namespace Messaging { namespace EventHubs {
         options,
         m_consumerClientOptions.RetryOptions,
         context);
+#endif
   }
 
   Models::EventHubProperties ConsumerClient::GetEventHubProperties(Core::Context const& context)