startop: Rewrite the perfetto trace collection part.

Test: pytest perfetto_trace_collector_test.py
Bug: 138233615
Change-Id: If13d895029e734a5e52bed73c5f870bb3f036c2f

3 files changed, 278 insertions, 1 deletions

diff --git a/startop/scripts/app_startup/lib/adb_utils.py b/startop/scripts/app_startup/lib/adb_utils.py
index 0e0065defd7f..1c60a17ada1b 100644
--- a/startop/scripts/app_startup/lib/adb_utils.py
+++ b/startop/scripts/app_startup/lib/adb_utils.py
@@ -104,4 +104,21 @@ def blocking_wait_for_logcat_displayed_time(timestamp: datetime.datetime,
     return None
   displayed_time = result[result.rfind('+'):]
 
-  return parse_time_to_milliseconds(displayed_time)
\ No newline at end of file
+  return parse_time_to_milliseconds(displayed_time)
+
+def delete_file_on_device(file_path: str) -> None:
+  """ Deletes a file on the device. """
+  cmd_utils.run_adb_shell_command(
+    "[[ -f '{file_path}' ]] && rm -f '{file_path}' || "
+    "exit 0".format(file_path=file_path))
+
+def set_prop(property: str, value: str) -> None:
+  """ Sets property using adb shell. """
+  cmd_utils.run_adb_shell_command('setprop "{property}" "{value}"'.format(
+      property=property, value=value))
+
+def pull_file(device_file_path: str, output_file_path: str) -> None:
+  """ Pulls file from device to output """
+  cmd_utils.run_shell_command('adb pull "{device_file_path}" "{output_file_path}"'.
+      format(device_file_path=device_file_path,
+             output_file_path=output_file_path))
diff --git a/startop/scripts/app_startup/lib/perfetto_trace_collector.py b/startop/scripts/app_startup/lib/perfetto_trace_collector.py
new file mode 100644
index 000000000000..20a6c5280c81
--- /dev/null
+++ b/startop/scripts/app_startup/lib/perfetto_trace_collector.py
@@ -0,0 +1,159 @@
+# Copyright 2019, 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.
+
+"""Class to collector perfetto trace."""
+import datetime
+from datetime import timedelta
+import os
+import re
+import sys
+import time
+from typing import Optional, List, Tuple
+
+# global variables
+DIR = os.path.abspath(os.path.dirname(__file__))
+
+sys.path.append(os.path.dirname(os.path.dirname(DIR)))
+
+import app_startup.lib.adb_utils as adb_utils
+from app_startup.lib.app_runner import AppRunner, AppRunnerListener
+import lib.print_utils as print_utils
+import lib.logcat_utils as logcat_utils
+import iorap.lib.iorapd_utils as iorapd_utils
+
+class PerfettoTraceCollector(AppRunnerListener):
+  """ Class to collect perfetto trace.
+
+      To set trace duration of perfetto, change the 'trace_duration_ms'.
+      To pull the generated perfetto trace on device, set the 'output'.
+  """
+  TRACE_FILE_SUFFIX = 'perfetto_trace.pb'
+  TRACE_DURATION_PROP = 'iorapd.perfetto.trace_duration_ms'
+  SECONDS_TO_MILLISECONDS = 1000
+
+  def __init__(self,
+               package: str,
+               activity: Optional[str],
+               compiler_filter: Optional[str],
+               timeout: Optional[int],
+               simulate: bool,
+               trace_duration: timedelta = timedelta(milliseconds=5000),
+               save_destination_file_path: Optional[str] = None):
+    """ Initialize the perfetto trace collector. """
+    self.app_runner = AppRunner(package,
+                                activity,
+                                compiler_filter,
+                                timeout,
+                                simulate)
+    self.app_runner.add_callbacks(self)
+
+    self.trace_duration = trace_duration
+    self.save_destination_file_path = save_destination_file_path
+
+  def purge_file(self, suffix: str) -> None:
+    print_utils.debug_print('iorapd-perfetto: purge file in ' +
+                            self._get_remote_path())
+    adb_utils.delete_file_on_device(self._get_remote_path())
+
+  def run(self) -> Optional[List[Tuple[str]]]:
+    """Runs an app.
+
+    Returns:
+      A list of (metric, value) tuples.
+    """
+    return self.app_runner.run()
+
+  def preprocess(self):
+    # Sets up adb environment.
+    adb_utils.root()
+    adb_utils.disable_selinux()
+    time.sleep(1)
+
+    # Kill any existing process of this app
+    adb_utils.pkill(self.app_runner.package)
+
+    # Remove existing trace and compiler files
+    self.purge_file(PerfettoTraceCollector.TRACE_FILE_SUFFIX)
+
+    # Set perfetto trace duration prop to milliseconds.
+    adb_utils.set_prop(PerfettoTraceCollector.TRACE_DURATION_PROP,
+                       int(self.trace_duration.total_seconds()*
+                           PerfettoTraceCollector.SECONDS_TO_MILLISECONDS))
+
+    if not iorapd_utils.stop_iorapd():
+      raise RuntimeError('Cannot stop iorapd!')
+
+    if not iorapd_utils.enable_iorapd_perfetto():
+      raise RuntimeError('Cannot enable perfetto!')
+
+    if not iorapd_utils.disable_iorapd_readahead():
+      raise RuntimeError('Cannot disable readahead!')
+
+    if not iorapd_utils.start_iorapd():
+      raise RuntimeError('Cannot start iorapd!')
+
+    # Drop all caches to get cold starts.
+    adb_utils.vm_drop_cache()
+
+  def postprocess(self, pre_launch_timestamp: str):
+    # Kill any existing process of this app
+    adb_utils.pkill(self.app_runner.package)
+
+    iorapd_utils.disable_iorapd_perfetto()
+
+    if self.save_destination_file_path is not None:
+      adb_utils.pull_file(self._get_remote_path(),
+                          self.save_destination_file_path)
+
+  def metrics_selector(self, am_start_output: str,
+                       pre_launch_timestamp: str) -> str:
+    """Parses the metric after app startup by reading from logcat in a blocking
+    manner until all metrics have been found".
+
+    Returns:
+      An empty string.
+    """
+    if not self._wait_for_perfetto_trace(pre_launch_timestamp):
+      raise RuntimeError('Could not save perfetto app trace file!')
+
+    return ''
+
+  def _wait_for_perfetto_trace(self, pre_launch_timestamp) -> Optional[str]:
+    """ Waits for the perfetto trace being saved to file.
+
+    The string is in the format of r".*Perfetto TraceBuffer saved to file:
+    <file path>.*"
+
+    Returns:
+      the string what the program waits for. If the string doesn't show up,
+      return None.
+    """
+    pattern = re.compile(r'.*Perfetto TraceBuffer saved to file: {}.*'.
+                         format(self._get_remote_path()))
+
+    # The pre_launch_timestamp is longer than what the datetime can parse. Trim
+    # last three digits to make them align.
+    timestamp = datetime.datetime.strptime(pre_launch_timestamp[:-3],
+                                           '%Y-%m-%d %H:%M:%S.%f')
+    timeout_dt = timestamp + datetime.timedelta(0, self.app_runner.timeout)
+
+    return logcat_utils.blocking_wait_for_logcat_pattern(timestamp,
+                                                         pattern,
+                                                         timeout_dt)
+
+  def _get_remote_path(self):
+    # For example: android.music%2Fmusic.TopLevelActivity.perfetto_trace.pb
+    return iorapd_utils._iorapd_path_to_data_file(self.app_runner.package,
+                                                  self.app_runner.activity,
+                                                  PerfettoTraceCollector.TRACE_FILE_SUFFIX)
diff --git a/startop/scripts/app_startup/lib/perfetto_trace_collector_test.py b/startop/scripts/app_startup/lib/perfetto_trace_collector_test.py
new file mode 100644
index 000000000000..8d94fc58bede
--- /dev/null
+++ b/startop/scripts/app_startup/lib/perfetto_trace_collector_test.py
@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+#
+# Copyright 2019, 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.
+#
+
+"""Unit tests for the data_frame.py script."""
+import os
+import sys
+from pathlib import Path
+from datetime import timedelta
+
+from mock import call, patch
+from perfetto_trace_collector import PerfettoTraceCollector
+
+sys.path.append(Path(os.path.realpath(__file__)).parents[2])
+from app_startup.lib.app_runner import AppRunner
+
+RUNNER = PerfettoTraceCollector(package='music',
+                                activity='MainActivity',
+                                compiler_filter=None,
+                                timeout=10,
+                                simulate=False,
+                                trace_duration = timedelta(milliseconds=1000),
+                                # No actual file will be created. Just to
+                                # check the command.
+                                save_destination_file_path='/tmp/trace.pb')
+
+def _mocked_run_shell_command(*args, **kwargs):
+  if args[0] == 'adb shell ps | grep "music" | awk \'{print $2;}\'':
+    return (True, '9999')
+  else:
+    return (True, '')
+
+@patch('lib.logcat_utils.blocking_wait_for_logcat_pattern')
+@patch('lib.cmd_utils.run_shell_command')
+def test_perfetto_trace_collector_preprocess(mock_run_shell_command,
+                                             mock_blocking_wait_for_logcat_pattern):
+  mock_run_shell_command.side_effect = _mocked_run_shell_command
+  mock_blocking_wait_for_logcat_pattern.return_value = "Succeed!"
+
+  RUNNER.preprocess()
+
+  calls = [call('adb root'),
+           call('adb shell "getenforce"'),
+           call('adb shell "setenforce 0"'),
+           call('adb shell "stop"'),
+           call('adb shell "start"'),
+           call('adb wait-for-device'),
+           call('adb shell ps | grep "music" | awk \'{print $2;}\''),
+           call('adb shell "kill 9999"'),
+           call(
+               'adb shell "[[ -f \'/data/misc/iorapd/music%2FMainActivity.perfetto_trace.pb\' ]] '
+               '&& rm -f \'/data/misc/iorapd/music%2FMainActivity.perfetto_trace.pb\' || exit 0"'),
+           call('adb shell "setprop "iorapd.perfetto.trace_duration_ms" "1000""'),
+           call(
+               'bash -c "source {}; iorapd_stop"'.format(
+                   AppRunner.IORAP_COMMON_BASH_SCRIPT)),
+           call(
+               'bash -c "source {}; iorapd_perfetto_enable"'.format(
+                   AppRunner.IORAP_COMMON_BASH_SCRIPT)),
+           call(
+               'bash -c "source {}; iorapd_readahead_disable"'.format(
+                   AppRunner.IORAP_COMMON_BASH_SCRIPT)),
+           call(
+               'bash -c "source {}; iorapd_start"'.format(
+                   AppRunner.IORAP_COMMON_BASH_SCRIPT)),
+           call('adb shell "echo 3 > /proc/sys/vm/drop_caches"')]
+
+  mock_run_shell_command.assert_has_calls(calls)
+
+@patch('lib.logcat_utils.blocking_wait_for_logcat_pattern')
+@patch('lib.cmd_utils.run_shell_command')
+def test_perfetto_trace_collector_postprocess(mock_run_shell_command,
+                                              mock_blocking_wait_for_logcat_pattern):
+  mock_run_shell_command.side_effect = _mocked_run_shell_command
+  mock_blocking_wait_for_logcat_pattern.return_value = "Succeed!"
+
+  RUNNER.postprocess('2019-07-02 23:20:06.972674825')
+
+  calls = [call('adb shell ps | grep "music" | awk \'{print $2;}\''),
+           call('adb shell "kill 9999"'),
+           call(
+               'bash -c "source {}; iorapd_perfetto_disable"'.format(
+                   AppRunner.IORAP_COMMON_BASH_SCRIPT)),
+           call('adb pull '
+                '"/data/misc/iorapd/music%2FMainActivity.perfetto_trace.pb" '
+                '"/tmp/trace.pb"')]
+
+  mock_run_shell_command.assert_has_calls(calls)