address PR feedback

Author: ledjon-behluli

src/Orleans.Runtime/Configuration/Options/GrainDirectoryOptions.cs

@@ -106,14 +106,12 @@ public enum CachingStrategyType
     /// Depending on the value of this, the duration is understood as:
     /// <list type="bullet">
     /// <item><c>SafetyLeaseHoldDuration > TimeSpan.Zero</c> - The lease duration is explicitly controlled by the user.</item>
-    /// <item><c>SafetyLeaseHoldDuration &lt;= TimeSpan.Zero</c> - The system computes a lease duration as:
-    /// <c>2 × <see cref="ClusterMembershipOptions.ProbeTimeout"/> × <see cref="ClusterMembershipOptions.NumMissedProbesLimit"/></c>.</item>
+    /// <item><c>SafetyLeaseHoldDuration = TimeSpan.Zero</c>. No leases are placed placed at all, effectively nullifying this safety option.</item>
+    /// <item><c>SafetyLeaseHoldDuration = null</c> - The system computes a lease duration as:
+    /// <c>2 × <see cref="ClusterMembershipOptions.ProbeTimeout"/> × <see cref="ClusterMembershipOptions.NumMissedProbesLimit"/></c>.
+    /// This is the default value, and is designed to be long enough to allow for failure detection and cluster stabilization.
+    /// </item>
     /// </list>
     /// </remarks>
-    public TimeSpan SafetyLeaseHoldDuration { get; set; }
-
-    /// <summary>
-    /// The default value for <see cref="SafetyLeaseHoldDuration"/>
-    /// </summary>
-    public static readonly TimeSpan DEFAULT_SAFETY_LEASE_HOLD_DURATION = TimeSpan.Zero;
+    public TimeSpan? SafetyLeaseHoldDuration { get; set; }
 }

src/Orleans.Runtime/GrainDirectory/DistributedGrainDirectory.cs

@@ -1,4 +1,5 @@
 using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Runtime.CompilerServices;
 using System.Timers;
 using Microsoft.Extensions.DependencyInjection;
@@ -93,12 +94,13 @@ public DistributedGrainDirectory(
         _serviceProvider = serviceProvider;
         _membershipService = membershipService;
         _logger = logger;
-        _leaseHoldDuration = directoryOptions.Value.SafetyLeaseHoldDuration;
 
-        if (_leaseHoldDuration <= TimeSpan.Zero)
+        _leaseHoldDuration = directoryOptions.Value.SafetyLeaseHoldDuration switch
         {
-            _leaseHoldDuration = 2 * membershipOptions.Value.ProbeTimeout * membershipOptions.Value.NumMissedProbesLimit;
-        }
+            null => 2 * membershipOptions.Value.ProbeTimeout * membershipOptions.Value.NumMissedProbesLimit,
+            TimeSpan duration when duration >= TimeSpan.Zero => duration,
+            _ => throw new InvalidOperationException("Lease hold duration must be non-negative.")
+        };
 
         var partitions = ImmutableArray.CreateBuilder<GrainDirectoryPartition>(DirectoryMembershipSnapshot.PartitionsPerSilo);
 
@@ -327,7 +329,11 @@ Task OnRuntimeInitializeStart(CancellationToken cancellationToken)
             WorkItemGroup.QueueAction(() =>
             {
                 _runTask = ProcessMembershipUpdates();
-                _leaseCleanupTask = RequestExpiredLeaseCleanups();
+
+                if (_leaseHoldDuration > TimeSpan.Zero)
+                {
+                    _leaseCleanupTask = RequestExpiredLeaseCleanups();
+                }
             });
 
             return Task.CompletedTask;
@@ -418,6 +424,8 @@ private async Task ProcessMembershipUpdates()
 
     private async Task RequestExpiredLeaseCleanups()
     {
+        Debug.Assert(_leaseHoldDuration > TimeSpan.Zero);
+
         // We request cleanups periodically to not let expired leases linger in the directory for too long.
         // We do it here as opposed to in the partitions to avoid having 30 (by default, maybe more) timers.
 

src/Orleans.Runtime/GrainDirectory/GrainDirectoryPartition.Interface.cs

@@ -1,3 +1,4 @@
+using System.Diagnostics;
 using System.Runtime.InteropServices;
 using Microsoft.Extensions.Logging;
 
@@ -19,46 +20,44 @@ async ValueTask<DirectoryResult<GrainAddress>> IGrainDirectoryPartition.Register
 
         DebugAssertOwnership(address.GrainId);
 
-        var utcNow = _timeProvider.GetUtcNow().UtcDateTime;
-        var rangeHash = address.GrainId.GetUniformHashCode();
-
-        // Range lease holds
-        for (var i = _rangeLeaseHolds.Count - 1; i >= 0; i--)
+        if (_leaseHoldDuration > TimeSpan.Zero)
         {
-            var (lockedRange, expiration) = _rangeLeaseHolds[i];
+            var utcNow = _timeProvider.GetUtcNow().UtcDateTime;
+            var rangeHash = address.GrainId.GetUniformHashCode();
 
-            if (utcNow >= expiration)
+            // Range lease holds
+            for (var i = _rangeLeaseHolds.Count - 1; i >= 0; i--)
             {
-                // We use this opportunity to cleanup this expired range lease hold.
-                _rangeLeaseHolds.RemoveAt(i);
-                continue;
-            }
+                var (lockedRange, expiration) = _rangeLeaseHolds[i];
 
-            // If it is still active, does it block this request?
-            if (lockedRange.Contains(rangeHash))
-            {
-                // We reject, the client should retry!
-                throw new DirectoryLeaseHoldException($"Range {lockedRange} is under a lease hold until {expiration - utcNow}.");
-            }
-        }
+                if (utcNow >= expiration)
+                {
+                    // We use this opportunity to cleanup this expired range lease hold.
+                    _rangeLeaseHolds.RemoveAt(i);
+                    continue;
+                }
 
-        // Grain lease holds
-        if (_grainLeaseHolds.TryGetValue(address.GrainId, out var tombstone))
-        {
-            if (utcNow >= tombstone.LeaseExpiration)
-            {
-                // We use this opportunity to cleanup this expired grain-specific lease hold.
-                _grainLeaseHolds.Remove(address.GrainId);
+                // If it is still active, does it block this request?
+                if (lockedRange.Contains(rangeHash))
+                {
+                    // We reject, the client should retry!
+                    throw new DirectoryLeaseHoldException($"Range {lockedRange} is under a lease hold until {expiration - utcNow}.");
+                }
             }
-            else
+
+            // Grain lease holds
+            if (_directory.TryGetValue(address.GrainId, out var existingActivation))
             {
-                // Is the new registration trying to point to the same dead silo?
-                // If yes, it is consistent, but dead; otherwise we must block.
-                if (!tombstone.DeadSilo.Equals(address.SiloAddress))
+                if (_siloLeaseHolds.TryGetValue(existingActivation.SiloAddress!, out var expiration) && utcNow < expiration)
                 {
-                    // The previous owner is dead, but the lease hasnt expired yet, we must reject. 
-                    // We can not guarantee the old activation is gone yet. The client should retry!
-                    throw new DirectoryLeaseHoldException($"Grain {address.GrainId} is under a lease hold until {tombstone.LeaseExpiration - utcNow}.");
+                    // This grain belongs to this parition, and the activation is sitting on a silo that has an active lease hold.
+                    // We need to check if the request include sthe previous activation id, and if it does its a valid update/override,
+                    // otherwise it's a new activation trying to "steal" the id while the lease is active, so we reject it!
+
+                    if (currentRegistration is null || !existingActivation.Matches(currentRegistration))
+                    {
+                        throw new DirectoryLeaseHoldException($"Silo {existingActivation.SiloAddress} is under a lease hold until {expiration - utcNow}.");
+                    }
                 }
             }
         }

src/Orleans.Runtime/GrainDirectory/GrainDirectoryPartition.cs

@@ -46,7 +46,7 @@ internal sealed partial class GrainDirectoryPartition : SystemTarget, IGrainDire
 
     private readonly TimeSpan _leaseHoldDuration;
     private readonly List<(RingRange Range, DateTime Expiration)> _rangeLeaseHolds = [];
-    private readonly Dictionary<GrainId, (SiloAddress DeadSilo, DateTime LeaseExpiration)> _grainLeaseHolds = [];
+    private readonly Dictionary<SiloAddress, DateTime> _siloLeaseHolds = [];
 
     /// <param name="partitionIndex">The index of this partition on this silo. Each silo hosts a fixed number of dynamically sized partitions.</param>
     public GrainDirectoryPartition(
@@ -258,37 +258,38 @@ private void OnSiloRemovedFromCluster(ClusterMember change, SiloStatus previousS
         // If it was ShuttingDown, it surrendered its ownership gracefully.
         // If it was Active (or Joining) and suddenly became Dead, it crashed.
 
-        var isUngraceful = previousStatus is not SiloStatus.ShuttingDown;
-        var expiration = _timeProvider.GetUtcNow().UtcDateTime.Add(_leaseHoldDuration);
-        var toRemove = new List<GrainAddress>();
+        if (previousStatus is not SiloStatus.ShuttingDown && _leaseHoldDuration > TimeSpan.Zero)
+        {
+            // Instead of just deleting, we mark it as tombstoned.
+            // This prevents a new activation on a healthy silo from registering 
+            // until we are sure the dead silo has actually stopped processing.
+
+            var expiration = _timeProvider.GetUtcNow().UtcDateTime.Add(_leaseHoldDuration);
 
-        foreach (var entry in _directory)
+            _siloLeaseHolds[change.SiloAddress] = expiration;
+
+            LogDebugLeaseHoldForSilo(_logger, change.SiloAddress, expiration);
+        }
+        else
         {
-            if (change.SiloAddress.Equals(entry.Value.SiloAddress))
-            {
-                toRemove.Add(entry.Value);
+            var toRemove = new List<GrainAddress>();
 
-                if (isUngraceful)
+            foreach (var entry in _directory)
+            {
+                if (change.SiloAddress.Equals(entry.Value.SiloAddress))
                 {
-                    // Instead of just deleting, we mark it as tombstoned.
-                    // This prevents a new activation on a healthy silo from registering 
-                    // until we are sure the dead silo has actually stopped processing.
-                    // So we block re-registration of this specific grain id if we did not already,
-                    // or extend based on the new expiration time.
-
-                    _grainLeaseHolds[entry.Key] = (change.SiloAddress, expiration);
-                    LogDebugLeaseHoldForGrain(_logger, entry.Key, change.SiloAddress, expiration);
+                    toRemove.Add(entry.Value);
                 }
             }
-        }
-
-        if (toRemove.Count > 0)
-        {
-            LogDebugDeletingEntries(_logger, toRemove.Count, change.SiloAddress);
 
-            foreach (var grainAddress in toRemove)
+            if (toRemove.Count > 0)
             {
-                DeregisterCore(grainAddress);
+                LogDebugDeletingEntries(_logger, toRemove.Count, change.SiloAddress);
+
+                foreach (var grainAddress in toRemove)
+                {
+                    DeregisterCore(grainAddress);
+                }
             }
         }
 
@@ -490,10 +491,13 @@ private async Task AcquireRangeAsync(DirectoryMembershipSnapshot previous, Direc
             var recovered = false;
             if (!success)
             {
-                // We pessimistically asssume if snapshot transfer failed, than safety is needed.
-                var expiration = _timeProvider.GetUtcNow().UtcDateTime.Add(_leaseHoldDuration);
-                _rangeLeaseHolds.Add((addedRange, expiration));
-                LogWarningLeaseHoldForRange(_logger, addedRange, expiration);
+                if (_leaseHoldDuration > TimeSpan.Zero)
+                {
+                    // We pessimistically asssume if snapshot transfer failed, than safety is needed.
+                    var expiration = _timeProvider.GetUtcNow().UtcDateTime.Add(_leaseHoldDuration);
+                    _rangeLeaseHolds.Add((addedRange, expiration));
+                    LogWarningLeaseHoldForRange(_logger, addedRange, expiration);
+                }
 
                 // Wait for previous versions to be unlocked before proceeding.
                 await WaitForRange(addedRange, previous.Version);
@@ -804,21 +808,31 @@ private void CleanupExpiredLeasesCore()
                 }
             }
 
-            if (_grainLeaseHolds.Count > 0)
+            if (_siloLeaseHolds.Count > 0)
             {
-                var expiredKeys = _grainLeaseHolds
-                    .Where(kvp => utcNow >= kvp.Value.LeaseExpiration)
+                var expiredSilos = _siloLeaseHolds
+                    .Where(kvp => utcNow >= kvp.Value)
                     .Select(kvp => kvp.Key)
                     .ToList();
 
-                foreach (var key in expiredKeys)
+                if (expiredSilos.Count > 0)
                 {
-                    _grainLeaseHolds.Remove(key);
-                }
+                    // These are the grains which we were supposed to have removed when the silo was marked as dead,
+                    // but we kept them around until we were sure the silo was actually dead.
 
-                if (expiredKeys.Count > 0)
-                {
-                    LogDebugPrunedExpiredGrainLeaseHolds(_logger, expiredKeys.Count);
+                    var toRemove = _directory.Where(kvp => expiredSilos.Contains(kvp.Value.SiloAddress!)).ToList();
+
+                    foreach (var kvp in toRemove)
+                    {
+                        _directory.Remove(kvp.Key);
+                    }
+
+                    foreach (var silo in expiredSilos)
+                    {
+                        _siloLeaseHolds.Remove(silo);
+                    }
+
+                    LogDebugPrunedExpiredSiloLeaseHolds(_logger, expiredSilos.Count, toRemove.Count);
                 }
             }
         }
@@ -834,21 +848,20 @@ private sealed record class PartitionSnapshotState(
 
     [LoggerMessage(
       Level = LogLevel.Debug,
-      Message = "Grain {GrainId} from silo {Silo} has been put under a lease until {Expiration}."
-    )]
-    private static partial void LogDebugLeaseHoldForGrain(ILogger logger, GrainId grainId, SiloAddress silo, DateTime expiration);
+      Message = "Placed lease hold on dead silo {SiloAddress} until {Expiration}.")]
+    private static partial void LogDebugLeaseHoldForSilo(ILogger logger, SiloAddress siloAddress, DateTime expiration);
+
+    [LoggerMessage(
+      Level = LogLevel.Debug,
+      Message = "Pruned {SiloCount} expired silo lease holds, removing {GrainCount} dead grain activations from the directory.")]
+    private static partial void LogDebugPrunedExpiredSiloLeaseHolds(ILogger logger, int siloCount, int grainCount);
 
     [LoggerMessage(
       Level = LogLevel.Warning,
       Message = "Grains in the range {Range} have been put under a lease until {Expiration}."
     )]
     private static partial void LogWarningLeaseHoldForRange(ILogger logger, RingRange range, DateTime expiration);
 
-    [LoggerMessage(
-        Level = LogLevel.Debug,
-        Message = "Pruned {Count} expired grain lease holds."
-    )]
-    private static partial void LogDebugPrunedExpiredGrainLeaseHolds(ILogger logger, int count);
 
     [LoggerMessage(
       Level = LogLevel.Debug,

test/Orleans.GrainDirectory.Tests/GrainDirectory/GrainDirectoryLeaseTests.cs

@@ -24,7 +24,8 @@ public class LeaseTestGrain : Grain, ILeaseTestGrain
 public class GrainDirectoryLeaseTests
 {
     public static readonly FakeTimeProvider TimeProvider = new(DateTime.UtcNow);
-    public static readonly TimeSpan LeaseHoldDuration = TimeSpan.FromSeconds(5); // The value doesnt matter we will advance time manually.
+    // The value doesnt matter we will advance time manually.
+    public static readonly TimeSpan LeaseHoldDuration = TimeSpan.FromSeconds(5);
 
     [Fact]
     public async Task BlocksReactivations_AfterUngracefulShutdown()