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);
         }
     }
 }