Avoid NRE and disposed exceptions when closing client twice (#3312)

Azure · Apr 28, 2023 · 06d0936 · 06d0936
1 parent 1260326
commit 06d0936
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 12 deletions.
diff --git a/e2e/test/iothub/DeviceClientE2eTests.cs b/e2e/test/iothub/DeviceClientE2eTests.cs
@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.Azure.Devices.E2ETests.Helpers;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+namespace Microsoft.Azure.Devices.E2ETests
+{
+    [TestClass]
+    [TestCategory("E2E")]
+    [TestCategory("IoTHub")]
+    public class DeviceClientE2eTests
+    {
+        [TestMethod]
+        public async Task DeviceClient_CloseAsync_CanBeCalledTwice()
+        {
+            // arrange
+
+            using TestDevice testDevice = await TestDevice.GetTestDeviceAsync(nameof(DeviceClient_CloseAsync_CanBeCalledTwice));
+
+            try
+            {
+                using Client.DeviceClient client = testDevice.CreateDeviceClient(Client.TransportType.Amqp_Tcp_Only);
+                await client.OpenAsync().ConfigureAwait(false);
+                await client.CloseAsync().ConfigureAwait(false);
+
+                // act
+                Func<Task> act = () => client.CloseAsync();
+
+                // assert
+                await act.Should().NotThrowAsync().ConfigureAwait(false);
+            }
+            finally
+            {
+                await testDevice.RemoveDeviceAsync().ConfigureAwait(false);
+                testDevice.Dispose();
+            }
+
+        }
+    }
+}
diff --git a/iothub/device/src/Pipeline/RetryDelegatingHandler.cs b/iothub/device/src/Pipeline/RetryDelegatingHandler.cs
@@ -784,16 +784,17 @@ public override async Task CloseAsync(CancellationToken cancellationToken)
         {
             try
             {
-                if (!_openCalled)
+                if (!_openCalled
+                    || GetClientTransportStatus() == ClientTransportStatus.Closed)
                 {
                     return;
                 }
 
                 if (Logging.IsEnabled)
                     Logging.Enter(this, cancellationToken, nameof(CloseAsync));
 
-                _handleDisconnectCts.Cancel();
-                _cancelPendingOperationsCts.Cancel();
+                _handleDisconnectCts?.Cancel();
+                _cancelPendingOperationsCts?.Cancel();
                 await base.CloseAsync(cancellationToken).ConfigureAwait(false);
             }
             finally
@@ -859,7 +860,7 @@ private async Task EnsureOpenedAsync(bool withRetry, CancellationToken cancellat
                         if (Logging.IsEnabled)
                             Logging.Info(this, "Race condition: Disposed during opening.", nameof(EnsureOpenedAsync));
 
-                        _handleDisconnectCts.Cancel();
+                        _handleDisconnectCts?.Cancel();
                     }
                 }
             }
@@ -930,7 +931,7 @@ private async Task EnsureOpenedAsync(bool withRetry, TimeoutHelper timeoutHelper
                         if (Logging.IsEnabled)
                             Logging.Info(this, "Race condition: Disposed during opening.", nameof(EnsureOpenedAsync));
 
-                        _handleDisconnectCts.Cancel();
+                        _handleDisconnectCts?.Cancel();
                     }
                 }
             }
@@ -1071,7 +1072,7 @@ private async Task HandleDisconnectAsync()
                 if (Logging.IsEnabled)
                     Logging.Info(this, "Disposed during disconnection.", nameof(HandleDisconnectAsync));
 
-                _handleDisconnectCts.Cancel();
+                _handleDisconnectCts?.Cancel();
             }
 
             try
@@ -1114,7 +1115,7 @@ private async Task HandleDisconnectAsync()
 
                 // always reconnect.
                 _onConnectionStatusChanged(ConnectionStatus.Disconnected_Retrying, ConnectionStatusChangeReason.Communication_Error);
-                CancellationToken cancellationToken = _handleDisconnectCts.Token;
+                CancellationToken cancellationToken = _handleDisconnectCts?.Token ?? CancellationToken.None;
 
                 // This will recover to the state before the disconnect.
                 await _internalRetryPolicy.RunWithRetryAsync(async () =>
@@ -1253,9 +1254,7 @@ protected override void Dispose(bool disposing)
             try
             {
                 if (Logging.IsEnabled)
-                {
                     Logging.Enter(this, $"{nameof(DefaultDelegatingHandler)}.Disposed={_isDisposed}; disposing={disposing}", $"{nameof(RetryDelegatingHandler)}.{nameof(Dispose)}");
-                }
 
                 if (!_isDisposed)
                 {
@@ -1266,6 +1265,9 @@ protected override void Dispose(bool disposing)
 
                     if (disposing)
                     {
+                        _handleDisconnectCts?.Cancel();
+                        _cancelPendingOperationsCts?.Cancel();
+
                         var disposables = new List<IDisposable>
                         {
                             _handleDisconnectCts,
@@ -1277,9 +1279,6 @@ protected override void Dispose(bool disposing)
                             _twinEventsSubscriptionSemaphore,
                         };
 
-                        _handleDisconnectCts?.Cancel();
-                        _cancelPendingOperationsCts?.Cancel();
-
                         foreach (IDisposable disposable in disposables)
                         {
                             try