perf(runtime): reduce activation working set entry overhead

src/Orleans.Runtime/Catalog/ActivationData.cs

@@ -32,7 +32,7 @@ internal sealed partial class ActivationData :
     IGrainContext,
     ICollectibleGrainContext,
     IGrainExtensionBinder,
-    IActivationWorkingSetMember,
+    IActivationWorkingSetMemberStatus,
     IGrainTimerRegistry,
     IGrainManagementExtension,
     IGrainCallCancellationExtension,
@@ -41,6 +41,11 @@ internal sealed partial class ActivationData :
     IDisposable
 {
     private const string GrainAddressMigrationContextKey = "sys.addr";
+    // Activation lifecycle and working-set CLOCK state share one byte. All writes occur while holding the activation lock.
+    private const byte ActivationStateMask = 0b0000_0111;
+    private const byte IsInWorkingSetMask = 0b0000_1000;
+    private const byte IsIdleInWorkingSetMask = 0b0001_0000;
+    private const byte WasRemovedByCollectionMask = 0b0010_0000;
     private readonly GrainTypeSharedContext _shared;
     private readonly IServiceScope _serviceScope;
     private readonly WorkItemGroup _workItemGroup;
@@ -50,7 +55,7 @@ internal sealed partial class ActivationData :
     private GrainLifecycle? _lifecycle;
     private Queue<object>? _pendingOperations;
     private Message? _blockingRequest;
-    private bool _isInWorkingSet = true;
+    private byte _status = IsInWorkingSetMask;
     private CoarseStopwatch _busyDuration;
     private CoarseStopwatch _idleDuration;
     private GrainReference? _selfReference;
@@ -146,7 +151,7 @@ public void Start(IGrainActivator grainActivator)
     public object? GrainInstance { get; private set; }
     public GrainAddress Address { get; private set; }
     public GrainReference GrainReference => _selfReference ??= _shared.GrainReferenceActivator.CreateReference(GrainId, default);
-    public ActivationState State { get; private set; } = ActivationState.Creating;
+    public ActivationState State => (ActivationState)(Volatile.Read(ref _status) & ActivationStateMask);
     public PlacementStrategy PlacementStrategy => _shared.PlacementStrategy;
     public DateTime CollectionTicket { get; set; }
     public IServiceProvider ActivationServices => _serviceScope.ServiceProvider;
@@ -167,6 +172,24 @@ public IGrainLifecycle ObservableLifecycle
     public DateTime KeepAliveUntil { get; set; } = DateTime.MinValue;
     public bool IsValid => State is ActivationState.Valid;
 
+    private bool IsInWorkingSet
+    {
+        get => (Volatile.Read(ref _status) & IsInWorkingSetMask) != 0;
+        set => SetStatusFlag(IsInWorkingSetMask, value);
+    }
+
+    private bool IsIdleInWorkingSet
+    {
+        get => (Volatile.Read(ref _status) & IsIdleInWorkingSetMask) != 0;
+        set => SetStatusFlag(IsIdleInWorkingSetMask, value);
+    }
+
+    private bool WasRemovedByCollection
+    {
+        get => (Volatile.Read(ref _status) & WasRemovedByCollectionMask) != 0;
+        set => SetStatusFlag(WasRemovedByCollectionMask, value);
+    }
+
     // Currently, the only supported multi-activation grain is one using the StatelessWorkerPlacement strategy.
     internal bool IsStatelessWorker => PlacementStrategy is StatelessWorkerPlacement;
 
@@ -404,7 +427,14 @@ internal void SetGrainInstance(object grainInstance)
 
     public void SetState(ActivationState state)
     {
-        State = state;
+        Debug.Assert(Monitor.IsEntered(this));
+        _status = (byte)((_status & ~ActivationStateMask) | (byte)state);
+    }
+
+    private void SetStatusFlag(byte mask, bool value)
+    {
+        Debug.Assert(Monitor.IsEntered(this));
+        _status = value ? (byte)(_status | mask) : (byte)(_status & ~mask);
     }
 
     /// <summary>
@@ -966,13 +996,34 @@ public TExtensionInterface GetExtension<TExtensionInterface>()
     bool IActivationWorkingSetMember.IsCandidateForRemoval(bool wouldRemove)
     {
         const int IdlenessLowerBound = 10_000;
-        lock (this)
+        Debug.Assert(Monitor.IsEntered(this));
+        var inactive = IsInactive && _idleDuration.ElapsedMilliseconds > IdlenessLowerBound;
+
+        WasRemovedByCollection = wouldRemove && inactive;
+        return inactive;
+    }
+
+    bool IActivationWorkingSetMemberStatus.IsInWorkingSet
+    {
+        get => IsInWorkingSet;
+        set
         {
-            var inactive = IsInactive && _idleDuration.ElapsedMilliseconds > IdlenessLowerBound;
+            Debug.Assert(Monitor.IsEntered(this));
+            IsInWorkingSet = value;
+            if (value)
+            {
+                WasRemovedByCollection = false;
+            }
+        }
+    }
 
-            // This instance will remain in the working set if it is either not pending removal or if it is currently active.
-            _isInWorkingSet = !wouldRemove || !inactive;
-            return inactive;
+    bool IActivationWorkingSetMemberStatus.IsIdle
+    {
+        get => IsIdleInWorkingSet;
+        set
+        {
+            Debug.Assert(Monitor.IsEntered(this));
+            IsIdleInWorkingSet = value;
         }
     }
 
@@ -1482,13 +1533,14 @@ private void OnCompletedRequest(Message message)
 
             // If the message is meant to keep the activation active, reset the idle timer and ensure the activation
             // is in the activation working set.
-            if (message.IsKeepAlive)
+            if (message.IsKeepAlive && State is ActivationState.Valid)
             {
                 _idleDuration = CoarseStopwatch.StartNew();
+                IsIdleInWorkingSet = false;
 
-                if (!_isInWorkingSet)
+                if (!IsInWorkingSet)
                 {
-                    _isInWorkingSet = true;
+                    IsInWorkingSet = true;
                     _shared.InternalRuntime.ActivationWorkingSet.OnActive(this);
                 }
             }
@@ -2034,7 +2086,7 @@ private async Task FinishDeactivating(Command.Deactivate deactivateCommand, Canc
                 deactivationMetrics = deactivationMetrics.Migration();
                 _shared.CatalogInstruments.ActivationShutdownViaMigration();
             }
-            else if (_isInWorkingSet)
+            else if (!WasRemovedByCollection)
             {
                 deactivationMetrics = deactivationMetrics.DeactivateOnIdle();
                 _shared.CatalogInstruments.ActivationShutdownViaDeactivateOnIdle();