Skip to content

Commit

Permalink
ledger: Add initial connection counter.
Browse files Browse the repository at this point in the history
  • Loading branch information
ceyhunsen committed Sep 5, 2024
1 parent ac603c4 commit 63100d0
Showing 1 changed file with 63 additions and 13 deletions.
76 changes: 63 additions & 13 deletions src/ledger/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,9 @@
//! immutable nature.
use crate::utils;
use rusqlite::Connection;
use rusqlite::{params, Connection};
use std::{
env,
process::Command,
sync::{Arc, Mutex},
};

Expand Down Expand Up @@ -44,24 +43,16 @@ impl Ledger {
let path = Ledger::get_database_path(path);
let _ = utils::initialize_logger();

// Check if database has another connections.
let is_open = {
let ret = Command::new("lsof")
.args([&path])
.output()
.expect("failed to execute process");

!ret.stdout.is_empty()
};

let database = Connection::open(path.clone()).unwrap();

// If database has another connections, skip clearing.
if !is_open {
if Ledger::get_database_connection_count(&database) == 0 {
tracing::trace!("Creating new database at path {path}");

Ledger::drop_tables(&database).unwrap();
Ledger::create_tables(&database).unwrap();
}
Ledger::increment_connection_count(&database);

tracing::trace!("Database connection to {path} is established");

Expand Down Expand Up @@ -95,13 +86,51 @@ impl Ledger {
}
}

/// Returns current connection count to the database. If not zero
fn get_database_connection_count(database: &Connection) -> i64 {
let count = database.query_row("SELECT count FROM connection_info", params![], |row| {
Ok(row.get::<_, i64>(0).unwrap())
});

let count = match count {
Ok(count) => count,
Err(_) => 0,
};
tracing::trace!("Current connection count: {count}");

count
}

/// Increments connection count.
fn increment_connection_count(database: &Connection) {
let count = Self::get_database_connection_count(database) + 1;
tracing::trace!("Incrementing connection count to {count}...");

database
.execute("UPDATE connection_info SET count = ?1", params![count])
.unwrap();
}

/// Decrements connection count.
fn decrement_connection_count(&self) {
let count = Self::get_database_connection_count(&self.database.lock().unwrap()) - 1;
tracing::trace!("Decrementing connection count to {count}...");

self.database
.lock()
.unwrap()
.execute("UPDATE connection_info SET count = ?1", params![count])
.unwrap();
}

fn get_database_path(path: &str) -> String {
env::temp_dir().to_str().unwrap().to_owned() + "/bitcoin_mock_rpc_" + path
}

fn drop_tables(database: &Connection) -> Result<(), rusqlite::Error> {
database.execute_batch(
"
DROP TABLE IF EXISTS connection_info;
DROP TABLE IF EXISTS blocks;
DROP TABLE IF EXISTS mempool;
DROP TABLE IF EXISTS transactions;
Expand All @@ -117,6 +146,14 @@ impl Ledger {
fn create_tables(database: &Connection) -> Result<(), rusqlite::Error> {
database.execute_batch(
"
CREATE TABLE connection_info
(
count INTEGER NOT NULL
CONSTRAINT count PRIMARY KEY
);
INSERT INTO connection_info (count) VALUES (0);
CREATE TABLE blocks
(
height INTEGER NOT NULL,
Expand Down Expand Up @@ -155,6 +192,12 @@ impl Ledger {
}
}

impl Drop for Ledger {
fn drop(&mut self) {
self.decrement_connection_count();
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand All @@ -163,4 +206,11 @@ mod tests {
fn new() {
let _should_not_panic = Ledger::new("ledger_new");
}

#[test]
fn concurrent_connections() {
let _ledger = Ledger::new("concurrent_connections");

let _ledger2 = Ledger::new("concurrent_connections");
}
}

0 comments on commit 63100d0

Please sign in to comment.