async def test_emit_recurrent_events() -> None:
message_bus = MessageBus()
sender = InMemoryEventSender()
model = ModelHandlerStatic(Model.empty())
sub = SubscriptionHandler(
InMemoryDb[Subscriber](Subscriber, lambda x: x.id), message_bus)
queue = WorkerTaskQueue()
fast = timedelta(seconds=0.001)
periodic = emit_recurrent_events(sender, model, sub, queue, message_bus,
fast, fast)
await periodic.start()
while len(sender.events) < 3:
await asyncio.sleep(0.01)
await periodic.stop()
model_event, subscriber_event, worker_event = sender.events[0:3]
assert model_event.kind == CoreEvent.ModelInfo
assert model_event.counters["model_count"] == 0
assert subscriber_event.kind == CoreEvent.SubscriberInfo
assert subscriber_event.counters["subscriber_count"] == 0
assert worker_event.kind == CoreEvent.WorkerQueueInfo
assert worker_event.counters["worker_count"] == 0
def message_bus() -> MessageBus:
return MessageBus()
def with_config(created: bool, system_data: SystemData, sdb: StandardDatabase,
cert_handler: CertificateHandler, config: CoreConfig) -> None:
reconfigure_logging(
config) # based on the config, logging might have changed
# only lg the editable config - to not log any passwords
log.debug(f"Starting with config: {config.editable}")
info = system_info()
event_sender = PostHogEventSender(
system_data) if config.runtime.usage_metrics else NoEventSender()
db = db_access(config, sdb, event_sender)
message_bus = MessageBus()
scheduler = Scheduler()
worker_task_queue = WorkerTaskQueue()
model = ModelHandlerDB(db.get_model_db(), config.runtime.plantuml_server)
template_expander = DBTemplateExpander(db.template_entity_db)
config_handler = ConfigHandlerService(db.config_entity_db,
db.config_validation_entity_db,
db.configs_model_db,
worker_task_queue, message_bus)
log_ship = event_stream(config, cert_handler.client_context)
cli_deps = CLIDependencies(
message_bus=message_bus,
event_sender=event_sender,
db_access=db,
model_handler=model,
worker_task_queue=worker_task_queue,
config=config,
template_expander=template_expander,
config_handler=config_handler,
cert_handler=cert_handler,
)
default_env = {
"graph": config.cli.default_graph,
"section": config.cli.default_section
}
cli = CLI(cli_deps, all_commands(cli_deps), default_env, alias_names())
subscriptions = SubscriptionHandler(db.subscribers_db, message_bus)
task_handler = TaskHandlerService(db.running_task_db, db.job_db,
message_bus, event_sender, subscriptions,
scheduler, cli, config)
core_config_handler = CoreConfigHandler(config, message_bus,
worker_task_queue, config_handler)
merge_outer_edges_handler = MergeOuterEdgesHandler(message_bus,
subscriptions,
task_handler)
cli_deps.extend(task_handler=task_handler)
api = Api(
db,
model,
subscriptions,
task_handler,
message_bus,
event_sender,
worker_task_queue,
cert_handler,
config_handler,
cli,
template_expander,
config,
)
event_emitter = emit_recurrent_events(event_sender, model, subscriptions,
worker_task_queue, message_bus,
timedelta(hours=1),
timedelta(hours=1))
async def on_start() -> None:
# queue must be created inside an async function!
cli_deps.extend(forked_tasks=Queue())
await db.start()
await event_sender.start()
await subscriptions.start()
await scheduler.start()
await worker_task_queue.start()
await event_emitter.start()
await cli.start()
await task_handler.start()
await core_config_handler.start()
await merge_outer_edges_handler.start()
await cert_handler.start()
await log_ship.start()
await api.start()
if created:
await event_sender.core_event(CoreEvent.SystemInstalled)
await event_sender.core_event(
CoreEvent.SystemStarted,
{
"version": version(),
"created_at": to_json(system_data.created_at),
"system": platform.system(),
"platform": platform.platform(),
"inside_docker": info.inside_docker,
},
cpu_count=info.cpus,
mem_total=info.mem_total,
mem_available=info.mem_available,
)
async def on_stop() -> None:
duration = utc() - info.started_at
await api.stop()
await log_ship.stop()
await cert_handler.stop()
await core_config_handler.stop()
await merge_outer_edges_handler.stop()
await task_handler.stop()
await cli.stop()
await event_sender.core_event(CoreEvent.SystemStopped,
total_seconds=int(
duration.total_seconds()))
await event_emitter.stop()
await worker_task_queue.stop()
await scheduler.stop()
await subscriptions.stop()
await db.stop()
await event_sender.stop()
async def async_initializer() -> Application:
async def clean_all_tasks() -> None:
log.info("Clean up all running tasks.")
for task in asyncio.all_tasks():
with suppress(asyncio.CancelledError):
if not task.done() or not task.cancelled():
task.cancel()
log.debug(f"Wait for task: {task}")
await task
async def on_start_stop(_: Application) -> AsyncIterator[None]:
await on_start()
log.info("Initialization done. Starting API.")
yield
log.info("Shutdown initiated. Stop all tasks.")
await on_stop()
await clean_all_tasks()
api.app.cleanup_ctx.append(on_start_stop)
return api.app
runner.run_app(
async_initializer(),
api.stop,
host=config.api.web_hosts,
port=config.api.web_port,
ssl_context=cert_handler.host_context,
)
def run(arguments: List[str]) -> None:
"""
Run application. When this method returns, the process is done.
:param arguments: the arguments provided to this process.
Note: this method is used in tests to specify arbitrary arguments.
"""
args = parse_args(arguments)
setup_process(args)
# after setup, logging is possible
info = system_info()
log.info(
f"Starting up version={info.version} on system with cpus={info.cpus}, "
f"available_mem={info.mem_available}, total_mem={info.mem_total}")
# wait here for an initial connection to the database before we continue. blocking!
created, system_data, sdb = DbAccess.connect(args, timedelta(seconds=60))
event_sender = NoEventSender(
) if args.analytics_opt_out else PostHogEventSender(system_data)
db = db_access(args, sdb, event_sender)
cert_handler = CertificateHandler.lookup(args, sdb)
message_bus = MessageBus()
scheduler = Scheduler()
worker_task_queue = WorkerTaskQueue()
model = ModelHandlerDB(db.get_model_db(), args.plantuml_server)
template_expander = DBTemplateExpander(db.template_entity_db)
config_handler = ConfigHandlerService(
db.config_entity_db,
db.config_validation_entity_db,
db.configs_model_db,
worker_task_queue,
message_bus,
)
cli_deps = CLIDependencies(
message_bus=message_bus,
event_sender=event_sender,
db_access=db,
model_handler=model,
worker_task_queue=worker_task_queue,
args=args,
template_expander=template_expander,
config_handler=config_handler,
)
default_env = {
"graph": args.cli_default_graph,
"section": args.cli_default_section
}
cli = CLI(cli_deps, all_commands(cli_deps), default_env, aliases())
subscriptions = SubscriptionHandler(db.subscribers_db, message_bus)
task_handler = TaskHandlerService(db.running_task_db, db.job_db,
message_bus, event_sender, subscriptions,
scheduler, cli, args)
cli_deps.extend(task_handler=task_handler)
api = Api(
db,
model,
subscriptions,
task_handler,
message_bus,
event_sender,
worker_task_queue,
cert_handler,
config_handler,
cli,
template_expander,
args,
)
event_emitter = emit_recurrent_events(event_sender, model, subscriptions,
worker_task_queue, message_bus,
timedelta(hours=1),
timedelta(hours=1))
async def on_start() -> None:
# queue must be created inside an async function!
cli_deps.extend(forked_tasks=Queue())
await db.start()
await event_sender.start()
await subscriptions.start()
await scheduler.start()
await worker_task_queue.start()
await event_emitter.start()
await cli.start()
await task_handler.start()
await api.start()
if created:
await event_sender.core_event(CoreEvent.SystemInstalled)
await event_sender.core_event(
CoreEvent.SystemStarted,
{
"version": version(),
"created_at": to_json(system_data.created_at),
"system": platform.system(),
"platform": platform.platform(),
"inside_docker": info.inside_docker,
},
cpu_count=info.cpus,
mem_total=info.mem_total,
mem_available=info.mem_available,
)
async def on_stop() -> None:
duration = utc() - info.started_at
await api.stop()
await task_handler.stop()
await cli.stop()
await event_sender.core_event(CoreEvent.SystemStopped,
total_seconds=int(
duration.total_seconds()))
await event_emitter.stop()
await worker_task_queue.stop()
await scheduler.stop()
await subscriptions.stop()
await db.stop()
await event_sender.stop()
async def async_initializer() -> Application:
async def on_start_stop(_: Application) -> AsyncIterator[None]:
await on_start()
log.info("Initialization done. Starting API.")
yield
log.info("Shutdown initiated. Stop all tasks.")
await on_stop()
api.app.cleanup_ctx.append(on_start_stop)
return api.app
tls_context: Optional[SSLContext] = None
if args.tls_cert:
tls_context = SSLContext(ssl.PROTOCOL_TLS)
tls_context.load_cert_chain(args.tls_cert, args.tls_key,
args.tls_password)
runner.run_app(async_initializer(),
api.stop,
host=args.host,
port=args.port,
ssl_context=tls_context)
async def handler(in_mem_db: SubscriberDb) -> SubscriptionHandler:
result = SubscriptionHandler(in_mem_db, MessageBus())
await result.add_subscription("sub_1", "test", True, timedelta(seconds=3))
return result