Merge changes from topic "perfetto_test" into main

* changes:
  Added Perfetto tests to TracePerfTest
  Expose perfetto start and stop sessions

10 files changed, 174 insertions, 192 deletions

diff --git a/apct-tests/perftests/core/Android.bp b/apct-tests/perftests/core/Android.bp
index 1e299cdf8002..f16f2caccd49 100644
--- a/apct-tests/perftests/core/Android.bp
+++ b/apct-tests/perftests/core/Android.bp
@@ -50,6 +50,7 @@ android_test {
         "junit-params",
         "core-tests-support",
         "guava",
+        "perfetto_trace_java_protos",
     ],
 
     libs: ["android.test.base.stubs.system"],
diff --git a/apct-tests/perftests/core/src/android/os/TracePerfTest.java b/apct-tests/perftests/core/src/android/os/TracePerfTest.java
index 0d64c390f4c2..bf7c96a3cb85 100644
--- a/apct-tests/perftests/core/src/android/os/TracePerfTest.java
+++ b/apct-tests/perftests/core/src/android/os/TracePerfTest.java
@@ -17,6 +17,8 @@
 
 package android.os;
 
+import static android.os.PerfettoTrace.Category;
+
 import android.perftests.utils.BenchmarkState;
 import android.perftests.utils.PerfStatusReporter;
 import android.perftests.utils.ShellHelper;
@@ -31,19 +33,35 @@ import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import perfetto.protos.DataSourceConfigOuterClass.DataSourceConfig;
+import perfetto.protos.TraceConfigOuterClass.TraceConfig;
+import perfetto.protos.TraceConfigOuterClass.TraceConfig.BufferConfig;
+import perfetto.protos.TraceConfigOuterClass.TraceConfig.DataSource;
+import perfetto.protos.TrackEventConfigOuterClass.TrackEventConfig;
+
 @RunWith(AndroidJUnit4.class)
 public class TracePerfTest {
     @Rule
     public PerfStatusReporter mPerfStatusReporter = new PerfStatusReporter();
 
+    private static final String FOO = "foo";
+    private static final Category FOO_CATEGORY = new Category(FOO);
+    private static PerfettoTrace.Session sPerfettoSession;
+
     @BeforeClass
     public static void startTracing() {
         ShellHelper.runShellCommandRaw("atrace -c --async_start -a *");
+        PerfettoTrace.register(false /* isBackendInProcess */);
+        FOO_CATEGORY.register();
+        sPerfettoSession = new PerfettoTrace.Session(false /* isBackendInProcess */,
+                                                      getTraceConfig(FOO).toByteArray());
     }
 
     @AfterClass
     public static void endTracing() {
         ShellHelper.runShellCommandRaw("atrace --async_stop");
+        FOO_CATEGORY.unregister();
+        sPerfettoSession.close();
     }
 
     @Before
@@ -84,4 +102,61 @@ public class TracePerfTest {
             Trace.setCounter("testCounter", 123);
         }
     }
+
+    @Test
+    public void testInstant() {
+        Trace.instant(Trace.TRACE_TAG_APP, "testInstantA");
+
+        BenchmarkState state = mPerfStatusReporter.getBenchmarkState();
+        while (state.keepRunning()) {
+            Trace.instant(Trace.TRACE_TAG_APP, "testInstantA");
+        }
+    }
+
+    @Test
+    public void testInstantPerfetto() {
+        PerfettoTrace.instant(FOO_CATEGORY, "testInstantP").emit();
+
+        BenchmarkState state = mPerfStatusReporter.getBenchmarkState();
+        while (state.keepRunning()) {
+            PerfettoTrace.instant(FOO_CATEGORY, "testInstantP").emit();
+        }
+    }
+
+    @Test
+    public void testInstantPerfettoWithArgs() {
+        PerfettoTrace.instant(FOO_CATEGORY, "testInstantP")
+                .addArg("foo", "bar")
+                .addFlow(1)
+                .emit();
+
+        BenchmarkState state = mPerfStatusReporter.getBenchmarkState();
+        while (state.keepRunning()) {
+            PerfettoTrace.instant(FOO_CATEGORY, "testInstantP")
+                    .addArg("foo", "bar")
+                    .addFlow(1)
+                    .emit();
+        }
+    }
+
+    private static TraceConfig getTraceConfig(String cat) {
+        BufferConfig bufferConfig = BufferConfig.newBuilder().setSizeKb(1024).build();
+        TrackEventConfig trackEventConfig = TrackEventConfig
+                .newBuilder()
+                .addEnabledCategories(cat)
+                .build();
+        DataSourceConfig dsConfig = DataSourceConfig
+                .newBuilder()
+                .setName("track_event")
+                .setTargetBuffer(0)
+                .setTrackEventConfig(trackEventConfig)
+                .build();
+        DataSource ds = DataSource.newBuilder().setConfig(dsConfig).build();
+        TraceConfig traceConfig = TraceConfig
+                .newBuilder()
+                .addBuffers(bufferConfig)
+                .addDataSources(ds)
+                .build();
+        return traceConfig;
+    }
 }
diff --git a/core/java/android/os/PerfettoTrace.java b/core/java/android/os/PerfettoTrace.java
index e3f251e34b45..68f1570154ff 100644
--- a/core/java/android/os/PerfettoTrace.java
+++ b/core/java/android/os/PerfettoTrace.java
@@ -154,14 +154,44 @@ public final class PerfettoTrace {
         }
     }
 
+    /**
+     * Manages a perfetto tracing session.
+     * Constructing this object with a config automatically starts a tracing session. Each session
+     * must be closed after use and then the resulting trace bytes can be read.
+     *
+     * The session could be in process or system wide, depending on {@code isBackendInProcess}.
+     * This functionality is intended for testing.
+     */
+    public static final class Session {
+        private final long mPtr;
+
+        /**
+         * Session ctor.
+         */
+        public Session(boolean isBackendInProcess, byte[] config) {
+            mPtr = native_start_session(isBackendInProcess, config);
+        }
+
+        /**
+         * Closes the session and returns the trace.
+         */
+        public byte[] close() {
+            return native_stop_session(mPtr);
+        }
+    }
+
     @CriticalNative
     private static native long native_get_process_track_uuid();
-
     @CriticalNative
     private static native long native_get_thread_track_uuid(long tid);
 
     @FastNative
     private static native void native_activate_trigger(String name, int ttlMs);
+    @FastNative
+    private static native void native_register(boolean isBackendInProcess);
+
+    private static native long native_start_session(boolean isBackendInProcess, byte[] config);
+    private static native byte[] native_stop_session(long ptr);
 
     /**
      * Writes a trace message to indicate a given section of code was invoked.
@@ -307,7 +337,7 @@ public final class PerfettoTrace {
     /**
      * Registers the process with Perfetto.
      */
-    public static void register() {
-        Trace.registerWithPerfetto();
+    public static void register(boolean isBackendInProcess) {
+        native_register(isBackendInProcess);
     }
 }
diff --git a/core/java/android/os/Trace.java b/core/java/android/os/Trace.java
index 4a37e0a70443..09e6a45dc294 100644
--- a/core/java/android/os/Trace.java
+++ b/core/java/android/os/Trace.java
@@ -164,8 +164,6 @@ public final class Trace {
     private static native void nativeInstant(long tag, String name);
     @FastNative
     private static native void nativeInstantForTrack(long tag, String trackName, String name);
-    @FastNative
-    private static native void nativeRegisterWithPerfetto();
 
     private Trace() {
     }
@@ -545,6 +543,6 @@ public final class Trace {
      * @hide
      */
     public static void registerWithPerfetto() {
-        nativeRegisterWithPerfetto();
+        PerfettoTrace.register(false /* isBackendInProcess */);
     }
 }
diff --git a/core/jni/android_os_PerfettoTrace.cpp b/core/jni/android_os_PerfettoTrace.cpp
index 962aefc482e4..9bedfa27fa1a 100644
--- a/core/jni/android_os_PerfettoTrace.cpp
+++ b/core/jni/android_os_PerfettoTrace.cpp
@@ -24,9 +24,12 @@
 #include <nativehelper/scoped_primitive_array.h>
 #include <nativehelper/scoped_utf_chars.h>
 #include <nativehelper/utils.h>
+#include <tracing_perfetto.h>
 #include <tracing_sdk.h>
 
 namespace android {
+constexpr int kFlushTimeoutMs = 5000;
+
 template <typename T>
 inline static T* toPointer(jlong ptr) {
     return reinterpret_cast<T*>(static_cast<uintptr_t>(ptr));
@@ -51,6 +54,10 @@ static void android_os_PerfettoTrace_activate_trigger(JNIEnv* env, jclass, jstri
     tracing_perfetto::activate_trigger(name_chars.c_str(), static_cast<uint32_t>(ttl_ms));
 }
 
+void android_os_PerfettoTrace_register(bool is_backend_in_process) {
+    tracing_perfetto::registerWithPerfetto(is_backend_in_process);
+}
+
 static jlong android_os_PerfettoTraceCategory_init(JNIEnv* env, jclass, jstring name, jstring tag,
                                                    jstring severity) {
     ScopedUtfChars name_chars = GET_UTF_OR_RETURN(env, name);
@@ -85,6 +92,36 @@ static jlong android_os_PerfettoTraceCategory_get_extra_ptr(jlong ptr) {
     return toJLong(category->get());
 }
 
+static jlong android_os_PerfettoTrace_start_session(JNIEnv* env, jclass /* obj */,
+                                                    jboolean is_backend_in_process,
+                                                    jbyteArray config_bytes) {
+    jsize length = env->GetArrayLength(config_bytes);
+    std::vector<uint8_t> data;
+    data.reserve(length);
+    env->GetByteArrayRegion(config_bytes, 0, length, reinterpret_cast<jbyte*>(data.data()));
+
+    tracing_perfetto::Session* session =
+            new tracing_perfetto::Session(is_backend_in_process, data.data(), length);
+
+    return reinterpret_cast<long>(session);
+}
+
+static jbyteArray android_os_PerfettoTrace_stop_session([[maybe_unused]] JNIEnv* env,
+                                                        jclass /* obj */, jlong ptr) {
+    tracing_perfetto::Session* session = reinterpret_cast<tracing_perfetto::Session*>(ptr);
+
+    session->FlushBlocking(kFlushTimeoutMs);
+    session->StopBlocking();
+
+    std::vector<uint8_t> data = session->ReadBlocking();
+
+    delete session;
+
+    jbyteArray bytes = env->NewByteArray(data.size());
+    env->SetByteArrayRegion(bytes, 0, data.size(), reinterpret_cast<jbyte*>(data.data()));
+    return bytes;
+}
+
 static const JNINativeMethod gCategoryMethods[] = {
         {"native_init", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)J",
          (void*)android_os_PerfettoTraceCategory_init},
@@ -101,7 +138,10 @@ static const JNINativeMethod gTraceMethods[] =
          {"native_get_thread_track_uuid", "(J)J",
           (void*)android_os_PerfettoTrace_get_thread_track_uuid},
          {"native_activate_trigger", "(Ljava/lang/String;I)V",
-          (void*)android_os_PerfettoTrace_activate_trigger}};
+          (void*)android_os_PerfettoTrace_activate_trigger},
+         {"native_register", "(Z)V", (void*)android_os_PerfettoTrace_register},
+         {"native_start_session", "(Z[B)J", (void*)android_os_PerfettoTrace_start_session},
+         {"native_stop_session", "(J)[B", (void*)android_os_PerfettoTrace_stop_session}};
 
 int register_android_os_PerfettoTrace(JNIEnv* env) {
     int res = jniRegisterNativeMethods(env, "android/os/PerfettoTrace", gTraceMethods,
diff --git a/core/jni/android_os_Trace.cpp b/core/jni/android_os_Trace.cpp
index 21e056dfa12b..50618c5a07af 100644
--- a/core/jni/android_os_Trace.cpp
+++ b/core/jni/android_os_Trace.cpp
@@ -131,10 +131,6 @@ static jboolean android_os_Trace_nativeIsTagEnabled(jlong tag) {
     return tracing_perfetto::isTagEnabled(tag);
 }
 
-static void android_os_Trace_nativeRegisterWithPerfetto(JNIEnv* env) {
-    tracing_perfetto::registerWithPerfetto();
-}
-
 static const JNINativeMethod gTraceMethods[] = {
         /* name, signature, funcPtr */
         {"nativeSetAppTracingAllowed", "(Z)V", (void*)android_os_Trace_nativeSetAppTracingAllowed},
@@ -157,7 +153,6 @@ static const JNINativeMethod gTraceMethods[] = {
         {"nativeInstant", "(JLjava/lang/String;)V", (void*)android_os_Trace_nativeInstant},
         {"nativeInstantForTrack", "(JLjava/lang/String;Ljava/lang/String;)V",
          (void*)android_os_Trace_nativeInstantForTrack},
-        {"nativeRegisterWithPerfetto", "()V", (void*)android_os_Trace_nativeRegisterWithPerfetto},
 
         // ----------- @CriticalNative  ----------------
         {"nativeIsTagEnabled", "(J)Z", (void*)android_os_Trace_nativeIsTagEnabled},
diff --git a/core/tests/coretests/Android.bp b/core/tests/coretests/Android.bp
index 1b6746ca4b63..c06ad64cc0f5 100644
--- a/core/tests/coretests/Android.bp
+++ b/core/tests/coretests/Android.bp
@@ -122,7 +122,6 @@ android_test {
         "android.view.flags-aconfig-java",
     ],
     jni_libs: [
-        "libperfetto_trace_test_jni",
         "libpowermanagertest_jni",
         "libviewRootImplTest_jni",
         "libworksourceparceltest_jni",
diff --git a/core/tests/coretests/jni/Android.bp b/core/tests/coretests/jni/Android.bp
index 798ec90eb884..d6379ca8c3e6 100644
--- a/core/tests/coretests/jni/Android.bp
+++ b/core/tests/coretests/jni/Android.bp
@@ -111,27 +111,3 @@ cc_test_library {
     ],
     gtest: false,
 }
-
-cc_test_library {
-    name: "libperfetto_trace_test_jni",
-    srcs: [
-        "PerfettoTraceTest.cpp",
-    ],
-    static_libs: [
-        "perfetto_trace_protos",
-        "libtracing_perfetto_test_utils",
-    ],
-    shared_libs: [
-        "liblog",
-        "libnativehelper",
-        "libperfetto_c",
-        "libprotobuf-cpp-lite",
-        "libtracing_perfetto",
-    ],
-    stl: "libc++_static",
-    cflags: [
-        "-Werror",
-        "-Wall",
-    ],
-    gtest: false,
-}
diff --git a/core/tests/coretests/jni/PerfettoTraceTest.cpp b/core/tests/coretests/jni/PerfettoTraceTest.cpp
deleted file mode 100644
index 41d02ed70c9a..000000000000
--- a/core/tests/coretests/jni/PerfettoTraceTest.cpp
+++ /dev/null
@@ -1,117 +0,0 @@
-/*
- * Copyright (C) 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.
- */
-
-// #define LOG_NDEBUG 0
-#define LOG_TAG "PerfettoTraceTest"
-
-#include <nativehelper/JNIHelp.h>
-#include <utils/Log.h>
-
-#include "jni.h"
-#include "perfetto/public/abi/data_source_abi.h"
-#include "perfetto/public/abi/heap_buffer.h"
-#include "perfetto/public/abi/pb_decoder_abi.h"
-#include "perfetto/public/abi/tracing_session_abi.h"
-#include "perfetto/public/abi/track_event_abi.h"
-#include "perfetto/public/compiler.h"
-#include "perfetto/public/data_source.h"
-#include "perfetto/public/pb_decoder.h"
-#include "perfetto/public/producer.h"
-#include "perfetto/public/protos/config/trace_config.pzc.h"
-#include "perfetto/public/protos/trace/interned_data/interned_data.pzc.h"
-#include "perfetto/public/protos/trace/test_event.pzc.h"
-#include "perfetto/public/protos/trace/trace.pzc.h"
-#include "perfetto/public/protos/trace/trace_packet.pzc.h"
-#include "perfetto/public/protos/trace/track_event/debug_annotation.pzc.h"
-#include "perfetto/public/protos/trace/track_event/track_descriptor.pzc.h"
-#include "perfetto/public/protos/trace/track_event/track_event.pzc.h"
-#include "perfetto/public/protos/trace/trigger.pzc.h"
-#include "perfetto/public/te_category_macros.h"
-#include "perfetto/public/te_macros.h"
-#include "perfetto/public/track_event.h"
-#include "protos/perfetto/trace/interned_data/interned_data.pb.h"
-#include "protos/perfetto/trace/trace.pb.h"
-#include "protos/perfetto/trace/trace_packet.pb.h"
-#include "tracing_perfetto.h"
-#include "utils.h"
-
-namespace android {
-using ::perfetto::protos::EventCategory;
-using ::perfetto::protos::EventName;
-using ::perfetto::protos::FtraceEvent;
-using ::perfetto::protos::FtraceEventBundle;
-using ::perfetto::protos::InternedData;
-using ::perfetto::protos::Trace;
-using ::perfetto::protos::TracePacket;
-
-using ::perfetto::shlib::test_utils::TracingSession;
-
-struct TracingSessionHolder {
-    TracingSession tracing_session;
-};
-
-static void nativeRegisterPerfetto([[maybe_unused]] JNIEnv* env, jclass /* obj */) {
-    tracing_perfetto::registerWithPerfetto(false /* test */);
-}
-
-static jlong nativeStartTracing(JNIEnv* env, jclass /* obj */, jbyteArray configBytes) {
-    jsize length = env->GetArrayLength(configBytes);
-    std::vector<uint8_t> data;
-    data.reserve(length);
-    env->GetByteArrayRegion(configBytes, 0, length, reinterpret_cast<jbyte*>(data.data()));
-
-    TracingSession session = TracingSession::FromBytes(data.data(), length);
-    TracingSessionHolder* holder = new TracingSessionHolder(std::move(session));
-
-    return reinterpret_cast<long>(holder);
-}
-
-static jbyteArray nativeStopTracing([[maybe_unused]] JNIEnv* env, jclass /* obj */, jlong ptr) {
-    TracingSessionHolder* holder = reinterpret_cast<TracingSessionHolder*>(ptr);
-
-    // Stop
-    holder->tracing_session.FlushBlocking(5000);
-    holder->tracing_session.StopBlocking();
-
-    std::vector<uint8_t> data = holder->tracing_session.ReadBlocking();
-
-    delete holder;
-
-    jbyteArray bytes = env->NewByteArray(data.size());
-    env->SetByteArrayRegion(bytes, 0, data.size(), reinterpret_cast<jbyte*>(data.data()));
-    return bytes;
-}
-
-extern "C" jint JNI_OnLoad(JavaVM* vm, void* /* reserved */) {
-    JNIEnv* env;
-    const JNINativeMethod methodTable[] = {/* name, signature, funcPtr */
-                                           {"nativeStartTracing", "([B)J",
-                                            (void*)nativeStartTracing},
-                                           {"nativeStopTracing", "(J)[B", (void*)nativeStopTracing},
-                                           {"nativeRegisterPerfetto", "()V",
-                                            (void*)nativeRegisterPerfetto}};
-
-    if (vm->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6) != JNI_OK) {
-        return JNI_ERR;
-    }
-
-    jniRegisterNativeMethods(env, "android/os/PerfettoTraceTest", methodTable,
-                             sizeof(methodTable) / sizeof(JNINativeMethod));
-
-    return JNI_VERSION_1_6;
-}
-
-} /* namespace android */
diff --git a/core/tests/coretests/src/android/os/PerfettoTraceTest.java b/core/tests/coretests/src/android/os/PerfettoTraceTest.java
index ad28383689af..0b5a44665d2b 100644
--- a/core/tests/coretests/src/android/os/PerfettoTraceTest.java
+++ b/core/tests/coretests/src/android/os/PerfettoTraceTest.java
@@ -28,7 +28,6 @@ import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.util.ArraySet;
-import android.util.Log;
 
 import androidx.test.InstrumentationRegistry;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -84,19 +83,9 @@ public class PerfettoTraceTest {
     private final Set<String> mDebugAnnotationNames = new ArraySet<>();
     private final Set<String> mTrackNames = new ArraySet<>();
 
-    static {
-        try {
-            System.loadLibrary("perfetto_trace_test_jni");
-            Log.i(TAG, "Successfully loaded trace_test native library");
-        } catch (UnsatisfiedLinkError ule) {
-            Log.w(TAG, "Could not load trace_test native library");
-        }
-    }
-
     @Before
     public void setUp() {
-        PerfettoTrace.register();
-        nativeRegisterPerfetto();
+        PerfettoTrace.register(true);
         FOO_CATEGORY.register();
 
         mCategoryNames.clear();
@@ -110,7 +99,7 @@ public class PerfettoTraceTest {
     public void testDebugAnnotations() throws Exception {
         TraceConfig traceConfig = getTraceConfig(FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.instant(FOO_CATEGORY, "event")
                 .addFlow(2)
@@ -121,7 +110,7 @@ public class PerfettoTraceTest {
                 .addArg("string_val", FOO)
                 .emit();
 
-        byte[] traceBytes = nativeStopTracing(ptr);
+        byte[] traceBytes = session.close();
 
         Trace trace = Trace.parseFrom(traceBytes);
 
@@ -165,11 +154,11 @@ public class PerfettoTraceTest {
     public void testDebugAnnotationsWithLambda() throws Exception {
         TraceConfig traceConfig = getTraceConfig(FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.instant(FOO_CATEGORY, "event").addArg("long_val", 123L).emit();
 
-        byte[] traceBytes = nativeStopTracing(ptr);
+        byte[] traceBytes = session.close();
 
         Trace trace = Trace.parseFrom(traceBytes);
 
@@ -200,7 +189,7 @@ public class PerfettoTraceTest {
     public void testNamedTrack() throws Exception {
         TraceConfig traceConfig = getTraceConfig(FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.begin(FOO_CATEGORY, "event")
                 .usingNamedTrack(PerfettoTrace.getProcessTrackUuid(), FOO)
@@ -211,7 +200,7 @@ public class PerfettoTraceTest {
                 .usingNamedTrack(PerfettoTrace.getThreadTrackUuid(Process.myTid()), "bar")
                 .emit();
 
-        Trace trace = Trace.parseFrom(nativeStopTracing(ptr));
+        Trace trace = Trace.parseFrom(session.close());
 
         boolean hasTrackEvent = false;
         boolean hasTrackUuid = false;
@@ -248,7 +237,7 @@ public class PerfettoTraceTest {
     public void testProcessThreadNamedTrack() throws Exception {
         TraceConfig traceConfig = getTraceConfig(FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.begin(FOO_CATEGORY, "event")
                 .usingProcessNamedTrack(FOO)
@@ -259,7 +248,7 @@ public class PerfettoTraceTest {
                 .usingThreadNamedTrack(Process.myTid(), "%s-%s", "bar", "stool")
                 .emit();
 
-        Trace trace = Trace.parseFrom(nativeStopTracing(ptr));
+        Trace trace = Trace.parseFrom(session.close());
 
         boolean hasTrackEvent = false;
         boolean hasTrackUuid = false;
@@ -296,13 +285,13 @@ public class PerfettoTraceTest {
     public void testCounterSimple() throws Exception {
         TraceConfig traceConfig = getTraceConfig(FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.counter(FOO_CATEGORY, 16, FOO).emit();
 
         PerfettoTrace.counter(FOO_CATEGORY, 3.14, "bar").emit();
 
-        Trace trace = Trace.parseFrom(nativeStopTracing(ptr));
+        Trace trace = Trace.parseFrom(session.close());
 
         boolean hasTrackEvent = false;
         boolean hasCounterValue = false;
@@ -339,7 +328,7 @@ public class PerfettoTraceTest {
     public void testCounter() throws Exception {
         TraceConfig traceConfig = getTraceConfig(FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.counter(FOO_CATEGORY, 16)
                 .usingCounterTrack(PerfettoTrace.getProcessTrackUuid(), FOO).emit();
@@ -348,7 +337,7 @@ public class PerfettoTraceTest {
                 .usingCounterTrack(PerfettoTrace.getThreadTrackUuid(Process.myTid()),
                                    "%s-%s", "bar", "stool").emit();
 
-        Trace trace = Trace.parseFrom(nativeStopTracing(ptr));
+        Trace trace = Trace.parseFrom(session.close());
 
         boolean hasTrackEvent = false;
         boolean hasCounterValue = false;
@@ -385,14 +374,14 @@ public class PerfettoTraceTest {
     public void testProcessThreadCounter() throws Exception {
         TraceConfig traceConfig = getTraceConfig(FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.counter(FOO_CATEGORY, 16).usingProcessCounterTrack(FOO).emit();
 
         PerfettoTrace.counter(FOO_CATEGORY, 3.14)
                 .usingThreadCounterTrack(Process.myTid(), "%s-%s", "bar", "stool").emit();
 
-        Trace trace = Trace.parseFrom(nativeStopTracing(ptr));
+        Trace trace = Trace.parseFrom(session.close());
 
         boolean hasTrackEvent = false;
         boolean hasCounterValue = false;
@@ -429,7 +418,7 @@ public class PerfettoTraceTest {
     public void testProto() throws Exception {
         TraceConfig traceConfig = getTraceConfig(FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.instant(FOO_CATEGORY, "event_proto")
                 .beginProto()
@@ -441,7 +430,7 @@ public class PerfettoTraceTest {
                 .endProto()
                 .emit();
 
-        byte[] traceBytes = nativeStopTracing(ptr);
+        byte[] traceBytes = session.close();
 
         Trace trace = Trace.parseFrom(traceBytes);
 
@@ -477,7 +466,7 @@ public class PerfettoTraceTest {
     public void testProtoNested() throws Exception {
         TraceConfig traceConfig = getTraceConfig(FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.instant(FOO_CATEGORY, "event_proto_nested")
                 .beginProto()
@@ -494,7 +483,7 @@ public class PerfettoTraceTest {
                 .endProto()
                 .emit();
 
-        byte[] traceBytes = nativeStopTracing(ptr);
+        byte[] traceBytes = session.close();
 
         Trace trace = Trace.parseFrom(traceBytes);
 
@@ -538,13 +527,13 @@ public class PerfettoTraceTest {
     public void testActivateTrigger() throws Exception {
         TraceConfig traceConfig = getTriggerTraceConfig(FOO, FOO);
 
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.instant(FOO_CATEGORY, "event_trigger").emit();
 
         PerfettoTrace.activateTrigger(FOO, 1000);
 
-        byte[] traceBytes = nativeStopTracing(ptr);
+        byte[] traceBytes = session.close();
 
         Trace trace = Trace.parseFrom(traceBytes);
 
@@ -569,7 +558,7 @@ public class PerfettoTraceTest {
         TraceConfig traceConfig = getTraceConfig(BAR);
 
         Category barCategory = new Category(BAR);
-        long ptr = nativeStartTracing(traceConfig.toByteArray());
+        PerfettoTrace.Session session = new PerfettoTrace.Session(true, traceConfig.toByteArray());
 
         PerfettoTrace.instant(barCategory, "event")
                 .addArg("before", 1)
@@ -581,7 +570,7 @@ public class PerfettoTraceTest {
                 .addArg("after", 1)
                 .emit();
 
-        byte[] traceBytes = nativeStopTracing(ptr);
+        byte[] traceBytes = session.close();
 
         Trace trace = Trace.parseFrom(traceBytes);
 
@@ -603,10 +592,6 @@ public class PerfettoTraceTest {
         assertThat(mDebugAnnotationNames).doesNotContain("before");
     }
 
-    private static native long nativeStartTracing(byte[] config);
-    private static native void nativeRegisterPerfetto();
-    private static native byte[] nativeStopTracing(long ptr);
-
     private TrackEvent getTrackEvent(Trace trace, int idx) {
         int curIdx = 0;
         for (TracePacket packet: trace.getPacketList()) {