Do not assume there is only one crosvm process

Our benchmarks search for crosvm's PID by running 'ps -a | grep crosvm'.
This may work now but will fail if/when then system runs VMs
independently of the test process.
Add methods to find crosvm as the child process of virtmgr, which in
turn is the child process of the current PID.

Test: atest MicrodroidBenchmarks
Change-Id: I0dcd1f3aba6f69becccb8a1573ea71046e932de0
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index 4b11d77..9851a17 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -36,6 +36,7 @@
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineConfig;
 import android.system.virtualmachine.VirtualMachineException;
+import android.system.Os;
 import android.util.Log;
 
 import com.android.microdroid.test.common.MetricsProcessor;
@@ -347,16 +348,7 @@
 
         CrosvmStats(Function<String, String> shellExecutor) {
             try {
-                List<Integer> crosvmPids =
-                        ProcessUtil.getProcessMap(shellExecutor).entrySet().stream()
-                                .filter(e -> e.getValue().contains("crosvm"))
-                                .map(e -> e.getKey())
-                                .collect(java.util.stream.Collectors.toList());
-                if (crosvmPids.size() != 1) {
-                    throw new IllegalStateException(
-                            "expected to find exactly one crosvm processes, found "
-                                    + crosvmPids.size());
-                }
+                int crosvmPid = ProcessUtil.getCrosvmPid(Os.getpid(), shellExecutor);
 
                 long hostRss = 0;
                 long hostPss = 0;
@@ -364,7 +356,7 @@
                 long guestPss = 0;
                 boolean hasGuestMaps = false;
                 for (ProcessUtil.SMapEntry entry :
-                        ProcessUtil.getProcessSmaps(crosvmPids.get(0), shellExecutor)) {
+                        ProcessUtil.getProcessSmaps(crosvmPid, shellExecutor)) {
                     long rss = entry.metrics.get("Rss");
                     long pss = entry.metrics.get("Pss");
                     if (entry.name.contains("crosvm_guest")) {
diff --git a/tests/helper/src/java/com/android/microdroid/test/common/ProcessUtil.java b/tests/helper/src/java/com/android/microdroid/test/common/ProcessUtil.java
index 940ec9c..c72d91e 100644
--- a/tests/helper/src/java/com/android/microdroid/test/common/ProcessUtil.java
+++ b/tests/helper/src/java/com/android/microdroid/test/common/ProcessUtil.java
@@ -22,9 +22,12 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
+import java.util.stream.IntStream;
 
 /** This class provides process utility for both device tests and host tests. */
 public final class ProcessUtil {
+    private static final String CROSVM_BIN = "/apex/com.android.virt/bin/crosvm";
+    private static final String VIRTMGR_BIN = "/apex/com.android.virt/bin/virtmgr";
 
     /** A memory map entry from /proc/{pid}/smaps */
     public static class SMapEntry {
@@ -89,6 +92,35 @@
         return processMap;
     }
 
+    private static IntStream getChildProcesses(
+            int pid, String cmdlineFilter, Function<String, String> shellExecutor) {
+        String cmd = "pgrep -P " + pid;
+        if (cmdlineFilter != null) {
+            cmd += " -f " + cmdlineFilter;
+        }
+        return shellExecutor.apply(cmd).trim().lines().mapToInt(Integer::parseInt);
+    }
+
+    private static int getSingleChildProcess(
+            int parentPid, String cmdlineFilter, Function<String, String> shellExecutor) {
+        int[] pids = getChildProcesses(parentPid, cmdlineFilter, shellExecutor).toArray();
+        if (pids.length == 0) {
+            throw new IllegalStateException("No process found for " + cmdlineFilter);
+        } else if (pids.length > 1) {
+            throw new IllegalStateException("More than one process found for " + cmdlineFilter);
+        }
+        return pids[0];
+    }
+
+    public static int getVirtmgrPid(int parentPid, Function<String, String> shellExecutor) {
+        return getSingleChildProcess(parentPid, VIRTMGR_BIN, shellExecutor);
+    }
+
+    public static int getCrosvmPid(int parentPid, Function<String, String> shellExecutor) {
+        int virtmgrPid = getVirtmgrPid(parentPid, shellExecutor);
+        return getSingleChildProcess(virtmgrPid, CROSVM_BIN, shellExecutor);
+    }
+
     // To ensures that only one object is created at a time.
     private ProcessUtil() {}