Cache Active Broker In Memory (BrokerDiscoveryClient), Fixes AB#3429928 (#2817)

Author: rpdome

changelog.txt

@@ -1,5 +1,6 @@
 vNext
 ----------
+- [MINOR] Cache Active Broker In Memory (BrokerDiscoveryClient) (#2817)
 - [MINOR] Enable Broker Discovery by default in MSAL/Broker API (#2818)
 - [MINOR] Share SharedPreferencesInMemoryCache across instances of BrokerOAuth2TokenCache
 - [MINOR] Use SharedPreferencesInMemoryCache implementation in Broker (#2802)

common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/BrokerDiscoveryClient.kt

@@ -24,6 +24,7 @@ package com.microsoft.identity.common.internal.activebrokerdiscovery
 
 import android.content.Context
 import android.os.Bundle
+import androidx.annotation.VisibleForTesting
 import com.microsoft.identity.common.exception.BrokerCommunicationException
 import com.microsoft.identity.common.internal.broker.BrokerData
 import com.microsoft.identity.common.internal.broker.BrokerValidator
@@ -206,6 +207,21 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
         }
     }
 
+    data class CachedBrokerData (
+        val brokerData: BrokerData?
+    )
+
+    /**
+     * In-memory cache for the active broker data.
+     * There are 3 possible states:
+     * 1. null: the cache hasn't been initialized (needs to read from storage or query from broker first).
+     * 2. CachedBrokerData with null brokerData: no active broker found.
+     * 3. CachedBrokerData with non-null brokerData: found an active broker.
+     **/
+    @Volatile
+    @VisibleForTesting
+    var cachedData: CachedBrokerData? = null
+
     constructor(context: Context,
                 components: IPlatformComponents,
                 cache: IClientActiveBrokerCache) : this(
@@ -273,9 +289,59 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
         }
     }
 
+    override fun getActiveBrokerWithInMemoryCache(telemetryCallback: IBrokerDiscoveryClientTelemetryCallback?): BrokerData?{
+        cachedData?.let {
+            if (it.brokerData == null) {
+                return null
+            }
+            if (validateInMemoryCacheValue(it.brokerData, telemetryCallback)){
+                return it.brokerData
+            }
+
+            // Even if the value is not valid, we don't modify (invalidate) the class variable here since we're not in the lock.
+            // The variable will be updated in the block below.
+        }
+
+        return runBlocking {
+            val timeStartAcquiringLock = System.nanoTime()
+            classLevelLock.withLock {
+                telemetryCallback?.onLockAcquired(System.nanoTime() - timeStartAcquiringLock)
+
+                // just in case the value is already populated while waiting for the lock.
+                cachedData?.let {
+                    if (it.brokerData == null) {
+                        return@runBlocking null
+                    }
+                    if (validateInMemoryCacheValue(it.brokerData, telemetryCallback)){
+                        return@runBlocking it.brokerData
+                    }
+                }
+
+                val brokerData = getActiveBrokerAsync(shouldSkipCache = false, telemetryCallback)
+                cachedData = CachedBrokerData(brokerData)
+                return@runBlocking brokerData
+            }
+        }
+    }
+
+    /**
+     * Make sure the [BrokerData] that stays in the memory cache (if any) is valid.
+     **/
+    private fun validateInMemoryCacheValue(
+        data: BrokerData,
+        telemetryCallback: IBrokerDiscoveryClientTelemetryCallback?,
+    ): Boolean {
+        val timeStartIsValidBroker = System.nanoTime()
+        val isValidBroker = isValidBroker(data)
+        telemetryCallback?.onFinishCheckingIfValidBroker(System.nanoTime() - timeStartIsValidBroker)
+        return isValidBroker
+    }
+
     override fun getActiveBroker(shouldSkipCache: Boolean): BrokerData? {
         return runBlocking {
-            return@runBlocking getActiveBrokerAsync(shouldSkipCache, null)
+            classLevelLock.withLock {
+                return@runBlocking getActiveBrokerAsync(shouldSkipCache, null)
+            }
         }
     }
 
@@ -284,102 +350,88 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
         telemetryCallback: IBrokerDiscoveryClientTelemetryCallback
     ): BrokerData? {
         return runBlocking {
-            return@runBlocking getActiveBrokerAsync(shouldSkipCache, telemetryCallback)
+            val timeStartAcquiringLock = System.nanoTime()
+            classLevelLock.withLock {
+                telemetryCallback.onLockAcquired(System.nanoTime() - timeStartAcquiringLock)
+                return@runBlocking getActiveBrokerAsync(shouldSkipCache, telemetryCallback)
+            }
         }
     }
 
     private suspend fun getActiveBrokerAsync(shouldSkipCache:Boolean,
                                              telemetryCallback: IBrokerDiscoveryClientTelemetryCallback?): BrokerData?{
         val methodTag = "$TAG:getActiveBrokerAsync"
-
-        val timeStartAcquiringLock = System.nanoTime()
-        classLevelLock.withLock {
-            telemetryCallback?.onLockAcquired(System.nanoTime() - timeStartAcquiringLock)
-            if (!shouldSkipCache) {
-                if (cache.shouldUseAccountManager()) {
-                    telemetryCallback?.onUseAccountManager()
-                    return getActiveBrokerFromAccountManager()
+        if (!shouldSkipCache) {
+            if (cache.shouldUseAccountManager()) {
+                telemetryCallback?.onUseAccountManager()
+                return getActiveBrokerFromAccountManager()
+            }
+            val timeStartReadingFromCache = System.nanoTime()
+            cache.getCachedActiveBroker()?.let {
+                telemetryCallback?.onReadFromCache(System.nanoTime() - timeStartReadingFromCache)
+
+                val timeStartIsPackageInstalled = System.nanoTime()
+                val isPackageInstalled = isPackageInstalled(it)
+                telemetryCallback?.onFinishCheckingIfPackageIsInstalled(System.nanoTime() - timeStartIsPackageInstalled)
+                if (!isPackageInstalled) {
+                    Logger.info(
+                        methodTag,
+                        "There is a cached broker: $it, but the app is no longer installed."
+                    )
+                    cache.clearCachedActiveBroker()
+                    return@let
                 }
-                val timeStartReadingFromCache = System.nanoTime()
-                cache.getCachedActiveBroker()?.let {
-                    telemetryCallback?.onReadFromCache(System.nanoTime() - timeStartReadingFromCache)
-
-                    val timeStartIsPackageInstalled = System.nanoTime()
-                    val isPackageInstalled = isPackageInstalled(it)
-                    telemetryCallback?.onFinishCheckingIfPackageIsInstalled(System.nanoTime() - timeStartIsPackageInstalled)
-                    if (!isPackageInstalled) {
-                        Logger.info(
-                            methodTag,
-                            "There is a cached broker: $it, but the app is no longer installed."
-                        )
-                        cache.clearCachedActiveBroker()
-                        return@let
-                    }
 
-                    val timeStartIsValidBroker = System.nanoTime()
-                    val isValidBroker = isValidBroker(it)
-                    telemetryCallback?.onFinishCheckingIfValidBroker(System.nanoTime() - timeStartIsValidBroker)
-                    if (!isValidBroker) {
-                        Logger.info(
-                            methodTag,
-                            "Clearing cache as the installed app does not have a matching signature hash."
-                        )
-                        cache.clearCachedActiveBroker()
-                        return@let
-                    }
-
-                    val timeStartIsSupportedByTargetedBroker = System.nanoTime()
-                    val isSupportedByTargetedBroker =
-                        ipcStrategy.isSupportedByTargetedBroker(it.packageName)
-                    telemetryCallback?.onFinishCheckingIfSupportedByTargetedBroker(System.nanoTime() - timeStartIsSupportedByTargetedBroker)
-                    if(!isSupportedByTargetedBroker){
-                        Logger.info(
-                            methodTag,
-                            "Clearing cache as the installed app does not provide any IPC mechanism to communicate to. (e.g. the broker code isn't shipped with this apk)"
-                        )
-                        cache.clearCachedActiveBroker()
-                        return@let
-                    }
-
-                    Logger.info(methodTag, "Returning cached broker: $it")
-                    return it
+                val timeStartIsValidBroker = System.nanoTime()
+                val isValidBroker = isValidBroker(it)
+                telemetryCallback?.onFinishCheckingIfValidBroker(System.nanoTime() - timeStartIsValidBroker)
+                if (!isValidBroker) {
+                    Logger.info(
+                        methodTag,
+                        "Clearing cache as the installed app does not have a matching signature hash."
+                    )
+                    cache.clearCachedActiveBroker()
+                    return@let
                 }
-            }
 
-            val timeStartQueryFromBroker = System.nanoTime()
-            val brokerData = queryFromBroker(
-                brokerCandidates = brokerCandidates,
-                ipcStrategy = ipcStrategy,
-                isPackageInstalled = isPackageInstalled,
-                isValidBroker = isValidBroker
-            )
-            telemetryCallback?.onFinishQueryingResultFromBroker(System.nanoTime() - timeStartQueryFromBroker)
-
-            if (brokerData != null) {
-                cache.setCachedActiveBroker(brokerData)
-                return brokerData
+                Logger.info(methodTag, "Returning cached broker: $it")
+                return it
             }
+        }
 
-            Logger.info(
-                methodTag,
-                "Will skip broker discovery via IPC and fall back to AccountManager " +
-                        "for the next 60 minutes."
-            )
-            cache.clearCachedActiveBroker()
-            cache.setShouldUseAccountManagerForTheNextMilliseconds(
-                TimeUnit.MINUTES.toMillis(
-                    60
-                )
-            )
+        val timeStartQueryFromBroker = System.nanoTime()
+        val brokerData = queryFromBroker(
+            brokerCandidates = brokerCandidates,
+            ipcStrategy = ipcStrategy,
+            isPackageInstalled = isPackageInstalled,
+            isValidBroker = isValidBroker
+        )
+        telemetryCallback?.onFinishQueryingResultFromBroker(System.nanoTime() - timeStartQueryFromBroker)
+
+        if (brokerData != null) {
+            cache.setCachedActiveBroker(brokerData)
+            return brokerData
+        }
 
-            telemetryCallback?.onUseAccountManager()
-            val accountManagerResult = getActiveBrokerFromAccountManager()
-            Logger.info(
-                methodTag, "Tried getting active broker from account manager, " +
-                        "get ${accountManagerResult?.packageName}."
+        Logger.info(
+            methodTag,
+            "Will skip broker discovery via IPC and fall back to AccountManager " +
+                    "for the next 60 minutes."
+        )
+        cache.clearCachedActiveBroker()
+        cache.setShouldUseAccountManagerForTheNextMilliseconds(
+            TimeUnit.MINUTES.toMillis(
+                60
             )
+        )
 
-            return accountManagerResult
-        }
+        telemetryCallback?.onUseAccountManager()
+        val accountManagerResult = getActiveBrokerFromAccountManager()
+        Logger.info(
+            methodTag, "Tried getting active broker from account manager, " +
+                    "get ${accountManagerResult?.packageName}."
+        )
+
+        return accountManagerResult
     }
 }

common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/IBrokerDiscoveryClient.kt

@@ -65,4 +65,15 @@ interface IBrokerDiscoveryClient {
      * **/
     @kotlin.jvm.Throws(ClientException::class)
     fun forceBrokerRediscovery(brokerCandidate: BrokerData): BrokerData
+
+    /**
+     * This method will return the cached [BrokerData] in the memory (if any).
+     * If the cached value is invalid, it will then proceed with the regular [getActiveBroker] flow.
+     *
+     * The cached value in memory will remain the same during the app process lifetime.
+     *
+     * @param telemetryCallback     callback with telemetry data.
+     * @return BrokerData package name and signature hash of the targeted app.
+     * */
+    fun getActiveBrokerWithInMemoryCache(telemetryCallback: IBrokerDiscoveryClientTelemetryCallback?): BrokerData?
 }