def _worker_done(num, future):
teardown = get_fixture('teardown')
if teardown is not None:
try:
teardown(num)
except Exception as e:
# we can't stop the teardown process
log(e)
def _process(args):
global _STOP
if args.processes > 1:
signal.signal(signal.SIGINT, _shutdown)
signal.signal(signal.SIGTERM, _shutdown)
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
else:
loop = asyncio.get_event_loop()
if args.debug:
log('**** RUNNING IN DEBUG MODE == SLOW ****')
loop.set_debug(True)
results = {'OK': 0, 'FAILED': 0}
stream = asyncio.Queue(loop=loop)
co_tasks = []
if args.statsd:
statsd = get_statsd_client(args.statsd_server, args.statsd_port)
stastd_task = asyncio.ensure_future(statsd.run())
else:
statsd = stastd_task = None
consumer = asyncio.ensure_future(
consume(stream, args.workers, args.console, args.verbose))
co_tasks.append(consumer)
_TASKS.extend(co_tasks)
co_tasks = asyncio.gather(*co_tasks, loop=loop, return_exceptions=True)
workers = _runner(loop, args, results, stream, statsd)
run_task = asyncio.gather(*workers, loop=loop, return_exceptions=True)
_TASKS.extend(workers)
_STATSD.append((statsd, stastd_task))
try:
loop.run_until_complete(run_task)
except asyncio.CancelledError:
_STOP = True
co_tasks.cancel()
loop.run_until_complete(co_tasks)
run_task.cancel()
loop.run_until_complete(run_task)
finally:
_stop_statsd()
for task in _TASKS:
del task
loop.close()
return results
async def worker(num, loop, results, args, stream, statsd):
global _STOP
quiet = args.quiet
duration = args.duration
verbose = args.verbose
exception = args.exception
count = 1
start = _now()
howlong = 0
setup = get_fixture('setup')
if setup is not None:
try:
options = await setup(num, args)
except Exception as e:
log(e)
await stream.put('WORKER_STOPPED')
return
if options is None:
options = {}
elif not isinstance(options, dict):
log('The setup function needs to return a dict')
await stream.put('WORKER_STOPPED')
return
else:
options = {}
async with Session(loop, stream, verbose, statsd, **options) as session:
while howlong < duration and not _STOP:
howlong = _now() - start
result = await step(session, quiet, verbose, stream)
if result == 1:
results['OK'] += 1
elif result == -1:
results['FAILED'] += 1
if exception:
await stream.put('WORKER_STOPPED')
_STOP = True
elif result == 0:
break
count += 1
if not _STOP:
await stream.put('WORKER_STOPPED')
def runner(args, screen=None):
global_setup = get_fixture('global_setup')
if global_setup is not None:
try:
global_setup(args)
except Exception:
log("The global_setup() fixture failed")
raise
try:
return _launch_processes(args, screen)
finally:
global_teardown = get_fixture('global_teardown')
if global_teardown is not None:
try:
global_teardown()
except Exception as e:
# we can't stop the teardown process
log(e)
def run(args):
if not args.quiet:
log('**** Molotov v%s. Happy breaking! ****' % __version__, pid=None)
if os.path.exists(args.scenario):
spec = spec_from_file_location("loadtest", args.scenario)
module = module_from_spec(spec)
spec.loader.exec_module(module)
else:
try:
import_module(args.scenario)
except (ImportError, ValueError):
print('Cannot import %r' % args.scenario)
sys.exit(1)
if len(get_scenarios()) == 0:
print('You need at least one scenario. No scenario was found.')
print('A scenario with a weight of 0 is ignored')
sys.exit(1)
if args.verbose and args.quiet:
print("You can't use -q and -v at the same time")
sys.exit(1)
if args.verbose and not args.console:
print("You have to be in console mode (-c) to use -v")
sys.exit(1)
res = runner(args, screen=ui.init_screen)
if not args.quiet:
log('', pid=False)
log('%(OK)d OK, %(FAILED)d Failed' % res, pid=False)
def _launch_processes(args, screen):
results = {'FAILED': 0, 'OK': 0}
signal.signal(signal.SIGINT, _shutdown)
signal.signal(signal.SIGTERM, _shutdown)
if args.processes > 1:
if not args.quiet:
log('Forking %d processes' % args.processes, pid=False)
result_queue = multiprocessing.Queue()
ui = None
loop = asyncio.get_event_loop()
def _pprocess(result_queue):
result_queue.put(_process(args))
jobs = []
for i in range(args.processes):
p = multiprocessing.Process(target=_pprocess,
args=(result_queue,))
jobs.append(p)
p.start()
for job in jobs:
_PROCESSES.append(job)
if screen is not None and not args.console:
if args.processes == 1:
pids = [os.getpid()]
else:
pids = [job.pid for job in jobs]
ui = screen(pids, get_live_results)
def check_procs(*args):
dead = [not p.is_alive() for p in _PROCESSES]
if all(dead):
quit_screen()
ui.set_alarm_in(1, check_procs)
ui.run()
async def run(loop, quiet, console):
while len(_PROCESSES) > 0:
if not quiet and console:
try:
print(get_live_results(), end='\r')
except ClosedError:
# finished
return
for job in jobs:
if job.exitcode is not None and job in _PROCESSES:
_PROCESSES.remove(job)
await asyncio.sleep(.2)
loop.run_until_complete(asyncio.ensure_future(run(loop, args.quiet,
args.console)))
for job in jobs:
proc_result = result_queue.get()
results['FAILED'] += proc_result['FAILED']
results['OK'] += proc_result['OK']
else:
loop = asyncio.get_event_loop()
if screen is not None and not args.console:
ui = screen([os.getpid()], get_live_results, loop)
ui.start()
else:
ui = None
if not args.quiet and args.console and not args.verbose:
def _display(loop):
try:
print(get_live_results(), end='\r')
except ClosedError:
return
loop.call_later(_REFRESH, _display, loop)
loop.call_soon(_display, loop)
results = _process(args)
if ui is not None:
ui.stop()
return results
async def worker(num, loop, results, args, stream, statsd):
global _STOP
quiet = args.quiet
duration = args.duration
verbose = args.verbose
exception = args.exception
if args.single_mode:
single = get_scenario(args.single_mode)
else:
single = None
count = 1
start = _now()
howlong = 0
setup = get_fixture('setup')
if setup is not None:
try:
options = await setup(num, args)
except Exception as e:
log(e)
await stream.put('WORKER_STOPPED')
return
if options is None:
options = {}
elif not isinstance(options, dict):
log('The setup function needs to return a dict')
await stream.put('WORKER_STOPPED')
return
else:
options = {}
ssetup = get_fixture('setup_session')
steardown = get_fixture('teardown_session')
async with Session(loop, stream, verbose, statsd, **options) as session:
session.args = args
session.worker_id = num
if ssetup is not None:
await ssetup(num, session)
while howlong < duration and not _STOP:
if args.max_runs and count > args.max_runs:
break
howlong = _now() - start
session.step = count
result = await step(session, quiet, verbose, stream, single)
if result == 1:
results['OK'] += 1
elif result == -1:
results['FAILED'] += 1
if exception:
await stream.put('WORKER_STOPPED')
_STOP = True
elif result == 0:
break
count += 1
if steardown is not None:
try:
await steardown(num, session)
except Exception as e:
# we can't stop the teardown process
log(e)
if not _STOP:
await stream.put('WORKER_STOPPED')