From 1137b137ddbd1cbe139aad39f521fce8d84912eb Mon Sep 17 00:00:00 2001
From: Ming-Shin Lu <lumark@google.com>
Date: Thu, 24 Mar 2022 20:24:22 +0800
Subject: [PATCH] Fix IME snapshot missing to remove by 2 consecutive creation
 calls

There is a timing could happen to create another new IME snapshot
before recents animation finish, so system lost the previous
one's reference and never remove it.

Ensure removing the obsoleted IME snapshot before we create a new one.

Fix: 224664116
Test: atest DisplayContentTests#\
    testShowImeScreenshot_removeCurSnapshotBeforeCreateNext
Change-Id: I3927557453317e9beaab0d5995d6d415d09bb05d
---
 .../com/android/server/wm/DisplayContent.java |  4 +++
 .../server/wm/DisplayContentTests.java        | 26 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index f5ace6c78288c..f16e7ffbdd803 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -4143,6 +4143,10 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         final SurfaceControl.Transaction t = getPendingTransaction();
         // Prepare IME screenshot for the target if it allows to attach into.
         if (mInputMethodWindow != null && mInputMethodWindow.isVisible()) {
+            // Remove the obsoleted IME snapshot first in case the new snapshot happens to
+            // override the current one before the transition finish and the surface never be
+            // removed on the task.
+            removeImeSurfaceImmediately();
             mImeScreenshot = new ImeScreenshot(
                     mWmService.mSurfaceControlFactory.apply(null), mImeLayeringTarget);
             mImeScreenshot.attachAndShow(t);
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index 40b460157bc29..c18f4873c152d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -2023,6 +2023,32 @@ public class DisplayContentTests extends WindowTestsBase {
         verify(mDisplayContent, never()).showImeScreenshot();
     }
 
+    @UseTestDisplay(addWindows = {W_INPUT_METHOD})
+    @Test
+    public void testShowImeScreenshot_removeCurSnapshotBeforeCreateNext() {
+        final Task rootTask = createTask(mDisplayContent);
+        final Task task = createTaskInRootTask(rootTask, 0 /* userId */);
+        final ActivityRecord activity = createActivityRecord(mDisplayContent, task);
+        final WindowState win = createWindow(null, TYPE_BASE_APPLICATION, activity, "win");
+
+        mDisplayContent.setImeLayeringTarget(win);
+        mDisplayContent.setImeInputTarget(win);
+        spyOn(mDisplayContent);
+        spyOn(mDisplayContent.mInputMethodWindow);
+        doReturn(true).when(mDisplayContent.mInputMethodWindow).isVisible();
+        mDisplayContent.getInsetsStateController().getImeSourceProvider().setImeShowing(true);
+
+        // Verify when the timing of 2 showImeScreenshot invocations are very close, will first
+        // detach the current snapshot then create the next one.
+        mDisplayContent.showImeScreenshot();
+        DisplayContent.ImeScreenshot curSnapshot = mDisplayContent.mImeScreenshot;
+        spyOn(curSnapshot);
+        mDisplayContent.showImeScreenshot();
+        verify(curSnapshot).detach(any());
+        assertNotNull(mDisplayContent.mImeScreenshot);
+        assertNotEquals(curSnapshot, mDisplayContent.mImeScreenshot);
+    }
+
     @Test
     public void testRotateBounds_keepSamePhysicalPosition() {
         final DisplayContent dc =