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 = os.path.join(self._base_path_in_bucket, 'traces')
+    self._trace_database_path = os.path.join(
+        self._traces_dir,
+        config.get('trace_database_filename', 'trace_database.json'))
+
+    # 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."""
+    self._logger.info('Downloading trace database')
+    trace_database_string = self._google_storage_accessor.DownloadAsString(
+        self._trace_database_path) or '{}'
+    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')
+    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
@@ skipping 54 matching lines @@
       self._logger.error('Unsupported task action: %s' % clovis_task.Action())
       return
 
     # Extract the task parameters.
     params = clovis_task.Params()
     urls = params['urls']
     repeat_count = params.get('repeat_count', 1)
     emulate_device = params.get('emulate_device')
     emulate_network = params.get('emulate_network')
 
-    failures_dir = self._base_path_in_bucket + 'failures/'
+    failures_dir = os.path.join(self._base_path_in_bucket, 'failures')
     # TODO(blundell): Fix this up.
-    logs_dir = self._base_path_in_bucket + 'analyze_logs/'
+    logs_dir = os.path.join(self._base_path_in_bucket, 'analyze_logs')
     log_filename = 'analyze.log'
     # Avoid special characters in storage object names
     pattern = re.compile(r"[#\?\[\]\*/]")
 
     while len(urls) > 0:
       url = urls.pop()
       local_filename = pattern.sub('_', url)
       for repeat in range(repeat_count):
         self._logger.debug('Generating trace for URL: %s' % url)
-        remote_filename = local_filename + '/' + str(repeat)
+        remote_filename = os.path.join(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
+          remote_trace_location = os.path.join(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)
+          full_cloud_storage_path = os.path.join('gs://' + self._bucket_name,
+                                                 remote_trace_location)
+          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(log_filename,
-                                                 logs_dir + remote_filename)
-
+            self._google_storage_accessor.UploadFile(
+                local_filename, os.path.join(failures_dir, remote_filename))
+        self._logger.debug('Uploading analyze log')
+        self._google_storage_accessor.UploadFile(
+            log_filename, os.path.join(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()