Merge pull request #430 from alrz/use-typemanifestoptions

Avoid loading all types when collecting grain types
OrleansContrib · Jul 24, 2024 · ef58c7a · ef58c7a
2 parents 59ba91b + 20676af
commit ef58c7a
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 24 deletions.
diff --git a/OrleansDashboard/Implementation/Grains/DashboardGrain.cs b/OrleansDashboard/Implementation/Grains/DashboardGrain.cs
@@ -9,6 +9,7 @@
 using Orleans;
 using Orleans.Concurrency;
 using Orleans.Runtime;
+using Orleans.Serialization.Configuration;
 using OrleansDashboard.Implementation.Helpers;
 using OrleansDashboard.Metrics;
 using OrleansDashboard.Metrics.Details;
@@ -27,6 +28,7 @@ public class DashboardGrain : Grain, IDashboardGrain
         private readonly ISiloGrainClient siloGrainClient;
         private readonly DashboardCounters counters;
         private readonly GrainProfilerOptions grainProfilerOptions;
+        private readonly TypeManifestOptions typeManifestOptions;
         private readonly TimeSpan updateInterval;
         private bool isUpdating;
         private DateTime startTime = DateTime.UtcNow;
@@ -37,6 +39,7 @@ public class DashboardGrain : Grain, IDashboardGrain
         public DashboardGrain(
             IOptions<DashboardOptions> options,
             IOptions<GrainProfilerOptions> grainProfilerOptions,
+            IOptions<TypeManifestOptions> typeManifestOptions,
             ISiloDetailsProvider siloDetailsProvider,
             ISiloGrainClient siloGrainClient)
         {
@@ -45,6 +48,7 @@ public DashboardGrain(
 
             // Store the options to bypass the broadcase of the isEnabled flag.
             this.grainProfilerOptions = grainProfilerOptions.Value;
+            this.typeManifestOptions = typeManifestOptions.Value;
 
             // Do not allow smaller timers than 1000ms = 1sec.
             updateInterval = TimeSpan.FromMilliseconds(Math.Max(options.Value.CounterUpdateIntervalMs, 1000));
@@ -268,7 +272,7 @@ public async Task<Immutable<string>> GetGrainState(string id, string grainType)
 
             try
             {
-                var implementationType = GrainStateHelper.GetGrainType(grainType);
+                var implementationType = GrainStateHelper.GetGrainType(grainType, typeManifestOptions);
 
                 var mappedGrainId = GrainStateHelper.GetGrainId(id, implementationType);
                 object grainId = mappedGrainId.Item1;
@@ -343,8 +347,8 @@ public async Task<Immutable<string>> GetGrainState(string id, string grainType)
 
         public Task<Immutable<string[]>> GetGrainTypes()
         {
-            return Task.FromResult(GrainStateHelper.GetGrainTypes()
-                .Select(s => s.Namespace + "." + s.Name)
+            return Task.FromResult(typeManifestOptions.InterfaceImplementations
+                .Select(s => s.FullName)
                 .ToArray()
                 .AsImmutable());
         }

diff --git a/OrleansDashboard/Implementation/Helpers/GrainStateHelper.cs b/OrleansDashboard/Implementation/Helpers/GrainStateHelper.cs
@@ -4,21 +4,12 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Reflection;
+using Orleans.Serialization.Configuration;
 
 namespace OrleansDashboard.Implementation.Helpers
 {
     internal static class GrainStateHelper
     {
-        public static IEnumerable<Type> GetGrainTypes()
-        {
-            return AppDomain.CurrentDomain.GetAssemblies()
-                                     .SelectMany(s => s.GetTypes())
-                                     .Where(w => w.IsAssignableTo(typeof(IGrain))
-                                               && !w.Namespace.StartsWith("Orleans")
-                                               && w.IsClass
-                                               && !w.IsGenericType);
-        }
-
         public static (object, string) GetGrainId(string id, Type implementationType)
         {
             object grainId = null;
@@ -102,14 +93,10 @@ public static MethodInfo GenerateGetGrainMethod(IGrainFactory grainFactory, obje
             }
         }
 
-        public static Type GetGrainType(string grainType)
+        public static Type GetGrainType(string grainType, TypeManifestOptions typeManifestOptions)
         {
-            var _grainType = grainType.Split(".").Last();
-
-            return AppDomain.CurrentDomain.GetAssemblies()
-                                .SelectMany(s => s.GetTypes())
-                                .Where(w => w.Name.Equals(_grainType))
-                                .FirstOrDefault();
+            return typeManifestOptions.InterfaceImplementations
+                .FirstOrDefault(w => w.FullName.Equals(grainType));
         }
     }
 }
diff --git a/Tests/UnitTests/GrainStateTests.cs b/Tests/UnitTests/GrainStateTests.cs
@@ -30,12 +30,12 @@ public void Dispose()
         }
 
         [Fact]
-        public void TestGetGrainsTypes()
+        public async Task TestGetGrainsTypes()
         {
-            var types = GrainStateHelper.GetGrainTypes()
-                        .Select(s => s.Namespace + "." + s.Name);
+            var dashboardGrain = _cluster.GrainFactory.GetGrain<IDashboardGrain>(1);
+            var types = await dashboardGrain.GetGrainTypes();
 
-            Assert.Contains("TestGrains.TestStateInMemoryGrain", types);
+            Assert.Contains("TestGrains.TestStateInMemoryGrain", types.Value);
         }
 
         [Fact]