Caching widget labels in icon cache to avoid lookup at startup

Change-Id: Ie026ee47905454bd70e774d422cd7fe142aec7e2
diff --git a/src/com/android/launcher3/LauncherAppWidgetProviderInfo.java b/src/com/android/launcher3/LauncherAppWidgetProviderInfo.java
index 80758c9..b2b05b1 100644
--- a/src/com/android/launcher3/LauncherAppWidgetProviderInfo.java
+++ b/src/com/android/launcher3/LauncherAppWidgetProviderInfo.java
@@ -2,11 +2,15 @@
 
 import android.appwidget.AppWidgetHostView;
 import android.appwidget.AppWidgetProviderInfo;
+import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.Parcel;
+import android.os.UserHandle;
+
+import com.android.launcher3.icons.ComponentWithLabel;
 
 /**
  * This class is a thin wrapper around the framework AppWidgetProviderInfo class. This class affords
@@ -14,7 +18,8 @@
  * (who's implementation is owned by the launcher). This object represents a widget type / class,
  * as opposed to a widget instance, and so should not be confused with {@link LauncherAppWidgetInfo}
  */
-public class LauncherAppWidgetProviderInfo extends AppWidgetProviderInfo {
+public class LauncherAppWidgetProviderInfo extends AppWidgetProviderInfo
+        implements ComponentWithLabel {
 
     public static final String CLS_CUSTOM_WIDGET_PREFIX = "#custom-widget-";
 
@@ -102,4 +107,14 @@
             return 0;
         }
     }
- }
+
+    @Override
+    public final ComponentName getComponent() {
+        return provider;
+    }
+
+    @Override
+    public final UserHandle getUser() {
+        return getProfile();
+    }
+}
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 68abd68..5e09f8a 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -599,6 +599,19 @@
                 CacheDataUpdatedTask.OP_CACHE_UPDATE, user, updatedPackages));
     }
 
+    /**
+     * Called when the labels for the widgets has updated in the icon cache.
+     */
+    public void onWidgetLabelsUpdated(HashSet<String> updatedPackages, UserHandle user) {
+        enqueueModelUpdateTask(new BaseModelUpdateTask() {
+            @Override
+            public void execute(LauncherAppState app, BgDataModel dataModel, AllAppsList apps) {
+                dataModel.widgetsModel.onPackageIconsUpdated(updatedPackages, user, app);
+                bindUpdatedWidgets(dataModel);
+            }
+        });
+    }
+
     public void enqueueModelUpdateTask(ModelUpdateTask task) {
         task.init(mApp, this, sBgDataModel, mBgAllAppsList, mUiExecutor);
         runOnWorkerThread(task);
diff --git a/src/com/android/launcher3/compat/LauncherAppsCompatVL.java b/src/com/android/launcher3/compat/LauncherAppsCompatVL.java
index ba0ac6e..b641391 100644
--- a/src/com/android/launcher3/compat/LauncherAppsCompatVL.java
+++ b/src/com/android/launcher3/compat/LauncherAppsCompatVL.java
@@ -200,7 +200,7 @@
                 pm.queryIntentActivities(new Intent(Intent.ACTION_CREATE_SHORTCUT), 0)) {
             if (packageUser == null || packageUser.mPackageName
                     .equals(info.activityInfo.packageName)) {
-                result.add(new ShortcutConfigActivityInfoVL(info.activityInfo, pm));
+                result.add(new ShortcutConfigActivityInfoVL(info.activityInfo));
             }
         }
         return result;
diff --git a/src/com/android/launcher3/compat/ShortcutConfigActivityInfo.java b/src/com/android/launcher3/compat/ShortcutConfigActivityInfo.java
index d260e24..76eec6d 100644
--- a/src/com/android/launcher3/compat/ShortcutConfigActivityInfo.java
+++ b/src/com/android/launcher3/compat/ShortcutConfigActivityInfo.java
@@ -32,6 +32,7 @@
 import android.util.Log;
 import android.widget.Toast;
 
+import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.R;
@@ -40,7 +41,7 @@
 /**
  * Wrapper class for representing a shortcut configure activity.
  */
-public abstract class ShortcutConfigActivityInfo {
+public abstract class ShortcutConfigActivityInfo implements ComponentWithLabel {
 
     private static final String TAG = "SCActivityInfo";
 
@@ -52,10 +53,12 @@
         mUser = user;
     }
 
+    @Override
     public ComponentName getComponent() {
         return mCn;
     }
 
+    @Override
     public UserHandle getUser() {
         return mUser;
     }
@@ -64,8 +67,6 @@
         return LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
     }
 
-    public abstract CharSequence getLabel();
-
     public abstract Drawable getFullResIcon(IconCache cache);
 
     /**
@@ -94,8 +95,8 @@
     }
 
     /**
-     * Returns true if various properties ({@link #getLabel()}, {@link #getFullResIcon}) can
-     * be safely persisted.
+     * Returns true if various properties ({@link #getLabel(PackageManager)},
+     * {@link #getFullResIcon}) can be safely persisted.
      */
     public boolean isPersistable() {
         return true;
@@ -104,18 +105,15 @@
     static class ShortcutConfigActivityInfoVL extends ShortcutConfigActivityInfo {
 
         private final ActivityInfo mInfo;
-        private final PackageManager mPm;
 
-
-        public ShortcutConfigActivityInfoVL(ActivityInfo info, PackageManager pm) {
+        public ShortcutConfigActivityInfoVL(ActivityInfo info) {
             super(new ComponentName(info.packageName, info.name), Process.myUserHandle());
             mInfo = info;
-            mPm = pm;
         }
 
         @Override
-        public CharSequence getLabel() {
-            return mInfo.loadLabel(mPm);
+        public CharSequence getLabel(PackageManager pm) {
+            return mInfo.loadLabel(pm);
         }
 
         @Override
@@ -135,7 +133,7 @@
         }
 
         @Override
-        public CharSequence getLabel() {
+        public CharSequence getLabel(PackageManager pm) {
             return mInfo.getLabel();
         }
 
diff --git a/src/com/android/launcher3/dragndrop/AddItemActivity.java b/src/com/android/launcher3/dragndrop/AddItemActivity.java
index 278eefd..daf7dc6 100644
--- a/src/com/android/launcher3/dragndrop/AddItemActivity.java
+++ b/src/com/android/launcher3/dragndrop/AddItemActivity.java
@@ -32,6 +32,7 @@
 import android.graphics.Point;
 import android.graphics.PointF;
 import android.graphics.Rect;
+import android.os.AsyncTask;
 import android.os.Build;
 import android.os.Bundle;
 import android.view.MotionEvent;
@@ -46,6 +47,7 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherAppWidgetHost;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.LauncherModel;
 import com.android.launcher3.R;
 import com.android.launcher3.compat.AppWidgetManagerCompat;
 import com.android.launcher3.compat.LauncherAppsCompatVO;
@@ -54,6 +56,8 @@
 import com.android.launcher3.userevent.nano.LauncherLogProto.Action;
 import com.android.launcher3.userevent.nano.LauncherLogProto.ContainerType;
 import com.android.launcher3.util.InstantAppResolver;
+import com.android.launcher3.util.LooperExecutor;
+import com.android.launcher3.util.Provider;
 import com.android.launcher3.widget.PendingAddShortcutInfo;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 import com.android.launcher3.widget.WidgetHostViewLoader;
@@ -78,7 +82,6 @@
     // Widget request specific options.
     private LauncherAppWidgetHost mAppWidgetHost;
     private AppWidgetManagerCompat mAppWidgetManager;
-    private PendingAddWidgetInfo mPendingWidgetInfo;
     private int mPendingBindWidgetId;
     private Bundle mWidgetOptions;
 
@@ -189,10 +192,9 @@
     private void setupShortcut() {
         PinShortcutRequestActivityInfo shortcutInfo =
                 new PinShortcutRequestActivityInfo(mRequest, this);
-        WidgetItem item = new WidgetItem(shortcutInfo);
         mWidgetCell.getWidgetView().setTag(new PendingAddShortcutInfo(shortcutInfo));
-        mWidgetCell.applyFromCellItem(item, mApp.getWidgetCache());
-        mWidgetCell.ensurePreview();
+        applyWidgetItemAsync(
+                () -> new WidgetItem(shortcutInfo, mApp.getIconCache(), getPackageManager()));
     }
 
     private boolean setupWidget() {
@@ -207,18 +209,32 @@
         mAppWidgetManager = AppWidgetManagerCompat.getInstance(this);
         mAppWidgetHost = new LauncherAppWidgetHost(this);
 
-        mPendingWidgetInfo = new PendingAddWidgetInfo(widgetInfo);
-        mPendingWidgetInfo.spanX = Math.min(mIdp.numColumns, widgetInfo.spanX);
-        mPendingWidgetInfo.spanY = Math.min(mIdp.numRows, widgetInfo.spanY);
-        mWidgetOptions = WidgetHostViewLoader.getDefaultOptionsForWidget(this, mPendingWidgetInfo);
+        PendingAddWidgetInfo pendingInfo = new PendingAddWidgetInfo(widgetInfo);
+        pendingInfo.spanX = Math.min(mIdp.numColumns, widgetInfo.spanX);
+        pendingInfo.spanY = Math.min(mIdp.numRows, widgetInfo.spanY);
+        mWidgetOptions = WidgetHostViewLoader.getDefaultOptionsForWidget(this, pendingInfo);
+        mWidgetCell.getWidgetView().setTag(pendingInfo);
 
-        WidgetItem item = new WidgetItem(widgetInfo, getPackageManager(), mIdp);
-        mWidgetCell.getWidgetView().setTag(mPendingWidgetInfo);
-        mWidgetCell.applyFromCellItem(item, mApp.getWidgetCache());
-        mWidgetCell.ensurePreview();
+        applyWidgetItemAsync(() -> new WidgetItem(widgetInfo, mIdp, mApp.getIconCache()));
         return true;
     }
 
+    private void applyWidgetItemAsync(final Provider<WidgetItem> itemProvider) {
+        new AsyncTask<Void, Void, WidgetItem>() {
+            @Override
+            protected WidgetItem doInBackground(Void... voids) {
+                return itemProvider.get();
+            }
+
+            @Override
+            protected void onPostExecute(WidgetItem item) {
+                mWidgetCell.applyFromCellItem(item, mApp.getWidgetCache());
+                mWidgetCell.ensurePreview();
+            }
+        }.executeOnExecutor(new LooperExecutor(LauncherModel.getWorkerLooper()));
+        // TODO: Create a worker looper executor and reuse that everywhere.
+    }
+
     /**
      * Called when the cancel button is clicked.
      */
diff --git a/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java b/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java
index bd919bc..64655cc 100644
--- a/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java
+++ b/src/com/android/launcher3/dragndrop/PinShortcutRequestActivityInfo.java
@@ -22,6 +22,7 @@
 import android.content.Context;
 import android.content.pm.LauncherApps;
 import android.content.pm.LauncherApps.PinItemRequest;
+import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
 import android.graphics.drawable.Drawable;
 import android.os.Build;
@@ -65,7 +66,7 @@
     }
 
     @Override
-    public CharSequence getLabel() {
+    public CharSequence getLabel(PackageManager pm) {
         return mInfo.getShortLabel();
     }
 
diff --git a/src/com/android/launcher3/icons/BaseIconCache.java b/src/com/android/launcher3/icons/BaseIconCache.java
index f7e2be1..6433103 100644
--- a/src/com/android/launcher3/icons/BaseIconCache.java
+++ b/src/com/android/launcher3/icons/BaseIconCache.java
@@ -226,12 +226,12 @@
             entry = new CacheEntry();
             cachingLogic.loadIcon(mContext, this, object, entry);
         }
-        entry.title = cachingLogic.getLabel(object);
+        entry.title = cachingLogic.getLabel(object, mPackageManager);
         entry.contentDescription = mUserManager.getBadgedLabelForUser(entry.title, user);
         mCache.put(key, entry);
 
-        ContentValues values = newContentValues(entry.icon, entry.color,
-                entry.title.toString(), componentName.getPackageName());
+        ContentValues values = newContentValues(entry, entry.title.toString(),
+                componentName.getPackageName());
         addIconToDB(values, componentName, info, userSerial);
     }
 
@@ -280,16 +280,25 @@
      * This method is not thread safe, it must be called from a synchronized method.
      */
     protected <T> CacheEntry cacheLocked(
-            @NonNull ComponentName componentName,
-            @NonNull Provider<T> infoProvider,
-            @NonNull CachingLogic<T> cachingLogic,
-            UserHandle user, boolean usePackageIcon, boolean useLowResIcon) {
+            @NonNull ComponentName componentName, @NonNull UserHandle user,
+            @NonNull Provider<T> infoProvider, @NonNull CachingLogic<T> cachingLogic,
+            boolean usePackageIcon, boolean useLowResIcon) {
+        return cacheLocked(componentName, user, infoProvider, cachingLogic, usePackageIcon,
+                useLowResIcon, true);
+    }
+
+    protected <T> CacheEntry cacheLocked(
+            @NonNull ComponentName componentName, @NonNull UserHandle user,
+            @NonNull Provider<T> infoProvider, @NonNull CachingLogic<T> cachingLogic,
+            boolean usePackageIcon, boolean useLowResIcon, boolean addToMemCache) {
         Preconditions.assertWorkerThread();
         ComponentKey cacheKey = new ComponentKey(componentName, user);
         CacheEntry entry = mCache.get(cacheKey);
         if (entry == null || (entry.isLowRes() && !useLowResIcon)) {
             entry = new CacheEntry();
-            mCache.put(cacheKey, entry);
+            if (addToMemCache) {
+                mCache.put(cacheKey, entry);
+            }
 
             // Check the DB first.
             T object = null;
@@ -327,7 +336,7 @@
                     providerFetchedOnce = true;
                 }
                 if (object != null) {
-                    entry.title = cachingLogic.getLabel(object);
+                    entry.title = cachingLogic.getLabel(object, mPackageManager);
                     entry.contentDescription = mUserManager.getBadgedLabelForUser(entry.title, user);
                 }
             }
@@ -413,8 +422,8 @@
 
                     // Add the icon in the DB here, since these do not get written during
                     // package updates.
-                    ContentValues values = newContentValues(iconInfo.icon, entry.color,
-                            entry.title.toString(), packageName);
+                    ContentValues values = newContentValues(
+                            iconInfo, entry.title.toString(), packageName);
                     addIconToDB(values, cacheKey.componentName, info,
                             mUserManager.getSerialNumberForUser(user));
 
@@ -515,11 +524,11 @@
         }
     }
 
-    private ContentValues newContentValues(Bitmap icon, int iconColor, String label,
-            String packageName) {
+    private ContentValues newContentValues(BitmapInfo bitmapInfo, String label, String packageName) {
         ContentValues values = new ContentValues();
-        values.put(IconDB.COLUMN_ICON, Utilities.flattenBitmap(icon));
-        values.put(IconDB.COLUMN_ICON_COLOR, iconColor);
+        values.put(IconDB.COLUMN_ICON,
+                bitmapInfo.isLowRes() ? null : Utilities.flattenBitmap(bitmapInfo.icon));
+        values.put(IconDB.COLUMN_ICON_COLOR, bitmapInfo.color);
 
         values.put(IconDB.COLUMN_LABEL, label);
         values.put(IconDB.COLUMN_SYSTEM_STATE, mIconProvider.getIconSystemState(packageName));
diff --git a/src/com/android/launcher3/icons/CachingLogic.java b/src/com/android/launcher3/icons/CachingLogic.java
index 13fdc6a..24186ef 100644
--- a/src/com/android/launcher3/icons/CachingLogic.java
+++ b/src/com/android/launcher3/icons/CachingLogic.java
@@ -18,6 +18,7 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.LauncherActivityInfo;
+import android.content.pm.PackageManager;
 import android.os.UserHandle;
 
 public interface CachingLogic<T> {
@@ -26,7 +27,7 @@
 
     UserHandle getUser(T object);
 
-    CharSequence getLabel(T object);
+    CharSequence getLabel(T object, PackageManager pm);
 
     void loadIcon(Context context, BaseIconCache cache, T object, BitmapInfo target);
 
@@ -44,7 +45,7 @@
         }
 
         @Override
-        public CharSequence getLabel(LauncherActivityInfo object) {
+        public CharSequence getLabel(LauncherActivityInfo object, PackageManager pm) {
             return object.getLabel();
         }
 
@@ -57,4 +58,30 @@
             li.recycle();
         }
     };
+
+    CachingLogic<ComponentWithLabel> COMPONENT_WITH_LABEL =
+            new CachingLogic<ComponentWithLabel>() {
+
+        @Override
+        public ComponentName getComponent(ComponentWithLabel object) {
+            return object.getComponent();
+        }
+
+        @Override
+        public UserHandle getUser(ComponentWithLabel object) {
+            return object.getUser();
+        }
+
+        @Override
+        public CharSequence getLabel(ComponentWithLabel object, PackageManager pm) {
+            return object.getLabel(pm);
+        }
+
+        @Override
+        public void loadIcon(Context context, BaseIconCache cache,
+                ComponentWithLabel object, BitmapInfo target) {
+            // Do not load icon.
+            target.icon = BitmapInfo.LOW_RES_ICON;
+        }
+    };
 }
diff --git a/src/com/android/launcher3/icons/ComponentWithLabel.java b/src/com/android/launcher3/icons/ComponentWithLabel.java
new file mode 100644
index 0000000..2badb4c
--- /dev/null
+++ b/src/com/android/launcher3/icons/ComponentWithLabel.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2018 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.icons;
+
+import android.content.ComponentName;
+import android.content.pm.PackageManager;
+import android.os.UserHandle;
+
+public interface ComponentWithLabel {
+
+    ComponentName getComponent();
+
+    UserHandle getUser();
+
+    CharSequence getLabel(PackageManager pm);
+}
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index eae6f01..4349455 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.icons;
 
+import static com.android.launcher3.icons.CachingLogic.COMPONENT_WITH_LABEL;
 import static com.android.launcher3.icons.CachingLogic.LAUNCHER_ACTIVITY_INFO;
 
 import android.content.Context;
@@ -114,8 +115,8 @@
      */
     public synchronized void updateTitleAndIcon(AppInfo application) {
         CacheEntry entry = cacheLocked(application.componentName,
-                Provider.of(null), LAUNCHER_ACTIVITY_INFO,
-                application.user, false, application.usingLowResIcon());
+                application.user, Provider.of(null), LAUNCHER_ACTIVITY_INFO,
+                false, application.usingLowResIcon());
         if (entry.icon != null && !isDefaultIcon(entry.icon, application.user)) {
             applyCacheEntry(entry, application);
         }
@@ -148,6 +149,13 @@
         }
     }
 
+    public synchronized String getTitleNoCache(ComponentWithLabel info) {
+        CacheEntry entry = cacheLocked(info.getComponent(), info.getUser(), Provider.of(info),
+                COMPONENT_WITH_LABEL, false /* usePackageIcon */, true /* useLowResIcon */,
+                false /* addToMemCache */);
+        return Utilities.trim(entry.title);
+    }
+
     /**
      * Fill in {@param shortcutInfo} with the icon and label for {@param info}
      */
@@ -155,8 +163,9 @@
             @NonNull ItemInfoWithIcon infoInOut,
             @NonNull Provider<LauncherActivityInfo> activityInfoProvider,
             boolean usePkgIcon, boolean useLowResIcon) {
-        CacheEntry entry = cacheLocked(infoInOut.getTargetComponent(), activityInfoProvider,
-                LAUNCHER_ACTIVITY_INFO, infoInOut.user, usePkgIcon, useLowResIcon);
+        CacheEntry entry = cacheLocked(infoInOut.getTargetComponent(), infoInOut.user,
+                activityInfoProvider,
+                LAUNCHER_ACTIVITY_INFO, usePkgIcon, useLowResIcon);
         applyCacheEntry(entry, infoInOut);
     }
 
diff --git a/src/com/android/launcher3/icons/IconCacheUpdateHandler.java b/src/com/android/launcher3/icons/IconCacheUpdateHandler.java
index 64e3fbf..8b3af01 100644
--- a/src/com/android/launcher3/icons/IconCacheUpdateHandler.java
+++ b/src/com/android/launcher3/icons/IconCacheUpdateHandler.java
@@ -25,15 +25,17 @@
 import android.os.UserHandle;
 import android.text.TextUtils;
 import android.util.Log;
+import android.util.SparseBooleanArray;
 
-import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.icons.BaseIconCache.IconDB;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map.Entry;
 import java.util.Set;
 import java.util.Stack;
 
@@ -44,12 +46,28 @@
 
     private static final String TAG = "IconCacheUpdateHandler";
 
+    /**
+     * In this mode, all invalid icons are marked as to-be-deleted in {@link #mItemsToDelete}.
+     * This mode is used for the first run.
+     */
+    private static final boolean MODE_SET_INVALID_ITEMS = true;
+
+    /**
+     * In this mode, any valid icon is removed from {@link #mItemsToDelete}. This is used for all
+     * subsequent runs, which essentially acts as set-union of all valid items.
+     */
+    private static final boolean MODE_CLEAR_VALID_ITEMS = false;
+
     private static final Object ICON_UPDATE_TOKEN = new Object();
 
     private final HashMap<String, PackageInfo> mPkgInfoMap;
     private final BaseIconCache mIconCache;
+
     private final HashMap<UserHandle, Set<String>> mPackagesToIgnore = new HashMap<>();
 
+    private final SparseBooleanArray mItemsToDelete = new SparseBooleanArray();
+    private boolean mFilterMode = MODE_SET_INVALID_ITEMS;
+
     IconCacheUpdateHandler(BaseIconCache cache) {
         mIconCache = cache;
 
@@ -77,24 +95,43 @@
      * the DB and are updated.
      * @return The set of packages for which icons have updated.
      */
-    public <T> void updateIcons(List<T> apps, CachingLogic<T> cachingLogic) {
-        if (apps.isEmpty()) {
-            return;
+    public <T> void updateIcons(List<T> apps, CachingLogic<T> cachingLogic,
+            OnUpdateCallback onUpdateCallback) {
+        // Filter the list per user
+        HashMap<UserHandle, HashMap<ComponentName, T>> userComponentMap = new HashMap<>();
+        int count = apps.size();
+        for (int i = 0; i < count; i++) {
+            T app = apps.get(i);
+            UserHandle userHandle = cachingLogic.getUser(app);
+            HashMap<ComponentName, T> componentMap = userComponentMap.get(userHandle);
+            if (componentMap == null) {
+                componentMap = new HashMap<>();
+                userComponentMap.put(userHandle, componentMap);
+            }
+            componentMap.put(cachingLogic.getComponent(app), app);
         }
-        UserHandle user = cachingLogic.getUser(apps.get(0));
 
+        for (Entry<UserHandle, HashMap<ComponentName, T>> entry : userComponentMap.entrySet()) {
+            updateIconsPerUser(entry.getKey(), entry.getValue(), cachingLogic, onUpdateCallback);
+        }
+
+        // From now on, clear every valid item from the global valid map.
+        mFilterMode = MODE_CLEAR_VALID_ITEMS;
+    }
+
+    /**
+     * Updates the persistent DB, such that only entries corresponding to {@param apps} remain in
+     * the DB and are updated.
+     * @return The set of packages for which icons have updated.
+     */
+    private <T> void updateIconsPerUser(UserHandle user, HashMap<ComponentName, T> componentMap,
+            CachingLogic<T> cachingLogic, OnUpdateCallback onUpdateCallback) {
         Set<String> ignorePackages = mPackagesToIgnore.get(user);
         if (ignorePackages == null) {
             ignorePackages = Collections.emptySet();
         }
-
         long userSerial = mIconCache.mUserManager.getSerialNumberForUser(user);
-        HashMap<ComponentName, T> componentMap = new HashMap<>();
-        for (T app : apps) {
-            componentMap.put(cachingLogic.getComponent(app), app);
-        }
 
-        HashSet<Integer> itemsToRemove = new HashSet<>();
         Stack<T> appsToUpdate = new Stack<>();
 
         try (Cursor c = mIconCache.mIconDb.query(
@@ -114,10 +151,15 @@
                 String cn = c.getString(indexComponent);
                 ComponentName component = ComponentName.unflattenFromString(cn);
                 PackageInfo info = mPkgInfoMap.get(component.getPackageName());
+
+                int rowId = c.getInt(rowIndex);
                 if (info == null) {
                     if (!ignorePackages.contains(component.getPackageName())) {
-                        mIconCache.remove(component, user);
-                        itemsToRemove.add(c.getInt(rowIndex));
+
+                        if (mFilterMode == MODE_SET_INVALID_ITEMS) {
+                            mIconCache.remove(component, user);
+                            mItemsToDelete.put(rowId, true);
+                        }
                     }
                     continue;
                 }
@@ -132,11 +174,17 @@
                 if (version == info.versionCode && updateTime == info.lastUpdateTime &&
                         TextUtils.equals(c.getString(systemStateIndex),
                                 mIconCache.mIconProvider.getIconSystemState(info.packageName))) {
+
+                    if (mFilterMode == MODE_CLEAR_VALID_ITEMS) {
+                        mItemsToDelete.put(rowId, false);
+                    }
                     continue;
                 }
                 if (app == null) {
-                    mIconCache.remove(component, user);
-                    itemsToRemove.add(c.getInt(rowIndex));
+                    if (mFilterMode == MODE_SET_INVALID_ITEMS) {
+                        mIconCache.remove(component, user);
+                        mItemsToDelete.put(rowId, true);
+                    }
                 } else {
                     appsToUpdate.add(app);
                 }
@@ -145,17 +193,28 @@
             Log.d(TAG, "Error reading icon cache", e);
             // Continue updating whatever we have read so far
         }
-        if (!itemsToRemove.isEmpty()) {
-            mIconCache.mIconDb.delete(
-                    Utilities.createDbSelectionQuery(IconDB.COLUMN_ROWID, itemsToRemove), null);
-        }
 
         // Insert remaining apps.
         if (!componentMap.isEmpty() || !appsToUpdate.isEmpty()) {
             Stack<T> appsToAdd = new Stack<>();
             appsToAdd.addAll(componentMap.values());
-            new SerializedIconUpdateTask(userSerial, user, appsToAdd, appsToUpdate, cachingLogic)
-                    .scheduleNext();
+            new SerializedIconUpdateTask(userSerial, user, appsToAdd, appsToUpdate, cachingLogic,
+                    onUpdateCallback).scheduleNext();
+        }
+    }
+
+    public void finish() {
+        // Commit all deletes
+        ArrayList<Integer> deleteIds = new ArrayList<>();
+        int count = mItemsToDelete.size();
+        for (int i = 0;  i < count; i++) {
+            if (mItemsToDelete.valueAt(i)) {
+                deleteIds.add(mItemsToDelete.keyAt(i));
+            }
+        }
+        if (!deleteIds.isEmpty()) {
+            mIconCache.mIconDb.delete(
+                    Utilities.createDbSelectionQuery(IconDB.COLUMN_ROWID, deleteIds), null);
         }
     }
 
@@ -172,14 +231,17 @@
         private final Stack<T> mAppsToUpdate;
         private final CachingLogic<T> mCachingLogic;
         private final HashSet<String> mUpdatedPackages = new HashSet<>();
+        private final OnUpdateCallback mOnUpdateCallback;
 
         SerializedIconUpdateTask(long userSerial, UserHandle userHandle,
-                Stack<T> appsToAdd, Stack<T> appsToUpdate, CachingLogic<T> cachingLogic) {
+                Stack<T> appsToAdd, Stack<T> appsToUpdate, CachingLogic<T> cachingLogic,
+                OnUpdateCallback onUpdateCallback) {
             mUserHandle = userHandle;
             mUserSerial = userSerial;
             mAppsToAdd = appsToAdd;
             mAppsToUpdate = appsToUpdate;
             mCachingLogic = cachingLogic;
+            mOnUpdateCallback = onUpdateCallback;
         }
 
         @Override
@@ -193,9 +255,8 @@
                 mUpdatedPackages.add(pkg);
 
                 if (mAppsToUpdate.isEmpty() && !mUpdatedPackages.isEmpty()) {
-                    // No more app to update. Notify model.
-                    LauncherAppState.getInstance(mIconCache.mContext).getModel()
-                            .onPackageIconsUpdated(mUpdatedPackages, mUserHandle);
+                    // No more app to update. Notify callback.
+                    mOnUpdateCallback.onPackageIconsUpdated(mUpdatedPackages, mUserHandle);
                 }
 
                 // Let it run one more time.
@@ -221,4 +282,9 @@
                     SystemClock.uptimeMillis() + 1);
         }
     }
+
+    public interface OnUpdateCallback {
+
+        void onPackageIconsUpdated(HashSet<String> updatedPackages, UserHandle user);
+    }
 }
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index d6b7b0f..e0da6b1 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -20,6 +20,7 @@
 import static com.android.launcher3.ItemInfoWithIcon.FLAG_DISABLED_SAFEMODE;
 import static com.android.launcher3.ItemInfoWithIcon.FLAG_DISABLED_SUSPENDED;
 import static com.android.launcher3.folder.ClippedFolderIconLayoutRule.MAX_NUM_ITEMS_IN_PREVIEW;
+import static com.android.launcher3.icons.CachingLogic.COMPONENT_WITH_LABEL;
 import static com.android.launcher3.icons.CachingLogic.LAUNCHER_ACTIVITY_INFO;
 import static com.android.launcher3.model.LoaderResults.filterCurrentWorkspaceItems;
 
@@ -44,6 +45,7 @@
 import com.android.launcher3.AllAppsList;
 import com.android.launcher3.AppInfo;
 import com.android.launcher3.FolderInfo;
+import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCacheUpdateHandler;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.InstallShortcutReceiver;
@@ -184,7 +186,7 @@
 
             // second step
             TraceHelper.partitionSection(TAG, "step 2.1: loading all apps");
-            List<List<LauncherActivityInfo>> activityListPerUser = loadAllApps();
+            List<LauncherActivityInfo> allActivityList = loadAllApps();
 
             TraceHelper.partitionSection(TAG, "step 2.2: Binding all apps");
             verifyNotStopped();
@@ -194,7 +196,8 @@
             TraceHelper.partitionSection(TAG, "step 2.3: Update icon cache");
             IconCacheUpdateHandler updateHandler = mIconCache.getUpdateHandler();
             setIgnorePackages(updateHandler);
-            updateIconCacheForApps(updateHandler, activityListPerUser);
+            updateHandler.updateIcons(allActivityList, LAUNCHER_ACTIVITY_INFO,
+                    mApp.getModel()::onPackageIconsUpdated);
 
             // Take a break
             TraceHelper.partitionSection(TAG, "step 2 completed, wait for idle");
@@ -216,12 +219,21 @@
 
             // fourth step
             TraceHelper.partitionSection(TAG, "step 4.1: loading widgets");
-            mBgDataModel.widgetsModel.update(mApp, null);
+            List<ComponentWithLabel> allWidgetsList = mBgDataModel.widgetsModel.update(mApp, null);
 
             verifyNotStopped();
             TraceHelper.partitionSection(TAG, "step 4.2: Binding widgets");
             mResults.bindWidgets();
 
+            verifyNotStopped();
+            TraceHelper.partitionSection(TAG, "step 4.3: Update icon cache");
+            updateHandler.updateIcons(allWidgetsList, COMPONENT_WITH_LABEL,
+                    mApp.getModel()::onWidgetLabelsUpdated);
+
+            verifyNotStopped();
+            TraceHelper.partitionSection(TAG, "step 5: Finish icon cache update");
+            updateHandler.finish();
+
             transaction.commit();
         } catch (CancellationException e) {
             // Loader stopped, ignore
@@ -799,17 +811,9 @@
         updateHandler.setPackagesToIgnore(Process.myUserHandle(), packagesToIgnore);
     }
 
-    private void updateIconCacheForApps(IconCacheUpdateHandler updateHandler,
-            List<List<LauncherActivityInfo>> activityListPerUser) {
-        int userCount = activityListPerUser.size();
-        for (int i = 0; i < userCount; i++) {
-            updateHandler.updateIcons(activityListPerUser.get(i), LAUNCHER_ACTIVITY_INFO);
-        }
-    }
-
-    private List<List<LauncherActivityInfo>> loadAllApps() {
+    private List<LauncherActivityInfo> loadAllApps() {
         final List<UserHandle> profiles = mUserManager.getUserProfiles();
-        List<List<LauncherActivityInfo>> activityListPerUser = new ArrayList<>();
+        List<LauncherActivityInfo> allActivityList = new ArrayList<>();
         // Clear the list of apps
         mBgAllAppsList.clear();
         for (UserHandle user : profiles) {
@@ -818,7 +822,7 @@
             // Fail if we don't have any apps
             // TODO: Fix this. Only fail for the current user.
             if (apps == null || apps.isEmpty()) {
-                return activityListPerUser;
+                return allActivityList;
             }
             boolean quietMode = mUserManager.isQuietModeEnabled(user);
             // Create the ApplicationInfos
@@ -827,7 +831,7 @@
                 // This builds the icon bitmaps.
                 mBgAllAppsList.add(new AppInfo(app, user, quietMode), app);
             }
-            activityListPerUser.add(apps);
+            allActivityList.addAll(apps);
         }
 
         if (FeatureFlags.LAUNCHER3_PROMISE_APPS_IN_ALL_APPS) {
@@ -840,7 +844,7 @@
         }
 
         mBgAllAppsList.added = new ArrayList<>();
-        return activityListPerUser;
+        return allActivityList;
     }
 
     private void loadDeepShortcuts() {
diff --git a/src/com/android/launcher3/model/WidgetItem.java b/src/com/android/launcher3/model/WidgetItem.java
index 1e96dec..e38529b 100644
--- a/src/com/android/launcher3/model/WidgetItem.java
+++ b/src/com/android/launcher3/model/WidgetItem.java
@@ -9,6 +9,7 @@
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.compat.ShortcutConfigActivityInfo;
+import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.util.ComponentKey;
 
 import java.text.Collator;
@@ -29,11 +30,11 @@
     public final String label;
     public final int spanX, spanY;
 
-    public WidgetItem(LauncherAppWidgetProviderInfo info, PackageManager pm,
-            InvariantDeviceProfile idp) {
+    public WidgetItem(LauncherAppWidgetProviderInfo info,
+            InvariantDeviceProfile idp, IconCache iconCache) {
         super(info.provider, info.getProfile());
 
-        label = Utilities.trim(info.getLabel(pm));
+        label = iconCache.getTitleNoCache(info);
         widgetInfo = info;
         activityInfo = null;
 
@@ -41,9 +42,10 @@
         spanY = Math.min(info.spanY, idp.numRows);
     }
 
-    public WidgetItem(ShortcutConfigActivityInfo info) {
+    public WidgetItem(ShortcutConfigActivityInfo info, IconCache iconCache, PackageManager pm) {
         super(info.getComponent(), info.getUser());
-        label = Utilities.trim(info.getLabel());
+        label = info.isPersistable() ? iconCache.getTitleNoCache(info) :
+                Utilities.trim(info.getLabel(pm));
         widgetInfo = null;
         activityInfo = info;
         spanX = spanY = 1;
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index 82f4fe1..9a17ec6 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -11,6 +11,7 @@
 import android.util.Log;
 
 import com.android.launcher3.AppFilter;
+import com.android.launcher3.icons.ComponentWithLabel;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
@@ -31,7 +32,10 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
 
 import androidx.annotation.Nullable;
 
@@ -76,26 +80,32 @@
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
      */
-    public void update(LauncherAppState app, @Nullable PackageUserKey packageUser) {
+    public List<ComponentWithLabel> update(LauncherAppState app, @Nullable PackageUserKey packageUser) {
         Preconditions.assertWorkerThread();
 
         Context context = app.getContext();
         final ArrayList<WidgetItem> widgetsAndShortcuts = new ArrayList<>();
+        List<ComponentWithLabel> updatedItems = new ArrayList<>();
         try {
-            PackageManager pm = context.getPackageManager();
             InvariantDeviceProfile idp = app.getInvariantDeviceProfile();
+            PackageManager pm = app.getContext().getPackageManager();
 
             // Widgets
             AppWidgetManagerCompat widgetManager = AppWidgetManagerCompat.getInstance(context);
             for (AppWidgetProviderInfo widgetInfo : widgetManager.getAllProviders(packageUser)) {
-                widgetsAndShortcuts.add(new WidgetItem(LauncherAppWidgetProviderInfo
-                        .fromProviderInfo(context, widgetInfo), pm, idp));
+                LauncherAppWidgetProviderInfo launcherWidgetInfo =
+                        LauncherAppWidgetProviderInfo.fromProviderInfo(context, widgetInfo);
+
+                widgetsAndShortcuts.add(new WidgetItem(
+                        launcherWidgetInfo, idp, app.getIconCache()));
+                updatedItems.add(launcherWidgetInfo);
             }
 
             // Shortcuts
             for (ShortcutConfigActivityInfo info : LauncherAppsCompat.getInstance(context)
                     .getCustomShortcutActivityList(packageUser)) {
-                widgetsAndShortcuts.add(new WidgetItem(info));
+                widgetsAndShortcuts.add(new WidgetItem(info, app.getIconCache(), pm));
+                updatedItems.add(info);
             }
             setWidgetsAndShortcuts(widgetsAndShortcuts, app, packageUser);
         } catch (Exception e) {
@@ -110,6 +120,7 @@
         }
 
         app.getWidgetCache().removeObsoletePreviews(widgetsAndShortcuts, packageUser);
+        return updatedItems;
     }
 
     private synchronized void setWidgetsAndShortcuts(ArrayList<WidgetItem> rawWidgetsShortcuts,
@@ -204,4 +215,26 @@
             iconCache.getTitleAndIconForApp(p, true /* userLowResIcon */);
         }
     }
+
+    public void onPackageIconsUpdated(Set<String> packageNames, UserHandle user,
+            LauncherAppState app) {
+        for (Entry<PackageItemInfo, ArrayList<WidgetItem>> entry : mWidgetsList.entrySet()) {
+            if (packageNames.contains(entry.getKey().packageName)) {
+                ArrayList<WidgetItem> items = entry.getValue();
+                int count = items.size();
+                for (int i = 0; i < count; i++) {
+                    WidgetItem item = items.get(i);
+                    if (item.user.equals(user)) {
+                        if (item.activityInfo != null) {
+                            items.set(i, new WidgetItem(item.activityInfo, app.getIconCache(),
+                                    app.getContext().getPackageManager()));
+                        } else {
+                            items.set(i, new WidgetItem(item.widgetInfo,
+                                    app.getInvariantDeviceProfile(), app.getIconCache()));
+                        }
+                    }
+                }
+            }
+        }
+    }
 }
\ No newline at end of file
diff --git a/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java b/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
index 1e56439..8503547 100644
--- a/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
+++ b/tests/src/com/android/launcher3/model/BaseModelUpdateTaskTestCase.java
@@ -205,9 +205,9 @@
         @Override
         protected <T> CacheEntry cacheLocked(
                 @NonNull ComponentName componentName,
-                @NonNull Provider<T> infoProvider,
+                UserHandle user, @NonNull Provider<T> infoProvider,
                 @NonNull CachingLogic<T> cachingLogic,
-                UserHandle user, boolean usePackageIcon, boolean useLowResIcon) {
+                boolean usePackageIcon, boolean useLowResIcon) {
             CacheEntry entry = mCache.get(new ComponentKey(componentName, user));
             if (entry == null) {
                 entry = new CacheEntry();
diff --git a/tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java b/tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
index 307a53e..a31d8a6 100644
--- a/tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
+++ b/tests/src/com/android/launcher3/widget/WidgetsListAdapterTest.java
@@ -22,7 +22,6 @@
 
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.graphics.Bitmap;
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
@@ -129,11 +128,10 @@
         if (num <= 0) return result;
 
         MultiHashMap<PackageItemInfo, WidgetItem> newMap = new MultiHashMap();
-        PackageManager pm = mContext.getPackageManager();
         AppWidgetManagerCompat widgetManager = AppWidgetManagerCompat.getInstance(mContext);
         for (AppWidgetProviderInfo widgetInfo : widgetManager.getAllProviders(null)) {
             WidgetItem wi = new WidgetItem(LauncherAppWidgetProviderInfo
-                    .fromProviderInfo(mContext, widgetInfo), pm, mTestProfile);
+                    .fromProviderInfo(mContext, widgetInfo), mTestProfile, mIconCache);
 
             PackageItemInfo pInfo = new PackageItemInfo(wi.componentName.getPackageName());
             pInfo.title = pInfo.packageName;