Persist task snapshots to disk

So they can be used again after rebooting or when the process gets
killed, but the snapshot is still used for recents.

Also implement TaskSnapshotLoader, to restore it from disk. The
infrastructure around restoring and caching snapshots for recents
will be implemented in the next CL.

Test: runtest frameworks-services -c
com.android.server.wm.TaskSnapshotPersisterLoaderTest

Bug: 31339431
Change-Id: Iaec03c4cc92e04b6dd7e623bca755ddc92613bce
diff --git a/proto/src/task_snapshot.proto b/proto/src/task_snapshot.proto
new file mode 100644
index 0000000..c9d5c27
--- /dev/null
+++ b/proto/src/task_snapshot.proto
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+ syntax = "proto3";
+
+ package com.android.server.wm;
+
+ option java_package = "com.android.server.wm";
+ option java_outer_classname = "WindowManagerProtos";
+
+ message TaskSnapshotProto {
+     int32 orientation = 1;
+     int32 inset_left = 2;
+     int32 inset_top = 3;
+     int32 inset_right = 4;
+     int32 inset_bottom = 5;
+ }
\ No newline at end of file
diff --git a/services/core/java/com/android/server/am/TaskPersister.java b/services/core/java/com/android/server/am/TaskPersister.java
index 7a62f2c..9dde39e 100644
--- a/services/core/java/com/android/server/am/TaskPersister.java
+++ b/services/core/java/com/android/server/am/TaskPersister.java
@@ -651,6 +651,8 @@
                                         "omitting from persistentTaskIds task=" + task);
                             }
                         }
+                        mService.mWindowManager.removeObsoleteTaskFiles(persistentTaskIds,
+                                mRecentTasks.usersWithRecentsLoadedLocked());
                     }
                     removeObsoleteFiles(persistentTaskIds);
                 }
diff --git a/services/core/java/com/android/server/am/TaskRecord.java b/services/core/java/com/android/server/am/TaskRecord.java
index a72a958..db6c0f7 100644
--- a/services/core/java/com/android/server/am/TaskRecord.java
+++ b/services/core/java/com/android/server/am/TaskRecord.java
@@ -788,6 +788,9 @@
             inRecents = false;
             mService.notifyTaskPersisterLocked(this, false);
         }
+
+        // TODO: Use window container controller once tasks are better synced between AM and WM
+        mService.mWindowManager.notifyTaskRemovedFromRecents(taskId, userId);
     }
 
     void setTaskToAffiliateWith(TaskRecord taskToAffiliateWith) {
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index df8679d..10ecf3b 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -22,11 +22,13 @@
 import android.app.ActivityManager.StackId;
 import android.app.ActivityManager.TaskSnapshot;
 import android.graphics.GraphicBuffer;
+import android.os.Environment;
 import android.util.ArraySet;
 import android.view.WindowManagerPolicy.StartingSurface;
 
 import com.android.internal.annotations.VisibleForTesting;
 
+import java.io.File;
 import java.io.PrintWriter;
 
 /**
@@ -45,14 +47,21 @@
 class TaskSnapshotController {
 
     private final WindowManagerService mService;
-    private final TaskSnapshotCache mCache = new TaskSnapshotCache();
 
+    private final TaskSnapshotCache mCache = new TaskSnapshotCache();
+    private final TaskSnapshotPersister mPersister = new TaskSnapshotPersister(
+            Environment::getDataSystemCeDirectory);
+    private final TaskSnapshotLoader mLoader = new TaskSnapshotLoader(mPersister);
     private final ArraySet<Task> mTmpTasks = new ArraySet<>();
 
     TaskSnapshotController(WindowManagerService service) {
         mService = service;
     }
 
+    void systemReady() {
+        mPersister.start();
+    }
+
     void onTransitionStarting() {
         if (!ENABLE_TASK_SNAPSHOTS) {
             return;
@@ -69,6 +78,7 @@
             final TaskSnapshot snapshot = snapshotTask(task);
             if (snapshot != null) {
                 mCache.putSnapshot(task, snapshot);
+                mPersister.persistSnapshot(task.mTaskId, task.mUserId, snapshot);
                 if (task.getController() != null) {
                     task.getController().reportSnapshotChanged(snapshot);
                 }
@@ -141,6 +151,17 @@
         mCache.cleanCache(wtoken);
     }
 
+    void notifyTaskRemovedFromRecents(int taskId, int userId) {
+        mPersister.onTaskRemovedFromRecents(taskId, userId);
+    }
+
+    /**
+     * See {@link TaskSnapshotPersister#removeObsoleteFiles}
+     */
+    void removeObsoleteTaskFiles(ArraySet<Integer> persistentTaskIds, int[] runningUserIds) {
+        mPersister.removeObsoleteFiles(persistentTaskIds, runningUserIds);
+    }
+
     void dump(PrintWriter pw, String prefix) {
         mCache.dump(pw, prefix);
     }
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotLoader.java b/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
new file mode 100644
index 0000000..4340822
--- /dev/null
+++ b/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
@@ -0,0 +1,93 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+
+package com.android.server.wm;
+
+import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
+import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
+
+import android.app.ActivityManager.TaskSnapshot;
+import android.graphics.Bitmap;
+import android.graphics.Bitmap.Config;
+import android.graphics.BitmapFactory;
+import android.graphics.BitmapFactory.Options;
+import android.graphics.GraphicBuffer;
+import android.graphics.Rect;
+import android.util.Slog;
+
+import com.android.server.wm.nano.WindowManagerProtos.TaskSnapshotProto;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+
+/**
+ * Loads a persisted {@link TaskSnapshot} from disk.
+ * <p>
+ * Do not hold the window manager lock when accessing this class.
+ * <p>
+ * Test class: {@link TaskSnapshotPersisterLoaderTest}
+ */
+class TaskSnapshotLoader {
+
+    private static final String TAG = TAG_WITH_CLASS_NAME ? "TaskSnapshotLoader" : TAG_WM;
+
+    private final TaskSnapshotPersister mPersister;
+
+    TaskSnapshotLoader(TaskSnapshotPersister persister) {
+        mPersister = persister;
+    }
+
+    /**
+     * Loads a task from the disk.
+     * <p>
+     * Do not hold the window manager lock when calling this method, as we directly read data from
+     * disk here, which might be slow.
+     *
+     * @param taskId The id of the task to load.
+     * @param userId The id of the user the task belonged to.
+     * @return The loaded {@link TaskSnapshot} or {@code null} if it couldn't be loaded.
+     */
+    TaskSnapshot loadTask(int taskId, int userId) {
+        final File protoFile = mPersister.getProtoFile(taskId, userId);
+        final File bitmapFile = mPersister.getBitmapFile(taskId, userId);
+        if (!protoFile.exists() || !bitmapFile.exists()) {
+            return null;
+        }
+        try {
+            final byte[] bytes = Files.readAllBytes(protoFile.toPath());
+            final TaskSnapshotProto proto = TaskSnapshotProto.parseFrom(bytes);
+            final Options options = new Options();
+            options.inPreferredConfig = Config.HARDWARE;
+            final Bitmap bitmap = BitmapFactory.decodeFile(bitmapFile.getPath(), options);
+            if (bitmap == null) {
+                Slog.w(TAG, "Failed to load bitmap: " + bitmapFile.getPath());
+                return null;
+            }
+            final GraphicBuffer buffer = bitmap.createGraphicBufferHandle();
+            if (buffer == null) {
+                Slog.w(TAG, "Failed to retrieve gralloc buffer for bitmap: "
+                        + bitmapFile.getPath());
+                return null;
+            }
+            return new TaskSnapshot(buffer, proto.orientation,
+                    new Rect(proto.insetLeft, proto.insetTop, proto.insetRight, proto.insetBottom));
+        } catch (IOException e) {
+            Slog.w(TAG, "Unable to load task snapshot data for taskId=" + taskId);
+            return null;
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotPersister.java b/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
new file mode 100644
index 0000000..3a06c38
--- /dev/null
+++ b/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
@@ -0,0 +1,335 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+
+package com.android.server.wm;
+
+import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
+import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
+
+import android.annotation.TestApi;
+import android.app.ActivityManager.TaskSnapshot;
+import android.graphics.Bitmap;
+import android.graphics.Bitmap.CompressFormat;
+import android.os.Process;
+import android.os.SystemClock;
+import android.util.ArraySet;
+import android.util.Slog;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.os.AtomicFile;
+import com.android.server.wm.nano.WindowManagerProtos.TaskSnapshotProto;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.util.ArrayDeque;
+
+/**
+ * Persists {@link TaskSnapshot}s to disk.
+ * <p>
+ * Test class: {@link TaskSnapshotPersisterLoaderTest}
+ */
+class TaskSnapshotPersister {
+
+    private static final String TAG = TAG_WITH_CLASS_NAME ? "TaskSnapshotPersister" : TAG_WM;
+    private static final String SNAPSHOTS_DIRNAME = "snapshots";
+    private static final long DELAY_MS = 100;
+    private static final String PROTO_EXTENSION = ".proto";
+    private static final String BITMAP_EXTENSION = ".png";
+
+    @GuardedBy("mLock")
+    private final ArrayDeque<WriteQueueItem> mWriteQueue = new ArrayDeque<>();
+    @GuardedBy("mLock")
+    private boolean mQueueIdling;
+    private boolean mStarted;
+    private final Object mLock = new Object();
+    private final DirectoryResolver mDirectoryResolver;
+
+    /**
+     * The list of ids of the tasks that have been persisted since {@link #removeObsoleteFiles} was
+     * called.
+     */
+    @GuardedBy("mLock")
+    private final ArraySet<Integer> mPersistedTaskIdsSinceLastRemoveObsolete = new ArraySet<>();
+
+    TaskSnapshotPersister(DirectoryResolver resolver) {
+        mDirectoryResolver = resolver;
+    }
+
+    /**
+     * Starts persisting.
+     */
+    void start() {
+        if (!mStarted) {
+            mStarted = true;
+            mPersister.start();
+        }
+    }
+
+    /**
+     * Persists a snapshot of a task to disk.
+     *
+     * @param taskId The id of the task that needs to be persisted.
+     * @param userId The id of the user this tasks belongs to.
+     * @param snapshot The snapshot to persist.
+     */
+    void persistSnapshot(int taskId, int userId, TaskSnapshot snapshot) {
+        synchronized (mLock) {
+            mPersistedTaskIdsSinceLastRemoveObsolete.add(taskId);
+            sendToQueueLocked(new StoreWriteQueueItem(taskId, userId, snapshot));
+        }
+    }
+
+    /**
+     * Callend when a task has been removed.
+     *
+     * @param taskId The id of task that has been removed.
+     * @param userId The id of the user the task belonged to.
+     */
+    void onTaskRemovedFromRecents(int taskId, int userId) {
+        synchronized (mLock) {
+            mPersistedTaskIdsSinceLastRemoveObsolete.remove(taskId);
+            sendToQueueLocked(new DeleteWriteQueueItem(taskId, userId));
+        }
+    }
+
+    /**
+     * In case a write/delete operation was lost because the system crashed, this makes sure to
+     * clean up the directory to remove obsolete files.
+     *
+     * @param persistentTaskIds A set of task ids that exist in our in-memory model.
+     * @param runningUserIds The ids of the list of users that have tasks loaded in our in-memory
+     *                       model.
+     */
+    void removeObsoleteFiles(ArraySet<Integer> persistentTaskIds, int[] runningUserIds) {
+        synchronized (mLock) {
+            mPersistedTaskIdsSinceLastRemoveObsolete.clear();
+            sendToQueueLocked(new RemoveObsoleteFilesQueueItem(persistentTaskIds, runningUserIds));
+        }
+    }
+
+    @TestApi
+    void waitForQueueEmpty() {
+        while (true) {
+            synchronized (mLock) {
+                if (mWriteQueue.isEmpty() && mQueueIdling) {
+                    return;
+                }
+            }
+            SystemClock.sleep(100);
+        }
+    }
+
+    @GuardedBy("mLock")
+    private void sendToQueueLocked(WriteQueueItem item) {
+        mWriteQueue.offer(item);
+        mLock.notifyAll();
+    }
+
+    private File getDirectory(int userId) {
+        return new File(mDirectoryResolver.getSystemDirectoryForUser(userId), SNAPSHOTS_DIRNAME);
+    }
+
+    File getProtoFile(int taskId, int userId) {
+        return new File(getDirectory(userId), taskId + PROTO_EXTENSION);
+    }
+
+    File getBitmapFile(int taskId, int userId) {
+        return new File(getDirectory(userId), taskId + BITMAP_EXTENSION);
+    }
+
+    private boolean createDirectory(int userId) {
+        final File dir = getDirectory(userId);
+        return dir.exists() || dir.mkdirs();
+    }
+
+    private void deleteSnapshot(int taskId, int userId) {
+        final File protoFile = getProtoFile(taskId, userId);
+        final File bitmapFile = getBitmapFile(taskId, userId);
+        protoFile.delete();
+        bitmapFile.delete();
+    }
+
+    interface DirectoryResolver {
+        File getSystemDirectoryForUser(int userId);
+    }
+
+    private Thread mPersister = new Thread("TaskSnapshotPersister") {
+        public void run() {
+            android.os.Process.setThreadPriority(Process.THREAD_PRIORITY_BACKGROUND);
+            while (true) {
+                WriteQueueItem next;
+                synchronized (mLock) {
+                    next = mWriteQueue.poll();
+                }
+                if (next != null) {
+                    next.write();
+                    SystemClock.sleep(DELAY_MS);
+                }
+                synchronized (mLock) {
+                    if (!mWriteQueue.isEmpty()) {
+                        continue;
+                    }
+                    try {
+                        mQueueIdling = true;
+                        mLock.wait();
+                        mQueueIdling = false;
+                    } catch (InterruptedException e) {
+                    }
+                }
+            }
+        }
+    };
+
+    private abstract class WriteQueueItem {
+        abstract void write();
+    }
+
+    private class StoreWriteQueueItem extends WriteQueueItem {
+        private final int mTaskId;
+        private final int mUserId;
+        private final TaskSnapshot mSnapshot;
+
+        StoreWriteQueueItem(int taskId, int userId, TaskSnapshot snapshot) {
+            mTaskId = taskId;
+            mUserId = userId;
+            mSnapshot = snapshot;
+        }
+
+        @Override
+        void write() {
+            if (!createDirectory(mUserId)) {
+                Slog.e(TAG, "Unable to create snapshot directory for user dir="
+                        + getDirectory(mUserId));
+            }
+            boolean failed = false;
+            if (!writeProto()) {
+                failed = true;
+            }
+            if (!writeBuffer()) {
+                writeBuffer();
+                failed = true;
+            }
+            if (failed) {
+                deleteSnapshot(mTaskId, mUserId);
+            }
+        }
+
+        boolean writeProto() {
+            final TaskSnapshotProto proto = new TaskSnapshotProto();
+            proto.orientation = mSnapshot.getOrientation();
+            proto.insetLeft = mSnapshot.getContentInsets().left;
+            proto.insetTop = mSnapshot.getContentInsets().top;
+            proto.insetRight = mSnapshot.getContentInsets().right;
+            proto.insetBottom = mSnapshot.getContentInsets().bottom;
+            final byte[] bytes = TaskSnapshotProto.toByteArray(proto);
+            final File file = getProtoFile(mTaskId, mUserId);
+            final AtomicFile atomicFile = new AtomicFile(file);
+            FileOutputStream fos = null;
+            try {
+                fos = atomicFile.startWrite();
+                fos.write(bytes);
+                atomicFile.finishWrite(fos);
+            } catch (IOException e) {
+                atomicFile.failWrite(fos);
+                Slog.e(TAG, "Unable to open " + file + " for persisting. " + e);
+                return false;
+            }
+            return true;
+        }
+
+        boolean writeBuffer() {
+            final File file = getBitmapFile(mTaskId, mUserId);
+            final Bitmap bitmap = Bitmap.createHardwareBitmap(mSnapshot.getSnapshot());
+            try {
+                FileOutputStream fos = new FileOutputStream(file);
+                bitmap.compress(CompressFormat.PNG, 0 /* quality */, fos);
+                fos.close();
+            } catch (IOException e) {
+                Slog.e(TAG, "Unable to open " + file + " for persisting. " + e);
+                return false;
+            }
+            return true;
+        }
+    }
+
+    private class DeleteWriteQueueItem extends WriteQueueItem {
+        private final int mTaskId;
+        private final int mUserId;
+
+        DeleteWriteQueueItem(int taskId, int userId) {
+            mTaskId = taskId;
+            mUserId = userId;
+        }
+
+        @Override
+        void write() {
+            deleteSnapshot(mTaskId, mUserId);
+        }
+    }
+
+    @VisibleForTesting
+    class RemoveObsoleteFilesQueueItem extends WriteQueueItem {
+        private final ArraySet<Integer> mPersistentTaskIds;
+        private final int[] mRunningUserIds;
+
+        @VisibleForTesting
+        RemoveObsoleteFilesQueueItem(ArraySet<Integer> persistentTaskIds,
+                int[] runningUserIds) {
+            mPersistentTaskIds = persistentTaskIds;
+            mRunningUserIds = runningUserIds;
+        }
+
+        @Override
+        void write() {
+            final ArraySet<Integer> newPersistedTaskIds;
+            synchronized (mLock) {
+                newPersistedTaskIds = new ArraySet<>(mPersistedTaskIdsSinceLastRemoveObsolete);
+            }
+            for (int userId : mRunningUserIds) {
+                final File dir = getDirectory(userId);
+                final String[] files = dir.list();
+                if (files == null) {
+                    continue;
+                }
+                for (String file : files) {
+                    final int taskId = getTaskId(file);
+                    if (!mPersistentTaskIds.contains(taskId)
+                            && !newPersistedTaskIds.contains(taskId)) {
+                        new File(dir, file).delete();
+                    }
+                }
+            }
+        }
+
+        @VisibleForTesting
+        int getTaskId(String fileName) {
+            if (!fileName.endsWith(PROTO_EXTENSION) && !fileName.endsWith(BITMAP_EXTENSION)) {
+                return -1;
+            }
+            final int end = fileName.lastIndexOf('.');
+            if (end == -1) {
+                return -1;
+            }
+            try {
+                return Integer.parseInt(fileName.substring(0, end));
+            } catch (NumberFormatException e) {
+                return -1;
+            }
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index dcc0c6f..b0d22b5 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -3902,6 +3902,20 @@
     }
 
     /**
+     * In case a task write/delete operation was lost because the system crashed, this makes sure to
+     * clean up the directory to remove obsolete files.
+     *
+     * @param persistentTaskIds A set of task ids that exist in our in-memory model.
+     * @param runningUserIds The ids of the list of users that have tasks loaded in our in-memory
+     *                       model.
+     */
+    public void removeObsoleteTaskFiles(ArraySet<Integer> persistentTaskIds, int[] runningUserIds) {
+        synchronized (mWindowMap) {
+            mTaskSnapshotController.removeObsoleteTaskFiles(persistentTaskIds, runningUserIds);
+        }
+    }
+
+    /**
      * Takes a snapshot of the screen.  In landscape mode this grabs the whole screen.
      * In portrait mode, it grabs the full screenshot.
      *
@@ -5346,6 +5360,7 @@
 
     public void systemReady() {
         mPolicy.systemReady();
+        mTaskSnapshotController.systemReady();
     }
 
     // -------------------------------------------------------------
@@ -6985,6 +7000,18 @@
         }
     }
 
+    /**
+     * Called when a task has been removed from the recent tasks list.
+     * <p>
+     * Note: This doesn't go through {@link TaskWindowContainerController} yet as the window
+     * container may not exist when this happens.
+     */
+    public void notifyTaskRemovedFromRecents(int taskId, int userId) {
+        synchronized (mWindowMap) {
+            mTaskSnapshotController.notifyTaskRemovedFromRecents(taskId, userId);
+        }
+    }
+
     @Override
     public int getDockedDividerInsetsLw() {
         return getDefaultDisplayContentLocked().getDockedDividerController().getContentInsets();
diff --git a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotPersisterLoaderTest.java b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotPersisterLoaderTest.java
new file mode 100644
index 0000000..8d6d2da
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotPersisterLoaderTest.java
@@ -0,0 +1,222 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+
+package com.android.server.wm;
+
+import static android.graphics.GraphicBuffer.USAGE_HW_TEXTURE;
+import static android.graphics.GraphicBuffer.USAGE_SW_READ_RARELY;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import android.app.ActivityManager.TaskSnapshot;
+import android.content.pm.UserInfo;
+import android.content.res.Configuration;
+import android.graphics.Bitmap;
+import android.graphics.Canvas;
+import android.graphics.Color;
+import android.graphics.GraphicBuffer;
+import android.graphics.PixelFormat;
+import android.graphics.Rect;
+import android.os.SystemClock;
+import android.os.UserManager;
+import android.platform.test.annotations.Presubmit;
+import android.support.test.InstrumentationRegistry;
+import android.support.test.filters.MediumTest;
+import android.support.test.runner.AndroidJUnit4;
+import android.util.ArraySet;
+
+import com.android.internal.util.Predicate;
+import com.android.server.wm.TaskSnapshotPersister.RemoveObsoleteFilesQueueItem;
+
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+
+/**
+ * Test class for {@link TaskSnapshotPersister} and {@link TaskSnapshotLoader}
+ *
+ * runtest frameworks-services -c com.android.server.wm.TaskSnapshotPersisterLoaderTest
+ */
+@MediumTest
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+public class TaskSnapshotPersisterLoaderTest {
+
+    private static final String TEST_USER_NAME = "TaskSnapshotPersisterTest User";
+    private static final Rect TEST_INSETS = new Rect(10, 20, 30, 40);
+
+    private TaskSnapshotPersister mPersister;
+    private TaskSnapshotLoader mLoader;
+    private static int sTestUserId;
+    private static File sFilesDir;
+    private static UserManager sUserManager;
+
+    @BeforeClass
+    public static void setUpUser() {
+        sUserManager = UserManager.get(InstrumentationRegistry.getContext());
+        sTestUserId = createUser(TEST_USER_NAME, 0);
+        sFilesDir = InstrumentationRegistry.getContext().getFilesDir();
+    }
+
+    @AfterClass
+    public static void tearDownUser() {
+        removeUser(sTestUserId);
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        mPersister = new TaskSnapshotPersister(
+                userId -> sFilesDir);
+        mLoader = new TaskSnapshotLoader(mPersister);
+        mPersister.start();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        cleanDirectory();
+    }
+
+    @Test
+    public void testPersistAndLoadSnapshot() {
+        mPersister.persistSnapshot(1 , sTestUserId, createSnapshot());
+        mPersister.waitForQueueEmpty();
+        final File[] files = new File[] { new File(sFilesDir.getPath() + "/snapshots/1.proto"),
+                new File(sFilesDir.getPath() + "/snapshots/1.png") };
+        assertTrueForFiles(files, File::exists, " must exist");
+        final TaskSnapshot snapshot = mLoader.loadTask(1, sTestUserId);
+        assertNotNull(snapshot);
+        assertEquals(TEST_INSETS, snapshot.getContentInsets());
+        assertNotNull(snapshot.getSnapshot());
+        assertEquals(Configuration.ORIENTATION_PORTRAIT, snapshot.getOrientation());
+    }
+
+    private void assertTrueForFiles(File[] files, Predicate<File> predicate, String message) {
+        for (File file : files) {
+            assertTrue(file.getName() + message, predicate.apply(file));
+        }
+    }
+
+    @Test
+    public void testTaskRemovedFromRecents() {
+        mPersister.persistSnapshot(1, sTestUserId, createSnapshot());
+        mPersister.onTaskRemovedFromRecents(1, sTestUserId);
+        mPersister.waitForQueueEmpty();
+        assertFalse(new File(sFilesDir.getPath() + "/snapshots/1.proto").exists());
+        assertFalse(new File(sFilesDir.getPath() + "/snapshots/1.png").exists());
+    }
+
+    /**
+     * Tests that persisting a couple of snapshots is being throttled.
+     */
+    @Test
+    public void testThrottling() {
+        long ms = SystemClock.elapsedRealtime();
+        mPersister.persistSnapshot(1, sTestUserId, createSnapshot());
+        mPersister.persistSnapshot(2, sTestUserId, createSnapshot());
+        mPersister.persistSnapshot(3, sTestUserId, createSnapshot());
+        mPersister.persistSnapshot(4, sTestUserId, createSnapshot());
+        mPersister.persistSnapshot(5, sTestUserId, createSnapshot());
+        mPersister.persistSnapshot(6, sTestUserId, createSnapshot());
+        mPersister.waitForQueueEmpty();
+        assertTrue(SystemClock.elapsedRealtime() - ms > 500);
+    }
+
+    @Test
+    public void testGetTaskId() {
+        RemoveObsoleteFilesQueueItem removeObsoleteFilesQueueItem =
+                mPersister.new RemoveObsoleteFilesQueueItem(new ArraySet<>(), new int[] {});
+        assertEquals(-1, removeObsoleteFilesQueueItem.getTaskId("blablablulp"));
+        assertEquals(-1, removeObsoleteFilesQueueItem.getTaskId("nothing.err"));
+        assertEquals(-1, removeObsoleteFilesQueueItem.getTaskId("/invalid/"));
+        assertEquals(12, removeObsoleteFilesQueueItem.getTaskId("12.png"));
+        assertEquals(12, removeObsoleteFilesQueueItem.getTaskId("12.proto"));
+        assertEquals(1, removeObsoleteFilesQueueItem.getTaskId("1.png"));
+    }
+
+    @Test
+    public void testRemoveObsoleteFiles() {
+        mPersister.persistSnapshot(1, sTestUserId, createSnapshot());
+        mPersister.persistSnapshot(2, sTestUserId, createSnapshot());
+        final ArraySet<Integer> taskIds = new ArraySet<>();
+        taskIds.add(1);
+        mPersister.removeObsoleteFiles(taskIds, new int[] { sTestUserId });
+        mPersister.waitForQueueEmpty();
+        final File[] existsFiles = new File[] {
+                new File(sFilesDir.getPath() + "/snapshots/1.proto"),
+                new File(sFilesDir.getPath() + "/snapshots/1.png") };
+        final File[] nonExistsFiles = new File[] {
+                new File(sFilesDir.getPath() + "/snapshots/2.proto"),
+                new File(sFilesDir.getPath() + "/snapshots/2.png") };
+        assertTrueForFiles(existsFiles, File::exists, " must exist");
+        assertTrueForFiles(nonExistsFiles, file -> !file.exists(), " must not exist");
+    }
+
+    @Test
+    public void testRemoveObsoleteFiles_addedOneInTheMeantime() {
+        mPersister.persistSnapshot(1, sTestUserId, createSnapshot());
+        final ArraySet<Integer> taskIds = new ArraySet<>();
+        taskIds.add(1);
+        mPersister.removeObsoleteFiles(taskIds, new int[] { sTestUserId });
+        mPersister.persistSnapshot(2, sTestUserId, createSnapshot());
+        mPersister.waitForQueueEmpty();
+        final File[] existsFiles = new File[] {
+                new File(sFilesDir.getPath() + "/snapshots/1.proto"),
+                new File(sFilesDir.getPath() + "/snapshots/1.png"),
+                new File(sFilesDir.getPath() + "/snapshots/2.proto"),
+                new File(sFilesDir.getPath() + "/snapshots/2.png") };
+        assertTrueForFiles(existsFiles, File::exists, " must exist");
+    }
+
+    private TaskSnapshot createSnapshot() {
+        GraphicBuffer buffer = GraphicBuffer.create(100, 100, PixelFormat.RGBA_8888,
+                USAGE_HW_TEXTURE | USAGE_SW_READ_RARELY | USAGE_SW_READ_RARELY);
+        Canvas c = buffer.lockCanvas();
+        c.drawColor(Color.RED);
+        buffer.unlockCanvasAndPost(c);
+        return new TaskSnapshot(buffer, Configuration.ORIENTATION_PORTRAIT, TEST_INSETS);
+    }
+
+    private static int createUser(String name, int flags) {
+        UserInfo user = sUserManager.createUser(name, flags);
+        if (user == null) {
+            Assert.fail("Error while creating the test user: " + TEST_USER_NAME);
+        }
+        return user.id;
+    }
+
+    private static void removeUser(int userId) {
+        if (!sUserManager.removeUser(userId)) {
+            Assert.fail("Error while removing the test user: " + TEST_USER_NAME);
+        }
+    }
+
+    private void cleanDirectory() {
+        for (File file : new File(sFilesDir, "snapshots").listFiles()) {
+            if (!file.isDirectory()) {
+                file.delete();
+            }
+        }
+    }
+}