diff --git a/Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/FeedManagement/PartitionControllerCore.cs b/Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/FeedManagement/PartitionControllerCore.cs index e181dadbec..7407477c21 100644 --- a/Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/FeedManagement/PartitionControllerCore.cs +++ b/Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/FeedManagement/PartitionControllerCore.cs @@ -89,8 +89,7 @@ public override async Task AddOrUpdateLeaseAsync(DocumentServiceLease lease) throw; } - PartitionSupervisor supervisor = this.partitionSupervisorFactory.Create(lease); - this.ProcessPartitionAsync(supervisor, lease).LogException(); + this.ProcessPartitionAsync(lease).LogException(); } public override async Task ShutdownAsync() @@ -146,8 +145,10 @@ private async Task RemoveLeaseAsync(DocumentServiceLease lease, bool wasAcquired } } - private async Task ProcessPartitionAsync(PartitionSupervisor partitionSupervisor, DocumentServiceLease lease) + private async Task ProcessPartitionAsync(DocumentServiceLease lease) { + using PartitionSupervisor partitionSupervisor = this.partitionSupervisorFactory.Create(lease); + try { await partitionSupervisor.RunAsync(this.shutdownCts.Token).ConfigureAwait(false); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/PartitionControllerSplitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/PartitionControllerSplitTests.cs index 25704e05c8..ba56424ee0 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/PartitionControllerSplitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/PartitionControllerSplitTests.cs @@ -48,6 +48,9 @@ public async Task Controller_ShouldSignalSynchronizerSplitPartition_IfPartitionS await sut.ShutdownAsync().ConfigureAwait(false); Mock.Get(synchronizer).VerifyAll(); + + Mock.Get(partitionSupervisor) + .Verify(s => s.Dispose(), Times.Once); } [TestMethod] @@ -76,6 +79,9 @@ public async Task Controller_ShouldPassLastKnownContinuationTokenToSynchronizer_ await sut.ShutdownAsync().ConfigureAwait(false); Mock.Get(synchronizer).VerifyAll(); + + Mock.Get(partitionSupervisor) + .Verify(s => s.Dispose(), Times.Once); } [TestMethod] @@ -108,6 +114,9 @@ public async Task Controller_ShouldCopyParentLeaseProperties_IfPartitionSplitHap .VerifySet(l => l.Properties = customProperties, Times.Once); Mock.Get(leaseChild2) .VerifySet(l => l.Properties = customProperties, Times.Once); + + Mock.Get(partitionSupervisor) + .Verify(s => s.Dispose(), Times.Once); } [TestMethod] @@ -132,6 +141,9 @@ public async Task Controller_ShouldKeepParentLease_IfSplitThrows() await sut.ShutdownAsync().ConfigureAwait(false); Mock.Get(leaseManager).Verify(manager => manager.DeleteAsync(lease), Times.Never); + + Mock.Get(partitionSupervisor) + .Verify(s => s.Dispose(), Times.Once); } [TestMethod] @@ -180,6 +192,13 @@ public async Task Controller_ShouldRunProcessingOnChildPartitions_IfHappyPath() monitor.Verify(m => m.NotifyLeaseAcquireAsync(leaseChild1.CurrentLeaseToken), Times.Once); monitor.Verify(m => m.NotifyLeaseReleaseAsync(leaseChild2.CurrentLeaseToken), Times.Once); + + Mock.Get(partitionSupervisor) + .Verify(s => s.Dispose(), Times.Once); + Mock.Get(partitionSupervisor1) + .Verify(s => s.Dispose(), Times.Once); + Mock.Get(partitionSupervisor2) + .Verify(s => s.Dispose(), Times.Once); } [TestMethod]