Make StorageCredentialCache safe for mutli-realm usage

Author: XN137

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java

@@ -224,8 +224,7 @@ public synchronized StorageCredentialCache getOrCreateStorageCredentialCache(
       RealmContext realmContext) {
     if (!storageCredentialCacheMap.containsKey(realmContext.getRealmIdentifier())) {
       storageCredentialCacheMap.put(
-          realmContext.getRealmIdentifier(),
-          new StorageCredentialCache(realmContext, configurationStore));
+          realmContext.getRealmIdentifier(), new StorageCredentialCache(configurationStore));
     }
 
     return storageCredentialCacheMap.get(realmContext.getRealmIdentifier());

polaris-core/src/main/java/org/apache/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java

@@ -182,8 +182,7 @@ public synchronized StorageCredentialCache getOrCreateStorageCredentialCache(
       RealmContext realmContext) {
     if (!storageCredentialCacheMap.containsKey(realmContext.getRealmIdentifier())) {
       storageCredentialCacheMap.put(
-          realmContext.getRealmIdentifier(),
-          new StorageCredentialCache(realmContext, configurationStore));
+          realmContext.getRealmIdentifier(), new StorageCredentialCache(configurationStore));
     }
 
     return storageCredentialCacheMap.get(realmContext.getRealmIdentifier());

polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java

@@ -50,13 +50,10 @@ public class StorageCredentialCache {
   private static final long CACHE_MAX_NUMBER_OF_ENTRIES = 10_000L;
 
   private final LoadingCache<StorageCredentialCacheKey, StorageCredentialCacheEntry> cache;
-  private final RealmContext realmContext;
   private final PolarisConfigurationStore configurationStore;
 
   /** Initialize the creds cache */
-  public StorageCredentialCache(
-      RealmContext realmContext, PolarisConfigurationStore configurationStore) {
-    this.realmContext = realmContext;
+  public StorageCredentialCache(PolarisConfigurationStore configurationStore) {
     this.configurationStore = configurationStore;
     cache =
         Caffeine.newBuilder()
@@ -69,7 +66,7 @@ public StorageCredentialCache(
                               0,
                               Math.min(
                                   (entry.getExpirationTime() - System.currentTimeMillis()) / 2,
-                                  this.maxCacheDurationMs()));
+                                  entry.getMaxCacheDurationMs()));
                       return Duration.ofMillis(expireAfterMillis);
                     }))
             .build(
@@ -80,7 +77,7 @@ public StorageCredentialCache(
   }
 
   /** How long credentials should remain in the cache. */
-  private long maxCacheDurationMs() {
+  private long maxCacheDurationMs(RealmContext realmContext) {
     var cacheDurationSeconds =
         configurationStore.getConfiguration(
             realmContext, FeatureConfiguration.STORAGE_CREDENTIAL_CACHE_DURATION_SECONDS);
@@ -123,26 +120,27 @@ public AccessConfig getOrGenerateSubScopeCreds(
     }
     StorageCredentialCacheKey key =
         new StorageCredentialCacheKey(
+            callCtx.getRealmContext().getRealmIdentifier(),
             polarisEntity,
             allowListOperation,
             allowedReadLocations,
-            allowedWriteLocations,
-            callCtx);
+            allowedWriteLocations);
     LOGGER.atDebug().addKeyValue("key", key).log("subscopedCredsCache");
     Function<StorageCredentialCacheKey, StorageCredentialCacheEntry> loader =
         k -> {
           LOGGER.atDebug().log("StorageCredentialCache::load");
           ScopedCredentialsResult scopedCredentialsResult =
               credentialVendor.getSubscopedCredsForEntity(
-                  k.getCallContext(),
+                  callCtx,
                   k.getCatalogId(),
                   k.getEntityId(),
                   polarisEntity.getType(),
                   k.isAllowedListAction(),
                   k.getAllowedReadLocations(),
                   k.getAllowedWriteLocations());
           if (scopedCredentialsResult.isSuccess()) {
-            return new StorageCredentialCacheEntry(scopedCredentialsResult);
+            long maxCacheDurationMs = maxCacheDurationMs(callCtx.getRealmContext());
+            return new StorageCredentialCacheEntry(scopedCredentialsResult, maxCacheDurationMs);
           }
           LOGGER
               .atDebug()

polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheEntry.java

@@ -32,12 +32,19 @@ public class StorageCredentialCacheEntry {
   public final EnumMap<StorageAccessProperty, String> credsMap;
 
   private final ScopedCredentialsResult scopedCredentialsResult;
+  private final long maxCacheDurationMs;
 
-  public StorageCredentialCacheEntry(ScopedCredentialsResult scopedCredentialsResult) {
+  public StorageCredentialCacheEntry(
+      ScopedCredentialsResult scopedCredentialsResult, long maxCacheDurationMs) {
     this.scopedCredentialsResult = scopedCredentialsResult;
+    this.maxCacheDurationMs = maxCacheDurationMs;
     this.credsMap = scopedCredentialsResult.getCredentials();
   }
 
+  public long getMaxCacheDurationMs() {
+    return maxCacheDurationMs;
+  }
+
   /** Get the expiration time in millisecond for the cached entry */
   public long getExpirationTime() {
     if (credsMap.containsKey(StorageAccessProperty.GCS_ACCESS_TOKEN_EXPIRES_AT)) {

polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheKey.java

@@ -18,16 +18,14 @@
  */
 package org.apache.polaris.core.storage.cache;
 
-import jakarta.annotation.Nullable;
 import java.util.Objects;
 import java.util.Set;
-import org.apache.polaris.core.PolarisCallContext;
-import org.apache.polaris.core.context.CallContext;
 import org.apache.polaris.core.entity.PolarisEntity;
 import org.apache.polaris.core.entity.PolarisEntityConstants;
 
 public class StorageCredentialCacheKey {
 
+  private final String realmId;
   private final long catalogId;
 
   /** The serialized string of the storage config. */
@@ -44,18 +42,13 @@ public class StorageCredentialCacheKey {
 
   private final Set<String> allowedWriteLocations;
 
-  /**
-   * The callContext is passed to be used to fetch subscoped creds, but is not used to hash/equals
-   * as part of the cache key.
-   */
-  private @Nullable PolarisCallContext callContext;
-
   public StorageCredentialCacheKey(
+      String realmId,
       PolarisEntity entity,
       boolean allowedListAction,
       Set<String> allowedReadLocations,
-      Set<String> allowedWriteLocations,
-      @Nullable PolarisCallContext callContext) {
+      Set<String> allowedWriteLocations) {
+    this.realmId = realmId;
     this.catalogId = entity.getCatalogId();
     this.storageConfigSerializedStr =
         entity
@@ -65,10 +58,10 @@ public StorageCredentialCacheKey(
     this.allowedListAction = allowedListAction;
     this.allowedReadLocations = allowedReadLocations;
     this.allowedWriteLocations = allowedWriteLocations;
-    this.callContext = callContext;
-    if (this.callContext == null) {
-      this.callContext = CallContext.getCurrentContext().getPolarisCallContext();
-    }
+  }
+
+  public String getRealmId() {
+    return realmId;
   }
 
   public long getCatalogId() {
@@ -95,16 +88,13 @@ public Set<String> getAllowedWriteLocations() {
     return allowedWriteLocations;
   }
 
-  public @Nullable PolarisCallContext getCallContext() {
-    return callContext;
-  }
-
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;
     if (o == null || getClass() != o.getClass()) return false;
     StorageCredentialCacheKey cacheKey = (StorageCredentialCacheKey) o;
-    return catalogId == cacheKey.getCatalogId()
+    return Objects.equals(realmId, cacheKey.getRealmId())
+        && catalogId == cacheKey.getCatalogId()
         && Objects.equals(storageConfigSerializedStr, cacheKey.getStorageConfigSerializedStr())
         && allowedListAction == cacheKey.allowedListAction
         && Objects.equals(allowedReadLocations, cacheKey.allowedReadLocations)
@@ -114,6 +104,7 @@ public boolean equals(Object o) {
   @Override
   public int hashCode() {
     return Objects.hash(
+        realmId,
         catalogId,
         storageConfigSerializedStr,
         allowedListAction,
@@ -124,7 +115,9 @@ public int hashCode() {
   @Override
   public String toString() {
     return "StorageCredentialCacheKey{"
-        + "catalogId="
+        + "realmId="
+        + realmId
+        + ", catalogId="
         + catalogId
         + ", storageConfigSerializedStr='"
         + storageConfigSerializedStr