Drawing app theme snapshot without hold wm lock.

A potential deadlock could happen when hold wm lock and require
hardware buffer from createHardwareBitmap, the rendering thread can be
blocked because it is waiting for dequeue buffer.
Defer the create snapshot process in a supplier, to escape the wm lock
if the snapshot should be draw by render thread.
Also add a parameter allowAppTheme in captureSnapshot, then the caller
can decide to not capture fake snapshot.

Flag: com.android.window.flags.exclude_drawing_app_theme_snapshot_from_lock
Bug: 373502791
Test: atest ActivitySnapshotControllerTests TransitionTests
TaskSnapshotControllerTest
Test: verify there won't block render thread while capture app snapshot
for app theme mode.

Change-Id: I7add8baecb0d068bac485fe9eee0c90f23d2f823

9 files changed, 233 insertions, 63 deletions

diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index 9d11d149b0ed..80f42e587c50 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -430,4 +430,15 @@ flag {
   description: "Support insets definition and calculation relative to task bounds."
   bug: "277292497"
   is_fixed_read_only: true
+}
+
+flag {
+    name: "exclude_drawing_app_theme_snapshot_from_lock"
+    namespace: "windowing_frontend"
+    description: "Do not hold wm lock when drawing app theme snapshot."
+    is_fixed_read_only: true
+    bug: "373502791"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
 }
\ No newline at end of file
diff --git a/services/core/java/com/android/server/wm/AbsAppSnapshotController.java b/services/core/java/com/android/server/wm/AbsAppSnapshotController.java
index 19eba5fe5755..90c2216f6b22 100644
--- a/services/core/java/com/android/server/wm/AbsAppSnapshotController.java
+++ b/services/core/java/com/android/server/wm/AbsAppSnapshotController.java
@@ -51,8 +51,11 @@ import android.window.TaskSnapshot;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.graphics.ColorUtils;
 import com.android.server.wm.utils.InsetUtils;
+import com.android.window.flags.Flags;
 
 import java.io.PrintWriter;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 /**
  * Base class for a Snapshot controller
@@ -148,43 +151,60 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
     protected abstract Rect getLetterboxInsets(ActivityRecord topActivity);
 
     /**
-     * This is different than {@link #recordSnapshotInner(TYPE)} because it doesn't store
-     * the snapshot to the cache and returns the TaskSnapshot immediately.
-     *
-     * This is only used for testing so the snapshot content can be verified.
+     * This is different than {@link #recordSnapshotInner(TYPE, boolean, Consumer)}  because it
+     * doesn't store the snapshot to the cache and returns the TaskSnapshot immediately.
      */
     @VisibleForTesting
-    TaskSnapshot captureSnapshot(TYPE source) {
-        final TaskSnapshot snapshot;
+    SnapshotSupplier captureSnapshot(TYPE source, boolean allowAppTheme) {
+        final SnapshotSupplier supplier = new SnapshotSupplier();
         switch (getSnapshotMode(source)) {
-            case SNAPSHOT_MODE_NONE:
-                return null;
             case SNAPSHOT_MODE_APP_THEME:
-                snapshot = drawAppThemeSnapshot(source);
+                Trace.traceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "drawAppThemeSnapshot");
+                if (Flags.excludeDrawingAppThemeSnapshotFromLock()) {
+                    if (allowAppTheme) {
+                        supplier.setSupplier(drawAppThemeSnapshot(source));
+                    }
+                } else {
+                    final Supplier<TaskSnapshot> original = drawAppThemeSnapshot(source);
+                    final TaskSnapshot snapshot = original != null ? original.get() : null;
+                    supplier.setSnapshot(snapshot);
+                }
+                Trace.traceEnd(Trace.TRACE_TAG_WINDOW_MANAGER);
                 break;
             case SNAPSHOT_MODE_REAL:
-                snapshot = snapshot(source);
+                supplier.setSnapshot(snapshot(source));
                 break;
             default:
-                snapshot = null;
                 break;
         }
-        return snapshot;
+        return supplier;
     }
 
-    final TaskSnapshot recordSnapshotInner(TYPE source) {
+    /**
+     * @param allowAppTheme If true, allows to draw app theme snapshot when it's not allowed to take
+     *                      a real screenshot, but create a fake representation of the app.
+     * @param inLockConsumer Extra task to do in WM lock when first get the snapshot object.
+     */
+    final SnapshotSupplier recordSnapshotInner(TYPE source, boolean allowAppTheme,
+            @Nullable Consumer<TaskSnapshot> inLockConsumer) {
         if (shouldDisableSnapshots()) {
             return null;
         }
-        final TaskSnapshot snapshot = captureSnapshot(source);
-        if (snapshot == null) {
-            return null;
-        }
-        mCache.putSnapshot(source, snapshot);
-        return snapshot;
+        final SnapshotSupplier supplier = captureSnapshot(source, allowAppTheme);
+        supplier.setConsumer(t -> {
+            synchronized (mService.mGlobalLock) {
+                if (!source.isAttached()) {
+                    return;
+                }
+                mCache.putSnapshot(source, t);
+                if (inLockConsumer != null) {
+                    inLockConsumer.accept(t);
+                }
+            }
+        });
+        return supplier;
     }
 
-    @VisibleForTesting
     int getSnapshotMode(TYPE source) {
         final int type = source.getActivityType();
         if (type == ACTIVITY_TYPE_RECENTS || type == ACTIVITY_TYPE_DREAM) {
@@ -400,7 +420,7 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
      * If we are not allowed to take a real screenshot, this attempts to represent the app as best
      * as possible by using the theme's window background.
      */
-    private TaskSnapshot drawAppThemeSnapshot(TYPE source) {
+    private Supplier<TaskSnapshot> drawAppThemeSnapshot(TYPE source) {
         final ActivityRecord topActivity = getTopActivity(source);
         if (topActivity == null) {
             return null;
@@ -432,26 +452,46 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
         decorPainter.setInsets(systemBarInsets);
         decorPainter.drawDecors(c /* statusBarExcludeFrame */, null /* alreadyDrawFrame */);
         node.end(c);
-        final Bitmap hwBitmap = ThreadedRenderer.createHardwareBitmap(node, width, height);
-        if (hwBitmap == null) {
-            return null;
-        }
+
         final Rect contentInsets = new Rect(systemBarInsets);
         final Rect letterboxInsets = getLetterboxInsets(topActivity);
         InsetUtils.addInsets(contentInsets, letterboxInsets);
-        // Note, the app theme snapshot is never translucent because we enforce a non-translucent
-        // color above
-        final TaskSnapshot taskSnapshot = new TaskSnapshot(
-                System.currentTimeMillis() /* id */,
-                SystemClock.elapsedRealtimeNanos() /* captureTime */,
-                topActivity.mActivityComponent, hwBitmap.getHardwareBuffer(),
-                hwBitmap.getColorSpace(), mainWindow.getConfiguration().orientation,
-                mainWindow.getWindowConfiguration().getRotation(), new Point(taskWidth, taskHeight),
-                contentInsets, letterboxInsets, false /* isLowResolution */,
-                false /* isRealSnapshot */, source.getWindowingMode(),
-                attrs.insetsFlags.appearance, false /* isTranslucent */, false /* hasImeSurface */,
-                topActivity.getConfiguration().uiMode /* uiMode */);
-        return validateSnapshot(taskSnapshot);
+
+        final TaskSnapshot.Builder builder = new TaskSnapshot.Builder();
+        builder.setIsRealSnapshot(false);
+        builder.setId(System.currentTimeMillis());
+        builder.setContentInsets(contentInsets);
+        builder.setLetterboxInsets(letterboxInsets);
+
+        builder.setTopActivityComponent(topActivity.mActivityComponent);
+        // Note, the app theme snapshot is never translucent because we enforce a
+        // non-translucent color above.
+        builder.setIsTranslucent(false);
+        builder.setWindowingMode(source.getWindowingMode());
+        builder.setAppearance(attrs.insetsFlags.appearance);
+        builder.setUiMode(topActivity.getConfiguration().uiMode);
+
+        builder.setRotation(mainWindow.getWindowConfiguration().getRotation());
+        builder.setOrientation(mainWindow.getConfiguration().orientation);
+        builder.setTaskSize(new Point(taskWidth, taskHeight));
+        builder.setCaptureTime(SystemClock.elapsedRealtimeNanos());
+
+        return () -> {
+            try {
+                Trace.traceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "drawAppThemeSnapshot_acquire");
+                // Do not hold WM lock when calling to render thread.
+                final Bitmap hwBitmap = ThreadedRenderer.createHardwareBitmap(node, width,
+                        height);
+                if (hwBitmap == null) {
+                    return null;
+                }
+                builder.setSnapshot(hwBitmap.getHardwareBuffer());
+                builder.setColorSpace(hwBitmap.getColorSpace());
+                return validateSnapshot(builder.build());
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_WINDOW_MANAGER);
+            }
+        };
     }
 
     static Rect getSystemBarInsets(Rect frame, InsetsState state) {
@@ -482,4 +522,45 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
         pw.println(prefix + "mSnapshotEnabled=" + mSnapshotEnabled);
         mCache.dump(pw, prefix);
     }
+
+    static class SnapshotSupplier implements Supplier<TaskSnapshot> {
+
+        private TaskSnapshot mSnapshot;
+        private boolean mHasSet;
+        private Consumer<TaskSnapshot> mConsumer;
+        private Supplier<TaskSnapshot> mSupplier;
+
+        /** Callback when the snapshot is get for the first time. */
+        void setConsumer(@NonNull Consumer<TaskSnapshot> consumer) {
+            mConsumer = consumer;
+        }
+
+        void setSupplier(@NonNull Supplier<TaskSnapshot> createSupplier) {
+            mSupplier = createSupplier;
+        }
+
+        void setSnapshot(TaskSnapshot snapshot) {
+            mSnapshot = snapshot;
+        }
+
+        void handleSnapshot() {
+            if (mHasSet) {
+                return;
+            }
+            mHasSet = true;
+            if (mSnapshot == null) {
+                mSnapshot = mSupplier != null ? mSupplier.get() : null;
+            }
+            if (mConsumer != null && mSnapshot != null) {
+                mConsumer.accept(mSnapshot);
+            }
+        }
+
+        @Override
+        @Nullable
+        public TaskSnapshot get() {
+            handleSnapshot();
+            return mSnapshot;
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/wm/ActivitySnapshotController.java b/services/core/java/com/android/server/wm/ActivitySnapshotController.java
index 9aaa0e1cfd6b..cfd324830db5 100644
--- a/services/core/java/com/android/server/wm/ActivitySnapshotController.java
+++ b/services/core/java/com/android/server/wm/ActivitySnapshotController.java
@@ -38,6 +38,7 @@ import com.android.server.wm.BaseAppSnapshotPersister.PersistInfoProvider;
 import java.io.File;
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.function.Supplier;
 
 /**
  * When an app token becomes invisible, we take a snapshot (bitmap) and put it into our cache.
@@ -355,7 +356,9 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
         final int[] mixedCode = new int[size];
         if (size == 1) {
             final ActivityRecord singleActivity = activity.get(0);
-            final TaskSnapshot snapshot = recordSnapshotInner(singleActivity);
+            final Supplier<TaskSnapshot> supplier = recordSnapshotInner(singleActivity,
+                    false /* allowAppTheme */, null /* inLockConsumer */);
+            final TaskSnapshot snapshot = supplier != null ? supplier.get() : null;
             if (snapshot != null) {
                 mixedCode[0] = getSystemHashCode(singleActivity);
                 addUserSavedFile(singleActivity.mUserId, snapshot, mixedCode);
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 5eee8ece6a67..290f155bb4cd 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -314,6 +314,7 @@ import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.Supplier;
 
 /**
  * System service for managing activities and their containers (task, displays,... ).
@@ -4038,6 +4039,7 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
         mAmInternal.enforceCallingPermission(READ_FRAME_BUFFER, "takeTaskSnapshot()");
         final long ident = Binder.clearCallingIdentity();
         try {
+            final Supplier<TaskSnapshot> supplier;
             synchronized (mGlobalLock) {
                 final Task task = mRootWindowContainer.anyTaskForId(taskId,
                         MATCH_ATTACHED_TASK_OR_RECENT_TASKS);
@@ -4050,11 +4052,13 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
                 // be retrieved by recents. While if updateCache is false, the real snapshot will
                 // always be taken and the snapshot won't be put into SnapshotPersister.
                 if (updateCache) {
-                    return mWindowManager.mTaskSnapshotController.recordSnapshot(task);
+                    supplier = mWindowManager.mTaskSnapshotController
+                            .getRecordSnapshotSupplier(task);
                 } else {
                     return mWindowManager.mTaskSnapshotController.snapshot(task);
                 }
             }
+            return supplier != null ? supplier.get() : null;
         } finally {
             Binder.restoreCallingIdentity(ident);
         }
@@ -6403,6 +6407,7 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
         @Override
         public boolean shuttingDown(boolean booted, int timeout) {
             mShuttingDown = true;
+            mWindowManager.mSnapshotController.mTaskSnapshotController.prepareShutdown();
             synchronized (mGlobalLock) {
                 mRootWindowContainer.prepareForShutdown();
                 updateEventDispatchingLocked(booted);
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 3d2868540334..865d5facc4a4 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2926,7 +2926,6 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     }
 
     void prepareForShutdown() {
-        mWindowManager.mSnapshotController.mTaskSnapshotController.prepareShutdown();
         for (int i = 0; i < getChildCount(); i++) {
             createSleepToken("shutdown", getChildAt(i).mDisplayId);
         }
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 38a2ebeba332..7d300e98f44b 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -36,7 +36,9 @@ import android.window.TaskSnapshot;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.policy.WindowManagerPolicy.ScreenOffListener;
 import com.android.server.wm.BaseAppSnapshotPersister.PersistInfoProvider;
+import com.android.window.flags.Flags;
 
+import java.util.ArrayList;
 import java.util.Set;
 
 /**
@@ -154,6 +156,8 @@ class TaskSnapshotController extends AbsAppSnapshotController<Task, TaskSnapshot
      * The attributes of task snapshot are based on task configuration. But sometimes the
      * configuration may have been changed during a transition, so supply the ChangeInfo that
      * stored the previous appearance of the closing task.
+     *
+     * The snapshot won't be created immediately if it should be captured as fake snapshot.
      */
     void recordSnapshot(Task task, Transition.ChangeInfo changeInfo) {
         mCurrentChangeInfo = changeInfo;
@@ -164,13 +168,35 @@ class TaskSnapshotController extends AbsAppSnapshotController<Task, TaskSnapshot
         }
     }
 
-    TaskSnapshot recordSnapshot(Task task) {
-        final TaskSnapshot snapshot = recordSnapshotInner(task);
-        if (snapshot != null && !task.isActivityTypeHome()) {
-            mPersister.persistSnapshot(task.mTaskId, task.mUserId, snapshot);
-            task.onSnapshotChanged(snapshot);
+    void recordSnapshot(Task task) {
+        if (shouldDisableSnapshots()) {
+            return;
+        }
+        final SnapshotSupplier supplier = getRecordSnapshotSupplier(task);
+        if (supplier == null) {
+            return;
         }
-        return snapshot;
+        final int mode = getSnapshotMode(task);
+        if (Flags.excludeDrawingAppThemeSnapshotFromLock() && mode == SNAPSHOT_MODE_APP_THEME) {
+            mService.mH.post(supplier::handleSnapshot);
+        } else {
+            supplier.handleSnapshot();
+        }
+    }
+
+    /**
+     * Note that the snapshot is not created immediately, if the returned supplier is non-null, the
+     * caller must call {@link AbsAppSnapshotController.SnapshotSupplier#get} or
+     * {@link AbsAppSnapshotController.SnapshotSupplier#handleSnapshot} to complete the entire
+     * record request.
+     */
+    SnapshotSupplier getRecordSnapshotSupplier(Task task) {
+        return recordSnapshotInner(task, true /* allowAppTheme */, snapshot -> {
+            if (!task.isActivityTypeHome()) {
+                mPersister.persistSnapshot(task.mTaskId, task.mUserId, snapshot);
+                task.onSnapshotChanged(snapshot);
+            }
+        });
     }
 
     /**
@@ -328,27 +354,38 @@ class TaskSnapshotController extends AbsAppSnapshotController<Task, TaskSnapshot
      * Record task snapshots before shutdown.
      */
     void prepareShutdown() {
-        if (!com.android.window.flags.Flags.recordTaskSnapshotsBeforeShutdown()) {
+        if (!Flags.recordTaskSnapshotsBeforeShutdown()) {
             return;
         }
-        // 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);
+        final ArrayList<SnapshotSupplier> supplierArrayList = new ArrayList<>();
+        synchronized (mService.mGlobalLock) {
+            // 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 SnapshotSupplier supplier = captureSnapshot(task,
+                                true /* allowAppTheme */);
+                        if (supplier != null) {
+                            supplier.setConsumer(t ->
+                                    mPersister.persistSnapshot(task.mTaskId, task.mUserId, t));
+                            supplierArrayList.add(supplier);
+                        }
                     }
-                }
-            }, true /* traverseTopToBottom */);
+                }, true /* traverseTopToBottom */);
+            }
+        }
+        for (int i = supplierArrayList.size() - 1; i >= 0; --i) {
+            supplierArrayList.get(i).handleSnapshot();
+        }
+        synchronized (mService.mGlobalLock) {
+            mPersister.mSnapshotPersistQueue.setPaused(false);
         }
-        mPersister.mSnapshotPersistQueue.setPaused(false);
     }
 
     void waitFlush(long timeout) {
-        if (!com.android.window.flags.Flags.recordTaskSnapshotsBeforeShutdown()) {
+        if (!Flags.recordTaskSnapshotsBeforeShutdown()) {
             return;
         }
         mPersister.mSnapshotPersistQueue.waitFlush(timeout);
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 9d9c53dfe0f4..ac368d87877f 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -10084,14 +10084,16 @@ public class WindowManagerService extends IWindowManager.Stub
         TaskSnapshot taskSnapshot;
         final long token = Binder.clearCallingIdentity();
         try {
+            final Supplier<TaskSnapshot> supplier;
             synchronized (mGlobalLock) {
                 Task task = mRoot.anyTaskForId(taskId, MATCH_ATTACHED_TASK_OR_RECENT_TASKS);
                 if (task == null) {
                     throw new IllegalArgumentException(
                             "Failed to find matching task for taskId=" + taskId);
                 }
-                taskSnapshot = mTaskSnapshotController.captureSnapshot(task);
+                supplier = mTaskSnapshotController.captureSnapshot(task, true /* allowAppTheme */);
             }
+            taskSnapshot = supplier != null ? supplier.get() : null;
         } finally {
             Binder.restoreCallingIdentity(token);
         }
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java
index a7fc10f2fcc5..948371f74a9c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java
@@ -29,6 +29,7 @@ import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.never;
 
@@ -253,7 +254,11 @@ public class ActivitySnapshotControllerTests extends TaskSnapshotPersisterTestBa
      */
     @Test
     public void testSkipRecordActivity() {
-        doReturn(createSnapshot()).when(mActivitySnapshotController).recordSnapshotInner(any());
+        final AbsAppSnapshotController.SnapshotSupplier supplier =
+                new AbsAppSnapshotController.SnapshotSupplier();
+        supplier.setSupplier(this::createSnapshot);
+        doReturn(supplier).when(mActivitySnapshotController).recordSnapshotInner(
+                any(), anyBoolean(), any());
         final Task task = createTask(mDisplayContent);
 
         mSnapshotPersistQueue.setPaused(true);
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
index 6655932b060b..c6b2a6b8d42f 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
@@ -33,6 +33,7 @@ import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -45,12 +46,15 @@ import android.graphics.PixelFormat;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.util.ArraySet;
 import android.window.TaskSnapshot;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.window.flags.Flags;
+
 import com.google.android.collect.Sets;
 
 import org.junit.Test;
@@ -285,4 +289,27 @@ public class TaskSnapshotControllerTest extends WindowTestsBase {
 
         assertFalse(success);
     }
+
+    @Test
+    @EnableFlags(Flags.FLAG_EXCLUDE_DRAWING_APP_THEME_SNAPSHOT_FROM_LOCK)
+    public void testRecordTaskSnapshot() {
+        spyOn(mWm.mTaskSnapshotController.mCache);
+        spyOn(mWm.mTaskSnapshotController);
+        doReturn(false).when(mWm.mTaskSnapshotController).shouldDisableSnapshots();
+
+        final WindowState normalWindow = createWindow(null,
+                FIRST_APPLICATION_WINDOW, mDisplayContent, "normalWindow");
+        final TaskSnapshot snapshot = new TaskSnapshotPersisterTestBase.TaskSnapshotBuilder()
+                .setTopActivityComponent(normalWindow.mActivityRecord.mActivityComponent).build();
+        doReturn(snapshot).when(mWm.mTaskSnapshotController).snapshot(any());
+        final Task task = normalWindow.mActivityRecord.getTask();
+        mWm.mTaskSnapshotController.recordSnapshot(task);
+        verify(mWm.mTaskSnapshotController.mCache).putSnapshot(eq(task), any());
+        clearInvocations(mWm.mTaskSnapshotController.mCache);
+
+        normalWindow.mAttrs.flags |= FLAG_SECURE;
+        mWm.mTaskSnapshotController.recordSnapshot(task);
+        waitHandlerIdle(mWm.mH);
+        verify(mWm.mTaskSnapshotController.mCache).putSnapshot(eq(task), any());
+    }
 }