Block the shutdown process until the snapshots are written to disk.

Reuse existing timeout to ensure task snapshots are written to disk
before device shuts down.

Flag: com.android.window.flags.record_task_snapshots_before_shutdown
Bug: 376821232
Test: launch multiple tasks into desktop mode, reboot then verify the
tasksnapshots were written to disk.

Change-Id: I36578099810a581af8f44ffda2845aad1c3547c4

4 files changed, 63 insertions, 20 deletions

diff --git a/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java b/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java
index 70a8f563275f..a077a0b9a2ca 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java
@@ -2054,6 +2054,8 @@ public class ActivityTaskSupervisor implements RecentTasks.Callbacks {
                 break;
             }
         }
+        long timeRemaining = endTime - System.currentTimeMillis();
+        mWindowManager.mSnapshotController.mTaskSnapshotController.waitFlush(timeRemaining);
 
         // Force checkReadyForSleep to complete.
         checkReadyForSleepLocked(false /* allowDelay */);
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index c89feb41e723..46312aff1fb6 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2853,11 +2853,9 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     }
 
     void prepareForShutdown() {
+        mWindowManager.mSnapshotController.mTaskSnapshotController.prepareShutdown();
         for (int i = 0; i < getChildCount(); i++) {
-            final int displayId = getChildAt(i).mDisplayId;
-            mWindowManager.mSnapshotController.mTaskSnapshotController
-                    .snapshotForShutdown(displayId);
-            createSleepToken("shutdown", displayId);
+            createSleepToken("shutdown", getChildAt(i).mDisplayId);
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/SnapshotPersistQueue.java b/services/core/java/com/android/server/wm/SnapshotPersistQueue.java
index bd8e8f4008de..8b63ecf7135f 100644
--- a/services/core/java/com/android/server/wm/SnapshotPersistQueue.java
+++ b/services/core/java/com/android/server/wm/SnapshotPersistQueue.java
@@ -103,12 +103,42 @@ class SnapshotPersistQueue {
     }
 
     /**
-     * Write out everything in the queue because of shutdown.
+     * Prepare to enqueue all visible task snapshots because of shutdown.
      */
-    void shutdown() {
+    void prepareShutdown() {
         synchronized (mLock) {
             mShutdown = true;
-            mLock.notifyAll();
+        }
+    }
+
+    private boolean isQueueEmpty() {
+        synchronized (mLock) {
+            return mWriteQueue.isEmpty() || mQueueIdling || mPaused;
+        }
+    }
+
+    void waitFlush(long timeout) {
+        if (timeout <= 0) {
+            return;
+        }
+        final long endTime = System.currentTimeMillis() + timeout;
+        while (true) {
+            if (!isQueueEmpty()) {
+                long timeRemaining = endTime - System.currentTimeMillis();
+                if (timeRemaining > 0) {
+                    synchronized (mLock) {
+                        try {
+                            mLock.wait(timeRemaining);
+                        } catch (InterruptedException e) {
+                        }
+                    }
+                } else {
+                    Slog.w(TAG, "Snapshot Persist Queue flush timed out");
+                    break;
+                }
+            } else {
+                break;
+            }
         }
     }
 
@@ -139,7 +169,9 @@ class SnapshotPersistQueue {
             mWriteQueue.addLast(item);
         }
         item.onQueuedLocked();
-        ensureStoreQueueDepthLocked();
+        if (!mShutdown) {
+            ensureStoreQueueDepthLocked();
+        }
         if (!mPaused) {
             mLock.notifyAll();
         }
@@ -213,6 +245,9 @@ class SnapshotPersistQueue {
                     if (!writeQueueEmpty && !mPaused) {
                         continue;
                     }
+                    if (mShutdown && writeQueueEmpty) {
+                        mLock.notifyAll();
+                    }
                     try {
                         mQueueIdling = writeQueueEmpty;
                         mLock.wait();
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 9fe3f7563902..c130931277fe 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -309,23 +309,31 @@ class TaskSnapshotController extends AbsAppSnapshotController<Task, TaskSnapshot
     /**
      * Record task snapshots before shutdown.
      */
-    void snapshotForShutdown(int displayId) {
+    void prepareShutdown() {
         if (!com.android.window.flags.Flags.recordTaskSnapshotsBeforeShutdown()) {
             return;
         }
-        final DisplayContent displayContent = mService.mRoot.getDisplayContent(displayId);
-        if (displayContent == null) {
+        // Make write items run in a batch.
+        mPersister.mSnapshotPersistQueue.setPaused(true);
+        mPersister.mSnapshotPersistQueue.prepareShutdown();
+        for (int i = 0; i < mService.mRoot.getChildCount(); i++) {
+            mService.mRoot.getChildAt(i).forAllLeafTasks(task -> {
+                if (task.isVisible() && !task.isActivityTypeHome()) {
+                    final TaskSnapshot snapshot = captureSnapshot(task);
+                    if (snapshot != null) {
+                        mPersister.persistSnapshot(task.mTaskId, task.mUserId, snapshot);
+                    }
+                }
+            }, true /* traverseTopToBottom */);
+        }
+        mPersister.mSnapshotPersistQueue.setPaused(false);
+    }
+
+    void waitFlush(long timeout) {
+        if (!com.android.window.flags.Flags.recordTaskSnapshotsBeforeShutdown()) {
             return;
         }
-        displayContent.forAllLeafTasks(task -> {
-            if (task.isVisible() && !task.isActivityTypeHome()) {
-                final TaskSnapshot snapshot = captureSnapshot(task);
-                if (snapshot != null) {
-                    mPersister.persistSnapshot(task.mTaskId, task.mUserId, snapshot);
-                }
-            }
-        }, true /* traverseTopToBottom */);
-        mPersister.mSnapshotPersistQueue.shutdown();
+        mPersister.mSnapshotPersistQueue.waitFlush(timeout);
     }
 
     /**