MK Ryu | 7911ad5 | 2015-12-18 11:40:04 -0800 | [diff] [blame] | 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 |