From 574f030b3eaad227e33b65a8e547c4a778d2be9f Mon Sep 17 00:00:00 2001
From: Pinyao Ting <pinyaoting@google.com>
Date: Mon, 31 Jan 2022 19:31:45 +0000
Subject: [PATCH] Refrain from calling AppSearchUserInstance after user is
 stopped

Bug: 217173331
Test: manual
Change-Id: Id4a9972046de73235ec69a26096120100f282963
---
 .../android/server/pm/ShortcutService.java    |  4 ++++
 .../com/android/server/pm/ShortcutUser.java   | 21 +++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/services/core/java/com/android/server/pm/ShortcutService.java b/services/core/java/com/android/server/pm/ShortcutService.java
index 27605789891e3..1cf2dc52e4307 100644
--- a/services/core/java/com/android/server/pm/ShortcutService.java
+++ b/services/core/java/com/android/server/pm/ShortcutService.java
@@ -735,6 +735,9 @@ public class ShortcutService extends IShortcutService.Stub {
         if (DEBUG || DEBUG_REBOOT) {
             Slog.d(TAG, "unloadUserLocked: user=" + userId);
         }
+        // Cancel any ongoing background tasks.
+        getUserShortcutsLocked(userId).cancelAllInFlightTasks();
+
         // Save all dirty information.
         saveDirtyInfo(false);
 
@@ -3736,6 +3739,7 @@ public class ShortcutService extends IShortcutService.Stub {
             synchronized (mLock) {
                 if (mHandler.hasCallbacks(mSaveDirtyInfoRunner)) {
                     mHandler.removeCallbacks(mSaveDirtyInfoRunner);
+                    forEachLoadedUserLocked(ShortcutUser::cancelAllInFlightTasks);
                     saveDirtyInfo(false);
                 }
                 mShutdown.set(true);
diff --git a/services/core/java/com/android/server/pm/ShortcutUser.java b/services/core/java/com/android/server/pm/ShortcutUser.java
index 408f045f47b8c..4bb5dcfa4b269 100644
--- a/services/core/java/com/android/server/pm/ShortcutUser.java
+++ b/services/core/java/com/android/server/pm/ShortcutUser.java
@@ -33,6 +33,7 @@ import android.util.Slog;
 import android.util.TypedXmlPullParser;
 import android.util.TypedXmlSerializer;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.infra.AndroidFuture;
 import com.android.internal.logging.MetricsLogger;
@@ -50,7 +51,9 @@ import org.xmlpull.v1.XmlPullParserException;
 import java.io.File;
 import java.io.IOException;
 import java.io.PrintWriter;
+import java.util.ArrayList;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
@@ -138,6 +141,11 @@ class ShortcutUser {
     private String mLastAppScanOsFingerprint;
     private String mRestoreFromOsFingerprint;
 
+    private final Object mLock = new Object();
+
+    @GuardedBy("mLock")
+    private final ArrayList<AndroidFuture<AppSearchSession>> mInFlightSessions = new ArrayList<>();
+
     public ShortcutUser(ShortcutService service, int userId) {
         mService = service;
         mUserId = userId;
@@ -718,6 +726,10 @@ class ShortcutUser {
     AndroidFuture<AppSearchSession> getAppSearch(
             @NonNull final AppSearchManager.SearchContext searchContext) {
         final AndroidFuture<AppSearchSession> future = new AndroidFuture<>();
+        synchronized (mLock) {
+            mInFlightSessions.removeIf(CompletableFuture::isDone);
+            mInFlightSessions.add(future);
+        }
         if (mAppSearchManager == null) {
             future.completeExceptionally(new RuntimeException("app search manager is null"));
             return future;
@@ -743,4 +755,13 @@ class ShortcutUser {
         }
         return future;
     }
+
+    void cancelAllInFlightTasks() {
+        synchronized (mLock) {
+            for (AndroidFuture<AppSearchSession> session : mInFlightSessions) {
+                session.cancel(true);
+            }
+            mInFlightSessions.clear();
+        }
+    }
 }