re-uploading work

crates/worker/src/heartbeat.rs

@@ -0,0 +1,140 @@
+//! Heartbeat Module
+//!
+//! Manages periodic heartbeat updates to keep the worker's status fresh in the database.
+
+use attune_common::error::Result;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::RwLock;
+use tokio::time;
+use tracing::{debug, error, info, warn};
+
+use crate::registration::WorkerRegistration;
+
+/// Heartbeat manager for worker status updates
+pub struct HeartbeatManager {
+    registration: Arc<RwLock<WorkerRegistration>>,
+    interval: Duration,
+    running: Arc<RwLock<bool>>,
+}
+
+impl HeartbeatManager {
+    /// Create a new heartbeat manager
+    ///
+    /// # Arguments
+    /// * `registration` - Worker registration instance
+    /// * `interval_secs` - Heartbeat interval in seconds
+    pub fn new(registration: Arc<RwLock<WorkerRegistration>>, interval_secs: u64) -> Self {
+        Self {
+            registration,
+            interval: Duration::from_secs(interval_secs),
+            running: Arc::new(RwLock::new(false)),
+        }
+    }
+
+    /// Start the heartbeat loop
+    ///
+    /// This spawns a background task that periodically updates the worker's heartbeat
+    /// in the database. The task will continue running until `stop()` is called.
+    pub async fn start(&self) -> Result<()> {
+        let mut running = self.running.write().await;
+        if *running {
+            warn!("Heartbeat manager is already running");
+            return Ok(());
+        }
+        *running = true;
+        drop(running);
+
+        info!(
+            "Starting heartbeat manager with interval: {:?}",
+            self.interval
+        );
+
+        let registration = self.registration.clone();
+        let interval = self.interval;
+        let running = self.running.clone();
+
+        tokio::spawn(async move {
+            let mut ticker = time::interval(interval);
+
+            loop {
+                ticker.tick().await;
+
+                // Check if we should stop
+                {
+                    let is_running = running.read().await;
+                    if !*is_running {
+                        info!("Heartbeat manager stopping");
+                        break;
+                    }
+                }
+
+                // Send heartbeat
+                let reg = registration.read().await;
+                match reg.update_heartbeat().await {
+                    Ok(_) => {
+                        debug!("Heartbeat sent successfully");
+                    }
+                    Err(e) => {
+                        error!("Failed to send heartbeat: {}", e);
+                        // Continue trying - don't break the loop on transient errors
+                    }
+                }
+            }
+
+            info!("Heartbeat manager stopped");
+        });
+
+        Ok(())
+    }
+
+    /// Stop the heartbeat loop
+    pub async fn stop(&self) {
+        info!("Stopping heartbeat manager");
+        let mut running = self.running.write().await;
+        *running = false;
+    }
+
+    /// Check if the heartbeat manager is running
+    pub async fn is_running(&self) -> bool {
+        let running = self.running.read().await;
+        *running
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::registration::WorkerRegistration;
+    use attune_common::config::Config;
+    use attune_common::db::Database;
+
+    #[tokio::test]
+    #[ignore] // Requires database
+    async fn test_heartbeat_manager() {
+        let config = Config::load().unwrap();
+        let db = Database::new(&config.database).await.unwrap();
+        let pool = db.pool().clone();
+        let mut registration = WorkerRegistration::new(pool, &config);
+        registration.register().await.unwrap();
+
+        let registration = Arc::new(RwLock::new(registration));
+        let manager = HeartbeatManager::new(registration.clone(), 1);
+
+        // Start heartbeat
+        manager.start().await.unwrap();
+        assert!(manager.is_running().await);
+
+        // Wait for a few heartbeats
+        tokio::time::sleep(Duration::from_secs(3)).await;
+
+        // Stop heartbeat
+        manager.stop().await;
+        tokio::time::sleep(Duration::from_millis(100)).await;
+        assert!(!manager.is_running().await);
+
+        // Deregister worker
+        let reg = registration.read().await;
+        reg.deregister().await.unwrap();
+    }
+}