Ensure no thread leak after NetworkStatsServiceTest

This change includes:
 1. A mechanism in DevSdkIngoreRunnrer to dump thread counts
    before/after tests, and compare thread counts to detect leaks.
 2. Add an annotation @MonitorThreadLeak for test classes to
    annotate the classes which should monitor thread leaks.
 3. Annotated NetworkStatsServiceTest to apply the enforcement.

Sample output:
[1/2] android.net.connectivity.com.android.server.net.NetworkStatsServiceTest#testDumpStatsMap: PASSED (1.187s)
[2/2] android.net.connectivity.com.android.server.net.NetworkStatsServiceTest#ThreadLeakMonitor: FAILED (7ms)

STACKTRACE:
java.lang.IllegalStateException: Expected threads: {binder:26055_3=1, Instr: androidx.test.runner.AndroidJUnitRunner=1, main=1, InstrumentationConnectionThread=1, binder:26055_2=1, binder:26055_1=1} but got: {binder:26055_3=1, NetworkStatsObservers=1, Instr: androidx.test.runner.AndroidJUnitRunner=1, main=1, InstrumentationConnectionThread=1, binder:26055_2=1, binder:26055_1=1}

Test: atest ConnectivityCoverageTests:android.net.connectivity.com.android.server.net.NetworkStatsServiceTest \
      --rerun-until-failure 100
Bug: 308544001
Bug: 307693729
Change-Id: Ia0bccb82c5985df608b8402009b32626b6b17c5a
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt b/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt
index 2d281fd..1ba83ca 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/DevSdkIgnoreRunner.kt
@@ -19,6 +19,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.testutils.DevSdkIgnoreRule.IgnoreAfter
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
+import java.lang.IllegalStateException
 import java.lang.reflect.Modifier
 import org.junit.runner.Description
 import org.junit.runner.Runner
@@ -27,6 +28,7 @@
 import org.junit.runner.manipulation.NoTestsRemainException
 import org.junit.runner.manipulation.Sortable
 import org.junit.runner.manipulation.Sorter
+import org.junit.runner.notification.Failure
 import org.junit.runner.notification.RunNotifier
 import org.junit.runners.Parameterized
 
@@ -52,6 +54,9 @@
  *     class MyTestClass { ... }
  */
 class DevSdkIgnoreRunner(private val klass: Class<*>) : Runner(), Filterable, Sortable {
+    private val leakMonitorDesc = Description.createTestDescription(klass, "ThreadLeakMonitor")
+    private val shouldThreadLeakFailTest = klass.isAnnotationPresent(MonitorThreadLeak::class.java)
+
     // Inference correctly infers Runner & Filterable & Sortable for |baseRunner|, but the
     // Java bytecode doesn't have a way to express this. Give this type a name by wrapping it.
     private class RunnerWrapper<T>(private val wrapped: T) :
@@ -61,6 +66,10 @@
         override fun run(notifier: RunNotifier?) = wrapped.run(notifier)
     }
 
+    // Annotation for test classes to indicate the test runner should monitor thread leak.
+    // TODO(b/307693729): Remove this annotation and monitor thread leak by default.
+    annotation class MonitorThreadLeak
+
     private val baseRunner: RunnerWrapper<*>? = klass.let {
         val ignoreAfter = it.getAnnotation(IgnoreAfter::class.java)
         val ignoreUpTo = it.getAnnotation(IgnoreUpTo::class.java)
@@ -81,20 +90,52 @@
                 it.isAnnotationPresent(Parameterized.Parameters::class.java) }
 
     override fun run(notifier: RunNotifier) {
-        if (baseRunner != null) {
+        if (baseRunner == null) {
+            // Report a single, skipped placeholder test for this class, as the class is expected to
+            // report results when run. In practice runners that apply the Filterable implementation
+            // would see a NoTestsRemainException and not call the run method.
+            notifier.fireTestIgnored(
+                    Description.createTestDescription(klass, "skippedClassForDevSdkMismatch"))
+            return
+        }
+        if (!shouldThreadLeakFailTest) {
             baseRunner.run(notifier)
             return
         }
 
-        // Report a single, skipped placeholder test for this class, as the class is expected to
-        // report results when run. In practice runners that apply the Filterable implementation
-        // would see a NoTestsRemainException and not call the run method.
-        notifier.fireTestIgnored(
-                Description.createTestDescription(klass, "skippedClassForDevSdkMismatch"))
+        // Dump threads as a baseline to monitor thread leaks.
+        val threadCountsBeforeTest = getAllThreadNameCounts()
+
+        baseRunner.run(notifier)
+
+        notifier.fireTestStarted(leakMonitorDesc)
+        val threadCountsAfterTest = getAllThreadNameCounts()
+        if (threadCountsBeforeTest != threadCountsAfterTest) {
+            notifier.fireTestFailure(Failure(leakMonitorDesc,
+                    IllegalStateException("Expected threads: $threadCountsBeforeTest " +
+                            "but got: $threadCountsAfterTest")))
+        }
+        notifier.fireTestFinished(leakMonitorDesc)
+    }
+
+    private fun getAllThreadNameCounts(): Map<String, Int> {
+        // Get the counts of threads in the group per name.
+        // Filter system thread groups.
+        return Thread.getAllStackTraces().keys
+                .filter { it.threadGroup?.name != "system" }
+                .groupingBy { it.name }.eachCount()
     }
 
     override fun getDescription(): Description {
-        return baseRunner?.description ?: Description.createSuiteDescription(klass)
+        if (baseRunner == null) {
+            return Description.createSuiteDescription(klass)
+        }
+
+        return baseRunner.description.also {
+            if (shouldThreadLeakFailTest) {
+                it.addChild(leakMonitorDesc)
+            }
+        }
     }
 
     /**
@@ -102,7 +143,9 @@
      */
     override fun testCount(): Int {
         // When ignoring the tests, a skipped placeholder test is reported, so test count is 1.
-        return baseRunner?.testCount() ?: 1
+        if (baseRunner == null) return 1
+
+        return baseRunner.testCount() + if (shouldThreadLeakFailTest) 1 else 0
     }
 
     @Throws(NoTestsRemainException::class)
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 92a5b64..7a4dfed 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -193,6 +193,7 @@
  * TODO: This test used to be really brittle because it used Easymock - it uses Mockito now, but
  * still uses the Easymock structure, which could be simplified.
  */
+@DevSdkIgnoreRunner.MonitorThreadLeak
 @RunWith(DevSdkIgnoreRunner.class)
 @SmallTest
 // NetworkStatsService is not updatable before T, so tests do not need to be backwards compatible