Migrated db connections to pool

backend/src/db.rs

@@ -1,50 +1,63 @@
 pub mod devices;
 
 use include_dir::{Dir, include_dir};
+use lazy_static::lazy_static;
 use log::{debug, error};
-use rusqlite::Connection;
+use r2d2::PooledConnection;
+use r2d2_sqlite::SqliteConnectionManager;
 use rusqlite_migration::Migrations;
-use std::result::Result;
-use std::sync::LazyLock;
+use tokio::sync::Mutex;
 
 use crate::settings;
 
 static MIGRATIONS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/database_migrations");
+static INITIALISED: Mutex<bool> = Mutex::const_new(false);
 
-// Define migrations. These are applied atomically.
-static MIGRATIONS: LazyLock<Migrations<'static>> =
-    LazyLock::new(|| Migrations::from_directory(&MIGRATIONS_DIR).unwrap());
+lazy_static! {
+    // Define migrations. These are applied atomically.
+    static ref MIGRATIONS: Migrations<'static> =
+        Migrations::from_directory(&MIGRATIONS_DIR).unwrap();
 
-pub fn init_db() -> Result<Connection, String> {
-    if settings::CONFIG.database.path.is_empty() {
-        error!("Database path not set. Make sure to define database.path in your config file.");
-        Err(format!(
-            "Database path not set. Make sure to define database.path in your config file."
-        ))
-    } else {
-        debug!("Opening database at {}.", settings::CONFIG.database.path);
+    // TODO : Move pool size to configuration file
+    static ref POOL: r2d2::Pool<SqliteConnectionManager> = r2d2::Pool::builder().max_size(10).build(r2d2_sqlite::SqliteConnectionManager::file(settings::CONFIG.database.path.as_str())).unwrap();
+}
 
-        let database_path = settings::CONFIG.database.path.clone();
+pub fn get_db_connection() -> PooledConnection<SqliteConnectionManager> {
+    let result = POOL.get();
 
-        let mut conn = match Connection::open(database_path) {
-            Ok(value) => value,
-            Err(error) => {
-                error!("Error opening database (oott.db): {error}");
-                return Err(format!("Error opening database (oott.db): {error}"));
-            }
-        };
-
-        debug!("Database open, executing migrations if needed.");
-        // Update the database schema, atomically
-        match MIGRATIONS.to_latest(&mut conn) {
-            Ok(_) => {
-                debug!("Database up to date.");
-                Ok(conn)
-            }
-            Err(error) => {
-                error!("Error updating database: {error}");
-                Err(format!("Error updating database: {error}"))
-            }
+    match result {
+        Ok(value) => value,
+        Err(error) => {
+            error!("Error obtaining database connection from the pool: {error}");
+            panic!("Error obtaining database connection from the pool: {error}");
         }
     }
 }
+
+pub async fn init_db() -> Result<(), String> {
+    let mut initialised = INITIALISED.lock().await;
+    if *initialised {
+        return Ok(());
+    }
+
+    debug!("Getting database connection");
+
+    let mut conn = get_db_connection();
+
+    debug!("Executing database migrations if needed.");
+
+    let result = match MIGRATIONS.to_latest(&mut conn) {
+        Ok(_) => {
+            debug!("Database up to date.");
+            Ok(())
+        }
+        Err(error) => {
+            error!("Error updating database: {error}");
+            Err(format!("Error updating database: {error}"))
+        }
+    };
+
+    *initialised = true;
+
+    result
+}

backend/src/db/devices.rs

@@ -1,11 +1,12 @@
 use log::{debug, error};
-use rusqlite::{Connection, params};
-use std::sync::MutexGuard;
+use rusqlite::params;
 
-use crate::model::devices::Device;
+use crate::{db, model::devices::Device};
 
 // Read device from its MAC address
-pub fn read(conn: MutexGuard<Connection>, mac_address: String) -> Option<Device> {
+pub fn read(mac_address: String) -> Option<Device> {
+    let conn = db::get_db_connection();
+
     let result: Result<Device, rusqlite::Error> = conn.query_one(
         "SELECT mac_address, ipv4_address, vendor, last_seen FROM devices WHERE mac_address=?1",
         params![mac_address],
@@ -33,7 +34,9 @@ pub fn read(conn: MutexGuard<Connection>, mac_address: String) -> Option<Device>
     }
 }
 
-pub fn insert(conn: MutexGuard<Connection>, device: Device) -> Result<(), String> {
+pub fn insert(device: Device) -> Result<(), String> {
+    let conn = db::get_db_connection();
+
     match conn.execute(
         "INSERT INTO devices (mac_address, ipv4_address, vendor, last_seen) VALUES (?1, ?2, ?3, ?4)",
         params![device.mac_address, device.ipv4_address, device.vendor, device.last_seen]) {
@@ -48,7 +51,8 @@ pub fn insert(conn: MutexGuard<Connection>, device: Device) -> Result<(), String
     }
 }
 
-pub fn update(conn: MutexGuard<Connection>, device: Device) -> Result<(), String> {
+pub fn update(device: Device) -> Result<(), String> {
+    let conn = db::get_db_connection();
     match conn.execute(
         "UPDATE devices SET ipv4_address=?1, vendor=?2, last_seen=?3 WHERE mac_address=?4",
         params![

backend/src/main.rs

@@ -31,6 +31,9 @@ async fn main() -> Result<(), String> {
     // Now onto the important stuff
     info!("Starting up oott");
 
+    // Initialize database
+    db::init_db().await?;
+
     // Start the device scanner and web server (for API and UI) in parallel
     tokio::join!(scanner::scan(), web_server::serve()).0
 }

backend/src/scanner.rs

@@ -3,13 +3,9 @@ use crate::device_finders;
 use crate::events;
 use crate::settings::CONFIG;
 use log::{debug, info};
-use std::sync::{Arc, Mutex};
 use tokio::time::{Duration, sleep};
 
 pub async fn scan() -> Result<(), String> {
-    // Get database connection - thread protected
-    let db_conn = Arc::new(Mutex::new(db::init_db().unwrap()));
-
     loop {
         // Find online devices via ARP
         let devices = device_finders::arp::find(CONFIG.networking.interface.to_string()).await?;
@@ -21,13 +17,8 @@ pub async fn scan() -> Result<(), String> {
         for device in devices.iter() {
             debug!("Online device found {}", device);
 
-            // Using just one connection for now, need to update if moved DB portion to multi-thread
-            // need to change to a clone if we need multiple threads
-            let db_conn_clone = Arc::clone(&db_conn);
-
             // Read device from database
-            let recorded_device_result =
-                db::devices::read(db_conn_clone.lock().unwrap(), device.mac_address.clone());
+            let recorded_device_result = db::devices::read(device.mac_address.clone());
 
             match recorded_device_result {
                 Some(recorded_device) => {
@@ -36,7 +27,7 @@ pub async fn scan() -> Result<(), String> {
                         "Device found in database {}. Updating to {}.",
                         recorded_device, device
                     );
-                    db::devices::update(db_conn_clone.lock().unwrap(), device.clone())?;
+                    db::devices::update(device.clone())?;
                     events::trigger_existing_device(recorded_device, device.clone()).ok(); // Ignoring errors here, do not stop loop if notification delivery fails
                 }
                 None => {
@@ -46,7 +37,7 @@ pub async fn scan() -> Result<(), String> {
                         device.mac_address
                     );
 
-                    db::devices::insert(db_conn_clone.lock().unwrap(), device.clone())?;
+                    db::devices::insert(device.clone())?;
                     events::trigger_new_device(device.clone()).ok(); // Ignoring errors here, do not stop loop if notification delivery fails
                 }
             };