Fix grains not recycling after cancelling DelayDeactivation

src/Orleans.Core.Abstractions/Core/Grain.cs

@@ -135,10 +135,11 @@ protected void MigrateOnIdle()
 
     /// <summary>
     /// Delay Deactivation of this activation at least for the specified time duration.
-    /// A positive <c>timeSpan</c> value means “prevent GC of this activation for that time span”.
-    /// A negative <c>timeSpan</c> value means “cancel the previous setting of the DelayDeactivation call and make this activation behave based on the regular Activation Garbage Collection settings”.
-    /// DeactivateOnIdle method would undo / override any current “keep alive” setting,
-    /// making this grain immediately available for deactivation.
+    /// <para>A positive <c>timeSpan</c> value means “prevent GC of this activation for that time span”.</para>
+    /// <para>A <see cref="TimeSpan.Zero"/> value means “cancel the previous setting of the <see cref="DelayDeactivation(TimeSpan)"/> call and make this activation behave based on the regular Activation Garbage Collection settings”.</para>
+    /// <para>A <see cref="Timeout.InfiniteTimeSpan"/> value means “delay deactivation indefinitely”.</para>
+    /// <para><see cref="DeactivateOnIdle"/> would undo / override any current “keep alive” setting,
+    /// making this grain immediately available for deactivation.</para>
     /// </summary>
     protected void DelayDeactivation(TimeSpan timeSpan)
     {

src/Orleans.Runtime/Catalog/ActivationData.cs

@@ -482,13 +482,20 @@ public void DelayDeactivation(TimeSpan timespan)
     {
         if (timespan == TimeSpan.MaxValue || timespan == Timeout.InfiniteTimeSpan)
         {
-            // otherwise creates negative time.
+            // Adding these values to the current time would overflow, so use DateTime.MaxValue directly.
             KeepAliveUntil = DateTime.MaxValue;
         }
         else if (timespan <= TimeSpan.Zero)
         {
-            // reset any current keepAliveUntil
+            // Cancel the previous DelayDeactivation and revert to normal collection behavior.
+            // If there was an active keep-alive, reschedule collection so the grain can be collected
+            // after CollectionAgeLimit rather than waiting for the previously scheduled far-future time.
+            var hadActiveKeepAlive = KeepAliveUntil > GrainRuntime.TimeProvider.GetUtcNow().UtcDateTime;
             ResetKeepAliveRequest();
+            if (hadActiveKeepAlive)
+            {
+                _shared.InternalRuntime.ActivationCollector.TryRescheduleCollection(this);
+            }
         }
         else
         {

src/Orleans.Runtime/Core/InternalGrainRuntime.cs

@@ -19,7 +19,8 @@ internal class InternalGrainRuntime(
         CompatibilityDirectorManager compatibilityDirectorManager,
         IOptions<GrainCollectionOptions> collectionOptions,
         ILocalGrainDirectory localGrainDirectory,
-        IActivationWorkingSet activationWorkingSet)
+        IActivationWorkingSet activationWorkingSet,
+        ActivationCollector activationCollector)
     {
         public InsideRuntimeClient RuntimeClient { get; } = catalog.RuntimeClient;
         public MessageCenter MessageCenter { get; } = messageCenter;
@@ -31,5 +32,6 @@ internal class InternalGrainRuntime(
         public IOptions<GrainCollectionOptions> CollectionOptions { get; } = collectionOptions;
         public ILocalGrainDirectory LocalGrainDirectory { get; } = localGrainDirectory;
         public IActivationWorkingSet ActivationWorkingSet { get; } = activationWorkingSet;
+        public ActivationCollector ActivationCollector { get; } = activationCollector;
     }
 }

test/Grains/TestInternalGrainInterfaces/ActivationGCTestGrainInterfaces.cs

@@ -38,4 +38,10 @@ public interface IStatelessWorkerActivationCollectorTestGrain1 : IGrainWithGuidK
         Task Delay(TimeSpan dt);
         Task<string> IdentifyActivation();
     }
+
+    public interface IKeepAliveActivationGcTestGrain : IGrainWithGuidKey
+    {
+        Task SetKeepAlive(TimeSpan keepAlive);
+        Task CancelKeepAlive();
+    }
 }

test/Grains/TestInternalGrains/ActivationGCTestGrains.cs

@@ -94,4 +94,19 @@ public Task<string> IdentifyActivation()
         }
 
     }
+
+    public class KeepAliveActivationGcTestGrain : Grain, IKeepAliveActivationGcTestGrain
+    {
+        public Task SetKeepAlive(TimeSpan keepAlive)
+        {
+            DelayDeactivation(keepAlive);
+            return Task.CompletedTask;
+        }
+
+        public Task CancelKeepAlive()
+        {
+            DelayDeactivation(TimeSpan.Zero);
+            return Task.CompletedTask;
+        }
+    }
 }

test/Orleans.Runtime.Internal.Tests/ActivationsLifeCycleTests/ActivationCollectorTests.cs

@@ -529,6 +529,36 @@ public async Task ActivationCollectorShouldCollectByCollectionSpecificAgeLimitFo
             Assert.Equal(0, activationsNotCollected);
         }
 
+        [Fact, TestCategory("ActivationCollector"), TestCategory("Functional")]
+        public async Task ActivationCollectorShouldCollectAfterCancellingKeepAlive()
+        {
+            await Initialize(DEFAULT_IDLE_TIMEOUT);
+
+            var fullGrainTypeName = RuntimeTypeNameFormatter.Format(typeof(KeepAliveActivationGcTestGrain));
+
+            // Activate the grain and extend its lifetime to 5 minutes.
+            var grain = this.testCluster.GrainFactory.GetGrain<IKeepAliveActivationGcTestGrain>(Guid.NewGuid());
+            await grain.SetKeepAlive(TimeSpan.FromMinutes(5));
+
+            // Verify the grain is not collected while it has an active keep-alive.
+            await Task.Delay(WAIT_TIME);
+            int activationsWithKeepAlive = await TestUtils.GetActivationCount(this.testCluster.GrainFactory, fullGrainTypeName);
+            Assert.Equal(1, activationsWithKeepAlive);
+
+            // Cancel the keep-alive. The grain should now be collectable after the standard idle timeout.
+            await grain.CancelKeepAlive();
+
+            // Wait for the grain to become idle past the collection age limit.
+            await Task.Delay(DEFAULT_IDLE_TIMEOUT + DEFAULT_COLLECTION_QUANTUM);
+
+            // Force collection to deterministically verify the grain is now collectable.
+            var mgmtGrain = this.testCluster.GrainFactory.GetGrain<IManagementGrain>(0);
+            await mgmtGrain.ForceActivationCollection(DEFAULT_IDLE_TIMEOUT);
+
+            int activationsNotCollected = await TestUtils.GetActivationCount(this.testCluster.GrainFactory, fullGrainTypeName);
+            Assert.Equal(0, activationsNotCollected);
+        }
+
         [Fact, TestCategory("SlowBVT"), TestCategory("Timers")]
         public async Task NonReentrantGrainTimer_NoKeepAlive_Test()
         {