When app dies, destroy snapshot

Also destroy snapshot when we remove the AppWindowToken.

Test: runtest frameworks-services -c
com.android.server.wm.SnapshotCacheTest
Test: Open app, go home, kill app, make sure snapshots are
destroyed.

Change-Id: I532c2d7499a86164175f9fcbc8b77c6eb6bfeae6
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotCache.java b/services/core/java/com/android/server/wm/TaskSnapshotCache.java
index 994a155..c86229b 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotCache.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotCache.java
@@ -20,6 +20,8 @@
 import android.app.ActivityManager.TaskSnapshot;
 import android.util.ArrayMap;
 
+import java.io.PrintWriter;
+
 /**
  * Caches snapshots. See {@link TaskSnapshotController}.
  * <p>
@@ -27,13 +29,65 @@
  */
 class TaskSnapshotCache {
 
-    private final ArrayMap<Task, TaskSnapshot> mCache = new ArrayMap<>();
+    private final ArrayMap<AppWindowToken, Task> mAppTaskMap = new ArrayMap<>();
+    private final ArrayMap<Task, CacheEntry> mCache = new ArrayMap<>();
 
     void putSnapshot(Task task, TaskSnapshot snapshot) {
-        mCache.put(task, snapshot);
+        final CacheEntry entry = mCache.get(task);
+        if (entry != null) {
+            mAppTaskMap.remove(entry.topApp);
+        }
+        final AppWindowToken top = task.getTopChild();
+        mAppTaskMap.put(top, task);
+        mCache.put(task, new CacheEntry(snapshot, task.getTopChild()));
     }
 
     @Nullable TaskSnapshot getSnapshot(Task task) {
-        return mCache.get(task);
+        final CacheEntry entry = mCache.get(task);
+        return entry != null ? entry.snapshot : null;
+    }
+
+    /**
+     * Cleans the cache after an app window token's process died.
+     */
+    void cleanCache(AppWindowToken wtoken) {
+        final Task task = mAppTaskMap.get(wtoken);
+        if (task != null) {
+            removeEntry(task);
+        }
+    }
+
+    private void removeEntry(Task task) {
+        final CacheEntry entry = mCache.get(task);
+        if (entry != null) {
+            mAppTaskMap.remove(entry.topApp);
+            mCache.remove(task);
+        }
+    }
+
+    void dump(PrintWriter pw, String prefix) {
+        final String doublePrefix = prefix + "  ";
+        final String triplePrefix = doublePrefix + "  ";
+        pw.println(prefix + "SnapshotCache");
+        for (int i = mCache.size() - 1; i >= 0; i--) {
+            final CacheEntry entry = mCache.valueAt(i);
+            pw.println(doublePrefix + "Entry taskId=" + mCache.keyAt(i).mTaskId);
+            pw.println(triplePrefix + "topApp=" + entry.topApp);
+            pw.println(triplePrefix + "snapshot=" + entry.snapshot);
+        }
+    }
+
+    private static final class CacheEntry {
+
+        /** The snapshot. */
+        final TaskSnapshot snapshot;
+
+        /** The app token that was on top of the task when the snapshot was taken */
+        final AppWindowToken topApp;
+
+        CacheEntry(TaskSnapshot snapshot, AppWindowToken topApp) {
+            this.snapshot = snapshot;
+            this.topApp = topApp;
+        }
     }
 }