Add scripts for creating dirty-image-objects

Test: ./art/imgdiag/run_imgdiag.py
Test: ./art/imgdiag/create_dirty_image_objects.py ./imgdiag_*
Change-Id: I7ff85e5998973544e89262894746c50b39a5c081

3 files changed, 313 insertions, 0 deletions

diff --git a/imgdiag/create_dirty_image_objects.py b/imgdiag/create_dirty_image_objects.py
new file mode 100755
index 0000000000..d090fcc2d1
--- /dev/null
+++ b/imgdiag/create_dirty_image_objects.py
@@ -0,0 +1,131 @@
+#! /usr/bin/env python3
+#
+# Copyright 2023, The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from collections import defaultdict
+import os
+
+
+def process_dirty_entries(entries, with_sort):
+  mark_counts = defaultdict(int)
+  dirty_image_objects = []
+
+  union = set()
+  for v in entries.values():
+    union = union.union(v)
+
+  for obj in union:
+    str_marker = ''
+    marker = 0
+    # Sort marker is uint32_t, where Nth bit is set if Nth process has this object dirty.
+    for idx, v in enumerate(entries.values()):
+      if obj in v:
+        str_marker += chr(ord('A') + idx)
+        marker = (marker << 1) | 1
+      else:
+        str_marker += '_'
+        marker = marker << 1
+
+    if with_sort:
+      dirty_image_objects.append(obj + ' ' + str(marker) + '\n')
+    else:
+      dirty_image_objects.append(obj + '\n')
+
+    mark_counts[str_marker] += 1
+
+  return (dirty_image_objects, mark_counts)
+
+
+def main():
+  parser = argparse.ArgumentParser(
+      description=(
+          'Create dirty-image-objects file from specified imgdiag output files.'
+      ),
+      formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+  )
+  parser.add_argument(
+      'imgdiag_files',
+      nargs='+',
+      help='imgdiag files to use.',
+  )
+  parser.add_argument(
+      '--sort-objects',
+      action=argparse.BooleanOptionalAction,
+      default=True,
+      help='Use object sorting.',
+  )
+  parser.add_argument(
+      '--output-filename',
+      default='dirty-image-objects.txt',
+      help='Output file for dirty image objects.',
+  )
+  parser.add_argument(
+      '--print-stats',
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help='Print dirty object stats.',
+  )
+
+  args = parser.parse_args()
+
+  entries = dict()
+  for path in args.imgdiag_files:
+    with open(path) as f:
+      lines = f.readlines()
+    prefix = 'dirty_obj: '
+    lines = [l.strip().removeprefix(prefix) for l in lines if prefix in l]
+    entries[path] = set(lines)
+
+  if args.sort_objects and len(entries) > 32:
+    print(
+        'WARNING: too many processes for sorting, using top 32 by number of'
+        ' dirty objects.'
+    )
+    entries_list = sorted(
+        list(entries.items()), reverse=True, key=lambda x: len(x[1])
+    )
+    entries_list = entries_list[0:32]
+    entries = {k: v for (k, v) in entries_list}
+
+  print('Using processes:')
+  for k, v in sorted(entries.items(), key=lambda x: len(x[1])):
+    print(f'{k}: {len(v)}')
+  print()
+
+  dirty_image_objects, mark_counts = process_dirty_entries(
+      entries=entries, with_sort=args.sort_objects
+  )
+
+  with open(args.output_filename, 'w') as f:
+    f.writelines(dirty_image_objects)
+
+  if args.print_stats:
+    mark_counts = sorted(
+        list(mark_counts.items()), key=lambda x: x[1], reverse=True
+    )
+
+    for i, path in enumerate(entries.keys()):
+      print(path, chr(ord('A') + i))
+
+    total_count = 0
+    for marker, count in mark_counts:
+      print(marker, count)
+      total_count += count
+    print('total: ', total_count)
+
+
+if __name__ == '__main__':
+  main()
diff --git a/imgdiag/dirty_image_objects.md b/imgdiag/dirty_image_objects.md
new file mode 100644
index 0000000000..bf020691f9
--- /dev/null
+++ b/imgdiag/dirty_image_objects.md
@@ -0,0 +1,62 @@
+# How to update dirty-image-objects
+
+1. Add `imgdiag` to ART APEX.
+
+The easiest way is to modify `art/build/apex/Android.bp` like this:
+```
+ art_runtime_binaries_both = [
+     "dalvikvm",
+     "dex2oat",
++    "imgdiag",
+ ]
+```
+
+2. Install ART APEX and reboot, e.g.:
+
+```
+m apps_only dist
+adb install out/dist/com.android.art.apex
+adb reboot
+```
+
+3. Collect imgdiag output.
+
+```
+# To see all options check: art/imgdiag/run_imgdiag.py -h
+
+art/imgdiag/run_imgdiag.py
+```
+
+4. Create new dirty-image-objects.
+
+```
+# To see all options check: art/imgdiag/create_dirty_image_objects.py -h
+
+# Using all imgdiag files:
+art/imgdiag/create_dirty_image_objects.py ./imgdiag_*
+
+# Or using only specified files:
+art/imgdiag/create_dirty_image_objects.py \
+  ./imgdiag_system_server.txt \
+  ./imgdiag_com.android.systemui.txt \
+  ./imgdiag_com.google.android.gms.txt \
+  ./imgdiag_com.google.android.gms.persistent.txt \
+  ./imgdiag_com.google.android.gms.ui.txt \
+  ./imgdiag_com.google.android.gms.unstable.txt
+```
+
+5. Push new dirty-image-objects to the device.
+
+```
+adb push dirty-image-objects.txt /etc/dirty-image-objects
+```
+
+6. Reinstall ART APEX to update the boot image.
+
+```
+adb install out/dist/com.android.art.apex
+adb reboot
+```
+
+At this point the device should have new `boot.art` with optimized dirty object layout.
+This can be checked by collecting imgdiag output again and comparing dirty page counts to the previous run.
diff --git a/imgdiag/run_imgdiag.py b/imgdiag/run_imgdiag.py
new file mode 100755
index 0000000000..6d035540be
--- /dev/null
+++ b/imgdiag/run_imgdiag.py
@@ -0,0 +1,120 @@
+#! /usr/bin/env python3
+#
+# Copyright 2023, The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from collections import namedtuple
+import subprocess
+
+try:
+  from tqdm import tqdm
+except:
+  def tqdm(x):
+    return x
+
+
+ProcEntry = namedtuple('ProcEntry', 'pid, ppid, cmd, name, etc_args')
+
+def get_mem_stats(zygote_pid, target_pid, target_name, imgdiag_path, boot_image, device_out_dir):
+  imgdiag_output_path = f'{device_out_dir}/imgdiag_{target_name}.txt'
+  cmd_collect = (
+      'adb shell '
+      f'"{imgdiag_path} --zygote-diff-pid={zygote_pid} --image-diff-pid={target_pid} '
+      f'--output={imgdiag_output_path} --boot-image={boot_image} --dump-dirty-objects"'
+  )
+
+  try:
+    subprocess.run(cmd_collect, shell=True, check=True)
+  except:
+    print('imgdiag call failed on:', target_pid, target_name)
+    return
+
+  cmd_pull = f'adb pull {imgdiag_output_path} ./'
+  subprocess.run(cmd_pull, shell=True, check=True, capture_output=True)
+
+
+def main():
+  parser = argparse.ArgumentParser(
+      description='Run imgdiag on selected processes and pull results from the device.',
+      formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+  )
+  parser.add_argument(
+      'process_names',
+      nargs='*',
+      help='Process names to use. If none - dump all zygote children.',
+  )
+  parser.add_argument(
+      '--boot-image',
+      dest='boot_image',
+      default='/data/misc/apexdata/com.android.art/dalvik-cache/boot.art',
+      help='Path to boot.art',
+  )
+  parser.add_argument(
+      '--zygote',
+      default='zygote64',
+      help='Zygote process name',
+  )
+  parser.add_argument(
+      '--imgdiag',
+      default='/apex/com.android.art/bin/imgdiag64',
+      help='Path to imgdiag binary.',
+  )
+  parser.add_argument(
+      '--device-out-dir',
+      default='/data/local/tmp/imgdiag_out',
+      help='Directory for imgdiag output files on the device.',
+  )
+
+  args = parser.parse_args()
+
+  res = subprocess.run(
+      args='adb shell ps -o pid:1,ppid:1,cmd:1,args:1',
+      capture_output=True,
+      shell=True,
+      check=True,
+      text=True,
+  )
+
+  proc_entries = []
+  for line in res.stdout.splitlines()[1:]:  # skip header
+    pid, ppid, cmd, name, *etc_args = line.split(' ')
+    entry = ProcEntry(int(pid), int(ppid), cmd, name, etc_args)
+    proc_entries.append(entry)
+
+  zygote_entry = next(e for e in proc_entries if e.name == args.zygote)
+  zygote_children = [e for e in proc_entries if e.ppid == zygote_entry.pid]
+
+  if args.process_names:
+    zygote_children = [e for e in proc_entries if e.name in args.process_names]
+
+  print('\n'.join(str(e.pid) + ' ' + e.name for e in zygote_children))
+
+  subprocess.run(
+      args=f'adb shell "mkdir -p {args.device_out_dir}"', check=True, shell=True
+  )
+
+  for entry in tqdm(zygote_children):
+    get_mem_stats(
+        zygote_pid=entry.ppid,
+        target_pid=entry.pid,
+        target_name=entry.name,
+        imgdiag_path=args.imgdiag,
+        boot_image=args.boot_image,
+        device_out_dir=args.device_out_dir,
+    )
+
+
+if __name__ == '__main__':
+    main()