Makes ImmovableAttribute configurable (#9205)

---------

Co-authored-by: Ledjon Behluli <Ledjon@notiphy.io>

src/Orleans.Core.Abstractions/Placement/PlacementAttribute.cs

@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using Orleans.Metadata;
 using Orleans.Runtime;
+using static Orleans.Placement.ImmovableAttribute;
 
 namespace Orleans.Placement
 {
@@ -113,14 +114,40 @@ public ResourceOptimizedPlacementAttribute() :
     }
 
     /// <summary>
-    /// Ensures that when active-rebalancing is enabled, activations of this grain type will not be migrated automatically.
+    /// Ensures that activations of this grain type will not be migrated automatically.
     /// </summary>
     /// <remarks>Activations can still be migrated by user initiated code.</remarks>
     [AttributeUsage(AttributeTargets.Class, AllowMultiple = false)]
-    public sealed class ImmovableAttribute : Attribute, IGrainPropertiesProviderAttribute
+    public sealed class ImmovableAttribute(ImmovableKind kind = ImmovableKind.Any)
+        : Attribute, IGrainPropertiesProviderAttribute
     {
+        /// <summary>
+        /// The kind of immovability.
+        /// </summary>
+        public ImmovableKind Kind { get; } = kind;
+
         /// <inheritdoc/>
         public void Populate(IServiceProvider services, Type grainClass, GrainType grainType, Dictionary<string, string> properties)
-            => properties[WellKnownGrainTypeProperties.Immovable] = "true";
+            => properties[WellKnownGrainTypeProperties.Immovable] = ((byte)Kind).ToString();
+    }
+
+    /// <summary>
+    /// Emphasizes that immovability is restricted to certain components.
+    /// </summary>
+    [Flags]
+    public enum ImmovableKind : byte
+    {
+        /// <summary>
+        /// Activations of this grain type will not be migrated by the repartitioner.
+        /// </summary>
+        Repartitioner = 1,
+        /// <summary>
+        /// Activations of this grain type will not be migrated by the rebalancer.
+        /// </summary>
+        Rebalancer = 2,
+        /// <summary>
+        /// Activations of this grain type will not be migrated by anything.
+        /// </summary>
+        Any = Repartitioner | Rebalancer
     }
 }

src/Orleans.Runtime/Placement/GrainMigratabilityChecker.cs

@@ -1,4 +1,5 @@
 using Orleans.Metadata;
+using Orleans.Placement;
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Frozen;
@@ -13,31 +14,39 @@ internal sealed class GrainMigratabilityChecker(
     IClusterManifestProvider clusterManifestProvider,
     TimeProvider timeProvider)
 {
+    // We override equality and hashcode as this type is used as the dictionary key,
+    // and record structs use default equality comparer, which for an enum is not that great for performance.
+    private readonly record struct StatusKey(GrainType Type, ImmovableKind Kind) : IEquatable<StatusKey>
+    {
+        public bool Equals(StatusKey other) => Type == other.Type && Kind == other.Kind;
+        public override int GetHashCode() => HashCode.Combine(Type.GetUniformHashCode(), Kind);
+    }
+
     private readonly GrainManifest _localManifest = clusterManifestProvider.LocalGrainManifest;
     private readonly PlacementStrategyResolver _strategyResolver = strategyResolver;
     private readonly TimeProvider _timeProvider = timeProvider;
-    private readonly ConcurrentDictionary<uint, bool> _migratableStatuses = new();
-    private FrozenDictionary<uint, bool>? _migratableStatusesCache;
+    private readonly ConcurrentDictionary<StatusKey, bool> _migratableStatuses = new();
+    private FrozenDictionary<StatusKey, bool>? _migratableStatusesCache;
     private long _lastRegeneratedCacheTimestamp = timeProvider.GetTimestamp();
 
-    public bool IsMigratable(GrainType grainType)
+    public bool IsMigratable(GrainType grainType, ImmovableKind expectedKind)
     {
-        var hash = grainType.GetUniformHashCode();
-        if (_migratableStatusesCache is { } cache && cache.TryGetValue(hash, out var isMigratable))
+        var statusKey = new StatusKey(grainType, expectedKind);
+        if (_migratableStatusesCache is { } cache && cache.TryGetValue(statusKey, out var isMigratable))
         {
             return isMigratable;
         }
 
-        return IsMigratableRare(grainType, hash);
+        return IsMigratableRare(grainType, statusKey);
 
-        bool IsMigratableRare(GrainType grainType, uint hash)
+        bool IsMigratableRare(GrainType grainType, StatusKey statusKey)
         {
             // _migratableStatuses holds statuses for each grain type if its migratable type or not, so we can make fast lookups.
             // since we don't anticipate a huge number of grain *types*, i think its just fine to have this in place as fast-check.
-            if (!_migratableStatuses.TryGetValue(hash, out var isMigratable))
+            if (!_migratableStatuses.TryGetValue(statusKey, out var isMigratable))
             {
                 isMigratable = !(grainType.IsClient() || grainType.IsSystemTarget() || grainType.IsGrainService() || IsStatelessWorker(grainType) || IsImmovable(grainType));
-                _migratableStatuses.TryAdd(hash, isMigratable);
+                _migratableStatuses.TryAdd(statusKey, isMigratable);
             }
 
             // Regenerate the cache periodically.
@@ -61,9 +70,19 @@ bool IsImmovable(GrainType grainType)
             if (_localManifest.Grains.TryGetValue(grainType, out var props))
             {
                 // If there is no 'Immovable' property, it is not immovable.
+                if (!props.Properties.TryGetValue(WellKnownGrainTypeProperties.Immovable, out var value))
+                {
+                    return false;
+                }
+
                 // If the value fails to parse, assume it's immovable.
-                // If the value is true, it's immovable.
-                return props.Properties.TryGetValue(WellKnownGrainTypeProperties.Immovable, out var value) && (!bool.TryParse(value, out var result) || result);
+                if (!byte.TryParse(value, out var actualKindValue))
+                {
+                    return true;
+                }
+
+                // It is immovable, but does the kind match with the parameter.
+                return ((ImmovableKind)actualKindValue & expectedKind) == expectedKind;
             }
 
             // Assume unknown grains are immovable.

src/Orleans.Runtime/Placement/Repartitioning/RepartitionerMessageFilter.cs

@@ -1,5 +1,7 @@
 #nullable enable
 
+using Orleans.Placement;
+
 namespace Orleans.Runtime.Placement.Repartitioning;
 
 internal interface IRepartitionerMessageFilter
@@ -22,8 +24,8 @@ public bool IsAcceptable(Message message, out bool isSenderMigratable, out bool
             return false;
         }
 
-        isSenderMigratable = checker.IsMigratable(message.SendingGrain.Type);
-        isTargetMigratable = checker.IsMigratable(message.TargetGrain.Type);
+        isSenderMigratable = checker.IsMigratable(message.SendingGrain.Type, ImmovableKind.Repartitioner);
+        isTargetMigratable = checker.IsMigratable(message.TargetGrain.Type, ImmovableKind.Repartitioner);
 
         // If both are not migratable types we ignore this. But if one of them is not, then we allow passing, as we wish to move grains closer to them, as with any type of grain.
         return isSenderMigratable || isTargetMigratable;

src/Orleans.Runtime/Silo/SiloControl.cs

@@ -10,6 +10,7 @@
 using Orleans.Configuration;
 using Orleans.GrainDirectory;
 using Orleans.Metadata;
+using Orleans.Placement;
 using Orleans.Providers;
 using Orleans.Runtime.GrainDirectory;
 using Orleans.Runtime.Placement;
@@ -307,7 +308,7 @@ public Task MigrateRandomActivations(SiloAddress target, int count)
             var remainingCount = count;
             foreach (var (grainId, grainContext) in activationDirectory)
             {
-                if (!_migratabilityChecker.IsMigratable(grainId.Type))
+                if (!_migratabilityChecker.IsMigratable(grainId.Type, ImmovableKind.Rebalancer))
                 {
                     continue;
                 }

test/TesterInternal/ActivationRepartitioningTests/CustomToleranceTests.cs

@@ -213,7 +213,7 @@ public override Task OnDeactivateAsync(DeactivationReason reason, CancellationTo
     /// This is simply to achieve initial balance between the 2 silos, as by default the primary
     /// will have 1 more activation than the secondary. That activations is 'sys.svc.clustering.dev'
     /// </summary>
-    [Immovable]
+    [Immovable(ImmovableKind.Repartitioner)]
     public class X : Grain, IX
     {
         public Task Ping() => Task.CompletedTask;

test/TesterInternal/ActivationRepartitioningTests/DefaultToleranceTests.cs

@@ -461,7 +461,7 @@ public async Task Ping(Scenario scenario, SiloAddress silo2)
         }
     }
 
-    [Immovable]
+    [Immovable(ImmovableKind.Repartitioner)]
     public class CImmovable : GrainBase, ICImmovable
     {
         public Task Ping(Scenario scenario) =>
@@ -482,7 +482,7 @@ public Task Ping(Scenario scenario) =>
             };
     }
 
-    [Immovable]
+    [Immovable(ImmovableKind.Repartitioner)]
     public class SP : GrainBase, ISP
     {
         // We are just 'Immovable' on this type, because we just want it to push messages to the stream,