diff --git a/protos/launcher_atom.proto b/protos/launcher_atom.proto
index cf854ed..10eedc8 100644
--- a/protos/launcher_atom.proto
+++ b/protos/launcher_atom.proto
@@ -76,6 +76,9 @@
 
 // Represents the apps list sorted alphabetically inside the all-apps view.
 message AllAppsContainer {
+  oneof ParentContainer {
+    TaskBarContainer taskbar_container = 1;
+  }
 }
 
 message WidgetsContainer {
@@ -83,6 +86,9 @@
 
 // Represents the predicted apps row(top row) in the all-apps view.
 message PredictionContainer {
+  oneof ParentContainer {
+    TaskBarContainer taskbar_container = 1;
+  }
 }
 
 // Represents the apps container within search results.
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 89e54b8..ef9c62d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -380,6 +380,14 @@
             folderBuilder.clearHotseat();
             itemInfoBuilder.setContainerInfo(LauncherAtom.ContainerInfo.newBuilder()
                     .setFolder(folderBuilder));
+        } else if (oldContainer.hasAllAppsContainer()) {
+            itemInfoBuilder.setContainerInfo(LauncherAtom.ContainerInfo.newBuilder()
+                    .setAllAppsContainer(oldContainer.getAllAppsContainer().toBuilder()
+                            .setTaskbarContainer(LauncherAtom.TaskBarContainer.newBuilder())));
+        } else if (oldContainer.hasPredictionContainer()) {
+            itemInfoBuilder.setContainerInfo(LauncherAtom.ContainerInfo.newBuilder()
+                    .setPredictionContainer(oldContainer.getPredictionContainer().toBuilder()
+                            .setTaskbarContainer(LauncherAtom.TaskBarContainer.newBuilder())));
         }
     }
 
@@ -703,6 +711,7 @@
             }
         } else if (tag instanceof AppInfo) {
             startItemInfoActivity((AppInfo) tag);
+            mControllers.uiController.onTaskbarIconLaunched((AppInfo) tag);
         } else {
             Log.e(TAG, "Unknown type clicked: " + tag);
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
index e1ce898..6416720 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
@@ -18,6 +18,7 @@
 import static com.android.launcher3.LauncherAnimUtils.SCALE_PROPERTY;
 import static com.android.launcher3.Utilities.squaredHypot;
 import static com.android.launcher3.anim.Interpolators.LINEAR;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_ALLAPPS_BUTTON_TAP;
 import static com.android.quickstep.AnimatedFloat.VALUE;
 
 import android.annotation.NonNull;
@@ -341,7 +342,10 @@
         }
 
         public View.OnClickListener getAllAppsButtonClickListener() {
-            return v -> mControllers.taskbarAllAppsController.show();
+            return v -> {
+                mActivity.getStatsLogManager().logger().log(LAUNCHER_TASKBAR_ALLAPPS_BUTTON_TAP);
+                mControllers.taskbarAllAppsController.show();
+            };
         }
 
         public View.OnLongClickListener getIconOnLongClickListener() {
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index 13007ea..10c56c9 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -20,6 +20,7 @@
 import static androidx.core.util.Preconditions.checkState;
 
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_NON_ACTIONABLE;
+import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.ALL_APPS_CONTAINER;
 import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.EXTENDED_CONTAINERS;
 import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.FOLDER;
 import static com.android.launcher3.logger.LauncherAtom.ContainerInfo.ContainerCase.SEARCH_RESULT_CONTAINER;
@@ -92,6 +93,7 @@
     private static final int FOLDER_HIERARCHY_OFFSET = 100;
     private static final int SEARCH_RESULT_HIERARCHY_OFFSET = 200;
     private static final int EXTENDED_CONTAINERS_HIERARCHY_OFFSET = 300;
+    private static final int ALL_APPS_HIERARCHY_OFFSET = 400;
 
     /**
      * Flags for converting SearchAttribute to integer value.
@@ -632,6 +634,9 @@
         } else if (info.getContainerInfo().getContainerCase() == EXTENDED_CONTAINERS) {
             return info.getContainerInfo().getExtendedContainers().getContainerCase().getNumber()
                     + EXTENDED_CONTAINERS_HIERARCHY_OFFSET;
+        } else if (info.getContainerInfo().getContainerCase() == ALL_APPS_CONTAINER) {
+            return info.getContainerInfo().getAllAppsContainer().getParentContainerCase()
+                    .getNumber() + ALL_APPS_HIERARCHY_OFFSET;
         } else {
             return info.getContainerInfo().getContainerCase().getNumber();
         }
diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index ca6712f..f7600ff 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -57,6 +57,7 @@
 import com.android.launcher3.testcomponent.TestCommandReceiver;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.FailureWatcher;
+import com.android.launcher3.util.rule.SamplerRule;
 import com.android.launcher3.util.rule.ScreenRecordRule;
 import com.android.quickstep.views.RecentsView;
 
@@ -111,7 +112,8 @@
         }
 
         mOrderSensitiveRules = RuleChain
-                .outerRule(new NavigationModeSwitchRule(mLauncher))
+                .outerRule(new SamplerRule())
+                .around(new NavigationModeSwitchRule(mLauncher))
                 .around(new FailureWatcher(mDevice, mLauncher));
 
         mOtherLauncherActivity = context.getPackageManager().queryIntentActivities(
diff --git a/res/drawable/gm_edit_24.xml b/res/drawable/gm_edit_24.xml
index 59a0dc2..f741333 100644
--- a/res/drawable/gm_edit_24.xml
+++ b/res/drawable/gm_edit_24.xml
@@ -5,6 +5,6 @@
     android:viewportHeight="24"
     android:tint="?attr/colorControlNormal">
   <path
-      android:fillColor="@android:color/white"
+      android:fillColor="?android:attr/textColorPrimaryInverse"
       android:pathData="M20.41,4.94l-1.35,-1.35c-0.78,-0.78 -2.05,-0.78 -2.83,0L3,16.82L3,21h4.18L20.41,7.77c0.79,-0.78 0.79,-2.05 0,-2.83zM6.41,19.06L5,19v-1.36l9.82,-9.82 1.41,1.41 -9.82,9.83z"/>
 </vector>
diff --git a/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java b/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java
index 462daf5..12eb837 100644
--- a/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java
+++ b/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java
@@ -313,7 +313,7 @@
      * Find empty space on the workspace and returns the screenId.
      */
     protected int findSpaceOnWorkspace(ItemInfo info, int[] outCoordinates) {
-        Workspace workspace = mContext.getWorkspace();
+        Workspace<?> workspace = mContext.getWorkspace();
         IntArray workspaceScreens = workspace.getScreenOrder();
         int screenId;
 
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index 7ba2317..fe9b633 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -42,8 +42,10 @@
 import android.os.UserHandle;
 import android.text.TextUtils;
 import android.util.Log;
+import android.util.SparseArray;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.core.util.Pair;
 
 import com.android.launcher3.InvariantDeviceProfile;
@@ -94,6 +96,8 @@
     private final InstantAppResolver mInstantAppResolver;
     private final IconProvider mIconProvider;
 
+    private final SparseArray<BitmapInfo> mWidgetCategoryBitmapInfos;
+
     private int mPendingIconRequestCount = 0;
 
     public IconCache(Context context, InvariantDeviceProfile idp) {
@@ -111,6 +115,7 @@
         mUserManager = UserCache.INSTANCE.get(mContext);
         mInstantAppResolver = InstantAppResolver.newInstance(mContext);
         mIconProvider = iconProvider;
+        mWidgetCategoryBitmapInfos = new SparseArray<>();
     }
 
     @Override
@@ -477,13 +482,39 @@
         CacheEntry entry = getEntryForPackageLocked(
                 infoInOut.packageName, infoInOut.user, useLowResIcon);
         applyCacheEntry(entry, infoInOut);
-        if (infoInOut.widgetCategory != NO_CATEGORY) {
-            WidgetSection widgetSection = WidgetSections.getWidgetSections(mContext)
-                    .get(infoInOut.widgetCategory);
-            infoInOut.title = mContext.getString(widgetSection.mSectionTitle);
-            infoInOut.contentDescription = mPackageManager.getUserBadgedLabel(
-                    infoInOut.title, infoInOut.user);
+        if (infoInOut.widgetCategory == NO_CATEGORY) {
+            return;
         }
+
+        WidgetSection widgetSection = WidgetSections.getWidgetSections(mContext)
+                .get(infoInOut.widgetCategory);
+        infoInOut.title = mContext.getString(widgetSection.mSectionTitle);
+        infoInOut.contentDescription = mPackageManager.getUserBadgedLabel(
+                infoInOut.title, infoInOut.user);
+        final BitmapInfo cachedBitmap = mWidgetCategoryBitmapInfos.get(infoInOut.widgetCategory);
+        if (cachedBitmap != null) {
+            infoInOut.bitmap = getBadgedIcon(cachedBitmap, infoInOut.user);
+            return;
+        }
+
+        try (LauncherIcons li = LauncherIcons.obtain(mContext)) {
+            final BitmapInfo tempBitmap = li.createBadgedIconBitmap(
+                    mContext.getDrawable(widgetSection.mSectionDrawable),
+                    new BaseIconFactory.IconOptions().setShrinkNonAdaptiveIcons(false));
+            mWidgetCategoryBitmapInfos.put(infoInOut.widgetCategory, tempBitmap);
+            infoInOut.bitmap = getBadgedIcon(tempBitmap, infoInOut.user);
+        } catch (Exception e) {
+            Log.e(TAG, "Error initializing bitmap for icons with widget category", e);
+        }
+
+    }
+
+    private synchronized BitmapInfo getBadgedIcon(@Nullable final BitmapInfo bitmap,
+            @NonNull final UserHandle user) {
+        if (bitmap == null) {
+            return getDefaultIcon(user);
+        }
+        return bitmap.withFlags(getUserFlagOpLocked(user));
     }
 
     protected void applyCacheEntry(CacheEntry entry, ItemInfoWithIcon info) {
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index 9af72c3..2b6e426 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -586,6 +586,9 @@
 
         @UiEvent(doc = "User clicked on IME quicksearch button.")
         LAUNCHER_ALLAPPS_QUICK_SEARCH_WITH_IME(1047),
+
+        @UiEvent(doc = "User tapped taskbar All Apps button.")
+        LAUNCHER_TASKBAR_ALLAPPS_BUTTON_TAP(1057),
         ;
 
         // ADD MORE
diff --git a/src/com/android/launcher3/popup/ArrowPopup.java b/src/com/android/launcher3/popup/ArrowPopup.java
index df8a3e2..196cc56 100644
--- a/src/com/android/launcher3/popup/ArrowPopup.java
+++ b/src/com/android/launcher3/popup/ArrowPopup.java
@@ -602,6 +602,7 @@
         mIsAboveIcon = y > dragLayer.getTop() + insets.top;
         if (!mIsAboveIcon) {
             y = mTempRect.top + iconHeight + extraVerticalSpace;
+            height -= extraVerticalSpace;
         }
 
         // Insets are added later, so subtract them now.
@@ -609,7 +610,7 @@
         y -= insets.top;
 
         mGravity = 0;
-        if (y + height > dragLayer.getBottom() - insets.bottom) {
+        if ((insets.top + y + height) > (dragLayer.getBottom() - insets.bottom)) {
             // The container is opening off the screen, so just center it in the drag layer instead.
             mGravity = Gravity.CENTER_VERTICAL;
             // Put the container next to the icon, preferring the right side in ltr (left in rtl).
diff --git a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
index 1f6551e..130ee3a 100644
--- a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
@@ -18,7 +18,6 @@
 
 import static com.android.launcher3.graphics.PreloadIconDrawable.newPendingIcon;
 import static com.android.launcher3.icons.FastBitmapDrawable.getDisabledColorFilter;
-import static com.android.launcher3.widget.WidgetSections.getWidgetSections;
 
 import android.content.Context;
 import android.graphics.Canvas;
@@ -341,8 +340,6 @@
         if (mInfo.pendingItemInfo.widgetCategory == WidgetSections.NO_CATEGORY) {
             return null;
         }
-        Context context = getContext();
-        return context.getDrawable(getWidgetSections(context).get(
-                mInfo.pendingItemInfo.widgetCategory).mSectionDrawable);
+        return mInfo.pendingItemInfo.newIcon(getContext());
     }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
index 932e06d..b0e2ec1 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListHeader.java
@@ -15,7 +15,6 @@
  */
 package com.android.launcher3.widget.picker;
 
-import static com.android.launcher3.widget.WidgetSections.NO_CATEGORY;
 
 import android.content.Context;
 import android.content.res.Resources;
@@ -43,8 +42,6 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.PluralMessageFormat;
 import com.android.launcher3.views.ActivityContext;
-import com.android.launcher3.widget.WidgetSections;
-import com.android.launcher3.widget.WidgetSections.WidgetSection;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.model.WidgetsListSearchHeaderEntry;
 
@@ -177,13 +174,7 @@
 
     private void setIcon(PackageItemInfo info) {
         Drawable icon;
-        if (info.widgetCategory == NO_CATEGORY) {
-            icon = info.newIcon(getContext());
-        } else {
-            WidgetSection widgetSection = WidgetSections.getWidgetSections(getContext())
-                    .get(info.widgetCategory);
-            icon = getContext().getDrawable(widgetSection.mSectionDrawable);
-        }
+        icon = info.newIcon(getContext());
         applyDrawables(icon);
         mIconDrawable = icon;
         if (mIconDrawable != null) {
diff --git a/tests/Android.bp b/tests/Android.bp
index 7542d04..54cded0 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -37,6 +37,7 @@
       "src/com/android/launcher3/util/WidgetUtils.java",
       "src/com/android/launcher3/util/rule/FailureWatcher.java",
       "src/com/android/launcher3/util/rule/LauncherActivityRule.java",
+      "src/com/android/launcher3/util/rule/SamplerRule.java",
       "src/com/android/launcher3/util/rule/ScreenRecordRule.java",
       "src/com/android/launcher3/util/rule/ShellCommandRule.java",
       "src/com/android/launcher3/util/rule/SimpleActivityRule.java",
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 136f115..7080c85 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -68,6 +68,7 @@
 import com.android.launcher3.util.WidgetUtils;
 import com.android.launcher3.util.rule.FailureWatcher;
 import com.android.launcher3.util.rule.LauncherActivityRule;
+import com.android.launcher3.util.rule.SamplerRule;
 import com.android.launcher3.util.rule.ScreenRecordRule;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.util.rule.TestStabilityRule;
@@ -227,7 +228,8 @@
 
     @Rule
     public TestRule mOrderSensitiveRules = RuleChain
-            .outerRule(new TestStabilityRule())
+            .outerRule(new SamplerRule())
+            .around(new TestStabilityRule())
             .around(mActivityMonitor)
             .around(getRulesInsideActivityMonitor());
 
diff --git a/tests/src/com/android/launcher3/util/rule/SamplerRule.java b/tests/src/com/android/launcher3/util/rule/SamplerRule.java
new file mode 100644
index 0000000..6125f2a
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/rule/SamplerRule.java
@@ -0,0 +1,129 @@
+/*
+ * Copyright (C) 2022 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.util.rule;
+
+import android.os.SystemClock;
+
+import androidx.test.InstrumentationRegistry;
+
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+import java.io.BufferedOutputStream;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.OutputStreamWriter;
+import java.text.SimpleDateFormat;
+import java.util.Date;
+
+/**
+ * A rule that generates a file that helps diagnosing cases when the test process was terminated
+ * because the test execution took too long, and tests that ran for too long even without being
+ * terminated. If the process was terminated or the test was long, the test leaves an artifact with
+ * stack traces of all threads, every SAMPLE_INTERVAL_MS. This will help understanding where we
+ * stuck.
+ */
+public class SamplerRule implements TestRule {
+    private static final int TOO_LONG_TEST_MS = 180000;
+    private static final int SAMPLE_INTERVAL_MS = 3000;
+
+    public static Thread startThread(Description description) {
+        Thread thread =
+                new Thread() {
+                    @Override
+                    public void run() {
+                        // Write all-threads stack stace every SAMPLE_INTERVAL_MS while the test
+                        // is running.
+                        // After the test finishes, delete that file. If the test process is
+                        // terminated due to timeout, the trace file won't be deleted.
+                        final File file = getFile();
+
+                        final long startTime = SystemClock.elapsedRealtime();
+                        try (OutputStreamWriter outputStreamWriter =
+                                     new OutputStreamWriter(
+                                             new BufferedOutputStream(
+                                                     new FileOutputStream(file)))) {
+                            writeSamples(outputStreamWriter);
+                        } catch (IOException | InterruptedException e) {
+                            // Simply suppressing the exceptions, nothing to do here.
+                        } finally {
+                            // If the process is not killed, then there was no test timeout, and
+                            // we are not interested in the trace file, unless the test ran too
+                            // long.
+                            if (SystemClock.elapsedRealtime() - startTime < TOO_LONG_TEST_MS) {
+                                file.delete();
+                            }
+                        }
+                    }
+
+                    private File getFile() {
+                        final String strDate = new SimpleDateFormat("HH:mm:ss").format(new Date());
+
+                        final String descStr = description.getTestClass().getSimpleName() + "."
+                                + description.getMethodName();
+                        return artifactFile(
+                                "ThreadStackSamples-" + strDate + "-" + descStr + ".txt");
+                    }
+
+                    private void writeSamples(OutputStreamWriter writer)
+                            throws IOException, InterruptedException {
+                        int count = 0;
+                        while (true) {
+                            writer.write(
+                                    "#"
+                                            + (count++)
+                                            + " =============================================\r\n");
+                            for (StackTraceElement[] stack : getAllStackTraces().values()) {
+                                writer.write("---------------------\r\n");
+                                for (StackTraceElement frame : stack) {
+                                    writer.write(frame.toString() + "\r\n");
+                                }
+                            }
+                            writer.flush();
+
+                            sleep(SAMPLE_INTERVAL_MS);
+                        }
+                    }
+                };
+
+        thread.start();
+        return thread;
+    }
+
+    @Override
+    public Statement apply(Statement base, Description description) {
+        return new Statement() {
+            @Override
+            public void evaluate() throws Throwable {
+                final Thread traceThread = startThread(description);
+                try {
+                    base.evaluate();
+                } finally {
+                    traceThread.interrupt();
+                    traceThread.join();
+                }
+            }
+        };
+    }
+
+    private static File artifactFile(String fileName) {
+        return new File(
+                InstrumentationRegistry.getInstrumentation().getTargetContext().getFilesDir(),
+                fileName);
+    }
+}