diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java
index 556f481..0eafb44 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java
@@ -417,7 +417,8 @@
             return;
         }
 
-        TraceHelper.INSTANCE.beginFlagsOverride(TraceHelper.FLAG_ALLOW_BINDER_TRACKING);
+        Object traceToken = TraceHelper.INSTANCE.beginFlagsOverride(
+                TraceHelper.FLAG_ALLOW_BINDER_TRACKING);
         MotionEvent event = (MotionEvent) ev;
         if (event.getAction() == ACTION_DOWN) {
             GestureState newGestureState = new GestureState(
@@ -446,7 +447,7 @@
 
         ActiveGestureLog.INSTANCE.addLog("onMotionEvent", event.getActionMasked());
         mUncheckedConsumer.onMotionEvent(event);
-        TraceHelper.INSTANCE.endFlagsOverride();
+        TraceHelper.INSTANCE.endFlagsOverride(traceToken);
     }
 
     private InputConsumer newConsumer(GestureState gestureState, boolean useSharedState,
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
index 8dd1fff..1168758 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
@@ -367,13 +367,19 @@
         if (mWasLauncherAlreadyVisible) {
             mStateCallback.setState(STATE_LAUNCHER_DRAWN);
         } else {
-            TraceHelper.INSTANCE.beginSection("WTS-init");
+            Object traceToken = TraceHelper.INSTANCE.beginSection("WTS-init");
             View dragLayer = activity.getDragLayer();
             dragLayer.getViewTreeObserver().addOnDrawListener(new OnDrawListener() {
+                boolean mHandled = false;
 
                 @Override
                 public void onDraw() {
-                    TraceHelper.INSTANCE.endSection();
+                    if (mHandled) {
+                        return;
+                    }
+                    mHandled = true;
+
+                    TraceHelper.INSTANCE.endSection(traceToken);
                     dragLayer.post(() ->
                             dragLayer.getViewTreeObserver().removeOnDrawListener(this));
                     if (activity != mActivity) {
@@ -416,13 +422,14 @@
     private void initializeLauncherAnimationController() {
         buildAnimationController();
 
-        TraceHelper.INSTANCE.beginSection("logToggleRecents", TraceHelper.FLAG_IGNORE_BINDERS);
+        Object traceToken = TraceHelper.INSTANCE.beginSection("logToggleRecents",
+                TraceHelper.FLAG_IGNORE_BINDERS);
         // Only used in debug builds
         if (LatencyTrackerCompat.isEnabled(mContext)) {
             LatencyTrackerCompat.logToggleRecents(
                     (int) (mLauncherFrameDrawnTime - mTouchTimeMs));
         }
-        TraceHelper.INSTANCE.endSection();
+        TraceHelper.INSTANCE.endSection(traceToken);
 
         // This method is only called when STATE_GESTURE_STARTED is set, so we can enable the
         // high-res thumbnail loader here once we are sure that we will end up in an overview state
@@ -1144,10 +1151,10 @@
             }
             if (!finishTransitionPosted) {
                 // If we haven't posted a draw callback, set the state immediately.
-                TraceHelper.INSTANCE.beginSection(SCREENSHOT_CAPTURED_EVT,
+                Object traceToken = TraceHelper.INSTANCE.beginSection(SCREENSHOT_CAPTURED_EVT,
                         TraceHelper.FLAG_CHECK_FOR_RACE_CONDITIONS);
                 setStateOnUiThread(STATE_SCREENSHOT_CAPTURED);
-                TraceHelper.INSTANCE.endSection();
+                TraceHelper.INSTANCE.endSection(traceToken);
             }
         }
     }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
index d231d51..02f4c40 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/OtherActivityInputConsumer.java
@@ -198,7 +198,8 @@
 
         switch (ev.getActionMasked()) {
             case ACTION_DOWN: {
-                TraceHelper.INSTANCE.beginSection(DOWN_EVT, FLAG_CHECK_FOR_RACE_CONDITIONS);
+                Object traceToken = TraceHelper.INSTANCE.beginSection(DOWN_EVT,
+                        FLAG_CHECK_FOR_RACE_CONDITIONS);
                 mActivePointerId = ev.getPointerId(0);
                 mDownPos.set(ev.getX(), ev.getY());
                 mLastPos.set(mDownPos);
@@ -209,7 +210,7 @@
                     startTouchTrackingForWindowAnimation(ev.getEventTime(), false);
                 }
 
-                TraceHelper.INSTANCE.endSection();
+                TraceHelper.INSTANCE.endSection(traceToken);
                 break;
             }
             case ACTION_POINTER_DOWN: {
@@ -355,7 +356,8 @@
      * the animation can still be running.
      */
     private void finishTouchTracking(MotionEvent ev) {
-        TraceHelper.INSTANCE.beginSection(UP_EVT, FLAG_CHECK_FOR_RACE_CONDITIONS);
+        Object traceToken = TraceHelper.INSTANCE.beginSection(UP_EVT,
+                FLAG_CHECK_FOR_RACE_CONDITIONS);
 
         if (mPassedWindowMoveSlop && mInteractionHandler != null) {
             if (ev.getActionMasked() == ACTION_CANCEL) {
@@ -389,7 +391,7 @@
         mVelocityTracker.recycle();
         mVelocityTracker = null;
         mMotionPauseDetector.clear();
-        TraceHelper.INSTANCE.endSection();
+        TraceHelper.INSTANCE.endSection(traceToken);
     }
 
     @Override
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 2951d89..4b4d793 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -313,7 +313,8 @@
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
-        TraceHelper.INSTANCE.beginSection(ON_CREATE_EVT, TraceHelper.FLAG_UI_EVENT);
+        Object traceToken = TraceHelper.INSTANCE.beginSection(ON_CREATE_EVT,
+                TraceHelper.FLAG_UI_EVENT);
         if (DEBUG_STRICT_MODE) {
             StrictMode.setThreadPolicy(new StrictMode.ThreadPolicy.Builder()
                     .detectDiskReads()
@@ -431,7 +432,7 @@
             }
         });
 
-        TraceHelper.INSTANCE.endSection();
+        TraceHelper.INSTANCE.endSection(traceToken);
     }
 
     protected LauncherOverlayManager getDefaultOverlay() {
@@ -936,14 +937,15 @@
 
     @Override
     protected void onStart() {
-        TraceHelper.INSTANCE.beginSection(ON_START_EVT, TraceHelper.FLAG_UI_EVENT);
+        Object traceToken = TraceHelper.INSTANCE.beginSection(ON_START_EVT,
+                TraceHelper.FLAG_UI_EVENT);
         super.onStart();
         if (!mDeferOverlayCallbacks) {
             mOverlayManager.onActivityStarted(this);
         }
 
         mAppWidgetHost.setListenIfResumed(true);
-        TraceHelper.INSTANCE.endSection();
+        TraceHelper.INSTANCE.endSection(traceToken);
     }
 
     private void handleDeferredResume() {
@@ -1038,7 +1040,8 @@
 
     @Override
     protected void onResume() {
-        TraceHelper.INSTANCE.beginSection(ON_RESUME_EVT, TraceHelper.FLAG_UI_EVENT);
+        Object traceToken = TraceHelper.INSTANCE.beginSection(ON_RESUME_EVT,
+                TraceHelper.FLAG_UI_EVENT);
         super.onResume();
 
         mHandler.removeCallbacks(mHandleDeferredResume);
@@ -1059,7 +1062,7 @@
             mOverlayManager.onActivityResumed(this);
         }
 
-        TraceHelper.INSTANCE.endSection();
+        TraceHelper.INSTANCE.endSection(traceToken);
     }
 
     @Override
@@ -1423,7 +1426,7 @@
 
     @Override
     protected void onNewIntent(Intent intent) {
-        TraceHelper.INSTANCE.beginSection(ON_NEW_INTENT_EVT);
+        Object traceToken = TraceHelper.INSTANCE.beginSection(ON_NEW_INTENT_EVT);
         super.onNewIntent(intent);
 
         boolean alreadyOnHome = hasWindowFocus() && ((intent.getFlags() &
@@ -1475,7 +1478,7 @@
             mOverlayManager.hideOverlay(isStarted() && !isForceInvisible());
         }
 
-        TraceHelper.INSTANCE.endSection();
+        TraceHelper.INSTANCE.endSection(traceToken);
     }
 
     @Override
@@ -1978,7 +1981,7 @@
      * Implementation of the method from LauncherModel.Callbacks.
      */
     public void startBinding() {
-        TraceHelper.INSTANCE.beginSection("startBinding");
+        Object traceToken = TraceHelper.INSTANCE.beginSection("startBinding");
         // Floating panels (except the full widget sheet) are associated with individual icons. If
         // we are starting a fresh bind, close all such panels as all the icons are about
         // to go away.
@@ -1996,7 +1999,7 @@
         if (mHotseat != null) {
             mHotseat.resetLayout(getWallpaperDeviceProfile().isVerticalBarLayout());
         }
-        TraceHelper.INSTANCE.endSection();
+        TraceHelper.INSTANCE.endSection(traceToken);
     }
 
     @Override
@@ -2196,7 +2199,7 @@
             return view;
         }
 
-        TraceHelper.INSTANCE.beginSection("BIND_WIDGET_id=" + item.appWidgetId);
+        Object traceToken = TraceHelper.INSTANCE.beginSection("BIND_WIDGET_id=" + item.appWidgetId);
 
         try {
             final LauncherAppWidgetProviderInfo appWidgetInfo;
@@ -2293,7 +2296,7 @@
             }
             prepareAppWidget(view, item);
         } finally {
-            TraceHelper.INSTANCE.endSection();
+            TraceHelper.INSTANCE.endSection(traceToken);
         }
 
         return view;
@@ -2373,7 +2376,7 @@
      * Implementation of the method from LauncherModel.Callbacks.
      */
     public void finishBindingItems(int pageBoundFirst) {
-        TraceHelper.INSTANCE.beginSection("finishBindingItems");
+        Object traceToken = TraceHelper.INSTANCE.beginSection("finishBindingItems");
         mWorkspace.restoreInstanceStateForRemainingPages();
 
         setWorkspaceLoading(false);
@@ -2397,7 +2400,7 @@
                 mDeviceProfile.inv.numFolderColumns * mDeviceProfile.inv.numFolderRows);
         getViewCache().setCacheSize(R.layout.folder_page, 2);
 
-        TraceHelper.INSTANCE.endSection();
+        TraceHelper.INSTANCE.endSection(traceToken);
     }
 
     private boolean canRunNewAppsAnimation() {
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 5582363..81b701d 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -180,7 +180,7 @@
                     "LoaderTask2 " + this);
         }
 
-        TraceHelper.INSTANCE.beginSection(TAG);
+        Object traceToken = TraceHelper.INSTANCE.beginSection(TAG);
         TimingLogger logger = TestProtocol.sDebugTracing ?
                 new TimingLogger(TAG, "run") {
                     @Override
@@ -281,7 +281,7 @@
         } finally {
             logger.dumpToLog();
         }
-        TraceHelper.INSTANCE.endSection();
+        TraceHelper.INSTANCE.endSection(traceToken);
     }
 
     public synchronized void stopLocked() {
diff --git a/src/com/android/launcher3/util/TraceHelper.java b/src/com/android/launcher3/util/TraceHelper.java
index 073fb6a..168227d 100644
--- a/src/com/android/launcher3/util/TraceHelper.java
+++ b/src/com/android/launcher3/util/TraceHelper.java
@@ -46,35 +46,44 @@
      */
     public static TraceHelper INSTANCE = new TraceHelper();
 
-    public void beginSection(String sectionName) {
-        beginSection(sectionName, 0);
+    /**
+     * @return a token to pass into {@link #endSection(Object)}.
+     */
+    public Object beginSection(String sectionName) {
+        return beginSection(sectionName, 0);
     }
 
-    public void beginSection(String sectionName, int flags) {
+    public Object beginSection(String sectionName, int flags) {
         Trace.beginSection(sectionName);
+        return null;
     }
 
-    public void endSection() {
+    /**
+     * @param token the token returned from {@link #beginSection(String, int)}
+     */
+    public void endSection(Object token) {
         Trace.endSection();
     }
 
     /**
      * Similar to {@link #beginSection} but doesn't add a trace section.
      */
-    public void beginFlagsOverride(int flags) { }
+    public Object beginFlagsOverride(int flags) {
+        return null;
+    }
 
-    public void endFlagsOverride() { }
+    public void endFlagsOverride(Object token) { }
 
     /**
      * Temporarily ignore blocking binder calls for the duration of this {@link Supplier}.
      */
     @MainThread
     public static <T> T whitelistIpcs(String rpcName, Supplier<T> supplier) {
-        INSTANCE.beginSection(rpcName, FLAG_IGNORE_BINDERS);
+        Object traceToken = INSTANCE.beginSection(rpcName, FLAG_IGNORE_BINDERS);
         try {
             return supplier.get();
         } finally {
-            INSTANCE.endSection();
+            INSTANCE.endSection(traceToken);
         }
     }
 }
diff --git a/tests/src/com/android/launcher3/util/TraceHelperForTest.java b/tests/src/com/android/launcher3/util/TraceHelperForTest.java
index 9125f5f..61e9520 100644
--- a/tests/src/com/android/launcher3/util/TraceHelperForTest.java
+++ b/tests/src/com/android/launcher3/util/TraceHelperForTest.java
@@ -41,9 +41,10 @@
     private TraceHelperForTest() { }
 
     @Override
-    public void beginSection(String sectionName, int flags) {
+    public Object beginSection(String sectionName, int flags) {
         LinkedList<TraceInfo> stack = mStack.get();
-        stack.add(new TraceInfo(sectionName, flags));
+        TraceInfo info = new TraceInfo(sectionName, flags);
+        stack.add(info);
 
         if ((flags & TraceHelper.FLAG_CHECK_FOR_RACE_CONDITIONS) != 0
                  && mRaceConditionReproducer != null) {
@@ -52,39 +53,49 @@
         updateBinderTracking(stack);
 
         super.beginSection(sectionName, flags);
+        return info;
     }
 
     @Override
-    public void endSection() {
+    public void endSection(Object token) {
         LinkedList<TraceInfo> stack = mStack.get();
-        TraceInfo info = stack.pollLast();
+        if (stack.size() == 0) {
+            new Throwable().printStackTrace();
+        }
+        TraceInfo info = (TraceInfo) token;
+        stack.remove(info);
         if ((info.flags & TraceHelper.FLAG_CHECK_FOR_RACE_CONDITIONS) != 0
                 && mRaceConditionReproducer != null) {
             mRaceConditionReproducer.onEvent(RaceConditionReproducer.exitEvt(info.sectionName));
         }
         updateBinderTracking(stack);
 
-        super.endSection();
+        super.endSection(token);
     }
 
     @Override
-    public void beginFlagsOverride(int flags) {
+    public Object beginFlagsOverride(int flags) {
         LinkedList<TraceInfo> stack = mStack.get();
-        stack.push(new TraceInfo(null, flags));
+        TraceInfo info = new TraceInfo(null, flags);
+        stack.add(info);
         updateBinderTracking(stack);
         super.beginFlagsOverride(flags);
+        return info;
     }
 
     @Override
-    public void endFlagsOverride() {
-        super.endFlagsOverride();
-        updateBinderTracking(mStack.get());
+    public void endFlagsOverride(Object token) {
+        super.endFlagsOverride(token);
+        LinkedList<TraceInfo> stack = mStack.get();
+        TraceInfo info = (TraceInfo) token;
+        stack.remove(info);
+        updateBinderTracking(stack);
     }
 
     private void updateBinderTracking(LinkedList<TraceInfo> stack) {
         if (mFlagsChangeListener != null) {
             mFlagsChangeListener.accept(stack.stream()
-                    .mapToInt(s -> s.flags).reduce(0, (a, b) -> a | b));
+                    .mapToInt(info -> info.flags).reduce(0, (a, b) -> a | b));
         }
     }