implement storage + improve code

server/src/storage/memory.rs

@@ -0,0 +1,59 @@
+use std::collections::HashMap;
+use chrono::Utc;
+use tokio::sync::Mutex;
+use async_trait::async_trait;
+
+use super::{StorageRepository, UptimeMessage, UptimeModel};
+
+pub struct InMemoryRepository {
+    agents: Mutex<HashMap<String, UptimeModel>>
+}
+
+impl InMemoryRepository {
+    pub fn new() -> Self {
+        Self {
+            agents: Default::default(),
+        }
+    }
+}
+
+#[async_trait]
+impl StorageRepository for InMemoryRepository {
+    async fn record_message(&self, agent: &str) -> anyhow::Result<()> {
+        let mut agents = self.agents.lock().await;
+        let now = Utc::now();
+
+        agents.entry(agent.to_string())
+            .and_modify(|a| {
+                a.last_seen = now;
+                a.message_count += 1;
+            })
+            .or_insert_with(|| UptimeModel {
+                id: agent.to_string(),
+                first_seen: now,
+                last_seen: now,
+                message_count: 1,
+            });
+
+        Ok(())
+    }
+
+    /*
+    async fn get_uptime(&self, agent: &str) -> anyhow::Result<Option<UptimeMessage>> {
+        let agents = self.agents.lock().await;
+
+        match agents.get(agent) {
+            Some(data) => {
+                Ok(Some(data.clone().into()))
+            }
+            None => Ok(None),
+        }
+    }
+    */
+
+    async fn get_agents(&self) -> anyhow::Result<Vec<UptimeMessage>> {
+        let agents = self.agents.lock().await;
+
+        Ok(agents.values().cloned().map(Into::into).collect())
+    }
+}

server/src/storage/mod.rs

@@ -0,0 +1,70 @@
+use serde::{Deserialize, Serialize};
+use chrono::{DateTime, Utc};
+use sqlite::SQLiteRepository;
+use std::sync::Arc;
+use memory::InMemoryRepository;
+use async_trait::async_trait;
+use common::MQTT_SEND_INTERVAL;
+
+pub mod memory;
+pub mod sqlite;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct UptimeModel {
+    pub id: String,
+    pub first_seen: DateTime<Utc>,
+    pub last_seen: DateTime<Utc>,
+    pub message_count: u64,
+}
+
+impl Into<UptimeMessage> for UptimeModel {
+    fn into(self) -> UptimeMessage {
+        let duration = Utc::now().signed_duration_since(self.first_seen);
+        let expected_messages = duration.num_seconds() as f64 / MQTT_SEND_INTERVAL as f64;
+
+        let uptime_pct = (self.message_count as f64 / expected_messages * 100.0).min(100.0);
+
+        UptimeMessage {
+            agent: self.id,
+            uptime: uptime_pct,
+            last_seen: self.last_seen,
+        }
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct UptimeMessage {
+    pub agent: String,
+    pub uptime: f64,
+    pub last_seen: DateTime<Utc>,
+}
+
+#[async_trait]
+pub trait StorageRepository: Send + Sync {
+    async fn record_message(&self, agent: &str) -> anyhow::Result<()>;
+    async fn get_agents(&self) -> anyhow::Result<Vec<UptimeMessage>>;
+}
+
+pub enum StorageStrategy {
+    InMemory,
+    SQLite(String),
+}
+
+pub struct StorageRepositoryImpl {
+    inner: Arc<dyn StorageRepository>
+}
+
+impl StorageRepositoryImpl {
+    pub fn new(strategy: StorageStrategy) -> Self {
+        let inner: Arc<dyn StorageRepository> = match strategy {
+            StorageStrategy::InMemory => Arc::new(InMemoryRepository::new()),
+            StorageStrategy::SQLite(path) => Arc::new(SQLiteRepository::new(path)),
+        };
+
+        Self { inner }
+    }
+
+    pub fn inner(&self) -> Arc<dyn StorageRepository> {
+        self.inner.clone()
+    }
+}

server/src/storage/sqlite.rs

@@ -0,0 +1,79 @@
+use async_trait::async_trait;
+use rusqlite::Connection;
+use tokio::sync::Mutex;
+use std::path::Path;
+use chrono::{DateTime, Utc};
+
+use super::{StorageRepository, UptimeMessage, UptimeModel};
+
+pub struct SQLiteRepository {
+    conn: Mutex<Connection>
+}
+
+impl SQLiteRepository {
+    pub fn new<P: AsRef<Path>>(path: P) -> Self {
+        let conn = Connection::open(path).unwrap();
+
+        conn.pragma_update(None, "journal_mode", "WAL").unwrap();
+        conn.pragma_update(None, "foreign_keys", "ON").unwrap();
+        conn.pragma_update(None, "synchronous", "NORMAL").unwrap();
+
+        conn.execute_batch(
+            r#"
+            CREATE TABLE IF NOT EXISTS agents (
+                id TEXT PRIMARY KEY,
+                first_seen TEXT NOT NULL,
+                last_seen TEXT NOT NULL,
+                message_count INTEGER NOT NULL DEFAULT 0
+            ) STRICT;
+
+            CREATE INDEX IF NOT EXISTS idx_agents_id ON agents(id);
+            CREATE INDEX IF NOT EXISTS idx_agents_times ON agents(first_seen, last_seen);
+            "#
+        ).unwrap();
+
+        Self {
+            conn: Mutex::new(conn)
+        }
+    }
+}
+
+#[async_trait]
+impl StorageRepository for SQLiteRepository {
+    async fn record_message(&self, agent: &str) -> anyhow::Result<()> {
+        let conn = self.conn.lock().await;
+        let now = Utc::now().to_rfc3339();
+
+        conn.execute(r#"
+            INSERT INTO agents (id, first_seen, last_seen, message_count)
+            VALUES (?1, ?2, ?2, 1)
+            ON CONFLICT (id) DO UPDATE SET
+                last_seen = excluded.last_seen,
+                message_count = message_count + 1;
+            "#, [agent, &now]
+        )?;
+
+        Ok(())
+    }
+
+    async fn get_agents(&self) -> anyhow::Result<Vec<UptimeMessage>> {
+        let conn = self.conn.lock().await;
+        let mut stmt = conn.prepare("SELECT id, first_seen, last_seen, message_count FROM agents")?;
+
+        let result = stmt.query_map([], |row| {
+            let first_seen: DateTime<Utc> = row.get::<_, String>(1)?.parse().unwrap();
+            let last_seen: DateTime<Utc> = row.get::<_, String>(2)?.parse().unwrap();
+
+            Ok(UptimeModel {
+                id: row.get(0)?,
+                first_seen,
+                last_seen,
+                message_count: row.get(3)?,
+            })
+        })?;
+
+        let models: Result<Vec<UptimeModel>, _> = result.collect();
+
+        Ok(models?.into_iter().map(Into::into).collect())
+    }
+}