def test_task_queue(self):
""" Test the task queue """
task_queue = TaskQueue()
counters = [Counter() for _ in xrange(128)]
tasks = [task_queue.put(counters[i].add, i) for i in xrange(128)]
for task in tasks:
task.join()
for i in xrange(128):
self.assertEqual(counters[i].value, i)
def test_task_queue(self):
""" Test the task queue """
task_queue = TaskQueue()
counters = [Counter() for _ in xrange(128)]
tasks = [task_queue.put(counters[i].add, i) for i in xrange(128)]
for task in tasks:
task.join()
for i in xrange(128):
self.assertEqual(counters[i].value, i)
def test_sequential_delay(self):
""" Test that a single thread delays between tasks """
task_queue = TaskQueue(threads=1, delay=0.25)
counters = [Counter() for x in xrange(8)]
start = time.time()
tasks = [task_queue.put(x.add, 1) for x in counters]
for task in tasks:
task.join()
finish = time.time()
self.assertGreaterEqual(finish - start, 2.0)
def test_sequential_delay(self):
""" Test that a single thread delays between tasks """
task_queue = TaskQueue(threads=1, delay=0.25)
counters = [Counter() for x in xrange(8)]
start = time.time()
tasks = [task_queue.put(x.add, 1) for x in counters]
for task in tasks:
task.join()
finish = time.time()
self.assertGreaterEqual(finish - start, 2.0)
def run_command(args, nodes, name, *method_args):
"""Runs a command on nodes."""
if args.threads != None:
task_queue = TaskQueue(threads=args.threads, delay=args.command_delay)
else:
task_queue = TaskQueue(delay=args.command_delay)
tasks = {}
for node in nodes:
target = node
for member in name.split("."):
target = getattr(target, member)
tasks[node] = task_queue.put(target, *method_args)
results = {}
errors = {}
try:
counter = 0
while any(x.is_alive() for x in tasks.values()):
if not args.quiet:
_print_command_status(tasks, counter)
counter += 1
time.sleep(0.25)
for node, task in tasks.iteritems():
if task.status == "Completed":
results[node] = task.result
else:
errors[node] = task.error
except KeyboardInterrupt:
args.retry = 0
for node, task in tasks.iteritems():
if task.status == "Completed":
results[node] = task.result
elif task.status == "Failed":
errors[node] = task.error
else:
errors[node] = KeyboardInterrupt(
"Aborted by keyboard interrupt"
)
if not args.quiet:
_print_command_status(tasks, counter)
print("\n")
# Handle errors
should_retry = False
if errors:
_print_errors(args, nodes, errors)
if args.retry == None:
sys.stdout.write("Retry command on failed hosts? (y/n): ")
sys.stdout.flush()
while True:
command = raw_input().strip().lower()
if command in ['y', 'yes']:
should_retry = True
break
elif command in ['n', 'no']:
print
break
elif args.retry >= 1:
should_retry = True
if args.retry == 1:
print("Retrying command 1 more time...")
elif args.retry > 1:
print("Retrying command %i more times..." % args.retry)
args.retry -= 1
if should_retry:
nodes = [x for x in nodes if x in errors]
new_results, errors = run_command(args, nodes, name, *method_args)
results.update(new_results)
return results, errors
