1#! /usr/bin/env python3
2#
3# Copyright 2023, The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17import argparse
18from collections import namedtuple
19import subprocess
20
21try:
22  from tqdm import tqdm
23except:
24
25  def tqdm(x):
26    return x
27
28
29ProcEntry = namedtuple('ProcEntry', 'pid, ppid, cmd, name, etc_args')
30
31
32def get_mem_stats(
33    zygote_pid,
34    target_pid,
35    target_name,
36    imgdiag_path,
37    boot_image,
38    device_out_dir,
39    host_out_dir,
40):
41  imgdiag_output_path = (
42      f'{device_out_dir}/imgdiag_{target_name}_{target_pid}.txt'
43  )
44  cmd_collect = (
45      'adb shell '
46      f'"{imgdiag_path} --zygote-diff-pid={zygote_pid} --image-diff-pid={target_pid} '
47      f'--output={imgdiag_output_path} --boot-image={boot_image} --dump-dirty-objects"'
48  )
49
50  try:
51    subprocess.run(cmd_collect, shell=True, check=True)
52  except:
53    print('imgdiag call failed on:', target_pid, target_name)
54    return
55
56  cmd_pull = f'adb pull {imgdiag_output_path} {host_out_dir}'
57  subprocess.run(cmd_pull, shell=True, check=True, capture_output=True)
58
59
60def main():
61  parser = argparse.ArgumentParser(
62      description=(
63          'Run imgdiag on selected processes and pull results from the device.'
64      ),
65      formatter_class=argparse.ArgumentDefaultsHelpFormatter,
66  )
67  parser.add_argument(
68      'process_names',
69      nargs='*',
70      help='Process names to use. If none - dump all zygote children.',
71  )
72  parser.add_argument(
73      '--boot-image',
74      dest='boot_image',
75      default='/data/misc/apexdata/com.android.art/dalvik-cache/boot.art',
76      help='Path to boot.art',
77  )
78  parser.add_argument(
79      '--zygote',
80      default='zygote64',
81      help='Zygote process name',
82  )
83  parser.add_argument(
84      '--imgdiag',
85      default='/apex/com.android.art/bin/imgdiag64',
86      help='Path to imgdiag binary.',
87  )
88  parser.add_argument(
89      '--device-out-dir',
90      default='/data/local/tmp/imgdiag_out',
91      help='Directory for imgdiag output files on the device.',
92  )
93  parser.add_argument(
94      '--host-out-dir',
95      default='./',
96      help='Directory for imgdiag output files on the host.',
97  )
98
99  args = parser.parse_args()
100
101  res = subprocess.run(
102      args='adb shell ps -o pid:1,ppid:1,cmd:1,args:1',
103      capture_output=True,
104      shell=True,
105      check=True,
106      text=True,
107  )
108
109  proc_entries = []
110  for line in res.stdout.splitlines()[1:]:  # skip header
111    pid, ppid, cmd, name, *etc_args = line.split(' ')
112    entry = ProcEntry(int(pid), int(ppid), cmd, name, etc_args)
113    proc_entries.append(entry)
114
115  zygote_entry = next(e for e in proc_entries if e.name == args.zygote)
116  zygote_children = [e for e in proc_entries if e.ppid == zygote_entry.pid]
117
118  if args.process_names:
119    zygote_children = [e for e in proc_entries if e.name in args.process_names]
120
121  print('\n'.join(str(e.pid) + ' ' + e.name for e in zygote_children))
122
123  subprocess.run(
124      args=f'adb shell "mkdir -p {args.device_out_dir}"', check=True, shell=True
125  )
126  subprocess.run(args=f'mkdir -p {args.host_out_dir}', check=True, shell=True)
127
128  for entry in tqdm(zygote_children):
129    get_mem_stats(
130        zygote_pid=entry.ppid,
131        target_pid=entry.pid,
132        target_name=entry.name,
133        imgdiag_path=args.imgdiag,
134        boot_image=args.boot_image,
135        device_out_dir=args.device_out_dir,
136        host_out_dir=args.host_out_dir,
137    )
138
139
140if __name__ == '__main__':
141  main()
142