From dabeba43e20c213173a0190fc29817f602ff60fb Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Sun, 22 Mar 2020 18:21:40 -0500
Subject: [PATCH] Make connections configurable

---
 README.md     | 2 ++
 src/config.rs | 6 ++++++
 src/db.rs     | 8 ++++++--
 src/main.rs   | 8 ++------
 src/notify.rs | 7 +++++--
 5 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 462b5b8..bbdf73a 100644
--- a/README.md
+++ b/README.md
@@ -63,6 +63,7 @@ HTTPS=false
 DATABASE_URL=
 PRETTY_LOG=true
 PUBLISH_BLOCKS=false
+CONNECTIONS_PER_CORE=2 # how many postgres connections should be made for each core on the host
 ```
 To run this server in production, you'll likely want to set most of them
 ```env
@@ -76,6 +77,7 @@ HTTPS=true
 DATABASE_URL=postgres://pg_user:pg_pass@pg_host:pg_port/pg_database
 PRETTY_LOG=false
 PUBLISH_BLOCKS=true
+CONNECTIONS_PER_CORE=4
 ```
 
 ### Contributing
diff --git a/src/config.rs b/src/config.rs
index 6483c1a..18b7395 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -17,6 +17,7 @@ pub struct Config {
     database_url: String,
     pretty_log: bool,
     publish_blocks: bool,
+    connections_per_core: u32,
 }
 
 pub enum UrlKind {
@@ -44,6 +45,7 @@ impl Config {
             .set_default("https", false)?
             .set_default("pretty_log", true)?
             .set_default("publish_blocks", false)?
+            .set_default("connections_per_core", 2)?
             .merge(Environment::new())?;
 
         Ok(config.try_into()?)
@@ -53,6 +55,10 @@ impl Config {
         self.pretty_log
     }
 
+    pub fn connections_per_core(&self) -> u32 {
+        self.connections_per_core
+    }
+
     pub fn validate_signatures(&self) -> bool {
         self.validate_signatures
     }
diff --git a/src/db.rs b/src/db.rs
index 6c663e5..c5e7248 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -22,11 +22,15 @@ pub struct Db {
 }
 
 impl Db {
-    pub async fn build(config: Config) -> Result<Self, MyError> {
+    pub async fn build(config: &crate::config::Config) -> Result<Self, MyError> {
+        let cpus: u32 = num_cpus::get().try_into()?;
+        let max_conns = cpus * config.connections_per_core();
+
+        let config: Config = config.database_url().parse()?;
         let manager = PostgresConnectionManager::new(config, NoTls);
 
         let pool = bb8::Pool::builder()
-            .max_size((num_cpus::get() * 4).try_into()?)
+            .max_size(max_conns)
             .build(manager)
             .await?;
 
diff --git a/src/main.rs b/src/main.rs
index 8b2f4b9..a9e8bf0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,7 +3,6 @@ use actix_web::{
     middleware::Logger,
     web, App, HttpResponse, HttpServer,
 };
-use bb8_postgres::tokio_postgres;
 use log::error;
 use std::{
     io::BufWriter,
@@ -33,7 +32,6 @@ use self::{
     db::Db,
     error::MyError,
     jobs::{create_server, create_workers},
-    notify::notify_loop,
     state::State,
     templates::statics::StaticFile,
     webfinger::RelayResolver,
@@ -90,8 +88,7 @@ async fn main() -> Result<(), anyhow::Error> {
         env_logger::init();
     }
 
-    let pg_config: tokio_postgres::Config = config.database_url().parse()?;
-    let db = Db::build(pg_config.clone()).await?;
+    let db = Db::build(&config).await?;
 
     let args = Args::new();
 
@@ -109,11 +106,10 @@ async fn main() -> Result<(), anyhow::Error> {
     let state = State::hydrate(config.clone(), &db).await?;
 
     rehydrate::spawn(db.clone(), state.clone());
+    notify::spawn(state.clone(), &config)?;
 
     let job_server = create_server(db.clone());
 
-    let _ = notify_loop(state.clone(), pg_config.clone());
-
     let bind_address = config.bind_address();
     HttpServer::new(move || {
         create_workers(state.clone(), job_server.clone());
diff --git a/src/notify.rs b/src/notify.rs
index b69fe28..6db4ec6 100644
--- a/src/notify.rs
+++ b/src/notify.rs
@@ -1,4 +1,4 @@
-use crate::{db::listen, state::State};
+use crate::{db::listen, error::MyError, state::State};
 use activitystreams::primitives::XsdAnyUri;
 use actix::clock::{delay_for, Duration};
 use bb8_postgres::tokio_postgres::{tls::NoTls, AsyncMessage, Config, Notification};
@@ -43,7 +43,9 @@ async fn handle_notification(state: &State, notif: Notification) {
     };
 }
 
-pub fn notify_loop(state: State, config: Config) {
+pub fn spawn(state: State, config: &crate::config::Config) -> Result<(), MyError> {
+    let config: Config = config.database_url().parse()?;
+
     actix::spawn(async move {
         let mut client;
 
@@ -93,4 +95,5 @@ pub fn notify_loop(state: State, config: Config) {
             warn!("Restarting listener task");
         }
     });
+    Ok(())
 }