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
5import logging
6import os
7import shutil
8import sys
9import tempfile
10import zipfile
11
12from catapult_base import cloud_storage  # pylint: disable=import-error
13
14from telemetry.core import exceptions
15from telemetry.core import util
16from telemetry import decorators
17from telemetry.internal.browser import browser_finder
18from telemetry.internal.browser import browser_finder_exceptions
19from telemetry.internal.browser import browser_info as browser_info_module
20from telemetry.internal.platform.profiler import profiler_finder
21from telemetry.internal.util import exception_formatter
22from telemetry.internal.util import file_handle
23from telemetry.page import page_test
24from telemetry import story
25from telemetry.util import image_util
26from telemetry.util import wpr_modes
27from telemetry.web_perf import timeline_based_measurement
28
29
30def _PrepareFinderOptions(finder_options, test, device_type):
31  browser_options = finder_options.browser_options
32  # Set up user agent.
33  browser_options.browser_user_agent_type = device_type
34
35  test.CustomizeBrowserOptions(finder_options.browser_options)
36  if finder_options.profiler:
37    profiler_class = profiler_finder.FindProfiler(finder_options.profiler)
38    profiler_class.CustomizeBrowserOptions(browser_options.browser_type,
39                                           finder_options)
40
41
42class SharedPageState(story.SharedState):
43  """
44  This class contains all specific logic necessary to run a Chrome browser
45  benchmark.
46  """
47
48  _device_type = None
49
50  def __init__(self, test, finder_options, story_set):
51    super(SharedPageState, self).__init__(test, finder_options, story_set)
52    if isinstance(test, timeline_based_measurement.TimelineBasedMeasurement):
53      assert not finder_options.profiler, (
54          'This is a Timeline Based Measurement benchmark. You cannot run it '
55          'with the --profiler flag. If you need trace data, tracing is always '
56          ' enabled in Timeline Based Measurement benchmarks and you can get '
57          'the trace data by using --output-format=json.')
58      # This is to avoid the cyclic-import caused by timeline_based_page_test.
59      from telemetry.web_perf import timeline_based_page_test
60      self._test = timeline_based_page_test.TimelineBasedPageTest(test)
61    else:
62      self._test = test
63    device_type = self._device_type
64    # TODO(aiolos, nednguyen): Remove this logic of pulling out user_agent_type
65    # from story_set once all page_set are converted to story_set
66    # (crbug.com/439512).
67
68    def _IsPageSetInstance(s):
69      # This is needed to avoid importing telemetry.page.page_set which will
70      # cause cyclic import.
71      return 'PageSet' == s.__class__.__name__ or 'PageSet' in (
72          list(c.__name__ for c in s.__class__.__bases__))
73    if not device_type and _IsPageSetInstance(story_set):
74      device_type = story_set.user_agent_type
75    _PrepareFinderOptions(finder_options, self._test, device_type)
76    self._browser = None
77    self._finder_options = finder_options
78    self._possible_browser = self._GetPossibleBrowser(
79        self._test, finder_options)
80
81    self._first_browser = True
82    self._did_login_for_current_page = False
83    self._current_page = None
84    self._current_tab = None
85    self._migrated_profile = None
86
87    self._pregenerated_profile_archive_dir = None
88    self._test.SetOptions(self._finder_options)
89
90    # TODO(crbug/404771): Move network controller options out of
91    # browser_options and into finder_options.
92    browser_options = self._finder_options.browser_options
93    if self._finder_options.use_live_sites:
94      wpr_mode = wpr_modes.WPR_OFF
95    elif browser_options.wpr_mode == wpr_modes.WPR_RECORD:
96      wpr_mode = wpr_modes.WPR_RECORD
97    else:
98      wpr_mode = wpr_modes.WPR_REPLAY
99
100    self.platform.network_controller.Open(wpr_mode,
101                                          browser_options.extra_wpr_args)
102
103
104  @property
105  def browser(self):
106    return self._browser
107
108  def _FindBrowser(self, finder_options):
109    possible_browser = browser_finder.FindBrowser(finder_options)
110    if not possible_browser:
111      raise browser_finder_exceptions.BrowserFinderException(
112          'No browser found.\n\nAvailable browsers:\n%s\n' %
113          '\n'.join(browser_finder.GetAllAvailableBrowserTypes(finder_options)))
114    return possible_browser
115
116  def _GetPossibleBrowser(self, test, finder_options):
117    """Return a possible_browser with the given options for |test|. """
118    possible_browser = self._FindBrowser(finder_options)
119    finder_options.browser_options.browser_type = (
120        possible_browser.browser_type)
121
122    enabled, msg = decorators.IsEnabled(test, possible_browser)
123    if not enabled and not finder_options.run_disabled_tests:
124      logging.warning(msg)
125      logging.warning('You are trying to run a disabled test.')
126      logging.warning(
127          'Pass --also-run-disabled-tests to squelch this message.')
128      sys.exit(0)
129
130    if possible_browser.IsRemote():
131      possible_browser.RunRemote()
132      sys.exit(0)
133    return possible_browser
134
135  def _TryCaptureScreenShot(self, page, tab, results):
136    try:
137      # TODO(nednguyen): once all platforms support taking screenshot,
138      # remove the tab checking logic and consider moving this to story_runner.
139      # (crbug.com/369490)
140      if tab.browser.platform.CanTakeScreenshot():
141        tf = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
142        tf.close()
143        tab.browser.platform.TakeScreenshot(tf.name)
144        results.AddProfilingFile(page, file_handle.FromTempFile(tf))
145      elif tab.IsAlive() and tab.screenshot_supported:
146        tf = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
147        tf.close()
148        image = tab.Screenshot()
149        image_util.WritePngFile(image, tf.name)
150        results.AddProfilingFile(page, file_handle.FromTempFile(tf))
151      else:
152        logging.warning(
153            'Either tab has crashed or browser does not support taking tab '
154            'screenshot. Skip taking screenshot on failure.')
155    except Exception as e:
156      logging.warning('Exception when trying to capture screenshot: %s',
157                      repr(e))
158
159  def DidRunStory(self, results):
160    if self._finder_options.profiler:
161      self._StopProfiling(results)
162    # We might hang while trying to close the connection, and need to guarantee
163    # the page will get cleaned up to avoid future tests failing in weird ways.
164    try:
165      if self._current_tab and self._current_tab.IsAlive():
166        self._current_tab.CloseConnections()
167    except Exception:
168      if self._current_tab:
169        self._current_tab.Close()
170    finally:
171      if self._current_page.credentials and self._did_login_for_current_page:
172        self.browser.credentials.LoginNoLongerNeeded(
173            self._current_tab, self._current_page.credentials)
174      if self._test.StopBrowserAfterPage(self.browser, self._current_page):
175        self._StopBrowser()
176      self._current_page = None
177      self._current_tab = None
178
179  @property
180  def platform(self):
181    return self._possible_browser.platform
182
183  def _StartBrowser(self, page):
184    assert self._browser is None
185    self._possible_browser.SetCredentialsPath(page.credentials_path)
186
187    self._test.WillStartBrowser(self.platform)
188    if page.startup_url:
189      self._finder_options.browser_options.startup_url = page.startup_url
190    self._browser = self._possible_browser.Create(self._finder_options)
191    self._test.DidStartBrowser(self.browser)
192
193    if self._first_browser:
194      self._first_browser = False
195      self.browser.credentials.WarnIfMissingCredentials(page)
196      logging.info('OS: %s %s',
197                   self.platform.GetOSName(),
198                   self.platform.GetOSVersionName())
199      if self.browser.supports_system_info:
200        system_info = self.browser.GetSystemInfo()
201        if system_info.model_name:
202          logging.info('Model: %s', system_info.model_name)
203        if system_info.gpu:
204          for i, device in enumerate(system_info.gpu.devices):
205            logging.info('GPU device %d: %s', i, device)
206          if system_info.gpu.aux_attributes:
207            logging.info('GPU Attributes:')
208            for k, v in sorted(system_info.gpu.aux_attributes.iteritems()):
209              logging.info('  %-20s: %s', k, v)
210          if system_info.gpu.feature_status:
211            logging.info('Feature Status:')
212            for k, v in sorted(system_info.gpu.feature_status.iteritems()):
213              logging.info('  %-20s: %s', k, v)
214          if system_info.gpu.driver_bug_workarounds:
215            logging.info('Driver Bug Workarounds:')
216            for workaround in system_info.gpu.driver_bug_workarounds:
217              logging.info('  %s', workaround)
218        else:
219          logging.info('No GPU devices')
220      else:
221        logging.warning('System info not supported')
222
223  def WillRunStory(self, page):
224    if not self.platform.tracing_controller.is_tracing_running:
225      # For TimelineBasedMeasurement benchmarks, tracing has already started.
226      # For PageTest benchmarks, tracing has not yet started. We need to make
227      # sure no tracing state is left before starting the browser for PageTest
228      # benchmarks.
229      self.platform.tracing_controller.ClearStateIfNeeded()
230
231    if self._ShouldDownloadPregeneratedProfileArchive():
232      self._DownloadPregeneratedProfileArchive()
233
234      if self._ShouldMigrateProfile():
235        self._MigratePregeneratedProfile()
236
237    page_set = page.page_set
238    self._current_page = page
239    if self._browser and (self._test.RestartBrowserBeforeEachPage()
240                          or page.startup_url):
241      assert not self.platform.tracing_controller.is_tracing_running, (
242          'Should not restart browser when tracing is already running. For '
243          'TimelineBasedMeasurement (TBM) benchmarks, you should not use '
244          'startup_url. Use benchmark.ShouldTearDownStateAfterEachStoryRun '
245          'instead.')
246      self._StopBrowser()
247    started_browser = not self.browser
248
249    archive_path = page_set.WprFilePathForStory(page)
250    # TODO(nednguyen, perezju): Ideally we should just let the network
251    # controller raise an exception when the archive_path is not found.
252    if archive_path is not None and not os.path.isfile(archive_path):
253      logging.warning('WPR archive missing: %s', archive_path)
254      archive_path = None
255    self.platform.network_controller.StartReplay(
256        archive_path, page.make_javascript_deterministic)
257
258    if self.browser:
259      # Set new credential path for browser.
260      self.browser.credentials.credentials_path = page.credentials_path
261    else:
262      self._StartBrowser(page)
263    if self.browser.supports_tab_control and self._test.close_tabs_before_run:
264      # Create a tab if there's none.
265      if len(self.browser.tabs) == 0:
266        self.browser.tabs.New()
267
268      # Ensure only one tab is open, unless the test is a multi-tab test.
269      if not self._test.is_multi_tab_test:
270        while len(self.browser.tabs) > 1:
271          self.browser.tabs[-1].Close()
272
273      # Must wait for tab to commit otherwise it can commit after the next
274      # navigation has begun and RenderFrameHostManager::DidNavigateMainFrame()
275      # will cancel the next navigation because it's pending. This manifests as
276      # the first navigation in a PageSet freezing indefinitely because the
277      # navigation was silently canceled when |self.browser.tabs[0]| was
278      # committed. Only do this when we just started the browser, otherwise
279      # there are cases where previous pages in a PageSet never complete
280      # loading so we'll wait forever.
281      if started_browser:
282        self.browser.tabs[0].WaitForDocumentReadyStateToBeComplete()
283
284    # Start profiling if needed.
285    if self._finder_options.profiler:
286      self._StartProfiling(self._current_page)
287
288  def CanRunStory(self, page):
289    return self.CanRunOnBrowser(browser_info_module.BrowserInfo(self.browser),
290                                page)
291
292  def CanRunOnBrowser(self, browser_info,
293                      page):  # pylint: disable=unused-argument
294    """Override this to return whether the browser brought up by this state
295    instance is suitable for running the given page.
296
297    Args:
298      browser_info: an instance of telemetry.core.browser_info.BrowserInfo
299      page: an instance of telemetry.page.Page
300    """
301    del browser_info, page  # unused
302    return True
303
304  def _PreparePage(self):
305    self._current_tab = self._test.TabForPage(self._current_page, self.browser)
306    if self._current_page.is_file:
307      self.platform.SetHTTPServerDirectories(
308          self._current_page.page_set.serving_dirs |
309          set([self._current_page.serving_dir]))
310
311    if self._current_page.credentials:
312      if not self.browser.credentials.LoginNeeded(
313          self._current_tab, self._current_page.credentials):
314        raise page_test.Failure(
315            'Login as ' + self._current_page.credentials + ' failed')
316      self._did_login_for_current_page = True
317
318    if self._test.clear_cache_before_each_run:
319      self._current_tab.ClearCache(force=True)
320
321  @property
322  def current_page(self):
323    return self._current_page
324
325  @property
326  def current_tab(self):
327    return self._current_tab
328
329  @property
330  def page_test(self):
331    return self._test
332
333  def RunStory(self, results):
334    try:
335      self._PreparePage()
336      self._current_page.Run(self)
337      self._test.ValidateAndMeasurePage(
338          self._current_page, self._current_tab, results)
339    except exceptions.Error:
340      if self._finder_options.browser_options.take_screenshot_for_failed_page:
341        self._TryCaptureScreenShot(self._current_page, self._current_tab,
342                                   results)
343      if self._test.is_multi_tab_test:
344        # Avoid trying to recover from an unknown multi-tab state.
345        exception_formatter.PrintFormattedException(
346            msg='Telemetry Error during multi tab test:')
347        raise page_test.MultiTabTestAppCrashError
348      raise
349    except Exception:
350      if self._finder_options.browser_options.take_screenshot_for_failed_page:
351        self._TryCaptureScreenShot(self._current_page, self._current_tab,
352                                   results)
353      raise
354
355  def TearDownState(self):
356    if self._migrated_profile:
357      shutil.rmtree(self._migrated_profile)
358      self._migrated_profile = None
359
360    self._StopBrowser()
361    self.platform.StopAllLocalServers()
362    self.platform.network_controller.Close()
363
364  def _StopBrowser(self):
365    if self._browser:
366      self._browser.Close()
367      self._browser = None
368
369  def _StartProfiling(self, page):
370    output_file = os.path.join(self._finder_options.output_dir,
371                               page.file_safe_name)
372    is_repeating = (self._finder_options.page_repeat != 1 or
373                    self._finder_options.pageset_repeat != 1)
374    if is_repeating:
375      output_file = util.GetSequentialFileName(output_file)
376    self.browser.profiling_controller.Start(
377        self._finder_options.profiler, output_file)
378
379  def _StopProfiling(self, results):
380    if self.browser:
381      profiler_files = self.browser.profiling_controller.Stop()
382      for f in profiler_files:
383        if os.path.isfile(f):
384          results.AddProfilingFile(self._current_page,
385                                   file_handle.FromFilePath(f))
386
387  def _ShouldMigrateProfile(self):
388    return not self._migrated_profile
389
390  def _MigrateProfile(self, finder_options, found_browser,
391                      initial_profile, final_profile):
392    """Migrates a profile to be compatible with a newer version of Chrome.
393
394    Launching Chrome with the old profile will perform the migration.
395    """
396    # Save the current input and output profiles.
397    saved_input_profile = finder_options.browser_options.profile_dir
398    saved_output_profile = finder_options.output_profile_path
399
400    # Set the input and output profiles.
401    finder_options.browser_options.profile_dir = initial_profile
402    finder_options.output_profile_path = final_profile
403
404    # Launch the browser, then close it.
405    browser = found_browser.Create(finder_options)
406    browser.Close()
407
408    # Load the saved input and output profiles.
409    finder_options.browser_options.profile_dir = saved_input_profile
410    finder_options.output_profile_path = saved_output_profile
411
412  def _MigratePregeneratedProfile(self):
413    """Migrates the pre-generated profile by launching Chrome with it.
414
415    On success, updates self._migrated_profile and
416    self._finder_options.browser_options.profile_dir with the directory of the
417    migrated profile.
418    """
419    self._migrated_profile = tempfile.mkdtemp()
420    logging.info("Starting migration of pre-generated profile to %s",
421                 self._migrated_profile)
422    pregenerated_profile = self._finder_options.browser_options.profile_dir
423
424    possible_browser = self._FindBrowser(self._finder_options)
425    self._MigrateProfile(self._finder_options, possible_browser,
426                         pregenerated_profile, self._migrated_profile)
427    self._finder_options.browser_options.profile_dir = self._migrated_profile
428    logging.info("Finished migration of pre-generated profile to %s",
429                 self._migrated_profile)
430
431  def GetPregeneratedProfileArchiveDir(self):
432    return self._pregenerated_profile_archive_dir
433
434  def SetPregeneratedProfileArchiveDir(self, archive_path):
435    """
436    Benchmarks can set a pre-generated profile archive to indicate that when
437    Chrome is launched, it should have a --user-data-dir set to the
438    pre-generated profile, rather than to an empty profile.
439
440    If the benchmark is invoked with the option --profile-dir=<dir>, that
441    option overrides this value.
442    """
443    self._pregenerated_profile_archive_dir = archive_path
444
445  def _ShouldDownloadPregeneratedProfileArchive(self):
446    """Whether to download a pre-generated profile archive."""
447    # There is no pre-generated profile archive.
448    if not self.GetPregeneratedProfileArchiveDir():
449      return False
450
451    # If profile dir is specified on command line, use that instead.
452    if self._finder_options.browser_options.profile_dir:
453      logging.warning("Profile directory specified on command line: %s, this"
454                      "overrides the benchmark's default profile directory.",
455                      self._finder_options.browser_options.profile_dir)
456      return False
457
458    # If the browser is remote, a local download has no effect.
459    if self._possible_browser.IsRemote():
460      return False
461
462    return True
463
464  def _DownloadPregeneratedProfileArchive(self):
465    """Download and extract the profile directory archive if one exists.
466
467    On success, updates self._finder_options.browser_options.profile_dir with
468    the directory of the extracted profile.
469    """
470    # Download profile directory from cloud storage.
471    generated_profile_archive_path = self.GetPregeneratedProfileArchiveDir()
472
473    try:
474      cloud_storage.GetIfChanged(generated_profile_archive_path,
475                                 cloud_storage.PUBLIC_BUCKET)
476    except (cloud_storage.CredentialsError,
477            cloud_storage.PermissionError) as e:
478      if os.path.exists(generated_profile_archive_path):
479        # If the profile directory archive exists, assume the user has their
480        # own local copy simply warn.
481        logging.warning('Could not download Profile archive: %s',
482                        generated_profile_archive_path)
483      else:
484        # If the archive profile directory doesn't exist, this is fatal.
485        logging.error('Can not run without required profile archive: %s. '
486                      'If you believe you have credentials, follow the '
487                      'instructions below.',
488                      generated_profile_archive_path)
489        logging.error(str(e))
490        sys.exit(-1)
491
492    # Check to make sure the zip file exists.
493    if not os.path.isfile(generated_profile_archive_path):
494      raise Exception("Profile directory archive not downloaded: ",
495                      generated_profile_archive_path)
496
497    # The location to extract the profile into.
498    extracted_profile_dir_path = (
499        os.path.splitext(generated_profile_archive_path)[0])
500
501    # Unzip profile directory.
502    with zipfile.ZipFile(generated_profile_archive_path) as f:
503      try:
504        f.extractall(os.path.dirname(generated_profile_archive_path))
505      except e:
506        # Cleanup any leftovers from unzipping.
507        if os.path.exists(extracted_profile_dir_path):
508          shutil.rmtree(extracted_profile_dir_path)
509        logging.error("Error extracting profile directory zip file: %s", e)
510        sys.exit(-1)
511
512    # Run with freshly extracted profile directory.
513    logging.info("Using profile archive directory: %s",
514                 extracted_profile_dir_path)
515    self._finder_options.browser_options.profile_dir = (
516        extracted_profile_dir_path)
517
518
519class SharedMobilePageState(SharedPageState):
520  _device_type = 'mobile'
521
522
523class SharedDesktopPageState(SharedPageState):
524  _device_type = 'desktop'
525
526
527class SharedTabletPageState(SharedPageState):
528  _device_type = 'tablet'
529
530
531class Shared10InchTabletPageState(SharedPageState):
532  _device_type = 'tablet_10_inch'
533