Merge "Add tool for injecting tracing code into a method." am: a7006818d3 am: 3e53e056ee

Original change: https://android-review.googlesource.com/c/platform/frameworks/base/+/2016118

Change-Id: If83cc7c52c09e0a226e759755e60ccc9a3412951

8 files changed, 715 insertions, 0 deletions

diff --git a/tools/traceinjection/Android.bp b/tools/traceinjection/Android.bp
new file mode 100644
index 000000000000..1395c5f2e635
--- /dev/null
+++ b/tools/traceinjection/Android.bp
@@ -0,0 +1,49 @@
+package {
+    // See: http://go/android-license-faq
+    // A large-scale-change added 'default_applicable_licenses' to import
+    // all of the 'license_kinds' from "frameworks_base_license"
+    // to get the below license kinds:
+    //   SPDX-license-identifier-Apache-2.0
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+java_binary_host {
+    name: "traceinjection",
+    manifest: "manifest.txt",
+    srcs: ["src/**/*.java"],
+    static_libs: [
+        "asm-7.0",
+        "asm-commons-7.0",
+        "asm-tree-7.0",
+        "asm-analysis-7.0",
+        "guava-21.0",
+    ],
+}
+
+java_library_host {
+    name: "TraceInjectionTests-Uninjected",
+    srcs: ["test/**/*.java"],
+    static_libs: [
+        "junit",
+    ],
+}
+
+java_genrule_host {
+    name: "TraceInjectionTests-Injected",
+    srcs: [":TraceInjectionTests-Uninjected"],
+    tools: ["traceinjection"],
+    cmd: "$(location traceinjection) " +
+        "  --annotation \"com/android/traceinjection/Trace\"" +
+        "  --start \"com/android/traceinjection/InjectionTests.traceStart\"" +
+        "  --end \"com/android/traceinjection/InjectionTests.traceEnd\"" +
+        "  -o $(out) " +
+        "  -i $(in)",
+    out: ["TraceInjectionTests-Injected.jar"],
+}
+
+java_test_host {
+    name: "TraceInjectionTests",
+    static_libs: [
+        "TraceInjectionTests-Injected",
+    ],
+}
diff --git a/tools/traceinjection/manifest.txt b/tools/traceinjection/manifest.txt
new file mode 100644
index 000000000000..7f4ee1d617fa
--- /dev/null
+++ b/tools/traceinjection/manifest.txt
@@ -0,0 +1 @@
+Main-Class: com.android.traceinjection.Main
diff --git a/tools/traceinjection/src/com/android/traceinjection/Main.java b/tools/traceinjection/src/com/android/traceinjection/Main.java
new file mode 100644
index 000000000000..190df819dd64
--- /dev/null
+++ b/tools/traceinjection/src/com/android/traceinjection/Main.java
@@ -0,0 +1,121 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.traceinjection;
+
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassWriter;
+
+import java.io.BufferedInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.Enumeration;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipFile;
+import java.util.zip.ZipOutputStream;
+
+public class Main {
+    public static void main(String[] args) throws IOException {
+        String inJar = null;
+        String outJar = null;
+        String annotation = null;
+        String traceStart = null;
+        String traceEnd = null;
+
+        // All arguments require a value currently, so just make sure we have an even number and
+        // then process them all two at a time.
+        if (args.length % 2 != 0) {
+            throw new IllegalArgumentException("Argument is missing corresponding value");
+        }
+        for (int i = 0; i < args.length - 1; i += 2) {
+            final String arg = args[i].trim();
+            final String argValue = args[i + 1].trim();
+            if ("-i".equals(arg)) {
+                inJar = argValue;
+            } else if ("-o".equals(arg)) {
+                outJar = argValue;
+            } else if ("--annotation".equals(arg)) {
+                annotation = argValue;
+            } else if ("--start".equals(arg)) {
+                traceStart = argValue;
+            } else if ("--end".equals(arg)) {
+                traceEnd = argValue;
+            } else {
+                throw new IllegalArgumentException("Unknown argument: " + arg);
+            }
+        }
+
+        if (inJar == null) {
+            throw new IllegalArgumentException("input jar is required");
+        }
+
+        if (outJar == null) {
+            throw new IllegalArgumentException("output jar is required");
+        }
+
+        if (annotation == null) {
+            throw new IllegalArgumentException("trace annotation is required");
+        }
+
+        if (traceStart == null) {
+            throw new IllegalArgumentException("start trace method is required");
+        }
+
+        if (traceEnd == null) {
+            throw new IllegalArgumentException("end trace method is required");
+        }
+
+        TraceInjectionConfiguration params =
+                new TraceInjectionConfiguration(annotation, traceStart, traceEnd);
+
+        try (
+                ZipFile zipSrc = new ZipFile(inJar);
+                ZipOutputStream zos = new ZipOutputStream(new FileOutputStream(outJar));
+        ) {
+            Enumeration<? extends ZipEntry> srcEntries = zipSrc.entries();
+            while (srcEntries.hasMoreElements()) {
+                ZipEntry entry = srcEntries.nextElement();
+                ZipEntry newEntry = new ZipEntry(entry.getName());
+                newEntry.setTime(entry.getTime());
+                zos.putNextEntry(newEntry);
+                BufferedInputStream bis = new BufferedInputStream(zipSrc.getInputStream(entry));
+
+                if (entry.getName().endsWith(".class")) {
+                    convert(bis, zos, params);
+                } else {
+                    while (bis.available() > 0) {
+                        zos.write(bis.read());
+                    }
+                    zos.closeEntry();
+                    bis.close();
+                }
+            }
+            zos.finish();
+        }
+    }
+
+    private static void convert(InputStream in, OutputStream out,
+            TraceInjectionConfiguration params) throws IOException {
+        ClassReader cr = new ClassReader(in);
+        ClassWriter cw = new ClassWriter(0);
+        TraceInjectionClassVisitor cv = new TraceInjectionClassVisitor(cw, params);
+        cr.accept(cv, ClassReader.EXPAND_FRAMES);
+        byte[] data = cw.toByteArray();
+        out.write(data);
+    }
+}
diff --git a/tools/traceinjection/src/com/android/traceinjection/TraceInjectionClassVisitor.java b/tools/traceinjection/src/com/android/traceinjection/TraceInjectionClassVisitor.java
new file mode 100644
index 000000000000..863f976b8aff
--- /dev/null
+++ b/tools/traceinjection/src/com/android/traceinjection/TraceInjectionClassVisitor.java
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.traceinjection;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+/**
+ * {@link ClassVisitor} that injects tracing code to methods annotated with the configured
+ * annotation.
+ */
+public class TraceInjectionClassVisitor extends ClassVisitor {
+    private final TraceInjectionConfiguration mParams;
+    public TraceInjectionClassVisitor(ClassVisitor classVisitor,
+            TraceInjectionConfiguration params) {
+        super(Opcodes.ASM7, classVisitor);
+        mParams = params;
+    }
+
+    @Override
+    public MethodVisitor visitMethod(int access, String name, String desc, String signature,
+            String[] exceptions) {
+        MethodVisitor chain = super.visitMethod(access, name, desc, signature, exceptions);
+        return new TraceInjectionMethodAdapter(chain, access, name, desc, mParams);
+    }
+}
diff --git a/tools/traceinjection/src/com/android/traceinjection/TraceInjectionConfiguration.java b/tools/traceinjection/src/com/android/traceinjection/TraceInjectionConfiguration.java
new file mode 100644
index 000000000000..f9595bdad9cf
--- /dev/null
+++ b/tools/traceinjection/src/com/android/traceinjection/TraceInjectionConfiguration.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.traceinjection;
+
+/**
+ * Configuration data for trace method injection.
+ */
+public class TraceInjectionConfiguration {
+    public final String annotation;
+    public final String startMethodClass;
+    public final String startMethodName;
+    public final String endMethodClass;
+    public final String endMethodName;
+
+    public TraceInjectionConfiguration(String annotation, String startMethod, String endMethod) {
+        this.annotation = annotation;
+        String[] startMethodComponents = parseMethod(startMethod);
+        String[] endMethodComponents = parseMethod(endMethod);
+        startMethodClass = startMethodComponents[0];
+        startMethodName = startMethodComponents[1];
+        endMethodClass = endMethodComponents[0];
+        endMethodName = endMethodComponents[1];
+    }
+
+    public String toString() {
+        return "TraceInjectionParams{annotation=" + annotation
+                + ", startMethod=" + startMethodClass + "." + startMethodName
+                + ", endMethod=" + endMethodClass + "." + endMethodName + "}";
+    }
+
+    private static String[] parseMethod(String method) {
+        String[] methodComponents = method.split("\\.");
+        if (methodComponents.length != 2) {
+            throw new IllegalArgumentException("Invalid method descriptor: " + method);
+        }
+        return methodComponents;
+    }
+}
diff --git a/tools/traceinjection/src/com/android/traceinjection/TraceInjectionMethodAdapter.java b/tools/traceinjection/src/com/android/traceinjection/TraceInjectionMethodAdapter.java
new file mode 100644
index 000000000000..c2bbddcb5668
--- /dev/null
+++ b/tools/traceinjection/src/com/android/traceinjection/TraceInjectionMethodAdapter.java
@@ -0,0 +1,183 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.traceinjection;
+
+import org.objectweb.asm.AnnotationVisitor;
+import org.objectweb.asm.Label;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+import org.objectweb.asm.commons.AdviceAdapter;
+import org.objectweb.asm.commons.Method;
+
+/**
+ * Adapter that injects tracing code to methods annotated with the configured annotation.
+ *
+ * Assuming the configured annotation is {@code @Trace} and the configured methods are
+ * {@code Tracing.begin()} and {@code Tracing.end()}, it effectively transforms:
+ *
+ * <pre>{@code
+ * @Trace
+ * void method() {
+ *     doStuff();
+ * }
+ * }</pre>
+ *
+ * into:
+ * <pre>{@code
+ * @Trace
+ * void method() {
+ *     Tracing.begin();
+ *     try {
+ *         doStuff();
+ *     } finally {
+ *         Tracing.end();
+ *     }
+ * }
+ * }</pre>
+ */
+public class TraceInjectionMethodAdapter extends AdviceAdapter {
+    private final TraceInjectionConfiguration mParams;
+    private final Label mStartFinally = newLabel();
+    private final boolean mIsConstructor;
+
+    private boolean mShouldTrace;
+    private long mTraceId;
+    private String mTraceLabel;
+
+    public TraceInjectionMethodAdapter(MethodVisitor methodVisitor, int access,
+            String name, String descriptor, TraceInjectionConfiguration params) {
+        super(Opcodes.ASM7, methodVisitor, access, name, descriptor);
+        mParams = params;
+        mIsConstructor = "<init>".equals(name);
+    }
+
+    @Override
+    public void visitCode() {
+        super.visitCode();
+        if (mShouldTrace) {
+            visitLabel(mStartFinally);
+        }
+    }
+
+    @Override
+    protected void onMethodEnter() {
+        if (!mShouldTrace) {
+            return;
+        }
+        Type type = Type.getType(toJavaSpecifier(mParams.startMethodClass));
+        Method trace = Method.getMethod("void " + mParams.startMethodName + " (long, String)");
+        push(mTraceId);
+        push(getTraceLabel());
+        invokeStatic(type, trace);
+    }
+
+    private String getTraceLabel() {
+        return !isEmpty(mTraceLabel) ? mTraceLabel : getName();
+    }
+
+    @Override
+    protected void onMethodExit(int opCode) {
+        // Any ATHROW exits will be caught as part of our exception-handling block, so putting it
+        // here would cause us to call the end trace method multiple times.
+        if (opCode != ATHROW) {
+            onFinally();
+        }
+    }
+
+    private void onFinally() {
+        if (!mShouldTrace) {
+            return;
+        }
+        Type type = Type.getType(toJavaSpecifier(mParams.endMethodClass));
+        Method trace = Method.getMethod("void " + mParams.endMethodName + " (long)");
+        push(mTraceId);
+        invokeStatic(type, trace);
+    }
+
+    @Override
+    public void visitMaxs(int maxStack, int maxLocals) {
+        final int minStackSize;
+        if (mShouldTrace) {
+            Label endFinally = newLabel();
+            visitLabel(endFinally);
+            catchException(mStartFinally, endFinally, null);
+            // The stack will always contain exactly one element: the exception we caught
+            final Object[] stack = new Object[]{ "java/lang/Throwable"};
+            // Because we use EXPAND_FRAMES, the frame type must always be F_NEW.
+            visitFrame(F_NEW, /* numLocal= */ 0, /* local= */ null, stack.length, stack);
+            onFinally();
+            // Rethrow the exception that we caught in the finally block.
+            throwException();
+
+            // Make sure we have at least enough stack space to push the trace arguments
+            // (long, String)
+            minStackSize = Type.LONG_TYPE.getSize() + Type.getType(String.class).getSize();
+        } else {
+            // We didn't inject anything, so no need for additional stack space.
+            minStackSize = 0;
+        }
+
+        super.visitMaxs(Math.max(minStackSize, maxStack), maxLocals);
+    }
+
+    @Override
+    public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
+        AnnotationVisitor av = super.visitAnnotation(descriptor, visible);
+        if (descriptor.equals(toJavaSpecifier(mParams.annotation))) {
+            if (mIsConstructor) {
+                // TODO: Support constructor tracing. At the moment, constructors aren't supported
+                //  because you can't put an exception handler around a super() call within the
+                //  constructor itself.
+                throw new IllegalStateException("Cannot trace constructors");
+            }
+            av = new TracingAnnotationVisitor(av);
+        }
+        return av;
+    }
+
+    /**
+     * An AnnotationVisitor that pulls the trace ID and label information from the configured
+     * annotation.
+     */
+    class TracingAnnotationVisitor extends AnnotationVisitor {
+
+        TracingAnnotationVisitor(AnnotationVisitor annotationVisitor) {
+            super(Opcodes.ASM7, annotationVisitor);
+        }
+
+        @Override
+        public void visit(String name, Object value) {
+            if ("tag".equals(name)) {
+                mTraceId = (long) value;
+                // If we have a trace annotation and ID, then we have everything we need to trace
+                mShouldTrace = true;
+            } else if ("label".equals(name)) {
+                mTraceLabel = (String) value;
+            }
+            super.visit(name, value);
+        }
+    }
+
+    private static String toJavaSpecifier(String klass) {
+        return "L" + klass + ";";
+    }
+
+    private static boolean isEmpty(String str) {
+        return str == null || "".equals(str);
+    }
+}
diff --git a/tools/traceinjection/test/com/android/traceinjection/InjectionTests.java b/tools/traceinjection/test/com/android/traceinjection/InjectionTests.java
new file mode 100644
index 000000000000..81bf235fe0a6
--- /dev/null
+++ b/tools/traceinjection/test/com/android/traceinjection/InjectionTests.java
@@ -0,0 +1,246 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.traceinjection;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+@RunWith(JUnit4.class)
+public class InjectionTests {
+    public static final int TRACE_TAG = 42;
+    public static final String CUSTOM_TRACE_NAME = "Custom";
+
+    public static final TraceTracker TRACKER = new TraceTracker();
+
+    @After
+    public void tearDown() {
+        TRACKER.reset();
+    }
+
+    @Test
+    public void testDefaultLabel() {
+        assertTraces(this::tracedMethod, "tracedMethod");
+        tracedMethodThrowsAndCatches();
+    }
+
+    @Test
+    public void testCustomLabel() {
+        assertTraces(this::tracedMethodHasCustomName, CUSTOM_TRACE_NAME);
+    }
+
+    @Test
+    public void testTracedMethodsStillThrow() {
+        assertTraces(() -> assertThrows(IllegalArgumentException.class, this::tracedMethodThrows),
+                "tracedMethodThrows");
+        // Also test that we rethrow exceptions from method calls. This is slightly different from
+        // the previous case because the ATHROW instruction is not actually present at all in the
+        // bytecode of the instrumented method.
+        TRACKER.reset();
+        assertTraces(() -> assertThrows(NullPointerException.class,
+                        this::tracedMethodCallsThrowingMethod),
+                "tracedMethodCallsThrowingMethod");
+    }
+
+    @Test
+    public void testNestedTracedMethods() {
+        assertTraces(this::outerTracedMethod, "outerTracedMethod", "innerTracedMethod");
+    }
+
+    @Test
+    public void testTracedMethodWithCatchBlock() {
+        assertTraces(this::tracedMethodThrowsAndCatches, "tracedMethodThrowsAndCatches");
+    }
+
+    @Test
+    public void testTracedMethodWithFinallyBlock() {
+        assertTraces(() -> assertThrows(IllegalArgumentException.class,
+                this::tracedMethodThrowWithFinally), "tracedMethodThrowWithFinally");
+    }
+
+    @Test
+    public void testNonVoidMethod() {
+        assertTraces(this::tracedNonVoidMethod, "tracedNonVoidMethod");
+    }
+
+    @Test
+    public void testNonVoidMethodReturnsWithinCatches() {
+        assertTraces(this::tracedNonVoidMethodReturnsWithinCatches,
+                "tracedNonVoidMethodReturnsWithinCatches");
+    }
+
+    @Test
+    public void testNonVoidMethodReturnsWithinFinally() {
+        assertTraces(this::tracedNonVoidMethodReturnsWithinFinally,
+                "tracedNonVoidMethodReturnsWithinFinally");
+    }
+
+    @Test
+    public void testTracedStaticMethod() {
+        assertTraces(InjectionTests::tracedStaticMethod, "tracedStaticMethod");
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public void tracedMethod() {
+        assertEquals(1, TRACKER.getTraceCount(TRACE_TAG));
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public void tracedMethodThrows() {
+        throw new IllegalArgumentException();
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public void tracedMethodCallsThrowingMethod() {
+        throwingMethod();
+    }
+
+    private void throwingMethod() {
+        throw new NullPointerException();
+    }
+
+
+    @Trace(tag = TRACE_TAG)
+    public void tracedMethodThrowsAndCatches() {
+        try {
+            throw new IllegalArgumentException();
+        } catch (IllegalArgumentException ignored) {
+            assertEquals(1, TRACKER.getTraceCount(TRACE_TAG));
+        }
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public void tracedMethodThrowWithFinally() {
+        try {
+            throw new IllegalArgumentException();
+        } finally {
+            assertEquals(1, TRACKER.getTraceCount(TRACE_TAG));
+        }
+    }
+
+    @Trace(tag = TRACE_TAG, label = CUSTOM_TRACE_NAME)
+    public void tracedMethodHasCustomName() {
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public void outerTracedMethod() {
+        innerTracedMethod();
+        assertEquals(1, TRACKER.getTraceCount(TRACE_TAG));
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public void innerTracedMethod() {
+        assertEquals(2, TRACKER.getTraceCount(TRACE_TAG));
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public int tracedNonVoidMethod() {
+        assertEquals(1, TRACKER.getTraceCount(TRACE_TAG));
+        return 0;
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public int tracedNonVoidMethodReturnsWithinCatches() {
+        try {
+            throw new IllegalArgumentException();
+        } catch (IllegalArgumentException ignored) {
+            assertEquals(1, TRACKER.getTraceCount(TRACE_TAG));
+            return 0;
+        }
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public int tracedNonVoidMethodReturnsWithinFinally() {
+        try {
+            throw new IllegalArgumentException();
+        } finally {
+            assertEquals(1, TRACKER.getTraceCount(TRACE_TAG));
+            return 0;
+        }
+    }
+
+    @Trace(tag = TRACE_TAG)
+    public static void tracedStaticMethod() {
+        assertEquals(1, TRACKER.getTraceCount(TRACE_TAG));
+    }
+
+    public void assertTraces(Runnable r, String... traceLabels) {
+        r.run();
+        assertEquals(Arrays.asList(traceLabels), TRACKER.getTraceLabels(TRACE_TAG));
+        TRACKER.assertAllTracesClosed();
+    }
+
+    public static void traceStart(long tag, String name) {
+        TRACKER.onTraceStart(tag, name);
+    }
+
+    public static void traceEnd(long tag) {
+        TRACKER.onTraceEnd(tag);
+    }
+
+    static class TraceTracker {
+        private final Map<Long, List<String>> mTraceLabelsByTag = new HashMap<>();
+        private final Map<Long, Integer> mTraceCountsByTag = new HashMap<>();
+
+        public void onTraceStart(long tag, String name) {
+            getTraceLabels(tag).add(name);
+            mTraceCountsByTag.put(tag, mTraceCountsByTag.getOrDefault(tag, 0) + 1);
+        }
+
+        public void onTraceEnd(long tag) {
+            final int newCount = getTraceCount(tag) - 1;
+            if (newCount < 0) {
+                throw new IllegalStateException("Trace count has gone negative for tag " + tag);
+            }
+            mTraceCountsByTag.put(tag, newCount);
+        }
+
+        public void reset() {
+            mTraceLabelsByTag.clear();
+            mTraceCountsByTag.clear();
+        }
+
+        public List<String> getTraceLabels(long tag) {
+            if (!mTraceLabelsByTag.containsKey(tag)) {
+                mTraceLabelsByTag.put(tag, new ArrayList<>());
+            }
+            return mTraceLabelsByTag.get(tag);
+        }
+
+        public int getTraceCount(long tag) {
+            return mTraceCountsByTag.getOrDefault(tag, 0);
+        }
+
+        public void assertAllTracesClosed() {
+            for (Map.Entry<Long, Integer> count: mTraceCountsByTag.entrySet()) {
+                final String errorMsg = "Tag " + count.getKey() + " is not fully closed (count="
+                        + count.getValue() + ")";
+                assertEquals(errorMsg, 0, (int) count.getValue());
+            }
+        }
+    }
+}
diff --git a/tools/traceinjection/test/com/android/traceinjection/Trace.java b/tools/traceinjection/test/com/android/traceinjection/Trace.java
new file mode 100644
index 000000000000..9e1c545673e8
--- /dev/null
+++ b/tools/traceinjection/test/com/android/traceinjection/Trace.java
@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.traceinjection;
+
+public @interface Trace {
+    long tag();
+    String label() default "";
+}