http_request action working nicely

This commit is contained in:
2026-02-09 23:21:23 -06:00
parent e31ecb781b
commit 966a5af188
18 changed files with 720 additions and 395 deletions

View File

@@ -8,6 +8,7 @@ use anyhow::Result;
use attune_common::config::Config;
use attune_sensor::service::SensorService;
use clap::Parser;
use tokio::signal::unix::{signal, SignalKind};
use tracing::{error, info};
#[derive(Parser, Debug)]
@@ -56,32 +57,38 @@ async fn main() -> Result<()> {
info!("Message Queue: {}", mask_connection_string(&mq_config.url));
}
// Create sensor service
// Create and start 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
// Start the service (spawns background tasks and returns)
info!("Starting Sensor Service components...");
if let Err(e) = service.start().await {
error!("Sensor Service error: {}", e);
return Err(e);
service.start().await?;
info!("Attune Sensor Service is ready");
// Setup signal handlers for graceful shutdown
let mut sigint = signal(SignalKind::interrupt())?;
let mut sigterm = signal(SignalKind::terminate())?;
tokio::select! {
_ = sigint.recv() => {
info!("Received SIGINT signal");
}
_ = sigterm.recv() => {
info!("Received SIGTERM signal");
}
}
info!("Sensor Service has shut down gracefully");
info!("Shutting down gracefully...");
// Stop the service: deregister worker, stop sensors, clean up connections
if let Err(e) = service.stop().await {
error!("Error during shutdown: {}", e);
}
info!("Attune Sensor Service shutdown complete");
Ok(())
}

View File

@@ -12,6 +12,7 @@ use serde_json::Value as JsonValue;
use sqlx::PgPool;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
use tracing::{error, info, warn};
use crate::sensor_manager::SensorManager;
@@ -22,6 +23,7 @@ pub struct RuleLifecycleListener {
connection: Connection,
sensor_manager: Arc<SensorManager>,
consumer: Arc<RwLock<Option<Consumer>>>,
task_handle: RwLock<Option<JoinHandle<()>>>,
}
impl RuleLifecycleListener {
@@ -32,6 +34,7 @@ impl RuleLifecycleListener {
connection,
sensor_manager,
consumer: Arc::new(RwLock::new(None)),
task_handle: RwLock::new(None),
}
}
@@ -88,19 +91,20 @@ impl RuleLifecycleListener {
);
}
// Store consumer
// Store consumer reference (for cleanup on drop)
*self.consumer.write().await = Some(consumer);
// Clone self for async handler
// Clone references for the spawned task
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() {
// Start consuming messages in a background task.
// Take the consumer out of the Arc<RwLock> so we don't hold the read lock
// for the entire duration of consume_with_handler (which would deadlock stop()).
let handle = tokio::spawn(async move {
let consumer = consumer_ref.write().await.take();
if let Some(consumer) = consumer {
let result = consumer
.consume_with_handler::<JsonValue, _, _>(move |envelope| {
let db = db.clone();
@@ -129,6 +133,8 @@ impl RuleLifecycleListener {
}
});
*self.task_handle.write().await = Some(handle);
info!("Rule lifecycle listener started");
Ok(())
@@ -138,8 +144,15 @@ impl RuleLifecycleListener {
pub async fn stop(&self) -> Result<()> {
info!("Stopping rule lifecycle listener");
// Abort the consumer task first — this ends the consume_with_handler loop
// and drops the Consumer (and its channel) inside the task.
if let Some(handle) = self.task_handle.write().await.take() {
handle.abort();
let _ = handle.await; // wait for abort to complete
}
// Clean up any consumer that wasn't taken by the task (e.g. if task never started)
if let Some(consumer) = self.consumer.write().await.take() {
// Consumer will be dropped and connection closed
drop(consumer);
}

View File

@@ -2,6 +2,12 @@
//!
//! Main service orchestrator that coordinates sensor management
//! and rule lifecycle listening.
//!
//! Shutdown follows the same pattern as the worker service:
//! 1. Deregister worker (mark inactive, stop receiving new work)
//! 2. Stop heartbeat
//! 3. Stop sensor processes with configurable timeout
//! 4. Close MQ and DB connections
use crate::rule_lifecycle_listener::RuleLifecycleListener;
use crate::sensor_manager::SensorManager;
@@ -12,6 +18,7 @@ use attune_common::db::Database;
use attune_common::mq::MessageQueue;
use sqlx::PgPool;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{error, info, warn};
@@ -29,7 +36,7 @@ struct SensorServiceInner {
rule_lifecycle_listener: Arc<RuleLifecycleListener>,
sensor_worker_registration: Arc<RwLock<SensorWorkerRegistration>>,
heartbeat_interval: u64,
running: Arc<RwLock<bool>>,
heartbeat_running: Arc<RwLock<bool>>,
}
impl SensorService {
@@ -112,18 +119,18 @@ impl SensorService {
rule_lifecycle_listener,
sensor_worker_registration: Arc::new(RwLock::new(sensor_worker_registration)),
heartbeat_interval,
running: Arc::new(RwLock::new(false)),
heartbeat_running: Arc::new(RwLock::new(false)),
}),
})
}
/// Start the sensor service
///
/// Spawns background tasks (heartbeat, rule listener, sensor manager) and returns.
/// The caller is responsible for blocking on shutdown signals and calling `stop()`.
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
@@ -152,37 +159,48 @@ impl SensorService {
info!("Sensor manager started");
// Start heartbeat loop
*self.inner.heartbeat_running.write().await = true;
let registration = self.inner.sensor_worker_registration.clone();
let heartbeat_interval = self.inner.heartbeat_interval;
let running = self.inner.running.clone();
let heartbeat_running = self.inner.heartbeat_running.clone();
tokio::spawn(async move {
while *running.read().await {
tokio::time::sleep(tokio::time::Duration::from_secs(heartbeat_interval)).await;
let mut ticker = tokio::time::interval(Duration::from_secs(heartbeat_interval));
loop {
ticker.tick().await;
if !*heartbeat_running.read().await {
info!("Heartbeat loop stopping");
break;
}
if let Err(e) = registration.read().await.heartbeat().await {
error!("Failed to send sensor worker heartbeat: {}", e);
}
}
info!("Heartbeat loop stopped");
});
// Wait until stopped
while *self.inner.running.read().await {
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
}
info!("Sensor Service stopped");
info!("Sensor Service started successfully");
Ok(())
}
/// Stop the sensor service
/// Stop the sensor service gracefully
///
/// Shutdown order (mirrors worker service pattern):
/// 1. Deregister worker (mark inactive to stop being scheduled for new work)
/// 2. Stop heartbeat
/// 3. Stop sensor processes with timeout
/// 4. Stop rule lifecycle listener
/// 5. Close MQ and DB connections
pub async fn stop(&self) -> Result<()> {
info!("Stopping Sensor Service");
info!("Stopping Sensor Service - initiating graceful shutdown");
// Mark as not running
*self.inner.running.write().await = false;
// Deregister sensor worker
info!("Deregistering sensor worker...");
// 1. Deregister sensor worker first to stop receiving new work
info!("Marking sensor worker as inactive to stop receiving new work");
if let Err(e) = self
.inner
.sensor_worker_registration
@@ -194,25 +212,51 @@ impl SensorService {
error!("Failed to deregister sensor worker: {}", e);
}
// Stop rule lifecycle listener
// 2. Stop heartbeat
info!("Stopping heartbeat updates");
*self.inner.heartbeat_running.write().await = false;
// Wait a bit for heartbeat loop to notice the flag
tokio::time::sleep(Duration::from_millis(100)).await;
// 3. Stop sensor processes with timeout
let shutdown_timeout = self
.inner
.config
.sensor
.as_ref()
.map(|s| s.shutdown_timeout)
.unwrap_or(30);
info!(
"Waiting up to {} seconds for sensor processes to stop",
shutdown_timeout
);
let sensor_manager = self.inner.sensor_manager.clone();
let timeout_duration = Duration::from_secs(shutdown_timeout);
match tokio::time::timeout(timeout_duration, sensor_manager.stop()).await {
Ok(Ok(_)) => info!("All sensor processes stopped"),
Ok(Err(e)) => error!("Error stopping sensor processes: {}", e),
Err(_) => warn!(
"Shutdown timeout reached ({} seconds) - some sensor processes may have been interrupted",
shutdown_timeout
),
}
// 4. 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
// 5. 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
// 6. Close database connection
info!("Closing database connection...");
self.inner.db.close().await;
@@ -221,11 +265,6 @@ impl SensorService {
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
@@ -243,11 +282,6 @@ impl SensorService {
/// 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));