Use nested scope for ShortcutLoader

Use a nested coroutine scope for each ShortcutLoader instance and close
it when the instance is stopped being used.

Fix: 358135601
Test: atest IntentResolver-tests-unit
Test: manual testing checking the logs
Flag: com.android.intentresolver.fix_shortcut_loader_job_leak

Change-Id: I84bbca75612f153193b03e24bfc0cb8842a8d3e2

4 files changed, 82 insertions, 17 deletions

diff --git a/aconfig/FeatureFlags.aconfig b/aconfig/FeatureFlags.aconfig
index 87a584c9..07ef198f 100644
--- a/aconfig/FeatureFlags.aconfig
+++ b/aconfig/FeatureFlags.aconfig
@@ -77,6 +77,16 @@ flag {
 }
 
 flag {
+  name: "fix_shortcut_loader_job_leak"
+  namespace: "intentresolver"
+  description: "User a nested coroutine scope for shortcut loader instances"
+  bug: "358135601"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
+
+flag {
   name: "preview_image_loader"
   namespace: "intentresolver"
   description: "Use the unified preview image loader for all preview variations; support variable preview sizes."
diff --git a/java/src/com/android/intentresolver/ChooserActivity.java b/java/src/com/android/intentresolver/ChooserActivity.java
index 8c6c7b7f..d3c21f3e 100644
--- a/java/src/com/android/intentresolver/ChooserActivity.java
+++ b/java/src/com/android/intentresolver/ChooserActivity.java
@@ -2701,6 +2701,9 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
             if (appPredictor != null) {
                 appPredictor.destroy();
             }
+            if (shortcutLoader != null) {
+                shortcutLoader.destroy();
+            }
         }
     }
 }
diff --git a/java/src/com/android/intentresolver/shortcuts/ShortcutLoader.kt b/java/src/com/android/intentresolver/shortcuts/ShortcutLoader.kt
index 08230d90..68412256 100644
--- a/java/src/com/android/intentresolver/shortcuts/ShortcutLoader.kt
+++ b/java/src/com/android/intentresolver/shortcuts/ShortcutLoader.kt
@@ -35,6 +35,7 @@ import androidx.annotation.MainThread
 import androidx.annotation.OpenForTesting
 import androidx.annotation.VisibleForTesting
 import androidx.annotation.WorkerThread
+import com.android.intentresolver.Flags.fixShortcutLoaderJobLeak
 import com.android.intentresolver.chooser.DisplayResolveInfo
 import com.android.intentresolver.measurements.Tracer
 import com.android.intentresolver.measurements.runTracing
@@ -43,7 +44,9 @@ import java.util.function.Consumer
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.asExecutor
+import kotlinx.coroutines.cancel
 import kotlinx.coroutines.channels.BufferOverflow
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.combine
@@ -65,7 +68,7 @@ open class ShortcutLoader
 @VisibleForTesting
 constructor(
     private val context: Context,
-    private val scope: CoroutineScope,
+    parentScope: CoroutineScope,
     private val appPredictor: AppPredictorProxy?,
     private val userHandle: UserHandle,
     private val isPersonalProfile: Boolean,
@@ -73,6 +76,8 @@ constructor(
     private val dispatcher: CoroutineDispatcher,
     private val callback: Consumer<Result>
 ) {
+    private val scope =
+        if (fixShortcutLoaderJobLeak()) parentScope.createChildScope() else parentScope
     private val shortcutToChooserTargetConverter = ShortcutToChooserTargetConverter()
     private val userManager = context.getSystemService(Context.USER_SERVICE) as UserManager
     private val appPredictorCallback =
@@ -88,6 +93,9 @@ constructor(
     private val isDestroyed
         get() = !scope.isActive
 
+    private val id
+        get() = System.identityHashCode(this).toString(Character.MAX_RADIX)
+
     @MainThread
     constructor(
         context: Context,
@@ -132,7 +140,7 @@ constructor(
             }
             .invokeOnCompletion {
                 runCatching { appPredictor?.unregisterPredictionUpdates(appPredictorCallback) }
-                Log.d(TAG, "destroyed, user: $userHandle")
+                Log.d(TAG, "[$id] destroyed, user: $userHandle")
             }
         reset()
     }
@@ -140,7 +148,7 @@ constructor(
     /** Clear application targets (see [updateAppTargets] and initiate shortcuts loading. */
     @OpenForTesting
     open fun reset() {
-        Log.d(TAG, "reset shortcut loader for user $userHandle")
+        Log.d(TAG, "[$id] reset shortcut loader for user $userHandle")
         appTargetSource.tryEmit(null)
         shortcutSource.tryEmit(null)
         scope.launch(dispatcher) { loadShortcuts() }
@@ -155,14 +163,21 @@ constructor(
         appTargetSource.tryEmit(appTargets)
     }
 
+    @OpenForTesting
+    open fun destroy() {
+        if (fixShortcutLoaderJobLeak()) {
+            scope.cancel()
+        }
+    }
+
     @WorkerThread
     private fun loadShortcuts() {
         // no need to query direct share for work profile when its locked or disabled
         if (!shouldQueryDirectShareTargets()) {
-            Log.d(TAG, "skip shortcuts loading for user $userHandle")
+            Log.d(TAG, "[$id] skip shortcuts loading for user $userHandle")
             return
         }
-        Log.d(TAG, "querying direct share targets for user $userHandle")
+        Log.d(TAG, "[$id] querying direct share targets for user $userHandle")
         queryDirectShareTargets(false)
     }
 
@@ -170,7 +185,7 @@ constructor(
     private fun queryDirectShareTargets(skipAppPredictionService: Boolean) {
         if (!skipAppPredictionService && appPredictor != null) {
             try {
-                Log.d(TAG, "query AppPredictor for user $userHandle")
+                Log.d(TAG, "[$id] query AppPredictor for user $userHandle")
                 Tracer.beginAppPredictorQueryTrace(userHandle)
                 appPredictor.requestPredictionUpdate()
                 return
@@ -180,12 +195,12 @@ constructor(
                 if (isDestroyed) {
                     return
                 }
-                Log.e(TAG, "Failed to query AppPredictor for user $userHandle", e)
+                Log.e(TAG, "[$id] failed to query AppPredictor for user $userHandle", e)
             }
         }
         // Default to just querying ShortcutManager if AppPredictor not present.
         if (targetIntentFilter == null) {
-            Log.d(TAG, "skip querying ShortcutManager for $userHandle")
+            Log.d(TAG, "[$id] skip querying ShortcutManager for $userHandle")
             sendShareShortcutInfoList(
                 emptyList(),
                 isFromAppPredictor = false,
@@ -193,12 +208,12 @@ constructor(
             )
             return
         }
-        Log.d(TAG, "query ShortcutManager for user $userHandle")
+        Log.d(TAG, "[$id] query ShortcutManager for user $userHandle")
         val shortcuts =
             runTracing("shortcut-mngr-${userHandle.identifier}") {
                 queryShortcutManager(targetIntentFilter)
             }
-        Log.d(TAG, "receive shortcuts from ShortcutManager for user $userHandle")
+        Log.d(TAG, "[$id] receive shortcuts from ShortcutManager for user $userHandle")
         sendShareShortcutInfoList(shortcuts, false, null)
     }
 
@@ -210,14 +225,13 @@ constructor(
         val pm = context.createContextAsUser(userHandle, 0 /* flags */).packageManager
         return sm?.getShareTargets(targetIntentFilter)?.filter {
             pm.isPackageEnabled(it.targetComponent.packageName)
-        }
-            ?: emptyList()
+        } ?: emptyList()
     }
 
     @WorkerThread
     private fun onAppPredictorCallback(appPredictorTargets: List<AppTarget>) {
         endAppPredictorQueryTrace(userHandle)
-        Log.d(TAG, "receive app targets from AppPredictor")
+        Log.d(TAG, "[$id] receive app targets from AppPredictor")
         if (appPredictorTargets.isEmpty() && shouldQueryDirectShareTargets()) {
             // APS may be disabled, so try querying targets ourselves.
             queryDirectShareTargets(true)
@@ -330,6 +344,11 @@ constructor(
         val directShareShortcutInfoCache: Map<ChooserTarget, ShortcutInfo>
     )
 
+    private fun endAppPredictorQueryTrace(userHandle: UserHandle) {
+        val duration = Tracer.endAppPredictorQueryTrace(userHandle)
+        Log.d(TAG, "[$id] AppPredictor query duration for user $userHandle: $duration ms")
+    }
+
     /** Shortcuts grouped by app. */
     class ShortcutResultInfo(
         val appTarget: DisplayResolveInfo,
@@ -378,9 +397,12 @@ constructor(
                 .getOrDefault(false)
         }
 
-        private fun endAppPredictorQueryTrace(userHandle: UserHandle) {
-            val duration = Tracer.endAppPredictorQueryTrace(userHandle)
-            Log.d(TAG, "AppPredictor query duration for user $userHandle: $duration ms")
-        }
+        /**
+         * Creates a new coroutine scope and makes its job a child of the given, `this`, coroutine
+         * scope's job. This ensures that the new scope will be canceled when the parent scope is
+         * canceled (but not vice versa).
+         */
+        private fun CoroutineScope.createChildScope() =
+            CoroutineScope(coroutineContext + Job(parent = coroutineContext[Job]))
     }
 }
diff --git a/tests/unit/src/com/android/intentresolver/shortcuts/ShortcutLoaderTest.kt b/tests/unit/src/com/android/intentresolver/shortcuts/ShortcutLoaderTest.kt
index fbdc062b..9c84cdcf 100644
--- a/tests/unit/src/com/android/intentresolver/shortcuts/ShortcutLoaderTest.kt
+++ b/tests/unit/src/com/android/intentresolver/shortcuts/ShortcutLoaderTest.kt
@@ -26,7 +26,10 @@ import android.content.pm.PackageManager.ApplicationInfoFlags
 import android.content.pm.ShortcutManager
 import android.os.UserHandle
 import android.os.UserManager
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.filters.SmallTest
+import com.android.intentresolver.Flags.FLAG_FIX_SHORTCUT_LOADER_JOB_LEAK
 import com.android.intentresolver.chooser.DisplayResolveInfo
 import com.android.intentresolver.createAppTarget
 import com.android.intentresolver.createShareShortcutInfo
@@ -42,6 +45,7 @@ import org.junit.Assert.assertArrayEquals
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
 import org.junit.Assert.assertTrue
+import org.junit.Rule
 import org.junit.Test
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
@@ -56,6 +60,8 @@ import org.mockito.kotlin.whenever
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 class ShortcutLoaderTest {
+    @get:Rule val flagRule = SetFlagsRule()
+
     private val appInfo =
         ApplicationInfo().apply {
             enabled = true
@@ -465,6 +471,30 @@ class ShortcutLoaderTest {
         testAlwaysCallSystemForMainProfile(isQuietModeEnabled = true)
     }
 
+    @Test
+    @EnableFlags(FLAG_FIX_SHORTCUT_LOADER_JOB_LEAK)
+    fun test_ShortcutLoaderDestroyed_appPredictorCallbackUnregisteredAndWatchdogCancelled() {
+        scope.runTest {
+            val testSubject =
+                ShortcutLoader(
+                    context,
+                    backgroundScope,
+                    appPredictor,
+                    UserHandle.of(0),
+                    true,
+                    intentFilter,
+                    dispatcher,
+                    callback
+                )
+
+            testSubject.updateAppTargets(appTargets)
+            testSubject.destroy()
+
+            verify(appPredictor, times(1)).registerPredictionUpdates(any(), any())
+            verify(appPredictor, times(1)).unregisterPredictionUpdates(any())
+        }
+    }
+
     private fun testDisabledWorkProfileDoNotCallSystem(
         isUserRunning: Boolean = true,
         isUserUnlocked: Boolean = true,