Prevent sleep tokens from being created recursively.

Sometimes a keyguard sleep token can be created with recursively call
because we use the local object existing as a condition to create it.
The code flow looks like:

KeyguardController#updateSleepToken
if (mSleepToken == null) then a = acquireSleepToken();
acquireSleepToken
  updateSleepIfNeededLocked
    ensureActivitiesVisible
      keyguardController#visibilitiesUpdated
        updateSleepToken
Then since the local mSleepToken haven't been assigned,
acquireSleepToken would be called again. Which would create another
sleep token and leave the previous one in RootWindowContainer.

To prevent this from happening again, separate updateSleepIfNeeded
from acquireSleepToken, and store sleep token as a map, so no more
sleep token can be created with same key.

Bug: 160351101
Test: atest MultiDisplayKeyguardTests MultiDisplayLockedKeyguardTests/
KeyguardTests KeyguardInputTests KeyguardLockedTests

Change-Id: I168dfe0fa101cb65eb676ca256c03439c8e4193d

6 files changed, 110 insertions, 73 deletions

diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index e75dab73478e..03868e922bdd 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -209,7 +209,6 @@ import com.android.server.policy.keyguard.KeyguardStateMonitor.StateCallback;
 import com.android.server.statusbar.StatusBarManagerInternal;
 import com.android.server.vr.VrManagerInternal;
 import com.android.server.wm.ActivityTaskManagerInternal;
-import com.android.server.wm.ActivityTaskManagerInternal.SleepToken;
 import com.android.server.wm.AppTransition;
 import com.android.server.wm.DisplayPolicy;
 import com.android.server.wm.DisplayRotation;
@@ -491,7 +490,7 @@ public class PhoneWindowManager implements WindowManagerPolicy {
     private boolean mPendingKeyguardOccluded;
     private boolean mKeyguardOccludedChanged;
 
-    SleepToken mScreenOffSleepToken;
+    private ActivityTaskManagerInternal.SleepTokenAcquirer mScreenOffSleepTokenAcquirer;
     volatile boolean mKeyguardOccluded;
     Intent mHomeIntent;
     Intent mCarDockIntent;
@@ -1741,6 +1740,9 @@ public class PhoneWindowManager implements WindowManagerPolicy {
                 new AccessibilityShortcutController(mContext, new Handler(), mCurrentUserId);
         mLogger = new MetricsLogger();
 
+        mScreenOffSleepTokenAcquirer = mActivityTaskManagerInternal
+                .createSleepTokenAcquirer("ScreenOff");
+
         Resources res = mContext.getResources();
         mWakeOnDpadKeyPress =
                 res.getBoolean(com.android.internal.R.bool.config_wakeOnDpadKeyPress);
@@ -4984,15 +4986,9 @@ public class PhoneWindowManager implements WindowManagerPolicy {
     // TODO (multidisplay): Support multiple displays in WindowManagerPolicy.
     private void updateScreenOffSleepToken(boolean acquire) {
         if (acquire) {
-            if (mScreenOffSleepToken == null) {
-                mScreenOffSleepToken = mActivityTaskManagerInternal.acquireSleepToken(
-                        "ScreenOff", DEFAULT_DISPLAY);
-            }
+            mScreenOffSleepTokenAcquirer.acquire(DEFAULT_DISPLAY);
         } else {
-            if (mScreenOffSleepToken != null) {
-                mScreenOffSleepToken.release();
-                mScreenOffSleepToken = null;
-            }
+            mScreenOffSleepTokenAcquirer.release(DEFAULT_DISPLAY);
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java b/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
index d5df9068e81d..a903bcd3d728 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
@@ -125,20 +125,30 @@ public abstract class ActivityTaskManagerInternal {
      * Sleep tokens cause the activity manager to put the top activity to sleep.
      * They are used by components such as dreams that may hide and block interaction
      * with underlying activities.
+     * The Acquirer provides an interface that encapsulates the underlying work, so the user does
+     * not need to handle the token by him/herself.
      */
-    public static abstract class SleepToken {
+    public interface SleepTokenAcquirer {
 
-        /** Releases the sleep token. */
-        public abstract void release();
+        /**
+         * Acquires a sleep token.
+         * @param displayId The display to apply to.
+         */
+        void acquire(int displayId);
+
+        /**
+         * Releases the sleep token.
+         * @param displayId The display to apply to.
+         */
+        void release(int displayId);
     }
 
     /**
-     * Acquires a sleep token for the specified display with the specified tag.
+     * Creates a sleep token acquirer for the specified display with the specified tag.
      *
-     * @param tag A string identifying the purpose of the token (eg. "Dream").
-     * @param displayId The display to apply the sleep token to.
+     * @param tag A string identifying the purpose (eg. "Dream").
      */
-    public abstract SleepToken acquireSleepToken(@NonNull String tag, int displayId);
+    public abstract SleepTokenAcquirer createSleepTokenAcquirer(@NonNull String tag);
 
     /**
      * Returns home activity for the specified user.
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 4e1d789bebd8..df8ac572384a 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -5070,7 +5070,10 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
         final long sleepToken = proto.start(ActivityManagerServiceDumpProcessesProto.SLEEP_STATUS);
         proto.write(ActivityManagerServiceDumpProcessesProto.SleepStatus.WAKEFULNESS,
                 PowerManagerInternal.wakefulnessToProtoEnum(wakeFullness));
-        for (ActivityTaskManagerInternal.SleepToken st : mRootWindowContainer.mSleepTokens) {
+        final int tokenSize = mRootWindowContainer.mSleepTokens.size();
+        for (int i = 0; i < tokenSize; i++) {
+            final RootWindowContainer.SleepToken st =
+                    mRootWindowContainer.mSleepTokens.valueAt(i);
             proto.write(ActivityManagerServiceDumpProcessesProto.SleepStatus.SLEEP_TOKENS,
                     st.toString());
         }
@@ -5504,12 +5507,35 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
                 reason);
     }
 
-    ActivityTaskManagerInternal.SleepToken acquireSleepToken(String tag, int displayId) {
-        synchronized (mGlobalLock) {
-            final ActivityTaskManagerInternal.SleepToken token =
-                    mRootWindowContainer.createSleepToken(tag, displayId);
-            updateSleepIfNeededLocked();
-            return token;
+    final class SleepTokenAcquirerImpl implements ActivityTaskManagerInternal.SleepTokenAcquirer {
+        private final String mTag;
+        private final SparseArray<RootWindowContainer.SleepToken> mSleepTokens =
+                new SparseArray<>();
+
+        SleepTokenAcquirerImpl(@NonNull String tag) {
+            mTag = tag;
+        }
+
+        @Override
+        public void acquire(int displayId) {
+            synchronized (mGlobalLock) {
+                if (!mSleepTokens.contains(displayId)) {
+                    mSleepTokens.append(displayId,
+                            mRootWindowContainer.createSleepToken(mTag, displayId));
+                    updateSleepIfNeededLocked();
+                }
+            }
+        }
+
+        @Override
+        public void release(int displayId) {
+            synchronized (mGlobalLock) {
+                final RootWindowContainer.SleepToken token = mSleepTokens.get(displayId);
+                if (token != null) {
+                    mRootWindowContainer.removeSleepToken(token);
+                    mSleepTokens.remove(displayId);
+                }
+            }
         }
     }
 
@@ -6068,9 +6094,9 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
 
     final class LocalService extends ActivityTaskManagerInternal {
         @Override
-        public SleepToken acquireSleepToken(String tag, int displayId) {
+        public SleepTokenAcquirer createSleepTokenAcquirer(@NonNull String tag) {
             Objects.requireNonNull(tag);
-            return ActivityTaskManagerService.this.acquireSleepToken(tag, displayId);
+            return new SleepTokenAcquirerImpl(tag);
         }
 
         @Override
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 8fff81a1e8c4..609672397024 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -585,9 +585,9 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     private IntArray mDisplayAccessUIDs = new IntArray();
 
     /** All tokens used to put activities on this stack to sleep (including mOffToken) */
-    final ArrayList<ActivityTaskManagerInternal.SleepToken> mAllSleepTokens = new ArrayList<>();
-    /** The token acquired by ActivityStackSupervisor to put stacks on the display to sleep */
-    ActivityTaskManagerInternal.SleepToken mOffToken;
+    final ArrayList<RootWindowContainer.SleepToken> mAllSleepTokens = new ArrayList<>();
+    /** The token acquirer to put stacks on the display to sleep */
+    private final ActivityTaskManagerInternal.SleepTokenAcquirer mOffTokenAcquirer;
 
     private boolean mSleeping;
 
@@ -923,6 +923,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         mAtmService = mWmService.mAtmService;
         mDisplay = display;
         mDisplayId = display.getDisplayId();
+        mOffTokenAcquirer = mRootWindowContainer.mDisplayOffTokenAcquirer;
         mWallpaperController = new WallpaperController(mWmService, this);
         display.getDisplayInfo(mDisplayInfo);
         display.getMetrics(mDisplayMetrics);
@@ -4931,11 +4932,10 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         final int displayId = mDisplay.getDisplayId();
         if (displayId != DEFAULT_DISPLAY) {
             final int displayState = mDisplay.getState();
-            if (displayState == Display.STATE_OFF && mOffToken == null) {
-                mOffToken = mAtmService.acquireSleepToken("Display-off", displayId);
-            } else if (displayState == Display.STATE_ON && mOffToken != null) {
-                mOffToken.release();
-                mOffToken = null;
+            if (displayState == Display.STATE_OFF) {
+                mOffTokenAcquirer.acquire(mDisplayId);
+            } else if (displayState == Display.STATE_ON) {
+                mOffTokenAcquirer.release(mDisplayId);
             }
         }
         mWmService.requestTraversal();
@@ -5175,7 +5175,8 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         mDisplayPolicy.release();
 
         if (!mAllSleepTokens.isEmpty()) {
-            mRootWindowContainer.mSleepTokens.removeAll(mAllSleepTokens);
+            mAllSleepTokens.forEach(token ->
+                    mRootWindowContainer.mSleepTokens.remove(token.mHashKey));
             mAllSleepTokens.clear();
             mAtmService.updateSleepIfNeededLocked();
         }
diff --git a/services/core/java/com/android/server/wm/KeyguardController.java b/services/core/java/com/android/server/wm/KeyguardController.java
index 76f236534b69..4f4f8dc2c5d2 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -50,7 +50,6 @@ import android.util.proto.ProtoOutputStream;
 
 import com.android.internal.policy.IKeyguardDismissCallback;
 import com.android.server.policy.WindowManagerPolicy;
-import com.android.server.wm.ActivityTaskManagerInternal.SleepToken;
 
 import java.io.PrintWriter;
 
@@ -74,11 +73,14 @@ class KeyguardController {
     private final SparseArray<KeyguardDisplayState> mDisplayStates = new SparseArray<>();
     private final ActivityTaskManagerService mService;
     private RootWindowContainer mRootWindowContainer;
+    private final ActivityTaskManagerInternal.SleepTokenAcquirer mSleepTokenAcquirer;
+
 
     KeyguardController(ActivityTaskManagerService service,
             ActivityStackSupervisor stackSupervisor) {
         mService = service;
         mStackSupervisor = stackSupervisor;
+        mSleepTokenAcquirer = mService.new SleepTokenAcquirerImpl("keyguard");
     }
 
     void setWindowManager(WindowManagerService windowManager) {
@@ -412,17 +414,17 @@ class KeyguardController {
 
     private void updateKeyguardSleepToken(int displayId) {
         final KeyguardDisplayState state = getDisplay(displayId);
-        if (isKeyguardUnoccludedOrAodShowing(displayId) && state.mSleepToken == null) {
-            state.acquiredSleepToken();
-        } else if (!isKeyguardUnoccludedOrAodShowing(displayId) && state.mSleepToken != null) {
-            state.releaseSleepToken();
+        if (isKeyguardUnoccludedOrAodShowing(displayId)) {
+            state.mSleepTokenAcquirer.acquire(displayId);
+        } else if (!isKeyguardUnoccludedOrAodShowing(displayId)) {
+            state.mSleepTokenAcquirer.release(displayId);
         }
     }
 
     private KeyguardDisplayState getDisplay(int displayId) {
         KeyguardDisplayState state = mDisplayStates.get(displayId);
         if (state == null) {
-            state = new KeyguardDisplayState(mService, displayId);
+            state = new KeyguardDisplayState(mService, displayId, mSleepTokenAcquirer);
             mDisplayStates.append(displayId, state);
         }
         return state;
@@ -443,29 +445,18 @@ class KeyguardController {
         private ActivityRecord mDismissingKeyguardActivity;
         private boolean mRequestDismissKeyguard;
         private final ActivityTaskManagerService mService;
-        private SleepToken mSleepToken;
+        private final ActivityTaskManagerInternal.SleepTokenAcquirer mSleepTokenAcquirer;
 
-        KeyguardDisplayState(ActivityTaskManagerService service, int displayId) {
+        KeyguardDisplayState(ActivityTaskManagerService service, int displayId,
+                ActivityTaskManagerInternal.SleepTokenAcquirer acquirer) {
             mService = service;
             mDisplayId = displayId;
+            mSleepTokenAcquirer = acquirer;
         }
 
         void onRemoved() {
             mDismissingKeyguardActivity = null;
-            releaseSleepToken();
-        }
-
-        void acquiredSleepToken() {
-            if (mSleepToken == null) {
-                mSleepToken = mService.acquireSleepToken("keyguard", mDisplayId);
-            }
-        }
-
-        void releaseSleepToken() {
-            if (mSleepToken != null) {
-                mSleepToken.release();
-                mSleepToken = null;
-            }
+            mSleepTokenAcquirer.release(mDisplayId);
         }
 
         void visibilitiesUpdated(KeyguardController controller, DisplayContent display) {
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 50ae4ea2f627..f075fa519666 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -220,6 +220,9 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     // transaction from the global transaction.
     private final SurfaceControl.Transaction mDisplayTransaction;
 
+    /** The token acquirer to put stacks on the displays to sleep */
+    final ActivityTaskManagerInternal.SleepTokenAcquirer mDisplayOffTokenAcquirer;
+
     /**
      * The modes which affect which tasks are returned when calling
      * {@link RootWindowContainer#anyTaskForId(int)}.
@@ -259,7 +262,7 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
      * They are used by components that may hide and block interaction with underlying
      * activities.
      */
-    final ArrayList<ActivityTaskManagerInternal.SleepToken> mSleepTokens = new ArrayList<>();
+    final SparseArray<SleepToken> mSleepTokens = new SparseArray<>();
 
     /** Set when a power mode launch has started, but not ended. */
     private boolean mPowerModeLaunchStarted;
@@ -444,6 +447,7 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
         mService = service.mAtmService;
         mStackSupervisor = mService.mStackSupervisor;
         mStackSupervisor.mRootWindowContainer = this;
+        mDisplayOffTokenAcquirer = mService.new SleepTokenAcquirerImpl("Display-off");
     }
 
     boolean updateFocusedWindowLocked(int mode, boolean updateInputWindows) {
@@ -2619,20 +2623,29 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
         }
     }
 
-    ActivityTaskManagerInternal.SleepToken createSleepToken(String tag, int displayId) {
+    SleepToken createSleepToken(String tag, int displayId) {
         final DisplayContent display = getDisplayContent(displayId);
         if (display == null) {
             throw new IllegalArgumentException("Invalid display: " + displayId);
         }
 
-        final SleepTokenImpl token = new SleepTokenImpl(tag, displayId);
-        mSleepTokens.add(token);
-        display.mAllSleepTokens.add(token);
+        final int tokenKey = makeSleepTokenKey(tag, displayId);
+        SleepToken token = mSleepTokens.get(tokenKey);
+        if (token == null) {
+            token = new SleepToken(tag, displayId);
+            mSleepTokens.put(tokenKey, token);
+            display.mAllSleepTokens.add(token);
+        } else {
+            throw new RuntimeException("Create the same sleep token twice: " + token);
+        }
         return token;
     }
 
-    private void removeSleepToken(SleepTokenImpl token) {
-        mSleepTokens.remove(token);
+    void removeSleepToken(SleepToken token) {
+        if (!mSleepTokens.contains(token.mHashKey)) {
+            Slog.d(TAG, "Remove non-exist sleep token: " + token + " from " + Debug.getCallers(6));
+        }
+        mSleepTokens.remove(token.mHashKey);
 
         final DisplayContent display = getDisplayContent(token.mDisplayId);
         if (display != null) {
@@ -3613,22 +3626,22 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
         return printed[0];
     }
 
-    private final class SleepTokenImpl extends ActivityTaskManagerInternal.SleepToken {
+    private static int makeSleepTokenKey(String tag, int displayId) {
+        final String tokenKey = tag + displayId;
+        return tokenKey.hashCode();
+    }
+
+    static final class SleepToken {
         private final String mTag;
         private final long mAcquireTime;
         private final int mDisplayId;
+        final int mHashKey;
 
-        public SleepTokenImpl(String tag, int displayId) {
+        SleepToken(String tag, int displayId) {
             mTag = tag;
             mDisplayId = displayId;
             mAcquireTime = SystemClock.uptimeMillis();
-        }
-
-        @Override
-        public void release() {
-            synchronized (mService.mGlobalLock) {
-                removeSleepToken(this);
-            }
+            mHashKey = makeSleepTokenKey(mTag, mDisplayId);
         }
 
         @Override