Merge "Add tool for injecting tracing code into a method."
diff --git a/tools/traceinjection/Android.bp b/tools/traceinjection/Android.bp
new file mode 100644
index 0000000..1395c5f
--- /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 0000000..7f4ee1d
--- /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 0000000..190df81
--- /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 0000000..863f976
--- /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 0000000..f9595bd
--- /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 0000000..c2bbddc
--- /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 0000000..81bf235
--- /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 0000000..9e1c545
--- /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 "";
+}