tools/perf/profile_creators/profile_generator.py

   1 # Copyright 2014 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Handles generating profiles and transferring them to/from mobile devices."""
   6
   7 import logging
   8 import optparse
   9 import os
  10 import shutil
  11 import stat
  12 import sys
  13 import tempfile
  14
  15 from profile_creators import profile_extender
  16 from telemetry.core import discover
  17 from telemetry.core import util
  18 from telemetry.internal.browser import browser_finder
  19 from telemetry.internal.browser import browser_options
  20 from telemetry.internal import story_runner
  21
  22
  23 def _DiscoverProfileExtenderClasses():
  24   profile_extenders_dir = os.path.abspath(os.path.join(util.GetBaseDir(),
  25       os.pardir, 'perf', 'profile_creators'))
  26   base_dir = os.path.abspath(os.path.join(profile_extenders_dir, os.pardir))
  27
  28   profile_extenders_unfiltered = discover.DiscoverClasses(
  29       profile_extenders_dir, base_dir, profile_extender.ProfileExtender)
  30
  31   # Remove 'extender' suffix from keys.
  32   profile_extenders = {}
  33   for test_name, test_class in profile_extenders_unfiltered.iteritems():
  34     assert test_name.endswith('_extender')
  35     test_name = test_name[:-len('_extender')]
  36     profile_extenders[test_name] = test_class
  37   return profile_extenders
  38
  39
  40 def _IsPseudoFile(directory, paths):
  41   """Filter function for shutil.copytree() to reject socket files and symlinks
  42   since those can't be copied around on bots."""
  43   def IsSocket(full_path):
  44     """Check if a file at a given path is a socket."""
  45     try:
  46       if stat.S_ISSOCK(os.stat(full_path).st_mode):
  47         return True
  48     except OSError:
  49       # Thrown if we encounter a broken symlink.
  50       pass
  51     return False
  52
  53   ignore_list = []
  54   for path in paths:
  55     full_path = os.path.join(directory, path)
  56
  57     if os.path.isdir(full_path):
  58       continue
  59     if not IsSocket(full_path) and not os.path.islink(full_path):
  60       continue
  61
  62     logging.warning('Ignoring pseudo file: %s' % full_path)
  63     ignore_list.append(path)
  64
  65   return ignore_list
  66
  67
  68 class ProfileGenerator(object):
  69   """Generate profile.
  70
  71   On desktop the generated profile is copied to the specified location so later
  72   runs can reuse it.
  73   On CrOS profile resides on cryptohome and there is no easy way to
  74   override it before user login. So for CrOS we just generate the profile
  75   every time when the benchmark starts to run.
  76   """
  77   def __init__(self, profile_extender_class, profile_name):
  78     self._profile_extender_class = profile_extender_class
  79     self._profile_name = profile_name
  80
  81   def Run(self, options):
  82     """Kick off the process.
  83
  84     Args:
  85       options: Instance of BrowserFinderOptions to search for proper browser.
  86
  87     Returns:
  88       The path of generated profile or existing profile if --profile-dir
  89       is given. Could be None if it's generated on default location
  90       (e.g., crytohome on CrOS).
  91     """
  92     possible_browser = browser_finder.FindBrowser(options)
  93     is_cros = possible_browser.browser_type.startswith('cros')
  94
  95     out_dir = None
  96     if is_cros:
  97       self.Create(options, None)
  98     else:
  99       # Decide profile output path.
 100       out_dir = (options.browser_options.profile_dir or
 101           os.path.abspath(os.path.join(
 102               tempfile.gettempdir(), self._profile_name, self._profile_name)))
 103       if not os.path.exists(out_dir):
 104         self.Create(options, out_dir)
 105
 106     return out_dir
 107
 108   def Create(self, options, out_dir):
 109     """Generate profile.
 110
 111     If out_dir is given, copy the generated profile to out_dir.
 112     Otherwise the profile is generated to its default position
 113     (e.g., cryptohome on CrOS).
 114     """
 115
 116     # Leave the global options intact.
 117     creator_options = options.Copy()
 118
 119     if out_dir:
 120       sys.stderr.write('Generating profile to: %s \n' % out_dir)
 121       # The genrated profile is copied to out_dir only if the generation is
 122       # successful. In the generation process a temp directory is used so
 123       # the default profile is not polluted on failure.
 124       tmp_profile_path = tempfile.mkdtemp()
 125       creator_options.output_profile_path = tmp_profile_path
 126
 127     creator = self._profile_extender_class(creator_options)
 128
 129     try:
 130       creator.Run()
 131     except Exception as e:
 132       logging.exception('Profile creation failed.')
 133       raise e
 134     else:
 135       sys.stderr.write('SUCCESS: Profile generated.\n')
 136
 137       # Copy generated profile to final destination if out_dir is given.
 138       if out_dir:
 139         if os.path.exists(out_dir):
 140           shutil.rmtree(out_dir)
 141         shutil.copytree(tmp_profile_path,
 142                         out_dir, ignore=_IsPseudoFile)
 143         sys.stderr.write(
 144             "SUCCESS: Generated profile copied to: '%s'.\n" % out_dir)
 145     finally:
 146       if out_dir:
 147         shutil.rmtree(tmp_profile_path)
 148
 149
 150 def AddCommandLineArgs(parser):
 151   story_runner.AddCommandLineArgs(parser)
 152
 153   profile_extenders = _DiscoverProfileExtenderClasses().keys()
 154   legal_profile_creators = '|'.join(profile_extenders)
 155   group = optparse.OptionGroup(parser, 'Profile generation options')
 156   group.add_option('--profile-type-to-generate',
 157       dest='profile_type_to_generate',
 158       default=None,
 159       help='Type of profile to generate. '
 160            'Supported values: %s' % legal_profile_creators)
 161   parser.add_option_group(group)
 162
 163
 164 def ProcessCommandLineArgs(parser, args):
 165   story_runner.ProcessCommandLineArgs(parser, args)
 166
 167   if not args.profile_type_to_generate:
 168     parser.error("Must specify --profile-type-to-generate option.")
 169
 170   profile_extenders = _DiscoverProfileExtenderClasses().keys()
 171   if args.profile_type_to_generate not in profile_extenders:
 172     legal_profile_creators = '|'.join(profile_extenders)
 173     parser.error("Invalid profile type, legal values are: %s." %
 174         legal_profile_creators)
 175
 176   if not args.browser_type:
 177     parser.error("Must specify --browser option.")
 178
 179   if not args.output_dir:
 180     parser.error("Must specify --output-dir option.")
 181
 182   if args.browser_options.dont_override_profile:
 183     parser.error("Can't use existing profile when generating profile.")
 184
 185
 186 def Main():
 187   options = browser_options.BrowserFinderOptions()
 188   parser = options.CreateParser(
 189       "%%prog <--profile-type-to-generate=...> <--browser=...> <--output-dir>")
 190   AddCommandLineArgs(parser)
 191   _, _ = parser.parse_args()
 192   ProcessCommandLineArgs(parser, options)
 193
 194   # Generate profile.
 195   profile_extenders = _DiscoverProfileExtenderClasses()
 196   profile_extender_class = profile_extenders[options.profile_type_to_generate]
 197
 198   generator = ProfileGenerator(profile_extender_class,
 199       options.profile_type_to_generate)
 200   generator.Create(options, options.output_dir)
 201   return 0