diff --git a/Microsoft.Azure.Cosmos/src/DocumentClient.cs b/Microsoft.Azure.Cosmos/src/DocumentClient.cs index 02456b67c1..c70b0b8668 100644 --- a/Microsoft.Azure.Cosmos/src/DocumentClient.cs +++ b/Microsoft.Azure.Cosmos/src/DocumentClient.cs @@ -1405,6 +1405,12 @@ public void Dispose() this.cosmosAuthorization.Dispose(); } + if (this.PartitionKeyRangeLocation != null) + { + (this.PartitionKeyRangeLocation as IDisposable)?.Dispose(); + this.PartitionKeyRangeLocation = null; + } + if (this.GlobalEndpointManager != null) { this.GlobalEndpointManager.OnEnablePartitionLevelFailoverConfigChanged -= this.UpdatePartitionLevelFailoverConfigWithAccountRefresh; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeFailoverTests/GlobalPartitionEndpointManagerUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeFailoverTests/GlobalPartitionEndpointManagerUnitTests.cs index dc702d37f4..90ae13ab13 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeFailoverTests/GlobalPartitionEndpointManagerUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/PartitionKeyRangeFailoverTests/GlobalPartitionEndpointManagerUnitTests.cs @@ -8,6 +8,7 @@ namespace Microsoft.Azure.Cosmos.Tests using System.Collections.Generic; using System.Collections.ObjectModel; using System.Linq; + using System.Reflection; using System.Threading.Tasks; using System.Threading; using Microsoft.Azure.Cosmos.Routing; @@ -394,6 +395,113 @@ private async Task OpenConnectionToUnhealthyEndpointsAsync( } } + /// + /// Verifies that DocumentClient.Dispose() disposes PartitionKeyRangeLocation + /// when the implementation is IDisposable (GlobalPartitionEndpointManagerCore) + /// and sets it to null. + /// Regression test for: https://github.com/Azure/azure-cosmos-dotnet-v3/pull/5777 + /// + [TestMethod] + [Timeout(10000)] + public void Dispose_DisposesPartitionKeyRangeLocationWhenIDisposable() + { + using MockDocumentClient documentClient = new MockDocumentClient(); + + Mock mockEndpointManager = new Mock(MockBehavior.Loose); + GlobalPartitionEndpointManagerCore manager = new GlobalPartitionEndpointManagerCore( + mockEndpointManager.Object, + isPartitionLevelFailoverEnabled: false, + isPartitionLevelCircuitBreakerEnabled: false); + + PropertyInfo property = typeof(DocumentClient).GetProperty( + nameof(DocumentClient.PartitionKeyRangeLocation), + BindingFlags.Instance | BindingFlags.NonPublic); + + Assert.IsNotNull(property, "Could not find PartitionKeyRangeLocation property on DocumentClient."); + property.SetValue(documentClient, manager); + Assert.IsNotNull(documentClient.PartitionKeyRangeLocation); + Assert.IsInstanceOfType(documentClient.PartitionKeyRangeLocation, typeof(IDisposable)); + + documentClient.Dispose(); + + Assert.IsNull(documentClient.PartitionKeyRangeLocation, "PartitionKeyRangeLocation should be null after Dispose."); + + // Verify the manager was actually disposed: after disposal the cancellation token + // is cancelled, so re-initialization of the background loop is a no-op. + manager.InitializeAndStartCircuitBreakerFailbackBackgroundRefresh(); + } + + [TestMethod] + [Timeout(10000)] + public async Task Dispose_StopsBackgroundFailbackLoop() + { + Environment.SetEnvironmentVariable(ConfigurationManager.StalePartitionUnavailabilityRefreshIntervalInSeconds, "1"); + Environment.SetEnvironmentVariable(ConfigurationManager.AllowedPartitionUnavailabilityDurationInSeconds, "1"); + try + { + Mock mockEndpointManager = new Mock(MockBehavior.Strict); + + List readRegions = new(); + for (int i = 1; i <= 3; i++) + { + readRegions.Add(new Uri($"https://localhost:{i}/")); + } + + mockEndpointManager.Setup(x => x.ReadEndpoints).Returns(() => new ReadOnlyCollection(readRegions)); + mockEndpointManager.Setup(x => x.AccountReadEndpoints).Returns(() => new ReadOnlyCollection(readRegions)); + mockEndpointManager.Setup(x => x.WriteEndpoints).Returns(() => new ReadOnlyCollection(readRegions)); + mockEndpointManager.Setup(x => x.CanSupportMultipleWriteLocations(ResourceType.Document, OperationType.Create)).Returns(true); + + int callbackInvocationCount = 0; + + GlobalPartitionEndpointManagerCore manager = new GlobalPartitionEndpointManagerCore( + mockEndpointManager.Object, + isPartitionLevelFailoverEnabled: false, + isPartitionLevelCircuitBreakerEnabled: true); + + manager.SetBackgroundConnectionPeriodicRefreshTask( + async (pkRangeUriMappings) => + { + Interlocked.Increment(ref callbackInvocationCount); + await Task.CompletedTask; + }); + + PartitionKeyRange partitionKeyRange = new PartitionKeyRange() + { + Id = "0", + MinInclusive = "", + MaxExclusive = "FF" + }; + + Uri routeToLocation = new Uri("https://localhost:0/"); + + using DocumentServiceRequest createRequest = DocumentServiceRequest.Create(OperationType.Create, ResourceType.Document, AuthorizationTokenType.PrimaryMasterKey); + createRequest.RequestContext.ResolvedPartitionKeyRange = partitionKeyRange; + createRequest.RequestContext.RouteToLocation(routeToLocation); + createRequest.RequestContext.ResolvedCollectionRid = "test-collection"; + + GlobalPartitionEndpointManagerUnitTests.SimulateConsecutiveFailures(manager, createRequest); + + Assert.IsTrue(manager.TryMarkEndpointUnavailableForPartitionKeyRange(createRequest)); + + // Dispose should cancel the background loop. + manager.Dispose(); + + int countAfterDispose = callbackInvocationCount; + + // Wait long enough for the background loop to have triggered if it were still running. + await Task.Delay(TimeSpan.FromSeconds(3)); + + // The callback should not be invoked after dispose. + Assert.AreEqual(countAfterDispose, callbackInvocationCount, "Background failback loop should not invoke callback after Dispose."); + } + finally + { + Environment.SetEnvironmentVariable(ConfigurationManager.AllowedPartitionUnavailabilityDurationInSeconds, null); + Environment.SetEnvironmentVariable(ConfigurationManager.StalePartitionUnavailabilityRefreshIntervalInSeconds, null); + } + } + private static void SimulateConsecutiveFailures( GlobalPartitionEndpointManagerCore failoverManager, DocumentServiceRequest requestMessage)