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(
73             Function<String, String> shellExecutor) 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 = parseMemoryInfo(header
78                 + 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, Function<String, String> shellExecutor)130     public static int getCrosvmPid(int parentPid, Function<String, String> shellExecutor) {
131         int virtmgrPid = getVirtmgrPid(parentPid, shellExecutor);
132         return getSingleChildProcess(virtmgrPid, CROSVM_BIN, shellExecutor);
133     }
134 
135     // To ensures that only one object is created at a time.
ProcessUtil()136     private ProcessUtil() {}
137 
parseMemoryInfo(String file)138     private static List<SMapEntry> parseMemoryInfo(String file) {
139         List<SMapEntry> entries = new ArrayList<SMapEntry>();
140         for (String line : file.split("\n")) {
141             line = line.trim();
142             if (line.length() == 0) {
143                 continue;
144             }
145             // Each line is '<metrics>:        <number> kB'.
146             // EX : Pss_Anon:        70712 kB
147             // EX : Active(file):     5792 kB
148             // EX : ProtectionKey:       0
149             if (line.matches("[\\w()]+:\\s+.*")) {
150                 if (entries.size() == 0) {
151                     throw new RuntimeException("unexpected line: " + line);
152                 }
153                 if (line.endsWith(" kB")) line = line.substring(0, line.length() - 3);
154                 String[] elems = line.split(":");
155                 String name = elems[0].trim();
156                 try {
157                     entries.get(entries.size() - 1)
158                             .metrics
159                             .put(name, Long.parseLong(elems[1].trim()));
160                 } catch (java.lang.NumberFormatException e) {
161                     // Some entries, like "VmFlags", aren't numbers, just ignore.
162                 }
163                 continue;
164             }
165             // Parse the header and create a new entry for it.
166             // Some header examples:
167             //     7f644098a000-7f644098c000 rw-p 00000000 00:00 0
168             //     00400000-0048a000 r-xp 00000000 fd:03 960637   /bin/bash
169             //     75e42af000-75f42af000 rw-s 00000000 00:01 235  /memfd:crosvm_guest (deleted)
170             SMapEntry entry = new SMapEntry();
171             String[] parts = line.split("\\s+", 6);
172             if (parts.length >= 6) {
173                 entry.name = parts[5];
174             } else {
175                 entry.name = "";
176             }
177             entry.metrics = new HashMap<String, Long>();
178             entries.add(entry);
179         }
180         return entries;
181     }
182 
skipFirstLine(String str)183     private static String skipFirstLine(String str) {
184         int index = str.indexOf("\n");
185         return (index < 0) ? "" : str.substring(index + 1);
186     }
187 }
188