Merge "Fixes the crash issue that occurs after scrolling to page beyond the previous PDFs total pages." into androidx-main
diff --git a/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListItemPlacementAnimationTest.kt b/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListItemPlacementAnimationTest.kt
index e3e55ef..f1f79b5 100644
--- a/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListItemPlacementAnimationTest.kt
+++ b/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListItemPlacementAnimationTest.kt
@@ -543,6 +543,39 @@
     }
 
     @Test
+    fun movingAwayItem_itemWithMoreChildren_crossAxisAlignmentDefined_shouldNotCrash() {
+        var list by mutableStateOf(listOf(0, 1, 2, 3, 4, 5))
+        val listSize = itemSize * 3
+        val listSizeDp = with(rule.density) { listSize.toDp() }
+        rule.setContent {
+            LazyList(
+                maxSize = listSizeDp,
+                startIndex = 3,
+                crossAxisAlignment = CrossAxisAlignment.Center
+            ) {
+                items(list, key = { it }) {
+                    Item(it)
+                    if (it != list.last()) {
+                        Box(modifier = Modifier)
+                    }
+                }
+            }
+        }
+
+        assertPositions(3 to 0f, 4 to itemSize, 5 to itemSize * 2)
+
+        // move item 5 out of bounds
+        rule.runOnUiThread { list = listOf(5, 0, 1, 2, 3, 4) }
+
+        // should not crash
+        onAnimationFrame { fraction ->
+            if (fraction == 1.0f) {
+                assertPositions(2 to 0f, 3 to itemSize, 4 to itemSize * 2)
+            }
+        }
+    }
+
+    @Test
     fun moveItemToTheBottomOutsideOfBounds_withSpacing() {
         var list by mutableStateOf(listOf(0, 1, 2, 3, 4, 5))
         val listSize = itemSize * 3 + spacing * 2
diff --git a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListMeasuredItem.kt b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListMeasuredItem.kt
index d3c64fc..73a509f 100644
--- a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListMeasuredItem.kt
+++ b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListMeasuredItem.kt
@@ -113,10 +113,6 @@
         layoutWidth: Int,
         layoutHeight: Int
     ) {
-        require(crossAxisOffset == 0) {
-            "positioning a list item with non zero crossAxisOffset is not supported." +
-                "$crossAxisOffset was passed."
-        }
         position(mainAxisOffset, layoutWidth, layoutHeight)
     }
 
diff --git a/pdf/pdf-viewer-fragment/src/main/java/androidx/pdf/viewer/fragment/PdfViewerFragment.kt b/pdf/pdf-viewer-fragment/src/main/java/androidx/pdf/viewer/fragment/PdfViewerFragment.kt
index 8d0f787..dbd7846 100644
--- a/pdf/pdf-viewer-fragment/src/main/java/androidx/pdf/viewer/fragment/PdfViewerFragment.kt
+++ b/pdf/pdf-viewer-fragment/src/main/java/androidx/pdf/viewer/fragment/PdfViewerFragment.kt
@@ -23,6 +23,7 @@
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.view.WindowInsets
 import android.widget.FrameLayout
 import androidx.core.os.BundleCompat
 import androidx.core.view.ViewCompat
@@ -262,8 +263,6 @@
             loadingView?.showLoadingView()
         }
 
-        adjustInsetsForSearchMenu(findInFileView!!)
-
         pdfLoaderCallbacks =
             PdfLoaderCallbacksImpl(
                 requireContext(),
@@ -289,16 +288,18 @@
                     if (shouldRedrawOnDocumentLoaded) {
                         shouldRedrawOnDocumentLoaded = false
                     }
-
-                    if (annotationButton != null && isAnnotationIntentResolvable) {
-                        annotationButton?.visibility = View.VISIBLE
-                    }
                 },
                 onDocumentLoadFailure = { thrown -> onLoadDocumentError(thrown) }
             )
 
         setUpEditFab()
 
+        // Need to adjust the view only after the layout phase is completed for the views to
+        // accurately calculate the height of the view
+        pdfViewer?.viewTreeObserver?.addOnGlobalLayoutListener {
+            adjustInsetsForSearchMenu(findInFileView!!)
+        }
+
         viewModel.pdfLoaderStateFlow.value?.let { loader ->
             pdfLoader = loader
             refreshContentAndModels(loader)
@@ -350,18 +351,24 @@
      * mode but this is not required in landscape mode as the keyboard is detached from the bottom
      * of the view.
      */
-    private fun adjustInsetsForSearchMenu(view: FindInFileView) {
+    private fun adjustInsetsForSearchMenu(findInFileView: FindInFileView) {
         if (resources.configuration.orientation == Configuration.ORIENTATION_PORTRAIT) {
             WindowCompat.setDecorFitsSystemWindows(
                 requireActivity().window,
                 /* decorFitsSystemWindows= */ false
             )
 
+            val screenHeight = requireActivity().resources.displayMetrics.heightPixels
+            val height = pdfViewer?.findViewById<FrameLayout>(R.id.parent_pdf_container)!!.height
+
             // Set the listener to handle window insets
-            ViewCompat.setOnApplyWindowInsetsListener(view) { v, insets ->
+            ViewCompat.setOnApplyWindowInsetsListener(findInFileView) { view, insets ->
                 val imeInsets = insets.getInsets(WindowInsetsCompat.Type.ime())
-                v.updateLayoutParams<ViewGroup.MarginLayoutParams> {
-                    bottomMargin = imeInsets.bottom
+                view.updateLayoutParams<ViewGroup.MarginLayoutParams> {
+                    bottomMargin = 0
+                    if (!isKeyboardCollapsed(view)) {
+                        bottomMargin = imeInsets.bottom - (screenHeight - height)
+                    }
                 }
 
                 // Consume only the IME insets
@@ -370,6 +377,12 @@
         }
     }
 
+    private fun isKeyboardCollapsed(view: View): Boolean {
+        val windowInsets = view.rootWindowInsets
+        val imeVisible = windowInsets?.isVisible(WindowInsets.Type.ime()) ?: false
+        return !imeVisible
+    }
+
     /** Called after this viewer enters the screen and becomes visible. */
     private fun onEnter() {
         participateInAccessibility(true)
@@ -415,14 +428,17 @@
 
     /**
      * Posts a [.onContentsAvailable] method to be run as soon as permitted (when this Viewer has
-     * its view hierarchy built up and [.onCreateView] has finished). It might run right now if the
+     * its view hierarchy built up and [onCreateView] has finished). It might run right now if the
      * Viewer is currently started.
      */
-    private fun postContentsAvailable(contents: DisplayData) {
+    private fun postContentsAvailable(
+        contents: DisplayData,
+        showAnnotationButton: Boolean = false
+    ) {
         Preconditions.checkState(delayedContentsAvailable == null, "Already waits for contents")
 
         if (isStarted()) {
-            onContentsAvailable(contents)
+            onContentsAvailable(contents, showAnnotationButton)
             hasContents = true
         } else {
             delayedContentsAvailable = Runnable {
@@ -430,14 +446,14 @@
                     !hasContents,
                     "Received contents while restoring another copy"
                 )
-                onContentsAvailable(contents)
+                onContentsAvailable(contents, showAnnotationButton)
                 delayedContentsAvailable = null
                 hasContents = true
             }
         }
     }
 
-    private fun onContentsAvailable(contents: DisplayData) {
+    private fun onContentsAvailable(contents: DisplayData, showAnnotationButton: Boolean) {
         fileData = contents
 
         // Update the PdfLoader in the ViewModel with the new DisplayData
@@ -446,6 +462,26 @@
             contents,
             pdfLoaderCallbacks!!
         )
+        setAnnotationIntentResolvability()
+        setAnnotationButtonVisibility(showAnnotationButton)
+    }
+
+    private fun setAnnotationIntentResolvability() {
+        isAnnotationIntentResolvable =
+            AnnotationUtils.resolveAnnotationIntent(requireContext(), documentUri!!)
+        singleTapHandler?.setAnnotationIntentResolvable(isAnnotationIntentResolvable)
+        findInFileView!!.setAnnotationIntentResolvable(isAnnotationIntentResolvable)
+        (zoomScrollObserver as? ZoomScrollValueObserver)?.setAnnotationIntentResolvable(
+            isAnnotationIntentResolvable
+        )
+    }
+
+    private fun setAnnotationButtonVisibility(showAnnotationButton: Boolean) {
+        annotationButton?.let { button ->
+            if (showAnnotationButton && isAnnotationIntentResolvable) {
+                button.visibility = View.VISIBLE
+            }
+        }
     }
 
     /**
@@ -536,7 +572,8 @@
                 pdfLoader,
                 paginatedView!!,
                 zoomView!!,
-                singleTapHandler!!
+                singleTapHandler!!,
+                findInFileView!!
             )
         updatePageViewFactory(pageViewFactory!!)
     }
@@ -726,14 +763,6 @@
             annotationButton?.visibility = View.GONE
         }
         localUri = fileUri
-        isAnnotationIntentResolvable =
-            AnnotationUtils.resolveAnnotationIntent(requireContext(), localUri!!)
-        singleTapHandler?.setAnnotationIntentResolvable(isAnnotationIntentResolvable)
-        findInFileView!!.setAnnotationIntentResolvable(isAnnotationIntentResolvable)
-        findInFileView!!.resetFindInFile()
-        (zoomScrollObserver as? ZoomScrollValueObserver)?.setAnnotationIntentResolvable(
-            isAnnotationIntentResolvable
-        )
     }
 
     private fun validateFileUri(fileUri: Uri) {
@@ -794,7 +823,7 @@
     /** Feed this Viewer with contents to be displayed. */
     private fun feed(contents: DisplayData?): PdfViewerFragment {
         if (contents != null) {
-            postContentsAvailable(contents)
+            postContentsAvailable(contents, true)
         }
         return this
     }
diff --git a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PageTouchListener.java b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PageTouchListener.java
index 6a802cd..518aef2 100644
--- a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PageTouchListener.java
+++ b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PageTouchListener.java
@@ -21,6 +21,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.RestrictTo;
+import androidx.pdf.find.FindInFileView;
 import androidx.pdf.models.SelectionBoundary;
 import androidx.pdf.util.GestureTracker;
 import androidx.pdf.viewer.loader.PdfLoader;
@@ -33,15 +34,18 @@
 
     private final PdfLoader mPdfLoader;
 
+    private final FindInFileView mFindInFileView;
 
     private final SingleTapHandler mSingleTapHandler;
 
     PageTouchListener(@NonNull PageViewFactory.PageView pageView,
             @NonNull PdfLoader pdfLoader,
-            @NonNull SingleTapHandler singleTapHandler) {
+            @NonNull SingleTapHandler singleTapHandler,
+            @NonNull FindInFileView findInFileView) {
         this.mPageView = pageView;
         this.mPdfLoader = pdfLoader;
         this.mSingleTapHandler = singleTapHandler;
+        this.mFindInFileView = findInFileView;
     }
 
     @Override
@@ -57,6 +61,7 @@
 
     @Override
     public void onLongPress(MotionEvent e) {
+        mFindInFileView.resetFindInFile();
         SelectionBoundary boundary =
                 SelectionBoundary.atPoint(new Point((int) e.getX(), (int) e.getY()));
         mPdfLoader.selectPageText(mPageView.getPageNum(), boundary, boundary);
diff --git a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PageViewFactory.java b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PageViewFactory.java
index 4d1014f..bbc9b59 100644
--- a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PageViewFactory.java
+++ b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PageViewFactory.java
@@ -24,6 +24,7 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.RestrictTo;
 import androidx.annotation.VisibleForTesting;
+import androidx.pdf.find.FindInFileView;
 import androidx.pdf.models.Dimensions;
 import androidx.pdf.models.GotoLink;
 import androidx.pdf.models.LinkRects;
@@ -54,17 +55,20 @@
     private final PaginatedView mPaginatedView;
     private final ZoomView mZoomView;
     private final SingleTapHandler mSingleTapHandler;
+    private final FindInFileView mFindInFileView;
 
     public PageViewFactory(@NonNull Context context,
             @NonNull PdfLoader pdfLoader,
             @NonNull PaginatedView paginatedView,
             @NonNull ZoomView zoomView,
-            @NonNull SingleTapHandler singleTapHandler) {
+            @NonNull SingleTapHandler singleTapHandler,
+            @NonNull FindInFileView findInFileView) {
         this.mContext = context;
         this.mPdfLoader = pdfLoader;
         this.mPaginatedView = paginatedView;
         this.mZoomView = zoomView;
         this.mSingleTapHandler = singleTapHandler;
+        this.mFindInFileView = findInFileView;
     }
 
     /**
@@ -182,7 +186,7 @@
 
         GestureTracker gestureTracker = new GestureTracker(mContext);
         gestureTracker.setDelegateHandler(new PageTouchListener(pageView, mPdfLoader,
-                mSingleTapHandler));
+                mSingleTapHandler, mFindInFileView));
         pageView.asView().setOnTouchListener(gestureTracker);
 
         PageMosaicView pageMosaicView = pageView.getPageView();
diff --git a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PdfViewer.java b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PdfViewer.java
index 58f6a6b..a6b4e01 100644
--- a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PdfViewer.java
+++ b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/PdfViewer.java
@@ -282,7 +282,7 @@
         mSingleTapHandler = new SingleTapHandler(getContext(), mAnnotationButton,
                 mFindInFileView, mZoomView, mSelectionModel, mPaginationModel, mLayoutHandler);
         mPageViewFactory = new PageViewFactory(requireContext(), mPdfLoader,
-                mPaginatedView, mZoomView, mSingleTapHandler);
+                mPaginatedView, mZoomView, mSingleTapHandler, mFindInFileView);
         mPaginatedView.setPageViewFactory(mPageViewFactory);
 
         mSelectionObserver =
diff --git a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/ZoomScrollValueObserver.java b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/ZoomScrollValueObserver.java
index 2f8b034..1baeb08 100644
--- a/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/ZoomScrollValueObserver.java
+++ b/pdf/pdf-viewer/src/main/java/androidx/pdf/viewer/ZoomScrollValueObserver.java
@@ -73,6 +73,13 @@
             return;
         }
         loadPageAssets(position);
+
+        if (oldPosition.scrollY > position.scrollY) {
+            mIsPageScrollingUp = true;
+        } else if (oldPosition.scrollY < position.scrollY) {
+            mIsPageScrollingUp = false;
+        }
+
         if (mIsAnnotationIntentResolvable) {
 
             if (!isAnnotationButtonVisible() && position.scrollY == 0
diff --git a/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/MockPageViewAccessbilityDisabledFactory.java b/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/MockPageViewAccessbilityDisabledFactory.java
index 7b070fc..1589951 100644
--- a/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/MockPageViewAccessbilityDisabledFactory.java
+++ b/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/MockPageViewAccessbilityDisabledFactory.java
@@ -19,6 +19,7 @@
 import android.content.Context;
 
 import androidx.annotation.NonNull;
+import androidx.pdf.find.FindInFileView;
 import androidx.pdf.models.Dimensions;
 import androidx.pdf.util.TileBoard;
 import androidx.pdf.viewer.loader.PdfLoader;
@@ -30,8 +31,9 @@
             @NonNull PdfLoader pdfLoader,
             @NonNull PaginatedView paginatedView,
             @NonNull ZoomView zoomView,
-            @NonNull SingleTapHandler singleTapHandler) {
-        super(context, pdfLoader, paginatedView, zoomView, singleTapHandler);
+            @NonNull SingleTapHandler singleTapHandler,
+            @NonNull FindInFileView findInFileView) {
+        super(context, pdfLoader, paginatedView, zoomView, singleTapHandler, findInFileView);
     }
 
     @NonNull
diff --git a/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/MockPageViewAccessbilityEnabledFactory.java b/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/MockPageViewAccessbilityEnabledFactory.java
index cb51642..3bb879b 100644
--- a/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/MockPageViewAccessbilityEnabledFactory.java
+++ b/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/MockPageViewAccessbilityEnabledFactory.java
@@ -19,6 +19,7 @@
 import android.content.Context;
 
 import androidx.annotation.NonNull;
+import androidx.pdf.find.FindInFileView;
 import androidx.pdf.models.Dimensions;
 import androidx.pdf.util.TileBoard;
 import androidx.pdf.viewer.loader.PdfLoader;
@@ -32,8 +33,9 @@
             @NonNull PdfLoader pdfLoader,
             @NonNull PaginatedView paginatedView,
             @NonNull ZoomView zoomView,
-            @NonNull SingleTapHandler singleTapHandler) {
-        super(context, pdfLoader, paginatedView, zoomView, singleTapHandler);
+            @NonNull SingleTapHandler singleTapHandler,
+            @NonNull FindInFileView findInFileView) {
+        super(context, pdfLoader, paginatedView, zoomView, singleTapHandler, findInFileView);
     }
 
     @NonNull
diff --git a/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/PageViewFactoryTest.java b/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/PageViewFactoryTest.java
index c3869d4..f8e3651 100644
--- a/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/PageViewFactoryTest.java
+++ b/pdf/pdf-viewer/src/test/java/androidx/pdf/viewer/PageViewFactoryTest.java
@@ -27,6 +27,7 @@
 import android.graphics.drawable.ColorDrawable;
 import android.view.View;
 
+import androidx.pdf.find.FindInFileView;
 import androidx.pdf.models.Dimensions;
 import androidx.pdf.util.ObservableValue;
 import androidx.pdf.viewer.loader.PdfLoader;
@@ -52,6 +53,8 @@
 
     private final SingleTapHandler mMockSingleTapHandler = mock(SingleTapHandler.class);
 
+    private final FindInFileView mMockFindInFileView = mock(FindInFileView.class);
+
     @Before
     public void setup() {
         when(mMockZoomView.zoomScroll()).thenReturn(mock(ObservableValue.class));
@@ -68,7 +71,7 @@
                 PageViewFactory.PageView.class);
         PageViewFactory mockPageViewFactory = new MockPageViewAccessbilityDisabledFactory(
                 ApplicationProvider.getApplicationContext(), mMockPdfLoader, mMockPaginatedView,
-                mMockZoomView, mMockSingleTapHandler
+                mMockZoomView, mMockSingleTapHandler, mMockFindInFileView
         );
 
         // Act
@@ -99,7 +102,7 @@
                 PageViewFactory.PageView.class);
         PageViewFactory mockPageViewFactory = new MockPageViewAccessbilityEnabledFactory(
                 ApplicationProvider.getApplicationContext(), mMockPdfLoader, mMockPaginatedView,
-                mMockZoomView, mMockSingleTapHandler
+                mMockZoomView, mMockSingleTapHandler, mMockFindInFileView
         );
 
         // Act
diff --git a/tracing/tracing-ktx/api/current.ignore b/tracing/tracing-ktx/api/current.ignore
new file mode 100644
index 0000000..ac9c9b7
--- /dev/null
+++ b/tracing/tracing-ktx/api/current.ignore
@@ -0,0 +1,3 @@
+// Baseline format: 1.0
+RemovedPackage: androidx.tracing:
+    Removed package androidx.tracing
diff --git a/tracing/tracing-ktx/api/current.txt b/tracing/tracing-ktx/api/current.txt
index 9ea7a34..e6f50d0 100644
--- a/tracing/tracing-ktx/api/current.txt
+++ b/tracing/tracing-ktx/api/current.txt
@@ -1,12 +1 @@
 // Signature format: 4.0
-package androidx.tracing {
-
-  public final class TraceKt {
-    method public static inline <T> T trace(String label, kotlin.jvm.functions.Function0<? extends T> block);
-    method public static inline <T> T trace(kotlin.jvm.functions.Function0<java.lang.String> lazyLabel, kotlin.jvm.functions.Function0<? extends T> block);
-    method public static suspend inline <T> Object? traceAsync(String methodName, int cookie, kotlin.jvm.functions.Function1<? super kotlin.coroutines.Continuation<? super T>,? extends java.lang.Object?> block, kotlin.coroutines.Continuation<? super T>);
-    method public static inline <T> T traceAsync(kotlin.jvm.functions.Function0<java.lang.String> lazyMethodName, kotlin.jvm.functions.Function0<java.lang.Integer> lazyCookie, kotlin.jvm.functions.Function0<? extends T> block);
-  }
-
-}
-
diff --git a/tracing/tracing-ktx/api/restricted_current.ignore b/tracing/tracing-ktx/api/restricted_current.ignore
new file mode 100644
index 0000000..ac9c9b7
--- /dev/null
+++ b/tracing/tracing-ktx/api/restricted_current.ignore
@@ -0,0 +1,3 @@
+// Baseline format: 1.0
+RemovedPackage: androidx.tracing:
+    Removed package androidx.tracing
diff --git a/tracing/tracing-ktx/api/restricted_current.txt b/tracing/tracing-ktx/api/restricted_current.txt
index 9ea7a34..e6f50d0 100644
--- a/tracing/tracing-ktx/api/restricted_current.txt
+++ b/tracing/tracing-ktx/api/restricted_current.txt
@@ -1,12 +1 @@
 // Signature format: 4.0
-package androidx.tracing {
-
-  public final class TraceKt {
-    method public static inline <T> T trace(String label, kotlin.jvm.functions.Function0<? extends T> block);
-    method public static inline <T> T trace(kotlin.jvm.functions.Function0<java.lang.String> lazyLabel, kotlin.jvm.functions.Function0<? extends T> block);
-    method public static suspend inline <T> Object? traceAsync(String methodName, int cookie, kotlin.jvm.functions.Function1<? super kotlin.coroutines.Continuation<? super T>,? extends java.lang.Object?> block, kotlin.coroutines.Continuation<? super T>);
-    method public static inline <T> T traceAsync(kotlin.jvm.functions.Function0<java.lang.String> lazyMethodName, kotlin.jvm.functions.Function0<java.lang.Integer> lazyCookie, kotlin.jvm.functions.Function0<? extends T> block);
-  }
-
-}
-
diff --git a/tracing/tracing-ktx/build.gradle b/tracing/tracing-ktx/build.gradle
index 6add757..7fa130f 100644
--- a/tracing/tracing-ktx/build.gradle
+++ b/tracing/tracing-ktx/build.gradle
@@ -31,13 +31,6 @@
 
 dependencies {
     api(project(":tracing:tracing"))
-    api(libs.kotlinStdlib)
-
-    androidTestImplementation(libs.testExtJunit)
-    androidTestImplementation(libs.testCore)
-    androidTestImplementation(libs.testRunner)
-    androidTestImplementation(libs.kotlinCoroutinesAndroid)
-    testImplementation(libs.junit)
 }
 
 androidx {
diff --git a/tracing/tracing-ktx/src/main/java/androidx/tracing/Trace.kt b/tracing/tracing-ktx/src/main/java/androidx/tracing/Trace.kt
deleted file mode 100644
index d5e6d33..0000000
--- a/tracing/tracing-ktx/src/main/java/androidx/tracing/Trace.kt
+++ /dev/null
@@ -1,107 +0,0 @@
-/*
- * Copyright 2020 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 androidx.tracing
-
-/**
- * Wrap the specified [block] in calls to [Trace.beginSection] (with the supplied [label]) and
- * [Trace.endSection].
- *
- * @param label A name of the code section to appear in the trace.
- * @param block A block of code which is being traced.
- */
-public inline fun <T> trace(label: String, block: () -> T): T {
-    Trace.beginSection(label)
-    try {
-        return block()
-    } finally {
-        Trace.endSection()
-    }
-}
-
-/**
- * Wrap the specified [block] in calls to [Trace.beginSection] (with a lazy-computed [lazyLabel],
- * only if tracing is enabled - [Trace.isEnabled]) and [Trace.endSection].
- *
- * This variant allows you to build a dynamic label, but only when tracing is enabled, avoiding the
- * cost of String construction otherwise.
- *
- * @param lazyLabel A name of the code section to appear in the trace, computed lazily if needed.
- * @param block A block of code which is being traced.
- */
-public inline fun <T> trace(lazyLabel: () -> String, block: () -> T): T {
-    val isEnabled = Trace.isEnabled()
-    if (isEnabled) {
-        Trace.beginSection(lazyLabel())
-    }
-    try {
-        return block()
-    } finally {
-        if (isEnabled) {
-            Trace.endSection()
-        }
-    }
-}
-
-/**
- * Wrap the specified [block] in calls to [Trace.beginAsyncSection] (with the supplied [methodName]
- * and [cookie]) and [Trace.endAsyncSection].
- *
- * @param methodName The method name to appear in the trace.
- * @param cookie Unique identifier for distinguishing simultaneous events
- */
-public suspend inline fun <T> traceAsync(
-    methodName: String,
-    cookie: Int,
-    crossinline block: suspend () -> T
-): T {
-    Trace.beginAsyncSection(methodName, cookie)
-    try {
-        return block()
-    } finally {
-        Trace.endAsyncSection(methodName, cookie)
-    }
-}
-
-/**
- * Wrap the specified [block] in calls to [Trace.beginAsyncSection] and [Trace.endAsyncSection],
- * with a lazy-computed [lazyMethodName] and [lazyCookie], only if tracing is
- * enabled - [Trace.isEnabled].
- *
- * @param lazyMethodName The method name to appear in the trace, computed lazily if needed.
- * @param lazyCookie Unique identifier for distinguishing simultaneous events, computed lazily if
- *   needed.
- */
-public inline fun <T> traceAsync(
-    lazyMethodName: () -> String,
-    lazyCookie: () -> Int,
-    block: () -> T
-): T {
-    var methodName: String? = null
-    var cookie = 0
-    if (Trace.isEnabled()) {
-        methodName = lazyMethodName()
-        cookie = lazyCookie()
-        Trace.beginAsyncSection(methodName, cookie)
-    }
-    try {
-        return block()
-    } finally {
-        if (methodName != null) {
-            Trace.endAsyncSection(methodName, cookie)
-        }
-    }
-}
diff --git a/tracing/tracing/api/current.txt b/tracing/tracing/api/current.txt
index c883da2..7d733c5 100644
--- a/tracing/tracing/api/current.txt
+++ b/tracing/tracing/api/current.txt
@@ -2,13 +2,21 @@
 package androidx.tracing {
 
   public final class Trace {
-    method public static void beginAsyncSection(String, int);
-    method public static void beginSection(String);
-    method public static void endAsyncSection(String, int);
+    method public static void beginAsyncSection(String methodName, int cookie);
+    method public static void beginSection(String label);
+    method public static void endAsyncSection(String methodName, int cookie);
     method public static void endSection();
     method public static void forceEnableAppTracing();
     method public static boolean isEnabled();
-    method public static void setCounter(String, int);
+    method public static void setCounter(String counterName, int counterValue);
+    field public static final androidx.tracing.Trace INSTANCE;
+  }
+
+  public final class TraceKt {
+    method public static inline <T> T trace(String label, kotlin.jvm.functions.Function0<? extends T> block);
+    method public static inline <T> T trace(kotlin.jvm.functions.Function0<java.lang.String> lazyLabel, kotlin.jvm.functions.Function0<? extends T> block);
+    method public static suspend inline <T> Object? traceAsync(String methodName, int cookie, kotlin.jvm.functions.Function1<? super kotlin.coroutines.Continuation<? super T>,? extends java.lang.Object?> block, kotlin.coroutines.Continuation<? super T>);
+    method public static inline <T> T traceAsync(kotlin.jvm.functions.Function0<java.lang.String> lazyMethodName, kotlin.jvm.functions.Function0<java.lang.Integer> lazyCookie, kotlin.jvm.functions.Function0<? extends T> block);
   }
 
 }
diff --git a/tracing/tracing/api/restricted_current.txt b/tracing/tracing/api/restricted_current.txt
index c883da2..7d733c5 100644
--- a/tracing/tracing/api/restricted_current.txt
+++ b/tracing/tracing/api/restricted_current.txt
@@ -2,13 +2,21 @@
 package androidx.tracing {
 
   public final class Trace {
-    method public static void beginAsyncSection(String, int);
-    method public static void beginSection(String);
-    method public static void endAsyncSection(String, int);
+    method public static void beginAsyncSection(String methodName, int cookie);
+    method public static void beginSection(String label);
+    method public static void endAsyncSection(String methodName, int cookie);
     method public static void endSection();
     method public static void forceEnableAppTracing();
     method public static boolean isEnabled();
-    method public static void setCounter(String, int);
+    method public static void setCounter(String counterName, int counterValue);
+    field public static final androidx.tracing.Trace INSTANCE;
+  }
+
+  public final class TraceKt {
+    method public static inline <T> T trace(String label, kotlin.jvm.functions.Function0<? extends T> block);
+    method public static inline <T> T trace(kotlin.jvm.functions.Function0<java.lang.String> lazyLabel, kotlin.jvm.functions.Function0<? extends T> block);
+    method public static suspend inline <T> Object? traceAsync(String methodName, int cookie, kotlin.jvm.functions.Function1<? super kotlin.coroutines.Continuation<? super T>,? extends java.lang.Object?> block, kotlin.coroutines.Continuation<? super T>);
+    method public static inline <T> T traceAsync(kotlin.jvm.functions.Function0<java.lang.String> lazyMethodName, kotlin.jvm.functions.Function0<java.lang.Integer> lazyCookie, kotlin.jvm.functions.Function0<? extends T> block);
   }
 
 }
diff --git a/tracing/tracing/build.gradle b/tracing/tracing/build.gradle
index 2c1207d..47634ef 100644
--- a/tracing/tracing/build.gradle
+++ b/tracing/tracing/build.gradle
@@ -31,7 +31,7 @@
 
 dependencies {
     implementation("androidx.annotation:annotation:1.8.1")
-    androidTestImplementation(libs.kotlinStdlib)
+    api(libs.kotlinStdlib)
     androidTestImplementation(libs.testExtJunit)
     androidTestImplementation(libs.testCore)
     androidTestImplementation(libs.testRunner)
diff --git a/tracing/tracing-ktx/src/androidTest/java/androidx/tracing/TraceTestKt.kt b/tracing/tracing/src/androidTest/java/androidx/tracing/TraceKtTest.kt
similarity index 98%
rename from tracing/tracing-ktx/src/androidTest/java/androidx/tracing/TraceTestKt.kt
rename to tracing/tracing/src/androidTest/java/androidx/tracing/TraceKtTest.kt
index c99120c..1d8b29d 100644
--- a/tracing/tracing-ktx/src/androidTest/java/androidx/tracing/TraceTestKt.kt
+++ b/tracing/tracing/src/androidTest/java/androidx/tracing/TraceKtTest.kt
@@ -27,7 +27,7 @@
 
 @RunWith(AndroidJUnit4::class)
 @MediumTest
-class TraceTestKt {
+class TraceKtTest {
     @Test
     fun traceTest() {
         val x = trace("Test") { 10 }
diff --git a/tracing/tracing/src/androidTest/java/androidx/tracing/TraceTest.java b/tracing/tracing/src/androidTest/java/androidx/tracing/TraceTest.java
deleted file mode 100644
index 05cfb2e..0000000
--- a/tracing/tracing/src/androidTest/java/androidx/tracing/TraceTest.java
+++ /dev/null
@@ -1,197 +0,0 @@
-/*
- * Copyright 2020 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 androidx.tracing;
-
-import static androidx.tracing.Trace.MAX_TRACE_LABEL_LENGTH;
-
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-
-import static java.nio.charset.StandardCharsets.UTF_8;
-
-import android.app.UiAutomation;
-import android.os.Build;
-import android.os.ParcelFileDescriptor;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-import androidx.test.core.app.ApplicationProvider;
-import androidx.test.filters.LargeTest;
-import androidx.test.filters.SdkSuppress;
-import androidx.test.filters.SmallTest;
-import androidx.test.platform.app.InstrumentationRegistry;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Ignore;
-import org.junit.Test;
-
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-
-@LargeTest
-@SdkSuppress(minSdkVersion = 21) // Required for UiAutomation#executeShellCommand()
-public final class TraceTest {
-
-    private static final int TRACE_BUFFER_SIZE = 8192;
-    private ByteArrayOutputStream mByteArrayOutputStream;
-
-    @Before
-    public void setUp() {
-        mByteArrayOutputStream = new ByteArrayOutputStream();
-    }
-
-    @After
-    public void stopAtrace() throws IOException {
-        // Since API 23, 'async_stop' will work. On lower API levels it was broken (see aosp/157142)
-        if (Build.VERSION.SDK_INT >= 23) {
-            executeCommand("atrace --async_stop");
-        } else {
-            // Ensure tracing is not currently running by performing a short synchronous trace.
-            executeCommand("atrace -t 0");
-        }
-    }
-
-    @Test
-    @Ignore("b/280041271")
-    public void beginAndEndSection() throws IOException {
-        startTrace();
-        Trace.beginSection("beginAndEndSection");
-        Trace.endSection();
-        dumpTrace();
-
-        assertTraceContains("tracing_mark_write:\\ B\\|.*\\|beginAndEndSection");
-        assertTraceContains("tracing_mark_write:\\ E");
-    }
-
-    @Test
-    @Ignore("b/280041271")
-    public void beginAndEndTraceSectionLongLabel() throws IOException {
-        StringBuilder builder = new StringBuilder();
-        for (int i = 0; i < 20; i++) {
-            builder.append("longLabel");
-        }
-        startTrace();
-        Trace.beginSection(builder.toString());
-        Trace.endSection();
-        dumpTrace();
-        assertTraceContains(
-                "tracing_mark_write:\\ B\\|.*\\|" + builder.substring(0, MAX_TRACE_LABEL_LENGTH));
-        assertTraceContains("tracing_mark_write:\\ E");
-    }
-
-    @Test
-    @SdkSuppress(minSdkVersion = 29) // SELinux
-    public void beginAndEndSectionAsync() throws IOException {
-        startTrace();
-        Trace.beginAsyncSection("beginAndEndSectionAsync", /*cookie=*/5099);
-        Trace.endAsyncSection("beginAndEndSectionAsync", /*cookie=*/5099);
-        dumpTrace();
-
-        assertTraceContains("tracing_mark_write:\\ S\\|.*\\|beginAndEndSectionAsync\\|5099");
-        assertTraceContains("tracing_mark_write:\\ F\\|.*\\|beginAndEndSectionAsync\\|5099");
-    }
-
-    @Test
-    @SdkSuppress(minSdkVersion = 29) // SELinux
-    public void setCounter() throws IOException {
-        startTrace();
-        assertTrue("Checking that tracing is enabled", Trace.isEnabled());
-        Trace.beginSection("setting counters");
-        Trace.setCounter("counterName", 42);
-        Trace.setCounter("counterName", 47);
-        Trace.setCounter("counterName", 9787);
-        Trace.endSection();
-        assertTrue("Checking that tracing is enabled", Trace.isEnabled());
-        dumpTrace();
-        assertTraceContains("setting counters");
-        assertTraceContains("tracing_mark_write:\\ C\\|.*\\|counterName\\|42");
-        assertTraceContains("tracing_mark_write:\\ C\\|.*\\|counterName\\|47");
-        assertTraceContains("tracing_mark_write:\\ C\\|.*\\|counterName\\|9787");
-    }
-
-    @Test
-    @Ignore("b/280041271")
-    public void isEnabledDuringTrace() throws IOException {
-        startTrace();
-        boolean enabled = Trace.isEnabled();
-        dumpTrace();
-        assertTrue(enabled);
-    }
-
-    @SmallTest
-    @Test
-    public void isNotEnabledWhenNotTracing() {
-        assertFalse(Trace.isEnabled());
-    }
-
-    private void startTrace() throws IOException {
-        String processName =
-                ApplicationProvider.getApplicationContext().getApplicationInfo().processName;
-
-        // Write the "async_start" status to the byte array to ensure atrace has fully started
-        // before issuing any trace commands. This will also capture any errors that occur during
-        // start so they can be added to the assertion error's message.
-        executeCommand(
-                String.format("atrace --async_start -b %d -a %s", TRACE_BUFFER_SIZE, processName));
-    }
-
-    private void dumpTrace() throws IOException {
-        // On older versions of atrace, the -b option is required when dumping the trace so the
-        // trace buffer doesn't get cleared before being dumped.
-        executeCommand(
-                String.format("atrace --async_dump -b %d", TRACE_BUFFER_SIZE),
-                mByteArrayOutputStream);
-    }
-
-    private static void executeCommand(@NonNull String command) throws IOException {
-        executeCommand(command, null);
-    }
-
-    private static void executeCommand(@NonNull String command,
-            @Nullable ByteArrayOutputStream outputStream) throws IOException {
-        UiAutomation automation = InstrumentationRegistry.getInstrumentation().getUiAutomation();
-
-        try (ParcelFileDescriptor pfDescriptor = automation.executeShellCommand(command);
-             ParcelFileDescriptor.AutoCloseInputStream inputStream =
-                     new ParcelFileDescriptor.AutoCloseInputStream(
-                             pfDescriptor)) {
-            byte[] buffer = new byte[1024];
-
-            int length;
-            while ((length = inputStream.read(buffer)) >= 0) {
-                if (outputStream != null) {
-                    outputStream.write(buffer, 0, length);
-                }
-            }
-        }
-    }
-
-    private void assertTraceContains(@NonNull String contentRegex) {
-        String traceString = new String(mByteArrayOutputStream.toByteArray(), UTF_8);
-        Pattern pattern = Pattern.compile(contentRegex);
-        Matcher matcher = pattern.matcher(traceString);
-
-        if (!matcher.find()) {
-            throw new AssertionError(
-                    String.format("Trace does not contain requested regex: %s\n%s", contentRegex,
-                            traceString));
-        }
-    }
-}
diff --git a/tracing/tracing/src/androidTest/java/androidx/tracing/TraceTest.kt b/tracing/tracing/src/androidTest/java/androidx/tracing/TraceTest.kt
new file mode 100644
index 0000000..ebf4ba5
--- /dev/null
+++ b/tracing/tracing/src/androidTest/java/androidx/tracing/TraceTest.kt
@@ -0,0 +1,165 @@
+/*
+ * Copyright 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 androidx.tracing
+
+import android.content.Context
+import android.os.Build
+import android.os.ParcelFileDescriptor
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.filters.LargeTest
+import androidx.test.filters.SdkSuppress
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import androidx.tracing.Trace.MAX_TRACE_LABEL_LENGTH
+import java.io.ByteArrayOutputStream
+import java.nio.charset.StandardCharsets
+import org.junit.After
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
+import org.junit.Ignore
+import org.junit.Test
+
+@LargeTest
+@SdkSuppress(minSdkVersion = 21) // Required for UiAutomation#executeShellCommand()
+class TraceTest {
+    private var byteArrayOutputStream = ByteArrayOutputStream()
+
+    @After
+    fun stopAtrace() {
+        // Since API 23, 'async_stop' will work. On lower API levels it was broken (see aosp/157142)
+        if (Build.VERSION.SDK_INT >= 23) {
+            executeCommand("atrace --async_stop")
+        } else {
+            // Ensure tracing is not currently running by performing a short synchronous trace.
+            executeCommand("atrace -t 0")
+        }
+    }
+
+    @Test
+    @Ignore("b/280041271")
+    fun beginAndEndSection() {
+        startTrace()
+        Trace.beginSection("beginAndEndSection")
+        Trace.endSection()
+        dumpTrace()
+
+        assertTraceContains("tracing_mark_write:\\ B\\|.*\\|beginAndEndSection")
+        assertTraceContains("tracing_mark_write:\\ E")
+    }
+
+    @Test
+    @Ignore("b/280041271")
+    fun beginAndEndTraceSectionLongLabel() {
+        val builder = StringBuilder()
+        for (i in 0..19) {
+            builder.append("longLabel")
+        }
+        startTrace()
+        Trace.beginSection(label = builder.toString())
+        Trace.endSection()
+        dumpTrace()
+        assertTraceContains(
+            "tracing_mark_write:\\ B\\|.*\\|" + builder.substring(0, MAX_TRACE_LABEL_LENGTH)
+        )
+        assertTraceContains("tracing_mark_write:\\ E")
+    }
+
+    @Test
+    @SdkSuppress(minSdkVersion = 29) // SELinux
+    fun beginAndEndSectionAsync() {
+        startTrace()
+        Trace.beginAsyncSection(methodName = "beginAndEndSectionAsync", cookie = 5099)
+        Trace.endAsyncSection(methodName = "beginAndEndSectionAsync", cookie = 5099)
+        dumpTrace()
+
+        assertTraceContains("tracing_mark_write:\\ S\\|.*\\|beginAndEndSectionAsync\\|5099")
+        assertTraceContains("tracing_mark_write:\\ F\\|.*\\|beginAndEndSectionAsync\\|5099")
+    }
+
+    @Test
+    @SdkSuppress(minSdkVersion = 29) // SELinux
+    fun setCounter() {
+        startTrace()
+        assertTrue("Checking that tracing is enabled", Trace.isEnabled())
+        Trace.beginSection(label = "setting counters")
+        Trace.setCounter(counterName = "counterName", counterValue = 42)
+        Trace.setCounter(counterName = "counterName", counterValue = 47)
+        Trace.setCounter(counterName = "counterName", counterValue = 9787)
+        Trace.endSection()
+        assertTrue("Checking that tracing is enabled", Trace.isEnabled())
+        dumpTrace()
+        assertTraceContains("setting counters")
+        assertTraceContains("tracing_mark_write:\\ C\\|.*\\|counterName\\|42")
+        assertTraceContains("tracing_mark_write:\\ C\\|.*\\|counterName\\|47")
+        assertTraceContains("tracing_mark_write:\\ C\\|.*\\|counterName\\|9787")
+    }
+
+    @Ignore("b/280041271")
+    @Test
+    fun isEnabledDuringTrace() {
+        startTrace()
+        val enabled = Trace.isEnabled()
+        dumpTrace()
+        assertTrue(enabled)
+    }
+
+    @Test
+    @SmallTest
+    fun isNotEnabledWhenNotTracing() {
+        assertFalse(Trace.isEnabled())
+    }
+
+    private fun startTrace() {
+        val processName = getApplicationContext<Context>().applicationInfo.processName
+
+        // Write the "async_start" status to the byte array to ensure atrace has fully started
+        // before issuing any trace commands. This will also capture any errors that occur during
+        // start so they can be added to the assertion error's message.
+        executeCommand("atrace --async_start -b $TRACE_BUFFER_SIZE -a $processName")
+    }
+
+    private fun dumpTrace() {
+        // On older versions of atrace, the -b option is required when dumping the trace so the
+        // trace buffer doesn't get cleared before being dumped.
+        executeCommand("atrace --async_dump -b $TRACE_BUFFER_SIZE", byteArrayOutputStream)
+    }
+
+    private fun assertTraceContains(contentRegex: String) {
+        val traceString = byteArrayOutputStream.toByteArray().toString(StandardCharsets.UTF_8)
+
+        if (!contentRegex.toRegex().containsMatchIn(traceString)) {
+            throw AssertionError(
+                "Trace does not contain requested regex: $contentRegex\n$traceString"
+            )
+        }
+    }
+}
+
+private const val TRACE_BUFFER_SIZE = 8192
+
+private fun executeCommand(command: String, outputStream: ByteArrayOutputStream? = null) {
+    val automation = InstrumentationRegistry.getInstrumentation().uiAutomation
+
+    automation.executeShellCommand(command).use { pfDescriptor ->
+        ParcelFileDescriptor.AutoCloseInputStream(pfDescriptor).use { inputStream ->
+            val buffer = ByteArray(1024)
+            var length: Int
+            while (inputStream.read(buffer).also { length = it } >= 0) {
+                outputStream?.write(buffer, 0, length)
+            }
+        }
+    }
+}
diff --git a/tracing/tracing/src/main/java/androidx/tracing/Trace.java b/tracing/tracing/src/main/java/androidx/tracing/Trace.java
deleted file mode 100644
index d2eef10..0000000
--- a/tracing/tracing/src/main/java/androidx/tracing/Trace.java
+++ /dev/null
@@ -1,302 +0,0 @@
-/*
- * Copyright 2020 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 androidx.tracing;
-
-import android.os.Build;
-import android.util.Log;
-
-import androidx.annotation.NonNull;
-
-import java.lang.reflect.Field;
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
-
-/**
- * Writes trace events to the system trace buffer.
- *
- * <p>These trace events can be collected and visualized using the Android Studio System
- * Trace, Perfetto, and Systrace tools.
- *
- * <p>Tracing should generally be performed in a non-debuggable app for more accurate
- * measurements, representative of real user experience. In a non-debuggable app, tracing is
- * {@link #isEnabled() enabled} if a trace is currently being captured, as well as one of the
- * following:
- * <ul>
- *   <li>Android 12 (API 31) or greater: On by default, unless
- *     <pre>&lt;profileable enabled=false/&gt;</pre>
- *     or <pre>&lt;profileable shell=false/&gt;</pre> is set in the manifest.</li>
- *   <li>Android 10 or 11 (API 29 or 30): <pre>&lt;profileable shell=true/&gt;</pre> is set in the
- *     manifest, or {@link #forceEnableAppTracing()} has been called</li>
- *   <li>JellyBean through Android 11 (API 18 through API 28): {@link #forceEnableAppTracing()} has
- *     been called</li>
- * </ul>
- *
- * <p>This tracing mechanism is independent of the method tracing mechanism offered by
- * {@link android.os.Debug#startMethodTracing}.  In particular, it enables tracing of events that
- * occur across multiple processes.
- *
- * <p>For information see
- * <a href="{@docRoot}studio/profile/systrace/">Overview of system tracing</a>.
- */
-public final class Trace {
-    static final String TAG = "Trace";
-    static final int MAX_TRACE_LABEL_LENGTH = 127;
-
-    private static long sTraceTagApp;
-    private static Method sIsTagEnabledMethod;
-    private static Method sAsyncTraceBeginMethod;
-    private static Method sAsyncTraceEndMethod;
-    private static Method sTraceCounterMethod;
-    private static boolean sHasAppTracingEnabled;
-
-    /**
-     * Checks whether or not tracing is currently enabled.
-     *
-     * <p>This is useful to avoid intermediate string creation for trace sections that require
-     * formatting. It is not necessary to guard all Trace method calls as they internally already
-     * check this. However it is recommended to use this to prevent creating any temporary
-     * objects that would then be passed to those methods to reduce runtime cost when tracing
-     * isn't enabled.
-     *
-     * @return true if tracing is currently enabled, false otherwise
-     */
-    public static boolean isEnabled() {
-        if (Build.VERSION.SDK_INT >= 29) {
-            return TraceApi29Impl.isEnabled();
-        }
-        return isEnabledFallback();
-    }
-
-    /**
-     * Enables the app tracing tag in a non-debuggable process.
-     *
-     * Beginning in Android 12 (API 31), app tracing - custom tracing performed by app code via
-     * this class or android.os.Trace - is always enabled in all apps. Prior to this, app tracing
-     * was only enabled in debuggable apps (as well as profileable apps, on API 29/30).
-     *
-     * Calling this method enables the app to record custom trace content without debuggable=true
-     * on any platform version that supports tracing. Tracing of non-debuggable apps is highly
-     * recommended, to ensure accurate performance measurements.
-     *
-     * As app tracing is always enabled on Android 12 (API 31) and above, this does nothing after
-     * API 31.
-     */
-    public static void forceEnableAppTracing() {
-        if (Build.VERSION.SDK_INT < 31) {
-            try {
-                if (!sHasAppTracingEnabled) {
-                    sHasAppTracingEnabled = true; // only attempt once
-                    @SuppressWarnings("JavaReflectionMemberAccess")
-                    Method setAppTracingAllowed = android.os.Trace.class.getMethod(
-                            "setAppTracingAllowed",
-                            boolean.class
-                    );
-                    setAppTracingAllowed.invoke(null, true);
-                }
-            } catch (Exception exception) {
-                handleException("setAppTracingAllowed", exception);
-            }
-        }
-    }
-
-    /**
-     * Writes a trace message to indicate that a given section of code has begun.
-     *
-     * <p>This call must be followed by a corresponding call to {@link #endSection()} on the same
-     * thread.
-     *
-     * <p class="note"> At this time the vertical bar character '|', newline character '\n', and
-     * null character '\0' are used internally by the tracing mechanism.  If sectionName contains
-     * these characters they will be replaced with a space character in the trace.
-     *
-     * @param label The name of the code section to appear in the trace.
-     */
-    public static void beginSection(@NonNull String label) {
-        android.os.Trace.beginSection(truncatedTraceSectionLabel(label));
-    }
-
-    /**
-     * Writes a trace message to indicate that a given section of code has ended.
-     *
-     * <p>This call must be preceded by a corresponding call to {@link #beginSection(String)}.
-     * Calling this method will mark the end of the most recently begun section of code, so care
-     * must be taken to ensure that beginSection / endSection pairs are properly nested and
-     * called from the same thread.
-     */
-    public static void endSection() {
-        android.os.Trace.endSection();
-    }
-
-    /**
-     * Writes a trace message to indicate that a given section of code has begun.
-     *
-     * <p>Must be followed by a call to {@link #endAsyncSection(String, int)} with the same
-     * methodName and cookie. Unlike {@link #beginSection(String)} and {@link #endSection()},
-     * asynchronous events do not need to be nested. The name and cookie used to begin an event
-     * must be used to end it.
-     *
-     * The cookie must be unique to any overlapping events. If events don't overlap, you can
-     * simply always pass the same integer (e.g. `0`). If they do overlap, the cookie is used to
-     * disambiguate between overlapping events, like the following scenario:
-     * <pre>
-     * [==========================]
-     *           [=====================================]
-     *                                      [====]
-     * </pre>
-     * Without unique cookies, these start/stop timestamps could be misinterpreted by the trace
-     * display like the following, to show very different ranges:
-     * <pre>
-     * [=========================================]
-     *           [================]
-     *                                      [==========]
-     * </pre>
-     *
-     * @param methodName The method name to appear in the trace.
-     * @param cookie     Unique identifier for distinguishing simultaneous events with the same
-     *                   methodName
-     * @see #endAsyncSection
-     */
-    public static void beginAsyncSection(@NonNull String methodName, int cookie) {
-        if (Build.VERSION.SDK_INT >= 29) {
-            TraceApi29Impl.beginAsyncSection(truncatedTraceSectionLabel(methodName), cookie);
-        } else {
-            beginAsyncSectionFallback(truncatedTraceSectionLabel(methodName), cookie);
-        }
-    }
-
-    /**
-     * Writes a trace message to indicate that the current method has ended.
-     *
-     * <p>Must be called exactly once for each call to {@link #beginAsyncSection(String, int)}
-     * using the same name and cookie.
-     *
-     * @param methodName The method name to appear in the trace.
-     * @param cookie     Unique identifier for distinguishing simultaneous events with the same
-     *                   methodName
-     * @see #beginAsyncSection
-     */
-    public static void endAsyncSection(@NonNull String methodName, int cookie) {
-        if (Build.VERSION.SDK_INT >= 29) {
-            TraceApi29Impl.endAsyncSection(truncatedTraceSectionLabel(methodName), cookie);
-        } else {
-            endAsyncSectionFallback(truncatedTraceSectionLabel(methodName), cookie);
-        }
-    }
-
-    /**
-     * Writes trace message to indicate the value of a given counter.
-     *
-     * @param counterName  The counter name to appear in the trace.
-     * @param counterValue The counter value.
-     */
-    public static void setCounter(@NonNull String counterName, int counterValue) {
-        if (Build.VERSION.SDK_INT >= 29) {
-            TraceApi29Impl.setCounter(truncatedTraceSectionLabel(counterName), counterValue);
-        } else {
-            setCounterFallback(truncatedTraceSectionLabel(counterName), counterValue);
-        }
-    }
-
-    @SuppressWarnings({"JavaReflectionMemberAccess", "BanUncheckedReflection"})
-    private static boolean isEnabledFallback() {
-        try {
-            if (sIsTagEnabledMethod == null) {
-                Field traceTagAppField = android.os.Trace.class.getField("TRACE_TAG_APP");
-                sTraceTagApp = traceTagAppField.getLong(null);
-                sIsTagEnabledMethod =
-                        android.os.Trace.class.getMethod("isTagEnabled", long.class);
-            }
-            return (boolean) sIsTagEnabledMethod.invoke(null, sTraceTagApp);
-        } catch (Exception exception) {
-            handleException("isTagEnabled", exception);
-        }
-        // Never enabled on < API 18
-        return false;
-    }
-
-    @SuppressWarnings({"JavaReflectionMemberAccess", "BanUncheckedReflection"})
-    private static void beginAsyncSectionFallback(@NonNull String methodName, int cookie) {
-        try {
-            if (sAsyncTraceBeginMethod == null) {
-                sAsyncTraceBeginMethod = android.os.Trace.class.getMethod(
-                        "asyncTraceBegin",
-                        long.class,
-                        String.class, int.class
-                );
-            }
-            sAsyncTraceBeginMethod.invoke(null, sTraceTagApp, methodName, cookie);
-        } catch (Exception exception) {
-            handleException("asyncTraceBegin", exception);
-        }
-    }
-
-    @SuppressWarnings({"JavaReflectionMemberAccess", "BanUncheckedReflection"})
-    private static void endAsyncSectionFallback(@NonNull String methodName, int cookie) {
-        try {
-            if (sAsyncTraceEndMethod == null) {
-                sAsyncTraceEndMethod = android.os.Trace.class.getMethod(
-                        "asyncTraceEnd",
-                        long.class,
-                        String.class, int.class
-                );
-            }
-            sAsyncTraceEndMethod.invoke(null, sTraceTagApp, methodName, cookie);
-        } catch (Exception exception) {
-            handleException("asyncTraceEnd", exception);
-        }
-    }
-
-    @SuppressWarnings({"JavaReflectionMemberAccess", "BanUncheckedReflection"})
-    private static void setCounterFallback(@NonNull String counterName, int counterValue) {
-        try {
-            if (sTraceCounterMethod == null) {
-                sTraceCounterMethod = android.os.Trace.class.getMethod(
-                        "traceCounter",
-                        long.class,
-                        String.class,
-                        int.class
-                );
-            }
-            sTraceCounterMethod.invoke(null, sTraceTagApp, counterName, counterValue);
-        } catch (Exception exception) {
-            handleException("traceCounter", exception);
-        }
-    }
-
-    private static void handleException(@NonNull String methodName, @NonNull Exception exception) {
-        if (exception instanceof InvocationTargetException) {
-            Throwable cause = exception.getCause();
-            if (cause instanceof RuntimeException) {
-                throw (RuntimeException) cause;
-            } else {
-                throw new RuntimeException(cause);
-            }
-        }
-        Log.v(TAG, "Unable to call " + methodName + " via reflection", exception);
-    }
-
-    @NonNull
-    private static String truncatedTraceSectionLabel(@NonNull String labelName) {
-        if (labelName.length() <= MAX_TRACE_LABEL_LENGTH) {
-            return labelName;
-        }
-        return labelName.substring(0, MAX_TRACE_LABEL_LENGTH);
-    }
-
-    private Trace() {
-    }
-}
diff --git a/tracing/tracing/src/main/java/androidx/tracing/Trace.kt b/tracing/tracing/src/main/java/androidx/tracing/Trace.kt
new file mode 100644
index 0000000..07ab30f
--- /dev/null
+++ b/tracing/tracing/src/main/java/androidx/tracing/Trace.kt
@@ -0,0 +1,398 @@
+/*
+ * Copyright 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 androidx.tracing
+
+import android.os.Build
+import android.os.Trace
+import android.util.Log
+import java.lang.reflect.InvocationTargetException
+import java.lang.reflect.Method
+
+/**
+ * Writes trace events to the system trace buffer.
+ *
+ * These trace events can be collected and visualized using the Android Studio System Trace,
+ * Perfetto, and Systrace tools.
+ *
+ * Tracing should generally be performed in a non-debuggable app for more accurate measurements,
+ * representative of real user experience. In a non-debuggable app, tracing is [enabled][.isEnabled]
+ * if a trace is currently being captured, as well as one of the following:
+ * * Android 12 (API 31) or greater: On by default, unless
+ * <pre>&lt;profileable enabled=false/&gt;</pre>
+ *
+ * or <pre>&lt;profileable shell=false/&gt;</pre> is set in the manifest.
+ * * Android 10 or 11 (API 29 or 30): <pre>&lt;profileable shell=true/&gt;</pre> is set in the
+ *   manifest, or [.forceEnableAppTracing] has been called
+ * * JellyBean through Android 11 (API 18 through API 28): [.forceEnableAppTracing] has been called
+ *
+ * This tracing mechanism is independent of the method tracing mechanism offered by
+ * [android.os.Debug.startMethodTracing]. In particular, it enables tracing of events that occur
+ * across multiple processes.
+ *
+ * For information see [Overview of system tracing]({@docRoot}studio/profile/systrace/).
+ */
+object Trace {
+    private const val TAG: String = "Trace"
+    internal const val MAX_TRACE_LABEL_LENGTH: Int = 127
+
+    private var traceTagApp = 0L
+    private var isTagEnabledMethod: Method? = null
+    private var asyncTraceBeginMethod: Method? = null
+    private var asyncTraceEndMethod: Method? = null
+    private var traceCounterMethod: Method? = null
+    private var hasAppTracingEnabled = false
+
+    /**
+     * Checks whether or not tracing is currently enabled.
+     *
+     * This is useful to avoid intermediate string creation for trace sections that require
+     * formatting. It is not necessary to guard all Trace method calls as they internally already
+     * check this. However it is recommended to use this to prevent creating any temporary objects
+     * that would then be passed to those methods to reduce runtime cost when tracing isn't enabled.
+     *
+     * @return `true` if tracing is currently enabled, `false` otherwise.
+     */
+    @JvmStatic // A function (not a property) for source compatibility with Kotlin callers
+    fun isEnabled(): Boolean =
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
+            TraceApi29Impl.isEnabled
+        } else {
+            isEnabledFallback
+        }
+
+    /**
+     * Enables the app tracing tag in a non-debuggable process.
+     *
+     * Beginning in Android 12 (API 31), app tracing - custom tracing performed by app code via this
+     * class or android.os.Trace - is always enabled in all apps. Prior to this, app tracing was
+     * only enabled in debuggable apps (as well as profileable apps, on API 29/30).
+     *
+     * Calling this method enables the app to record custom trace content without debuggable=true on
+     * any platform version that supports tracing. Tracing of non-debuggable apps is highly
+     * recommended, to ensure accurate performance measurements.
+     *
+     * As app tracing is always enabled on Android 12 (API 31) and above, this does nothing after
+     * API 31.
+     */
+    @JvmStatic
+    fun forceEnableAppTracing() {
+        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.S) {
+            try {
+                if (!hasAppTracingEnabled) {
+                    hasAppTracingEnabled = true // only attempt once
+                    val setAppTracingAllowed =
+                        Trace::class
+                            .java
+                            .getMethod(
+                                "setAppTracingAllowed",
+                                Boolean::class.javaPrimitiveType,
+                            )
+                    setAppTracingAllowed.invoke(null, true)
+                }
+            } catch (exception: Exception) {
+                handleException("setAppTracingAllowed", exception)
+            }
+        }
+    }
+
+    /**
+     * Writes a trace message to indicate that a given section of code has begun.
+     *
+     * This call must be followed by a corresponding call to [endSection] on the same thread.
+     *
+     * At this time the vertical bar character '|', newline character '\n', and null character '\0'
+     * are used internally by the tracing mechanism. If sectionName contains these characters they
+     * will be replaced with a space character in the trace.
+     *
+     * @param label The name of the code section to appear in the trace.
+     */
+    @JvmStatic
+    fun beginSection(label: String) {
+        Trace.beginSection(label.truncatedTraceSectionLabel())
+    }
+
+    /**
+     * Writes a trace message to indicate that a given section of code has ended.
+     *
+     * This call must be preceded by a corresponding call to [beginSection]. Calling this method
+     * will mark the end of the most recently begun section of code, so care must be taken to ensure
+     * that beginSection / endSection pairs are properly nested and called from the same thread.
+     */
+    @JvmStatic
+    fun endSection() {
+        Trace.endSection()
+    }
+
+    /**
+     * Writes a trace message to indicate that a given section of code has begun.
+     *
+     * Must be followed by a call to [endAsyncSection] with the same methodName and cookie. Unlike
+     * [beginSection] and [endSection], asynchronous events do not need to be nested. The name and
+     * cookie used to begin an event must be used to end it.
+     *
+     * The cookie must be unique to any overlapping events. If events don't overlap, you can simply
+     * always pass the same integer (e.g. `0`). If they do overlap, the cookie is used to
+     * disambiguate between overlapping events, like the following scenario:
+     * ```
+     * [==========================]
+     *           [=====================================]
+     *                                      [====]
+     * ```
+     *
+     * Without unique cookies, these start/stop timestamps could be misinterpreted by the trace
+     * display like the following, to show very different ranges:
+     * ```
+     * [=========================================]
+     *           [================]
+     *                                      [==========]
+     * ```
+     *
+     * @param methodName The method name to appear in the trace.
+     * @param cookie Unique identifier for distinguishing simultaneous events with the same
+     *   methodName.
+     * @see endAsyncSection
+     */
+    @JvmStatic
+    fun beginAsyncSection(methodName: String, cookie: Int) {
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
+            TraceApi29Impl.beginAsyncSection(methodName.truncatedTraceSectionLabel(), cookie)
+        } else {
+            beginAsyncSectionFallback(methodName.truncatedTraceSectionLabel(), cookie)
+        }
+    }
+
+    /**
+     * Writes a trace message to indicate that the current method has ended.
+     *
+     * Must be called exactly once for each call to [beginAsyncSection] using the same name and
+     * cookie.
+     *
+     * @param methodName The method name to appear in the trace.
+     * @param cookie Unique identifier for distinguishing simultaneous events with the same
+     *   methodName.
+     * @see beginAsyncSection
+     */
+    @JvmStatic
+    fun endAsyncSection(methodName: String, cookie: Int) {
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
+            TraceApi29Impl.endAsyncSection(methodName.truncatedTraceSectionLabel(), cookie)
+        } else {
+            endAsyncSectionFallback(methodName.truncatedTraceSectionLabel(), cookie)
+        }
+    }
+
+    /**
+     * Writes trace message to indicate the value of a given counter.
+     *
+     * @param counterName The counter name to appear in the trace.
+     * @param counterValue The counter value.
+     */
+    @JvmStatic
+    fun setCounter(counterName: String, counterValue: Int) {
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
+            TraceApi29Impl.setCounter(counterName.truncatedTraceSectionLabel(), counterValue)
+        } else {
+            setCounterFallback(counterName.truncatedTraceSectionLabel(), counterValue)
+        }
+    }
+
+    @Suppress("JavaReflectionMemberAccess", "BanUncheckedReflection")
+    private val isEnabledFallback: Boolean
+        get() {
+            try {
+                if (isTagEnabledMethod == null) {
+                    val traceTagAppField = Trace::class.java.getField("TRACE_TAG_APP")
+                    traceTagApp = traceTagAppField.getLong(null)
+                    isTagEnabledMethod =
+                        Trace::class.java.getMethod("isTagEnabled", Long::class.javaPrimitiveType)
+                }
+                return requireNotNull(isTagEnabledMethod).invoke(null, traceTagApp) as Boolean
+            } catch (exception: Exception) {
+                handleException("isTagEnabled", exception)
+            }
+            // Never enabled on < API 18
+            return false
+        }
+
+    @Suppress("JavaReflectionMemberAccess", "BanUncheckedReflection")
+    private fun beginAsyncSectionFallback(methodName: String, cookie: Int) {
+        try {
+            if (asyncTraceBeginMethod == null) {
+                asyncTraceBeginMethod =
+                    Trace::class
+                        .java
+                        .getMethod(
+                            "asyncTraceBegin",
+                            Long::class.javaPrimitiveType,
+                            String::class.java,
+                            Int::class.javaPrimitiveType,
+                        )
+            }
+            requireNotNull(asyncTraceBeginMethod).invoke(null, traceTagApp, methodName, cookie)
+        } catch (exception: Exception) {
+            handleException("asyncTraceBegin", exception)
+        }
+    }
+
+    @Suppress("JavaReflectionMemberAccess", "BanUncheckedReflection")
+    private fun endAsyncSectionFallback(methodName: String, cookie: Int) {
+        try {
+            if (asyncTraceEndMethod == null) {
+                asyncTraceEndMethod =
+                    Trace::class
+                        .java
+                        .getMethod(
+                            "asyncTraceEnd",
+                            Long::class.javaPrimitiveType,
+                            String::class.java,
+                            Int::class.javaPrimitiveType,
+                        )
+            }
+            requireNotNull(asyncTraceEndMethod).invoke(null, traceTagApp, methodName, cookie)
+        } catch (exception: Exception) {
+            handleException("asyncTraceEnd", exception)
+        }
+    }
+
+    @Suppress("JavaReflectionMemberAccess", "BanUncheckedReflection")
+    private fun setCounterFallback(counterName: String, counterValue: Int) {
+        try {
+            if (traceCounterMethod == null) {
+                traceCounterMethod =
+                    Trace::class
+                        .java
+                        .getMethod(
+                            "traceCounter",
+                            Long::class.javaPrimitiveType,
+                            String::class.java,
+                            Int::class.javaPrimitiveType,
+                        )
+            }
+            requireNotNull(traceCounterMethod).invoke(null, traceTagApp, counterName, counterValue)
+        } catch (exception: Exception) {
+            handleException("traceCounter", exception)
+        }
+    }
+
+    private fun handleException(methodName: String, exception: Exception) {
+        if (exception is InvocationTargetException) {
+            val cause = exception.cause
+            if (cause is RuntimeException) {
+                throw cause
+            } else {
+                throw RuntimeException(cause)
+            }
+        }
+        Log.v(TAG, "Unable to call $methodName via reflection", exception)
+    }
+
+    private fun String.truncatedTraceSectionLabel(): String =
+        takeIf { it.length <= MAX_TRACE_LABEL_LENGTH } ?: substring(0, MAX_TRACE_LABEL_LENGTH)
+}
+
+/**
+ * Wrap the specified [block] in calls to [Trace.beginSection] (with the supplied [label]) and
+ * [Trace.endSection].
+ *
+ * @param label A name of the code section to appear in the trace.
+ * @param block A block of code which is being traced.
+ */
+public inline fun <T> trace(label: String, block: () -> T): T {
+    androidx.tracing.Trace.beginSection(label)
+    try {
+        return block()
+    } finally {
+        androidx.tracing.Trace.endSection()
+    }
+}
+
+/**
+ * Wrap the specified [block] in calls to [Trace.beginSection] (with a lazy-computed [lazyLabel],
+ * only if tracing is enabled - [Trace.isEnabled]) and [Trace.endSection].
+ *
+ * This variant allows you to build a dynamic label, but only when tracing is enabled, avoiding the
+ * cost of String construction otherwise.
+ *
+ * @param lazyLabel A name of the code section to appear in the trace, computed lazily if needed.
+ * @param block A block of code which is being traced.
+ */
+public inline fun <T> trace(lazyLabel: () -> String, block: () -> T): T {
+    val isEnabled = androidx.tracing.Trace.isEnabled()
+    if (isEnabled) {
+        androidx.tracing.Trace.beginSection(lazyLabel())
+    }
+    try {
+        return block()
+    } finally {
+        if (isEnabled) {
+            androidx.tracing.Trace.endSection()
+        }
+    }
+}
+
+/**
+ * Wrap the specified [block] in calls to [Trace.beginAsyncSection] (with the supplied [methodName]
+ * and [cookie]) and [Trace.endAsyncSection].
+ *
+ * @param methodName The method name to appear in the trace.
+ * @param cookie Unique identifier for distinguishing simultaneous events
+ * @param block A code block to be wrapped between [Trace.beginAsyncSection] and
+ *   [Trace.endAsyncSection].
+ */
+public suspend inline fun <T> traceAsync(
+    methodName: String,
+    cookie: Int,
+    crossinline block: suspend () -> T
+): T {
+    androidx.tracing.Trace.beginAsyncSection(methodName, cookie)
+    try {
+        return block()
+    } finally {
+        androidx.tracing.Trace.endAsyncSection(methodName, cookie)
+    }
+}
+
+/**
+ * Wrap the specified [block] in calls to [Trace.beginAsyncSection] and [Trace.endAsyncSection],
+ * with a lazy-computed [lazyMethodName] and [lazyCookie], only if tracing is
+ * enabled - [Trace.isEnabled].
+ *
+ * @param lazyMethodName The method name to appear in the trace, computed lazily if needed.
+ * @param lazyCookie Unique identifier for distinguishing simultaneous events, computed lazily if
+ *   needed.
+ * @param block a code block to be wrapped between [Trace.beginAsyncSection] and
+ *   [Trace.endAsyncSection].
+ */
+public inline fun <T> traceAsync(
+    lazyMethodName: () -> String,
+    lazyCookie: () -> Int,
+    block: () -> T
+): T {
+    var methodName: String? = null
+    var cookie = 0
+    if (androidx.tracing.Trace.isEnabled()) {
+        methodName = lazyMethodName()
+        cookie = lazyCookie()
+        androidx.tracing.Trace.beginAsyncSection(methodName, cookie)
+    }
+    try {
+        return block()
+    } finally {
+        if (methodName != null) {
+            androidx.tracing.Trace.endAsyncSection(methodName, cookie)
+        }
+    }
+}
diff --git a/tracing/tracing/src/main/java/androidx/tracing/TraceApi29Impl.java b/tracing/tracing/src/main/java/androidx/tracing/TraceApi29Impl.java
deleted file mode 100644
index 6c51b18..0000000
--- a/tracing/tracing/src/main/java/androidx/tracing/TraceApi29Impl.java
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * Copyright 2020 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 androidx.tracing;
-
-import androidx.annotation.NonNull;
-import androidx.annotation.RequiresApi;
-
-/**
- * This is a helper class that handles {@link android.os.Trace} functionality in API >= 29.
- * <p>
- * This class is being defined separately to avoid class verification failures.
- * For more information read https://chromium.googlesource
- * .com/chromium/src/build/+/refs/heads/master/android/docs/class_verification_failures
- * .md#understanding-the-reason-for-the-failure
- */
-@RequiresApi(29)
-final class TraceApi29Impl {
-
-    private TraceApi29Impl() {
-        // Does nothing
-    }
-
-    /**
-     * Checks whether or not tracing is currently enabled.
-     */
-    public static boolean isEnabled() {
-        return android.os.Trace.isEnabled();
-    }
-
-    /**
-     * Writes a trace message to indicate that a given section of code has
-     * begun. Must be followed by a call to {@link #endAsyncSection(String, int)} with the same
-     * methodName and cookie.
-     *
-     * @param methodName The method name to appear in the trace.
-     * @param cookie     Unique identifier for distinguishing simultaneous events
-     */
-    public static void beginAsyncSection(@NonNull String methodName, int cookie) {
-        android.os.Trace.beginAsyncSection(methodName, cookie);
-    }
-
-    /**
-     * Writes a trace message to indicate that the current method has ended.
-     * Must be called exactly once for each call to {@link #beginAsyncSection(String, int)}
-     * using the same name and cookie.
-     *
-     * @param methodName The method name to appear in the trace.
-     * @param cookie     Unique identifier for distinguishing simultaneous events
-     */
-    public static void endAsyncSection(@NonNull String methodName, int cookie) {
-        android.os.Trace.endAsyncSection(methodName, cookie);
-    }
-
-    /**
-     * Writes trace message to indicate the value of a given counter.
-     *
-     * @param counterName  The counter name to appear in the trace.
-     * @param counterValue The counter value.
-     */
-    public static void setCounter(@NonNull String counterName, int counterValue) {
-        android.os.Trace.setCounter(counterName, counterValue);
-    }
-}
diff --git a/tracing/tracing/src/main/java/androidx/tracing/TraceApi29Impl.kt b/tracing/tracing/src/main/java/androidx/tracing/TraceApi29Impl.kt
new file mode 100644
index 0000000..00e7831
--- /dev/null
+++ b/tracing/tracing/src/main/java/androidx/tracing/TraceApi29Impl.kt
@@ -0,0 +1,67 @@
+/*
+ * Copyright 2020 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 androidx.tracing
+
+import android.os.Build
+import androidx.annotation.RequiresApi
+
+/**
+ * This is a helper class that handles [android.os.Trace] functionality in API >= 29.
+ *
+ * This class is being defined separately to avoid class verification failures. For more information
+ * read https://chromium.googlesource
+ * .com/chromium/src/build/+/refs/heads/master/android/docs/class_verification_failures
+ * .md#understanding-the-reason-for-the-failure
+ */
+@RequiresApi(Build.VERSION_CODES.Q)
+internal object TraceApi29Impl {
+
+    /** Checks whether or not tracing is currently enabled. */
+    val isEnabled: Boolean
+        get() = android.os.Trace.isEnabled()
+
+    /**
+     * Writes a trace message to indicate that a given section of code has begun. Must be followed
+     * by a call to [endAsyncSection] with the same methodName and cookie.
+     *
+     * @param methodName The method name to appear in the trace.
+     * @param cookie Unique identifier for distinguishing simultaneous events.
+     */
+    fun beginAsyncSection(methodName: String, cookie: Int) {
+        android.os.Trace.beginAsyncSection(methodName, cookie)
+    }
+
+    /**
+     * Writes a trace message to indicate that the current method has ended. Must be called exactly
+     * once for each call to [beginAsyncSection] using the same name and cookie.
+     *
+     * @param methodName The method name to appear in the trace.
+     * @param cookie Unique identifier for distinguishing simultaneous events.
+     */
+    fun endAsyncSection(methodName: String, cookie: Int) {
+        android.os.Trace.endAsyncSection(methodName, cookie)
+    }
+
+    /**
+     * Writes trace message to indicate the value of a given counter.
+     *
+     * @param counterName The counter name to appear in the trace.
+     * @param counterValue The counter value.
+     */
+    fun setCounter(counterName: String, counterValue: Int) {
+        android.os.Trace.setCounter(counterName, counterValue.toLong())
+    }
+}