From 303e3ffd4e31e7bad88006a14964504b151ee234 Mon Sep 17 00:00:00 2001
From: Chewing_Bever <roosensjef@gmail.com>
Date: Tue, 16 May 2023 09:00:12 +0200
Subject: [PATCH] feat: add database connection pooling

---
 Cargo.lock            | 21 +++++++++++++++++++++
 Cargo.toml            |  2 +-
 src/commands/users.rs |  4 ++--
 src/main.rs           | 14 ++++++++++----
 4 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c88ad20..c9d2c0a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -422,6 +422,7 @@ checksum = "72eb77396836a4505da85bae0712fa324b74acfe1876d7c2f7e694ef3d0ee373"
 dependencies = [
  "diesel_derives",
  "libsqlite3-sys",
+ "r2d2",
 ]
 
 [[package]]
@@ -1073,6 +1074,17 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "r2d2"
+version = "0.8.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
+dependencies = [
+ "log",
+ "parking_lot",
+ "scheduled-thread-pool",
+]
+
 [[package]]
 name = "rand"
 version = "0.8.5"
@@ -1231,6 +1243,15 @@ version = "1.0.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f91339c0467de62360649f8d3e185ca8de4224ff281f66000de5eb2a77a79041"
 
+[[package]]
+name = "scheduled-thread-pool"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
+dependencies = [
+ "parking_lot",
+]
+
 [[package]]
 name = "scopeguard"
 version = "1.1.0"
diff --git a/Cargo.toml b/Cargo.toml
index a4d6e6f..1faac2a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,4 +22,4 @@ chrono = "*"
 uuid = "*"
 poise = "0.5.5"
 async-minecraft-ping = "0.8.0"
-diesel = { version = "2.0.4", features = ["sqlite", "returning_clauses_for_sqlite_3_35"] }
+diesel = { version = "2.0.4", features = ["sqlite", "returning_clauses_for_sqlite_3_35", "r2d2"] }
diff --git a/src/commands/users.rs b/src/commands/users.rs
index df0e7e6..e164d9a 100644
--- a/src/commands/users.rs
+++ b/src/commands/users.rs
@@ -16,7 +16,7 @@ pub async fn register(
     };
 
     {
-        let mut conn = ctx.data().conn.lock().unwrap();
+        let mut conn = ctx.data().pool.get()?;
         user_insert(&mut conn, &user);
     }
 
@@ -26,7 +26,7 @@ pub async fn register(
 #[poise::command(prefix_command, slash_command)]
 pub async fn registered(ctx: Context<'_>) -> Result<(), Error> {
     let users = {
-        let mut conn = ctx.data().conn.lock().unwrap();
+        let mut conn = ctx.data().pool.get()?;
         user_all(&mut conn)
     };
 
diff --git a/src/main.rs b/src/main.rs
index e4364d8..2bc1ab9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,6 +6,8 @@ use diesel::Connection;
 use poise::serenity_prelude as serenity;
 use std::sync::Mutex;
 use std::{env::var, time::Duration};
+use diesel::r2d2::{ConnectionManager, Pool};
+use diesel::sqlite::SqliteConnection;
 
 // Types used by all command functions
 type Error = Box<dyn std::error::Error + Send + Sync>;
@@ -14,7 +16,7 @@ type Context<'a> = poise::Context<'a, Data, Error>;
 // Custom user data passed to all command functions
 pub struct Data {
     client: AffluencesClient,
-    conn: Mutex<diesel::sqlite::SqliteConnection>,
+    pool: Pool<ConnectionManager<SqliteConnection>>
 }
 
 async fn on_error(error: poise::FrameworkError<'_, Data, Error>) {
@@ -34,6 +36,12 @@ async fn on_error(error: poise::FrameworkError<'_, Data, Error>) {
     }
 }
 
+fn db_connection_pool(url: &str) -> Pool<ConnectionManager<SqliteConnection>> {
+    let manager = ConnectionManager::new(url);
+
+    Pool::builder().test_on_check_out(true).build(manager).expect("oops")
+}
+
 #[tokio::main]
 async fn main() {
     // FrameworkOptions contains all of poise's configuration option in one struct
@@ -95,9 +103,7 @@ async fn main() {
                 poise::builtins::register_globally(ctx, &framework.options().commands).await?;
                 Ok(Data {
                     client: AffluencesClient::new(),
-                    conn: Mutex::new(
-                        diesel::sqlite::SqliteConnection::establish("affy.db").unwrap(),
-                    ),
+                    pool: db_connection_pool("affy.db")
                 })
             })
         })