add callback for oneauth

rpdome · rpdome · commit ecd94350f58d · 2025-11-03T20:37:55.000-08:00
diff --git a/changelog.txt b/changelog.txt
@@ -1,5 +1,6 @@
 vNext
 ----------
+- [MINOR] Add callback for OneAuth for measuring Broker Discovery Client Perf (#2796)
 - [MINOR] Add new span name for DELEGATION_CERT_INSTALL's telemetry (#2790)
 - [MINOR] Refactor getAccountByLocalAccountId (#2781)
 - [MINOR] Add OTel Benchmarker (#2786)
diff --git a/common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/BrokerDiscoveryClient.kt b/common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/BrokerDiscoveryClient.kt
@@ -22,6 +22,7 @@
 //  THE SOFTWARE.
 package com.microsoft.identity.common.internal.activebrokerdiscovery
 
+import IBrokerDiscoveryClientTelemetryCallback
 import android.content.Context
 import android.os.Bundle
 import com.microsoft.identity.common.exception.BrokerCommunicationException
@@ -66,15 +67,6 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
                             private val isPackageInstalled: (BrokerData) -> Boolean,
                             private val isValidBroker: (BrokerData) -> Boolean) : IBrokerDiscoveryClient {
 
-    interface IBrokerDiscoveryClientTelemetryCallback {
-        fun onReadFromCache()
-        fun onShouldUseAccountManager()
-        fun onFinishCheckingIfPackageIsInstalled(timeSpent: Long)
-        fun onFinishCheckingIfSupportedByTargetedBroker(timeSpent: Long)
-        fun onFinishQueryingResultFromBroker(timeSpent: Long)
-        fun onFinishCheckingIfValidBroker(timeSpent: Long)
-    }
-
     companion object {
         val TAG = BrokerDiscoveryClient::class.simpleName
 
@@ -108,8 +100,6 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
 
         const val ERROR_BUNDLE_KEY = "ERROR_BUNDLE_KEY"
 
-        public val sTelemetryCallback: IBrokerDiscoveryClientTelemetryCallback? = null
-
         /**
          * Per-process Thread-safe, coroutine-safe Mutex of this class.
          * This is to prevent the IPC mechanism from being unnecessarily triggered due to race condition.
@@ -286,24 +276,35 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
 
     override fun getActiveBroker(shouldSkipCache: Boolean): BrokerData? {
         return runBlocking {
-            return@runBlocking getActiveBrokerAsync(shouldSkipCache)
+            return@runBlocking getActiveBrokerAsync(shouldSkipCache, null)
+        }
+    }
+
+    override fun getActiveBroker(
+        shouldSkipCache: Boolean,
+        telemetryCallback: IBrokerDiscoveryClientTelemetryCallback
+    ): BrokerData? {
+        return runBlocking {
+            return@runBlocking getActiveBrokerAsync(shouldSkipCache, telemetryCallback)
         }
     }
 
-    private suspend fun getActiveBrokerAsync(shouldSkipCache:Boolean): BrokerData?{
+    private suspend fun getActiveBrokerAsync(shouldSkipCache:Boolean,
+                                             telemetryCallback: IBrokerDiscoveryClientTelemetryCallback?): BrokerData?{
         val methodTag = "$TAG:getActiveBrokerAsync"
         classLevelLock.withLock {
             if (!shouldSkipCache) {
                 if (cache.shouldUseAccountManager()) {
-                    sTelemetryCallback?.onShouldUseAccountManager()
+                    telemetryCallback?.onUseAccountManager()
                     return getActiveBrokerFromAccountManager()
                 }
+                val timeStartReadingFromCache = System.nanoTime()
                 cache.getCachedActiveBroker()?.let {
-                    sTelemetryCallback?.onReadFromCache()
+                    telemetryCallback?.onReadFromCache(System.nanoTime() - timeStartReadingFromCache)
 
-                    val startTime1 = System.nanoTime()
+                    val timeStartIsPackageInstalled = System.nanoTime()
                     val isPackageInstalled= isPackageInstalled(it)
-                    sTelemetryCallback?.onFinishCheckingIfPackageIsInstalled(System.nanoTime() - startTime1)
+                    telemetryCallback?.onFinishCheckingIfPackageIsInstalled(System.nanoTime() - timeStartIsPackageInstalled)
                     if (!isPackageInstalled) {
                         Logger.info(
                             methodTag,
@@ -313,9 +314,9 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
                         return@let
                     }
 
-                    val startTime2 = System.nanoTime()
+                    val timeStartIsValidBroker = System.nanoTime()
                     val isValidBroker= isValidBroker(it)
-                    sTelemetryCallback?.onFinishCheckingIfValidBroker(System.nanoTime() - startTime2)
+                    telemetryCallback?.onFinishCheckingIfValidBroker(System.nanoTime() - timeStartIsValidBroker)
                     if (!isValidBroker) {
                         Logger.info(
                             methodTag,
@@ -325,10 +326,10 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
                         return@let
                     }
 
-                    val startTime3 = System.nanoTime()
+                    val timeStartIsSupportedByTargetedBroker = System.nanoTime()
                     val isSupportedByTargetedBroker =
                         ipcStrategy.isSupportedByTargetedBroker(it.packageName)
-                    sTelemetryCallback?.onFinishCheckingIfSupportedByTargetedBroker(System.nanoTime() - startTime3)
+                    telemetryCallback?.onFinishCheckingIfSupportedByTargetedBroker(System.nanoTime() - timeStartIsSupportedByTargetedBroker)
                     if(!isSupportedByTargetedBroker){
                         Logger.info(
                             methodTag,
@@ -343,14 +344,14 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
                 }
             }
 
-            val startTime4 = System.nanoTime()
+            val timeStartQueryFromBroker = System.nanoTime()
             val brokerData = queryFromBroker(
                 brokerCandidates = brokerCandidates,
                 ipcStrategy = ipcStrategy,
                 isPackageInstalled = isPackageInstalled,
                 isValidBroker = isValidBroker
             )
-            sTelemetryCallback?.onFinishCheckingIfPackageIsInstalled(System.nanoTime() - startTime4)
+            telemetryCallback?.onFinishQueryingResultFromBroker(System.nanoTime() - timeStartQueryFromBroker)
 
             if (brokerData != null) {
                 cache.setCachedActiveBroker(brokerData)
@@ -369,6 +370,7 @@ class BrokerDiscoveryClient(private val brokerCandidates: Set<BrokerData>,
                 )
             )
 
+            telemetryCallback?.onUseAccountManager()
             val accountManagerResult = getActiveBrokerFromAccountManager()
             Logger.info(
                 methodTag, "Tried getting active broker from account manager, " +
diff --git a/common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/IBrokerDiscoveryClient.kt b/common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/IBrokerDiscoveryClient.kt
@@ -22,6 +22,7 @@
 //  THE SOFTWARE.
 package com.microsoft.identity.common.internal.activebrokerdiscovery
 
+import IBrokerDiscoveryClientTelemetryCallback
 import com.microsoft.identity.common.internal.broker.BrokerData
 import com.microsoft.identity.common.java.exception.ClientException
 
@@ -37,6 +38,20 @@ interface IBrokerDiscoveryClient {
      * **/
     fun getActiveBroker(shouldSkipCache: Boolean = false) : BrokerData?
 
+
+    /**
+     * Performs a discovery to figure out which broker app the SDK (MSAL/OneAuth)
+     * has to send its request to.
+     *
+     * @param shouldSkipCache       If true, this will skip cached value (if any)
+     *                              and always query the broker for the result.
+     * @param telemetryCallback     callback with telemetry data.
+     * @return BrokerData package name and signature hash of the targeted app.
+     * **/
+    fun getActiveBroker(shouldSkipCache: Boolean = false,
+                        telemetryCallback: IBrokerDiscoveryClientTelemetryCallback) : BrokerData?
+
+
     /**
      * Force a broker app to figure out which broker app the SDK (MSAL/OneAuth)
      * has to send its request to.
diff --git a/common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/IBrokerDiscoveryClientTelemetryCallback.kt b/common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/IBrokerDiscoveryClientTelemetryCallback.kt
@@ -0,0 +1,72 @@
+//  Copyright (c) Microsoft Corporation.
+//  All rights reserved.
+//
+//  This code is licensed under the MIT License.
+//
+//  Permission is hereby granted, free of charge, to any person obtaining a copy
+//  of this software and associated documentation files(the "Software"), to deal
+//  in the Software without restriction, including without limitation the rights
+//  to use, copy, modify, merge, publish, distribute, sublicense, and / or sell
+//  copies of the Software, and to permit persons to whom the Software is
+//  furnished to do so, subject to the following conditions :
+//
+//  The above copyright notice and this permission notice shall be included in
+//  all copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+//  THE SOFTWARE.
+
+interface IBrokerDiscoveryClientTelemetryCallback {
+    /**
+     * If triggered, the broker discovery client is using Account Manager
+     * (e.g. one of the installed broker is too old and doesn't support this new mechanism).
+     **/
+    fun onUseAccountManager()
+
+    /**
+     * Triggered when BrokerDiscoveryClient finishes reading from its cache.
+     *
+     * @return timeSpentInNanoSeconds Time spent reading the cache (in nanoseconds).
+     **/
+    fun onReadFromCache(timeSpentInNanoSeconds: Long)
+
+    /**
+     * Triggered when BrokerDiscoveryClient finishes validating if the cached broker package is installed.
+     *
+     * @return timeSpentInNanoSeconds Time spent to validate if the package is installed (in nanoseconds).
+     **/
+    fun onFinishCheckingIfPackageIsInstalled(timeSpentInNanoSeconds: Long)
+
+    /**
+     * Triggered when BrokerDiscoveryClient finishes validating if the cached broker package supports the given IPC strategy.
+     *
+     * @return timeSpentInNanoSeconds Time spent to validate if the package supports the IPC strategy (in nanoseconds).
+     **/
+    fun onFinishCheckingIfSupportedByTargetedBroker(timeSpentInNanoSeconds: Long)
+
+    /**
+     * Triggered when BrokerDiscoveryClient finishes validating if the cached broker package is a valid broker.
+     *
+     * @return timeSpentInNanoSeconds Time spent to validate if the package is a valid broker (in nanoseconds).
+     **/
+    fun onFinishCheckingIfValidBroker(timeSpentInNanoSeconds: Long)
+
+    /**
+     * Triggered when BrokerDiscoveryClient finishes querying active broker from one of the broker apps.
+     *
+     * There are 2 possible cases when this is triggered.
+     * 1. Both Broker and this MSAL/OneAuth app are freshly installed.
+     * This will trigger the broker discovery in broker, and will take longer time.
+     *
+     * 2. Broker has already done broker discovery before (e.g. by other MSAL/OneAuth app). Only this MSAL/OneAuth app is freshly installed.
+     * When the request reaches the broker, the broker would just return the cached value, which will be faster.
+     *
+     * @return timeSpentInNanoSeconds Time spent to query the active broker (in nanoseconds).
+     **/
+    fun onFinishQueryingResultFromBroker(timeSpentInNanoSeconds: Long)
+}
diff --git a/common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/LegacyBrokerDiscoveryClient.kt b/common/src/main/java/com/microsoft/identity/common/internal/activebrokerdiscovery/LegacyBrokerDiscoveryClient.kt
@@ -22,6 +22,7 @@
 //  THE SOFTWARE.
 package com.microsoft.identity.common.internal.activebrokerdiscovery
 
+import IBrokerDiscoveryClientTelemetryCallback
 import android.content.Context
 import com.microsoft.identity.common.internal.broker.BrokerData
 
@@ -35,6 +36,15 @@ class LegacyBrokerDiscoveryClient(val context: Context): IBrokerDiscoveryClient
             .getActiveBrokerFromAccountManager()
     }
 
+    override fun getActiveBroker(
+        shouldSkipCache: Boolean,
+        telemetryCallback: IBrokerDiscoveryClientTelemetryCallback
+    ): BrokerData? {
+        // LegacyBrokerDiscoveryClient does not support telemetry callback.
+        return AccountManagerBrokerDiscoveryUtil(context)
+            .getActiveBrokerFromAccountManager()
+    }
+
     override fun forceBrokerRediscovery(brokerCandidate: BrokerData): BrokerData {
         throw UnsupportedOperationException("LegacyBrokerDiscoveryClient does not support forceBrokerRediscovery.")
     }
