1use crate::{Context, HealthStatus, NodeState, Result};
2use async_trait::async_trait;
3use schemars::JsonSchema;
4use serde::de::DeserializeOwned;
5use serde::{Deserialize, Serialize};
6use std::fmt::Debug;
7use std::sync::Arc;
8use std::time::Duration;
9use tokio::sync::Mutex;
10use tracing::{debug, error, info, warn};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
13#[serde(rename_all = "lowercase")]
14pub enum HealthPriority {
15 Critical,
16
17 #[default]
18 Normal,
19
20 Background,
21}
22
23impl HealthPriority {
24 const fn profile(&self) -> (Duration, u64, Duration) {
30 match self {
31 Self::Critical => (Duration::from_secs(1), 5, Duration::from_secs(3)),
32 Self::Normal => (Duration::from_secs(5), 60, Duration::from_secs(15)),
33 Self::Background => (Duration::from_secs(30), 120, Duration::from_secs(90)),
34 }
35 }
36
37 pub const fn interval(&self) -> Duration {
39 self.profile().0
40 }
41
42 pub const fn ttl_seconds(&self) -> u64 {
47 self.profile().1
48 }
49
50 pub const fn timeout(&self) -> Duration {
55 self.profile().2
56 }
57
58 pub const fn timeout_micros(&self) -> u64 {
60 self.timeout().as_micros() as u64
61 }
62}
63
64impl std::fmt::Display for HealthPriority {
65 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66 match self {
67 Self::Critical => write!(f, "critical"),
68 Self::Normal => write!(f, "normal"),
69 Self::Background => write!(f, "background"),
70 }
71 }
72}
73
74#[derive(Debug, Clone)]
75pub struct HealthReportingConfig {
76 pub enabled: bool,
77
78 pub priority: HealthPriority,
79
80 pub store_history: bool,
81}
82
83impl Default for HealthReportingConfig {
84 fn default() -> Self {
85 Self {
86 enabled: true,
87 priority: HealthPriority::Normal,
88 store_history: true,
89 }
90 }
91}
92
93impl HealthReportingConfig {
94 fn with_priority(priority: HealthPriority, store_history: bool) -> Self {
97 Self {
98 enabled: true,
99 priority,
100 store_history,
101 }
102 }
103
104 pub fn critical() -> Self {
106 Self::with_priority(HealthPriority::Critical, false)
107 }
108
109 pub fn normal() -> Self {
111 Self::default()
112 }
113
114 pub fn background() -> Self {
116 Self::with_priority(HealthPriority::Background, true)
117 }
118
119 pub fn interval(&self) -> Duration {
121 self.priority.interval()
122 }
123
124 pub fn ttl_seconds(&self) -> u64 {
126 self.priority.ttl_seconds()
127 }
128}
129
130pub trait RateConfig {
131 fn rate_hz(&self) -> f32;
133
134 fn create_interval(&self) -> tokio::time::Interval {
139 let duration = Duration::from_secs_f32(1.0 / self.rate_hz());
140 tokio::time::interval(duration)
141 }
142
143 fn period(&self) -> Duration {
145 Duration::from_secs_f32(1.0 / self.rate_hz())
146 }
147}
148
149pub trait HealthStatusExt {
150 fn healthy(message: impl Into<String>) -> Self;
152 fn unhealthy(message: impl Into<String>) -> Self;
154}
155
156impl HealthStatus {
157 fn with_check(healthy: bool, failure_count: u32, message: impl Into<String>) -> Self {
161 Self {
162 healthy,
163 last_check: crate::prelude::now_micros(),
164 failure_count,
165 message: Some(message.into()),
166 }
167 }
168}
169
170impl HealthStatusExt for HealthStatus {
171 fn healthy(message: impl Into<String>) -> Self {
172 Self::with_check(true, 0, message)
173 }
174
175 fn unhealthy(message: impl Into<String>) -> Self {
176 Self::with_check(false, 1, message)
177 }
178}
179
180#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
181#[serde(rename_all = "kebab-case")]
182pub enum NodeType {
183 Imu,
184 Camera,
185 ImageClassifier,
186 ObjectDetector,
187 BehaviorExecutor,
188 LlmCommand,
189 Teleop,
190 Motor,
191 WebsocketBridge,
192 VideoStreamer,
193 Simulator,
194 Speaker,
195 Listener,
196 Rl,
197 InstanceTrack,
198 Diagnostics,
199 Servo,
200 PanTilt,
201}
202
203impl NodeType {
204 pub const ALL: &'static [NodeType] = &[
205 NodeType::Imu,
206 NodeType::Camera,
207 NodeType::ImageClassifier,
208 NodeType::ObjectDetector,
209 NodeType::BehaviorExecutor,
210 NodeType::LlmCommand,
211 NodeType::Teleop,
212 NodeType::Motor,
213 NodeType::WebsocketBridge,
214 NodeType::VideoStreamer,
215 NodeType::Simulator,
216 NodeType::Speaker,
217 NodeType::Listener,
218 NodeType::Rl,
219 NodeType::InstanceTrack,
220 NodeType::Diagnostics,
221 NodeType::Servo,
222 NodeType::PanTilt,
223 ];
224
225 pub fn as_str(&self) -> &'static str {
232 match self {
233 NodeType::Imu => "imu",
234 NodeType::Camera => "camera",
235 NodeType::ImageClassifier => "image-classifier",
236 NodeType::ObjectDetector => "object-detector",
237 NodeType::BehaviorExecutor => "behavior-executor",
238 NodeType::LlmCommand => "llm-command",
239 NodeType::Teleop => "teleop",
240 NodeType::Motor => "motor",
241 NodeType::WebsocketBridge => "websocket-bridge",
242 NodeType::VideoStreamer => "video-streamer",
243 NodeType::Simulator => "simulator",
244 NodeType::Speaker => "speaker",
245 NodeType::Listener => "listener",
246 NodeType::Rl => "rl",
247 NodeType::InstanceTrack => "instance-track",
248 NodeType::Diagnostics => "diagnostics",
249 NodeType::Servo => "servo",
250 NodeType::PanTilt => "pan-tilt",
251 }
252 }
253}
254
255impl std::fmt::Display for NodeType {
256 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
257 write!(f, "{}", self.as_str())
258 }
259}
260
261#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
262#[serde(rename_all = "kebab-case")]
263pub enum NodeTag {
264 Sensor,
265 Inertial,
266 Capture,
267 Ai,
268 Classifier,
269 Detector,
270 Vision,
271 Planner,
272 Control,
273 Nlp,
274 Input,
275 Actuator,
276 Motion,
277 Relay,
278 Websocket,
279 Streamer,
280 Video,
281 Simulator,
282 Logger,
283}
284
285#[async_trait]
286pub trait NodeImpl: Send + Sync + Debug + 'static {
287 type Config: DeserializeOwned + Send + Sync + Debug;
288
289 const NODE_TYPE: Option<NodeType> = None;
290
291 const NODE_TAGS: &'static [NodeTag] = &[];
292
293 const CRATE_NAME: Option<&'static str> = None;
301
302 const DESCRIPTION: &'static str = "";
305
306 const DEFAULT_RUN_TARGET: &'static str = "robot";
309
310 const CATALOG_TAGS: &'static [&'static str] = &[];
314
315 async fn init(config: Self::Config) -> Result<Self>
326 where
327 Self: Sized;
328
329 async fn run(&mut self, ctx: &Context) -> Result<()>;
345
346 async fn shutdown(&mut self) -> Result<()> {
360 Ok(())
361 }
362
363 async fn health_check(&self) -> HealthStatus {
372 HealthStatus {
373 healthy: true,
374 last_check: crate::prelude::now_micros(),
375 failure_count: 0,
376 message: Some("Node is running".to_string()),
377 }
378 }
379
380 fn health_diagnostics(&self) -> Option<std::collections::HashMap<String, String>> {
391 None
392 }
393
394 fn preflight(&self) -> Result<()> {
418 Ok(())
419 }
420}
421
422pub fn fmt_node_with_controller(
429 f: &mut std::fmt::Formatter<'_>,
430 struct_name: &str,
431 config: &impl Debug,
432 controller_name: &str,
433 update_counter: u64,
434) -> std::fmt::Result {
435 f.debug_struct(struct_name)
436 .field("config", config)
437 .field("controller", &controller_name)
438 .field("update_counter", &update_counter)
439 .finish()
440}
441
442#[derive(Clone)]
443pub struct Node<T: NodeImpl> {
444 node: Arc<Mutex<T>>,
445 state: Arc<Mutex<NodeState>>,
446 static_diagnostics: Arc<Option<std::collections::HashMap<String, String>>>,
456}
457
458impl<T: NodeImpl> Node<T> {
459 pub async fn from_config(config: T::Config) -> Result<Self> {
467 info!("Initializing node");
468
469 let node = T::init(config).await?;
470 let static_diagnostics = node.health_diagnostics();
471
472 Ok(Self {
473 node: Arc::new(Mutex::new(node)),
474 state: Arc::new(Mutex::new(NodeState::Starting)),
475 static_diagnostics: Arc::new(static_diagnostics),
476 })
477 }
478
479 pub fn static_diagnostics(&self) -> Option<std::collections::HashMap<String, String>> {
482 (*self.static_diagnostics).clone()
483 }
484
485 pub async fn start(&self, ctx: Context) -> Result<()> {
499 {
500 let node = self.node.lock().await;
501 node.preflight()?;
502 }
503
504 {
505 let mut state = self.state.lock().await;
506 *state = NodeState::Running;
507 info!("Node transitioned to Running state");
508 }
509
510 let result = {
511 let mut node = self.node.lock().await;
512 debug!("Starting node execution");
513 node.run(&ctx).await
514 };
515
516 match result {
517 Ok(()) => {
518 info!("Node completed successfully");
519 self.shutdown().await?;
520 }
521 Err(e) => {
522 error!("Node failed with error: {}", e);
523 {
524 let mut state = self.state.lock().await;
525 *state = NodeState::Crashed;
526 }
527 self.shutdown().await?;
528 return Err(e);
529 }
530 }
531
532 Ok(())
533 }
534
535 pub async fn shutdown(&self) -> Result<()> {
539 info!("Shutting down node");
540
541 {
542 let mut state = self.state.lock().await;
543 *state = NodeState::Stopping;
544 }
545
546 let result = {
547 let mut node = self.node.lock().await;
548 node.shutdown().await
549 };
550
551 {
552 let mut state = self.state.lock().await;
553 *state = NodeState::Stopped;
554 info!("Node stopped");
555 }
556
557 if let Err(e) = result {
558 warn!("Error during shutdown: {}", e);
559 return Err(e);
560 }
561
562 Ok(())
563 }
564
565 pub async fn state(&self) -> NodeState {
567 self.state.lock().await.clone()
568 }
569
570 pub async fn health_check(&self) -> HealthStatus {
572 let node = self.node.lock().await;
573 node.health_check().await
574 }
575
576 pub async fn is_healthy(&self) -> bool {
578 matches!(self.state().await, NodeState::Running | NodeState::Starting)
579 }
580}
581
582pub async fn run_node<T: NodeImpl>(
633 config: T::Config,
634 ctx: Context,
635 health_config: HealthReportingConfig,
636) -> Result<()> {
637 crate::prelude::init_logging();
638
639 #[cfg(feature = "messaging")]
640 crate::log_publisher::attach(ctx.node_id(), ctx.clone());
641
642 let node = Node::<T>::from_config(config).await?;
643
644 #[cfg(feature = "messaging")]
645 {
646 if health_config.enabled {
647 let node_for_health = node.node.clone();
648 let static_diagnostics_for_health = node.static_diagnostics();
649 let ctx_for_health = ctx.clone();
650 let priority = health_config.priority;
651 let report_interval = health_config.interval();
652 let ttl_seconds = health_config.ttl_seconds();
653 let store_history = health_config.store_history;
654
655 tokio::spawn(async move {
656 use crate::health::HealthReporter;
657
658 let reporter = match HealthReporter::with_config_and_type_tags(
659 &ctx_for_health,
660 priority,
661 ttl_seconds,
662 store_history,
663 T::NODE_TYPE,
664 T::NODE_TAGS,
665 )
666 .await
667 {
668 Ok(r) => r,
669 Err(e) => {
670 warn!("Health reporting disabled: {}", e);
671 return;
672 }
673 };
674
675 debug!(
676 "Health reporting enabled: priority={}, interval={:?}, ttl={}s, history={}",
677 priority, report_interval, ttl_seconds, store_history
678 );
679
680 let mut interval = tokio::time::interval(report_interval);
681
682 loop {
683 interval.tick().await;
684 debug!("Health tick - checking node health");
685
686 let msg_status = match node_for_health.try_lock() {
687 Ok(node) => {
688 let status = node.health_check().await;
689 let diagnostics = node.health_diagnostics();
690 crate::messages::system::HealthStatus {
691 timestamp: status.last_check,
692 level: if status.healthy {
693 crate::messages::system::HealthLevel::Ok
694 } else {
695 crate::messages::system::HealthLevel::Error
696 },
697 message: status.message.unwrap_or_else(|| "No message".to_string()),
698 diagnostics,
699 }
700 }
701 Err(_) => {
702 debug!("Node busy, reporting as healthy");
713 crate::messages::system::HealthStatus {
714 timestamp: crate::prelude::now_micros(),
715 level: crate::messages::system::HealthLevel::Ok,
716 message: "Node running (busy)".to_string(),
717 diagnostics: static_diagnostics_for_health.clone(),
718 }
719 }
720 };
721
722 if let Err(e) = reporter.report(msg_status).await {
723 warn!("Failed to report health: {}", e);
724 } else {
725 tracing::trace!("Health report sent successfully");
726 }
727 }
728 });
729 } else {
730 info!("Health reporting disabled");
731 }
732 }
733 #[cfg(not(feature = "messaging"))]
734 let _ = &health_config;
735
736 node.start(ctx).await
737}
738
739#[cfg(all(test, not(feature = "messaging")))]
740mod tests {
741 use super::*;
742 use crate::error::Mecha10Error;
743 use std::sync::atomic::{AtomicBool, Ordering};
744 use std::sync::Arc as StdArc;
745
746 #[derive(Debug, Default, serde::Deserialize)]
747 struct FakeConfig;
748
749 #[derive(Debug)]
750 struct OkNode {
751 ran: StdArc<AtomicBool>,
752 }
753
754 #[async_trait]
755 impl NodeImpl for OkNode {
756 type Config = FakeConfig;
757
758 async fn init(_config: Self::Config) -> Result<Self> {
759 Ok(Self {
760 ran: StdArc::new(AtomicBool::new(false)),
761 })
762 }
763
764 async fn run(&mut self, _ctx: &Context) -> Result<()> {
765 self.ran.store(true, Ordering::SeqCst);
766 Ok(())
767 }
768 }
769
770 #[derive(Debug)]
771 struct RejectingPreflightNode {
772 ran: StdArc<AtomicBool>,
773 }
774
775 #[async_trait]
776 impl NodeImpl for RejectingPreflightNode {
777 type Config = FakeConfig;
778
779 async fn init(_config: Self::Config) -> Result<Self> {
780 Ok(Self {
781 ran: StdArc::new(AtomicBool::new(false)),
782 })
783 }
784
785 async fn run(&mut self, _ctx: &Context) -> Result<()> {
786 self.ran.store(true, Ordering::SeqCst);
787 Ok(())
788 }
789
790 fn preflight(&self) -> Result<()> {
791 Err(Mecha10Error::Configuration("boom: bad node content".to_string()))
792 }
793 }
794
795 #[tokio::test]
796 async fn default_preflight_is_a_noop_and_run_still_executes() {
797 let node = Node::<OkNode>::from_config(FakeConfig).await.unwrap();
798 let ctx = Context::new("test-node-ok").await.unwrap();
799
800 let ran_handle = { node.node.lock().await.ran.clone() };
801
802 node.start(ctx).await.unwrap();
803
804 assert!(ran_handle.load(Ordering::SeqCst), "run() should have executed");
805 }
806
807 #[tokio::test]
808 async fn start_fails_fast_when_preflight_rejects_config_and_never_calls_run() {
809 let node = Node::<RejectingPreflightNode>::from_config(FakeConfig).await.unwrap();
810 let ctx = Context::new("test-node-rejecting").await.unwrap();
811
812 let ran_handle = { node.node.lock().await.ran.clone() };
813
814 let result = node.start(ctx).await;
815
816 assert!(result.is_err());
817 assert!(result.unwrap_err().to_string().contains("boom: bad node content"));
818 assert!(
819 !ran_handle.load(Ordering::SeqCst),
820 "run() must not execute when preflight() rejects the node"
821 );
822 }
823
824 #[test]
825 fn node_type_kebab_case_matches_existing_builtin_node_types_strings() {
826 let expected: &[(NodeType, &str)] = &[
827 (NodeType::Imu, "imu"),
828 (NodeType::Camera, "camera"),
829 (NodeType::ImageClassifier, "image-classifier"),
830 (NodeType::ObjectDetector, "object-detector"),
831 (NodeType::BehaviorExecutor, "behavior-executor"),
832 (NodeType::LlmCommand, "llm-command"),
833 (NodeType::Teleop, "teleop"),
834 (NodeType::Motor, "motor"),
835 (NodeType::WebsocketBridge, "websocket-bridge"),
836 (NodeType::VideoStreamer, "video-streamer"),
837 (NodeType::Simulator, "simulator"),
838 (NodeType::Speaker, "speaker"),
839 (NodeType::Listener, "listener"),
840 (NodeType::Rl, "rl"),
841 (NodeType::InstanceTrack, "instance-track"),
842 (NodeType::Diagnostics, "diagnostics"),
843 ];
844
845 assert_eq!(expected.len(), NodeType::ALL.len());
846
847 for (variant, kebab) in expected {
848 let json = serde_json::to_string(variant).unwrap();
849 assert_eq!(json, format!("\"{}\"", kebab));
850 assert_eq!(variant.as_str(), *kebab);
851
852 let round_tripped: NodeType = serde_json::from_str(&json).unwrap();
853 assert_eq!(round_tripped, *variant);
854 }
855 }
856
857 #[test]
858 fn node_tag_kebab_case_matches_expected_strings() {
859 let expected: &[(NodeTag, &str)] = &[
860 (NodeTag::Sensor, "sensor"),
861 (NodeTag::Inertial, "inertial"),
862 (NodeTag::Capture, "capture"),
863 (NodeTag::Ai, "ai"),
864 (NodeTag::Classifier, "classifier"),
865 (NodeTag::Detector, "detector"),
866 (NodeTag::Vision, "vision"),
867 (NodeTag::Planner, "planner"),
868 (NodeTag::Control, "control"),
869 (NodeTag::Nlp, "nlp"),
870 (NodeTag::Input, "input"),
871 (NodeTag::Actuator, "actuator"),
872 (NodeTag::Motion, "motion"),
873 (NodeTag::Relay, "relay"),
874 (NodeTag::Websocket, "websocket"),
875 (NodeTag::Streamer, "streamer"),
876 (NodeTag::Video, "video"),
877 (NodeTag::Simulator, "simulator"),
878 (NodeTag::Logger, "logger"),
879 ];
880
881 for (variant, kebab) in expected {
882 let json = serde_json::to_string(variant).unwrap();
883 assert_eq!(json, format!("\"{}\"", kebab));
884
885 let round_tripped: NodeTag = serde_json::from_str(&json).unwrap();
886 assert_eq!(round_tripped, *variant);
887 }
888 }
889
890 #[test]
891 fn node_impl_defaults_node_type_to_none_and_tags_to_empty() {
892 assert_eq!(OkNode::NODE_TYPE, None);
893 assert!(OkNode::NODE_TAGS.is_empty());
894 }
895
896 #[test]
897 fn node_impl_defaults_registration_consts() {
898 assert_eq!(OkNode::CRATE_NAME, None);
899 assert_eq!(OkNode::DESCRIPTION, "");
900 assert_eq!(OkNode::DEFAULT_RUN_TARGET, "robot");
901 assert!(OkNode::CATALOG_TAGS.is_empty());
902 }
903}