Add starting windows while unlocking

- Modify screen capturing logic a bit such that it can also take a
  screenshot when the display is off.
- Also take snapshot when app visibility is changing without
  proper app transition.
- When unlocking, add strating windows for all visible apps.

Test: Unlock phone.
Test: Have an app that sleeps in onStart, make sure unlocking is
instant.
Bug: 31339431
Change-Id: I953ab6cb30d0d264554fd49a46bdc56e23356d13
diff --git a/services/core/java/com/android/server/am/ActivityStack.java b/services/core/java/com/android/server/am/ActivityStack.java
index d7b3728e..963a9dc 100644
--- a/services/core/java/com/android/server/am/ActivityStack.java
+++ b/services/core/java/com/android/server/am/ActivityStack.java
@@ -1795,6 +1795,12 @@
         }
     }
 
+    void addStartingWindowsForVisibleActivities(boolean taskSwitch) {
+        for (int taskNdx = mTaskHistory.size() - 1; taskNdx >= 0; --taskNdx) {
+            mTaskHistory.get(taskNdx).addStartingWindowsForVisibleActivities(taskSwitch);
+        }
+    }
+
     /**
      * @return true if the top visible activity wants to occlude the Keyguard, false otherwise
      */
diff --git a/services/core/java/com/android/server/am/ActivityStackSupervisor.java b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
index 14899b4..2c1c3a1 100644
--- a/services/core/java/com/android/server/am/ActivityStackSupervisor.java
+++ b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
@@ -3229,6 +3229,17 @@
         }
     }
 
+    void addStartingWindowsForVisibleActivities(boolean taskSwitch) {
+        for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
+            final ArrayList<ActivityStack> stacks = mActivityDisplays.valueAt(displayNdx).mStacks;
+            final int topStackNdx = stacks.size() - 1;
+            for (int stackNdx = topStackNdx; stackNdx >= 0; --stackNdx) {
+                final ActivityStack stack = stacks.get(stackNdx);
+                stack.addStartingWindowsForVisibleActivities(taskSwitch);
+            }
+        }
+    }
+
     void invalidateTaskLayers() {
         mTaskLayersChanged = true;
     }
diff --git a/services/core/java/com/android/server/am/KeyguardController.java b/services/core/java/com/android/server/am/KeyguardController.java
index cfe2eb0..b0a4746 100644
--- a/services/core/java/com/android/server/am/KeyguardController.java
+++ b/services/core/java/com/android/server/am/KeyguardController.java
@@ -29,6 +29,7 @@
 import static com.android.server.wm.AppTransition.TRANSIT_KEYGUARD_GOING_AWAY;
 import static com.android.server.wm.AppTransition.TRANSIT_KEYGUARD_OCCLUDE;
 import static com.android.server.wm.AppTransition.TRANSIT_KEYGUARD_UNOCCLUDE;
+import static com.android.server.wm.AppTransition.TRANSIT_NONE;
 import static com.android.server.wm.AppTransition.TRANSIT_UNSET;
 
 import android.os.IBinder;
@@ -120,6 +121,7 @@
 
                 // Some stack visibility might change (e.g. docked stack)
                 mStackSupervisor.ensureActivitiesVisibleLocked(null, 0, !PRESERVE_WINDOWS);
+                mStackSupervisor.addStartingWindowsForVisibleActivities(true /* taskSwitch */);
                 mWindowManager.executeAppTransition();
             } finally {
                 mWindowManager.continueSurfaceLayout();
diff --git a/services/core/java/com/android/server/am/TaskRecord.java b/services/core/java/com/android/server/am/TaskRecord.java
index 1f5152a..2745fe5 100644
--- a/services/core/java/com/android/server/am/TaskRecord.java
+++ b/services/core/java/com/android/server/am/TaskRecord.java
@@ -1981,6 +1981,15 @@
         return rootAffinity != null && getStackId() != PINNED_STACK_ID;
     }
 
+    void addStartingWindowsForVisibleActivities(boolean taskSwitch) {
+        for (int activityNdx = mActivities.size() - 1; activityNdx >= 0; --activityNdx) {
+            final ActivityRecord r = mActivities.get(activityNdx);
+            if (r.visible) {
+                r.showStartingWindow(null /* prev */, false /* newTask */, taskSwitch);
+            }
+        }
+    }
+
     void dump(PrintWriter pw, String prefix) {
         pw.print(prefix); pw.print("userId="); pw.print(userId);
                 pw.print(" effectiveUid="); UserHandle.formatUid(pw, effectiveUid);
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index ac9859d..29c685d 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -375,6 +375,7 @@
                 // affected.
                 mService.getDefaultDisplayContentLocked().getDockedDividerController()
                         .notifyAppVisibilityChanged();
+                mService.mTaskSnapshotController.notifyAppVisibilityChanged(this, visible);
             }
         }
 
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 3958510..73e88c1 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -2282,7 +2282,7 @@
                 final boolean foundTargetWs =
                         (w.mAppToken != null && w.mAppToken.token == appToken)
                                 || (mScreenshotApplicationState.appWin != null && wallpaperOnly);
-                if (foundTargetWs && w.isDisplayedLw() && winAnim.getShown()) {
+                if (foundTargetWs && w.isVisible() && winAnim.getShown()) {
                     mScreenshotApplicationState.screenshotReady = true;
                 }
 
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 15878f6..2b74f84 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -26,6 +26,8 @@
 import android.util.ArraySet;
 import android.view.WindowManagerPolicy.StartingSurface;
 
+import com.google.android.collect.Sets;
+
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.io.PrintWriter;
@@ -66,10 +68,27 @@
         if (!ENABLE_TASK_SNAPSHOTS) {
             return;
         }
+        handleClosingApps(mService.mClosingApps);
+    }
+
+
+    /**
+     * Called when the visibility of an app changes outside of the regular app transition flow.
+     */
+    void notifyAppVisibilityChanged(AppWindowToken appWindowToken, boolean visible) {
+        if (!ENABLE_TASK_SNAPSHOTS) {
+            return;
+        }
+        if (!visible) {
+            handleClosingApps(Sets.newArraySet(appWindowToken));
+        }
+    }
+
+    private void handleClosingApps(ArraySet<AppWindowToken> closingApps) {
 
         // We need to take a snapshot of the task if and only if all activities of the task are
         // either closing or hidden.
-        getClosingTasks(mService.mClosingApps, mTmpTasks);
+        getClosingTasks(closingApps, mTmpTasks);
         for (int i = mTmpTasks.size() - 1; i >= 0; i--) {
             final Task task = mTmpTasks.valueAt(i);
             if (!canSnapshotTask(task)) {