src/dctv/collect.py

# Copyright (C) 2020 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.
"""Collect a trace from a device

This command is intended as a "reference implementation" of DCTV trace
collection.  For production use, we'll collect traces using Perfetto
and Traceur, or maybe atrace on older devices.  This command exists so
we can collect some basic traces in *exactly* the way DCTV needs and
so we can isolate trace configuration problems from DCTV problems.
"""
import logging
import subprocess
import time
import re
from shlex import quote
import os
from os.path import dirname, join as pjoin
import json
from math import ceil

log = logging.getLogger(__name__)

# TODO(dancol): provide a way for atrace to export these event
# categories somehow dynamically instead of mirroring them here.

CATEGORIES = {
  "sched": [
    "cgroup_attach_task",
    "cgroup_destroy_root",
    "cgroup_mkdir",
    "cgroup_release",
    "cgroup_remount",
    "cgroup_rename",
    "cgroup_rmdir",
    "cgroup_setup_root",
    "cgroup_transfer_tasks",
    "oom_score_adj_update",
    "sched_blocked_reason",

    # TODO(dancol): conditionally enable the evetn
    # on devices that have it; also analyze it
    # intelligently.
    #"sched_cpu_hotplug",

    "sched_pi_setprio",
    "sched_process_exit",
    "sched_switch",
    # sched_wakeup is useless
    #"sched_wakeup",
    "sched_waking",
    "task_newtask",
    "task_rename",
  ],
  "irq": [
    "softireq_exit",
    "softirq_entry",
    "tasklet_entry",
    "tasklet_exit",
  ],
  "freq": [
    "clk_disable",
    "clk_enable",
    "clk_set_rate",
    "clock_disable",
    "clock_enable",
    "clock_set_rate",
    "cpu_frequency",
    "cpu_frequency_limits",
  ],
  "membus": [
    "memory_bus",
  ],
  "idle": [
    "cpu_idle",
  ],
  "disk": [
    "block_rq_complete",
    "block_rq_issue",
    "ext4_da_write_begin",
    "ext4_da_write_end",
    "f2fs_sync_file_enter",
    "f2fs_sync_file_exit",
    "f2fs_write_begin",
    "f2fs_write_end",
  ],
  "load": [
    "cpufreq_interactive",
  ],
  "workq": [
    "workqueue_activate_work",
    "workqueue_execute_end",
    "workqueue_execute_start",
    "workqueue_queue_work",
  ],
  "memreclaim": [
    "mm_vmscan_direct_reclaim_begin",
    "mm_vmscan_direct_reclaim_end",
    "mm_vmscan_kswapd_wake",
    "mm_vmscan_kswapd_sleep",
  ],
  "regulators": [
    "regulator_disable",
    "regulator_disable_complete",
    "regulator_enable",
    "regulator_enable_complete",
    "regulator_enable_delay",
    "regulator_set_voltage",
    "regulator_set_voltage_complete",
  ],
  "binder_driver": [
    "binder_set_priority",
    "binder_transaction",
    "binder_transaction_alloc_buf",
    "binder_transaction_received",
  ],
  "binder_lock": [
    "binder_lock",
    "binder_locked",
    "binder_unlock",
  ],
  "pagecache": [
    "mm_filemap_add_to_page_cache",
    "mm_filemap_delete_from_page_cache",
  ],
  "memory": [
    "rss_stat",
    "ion_heap_grow",
    "ion_heap_shrink",
  ],
}

EVENTLESS_CATEGORIES = {
  "adb",
  "aidl",
  "am",
  "audio",
  "bionic",
  "camera",
  "core_services",
  "dalvik",
  "database",
  "gfx",
  "hal",
  "i2c",
  "input",
  "ion",
  "network",
  "nnapi",
  "pdx",
  "pm",
  "power",
  "res",
  "rro",
  "rs",
  "sm",
  "ss",
  "sync",
  "vibrator",
  "video",
  "view",
  "webview",
  "wm",
}

for cat in EVENTLESS_CATEGORIES:
  CATEGORIES[cat] = ()

class AdbError(RuntimeError):
  """Exception raised when ADB fails"""

def run_adb_shell(*args):
  """Run an ADB shell command on device

  ARGS is a sequence of arguments, starting with the command name.
  Return the UTF-8 decoded standard output of the command with any
  final newline stripped, just like the shell does with $().  If the
  command fails, instead throw AdbError with the exception message
  incorporating the stderr output of the command.
  """
  args = ["adb", "shell"] + list(args)
  out = subprocess.run(
    args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  if out.returncode:
    raise AdbError(out.stderr.decode("UTF-8"))
  return out.stdout.decode("UTF-8").strip()

RESET_TRACE_STATE_SCRIPT = """\
set -e
cd /d/tracing
echo 0 > tracing_on
for option in options/*; do
  echo 0 > $option
done
echo nop > current_tracer
echo 0 > tracing_thresh
echo 0 > events/enable
echo ff > tracing_cpumask
echo 1 > buffer_size_kb
echo local > trace_clock
echo > trace
atrace --only_userspace --async_stop
"""

def reset_trace_state():
  """Reset device trace configuration to known state"""
  log.info("Resetting ftrace configuration to pristine")
  run_adb_shell(RESET_TRACE_STATE_SCRIPT)

def _make_start_tracing_script(events,
                               atrace_categories,
                               atrace_apps):
  atrace_args = []
  if atrace_apps:
    atrace_args.append("-a")
    atrace_args.append(",".join(atrace_apps))
  atrace_args.extend(atrace_categories)
  lines = [
    "set -e -x",
    "cd /d/tracing",
    "echo 1 > options/context-info",
    "echo 1 > options/markers",
    "echo 1 > options/trace_printk",
    "echo 1 > options/irq-info",
    "echo 128 > buffer_size_kb",
    "echo boot > trace_clock",
    "echo '{}' > set_event". format(" ".join(events)),
    # Need tracing_on before atrace so that trace_marker writes work
    "echo 1 > tracing_on",
    "atrace --only_userspace --async_start {}".format(
      " ".join(map(quote, atrace_args))),
  ]
  script = "\n".join(lines) + "\n"
  with open(pjoin(dirname(__file__), "snapshot.sh")) as snapshot_script:
    script += snapshot_script.read()
  return script

# We write this header to trace files to reserve space for the footer
# size and to mark the file incomplete in case of casual examination.
# We overwrite it the length of the footer on success.
INCOMPLETE_TRACE_HEADER = b"# DCTV INCOMPLETE TRACE"

def _collect_trace_1(enabled_events,
                     atrace_categories,
                     atrace_apps):
  log.info("Enabling tracing on device")
  log.info("Collecting snapshot...")
  snapshot = run_adb_shell(
    _make_start_tracing_script(enabled_events,
                               atrace_categories,
                               atrace_apps))
  log.info("... snapshot collection complete")
  log.info("Collecting system information...")
  sysinfo = {
    "cpuinfo": run_adb_shell("cat /proc/cpuinfo"),
    "kernel_version": run_adb_shell("cat /proc/version"),
  }
  log.info("... system information collection complete")
  try:
    log.info("Trace running: use control-C to stop trace")
    while True:
      time.sleep(10000)  # No support for pause(2)?
  except KeyboardInterrupt:
    log.info("Saw control-C: stopping trace")
  return snapshot, sysinfo

def _ensure_root():
  uid = int(run_adb_shell("id", "-u"))
  if uid:
    log.info("running adb root")
    subprocess.check_call(["adb", "root"])
    uid = int(run_adb_shell("id", "-u"))
    if uid:
      raise ValueError("adb root didn't work: uid={!r}".format(uid))

def _collect_trace(trace_file,
                   enabled_events,
                   atrace_categories,
                   atrace_apps):

  trace_file.write(INCOMPLETE_TRACE_HEADER + b"\n")
  trace_file.flush()
  copier = subprocess.Popen(["adb", "shell",
                             "cat", "/d/tracing/trace_pipe"],
                            stdout=trace_file)
  try:
    snapshot, sysinfo = _collect_trace_1(enabled_events,
                                         atrace_categories,
                                         atrace_apps)
  finally:
    copier.kill()
    copier.wait()
    reset_trace_state()
  log.info("Collected %dkB of raw trace data into %r",
           int(ceil(os.fstat(trace_file.fileno()).st_size / 1024)),
           trace_file.name)
  log.info("Appending trace metadata")
  # Make sure the blob has a terminating newline so that editors don't
  # munge the trace by helpfully inserting a trailing newline and
  # wrecking the footer length field at the top.
  metadata = {
    "version": "dctv-1",
    "enabled_events": sorted(enabled_events),
    "atrace_categories": sorted(atrace_categories),
    "atrace_apps": sorted(atrace_apps),
    "snapshots": [snapshot],
  }
  metadata.update(sysinfo)
  metadata_blob = b"# " + json.dumps(metadata).encode("UTF-8") + b"\n"
  trace_file.write(metadata_blob)
  log.info("Finalizing file")
  header = b"# DCTV %d" % len(metadata_blob)
  assert len(header) <= len(INCOMPLETE_TRACE_HEADER)
  trace_file.seek(0)
  trace_file.write(header.ljust(len(INCOMPLETE_TRACE_HEADER)))
  log.info("Success: trace file %r ready to load",
           trace_file.name)

def _get_available_events():
  ev_fix_re = re.compile(r"^[^:]*:")
  return set(
    ev_fix_re.sub("", e)
    for e in run_adb_shell(
      "cat", "/d/tracing/available_events").split())

def _get_available_categories():
  cat_fix_re = re.compile("^[ \t]*([a-zA-Z0-9_]+) - .*$")
  def _fix_cat(c):
    assert cat_fix_re.match(c), \
      "bad category line: {!r}".format(c)
    return cat_fix_re.sub("\\1", c)
  return set(
    _fix_cat(c)
    for c in run_adb_shell("atrace", "--list_categories").split("\n"))

def run_collect(args):
  """Trace collection subcommand"""
  available_events = _get_available_events()
  available_categories = _get_available_categories()
  unknown_categories = available_categories - set(CATEGORIES)
  if unknown_categories:
    log.warning("unknown categories: %r",
                sorted(available_categories -
                       set(CATEGORIES)))
  if args.list_categories:
    for catx in sorted(available_categories):
      print(catx)
    return

  enabled_events = set()
  atrace_categories = args.categories
  for category in atrace_categories:
    enabled_events.update(CATEGORIES[category])
  if not enabled_events.issubset(available_events):
    log.error("requested events unavailable: %s",
              enabled_events - available_events)
  _ensure_root()
  open_mode = "w+b" if args.force else "x+b"
  with open(args.output, open_mode) as trace_file:
    try:
      reset_trace_state()
      log.info("Recording trace to %r", trace_file.name)
      _collect_trace(trace_file,
                     enabled_events,
                     atrace_categories,
                     [app for app in args.apps.split(",") if app])
    except:
      log.warning("Failed to collect trace: deleting %r",
                  trace_file.name)
      os.unlink(trace_file.name)
      raise