From 9aab3f71046840fd26c436e88b18f46bd4f85587 Mon Sep 17 00:00:00 2001
From: Reuben Bond <203839+ReubenBond@users.noreply.github.com>
Date: Fri, 8 Mar 2024 09:43:40 -0800
Subject: [PATCH] Use GrainDirectoryCacheFactory to construct a
 IGrainDirectoryCache (#8844) (#8898)

---------

Co-authored-by: Mostafa Abdalla <mostafa.abdalla@cm.com>
Co-authored-by: Reuben Bond <203839+ReubenBond@users.noreply.github.com>
(cherry picked from commit ab644bf2da188819c8a3a01fb8c7573b4131a437)

Co-authored-by: Mostafa Goher <Mostafa.Goher@gmail.com>
---
 .../GrainDirectory/CachedGrainLocator.cs            |  7 +++++--
 .../GrainDirectory/GrainDirectoryCacheFactory.cs    | 13 +++++++++++++
 .../Directory/CachedGrainLocatorTests.cs            | 11 +++++++++--
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/Orleans.Runtime/GrainDirectory/CachedGrainLocator.cs b/src/Orleans.Runtime/GrainDirectory/CachedGrainLocator.cs
index 7ab2aa0ced..c476d002ed 100644
--- a/src/Orleans.Runtime/GrainDirectory/CachedGrainLocator.cs
+++ b/src/Orleans.Runtime/GrainDirectory/CachedGrainLocator.cs
@@ -3,6 +3,7 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Extensions.Options;
 using Orleans.Configuration;
 using Orleans.GrainDirectory;
 using Orleans.Internal;
@@ -30,12 +31,14 @@ internal interface ITestAccessor
         MembershipVersion ITestAccessor.LastMembershipVersion { get; set; }
 
         public CachedGrainLocator(
+            IServiceProvider serviceProvider,
             GrainDirectoryResolver grainDirectoryResolver,
-            IClusterMembershipService clusterMembershipService)
+            IClusterMembershipService clusterMembershipService,
+            IOptions<GrainDirectoryOptions> grainDirectoryOptions)
         {
             this.grainDirectoryResolver = grainDirectoryResolver;
             this.clusterMembershipService = clusterMembershipService;
-            this.cache = new LRUBasedGrainDirectoryCache(GrainDirectoryOptions.DEFAULT_CACHE_SIZE, GrainDirectoryOptions.DEFAULT_MAXIMUM_CACHE_TTL);
+            this.cache = GrainDirectoryCacheFactory.CreateCustomGrainDirectoryCache(serviceProvider, grainDirectoryOptions.Value);
         }
 
         public async ValueTask<GrainAddress> Lookup(GrainId grainId)
diff --git a/src/Orleans.Runtime/GrainDirectory/GrainDirectoryCacheFactory.cs b/src/Orleans.Runtime/GrainDirectory/GrainDirectoryCacheFactory.cs
index e6832ae27e..100982050f 100644
--- a/src/Orleans.Runtime/GrainDirectory/GrainDirectoryCacheFactory.cs
+++ b/src/Orleans.Runtime/GrainDirectory/GrainDirectoryCacheFactory.cs
@@ -36,6 +36,19 @@ public static IGrainDirectoryCache CreateGrainDirectoryCache(IServiceProvider se
             }
         }
 
+        internal static IGrainDirectoryCache CreateCustomGrainDirectoryCache(IServiceProvider services, GrainDirectoryOptions options)
+        {
+            var grainDirectoryCache = services.GetService<IGrainDirectoryCache>();
+            if (grainDirectoryCache != null)
+            {
+                return grainDirectoryCache;
+            }
+            else
+            {
+                return new LRUBasedGrainDirectoryCache(options.CacheSize, options.MaximumCacheTTL);
+            }
+        }
+
         internal static AdaptiveDirectoryCacheMaintainer CreateGrainDirectoryCacheMaintainer(
             LocalGrainDirectory router,
             IGrainDirectoryCache cache,
diff --git a/test/NonSilo.Tests/Directory/CachedGrainLocatorTests.cs b/test/NonSilo.Tests/Directory/CachedGrainLocatorTests.cs
index 94ba121a55..f0c978c674 100644
--- a/test/NonSilo.Tests/Directory/CachedGrainLocatorTests.cs
+++ b/test/NonSilo.Tests/Directory/CachedGrainLocatorTests.cs
@@ -1,9 +1,13 @@
+// Ignore Spelling: Locator
+
 using System.Collections.Immutable;
 using System.Net;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
 using NSubstitute;
 using NSubstitute.ReceivedExtensions;
+using Orleans.Configuration;
 using Orleans.GrainDirectory;
 using Orleans.Metadata;
 using Orleans.Runtime;
@@ -19,7 +23,7 @@ public class CachedGrainLocatorTests
     {
         private readonly LoggerFactory loggerFactory;
         private readonly SiloLifecycleSubject lifecycle;
-
+        private readonly IOptions<GrainDirectoryOptions> grainDirectoryOptions;
         private readonly IGrainDirectory grainDirectory;
         private readonly GrainDirectoryResolver grainDirectoryResolver;
         private readonly MockClusterMembershipService mockMembershipService;
@@ -42,9 +46,12 @@ public CachedGrainLocatorTests(ITestOutputHelper output)
                 Array.Empty<IGrainDirectoryResolver>());
             this.mockMembershipService = new MockClusterMembershipService();
 
+            grainDirectoryOptions = Options.Create(new GrainDirectoryOptions());
             this.grainLocator = new CachedGrainLocator(
+                services,
                 this.grainDirectoryResolver, 
-                this.mockMembershipService.Target);
+                this.mockMembershipService.Target,
+                grainDirectoryOptions);
 
             this.grainLocator.Participate(this.lifecycle);
         }