PR feedback 2

dotnet · Sep 21, 2024 · c86456c · c86456c
1 parent a2cba2c
commit c86456c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 23 deletions.
diff --git a/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancerWorker.cs b/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancerWorker.cs
@@ -150,8 +150,7 @@ public void RemoveSilo(SiloAddress silo)
 
     public void SiloStatisticsChangeNotification(SiloAddress address, SiloRuntimeStatistics statistics)
     {
-        ref var stats = ref CollectionsMarshal.GetValueRefOrAddDefault(_siloStatistics, address, out _);
-        stats = new(statistics.EnvironmentStatistics.MemoryUsageBytes, statistics.ActivationCount);
+        _siloStatistics[address] = new(statistics.EnvironmentStatistics.MemoryUsageBytes, statistics.ActivationCount);
     }
 
     public ValueTask<RebalancingReport> GetReport() => new(BuildReport());

diff --git a/src/Orleans.Runtime/Silo/SiloControl.cs b/src/Orleans.Runtime/Silo/SiloControl.cs
@@ -273,36 +273,31 @@ public Task<List<GrainId>> GetActiveGrains(GrainType grainType)
 
         public Task MigrateRandomActivations(SiloAddress target, int count)
         {
-            var statistics = GetDetailedGrainStatisticsCore();
-            var grainIds = statistics.Select(x => x.GrainId).ToList();
-            var idsToMigrate = new HashSet<GrainId>();
+            ArgumentNullException.ThrowIfNull(target);
+            ArgumentOutOfRangeException.ThrowIfNegative(count);
             var migrationContext = new Dictionary<string, object>()
             {
                 [IPlacementDirector.PlacementHintKey] = target
             };
 
-            // Loop until we've migrated the desired count of activations or run out of activations to try
-            while (idsToMigrate.Count < count && grainIds.Count > 0)
-            {   
-                var remainingCount = count - idsToMigrate.Count; // How many more we still need to migrate?
-                var candidates = Random.Shared.GetItems(grainIds.ToArray(), remainingCount); // Randomly select up to the remaining count
-
-                foreach (var grainId in candidates)
+            // Loop until we've migrated the desired count of activations or run out of activations to try.
+            // Note that we have a weak pseudorandom enumeration here, and lossy counting: this is not a precise
+            // or deterministic operation.
+            var remainingCount = count;
+            foreach (var (grainId, grainContext) in activationDirectory)
+            {
+                if (!_migratabilityChecker.IsMigratable(grainId.Type))
                 {
-                    grainIds.Remove(grainId); // Remove it from the list to avoid re-selection
-
-                    if (_migratabilityChecker.IsMigratable(grainId.Type) &&
-                        activationDirectory.FindTarget(grainId) is { } activation)
-                    {
-                        activation.Migrate(migrationContext);
-                        idsToMigrate.Add(grainId);
+                    continue;
+                }
 
-                        if (idsToMigrate.Count >= count)
-                            break;
-                    }
+                if (--remainingCount <= 0)
+                {
+                    break;
                 }
-            }
 
+                grainContext.Migrate(migrationContext);
+            }
 
             return Task.CompletedTask;
         }