Add tool for injecting tracing code into a method.

This tool rewrites the bytecode in the designated JAR files to produce
tracing calls on enter and exit, while making sure to close the tracing
span even on exceptions being thrown.

The idea is mostly to reduce the amount of noise within methods when
just trying to add some tracing.

Test: atest --host TraceInjectionTests
Change-Id: If6acb72f34cbb83d9b041a62ee3d8c2abf74b69e

2 files changed, 268 insertions, 0 deletions

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 "";
+}