Converting ModelTask to an interface instead of an abstract class

This allows extensibility and better use of lambdas

Bug: 338282246
Test: Presubmit
Flag: None

Change-Id: Ia41067f0068b3b631eeb4faf877dc77f8587e1f6
diff --git a/quickstep/src/com/android/launcher3/model/PredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/PredictionUpdateTask.java
index 2fcbe4e..d604742 100644
--- a/quickstep/src/com/android/launcher3/model/PredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/PredictionUpdateTask.java
@@ -15,8 +15,8 @@
  */
 package com.android.launcher3.model;
 
-import static com.android.launcher3.LauncherPrefs.nonRestorableItem;
 import static com.android.launcher3.EncryptionType.ENCRYPTED;
+import static com.android.launcher3.LauncherPrefs.nonRestorableItem;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.quickstep.InstantAppResolverImpl.COMPONENT_CLASS_MARKER;
 
@@ -32,6 +32,7 @@
 
 import com.android.launcher3.ConstantItem;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
@@ -47,7 +48,7 @@
 /**
  * Task to update model as a result of predicted apps update
  */
-public class PredictionUpdateTask extends BaseModelUpdateTask {
+public class PredictionUpdateTask implements ModelUpdateTask {
 
     public static final ConstantItem<Boolean> LAST_PREDICTION_ENABLED =
             nonRestorableItem("last_prediction_enabled_state", true, ENCRYPTED);
@@ -61,8 +62,9 @@
     }
 
     @Override
-    public void execute(@NonNull final LauncherAppState app, @NonNull final BgDataModel dataModel,
-            @NonNull final AllAppsList apps) {
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList apps) {
+        LauncherAppState app = taskController.getApp();
         Context context = app.getContext();
 
         // TODO: remove this
@@ -119,7 +121,7 @@
 
         FixedContainerItems fci = new FixedContainerItems(mPredictorState.containerId, items);
         dataModel.extraItems.put(fci.containerId, fci);
-        bindExtraContainerItems(fci);
+        taskController.bindExtraContainerItems(fci);
         usersForChangedShortcuts.forEach(
                 u -> dataModel.updateShortcutPinnedState(app.getContext(), u));
 
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index f4cbf17..39f2c00 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -24,7 +24,7 @@
 
 import androidx.annotation.NonNull;
 
-import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
 import com.android.launcher3.model.data.ItemInfo;
@@ -41,7 +41,7 @@
 import java.util.stream.Collectors;
 
 /** Task to update model as a result of predicted widgets update */
-public final class WidgetsPredictionUpdateTask extends BaseModelUpdateTask {
+public final class WidgetsPredictionUpdateTask implements ModelUpdateTask {
     private final PredictorState mPredictorState;
     private final List<AppTarget> mTargets;
 
@@ -58,8 +58,8 @@
      * workspace.
      */
     @Override
-    public void execute(@NonNull final LauncherAppState appState,
-            @NonNull final BgDataModel dataModel, @NonNull final AllAppsList apps) {
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList apps) {
         Set<ComponentKey> widgetsInWorkspace = dataModel.appWidgets.stream().map(
                 widget -> new ComponentKey(widget.providerName, widget.user)).collect(
                 Collectors.toSet());
@@ -98,7 +98,7 @@
 
         List<ItemInfo> items;
         if (enableCategorizedWidgetSuggestions()) {
-            Context context = appState.getContext();
+            Context context = taskController.getApp().getContext();
             WidgetRecommendationCategoryProvider categoryProvider =
                     WidgetRecommendationCategoryProvider.newInstance(context);
             items = servicePredictedItems.stream()
@@ -115,7 +115,7 @@
                 new FixedContainerItems(mPredictorState.containerId, items);
 
         dataModel.extraItems.put(mPredictorState.containerId, fixedContainerItems);
-        bindExtraContainerItems(fixedContainerItems);
+        taskController.bindExtraContainerItems(fixedContainerItems);
 
         // Don't store widgets prediction to disk because it is not used frequently.
     }
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index ed633df..e3e14ae 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -58,10 +58,6 @@
 import com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers;
 import com.android.launcher3.logging.InstanceId;
 import com.android.launcher3.logging.StatsLogManager;
-import com.android.launcher3.model.AllAppsList;
-import com.android.launcher3.model.BaseModelUpdateTask;
-import com.android.launcher3.model.BgDataModel;
-import com.android.launcher3.model.data.CollectionInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.LogConfig;
@@ -371,17 +367,9 @@
             if (mItemInfo.container < 0 || !LauncherAppState.INSTANCE.executeIfCreated(app -> {
                 // Item is inside a collection, fetch collection info in a BG thread
                 // and then write to StatsLog.
-                app.getModel().enqueueModelUpdateTask(
-                        new BaseModelUpdateTask() {
-                            @Override
-                            public void execute(@NonNull final LauncherAppState app,
-                                    @NonNull final BgDataModel dataModel,
-                                    @NonNull final AllAppsList apps) {
-                                CollectionInfo collectionInfo =
-                                        dataModel.collections.get(mItemInfo.container);
-                                write(event, applyOverwrites(mItemInfo.buildProto(collectionInfo)));
-                            }
-                        });
+                app.getModel().enqueueModelUpdateTask((taskController, dataModel, apps) ->
+                        write(event, applyOverwrites(mItemInfo.buildProto(
+                                dataModel.collections.get(mItemInfo.container)))));
             })) {
                 // Write log on the model thread so that logs do not go out of order
                 // (for eg: drop comes after drag)
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 9b0e0ec..be98589 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -48,7 +48,6 @@
 import com.android.launcher3.model.AddWorkspaceItemsTask;
 import com.android.launcher3.model.AllAppsList;
 import com.android.launcher3.model.BaseLauncherBinder;
-import com.android.launcher3.model.BaseModelUpdateTask;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.CacheDataUpdatedTask;
@@ -57,6 +56,7 @@
 import com.android.launcher3.model.ModelDbController;
 import com.android.launcher3.model.ModelDelegate;
 import com.android.launcher3.model.ModelLauncherCallbacks;
+import com.android.launcher3.model.ModelTaskController;
 import com.android.launcher3.model.ModelWriter;
 import com.android.launcher3.model.PackageInstallStateChangedTask;
 import com.android.launcher3.model.PackageUpdatedTask;
@@ -82,7 +82,6 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.concurrent.CancellationException;
-import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
@@ -426,13 +425,9 @@
     @Override
     public void onInstallSessionCreated(@NonNull final PackageInstallInfo sessionInfo) {
         if (FeatureFlags.PROMISE_APPS_IN_ALL_APPS.get()) {
-            enqueueModelUpdateTask(new BaseModelUpdateTask() {
-                @Override
-                public void execute(@NonNull final LauncherAppState app,
-                        @NonNull final BgDataModel dataModel, @NonNull final AllAppsList apps) {
-                    apps.addPromiseApp(app.getContext(), sessionInfo);
-                    bindApplicationsIfNeeded();
-                }
+            enqueueModelUpdateTask((taskController, dataModel, apps) -> {
+                apps.addPromiseApp(mApp.getContext(), sessionInfo);
+                taskController.bindApplicationsIfNeeded();
             });
         }
     }
@@ -440,60 +435,56 @@
     @Override
     public void onSessionFailure(@NonNull final String packageName,
             @NonNull final UserHandle user) {
-        enqueueModelUpdateTask(new BaseModelUpdateTask() {
-            @Override
-            public void execute(@NonNull final LauncherAppState app,
-                    @NonNull final BgDataModel dataModel, @NonNull final AllAppsList apps) {
-                IconCache iconCache = app.getIconCache();
-                final IntSet removedIds = new IntSet();
-                HashSet<WorkspaceItemInfo> archivedWorkspaceItemsToCacheRefresh = new HashSet<>();
-                boolean isAppArchived = PackageManagerHelper.INSTANCE.get(mApp.getContext())
-                        .isAppArchivedForUser(packageName, user);
-                synchronized (dataModel) {
-                    if (isAppArchived) {
-                        // Remove package icon cache entry for archived app in case of a session
-                        // failure.
-                        mApp.getIconCache().remove(
-                                new ComponentName(packageName, packageName + EMPTY_CLASS_NAME),
-                                user);
-                    }
+        enqueueModelUpdateTask((taskController, dataModel, apps) -> {
+            IconCache iconCache = mApp.getIconCache();
+            final IntSet removedIds = new IntSet();
+            HashSet<WorkspaceItemInfo> archivedWorkspaceItemsToCacheRefresh = new HashSet<>();
+            boolean isAppArchived = PackageManagerHelper.INSTANCE.get(mApp.getContext())
+                    .isAppArchivedForUser(packageName, user);
+            synchronized (dataModel) {
+                if (isAppArchived) {
+                    // Remove package icon cache entry for archived app in case of a session
+                    // failure.
+                    mApp.getIconCache().remove(
+                            new ComponentName(packageName, packageName + EMPTY_CLASS_NAME),
+                            user);
+                }
 
-                    for (ItemInfo info : dataModel.itemsIdMap) {
-                        if (info instanceof WorkspaceItemInfo
-                                && ((WorkspaceItemInfo) info).hasPromiseIconUi()
-                                && user.equals(info.user)
-                                && info.getIntent() != null) {
-                            if (TextUtils.equals(packageName, info.getIntent().getPackage())) {
-                                removedIds.add(info.id);
-                            }
-                            if (((WorkspaceItemInfo) info).isArchived()) {
-                                WorkspaceItemInfo workspaceItem = (WorkspaceItemInfo) info;
-                                // Refresh icons on the workspace for archived apps.
-                                iconCache.getTitleAndIcon(workspaceItem,
-                                        workspaceItem.usingLowResIcon());
-                                archivedWorkspaceItemsToCacheRefresh.add(workspaceItem);
-                            }
+                for (ItemInfo info : dataModel.itemsIdMap) {
+                    if (info instanceof WorkspaceItemInfo
+                            && ((WorkspaceItemInfo) info).hasPromiseIconUi()
+                            && user.equals(info.user)
+                            && info.getIntent() != null) {
+                        if (TextUtils.equals(packageName, info.getIntent().getPackage())) {
+                            removedIds.add(info.id);
+                        }
+                        if (((WorkspaceItemInfo) info).isArchived()) {
+                            WorkspaceItemInfo workspaceItem = (WorkspaceItemInfo) info;
+                            // Refresh icons on the workspace for archived apps.
+                            iconCache.getTitleAndIcon(workspaceItem,
+                                    workspaceItem.usingLowResIcon());
+                            archivedWorkspaceItemsToCacheRefresh.add(workspaceItem);
                         }
                     }
-
-                    if (isAppArchived) {
-                        apps.updateIconsAndLabels(new HashSet<>(List.of(packageName)), user);
-                    }
                 }
 
-                if (!removedIds.isEmpty()) {
-                    deleteAndBindComponentsRemoved(
-                            ItemInfoMatcher.ofItemIds(removedIds),
-                            "removed because install session failed");
-                }
-                if (!archivedWorkspaceItemsToCacheRefresh.isEmpty()) {
-                    bindUpdatedWorkspaceItems(
-                            archivedWorkspaceItemsToCacheRefresh.stream().toList());
-                }
                 if (isAppArchived) {
-                    bindApplicationsIfNeeded();
+                    apps.updateIconsAndLabels(new HashSet<>(List.of(packageName)), user);
                 }
             }
+
+            if (!removedIds.isEmpty()) {
+                taskController.deleteAndBindComponentsRemoved(
+                        ItemInfoMatcher.ofItemIds(removedIds),
+                        "removed because install session failed");
+            }
+            if (!archivedWorkspaceItemsToCacheRefresh.isEmpty()) {
+                taskController.bindUpdatedWorkspaceItems(
+                        archivedWorkspaceItemsToCacheRefresh.stream().toList());
+            }
+            if (isAppArchived) {
+                taskController.bindApplicationsIfNeeded();
+            }
         });
     }
 
@@ -583,13 +574,9 @@
      */
     public void onWidgetLabelsUpdated(@NonNull final HashSet<String> updatedPackages,
             @NonNull final UserHandle user) {
-        enqueueModelUpdateTask(new BaseModelUpdateTask() {
-            @Override
-            public void execute(@NonNull final LauncherAppState app,
-                    @NonNull final BgDataModel dataModel, @NonNull final AllAppsList apps) {
-                dataModel.widgetsModel.onPackageIconsUpdated(updatedPackages, user, app);
-                bindUpdatedWidgets(dataModel);
-            }
+        enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
+            dataModel.widgetsModel.onPackageIconsUpdated(updatedPackages, user, mApp);
+            taskController.bindUpdatedWidgets(dataModel);
         });
     }
 
@@ -597,8 +584,15 @@
         if (mModelDestroyed) {
             return;
         }
-        task.init(mApp, this, mBgDataModel, mBgAllAppsList, MAIN_EXECUTOR);
-        MODEL_EXECUTOR.execute(task);
+        MODEL_EXECUTOR.execute(() -> {
+            if (!isModelLoaded()) {
+                // Loader has not yet run.
+                return;
+            }
+            ModelTaskController controller = new ModelTaskController(
+                    mApp, mBgDataModel, mBgAllAppsList, this, MAIN_EXECUTOR);
+            task.execute(controller, mBgDataModel, mBgAllAppsList);
+        });
     }
 
     /**
@@ -610,18 +604,10 @@
         void execute(@NonNull Callbacks callbacks);
     }
 
-    /**
-     * A runnable which changes/updates the data model of the launcher based on certain events.
-     */
-    public interface ModelUpdateTask extends Runnable {
+    public interface ModelUpdateTask {
 
-        /**
-         * Called before the task is posted to initialize the internal state.
-         */
-        void init(@NonNull LauncherAppState app, @NonNull LauncherModel model,
-                @NonNull BgDataModel dataModel, @NonNull AllAppsList allAppsList,
-                @NonNull Executor uiExecutor);
-
+        void execute(@NonNull ModelTaskController taskController,
+                @NonNull BgDataModel dataModel, @NonNull AllAppsList apps);
     }
 
     public void updateAndBindWorkspaceItem(@NonNull final WorkspaceItemInfo si,
@@ -638,27 +624,19 @@
      */
     public void updateAndBindWorkspaceItem(
             @NonNull final Supplier<WorkspaceItemInfo> itemProvider) {
-        enqueueModelUpdateTask(new BaseModelUpdateTask() {
-            @Override
-            public void execute(@NonNull final LauncherAppState app,
-                    @NonNull final BgDataModel dataModel, @NonNull final AllAppsList apps) {
-                WorkspaceItemInfo info = itemProvider.get();
-                getModelWriter().updateItemInDatabase(info);
-                ArrayList<WorkspaceItemInfo> update = new ArrayList<>();
-                update.add(info);
-                bindUpdatedWorkspaceItems(update);
-            }
+        enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
+            WorkspaceItemInfo info = itemProvider.get();
+            taskController.getModelWriter().updateItemInDatabase(info);
+            ArrayList<WorkspaceItemInfo> update = new ArrayList<>();
+            update.add(info);
+            taskController.bindUpdatedWorkspaceItems(update);
         });
     }
 
     public void refreshAndBindWidgetsAndShortcuts(@Nullable final PackageUserKey packageUser) {
-        enqueueModelUpdateTask(new BaseModelUpdateTask() {
-            @Override
-            public void execute(@NonNull final LauncherAppState app,
-                    @NonNull final BgDataModel dataModel, @NonNull final AllAppsList apps) {
-                dataModel.widgetsModel.update(app, packageUser);
-                bindUpdatedWidgets(dataModel);
-            }
+        enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
+            dataModel.widgetsModel.update(taskController.getApp(), packageUser);
+            taskController.bindUpdatedWidgets(dataModel);
         });
     }
 
diff --git a/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java b/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
index ab47097..8121d2a 100644
--- a/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
+++ b/src/com/android/launcher3/allapps/search/DefaultAppSearchAlgorithm.java
@@ -22,13 +22,9 @@
 import android.os.Handler;
 
 import androidx.annotation.AnyThread;
-import androidx.annotation.NonNull;
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.allapps.BaseAllAppsAdapter.AdapterItem;
-import com.android.launcher3.model.AllAppsList;
-import com.android.launcher3.model.BaseModelUpdateTask;
-import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.search.SearchAlgorithm;
 import com.android.launcher3.search.SearchCallback;
@@ -67,16 +63,12 @@
 
     @Override
     public void doSearch(String query, SearchCallback<AdapterItem> callback) {
-        mAppState.getModel().enqueueModelUpdateTask(new BaseModelUpdateTask() {
-            @Override
-            public void execute(@NonNull final LauncherAppState app,
-                    @NonNull final BgDataModel dataModel, @NonNull final AllAppsList apps) {
-                ArrayList<AdapterItem> result = getTitleMatchResult(apps.data, query);
-                if (mAddNoResultsMessage && result.isEmpty()) {
-                    result.add(getEmptyMessageAdapterItem(query));
-                }
-                mResultHandler.post(() -> callback.onSearchResult(query, result));
+        mAppState.getModel().enqueueModelUpdateTask((taskController, dataModel, apps) ->  {
+            ArrayList<AdapterItem> result = getTitleMatchResult(apps.data, query);
+            if (mAddNoResultsMessage && result.isEmpty()) {
+                result.add(getEmptyMessageAdapterItem(query));
             }
+            mResultHandler.post(() -> callback.onSearchResult(query, result));
         });
     }
 
diff --git a/src/com/android/launcher3/folder/FolderNameProvider.java b/src/com/android/launcher3/folder/FolderNameProvider.java
index 5d2bb3a..be5f8f7 100644
--- a/src/com/android/launcher3/folder/FolderNameProvider.java
+++ b/src/com/android/launcher3/folder/FolderNameProvider.java
@@ -28,11 +28,12 @@
 import androidx.annotation.WorkerThread;
 
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.model.AllAppsList;
-import com.android.launcher3.model.BaseModelUpdateTask;
 import com.android.launcher3.model.BgDataModel;
+import com.android.launcher3.model.ModelTaskController;
 import com.android.launcher3.model.StringCache;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.CollectionInfo;
@@ -191,10 +192,11 @@
         nameInfos.setLabel(labels.length - 1, label, 1.0f);
     }
 
-    private class FolderNameWorker extends BaseModelUpdateTask {
+    private class FolderNameWorker implements ModelUpdateTask {
+
         @Override
-        public void execute(@NonNull final LauncherAppState app,
-                @NonNull final BgDataModel dataModel, @NonNull final AllAppsList apps) {
+        public void execute(@NonNull ModelTaskController taskController,
+                @NonNull BgDataModel dataModel, @NonNull AllAppsList apps) {
             mCollectionInfos = dataModel.collections.clone();
             mAppInfos = Arrays.asList(apps.copyData());
         }
diff --git a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
index 3fa6da4..427fb97 100644
--- a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
+++ b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
@@ -26,9 +26,10 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel.CallbackTask;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.data.AppInfo;
@@ -50,7 +51,7 @@
 /**
  * Task to add auto-created workspace items.
  */
-public class AddWorkspaceItemsTask extends BaseModelUpdateTask {
+public class AddWorkspaceItemsTask implements ModelUpdateTask {
 
     private static final String LOG = "AddWorkspaceItemsTask";
 
@@ -77,16 +78,17 @@
         mItemSpaceFinder = itemSpaceFinder;
     }
 
+
     @Override
-    public void execute(@NonNull final LauncherAppState app, @NonNull final BgDataModel dataModel,
-            @NonNull final AllAppsList apps) {
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList apps) {
         if (mItemList.isEmpty()) {
             return;
         }
 
         final ArrayList<ItemInfo> addedItemsFinal = new ArrayList<>();
         final IntArray addedWorkspaceScreensFinal = new IntArray();
-        final Context context = app.getContext();
+        final Context context = taskController.getApp().getContext();
 
         synchronized (dataModel) {
             IntArray workspaceScreens = dataModel.collectWorkspaceScreens();
@@ -128,8 +130,8 @@
 
             for (ItemInfo item : filteredItems) {
                 // Find appropriate space for the item.
-                int[] coords = mItemSpaceFinder.findSpaceForItem(app, dataModel, workspaceScreens,
-                        addedWorkspaceScreensFinal, item.spanX, item.spanY);
+                int[] coords = mItemSpaceFinder.findSpaceForItem(taskController.getApp(), dataModel,
+                        workspaceScreens, addedWorkspaceScreensFinal, item.spanX, item.spanY);
                 int screenId = coords[0];
 
                 ItemInfo itemInfo;
@@ -176,8 +178,8 @@
                     if (hasActivity) {
                         // App was installed while launcher was in the background,
                         // or app was already installed for another user.
-                        itemInfo = new AppInfo(app.getContext(), activities.get(0), item.user)
-                                .makeWorkspaceItem(app.getContext());
+                        itemInfo = new AppInfo(context, activities.get(0), item.user)
+                                .makeWorkspaceItem(context);
 
                         if (shortcutExists(dataModel, itemInfo.getIntent(), itemInfo.user)) {
                             // We need this additional check here since we treat all auto added
@@ -187,16 +189,17 @@
                             continue;
                         }
 
+                        IconCache cache = taskController.getApp().getIconCache();
                         WorkspaceItemInfo wii = (WorkspaceItemInfo) itemInfo;
                         wii.title = "";
-                        wii.bitmap = app.getIconCache().getDefaultIcon(item.user);
-                        app.getIconCache().getTitleAndIcon(wii,
+                        wii.bitmap = cache.getDefaultIcon(item.user);
+                        cache.getTitleAndIcon(wii,
                                 ((WorkspaceItemInfo) itemInfo).usingLowResIcon());
                     }
                 }
 
                 // Add the shortcut to the db
-                getModelWriter().addItemToDatabase(itemInfo,
+                taskController.getModelWriter().addItemToDatabase(itemInfo,
                         LauncherSettings.Favorites.CONTAINER_DESKTOP, screenId,
                         coords[1], coords[2]);
 
@@ -209,7 +212,7 @@
         }
 
         if (!addedItemsFinal.isEmpty()) {
-            scheduleCallbackTask(new CallbackTask() {
+            taskController.scheduleCallbackTask(new CallbackTask() {
                 @Override
                 public void execute(@NonNull Callbacks callbacks) {
                     final ArrayList<ItemInfo> addAnimated = new ArrayList<>();
diff --git a/src/com/android/launcher3/model/BaseModelUpdateTask.java b/src/com/android/launcher3/model/BaseModelUpdateTask.java
deleted file mode 100644
index 529c30a..0000000
--- a/src/com/android/launcher3/model/BaseModelUpdateTask.java
+++ /dev/null
@@ -1,160 +0,0 @@
-/*
- * Copyright (C) 2016 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.launcher3.model;
-
-import android.util.Log;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherModel;
-import com.android.launcher3.LauncherModel.CallbackTask;
-import com.android.launcher3.LauncherModel.ModelUpdateTask;
-import com.android.launcher3.celllayout.CellPosMapper;
-import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.BgDataModel.FixedContainerItems;
-import com.android.launcher3.model.data.AppInfo;
-import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.util.ComponentKey;
-import com.android.launcher3.util.PackageUserKey;
-import com.android.launcher3.widget.model.WidgetsListBaseEntry;
-
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.concurrent.Executor;
-import java.util.function.Predicate;
-import java.util.stream.Collectors;
-
-/**
- * Extension of {@link ModelUpdateTask} with some utility methods
- */
-public abstract class BaseModelUpdateTask implements ModelUpdateTask {
-
-    private static final boolean DEBUG_TASKS = false;
-    private static final String TAG = "BaseModelUpdateTask";
-
-    // Nullabilities are explicitly omitted here because these are late-init fields,
-    // They will be non-null after init(), which is always the case in enqueueModelUpdateTask().
-    private LauncherAppState mApp;
-    private LauncherModel mModel;
-    private BgDataModel mDataModel;
-    private AllAppsList mAllAppsList;
-    private Executor mUiExecutor;
-
-    public void init(@NonNull final LauncherAppState app, @NonNull final LauncherModel model,
-            @NonNull final BgDataModel dataModel, @NonNull final AllAppsList allAppsList,
-            @NonNull final Executor uiExecutor) {
-        mApp = app;
-        mModel = model;
-        mDataModel = dataModel;
-        mAllAppsList = allAppsList;
-        mUiExecutor = uiExecutor;
-    }
-
-    @Override
-    public final void run() {
-        boolean isModelLoaded = Objects.requireNonNull(mModel).isModelLoaded();
-        if (!isModelLoaded) {
-            if (DEBUG_TASKS) {
-                Log.d(TAG, "Ignoring model task since loader is pending=" + this);
-            }
-            // Loader has not yet run.
-            return;
-        }
-        execute(mApp, mDataModel, mAllAppsList);
-    }
-
-    /**
-     * Execute the actual task. Called on the worker thread.
-     */
-    public abstract void execute(@NonNull LauncherAppState app,
-            @NonNull BgDataModel dataModel, @NonNull AllAppsList apps);
-
-    /**
-     * Schedules a {@param task} to be executed on the current callbacks.
-     */
-    public final void scheduleCallbackTask(@NonNull final CallbackTask task) {
-        for (final Callbacks cb : mModel.getCallbacks()) {
-            mUiExecutor.execute(() -> task.execute(cb));
-        }
-    }
-
-    public ModelWriter getModelWriter() {
-        // Updates from model task, do not deal with icon position in hotseat. Also no need to
-        // verify changes as the ModelTasks always push the changes to callbacks
-        return mModel.getWriter(false /* verifyChanges */, CellPosMapper.DEFAULT, null);
-    }
-
-    public void bindUpdatedWorkspaceItems(@NonNull final List<WorkspaceItemInfo> allUpdates) {
-        // Bind workspace items
-        List<WorkspaceItemInfo> workspaceUpdates = allUpdates.stream()
-                .filter(info -> info.id != ItemInfo.NO_ID)
-                .collect(Collectors.toList());
-        if (!workspaceUpdates.isEmpty()) {
-            scheduleCallbackTask(c -> c.bindWorkspaceItemsChanged(workspaceUpdates));
-        }
-
-        // Bind extra items if any
-        allUpdates.stream()
-                .mapToInt(info -> info.container)
-                .distinct()
-                .mapToObj(mDataModel.extraItems::get)
-                .filter(Objects::nonNull)
-                .forEach(this::bindExtraContainerItems);
-    }
-
-    public void bindExtraContainerItems(@NonNull final FixedContainerItems item) {
-        scheduleCallbackTask(c -> c.bindExtraContainerItems(item));
-    }
-
-    public void bindDeepShortcuts(@NonNull final BgDataModel dataModel) {
-        final HashMap<ComponentKey, Integer> shortcutMapCopy =
-                new HashMap<>(dataModel.deepShortcutMap);
-        scheduleCallbackTask(callbacks -> callbacks.bindDeepShortcutMap(shortcutMapCopy));
-    }
-
-    public void bindUpdatedWidgets(@NonNull final BgDataModel dataModel) {
-        final ArrayList<WidgetsListBaseEntry> widgets =
-                dataModel.widgetsModel.getWidgetsListForPicker(mApp.getContext());
-        scheduleCallbackTask(c -> c.bindAllWidgets(widgets));
-    }
-
-    public void deleteAndBindComponentsRemoved(final Predicate<ItemInfo> matcher,
-            @Nullable final String reason) {
-        getModelWriter().deleteItemsFromDatabase(matcher, reason);
-
-        // Call the components-removed callback
-        scheduleCallbackTask(c -> c.bindWorkspaceComponentsRemoved(matcher));
-    }
-
-    public void bindApplicationsIfNeeded() {
-        if (mAllAppsList.getAndResetChangeFlag()) {
-            AppInfo[] apps = mAllAppsList.copyData();
-            int flags = mAllAppsList.getFlags();
-            Map<PackageUserKey, Integer> packageUserKeytoUidMap = Arrays.stream(apps).collect(
-                    Collectors.toMap(
-                            appInfo -> new PackageUserKey(appInfo.componentName.getPackageName(),
-                                    appInfo.user), appInfo -> appInfo.uid, (a, b) -> a));
-            scheduleCallbackTask(c -> c.bindAllApplications(apps, flags, packageUserKeytoUidMap));
-        }
-    }
-}
diff --git a/src/com/android/launcher3/model/CacheDataUpdatedTask.java b/src/com/android/launcher3/model/CacheDataUpdatedTask.java
index 57fefaa..66b4fd9 100644
--- a/src/com/android/launcher3/model/CacheDataUpdatedTask.java
+++ b/src/com/android/launcher3/model/CacheDataUpdatedTask.java
@@ -20,7 +20,7 @@
 
 import androidx.annotation.NonNull;
 
-import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
@@ -31,7 +31,7 @@
 /**
  * Handles changes due to cache updates.
  */
-public class CacheDataUpdatedTask extends BaseModelUpdateTask {
+public class CacheDataUpdatedTask implements ModelUpdateTask {
 
     public static final int OP_CACHE_UPDATE = 1;
     public static final int OP_SESSION_UPDATE = 2;
@@ -52,9 +52,9 @@
     }
 
     @Override
-    public void execute(@NonNull final LauncherAppState app, @NonNull final BgDataModel dataModel,
-            @NonNull final AllAppsList apps) {
-        IconCache iconCache = app.getIconCache();
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList apps) {
+        IconCache iconCache = taskController.getApp().getIconCache();
         ArrayList<WorkspaceItemInfo> updatedShortcuts = new ArrayList<>();
 
         synchronized (dataModel) {
@@ -69,8 +69,8 @@
             });
             apps.updateIconsAndLabels(mPackages, mUser);
         }
-        bindUpdatedWorkspaceItems(updatedShortcuts);
-        bindApplicationsIfNeeded();
+        taskController.bindUpdatedWorkspaceItems(updatedShortcuts);
+        taskController.bindApplicationsIfNeeded();
     }
 
     public boolean isValidShortcut(@NonNull final WorkspaceItemInfo si) {
diff --git a/src/com/android/launcher3/model/ModelTaskController.kt b/src/com/android/launcher3/model/ModelTaskController.kt
new file mode 100644
index 0000000..266ed0c
--- /dev/null
+++ b/src/com/android/launcher3/model/ModelTaskController.kt
@@ -0,0 +1,104 @@
+/*
+ * Copyright (C) 2024 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.launcher3.model
+
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherModel
+import com.android.launcher3.LauncherModel.CallbackTask
+import com.android.launcher3.celllayout.CellPosMapper
+import com.android.launcher3.model.BgDataModel.FixedContainerItems
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.util.PackageUserKey
+import java.util.Objects
+import java.util.concurrent.Executor
+import java.util.function.Predicate
+
+/** Class with utility methods and properties for running a LauncherModel Task */
+class ModelTaskController(
+    val app: LauncherAppState,
+    val dataModel: BgDataModel,
+    val allAppsList: AllAppsList,
+    private val model: LauncherModel,
+    private val uiExecutor: Executor
+) {
+
+    /** Schedules a {@param task} to be executed on the current callbacks. */
+    fun scheduleCallbackTask(task: CallbackTask) {
+        for (cb in model.callbacks) {
+            uiExecutor.execute { task.execute(cb) }
+        }
+    }
+
+    /**
+     * Updates from model task, do not deal with icon position in hotseat. Also no need to verify
+     * changes as the ModelTasks always push the changes to callbacks
+     */
+    fun getModelWriter() = model.getWriter(false /* verifyChanges */, CellPosMapper.DEFAULT, null)
+
+    fun bindUpdatedWorkspaceItems(allUpdates: List<WorkspaceItemInfo>) {
+        // Bind workspace items
+        val workspaceUpdates =
+            allUpdates.stream().filter { info -> info.id != ItemInfo.NO_ID }.toList()
+        if (workspaceUpdates.isNotEmpty()) {
+            scheduleCallbackTask { it.bindWorkspaceItemsChanged(workspaceUpdates) }
+        }
+
+        // Bind extra items if any
+        allUpdates
+            .stream()
+            .mapToInt { info: WorkspaceItemInfo -> info.container }
+            .distinct()
+            .mapToObj { dataModel.extraItems.get(it) }
+            .filter { Objects.nonNull(it) }
+            .forEach { bindExtraContainerItems(it) }
+    }
+
+    fun bindExtraContainerItems(item: FixedContainerItems) {
+        scheduleCallbackTask { it.bindExtraContainerItems(item) }
+    }
+
+    fun bindDeepShortcuts(dataModel: BgDataModel) {
+        val shortcutMapCopy = HashMap(dataModel.deepShortcutMap)
+        scheduleCallbackTask { it.bindDeepShortcutMap(shortcutMapCopy) }
+    }
+
+    fun bindUpdatedWidgets(dataModel: BgDataModel) {
+        val widgets = dataModel.widgetsModel.getWidgetsListForPicker(app.context)
+        scheduleCallbackTask { it.bindAllWidgets(widgets) }
+    }
+
+    fun deleteAndBindComponentsRemoved(matcher: Predicate<ItemInfo?>, reason: String?) {
+        getModelWriter().deleteItemsFromDatabase(matcher, reason)
+
+        // Call the components-removed callback
+        scheduleCallbackTask { it.bindWorkspaceComponentsRemoved(matcher) }
+    }
+
+    fun bindApplicationsIfNeeded() {
+        if (allAppsList.getAndResetChangeFlag()) {
+            val apps = allAppsList.copyData()
+            val flags = allAppsList.flags
+            val packageUserKeyToUidMap =
+                apps.associateBy(
+                    keySelector = { PackageUserKey(it.componentName!!.packageName, it.user) },
+                    valueTransform = { it.uid }
+                )
+            scheduleCallbackTask { it.bindAllApplications(apps, flags, packageUserKeyToUidMap) }
+        }
+    }
+}
diff --git a/src/com/android/launcher3/model/PackageIncrementalDownloadUpdatedTask.java b/src/com/android/launcher3/model/PackageIncrementalDownloadUpdatedTask.java
index b9fba9d..f924a9f 100644
--- a/src/com/android/launcher3/model/PackageIncrementalDownloadUpdatedTask.java
+++ b/src/com/android/launcher3/model/PackageIncrementalDownloadUpdatedTask.java
@@ -19,7 +19,7 @@
 
 import androidx.annotation.NonNull;
 
-import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
@@ -31,7 +31,7 @@
 /**
  * Handles updates due to incremental download progress updates.
  */
-public class PackageIncrementalDownloadUpdatedTask extends BaseModelUpdateTask {
+public class PackageIncrementalDownloadUpdatedTask implements ModelUpdateTask {
 
     @NonNull
     private final UserHandle mUser;
@@ -49,8 +49,8 @@
     }
 
     @Override
-    public void execute(@NonNull LauncherAppState app, @NonNull final BgDataModel dataModel,
-            @NonNull final AllAppsList appsList) {
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList appsList) {
         PackageInstallInfo downloadInfo = new PackageInstallInfo(
                 mPackageName,
                 PackageInstallInfo.STATUS_INSTALLED_DOWNLOADING,
@@ -62,11 +62,11 @@
             if (!updatedAppInfos.isEmpty()) {
                 for (AppInfo appInfo : updatedAppInfos) {
                     appInfo.runtimeStatusFlags &= ~ItemInfoWithIcon.FLAG_INSTALL_SESSION_ACTIVE;
-                    scheduleCallbackTask(
+                    taskController.scheduleCallbackTask(
                             c -> c.bindIncrementalDownloadProgressUpdated(appInfo));
                 }
             }
-            bindApplicationsIfNeeded();
+            taskController.bindApplicationsIfNeeded();
         }
 
         final ArrayList<WorkspaceItemInfo> updatedWorkspaceItems = new ArrayList<>();
@@ -79,6 +79,6 @@
                 }
             });
         }
-        bindUpdatedWorkspaceItems(updatedWorkspaceItems);
+        taskController.bindUpdatedWorkspaceItems(updatedWorkspaceItems);
     }
 }
diff --git a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
index 2457a42..d238213 100644
--- a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
+++ b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
@@ -15,12 +15,13 @@
  */
 package com.android.launcher3.model;
 
+import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 
 import androidx.annotation.NonNull;
 
-import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
@@ -33,7 +34,7 @@
 /**
  * Handles changes due to a sessions updates for a currently installing app.
  */
-public class PackageInstallStateChangedTask extends BaseModelUpdateTask {
+public class PackageInstallStateChangedTask implements ModelUpdateTask {
 
     @NonNull
     private final PackageInstallInfo mInstallInfo;
@@ -43,16 +44,17 @@
     }
 
     @Override
-    public void execute(@NonNull final LauncherAppState app, @NonNull final BgDataModel dataModel,
-            @NonNull final AllAppsList apps) {
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList apps) {
         if (mInstallInfo.state == PackageInstallInfo.STATUS_INSTALLED) {
             try {
                 // For instant apps we do not get package-add. Use setting events to update
                 // any pinned icons.
-                ApplicationInfo ai = app.getContext()
+                Context context = taskController.getApp().getContext();
+                ApplicationInfo ai = context
                         .getPackageManager().getApplicationInfo(mInstallInfo.packageName, 0);
-                if (InstantAppResolver.newInstance(app.getContext()).isInstantApp(ai)) {
-                    app.getModel().newModelCallbacks()
+                if (InstantAppResolver.newInstance(context).isInstantApp(ai)) {
+                    taskController.getApp().getModel().newModelCallbacks()
                             .onPackageAdded(ai.packageName, mInstallInfo.user);
                 }
             } catch (PackageManager.NameNotFoundException e) {
@@ -66,10 +68,11 @@
             List<AppInfo> updatedAppInfos = apps.updatePromiseInstallInfo(mInstallInfo);
             if (!updatedAppInfos.isEmpty()) {
                 for (AppInfo appInfo : updatedAppInfos) {
-                    scheduleCallbackTask(c -> c.bindIncrementalDownloadProgressUpdated(appInfo));
+                    taskController.scheduleCallbackTask(
+                            c -> c.bindIncrementalDownloadProgressUpdated(appInfo));
                 }
             }
-            bindApplicationsIfNeeded();
+            taskController.bindApplicationsIfNeeded();
         }
 
         synchronized (dataModel) {
@@ -90,7 +93,8 @@
             }
 
             if (!updates.isEmpty()) {
-                scheduleCallbackTask(callbacks -> callbacks.bindRestoreItemsChange(updates));
+                taskController.scheduleCallbackTask(
+                        callbacks -> callbacks.bindRestoreItemsChange(updates));
             }
         }
     }
diff --git a/src/com/android/launcher3/model/PackageUpdatedTask.java b/src/com/android/launcher3/model/PackageUpdatedTask.java
index 6432d33..802faae 100644
--- a/src/com/android/launcher3/model/PackageUpdatedTask.java
+++ b/src/com/android/launcher3/model/PackageUpdatedTask.java
@@ -36,9 +36,9 @@
 
 import com.android.launcher3.Flags;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
-import com.android.launcher3.Utilities;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.logging.FileLog;
@@ -71,7 +71,7 @@
  * or when a user availability changes.
  */
 @SuppressWarnings("NewApi")
-public class PackageUpdatedTask extends BaseModelUpdateTask {
+public class PackageUpdatedTask implements ModelUpdateTask {
 
     // TODO(b/290090023): Set to false after root causing is done.
     private static final boolean DEBUG = true;
@@ -102,8 +102,9 @@
     }
 
     @Override
-    public void execute(@NonNull final LauncherAppState app, @NonNull final BgDataModel dataModel,
-            @NonNull final AllAppsList appsList) {
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList appsList) {
+        final LauncherAppState app = taskController.getApp();
         final Context context = app.getContext();
         final IconCache iconCache = app.getIconCache();
 
@@ -192,7 +193,7 @@
                 break;
         }
 
-        bindApplicationsIfNeeded();
+        taskController.bindApplicationsIfNeeded();
 
         final IntSet removedShortcuts = new IntSet();
         // Shortcuts to keep even if the corresponding app was removed
@@ -305,7 +306,7 @@
                         updatedWorkspaceItems.add(si);
                     }
                     if (infoUpdated && si.id != ItemInfo.NO_ID) {
-                        getModelWriter().updateItemInDatabase(si);
+                        taskController.getModelWriter().updateItemInDatabase(si);
                     }
                 });
 
@@ -324,20 +325,20 @@
                         widgetInfo.restoreStatus |= LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
 
                         widgets.add(widgetInfo);
-                        getModelWriter().updateItemInDatabase(widgetInfo);
+                        taskController.getModelWriter().updateItemInDatabase(widgetInfo);
                     }
                 }
             }
 
-            bindUpdatedWorkspaceItems(updatedWorkspaceItems);
+            taskController.bindUpdatedWorkspaceItems(updatedWorkspaceItems);
             if (!removedShortcuts.isEmpty()) {
-                deleteAndBindComponentsRemoved(
+                taskController.deleteAndBindComponentsRemoved(
                         ItemInfoMatcher.ofItemIds(removedShortcuts),
                         "removed because the target component is invalid");
             }
 
             if (!widgets.isEmpty()) {
-                scheduleCallbackTask(c -> c.bindWidgetsRestored(widgets));
+                taskController.scheduleCallbackTask(c -> c.bindWidgetsRestored(widgets));
             }
         }
 
@@ -363,7 +364,7 @@
                     ItemInfoMatcher.ofPackages(removedPackages, mUser)
                             .or(ItemInfoMatcher.ofComponents(removedComponents, mUser))
                             .and(ItemInfoMatcher.ofItemIds(forceKeepShortcuts).negate());
-            deleteAndBindComponentsRemoved(removeMatch,
+            taskController.deleteAndBindComponentsRemoved(removeMatch,
                     "removed because the corresponding package or component is removed. "
                             + "mOp=" + mOp + " removedPackages=" + removedPackages.stream().collect(
                                     Collectors.joining(",", "[", "]"))
@@ -382,7 +383,7 @@
             for (int i = 0; i < N; i++) {
                 dataModel.widgetsModel.update(app, new PackageUserKey(packages[i], mUser));
             }
-            bindUpdatedWidgets(dataModel);
+            taskController.bindUpdatedWidgets(dataModel);
         }
     }
 
diff --git a/src/com/android/launcher3/model/ReloadStringCacheTask.java b/src/com/android/launcher3/model/ReloadStringCacheTask.java
index 34f7057..3d974d6 100644
--- a/src/com/android/launcher3/model/ReloadStringCacheTask.java
+++ b/src/com/android/launcher3/model/ReloadStringCacheTask.java
@@ -17,13 +17,13 @@
 
 import androidx.annotation.NonNull;
 
-import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 
 /**
  * Handles updates due to changes in Device Policy Management resources triggered by
  * {@link android.app.admin.DevicePolicyManager#ACTION_DEVICE_POLICY_RESOURCE_UPDATED}.
  */
-public class ReloadStringCacheTask extends BaseModelUpdateTask {
+public class ReloadStringCacheTask implements ModelUpdateTask {
 
     @NonNull
     private ModelDelegate mModelDelegate;
@@ -33,12 +33,12 @@
     }
 
     @Override
-    public void execute(@NonNull final LauncherAppState app, @NonNull final BgDataModel dataModel,
-            @NonNull final AllAppsList appsList) {
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList apps) {
         synchronized (dataModel) {
             mModelDelegate.loadStringCache(dataModel.stringCache);
             StringCache cloneSC = dataModel.stringCache.clone();
-            scheduleCallbackTask(c -> c.bindStringCache(cloneSC));
+            taskController.scheduleCallbackTask(c -> c.bindStringCache(cloneSC));
         }
     }
 }
diff --git a/src/com/android/launcher3/model/ShortcutsChangedTask.java b/src/com/android/launcher3/model/ShortcutsChangedTask.java
index 1cb5215..1916d23 100644
--- a/src/com/android/launcher3/model/ShortcutsChangedTask.java
+++ b/src/com/android/launcher3/model/ShortcutsChangedTask.java
@@ -22,6 +22,7 @@
 import androidx.annotation.NonNull;
 
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.shortcuts.ShortcutKey;
@@ -38,7 +39,7 @@
 /**
  * Handles changes due to shortcut manager updates (deep shortcut changes)
  */
-public class ShortcutsChangedTask extends BaseModelUpdateTask {
+public class ShortcutsChangedTask implements ModelUpdateTask {
 
     @NonNull
     private final String mPackageName;
@@ -61,8 +62,9 @@
     }
 
     @Override
-    public void execute(@NonNull final LauncherAppState app, @NonNull final BgDataModel dataModel,
-            @NonNull final AllAppsList apps) {
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList apps) {
+        final LauncherAppState app = taskController.getApp();
         final Context context = app.getContext();
         // Find WorkspaceItemInfo's that have changed on the workspace.
         ArrayList<WorkspaceItemInfo> matchingWorkspaceItems = new ArrayList<>();
@@ -78,8 +80,8 @@
 
         if (!matchingWorkspaceItems.isEmpty()) {
             if (mShortcuts.isEmpty()) {
-                PackageManagerHelper packageManagerHelper = PackageManagerHelper.INSTANCE.get(
-                        app.getContext());
+                PackageManagerHelper packageManagerHelper =
+                        PackageManagerHelper.INSTANCE.get(context);
                 // Verify that the app is indeed installed.
                 if (!packageManagerHelper.isAppInstalled(mPackageName, mUser)
                         && !packageManagerHelper.isAppArchivedForUser(mPackageName, mUser)) {
@@ -115,9 +117,9 @@
                         });
             }
 
-            bindUpdatedWorkspaceItems(updatedWorkspaceItemInfos);
+            taskController.bindUpdatedWorkspaceItems(updatedWorkspaceItemInfos);
             if (!nonPinnedIds.isEmpty()) {
-                deleteAndBindComponentsRemoved(ItemInfoMatcher.ofShortcutKeys(
+                taskController.deleteAndBindComponentsRemoved(ItemInfoMatcher.ofShortcutKeys(
                         nonPinnedIds.stream()
                                 .map(id -> new ShortcutKey(mPackageName, mUser, id))
                                 .collect(Collectors.toSet())),
@@ -128,7 +130,7 @@
         if (mUpdateIdMap) {
             // Update the deep shortcut map if the list of ids has changed for an activity.
             dataModel.updateDeepShortcutCounts(mPackageName, mUser, mShortcuts);
-            bindDeepShortcuts(dataModel);
+            taskController.bindDeepShortcuts(dataModel);
         }
     }
 }
diff --git a/src/com/android/launcher3/model/UserLockStateChangedTask.java b/src/com/android/launcher3/model/UserLockStateChangedTask.java
index 63ca35b..3dc5ff3 100644
--- a/src/com/android/launcher3/model/UserLockStateChangedTask.java
+++ b/src/com/android/launcher3/model/UserLockStateChangedTask.java
@@ -24,6 +24,7 @@
 import androidx.annotation.NonNull;
 
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherModel.ModelUpdateTask;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.shortcuts.ShortcutKey;
@@ -40,7 +41,7 @@
 /**
  * Task to handle changing of lock state of the user
  */
-public class UserLockStateChangedTask extends BaseModelUpdateTask {
+public class UserLockStateChangedTask implements ModelUpdateTask {
 
     @NonNull
     private final UserHandle mUser;
@@ -52,8 +53,9 @@
     }
 
     @Override
-    public void execute(@NonNull final LauncherAppState app, @NonNull final BgDataModel dataModel,
-            @NonNull final AllAppsList apps) {
+    public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
+            @NonNull AllAppsList apps) {
+        LauncherAppState app = taskController.getApp();
         Context context = app.getContext();
 
         HashMap<ShortcutKey, ShortcutInfo> pinnedShortcuts = new HashMap<>();
@@ -98,9 +100,10 @@
                 }
             });
         }
-        bindUpdatedWorkspaceItems(updatedWorkspaceItemInfos);
+        taskController.bindUpdatedWorkspaceItems(updatedWorkspaceItemInfos);
         if (!removedKeys.isEmpty()) {
-            deleteAndBindComponentsRemoved(ItemInfoMatcher.ofShortcutKeys(removedKeys),
+            taskController.deleteAndBindComponentsRemoved(
+                    ItemInfoMatcher.ofShortcutKeys(removedKeys),
                     "removed during unlock because it's no longer available"
                             + " (possibly due to clear data)");
         }
@@ -118,6 +121,6 @@
                     null, mUser,
                     new ShortcutRequest(context, mUser).query(ShortcutRequest.ALL));
         }
-        bindDeepShortcuts(dataModel);
+        taskController.bindDeepShortcuts(dataModel);
     }
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
index 002f496..f18c02b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -44,15 +44,12 @@
 import android.test.mock.MockContentResolver;
 import android.util.ArrayMap;
 
-import androidx.annotation.NonNull;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.uiautomator.UiDevice;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
-import com.android.launcher3.LauncherModel.ModelUpdateTask;
-import com.android.launcher3.model.AllAppsList;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.testing.TestInformationProvider;
@@ -66,7 +63,6 @@
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Executor;
 
 /**
  * Utility class to help manage Launcher Model and related objects for test.
@@ -115,17 +111,9 @@
 
     public synchronized BgDataModel getBgDataModel() {
         if (mDataModel == null) {
-            getModel().enqueueModelUpdateTask(new ModelUpdateTask() {
-                @Override
-                public void init(@NonNull LauncherAppState app, @NonNull LauncherModel model,
-                        @NonNull BgDataModel dataModel, @NonNull AllAppsList allAppsList,
-                        @NonNull Executor uiExecutor) {
-                    mDataModel = dataModel;
-                }
-
-                @Override
-                public void run() { }
-            });
+            getModel().enqueueModelUpdateTask((taskController, dataModel, apps) ->
+                    mDataModel = dataModel);
+            runOnExecutorSync(Executors.MODEL_EXECUTOR, () -> { });
         }
         return mDataModel;
     }