1# Copyright 2015 The Chromium OS 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 6 7from chromite.lib import metrics 8 9 10DRONE_ACCESSIBILITY_METRIC = metrics.Boolean( 11 'chromeos/autotest/scheduler/drone_accessibility') 12 13class DroneTaskQueueException(Exception): 14 """Generic task queue exception.""" 15 pass 16 17 18class DroneTaskQueue(object): 19 """A manager to run queued tasks in drones and gather results from them.""" 20 21 def __init__(self): 22 self.results = dict() 23 24 25 def get_results(self): 26 """Get a results dictionary keyed on drones. 27 28 @return: A dictionary of return values from drones. 29 """ 30 results_copy = self.results.copy() 31 self.results.clear() 32 return results_copy 33 34 35 def execute(self, drones, wait=True): 36 """Invoke methods via SSH to a drone. 37 38 @param drones: A list of drones with calls to execute. 39 @param wait: If True, this method will only return when all the drones 40 have returned the result of their respective invocations of 41 drone_utility. The `results` map will be cleared. 42 If False, the caller must clear the map before the next invocation 43 of `execute`, by calling `get_results`. 44 45 @return: A dictionary keyed on the drones, containing a list of return 46 values from the execution of drone_utility. 47 48 @raises DroneTaskQueueException: If the results map isn't empty at the 49 time of invocation. 50 """ 51 if self.results: 52 raise DroneTaskQueueException( 53 'Cannot clobber results map: %s, it should be cleared ' 54 'through get_results.' % self.results) 55 for drone in drones: 56 if not drone.get_calls(): 57 logging.debug("Drone %s has no work, skipping. crbug.com/853861" 58 , drone) 59 continue 60 metric_fields = { 61 'drone_hostname': drone.hostname, 62 'call_count': len(drone.get_calls()) 63 } 64 drone_reachable = True 65 try: 66 drone_results = drone.execute_queued_calls() 67 logging.debug("Drone %s scheduled. crbug.com/853861", drone) 68 except IOError: 69 drone_reachable = False 70 logging.error( 71 "Drone %s is not reachable by the scheduler.", drone) 72 continue 73 finally: 74 DRONE_ACCESSIBILITY_METRIC.set( 75 drone_reachable, fields=metric_fields) 76 if drone in self.results: 77 raise DroneTaskQueueException( 78 'Task queue has recorded results for drone %s: %s' % 79 (drone, self.results)) 80 self.results[drone] = drone_results 81 return self.get_results() if wait else None 82