1use axum::Router;
25use std::collections::HashSet;
26use std::sync::Arc;
27
28use tokio_util::sync::CancellationToken;
29use uuid::Uuid;
30
31use crate::backends::OopSpawnConfig;
32use crate::client_hub::ClientHub;
33use crate::config::ConfigProvider;
34use crate::context::GearContextBuilder;
35use crate::registry::{
36 ApiGatewayCap, GearEntry, GearRegistry, GrpcHubCap, RegistryError, RestApiCap, RunnableCap,
37 SystemCap,
38};
39use crate::runtime::{GearManager, GrpcInstallerStore, OopSpawnOptions, SystemContext};
40
41#[cfg(feature = "db")]
42use crate::registry::DatabaseCap;
43
44#[derive(Clone)]
46pub enum DbOptions {
47 None,
49 #[cfg(feature = "db")]
51 Manager(Arc<toolkit_db::DbManager>),
52}
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
56pub enum RunMode {
57 Full,
59 MigrateOnly,
61}
62
63pub const TOOLKIT_DIRECTORY_ENDPOINT_ENV: &str = "TOOLKIT_DIRECTORY_ENDPOINT";
65
66pub const TOOLKIT_MODULE_CONFIG_ENV: &str = "TOOLKIT_MODULE_CONFIG";
68
69pub const DEFAULT_SHUTDOWN_DEADLINE: std::time::Duration = std::time::Duration::from_secs(35);
75
76fn static_endpoint_override(
84 cfg: &dyn ConfigProvider,
85 owner_gear: &str,
86 dep_gear: &str,
87) -> Option<String> {
88 cfg.get_gear_config(owner_gear)?
89 .get("config")?
90 .get("consumer_wiring")?
91 .get(dep_gear)?
92 .as_str()
93 .map(str::to_owned)
94}
95
96pub struct HostRuntime {
97 registry: GearRegistry,
98 ctx_builder: GearContextBuilder,
99 instance_id: Uuid,
100 gear_manager: Arc<GearManager>,
101 grpc_installers: Arc<GrpcInstallerStore>,
102 client_hub: Arc<ClientHub>,
103 gears_cfg: Arc<dyn ConfigProvider>,
106 dep_checker: Arc<super::readiness::DependencyChecker>,
112 rest_providers_registered: std::sync::atomic::AtomicBool,
118 cancel: CancellationToken,
119 #[allow(dead_code)]
120 db_options: DbOptions,
121 oop_options: Option<OopSpawnOptions>,
123 shutdown_deadline: std::time::Duration,
125}
126
127impl HostRuntime {
128 pub fn new(
132 registry: GearRegistry,
133 gears_cfg: Arc<dyn ConfigProvider>,
134 db_options: DbOptions,
135 client_hub: Arc<ClientHub>,
136 cancel: CancellationToken,
137 instance_id: Uuid,
138 oop_options: Option<OopSpawnOptions>,
139 ) -> Self {
140 let gear_manager = Arc::new(GearManager::new());
142 let grpc_installers = Arc::new(GrpcInstallerStore::new());
143
144 let dep_checker = Arc::new(super::readiness::DependencyChecker::new());
148 client_hub.register::<super::readiness::DependencyChecker>(dep_checker.clone());
149
150 let ctx_builder = GearContextBuilder::new(
152 instance_id,
153 gears_cfg.clone(),
154 client_hub.clone(),
155 cancel.clone(),
156 );
157 #[cfg(feature = "db")]
158 let ctx_builder = match &db_options {
159 DbOptions::Manager(mgr) => ctx_builder.with_db_manager(mgr.clone()),
160 DbOptions::None => ctx_builder,
161 };
162
163 Self {
164 registry,
165 ctx_builder,
166 instance_id,
167 gear_manager,
168 grpc_installers,
169 client_hub,
170 gears_cfg,
171 dep_checker,
172 rest_providers_registered: std::sync::atomic::AtomicBool::new(false),
173 cancel,
174 db_options,
175 oop_options,
176 shutdown_deadline: DEFAULT_SHUTDOWN_DEADLINE,
177 }
178 }
179
180 #[must_use]
197 pub fn with_shutdown_deadline(mut self, deadline: std::time::Duration) -> Self {
198 self.shutdown_deadline = deadline;
199 self
200 }
201
202 #[must_use]
208 pub fn with_internal_token_provider(
209 mut self,
210 provider: Option<toolkit_contract::runtime::config::InternalTokenProvider>,
211 ) -> Self {
212 self.ctx_builder = self.ctx_builder.with_internal_token_provider(provider);
213 self
214 }
215
216 pub fn run_pre_init_phase(&self) -> Result<(), RegistryError> {
223 tracing::info!("Phase: pre_init");
224
225 let sys_ctx = SystemContext::new(
226 self.instance_id,
227 Arc::clone(&self.gear_manager),
228 Arc::clone(&self.grpc_installers),
229 );
230
231 for entry in self.registry.gears() {
232 if self.cancel.is_cancelled() {
234 tracing::warn!("Pre-init phase cancelled by signal");
235 return Err(RegistryError::Cancelled);
236 }
237
238 if let Some(sys_mod) = entry.caps.query::<SystemCap>() {
239 tracing::debug!(gear = entry.name, "Running system pre_init");
240 sys_mod
241 .pre_init(&sys_ctx)
242 .map_err(|e| RegistryError::PreInit {
243 gear: entry.name,
244 source: e,
245 })?;
246 }
247 }
248
249 Ok(())
250 }
251
252 #[cfg(feature = "db")]
254 async fn gear_context(
255 &self,
256 gear_name: &'static str,
257 ) -> Result<crate::context::GearCtx, RegistryError> {
258 self.ctx_builder
259 .for_gear(gear_name)
260 .await
261 .map_err(|e| RegistryError::DbMigrate {
262 gear: gear_name,
263 source: e,
264 })
265 }
266
267 #[cfg(feature = "db")]
269 async fn db_migration_target(
270 &self,
271 gear_name: &'static str,
272 ctx: &crate::context::GearCtx,
273 db_gear: Option<Arc<dyn crate::contracts::DatabaseCapability>>,
274 ) -> Result<
275 Option<(
276 toolkit_db::Db,
277 Arc<dyn crate::contracts::DatabaseCapability>,
278 )>,
279 RegistryError,
280 > {
281 let Some(dbm) = db_gear else {
282 return Ok(None);
283 };
284
285 let db = match &self.db_options {
289 DbOptions::None => None,
290 #[cfg(feature = "db")]
291 DbOptions::Manager(mgr) => {
292 mgr.get(gear_name)
293 .await
294 .map_err(|e| RegistryError::DbMigrate {
295 gear: gear_name,
296 source: e.into(),
297 })?
298 }
299 };
300
301 _ = ctx; Ok(db.map(|db| (db, dbm)))
303 }
304
305 #[cfg(feature = "db")]
310 async fn migrate_gear(
311 gear_name: &'static str,
312 db: &toolkit_db::Db,
313 db_gear: Arc<dyn crate::contracts::DatabaseCapability>,
314 ) -> Result<(), RegistryError> {
315 let migrations = db_gear.migrations();
317
318 if migrations.is_empty() {
319 tracing::debug!(gear = gear_name, "No migrations to run");
320 return Ok(());
321 }
322
323 tracing::debug!(
324 gear = gear_name,
325 count = migrations.len(),
326 "Running DB migrations"
327 );
328
329 let result =
331 toolkit_db::migration_runner::run_migrations_for_gear(db, gear_name, migrations)
332 .await
333 .map_err(|e| RegistryError::DbMigrate {
334 gear: gear_name,
335 source: anyhow::Error::new(e),
336 })?;
337
338 tracing::info!(
339 gear = gear_name,
340 applied = result.applied,
341 skipped = result.skipped,
342 "DB migrations completed"
343 );
344
345 Ok(())
346 }
347
348 #[cfg(feature = "db")]
357 async fn run_db_phase(&self) -> Result<(), RegistryError> {
358 tracing::info!("Phase: db (before init)");
359
360 for entry in self.registry.gears_by_system_priority() {
361 if self.cancel.is_cancelled() {
363 tracing::warn!("DB migration phase cancelled by signal");
364 return Err(RegistryError::Cancelled);
365 }
366
367 let ctx = self.gear_context(entry.name).await?;
368 let db_gear = entry.caps.query::<DatabaseCap>();
369
370 match self
371 .db_migration_target(entry.name, &ctx, db_gear.clone())
372 .await?
373 {
374 Some((db, dbm)) => {
375 Self::migrate_gear(entry.name, &db, dbm).await?;
376 }
377 None if db_gear.is_some() => {
378 tracing::debug!(
379 gear = entry.name,
380 "Gear has DbGear trait but no DB handle (no config)"
381 );
382 }
383 None => {}
384 }
385 }
386
387 Ok(())
388 }
389
390 async fn run_init_phase(&self) -> Result<(), RegistryError> {
394 tracing::info!("Phase: init");
395
396 for entry in self.registry.gears_by_system_priority() {
397 let ctx =
398 self.ctx_builder
399 .for_gear(entry.name)
400 .await
401 .map_err(|e| RegistryError::Init {
402 gear: entry.name,
403 source: e,
404 })?;
405 tracing::info!(gear = entry.name, "Initializing a gear...");
406 entry
407 .core
408 .init(&ctx)
409 .await
410 .map_err(|e| RegistryError::Init {
411 gear: entry.name,
412 source: e,
413 })?;
414 tracing::info!(gear = entry.name, "Initialized a gear.");
415 }
416
417 Ok(())
418 }
419
420 #[allow(
435 clippy::unused_async,
436 reason = "kept async for symmetry with the other `run_*_phase` steps awaited in sequence by `run_gear_phases`; the awaited work runs in a spawned readiness-probe task"
437 )]
438 async fn run_proxy_wiring_phase(&self) -> Result<(), RegistryError> {
439 use crate::discovery::{
440 ConsumerRegistration, DirectoryEndpointResolver, NullEndpointResolver,
441 };
442 use toolkit_contract::runtime::resolving::EndpointResolver;
443
444 let regs: Vec<&ConsumerRegistration> = inventory::iter::<ConsumerRegistration>
445 .into_iter()
446 .collect();
447 if regs.is_empty() {
448 return Ok(());
449 }
450 tracing::info!(
451 count = regs.len(),
452 "Phase: proxy-wiring (consumer discovery)"
453 );
454
455 let (resolver, have_directory): (Arc<dyn EndpointResolver>, bool) =
462 if let Ok(dir) = self.client_hub.get::<dyn crate::DirectoryClient>() {
463 (Arc::new(DirectoryEndpointResolver::new(dir)), true)
464 } else {
465 tracing::error!(
466 consumers = regs.len(),
467 "proxy-wiring: no DirectoryClient in ClientHub; remote consumer \
468 dependencies cannot be resolved and will gate /readyz (503). \
469 Co-located (local) dependencies are unaffected."
470 );
471 (Arc::new(NullEndpointResolver), false)
472 };
473
474 let known_gears: std::collections::HashSet<&str> =
486 self.registry.gears().iter().map(GearEntry::name).collect();
487 for reg in ®s {
488 if !known_gears.contains(reg.owner_gear) {
489 tracing::warn!(
490 owner = reg.owner_gear,
491 dep = reg.dep_gear,
492 "proxy-wiring: consumer's owner gear name does not match any registered gear; \
493 the `gears.{}.config.consumer_wiring.{}` static override will never resolve. \
494 Rename the gear to the kebab-case of its struct ident.",
495 reg.owner_gear,
496 reg.dep_gear,
497 );
498 }
499 }
500
501 let mut remote_deps: Vec<String> = Vec::new();
502 for reg in ®s {
503 let static_override =
509 static_endpoint_override(self.gears_cfg.as_ref(), reg.owner_gear, reg.dep_gear);
510 let (reg_resolver, is_static): (Arc<dyn EndpointResolver>, bool) =
511 if let Some(endpoint) = &static_override {
512 tracing::warn!(
513 owner = reg.owner_gear,
514 dep = reg.dep_gear,
515 endpoint = %endpoint,
516 "proxy-wiring: STATIC endpoint override in use (ADR-0004 dev/test \
517 escape hatch) - bypasses service discovery; MUST NOT be used in \
518 production"
519 );
520 (
521 Arc::new(crate::discovery::StaticEndpointResolver::new(
522 endpoint.clone(),
523 )),
524 true,
525 )
526 } else {
527 (Arc::clone(&resolver), false)
528 };
529
530 let outcome = (reg.wire)(
536 &self.client_hub,
537 reg_resolver,
538 self.ctx_builder.internal_token_provider(),
539 )
540 .map_err(|source| RegistryError::ProxyWiring {
541 gear: reg.owner_gear,
542 source,
543 })?;
544 self.dep_checker.register_dep(reg.dep_gear.to_owned());
545 match outcome {
546 crate::discovery::WireOutcome::Local => {
548 self.dep_checker.mark_resolved(reg.dep_gear);
549 }
550 crate::discovery::WireOutcome::Remote if is_static => {
552 self.dep_checker.mark_resolved(reg.dep_gear);
553 }
554 crate::discovery::WireOutcome::Remote => remote_deps.push(reg.dep_gear.to_owned()),
556 }
557 tracing::debug!(
558 owner = reg.owner_gear,
559 dep = reg.dep_gear,
560 outcome = ?outcome,
561 static_override = is_static,
562 "wired consumer contract"
563 );
564 }
565
566 if !have_directory || remote_deps.is_empty() {
570 return Ok(());
571 }
572
573 let readiness = Arc::clone(&self.dep_checker);
574 let cancel = self.cancel.clone();
575 tokio::spawn(async move {
576 const BASE: std::time::Duration = std::time::Duration::from_millis(100);
577 const MAX: std::time::Duration = std::time::Duration::from_secs(30);
578 let mut pending = remote_deps;
579 let mut backoff = BASE;
580 while !pending.is_empty() {
581 let mut still_pending = Vec::new();
582 for dep in pending {
583 match resolver.resolve_endpoint(&dep).await {
584 Ok(Some(_)) => {
585 readiness.mark_resolved(&dep);
586 tracing::info!(dep = %dep, "readiness: dependency resolved");
587 }
588 Ok(None) => still_pending.push(dep),
590 Err(e) => {
593 tracing::warn!(dep = %dep, error = %e, "readiness: directory lookup failed");
594 still_pending.push(dep);
595 }
596 }
597 }
598 pending = still_pending;
599 if pending.is_empty() {
600 break;
601 }
602 tokio::select! {
603 () = cancel.cancelled() => break,
604 () = tokio::time::sleep(backoff) => {}
605 }
606 backoff = (backoff * 2).min(MAX);
607 }
608 });
609
610 Ok(())
611 }
612
613 async fn run_init_wiring_post_init(&self) -> Result<(), RegistryError> {
632 self.run_init_phase().await?;
633
634 self.run_proxy_wiring_phase().await?;
635
636 self.run_post_init_phase().await
637 }
638
639 async fn run_post_init_phase(&self) -> Result<(), RegistryError> {
642 tracing::info!("Phase: post_init");
643
644 let sys_ctx = SystemContext::new(
645 self.instance_id,
646 Arc::clone(&self.gear_manager),
647 Arc::clone(&self.grpc_installers),
648 );
649
650 for entry in self.registry.gears_by_system_priority() {
651 if let Some(sys_mod) = entry.caps.query::<SystemCap>() {
652 sys_mod
653 .post_init(&sys_ctx)
654 .await
655 .map_err(|e| RegistryError::PostInit {
656 gear: entry.name,
657 source: e,
658 })?;
659 }
660 }
661
662 Ok(())
663 }
664
665 async fn run_rest_phase(&self) -> Result<Router, RegistryError> {
672 tracing::info!("Phase: rest (sync)");
673
674 let mut router = Router::new();
675
676 let host_count = self
678 .registry
679 .gears()
680 .iter()
681 .filter(|e| e.caps.has::<ApiGatewayCap>())
682 .count();
683
684 match host_count {
685 0 => {
686 return if self
687 .registry
688 .gears()
689 .iter()
690 .any(|e| e.caps.has::<RestApiCap>())
691 {
692 Err(RegistryError::RestRequiresHost)
693 } else {
694 Ok(router)
695 };
696 }
697 1 => { }
698 _ => return Err(RegistryError::MultipleRestHosts),
699 }
700
701 let host_idx = self
703 .registry
704 .gears()
705 .iter()
706 .position(|e| e.caps.has::<ApiGatewayCap>())
707 .ok_or(RegistryError::RestHostNotFoundAfterValidation)?;
708 let host_entry = &self.registry.gears()[host_idx];
709 let Some(host) = host_entry.caps.query::<ApiGatewayCap>() else {
710 return Err(RegistryError::RestHostMissingFromEntry);
711 };
712 let host_ctx = self
713 .ctx_builder
714 .for_gear(host_entry.name)
715 .await
716 .map_err(|e| RegistryError::RestPrepare {
717 gear: host_entry.name,
718 source: e,
719 })?;
720
721 let registry: &dyn crate::contracts::OpenApiRegistry = host.as_registry();
723
724 let hc_registry = Arc::new(
728 crate::healthcheck::RestHealthcheckRegistry::with_cancellation(
729 host_ctx.cancellation_token().clone(),
730 ),
731 );
732
733 hc_registry.register(
739 "readiness",
740 Arc::new(super::readiness::ReadinessHealthcheck::new(
741 self.dep_checker.clone(),
742 )),
743 );
744
745 router = host
747 .rest_prepare(&host_ctx, router, hc_registry.clone())
748 .map_err(|source| RegistryError::RestPrepare {
749 gear: host_entry.name,
750 source,
751 })?;
752
753 for e in self.registry.gears() {
755 if let Some(rest) = e.caps.query::<RestApiCap>() {
756 let ctx = self.ctx_builder.for_gear(e.name).await.map_err(|err| {
757 RegistryError::RestRegister {
758 gear: e.name,
759 source: err,
760 }
761 })?;
762
763 router = rest
764 .register_rest(&ctx, router, registry)
765 .map_err(|source| RegistryError::RestRegister {
766 gear: e.name,
767 source,
768 })?;
769
770 if let Some(hc) = rest.healthcheck(&ctx) {
772 hc_registry.register(e.name, hc);
773 }
774 }
775 }
776
777 router = host
779 .rest_finalize(&host_ctx, router, hc_registry)
780 .map_err(|source| RegistryError::RestFinalize {
781 gear: host_entry.name,
782 source,
783 })?;
784
785 Ok(router)
786 }
787
788 async fn run_grpc_phase(&self) -> Result<(), RegistryError> {
792 tracing::info!("Phase: grpc (registration)");
793
794 if self.registry.grpc_hub.is_none() && self.registry.grpc_services.is_empty() {
796 return Ok(());
797 }
798
799 if self.registry.grpc_hub.is_none() && !self.registry.grpc_services.is_empty() {
801 return Err(RegistryError::GrpcRequiresHub);
802 }
803
804 if let Some(hub_name) = &self.registry.grpc_hub {
806 let mut gears_data = Vec::new();
807 let mut seen = HashSet::new();
808
809 for (gear_name, service_gear) in &self.registry.grpc_services {
811 let ctx = self.ctx_builder.for_gear(gear_name).await.map_err(|err| {
812 RegistryError::GrpcRegister {
813 gear: gear_name.clone(),
814 source: err,
815 }
816 })?;
817
818 let installers = service_gear
819 .get_grpc_services(&ctx)
820 .await
821 .map_err(|source| RegistryError::GrpcRegister {
822 gear: gear_name.clone(),
823 source,
824 })?;
825
826 for reg in &installers {
827 if !seen.insert(reg.service_name) {
828 return Err(RegistryError::GrpcRegister {
829 gear: gear_name.clone(),
830 source: anyhow::anyhow!(
831 "Duplicate gRPC service name: {}",
832 reg.service_name
833 ),
834 });
835 }
836 }
837
838 gears_data.push(crate::runtime::GearInstallers {
839 gear_name: gear_name.clone(),
840 installers,
841 });
842 }
843
844 self.grpc_installers
845 .set(crate::runtime::GrpcInstallerData { gears: gears_data })
846 .map_err(|source| RegistryError::GrpcRegister {
847 gear: hub_name.clone(),
848 source,
849 })?;
850 }
851
852 Ok(())
853 }
854
855 async fn run_start_phase(&self) -> Result<(), RegistryError> {
859 tracing::info!("Phase: start");
860
861 for e in self.registry.gears_by_system_priority() {
862 if let Some(s) = e.caps.query::<RunnableCap>() {
863 tracing::debug!(
864 gear = e.name,
865 is_system = e.caps.has::<SystemCap>(),
866 "Starting stateful gear"
867 );
868 s.start(self.cancel.clone())
869 .await
870 .map_err(|source| RegistryError::Start {
871 gear: e.name,
872 source,
873 })?;
874 tracing::info!(gear = e.name, "Started gear");
875 }
876 }
877
878 Ok(())
879 }
880
881 async fn stop_one_gear(entry: &GearEntry, cancel: CancellationToken) {
883 if let Some(s) = entry.caps.query::<RunnableCap>() {
884 match s.stop(cancel).await {
885 Err(err) => {
886 tracing::warn!(gear = entry.name, error = %err, "Failed to stop gear");
887 }
888 _ => {
889 tracing::info!(gear = entry.name, "Stopped gear");
890 }
891 }
892 }
893 }
894
895 async fn run_stop_phase(&self) -> Result<(), RegistryError> {
920 tracing::info!("Phase: stop");
921
922 self.deregister_rest_providers().await;
925
926 let deadline = self.shutdown_deadline;
927
928 for e in self.registry.gears().iter().rev() {
930 let gear_name = e.name;
931
932 let deadline_token = CancellationToken::new();
935 let deadline_token_for_timeout = deadline_token.clone();
936
937 let deadline_task = tokio::spawn(async move {
939 tokio::time::sleep(deadline).await;
940 tracing::warn!(
941 gear = gear_name,
942 deadline_secs = deadline.as_secs(),
943 "Gear shutdown deadline reached, sending hard-stop signal"
944 );
945 deadline_token_for_timeout.cancel();
946 });
947
948 Self::stop_one_gear(e, deadline_token).await;
951
952 deadline_task.abort();
954 #[allow(clippy::let_underscore_must_use)]
955 let _ = deadline_task.await;
956 }
957
958 Ok(())
959 }
960
961 async fn run_stop_phase_guarded(&self) -> Result<(), RegistryError> {
966 let gear_count = u32::try_from(self.registry.gears().len().max(1)).unwrap_or(1);
967 let stop_timeout = self
968 .shutdown_deadline
969 .checked_mul(gear_count)
970 .and_then(|d| d.checked_add(std::time::Duration::from_secs(5)))
971 .unwrap_or(self.shutdown_deadline);
972
973 let (disarm_tx, disarm_rx) = std::sync::mpsc::channel::<()>();
977 std::thread::spawn(move || {
978 match disarm_rx.recv_timeout(stop_timeout) {
979 Ok(()) | Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
980 }
982 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
983 tracing::warn!(
984 timeout_secs = stop_timeout.as_secs(),
985 "shutdown: stop phase timed out, force exiting"
986 );
987 std::process::exit(1);
988 }
989 }
990 });
991
992 let stop_result = self.run_stop_phase().await;
993 let _ = disarm_tx.send(()).ok();
997
998 stop_result
999 }
1000
1001 async fn run_oop_spawn_phase(&self) -> Result<(), RegistryError> {
1006 let oop_opts = match &self.oop_options {
1007 Some(opts) if !opts.gears.is_empty() => opts,
1008 _ => return Ok(()),
1009 };
1010
1011 tracing::info!("Phase: oop_spawn");
1012
1013 let directory_endpoint = self.wait_for_grpc_hub_endpoint().await;
1015
1016 for gear_cfg in &oop_opts.gears {
1017 let mut env = gear_cfg.env.clone();
1020 env.insert(
1021 TOOLKIT_MODULE_CONFIG_ENV.to_owned(),
1022 gear_cfg.rendered_config_json.clone(),
1023 );
1024 if let Some(ref endpoint) = directory_endpoint {
1025 env.insert(TOOLKIT_DIRECTORY_ENDPOINT_ENV.to_owned(), endpoint.clone());
1026 }
1027
1028 let args = gear_cfg.args.clone();
1030
1031 let spawn_config = OopSpawnConfig {
1032 gear_name: gear_cfg.gear_name.clone(),
1033 binary: gear_cfg.binary.clone(),
1034 args,
1035 env,
1036 working_directory: gear_cfg.working_directory.clone(),
1037 };
1038
1039 oop_opts
1040 .backend
1041 .spawn(spawn_config)
1042 .await
1043 .map_err(|e| RegistryError::OopSpawn {
1044 gear: gear_cfg.gear_name.clone(),
1045 source: e,
1046 })?;
1047
1048 tracing::info!(
1049 gear = %gear_cfg.gear_name,
1050 directory_endpoint = ?directory_endpoint,
1051 "Spawned OoP gear via backend"
1052 );
1053 }
1054
1055 Ok(())
1056 }
1057
1058 async fn wait_for_grpc_hub_endpoint(&self) -> Option<String> {
1063 const POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(10);
1064 const MAX_WAIT: std::time::Duration = std::time::Duration::from_secs(5);
1065
1066 let grpc_hub = self
1068 .registry
1069 .gears()
1070 .iter()
1071 .find_map(|e| e.caps.query::<GrpcHubCap>());
1072
1073 let Some(hub) = grpc_hub else {
1074 return None; };
1076
1077 let start = std::time::Instant::now();
1078
1079 loop {
1080 if let Some(endpoint) = hub.bound_endpoint() {
1081 tracing::debug!(
1082 endpoint = %endpoint,
1083 elapsed_ms = start.elapsed().as_millis(),
1084 "gRPC hub endpoint available"
1085 );
1086 return Some(endpoint);
1087 }
1088
1089 if start.elapsed() > MAX_WAIT {
1090 tracing::warn!("Timed out waiting for gRPC hub to bind");
1091 return None;
1092 }
1093
1094 tokio::time::sleep(POLL_INTERVAL).await;
1095 }
1096 }
1097
1098 async fn wait_for_rest_endpoint(
1104 &self,
1105 host: &Arc<dyn crate::contracts::ApiGatewayCapability>,
1106 ) -> Option<String> {
1107 const POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(10);
1108 const MAX_WAIT: std::time::Duration = std::time::Duration::from_secs(5);
1109
1110 let start = std::time::Instant::now();
1111 loop {
1112 if let Some(endpoint) = host.bound_endpoint() {
1113 return Some(endpoint);
1114 }
1115 if start.elapsed() > MAX_WAIT {
1116 tracing::warn!("Timed out waiting for REST host to bind");
1117 return None;
1118 }
1119 tokio::time::sleep(POLL_INTERVAL).await;
1120 }
1121 }
1122
1123 async fn run_directory_register_phase(&self) -> Result<(), RegistryError> {
1136 let rest_gears = self.rest_provider_gears();
1137 if rest_gears.is_empty() {
1138 return Ok(());
1139 }
1140
1141 let Some(host) = self
1142 .registry
1143 .gears()
1144 .iter()
1145 .find_map(|e| e.caps.query::<ApiGatewayCap>())
1146 else {
1147 return Ok(()); };
1149
1150 let Some(endpoint) = self.wait_for_rest_endpoint(&host).await else {
1151 tracing::warn!(
1152 "directory-register: REST host endpoint unavailable; skipping REST provider registration"
1153 );
1154 return Ok(());
1155 };
1156
1157 let Ok(dir) = self.client_hub.get::<dyn crate::DirectoryClient>() else {
1158 tracing::debug!(
1159 "directory-register: no DirectoryClient in ClientHub; skipping REST provider registration"
1160 );
1161 return Ok(());
1162 };
1163
1164 let instance_id = self.instance_id.to_string();
1165 for gear in rest_gears {
1166 let (grpc_services, version) = match dir.list_instances(gear).await {
1180 Ok(insts) => insts
1181 .into_iter()
1182 .find(|i| i.instance_id == instance_id)
1183 .map(|i| (i.grpc_services, i.version))
1184 .unwrap_or_default(),
1185 Err(e) => {
1186 tracing::warn!(
1190 gear,
1191 error = %e,
1192 "directory-register: failed to read existing registration; \
1193 re-registering with empty grpc_services/version"
1194 );
1195 (Vec::new(), None)
1196 }
1197 };
1198 let mut info = crate::RegisterInstanceInfo::new(gear.to_owned(), instance_id.clone())
1202 .with_grpc_services(grpc_services)
1203 .with_rest_endpoint(crate::ServiceEndpoint::new(endpoint.clone()));
1204 if let Some(version) = version {
1205 info = info.with_version(version);
1206 }
1207 match dir.register_instance(info).await {
1208 Ok(()) => {
1209 tracing::info!(gear, endpoint = %endpoint, "registered REST provider in directory");
1210 }
1211 Err(e) => {
1212 tracing::warn!(gear, error = %e, "directory-register: failed to register REST provider");
1213 }
1214 }
1215 }
1216 self.rest_providers_registered
1220 .store(true, std::sync::atomic::Ordering::SeqCst);
1221 Ok(())
1222 }
1223
1224 fn rest_provider_gears(&self) -> Vec<&'static str> {
1229 self.registry
1230 .gears()
1231 .iter()
1232 .filter(|e| e.caps.has::<RestApiCap>() && !e.caps.has::<ApiGatewayCap>())
1233 .map(|e| e.name)
1234 .collect()
1235 }
1236
1237 async fn deregister_rest_providers(&self) {
1240 if !self
1244 .rest_providers_registered
1245 .load(std::sync::atomic::Ordering::SeqCst)
1246 {
1247 return;
1248 }
1249 let rest_gears = self.rest_provider_gears();
1250 if rest_gears.is_empty() {
1251 return;
1252 }
1253 let Ok(dir) = self.client_hub.get::<dyn crate::DirectoryClient>() else {
1254 return;
1255 };
1256 let instance_id = self.instance_id.to_string();
1257 for gear in rest_gears {
1258 if let Err(e) = dir.deregister_instance(gear, &instance_id).await {
1259 tracing::warn!(gear, error = %e, "directory-deregister: failed to deregister REST provider");
1260 }
1261 }
1262 }
1263
1264 pub async fn run_gear_phases(self) -> anyhow::Result<()> {
1273 self.run_phases_internal(RunMode::Full).await
1274 }
1275
1276 pub async fn run_migration_phases(self) -> anyhow::Result<()> {
1286 self.run_phases_internal(RunMode::MigrateOnly).await
1287 }
1288
1289 async fn run_phases_internal(self, mode: RunMode) -> anyhow::Result<()> {
1312 match mode {
1314 RunMode::Full => {
1315 tracing::info!("Running full lifecycle (all phases)");
1316 }
1317 RunMode::MigrateOnly => {
1318 tracing::info!("Running in migration mode (pre-init + db phases only)");
1319 }
1320 }
1321
1322 self.run_pre_init_phase()?;
1324
1325 #[cfg(feature = "db")]
1327 {
1328 self.run_db_phase().await?;
1329 }
1330 #[cfg(not(feature = "db"))]
1331 {
1332 }
1334
1335 if mode == RunMode::MigrateOnly {
1337 tracing::info!("Migration phases completed successfully");
1338 return Ok(());
1339 }
1340
1341 self.run_init_wiring_post_init().await?;
1343
1344 let _router = self.run_rest_phase().await?;
1346
1347 self.run_grpc_phase().await?;
1349
1350 self.run_start_phase().await?;
1352
1353 {
1357 let readiness = Arc::clone(&self.dep_checker);
1358 let cancel = self.cancel.clone();
1359 tokio::spawn(async move {
1360 cancel.cancelled().await;
1361 readiness.set_draining(true);
1362 });
1363 }
1364
1365 self.run_directory_register_phase().await?;
1368
1369 self.run_oop_spawn_phase().await?;
1371
1372 self.cancel.cancelled().await;
1374
1375 self.run_stop_phase_guarded().await?;
1379 Ok(())
1380 }
1381}
1382
1383#[cfg(feature = "bootstrap")]
1385impl HostRuntime {
1386 async fn compose_oop_router(
1392 &self,
1393 options: &crate::runtime::OopServeOptions,
1394 hc_registry: &Arc<crate::healthcheck::RestHealthcheckRegistry>,
1395 ) -> anyhow::Result<(Router, String)> {
1396 use crate::api::{OpenApiInfo, OpenApiRegistryImpl};
1397 use anyhow::Context as _;
1398
1399 let registry = OpenApiRegistryImpl::new();
1400 let mut router = Router::new();
1401
1402 for entry in self.registry.gears() {
1403 if let Some(rest) = entry.caps.query::<RestApiCap>() {
1404 let ctx = self
1405 .ctx_builder
1406 .for_gear(entry.name)
1407 .await
1408 .with_context(|| format!("OoP router: build context for '{}'", entry.name))?;
1409 router = rest
1410 .register_rest(&ctx, router, ®istry)
1411 .with_context(|| format!("OoP router: register_rest for '{}'", entry.name))?;
1412
1413 if let Some(hc) = rest.healthcheck(&ctx) {
1417 hc_registry.register(entry.name, hc);
1418 }
1419 }
1420 }
1421
1422 let info = OpenApiInfo {
1423 title: options.gear_name.clone(),
1424 version: options
1425 .version
1426 .clone()
1427 .unwrap_or_else(|| "0.0.0".to_owned()),
1428 description: None,
1429 servers: vec![],
1430 };
1431 let openapi = registry
1432 .build_openapi(&info)
1433 .context("OoP router: build OpenAPI document")?;
1434 let json = serde_json::to_string(&openapi).context("OoP router: serialize OpenAPI")?;
1435
1436 Ok((router, json))
1437 }
1438
1439 pub async fn run_oop_serving(
1447 self,
1448 options: crate::runtime::OopServeOptions,
1449 ) -> anyhow::Result<()> {
1450 use crate::runtime::ReadinessState;
1451
1452 tracing::info!("Running OoP serving lifecycle");
1453
1454 if self.client_hub.get::<dyn crate::DirectoryClient>().is_err() {
1459 self.client_hub
1460 .register::<dyn crate::DirectoryClient>(Arc::clone(&options.directory));
1461 }
1462
1463 let hc_registry = Arc::new(
1467 crate::healthcheck::RestHealthcheckRegistry::with_cancellation(self.cancel.clone()),
1468 );
1469
1470 let readiness = ReadinessState::from_checker(
1477 Arc::clone(&self.dep_checker),
1478 Arc::clone(&hc_registry),
1479 options.healthcheck_timeout,
1480 );
1481
1482 let mut server = super::oop_serve::OopHttpServer::start(
1487 Arc::clone(&readiness),
1488 options,
1489 self.cancel.clone(),
1490 )
1491 .await?;
1492
1493 let mut started = false;
1499 let composed: anyhow::Result<(Router, String)> = async {
1500 self.run_pre_init_phase()?;
1501 #[cfg(feature = "db")]
1502 self.run_db_phase().await?;
1503 self.run_init_wiring_post_init().await?;
1507 self.run_grpc_phase().await?;
1508 self.run_start_phase().await?;
1509 started = true;
1510 server.resolve_bearer_authenticator(&self.client_hub);
1516 self.compose_oop_router(server.options(), &hc_registry)
1517 .await
1518 }
1519 .await;
1520
1521 let serve_result = match composed {
1522 Ok((gear_router, openapi_json)) => {
1523 server.attach(gear_router, openapi_json);
1526 server.join().await
1528 }
1529 Err(e) => {
1530 tracing::error!(error = %e, "OoP startup failed before serving gear routes");
1531 self.cancel.cancel();
1533 if let Err(join_err) = server.join().await {
1534 tracing::warn!(error = %join_err, "OoP probe server teardown after startup failure errored");
1535 }
1536 Err(e)
1537 }
1538 };
1539
1540 if started && let Err(e) = self.run_stop_phase_guarded().await {
1543 tracing::warn!(error = %e, "OoP stop phase reported an error");
1544 }
1545
1546 serve_result
1547 }
1548}
1549
1550#[cfg(test)]
1551#[cfg(feature = "bootstrap")]
1552#[cfg_attr(coverage_nightly, coverage(off))]
1553#[path = "host_runtime_oop_tests.rs"]
1554mod host_runtime_oop_tests;
1555
1556#[cfg(test)]
1557#[cfg_attr(coverage_nightly, coverage(off))]
1558mod tests {
1559 use super::*;
1560 use crate::context::GearCtx;
1561 use crate::contracts::{Gear, RunnableCapability, SystemCapability};
1562 use crate::registry::RegistryBuilder;
1563 use std::sync::Arc;
1564 use std::sync::atomic::{AtomicUsize, Ordering};
1565 use tokio::sync::Mutex;
1566
1567 #[derive(Default)]
1568 #[allow(dead_code)]
1569 struct DummyCore;
1570 #[async_trait::async_trait]
1571 impl Gear for DummyCore {
1572 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1573 Ok(())
1574 }
1575 }
1576
1577 struct StopOrderTracker {
1578 my_order: usize,
1579 stop_order: Arc<AtomicUsize>,
1580 }
1581
1582 impl StopOrderTracker {
1583 fn new(counter: &Arc<AtomicUsize>, stop_order: Arc<AtomicUsize>) -> Self {
1584 let my_order = counter.fetch_add(1, Ordering::SeqCst);
1585 Self {
1586 my_order,
1587 stop_order,
1588 }
1589 }
1590 }
1591
1592 #[async_trait::async_trait]
1593 impl Gear for StopOrderTracker {
1594 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1595 Ok(())
1596 }
1597 }
1598
1599 #[async_trait::async_trait]
1600 impl RunnableCapability for StopOrderTracker {
1601 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1602 Ok(())
1603 }
1604 async fn stop(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1605 let order = self.stop_order.fetch_add(1, Ordering::SeqCst);
1606 tracing::info!(my_order = self.my_order, stop_order = order, "Gear stopped");
1607 Ok(())
1608 }
1609 }
1610
1611 #[tokio::test]
1612 async fn test_stop_phase_reverse_order() {
1613 let counter = Arc::new(AtomicUsize::new(0));
1614 let stop_order = Arc::new(AtomicUsize::new(0));
1615
1616 let gear_a = Arc::new(StopOrderTracker::new(&counter, stop_order.clone()));
1617 let gear_b = Arc::new(StopOrderTracker::new(&counter, stop_order.clone()));
1618 let gear_c = Arc::new(StopOrderTracker::new(&counter, stop_order.clone()));
1619
1620 let mut builder = RegistryBuilder::default();
1621 builder.register_core_with_meta("a", &[], gear_a.clone() as Arc<dyn Gear>);
1622 builder.register_core_with_meta("b", &["a"], gear_b.clone() as Arc<dyn Gear>);
1623 builder.register_core_with_meta("c", &["b"], gear_c.clone() as Arc<dyn Gear>);
1624
1625 builder.register_stateful_with_meta("a", gear_a.clone() as Arc<dyn RunnableCapability>);
1626 builder.register_stateful_with_meta("b", gear_b.clone() as Arc<dyn RunnableCapability>);
1627 builder.register_stateful_with_meta("c", gear_c.clone() as Arc<dyn RunnableCapability>);
1628
1629 let registry = builder.build_topo_sorted().unwrap();
1630
1631 let gear_names: Vec<_> = registry.gears().iter().map(|m| m.name).collect();
1633 assert_eq!(gear_names, vec!["a", "b", "c"]);
1634
1635 let client_hub = Arc::new(ClientHub::new());
1636 let cancel = CancellationToken::new();
1637 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1638
1639 let runtime = HostRuntime::new(
1640 registry,
1641 config_provider,
1642 DbOptions::None,
1643 client_hub,
1644 cancel.clone(),
1645 Uuid::new_v4(),
1646 None,
1647 );
1648
1649 runtime.run_stop_phase().await.unwrap();
1651
1652 assert_eq!(stop_order.load(Ordering::SeqCst), 3);
1656 }
1657
1658 #[tokio::test]
1659 async fn test_stop_phase_continues_on_error() {
1660 struct FailingGear {
1661 should_fail: bool,
1662 stopped: Arc<AtomicUsize>,
1663 }
1664
1665 #[async_trait::async_trait]
1666 impl Gear for FailingGear {
1667 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1668 Ok(())
1669 }
1670 }
1671
1672 #[async_trait::async_trait]
1673 impl RunnableCapability for FailingGear {
1674 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1675 Ok(())
1676 }
1677 async fn stop(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1678 self.stopped.fetch_add(1, Ordering::SeqCst);
1679 if self.should_fail {
1680 anyhow::bail!("Intentional failure")
1681 }
1682 Ok(())
1683 }
1684 }
1685
1686 let stopped = Arc::new(AtomicUsize::new(0));
1687 let gear_a = Arc::new(FailingGear {
1688 should_fail: false,
1689 stopped: stopped.clone(),
1690 });
1691 let gear_b = Arc::new(FailingGear {
1692 should_fail: true,
1693 stopped: stopped.clone(),
1694 });
1695 let gear_c = Arc::new(FailingGear {
1696 should_fail: false,
1697 stopped: stopped.clone(),
1698 });
1699
1700 let mut builder = RegistryBuilder::default();
1701 builder.register_core_with_meta("a", &[], gear_a.clone() as Arc<dyn Gear>);
1702 builder.register_core_with_meta("b", &["a"], gear_b.clone() as Arc<dyn Gear>);
1703 builder.register_core_with_meta("c", &["b"], gear_c.clone() as Arc<dyn Gear>);
1704
1705 builder.register_stateful_with_meta("a", gear_a.clone() as Arc<dyn RunnableCapability>);
1706 builder.register_stateful_with_meta("b", gear_b.clone() as Arc<dyn RunnableCapability>);
1707 builder.register_stateful_with_meta("c", gear_c.clone() as Arc<dyn RunnableCapability>);
1708
1709 let registry = builder.build_topo_sorted().unwrap();
1710
1711 let client_hub = Arc::new(ClientHub::new());
1712 let cancel = CancellationToken::new();
1713 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1714
1715 let runtime = HostRuntime::new(
1716 registry,
1717 config_provider,
1718 DbOptions::None,
1719 client_hub,
1720 cancel.clone(),
1721 Uuid::new_v4(),
1722 None,
1723 );
1724
1725 runtime.run_stop_phase().await.unwrap();
1727
1728 assert_eq!(stopped.load(Ordering::SeqCst), 3);
1730 }
1731
1732 struct EmptyConfigProvider;
1733 impl ConfigProvider for EmptyConfigProvider {
1734 fn get_gear_config(&self, _gear_name: &str) -> Option<&serde_json::Value> {
1735 None
1736 }
1737 }
1738
1739 #[test]
1740 fn static_endpoint_override_reads_nested_consumer_wiring_key() {
1741 struct MapCfg(std::collections::HashMap<String, serde_json::Value>);
1742 impl ConfigProvider for MapCfg {
1743 fn get_gear_config(&self, gear: &str) -> Option<&serde_json::Value> {
1744 self.0.get(gear)
1745 }
1746 }
1747 let mut map = std::collections::HashMap::new();
1748 map.insert(
1749 "orders".to_owned(),
1750 serde_json::json!({
1751 "config": { "consumer_wiring": { "billing": "http://localhost:8081" } }
1752 }),
1753 );
1754 let cfg = MapCfg(map);
1755
1756 assert_eq!(
1758 super::static_endpoint_override(&cfg, "orders", "billing").as_deref(),
1759 Some("http://localhost:8081")
1760 );
1761 assert_eq!(
1763 super::static_endpoint_override(&cfg, "orders", "inventory"),
1764 None
1765 );
1766 assert_eq!(
1767 super::static_endpoint_override(&cfg, "warehouse", "billing"),
1768 None
1769 );
1770 assert_eq!(
1771 super::static_endpoint_override(&EmptyConfigProvider, "orders", "billing"),
1772 None
1773 );
1774 }
1775
1776 #[test]
1782 fn static_endpoint_override_is_keyed_by_kebab_gear_name() {
1783 struct MapCfg(std::collections::HashMap<String, serde_json::Value>);
1784 impl ConfigProvider for MapCfg {
1785 fn get_gear_config(&self, gear: &str) -> Option<&serde_json::Value> {
1786 self.0.get(gear)
1787 }
1788 }
1789 let mut map = std::collections::HashMap::new();
1790 map.insert(
1791 "api-contracts-consumer".to_owned(),
1792 serde_json::json!({
1793 "config": { "consumer_wiring": { "api-contracts": "http://localhost:9099" } }
1794 }),
1795 );
1796 let cfg = MapCfg(map);
1797
1798 assert_eq!(
1799 super::static_endpoint_override(&cfg, "api-contracts-consumer", "api-contracts")
1800 .as_deref(),
1801 Some("http://localhost:9099"),
1802 );
1803 assert_eq!(
1805 super::static_endpoint_override(&cfg, "ApiContractsConsumer", "api-contracts"),
1806 None,
1807 );
1808 }
1809
1810 #[tokio::test]
1811 async fn test_post_init_runs_after_all_init_and_system_first() {
1812 #[derive(Clone)]
1813 struct TrackHooks {
1814 name: &'static str,
1815 events: Arc<Mutex<Vec<String>>>,
1816 }
1817
1818 #[async_trait::async_trait]
1819 impl Gear for TrackHooks {
1820 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1821 self.events.lock().await.push(format!("init:{}", self.name));
1822 Ok(())
1823 }
1824 }
1825
1826 #[async_trait::async_trait]
1827 impl SystemCapability for TrackHooks {
1828 fn pre_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
1829 Ok(())
1830 }
1831
1832 async fn post_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
1833 self.events
1834 .lock()
1835 .await
1836 .push(format!("post_init:{}", self.name));
1837 Ok(())
1838 }
1839 }
1840
1841 let events = Arc::new(Mutex::new(Vec::<String>::new()));
1842 let sys_a = Arc::new(TrackHooks {
1843 name: "sys_a",
1844 events: events.clone(),
1845 });
1846 let user_b = Arc::new(TrackHooks {
1847 name: "user_b",
1848 events: events.clone(),
1849 });
1850 let user_c = Arc::new(TrackHooks {
1851 name: "user_c",
1852 events: events.clone(),
1853 });
1854
1855 let mut builder = RegistryBuilder::default();
1856 builder.register_core_with_meta("sys_a", &[], sys_a.clone() as Arc<dyn Gear>);
1857 builder.register_core_with_meta("user_b", &["sys_a"], user_b.clone() as Arc<dyn Gear>);
1858 builder.register_core_with_meta("user_c", &["user_b"], user_c.clone() as Arc<dyn Gear>);
1859 builder.register_system_with_meta("sys_a", sys_a.clone() as Arc<dyn SystemCapability>);
1860
1861 let registry = builder.build_topo_sorted().unwrap();
1862
1863 let client_hub = Arc::new(ClientHub::new());
1864 let cancel = CancellationToken::new();
1865 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1866
1867 let runtime = HostRuntime::new(
1868 registry,
1869 config_provider,
1870 DbOptions::None,
1871 client_hub,
1872 cancel,
1873 Uuid::new_v4(),
1874 None,
1875 );
1876
1877 runtime.run_init_phase().await.unwrap();
1879 runtime.run_post_init_phase().await.unwrap();
1880
1881 let events = events.lock().await.clone();
1882 let first_post_init = events
1883 .iter()
1884 .position(|e| e.starts_with("post_init:"))
1885 .expect("expected post_init events");
1886 assert!(
1887 events[..first_post_init]
1888 .iter()
1889 .all(|e| e.starts_with("init:")),
1890 "expected all init events before post_init, got: {events:?}"
1891 );
1892
1893 assert_eq!(
1895 events,
1896 vec![
1897 "init:sys_a",
1898 "init:user_b",
1899 "init:user_c",
1900 "post_init:sys_a",
1901 ]
1902 );
1903 }
1904
1905 #[tokio::test]
1919 async fn init_wiring_post_init_runs_as_one_ordered_segment() {
1920 #[derive(Clone)]
1921 struct TrackHooks {
1922 events: Arc<Mutex<Vec<String>>>,
1923 }
1924
1925 #[async_trait::async_trait]
1926 impl Gear for TrackHooks {
1927 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1928 self.events.lock().await.push("init".to_owned());
1929 Ok(())
1930 }
1931 }
1932
1933 #[async_trait::async_trait]
1934 impl SystemCapability for TrackHooks {
1935 fn pre_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
1936 Ok(())
1937 }
1938
1939 async fn post_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
1940 self.events.lock().await.push("post_init".to_owned());
1941 Ok(())
1942 }
1943 }
1944
1945 let events = Arc::new(Mutex::new(Vec::<String>::new()));
1946 let gear = Arc::new(TrackHooks {
1947 events: events.clone(),
1948 });
1949
1950 let mut builder = RegistryBuilder::default();
1951 builder.register_core_with_meta("sys", &[], gear.clone() as Arc<dyn Gear>);
1952 builder.register_system_with_meta("sys", gear.clone() as Arc<dyn SystemCapability>);
1953 let registry = builder.build_topo_sorted().unwrap();
1954
1955 let runtime = HostRuntime::new(
1956 registry,
1957 Arc::new(EmptyConfigProvider) as Arc<dyn ConfigProvider>,
1958 DbOptions::None,
1959 Arc::new(ClientHub::new()),
1960 CancellationToken::new(),
1961 Uuid::new_v4(),
1962 None,
1963 );
1964
1965 runtime.run_init_wiring_post_init().await.unwrap();
1966
1967 assert_eq!(events.lock().await.clone(), vec!["init", "post_init"]);
1968 }
1969
1970 #[tokio::test]
1971 async fn test_stop_phase_provides_fresh_deadline_token() {
1972 use std::sync::atomic::AtomicBool;
1973
1974 struct TokenCheckGear {
1975 stop_was_called: AtomicBool,
1976 token_was_cancelled_on_entry: AtomicBool,
1977 }
1978
1979 #[async_trait::async_trait]
1980 impl Gear for TokenCheckGear {
1981 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1982 Ok(())
1983 }
1984 }
1985
1986 #[async_trait::async_trait]
1987 impl RunnableCapability for TokenCheckGear {
1988 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1989 Ok(())
1990 }
1991 async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
1992 self.stop_was_called.store(true, Ordering::SeqCst);
1994 self.token_was_cancelled_on_entry
1996 .store(deadline_token.is_cancelled(), Ordering::SeqCst);
1997 Ok(())
1998 }
1999 }
2000
2001 let gear = Arc::new(TokenCheckGear {
2002 stop_was_called: AtomicBool::new(false),
2003 token_was_cancelled_on_entry: AtomicBool::new(true),
2005 });
2006
2007 let mut builder = RegistryBuilder::default();
2008 builder.register_core_with_meta("test", &[], gear.clone() as Arc<dyn Gear>);
2009 builder.register_stateful_with_meta("test", gear.clone() as Arc<dyn RunnableCapability>);
2010
2011 let registry = builder.build_topo_sorted().unwrap();
2012 let client_hub = Arc::new(ClientHub::new());
2013 let cancel = CancellationToken::new();
2014 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
2015
2016 let runtime = HostRuntime::new(
2017 registry,
2018 config_provider,
2019 DbOptions::None,
2020 client_hub,
2021 cancel.clone(),
2022 Uuid::new_v4(),
2023 None,
2024 );
2025
2026 runtime.run_stop_phase().await.unwrap();
2028
2029 assert!(
2031 gear.stop_was_called.load(Ordering::SeqCst),
2032 "stop() was never called - gear may not have been registered correctly"
2033 );
2034
2035 assert!(
2038 !gear.token_was_cancelled_on_entry.load(Ordering::SeqCst),
2039 "deadline_token should NOT be cancelled when stop() is called - this enables graceful shutdown"
2040 );
2041 }
2042
2043 #[tokio::test]
2044 async fn test_stop_phase_graceful_shutdown_completes_before_deadline() {
2045 use std::sync::atomic::AtomicBool;
2046 use std::time::Duration;
2047
2048 struct GracefulGear {
2049 graceful_completed: AtomicBool,
2050 deadline_fired: AtomicBool,
2051 }
2052
2053 #[async_trait::async_trait]
2054 impl Gear for GracefulGear {
2055 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
2056 Ok(())
2057 }
2058 }
2059
2060 #[async_trait::async_trait]
2061 impl RunnableCapability for GracefulGear {
2062 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
2063 Ok(())
2064 }
2065 async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
2066 tokio::select! {
2068 () = tokio::time::sleep(Duration::from_millis(10)) => {
2069 self.graceful_completed.store(true, Ordering::SeqCst);
2070 }
2071 () = deadline_token.cancelled() => {
2072 self.deadline_fired.store(true, Ordering::SeqCst);
2073 }
2074 }
2075 Ok(())
2076 }
2077 }
2078
2079 let gear = Arc::new(GracefulGear {
2080 graceful_completed: AtomicBool::new(false),
2081 deadline_fired: AtomicBool::new(false),
2082 });
2083
2084 let mut builder = RegistryBuilder::default();
2085 builder.register_core_with_meta("test", &[], gear.clone() as Arc<dyn Gear>);
2086 builder.register_stateful_with_meta("test", gear.clone() as Arc<dyn RunnableCapability>);
2087
2088 let registry = builder.build_topo_sorted().unwrap();
2089 let client_hub = Arc::new(ClientHub::new());
2090 let cancel = CancellationToken::new();
2091 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
2092
2093 let runtime = HostRuntime::new(
2095 registry,
2096 config_provider,
2097 DbOptions::None,
2098 client_hub,
2099 cancel.clone(),
2100 Uuid::new_v4(),
2101 None,
2102 )
2103 .with_shutdown_deadline(Duration::from_secs(5));
2104
2105 runtime.run_stop_phase().await.unwrap();
2106
2107 assert!(
2109 gear.graceful_completed.load(Ordering::SeqCst),
2110 "graceful shutdown should complete"
2111 );
2112 assert!(
2114 !gear.deadline_fired.load(Ordering::SeqCst),
2115 "deadline should not fire when graceful shutdown completes quickly"
2116 );
2117 }
2118
2119 #[tokio::test]
2120 async fn test_stop_phase_deadline_fires_for_slow_gear() {
2121 use std::sync::atomic::AtomicBool;
2122 use std::time::Duration;
2123
2124 struct SlowGear {
2125 graceful_completed: AtomicBool,
2126 deadline_fired: AtomicBool,
2127 }
2128
2129 #[async_trait::async_trait]
2130 impl Gear for SlowGear {
2131 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
2132 Ok(())
2133 }
2134 }
2135
2136 #[async_trait::async_trait]
2137 impl RunnableCapability for SlowGear {
2138 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
2139 Ok(())
2140 }
2141 async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
2142 tokio::select! {
2144 () = tokio::time::sleep(Duration::from_secs(10)) => {
2145 self.graceful_completed.store(true, Ordering::SeqCst);
2146 }
2147 () = deadline_token.cancelled() => {
2148 self.deadline_fired.store(true, Ordering::SeqCst);
2149 }
2150 }
2151 Ok(())
2152 }
2153 }
2154
2155 let gear = Arc::new(SlowGear {
2156 graceful_completed: AtomicBool::new(false),
2157 deadline_fired: AtomicBool::new(false),
2158 });
2159
2160 let mut builder = RegistryBuilder::default();
2161 builder.register_core_with_meta("test", &[], gear.clone() as Arc<dyn Gear>);
2162 builder.register_stateful_with_meta("test", gear.clone() as Arc<dyn RunnableCapability>);
2163
2164 let registry = builder.build_topo_sorted().unwrap();
2165 let client_hub = Arc::new(ClientHub::new());
2166 let cancel = CancellationToken::new();
2167 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
2168
2169 let runtime = HostRuntime::new(
2171 registry,
2172 config_provider,
2173 DbOptions::None,
2174 client_hub,
2175 cancel.clone(),
2176 Uuid::new_v4(),
2177 None,
2178 )
2179 .with_shutdown_deadline(Duration::from_millis(100));
2180
2181 runtime.run_stop_phase().await.unwrap();
2182
2183 assert!(
2185 !gear.graceful_completed.load(Ordering::SeqCst),
2186 "graceful shutdown should not complete when deadline fires first"
2187 );
2188 assert!(
2190 gear.deadline_fired.load(Ordering::SeqCst),
2191 "deadline should fire for slow gears"
2192 );
2193 }
2194}