| # 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 random |
| import sys |
| import time |
| |
| from googleapiclient import discovery |
| from oauth2client.client import GoogleCredentials |
| |
| # NOTE: The parent directory needs to be first in sys.path to avoid conflicts |
| # with catapult modules that have colliding names, as catapult inserts itself |
| # into the path as the second element. This is an ugly and fragile hack. |
| _CLOUD_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), |
| os.pardir) |
| sys.path.insert(0, os.path.join(_CLOUD_DIR, os.pardir)) |
| # Add _CLOUD_DIR to the path to access common code through the same path as the |
| # frontend. |
| sys.path.append(_CLOUD_DIR) |
| |
| from common.clovis_task import ClovisTask |
| import common.google_bigquery_helper |
| from common.google_instance_helper import GoogleInstanceHelper |
| from clovis_task_handler import ClovisTaskHandler |
| from failure_database import FailureDatabase |
| from google_storage_accessor import GoogleStorageAccessor |
| |
| |
| 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._instance_name = config.get('instance_name') |
| self._worker_log_path = config.get('worker_log_path') |
| self._credentials = GoogleCredentials.get_application_default() |
| self._logger = logger |
| self._self_destruct = config.get('self_destruct') |
| if self._self_destruct and not self._instance_name: |
| self._logger.error('Self destruction requires an instance name.') |
| |
| # Separate the task storage path into the bucket and the base path under |
| # the bucket. |
| storage_path_components = config['task_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 += '/' |
| |
| self._google_storage_accessor = GoogleStorageAccessor( |
| credentials=self._credentials, project_name=self._project_name, |
| bucket_name=self._bucket_name) |
| |
| if self._instance_name: |
| failure_database_filename = \ |
| 'failure_database_%s.json' % self._instance_name |
| else: |
| failure_database_filename = 'failure_dabatase.json' |
| self._failure_database_path = os.path.join(self._base_path_in_bucket, |
| failure_database_filename) |
| |
| # Recover any existing failures in case the worker died. |
| self._failure_database = self._GetFailureDatabase() |
| |
| if self._failure_database.ToJsonDict(): |
| # Script is restarting after a crash, or there are already files from a |
| # previous run in the directory. |
| self._failure_database.AddFailure(FailureDatabase.DIRTY_STATE_ERROR, |
| 'failure_database') |
| |
| bigquery_service = common.google_bigquery_helper.GetBigQueryService( |
| self._credentials) |
| self._clovis_task_handler = ClovisTaskHandler( |
| self._project_name, self._base_path_in_bucket, self._failure_database, |
| self._google_storage_accessor, bigquery_service, |
| config['binaries_path'], config['ad_rules_filename'], |
| config['tracking_rules_filename'], self._logger, self._instance_name) |
| |
| self._UploadFailureDatabase() |
| |
| def Start(self): |
| """Main worker loop. |
| |
| Repeatedly pulls tasks from the task queue and processes them. Returns when |
| the queue is empty. |
| """ |
| task_api = discovery.build('taskqueue', 'v1beta2', |
| credentials=self._credentials) |
| queue_name = 'clovis-queue' |
| # Workaround for |
| # https://code.google.com/p/googleappengine/issues/detail?id=10199 |
| project = 's~' + self._project_name |
| |
| while True: |
| self._logger.debug('Fetching new task.') |
| (clovis_task, task_id) = self._FetchClovisTask(project, task_api, |
| queue_name) |
| if not clovis_task: |
| break |
| |
| self._logger.info('Processing task %s' % task_id) |
| self._clovis_task_handler.Run(clovis_task) |
| self._UploadFailureDatabase() |
| 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 _GetFailureDatabase(self): |
| """Downloads the failure database from CloudStorage.""" |
| self._logger.info('Downloading failure database') |
| failure_database_string = self._google_storage_accessor.DownloadAsString( |
| self._failure_database_path) |
| return FailureDatabase(failure_database_string) |
| |
| def _UploadFailureDatabase(self): |
| """Uploads the failure database to CloudStorage.""" |
| if not self._failure_database.is_dirty: |
| return |
| self._logger.info('Uploading failure database') |
| self._google_storage_accessor.UploadString( |
| self._failure_database.ToJsonString(), |
| self._failure_database_path) |
| self._failure_database.is_dirty = False |
| |
| 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=600, |
| groupByTag=True, tag=self._taskqueue_tag).execute() |
| if (not response.get('items')) or (len(response['items']) < 1): |
| return (None, None) # The task queue is empty. |
| |
| google_task = response['items'][0] |
| task_id = google_task['id'] |
| |
| # Delete the task without processing if it already failed multiple times. |
| # TODO(droger): This is a workaround for internal bug b/28442122, revisit |
| # once it is fixed. |
| retry_count = google_task['retry_count'] |
| max_retry_count = 3 |
| skip_task = retry_count >= max_retry_count |
| if skip_task: |
| task_api.tasks().delete(project=project_name, taskqueue=queue_name, |
| task=task_id).execute() |
| |
| clovis_task = ClovisTask.FromBase64(google_task['payloadBase64']) |
| |
| if retry_count > 0: |
| self._failure_database.AddFailure('task_queue_retry', |
| clovis_task.ToJsonString()) |
| self._UploadFailureDatabase() |
| |
| if skip_task: |
| return self._FetchClovisTask(project_name, task_api, queue_name) |
| |
| return (clovis_task, task_id) |
| |
| def _Finalize(self): |
| """Called before exiting.""" |
| self._logger.info('Done') |
| self._clovis_task_handler.Finalize() |
| # Upload the worker log. |
| if self._worker_log_path: |
| self._logger.info('Uploading worker log.') |
| remote_log_path = os.path.join(self._base_path_in_bucket, 'worker_log') |
| if self._instance_name: |
| remote_log_path += '_' + self._instance_name |
| self._google_storage_accessor.UploadFile(self._worker_log_path, |
| remote_log_path) |
| # Self destruct. |
| if self._self_destruct: |
| # Workaround for ComputeEngine internal bug b/28760288. |
| random_delay = random.random() * 600.0 # Up to 10 minutes. |
| self._logger.info( |
| 'Wait %.0fs to avoid load spikes on compute engine.' % random_delay) |
| time.sleep(random_delay) |
| |
| self._logger.info('Starting instance destruction: ' + self._instance_name) |
| google_instance_helper = GoogleInstanceHelper( |
| self._credentials, self._project_name, self._logger) |
| success = google_instance_helper.DeleteInstance(self._taskqueue_tag, |
| self._instance_name) |
| if not success: |
| self._logger.error('Self destruction failed.') |
| # Do not add anything after this line, as the instance might be killed at |
| # any time. |
| |
| 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, |
| format='[%(asctime)s][%(levelname)s] %(message)s', |
| datefmt='%y-%m-%d %H:%M:%S') |
| logging.Formatter.converter = time.gmtime |
| 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() |
