From c343888dcd04ce4ceffb454ae4d39c521de40aa0 Mon Sep 17 00:00:00 2001
From: Krisztian Kovacs <krisztian@equilibrium.co>
Date: Mon, 25 Apr 2022 17:22:23 +0200
Subject: [PATCH] feat(storage): use a connection pool

Instead of re-opening the Sqlite database for each request Storage now
contains a connection pool.
---
 Cargo.lock                       | 32 +++++++++++++++++++++++++++
 crates/pathfinder/Cargo.toml     |  2 ++
 crates/pathfinder/src/storage.rs | 37 +++++++++++++-------------------
 3 files changed, 49 insertions(+), 22 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7c624f6b75..2b5ad75149 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2091,6 +2091,8 @@ dependencies = [
  "mockall",
  "num-bigint 0.4.3",
  "pretty_assertions",
+ "r2d2",
+ "r2d2_sqlite",
  "reqwest",
  "rusqlite",
  "semver 1.0.7",
@@ -2391,6 +2393,27 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "r2d2"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "545c5bc2b880973c9c10e4067418407a0ccaa3091781d1671d46eb35107cb26f"
+dependencies = [
+ "log",
+ "parking_lot 0.11.2",
+ "scheduled-thread-pool",
+]
+
+[[package]]
+name = "r2d2_sqlite"
+version = "0.20.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6fdc8e4da70586127893be32b7adf21326a4c6b1aba907611edf467d13ffe895"
+dependencies = [
+ "r2d2",
+ "rusqlite",
+]
+
 [[package]]
 name = "radium"
 version = "0.6.2"
@@ -2676,6 +2699,15 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "scheduled-thread-pool"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc6f74fd1204073fa02d5d5d68bec8021be4c38690b61264b2fdb48083d0e7d7"
+dependencies = [
+ "parking_lot 0.11.2",
+]
+
 [[package]]
 name = "scoped-tls"
 version = "1.0.0"
diff --git a/crates/pathfinder/Cargo.toml b/crates/pathfinder/Cargo.toml
index 448dfa8792..a79c81ac5b 100644
--- a/crates/pathfinder/Cargo.toml
+++ b/crates/pathfinder/Cargo.toml
@@ -30,6 +30,8 @@ home = "0.5.3"
 jsonrpsee = { version = "0.11.0", features = ["server"] }
 lazy_static = "1.4.0"
 num-bigint = { version = "0.4.3", features = ["serde"] }
+r2d2 = "0.8.9"
+r2d2_sqlite = "0.20.0"
 reqwest = { version = "0.11.4", features = ["json"] }
 rusqlite = { version = "0.27.0", features = ["bundled"] }
 semver = "1.0.7"
diff --git a/crates/pathfinder/src/storage.rs b/crates/pathfinder/src/storage.rs
index 235c81db01..47c7e184e1 100644
--- a/crates/pathfinder/src/storage.rs
+++ b/crates/pathfinder/src/storage.rs
@@ -21,6 +21,8 @@ pub use state::{
 };
 
 use anyhow::Context;
+use r2d2::Pool;
+use r2d2_sqlite::SqliteConnectionManager;
 use rusqlite::Connection;
 
 /// Indicates database is non-existant.
@@ -30,6 +32,8 @@ const DB_VERSION_CURRENT: u32 = 12;
 /// Sqlite key used for the PRAGMA user version.
 const VERSION_KEY: &str = "user_version";
 
+type PooledConnection = r2d2::PooledConnection<SqliteConnectionManager>;
+
 /// Used to create [Connection's](Connection) to the pathfinder database.
 ///
 /// Intended usage:
@@ -38,17 +42,12 @@ const VERSION_KEY: &str = "user_version";
 /// - Use [Storage::connection] to create connection's to the database, which can in turn
 ///   be used to interact with the various [tables](self).
 #[derive(Clone)]
-pub struct Storage(std::sync::Arc<Inner>);
+pub struct Storage(Inner);
 
+#[derive(Clone)]
 struct Inner {
     database_path: PathBuf,
-    /// Required to keep the in-memory variant alive. Sqlite drops in-memory databases
-    /// as soon as all living connections are dropped, so we prevent this by storing
-    /// a keep-alive connection.
-    ///
-    /// [Connection] is !Sync so we wrap it in [Mutex] to get sync back.
-    #[cfg(test)]
-    _keep_alive: Mutex<Connection>,
+    pool: Pool<SqliteConnectionManager>,
 }
 
 impl Storage {
@@ -59,31 +58,25 @@ impl Storage {
     ///
     /// May be cloned safely.
     pub fn migrate(database_path: PathBuf) -> anyhow::Result<Self> {
-        let mut conn = Self::open_connection(&database_path)?;
+        let manager = SqliteConnectionManager::file(&database_path);
+        let pool = Pool::builder().build(manager)?;
+
+        let mut conn = pool.get()?;
         migrate_database(&mut conn).context("Migrate database")?;
 
-        #[cfg(not(test))]
-        let inner = Inner { database_path };
-        #[cfg(test)]
         let inner = Inner {
             database_path,
-            _keep_alive: Mutex::new(conn),
+            pool,
         };
 
-        let storage = Storage(std::sync::Arc::new(inner));
+        let storage = Storage(inner);
 
         Ok(storage)
     }
 
     /// Returns a new Sqlite [Connection] to the database.
-    pub fn connection(&self) -> anyhow::Result<Connection> {
-        Self::open_connection(&self.0.database_path)
-    }
-
-    /// Opens a connection the given database path.
-    fn open_connection(database_path: &Path) -> anyhow::Result<Connection> {
-        // TODO: think about flags?
-        let conn = Connection::open(database_path)?;
+    pub fn connection(&self) -> anyhow::Result<PooledConnection> {
+        let conn = self.0.pool.get()?;
         Ok(conn)
     }