Share SharedPreferencesInMemoryCache across cache instances Fixes AB#3428107 (#2813)

siddhijain · Copilot · web-flow · commit 228db7867cbd · 2025-11-14T20:05:12.000-06:00
Fixes [AB#3428107](https://identitydivision.visualstudio.com/fac9d424-53d2-45c0-91b5-ef6ba7a6bf26/_workitems/edit/3428107) Proposing to use static instance of in-memory cache layer for performance benefits. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
diff --git a/changelog.txt b/changelog.txt
@@ -1,5 +1,6 @@
 vNext
 ----------
+- [MINOR] Share SharedPreferencesInMemoryCache across instances of BrokerOAuth2TokenCache
 - [MINOR] Use SharedPreferencesInMemoryCache implementation in Broker (#2802)
 - [MINOR] Add OpenTelemetry support for passkey operations (#2795)
 - [MINOR] Add passkey registration support for WebView (#2769)
diff --git a/common/src/test/java/com/microsoft/identity/common/BrokerOAuth2TokenCacheTest.java b/common/src/test/java/com/microsoft/identity/common/BrokerOAuth2TokenCacheTest.java
@@ -42,16 +42,17 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.when;
 
-import android.content.Context;
 
+import androidx.annotation.NonNull;
 import androidx.test.core.app.ApplicationProvider;
 
-import com.microsoft.identity.common.components.AndroidPlatformComponentsFactory;
 import com.microsoft.identity.common.components.MockPlatformComponentsFactory;
 import com.microsoft.identity.common.internal.platform.AndroidPlatformUtil;
 import com.microsoft.identity.common.java.cache.BrokerApplicationMetadata;
@@ -66,10 +67,15 @@
 import com.microsoft.identity.common.java.cache.SharedPreferencesAccountCredentialCache;
 import com.microsoft.identity.common.java.cache.AccountDeletionRecord;
 import com.microsoft.identity.common.java.cache.ICacheRecord;
+import com.microsoft.identity.common.java.cache.SharedPreferencesAccountCredentialCacheWithMemoryCache;
 import com.microsoft.identity.common.java.dto.AccountRecord;
 import com.microsoft.identity.common.java.dto.Credential;
 import com.microsoft.identity.common.java.dto.CredentialType;
 import com.microsoft.identity.common.java.exception.ClientException;
+import com.microsoft.identity.common.java.flighting.CommonFlight;
+import com.microsoft.identity.common.java.flighting.CommonFlightsManager;
+import com.microsoft.identity.common.java.flighting.IFlightsManager;
+import com.microsoft.identity.common.java.flighting.IFlightsProvider;
 import com.microsoft.identity.common.java.interfaces.INameValueStorage;
 import com.microsoft.identity.common.java.interfaces.IPlatformComponents;
 import com.microsoft.identity.common.java.providers.microsoft.MicrosoftAccount;
@@ -79,14 +85,15 @@
 import com.microsoft.identity.common.java.providers.oauth2.OAuth2TokenCache;
 import com.microsoft.identity.common.shadows.ShadowAndroidSdkStorageEncryptionManager;
 
+import org.jetbrains.annotations.NotNull;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mockito;
 import org.powermock.api.mockito.PowerMockito;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.annotation.Config;
-import org.robolectric.shadows.ShadowSharedPreferences;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -236,6 +243,7 @@ public void tearDown() throws Exception {
         }
 
         mApplicationMetadataCache.clear();
+        CommonFlightsManager.INSTANCE.resetFlightsManager();
     }
 
     private void initOtherCaches(final IPlatformComponents components) {
@@ -1241,4 +1249,120 @@ public void testClearAll() throws ClientException {
             assertEquals(false, mBrokerOAuth2TokenCache.isClientIdKnownToCache(clientId));
         }
     }
+
+    @Test
+    public void testSingleCacheInstancePerStoreName_FlightEnabled() {
+        // Enable the flight
+        updateUseInMemoryCacheFlight(true);
+
+        final String storeName = "test_store_name";
+        final IPlatformComponents components1 = mPlatformComponents;
+        final IPlatformComponents components2 = mPlatformComponents;
+
+        // Call getCacheToBeUsed twice with the same storeName
+        final IAccountCredentialCache cache1 = BrokerOAuth2TokenCache.getCacheToBeUsed(components1, storeName);
+        final IAccountCredentialCache cache2 = BrokerOAuth2TokenCache.getCacheToBeUsed(components2, storeName);
+
+        // Verify both references point to the same instance
+        assertNotNull(cache1);
+        assertNotNull(cache2);
+        assertSame("Expected same cache instance for same storeName", cache1, cache2);
+        assertTrue("Cache should be of type SharedPreferencesAccountCredentialCacheWithMemoryCache",
+                cache1 instanceof SharedPreferencesAccountCredentialCacheWithMemoryCache);
+    }
+
+    @Test
+    public void testDifferentCacheInstancesPerStoreName_FlightEnabled() {
+        // Enable the flight
+        updateUseInMemoryCacheFlight(true);
+
+        final String storeName1 = "test_store_name_1";
+        final String storeName2 = "test_store_name_2";
+        final IPlatformComponents components = mPlatformComponents;
+
+        // Call getCacheToBeUsed with different storeNames
+        final IAccountCredentialCache cache1 = BrokerOAuth2TokenCache.getCacheToBeUsed(components, storeName1);
+        final IAccountCredentialCache cache2 = BrokerOAuth2TokenCache.getCacheToBeUsed(components, storeName2);
+
+        // Verify both are valid but different instances
+        assertNotNull(cache1);
+        assertNotNull(cache2);
+        assertNotSame("Expected different cache instances for different storeNames", cache1, cache2);
+        assertTrue("Cache should be of type SharedPreferencesAccountCredentialCacheWithMemoryCache",
+                cache1 instanceof SharedPreferencesAccountCredentialCacheWithMemoryCache);
+        assertTrue("Cache should be of type SharedPreferencesAccountCredentialCacheWithMemoryCache",
+                cache2 instanceof SharedPreferencesAccountCredentialCacheWithMemoryCache);
+    }
+
+    @Test
+    public void testCacheInstanceReusedAcrossMultipleBrokerTokenCaches_FlightEnabled() {
+        // Enable the flight
+        updateUseInMemoryCacheFlight(true);
+
+        final String storeName = getBrokerUidSequesteredFilename(TEST_APP_UID);
+
+        // Create multiple BrokerOAuth2TokenCache instances
+        final BrokerOAuth2TokenCache tokenCache1 = new BrokerOAuth2TokenCache
+                (mPlatformComponents,
+                        TEST_APP_UID,
+                        new NameValueStorageBrokerApplicationMetadataCache(mPlatformComponents));
+        final BrokerOAuth2TokenCache tokenCache2 = new BrokerOAuth2TokenCache(mPlatformComponents, TEST_APP_UID,
+                new NameValueStorageBrokerApplicationMetadataCache(mPlatformComponents));
+
+        // Get the underlying account credential caches
+        final IAccountCredentialCache cache1 = tokenCache1.getCacheToBeUsed(mPlatformComponents, storeName);
+        final IAccountCredentialCache cache2 = tokenCache2.getCacheToBeUsed(mPlatformComponents, storeName);
+
+        // Verify same instance is reused
+        assertNotNull(cache1);
+        assertNotNull(cache2);
+        assertSame(cache1, cache2);
+    }
+
+    @Test
+    public void testFociCacheInstanceReused_FlightEnabled() {
+        // Enable the flight
+        updateUseInMemoryCacheFlight(true);
+
+        final String fociStoreName = BROKER_FOCI_ACCOUNT_CREDENTIAL_SHARED_PREFERENCES;
+
+        // Call getCacheToBeUsed multiple times for FOCI cache
+        final IAccountCredentialCache fociCache1 = BrokerOAuth2TokenCache.getCacheToBeUsed(mPlatformComponents, fociStoreName);
+        final IAccountCredentialCache fociCache2 = BrokerOAuth2TokenCache.getCacheToBeUsed(mPlatformComponents, fociStoreName);
+
+        // Verify same FOCI cache instance is reused
+        assertNotNull(fociCache1);
+        assertNotNull(fociCache2);
+        assertSame(fociCache1, fociCache2);
+    }
+
+    private void updateUseInMemoryCacheFlight(boolean enabled) {
+        final IFlightsProvider mockFlightsProvider = Mockito.mock(IFlightsProvider.class);
+        Mockito.when(mockFlightsProvider.isFlightEnabled(CommonFlight.USE_IN_MEMORY_CACHE_FOR_ACCOUNTS_AND_CREDENTIALS))
+                .thenReturn(enabled);
+
+        // Create anonymous IFlightsManager
+        IFlightsManager anonymousFlightsManager = new IFlightsManager() {
+            @Override
+            public @NotNull IFlightsProvider getFlightsProvider(long waitForConfigsWithTimeoutInMs) {
+                return mockFlightsProvider;
+            }
+            @Override
+            public @NotNull IFlightsProvider getFlightsProviderForTenant(@NotNull String tenantId, long waitForConfigsWithTimeoutInMs) {
+                return mockFlightsProvider;
+            }
+            @Override
+            public @NotNull IFlightsProvider getFlightsProviderForTenant(@NotNull String tenantId) {
+                return mockFlightsProvider;
+            }
+            @NonNull
+            @Override
+            public IFlightsProvider getFlightsProvider() {
+                return mockFlightsProvider;
+            }
+        };
+
+        // Initialize CommonFlightsManager with the anonymous implementation
+        CommonFlightsManager.INSTANCE.initializeCommonFlightsManager(anonymousFlightsManager);
+    }
 }
diff --git a/common4j/src/main/com/microsoft/identity/common/java/cache/BrokerOAuth2TokenCache.java b/common4j/src/main/com/microsoft/identity/common/java/cache/BrokerOAuth2TokenCache.java
@@ -52,7 +52,9 @@
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 import edu.umd.cs.findbugs.annotations.Nullable;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
@@ -104,6 +106,20 @@ public class BrokerOAuth2TokenCache
     private final MicrosoftFamilyOAuth2TokenCache mFociCache;
     private final int mUid;
     private ProcessUidCacheFactory mDelegate = null;
+    /**
+     * Shared, process-wide registry of in-memory augmented account/credential caches keyed by
+     * storage (SharedPreferences) name.
+     * <p>
+     * Populated lazily via computeIfAbsent in getCacheToBeUsed(...). Allows multiple
+     * BrokerOAuth2TokenCache instances to reuse the same in-memory layer for the same underlying
+     * encrypted name-value store, reducing disk I/O and serialization overhead.
+     * <p>
+     * Thread-safety: ConcurrentHashMap ensures safe concurrent access and publication.
+     * Lifecycle: Entries are never removed for the lifetime of the process.
+     */
+    private static final Map<String, SharedPreferencesAccountCredentialCacheWithMemoryCache>
+            inMemoryCacheMapByStorage = new ConcurrentHashMap<>();
+
 
     /**
      * Constructs a new BrokerOAuth2TokenCache.
@@ -1337,6 +1353,7 @@ public void clearAll() {
 
         this.mFociCache.clearAll();
         this.mApplicationMetadataCache.clear();
+        inMemoryCacheMapByStorage.clear();
     }
 
     /**
@@ -1611,14 +1628,10 @@ private MsalOAuth2TokenCache initializeProcessUidCache(@NonNull final IPlatformC
             return mDelegate.getTokenCache(components, uid);
         }
 
-        final INameValueStorage<String> sharedPreferencesFileManager =
-                components.getStorageSupplier().getEncryptedNameValueStore(
-                        SharedPreferencesAccountCredentialCache
-                                .getBrokerUidSequesteredFilename(uid),
-                        String.class
-                );
+        final String storeName = SharedPreferencesAccountCredentialCache
+                .getBrokerUidSequesteredFilename(uid);
 
-        return getTokenCache(components, sharedPreferencesFileManager, false);
+        return getTokenCache(components, storeName, false);
     }
 
     private static MicrosoftFamilyOAuth2TokenCache initializeFociCache(@NonNull final IPlatformComponents components) {
@@ -1628,37 +1641,16 @@ private static MicrosoftFamilyOAuth2TokenCache initializeFociCache(@NonNull fina
                 "Initializing foci cache"
         );
 
-        final INameValueStorage<String> sharedPreferencesFileManager =
-                components.getStorageSupplier().getEncryptedNameValueStore(
-                        SharedPreferencesAccountCredentialCache.BROKER_FOCI_ACCOUNT_CREDENTIAL_SHARED_PREFERENCES,
-                        String.class
-                );
+        final String storeName = SharedPreferencesAccountCredentialCache.BROKER_FOCI_ACCOUNT_CREDENTIAL_SHARED_PREFERENCES;
 
-        return getTokenCache(components, sharedPreferencesFileManager, true);
+        return getTokenCache(components, storeName,true);
     }
 
     @SuppressWarnings(UNCHECKED)
     private static <T extends MsalOAuth2TokenCache> T getTokenCache(@NonNull final IPlatformComponents components,
-                                                                    @NonNull final INameValueStorage<String> spfm,
+                                                                    String storeName,
                                                                     boolean isFoci) {
-        final ICacheKeyValueDelegate cacheKeyValueDelegate = new CacheKeyValueDelegate();
-        final boolean isFlightEnabled = CommonFlightsManager.INSTANCE
-                .getFlightsProvider()
-                .isFlightEnabled(CommonFlight.USE_IN_MEMORY_CACHE_FOR_ACCOUNTS_AND_CREDENTIALS);
-        final IAccountCredentialCache accountCredentialCache;
-        if (isFlightEnabled) {
-            accountCredentialCache = new SharedPreferencesAccountCredentialCacheWithMemoryCache(
-                    cacheKeyValueDelegate,
-                    spfm
-            );
-            SpanExtension.current().setAttribute(AttributeName.in_memory_cache_used_for_accounts_and_credentials.name(), true);
-        } else {
-            accountCredentialCache = new SharedPreferencesAccountCredentialCache(
-                    cacheKeyValueDelegate,
-                    spfm
-            );
-            SpanExtension.current().setAttribute(AttributeName.in_memory_cache_used_for_accounts_and_credentials.name(), false);
-        }
+        final IAccountCredentialCache accountCredentialCache = getCacheToBeUsed(components, storeName);
         final MicrosoftStsAccountCredentialAdapter accountCredentialAdapter =
                 new MicrosoftStsAccountCredentialAdapter();
 
@@ -1678,6 +1670,54 @@ private static <T extends MsalOAuth2TokenCache> T getTokenCache(@NonNull final I
                 );
     }
 
+
+    /**
+     * Determines which cache implementation to use based on flighting.
+     * <p>
+     * When the flight {@code USE_IN_MEMORY_CACHE_FOR_ACCOUNTS_AND_CREDENTIALS} is enabled,
+     * returns a shared, cached instance of {@link SharedPreferencesAccountCredentialCacheWithMemoryCache}
+     * for the given storage, improving performance by reusing the in-memory cache layer across cache instances.
+     * When disabled, returns a new {@link SharedPreferencesAccountCredentialCache} instance.
+     * <p>
+     * Critical behavior: When flight is enabled, the same {@code SharedPreferencesAccountCredentialCacheWithMemoryCache}
+     * instance is returned for the same {@code spfm} reference, meaning the cache is shared across multiple
+     * {@code BrokerOAuth2TokenCache} instances.
+     * <p>
+     * Thread-safety: This method is thread-safe via {@code ConcurrentHashMap.computeIfAbsent}.
+     * <p>
+     * Lifecycle: Returned cached instances are never removed from the static map.
+     *
+     * @return A cached shared in-memory cache instance (flight enabled) or a new
+     *         non-cached instance (flight disabled).
+     */
+    public static IAccountCredentialCache getCacheToBeUsed(@NonNull final IPlatformComponents components,
+                                                           final String storeName) {
+        final boolean isFlightEnabled = CommonFlightsManager.INSTANCE
+                .getFlightsProvider()
+                .isFlightEnabled(CommonFlight.USE_IN_MEMORY_CACHE_FOR_ACCOUNTS_AND_CREDENTIALS);
+        SpanExtension.current().setAttribute(AttributeName.in_memory_cache_used_for_accounts_and_credentials.name(), isFlightEnabled);
+        if (isFlightEnabled) {
+            return inMemoryCacheMapByStorage.computeIfAbsent(storeName, s ->
+                    new SharedPreferencesAccountCredentialCacheWithMemoryCache(
+                            new CacheKeyValueDelegate(),
+                            components.getStorageSupplier().getEncryptedNameValueStore(
+                                    storeName,
+                                    String.class
+                    ))
+            );
+        } else {
+            final INameValueStorage<String> sharedPreferencesFileManager =
+                    components.getStorageSupplier().getEncryptedNameValueStore(
+                            storeName,
+                            String.class
+                    );
+            return new SharedPreferencesAccountCredentialCache(
+                    new CacheKeyValueDelegate(),
+                    sharedPreferencesFileManager
+            );
+        }
+    }
+
     @Nullable
     private MsalOAuth2TokenCache getTokenCacheForClient(@Nullable final BrokerApplicationMetadata metadata) {
         final String methodName = ":getTokenCacheForClient(bam)";
@@ -1736,53 +1776,4 @@ private MsalOAuth2TokenCache getTokenCacheForClient(@NonNull final String client
 
         return getTokenCacheForClient(metadata);
     }
-
-    /**
-     * Sets the SSO state for the supplied Account, relative to the provided uid.
-     *
-     * @param uidStr       The uid of the app whose SSO token is being inserted.
-     * @param account      The account for which the supplied token is being inserted.
-     * @param refreshToken The token to insert.
-     */
-    public void setSingleSignOnState(@NonNull final String uidStr,
-                                     @NonNull final GenericAccount account,
-                                     @NonNull final GenericRefreshToken refreshToken) {
-        final String methodName = ":setSingleSignOnState";
-
-        final boolean isFrt = refreshToken.getIsFamilyRefreshToken();
-
-        MsalOAuth2TokenCache targetCache;
-
-        final int uid = Integer.parseInt(uidStr);
-
-        if (isFrt) {
-            Logger.verbose(TAG + methodName, "Saving tokens to foci cache.");
-
-            targetCache = mFociCache;
-        } else {
-            // If there is an existing cache for this client id, use it. Otherwise, create a new
-            // one based on the supplied uid.
-            targetCache = initializeProcessUidCache(getComponents(), uid);
-        }
-        try {
-            targetCacheSetSingleSignOnState(account, refreshToken, targetCache);
-            updateApplicationMetadataCache(
-                    refreshToken.getClientId(),
-                    refreshToken.getEnvironment(),
-                    refreshToken.getFamilyId(),
-                    uid
-            );
-        } catch (ClientException e) {
-            Logger.warn(
-                    TAG + methodName,
-                    "Failed to save account/refresh token. Skipping."
-            );
-        }
-    }
-
-    // Suppressing unchecked warning as the generic type was not provided for targetCache
-    @SuppressWarnings(WarningType.unchecked_warning)
-    private void targetCacheSetSingleSignOnState(@NonNull GenericAccount account, @NonNull GenericRefreshToken refreshToken, MsalOAuth2TokenCache targetCache) throws ClientException {
-        targetCache.setSingleSignOnState(account, refreshToken);
-    }
 }
