feat: database reader connection pooling

Added connection pooling for queries, as well as basic configuration
options for min/max connections.

src/config.rs

@@ -22,6 +22,8 @@ pub struct Info {
 #[allow(unused)]
 pub struct Database {
     pub data_directory: String,
+    pub min_conn: u32,
+    pub max_conn: u32,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -93,6 +95,13 @@ impl Settings {
             // override with file contents
             .with_merged(config::File::with_name("config"))?
             .try_into()?;
+        // ensure connection pool size is logical
+        if settings.database.min_conn > settings.database.max_conn {
+            panic!(
+                "Database min_conn setting ({}) cannot exceed max_conn ({})",
+                settings.database.min_conn, settings.database.max_conn
+            );
+        }
         Ok(settings)
     }
 }
@@ -109,6 +118,8 @@ impl Default for Settings {
             },
             database: Database {
                 data_directory: ".".to_owned(),
+                min_conn: 4,
+                max_conn: 128,
             },
             network: Network {
                 port: 8080,

src/db.rs

@@ -1,4 +1,5 @@
 //! Event persistence and querying
+use crate::config;
 use crate::error::Result;
 use crate::event::Event;
 use crate::subscription::Subscription;
@@ -11,6 +12,8 @@ use rusqlite::Connection;
 use rusqlite::OpenFlags;
 //use std::num::NonZeroU32;
 use crate::config::SETTINGS;
+use r2d2;
+use r2d2_sqlite::SqliteConnectionManager;
 use rusqlite::limits::Limit;
 use rusqlite::types::ToSql;
 use std::path::Path;
@@ -18,6 +21,8 @@ use std::thread;
 use std::time::Instant;
 use tokio::task;
 
+pub type SqlitePool = r2d2::Pool<r2d2_sqlite::SqliteConnectionManager>;
+
 /// Database file
 const DB_FILE: &str = "nostr.db";
 
@@ -94,6 +99,23 @@ FOREIGN KEY(event_id) REFERENCES event(id) ON UPDATE RESTRICT ON DELETE CASCADE
 CREATE INDEX IF NOT EXISTS pubkey_ref_index ON pubkey_ref(referenced_pubkey);
 "##;
 
+pub fn build_read_pool() -> SqlitePool {
+    info!("Build a connection pool");
+    let config = config::SETTINGS.read().unwrap();
+    let db_dir = &config.database.data_directory;
+    let full_path = Path::new(db_dir).join(DB_FILE);
+    let manager = SqliteConnectionManager::file(&full_path)
+        .with_flags(OpenFlags::SQLITE_OPEN_READ_ONLY)
+        .with_init(|c| c.execute_batch(STARTUP_SQL));
+    let pool: SqlitePool = r2d2::Pool::builder()
+        .test_on_check_out(true) // no noticeable performance hit
+        .min_idle(Some(config.database.min_conn))
+        .max_size(config.database.max_conn)
+        .build(manager)
+        .unwrap();
+    return pool;
+}
+
 /// Upgrade DB to latest version, and execute pragma settings
 pub fn upgrade_db(conn: &mut Connection) -> Result<()> {
     // check the version.
@@ -615,22 +637,17 @@ fn query_from_sub(sub: &Subscription) -> (String, Vec<Box<dyn ToSql>>) {
 /// query is immediately aborted.
 pub async fn db_query(
     sub: Subscription,
+    conn: r2d2::PooledConnection<r2d2_sqlite::SqliteConnectionManager>,
     query_tx: tokio::sync::mpsc::Sender<QueryResult>,
     mut abandon_query_rx: tokio::sync::oneshot::Receiver<()>,
 ) {
     task::spawn_blocking(move || {
-        let config = SETTINGS.read().unwrap();
-        let db_dir = &config.database.data_directory;
-        let full_path = Path::new(db_dir).join(DB_FILE);
-
-        let conn = Connection::open_with_flags(&full_path, OpenFlags::SQLITE_OPEN_READ_ONLY)?;
-        debug!("opened database for reading");
         debug!("going to query for: {:?}", sub);
         let mut row_count: usize = 0;
         let start = Instant::now();
         // generate SQL query
         let (q, p) = query_from_sub(&sub);
-        // execute the query
+        // execute the query. Don't cache, since queries vary so much.
         let mut stmt = conn.prepare(&q)?;
         let mut event_rows = stmt.query(rusqlite::params_from_iter(p))?;
         while let Some(row) = event_rows.next()? {

src/main.rs

@@ -43,6 +43,7 @@ fn db_from_args(args: Vec<String>) -> Option<String> {
 /// Handle arbitrary HTTP requests, including for WebSocket upgrades.
 async fn handle_web_request(
     mut request: Request<Body>,
+    pool: db::SqlitePool,
     remote_addr: SocketAddr,
     broadcast: Sender<Event>,
     event_tx: tokio::sync::mpsc::Sender<Event>,
@@ -83,8 +84,9 @@ async fn handle_web_request(
                                     Some(config),
                                 )
                                 .await;
+
                                 tokio::spawn(nostr_server(
-                                    ws_stream, broadcast, event_tx, shutdown,
+                                    pool, ws_stream, broadcast, event_tx, shutdown,
                                 ));
                             }
                             Err(e) => println!(
@@ -188,6 +190,8 @@ fn main() -> Result<(), Error> {
     rt.block_on(async {
         let settings = config::SETTINGS.read().unwrap();
         info!("listening on: {}", socket_addr);
+        // build a connection pool for sqlite connections
+        let pool = db::build_read_pool();
         // all client-submitted valid events are broadcast to every
         // other client on this channel.  This should be large enough
         // to accomodate slower readers (messages are dropped if
@@ -214,6 +218,7 @@ fn main() -> Result<(), Error> {
         // A `Service` is needed for every connection, so this
         // creates one from our `handle_request` function.
         let make_svc = make_service_fn(|conn: &AddrStream| {
+            let svc_pool = pool.clone();
             let remote_addr = conn.remote_addr();
             let bcast = bcast_tx.clone();
             let event = event_tx.clone();
@@ -223,6 +228,7 @@ fn main() -> Result<(), Error> {
                 Ok::<_, Infallible>(service_fn(move |request: Request<Body>| {
                     handle_web_request(
                         request,
+                        svc_pool.clone(),
                         remote_addr,
                         bcast.clone(),
                         event.clone(),
@@ -246,6 +252,7 @@ fn main() -> Result<(), Error> {
 /// Handle new client connections.  This runs through an event loop
 /// for all client communication.
 async fn nostr_server(
+    pool: db::SqlitePool,
     ws_stream: WebSocketStream<Upgraded>,
     broadcast: Sender<Event>,
     event_tx: tokio::sync::mpsc::Sender<Event>,
@@ -337,7 +344,9 @@ async fn nostr_server(
                             Ok(()) => {
                                 running_queries.insert(s.id.to_owned(), abandon_query_tx);
                                 // start a database query
-                                db::db_query(s, query_tx.clone(), abandon_query_rx).await;
+                                // show pool stats
+                                debug!("DB pool stats: {:?}", pool.state());
+                                db::db_query(s, pool.get().expect("could not get connection"), query_tx.clone(), abandon_query_rx).await;
                             },
                             Err(e) => {
                                 info!("Subscription error: {}", e);