DocumentClient: Fixes ObjectDisposedException during Background Refresh by adding Cancellation Token

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -163,11 +163,13 @@ internal partial class DocumentClient : IDisposable, IAuthorizationTokenProvider
         //SessionContainer.
         internal ISessionContainer sessionContainer;
 
+        private CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
+
         private AsyncLazy<QueryPartitionProvider> queryPartitionProvider;
 
         private DocumentClientEventSource eventSource;
         private Func<bool, Task<bool>> initializeTaskFactory;
-        internal AsyncCacheNonBlocking<string, bool> initTaskCache = new AsyncCacheNonBlocking<string, bool>();
+        internal AsyncCacheNonBlocking<string, bool> initTaskCache;
 
         private JsonSerializerSettings serializerSettings;
         private event EventHandler<SendingRequestEventArgs> sendingRequest;
@@ -218,6 +220,8 @@ public DocumentClient(Uri serviceEndpoint,
             }
 
             this.Initialize(serviceEndpoint, connectionPolicy, desiredConsistencyLevel);
+            this.initTaskCache = new AsyncCacheNonBlocking<string, bool>(cancellationToken: this.cancellationTokenSource.Token);
+
         }
 
         /// <summary>
@@ -459,6 +463,7 @@ internal DocumentClient(Uri serviceEndpoint,
             this.cosmosAuthorization = cosmosAuthorization ?? throw new ArgumentNullException(nameof(cosmosAuthorization));
             this.transportClientHandlerFactory = transportClientHandlerFactory;
             this.IsLocalQuorumConsistency = isLocalQuorumConsistency;
+            this.initTaskCache = new AsyncCacheNonBlocking<string, bool>(cancellationToken: this.cancellationTokenSource.Token);
 
             this.Initialize(
                 serviceEndpoint: serviceEndpoint,
@@ -1200,6 +1205,11 @@ public void Dispose()
                 return;
             }
 
+            if (this.cancellationTokenSource != null && !this.cancellationTokenSource.IsCancellationRequested)
+            {
+                this.cancellationTokenSource.Cancel();
+            }
+
             if (this.StoreModel != null)
             {
                 this.StoreModel.Dispose();
@@ -6596,7 +6606,8 @@ private async Task InitializeGatewayConfigurationReaderAsync()
                     serviceEndpoint: this.ServiceEndpoint,
                     cosmosAuthorization: this.cosmosAuthorization,
                     connectionPolicy: this.ConnectionPolicy,
-                    httpClient: this.httpClient);
+                    httpClient: this.httpClient,
+                    cancellationToken: this.cancellationTokenSource.Token);
 
             this.accountServiceConfiguration = new CosmosAccountServiceConfiguration(accountReader.InitializeReaderAsync);
 

Microsoft.Azure.Cosmos/src/GatewayAccountReader.cs

@@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos
     using System;
     using System.Globalization;
     using System.Net.Http;
+    using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Resource.CosmosExceptions;
     using Microsoft.Azure.Cosmos.Routing;
@@ -21,17 +22,20 @@ internal sealed class GatewayAccountReader
         private readonly AuthorizationTokenProvider cosmosAuthorization;
         private readonly CosmosHttpClient httpClient;
         private readonly Uri serviceEndpoint;
+        private CancellationToken cancellationToken;
 
         // Backlog: Auth abstractions are spilling through. 4 arguments for this CTOR are result of it.
         public GatewayAccountReader(Uri serviceEndpoint,
                 AuthorizationTokenProvider cosmosAuthorization,
                 ConnectionPolicy connectionPolicy,
-                CosmosHttpClient httpClient)
+                CosmosHttpClient httpClient,
+                CancellationToken cancellationToken = default)
         {
             this.httpClient = httpClient;
             this.serviceEndpoint = serviceEndpoint;
             this.cosmosAuthorization = cosmosAuthorization ?? throw new ArgumentNullException(nameof(AuthorizationTokenProvider));
             this.connectionPolicy = connectionPolicy;
+            this.cancellationToken = cancellationToken;
         }
 
         private async Task<AccountProperties> GetDatabaseAccountAsync(Uri serviceEndpoint)
@@ -52,7 +56,7 @@ await this.cosmosAuthorization.AddAuthorizationHeaderAsync(
                     resourceType: ResourceType.DatabaseAccount,
                     timeoutPolicy: HttpTimeoutPolicyControlPlaneRead.Instance,
                     clientSideRequestStatistics: stats,
-                    cancellationToken: default))
+                    cancellationToken: this.cancellationToken))
                 {
                     using (DocumentServiceResponse documentServiceResponse = await ClientExtensions.ParseResponseAsync(responseMessage))
                     {
@@ -84,7 +88,7 @@ public async Task<AccountProperties> InitializeReaderAsync()
                 defaultEndpoint: this.serviceEndpoint,
                 locations: this.connectionPolicy.PreferredLocations,
                 getDatabaseAccountFn: this.GetDatabaseAccountAsync,
-                cancellationToken: default);
+                cancellationToken: this.cancellationToken);
 
             return databaseAccount;
         }

Microsoft.Azure.Cosmos/src/Routing/AsyncCacheNonBlocking.cs

@@ -21,7 +21,7 @@ namespace Microsoft.Azure.Cosmos
     /// </summary>
     internal sealed class AsyncCacheNonBlocking<TKey, TValue> : IDisposable
     {
-        private readonly CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
+        private readonly CancellationTokenSource cancellationTokenSource;
         private readonly ConcurrentDictionary<TKey, AsyncLazyWithRefreshTask<TValue>> values;
         private readonly Func<Exception, bool> removeFromCacheOnBackgroundRefreshException;
 
@@ -30,11 +30,15 @@ internal sealed class AsyncCacheNonBlocking<TKey, TValue> : IDisposable
 
         public AsyncCacheNonBlocking(
             Func<Exception, bool> removeFromCacheOnBackgroundRefreshException = null,
-            IEqualityComparer<TKey> keyEqualityComparer = null)
+            IEqualityComparer<TKey> keyEqualityComparer = null,
+            CancellationToken cancellationToken = default)
         {
             this.keyEqualityComparer = keyEqualityComparer ?? EqualityComparer<TKey>.Default;
             this.values = new ConcurrentDictionary<TKey, AsyncLazyWithRefreshTask<TValue>>(this.keyEqualityComparer);
             this.removeFromCacheOnBackgroundRefreshException = removeFromCacheOnBackgroundRefreshException ?? AsyncCacheNonBlocking<TKey, TValue>.RemoveNotFoundFromCacheOnException;
+            this.cancellationTokenSource = cancellationToken == default
+                ? new CancellationTokenSource()
+                : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
         }
 
         public AsyncCacheNonBlocking()

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosClientTests.cs

@@ -6,13 +6,15 @@ namespace Microsoft.Azure.Cosmos.Tests
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using System.Globalization;
     using System.Linq;
     using System.Net;
     using System.Net.Http;
     using System.Threading;
     using System.Threading.Tasks;
     using global::Azure.Core;
+    using Microsoft.Azure.Cosmos.Core.Trace;
     using Microsoft.Azure.Cosmos.Fluent;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
     using Moq;
@@ -230,5 +232,52 @@ public void Builder_ValidateHttpFactory()
                 WebProxy = null,
             };
         }
+
+        [TestMethod]
+        public void CosmosClientEarlyDisposeTest()
+        {
+            string disposeErrorMsg = "Cannot access a disposed object";
+            HashSet<string> errors = new HashSet<string>();
+
+            void TraceHandler(string message)
+            {
+                if (message.Contains(disposeErrorMsg))
+                {
+                    errors.Add(message);
+                }
+            }
+
+            DefaultTrace.TraceSource.Listeners.Add(new TestTraceListener { Callback = TraceHandler });
+            DefaultTrace.InitEventListener();
+
+            for (int z = 0; z < 100; ++z)
+            {
+                using CosmosClient cosmos = new(ConnectionString);
+            }
+
+            string assertMsg = "";
+
+            foreach (string s in errors)
+            {
+                assertMsg += s + "\n";
+            }
+
+            Assert.AreEqual(0, errors.Count, $"\nErrors found in trace:\n{assertMsg}");
+        }
+
+        private class TestTraceListener : TraceListener
+        {
+            public Action<string> Callback { get; set; }
+            public override bool IsThreadSafe => true;
+            public override void Write(string message)
+            {
+                this.Callback(message);
+            }
+
+            public override void WriteLine(string message)
+            {
+                this.Callback(message);
+            }
+        }
     }
 }