Add test infra to create profile scenerios

Currently, we only have a way to run profile tests on compilation, which
prevents us from creating more specific test scenerious.

This CL introduces a new rule that works with profile.sh to capture
profiling data from tests.

It is composed of two parts:
profile.sh: Previoulsy, profile.sh always attached to the gradle
compilation, which does not cover tests since they are run on a forked
process. Instead, now it has a new `target` parameter where the caller
can specifiy the target jvm process to attach (limited to gradle or
test).
ProfileRule: This is a test rule for tests to report trace sections.
When running compilation tests, most time is spent in actually compiling
the code and we don't care about that. ProfileRule adds a `trace` method
that can be used to mark trace sections. These trace call stacks are
automatically filtered by profile.sh to generate the report only
including them.

I've also added a simple profile test that obtains a raw type to
validate the performance improvements in the previous change.
(I792f643a80294a65bf55572cc4872b4a8af4e1b7).

Bug: 204923211
Test: RawTypeCreationScenarioTest
Change-Id: Id78d1266dc9815eb7eba2dbff4fd27b001e7b022
diff --git a/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/profiling/ProfileRule.kt b/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/profiling/ProfileRule.kt
new file mode 100644
index 0000000..c30c894
--- /dev/null
+++ b/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/profiling/ProfileRule.kt
@@ -0,0 +1,100 @@
+/*
+ * Copyright 2021 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.room.compiler.processing.profiling
+
+import org.junit.AssumptionViolatedException
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * Helper rule to run profiling tests.
+ *
+ * These tests are run along with `scripts/profile.sh` to build an async profile
+ * output based on a test scenario.
+ *
+ * If this rule is applied outside a profiling session, it will ignore the test.
+ */
+class ProfileRule : TestRule {
+    /**
+     * Runs the given block, repeatedly :).
+     *
+     * It will first run it [warmUps] times with a fake tracer. Then it will run
+     * the block [repeat] times with a real profiling scope that will be captured by
+     * profile.sh.
+     */
+    fun runRepeated(
+        warmUps: Int,
+        repeat: Int,
+        block: (ProfileScope) -> Unit
+    ) {
+        val warmUpScope = WarmUpProfileScope()
+        repeat(warmUps) {
+            block(warmUpScope)
+        }
+        val realProfileScope = RealProfileScope()
+        repeat(repeat) {
+            block(realProfileScope)
+        }
+    }
+
+    override fun apply(base: Statement, description: Description): Statement {
+        return object : Statement() {
+            override fun evaluate() {
+                assumeProfiling()
+                base.evaluate()
+            }
+        }
+    }
+
+    private fun assumeProfiling() {
+        if (!isProfilingEnabled) {
+            throw AssumptionViolatedException("No reason to run while not profiling")
+        }
+    }
+
+    interface ProfileScope {
+        /**
+         * Utility function for tests to mark certain areas of their code for tracking.
+         *
+         * This method is explicitly not marked as inline to ensure it shows up in the
+         * profiling output.
+         */
+        fun trace(block: () -> Unit)
+    }
+
+    private class RealProfileScope : ProfileScope {
+        override fun trace(block: () -> Unit) {
+            // this doesn't do anything but profile.sh trace profiler checks
+            // this class while filtering stacktraces
+            block()
+        }
+    }
+
+    private class WarmUpProfileScope : ProfileScope {
+        override fun trace(block: () -> Unit) {
+            block()
+        }
+    }
+
+    companion object {
+        val isProfilingEnabled by lazy {
+            // set by profile.sh
+            System.getenv("ANDROIDX_ROOM_ENABLE_PROFILE_TESTS") != null
+        }
+    }
+}
\ No newline at end of file
diff --git a/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/profiling/RawTypeCreationScenarioTest.kt b/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/profiling/RawTypeCreationScenarioTest.kt
new file mode 100644
index 0000000..83dabc8
--- /dev/null
+++ b/room/room-compiler-processing/src/test/java/androidx/room/compiler/processing/profiling/RawTypeCreationScenarioTest.kt
@@ -0,0 +1,58 @@
+/*
+ * Copyright 2021 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.room.compiler.processing.profiling
+
+import androidx.room.compiler.processing.util.Source
+import androidx.room.compiler.processing.util.runKspTest
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+class RawTypeCreationScenarioTest {
+    @get:Rule
+    val profileRule = ProfileRule()
+
+    @Test
+    fun profile() {
+        val classCount = 1000
+        val contents = buildString {
+            (0 until classCount).forEach { cnt ->
+                appendLine("class Sample$cnt")
+            }
+        }
+        val sources = Source.kotlin("Sample.kt", contents)
+        val classNames = (0 until classCount).map { "Sample$it" }
+        profileRule.runRepeated(
+            warmUps = 10,
+            repeat = 20
+        ) { profileScope ->
+            runKspTest(
+                sources = listOf(sources)
+            ) { invocation ->
+                profileScope.trace {
+                    classNames.forEach { className ->
+                        invocation.processingEnv.requireTypeElement(
+                            className
+                        ).type.rawType
+                    }
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/room/scripts/attach-async-profiler-to-tests-init-script.gradle b/room/scripts/attach-async-profiler-to-tests-init-script.gradle
new file mode 100644
index 0000000..97303cf
--- /dev/null
+++ b/room/scripts/attach-async-profiler-to-tests-init-script.gradle
@@ -0,0 +1,14 @@
+/**
+ * A gradle init script that will attach given Jvm Args to each Test task.
+ */
+def jvmArgs = startParameter.systemPropertiesArgs.get("androidx.room.testJvmArgs")
+taskGraph.addTaskExecutionGraphListener { graph ->
+    graph.beforeTask { task ->
+        if (task instanceof Test) {
+            task.jvmArgs(jvmArgs)
+            // this environment variable is used to avoid running profiling tests
+            // unless we are in a profiling execution
+            task.environment("ANDROIDX_ROOM_ENABLE_PROFILE_TESTS", "true")
+        }
+    }
+}
diff --git a/room/scripts/profile.sh b/room/scripts/profile.sh
index ce1acf6..a412414 100755
--- a/room/scripts/profile.sh
+++ b/room/scripts/profile.sh
@@ -5,8 +5,13 @@
 function usage {
     echo "usage: ./profile.sh <gradle tasks>
     ./profile.sh --name my_run \
+        --target gradle \
         :room:integration-tests:room-testapp-kotlin:kspWithKspDebugAndroidTestKotlin \
         --filter *K2JVMCompiler*
+    ./profile.sh --name my_run \
+        --target test \
+        :room:room-compiler-processing:test \
+        --tests RawTypeCreationScenarioTest
 
 
     Arguments:
@@ -15,6 +20,8 @@
     --filter <stacktrace include filter>: Regex to filter stacktraces. e.g. *room*
     --outputFile <file path>: Path to the output file
     --preset [kapt|ksp]: Predefined tasks to run ksp/kapt profile on the KotlinTestApp
+    --target [gradle|test]: The target process type to profile. Gradle for compilation,
+                            test for profile test tasks. Defaults to gradle
 
     It will run the task once without profiling, stop gradle daemon and re-run it again while also
     disabling up-to-date checks for the given tasks.
@@ -46,6 +53,7 @@
 REPORT_NAME="profile"
 GRADLE_ARGS=""
 ADDITIONAL_AGENT_PARAMS=""
+AGENT_TARGET="gradle"
 
 while [ $# -gt 0 ]; do
     case $1 in
@@ -64,14 +72,30 @@
             OUTPUT_FILE=$2
             shift
             ;;
+        --target)
+            if [ "$2" = "gradle" ]; then
+                AGENT_TARGET="gadle"
+            elif [ "$2" = "test" ]; then
+                AGENT_TARGET="test"
+                # add RealProfileScope to tracing automatically. Otherwise, it will be dominated
+                # by compilation.
+                ADDITIONAL_AGENT_PARAMS="$ADDITIONAL_AGENT_PARAMS,include=*RealProfileScope*"
+            else
+                echo "invalid target: $2"
+                exit
+            fi
+            shift
+            ;;
         --preset)
             if [ "$2" = "kapt" ]; then
                 GRADLE_ARGS=":room:integration-tests:room-testapp-kotlin:kaptGenerateStubsWithKaptDebugAndroidTestKotlin \
                     :room:integration-tests:room-testapp-kotlin:kaptWithKaptDebugAndroidTestKotlin"
                 ADDITIONAL_AGENT_PARAMS="$ADDITIONAL_AGENT_PARAMS,include=*AbstractKapt3Extension*"
+                AGENT_TARGET="gradle"
             elif [ "$2" = "ksp" ]; then
                 GRADLE_ARGS=":room:integration-tests:room-testapp-kotlin:kspWithKspDebugAndroidTestKotlin"
                 ADDITIONAL_AGENT_PARAMS="$ADDITIONAL_AGENT_PARAMS,include=*AbstractKotlinSymbolProcessingExtension*"
+                AGENT_TARGET="gradle"
             else
                 echo "invalid preset: $2"
                 exit
@@ -113,11 +137,16 @@
 
     local AGENT_PARAMS="file=${OUTPUT_FILE}${ADDITIONAL_AGENT_PARAMS}"
     $GRADLEW -p $PROJECT_DIR $GRADLE_ARGS
+    local AGENT_PARAMETER_NAME="-Dorg.gradle.jvmargs"
+    if [ $AGENT_TARGET = "test" ]; then
+        AGENT_PARAMETER_NAME="-Dandroidx.room.testJvmArgs"
+    fi
     $GRADLEW --no-daemon \
         --init-script $SCRIPT_DIR/rerun-requested-task-init-script.gradle \
+        --init-script $SCRIPT_DIR/attach-async-profiler-to-tests-init-script.gradle \
         -p $PROJECT_DIR $GRADLE_ARGS \
         -Dkotlin.compiler.execution.strategy="in-process"  \
-        -Dorg.gradle.jvmargs="-agentpath:$AGENT_PATH=start,event=cpu,$AGENT_PARAMS,interval=500000" #sample every .5 ms
+        $AGENT_PARAMETER_NAME="-agentpath:$AGENT_PATH=start,event=cpu,$AGENT_PARAMS,interval=500000" #sample every .5 ms
 }
 
 profile $GRADLE_ARGS, $REPORT_NAME
@@ -135,4 +164,4 @@
 }
 
 # open it in chrome once done
-openFileInChrome $OUTPUT_FILE
\ No newline at end of file
+openFileInChrome $OUTPUT_FILE