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: 162441580
Test: atest MultiDisplayKeyguardTests MultiDisplayLockedKeyguardTests/
KeyguardTests KeyguardInputTests KeyguardLockedTests

Change-Id: I168dfe0fa101cb65eb676ca256c03439c8e4193d
Merged-In: I168dfe0fa101cb65eb676ca256c03439c8e4193d
(cherry picked from commit 3f51a96792d3770d82efcad645714ea8445a445f)
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java b/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
index d5df906..a903bcd 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 @@
      * 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 205523b..a1d0e56 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -5087,7 +5087,10 @@
         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());
         }
@@ -5521,12 +5524,35 @@
                 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);
+                }
+            }
         }
     }
 
@@ -6085,9 +6111,9 @@
 
     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 e1693ba..fcbc756 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -596,9 +596,9 @@
     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;
 
@@ -934,6 +934,7 @@
         mAtmService = mWmService.mAtmService;
         mDisplay = display;
         mDisplayId = display.getDisplayId();
+        mOffTokenAcquirer = mRootWindowContainer.mDisplayOffTokenAcquirer;
         mWallpaperController = new WallpaperController(mWmService, this);
         display.getDisplayInfo(mDisplayInfo);
         display.getMetrics(mDisplayMetrics);
@@ -5168,11 +5169,10 @@
         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();
@@ -5424,7 +5424,8 @@
         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 9c535e4..8ba92b5 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -49,7 +49,6 @@
 
 import com.android.internal.policy.IKeyguardDismissCallback;
 import com.android.server.policy.WindowManagerPolicy;
-import com.android.server.wm.ActivityTaskManagerInternal.SleepToken;
 
 import java.io.PrintWriter;
 
@@ -73,11 +72,14 @@
     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) {
@@ -411,17 +413,17 @@
 
     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;
@@ -442,29 +444,18 @@
         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 84a3893..049e46f 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -219,6 +219,9 @@
     // 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)}.
@@ -258,7 +261,7 @@
      * 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 hint has started, but not ended. */
     private boolean mPowerHintSent;
@@ -443,6 +446,7 @@
         mService = service.mAtmService;
         mStackSupervisor = mService.mStackSupervisor;
         mStackSupervisor.mRootWindowContainer = this;
+        mDisplayOffTokenAcquirer = mService.new SleepTokenAcquirerImpl("Display-off");
     }
 
     boolean updateFocusedWindowLocked(int mode, boolean updateInputWindows) {
@@ -2639,20 +2643,29 @@
         }
     }
 
-    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) {
@@ -3661,22 +3674,22 @@
         return printed;
     }
 
-    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