1use axum::Router;
17use std::collections::HashSet;
18use std::sync::Arc;
19
20use tokio_util::sync::CancellationToken;
21use uuid::Uuid;
22
23use crate::backends::OopSpawnConfig;
24use crate::client_hub::ClientHub;
25use crate::config::ConfigProvider;
26use crate::context::GearContextBuilder;
27use crate::registry::{
28 ApiGatewayCap, GearEntry, GearRegistry, GrpcHubCap, RegistryError, RestApiCap, RunnableCap,
29 SystemCap,
30};
31use crate::runtime::{GearManager, GrpcInstallerStore, OopSpawnOptions, SystemContext};
32
33#[cfg(feature = "db")]
34use crate::registry::DatabaseCap;
35
36#[derive(Clone)]
38pub enum DbOptions {
39 None,
41 #[cfg(feature = "db")]
43 Manager(Arc<toolkit_db::DbManager>),
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum RunMode {
49 Full,
51 MigrateOnly,
53}
54
55pub const TOOLKIT_DIRECTORY_ENDPOINT_ENV: &str = "TOOLKIT_DIRECTORY_ENDPOINT";
57
58pub const TOOLKIT_MODULE_CONFIG_ENV: &str = "TOOLKIT_MODULE_CONFIG";
60
61pub const DEFAULT_SHUTDOWN_DEADLINE: std::time::Duration = std::time::Duration::from_secs(35);
67
68pub struct HostRuntime {
72 registry: GearRegistry,
73 ctx_builder: GearContextBuilder,
74 instance_id: Uuid,
75 gear_manager: Arc<GearManager>,
76 grpc_installers: Arc<GrpcInstallerStore>,
77 #[allow(dead_code)]
78 client_hub: Arc<ClientHub>,
79 cancel: CancellationToken,
80 #[allow(dead_code)]
81 db_options: DbOptions,
82 oop_options: Option<OopSpawnOptions>,
84 shutdown_deadline: std::time::Duration,
86}
87
88impl HostRuntime {
89 pub fn new(
93 registry: GearRegistry,
94 gears_cfg: Arc<dyn ConfigProvider>,
95 db_options: DbOptions,
96 client_hub: Arc<ClientHub>,
97 cancel: CancellationToken,
98 instance_id: Uuid,
99 oop_options: Option<OopSpawnOptions>,
100 ) -> Self {
101 let gear_manager = Arc::new(GearManager::new());
103 let grpc_installers = Arc::new(GrpcInstallerStore::new());
104
105 let ctx_builder =
107 GearContextBuilder::new(instance_id, gears_cfg, client_hub.clone(), cancel.clone());
108 #[cfg(feature = "db")]
109 let ctx_builder = match &db_options {
110 DbOptions::Manager(mgr) => ctx_builder.with_db_manager(mgr.clone()),
111 DbOptions::None => ctx_builder,
112 };
113
114 Self {
115 registry,
116 ctx_builder,
117 instance_id,
118 gear_manager,
119 grpc_installers,
120 client_hub,
121 cancel,
122 db_options,
123 oop_options,
124 shutdown_deadline: DEFAULT_SHUTDOWN_DEADLINE,
125 }
126 }
127
128 #[must_use]
144 pub fn with_shutdown_deadline(mut self, deadline: std::time::Duration) -> Self {
145 self.shutdown_deadline = deadline;
146 self
147 }
148
149 pub fn run_pre_init_phase(&self) -> Result<(), RegistryError> {
156 tracing::info!("Phase: pre_init");
157
158 let sys_ctx = SystemContext::new(
159 self.instance_id,
160 Arc::clone(&self.gear_manager),
161 Arc::clone(&self.grpc_installers),
162 );
163
164 for entry in self.registry.gears() {
165 if self.cancel.is_cancelled() {
167 tracing::warn!("Pre-init phase cancelled by signal");
168 return Err(RegistryError::Cancelled);
169 }
170
171 if let Some(sys_mod) = entry.caps.query::<SystemCap>() {
172 tracing::debug!(gear = entry.name, "Running system pre_init");
173 sys_mod
174 .pre_init(&sys_ctx)
175 .map_err(|e| RegistryError::PreInit {
176 gear: entry.name,
177 source: e,
178 })?;
179 }
180 }
181
182 Ok(())
183 }
184
185 #[cfg(feature = "db")]
187 async fn gear_context(
188 &self,
189 gear_name: &'static str,
190 ) -> Result<crate::context::GearCtx, RegistryError> {
191 self.ctx_builder
192 .for_gear(gear_name)
193 .await
194 .map_err(|e| RegistryError::DbMigrate {
195 gear: gear_name,
196 source: e,
197 })
198 }
199
200 #[cfg(feature = "db")]
202 async fn db_migration_target(
203 &self,
204 gear_name: &'static str,
205 ctx: &crate::context::GearCtx,
206 db_gear: Option<Arc<dyn crate::contracts::DatabaseCapability>>,
207 ) -> Result<
208 Option<(
209 toolkit_db::Db,
210 Arc<dyn crate::contracts::DatabaseCapability>,
211 )>,
212 RegistryError,
213 > {
214 let Some(dbm) = db_gear else {
215 return Ok(None);
216 };
217
218 let db = match &self.db_options {
222 DbOptions::None => None,
223 #[cfg(feature = "db")]
224 DbOptions::Manager(mgr) => {
225 mgr.get(gear_name)
226 .await
227 .map_err(|e| RegistryError::DbMigrate {
228 gear: gear_name,
229 source: e.into(),
230 })?
231 }
232 };
233
234 _ = ctx; Ok(db.map(|db| (db, dbm)))
236 }
237
238 #[cfg(feature = "db")]
243 async fn migrate_gear(
244 gear_name: &'static str,
245 db: &toolkit_db::Db,
246 db_gear: Arc<dyn crate::contracts::DatabaseCapability>,
247 ) -> Result<(), RegistryError> {
248 let migrations = db_gear.migrations();
250
251 if migrations.is_empty() {
252 tracing::debug!(gear = gear_name, "No migrations to run");
253 return Ok(());
254 }
255
256 tracing::debug!(
257 gear = gear_name,
258 count = migrations.len(),
259 "Running DB migrations"
260 );
261
262 let result =
264 toolkit_db::migration_runner::run_migrations_for_gear(db, gear_name, migrations)
265 .await
266 .map_err(|e| RegistryError::DbMigrate {
267 gear: gear_name,
268 source: anyhow::Error::new(e),
269 })?;
270
271 tracing::info!(
272 gear = gear_name,
273 applied = result.applied,
274 skipped = result.skipped,
275 "DB migrations completed"
276 );
277
278 Ok(())
279 }
280
281 #[cfg(feature = "db")]
290 async fn run_db_phase(&self) -> Result<(), RegistryError> {
291 tracing::info!("Phase: db (before init)");
292
293 for entry in self.registry.gears_by_system_priority() {
294 if self.cancel.is_cancelled() {
296 tracing::warn!("DB migration phase cancelled by signal");
297 return Err(RegistryError::Cancelled);
298 }
299
300 let ctx = self.gear_context(entry.name).await?;
301 let db_gear = entry.caps.query::<DatabaseCap>();
302
303 match self
304 .db_migration_target(entry.name, &ctx, db_gear.clone())
305 .await?
306 {
307 Some((db, dbm)) => {
308 Self::migrate_gear(entry.name, &db, dbm).await?;
309 }
310 None if db_gear.is_some() => {
311 tracing::debug!(
312 gear = entry.name,
313 "Gear has DbGear trait but no DB handle (no config)"
314 );
315 }
316 None => {}
317 }
318 }
319
320 Ok(())
321 }
322
323 async fn run_init_phase(&self) -> Result<(), RegistryError> {
327 tracing::info!("Phase: init");
328
329 for entry in self.registry.gears_by_system_priority() {
330 let ctx =
331 self.ctx_builder
332 .for_gear(entry.name)
333 .await
334 .map_err(|e| RegistryError::Init {
335 gear: entry.name,
336 source: e,
337 })?;
338 tracing::info!(gear = entry.name, "Initializing a gear...");
339 entry
340 .core
341 .init(&ctx)
342 .await
343 .map_err(|e| RegistryError::Init {
344 gear: entry.name,
345 source: e,
346 })?;
347 tracing::info!(gear = entry.name, "Initialized a gear.");
348 }
349
350 Ok(())
351 }
352
353 async fn run_post_init_phase(&self) -> Result<(), RegistryError> {
360 tracing::info!("Phase: post_init");
361
362 let sys_ctx = SystemContext::new(
363 self.instance_id,
364 Arc::clone(&self.gear_manager),
365 Arc::clone(&self.grpc_installers),
366 );
367
368 for entry in self.registry.gears_by_system_priority() {
369 if let Some(sys_mod) = entry.caps.query::<SystemCap>() {
370 sys_mod
371 .post_init(&sys_ctx)
372 .await
373 .map_err(|e| RegistryError::PostInit {
374 gear: entry.name,
375 source: e,
376 })?;
377 }
378 }
379
380 Ok(())
381 }
382
383 async fn run_rest_phase(&self) -> Result<Router, RegistryError> {
390 tracing::info!("Phase: rest (sync)");
391
392 let mut router = Router::new();
393
394 let host_count = self
396 .registry
397 .gears()
398 .iter()
399 .filter(|e| e.caps.has::<ApiGatewayCap>())
400 .count();
401
402 match host_count {
403 0 => {
404 return if self
405 .registry
406 .gears()
407 .iter()
408 .any(|e| e.caps.has::<RestApiCap>())
409 {
410 Err(RegistryError::RestRequiresHost)
411 } else {
412 Ok(router)
413 };
414 }
415 1 => { }
416 _ => return Err(RegistryError::MultipleRestHosts),
417 }
418
419 let host_idx = self
421 .registry
422 .gears()
423 .iter()
424 .position(|e| e.caps.has::<ApiGatewayCap>())
425 .ok_or(RegistryError::RestHostNotFoundAfterValidation)?;
426 let host_entry = &self.registry.gears()[host_idx];
427 let Some(host) = host_entry.caps.query::<ApiGatewayCap>() else {
428 return Err(RegistryError::RestHostMissingFromEntry);
429 };
430 let host_ctx = self
431 .ctx_builder
432 .for_gear(host_entry.name)
433 .await
434 .map_err(|e| RegistryError::RestPrepare {
435 gear: host_entry.name,
436 source: e,
437 })?;
438
439 let registry: &dyn crate::contracts::OpenApiRegistry = host.as_registry();
441
442 router =
444 host.rest_prepare(&host_ctx, router)
445 .map_err(|source| RegistryError::RestPrepare {
446 gear: host_entry.name,
447 source,
448 })?;
449
450 for e in self.registry.gears() {
452 if let Some(rest) = e.caps.query::<RestApiCap>() {
453 let ctx = self.ctx_builder.for_gear(e.name).await.map_err(|err| {
454 RegistryError::RestRegister {
455 gear: e.name,
456 source: err,
457 }
458 })?;
459
460 router = rest
461 .register_rest(&ctx, router, registry)
462 .map_err(|source| RegistryError::RestRegister {
463 gear: e.name,
464 source,
465 })?;
466 }
467 }
468
469 router = host.rest_finalize(&host_ctx, router).map_err(|source| {
471 RegistryError::RestFinalize {
472 gear: host_entry.name,
473 source,
474 }
475 })?;
476
477 Ok(router)
478 }
479
480 async fn run_grpc_phase(&self) -> Result<(), RegistryError> {
484 tracing::info!("Phase: grpc (registration)");
485
486 if self.registry.grpc_hub.is_none() && self.registry.grpc_services.is_empty() {
488 return Ok(());
489 }
490
491 if self.registry.grpc_hub.is_none() && !self.registry.grpc_services.is_empty() {
493 return Err(RegistryError::GrpcRequiresHub);
494 }
495
496 if let Some(hub_name) = &self.registry.grpc_hub {
498 let mut gears_data = Vec::new();
499 let mut seen = HashSet::new();
500
501 for (gear_name, service_gear) in &self.registry.grpc_services {
503 let ctx = self.ctx_builder.for_gear(gear_name).await.map_err(|err| {
504 RegistryError::GrpcRegister {
505 gear: gear_name.clone(),
506 source: err,
507 }
508 })?;
509
510 let installers = service_gear
511 .get_grpc_services(&ctx)
512 .await
513 .map_err(|source| RegistryError::GrpcRegister {
514 gear: gear_name.clone(),
515 source,
516 })?;
517
518 for reg in &installers {
519 if !seen.insert(reg.service_name) {
520 return Err(RegistryError::GrpcRegister {
521 gear: gear_name.clone(),
522 source: anyhow::anyhow!(
523 "Duplicate gRPC service name: {}",
524 reg.service_name
525 ),
526 });
527 }
528 }
529
530 gears_data.push(crate::runtime::GearInstallers {
531 gear_name: gear_name.clone(),
532 installers,
533 });
534 }
535
536 self.grpc_installers
537 .set(crate::runtime::GrpcInstallerData { gears: gears_data })
538 .map_err(|source| RegistryError::GrpcRegister {
539 gear: hub_name.clone(),
540 source,
541 })?;
542 }
543
544 Ok(())
545 }
546
547 async fn run_start_phase(&self) -> Result<(), RegistryError> {
551 tracing::info!("Phase: start");
552
553 for e in self.registry.gears_by_system_priority() {
554 if let Some(s) = e.caps.query::<RunnableCap>() {
555 tracing::debug!(
556 gear = e.name,
557 is_system = e.caps.has::<SystemCap>(),
558 "Starting stateful gear"
559 );
560 s.start(self.cancel.clone())
561 .await
562 .map_err(|source| RegistryError::Start {
563 gear: e.name,
564 source,
565 })?;
566 tracing::info!(gear = e.name, "Started gear");
567 }
568 }
569
570 Ok(())
571 }
572
573 async fn stop_one_gear(entry: &GearEntry, cancel: CancellationToken) {
575 if let Some(s) = entry.caps.query::<RunnableCap>() {
576 match s.stop(cancel).await {
577 Err(err) => {
578 tracing::warn!(gear = entry.name, error = %err, "Failed to stop gear");
579 }
580 _ => {
581 tracing::info!(gear = entry.name, "Stopped gear");
582 }
583 }
584 }
585 }
586
587 async fn run_stop_phase(&self) -> Result<(), RegistryError> {
612 tracing::info!("Phase: stop");
613
614 let deadline = self.shutdown_deadline;
615
616 for e in self.registry.gears().iter().rev() {
618 let gear_name = e.name;
619
620 let deadline_token = CancellationToken::new();
623 let deadline_token_for_timeout = deadline_token.clone();
624
625 let deadline_task = tokio::spawn(async move {
627 tokio::time::sleep(deadline).await;
628 tracing::warn!(
629 gear = gear_name,
630 deadline_secs = deadline.as_secs(),
631 "Gear shutdown deadline reached, sending hard-stop signal"
632 );
633 deadline_token_for_timeout.cancel();
634 });
635
636 Self::stop_one_gear(e, deadline_token).await;
639
640 deadline_task.abort();
642 #[allow(clippy::let_underscore_must_use)]
643 let _ = deadline_task.await;
644 }
645
646 Ok(())
647 }
648
649 async fn run_oop_spawn_phase(&self) -> Result<(), RegistryError> {
654 let oop_opts = match &self.oop_options {
655 Some(opts) if !opts.gears.is_empty() => opts,
656 _ => return Ok(()),
657 };
658
659 tracing::info!("Phase: oop_spawn");
660
661 let directory_endpoint = self.wait_for_grpc_hub_endpoint().await;
663
664 for gear_cfg in &oop_opts.gears {
665 let mut env = gear_cfg.env.clone();
668 env.insert(
669 TOOLKIT_MODULE_CONFIG_ENV.to_owned(),
670 gear_cfg.rendered_config_json.clone(),
671 );
672 if let Some(ref endpoint) = directory_endpoint {
673 env.insert(TOOLKIT_DIRECTORY_ENDPOINT_ENV.to_owned(), endpoint.clone());
674 }
675
676 let args = gear_cfg.args.clone();
678
679 let spawn_config = OopSpawnConfig {
680 gear_name: gear_cfg.gear_name.clone(),
681 binary: gear_cfg.binary.clone(),
682 args,
683 env,
684 working_directory: gear_cfg.working_directory.clone(),
685 };
686
687 oop_opts
688 .backend
689 .spawn(spawn_config)
690 .await
691 .map_err(|e| RegistryError::OopSpawn {
692 gear: gear_cfg.gear_name.clone(),
693 source: e,
694 })?;
695
696 tracing::info!(
697 gear = %gear_cfg.gear_name,
698 directory_endpoint = ?directory_endpoint,
699 "Spawned OoP gear via backend"
700 );
701 }
702
703 Ok(())
704 }
705
706 async fn wait_for_grpc_hub_endpoint(&self) -> Option<String> {
711 const POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(10);
712 const MAX_WAIT: std::time::Duration = std::time::Duration::from_secs(5);
713
714 let grpc_hub = self
716 .registry
717 .gears()
718 .iter()
719 .find_map(|e| e.caps.query::<GrpcHubCap>());
720
721 let Some(hub) = grpc_hub else {
722 return None; };
724
725 let start = std::time::Instant::now();
726
727 loop {
728 if let Some(endpoint) = hub.bound_endpoint() {
729 tracing::debug!(
730 endpoint = %endpoint,
731 elapsed_ms = start.elapsed().as_millis(),
732 "gRPC hub endpoint available"
733 );
734 return Some(endpoint);
735 }
736
737 if start.elapsed() > MAX_WAIT {
738 tracing::warn!("Timed out waiting for gRPC hub to bind");
739 return None;
740 }
741
742 tokio::time::sleep(POLL_INTERVAL).await;
743 }
744 }
745
746 pub async fn run_gear_phases(self) -> anyhow::Result<()> {
755 self.run_phases_internal(RunMode::Full).await
756 }
757
758 pub async fn run_migration_phases(self) -> anyhow::Result<()> {
768 self.run_phases_internal(RunMode::MigrateOnly).await
769 }
770
771 async fn run_phases_internal(self, mode: RunMode) -> anyhow::Result<()> {
794 match mode {
796 RunMode::Full => {
797 tracing::info!("Running full lifecycle (all phases)");
798 }
799 RunMode::MigrateOnly => {
800 tracing::info!("Running in migration mode (pre-init + db phases only)");
801 }
802 }
803
804 self.run_pre_init_phase()?;
806
807 #[cfg(feature = "db")]
809 {
810 self.run_db_phase().await?;
811 }
812 #[cfg(not(feature = "db"))]
813 {
814 }
816
817 if mode == RunMode::MigrateOnly {
819 tracing::info!("Migration phases completed successfully");
820 return Ok(());
821 }
822
823 self.run_init_phase().await?;
825
826 self.run_post_init_phase().await?;
828
829 let _router = self.run_rest_phase().await?;
831
832 self.run_grpc_phase().await?;
834
835 self.run_start_phase().await?;
837
838 self.run_oop_spawn_phase().await?;
840
841 self.cancel.cancelled().await;
843
844 let stop_timeout = std::time::Duration::from_secs(15);
850 let disarm = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
851 let disarm_clone = std::sync::Arc::clone(&disarm);
852 std::thread::spawn(move || {
853 std::thread::sleep(stop_timeout);
854 if !disarm_clone.load(std::sync::atomic::Ordering::Relaxed) {
855 tracing::warn!(
856 timeout_secs = stop_timeout.as_secs(),
857 "shutdown: stop phase timed out, force exiting"
858 );
859 std::process::exit(1);
860 }
861 });
862
863 self.run_stop_phase().await?;
864 disarm.store(true, std::sync::atomic::Ordering::Relaxed);
865
866 Ok(())
867 }
868}
869
870#[cfg(test)]
871#[cfg_attr(coverage_nightly, coverage(off))]
872mod tests {
873 use super::*;
874 use crate::context::GearCtx;
875 use crate::contracts::{Gear, RunnableCapability, SystemCapability};
876 use crate::registry::RegistryBuilder;
877 use std::sync::Arc;
878 use std::sync::atomic::{AtomicUsize, Ordering};
879 use tokio::sync::Mutex;
880
881 #[derive(Default)]
882 #[allow(dead_code)]
883 struct DummyCore;
884 #[async_trait::async_trait]
885 impl Gear for DummyCore {
886 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
887 Ok(())
888 }
889 }
890
891 struct StopOrderTracker {
892 my_order: usize,
893 stop_order: Arc<AtomicUsize>,
894 }
895
896 impl StopOrderTracker {
897 fn new(counter: &Arc<AtomicUsize>, stop_order: Arc<AtomicUsize>) -> Self {
898 let my_order = counter.fetch_add(1, Ordering::SeqCst);
899 Self {
900 my_order,
901 stop_order,
902 }
903 }
904 }
905
906 #[async_trait::async_trait]
907 impl Gear for StopOrderTracker {
908 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
909 Ok(())
910 }
911 }
912
913 #[async_trait::async_trait]
914 impl RunnableCapability for StopOrderTracker {
915 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
916 Ok(())
917 }
918 async fn stop(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
919 let order = self.stop_order.fetch_add(1, Ordering::SeqCst);
920 tracing::info!(my_order = self.my_order, stop_order = order, "Gear stopped");
921 Ok(())
922 }
923 }
924
925 #[tokio::test]
926 async fn test_stop_phase_reverse_order() {
927 let counter = Arc::new(AtomicUsize::new(0));
928 let stop_order = Arc::new(AtomicUsize::new(0));
929
930 let gear_a = Arc::new(StopOrderTracker::new(&counter, stop_order.clone()));
931 let gear_b = Arc::new(StopOrderTracker::new(&counter, stop_order.clone()));
932 let gear_c = Arc::new(StopOrderTracker::new(&counter, stop_order.clone()));
933
934 let mut builder = RegistryBuilder::default();
935 builder.register_core_with_meta("a", &[], gear_a.clone() as Arc<dyn Gear>);
936 builder.register_core_with_meta("b", &["a"], gear_b.clone() as Arc<dyn Gear>);
937 builder.register_core_with_meta("c", &["b"], gear_c.clone() as Arc<dyn Gear>);
938
939 builder.register_stateful_with_meta("a", gear_a.clone() as Arc<dyn RunnableCapability>);
940 builder.register_stateful_with_meta("b", gear_b.clone() as Arc<dyn RunnableCapability>);
941 builder.register_stateful_with_meta("c", gear_c.clone() as Arc<dyn RunnableCapability>);
942
943 let registry = builder.build_topo_sorted().unwrap();
944
945 let gear_names: Vec<_> = registry.gears().iter().map(|m| m.name).collect();
947 assert_eq!(gear_names, vec!["a", "b", "c"]);
948
949 let client_hub = Arc::new(ClientHub::new());
950 let cancel = CancellationToken::new();
951 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
952
953 let runtime = HostRuntime::new(
954 registry,
955 config_provider,
956 DbOptions::None,
957 client_hub,
958 cancel.clone(),
959 Uuid::new_v4(),
960 None,
961 );
962
963 runtime.run_stop_phase().await.unwrap();
965
966 assert_eq!(stop_order.load(Ordering::SeqCst), 3);
970 }
971
972 #[tokio::test]
973 async fn test_stop_phase_continues_on_error() {
974 struct FailingGear {
975 should_fail: bool,
976 stopped: Arc<AtomicUsize>,
977 }
978
979 #[async_trait::async_trait]
980 impl Gear for FailingGear {
981 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
982 Ok(())
983 }
984 }
985
986 #[async_trait::async_trait]
987 impl RunnableCapability for FailingGear {
988 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
989 Ok(())
990 }
991 async fn stop(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
992 self.stopped.fetch_add(1, Ordering::SeqCst);
993 if self.should_fail {
994 anyhow::bail!("Intentional failure")
995 }
996 Ok(())
997 }
998 }
999
1000 let stopped = Arc::new(AtomicUsize::new(0));
1001 let gear_a = Arc::new(FailingGear {
1002 should_fail: false,
1003 stopped: stopped.clone(),
1004 });
1005 let gear_b = Arc::new(FailingGear {
1006 should_fail: true,
1007 stopped: stopped.clone(),
1008 });
1009 let gear_c = Arc::new(FailingGear {
1010 should_fail: false,
1011 stopped: stopped.clone(),
1012 });
1013
1014 let mut builder = RegistryBuilder::default();
1015 builder.register_core_with_meta("a", &[], gear_a.clone() as Arc<dyn Gear>);
1016 builder.register_core_with_meta("b", &["a"], gear_b.clone() as Arc<dyn Gear>);
1017 builder.register_core_with_meta("c", &["b"], gear_c.clone() as Arc<dyn Gear>);
1018
1019 builder.register_stateful_with_meta("a", gear_a.clone() as Arc<dyn RunnableCapability>);
1020 builder.register_stateful_with_meta("b", gear_b.clone() as Arc<dyn RunnableCapability>);
1021 builder.register_stateful_with_meta("c", gear_c.clone() as Arc<dyn RunnableCapability>);
1022
1023 let registry = builder.build_topo_sorted().unwrap();
1024
1025 let client_hub = Arc::new(ClientHub::new());
1026 let cancel = CancellationToken::new();
1027 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1028
1029 let runtime = HostRuntime::new(
1030 registry,
1031 config_provider,
1032 DbOptions::None,
1033 client_hub,
1034 cancel.clone(),
1035 Uuid::new_v4(),
1036 None,
1037 );
1038
1039 runtime.run_stop_phase().await.unwrap();
1041
1042 assert_eq!(stopped.load(Ordering::SeqCst), 3);
1044 }
1045
1046 struct EmptyConfigProvider;
1047 impl ConfigProvider for EmptyConfigProvider {
1048 fn get_gear_config(&self, _gear_name: &str) -> Option<&serde_json::Value> {
1049 None
1050 }
1051 }
1052
1053 #[tokio::test]
1054 async fn test_post_init_runs_after_all_init_and_system_first() {
1055 #[derive(Clone)]
1056 struct TrackHooks {
1057 name: &'static str,
1058 events: Arc<Mutex<Vec<String>>>,
1059 }
1060
1061 #[async_trait::async_trait]
1062 impl Gear for TrackHooks {
1063 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1064 self.events.lock().await.push(format!("init:{}", self.name));
1065 Ok(())
1066 }
1067 }
1068
1069 #[async_trait::async_trait]
1070 impl SystemCapability for TrackHooks {
1071 fn pre_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
1072 Ok(())
1073 }
1074
1075 async fn post_init(&self, _sys: &crate::runtime::SystemContext) -> anyhow::Result<()> {
1076 self.events
1077 .lock()
1078 .await
1079 .push(format!("post_init:{}", self.name));
1080 Ok(())
1081 }
1082 }
1083
1084 let events = Arc::new(Mutex::new(Vec::<String>::new()));
1085 let sys_a = Arc::new(TrackHooks {
1086 name: "sys_a",
1087 events: events.clone(),
1088 });
1089 let user_b = Arc::new(TrackHooks {
1090 name: "user_b",
1091 events: events.clone(),
1092 });
1093 let user_c = Arc::new(TrackHooks {
1094 name: "user_c",
1095 events: events.clone(),
1096 });
1097
1098 let mut builder = RegistryBuilder::default();
1099 builder.register_core_with_meta("sys_a", &[], sys_a.clone() as Arc<dyn Gear>);
1100 builder.register_core_with_meta("user_b", &["sys_a"], user_b.clone() as Arc<dyn Gear>);
1101 builder.register_core_with_meta("user_c", &["user_b"], user_c.clone() as Arc<dyn Gear>);
1102 builder.register_system_with_meta("sys_a", sys_a.clone() as Arc<dyn SystemCapability>);
1103
1104 let registry = builder.build_topo_sorted().unwrap();
1105
1106 let client_hub = Arc::new(ClientHub::new());
1107 let cancel = CancellationToken::new();
1108 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1109
1110 let runtime = HostRuntime::new(
1111 registry,
1112 config_provider,
1113 DbOptions::None,
1114 client_hub,
1115 cancel,
1116 Uuid::new_v4(),
1117 None,
1118 );
1119
1120 runtime.run_init_phase().await.unwrap();
1122 runtime.run_post_init_phase().await.unwrap();
1123
1124 let events = events.lock().await.clone();
1125 let first_post_init = events
1126 .iter()
1127 .position(|e| e.starts_with("post_init:"))
1128 .expect("expected post_init events");
1129 assert!(
1130 events[..first_post_init]
1131 .iter()
1132 .all(|e| e.starts_with("init:")),
1133 "expected all init events before post_init, got: {events:?}"
1134 );
1135
1136 assert_eq!(
1138 events,
1139 vec![
1140 "init:sys_a",
1141 "init:user_b",
1142 "init:user_c",
1143 "post_init:sys_a",
1144 ]
1145 );
1146 }
1147
1148 #[tokio::test]
1149 async fn test_stop_phase_provides_fresh_deadline_token() {
1150 use std::sync::atomic::AtomicBool;
1151
1152 struct TokenCheckGear {
1153 stop_was_called: AtomicBool,
1154 token_was_cancelled_on_entry: AtomicBool,
1155 }
1156
1157 #[async_trait::async_trait]
1158 impl Gear for TokenCheckGear {
1159 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1160 Ok(())
1161 }
1162 }
1163
1164 #[async_trait::async_trait]
1165 impl RunnableCapability for TokenCheckGear {
1166 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1167 Ok(())
1168 }
1169 async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
1170 self.stop_was_called.store(true, Ordering::SeqCst);
1172 self.token_was_cancelled_on_entry
1174 .store(deadline_token.is_cancelled(), Ordering::SeqCst);
1175 Ok(())
1176 }
1177 }
1178
1179 let gear = Arc::new(TokenCheckGear {
1180 stop_was_called: AtomicBool::new(false),
1181 token_was_cancelled_on_entry: AtomicBool::new(true),
1183 });
1184
1185 let mut builder = RegistryBuilder::default();
1186 builder.register_core_with_meta("test", &[], gear.clone() as Arc<dyn Gear>);
1187 builder.register_stateful_with_meta("test", gear.clone() as Arc<dyn RunnableCapability>);
1188
1189 let registry = builder.build_topo_sorted().unwrap();
1190 let client_hub = Arc::new(ClientHub::new());
1191 let cancel = CancellationToken::new();
1192 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1193
1194 let runtime = HostRuntime::new(
1195 registry,
1196 config_provider,
1197 DbOptions::None,
1198 client_hub,
1199 cancel.clone(),
1200 Uuid::new_v4(),
1201 None,
1202 );
1203
1204 runtime.run_stop_phase().await.unwrap();
1206
1207 assert!(
1209 gear.stop_was_called.load(Ordering::SeqCst),
1210 "stop() was never called - gear may not have been registered correctly"
1211 );
1212
1213 assert!(
1216 !gear.token_was_cancelled_on_entry.load(Ordering::SeqCst),
1217 "deadline_token should NOT be cancelled when stop() is called - this enables graceful shutdown"
1218 );
1219 }
1220
1221 #[tokio::test]
1222 async fn test_stop_phase_graceful_shutdown_completes_before_deadline() {
1223 use std::sync::atomic::AtomicBool;
1224 use std::time::Duration;
1225
1226 struct GracefulGear {
1227 graceful_completed: AtomicBool,
1228 deadline_fired: AtomicBool,
1229 }
1230
1231 #[async_trait::async_trait]
1232 impl Gear for GracefulGear {
1233 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1234 Ok(())
1235 }
1236 }
1237
1238 #[async_trait::async_trait]
1239 impl RunnableCapability for GracefulGear {
1240 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1241 Ok(())
1242 }
1243 async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
1244 tokio::select! {
1246 () = tokio::time::sleep(Duration::from_millis(10)) => {
1247 self.graceful_completed.store(true, Ordering::SeqCst);
1248 }
1249 () = deadline_token.cancelled() => {
1250 self.deadline_fired.store(true, Ordering::SeqCst);
1251 }
1252 }
1253 Ok(())
1254 }
1255 }
1256
1257 let gear = Arc::new(GracefulGear {
1258 graceful_completed: AtomicBool::new(false),
1259 deadline_fired: AtomicBool::new(false),
1260 });
1261
1262 let mut builder = RegistryBuilder::default();
1263 builder.register_core_with_meta("test", &[], gear.clone() as Arc<dyn Gear>);
1264 builder.register_stateful_with_meta("test", gear.clone() as Arc<dyn RunnableCapability>);
1265
1266 let registry = builder.build_topo_sorted().unwrap();
1267 let client_hub = Arc::new(ClientHub::new());
1268 let cancel = CancellationToken::new();
1269 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1270
1271 let runtime = HostRuntime::new(
1273 registry,
1274 config_provider,
1275 DbOptions::None,
1276 client_hub,
1277 cancel.clone(),
1278 Uuid::new_v4(),
1279 None,
1280 )
1281 .with_shutdown_deadline(Duration::from_secs(5));
1282
1283 runtime.run_stop_phase().await.unwrap();
1284
1285 assert!(
1287 gear.graceful_completed.load(Ordering::SeqCst),
1288 "graceful shutdown should complete"
1289 );
1290 assert!(
1292 !gear.deadline_fired.load(Ordering::SeqCst),
1293 "deadline should not fire when graceful shutdown completes quickly"
1294 );
1295 }
1296
1297 #[tokio::test]
1298 async fn test_stop_phase_deadline_fires_for_slow_gear() {
1299 use std::sync::atomic::AtomicBool;
1300 use std::time::Duration;
1301
1302 struct SlowGear {
1303 graceful_completed: AtomicBool,
1304 deadline_fired: AtomicBool,
1305 }
1306
1307 #[async_trait::async_trait]
1308 impl Gear for SlowGear {
1309 async fn init(&self, _ctx: &GearCtx) -> anyhow::Result<()> {
1310 Ok(())
1311 }
1312 }
1313
1314 #[async_trait::async_trait]
1315 impl RunnableCapability for SlowGear {
1316 async fn start(&self, _cancel: CancellationToken) -> anyhow::Result<()> {
1317 Ok(())
1318 }
1319 async fn stop(&self, deadline_token: CancellationToken) -> anyhow::Result<()> {
1320 tokio::select! {
1322 () = tokio::time::sleep(Duration::from_secs(10)) => {
1323 self.graceful_completed.store(true, Ordering::SeqCst);
1324 }
1325 () = deadline_token.cancelled() => {
1326 self.deadline_fired.store(true, Ordering::SeqCst);
1327 }
1328 }
1329 Ok(())
1330 }
1331 }
1332
1333 let gear = Arc::new(SlowGear {
1334 graceful_completed: AtomicBool::new(false),
1335 deadline_fired: AtomicBool::new(false),
1336 });
1337
1338 let mut builder = RegistryBuilder::default();
1339 builder.register_core_with_meta("test", &[], gear.clone() as Arc<dyn Gear>);
1340 builder.register_stateful_with_meta("test", gear.clone() as Arc<dyn RunnableCapability>);
1341
1342 let registry = builder.build_topo_sorted().unwrap();
1343 let client_hub = Arc::new(ClientHub::new());
1344 let cancel = CancellationToken::new();
1345 let config_provider: Arc<dyn ConfigProvider> = Arc::new(EmptyConfigProvider);
1346
1347 let runtime = HostRuntime::new(
1349 registry,
1350 config_provider,
1351 DbOptions::None,
1352 client_hub,
1353 cancel.clone(),
1354 Uuid::new_v4(),
1355 None,
1356 )
1357 .with_shutdown_deadline(Duration::from_millis(100));
1358
1359 runtime.run_stop_phase().await.unwrap();
1360
1361 assert!(
1363 !gear.graceful_completed.load(Ordering::SeqCst),
1364 "graceful shutdown should not complete when deadline fires first"
1365 );
1366 assert!(
1368 gear.deadline_fired.load(Ordering::SeqCst),
1369 "deadline should fire for slow gears"
1370 );
1371 }
1372}