Merge "Drawing app theme snapshot without hold wm lock." into main

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 db62cebf7603..04650b9e0150 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());
+    }
 }