def _run_async(loop: asyncio.AbstractEventLoop, coro: Awaitable) -> Any:
"""Run a coroutine.
Does some special signal handling to perform a graceful exit.
Args:
loop: Event loop to run coroutine in.
coro: Coroutine to run
Returns:
Result of the coroutine.
"""
def graceful_exit() -> None:
pending = asyncio.Task.all_tasks()
for task in pending:
task.cancel()
loop.stop()
try:
loop.add_signal_handler(signal.SIGINT, graceful_exit)
loop.add_signal_handler(signal.SIGTERM, graceful_exit)
except NotImplementedError:
pass
try:
return loop.run_until_complete(coro)
except KeyboardInterrupt:
print("Exiting")
loop.close()
def __init__(self,
base_url: str,
username: str,
password: str,
tls: bool = True,
loop: asyncio.AbstractEventLoop = asyncio.get_event_loop()):
base_url = re.sub('http[s]?://', '', base_url)
if tls:
ws_url = f'wss://{base_url}/websocket'
rest_url = f'https://{base_url}'
else:
ws_url = f'ws://{base_url}/websocket'
rest_url = f'http://{base_url}'
self.ddp = client.DdpClient(ws_url, loop)
self.rest = client.RestClient(server_url=rest_url)
self._username = username
self._password = password
self._roomid_cache: Dict[str, m.Room] = {}
self._roomname_cache: Dict[str, m.Room] = {}
self._users_cache: Dict[str, m.UserRef] = {}
self.bots: List[b.BaseBot] = []
self._active_callbacks: List[asyncio.Task[Any]] = []
# Add signal handler for graceful stop
signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
for s in signals:
loop.add_signal_handler(
s, lambda s=s: asyncio.create_task(self.signal_handler(s)))
async def run(loop: asyncio.AbstractEventLoop,
server: FlagWriterServer) -> None:
await server.start()
for sig in [signal.SIGINT, signal.SIGTERM]:
loop.add_signal_handler(sig, lambda: on_shutdown(loop, server))
logger.info("Started flag writer server.")
await server.join()
def handle_error_signals(
self,
loop: asyncio.AbstractEventLoop,
error_signals: Optional[List[TSignal]] = None,
) -> None:
error_signals_to_handle = error_signals or _DEFAULT_ERROR_SIGNALS
for _signal in error_signals_to_handle:
loop.add_signal_handler(
_signal, lambda s=_signal: loop.create_task(self.shutdown(s)))
def _setup_termination(*, loop: asyncio.AbstractEventLoop):
def on_signal(signame):
_logger.info("Received signal %s. Exiting..." % signame)
loop.call_soon(lambda: loop.stop())
for signame in ('SIGINT', 'SIGTERM'):
loop.add_signal_handler(getattr(signal, signame),
functools.partial(on_signal, signame))
async def _do_async_fn(
async_fn: Callable[..., Coroutine[Any, Any, TReturn]],
args: Sequence[Any],
to_parent: BinaryIO,
loop: asyncio.AbstractEventLoop,
) -> TReturn:
# state: STARTED
update_state(to_parent, State.STARTED)
# A Future that will be set if any of the SHUTDOWN_SIGNALS signals are
# received causing _do_async_fn to raise a SystemExit
system_exit_signum: 'asyncio.Future[int]' = asyncio.Future()
# setup signal handlers.
for signum in SHUTDOWN_SIGNALS:
loop.add_signal_handler(
signum.value,
system_exit_signum.set_result,
signum,
)
# state: EXECUTING
update_state(to_parent, State.EXECUTING)
# Install a signal handler to set an asyncio.Event upon receiving a SIGINT
got_SIGINT = asyncio.Event()
loop.add_signal_handler(
signal.SIGINT,
got_SIGINT.set,
)
# First we need to generate a coroutine. We need this so we can throw
# exceptions into the running coroutine to allow it to handle keyboard
# interrupts.
async_fn_coro: Coroutine[Any, Any, TReturn] = async_fn(*args)
# The coroutine is then given to `_handle_coro` which waits for either the
# coroutine to finish, returning the result, or for a SIGINT signal at
# which point injects a `KeyboardInterrupt` into the running coroutine.
async_fn_task: 'asyncio.Future[TReturn]' = asyncio.ensure_future(
_handle_coro(async_fn_coro, got_SIGINT), )
# Now we wait for either a result from the coroutine or a SIGTERM which
# triggers immediate cancellation of the running coroutine.
done, pending = await asyncio.wait(
(async_fn_task, system_exit_signum),
return_when=asyncio.FIRST_COMPLETED,
)
# We prioritize the `SystemExit` case.
if system_exit_signum.done():
await _do_task_cleanup(async_fn_task)
raise SystemExit(system_exit_signum.result())
elif async_fn_task.done():
return async_fn_task.result()
else:
raise Exception("unreachable")
def _add_signal_handlers(loop: AbstractEventLoop) -> "Future[Any]":
"""Add signal handlers so shutdown can be handled normally, returning the stop future."""
log.info("Adding signal handlers...")
stop = loop.create_future()
for sig in SHUTDOWN_SIGNALS:
if sys.platform == "win32":
signal.signal(sig, lambda s, f: stop.set_result(None))
else:
loop.add_signal_handler(sig, stop.set_result, None)
return stop
def setup_graceful_shutdown(self, loop: asyncio.AbstractEventLoop):
# K8 uses SIGTERM on linux and SIGINT and windows
exit_signal = signal.SIGINT if platform.system(
) == 'Windows' else signal.SIGTERM
try:
loop.add_signal_handler(exit_signal, self.handle_signal)
except NotImplementedError:
# Disable graceful exit, but run anyway
logger.warning('%s does not support graceful shutdown',
platform.system())
def _add_signal_handlers(loop: AbstractEventLoop) -> None:
"""Add signal handlers so shutdown can be handled normally."""
log.info("Adding signal handlers...")
for sig in SHUTDOWN_SIGNALS:
if platform.system(
) == "Windows" or sys.platform == "win32": # stupid MyPy
signal.signal(sig, lambda s, f: asyncio.create_task(_terminate()))
else:
loop.add_signal_handler(sig,
lambda: asyncio.create_task(_terminate()))
async def exit_signal(loop: asyncio.AbstractEventLoop) -> AsyncGenerator[None, None]:
sigint_received = asyncio.Event()
for sig in [signal.SIGINT, signal.SIGTERM]:
# TODO also support Windows
loop.add_signal_handler(sig, sigint_received.set)
await sigint_received.wait()
try:
yield
finally:
loop.stop()
def run_sendria_servers(loop: asyncio.AbstractEventLoop) -> NoReturn:
# initialize db
loop.run_until_complete(db.setup(config.CONFIG.db))
# initialize and start webhooks
callbacks_enabled = callback.setup(
debug_mode=config.CONFIG.debug,
callback_webhook_url=config.CONFIG.callback_webhook_url,
callback_webhook_method=config.CONFIG.callback_webhook_method,
callback_webhook_auth=config.CONFIG.callback_webhook_auth,
)
if callbacks_enabled:
loop.create_task(callback.send_messages())
# initialize and start message saver
loop.create_task(db.message_saver())
# start smtp server
smtp.run(config.CONFIG.smtp_ip, config.CONFIG.smtp_port, config.CONFIG.smtp_auth, config.CONFIG.smtp_ident, config.CONFIG.debug)
logger.info('smtp server started', host=config.CONFIG.smtp_ip, port=config.CONFIG.smtp_port,
auth='enabled' if config.CONFIG.smtp_auth else 'disabled',
password_file=str(config.CONFIG.smtp_auth.path) if config.CONFIG.smtp_auth else None,
url=f'smtp://{config.CONFIG.smtp_ip}:{config.CONFIG.smtp_port}',
)
# initialize and start web server
app = http.setup()
runner = aiohttp.web.AppRunner(app)
loop.run_until_complete(runner.setup())
site = aiohttp.web.TCPSite(runner, host=config.CONFIG.http_ip, port=config.CONFIG.http_port)
server = site.start()
loop.run_until_complete(server)
logger.info('http server started',
host=config.CONFIG.http_ip, port=config.CONFIG.http_port,
url=f'http://{config.CONFIG.http_ip}:{config.CONFIG.http_port}',
auth='enabled' if config.CONFIG.http_auth else 'disabled',
password_file=str(config.CONFIG.http_auth.path) if config.CONFIG.http_auth else None,
)
# prepare for clean terminate
async def _initialize_aiohttp_services__stop() -> NoReturn:
for ws in set(app['websockets']):
await ws.close(code=aiohttp.WSCloseCode.GOING_AWAY, message='Server shutdown')
await app.shutdown()
SHUTDOWN.append(_initialize_aiohttp_services__stop())
signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
for s in signals:
loop.add_signal_handler(s, lambda s=s: asyncio.create_task(terminate_server(s, loop)))
def set_exit_handler(loop: asyncio.AbstractEventLoop, handler) -> None:
if IS_WINDOWS:
try:
import win32api
except ImportError:
version = ".".join(map(str, sys.version_info[:2]))
raise Exception(f"pypiwin32 is not installed for Python {version}")
else:
win32api.SetConsoleCtrlHandler(handler, True)
else:
import signal
loop.add_signal_handler(signal.SIGINT, handler)
loop.add_signal_handler(signal.SIGTERM, handler)
def _make_cancellation_event(loop: AbstractEventLoop) -> Event:
cancellation_event = Event()
def cancel(name, num):
msg = f'Received signal {name}'
if num == signal.SIGINT:
LOGGER.info(msg)
else:
LOGGER.warning(msg)
cancellation_event.set()
for signame in ['SIGINT', 'SIGTERM']:
signum = getattr(signal, signame)
loop.add_signal_handler(signum, cancel, signame, signum)
return cancellation_event
def run(loop: asyncio.AbstractEventLoop) -> None:
try:
loop.add_signal_handler(signal.SIGINT, loop.stop)
loop.add_signal_handler(signal.SIGTERM, loop.stop)
except (NotImplementedError, RuntimeError):
pass
asyncio.set_event_loop(loop)
try:
loop.run_forever()
except KeyboardInterrupt:
log.info('Received the signal to terminate the event loop.')
finally:
log.info('Cleaning up tasks.')
shutdown_loop(loop)
async def setup_server(loop: asyncio.AbstractEventLoop) -> None:
loop.create_task(heartbeat())
loop.add_reader(sys.stdin.fileno(), process_stdin)
loop.add_signal_handler(signal.SIGINT, ask_exit)
async def schedule_formatting(
sources: Set[Path],
fast: bool,
write_back: WriteBack,
mode: Mode,
report: "Report",
loop: asyncio.AbstractEventLoop,
executor: Executor,
) -> None:
"""Run formatting of `sources` in parallel using the provided `executor`.
(Use ProcessPoolExecutors for actual parallelism.)
`write_back`, `fast`, and `mode` options are passed to
:func:`format_file_in_place`.
"""
cache: Cache = {}
if write_back not in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
cache = read_cache(mode)
sources, cached = filter_cached(cache, sources)
for src in sorted(cached):
report.done(src, Changed.CACHED)
if not sources:
return
cancelled = []
sources_to_cache = []
lock = None
if write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
# For diff output, we need locks to ensure we don't interleave output
# from different processes.
manager = Manager()
lock = manager.Lock()
tasks = {
asyncio.ensure_future(
loop.run_in_executor(executor, format_file_in_place, src, fast,
mode, write_back, lock)): src
for src in sorted(sources)
}
pending = tasks.keys()
try:
loop.add_signal_handler(signal.SIGINT, cancel, pending)
loop.add_signal_handler(signal.SIGTERM, cancel, pending)
except NotImplementedError:
# There are no good alternatives for these on Windows.
pass
while pending:
done, _ = await asyncio.wait(pending,
return_when=asyncio.FIRST_COMPLETED)
for task in done:
src = tasks.pop(task)
if task.cancelled():
cancelled.append(task)
elif task.exception():
report.failed(src, str(task.exception()))
else:
changed = Changed.YES if task.result() else Changed.NO
# If the file was written back or was successfully checked as
# well-formatted, store this information in the cache.
if write_back is WriteBack.YES or (
write_back is WriteBack.CHECK
and changed is Changed.NO):
sources_to_cache.append(src)
report.done(src, changed)
if cancelled:
if sys.version_info >= (3, 7):
await asyncio.gather(*cancelled, return_exceptions=True)
else:
await asyncio.gather(*cancelled, loop=loop, return_exceptions=True)
if sources_to_cache:
write_cache(cache, sources_to_cache, mode)
async def _open_in_process(
async_fn: Callable[..., TReturn],
*args: Any,
loop: asyncio.AbstractEventLoop = None,
subprocess_kwargs: 'SubprocessKwargs' = None,
use_trio: bool = False,
) -> AsyncIterator[ProcessAPI[TReturn]]:
if use_trio and loop is not None:
raise ValueError("If using trio, cannot specify a loop")
proc: Process[TReturn] = Process(async_fn, args)
parent_r, child_w = os.pipe()
child_r, parent_w = os.pipe()
command = get_subprocess_command(child_r, child_w, use_trio)
sub_proc = await asyncio.create_subprocess_exec(
*command,
**_update_subprocess_kwargs(subprocess_kwargs, child_r, child_w),
)
if loop is None:
loop = asyncio.get_event_loop()
signal_queue: asyncio.Queue[signal.Signals] = asyncio.Queue()
for signum in RELAY_SIGNALS:
loop.add_signal_handler(
signum,
signal_queue.put_nowait,
signum,
)
# Monitoring
monitor_sub_proc_task = asyncio.ensure_future(_monitor_sub_proc(proc, sub_proc, parent_w))
relay_signals_task = asyncio.ensure_future(_relay_signals(proc, signal_queue))
monitor_state_task = asyncio.ensure_future(_monitor_state(proc, parent_r, child_w, loop))
startup_timeout = int(
os.getenv('ASYNCIO_RUN_IN_PROCESS_STARTUP_TIMEOUT', constants.STARTUP_TIMEOUT_SECONDS))
async with cleanup_tasks(monitor_sub_proc_task, relay_signals_task, monitor_state_task):
try:
await asyncio.wait_for(proc.wait_pid(), timeout=startup_timeout)
except asyncio.TimeoutError:
sub_proc.kill()
raise asyncio.TimeoutError(
f"{proc} took more than {startup_timeout} seconds to start up")
logger.debug(
"Got pid %d for %s, waiting for it to reach EXECUTING state before yielding",
proc.pid, proc)
# Wait until the child process has reached the EXECUTING
# state before yielding the context. This ensures that any
# calls to things like `terminate` or `kill` will be handled
# properly in the child process.
#
# The timeout ensures that if something is fundamentally wrong
# with the subprocess we don't hang indefinitely.
try:
logger.debug("Waiting for proc pid=%d to reach EXECUTING state", proc.pid)
await asyncio.wait_for(proc.wait_for_state(State.EXECUTING), timeout=startup_timeout)
except asyncio.TimeoutError:
sub_proc.kill()
raise asyncio.TimeoutError(
f"{proc} took more than {startup_timeout} seconds to start up")
try:
try:
yield proc
except KeyboardInterrupt as err:
# If a keyboard interrupt is encountered relay it to the
# child process and then give it a moment to cleanup before
# re-raising
logger.debug("Relaying SIGINT to pid=%d", sub_proc.pid)
try:
proc.send_signal(signal.SIGINT)
try:
await asyncio.wait_for(
proc.wait(), timeout=constants.SIGINT_TIMEOUT_SECONDS)
except asyncio.TimeoutError:
logger.debug(
"Timed out waiting for pid=%d to exit after relaying SIGINT",
sub_proc.pid,
)
except BaseException:
logger.exception(
"Unexpected error when terminating child; pid=%d", sub_proc.pid)
finally:
raise err
except asyncio.CancelledError as err:
# Send the child a SIGINT and wait SIGINT_TIMEOUT_SECONDS for it to terminate. If
# that times out, send a SIGTERM and wait SIGTERM_TIMEOUT_SECONDS before
# re-raising.
logger.debug(
"Got CancelledError while running subprocess pid=%d. Sending SIGINT.",
sub_proc.pid,
)
try:
proc.send_signal(signal.SIGINT)
try:
await asyncio.wait_for(
proc.wait(), timeout=constants.SIGINT_TIMEOUT_SECONDS)
except asyncio.TimeoutError:
logger.debug(
"Timed out waiting for pid=%d to exit after SIGINT, sending SIGTERM",
sub_proc.pid,
)
proc.terminate()
try:
await asyncio.wait_for(
proc.wait(), timeout=constants.SIGTERM_TIMEOUT_SECONDS)
except asyncio.TimeoutError:
logger.debug(
"Timed out waiting for pid=%d to exit after SIGTERM", sub_proc.pid)
except BaseException:
logger.exception(
"Unexpected error when terminating child; pid=%d", sub_proc.pid)
finally:
raise err
else:
# In the case that the yielded context block exits without an
# error we wait for the process to finish naturally. This can
# hang indefinitely.
logger.debug(
"Waiting for %s (pid=%d) to finish naturally, this can hang forever",
proc,
proc.pid,
)
await proc.wait()
finally:
if sub_proc.returncode is None:
# If the process has not returned at this stage we need to hard
# kill it to prevent it from hanging.
logger.warning(
"Child process pid=%d failed to exit cleanly. Sending SIGKILL",
sub_proc.pid,
# The `any` call is to include a stacktrace if this
# happened due to an exception but to omit it if this is
# somehow happening outside of an exception context.
exc_info=any(sys.exc_info()),
)
sub_proc.kill()
def _add_signal_handlers(loop: asyncio.AbstractEventLoop):
def handle_sigint():
logger.info("Caught SIGINT. Stopping ...")
loop.stop()
loop.add_signal_handler(signal.SIGINT, handle_sigint)
async def _open_in_process(
async_fn: Callable[..., TReturn],
*args: Any,
loop: asyncio.AbstractEventLoop = None,
) -> AsyncIterator[ProcessAPI[TReturn]]:
proc: Process[TReturn] = Process(async_fn, args)
parent_r, child_w = os.pipe()
child_r, parent_w = os.pipe()
parent_pid = os.getpid()
command = get_subprocess_command(child_r, child_w, parent_pid)
sub_proc = await asyncio.create_subprocess_exec(
*command,
# stdin=subprocess.PIPE,
# stdout=subprocess.PIPE,
# stderr=subprocess.PIPE,
pass_fds=(child_r, child_w),
)
if loop is None:
loop = asyncio.get_event_loop()
signal_queue: asyncio.Queue[signal.Signals] = asyncio.Queue()
for signum in RELAY_SIGNALS:
loop.add_signal_handler(
signum,
signal_queue.put_nowait,
signum,
)
# Monitoring
monitor_sub_proc_task = asyncio.ensure_future(
_monitor_sub_proc(proc, sub_proc, parent_w))
relay_signals_task = asyncio.ensure_future(
_relay_signals(proc, signal_queue))
monitor_state_task = asyncio.ensure_future(
_monitor_state(proc, parent_r, loop))
await proc.wait_pid()
# Wait until the child process has reached the STARTED
# state before yielding the context. This ensures that any
# calls to things like `terminate` or `kill` will be handled
# properly in the child process.
#
# The timeout ensures that if something is fundamentally wrong
# with the subprocess we don't hang indefinitely.
await proc.wait_for_state(State.STARTED)
try:
yield proc
except KeyboardInterrupt as err:
# If a keyboard interrupt is encountered relay it to the
# child process and then give it a moment to cleanup before
# re-raising
try:
proc.send_signal(signal.SIGINT)
try:
await asyncio.wait_for(proc.wait(), timeout=2)
except asyncio.TimeoutError:
pass
finally:
raise err
finally:
await proc.wait()
monitor_sub_proc_task.cancel()
try:
await monitor_sub_proc_task
except asyncio.CancelledError:
pass
monitor_state_task.cancel()
try:
await monitor_state_task
except asyncio.CancelledError:
pass
relay_signals_task.cancel()
try:
await relay_signals_task
except asyncio.CancelledError:
pass
def register_cancellation_event(cancellation_event: asyncio.Event,
loop: asyncio.AbstractEventLoop):
for signame in ('SIGHUP', 'SIGINT', 'SIGTERM'):
signum = getattr(signal, signame)
loop.add_signal_handler(signum, _cancel, signame, signum,
cancellation_event)
def create_tasks(
loop: asyncio.AbstractEventLoop,
lifecycle: Optional[Callable] = None,
registry: Optional[registries.BaseRegistry] = None,
standalone: bool = False,
priority: int = 0,
peering_name: str = peering.PEERING_DEFAULT_NAME,
namespace: Optional[str] = None,
):
"""
Create all the tasks needed to run the operator, but do not spawn/start them.
The tasks are properly inter-connected depending on the runtime specification.
They can be injected into any event loop as needed.
"""
# The freezer and the registry are scoped to this whole task-set, to sync them all.
lifecycle = lifecycle if lifecycle is not None else lifecycles.get_default_lifecycle(
)
registry = registry if registry is not None else registries.get_default_registry(
)
event_queue = asyncio.Queue(loop=loop)
freeze_flag = asyncio.Event(loop=loop)
should_stop = asyncio.Event(loop=loop)
tasks = []
# A top-level task for external stopping by setting a stop-flag. Once set,
# this task will exit, and thus all other top-level tasks will be cancelled.
tasks.extend([
loop.create_task(_stop_flag_checker(should_stop)),
])
# K8s-event posting. Events are queued in-memory and posted in the background.
# NB: currently, it is a global task, but can be made per-resource or per-object.
tasks.extend([
loop.create_task(posting.poster(event_queue=event_queue)),
])
# Monitor the peers, unless explicitly disabled.
ourselves: Optional[peering.Peer] = peering.Peer.detect(
id=peering.detect_own_id(),
priority=priority,
standalone=standalone,
namespace=namespace,
name=peering_name,
)
if ourselves:
tasks.extend([
loop.create_task(peering.peers_keepalive(ourselves=ourselves)),
loop.create_task(
queueing.watcher(
namespace=namespace,
resource=ourselves.resource,
handler=functools.partial(
peering.peers_handler,
ourselves=ourselves,
freeze=freeze_flag))), # freeze is set/cleared
])
# Resource event handling, only once for every known resource (de-duplicated).
for resource in registry.resources:
tasks.extend([
loop.create_task(
queueing.watcher(
namespace=namespace,
resource=resource,
handler=functools.partial(
handling.custom_object_handler,
lifecycle=lifecycle,
registry=registry,
resource=resource,
event_queue=event_queue,
freeze=freeze_flag))), # freeze is only checked
])
# On Ctrl+C or pod termination, cancel all tasks gracefully.
if threading.current_thread() is threading.main_thread():
loop.add_signal_handler(signal.SIGINT, should_stop.set)
loop.add_signal_handler(signal.SIGTERM, should_stop.set)
else:
logger.warning(
"OS signals are ignored: running not in the main thread.")
return tasks
def run(
func: Optional[Awaitable[None]] = None,
*,
finalize: Optional[Awaitable[None]] = None,
loop: AbstractEventLoop = None,
):
""" Configure the event loop to react to signals and exceptions then
run the provided coroutine and loop forever.
Shutdown the event loop when a signal or exception is received or the
supplied function explicitly requests the loop to stop.
This function provides some of the common boilerplate typically needed
when running asyncio applications. It registers signal handlers that
listen for SIGINT and SIGTERM that will stop the event loop and trigger
application shutdown actions. It registers a global exception handler
that will stop the loop and trigger application shutdown actions. This
helps catch a common problem that occurs in asyncio applications in
which an exception occurs in a task that was spun off but is not
reported until the event loop is stopped. This approach allows users
to be notified about these issues as soon as possible.
:param func: A coroutine to run before looping forever. This coroutine
is typically the "main" coroutine from which all other work is spawned.
The event loop will continue to run after the supplied coroutine
completes.
:param finalize: An optional coroutine to run when shutting down. Use this
to perform any graceful cleanup activities such as finalising log files,
disconnecting from services such as databases, etc.
:param loop: An optional event loop to run. If not supplied the default
event loop is used (i.e., whatever ``asyncio.get_event_loop()`` returns.
"""
logger.debug("Application runner starting")
if func:
if not (inspect.isawaitable(func)
or inspect.iscoroutinefunction(func)):
raise Exception("func must be a coroutine or a coroutine function "
f"that takes no arguments, got {func}")
if finalize:
if not (inspect.isawaitable(finalize)
or inspect.iscoroutinefunction(finalize)):
raise Exception(
"finalize must be a coroutine or a coroutine function "
f"that takes no arguments, got {finalize}")
# Use a supplied loop or the default event loop. If the loop is closed
# (which can happen in unit tests) then create a new event loop.
loop = loop or asyncio.get_event_loop()
if loop.is_closed():
loop = asyncio.new_event_loop()
def signal_handler(loop, sig):
logger.info(f"Caught {sig.name}, stopping.")
loop.call_soon(loop.stop)
loop.add_signal_handler(SIGINT, signal_handler, loop, SIGINT)
loop.add_signal_handler(SIGTERM, signal_handler, loop, SIGTERM)
def exception_handler(loop, context):
exc_msg = context["message"]
exc = context["exception"]
logger.exception(f"Caught exception: {exc_msg}", exc_info=exc)
loop.call_soon(loop.stop)
loop.set_exception_handler(exception_handler)
try:
if func:
if inspect.iscoroutinefunction(func):
func = func() # type: ignore
assert func is not None
loop.create_task(func)
loop.run_forever()
finally:
logger.debug("Application shutdown sequence starting")
if finalize:
if inspect.iscoroutinefunction(finalize):
finalize = finalize() # type: ignore
assert finalize is not None
loop.run_until_complete(finalize)
# Shutdown any outstanding tasks that are left running
pending_tasks = all_tasks(loop=loop)
if pending_tasks:
logger.debug(f"Cancelling {len(pending_tasks)} pending tasks.")
for task in pending_tasks:
logger.debug(f"Cancelling task: {task}")
task.cancel()
try:
loop.run_until_complete(asyncio.gather(*pending_tasks))
except asyncio.CancelledError:
pass
loop.run_until_complete(loop.shutdown_asyncgens())
logger.debug("Application shutdown sequence complete")
loop.close()
logger.debug("Application runner stopped")
def install_signal_handlers(self, loop: asyncio.AbstractEventLoop) -> None:
for signum, handler in self.signal_handlers.items():
loop.add_signal_handler(signum, handler)
def launch(loop: asyncio.AbstractEventLoop, workspace: Path, print_log: bool, print_metric_log: bool, verbose: bool):
config_file = workspace / 'config.json'
if not workspace.exists():
print(f'{workspace.resolve()} is not exist.', file=sys.stderr)
return False
elif not config_file.exists():
print(f'{config_file.resolve()} is not exist.', file=sys.stderr)
return False
with config_file.open() as local_config_fp, \
GLOBAL_CFG_PATH.open() as global_config_fp:
local_cfg_source: Dict[str, Any] = json.load(local_config_fp)
global_cfg_source: Dict[str, Any] = json.load(global_config_fp)
bench_cfgs = parse_workload_cfg(local_cfg_source['workloads'])
perf_cfg = parse_perf_cfg(global_cfg_source['perf'], local_cfg_source.get('perf', {'extra_events': []}))
rabbit_cfg = parse_rabbit_mq_cfg(global_cfg_source['rabbitMQ'])
launcher_cfg = parse_launcher_cfg(local_cfg_source.get('launcher'))
task_map: Dict[asyncio.Task, Benchmark] = dict()
was_successful = True
result_file = workspace / 'result.json'
if result_file.exists():
result_file.unlink()
def stop_all():
print('force stop all tasks...', file=sys.stderr)
for a_task in task_map:
if not a_task.done():
a_task.cancel()
nonlocal was_successful
was_successful = False
def store_runtime(a_task: asyncio.Task):
a_bench = task_map[a_task]
if result_file.exists():
with result_file.open('r+') as fp:
try:
original: Dict[str, Any] = json.load(fp)
if 'runtime' in original:
original['runtime'][a_bench.identifier] = a_bench.runtime
else:
original['runtime'] = {a_bench.identifier: a_bench.runtime}
fp.seek(0)
json.dump(original, fp, indent=4)
except json.decoder.JSONDecodeError:
fp.seek(0)
fp.truncate()
json.dump({'runtime': {a_bench.identifier: a_bench.runtime}}, fp, indent=4)
else:
with result_file.open('w') as fp:
json.dump({'runtime': {a_bench.identifier: a_bench.runtime}}, fp, indent=4)
a_task.remove_done_callback(store_runtime)
benches = tuple(create_benchmarks(bench_cfgs, perf_cfg, rabbit_cfg, workspace, verbose))
loop.add_signal_handler(signal.SIGHUP, stop_all)
loop.add_signal_handler(signal.SIGTERM, stop_all)
loop.add_signal_handler(signal.SIGINT, stop_all)
with power_monitor(workspace), hyper_threading_guard(loop, launcher_cfg.hyper_threading):
# invoke benchmark loaders in parallel and wait for launching actual benchmarks
loop.run_until_complete(asyncio.wait(tuple(bench.start_and_pause(print_log) for bench in benches)))
for bench in benches:
bench.resume()
for bench in benches:
task: asyncio.Task = loop.create_task(bench.monitor(print_metric_log))
task.add_done_callback(store_runtime)
task_map[task] = bench
# start monitoring
return_when = asyncio.FIRST_COMPLETED if launcher_cfg.stops_with_the_first else asyncio.ALL_COMPLETED
finished, unfinished = loop.run_until_complete(asyncio.wait(task_map.keys(), return_when=return_when))
for task in unfinished:
task.cancel()
loop.run_until_complete(asyncio.gather(*unfinished))
loop.remove_signal_handler(signal.SIGHUP)
loop.remove_signal_handler(signal.SIGTERM)
loop.remove_signal_handler(signal.SIGINT)
# run post scripts
for script in launcher_cfg.post_scripts:
name = str(script)
if name.endswith('.py'):
name = name[:-3]
script_module = importlib.import_module(name.replace('/', '.'))
script_module.run(workspace, GLOBAL_CFG_PATH)
return was_successful
def run_single(
app: Type[ASGIFramework],
config: Config,
*,
loop: asyncio.AbstractEventLoop,
sock: Optional[socket] = None,
is_child: bool = False,
) -> None:
"""Create a server to run the app on given the options.
Arguments:
app: The ASGI Framework to run.
config: The configuration that defines the server.
loop: Asyncio loop to create the server in, if None, take default one.
"""
if loop is None:
warnings.warn(
'Event loop is not specified, this can cause unexpected errors')
loop = asyncio.get_event_loop()
if config.pid_path is not None and not is_child:
_write_pid_file(config.pid_path)
loop.set_debug(config.debug)
if hasattr(app, 'startup'):
loop.run_until_complete(app.startup()) # type: ignore
if sock is not None:
create_server = loop.create_server(
lambda: Server(app, loop, config),
ssl=config.ssl,
sock=sock,
reuse_port=is_child,
)
elif config.file_descriptor is not None:
sock = socket_fromfd(config.file_descriptor, AF_UNIX, SOCK_STREAM)
create_server = loop.create_server(
lambda: Server(app, loop, config),
ssl=config.ssl,
sock=sock,
)
elif config.unix_domain is not None:
create_server = loop.create_unix_server(
lambda: Server(app, loop, config),
config.unix_domain,
ssl=config.ssl,
)
else:
create_server = loop.create_server(
lambda: Server(app, loop, config),
host=config.host,
port=config.port,
ssl=config.ssl,
reuse_port=is_child,
)
server = loop.run_until_complete(create_server)
if platform.system() == 'Windows':
loop.create_task(_windows_signal_support())
try:
loop.add_signal_handler(signal.SIGINT, _raise_shutdown)
loop.add_signal_handler(signal.SIGTERM, _raise_shutdown)
except NotImplementedError:
pass # Unix only
reload_ = False
try:
if config.use_reloader:
loop.run_until_complete(_observe_changes())
reload_ = True
else:
loop.run_forever()
except (SystemExit, KeyboardInterrupt):
pass
finally:
server.close()
loop.run_until_complete(server.wait_closed())
loop.run_until_complete(loop.shutdown_asyncgens())
try:
loop.remove_signal_handler(signal.SIGINT)
loop.remove_signal_handler(signal.SIGTERM)
except NotImplementedError:
pass # Unix only
if hasattr(app, 'cleanup'):
loop.run_until_complete(app.cleanup()) # type: ignore
loop.close()
if reload_:
# Restart this process (only safe for dev/debug)
os.execv(sys.executable, [sys.executable] + sys.argv)
def _setup_signal_handlers(loop: AbstractEventLoop) -> None:
loop.add_signal_handler(signal.SIGINT, loop.stop)
loop.add_signal_handler(signal.SIGTERM, loop.stop)