1 /*
2  * Copyright (C) 2022 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.microdroid.test.common;
18 
19 import java.io.IOException;
20 import java.util.ArrayList;
21 import java.util.HashMap;
22 import java.util.List;
23 import java.util.Map;
24 import java.util.function.Function;
25 import java.util.stream.IntStream;
26 
27 /** This class provides process utility for both device tests and host tests. */
28 public final class ProcessUtil {
29     private static final String CROSVM_BIN = "/apex/com.android.virt/bin/crosvm";
30     private static final String VIRTMGR_BIN = "/apex/com.android.virt/bin/virtmgr";
31 
32     /** A memory map entry from /proc/{pid}/smaps */
33     public static class SMapEntry {
34         public String name;
35         public Map<String, Long> metrics;
36 
37         @Override
toString()38         public String toString() {
39             StringBuilder sb = new StringBuilder();
40             sb.append("name: " + name + "\n");
41             metrics.forEach(
42                     (k, v) -> {
43                         sb.append("  " + k + ": " + v + "\n");
44                     });
45             return sb.toString();
46         }
47     }
48 
49     /** Gets metrics key and values mapping of specified process id */
getProcessSmaps(int pid, Function<String, String> shellExecutor)50     public static List<SMapEntry> getProcessSmaps(int pid, Function<String, String> shellExecutor)
51             throws IOException {
52         String path = "/proc/" + pid + "/smaps";
53         return parseMemoryInfo(shellExecutor.apply("cat " + path));
54     }
55 
56     /** Gets metrics key and values mapping of specified process id */
getProcessSmapsRollup( int pid, Function<String, String> shellExecutor)57     public static Map<String, Long> getProcessSmapsRollup(
58             int pid, Function<String, String> shellExecutor) throws IOException {
59         String path = "/proc/" + pid + "/smaps_rollup";
60         List<SMapEntry> entries = parseMemoryInfo(shellExecutor.apply("cat " + path + " || true"));
61         if (entries.size() > 1) {
62             throw new RuntimeException(
63                     "expected at most one entry in smaps_rollup, got " + entries.size());
64         }
65         if (entries.size() == 1) {
66             return entries.get(0).metrics;
67         }
68         return new HashMap<String, Long>();
69     }
70 
71     /** Gets global memory metrics key and values mapping */
getProcessMemoryMap(Function<String, String> shellExecutor)72     public static Map<String, Long> getProcessMemoryMap(Function<String, String> shellExecutor)
73             throws IOException {
74         // The input file of parseMemoryInfo need a header string as the key of output entries.
75         // /proc/meminfo doesn't have this line so add one as the key.
76         String header = "device memory info\n";
77         List<SMapEntry> entries =
78                 parseMemoryInfo(header + shellExecutor.apply("cat /proc/meminfo"));
79         if (entries.size() != 1) {
80             throw new RuntimeException(
81                     "expected one entry in /proc/meminfo, got " + entries.size());
82         }
83         return entries.get(0).metrics;
84     }
85 
86     /** Gets process id and process name mapping of the device */
getProcessMap(Function<String, String> shellExecutor)87     public static Map<Integer, String> getProcessMap(Function<String, String> shellExecutor)
88             throws IOException {
89         Map<Integer, String> processMap = new HashMap<>();
90         for (String ps : skipFirstLine(shellExecutor.apply("ps -Ao PID,NAME")).split("\n")) {
91             // Each line is '<pid> <name>'.
92             // EX : 11424 dex2oat64
93             ps = ps.trim();
94             if (ps.length() == 0) {
95                 continue;
96             }
97             int space = ps.indexOf(" ");
98             String pName = ps.substring(space + 1);
99             int pId = Integer.parseInt(ps.substring(0, space));
100             processMap.put(pId, pName);
101         }
102 
103         return processMap;
104     }
105 
getChildProcesses( int pid, String cmdlineFilter, Function<String, String> shellExecutor)106     private static IntStream getChildProcesses(
107             int pid, String cmdlineFilter, Function<String, String> shellExecutor) {
108         String cmd = "pgrep -P " + pid;
109         if (cmdlineFilter != null) {
110             cmd += " -f " + cmdlineFilter;
111         }
112         return shellExecutor.apply(cmd).trim().lines().mapToInt(Integer::parseInt);
113     }
114 
getSingleChildProcess( int parentPid, String cmdlineFilter, Function<String, String> shellExecutor)115     private static int getSingleChildProcess(
116             int parentPid, String cmdlineFilter, Function<String, String> shellExecutor) {
117         int[] pids = getChildProcesses(parentPid, cmdlineFilter, shellExecutor).toArray();
118         if (pids.length == 0) {
119             throw new IllegalStateException("No process found for " + cmdlineFilter);
120         } else if (pids.length > 1) {
121             throw new IllegalStateException("More than one process found for " + cmdlineFilter);
122         }
123         return pids[0];
124     }
125 
getVirtmgrPid(int parentPid, Function<String, String> shellExecutor)126     public static int getVirtmgrPid(int parentPid, Function<String, String> shellExecutor) {
127         return getSingleChildProcess(parentPid, VIRTMGR_BIN, shellExecutor);
128     }
129 
getCrosvmPid( int parentPid, String testName, Function<String, String> shellExecutor)130     public static int getCrosvmPid(
131             int parentPid, String testName, Function<String, String> shellExecutor) {
132         int virtmgrPid = getVirtmgrPid(parentPid, shellExecutor);
133         return getSingleChildProcess(virtmgrPid, "crosvm_" + testName, shellExecutor);
134     }
135 
136     // To ensures that only one object is created at a time.
ProcessUtil()137     private ProcessUtil() {}
138 
parseMemoryInfo(String file)139     private static List<SMapEntry> parseMemoryInfo(String file) {
140         List<SMapEntry> entries = new ArrayList<SMapEntry>();
141         for (String line : file.split("\n")) {
142             line = line.trim();
143             if (line.length() == 0) {
144                 continue;
145             }
146             // Each line is '<metrics>:        <number> kB'.
147             // EX : Pss_Anon:        70712 kB
148             // EX : Active(file):     5792 kB
149             // EX : ProtectionKey:       0
150             if (line.matches("[\\w()]+:\\s+.*")) {
151                 if (entries.size() == 0) {
152                     throw new RuntimeException("unexpected line: " + line);
153                 }
154                 if (line.endsWith(" kB")) line = line.substring(0, line.length() - 3);
155                 String[] elems = line.split(":");
156                 String name = elems[0].trim();
157                 try {
158                     entries.get(entries.size() - 1)
159                             .metrics
160                             .put(name, Long.parseLong(elems[1].trim()));
161                 } catch (java.lang.NumberFormatException e) {
162                     // Some entries, like "VmFlags", aren't numbers, just ignore.
163                 }
164                 continue;
165             }
166             // Parse the header and create a new entry for it.
167             // Some header examples:
168             //     7f644098a000-7f644098c000 rw-p 00000000 00:00 0
169             //     00400000-0048a000 r-xp 00000000 fd:03 960637   /bin/bash
170             //     75e42af000-75f42af000 rw-s 00000000 00:01 235  /memfd:crosvm_guest (deleted)
171             SMapEntry entry = new SMapEntry();
172             String[] parts = line.split("\\s+", 6);
173             if (parts.length >= 6) {
174                 entry.name = parts[5];
175             } else {
176                 entry.name = "";
177             }
178             entry.metrics = new HashMap<String, Long>();
179             entries.add(entry);
180         }
181         return entries;
182     }
183 
skipFirstLine(String str)184     private static String skipFirstLine(String str) {
185         int index = str.indexOf("\n");
186         return (index < 0) ? "" : str.substring(index + 1);
187     }
188 }
189