re-uploading work
This commit is contained in:
27
crates/sensor/.sqlx/query-5ef7e3bc2362b5b3da420e3913eaf3071100ab24f564b82799003ae9e27a6aed.json
generated
Normal file
27
crates/sensor/.sqlx/query-5ef7e3bc2362b5b3da420e3913eaf3071100ab24f564b82799003ae9e27a6aed.json
generated
Normal file
@@ -0,0 +1,27 @@
|
||||
{
|
||||
"db_name": "PostgreSQL",
|
||||
"query": "\n INSERT INTO event\n (trigger, trigger_ref, config, payload, source, source_ref)\n VALUES ($1, $2, $3, $4, $5, $6)\n RETURNING id\n ",
|
||||
"describe": {
|
||||
"columns": [
|
||||
{
|
||||
"ordinal": 0,
|
||||
"name": "id",
|
||||
"type_info": "Int8"
|
||||
}
|
||||
],
|
||||
"parameters": {
|
||||
"Left": [
|
||||
"Int8",
|
||||
"Text",
|
||||
"Jsonb",
|
||||
"Jsonb",
|
||||
"Int8",
|
||||
"Text"
|
||||
]
|
||||
},
|
||||
"nullable": [
|
||||
false
|
||||
]
|
||||
},
|
||||
"hash": "5ef7e3bc2362b5b3da420e3913eaf3071100ab24f564b82799003ae9e27a6aed"
|
||||
}
|
||||
71
crates/sensor/.sqlx/query-ddfd543c0ef1e25e0a2e5830faf285b02903258ef874b82bd0916b98114b8023.json
generated
Normal file
71
crates/sensor/.sqlx/query-ddfd543c0ef1e25e0a2e5830faf285b02903258ef874b82bd0916b98114b8023.json
generated
Normal file
@@ -0,0 +1,71 @@
|
||||
{
|
||||
"db_name": "PostgreSQL",
|
||||
"query": "\n SELECT\n id,\n trigger,\n trigger_ref,\n config,\n payload,\n source,\n source_ref,\n created,\n updated\n FROM event\n WHERE trigger_ref = $1\n ORDER BY created DESC\n LIMIT $2\n ",
|
||||
"describe": {
|
||||
"columns": [
|
||||
{
|
||||
"ordinal": 0,
|
||||
"name": "id",
|
||||
"type_info": "Int8"
|
||||
},
|
||||
{
|
||||
"ordinal": 1,
|
||||
"name": "trigger",
|
||||
"type_info": "Int8"
|
||||
},
|
||||
{
|
||||
"ordinal": 2,
|
||||
"name": "trigger_ref",
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"ordinal": 3,
|
||||
"name": "config",
|
||||
"type_info": "Jsonb"
|
||||
},
|
||||
{
|
||||
"ordinal": 4,
|
||||
"name": "payload",
|
||||
"type_info": "Jsonb"
|
||||
},
|
||||
{
|
||||
"ordinal": 5,
|
||||
"name": "source",
|
||||
"type_info": "Int8"
|
||||
},
|
||||
{
|
||||
"ordinal": 6,
|
||||
"name": "source_ref",
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"ordinal": 7,
|
||||
"name": "created",
|
||||
"type_info": "Timestamptz"
|
||||
},
|
||||
{
|
||||
"ordinal": 8,
|
||||
"name": "updated",
|
||||
"type_info": "Timestamptz"
|
||||
}
|
||||
],
|
||||
"parameters": {
|
||||
"Left": [
|
||||
"Text",
|
||||
"Int8"
|
||||
]
|
||||
},
|
||||
"nullable": [
|
||||
false,
|
||||
true,
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
false,
|
||||
false
|
||||
]
|
||||
},
|
||||
"hash": "ddfd543c0ef1e25e0a2e5830faf285b02903258ef874b82bd0916b98114b8023"
|
||||
}
|
||||
70
crates/sensor/.sqlx/query-e65380ade25b997bb41733d1b4fa510f1946fd2aa42de1c6f14e36a3f11b2aa1.json
generated
Normal file
70
crates/sensor/.sqlx/query-e65380ade25b997bb41733d1b4fa510f1946fd2aa42de1c6f14e36a3f11b2aa1.json
generated
Normal file
@@ -0,0 +1,70 @@
|
||||
{
|
||||
"db_name": "PostgreSQL",
|
||||
"query": "\n SELECT\n id,\n trigger,\n trigger_ref,\n config,\n payload,\n source,\n source_ref,\n created,\n updated\n FROM event\n WHERE id = $1\n ",
|
||||
"describe": {
|
||||
"columns": [
|
||||
{
|
||||
"ordinal": 0,
|
||||
"name": "id",
|
||||
"type_info": "Int8"
|
||||
},
|
||||
{
|
||||
"ordinal": 1,
|
||||
"name": "trigger",
|
||||
"type_info": "Int8"
|
||||
},
|
||||
{
|
||||
"ordinal": 2,
|
||||
"name": "trigger_ref",
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"ordinal": 3,
|
||||
"name": "config",
|
||||
"type_info": "Jsonb"
|
||||
},
|
||||
{
|
||||
"ordinal": 4,
|
||||
"name": "payload",
|
||||
"type_info": "Jsonb"
|
||||
},
|
||||
{
|
||||
"ordinal": 5,
|
||||
"name": "source",
|
||||
"type_info": "Int8"
|
||||
},
|
||||
{
|
||||
"ordinal": 6,
|
||||
"name": "source_ref",
|
||||
"type_info": "Text"
|
||||
},
|
||||
{
|
||||
"ordinal": 7,
|
||||
"name": "created",
|
||||
"type_info": "Timestamptz"
|
||||
},
|
||||
{
|
||||
"ordinal": 8,
|
||||
"name": "updated",
|
||||
"type_info": "Timestamptz"
|
||||
}
|
||||
],
|
||||
"parameters": {
|
||||
"Left": [
|
||||
"Int8"
|
||||
]
|
||||
},
|
||||
"nullable": [
|
||||
false,
|
||||
true,
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
false,
|
||||
false
|
||||
]
|
||||
},
|
||||
"hash": "e65380ade25b997bb41733d1b4fa510f1946fd2aa42de1c6f14e36a3f11b2aa1"
|
||||
}
|
||||
35
crates/sensor/Cargo.toml
Normal file
35
crates/sensor/Cargo.toml
Normal file
@@ -0,0 +1,35 @@
|
||||
[package]
|
||||
name = "attune-sensor"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
authors.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[lib]
|
||||
name = "attune_sensor"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[[bin]]
|
||||
name = "attune-sensor"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
attune-common = { path = "../common" }
|
||||
tokio = { workspace = true }
|
||||
sqlx = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tracing-subscriber = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
config = { workspace = true }
|
||||
|
||||
chrono = { workspace = true }
|
||||
clap = { workspace = true }
|
||||
lapin = { workspace = true }
|
||||
regex = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
cron = "0.15"
|
||||
reqwest = { workspace = true }
|
||||
hostname = "0.4"
|
||||
141
crates/sensor/src/api_client/mod.rs
Normal file
141
crates/sensor/src/api_client/mod.rs
Normal file
@@ -0,0 +1,141 @@
|
||||
//! API Client for Sensor Service
|
||||
//!
|
||||
//! This module provides an HTTP client for the sensor service to communicate
|
||||
//! with the Attune API for token provisioning and other operations.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// API client for sensor service
|
||||
#[derive(Clone)]
|
||||
pub struct ApiClient {
|
||||
base_url: String,
|
||||
client: Client,
|
||||
/// Optional admin token for authentication (if available)
|
||||
admin_token: Option<String>,
|
||||
}
|
||||
|
||||
/// Request to create a sensor token
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CreateSensorTokenRequest {
|
||||
pub sensor_ref: String,
|
||||
pub trigger_types: Vec<String>,
|
||||
pub ttl_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
/// Response from sensor token creation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SensorTokenResponse {
|
||||
pub identity_id: i64,
|
||||
pub sensor_ref: String,
|
||||
pub token: String,
|
||||
pub expires_at: String,
|
||||
pub trigger_types: Vec<String>,
|
||||
}
|
||||
|
||||
/// Wrapper for API responses
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ApiResponse<T> {
|
||||
pub data: T,
|
||||
}
|
||||
|
||||
impl ApiClient {
|
||||
/// Create a new API client
|
||||
pub fn new(base_url: String, admin_token: Option<String>) -> Self {
|
||||
Self {
|
||||
base_url,
|
||||
client: Client::new(),
|
||||
admin_token,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a sensor token via the API
|
||||
///
|
||||
/// This is used internally by the sensor service to provision tokens
|
||||
/// for standalone sensors when they are started.
|
||||
pub async fn create_sensor_token(
|
||||
&self,
|
||||
sensor_ref: &str,
|
||||
trigger_types: Vec<String>,
|
||||
ttl_seconds: Option<i64>,
|
||||
) -> Result<SensorTokenResponse> {
|
||||
let url = format!("{}/auth/internal/sensor-token", self.base_url);
|
||||
|
||||
let request = CreateSensorTokenRequest {
|
||||
sensor_ref: sensor_ref.to_string(),
|
||||
trigger_types,
|
||||
ttl_seconds,
|
||||
};
|
||||
|
||||
let mut req = self.client.post(&url).json(&request);
|
||||
|
||||
// Add authorization header if admin token is available
|
||||
if let Some(token) = &self.admin_token {
|
||||
req = req.header("Authorization", format!("Bearer {}", token));
|
||||
}
|
||||
|
||||
let response = req
|
||||
.send()
|
||||
.await
|
||||
.context("Failed to send sensor token creation request")?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(anyhow::anyhow!(
|
||||
"API request failed with status {}: {}",
|
||||
status,
|
||||
body
|
||||
));
|
||||
}
|
||||
|
||||
let api_response: ApiResponse<SensorTokenResponse> = response
|
||||
.json()
|
||||
.await
|
||||
.context("Failed to parse sensor token response")?;
|
||||
|
||||
Ok(api_response.data)
|
||||
}
|
||||
|
||||
/// Health check endpoint
|
||||
pub async fn health_check(&self) -> Result<()> {
|
||||
let url = format!("{}/health", self.base_url);
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.get(&url)
|
||||
.send()
|
||||
.await
|
||||
.context("Failed to send health check request")?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(anyhow::anyhow!(
|
||||
"Health check failed with status: {}",
|
||||
response.status()
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_api_client_creation() {
|
||||
let client = ApiClient::new("http://localhost:8080".to_string(), None);
|
||||
assert_eq!(client.base_url, "http://localhost:8080");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_api_client_with_token() {
|
||||
let client = ApiClient::new(
|
||||
"http://localhost:8080".to_string(),
|
||||
Some("test_token".to_string()),
|
||||
);
|
||||
assert_eq!(client.admin_token, Some("test_token".to_string()));
|
||||
}
|
||||
}
|
||||
17
crates/sensor/src/lib.rs
Normal file
17
crates/sensor/src/lib.rs
Normal file
@@ -0,0 +1,17 @@
|
||||
//! Attune Sensor Service Library
|
||||
//!
|
||||
//! This library provides the core functionality for the Attune Sensor Service,
|
||||
//! including event generation, rule matching, and template resolution.
|
||||
|
||||
pub mod api_client;
|
||||
pub mod rule_lifecycle_listener;
|
||||
pub mod sensor_manager;
|
||||
pub mod sensor_worker_registration;
|
||||
pub mod service;
|
||||
pub mod template_resolver;
|
||||
|
||||
// Re-export commonly used types
|
||||
pub use rule_lifecycle_listener::RuleLifecycleListener;
|
||||
pub use sensor_worker_registration::SensorWorkerRegistration;
|
||||
pub use service::SensorService;
|
||||
pub use template_resolver::{resolve_templates, TemplateContext};
|
||||
129
crates/sensor/src/main.rs
Normal file
129
crates/sensor/src/main.rs
Normal file
@@ -0,0 +1,129 @@
|
||||
//! Attune Sensor Service
|
||||
//!
|
||||
//! The Sensor Service monitors for trigger conditions and generates events.
|
||||
//! It executes custom sensor code, manages sensor lifecycle, and publishes
|
||||
//! events to the message queue for rule matching and enforcement creation.
|
||||
|
||||
use anyhow::Result;
|
||||
use attune_common::config::Config;
|
||||
use attune_sensor::service::SensorService;
|
||||
use clap::Parser;
|
||||
use tracing::{error, info};
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "attune-sensor")]
|
||||
#[command(about = "Attune Sensor Service - Event monitoring and generation", long_about = None)]
|
||||
struct Args {
|
||||
/// Path to configuration file
|
||||
#[arg(short, long)]
|
||||
config: Option<String>,
|
||||
|
||||
/// Log level (trace, debug, info, warn, error)
|
||||
#[arg(short, long, default_value = "info")]
|
||||
log_level: String,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
|
||||
// Initialize tracing with specified log level
|
||||
let log_level = args.log_level.parse().unwrap_or(tracing::Level::INFO);
|
||||
tracing_subscriber::fmt()
|
||||
.with_max_level(log_level)
|
||||
.with_target(false)
|
||||
.with_thread_ids(true)
|
||||
.with_file(true)
|
||||
.with_line_number(true)
|
||||
.init();
|
||||
|
||||
info!("Starting Attune Sensor Service");
|
||||
info!("Version: {}", env!("CARGO_PKG_VERSION"));
|
||||
|
||||
// Load configuration
|
||||
if let Some(config_path) = args.config {
|
||||
info!("Loading configuration from: {}", config_path);
|
||||
std::env::set_var("ATTUNE_CONFIG", config_path);
|
||||
}
|
||||
|
||||
let config = Config::load()?;
|
||||
config.validate()?;
|
||||
|
||||
info!("Configuration loaded successfully");
|
||||
info!("Environment: {}", config.environment);
|
||||
info!("Database: {}", mask_connection_string(&config.database.url));
|
||||
if let Some(ref mq_config) = config.message_queue {
|
||||
info!("Message Queue: {}", mask_connection_string(&mq_config.url));
|
||||
}
|
||||
|
||||
// Create sensor service
|
||||
let service = SensorService::new(config).await?;
|
||||
|
||||
info!("Sensor Service initialized successfully");
|
||||
|
||||
// Set up graceful shutdown handler
|
||||
let service_clone = service.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = tokio::signal::ctrl_c().await {
|
||||
error!("Failed to listen for shutdown signal: {}", e);
|
||||
} else {
|
||||
info!("Shutdown signal received");
|
||||
if let Err(e) = service_clone.stop().await {
|
||||
error!("Error during shutdown: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Start the service
|
||||
info!("Starting Sensor Service components...");
|
||||
if let Err(e) = service.start().await {
|
||||
error!("Sensor Service error: {}", e);
|
||||
return Err(e);
|
||||
}
|
||||
|
||||
info!("Sensor Service has shut down gracefully");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Mask sensitive parts of connection strings for logging
|
||||
fn mask_connection_string(url: &str) -> String {
|
||||
if let Some(at_pos) = url.find('@') {
|
||||
if let Some(proto_end) = url.find("://") {
|
||||
let protocol = &url[..proto_end + 3];
|
||||
let host_and_path = &url[at_pos..];
|
||||
return format!("{}***:***{}", protocol, host_and_path);
|
||||
}
|
||||
}
|
||||
"***:***@***".to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_mask_connection_string() {
|
||||
let url = "postgresql://user:password@localhost:5432/attune";
|
||||
let masked = mask_connection_string(url);
|
||||
assert!(!masked.contains("user"));
|
||||
assert!(!masked.contains("password"));
|
||||
assert!(masked.contains("@localhost"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mask_connection_string_no_credentials() {
|
||||
let url = "postgresql://localhost:5432/attune";
|
||||
let masked = mask_connection_string(url);
|
||||
assert_eq!(masked, "***:***@***");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mask_rabbitmq_connection() {
|
||||
let url = "amqp://admin:secret@rabbitmq:5672/%2F";
|
||||
let masked = mask_connection_string(url);
|
||||
assert!(!masked.contains("admin"));
|
||||
assert!(!masked.contains("secret"));
|
||||
assert!(masked.contains("@rabbitmq"));
|
||||
}
|
||||
}
|
||||
293
crates/sensor/src/rule_lifecycle_listener.rs
Normal file
293
crates/sensor/src/rule_lifecycle_listener.rs
Normal file
@@ -0,0 +1,293 @@
|
||||
//! Rule Lifecycle Listener
|
||||
//!
|
||||
//! This module listens for rule lifecycle events (created, enabled, disabled)
|
||||
//! and notifies the sensor manager to update sensor process lifecycles accordingly.
|
||||
|
||||
use anyhow::Result;
|
||||
use attune_common::mq::{
|
||||
Connection, Consumer, ConsumerConfig, MessageEnvelope, MessageType, RuleCreatedPayload,
|
||||
RuleDisabledPayload, RuleEnabledPayload,
|
||||
};
|
||||
use serde_json::Value as JsonValue;
|
||||
use sqlx::PgPool;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use crate::sensor_manager::SensorManager;
|
||||
|
||||
/// Rule lifecycle listener
|
||||
pub struct RuleLifecycleListener {
|
||||
db: PgPool,
|
||||
connection: Connection,
|
||||
sensor_manager: Arc<SensorManager>,
|
||||
consumer: Arc<RwLock<Option<Consumer>>>,
|
||||
}
|
||||
|
||||
impl RuleLifecycleListener {
|
||||
/// Create a new rule lifecycle listener
|
||||
pub fn new(db: PgPool, connection: Connection, sensor_manager: Arc<SensorManager>) -> Self {
|
||||
Self {
|
||||
db,
|
||||
connection,
|
||||
sensor_manager,
|
||||
consumer: Arc::new(RwLock::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Start listening for rule lifecycle events
|
||||
pub async fn start(&self) -> Result<()> {
|
||||
info!("Starting rule lifecycle listener");
|
||||
|
||||
// Create consumer configuration
|
||||
let consumer_config = ConsumerConfig {
|
||||
queue: "attune.rules.lifecycle.queue".to_string(),
|
||||
tag: "sensor-rule-lifecycle".to_string(),
|
||||
prefetch_count: 10,
|
||||
auto_ack: false,
|
||||
exclusive: false,
|
||||
};
|
||||
|
||||
// Create consumer
|
||||
let consumer = Consumer::new(&self.connection, consumer_config).await?;
|
||||
|
||||
// Bind queue to exchange with routing keys
|
||||
let exchange = "attune.events";
|
||||
let queue = "attune.rules.lifecycle.queue";
|
||||
|
||||
// Declare queue
|
||||
consumer
|
||||
.channel()
|
||||
.queue_declare(
|
||||
queue,
|
||||
lapin::options::QueueDeclareOptions {
|
||||
durable: true,
|
||||
exclusive: false,
|
||||
auto_delete: false,
|
||||
..Default::default()
|
||||
},
|
||||
lapin::types::FieldTable::default(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Bind to routing keys
|
||||
for routing_key in &["rule.created", "rule.enabled", "rule.disabled"] {
|
||||
consumer
|
||||
.channel()
|
||||
.queue_bind(
|
||||
queue,
|
||||
exchange,
|
||||
routing_key,
|
||||
lapin::options::QueueBindOptions::default(),
|
||||
lapin::types::FieldTable::default(),
|
||||
)
|
||||
.await?;
|
||||
info!(
|
||||
"Bound queue {} to exchange {} with routing key {}",
|
||||
queue, exchange, routing_key
|
||||
);
|
||||
}
|
||||
|
||||
// Store consumer
|
||||
*self.consumer.write().await = Some(consumer);
|
||||
|
||||
// Clone self for async handler
|
||||
let db = self.db.clone();
|
||||
let sensor_manager = self.sensor_manager.clone();
|
||||
let consumer_ref = self.consumer.clone();
|
||||
|
||||
// Start consuming messages
|
||||
tokio::spawn(async move {
|
||||
// Get consumer from the Arc<RwLock<Option<Consumer>>>
|
||||
let consumer_guard = consumer_ref.read().await;
|
||||
if let Some(consumer) = consumer_guard.as_ref() {
|
||||
let result = consumer
|
||||
.consume_with_handler::<JsonValue, _, _>(move |envelope| {
|
||||
let db = db.clone();
|
||||
let sensor_manager = sensor_manager.clone();
|
||||
|
||||
async move {
|
||||
if let Err(e) =
|
||||
Self::handle_message(&db, &sensor_manager, envelope).await
|
||||
{
|
||||
error!("Failed to handle rule lifecycle message: {}", e);
|
||||
return Err(attune_common::mq::MqError::Other(format!(
|
||||
"Handler error: {}",
|
||||
e
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
if let Err(e) = result {
|
||||
error!("Rule lifecycle listener stopped with error: {}", e);
|
||||
} else {
|
||||
info!("Rule lifecycle listener stopped");
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
info!("Rule lifecycle listener started");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Stop the listener
|
||||
pub async fn stop(&self) -> Result<()> {
|
||||
info!("Stopping rule lifecycle listener");
|
||||
|
||||
if let Some(consumer) = self.consumer.write().await.take() {
|
||||
// Consumer will be dropped and connection closed
|
||||
drop(consumer);
|
||||
}
|
||||
|
||||
info!("Rule lifecycle listener stopped");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Handle a rule lifecycle message
|
||||
async fn handle_message(
|
||||
db: &PgPool,
|
||||
sensor_manager: &Arc<SensorManager>,
|
||||
envelope: MessageEnvelope<JsonValue>,
|
||||
) -> Result<()> {
|
||||
match envelope.message_type {
|
||||
MessageType::RuleCreated => {
|
||||
let payload: RuleCreatedPayload = serde_json::from_value(envelope.payload)?;
|
||||
Self::handle_rule_created(db, sensor_manager, payload).await?;
|
||||
}
|
||||
MessageType::RuleEnabled => {
|
||||
let payload: RuleEnabledPayload = serde_json::from_value(envelope.payload)?;
|
||||
Self::handle_rule_enabled(db, sensor_manager, payload).await?;
|
||||
}
|
||||
MessageType::RuleDisabled => {
|
||||
let payload: RuleDisabledPayload = serde_json::from_value(envelope.payload)?;
|
||||
Self::handle_rule_disabled(sensor_manager, db, payload).await?;
|
||||
}
|
||||
_ => {
|
||||
warn!("Unexpected message type: {:?}", envelope.message_type);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Handle rule created event
|
||||
async fn handle_rule_created(
|
||||
_db: &PgPool,
|
||||
sensor_manager: &Arc<SensorManager>,
|
||||
payload: RuleCreatedPayload,
|
||||
) -> Result<()> {
|
||||
info!(
|
||||
"Handling RuleCreated: rule={}, trigger={}",
|
||||
payload.rule_ref, payload.trigger_ref
|
||||
);
|
||||
|
||||
// Notify sensor manager about rule change (may need to start sensors)
|
||||
if let Some(trigger_id) = payload.trigger_id {
|
||||
if let Err(e) = sensor_manager.handle_rule_change(trigger_id).await {
|
||||
error!(
|
||||
"Failed to handle sensor lifecycle for trigger {}: {}",
|
||||
trigger_id, e
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Handle rule enabled event
|
||||
async fn handle_rule_enabled(
|
||||
db: &PgPool,
|
||||
sensor_manager: &Arc<SensorManager>,
|
||||
payload: RuleEnabledPayload,
|
||||
) -> Result<()> {
|
||||
info!(
|
||||
"Handling RuleEnabled: rule={}, trigger={}",
|
||||
payload.rule_ref, payload.trigger_ref
|
||||
);
|
||||
|
||||
// Fetch trigger_id from database
|
||||
let trigger_id = match Self::get_trigger_id_for_rule(db, payload.rule_id).await {
|
||||
Ok(Some(id)) => id,
|
||||
Ok(None) => {
|
||||
warn!("Trigger not found for rule {}", payload.rule_id);
|
||||
return Ok(());
|
||||
}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"Failed to fetch trigger for rule {}: {}",
|
||||
payload.rule_id, e
|
||||
);
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
// Notify sensor manager about rule change (may need to start sensors)
|
||||
if let Err(e) = sensor_manager.handle_rule_change(trigger_id).await {
|
||||
error!(
|
||||
"Failed to handle sensor lifecycle for trigger {}: {}",
|
||||
trigger_id, e
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Handle rule disabled event
|
||||
async fn handle_rule_disabled(
|
||||
sensor_manager: &Arc<SensorManager>,
|
||||
db: &PgPool,
|
||||
payload: RuleDisabledPayload,
|
||||
) -> Result<()> {
|
||||
info!(
|
||||
"Handling RuleDisabled: rule={}, trigger={}",
|
||||
payload.rule_ref, payload.trigger_ref
|
||||
);
|
||||
|
||||
// Fetch trigger_id from database
|
||||
let trigger_id = match Self::get_trigger_id_for_rule(db, payload.rule_id).await {
|
||||
Ok(Some(id)) => id,
|
||||
Ok(None) => {
|
||||
warn!("Trigger not found for rule {}", payload.rule_id);
|
||||
return Ok(());
|
||||
}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"Failed to fetch trigger for rule {}: {}",
|
||||
payload.rule_id, e
|
||||
);
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
// Notify sensor manager about rule change (may need to stop sensors)
|
||||
if let Err(e) = sensor_manager.handle_rule_change(trigger_id).await {
|
||||
error!(
|
||||
"Failed to handle sensor lifecycle for trigger {}: {}",
|
||||
trigger_id, e
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Helper function to get trigger_id for a rule
|
||||
async fn get_trigger_id_for_rule(db: &PgPool, rule_id: i64) -> Result<Option<i64>> {
|
||||
let trigger_id = sqlx::query_scalar::<_, i64>(
|
||||
r#"
|
||||
SELECT trigger
|
||||
FROM rule
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(rule_id)
|
||||
.fetch_optional(db)
|
||||
.await?;
|
||||
|
||||
Ok(trigger_id)
|
||||
}
|
||||
}
|
||||
650
crates/sensor/src/sensor_manager.rs
Normal file
650
crates/sensor/src/sensor_manager.rs
Normal file
@@ -0,0 +1,650 @@
|
||||
//! Sensor Manager
|
||||
//!
|
||||
//! Manages the lifecycle of standalone sensor processes including loading,
|
||||
//! starting, stopping, and monitoring sensor instances.
|
||||
//!
|
||||
//! All sensors are independent processes that communicate with the API
|
||||
//! to create events. The sensor manager is responsible for:
|
||||
//! - Starting sensor processes when rules become active
|
||||
//! - Stopping sensor processes when no rules need them
|
||||
//! - Provisioning authentication tokens for sensor processes
|
||||
//! - Monitoring sensor health and restarting failed sensors
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use attune_common::models::{Id, Sensor, Trigger};
|
||||
use attune_common::repositories::{FindById, List};
|
||||
|
||||
use sqlx::{PgPool, Row};
|
||||
use std::collections::HashMap;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::{AsyncBufReadExt, BufReader};
|
||||
use tokio::process::{Child, Command};
|
||||
use tokio::sync::RwLock;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio::time::{interval, Duration};
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
use crate::api_client::ApiClient;
|
||||
|
||||
/// Sensor manager that coordinates all sensor instances
|
||||
#[derive(Clone)]
|
||||
pub struct SensorManager {
|
||||
inner: Arc<SensorManagerInner>,
|
||||
}
|
||||
|
||||
struct SensorManagerInner {
|
||||
db: PgPool,
|
||||
sensors: Arc<RwLock<HashMap<Id, SensorInstance>>>,
|
||||
running: Arc<RwLock<bool>>,
|
||||
packs_base_dir: String,
|
||||
api_client: ApiClient,
|
||||
api_url: String,
|
||||
mq_url: String,
|
||||
}
|
||||
|
||||
impl SensorManager {
|
||||
/// Create a new sensor manager
|
||||
pub fn new(db: PgPool) -> Self {
|
||||
// Get packs base directory from config or default
|
||||
let packs_base_dir =
|
||||
std::env::var("ATTUNE_PACKS_BASE_DIR").unwrap_or_else(|_| "./packs".to_string());
|
||||
|
||||
// Get API URL from config or default
|
||||
let api_url =
|
||||
std::env::var("ATTUNE_API_URL").unwrap_or_else(|_| "http://127.0.0.1:8080".to_string());
|
||||
|
||||
// Get MQ URL from config or default
|
||||
let mq_url = std::env::var("ATTUNE_MQ_URL")
|
||||
.unwrap_or_else(|_| "amqp://guest:guest@localhost:5672".to_string());
|
||||
|
||||
// Create API client for token provisioning (no admin token - uses internal endpoint)
|
||||
let api_client = ApiClient::new(api_url.clone(), None);
|
||||
|
||||
Self {
|
||||
inner: Arc::new(SensorManagerInner {
|
||||
db,
|
||||
sensors: Arc::new(RwLock::new(HashMap::new())),
|
||||
running: Arc::new(RwLock::new(false)),
|
||||
packs_base_dir,
|
||||
api_client,
|
||||
api_url,
|
||||
mq_url,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the sensor manager
|
||||
pub async fn start(&self) -> Result<()> {
|
||||
info!("Starting sensor manager");
|
||||
|
||||
// Mark as running
|
||||
*self.inner.running.write().await = true;
|
||||
|
||||
// Load and start all enabled sensors with active rules
|
||||
let sensors = self.load_enabled_sensors().await?;
|
||||
info!("Loaded {} enabled sensor(s)", sensors.len());
|
||||
|
||||
for sensor in sensors {
|
||||
// Only start sensors that have active rules
|
||||
match self.has_active_rules(sensor.trigger).await {
|
||||
Ok(true) => {
|
||||
let count = self
|
||||
.get_active_rule_count(sensor.trigger)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
info!(
|
||||
"Starting sensor {} - has {} active rule(s)",
|
||||
sensor.r#ref, count
|
||||
);
|
||||
if let Err(e) = self.start_sensor(sensor).await {
|
||||
error!("Failed to start sensor: {}", e);
|
||||
}
|
||||
}
|
||||
Ok(false) => {
|
||||
info!("Skipping sensor {} - no active rules", sensor.r#ref);
|
||||
}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"Failed to check active rules for sensor {}: {}",
|
||||
sensor.r#ref, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Start monitoring loop
|
||||
let manager = self.clone();
|
||||
tokio::spawn(async move {
|
||||
manager.monitoring_loop().await;
|
||||
});
|
||||
|
||||
info!("Sensor manager started");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Stop the sensor manager
|
||||
pub async fn stop(&self) -> Result<()> {
|
||||
info!("Stopping sensor manager");
|
||||
|
||||
// Mark as not running
|
||||
*self.inner.running.write().await = false;
|
||||
|
||||
// Collect sensor IDs to stop
|
||||
let sensor_ids: Vec<Id> = self.inner.sensors.read().await.keys().copied().collect();
|
||||
|
||||
// Stop all sensors
|
||||
for sensor_id in sensor_ids {
|
||||
info!("Stopping sensor {}", sensor_id);
|
||||
if let Err(e) = self.stop_sensor(sensor_id).await {
|
||||
error!("Failed to stop sensor {}: {}", sensor_id, e);
|
||||
}
|
||||
}
|
||||
|
||||
info!("Sensor manager stopped");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load all enabled sensors from the database
|
||||
async fn load_enabled_sensors(&self) -> Result<Vec<Sensor>> {
|
||||
use attune_common::repositories::SensorRepository;
|
||||
|
||||
let all_sensors = SensorRepository::list(&self.inner.db).await?;
|
||||
let enabled_sensors: Vec<Sensor> = all_sensors.into_iter().filter(|s| s.enabled).collect();
|
||||
Ok(enabled_sensors)
|
||||
}
|
||||
|
||||
/// Start a sensor instance
|
||||
async fn start_sensor(&self, sensor: Sensor) -> Result<()> {
|
||||
info!("Starting sensor {} ({})", sensor.r#ref, sensor.id);
|
||||
|
||||
// Load trigger information
|
||||
let trigger = self.load_trigger(sensor.trigger).await?;
|
||||
|
||||
// All sensors are now standalone processes
|
||||
let instance = self
|
||||
.start_standalone_sensor(sensor.clone(), trigger)
|
||||
.await?;
|
||||
|
||||
// Store instance
|
||||
self.inner.sensors.write().await.insert(sensor.id, instance);
|
||||
|
||||
info!("Sensor {} started successfully", sensor.r#ref);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Start a standalone sensor with token provisioning
|
||||
async fn start_standalone_sensor(
|
||||
&self,
|
||||
sensor: Sensor,
|
||||
trigger: Trigger,
|
||||
) -> Result<SensorInstance> {
|
||||
info!("Starting standalone sensor: {}", sensor.r#ref);
|
||||
|
||||
// Get trigger types
|
||||
let trigger_types = vec![trigger.r#ref.clone()];
|
||||
|
||||
// Provision sensor token via API
|
||||
info!("Provisioning token for sensor: {}", sensor.r#ref);
|
||||
let token_response = self
|
||||
.inner
|
||||
.api_client
|
||||
.create_sensor_token(&sensor.r#ref, trigger_types, Some(86400))
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to provision sensor token: {}", e))?;
|
||||
|
||||
info!(
|
||||
"Token provisioned for sensor {} (expires: {})",
|
||||
sensor.r#ref, token_response.expires_at
|
||||
);
|
||||
|
||||
// Build sensor script path
|
||||
let pack_ref = sensor
|
||||
.pack_ref
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow!("Sensor {} has no pack_ref", sensor.r#ref))?;
|
||||
|
||||
let sensor_script = format!(
|
||||
"{}/{}/sensors/{}",
|
||||
self.inner.packs_base_dir, pack_ref, sensor.entrypoint
|
||||
);
|
||||
|
||||
info!(
|
||||
"TRACE: Before fetching trigger instances for sensor {}",
|
||||
sensor.r#ref
|
||||
);
|
||||
info!("Starting standalone sensor process: {}", sensor_script);
|
||||
|
||||
// Fetch trigger instances (enabled rules with their trigger params)
|
||||
info!(
|
||||
"About to fetch trigger instances for sensor {} (trigger_id: {})",
|
||||
sensor.r#ref, sensor.trigger
|
||||
);
|
||||
let trigger_instances = match self.fetch_trigger_instances(sensor.trigger).await {
|
||||
Ok(instances) => {
|
||||
info!(
|
||||
"Fetched {} trigger instance(s) for sensor {}",
|
||||
instances.len(),
|
||||
sensor.r#ref
|
||||
);
|
||||
instances
|
||||
}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"Failed to fetch trigger instances for sensor {}: {}",
|
||||
sensor.r#ref, e
|
||||
);
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
let trigger_instances_json = serde_json::to_string(&trigger_instances)
|
||||
.map_err(|e| anyhow!("Failed to serialize trigger instances: {}", e))?;
|
||||
info!("Trigger instances JSON: {}", trigger_instances_json);
|
||||
|
||||
// Start the standalone sensor with token and configuration
|
||||
// Pass sensor ref (e.g., "core.interval_timer_sensor") for proper identification
|
||||
let mut child = Command::new(&sensor_script)
|
||||
.env("ATTUNE_API_URL", &self.inner.api_url)
|
||||
.env("ATTUNE_API_TOKEN", &token_response.token)
|
||||
.env("ATTUNE_SENSOR_REF", &sensor.r#ref)
|
||||
.env("ATTUNE_SENSOR_TRIGGERS", &trigger_instances_json)
|
||||
.env("ATTUNE_MQ_URL", &self.inner.mq_url)
|
||||
.env("ATTUNE_MQ_EXCHANGE", "attune.events")
|
||||
.env("ATTUNE_LOG_LEVEL", "info")
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.map_err(|e| anyhow!("Failed to start standalone sensor process: {}", e))?;
|
||||
|
||||
// Get stdout and stderr for logging (standalone sensors output JSON logs to stdout)
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| anyhow!("Failed to capture sensor stdout"))?;
|
||||
|
||||
let stderr = child
|
||||
.stderr
|
||||
.take()
|
||||
.ok_or_else(|| anyhow!("Failed to capture sensor stderr"))?;
|
||||
|
||||
// Spawn task to log stdout
|
||||
let sensor_ref_stdout = sensor.r#ref.clone();
|
||||
let stdout_handle = tokio::spawn(async move {
|
||||
let mut reader = BufReader::new(stdout).lines();
|
||||
|
||||
while let Ok(Some(line)) = reader.next_line().await {
|
||||
info!("Sensor {} stdout: {}", sensor_ref_stdout, line);
|
||||
}
|
||||
|
||||
info!("Sensor {} stdout stream closed", sensor_ref_stdout);
|
||||
});
|
||||
|
||||
// Spawn task to log stderr
|
||||
let sensor_ref_stderr = sensor.r#ref.clone();
|
||||
let stderr_handle = tokio::spawn(async move {
|
||||
let mut reader = BufReader::new(stderr).lines();
|
||||
|
||||
while let Ok(Some(line)) = reader.next_line().await {
|
||||
warn!("Sensor {} stderr: {}", sensor_ref_stderr, line);
|
||||
}
|
||||
|
||||
info!("Sensor {} stderr stream closed", sensor_ref_stderr);
|
||||
});
|
||||
|
||||
Ok(SensorInstance::new_standalone(
|
||||
child,
|
||||
stdout_handle,
|
||||
stderr_handle,
|
||||
))
|
||||
}
|
||||
|
||||
/// Load trigger information
|
||||
async fn load_trigger(&self, trigger_id: Id) -> Result<Trigger> {
|
||||
use attune_common::repositories::TriggerRepository;
|
||||
|
||||
TriggerRepository::find_by_id(&self.inner.db, trigger_id)
|
||||
.await?
|
||||
.ok_or_else(|| anyhow!("Trigger {} not found", trigger_id))
|
||||
}
|
||||
|
||||
/// Check if a trigger has any active/enabled rules
|
||||
async fn has_active_rules(&self, trigger_id: Id) -> Result<bool> {
|
||||
let count = sqlx::query_scalar::<_, i64>(
|
||||
r#"
|
||||
SELECT COUNT(*)
|
||||
FROM rule
|
||||
WHERE trigger = $1
|
||||
AND enabled = TRUE
|
||||
"#,
|
||||
)
|
||||
.bind(trigger_id)
|
||||
.fetch_one(&self.inner.db)
|
||||
.await?;
|
||||
|
||||
Ok(count > 0)
|
||||
}
|
||||
|
||||
/// Get count of active rules for a trigger
|
||||
async fn get_active_rule_count(&self, trigger_id: Id) -> Result<i64> {
|
||||
let count = sqlx::query_scalar::<_, i64>(
|
||||
r#"
|
||||
SELECT COUNT(*)
|
||||
FROM rule
|
||||
WHERE trigger = $1
|
||||
AND enabled = TRUE
|
||||
"#,
|
||||
)
|
||||
.bind(trigger_id)
|
||||
.fetch_one(&self.inner.db)
|
||||
.await?;
|
||||
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Fetch trigger instances (enabled rules with their trigger params) for a trigger
|
||||
async fn fetch_trigger_instances(&self, trigger_id: Id) -> Result<Vec<serde_json::Value>> {
|
||||
let rows = sqlx::query(
|
||||
r#"
|
||||
SELECT *
|
||||
FROM rule
|
||||
WHERE trigger = $1
|
||||
AND enabled = TRUE
|
||||
"#,
|
||||
)
|
||||
.bind(trigger_id)
|
||||
.fetch_all(&self.inner.db)
|
||||
.await?;
|
||||
|
||||
info!("Fetched {} rows from rule table", rows.len());
|
||||
|
||||
// Convert to the format expected by timer sensor
|
||||
let trigger_instances: Vec<serde_json::Value> = rows
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
let id: i64 = row.try_get("id").unwrap_or(0);
|
||||
let ref_str: String = row.try_get("ref").unwrap_or_default();
|
||||
let trigger_params: serde_json::Value = row
|
||||
.try_get("trigger_params")
|
||||
.unwrap_or(serde_json::json!({}));
|
||||
|
||||
info!(
|
||||
"Rule ID: {}, Ref: {}, Params: {}",
|
||||
id, ref_str, trigger_params
|
||||
);
|
||||
|
||||
serde_json::json!({
|
||||
"id": id,
|
||||
"ref": ref_str,
|
||||
"config": trigger_params
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(trigger_instances)
|
||||
}
|
||||
|
||||
/// Stop a sensor
|
||||
pub async fn stop_sensor(&self, sensor_id: Id) -> Result<()> {
|
||||
info!("Stopping sensor {}", sensor_id);
|
||||
|
||||
let mut sensors = self.inner.sensors.write().await;
|
||||
|
||||
if let Some(mut instance) = sensors.remove(&sensor_id) {
|
||||
instance.stop().await;
|
||||
info!("Sensor {} stopped", sensor_id);
|
||||
} else {
|
||||
warn!("Sensor {} not found in running instances", sensor_id);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Handle rule changes (created, enabled, disabled)
|
||||
pub async fn handle_rule_change(&self, trigger_id: Id) -> Result<()> {
|
||||
info!("Handling rule change for trigger {}", trigger_id);
|
||||
|
||||
// Find sensors for this trigger
|
||||
let sensors = sqlx::query_as::<_, Sensor>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
ref,
|
||||
pack,
|
||||
pack_ref,
|
||||
label,
|
||||
description,
|
||||
entrypoint,
|
||||
runtime,
|
||||
runtime_ref,
|
||||
trigger,
|
||||
trigger_ref,
|
||||
enabled,
|
||||
param_schema,
|
||||
config,
|
||||
created,
|
||||
updated
|
||||
FROM sensor
|
||||
WHERE trigger = $1
|
||||
AND enabled = TRUE
|
||||
"#,
|
||||
)
|
||||
.bind(trigger_id)
|
||||
.fetch_all(&self.inner.db)
|
||||
.await?;
|
||||
|
||||
for sensor in sensors {
|
||||
// Check if sensor is running
|
||||
let is_running = self.inner.sensors.read().await.contains_key(&sensor.id);
|
||||
|
||||
// Check if sensor should be running (has active rules)
|
||||
let should_run = self.has_active_rules(trigger_id).await?;
|
||||
|
||||
match (is_running, should_run) {
|
||||
(false, true) => {
|
||||
// Start sensor
|
||||
info!("Starting sensor {} due to rule change", sensor.r#ref);
|
||||
if let Err(e) = self.start_sensor(sensor).await {
|
||||
error!("Failed to start sensor: {}", e);
|
||||
}
|
||||
}
|
||||
(true, false) => {
|
||||
// Stop sensor
|
||||
info!("Stopping sensor {} - no active rules", sensor.r#ref);
|
||||
if let Err(e) = self.stop_sensor(sensor.id).await {
|
||||
error!("Failed to stop sensor: {}", e);
|
||||
}
|
||||
}
|
||||
(true, true) => {
|
||||
// Restart sensor to pick up new trigger instances
|
||||
info!(
|
||||
"Restarting sensor {} to update trigger instances",
|
||||
sensor.r#ref
|
||||
);
|
||||
if let Err(e) = self.stop_sensor(sensor.id).await {
|
||||
error!("Failed to stop sensor: {}", e);
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
if let Err(e) = self.start_sensor(sensor).await {
|
||||
error!("Failed to restart sensor: {}", e);
|
||||
}
|
||||
}
|
||||
(false, false) => {
|
||||
// No action needed
|
||||
debug!("Sensor {} - no action needed", sensor.r#ref);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Monitoring loop to check sensor health
|
||||
async fn monitoring_loop(&self) {
|
||||
let mut interval = interval(Duration::from_secs(60));
|
||||
|
||||
while *self.inner.running.read().await {
|
||||
interval.tick().await;
|
||||
|
||||
debug!("Sensor manager monitoring check");
|
||||
|
||||
let sensors = self.inner.sensors.read().await;
|
||||
for (sensor_id, instance) in sensors.iter() {
|
||||
let status = instance.status().await;
|
||||
|
||||
if status.failed {
|
||||
warn!(
|
||||
"Sensor {} has failed (failure_count: {})",
|
||||
sensor_id, status.failure_count
|
||||
);
|
||||
}
|
||||
|
||||
// Check if long-running process has died
|
||||
if let Some(ref _child) = instance.child_process {
|
||||
// Note: We can't easily check if child is still running without blocking
|
||||
// This would need enhancement with a better process management approach
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
info!("Sensor manager monitoring loop stopped");
|
||||
}
|
||||
|
||||
/// Get count of active sensors
|
||||
pub async fn active_count(&self) -> usize {
|
||||
let sensors = self.inner.sensors.read().await;
|
||||
let mut active = 0;
|
||||
|
||||
for instance in sensors.values() {
|
||||
let status = instance.status().await;
|
||||
if status.running && !status.failed {
|
||||
active += 1;
|
||||
}
|
||||
}
|
||||
|
||||
active
|
||||
}
|
||||
|
||||
/// Get count of failed sensors
|
||||
pub async fn failed_count(&self) -> usize {
|
||||
let sensors = self.inner.sensors.read().await;
|
||||
let mut failed = 0;
|
||||
|
||||
for instance in sensors.values() {
|
||||
let status = instance.status().await;
|
||||
if status.failed {
|
||||
failed += 1;
|
||||
}
|
||||
}
|
||||
|
||||
failed
|
||||
}
|
||||
|
||||
/// Get total count of sensors
|
||||
pub async fn total_count(&self) -> usize {
|
||||
self.inner.sensors.read().await.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Sensor instance managing a running sensor
|
||||
struct SensorInstance {
|
||||
status: Arc<RwLock<SensorStatus>>,
|
||||
child_process: Option<Child>,
|
||||
stderr_handle: Option<JoinHandle<()>>,
|
||||
stdout_handle: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl SensorInstance {
|
||||
/// Create a new standalone sensor instance
|
||||
fn new_standalone(
|
||||
child_process: Child,
|
||||
stdout_handle: JoinHandle<()>,
|
||||
stderr_handle: JoinHandle<()>,
|
||||
) -> Self {
|
||||
Self {
|
||||
status: Arc::new(RwLock::new(SensorStatus {
|
||||
running: true,
|
||||
failed: false,
|
||||
failure_count: 0,
|
||||
last_poll: Some(chrono::Utc::now()),
|
||||
})),
|
||||
child_process: Some(child_process),
|
||||
stderr_handle: Some(stderr_handle),
|
||||
stdout_handle: Some(stdout_handle),
|
||||
}
|
||||
}
|
||||
|
||||
/// Stop the sensor
|
||||
async fn stop(&mut self) {
|
||||
{
|
||||
let mut status = self.status.write().await;
|
||||
status.running = false;
|
||||
}
|
||||
|
||||
// Kill child process if exists
|
||||
if let Some(ref mut child) = self.child_process {
|
||||
if let Err(e) = child.start_kill() {
|
||||
error!("Failed to kill sensor process: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// Abort task handles
|
||||
if let Some(ref handle) = self.stdout_handle {
|
||||
handle.abort();
|
||||
}
|
||||
|
||||
if let Some(ref handle) = self.stderr_handle {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
/// Get sensor status
|
||||
async fn status(&self) -> SensorStatus {
|
||||
self.status.read().await.clone()
|
||||
}
|
||||
}
|
||||
|
||||
/// Sensor status information
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SensorStatus {
|
||||
/// Whether the sensor is running
|
||||
pub running: bool,
|
||||
|
||||
/// Whether the sensor has failed
|
||||
pub failed: bool,
|
||||
|
||||
/// Number of consecutive failures
|
||||
pub failure_count: u32,
|
||||
|
||||
/// Last successful poll time
|
||||
pub last_poll: Option<chrono::DateTime<chrono::Utc>>,
|
||||
}
|
||||
|
||||
impl Default for SensorStatus {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
running: false,
|
||||
failed: false,
|
||||
failure_count: 0,
|
||||
last_poll: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_sensor_status_default() {
|
||||
let status = SensorStatus::default();
|
||||
assert!(!status.running);
|
||||
assert!(!status.failed);
|
||||
assert_eq!(status.failure_count, 0);
|
||||
assert!(status.last_poll.is_none());
|
||||
}
|
||||
}
|
||||
354
crates/sensor/src/sensor_worker_registration.rs
Normal file
354
crates/sensor/src/sensor_worker_registration.rs
Normal file
@@ -0,0 +1,354 @@
|
||||
//! Sensor Worker Registration Module
|
||||
//!
|
||||
//! Handles sensor worker registration, discovery, and status management in the database.
|
||||
//! Similar to action worker registration but tailored for sensor service instances.
|
||||
//!
|
||||
//! Runtime detection uses the unified RuntimeDetector from common crate.
|
||||
|
||||
use attune_common::config::Config;
|
||||
use attune_common::error::Result;
|
||||
use attune_common::models::{Worker, WorkerRole, WorkerStatus, WorkerType};
|
||||
use attune_common::runtime_detection::RuntimeDetector;
|
||||
use chrono::Utc;
|
||||
use serde_json::json;
|
||||
use sqlx::{PgPool, Row};
|
||||
use std::collections::HashMap;
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// Sensor worker registration manager
|
||||
pub struct SensorWorkerRegistration {
|
||||
pool: PgPool,
|
||||
worker_id: Option<i64>,
|
||||
worker_name: String,
|
||||
host: Option<String>,
|
||||
capabilities: HashMap<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
impl SensorWorkerRegistration {
|
||||
/// Create a new sensor worker registration manager
|
||||
pub fn new(pool: PgPool, config: &Config) -> Self {
|
||||
let worker_name = config
|
||||
.sensor
|
||||
.as_ref()
|
||||
.and_then(|s| s.worker_name.clone())
|
||||
.unwrap_or_else(|| {
|
||||
format!(
|
||||
"sensor-{}",
|
||||
hostname::get()
|
||||
.unwrap_or_else(|_| "unknown".into())
|
||||
.to_string_lossy()
|
||||
)
|
||||
});
|
||||
|
||||
let host = config
|
||||
.sensor
|
||||
.as_ref()
|
||||
.and_then(|s| s.host.clone())
|
||||
.or_else(|| {
|
||||
hostname::get()
|
||||
.ok()
|
||||
.map(|h| h.to_string_lossy().to_string())
|
||||
});
|
||||
|
||||
// Initial capabilities (will be populated asynchronously)
|
||||
let mut capabilities = HashMap::new();
|
||||
|
||||
// Set max_concurrent_sensors from config
|
||||
let max_concurrent = config
|
||||
.sensor
|
||||
.as_ref()
|
||||
.and_then(|s| s.max_concurrent_sensors)
|
||||
.unwrap_or(10);
|
||||
capabilities.insert("max_concurrent_sensors".to_string(), json!(max_concurrent));
|
||||
|
||||
// Add sensor worker version metadata
|
||||
capabilities.insert(
|
||||
"sensor_version".to_string(),
|
||||
json!(env!("CARGO_PKG_VERSION")),
|
||||
);
|
||||
|
||||
// Placeholder for runtimes (will be detected asynchronously)
|
||||
capabilities.insert("runtimes".to_string(), json!(Vec::<String>::new()));
|
||||
|
||||
Self {
|
||||
pool,
|
||||
worker_id: None,
|
||||
worker_name,
|
||||
host,
|
||||
capabilities,
|
||||
}
|
||||
}
|
||||
|
||||
/// Register the sensor worker in the database
|
||||
pub async fn register(&mut self, config: &Config) -> Result<i64> {
|
||||
// Detect runtimes from database if not already configured
|
||||
self.detect_capabilities_async(config).await?;
|
||||
|
||||
info!("Registering sensor worker: {}", self.worker_name);
|
||||
|
||||
// Check if sensor worker with this name already exists
|
||||
let existing = sqlx::query_as::<_, Worker>(
|
||||
"SELECT * FROM worker WHERE name = $1 AND worker_role = 'sensor' ORDER BY created DESC LIMIT 1",
|
||||
)
|
||||
.bind(&self.worker_name)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
let worker_id = if let Some(existing_worker) = existing {
|
||||
info!(
|
||||
"Sensor worker '{}' already exists (ID: {}), updating status",
|
||||
self.worker_name, existing_worker.id
|
||||
);
|
||||
|
||||
// Update existing sensor worker to active status with new heartbeat
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE worker
|
||||
SET status = $1,
|
||||
capabilities = $2,
|
||||
last_heartbeat = $3,
|
||||
updated = $4,
|
||||
host = $5
|
||||
WHERE id = $6
|
||||
"#,
|
||||
)
|
||||
.bind(WorkerStatus::Active)
|
||||
.bind(serde_json::to_value(&self.capabilities)?)
|
||||
.bind(Utc::now())
|
||||
.bind(Utc::now())
|
||||
.bind(&self.host)
|
||||
.bind(existing_worker.id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
existing_worker.id
|
||||
} else {
|
||||
// Insert new sensor worker
|
||||
let row = sqlx::query(
|
||||
r#"
|
||||
INSERT INTO worker (name, worker_type, worker_role, host, status, capabilities, last_heartbeat)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
RETURNING id
|
||||
"#,
|
||||
)
|
||||
.bind(&self.worker_name)
|
||||
.bind(WorkerType::Local) // Sensor workers are always local
|
||||
.bind(WorkerRole::Sensor)
|
||||
.bind(&self.host)
|
||||
.bind(WorkerStatus::Active)
|
||||
.bind(serde_json::to_value(&self.capabilities)?)
|
||||
.bind(Utc::now())
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
let worker_id: i64 = row.get("id");
|
||||
info!("Sensor worker registered with ID: {}", worker_id);
|
||||
worker_id
|
||||
};
|
||||
|
||||
self.worker_id = Some(worker_id);
|
||||
Ok(worker_id)
|
||||
}
|
||||
|
||||
/// Send heartbeat to update last_heartbeat timestamp
|
||||
pub async fn heartbeat(&self) -> Result<()> {
|
||||
if let Some(worker_id) = self.worker_id {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE worker
|
||||
SET last_heartbeat = $1,
|
||||
status = $2,
|
||||
updated = $3
|
||||
WHERE id = $4
|
||||
"#,
|
||||
)
|
||||
.bind(Utc::now())
|
||||
.bind(WorkerStatus::Active)
|
||||
.bind(Utc::now())
|
||||
.bind(worker_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
debug!("Sensor worker heartbeat sent");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Mark sensor worker as inactive
|
||||
pub async fn deregister(&self) -> Result<()> {
|
||||
if let Some(worker_id) = self.worker_id {
|
||||
info!("Deregistering sensor worker: {}", self.worker_name);
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE worker
|
||||
SET status = $1,
|
||||
updated = $2
|
||||
WHERE id = $3
|
||||
"#,
|
||||
)
|
||||
.bind(WorkerStatus::Inactive)
|
||||
.bind(Utc::now())
|
||||
.bind(worker_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
info!("Sensor worker deregistered");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the registered sensor worker ID
|
||||
pub fn worker_id(&self) -> Option<i64> {
|
||||
self.worker_id
|
||||
}
|
||||
|
||||
/// Get the sensor worker name
|
||||
pub fn worker_name(&self) -> &str {
|
||||
&self.worker_name
|
||||
}
|
||||
|
||||
/// Add a capability to the sensor worker
|
||||
pub fn add_capability(&mut self, key: String, value: serde_json::Value) {
|
||||
self.capabilities.insert(key, value);
|
||||
}
|
||||
|
||||
/// Update sensor worker capabilities in the database
|
||||
pub async fn update_capabilities(&self) -> Result<()> {
|
||||
if let Some(worker_id) = self.worker_id {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE worker
|
||||
SET capabilities = $1,
|
||||
updated = $2
|
||||
WHERE id = $3
|
||||
"#,
|
||||
)
|
||||
.bind(serde_json::to_value(&self.capabilities)?)
|
||||
.bind(Utc::now())
|
||||
.bind(worker_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
|
||||
info!("Sensor worker capabilities updated");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Detect sensor worker capabilities based on database-driven runtime verification
|
||||
///
|
||||
/// This is a synchronous wrapper that should be called after pool is available.
|
||||
/// The actual detection happens in `detect_capabilities_async`.
|
||||
/// Detect available runtimes using the unified runtime detector
|
||||
pub async fn detect_capabilities_async(&mut self, config: &Config) -> Result<()> {
|
||||
info!("Detecting sensor worker capabilities...");
|
||||
|
||||
let detector = RuntimeDetector::new(self.pool.clone());
|
||||
|
||||
// Get config capabilities if available
|
||||
let config_capabilities = config.sensor.as_ref().and_then(|s| s.capabilities.as_ref());
|
||||
|
||||
// Detect capabilities with three-tier priority:
|
||||
// 1. ATTUNE_SENSOR_RUNTIMES env var
|
||||
// 2. Config file
|
||||
// 3. Database-driven detection
|
||||
let detected_capabilities = detector
|
||||
.detect_capabilities(config, "ATTUNE_SENSOR_RUNTIMES", config_capabilities)
|
||||
.await?;
|
||||
|
||||
// Merge detected capabilities with existing ones
|
||||
for (key, value) in detected_capabilities {
|
||||
self.capabilities.insert(key, value);
|
||||
}
|
||||
|
||||
info!(
|
||||
"Sensor worker capabilities detected: {:?}",
|
||||
self.capabilities
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for SensorWorkerRegistration {
|
||||
fn drop(&mut self) {
|
||||
// Note: We can't make this async, so we just log
|
||||
// The main service should call deregister() explicitly during shutdown
|
||||
if self.worker_id.is_some() {
|
||||
info!("SensorWorkerRegistration dropped - sensor worker should be deregistered");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // Requires database
|
||||
async fn test_database_driven_detection() {
|
||||
let config = Config::load().unwrap();
|
||||
let db = attune_common::db::Database::new(&config.database)
|
||||
.await
|
||||
.unwrap();
|
||||
let pool = db.pool().clone();
|
||||
let mut registration = SensorWorkerRegistration::new(pool, &config);
|
||||
|
||||
// Detect runtimes from database
|
||||
registration
|
||||
.detect_capabilities_async(&config)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Should have detected some runtimes
|
||||
let runtimes = registration.capabilities.get("runtimes").unwrap();
|
||||
let runtime_array = runtimes.as_array().unwrap();
|
||||
assert!(!runtime_array.is_empty());
|
||||
|
||||
println!("Detected runtimes: {:?}", runtime_array);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // Requires database
|
||||
async fn test_sensor_worker_registration() {
|
||||
let config = Config::load().unwrap();
|
||||
let db = attune_common::db::Database::new(&config.database)
|
||||
.await
|
||||
.unwrap();
|
||||
let pool = db.pool().clone();
|
||||
let mut registration = SensorWorkerRegistration::new(pool, &config);
|
||||
|
||||
// Test registration
|
||||
let worker_id = registration.register(&config).await.unwrap();
|
||||
assert!(worker_id > 0);
|
||||
assert_eq!(registration.worker_id(), Some(worker_id));
|
||||
|
||||
// Test heartbeat
|
||||
registration.heartbeat().await.unwrap();
|
||||
|
||||
// Test deregistration
|
||||
registration.deregister().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore] // Requires database
|
||||
async fn test_sensor_worker_capabilities() {
|
||||
let config = Config::load().unwrap();
|
||||
let db = attune_common::db::Database::new(&config.database)
|
||||
.await
|
||||
.unwrap();
|
||||
let pool = db.pool().clone();
|
||||
let mut registration = SensorWorkerRegistration::new(pool, &config);
|
||||
|
||||
registration.register(&config).await.unwrap();
|
||||
|
||||
// Add custom capability
|
||||
registration.add_capability("custom_feature".to_string(), json!(true));
|
||||
registration.update_capabilities().await.unwrap();
|
||||
|
||||
registration.deregister().await.unwrap();
|
||||
}
|
||||
}
|
||||
278
crates/sensor/src/service.rs
Normal file
278
crates/sensor/src/service.rs
Normal file
@@ -0,0 +1,278 @@
|
||||
//! Sensor Service
|
||||
//!
|
||||
//! Main service orchestrator that coordinates sensor management
|
||||
//! and rule lifecycle listening.
|
||||
|
||||
use crate::rule_lifecycle_listener::RuleLifecycleListener;
|
||||
use crate::sensor_manager::SensorManager;
|
||||
use crate::sensor_worker_registration::SensorWorkerRegistration;
|
||||
use anyhow::Result;
|
||||
use attune_common::config::Config;
|
||||
use attune_common::db::Database;
|
||||
use attune_common::mq::MessageQueue;
|
||||
use sqlx::PgPool;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
/// Sensor Service state
|
||||
#[derive(Clone)]
|
||||
pub struct SensorService {
|
||||
inner: Arc<SensorServiceInner>,
|
||||
}
|
||||
|
||||
struct SensorServiceInner {
|
||||
config: Config,
|
||||
db: PgPool,
|
||||
mq: MessageQueue,
|
||||
sensor_manager: Arc<SensorManager>,
|
||||
rule_lifecycle_listener: Arc<RuleLifecycleListener>,
|
||||
sensor_worker_registration: Arc<RwLock<SensorWorkerRegistration>>,
|
||||
heartbeat_interval: u64,
|
||||
running: Arc<RwLock<bool>>,
|
||||
}
|
||||
|
||||
impl SensorService {
|
||||
/// Create a new sensor service
|
||||
pub async fn new(config: Config) -> Result<Self> {
|
||||
info!("Initializing Sensor Service");
|
||||
|
||||
// Connect to database
|
||||
info!("Connecting to database...");
|
||||
let database = Database::new(&config.database).await?;
|
||||
let db = database.pool().clone();
|
||||
info!("Database connection established");
|
||||
|
||||
// Connect to message queue
|
||||
info!("Connecting to message queue...");
|
||||
let mq_config = config
|
||||
.message_queue
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("Message queue configuration is required"))?;
|
||||
let mq = MessageQueue::connect(&mq_config.url).await?;
|
||||
info!("Message queue connection established");
|
||||
|
||||
// Create service components
|
||||
info!("Creating service components...");
|
||||
|
||||
let sensor_manager = Arc::new(SensorManager::new(db.clone()));
|
||||
|
||||
// Create rule lifecycle listener
|
||||
let rule_lifecycle_listener = Arc::new(RuleLifecycleListener::new(
|
||||
db.clone(),
|
||||
mq.get_connection().clone(),
|
||||
sensor_manager.clone(),
|
||||
));
|
||||
|
||||
// Create sensor worker registration
|
||||
let sensor_worker_registration = SensorWorkerRegistration::new(db.clone(), &config);
|
||||
let heartbeat_interval = config
|
||||
.sensor
|
||||
.as_ref()
|
||||
.map(|s| s.heartbeat_interval)
|
||||
.unwrap_or(30);
|
||||
|
||||
Ok(Self {
|
||||
inner: Arc::new(SensorServiceInner {
|
||||
config,
|
||||
db,
|
||||
mq,
|
||||
sensor_manager,
|
||||
rule_lifecycle_listener,
|
||||
sensor_worker_registration: Arc::new(RwLock::new(sensor_worker_registration)),
|
||||
heartbeat_interval,
|
||||
running: Arc::new(RwLock::new(false)),
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
/// Start the sensor service
|
||||
pub async fn start(&self) -> Result<()> {
|
||||
info!("Starting Sensor Service");
|
||||
|
||||
// Mark as running
|
||||
*self.inner.running.write().await = true;
|
||||
|
||||
// Register sensor worker
|
||||
info!("Registering sensor worker...");
|
||||
let worker_id = self
|
||||
.inner
|
||||
.sensor_worker_registration
|
||||
.write()
|
||||
.await
|
||||
.register(&self.inner.config)
|
||||
.await?;
|
||||
info!("Sensor worker registered with ID: {}", worker_id);
|
||||
|
||||
// Start rule lifecycle listener
|
||||
info!("Starting rule lifecycle listener...");
|
||||
if let Err(e) = self.inner.rule_lifecycle_listener.start().await {
|
||||
error!("Failed to start rule lifecycle listener: {}", e);
|
||||
return Err(e);
|
||||
}
|
||||
info!("Rule lifecycle listener started");
|
||||
|
||||
// Start sensor manager
|
||||
info!("Starting sensor manager...");
|
||||
if let Err(e) = self.inner.sensor_manager.start().await {
|
||||
error!("Failed to start sensor manager: {}", e);
|
||||
return Err(e);
|
||||
}
|
||||
info!("Sensor manager started");
|
||||
|
||||
// Start heartbeat loop
|
||||
let registration = self.inner.sensor_worker_registration.clone();
|
||||
let heartbeat_interval = self.inner.heartbeat_interval;
|
||||
let running = self.inner.running.clone();
|
||||
tokio::spawn(async move {
|
||||
while *running.read().await {
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(heartbeat_interval)).await;
|
||||
if let Err(e) = registration.read().await.heartbeat().await {
|
||||
error!("Failed to send sensor worker heartbeat: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Wait until stopped
|
||||
while *self.inner.running.read().await {
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
info!("Sensor Service stopped");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Stop the sensor service
|
||||
pub async fn stop(&self) -> Result<()> {
|
||||
info!("Stopping Sensor Service");
|
||||
|
||||
// Mark as not running
|
||||
*self.inner.running.write().await = false;
|
||||
|
||||
// Deregister sensor worker
|
||||
info!("Deregistering sensor worker...");
|
||||
if let Err(e) = self
|
||||
.inner
|
||||
.sensor_worker_registration
|
||||
.read()
|
||||
.await
|
||||
.deregister()
|
||||
.await
|
||||
{
|
||||
error!("Failed to deregister sensor worker: {}", e);
|
||||
}
|
||||
|
||||
// Stop rule lifecycle listener
|
||||
info!("Stopping rule lifecycle listener...");
|
||||
if let Err(e) = self.inner.rule_lifecycle_listener.stop().await {
|
||||
error!("Failed to stop rule lifecycle listener: {}", e);
|
||||
}
|
||||
|
||||
// Stop sensor manager
|
||||
info!("Stopping sensor manager...");
|
||||
if let Err(e) = self.inner.sensor_manager.stop().await {
|
||||
error!("Failed to stop sensor manager: {}", e);
|
||||
}
|
||||
|
||||
// Close message queue connection
|
||||
info!("Closing message queue connection...");
|
||||
if let Err(e) = self.inner.mq.close().await {
|
||||
warn!("Error closing message queue: {}", e);
|
||||
}
|
||||
|
||||
// Close database connection
|
||||
info!("Closing database connection...");
|
||||
self.inner.db.close().await;
|
||||
|
||||
info!("Sensor Service stopped successfully");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check if service is running
|
||||
pub async fn is_running(&self) -> bool {
|
||||
*self.inner.running.read().await
|
||||
}
|
||||
|
||||
/// Get database pool
|
||||
pub fn db(&self) -> &PgPool {
|
||||
&self.inner.db
|
||||
}
|
||||
|
||||
/// Get message queue
|
||||
pub fn mq(&self) -> &MessageQueue {
|
||||
&self.inner.mq
|
||||
}
|
||||
|
||||
/// Get sensor manager
|
||||
pub fn sensor_manager(&self) -> Arc<SensorManager> {
|
||||
self.inner.sensor_manager.clone()
|
||||
}
|
||||
|
||||
/// Get health status
|
||||
pub async fn health_check(&self) -> HealthStatus {
|
||||
// Check if service is running
|
||||
if !*self.inner.running.read().await {
|
||||
return HealthStatus::Unhealthy("Service not running".to_string());
|
||||
}
|
||||
|
||||
// Check database connection
|
||||
if let Err(e) = sqlx::query("SELECT 1").execute(&self.inner.db).await {
|
||||
return HealthStatus::Unhealthy(format!("Database connection failed: {}", e));
|
||||
}
|
||||
|
||||
// Check sensor manager health
|
||||
let active_sensors = self.inner.sensor_manager.active_count().await;
|
||||
let failed_sensors = self.inner.sensor_manager.failed_count().await;
|
||||
|
||||
if active_sensors == 0 {
|
||||
return HealthStatus::Degraded("No active sensors".to_string());
|
||||
}
|
||||
|
||||
if failed_sensors > 10 {
|
||||
return HealthStatus::Degraded(format!("{} sensors have failed", failed_sensors));
|
||||
}
|
||||
|
||||
HealthStatus::Healthy
|
||||
}
|
||||
}
|
||||
|
||||
/// Health status enumeration
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum HealthStatus {
|
||||
/// Service is healthy
|
||||
Healthy,
|
||||
/// Service is degraded but operational
|
||||
Degraded(String),
|
||||
/// Service is unhealthy
|
||||
Unhealthy(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for HealthStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
HealthStatus::Healthy => write!(f, "healthy"),
|
||||
HealthStatus::Degraded(msg) => write!(f, "degraded: {}", msg),
|
||||
HealthStatus::Unhealthy(msg) => write!(f, "unhealthy: {}", msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_health_status_display() {
|
||||
assert_eq!(HealthStatus::Healthy.to_string(), "healthy");
|
||||
assert_eq!(
|
||||
HealthStatus::Degraded("test".to_string()).to_string(),
|
||||
"degraded: test"
|
||||
);
|
||||
assert_eq!(
|
||||
HealthStatus::Unhealthy("error".to_string()).to_string(),
|
||||
"unhealthy: error"
|
||||
);
|
||||
}
|
||||
}
|
||||
468
crates/sensor/src/template_resolver.rs
Normal file
468
crates/sensor/src/template_resolver.rs
Normal file
@@ -0,0 +1,468 @@
|
||||
//! Template Resolver
|
||||
//!
|
||||
//! Resolves template variables in rule action parameters using context from
|
||||
//! trigger payloads, pack configuration, and system variables.
|
||||
//!
|
||||
//! Supports template syntax: `{{ source.path.to.value }}`
|
||||
//!
|
||||
//! Example:
|
||||
//! ```rust
|
||||
//! use serde_json::json;
|
||||
//! use attune_sensor::template_resolver::{TemplateContext, resolve_templates};
|
||||
//!
|
||||
//! let params = json!({
|
||||
//! "message": "Error in {{ trigger.payload.service }}"
|
||||
//! });
|
||||
//!
|
||||
//! let context = TemplateContext {
|
||||
//! trigger_payload: json!({"service": "api-gateway"}),
|
||||
//! pack_config: json!({}),
|
||||
//! system_vars: json!({}),
|
||||
//! };
|
||||
//!
|
||||
//! let resolved = resolve_templates(¶ms, &context).unwrap();
|
||||
//! assert_eq!(resolved["message"], "Error in api-gateway");
|
||||
//! ```
|
||||
|
||||
use anyhow::Result;
|
||||
use regex::Regex;
|
||||
use serde_json::Value as JsonValue;
|
||||
use std::sync::LazyLock;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
/// Template context containing all available data sources
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TemplateContext {
|
||||
/// Event/trigger payload data
|
||||
pub trigger_payload: JsonValue,
|
||||
/// Pack configuration
|
||||
pub pack_config: JsonValue,
|
||||
/// System-provided variables
|
||||
pub system_vars: JsonValue,
|
||||
}
|
||||
|
||||
impl TemplateContext {
|
||||
/// Create a new template context
|
||||
pub fn new(trigger_payload: JsonValue, pack_config: JsonValue, system_vars: JsonValue) -> Self {
|
||||
Self {
|
||||
trigger_payload,
|
||||
pack_config,
|
||||
system_vars,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a value from the context using a dotted path
|
||||
///
|
||||
/// Supports paths like:
|
||||
/// - `trigger.payload.field`
|
||||
/// - `pack.config.setting`
|
||||
/// - `system.timestamp`
|
||||
pub fn get_value(&self, path: &str) -> Option<JsonValue> {
|
||||
let parts: Vec<&str> = path.split('.').collect();
|
||||
|
||||
if parts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Determine the root source
|
||||
let root = match parts[0] {
|
||||
"trigger" => {
|
||||
// trigger.payload.* paths
|
||||
if parts.len() < 2 || parts[1] != "payload" {
|
||||
warn!(
|
||||
"Invalid trigger path: {}, expected 'trigger.payload.*'",
|
||||
path
|
||||
);
|
||||
return None;
|
||||
}
|
||||
&self.trigger_payload
|
||||
}
|
||||
"pack" => {
|
||||
// pack.config.* paths
|
||||
if parts.len() < 2 || parts[1] != "config" {
|
||||
warn!("Invalid pack path: {}, expected 'pack.config.*'", path);
|
||||
return None;
|
||||
}
|
||||
&self.pack_config
|
||||
}
|
||||
"system" => &self.system_vars,
|
||||
_ => {
|
||||
warn!("Unknown template source: {}", parts[0]);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
// Navigate the path (skip the first 2 parts for trigger/pack, 1 for system)
|
||||
let skip_count = match parts[0] {
|
||||
"trigger" | "pack" => 2,
|
||||
"system" => 1,
|
||||
_ => return None,
|
||||
};
|
||||
|
||||
extract_nested_value(root, &parts[skip_count..])
|
||||
}
|
||||
}
|
||||
|
||||
/// Regex pattern to match template variables: {{ ... }}
|
||||
static TEMPLATE_REGEX: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r"\{\{\s*([^}]+?)\s*\}\}").expect("Failed to compile template regex")
|
||||
});
|
||||
|
||||
/// Resolve all template variables in a JSON value
|
||||
///
|
||||
/// Recursively processes objects and arrays, replacing template strings
|
||||
/// with values from the context.
|
||||
pub fn resolve_templates(value: &JsonValue, context: &TemplateContext) -> Result<JsonValue> {
|
||||
match value {
|
||||
JsonValue::String(s) => resolve_string_template(s, context),
|
||||
JsonValue::Object(map) => {
|
||||
let mut resolved = serde_json::Map::new();
|
||||
for (key, val) in map {
|
||||
resolved.insert(key.clone(), resolve_templates(val, context)?);
|
||||
}
|
||||
Ok(JsonValue::Object(resolved))
|
||||
}
|
||||
JsonValue::Array(arr) => {
|
||||
let resolved: Result<Vec<JsonValue>> =
|
||||
arr.iter().map(|v| resolve_templates(v, context)).collect();
|
||||
Ok(JsonValue::Array(resolved?))
|
||||
}
|
||||
// Pass through other types unchanged
|
||||
other => Ok(other.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve templates in a string value
|
||||
///
|
||||
/// If the string contains a single template that matches the entire string,
|
||||
/// returns the value with its original type (preserving numbers, booleans, etc).
|
||||
///
|
||||
/// If the string contains multiple templates or mixed content, performs
|
||||
/// string interpolation.
|
||||
fn resolve_string_template(s: &str, context: &TemplateContext) -> Result<JsonValue> {
|
||||
// Check if the entire string is a single template (for type preservation)
|
||||
if let Some(captures) = TEMPLATE_REGEX.captures(s) {
|
||||
let full_match = captures.get(0).unwrap();
|
||||
if full_match.start() == 0 && full_match.end() == s.len() {
|
||||
// Single template - preserve type
|
||||
let path = captures.get(1).unwrap().as_str().trim();
|
||||
debug!("Resolving single template: {}", path);
|
||||
|
||||
return match context.get_value(path) {
|
||||
Some(value) => {
|
||||
debug!("Resolved {} -> {:?}", path, value);
|
||||
Ok(value)
|
||||
}
|
||||
None => {
|
||||
warn!("Template variable not found: {}", path);
|
||||
Ok(JsonValue::Null)
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Multiple templates or mixed content - perform string interpolation
|
||||
let mut result = s.to_string();
|
||||
let mut any_replaced = false;
|
||||
|
||||
for captures in TEMPLATE_REGEX.captures_iter(s) {
|
||||
let full_match = captures.get(0).unwrap().as_str();
|
||||
let path = captures.get(1).unwrap().as_str().trim();
|
||||
|
||||
debug!("Resolving template in string: {}", path);
|
||||
|
||||
match context.get_value(path) {
|
||||
Some(value) => {
|
||||
let replacement = value_to_string(&value);
|
||||
debug!("Resolved {} -> {}", path, replacement);
|
||||
result = result.replace(full_match, &replacement);
|
||||
any_replaced = true;
|
||||
}
|
||||
None => {
|
||||
warn!("Template variable not found: {}", path);
|
||||
result = result.replace(full_match, "");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if any_replaced {
|
||||
debug!("String interpolation result: {}", result);
|
||||
}
|
||||
|
||||
Ok(JsonValue::String(result))
|
||||
}
|
||||
|
||||
/// Extract a nested value from JSON using a path
|
||||
fn extract_nested_value(root: &JsonValue, path: &[&str]) -> Option<JsonValue> {
|
||||
if path.is_empty() {
|
||||
return Some(root.clone());
|
||||
}
|
||||
|
||||
let mut current = root;
|
||||
|
||||
for part in path {
|
||||
match current {
|
||||
JsonValue::Object(map) => {
|
||||
current = map.get(*part)?;
|
||||
}
|
||||
JsonValue::Array(arr) => {
|
||||
// Try to parse part as array index
|
||||
if let Ok(index) = part.parse::<usize>() {
|
||||
current = arr.get(index)?;
|
||||
} else {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
_ => return None,
|
||||
}
|
||||
}
|
||||
|
||||
Some(current.clone())
|
||||
}
|
||||
|
||||
/// Convert a JSON value to a string for interpolation
|
||||
fn value_to_string(value: &JsonValue) -> String {
|
||||
match value {
|
||||
JsonValue::String(s) => s.clone(),
|
||||
JsonValue::Number(n) => n.to_string(),
|
||||
JsonValue::Bool(b) => b.to_string(),
|
||||
JsonValue::Null => String::new(),
|
||||
JsonValue::Array(_) | JsonValue::Object(_) => {
|
||||
// For complex types, serialize as JSON
|
||||
serde_json::to_string(value).unwrap_or_else(|_| String::new())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn create_test_context() -> TemplateContext {
|
||||
TemplateContext {
|
||||
trigger_payload: json!({
|
||||
"service": "api-gateway",
|
||||
"message": "Connection timeout",
|
||||
"severity": "critical",
|
||||
"count": 42,
|
||||
"enabled": true,
|
||||
"metadata": {
|
||||
"host": "web-01",
|
||||
"port": 8080
|
||||
},
|
||||
"tags": ["production", "backend"]
|
||||
}),
|
||||
pack_config: json!({
|
||||
"api_token": "secret123",
|
||||
"alert_channel": "#incidents",
|
||||
"timeout": 30
|
||||
}),
|
||||
system_vars: json!({
|
||||
"timestamp": "2026-01-17T15:30:00Z",
|
||||
"rule": {
|
||||
"id": 42,
|
||||
"ref": "test.rule"
|
||||
},
|
||||
"event": {
|
||||
"id": 123
|
||||
}
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_simple_string_substitution() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"message": "Hello {{ trigger.payload.service }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["message"], "Hello api-gateway");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_single_template_type_preservation() {
|
||||
let context = create_test_context();
|
||||
|
||||
// Number
|
||||
let template = json!({"count": "{{ trigger.payload.count }}"});
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["count"], 42);
|
||||
|
||||
// Boolean
|
||||
let template = json!({"enabled": "{{ trigger.payload.enabled }}"});
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["enabled"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nested_object_access() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"host": "{{ trigger.payload.metadata.host }}",
|
||||
"port": "{{ trigger.payload.metadata.port }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["host"], "web-01");
|
||||
assert_eq!(result["port"], 8080);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_array_access() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"first_tag": "{{ trigger.payload.tags.0 }}",
|
||||
"second_tag": "{{ trigger.payload.tags.1 }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["first_tag"], "production");
|
||||
assert_eq!(result["second_tag"], "backend");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pack_config_reference() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"token": "{{ pack.config.api_token }}",
|
||||
"channel": "{{ pack.config.alert_channel }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["token"], "secret123");
|
||||
assert_eq!(result["channel"], "#incidents");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_system_variables() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"timestamp": "{{ system.timestamp }}",
|
||||
"rule_id": "{{ system.rule.id }}",
|
||||
"event_id": "{{ system.event.id }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["timestamp"], "2026-01-17T15:30:00Z");
|
||||
assert_eq!(result["rule_id"], 42);
|
||||
assert_eq!(result["event_id"], 123);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_missing_value_returns_null() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"missing": "{{ trigger.payload.nonexistent }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert!(result["missing"].is_null());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multiple_templates_in_string() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"message": "Error in {{ trigger.payload.service }}: {{ trigger.payload.message }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(
|
||||
result["message"],
|
||||
"Error in api-gateway: Connection timeout"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_static_values_unchanged() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"static": "This is static",
|
||||
"number": 123,
|
||||
"boolean": false
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["static"], "This is static");
|
||||
assert_eq!(result["number"], 123);
|
||||
assert_eq!(result["boolean"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nested_objects_and_arrays() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"nested": {
|
||||
"field1": "{{ trigger.payload.service }}",
|
||||
"field2": "{{ pack.config.timeout }}"
|
||||
},
|
||||
"array": [
|
||||
"{{ trigger.payload.severity }}",
|
||||
"static value"
|
||||
]
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["nested"]["field1"], "api-gateway");
|
||||
assert_eq!(result["nested"]["field2"], 30);
|
||||
assert_eq!(result["array"][0], "critical");
|
||||
assert_eq!(result["array"][1], "static value");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_template_context() {
|
||||
let context = TemplateContext {
|
||||
trigger_payload: json!({}),
|
||||
pack_config: json!({}),
|
||||
system_vars: json!({}),
|
||||
};
|
||||
|
||||
let template = json!({
|
||||
"message": "{{ trigger.payload.missing }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert!(result["message"].is_null());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_whitespace_in_templates() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"message": "{{ trigger.payload.service }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["message"], "api-gateway");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_complex_real_world_example() {
|
||||
let context = create_test_context();
|
||||
let template = json!({
|
||||
"channel": "{{ pack.config.alert_channel }}",
|
||||
"message": "🚨 Error in {{ trigger.payload.service }}: {{ trigger.payload.message }}",
|
||||
"severity": "{{ trigger.payload.severity }}",
|
||||
"details": {
|
||||
"host": "{{ trigger.payload.metadata.host }}",
|
||||
"count": "{{ trigger.payload.count }}",
|
||||
"tags": "{{ trigger.payload.tags }}"
|
||||
},
|
||||
"timestamp": "{{ system.timestamp }}"
|
||||
});
|
||||
|
||||
let result = resolve_templates(&template, &context).unwrap();
|
||||
assert_eq!(result["channel"], "#incidents");
|
||||
assert_eq!(
|
||||
result["message"],
|
||||
"🚨 Error in api-gateway: Connection timeout"
|
||||
);
|
||||
assert_eq!(result["severity"], "critical");
|
||||
assert_eq!(result["details"]["host"], "web-01");
|
||||
assert_eq!(result["details"]["count"], 42);
|
||||
assert_eq!(result["timestamp"], "2026-01-17T15:30:00Z");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user