[Event Hubs] Fixed producer semaphore release (Azure#47965)

* [Event Hubs] Fixed producer semaphore release

The focus of these changes is to fox an obscure edge case in the
`EventHubBufferedProducer` client where an obscure race condition when
flushing/enqueuing events concurrently with disposing the producer could
cause a semaphore to be released inappropriately.  This error superseded
the `TaskCanceledException` that should have been surfaced.

Also moved to volatile reads for the total buffer count when making
state-related decisions, to ensure that the most recent value is used.

Author: jsquire

sdk/eventhub/Azure.Messaging.EventHubs/CHANGELOG.md

@@ -16,6 +16,8 @@ Thank you to our developer community members who helped to make the Event Hubs c
 
 - Querying runtime data and other management operations will now correctly guards against the race condition where an AMQP link is in the process of closing as the operation attempts to use it.  These errors will now properly be classified as retriable as they are for producer and consumer operations.
 
+- Fixed an obscure edge case in the `EventHubBufferedProducer` client where an obscure race condition when flushing/enqueuing events concurrently with disposing the producer could cause a semaphore to be released inappropriately.  This error superseded the `TaskCanceledException` that should have been surfaced.
+
 ### Other Changes
 
 - Added annotations to make the package compatible with trimming and native AOT compilation.

sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventHubBufferedProducerClient.cs

@@ -196,7 +196,7 @@ public virtual bool IsClosed
         ///   The total number of events that are currently buffered and waiting to be published, across all partitions.
         /// </summary>
         ///
-        public virtual int TotalBufferedEventCount => _totalBufferedEventCount;
+        public virtual int TotalBufferedEventCount => Volatile.Read(ref _totalBufferedEventCount);
 
         /// <summary>
         ///   The instance of <see cref="EventHubsEventSource" /> which can be mocked for testing.
@@ -603,7 +603,7 @@ private EventHubBufferedProducerClient(string fullyQualifiedNamespace,
             {
                 TokenCredential tokenCred => new EventHubProducerClient(fullyQualifiedNamespace, eventHubName, tokenCred, options),
                 AzureSasCredential sasCred => new EventHubProducerClient(fullyQualifiedNamespace, eventHubName, sasCred, options),
-                AzureNamedKeyCredential keyCred =>  new EventHubProducerClient(fullyQualifiedNamespace, eventHubName, keyCred, options),
+                AzureNamedKeyCredential keyCred => new EventHubProducerClient(fullyQualifiedNamespace, eventHubName, keyCred, options),
                 _ => throw new ArgumentException(Resources.UnsupportedCredential, nameof(credential))
             };
 
@@ -772,13 +772,13 @@ public virtual async Task<int> EnqueueEventAsync(EventData eventData,
 
                 if ((!IsPublishing) || (_producerManagementTask?.IsCompleted ?? false))
                 {
-                    try
+                    if (!_stateGuard.Wait(0, cancellationToken))
                     {
-                        if (!_stateGuard.Wait(0, cancellationToken))
-                        {
-                            await _stateGuard.WaitAsync(cancellationToken).ConfigureAwait(false);
-                        }
+                        await _stateGuard.WaitAsync(cancellationToken).ConfigureAwait(false);
+                    }
 
+                    try
+                    {
                         Argument.AssertNotClosed(_isClosed, nameof(EventHubBufferedProducerClient));
 
                         // StartPublishingAsync will verify that publishing is not already taking
@@ -935,13 +935,13 @@ public virtual async Task<int> EnqueueEventsAsync(IEnumerable<EventData> events,
 
                 if ((!IsPublishing) || (_producerManagementTask?.IsCompleted ?? false))
                 {
-                    try
+                    if (!_stateGuard.Wait(0, cancellationToken))
                     {
-                        if (!_stateGuard.Wait(0, cancellationToken))
-                        {
-                            await _stateGuard.WaitAsync(cancellationToken).ConfigureAwait(false);
-                        }
+                        await _stateGuard.WaitAsync(cancellationToken).ConfigureAwait(false);
+                    }
 
+                    try
+                    {
                         Argument.AssertNotClosed(_isClosed, nameof(EventHubBufferedProducerClient));
 
                         // StartPublishingAsync will verify that publishing is not already taking
@@ -1059,6 +1059,15 @@ public virtual async Task FlushAsync(CancellationToken cancellationToken = defau
 
                 await StopPublishingAsync(cancelActiveSendOperations: false, cancellationToken).ConfigureAwait(false);
                 await FlushInternalAsync(cancellationToken).ConfigureAwait(false);
+
+                // There's an unlikely race condition where it is possible that an event batch was being enqueued before publishing
+                // stopped and was not written to the partition buffer until after the flush for that partition completed.  To guard
+                // against this, restart publishing if any events are pending.
+
+                if (Volatile.Read(ref _totalBufferedEventCount) > 0)
+                {
+                    await StartPublishingAsync(cancellationToken).ConfigureAwait(false);
+                }
             }
             finally
             {
@@ -1092,15 +1101,13 @@ public virtual async Task CloseAsync(bool flush = true,
 
             var capturedExceptions = default(List<Exception>);
 
-            try
+            if (!_stateGuard.Wait(0, cancellationToken))
             {
-                if (!_stateGuard.Wait(0, cancellationToken))
-                {
-                    await _stateGuard.WaitAsync(cancellationToken).ConfigureAwait(false);
-                }
-
-                // If we've reached this point without an exception, the guard is held.
+                await _stateGuard.WaitAsync(cancellationToken).ConfigureAwait(false);
+            }
 
+            try
+            {
                 if (_isClosed)
                 {
                     return;
@@ -1675,7 +1682,7 @@ internal virtual async Task DrainAndPublishPartitionEvents(PartitionPublishingSt
                 }
 
                 var partitionStateDelta = (partitionState.BufferedEventCount * -1);
-                var totalDeltaZero = (_totalBufferedEventCount * -1);
+                var totalDeltaZero = (Volatile.Read(ref _totalBufferedEventCount) * -1);
 
                 Interlocked.Add(ref _totalBufferedEventCount, Math.Max(totalDeltaZero, partitionStateDelta));
                 Interlocked.Exchange(ref partitionState.BufferedEventCount, 0);
@@ -1840,7 +1847,7 @@ private async Task StartPublishingAsync(CancellationToken cancellationToken)
                 // If there is already a task running for the background management process,
                 // then no further initialization is needed.
 
-                if ((IsPublishing) && (!_producerManagementTask.IsCompleted))
+                if ((IsPublishing) && (!_producerManagementTask.IsCompleted) || (_isClosed))
                 {
                     return;
                 }
@@ -1954,13 +1961,12 @@ private async Task StopPublishingAsync(bool cancelActiveSendOperations,
 
                 _producerManagementTask = null;
             }
+            catch (OperationCanceledException ex) when (ex is not TaskCanceledException)
+            {
+                throw new TaskCanceledException(ex.Message, ex);
+            }
             catch (Exception ex)
             {
-                if (ex is OperationCanceledException opEx)
-                {
-                    throw new TaskCanceledException(opEx.Message, opEx);
-                }
-
                 Logger.BufferedProducerBackgroundProcessingStopError(Identifier, EventHubName, ex.Message);
                 (capturedExceptions ??= new List<Exception>()).Add(ex);
             }
@@ -2247,7 +2253,7 @@ private Task RunPublishingAsync(CancellationToken cancellationToken) =>
                             if ((!cancellationToken.IsCancellationRequested)
                                 && (_activePartitionStateMap.TryGetValue(partition, out var partitionState))
                                 && (partitionState.BufferedEventCount > 0)
-                                && ((partitionState.PartitionGuard.Wait(0, cancellationToken)) || (await partitionState.PartitionGuard.WaitAsync(PartitionPublishingGuardAcquireLimitMilliseconds, cancellationToken).ConfigureAwait((false)))))
+                                && ((partitionState.PartitionGuard.Wait(0, cancellationToken)) || (await partitionState.PartitionGuard.WaitAsync(PartitionPublishingGuardAcquireLimitMilliseconds, cancellationToken).ConfigureAwait(false))))
                             {
                                 // Responsibility for releasing the guard semaphore is passed to the task.
 
@@ -2257,7 +2263,7 @@ private Task RunPublishingAsync(CancellationToken cancellationToken) =>
                             // If there are no events in the buffer, avoid a tight loop by blocking to wait for events to be enqueued
                             // after a small delay.
 
-                            if (_totalBufferedEventCount == 0)
+                            if (Volatile.Read(ref _totalBufferedEventCount) == 0)
                             {
                                 // If completion source doesn't exist or was already set, then swap in a new completion source to be
                                 // set when an event is enqueued.  Allow the publishing loop to tick for one additional check of the
@@ -2281,10 +2287,15 @@ private Task RunPublishingAsync(CancellationToken cancellationToken) =>
 
                                     var idleWatch = ValueStopwatch.StartNew();
 
-                                    await _eventEnqueuedCompletionSource.Task.AwaitWithCancellation(cancellationToken);
-                                    _eventEnqueuedCompletionSource = null;
-
-                                    Logger.BufferedProducerIdleComplete(Identifier, EventHubName, operationId, idleWatch.GetElapsedTime().TotalSeconds);
+                                    try
+                                    {
+                                        await _eventEnqueuedCompletionSource.Task.AwaitWithCancellation(cancellationToken);
+                                        _eventEnqueuedCompletionSource = null;
+                                    }
+                                    finally
+                                    {
+                                        Logger.BufferedProducerIdleComplete(Identifier, EventHubName, operationId, idleWatch.GetElapsedTime().TotalSeconds);
+                                    }
                                 }
                             }
                         }
@@ -2332,7 +2343,7 @@ private Task RunPublishingAsync(CancellationToken cancellationToken) =>
         ///    set of Event Hub partitions has changed since they were last queried.
         /// </summary>
         ///
-        /// <param name="cancellationToken">A <see cref="CancellationToken" /> instance to signal the request to cancel the operation.</param>
+        /// v
         ///
         /// <remarks>
         ///   This method will potentially modify class state, overwriting the tracked set of partitions.
@@ -2468,12 +2479,12 @@ internal sealed class PartitionPublishingState : IDisposable
             /// <summary>The writer to use for enqueuing events to be published.</summary>
             public ChannelWriter<EventData> PendingEventsWriter => _pendingEvents.Writer;
 
-            /// <summary>The identifier of the partition that is being published.</summary>
-            public readonly string PartitionId;
-
             /// <summary>The primitive for synchronizing access for publishing to the partition.</summary>
             public readonly SemaphoreSlim PartitionGuard;
 
+            /// <summary>The identifier of the partition that is being published.</summary>
+            public readonly string PartitionId;
+
             /// <summary>The number of events that are currently buffered and waiting to be published for this partition.</summary>
             public int BufferedEventCount;
 

sdk/eventhub/Azure.Messaging.EventHubs/tests/Consumer/EventHubConsumerClientLiveTests.cs

@@ -1333,41 +1333,6 @@ public async Task ConsumerCannotReadFromInvalidConsumerGroup()
             }
         }
 
-        /// <summary>
-        ///   Verifies that the <see cref="EventHubConsumerClient" /> is able to
-        ///   connect to the Event Hubs service and perform operations.
-        /// </summary>
-        ///
-        [Test]
-        public async Task ConsumerCannotReadWithInvalidProxy()
-        {
-            await using (EventHubScope scope = await EventHubScope.CreateAsync(1))
-            {
-                using var cancellationSource = new CancellationTokenSource();
-                cancellationSource.CancelAfter(EventHubsTestEnvironment.Instance.TestExecutionTimeLimit);
-
-                var clientOptions = new EventHubConsumerClientOptions();
-                clientOptions.RetryOptions.MaximumRetries = 0;
-                clientOptions.RetryOptions.MaximumDelay = TimeSpan.FromMilliseconds(5);
-                clientOptions.RetryOptions.TryTimeout = TimeSpan.FromSeconds(45);
-                clientOptions.ConnectionOptions.Proxy = new WebProxy("http://1.2.3.4:9999");
-                clientOptions.ConnectionOptions.TransportType = EventHubsTransportType.AmqpWebSockets;
-
-                await using (var producer = new EventHubProducerClient(EventHubsTestEnvironment.Instance.FullyQualifiedNamespace, scope.EventHubName, EventHubsTestEnvironment.Instance.Credential))
-                await using (var invalidProxyConsumer = new EventHubConsumerClient(EventHubConsumerClient.DefaultConsumerGroupName, EventHubsTestEnvironment.Instance.FullyQualifiedNamespace, scope.EventHubName, EventHubsTestEnvironment.Instance.Credential, clientOptions))
-                {
-                    var partition = (await producer.GetPartitionIdsAsync(cancellationSource.Token)).First();
-
-                    // The sockets implementation in .NET Core on some platforms, such as Linux, does not trigger a specific socket exception and
-                    // will, instead, hang indefinitely.  The try timeout is intentionally set to a value smaller than the cancellation token to
-                    // invoke a timeout exception in these cases.
-
-                    Assert.That(async () => await ReadNothingAsync(invalidProxyConsumer, partition, cancellationSource.Token, iterationCount: 25), Throws.InstanceOf<WebSocketException>().Or.InstanceOf<TimeoutException>());
-                    Assert.That(cancellationSource.IsCancellationRequested, Is.False, "The cancellation token should not have been signaled.");
-                }
-            }
-        }
-
         /// <summary>
         ///   Verifies that the <see cref="EventHubConsumerClient" /> is able to
         ///   connect to the Event Hubs service and perform operations.
@@ -2200,53 +2165,6 @@ public async Task ConsumerCannotRetrievePartitionPropertiesWithInvalidPartitionI
             }
         }
 
-        /// <summary>
-        ///   Verifies that the <see cref="EventHubConsumerClient" /> is able to
-        ///   connect to the Event Hubs service.
-        /// </summary>
-        ///
-        [Test]
-        public async Task ConsumerCannotRetrieveMetadataWithInvalidProxy()
-        {
-            await using (EventHubScope scope = await EventHubScope.CreateAsync(1))
-            {
-                using var cancellationSource = new CancellationTokenSource();
-                cancellationSource.CancelAfter(EventHubsTestEnvironment.Instance.TestExecutionTimeLimit);
-
-                var invalidProxyOptions = new EventHubConsumerClientOptions();
-                invalidProxyOptions.RetryOptions.MaximumRetries = 0;
-                invalidProxyOptions.RetryOptions.MaximumDelay = TimeSpan.FromMilliseconds(5);
-                invalidProxyOptions.RetryOptions.TryTimeout = TimeSpan.FromSeconds(45);
-                invalidProxyOptions.ConnectionOptions.Proxy = new WebProxy("http://1.2.3.4:9999");
-                invalidProxyOptions.ConnectionOptions.TransportType = EventHubsTransportType.AmqpWebSockets;
-
-                await using (var consumer = new EventHubConsumerClient(
-                    EventHubConsumerClient.DefaultConsumerGroupName,
-                    EventHubsTestEnvironment.Instance.FullyQualifiedNamespace,
-                    scope.EventHubName,
-                    EventHubsTestEnvironment.Instance.Credential))
-
-                await using (var invalidProxyConsumer = new EventHubConsumerClient(
-                    EventHubConsumerClient.DefaultConsumerGroupName,
-                    EventHubsTestEnvironment.Instance.FullyQualifiedNamespace,
-                    scope.EventHubName,
-                    EventHubsTestEnvironment.Instance.Credential,
-                    invalidProxyOptions))
-                {
-                    var partition = (await consumer.GetPartitionIdsAsync(cancellationSource.Token)).First();
-
-                    // The sockets implementation in .NET Core on some platforms, such as Linux, does not trigger a specific socket exception and
-                    // will, instead, hang indefinitely.  The try timeout is intentionally set to a value smaller than the cancellation token to
-                    // invoke a timeout exception in these cases.
-
-                    Assert.That(async () => await invalidProxyConsumer.GetPartitionIdsAsync(cancellationSource.Token), Throws.InstanceOf<WebSocketException>().Or.InstanceOf<TimeoutException>());
-                    Assert.That(async () => await invalidProxyConsumer.GetEventHubPropertiesAsync(cancellationSource.Token), Throws.InstanceOf<WebSocketException>().Or.InstanceOf<TimeoutException>());
-                    Assert.That(async () => await invalidProxyConsumer.GetPartitionPropertiesAsync(partition, cancellationSource.Token), Throws.InstanceOf<WebSocketException>().Or.InstanceOf<TimeoutException>());
-                    Assert.That(cancellationSource.IsCancellationRequested, Is.False, "The cancellation token should not have been signaled.");
-                }
-            }
-        }
-
         /// <summary>
         ///   Verifies that the <see cref="EventHubConsumerClient" /> is able to
         ///   connect to the Event Hubs service and perform operations.