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 5 6 class DroneTaskQueueException(Exception): 7 """Generic task queue exception.""" 8 pass 9 10 11 class DroneTaskQueue(object): 12 """A manager to run queued tasks in drones and gather results from them.""" 13 14 def __init__(self): 15 self.results = dict() 16 17 18 def get_results(self): 19 """Get a results dictionary keyed on drones. 20 21 @return: A dictionary of return values from drones. 22 """ 23 results_copy = self.results.copy() 24 self.results.clear() 25 return results_copy 26 27 28 def execute(self, drones, wait=True): 29 """Invoke methods via SSH to a drone. 30 31 @param drones: A list of drones with calls to execute. 32 @param wait: If True, this method will only return when all the drones 33 have returned the result of their respective invocations of 34 drone_utility. The `results` map will be cleared. 35 If False, the caller must clear the map before the next invocation 36 of `execute`, by calling `get_results`. 37 38 @return: A dictionary keyed on the drones, containing a list of return 39 values from the execution of drone_utility. 40 41 @raises DroneTaskQueueException: If the results map isn't empty at the 42 time of invocation. 43 """ 44 if self.results: 45 raise DroneTaskQueueException( 46 'Cannot clobber results map: %s, it should be cleared ' 47 'through get_results.' % self.results) 48 for drone in drones: 49 if not drone.get_calls(): 50 continue 51 drone_results = drone.execute_queued_calls() 52 if drone in self.results: 53 raise DroneTaskQueueException( 54 'Task queue has recorded results for drone %s: %s' % 55 (drone, self.results)) 56 self.results[drone] = drone_results 57 return self.get_results() if wait else None 58
