Fix RealmEntityManagerFactory to reuse existing EntityCache

apache · Nov 26, 2024 · b9ded77 · b9ded77
1 parent 3475d24
commit b9ded77
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 13 deletions.
diff --git a/polaris-service/src/main/java/org/apache/polaris/service/PolarisApplication.java b/polaris-service/src/main/java/org/apache/polaris/service/PolarisApplication.java
@@ -84,6 +84,7 @@
 import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
 import org.apache.polaris.core.persistence.cache.EntityCache;
 import org.apache.polaris.core.persistence.cache.EntityCacheGrantManager;
+import org.apache.polaris.core.persistence.cache.RealmEntityCacheFactory;
 import org.apache.polaris.service.admin.PolarisServiceImpl;
 import org.apache.polaris.service.admin.api.PolarisCatalogsApi;
 import org.apache.polaris.service.admin.api.PolarisPrincipalRolesApi;
@@ -185,8 +186,14 @@ public void run(PolarisApplicationConfig configuration, Environment environment)
     if (metaStoreManagerFactory instanceof ConfigurationStoreAware) {
       ((ConfigurationStoreAware) metaStoreManagerFactory).setConfigurationStore(configurationStore);
     }
+    ConcurrentHashMap<RealmContext, EntityCache> realmEntityCache = new ConcurrentHashMap<>();
+    RealmEntityCacheFactory cacheFactory =
+        realm ->
+            realmEntityCache.computeIfAbsent(
+                realm,
+                (r) -> new EntityCache(metaStoreManagerFactory.getOrCreateMetaStoreManager(realm)));
     RealmEntityManagerFactory entityManagerFactory =
-        new RealmEntityManagerFactory(metaStoreManagerFactory);
+        new RealmEntityManagerFactory(metaStoreManagerFactory, cacheFactory);
     CallContextResolver callContextResolver = configuration.getCallContextResolver();
     callContextResolver.setMetaStoreManagerFactory(metaStoreManagerFactory);
     if (callContextResolver instanceof ConfigurationStoreAware csa) {
@@ -230,16 +237,9 @@ public void run(PolarisApplicationConfig configuration, Environment environment)
         new PolarisCallContextCatalogFactory(
             entityManagerFactory, metaStoreManagerFactory, taskExecutor, fileIOFactory);
 
-    ConcurrentHashMap<RealmContext, EntityCache> realmEntityCache = new ConcurrentHashMap<>();
     PolarisGrantManager.Factory factory =
         new EntityCacheGrantManager.EntityCacheGrantManagerFactory(
-            metaStoreManagerFactory,
-            realm ->
-                realmEntityCache.computeIfAbsent(
-                    realm,
-                    (r) ->
-                        new EntityCache(
-                            metaStoreManagerFactory.getOrCreateMetaStoreManager(realm))));
+            metaStoreManagerFactory, cacheFactory);
     PolarisAuthorizer authorizer = new PolarisAuthorizerImpl(configurationStore, factory);
     IcebergCatalogAdapter catalogAdapter =
         new IcebergCatalogAdapter(

diff --git a/...is-service/src/main/java/org/apache/polaris/service/config/RealmEntityManagerFactory.java b/...is-service/src/main/java/org/apache/polaris/service/config/RealmEntityManagerFactory.java
@@ -23,7 +23,8 @@
 import org.apache.polaris.core.context.RealmContext;
 import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
 import org.apache.polaris.core.persistence.PolarisEntityManager;
-import org.apache.polaris.core.persistence.cache.EntityCache;
+import org.apache.polaris.core.persistence.cache.RealmEntityCacheFactory;
+import org.jetbrains.annotations.NotNull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -34,28 +35,32 @@ public class RealmEntityManagerFactory {
 
   // Key: realmIdentifier
   private final Map<String, PolarisEntityManager> cachedEntityManagers = new ConcurrentHashMap<>();
+  private final RealmEntityCacheFactory entityCacheFactory;
 
   // Subclasses for test injection.
   protected RealmEntityManagerFactory() {
     this.metaStoreManagerFactory = null;
+    this.entityCacheFactory = null;
   }
 
-  public RealmEntityManagerFactory(MetaStoreManagerFactory metaStoreManagerFactory) {
+  public RealmEntityManagerFactory(
+      @NotNull MetaStoreManagerFactory metaStoreManagerFactory,
+      @NotNull RealmEntityCacheFactory entityCacheFactory) {
     this.metaStoreManagerFactory = metaStoreManagerFactory;
+    this.entityCacheFactory = entityCacheFactory;
   }
 
   public PolarisEntityManager getOrCreateEntityManager(RealmContext context) {
     String realm = context.getRealmIdentifier();
 
     LOGGER.debug("Looking up PolarisEntityManager for realm {}", realm);
-
     return cachedEntityManagers.computeIfAbsent(
         realm,
         r -> {
           LOGGER.info("Initializing new PolarisEntityManager for realm {}", r);
           return new PolarisEntityManager(
               metaStoreManagerFactory.getOrCreateMetaStoreManager(context),
-              new EntityCache(metaStoreManagerFactory.getOrCreateMetaStoreManager(context)),
+              entityCacheFactory.getOrCreateEntityCache(context),
               metaStoreManagerFactory.getOrCreateStorageCredentialCache(context));
         });
   }