tools/android/loading Loading trace database improvements

tools/android/loading/cloud/backend/worker.py

 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import argparse
 import json
 import logging
 import os
 import re
 import sys
@@ skipping 14 matching lines @@
 import loading_trace
 from loading_trace_database import LoadingTraceDatabase
 import options
 
 
 class Worker(object):
   def __init__(self, config, logger):
     """See README.md for the config format."""
     self._project_name = config['project_name']
     self._taskqueue_tag = config['taskqueue_tag']
+    self._src_path = config['src_path']
     self._credentials = GoogleCredentials.get_application_default()
     self._logger = logger
 
     # Separate the cloud storage path into the bucket and the base path under
     # the bucket.
     storage_path_components = config['cloud_storage_path'].split('/')
     self._bucket_name = storage_path_components[0]
     self._base_path_in_bucket = ''
     if len(storage_path_components) > 1:
       self._base_path_in_bucket = '/'.join(storage_path_components[1:])
       if not self._base_path_in_bucket.endswith('/'):
         self._base_path_in_bucket += '/'
 
-    # TODO: improve the trace database to support concurrent access.
-    self._traces_dir = self._base_path_in_bucket + 'traces/'
-    self._trace_database = LoadingTraceDatabase({})
-
-    self._src_path = config['src_path']
     self._google_storage_accessor = GoogleStorageAccessor(
         credentials=self._credentials, project_name=self._project_name,
         bucket_name=self._bucket_name)
 
+    self._traces_dir = self._base_path_in_bucket + 'traces/'

[Benoit L, 2016/04/20 15:02:22]

Generally speaking, I prefer to manipulate paths with os.path.join(), this
avoids having to care about whether all parts are neatly separated. It would
only blow up if the script is running on Windows and you want '/' as a
separator.

[droger, 2016/04/20 15:58:17]

Done.

+    self._trace_database_path = self._traces_dir + config.get(
+        'trace_database_filename', 'trace_database.json')

[Benoit L, 2016/04/20 15:02:22]

Is the trace dir per worker?
If not, then the trace DB is not per-worker by default, or am I missing
something?

[droger, 2016/04/20 15:58:17]

It is intended that all worker write to the same dir, but different files.
If you don't specify a 'trace_database_filename', then yes all workers would
write to the same file, which would not work.

I can make this key required instead of optional, but I don't think it would
solve the problem, since nothing is preventing several worker to use the same
file.

I'm adding a comment in the README to clarify.

+
+    # Recover any existing trace database in case the worker died.
+    self._DownloadTraceDatabase()
+
     # Initialize the global options that will be used during trace generation.
     options.OPTIONS.ParseArgs([])
     options.OPTIONS.local_binary = config['chrome_path']
 
   def Start(self):
     """Main worker loop.
 
     Repeatedly pulls tasks from the task queue and processes them. Returns when
     the queue is empty.
     """
@@ skipping 20 matching lines @@
           continue
 
       self._logger.info('Processing task %s' % task_id)
       self._ProcessClovisTask(clovis_task)
       self._logger.debug('Deleting task %s' % task_id)
       task_api.tasks().delete(project=project, taskqueue=queue_name,
                               task=task_id).execute()
       self._logger.info('Finished task %s' % task_id)
     self._Finalize()
 
+  def _DownloadTraceDatabase(self):
+    """Downloads the trace database from CloudStorage."""
+    trace_database_string = self._google_storage_accessor.DownloadAsString(

[Benoit L, 2016/04/20 15:02:22]

What about:

trace_db_str = self._google_storage_accessor...() or '{}'
self._trace_database = LoadingTraceDatabase(json.loads(trace_db_str))

Tangentially related: looks like we upload / download JSON quite often. Why not
add that to GoogleStorageAccessor?

[droger, 2016/04/20 15:58:17]

Done.
For me, what would make the most sense would be that the
LoadingTraceDatabase.ToJson() and FromJson() functions actually return JSON
(strings) instead of dictionaries.

+        self._trace_database_path)
+    trace_database_dict = {}
+    if trace_database_string:
+      self._logger.warning(
+          'Importing existing trace database %s' % self._trace_database_path)
+      trace_database_dict = json.loads(trace_database_string)
+    self._trace_database = LoadingTraceDatabase(trace_database_dict)
+
+  def _UploadTraceDatabase(self):
+    """Uploads the trace database to CloudStorage."""
+    self._logger.info('Uploading trace database')

[Benoit L, 2016/04/20 15:02:22]

tiny nit: Importing is "warning", uploading "info".

[droger, 2016/04/20 15:58:17]

Done.

+    self._google_storage_accessor.UploadString(
+        json.dumps(self._trace_database.ToJsonDict(), indent=2),
+        self._trace_database_path)
+
   def _FetchClovisTask(self, project_name, task_api, queue_name):
     """Fetches a ClovisTask from the task queue.
 
     Params:
       project_name(str): The name of the Google Cloud project.
       task_api: The TaskQueue service.
       queue_name(str): The name of the task queue.
 
     Returns:
       (ClovisTask, str): The fetched ClovisTask and its task ID, or (None, None)
                          if no tasks are found.
     """
     response = task_api.tasks().lease(
         project=project_name, taskqueue=queue_name, numTasks=1, leaseSecs=180,
         groupByTag=True, tag=self._taskqueue_tag).execute()
     if (not response.get('items')) or (len(response['items']) < 1):
       return (None, None)
 
     google_task = response['items'][0]
     task_id = google_task['id']
     clovis_task = ClovisTask.FromBase64(google_task['payloadBase64'])
     return (clovis_task, task_id)
 
   def _Finalize(self):
     """Called before exiting."""
-    self._logger.info('Uploading trace database')
-    self._google_storage_accessor.UploadString(
-        json.dumps(self._trace_database.ToJsonDict(), indent=2),
-        self._traces_dir + 'trace_database.json')
     # TODO(droger): Implement automatic instance destruction.
     self._logger.info('Done')
 
-
   def _GenerateTrace(self, url, emulate_device, emulate_network, filename,
                      log_filename):
     """ Generates a trace.
 
     Args:
       url: URL as a string.
       emulate_device: Name of the device to emulate. Empty for no emulation.
       emulate_network: Type of network emulation. Empty for no emulation.
       filename: Name of the file where the trace is saved.
       log_filename: Name of the file where standard output and errors are
                     logged.
 
     Returns:
       A dictionary of metadata about the trace, including a 'succeeded' field
       indicating whether the trace was successfully generated.
     """
     try:
       os.remove(filename)  # Remove any existing trace for this URL.
     except OSError:
       pass  # Nothing to remove.
 
     if not url.startswith('http') and not url.startswith('file'):
       url = 'http://' + url
 
-    old_stdout = sys.stdout
     old_stderr = sys.stderr
 
     trace_metadata = { 'succeeded' : False, 'url' : url }
     trace = None
     with open(log_filename, 'w') as sys.stdout:
       try:
         sys.stderr = sys.stdout
 
         # Set up the controller.
         chrome_ctl = controller.LocalChromeController()
@@ skipping 11 matching lines @@
               url, connection, chrome_ctl.ChromeMetadata())
           trace_metadata['succeeded'] = True
           trace_metadata.update(trace.ToJsonDict()[trace._METADATA_KEY])
       except Exception as e:
         sys.stderr.write(str(e))
 
       if trace:
         with open(filename, 'w') as f:
           json.dump(trace.ToJsonDict(), f, sort_keys=True, indent=2)
 
-    sys.stdout = old_stdout
     sys.stderr = old_stderr
 
     return trace_metadata
 
   def _ProcessClovisTask(self, clovis_task):
     """Processes one clovis_task."""
     if clovis_task.Action() != 'trace':
       self._logger.error('Unsupported task action: %s' % clovis_task.Action())
       return
 
@@ skipping 19 matching lines @@
         remote_filename = local_filename + '/' + str(repeat)
         trace_metadata = self._GenerateTrace(
             url, emulate_device, emulate_network, local_filename, log_filename)
         if trace_metadata['succeeded']:
           self._logger.debug('Uploading: %s' % remote_filename)
           remote_trace_location = self._traces_dir + remote_filename
           self._google_storage_accessor.UploadFile(local_filename,
                                                    remote_trace_location)
           full_cloud_storage_path = ('gs://' + self._bucket_name + '/' +
               remote_trace_location)
-          self._trace_database.AddTrace(full_cloud_storage_path, trace_metadata)
+          self._trace_database.SetTrace(full_cloud_storage_path, trace_metadata)
         else:
           self._logger.warning('Trace generation failed for URL: %s' % url)
-          # TODO: upload the failure
           if os.path.isfile(local_filename):
-            self._google_storage_accessor.UploadFile(local_filename,
-                                            failures_dir + remote_filename)
-            self._logger.debug('Uploading log')
+            self._google_storage_accessor.UploadFile(
+                local_filename, failures_dir + remote_filename)
+        self._logger.debug('Uploading analyze log')
         self._google_storage_accessor.UploadFile(log_filename,
                                                  logs_dir + remote_filename)
-
+    self._UploadTraceDatabase()
 
 if __name__ == '__main__':
   parser = argparse.ArgumentParser(
       description='ComputeEngine Worker for Clovis')
   parser.add_argument('--config', required=True,
                       help='Path to the configuration file.')
   args = parser.parse_args()
 
   # Configure logging.
   logging.basicConfig(level=logging.WARNING)
   worker_logger = logging.getLogger('worker')
   worker_logger.setLevel(logging.INFO)
 
   worker_logger.info('Reading configuration')
   with open(args.config) as config_json:
     worker = Worker(json.load(config_json), worker_logger)
     worker.Start()