Fixing task snapshots not getting updated in recents view

Bug: 124620962
Change-Id: I7c6b9d1183faa01a0836194e42a9cd8517c8b61a

4 files changed, 113 insertions, 55 deletions

diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
index 8faf95d5b..a7bf2c3ee 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
@@ -90,6 +90,7 @@ import com.android.launcher3.util.ViewPool;
 import com.android.quickstep.OverviewCallbacks;
 import com.android.quickstep.RecentsAnimationWrapper;
 import com.android.quickstep.RecentsModel;
+import com.android.quickstep.RecentsModel.TaskThumbnailChangeListener;
 import com.android.quickstep.TaskThumbnailCache;
 import com.android.quickstep.TaskUtils;
 import com.android.quickstep.util.ClipAnimationHelper;
@@ -113,7 +114,7 @@ import java.util.function.Consumer;
 @TargetApi(Build.VERSION_CODES.P)
 public abstract class RecentsView<T extends BaseActivity> extends PagedView implements Insettable,
         TaskThumbnailCache.HighResLoadingState.HighResLoadingStateChangedCallback,
-        InvariantDeviceProfile.OnIDPChangeListener {
+        InvariantDeviceProfile.OnIDPChangeListener, TaskThumbnailChangeListener {
 
     private static final String TAG = RecentsView.class.getSimpleName();
 
@@ -171,14 +172,6 @@ public abstract class RecentsView<T extends BaseActivity> extends PagedView impl
      */
     private final TaskStackChangeListener mTaskStackListener = new TaskStackChangeListener() {
         @Override
-        public void onTaskSnapshotChanged(int taskId, ThumbnailData snapshot) {
-            if (!mHandleTaskStackChanges) {
-                return;
-            }
-            updateThumbnail(taskId, snapshot);
-        }
-
-        @Override
         public void onActivityPinned(String packageName, int userId, int taskId, int stackId) {
             if (!mHandleTaskStackChanges) {
                 return;
@@ -262,7 +255,6 @@ public abstract class RecentsView<T extends BaseActivity> extends PagedView impl
 
     private boolean mOverviewStateEnabled;
     private boolean mHandleTaskStackChanges;
-    private Runnable mNextPageSwitchRunnable;
     private boolean mSwipeDownShouldLaunchApp;
     private boolean mTouchDownToStartHome;
     private final int mTouchSlop;
@@ -340,6 +332,19 @@ public abstract class RecentsView<T extends BaseActivity> extends PagedView impl
         return mIsRtl;
     }
 
+    @Override
+    public Task onTaskThumbnailChanged(int taskId, ThumbnailData thumbnailData) {
+        if (mHandleTaskStackChanges) {
+            TaskView taskView = getTaskView(taskId);
+            if (taskView != null) {
+                Task task = taskView.getTask();
+                taskView.getThumbnail().setThumbnail(task, thumbnailData);
+                return task;
+            }
+        }
+        return null;
+    }
+
     public TaskView updateThumbnail(int taskId, ThumbnailData thumbnailData) {
         TaskView taskView = getTaskView(taskId);
         if (taskView != null) {
@@ -371,6 +376,7 @@ public abstract class RecentsView<T extends BaseActivity> extends PagedView impl
         mActivity.addMultiWindowModeChangedListener(mMultiWindowModeChangedListener);
         ActivityManagerWrapper.getInstance().registerTaskStackListener(mTaskStackListener);
         mSyncTransactionApplier = new SyncRtSurfaceTransactionApplierCompat(this);
+        RecentsModel.INSTANCE.get(getContext()).addThumbnailChangeListener(this);
         mIdp.addOnChangeListener(this);
     }
 
@@ -382,6 +388,7 @@ public abstract class RecentsView<T extends BaseActivity> extends PagedView impl
         mActivity.removeMultiWindowModeChangedListener(mMultiWindowModeChangedListener);
         ActivityManagerWrapper.getInstance().unregisterTaskStackListener(mTaskStackListener);
         mSyncTransactionApplier = null;
+        RecentsModel.INSTANCE.get(getContext()).removeThumbnailChangeListener(this);
         mIdp.removeOnChangeListener(this);
     }
 
@@ -421,17 +428,9 @@ public abstract class RecentsView<T extends BaseActivity> extends PagedView impl
         updateTaskStackListenerState();
     }
 
-    public void setNextPageSwitchRunnable(Runnable r) {
-        mNextPageSwitchRunnable = r;
-    }
-
     @Override
     protected void onPageEndTransition() {
         super.onPageEndTransition();
-        if (mNextPageSwitchRunnable != null) {
-            mNextPageSwitchRunnable.run();
-            mNextPageSwitchRunnable = null;
-        }
         if (getNextPage() > 0) {
             setSwipeDownShouldLaunchApp(true);
         }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
index 942214152..fb58c243b 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
@@ -29,7 +29,6 @@ import android.animation.ObjectAnimator;
 import android.animation.TimeInterpolator;
 import android.app.ActivityOptions;
 import android.content.Context;
-import android.content.Intent;
 import android.content.res.Resources;
 import android.graphics.Outline;
 import android.graphics.drawable.Drawable;
@@ -215,6 +214,7 @@ public class TaskView extends FrameLayout implements PageCallbacks, Reusable {
      * Updates this task view to the given {@param task}.
      */
     public void bind(Task task) {
+        cancelPendingLoadTasks();
         mTask = task;
         mSnapshotView.bind(task);
     }
@@ -305,15 +305,15 @@ public class TaskView extends FrameLayout implements PageCallbacks, Reusable {
         if (mTask == null) {
             return;
         }
+        cancelPendingLoadTasks();
         if (visible) {
             // These calls are no-ops if the data is already loaded, try and load the high
             // resolution thumbnail if the state permits
             RecentsModel model = RecentsModel.INSTANCE.get(getContext());
             TaskThumbnailCache thumbnailCache = model.getThumbnailCache();
             TaskIconCache iconCache = model.getIconCache();
-            mThumbnailLoadRequest = thumbnailCache.updateThumbnailInBackground(mTask,
-                    !thumbnailCache.getHighResLoadingState().isEnabled() /* reducedResolution */,
-                    (task) -> mSnapshotView.setThumbnail(task, task.thumbnail));
+            mThumbnailLoadRequest = thumbnailCache.updateThumbnailInBackground(
+                    mTask, thumbnail -> mSnapshotView.setThumbnail(mTask, thumbnail));
             mIconLoadRequest = iconCache.updateIconInBackground(mTask,
                     (task) -> {
                         setIcon(task.icon);
@@ -325,17 +325,22 @@ public class TaskView extends FrameLayout implements PageCallbacks, Reusable {
                                 });
                     });
         } else {
-            if (mThumbnailLoadRequest != null) {
-                mThumbnailLoadRequest.cancel();
-            }
-            if (mIconLoadRequest != null) {
-                mIconLoadRequest.cancel();
-            }
             mSnapshotView.setThumbnail(null, null);
             setIcon(null);
         }
     }
 
+    private void cancelPendingLoadTasks() {
+        if (mThumbnailLoadRequest != null) {
+            mThumbnailLoadRequest.cancel();
+            mThumbnailLoadRequest = null;
+        }
+        if (mIconLoadRequest != null) {
+            mIconLoadRequest.cancel();
+            mIconLoadRequest = null;
+        }
+    }
+
     private boolean showTaskMenu() {
         getRecentsView().snapToPage(getRecentsView().indexOfChild(this));
         mMenuView = TaskMenuView.showForTask(this);
diff --git a/quickstep/src/com/android/quickstep/RecentsModel.java b/quickstep/src/com/android/quickstep/RecentsModel.java
index f3e1545e0..81a22a190 100644
--- a/quickstep/src/com/android/quickstep/RecentsModel.java
+++ b/quickstep/src/com/android/quickstep/RecentsModel.java
@@ -15,33 +15,29 @@
  */
 package com.android.quickstep;
 
+import static com.android.quickstep.TaskUtils.checkCurrentOrManagedUserId;
+
 import android.annotation.TargetApi;
 import android.app.ActivityManager;
 import android.content.ComponentCallbacks2;
 import android.content.Context;
 import android.os.Build;
-import android.os.Bundle;
 import android.os.HandlerThread;
 import android.os.Process;
 import android.os.RemoteException;
 import android.util.Log;
-import android.util.SparseArray;
 
-import com.android.launcher3.MainThreadExecutor;
 import com.android.launcher3.util.MainThreadInitializedObject;
-import com.android.launcher3.util.Preconditions;
 import com.android.systemui.shared.recents.ISystemUiProxy;
 import com.android.systemui.shared.recents.model.Task;
+import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.TaskStackChangeListener;
 
 import java.util.ArrayList;
+import java.util.List;
 import java.util.function.Consumer;
 
-import androidx.annotation.WorkerThread;
-
-import static com.android.quickstep.TaskUtils.checkCurrentOrManagedUserId;
-
 /**
  * Singleton class to load and manage recents model.
  */
@@ -54,8 +50,8 @@ public class RecentsModel extends TaskStackChangeListener {
     public static final MainThreadInitializedObject<RecentsModel> INSTANCE =
             new MainThreadInitializedObject<>(c -> new RecentsModel(c));
 
+    private final List<TaskThumbnailChangeListener> mThumbnailChangeListeners = new ArrayList<>();
     private final Context mContext;
-    private final MainThreadExecutor mMainThreadExecutor;
 
     private ISystemUiProxy mSystemUiProxy;
 
@@ -68,9 +64,6 @@ public class RecentsModel extends TaskStackChangeListener {
 
     private RecentsModel(Context context) {
         mContext = context;
-
-        mMainThreadExecutor = new MainThreadExecutor();
-
         HandlerThread loaderThread = new HandlerThread("TaskThumbnailIconCache",
                 Process.THREAD_PRIORITY_BACKGROUND);
         loaderThread.start();
@@ -168,6 +161,18 @@ public class RecentsModel extends TaskStackChangeListener {
         });
     }
 
+    @Override
+    public void onTaskSnapshotChanged(int taskId, ThumbnailData snapshot) {
+        mThumbnailCache.updateTaskSnapShot(taskId, snapshot);
+
+        for (int i = mThumbnailChangeListeners.size() - 1; i >= 0; i--) {
+            Task task = mThumbnailChangeListeners.get(i).onTaskThumbnailChanged(taskId, snapshot);
+            if (task != null) {
+                task.thumbnail = snapshot;
+            }
+        }
+    }
+
     public void setSystemUiProxy(ISystemUiProxy systemUiProxy) {
         mSystemUiProxy = systemUiProxy;
     }
@@ -239,4 +244,17 @@ public class RecentsModel extends TaskStackChangeListener {
                             + ": ", e);
         }
     }
+
+    public void addThumbnailChangeListener(TaskThumbnailChangeListener listener) {
+        mThumbnailChangeListeners.add(listener);
+    }
+
+    public void removeThumbnailChangeListener(TaskThumbnailChangeListener listener) {
+        mThumbnailChangeListeners.remove(listener);
+    }
+
+    public interface TaskThumbnailChangeListener {
+
+        Task onTaskThumbnailChanged(int taskId, ThumbnailData thumbnailData);
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/TaskThumbnailCache.java b/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
index 7a216ed2f..d05196bc8 100644
--- a/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
+++ b/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
@@ -26,6 +26,7 @@ import com.android.launcher3.Utilities;
 import com.android.launcher3.icons.cache.HandlerRunnable;
 import com.android.launcher3.util.Preconditions;
 import com.android.systemui.shared.recents.model.Task;
+import com.android.systemui.shared.recents.model.Task.TaskKey;
 import com.android.systemui.shared.recents.model.TaskKeyLruCache;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
@@ -38,7 +39,7 @@ public class TaskThumbnailCache {
     private final MainThreadExecutor mMainThreadExecutor;
 
     private final int mCacheSize;
-    private final TaskKeyLruCache<ThumbnailData> mCache;
+    private final ThumbnailCache mCache;
     private final HighResLoadingState mHighResLoadingState;
 
     public static class HighResLoadingState {
@@ -98,7 +99,7 @@ public class TaskThumbnailCache {
 
         Resources res = context.getResources();
         mCacheSize = res.getInteger(R.integer.recentsThumbnailCacheSize);
-        mCache = new TaskKeyLruCache<>(mCacheSize);
+        mCache = new ThumbnailCache(mCacheSize);
     }
 
     /**
@@ -106,13 +107,20 @@ public class TaskThumbnailCache {
      */
     public void updateThumbnailInCache(Task task) {
         Preconditions.assertUIThread();
-
         // Fetch the thumbnail for this task and put it in the cache
-        updateThumbnailInBackground(task, true /* reducedResolution */, (t) -> {
-            mCache.put(task.key, t.thumbnail);
-        });
+        if (task.thumbnail == null) {
+            updateThumbnailInBackground(task.key, true /* reducedResolution */,
+                    t -> task.thumbnail = t);
+        }
     }
 
+    /**
+     * Synchronously updates the thumbnail in the cache if it is already there.
+     */
+    public void updateTaskSnapShot(int taskId, ThumbnailData thumbnail) {
+        Preconditions.assertUIThread();
+        mCache.updateIfAlreadyInCache(taskId, thumbnail);
+    }
 
     /**
      * Asynchronously fetches the icon and other task data for the given {@param task}.
@@ -120,22 +128,33 @@ public class TaskThumbnailCache {
      * @param callback The callback to receive the task after its data has been populated.
      * @return A cancelable handle to the request
      */
-    public ThumbnailLoadRequest updateThumbnailInBackground(Task task, boolean reducedResolution,
-            Consumer<Task> callback) {
+    public ThumbnailLoadRequest updateThumbnailInBackground(
+            Task task, Consumer<ThumbnailData> callback) {
         Preconditions.assertUIThread();
 
+        boolean reducedResolution = !mHighResLoadingState.isEnabled();
         if (task.thumbnail != null && (!task.thumbnail.reducedResolution || reducedResolution)) {
             // Nothing to load, the thumbnail is already high-resolution or matches what the
             // request, so just callback
-            callback.accept(task);
+            callback.accept(task.thumbnail);
             return null;
         }
 
-        ThumbnailData cachedThumbnail = mCache.getAndInvalidateIfModified(task.key);
+
+        return updateThumbnailInBackground(task.key, !mHighResLoadingState.isEnabled(), t -> {
+            task.thumbnail = t;
+            callback.accept(t);
+        });
+    }
+
+    private ThumbnailLoadRequest updateThumbnailInBackground(TaskKey key, boolean reducedResolution,
+            Consumer<ThumbnailData> callback) {
+        Preconditions.assertUIThread();
+
+        ThumbnailData cachedThumbnail = mCache.getAndInvalidateIfModified(key);
         if (cachedThumbnail != null && (!cachedThumbnail.reducedResolution || reducedResolution)) {
             // Already cached, lets use that thumbnail
-            task.thumbnail = cachedThumbnail;
-            callback.accept(task);
+            callback.accept(cachedThumbnail);
             return null;
         }
 
@@ -144,14 +163,14 @@ public class TaskThumbnailCache {
             @Override
             public void run() {
                 ThumbnailData thumbnail = ActivityManagerWrapper.getInstance().getTaskThumbnail(
-                        task.key.id, reducedResolution);
+                        key.id, reducedResolution);
                 if (isCanceled()) {
                     // We don't call back to the provided callback in this case
                     return;
                 }
                 mMainThreadExecutor.execute(() -> {
-                    task.thumbnail = thumbnail;
-                    callback.accept(task);
+                    mCache.put(key, thumbnail);
+                    callback.accept(thumbnail);
                     onEnd();
                 });
             }
@@ -196,4 +215,21 @@ public class TaskThumbnailCache {
             this.reducedResolution = reducedResolution;
         }
     }
+
+    private static class ThumbnailCache extends TaskKeyLruCache<ThumbnailData> {
+
+        public ThumbnailCache(int cacheSize) {
+            super(cacheSize);
+        }
+
+        /**
+         * Updates the cache entry if it is already present in the cache
+         */
+        public void updateIfAlreadyInCache(int taskId, ThumbnailData thumbnailData) {
+            ThumbnailData oldData = getCacheEntry(taskId);
+            if (oldData != null) {
+                putCacheEntry(taskId, thumbnailData);
+            }
+        }
+    }
 }