feat(directory): add cancellation-aware operations

src/Orleans.Core.Abstractions/GrainDirectory/IGrainDirectory.cs

@@ -1,4 +1,5 @@
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using Orleans.Runtime;
 
@@ -18,6 +19,15 @@ public interface IGrainDirectory
         /// <returns>The <see cref="GrainAddress"/> that is effectively registered in the directory.</returns>
         Task<GrainAddress?> Register(GrainAddress address);
 
+        /// <summary>
+        /// Register a <see cref="GrainAddress"/> entry in the directory.
+        /// </summary>
+        /// <param name="address">The <see cref="GrainAddress"/> to register.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>The <see cref="GrainAddress"/> that is effectively registered in the directory.</returns>
+        Task<GrainAddress?> Register(GrainAddress address, CancellationToken cancellationToken) =>
+            Register(address).WaitAsync(cancellationToken);
+
         /// <summary>
         /// Register a <see cref="GrainAddress"/> entry in the directory.
         /// Only one <see cref="GrainAddress"/> per <see cref="GrainAddress.GrainId"/> can be registered. If there is already an
@@ -27,6 +37,16 @@ public interface IGrainDirectory
         /// <returns>The <see cref="GrainAddress"/> that is effectively registered in the directory.</returns>
         Task<GrainAddress?> Register(GrainAddress address, GrainAddress? previousAddress) => GrainDirectoryExtension.Register(this, address, previousAddress);
 
+        /// <summary>
+        /// Register a <see cref="GrainAddress"/> entry in the directory.
+        /// </summary>
+        /// <param name="address">The <see cref="GrainAddress"/> to register.</param>
+        /// <param name="previousAddress">The previous registration, if any.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>The <see cref="GrainAddress"/> that is effectively registered in the directory.</returns>
+        Task<GrainAddress?> Register(GrainAddress address, GrainAddress? previousAddress, CancellationToken cancellationToken) =>
+            Register(address, previousAddress).WaitAsync(cancellationToken);
+
         /// <summary>
         /// Unregisters the specified <see cref="GrainAddress"/> entry from the directory.
         /// </summary>
@@ -38,13 +58,31 @@ public interface IGrainDirectory
         /// </returns>
         Task Unregister(GrainAddress address);
 
+        /// <summary>
+        /// Unregisters the specified <see cref="GrainAddress"/> entry from the directory.
+        /// </summary>
+        /// <param name="address">The <see cref="GrainAddress"/> to unregister.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>A <see cref="Task"/> representing the operation.</returns>
+        Task Unregister(GrainAddress address, CancellationToken cancellationToken) =>
+            Unregister(address).WaitAsync(cancellationToken);
+
         /// <summary>
         /// Lookup for a <see cref="GrainAddress"/> for a given Grain ID.
         /// </summary>
         /// <param name="grainId">The Grain ID to lookup</param>
         /// <returns>The <see cref="GrainAddress"/> entry found in the directory, if any</returns>
         Task<GrainAddress?> Lookup(GrainId grainId);
 
+        /// <summary>
+        /// Lookup for a <see cref="GrainAddress"/> for a given Grain ID.
+        /// </summary>
+        /// <param name="grainId">The Grain ID to lookup.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>The <see cref="GrainAddress"/> entry found in the directory, if any.</returns>
+        Task<GrainAddress?> Lookup(GrainId grainId, CancellationToken cancellationToken) =>
+            Lookup(grainId).WaitAsync(cancellationToken);
+
         /// <summary>
         /// Unregisters all grain directory entries which point to any of the specified silos.
         /// </summary>
@@ -56,6 +94,15 @@ public interface IGrainDirectory
         /// A <see cref="Task"/> representing the operation.
         /// </returns>
         Task UnregisterSilos(List<SiloAddress> siloAddresses);
+
+        /// <summary>
+        /// Unregisters all grain directory entries which point to any of the specified silos.
+        /// </summary>
+        /// <param name="siloAddresses">The silos to be removed from the directory.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>A <see cref="Task"/> representing the operation.</returns>
+        Task UnregisterSilos(List<SiloAddress> siloAddresses, CancellationToken cancellationToken) =>
+            UnregisterSilos(siloAddresses).WaitAsync(cancellationToken);
     }
 
     internal static class GrainDirectoryExtension

src/Orleans.Runtime/Catalog/ActivationData.cs

@@ -1662,7 +1662,7 @@ private async Task ActivateAsync(Dictionary<string, object>? requestContextData,
                             LogRegisteringGrain(_shared.Logger, this, previousRegistration);
 
                             var result = await _shared.InternalRuntime.GrainLocator
-                                .Register(Address, previousRegistration).WaitAsync(cancellationToken);
+                                .Register(Address, previousRegistration, cancellationToken);
                             if (Address.Matches(result))
                             {
                                 Address = result;

src/Orleans.Runtime/GrainDirectory/CachedGrainLocator.cs

@@ -84,7 +84,13 @@ public CachedGrainLocator(
             return entry;
         }
 
-        public async Task<GrainAddress?> Register(GrainAddress address, GrainAddress? previousAddress)
+        public Task<GrainAddress?> Register(GrainAddress address, GrainAddress? previousAddress) =>
+            Register(address, previousAddress, CancellationToken.None);
+
+        internal async Task<GrainAddress?> Register(
+            GrainAddress address,
+            GrainAddress? previousAddress,
+            CancellationToken cancellationToken)
         {
             var grainType = address.GrainId.Type;
             if (grainType.IsClient() || grainType.IsSystemTarget())
@@ -100,7 +106,8 @@ public CachedGrainLocator(
                 MembershipVersion = clusterMembershipService.CurrentSnapshot.Version
             };
 
-            var result = await GetGrainDirectory(grainType).Register(address, previousAddress);
+            var directory = GetGrainDirectory(grainType);
+            var result = await directory.Register(address, previousAddress, cancellationToken);
 
             if (result is null)
             {
@@ -111,9 +118,8 @@ public CachedGrainLocator(
             // Check if the entry point to a dead silo
             if (IsKnownDeadSilo(result))
             {
-                // Remove outdated entry and retry to register
-                await GetGrainDirectory(grainType).Unregister(result);
-                result = await GetGrainDirectory(grainType).Register(address, previousAddress);
+                // Conditionally replace the outdated entry.
+                result = await directory.Register(address, result, cancellationToken);
             }
 
             // Cache update

src/Orleans.Runtime/GrainDirectory/DistributedGrainDirectory.cs

@@ -113,19 +113,51 @@ public DistributedGrainDirectory(
         DistributedRemoteGrainDirectory.Create(this, membershipService, shared);
     }
 
-    public async Task<GrainAddress?> Lookup(GrainId grainId) => await LookupAsync(grainId, _stoppedCts.Token);
+    public Task<GrainAddress?> Lookup(GrainId grainId) => Lookup(grainId, _stoppedCts.Token);
 
-    public async Task<GrainAddress?> Register(GrainAddress address) => await RegisterAsync(address, null, _stoppedCts.Token);
+    public async Task<GrainAddress?> Lookup(GrainId grainId, CancellationToken cancellationToken)
+    {
+        using var linkedCts = LinkWithStoppedToken(cancellationToken, out var effectiveCancellationToken);
+        return await LookupAsync(grainId, effectiveCancellationToken);
+    }
 
-    public async Task Unregister(GrainAddress address) => await InvokeAsync(
-        address.GrainId,
-        static (partition, version, address, cancellationToken) => partition.DeregisterAsync(version, address, cancellationToken),
-        address,
-        _stoppedCts.Token);
+    public Task<GrainAddress?> Register(GrainAddress address) => Register(address, _stoppedCts.Token);
+
+    public async Task<GrainAddress?> Register(GrainAddress address, CancellationToken cancellationToken)
+    {
+        using var linkedCts = LinkWithStoppedToken(cancellationToken, out var effectiveCancellationToken);
+        return await RegisterAsync(address, null, effectiveCancellationToken);
+    }
+
+    public Task Unregister(GrainAddress address) => Unregister(address, _stoppedCts.Token);
 
-    public async Task<GrainAddress?> Register(GrainAddress address, GrainAddress? previousAddress) => await RegisterAsync(address, previousAddress, _stoppedCts.Token);
+    public async Task Unregister(GrainAddress address, CancellationToken cancellationToken)
+    {
+        using var linkedCts = LinkWithStoppedToken(cancellationToken, out var effectiveCancellationToken);
+        await UnregisterAsync(address, effectiveCancellationToken);
+    }
 
-    public Task UnregisterSilos(List<SiloAddress> siloAddresses) => Task.CompletedTask;
+    public Task<GrainAddress?> Register(GrainAddress address, GrainAddress? previousAddress) =>
+        Register(address, previousAddress, _stoppedCts.Token);
+
+    public async Task<GrainAddress?> Register(
+        GrainAddress address,
+        GrainAddress? previousAddress,
+        CancellationToken cancellationToken)
+    {
+        using var linkedCts = LinkWithStoppedToken(cancellationToken, out var effectiveCancellationToken);
+        return await RegisterAsync(address, previousAddress, effectiveCancellationToken);
+    }
+
+    public Task UnregisterSilos(List<SiloAddress> siloAddresses) =>
+        UnregisterSilos(siloAddresses, _stoppedCts.Token);
+
+    public async Task UnregisterSilos(List<SiloAddress> siloAddresses, CancellationToken cancellationToken)
+    {
+        using var linkedCts = LinkWithStoppedToken(cancellationToken, out var effectiveCancellationToken);
+        effectiveCancellationToken.ThrowIfCancellationRequested();
+        await Task.CompletedTask;
+    }
 
     internal Task<GrainAddress?> LookupAsync(GrainId grainId, CancellationToken cancellationToken) => InvokeAsync(
         grainId,
@@ -145,6 +177,21 @@ internal Task UnregisterAsync(GrainAddress address, CancellationToken cancellati
         address,
         cancellationToken);
 
+    private CancellationTokenSource? LinkWithStoppedToken(
+        CancellationToken cancellationToken,
+        out CancellationToken effectiveCancellationToken)
+    {
+        if (cancellationToken == _stoppedCts.Token)
+        {
+            effectiveCancellationToken = cancellationToken;
+            return null;
+        }
+
+        var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _stoppedCts.Token);
+        effectiveCancellationToken = linkedCts.Token;
+        return linkedCts;
+    }
+
     private async Task<TResult?> InvokeAsync<TState, TResult>(
         GrainId grainId,
         Func<IGrainDirectoryPartition, MembershipVersion, TState, CancellationToken, ValueTask<DirectoryResult<TResult>>> func,

src/Orleans.Runtime/GrainDirectory/GrainLocator.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using System.Threading;
 using System.Threading.Tasks;
 using Orleans.GrainDirectory;
 
@@ -22,13 +23,18 @@ public GrainLocator(GrainLocatorResolver grainLocatorResolver, DirectoryInstrume
 
         public ValueTask<GrainAddress?> Lookup(GrainId grainId) => GetGrainLocator(grainId.Type).Lookup(grainId);
 
-        public async Task<GrainAddress?> Register(GrainAddress address, GrainAddress? previousRegistration)
+        public async Task<GrainAddress?> Register(
+            GrainAddress address,
+            GrainAddress? previousRegistration,
+            CancellationToken cancellationToken = default)
         {
             var grainLocator = GetGrainLocator(address.GrainId.Type);
             var metrics = RegistrationMetricTracker.Start(_directoryInstruments, grainLocator);
             try
             {
-                var result = await grainLocator.Register(address, previousRegistration);
+                var result = grainLocator is CachedGrainLocator cachedGrainLocator
+                    ? await cachedGrainLocator.Register(address, previousRegistration, cancellationToken)
+                    : await grainLocator.Register(address, previousRegistration).WaitAsync(cancellationToken);
                 metrics.RecordSucceeded();
                 return result;
             }

src/api/Orleans.Core.Abstractions/Orleans.Core.Abstractions.cs

@@ -1538,10 +1538,15 @@ public void Populate(System.IServiceProvider services, System.Type grainClass, R
 
     public partial interface IGrainDirectory
     {
+        System.Threading.Tasks.Task<Runtime.GrainAddress?> Lookup(Runtime.GrainId grainId, System.Threading.CancellationToken cancellationToken);
         System.Threading.Tasks.Task<Runtime.GrainAddress?> Lookup(Runtime.GrainId grainId);
+        System.Threading.Tasks.Task<Runtime.GrainAddress?> Register(Runtime.GrainAddress address, Runtime.GrainAddress? previousAddress, System.Threading.CancellationToken cancellationToken);
         System.Threading.Tasks.Task<Runtime.GrainAddress?> Register(Runtime.GrainAddress address, Runtime.GrainAddress? previousAddress);
+        System.Threading.Tasks.Task<Runtime.GrainAddress?> Register(Runtime.GrainAddress address, System.Threading.CancellationToken cancellationToken);
         System.Threading.Tasks.Task<Runtime.GrainAddress?> Register(Runtime.GrainAddress address);
+        System.Threading.Tasks.Task Unregister(Runtime.GrainAddress address, System.Threading.CancellationToken cancellationToken);
         System.Threading.Tasks.Task Unregister(Runtime.GrainAddress address);
+        System.Threading.Tasks.Task UnregisterSilos(System.Collections.Generic.List<Runtime.SiloAddress> siloAddresses, System.Threading.CancellationToken cancellationToken);
         System.Threading.Tasks.Task UnregisterSilos(System.Collections.Generic.List<Runtime.SiloAddress> siloAddresses);
     }
 }

test/Orleans.Core.Tests/Directory/CachedGrainLocatorTests.cs

@@ -41,6 +41,14 @@ public CachedGrainLocatorTests(ITestOutputHelper output)
             this.lifecycle = new SiloLifecycleSubject(this.loggerFactory.CreateLogger<SiloLifecycleSubject>());
 
             this.grainDirectory = Substitute.For<IGrainDirectory>();
+            this.grainDirectory
+                .Register(Arg.Any<GrainAddress>(), Arg.Any<GrainAddress?>(), Arg.Any<CancellationToken>())
+                .Returns(call => this.grainDirectory.Register(
+                    call.ArgAt<GrainAddress>(0),
+                    call.ArgAt<GrainAddress?>(1)));
+            this.grainDirectory
+                .Unregister(Arg.Any<GrainAddress>(), Arg.Any<CancellationToken>())
+                .Returns(call => this.grainDirectory.Unregister(call.ArgAt<GrainAddress>(0)));
             var services = new ServiceCollection()
                 .AddGrainDirectory(GrainDirectoryAttribute.DEFAULT_GRAIN_DIRECTORY, (sp, name) => this.grainDirectory)
                 .BuildServiceProvider();
@@ -444,13 +452,14 @@ public async Task RegisterWhenOtherEntryExistsButSiloIsDead()
             var expectedAddr = GenerateGrainAddress(expectedSilo);
             var outdatedAddr = GenerateGrainAddress(outdatedSilo);
 
-            // First returns the outdated entry, then the new one
-            this.grainDirectory.Register(expectedAddr, previousAddress: null).Returns(outdatedAddr, expectedAddr);
+            this.grainDirectory.Register(expectedAddr, previousAddress: null).Returns(outdatedAddr);
+            this.grainDirectory.Register(expectedAddr, previousAddress: outdatedAddr).Returns(expectedAddr);
 
             var actual = await this.grainLocator.Register(expectedAddr, previousAddress: null);
             Assert.Equal(expectedAddr, actual);
-            await this.grainDirectory.Received(2).Register(expectedAddr, previousAddress: null);
-            await this.grainDirectory.Received(1).Unregister(outdatedAddr);
+            await this.grainDirectory.Received(1).Register(expectedAddr, previousAddress: null);
+            await this.grainDirectory.Received(1).Register(expectedAddr, previousAddress: outdatedAddr);
+            await this.grainDirectory.DidNotReceive().Unregister(outdatedAddr);
 
             // Now should be in cache
             Assert.True(this.grainLocator.TryLookupInCache(expectedAddr.GrainId, out var result));

test/Orleans.Core.Tests/Directory/IGrainDirectoryTests.cs

@@ -0,0 +1,42 @@
+using Orleans.GrainDirectory;
+using Orleans.Runtime;
+using Xunit;
+
+namespace UnitTests.Directory;
+
+[TestCategory("BVT"), TestCategory("Directory")]
+public class IGrainDirectoryTests
+{
+    [Fact]
+    public async Task CancellationOverloads_DefaultToCancelableWait()
+    {
+        IGrainDirectory directory = new BlockingGrainDirectory();
+        var cancellationToken = new CancellationToken(canceled: true);
+
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(
+            () => directory.Register(default!, cancellationToken));
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(
+            () => directory.Register(default!, null, cancellationToken));
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(
+            () => directory.Unregister(default!, cancellationToken));
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(
+            () => directory.Lookup(default, cancellationToken));
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(
+            () => directory.UnregisterSilos([], cancellationToken));
+    }
+
+    private sealed class BlockingGrainDirectory : IGrainDirectory
+    {
+        public Task<GrainAddress?> Register(GrainAddress address) =>
+            new TaskCompletionSource<GrainAddress?>().Task;
+
+        public Task Unregister(GrainAddress address) =>
+            new TaskCompletionSource().Task;
+
+        public Task<GrainAddress?> Lookup(GrainId grainId) =>
+            new TaskCompletionSource<GrainAddress?>().Task;
+
+        public Task UnregisterSilos(List<SiloAddress> siloAddresses) =>
+            new TaskCompletionSource().Task;
+    }
+}