1use axum::Router;
17use std::collections::HashSet;
18use std::sync::Arc;
19
20use tokio_util::sync::CancellationToken;
21use uuid::Uuid;
22
23use crate::backends::OopSpawnConfig;
24use crate::client_hub::ClientHub;
25use crate::config::ConfigProvider;
26use crate::context::GearContextBuilder;
27use crate::registry::{
28 ApiGatewayCap, GearEntry, GearRegistry, GrpcHubCap, RegistryError, RestApiCap, RunnableCap,
29 SystemCap,
30};
31use crate::runtime::{GearManager, GrpcInstallerStore, OopSpawnOptions, SystemContext};
32
33#[cfg(feature = "db")]
34use crate::registry::DatabaseCap;
35
36#[derive(Clone)]
38pub enum DbOptions {
39 None,
41 #[cfg(feature = "db")]
43 Manager(Arc<toolkit_db::DbManager>),
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum RunMode {
49 Full,
51 MigrateOnly,
53}
54
55pub const TOOLKIT_DIRECTORY_ENDPOINT_ENV: &str = "TOOLKIT_DIRECTORY_ENDPOINT";
57
58pub const TOOLKIT_MODULE_CONFIG_ENV: &str = "TOOLKIT_MODULE_CONFIG";
60
61pub const DEFAULT_SHUTDOWN_DEADLINE: std::time::Duration = std::time::Duration::from_secs(35);
67
68pub struct HostRuntime {
72 registry: GearRegistry,
73 ctx_builder: GearContextBuilder,
74 instance_id: Uuid,
75 gear_manager: Arc<GearManager>,
76 grpc_installers: Arc<GrpcInstallerStore>,
77 #[allow(dead_code)]
78 client_hub: Arc<ClientHub>,
79 cancel: CancellationToken,
80 #[allow(dead_code)]
81 db_options: DbOptions,
82 oop_options: Option<OopSpawnOptions>,
84 shutdown_deadline: std::time::Duration,
86}
87
88impl HostRuntime {
89 pub fn new(
93 registry: GearRegistry,
94 gears_cfg: Arc<dyn ConfigProvider>,
95 db_options: DbOptions,
96 client_hub: Arc<ClientHub>,
97 cancel: CancellationToken,
98 instance_id: Uuid,
99 oop_options: Option<OopSpawnOptions>,
100 ) -> Self {
101 let gear_manager = Arc::new(GearManager::new());
103 let grpc_installers = Arc::new(GrpcInstallerStore::new());
104
105 let ctx_builder =
107 GearContextBuilder::new(instance_id, gears_cfg, client_hub.clone(), cancel.clone());
108 #[cfg(feature = "db")]
109 let ctx_builder = match &db_options {
110 DbOptions::Manager(mgr) => ctx_builder.with_db_manager(mgr.clone()),
111 DbOptions::None => ctx_builder,
112 };
113
114 Self {
115 registry,
116 ctx_builder,
117 instance_id,
118 gear_manager,
119 grpc_installers,
120 client_hub,
121 cancel,
122 db_options,
123 oop_options,
124 shutdown_deadline: DEFAULT_SHUTDOWN_DEADLINE,
125 }
126 }
127
128 #[must_use]
144 pub fn with_shutdown_deadline(mut self, deadline: std::time::Duration) -> Self {
145 self.shutdown_deadline = deadline;
146 self
147 }
148
149 pub fn run_pre_init_phase(&self) -> Result<(), RegistryError> {
156 tracing::info!("Phase: pre_init");
157
158 let sys_ctx = SystemContext::new(
159 self.instance_id,
160 Arc::clone(&self.gear_manager),
161 Arc::clone(&self.grpc_installers),
162 );
163
164 for entry in self.registry.gears() {
165 if self.cancel.is_cancelled() {
167 tracing::warn!("Pre-init phase cancelled by signal");
168 return Err(RegistryError::Cancelled);
169 }
170
171 if let Some(sys_mod) = entry.caps.query::<SystemCap>() {
172 tracing::debug!(gear = entry.name, "Running system pre_init");
173 sys_mod
174 .pre_init(&sys_ctx)
175 .map_err(|e| RegistryError::PreInit {
176 gear: entry.name,
177 source: e,
178 })?;
179 }
180 }
181
182 Ok(())
183 }
184
185 #[cfg(feature = "db")]
187 async fn gear_context(
188 &self,
189 gear_name: &'static str,
190 ) -> Result<crate::context::GearCtx, RegistryError> {
191 self.ctx_builder
192 .for_gear(gear_name)
193 .await
194 .map_err(|e| RegistryError::DbMigrate {
195 gear: gear_name,
196 source: e,
197 })
198 }
199
200 #[cfg(feature = "db")]
202 async fn db_migration_target(
203 &self,
204 gear_name: &'static str,
205 ctx: &crate::context::GearCtx,
206 db_gear: Option<Arc<dyn crate::contracts::DatabaseCapability>>,
207 ) -> Result<
208 Option<(
209 toolkit_db::Db,
210 Arc<dyn crate::contracts::DatabaseCapability>,
211 )>,
212 RegistryError,
213 > {
214 let Some(dbm) = db_gear else {
215 return Ok(None);
216 };
217
218 let db = match &self.db_options {
222 DbOptions::None => None,
223 #[cfg(feature = "db")]
224 DbOptions::Manager(mgr) => {
225 mgr.get(gear_name)
226 .await
227 .map_err(|e| RegistryError::DbMigrate {
228 gear: gear_name,
229 source: e.into(),
230 })?
231 }
232 };
233
234 _ = ctx; Ok(db.map(|db| (db, dbm)))
236 }
237
238 #[cfg(feature = "db")]
243 async fn migrate_gear(
244 gear_name: &'static str,
245 db: &toolkit_db::Db,
246 db_gear: Arc<dyn crate::contracts::DatabaseCapability>,
247 ) -> Result<(), RegistryError> {
248 let migrations = db_gear.migrations();
250
251 if migrations.is_empty() {
252 tracing::debug!(gear = gear_name, "No migrations to run");
253 return Ok(());
254 }
255
256 tracing::debug!(
257 gear = gear_name,
258 count = migrations.len(),
259 "Running DB migrations"
260 );
261
262 let result =
264 toolkit_db::migration_runner::run_migrations_for_gear(db, gear_name, migrations)
265 .await
266 .map_err(|e| RegistryError::DbMigrate {
267 gear: gear_name,
268 source: anyhow::Error::new(e),
269 })?;
270
271 tracing::info!(
272 gear = gear_name,
273 applied = result.applied,
274 skipped = result.skipped,
275 "DB migrations completed"
276 );
277
278 Ok(())
279 }
280
281 #[cfg(feature = "db")]
290 async fn run_db_phase(&self) -> Result<(), RegistryError> {
291 tracing::info!("Phase: db (before init)");
292
293 for entry in self.registry.gears_by_system_priority() {
294 if self.cancel.is_cancelled() {
296 tracing::warn!("DB migration phase cancelled by signal");
297 return Err(RegistryError::Cancelled);
298 }
299
300 let ctx = self.gear_context(entry.name).await?;
301 let db_gear = entry.caps.query::<DatabaseCap>();
302
303 match self
304 .db_migration_target(entry.name, &ctx, db_gear.clone())
305 .await?
306 {
307 Some((db, dbm)) => {
308 Self::migrate_gear(entry.name, &db, dbm).await?;
309 }
310 None if db_gear.is_some() => {
311 tracing::debug!(
312 gear = entry.name,
313 "Gear has DbGear trait but no DB handle (no config)"
314 );
315 }
316 None => {}
317 }
318 }
319
320 Ok(())
321 }
322
323 async fn run_init_phase(&self) -> Result<(), RegistryError> {
327 tracing::info!("Phase: init");
328
329 for entry in self.registry.gears_by_system_priority() {
330 let ctx =
331 self.ctx_builder
332 .for_gear(entry.name)
333 .await
334 .map_err(|e| RegistryError::Init {
335 gear: entry.name,
336 source: e,
337 })?;
338 tracing::info!(gear = entry.name, "Initializing a gear...");
339 entry
340 .core
341 .init(&ctx)
342 .await
343 .map_err(|e| RegistryError::Init {
344 gear: entry.name,
345 source: e,
346 })?;
347 tracing::info!(gear = entry.name, "Initialized a gear.");
348 }
349
350 Ok(())
351 }
352
353 async fn run_post_init_phase(&self) -> Result<(), RegistryError> {
360 tracing::info!("Phase: post_init");
361
362 let sys_ctx = SystemContext::new(
363 self.instance_id,
364 Arc::clone(&self.gear_manager),
365 Arc::clone(&self.grpc_installers),
366 );
367
368 for entry in self.registry.gears_by_system_priority() {
369 if let Some(sys_mod) = entry.caps.query::<SystemCap>() {
370 sys_mod
371 .post_init(&sys_ctx)
372 .await
373 .map_err(|e| RegistryError::PostInit {
374 gear: entry.name,
375 source: e,
376 })?;
377 }
378 }
379
380 Ok(())
381 }
382
383 async fn run_rest_phase(&self) -> Result<Router, RegistryError> {
390 tracing::info!("Phase: rest (sync)");
391
392 let mut router = Router::new();
393
394 let host_count = self
396 .registry
397 .gears()
398 .iter()
399 .filter(|e| e.caps.has::<ApiGatewayCap>())
400 .count();
401
402 match host_count {
403 0 => {
404 return if self
405 .registry
406 .gears()
407 .iter()
408 .any(|e| e.caps.has::<RestApiCap>())
409 {
410 Err(RegistryError::RestRequiresHost)
411 } else {
412 Ok(router)
413 };
414 }
415 1 => { }
416 _ => return Err(RegistryError::MultipleRestHosts),
417 }
418
419 let host_idx = self
421 .registry
422 .gears()
423 .iter()
424 .position(|e| e.caps.has::<ApiGatewayCap>())
425 .ok_or(RegistryError::RestHostNotFoundAfterValidation)?;
426 let host_entry = &self.registry.gears()[host_idx];
427 let Some(host) = host_entry.caps.query::<ApiGatewayCap>() else {
428 return Err(RegistryError::RestHostMissingFromEntry);
429 };
430 let host_ctx = self
431 .ctx_builder
432 .for_gear(host_entry.name)
433 .await
434 .map_err(|e| RegistryError::RestPrepare {
435 gear: host_entry.name,
436 source: e,
437 })?;
438
439 let registry: &dyn crate::contracts::OpenApiRegistry = host.as_registry();
441
442 let hc_registry = Arc::new(
446 crate::healthcheck::RestHealthcheckRegistry::with_cancellation(
447 host_ctx.cancellation_token().clone(),
448 ),
449 );
450
451 router = host
453 .rest_prepare(&host_ctx, router, hc_registry.clone())
454 .map_err(|source| RegistryError::RestPrepare {
455 gear: host_entry.name,
456 source,
457 })?;
458
459 for e in self.registry.gears() {
461 if let Some(rest) = e.caps.query::<RestApiCap>() {
462 let ctx = self.ctx_builder.for_gear(e.name).await.map_err(|err| {
463 RegistryError::RestRegister {
464 gear: e.name,
465 source: err,
466 }
467 })?;
468
469 router = rest
470 .register_rest(&ctx, router, registry)
471 .map_err(|source| RegistryError::RestRegister {
472 gear: e.name,
473 source,
474 })?;
475
476 if let Some(hc) = rest.healthcheck(&ctx) {
478 hc_registry.register(e.name, hc);
479 }
480 }
481 }
482
483 router = host
485 .rest_finalize(&host_ctx, router, hc_registry)
486 .map_err(|source| RegistryError::RestFinalize {
487 gear: host_entry.name,
488 source,
489 })?;
490
491 Ok(router)
492 }
493
494 async fn run_grpc_phase(&self) -> Result<(), RegistryError> {
498 tracing::info!("Phase: grpc (registration)");
499
500 if self.registry.grpc_hub.is_none() && self.registry.grpc_services.is_empty() {
502 return Ok(());
503 }
504
505 if self.registry.grpc_hub.is_none() && !self.registry.grpc_services.is_empty() {
507 return Err(RegistryError::GrpcRequiresHub);
508 }
509
510 if let Some(hub_name) = &self.registry.grpc_hub {
512 let mut gears_data = Vec::new();
513 let mut seen = HashSet::new();
514
515 for (gear_name, service_gear) in &self.registry.grpc_services {
517 let ctx = self.ctx_builder.for_gear(gear_name).await.map_err(|err| {
518 RegistryError::GrpcRegister {
519 gear: gear_name.clone(),
520 source: err,
521 }
522 })?;
523
524 let installers = service_gear
525 .get_grpc_services(&ctx)
526 .await
527 .map_err(|source| RegistryError::GrpcRegister {
528 gear: gear_name.clone(),
529 source,
530 })?;
531
532 for reg in &installers {
533 if !seen.insert(reg.service_name) {
534 return Err(RegistryError::GrpcRegister {
535 gear: gear_name.clone(),
536 source: anyhow::anyhow!(
537 "Duplicate gRPC service name: {}",
538 reg.service_name
539 ),
540 });
541 }
542 }
543
544 gears_data.push(crate::runtime::GearInstallers {
545 gear_name: gear_name.clone(),
546 installers,
547 });
548 }
549
550 self.grpc_installers
551 .set(crate::runtime::GrpcInstallerData { gears: gears_data })
552 .map_err(|source| RegistryError::GrpcRegister {
553 gear: hub_name.clone(),
554 source,
555 })?;
556 }
557
558 Ok(())
559 }
560
561 async fn run_start_phase(&self) -> Result<(), RegistryError> {
565 tracing::info!("Phase: start");
566
567 for e in self.registry.gears_by_system_priority() {
568 if let Some(s) = e.caps.query::<RunnableCap>() {
569 tracing::debug!(
570 gear = e.name,
571 is_system = e.caps.has::<SystemCap>(),
572 "Starting stateful gear"
573 );
574 s.start(self.cancel.clone())
575 .await
576 .map_err(|source| RegistryError::Start {
577 gear: e.name,
578 source,
579 })?;
580 tracing::info!(gear = e.name, "Started gear");
581 }
582 }
583
584 Ok(())
585 }
586
587 async fn stop_one_gear(entry: &GearEntry, cancel: CancellationToken) {
589 if let Some(s) = entry.caps.query::<RunnableCap>() {
590 match s.stop(cancel).await {
591 Err(err) => {
592 tracing::warn!(gear = entry.name, error = %err, "Failed to stop gear");
593 }
594 _ => {
595 tracing::info!(gear = entry.name, "Stopped gear");
596 }
597 }
598 }
599 }
600
601 async fn run_stop_phase(&self) -> Result<(), RegistryError> {
626 tracing::info!("Phase: stop");
627
628 let deadline = self.shutdown_deadline;
629
630 for e in self.registry.gears().iter().rev() {
632 let gear_name = e.name;
633
634 let deadline_token = CancellationToken::new();
637 let deadline_token_for_timeout = deadline_token.clone();
638
639 let deadline_task = tokio::spawn(async move {
641 tokio::time::sleep(deadline).await;
642 tracing::warn!(
643 gear = gear_name,
644 deadline_secs = deadline.as_secs(),
645 "Gear shutdown deadline reached, sending hard-stop signal"
646 );
647 deadline_token_for_timeout.cancel();
648 });
649
650 Self::stop_one_gear(e, deadline_token).await;
653
654 deadline_task.abort();
656 #[allow(clippy::let_underscore_must_use)]
657 let _ = deadline_task.await;
658 }
659
660 Ok(())
661 }
662
663 async fn run_stop_phase_guarded(&self) -> Result<(), RegistryError> {
668 let gear_count = u32::try_from(self.registry.gears().len().max(1)).unwrap_or(1);
669 let stop_timeout = self
670 .shutdown_deadline
671 .checked_mul(gear_count)
672 .and_then(|d| d.checked_add(std::time::Duration::from_secs(5)))
673 .unwrap_or(self.shutdown_deadline);
674
675 let (disarm_tx, disarm_rx) = std::sync::mpsc::channel::<()>();
679 std::thread::spawn(move || {
680 match disarm_rx.recv_timeout(stop_timeout) {
681 Ok(()) | Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
682 }
684 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
685 tracing::warn!(
686 timeout_secs = stop_timeout.as_secs(),
687 "shutdown: stop phase timed out, force exiting"
688 );
689 std::process::exit(1);
690 }
691 }
692 });
693
694 let stop_result = self.run_stop_phase().await;
695 let _ = disarm_tx.send(()).ok();
699
700 stop_result
701 }
702
703 async fn run_oop_spawn_phase(&self) -> Result<(), RegistryError> {
708 let oop_opts = match &self.oop_options {
709 Some(opts) if !opts.gears.is_empty() => opts,
710 _ => return Ok(()),
711 };
712
713 tracing::info!("Phase: oop_spawn");
714
715 let directory_endpoint = self.wait_for_grpc_hub_endpoint().await;
717
718 for gear_cfg in &oop_opts.gears {
719 let mut env = gear_cfg.env.clone();
722 env.insert(
723 TOOLKIT_MODULE_CONFIG_ENV.to_owned(),
724 gear_cfg.rendered_config_json.clone(),
725 );
726 if let Some(ref endpoint) = directory_endpoint {
727 env.insert(TOOLKIT_DIRECTORY_ENDPOINT_ENV.to_owned(), endpoint.clone());
728 }
729
730 let args = gear_cfg.args.clone();
732
733 let spawn_config = OopSpawnConfig {
734 gear_name: gear_cfg.gear_name.clone(),
735 binary: gear_cfg.binary.clone(),
736 args,
737 env,
738 working_directory: gear_cfg.working_directory.clone(),
739 };
740
741 oop_opts
742 .backend
743 .spawn(spawn_config)
744 .await
745 .map_err(|e| RegistryError::OopSpawn {
746 gear: gear_cfg.gear_name.clone(),
747 source: e,
748 })?;
749
750 tracing::info!(
751 gear = %gear_cfg.gear_name,
752 directory_endpoint = ?directory_endpoint,
753 "Spawned OoP gear via backend"
754 );
755 }
756
757 Ok(())
758 }
759
760 async fn wait_for_grpc_hub_endpoint(&self) -> Option<String> {
765 const POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(10);
766 const MAX_WAIT: std::time::Duration = std::time::Duration::from_secs(5);
767
768 let grpc_hub = self
770 .registry
771 .gears()
772 .iter()
773 .find_map(|e| e.caps.query::<GrpcHubCap>());
774
775 let Some(hub) = grpc_hub else {
776 return None; };
778
779 let start = std::time::Instant::now();
780
781 loop {
782 if let Some(endpoint) = hub.bound_endpoint() {
783 tracing::debug!(
784 endpoint = %endpoint,
785 elapsed_ms = start.elapsed().as_millis(),
786 "gRPC hub endpoint available"
787 );
788 return Some(endpoint);
789 }
790
791 if start.elapsed() > MAX_WAIT {
792 tracing::warn!("Timed out waiting for gRPC hub to bind");
793 return None;
794 }
795
796 tokio::time::sleep(POLL_INTERVAL).await;
797 }
798 }
799
800 pub async fn run_gear_phases(self) -> anyhow::Result<()> {
809 self.run_phases_internal(RunMode::Full).await
810 }
811
812 pub async fn run_migration_phases(self) -> anyhow::Result<()> {
822 self.run_phases_internal(RunMode::MigrateOnly).await
823 }
824
825 async fn run_phases_internal(self, mode: RunMode) -> anyhow::Result<()> {
848 match mode {
850 RunMode::Full => {
851 tracing::info!("Running full lifecycle (all phases)");
852 }
853 RunMode::MigrateOnly => {
854 tracing::info!("Running in migration mode (pre-init + db phases only)");
855 }
856 }
857
858 self.run_pre_init_phase()?;
860
861 #[cfg(feature = "db")]
863 {
864 self.run_db_phase().await?;
865 }
866 #[cfg(not(feature = "db"))]
867 {
868 }
870
871 if mode == RunMode::MigrateOnly {
873 tracing::info!("Migration phases completed successfully");
874 return Ok(());
875 }
876
877 self.run_init_phase().await?;
879
880 self.run_post_init_phase().await?;
882
883 let _router = self.run_rest_phase().await?;
885
886 self.run_grpc_phase().await?;
888
889 self.run_start_phase().await?;
891
892 self.run_oop_spawn_phase().await?;
894
895 self.cancel.cancelled().await;
897
898 self.run_stop_phase_guarded().await?;
902 Ok(())
903 }
904}
905
906#[cfg(feature = "bootstrap")]
908impl HostRuntime {
909 fn external_deps(&self) -> Vec<String> {
913 use std::collections::{BTreeSet, HashSet};
914
915 let present: HashSet<&str> = self.registry.gears().iter().map(|e| e.name).collect();
916 let mut deps = BTreeSet::new();
917 for entry in self.registry.gears() {
918 for dep in entry.deps() {
919 if !present.contains(dep) {
920 deps.insert((*dep).to_owned());
921 }
922 }
923 }
924 deps.into_iter().collect()
925 }
926
927 async fn compose_oop_router(
933 &self,
934 options: &crate::runtime::OopServeOptions,
935 hc_registry: &Arc<crate::healthcheck::RestHealthcheckRegistry>,
936 ) -> anyhow::Result<(Router, String)> {
937 use crate::api::{OpenApiInfo, OpenApiRegistryImpl};
938 use anyhow::Context as _;
939
940 let registry = OpenApiRegistryImpl::new();
941 let mut router = Router::new();
942
943 for entry in self.registry.gears() {
944 if let Some(rest) = entry.caps.query::<RestApiCap>() {
945 let ctx = self
946 .ctx_builder
947 .for_gear(entry.name)
948 .await
949 .with_context(|| format!("OoP router: build context for '{}'", entry.name))?;
950 router = rest
951 .register_rest(&ctx, router, ®istry)
952 .with_context(|| format!("OoP router: register_rest for '{}'", entry.name))?;
953
954 if let Some(hc) = rest.healthcheck(&ctx) {
958 hc_registry.register(entry.name, hc);
959 }
960 }
961 }
962
963 let info = OpenApiInfo {
964 title: options.gear_name.clone(),
965 version: options
966 .version
967 .clone()
968 .unwrap_or_else(|| "0.0.0".to_owned()),
969 description: None,
970 servers: vec![],
971 };
972 let openapi = registry
973 .build_openapi(&info)
974 .context("OoP router: build OpenAPI document")?;
975 let json = serde_json::to_string(&openapi).context("OoP router: serialize OpenAPI")?;
976
977 Ok((router, json))
978 }
979
980 pub async fn run_oop_serving(
988 self,
989 options: crate::runtime::OopServeOptions,
990 ) -> anyhow::Result<()> {
991 use crate::runtime::{ReadinessState, ResolvedRestEndpoints};
992
993 tracing::info!("Running OoP serving lifecycle");
994
995 let hc_registry = Arc::new(
999 crate::healthcheck::RestHealthcheckRegistry::with_cancellation(self.cancel.clone()),
1000 );
1001
1002 let deps = self.external_deps();
1005 let readiness = ReadinessState::with_check_timeout(
1006 deps.clone(),
1007 Arc::clone(&hc_registry),
1008 options.healthcheck_timeout,
1009 );
1010
1011 let resolved = Arc::new(ResolvedRestEndpoints::new());
1013 self.client_hub
1014 .register::<ResolvedRestEndpoints>(Arc::clone(&resolved));
1015
1016 let mut server = super::oop_serve::OopHttpServer::start(
1021 Arc::clone(&readiness),
1022 Arc::clone(&resolved),
1023 options,
1024 self.cancel.clone(),
1025 )
1026 .await?;
1027
1028 let mut started = false;
1032 let composed: anyhow::Result<(Router, String)> = async {
1033 self.run_pre_init_phase()?;
1034 #[cfg(feature = "db")]
1035 self.run_db_phase().await?;
1036 self.run_init_phase().await?;
1037 self.run_post_init_phase().await?;
1038 self.run_grpc_phase().await?;
1039 self.run_start_phase().await?;
1040 started = true;
1041 self.compose_oop_router(server.options(), &hc_registry)
1042 .await
1043 }
1044 .await;
1045
1046 let serve_result = match composed {
1047 Ok((gear_router, openapi_json)) => {
1048 server.attach(gear_router, openapi_json, deps);
1050 server.join().await
1052 }
1053 Err(e) => {
1054 tracing::error!(error = %e, "OoP startup failed before serving gear routes");
1055 self.cancel.cancel();
1057 if let Err(join_err) = server.join().await {
1058 tracing::warn!(error = %join_err, "OoP probe server teardown after startup failure errored");
1059 }
1060 Err(e)
1061 }
1062 };
1063
1064 if started && let Err(e) = self.run_stop_phase_guarded().await {
1067 tracing::warn!(error = %e, "OoP stop phase reported an error");
1068 }
1069
1070 serve_result
1071 }
1072}
1073
1074#[cfg(test)]
1075#[cfg(feature = "bootstrap")]
1076#[cfg_attr(coverage_nightly, coverage(off))]
1077#[path = "host_runtime_oop_tests.rs"]
1078mod host_runtime_oop_tests;
1079
1080#[cfg(test)]
1081#[cfg_attr(coverage_nightly, coverage(off))]
1082mod tests {
1083 use super::*;
1084 use crate::context::GearCtx;
1085 use crate::contracts::{Gear, RunnableCapability, SystemCapability};
1086 use crate::registry::RegistryBuilder;
1087 use std::sync::Arc;
1088 use std::sync::atomic::{AtomicUsize, Ordering};
1089 use tokio::sync::Mutex;
1090
1091 #[derive(Default)]
1092 #[allow(dead_code)]
1093 struct DummyCore;
1094 #[async_trait::async_trait]
1095 impl Gear for DummyCore {
1096 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1097 Ok(())
1098 }
1099 }
1100
1101 struct StopOrderTracker {
1102 my_order: usize,
1103 stop_order: Arc<AtomicUsize>,
1104 }
1105
1106 impl StopOrderTracker {
1107 fn new(counter: &Arc<AtomicUsize>, stop_order: Arc<AtomicUsize>) -> Self {
1108 let my_order = counter.fetch_add(1, Ordering::SeqCst);
1109 Self {
1110 my_order,
1111 stop_order,
1112 }
1113 }
1114 }
1115
1116 #[async_trait::async_trait]
1117 impl Gear for StopOrderTracker {
1118 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1119 Ok(())
1120 }
1121 }
1122
1123 #[async_trait::async_trait]
1124 impl RunnableCapability for StopOrderTracker {
1125 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1126 Ok(())
1127 }
1128 async fn stop(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1129 let order = self.stop_order.fetch_add(1, Ordering::SeqCst);
1130 tracing::info!(my_order = self.my_order, stop_order = order, "Gear stopped");
1131 Ok(())
1132 }
1133 }
1134
1135 #[tokio::test]
1136 async fn test_stop_phase_reverse_order() {
1137 let counter = Arc::new(AtomicUsize::new(0));
1138 let stop_order = Arc::new(AtomicUsize::new(0));
1139
1140 let gear_a = Arc::new(StopOrderTracker::new(&counter, stop_order.clone()));
1141 let gear_b = Arc::new(StopOrderTracker::new(&counter, stop_order.clone()));
1142 let gear_c = Arc::new(StopOrderTracker::new(&counter, stop_order.clone()));
1143
1144 let mut builder = RegistryBuilder::default();
1145 builder.register_core_with_meta("a", &[], gear_a.clone() as Arc<dyn Gear>);
1146 builder.register_core_with_meta("b", &["a"], gear_b.clone() as Arc<dyn Gear>);
1147 builder.register_core_with_meta("c", &["b"], gear_c.clone() as Arc<dyn Gear>);
1148
1149 builder.register_stateful_with_meta("a", gear_a.clone() as Arc<dyn RunnableCapability>);
1150 builder.register_stateful_with_meta("b", gear_b.clone() as Arc<dyn RunnableCapability>);
1151 builder.register_stateful_with_meta("c", gear_c.clone() as Arc<dyn RunnableCapability>);
1152
1153 let registry = builder.build_topo_sorted().unwrap();
1154
1155 let gear_names: Vec<_> = registry.gears().iter().map(|m| m.name).collect();
1157 assert_eq!(gear_names, vec!["a", "b", "c"]);
1158
1159 let client_hub = Arc::new(ClientHub::new());
1160 let cancel = CancellationToken::new();
1161 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1162
1163 let runtime = HostRuntime::new(
1164 registry,
1165 config_provider,
1166 DbOptions::None,
1167 client_hub,
1168 cancel.clone(),
1169 Uuid::new_v4(),
1170 None,
1171 );
1172
1173 runtime.run_stop_phase().await.unwrap();
1175
1176 assert_eq!(stop_order.load(Ordering::SeqCst), 3);
1180 }
1181
1182 #[tokio::test]
1183 async fn test_stop_phase_continues_on_error() {
1184 struct FailingGear {
1185 should_fail: bool,
1186 stopped: Arc<AtomicUsize>,
1187 }
1188
1189 #[async_trait::async_trait]
1190 impl Gear for FailingGear {
1191 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1192 Ok(())
1193 }
1194 }
1195
1196 #[async_trait::async_trait]
1197 impl RunnableCapability for FailingGear {
1198 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1199 Ok(())
1200 }
1201 async fn stop(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1202 self.stopped.fetch_add(1, Ordering::SeqCst);
1203 if self.should_fail {
1204 anyhow::bail!("Intentional failure")
1205 }
1206 Ok(())
1207 }
1208 }
1209
1210 let stopped = Arc::new(AtomicUsize::new(0));
1211 let gear_a = Arc::new(FailingGear {
1212 should_fail: false,
1213 stopped: stopped.clone(),
1214 });
1215 let gear_b = Arc::new(FailingGear {
1216 should_fail: true,
1217 stopped: stopped.clone(),
1218 });
1219 let gear_c = Arc::new(FailingGear {
1220 should_fail: false,
1221 stopped: stopped.clone(),
1222 });
1223
1224 let mut builder = RegistryBuilder::default();
1225 builder.register_core_with_meta("a", &[], gear_a.clone() as Arc<dyn Gear>);
1226 builder.register_core_with_meta("b", &["a"], gear_b.clone() as Arc<dyn Gear>);
1227 builder.register_core_with_meta("c", &["b"], gear_c.clone() as Arc<dyn Gear>);
1228
1229 builder.register_stateful_with_meta("a", gear_a.clone() as Arc<dyn RunnableCapability>);
1230 builder.register_stateful_with_meta("b", gear_b.clone() as Arc<dyn RunnableCapability>);
1231 builder.register_stateful_with_meta("c", gear_c.clone() as Arc<dyn RunnableCapability>);
1232
1233 let registry = builder.build_topo_sorted().unwrap();
1234
1235 let client_hub = Arc::new(ClientHub::new());
1236 let cancel = CancellationToken::new();
1237 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1238
1239 let runtime = HostRuntime::new(
1240 registry,
1241 config_provider,
1242 DbOptions::None,
1243 client_hub,
1244 cancel.clone(),
1245 Uuid::new_v4(),
1246 None,
1247 );
1248
1249 runtime.run_stop_phase().await.unwrap();
1251
1252 assert_eq!(stopped.load(Ordering::SeqCst), 3);
1254 }
1255
1256 struct EmptyConfigProvider;
1257 impl ConfigProvider for EmptyConfigProvider {
1258 fn get_gear_config(&self, _gear_name: &str) -> Option<&serde_json::Value> {
1259 None
1260 }
1261 }
1262
1263 #[tokio::test]
1264 async fn test_post_init_runs_after_all_init_and_system_first() {
1265 #[derive(Clone)]
1266 struct TrackHooks {
1267 name: &'static str,
1268 events: Arc<Mutex<Vec<String>>>,
1269 }
1270
1271 #[async_trait::async_trait]
1272 impl Gear for TrackHooks {
1273 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1274 self.events.lock().await.push(format!("init:{}", self.name));
1275 Ok(())
1276 }
1277 }
1278
1279 #[async_trait::async_trait]
1280 impl SystemCapability for TrackHooks {
1281 fn pre_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
1282 Ok(())
1283 }
1284
1285 async fn post_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
1286 self.events
1287 .lock()
1288 .await
1289 .push(format!("post_init:{}", self.name));
1290 Ok(())
1291 }
1292 }
1293
1294 let events = Arc::new(Mutex::new(Vec::<String>::new()));
1295 let sys_a = Arc::new(TrackHooks {
1296 name: "sys_a",
1297 events: events.clone(),
1298 });
1299 let user_b = Arc::new(TrackHooks {
1300 name: "user_b",
1301 events: events.clone(),
1302 });
1303 let user_c = Arc::new(TrackHooks {
1304 name: "user_c",
1305 events: events.clone(),
1306 });
1307
1308 let mut builder = RegistryBuilder::default();
1309 builder.register_core_with_meta("sys_a", &[], sys_a.clone() as Arc<dyn Gear>);
1310 builder.register_core_with_meta("user_b", &["sys_a"], user_b.clone() as Arc<dyn Gear>);
1311 builder.register_core_with_meta("user_c", &["user_b"], user_c.clone() as Arc<dyn Gear>);
1312 builder.register_system_with_meta("sys_a", sys_a.clone() as Arc<dyn SystemCapability>);
1313
1314 let registry = builder.build_topo_sorted().unwrap();
1315
1316 let client_hub = Arc::new(ClientHub::new());
1317 let cancel = CancellationToken::new();
1318 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1319
1320 let runtime = HostRuntime::new(
1321 registry,
1322 config_provider,
1323 DbOptions::None,
1324 client_hub,
1325 cancel,
1326 Uuid::new_v4(),
1327 None,
1328 );
1329
1330 runtime.run_init_phase().await.unwrap();
1332 runtime.run_post_init_phase().await.unwrap();
1333
1334 let events = events.lock().await.clone();
1335 let first_post_init = events
1336 .iter()
1337 .position(|e| e.starts_with("post_init:"))
1338 .expect("expected post_init events");
1339 assert!(
1340 events[..first_post_init]
1341 .iter()
1342 .all(|e| e.starts_with("init:")),
1343 "expected all init events before post_init, got: {events:?}"
1344 );
1345
1346 assert_eq!(
1348 events,
1349 vec![
1350 "init:sys_a",
1351 "init:user_b",
1352 "init:user_c",
1353 "post_init:sys_a",
1354 ]
1355 );
1356 }
1357
1358 #[tokio::test]
1359 async fn test_stop_phase_provides_fresh_deadline_token() {
1360 use std::sync::atomic::AtomicBool;
1361
1362 struct TokenCheckGear {
1363 stop_was_called: AtomicBool,
1364 token_was_cancelled_on_entry: AtomicBool,
1365 }
1366
1367 #[async_trait::async_trait]
1368 impl Gear for TokenCheckGear {
1369 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1370 Ok(())
1371 }
1372 }
1373
1374 #[async_trait::async_trait]
1375 impl RunnableCapability for TokenCheckGear {
1376 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1377 Ok(())
1378 }
1379 async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
1380 self.stop_was_called.store(true, Ordering::SeqCst);
1382 self.token_was_cancelled_on_entry
1384 .store(deadline_token.is_cancelled(), Ordering::SeqCst);
1385 Ok(())
1386 }
1387 }
1388
1389 let gear = Arc::new(TokenCheckGear {
1390 stop_was_called: AtomicBool::new(false),
1391 token_was_cancelled_on_entry: AtomicBool::new(true),
1393 });
1394
1395 let mut builder = RegistryBuilder::default();
1396 builder.register_core_with_meta("test", &[], gear.clone() as Arc<dyn Gear>);
1397 builder.register_stateful_with_meta("test", gear.clone() as Arc<dyn RunnableCapability>);
1398
1399 let registry = builder.build_topo_sorted().unwrap();
1400 let client_hub = Arc::new(ClientHub::new());
1401 let cancel = CancellationToken::new();
1402 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1403
1404 let runtime = HostRuntime::new(
1405 registry,
1406 config_provider,
1407 DbOptions::None,
1408 client_hub,
1409 cancel.clone(),
1410 Uuid::new_v4(),
1411 None,
1412 );
1413
1414 runtime.run_stop_phase().await.unwrap();
1416
1417 assert!(
1419 gear.stop_was_called.load(Ordering::SeqCst),
1420 "stop() was never called - gear may not have been registered correctly"
1421 );
1422
1423 assert!(
1426 !gear.token_was_cancelled_on_entry.load(Ordering::SeqCst),
1427 "deadline_token should NOT be cancelled when stop() is called - this enables graceful shutdown"
1428 );
1429 }
1430
1431 #[tokio::test]
1432 async fn test_stop_phase_graceful_shutdown_completes_before_deadline() {
1433 use std::sync::atomic::AtomicBool;
1434 use std::time::Duration;
1435
1436 struct GracefulGear {
1437 graceful_completed: AtomicBool,
1438 deadline_fired: AtomicBool,
1439 }
1440
1441 #[async_trait::async_trait]
1442 impl Gear for GracefulGear {
1443 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1444 Ok(())
1445 }
1446 }
1447
1448 #[async_trait::async_trait]
1449 impl RunnableCapability for GracefulGear {
1450 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1451 Ok(())
1452 }
1453 async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
1454 tokio::select! {
1456 () = tokio::time::sleep(Duration::from_millis(10)) => {
1457 self.graceful_completed.store(true, Ordering::SeqCst);
1458 }
1459 () = deadline_token.cancelled() => {
1460 self.deadline_fired.store(true, Ordering::SeqCst);
1461 }
1462 }
1463 Ok(())
1464 }
1465 }
1466
1467 let gear = Arc::new(GracefulGear {
1468 graceful_completed: AtomicBool::new(false),
1469 deadline_fired: AtomicBool::new(false),
1470 });
1471
1472 let mut builder = RegistryBuilder::default();
1473 builder.register_core_with_meta("test", &[], gear.clone() as Arc<dyn Gear>);
1474 builder.register_stateful_with_meta("test", gear.clone() as Arc<dyn RunnableCapability>);
1475
1476 let registry = builder.build_topo_sorted().unwrap();
1477 let client_hub = Arc::new(ClientHub::new());
1478 let cancel = CancellationToken::new();
1479 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1480
1481 let runtime = HostRuntime::new(
1483 registry,
1484 config_provider,
1485 DbOptions::None,
1486 client_hub,
1487 cancel.clone(),
1488 Uuid::new_v4(),
1489 None,
1490 )
1491 .with_shutdown_deadline(Duration::from_secs(5));
1492
1493 runtime.run_stop_phase().await.unwrap();
1494
1495 assert!(
1497 gear.graceful_completed.load(Ordering::SeqCst),
1498 "graceful shutdown should complete"
1499 );
1500 assert!(
1502 !gear.deadline_fired.load(Ordering::SeqCst),
1503 "deadline should not fire when graceful shutdown completes quickly"
1504 );
1505 }
1506
1507 #[tokio::test]
1508 async fn test_stop_phase_deadline_fires_for_slow_gear() {
1509 use std::sync::atomic::AtomicBool;
1510 use std::time::Duration;
1511
1512 struct SlowGear {
1513 graceful_completed: AtomicBool,
1514 deadline_fired: AtomicBool,
1515 }
1516
1517 #[async_trait::async_trait]
1518 impl Gear for SlowGear {
1519 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1520 Ok(())
1521 }
1522 }
1523
1524 #[async_trait::async_trait]
1525 impl RunnableCapability for SlowGear {
1526 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1527 Ok(())
1528 }
1529 async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
1530 tokio::select! {
1532 () = tokio::time::sleep(Duration::from_secs(10)) => {
1533 self.graceful_completed.store(true, Ordering::SeqCst);
1534 }
1535 () = deadline_token.cancelled() => {
1536 self.deadline_fired.store(true, Ordering::SeqCst);
1537 }
1538 }
1539 Ok(())
1540 }
1541 }
1542
1543 let gear = Arc::new(SlowGear {
1544 graceful_completed: AtomicBool::new(false),
1545 deadline_fired: AtomicBool::new(false),
1546 });
1547
1548 let mut builder = RegistryBuilder::default();
1549 builder.register_core_with_meta("test", &[], gear.clone() as Arc<dyn Gear>);
1550 builder.register_stateful_with_meta("test", gear.clone() as Arc<dyn RunnableCapability>);
1551
1552 let registry = builder.build_topo_sorted().unwrap();
1553 let client_hub = Arc::new(ClientHub::new());
1554 let cancel = CancellationToken::new();
1555 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1556
1557 let runtime = HostRuntime::new(
1559 registry,
1560 config_provider,
1561 DbOptions::None,
1562 client_hub,
1563 cancel.clone(),
1564 Uuid::new_v4(),
1565 None,
1566 )
1567 .with_shutdown_deadline(Duration::from_millis(100));
1568
1569 runtime.run_stop_phase().await.unwrap();
1570
1571 assert!(
1573 !gear.graceful_completed.load(Ordering::SeqCst),
1574 "graceful shutdown should not complete when deadline fires first"
1575 );
1576 assert!(
1578 gear.deadline_fired.load(Ordering::SeqCst),
1579 "deadline should fire for slow gears"
1580 );
1581 }
1582}