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