Structure config files

src/lib.rs

@@ -1,6 +1,5 @@
-use anyhow::Result;
+use anyhow::{Context, Result};
 use indoc::writedoc;
-use std::str::FromStr;
 use nix::errno::Errno;
 use nix::sys::signal::{self, Signal};
 use nix::unistd::Pid;
@@ -16,6 +15,7 @@ use std::io::Write;
 use std::os::unix::process::CommandExt;
 use std::path::{Path, PathBuf};
 use std::process::{Child, Command};
+use std::str::FromStr;
 use std::time::Duration;
 use std::time::SystemTime;
 use tempfile::NamedTempFile;
@@ -47,13 +47,102 @@ struct PgBouncerRow {
     peer_offset: i32,
 }
 
+struct DatabaseRow {
+    groupid: i32,
+    name: String,
+    host: String,
+    port: i32,
+}
+
 impl PgBouncerRow {
     fn shared_dir(&self) -> PathBuf {
         Path::new(BASE_DIR).join(self.groupid.to_string())
     }
 
-    fn dir(&self, index: i32) -> PathBuf {
-        self.shared_dir().join(index.to_string())
+    fn peers_ini(&self) -> PathBuf {
+        self.shared_dir().join("peers.ini")
+    }
+
+    fn databases_ini(&self) -> PathBuf {
+        self.shared_dir().join("databases.ini")
+    }
+
+    fn users_ini(&self) -> PathBuf {
+        self.shared_dir().join("users.ini")
+    }
+
+    fn settings_ini(&self) -> PathBuf {
+        self.shared_dir().join("settings.ini")
+    }
+
+    fn create_shared_configs(&self, processes: &[PgBouncer]) -> Result<()> {
+        let shared_dir = self.shared_dir();
+        create_dir_all(&shared_dir)?;
+
+        // peers
+        let mut f = NamedTempFile::new_in(TEMP_DIR)?;
+        for proc in processes.iter() {
+            let dir = proc.dir();
+            let dir_str: &str = dir.to_str().context("non utf8")?;
+            write!(
+                f,
+                "{} = host={} port={}\n",
+                proc.peer_id(),
+                dir_str,
+                proc.port
+            )?;
+        }
+        let peers_ini = self.peers_ini();
+        f.persist(&peers_ini)?;
+
+        // databases
+        // TODO: get from database
+        let databases = vec![DatabaseRow {
+            groupid: 1,
+            name: "*".into(),
+            // TODO: use unix_socket_directories value by default
+            host: "127.0.0.1".into(),
+            // TODO: Use postmaster port
+            port: 9700,
+        }];
+
+        let mut f = NamedTempFile::new_in(TEMP_DIR)?;
+        for db in databases {
+            write!(f, "{} = host={} port={}\n", db.name, db.host, db.port)?;
+        }
+        let databases_ini = self.databases_ini();
+        f.persist(&databases_ini)?;
+
+        // users
+        let mut f = NamedTempFile::new_in(TEMP_DIR)?;
+        write!(f, "\n")?;
+        let users_ini = self.users_ini();
+        f.persist(&users_ini)?;
+
+        // settings
+        let mut f = NamedTempFile::new_in(TEMP_DIR)?;
+        writedoc!(
+            f,
+            r#"
+            so_reuseport = 1
+            listen_addr = 127.0.0.1
+            listen_port = 6432
+
+            auth_type = trust
+            admin_users = jelte
+            auth_file = /home/jelte/opensource/pgbouncer/auth_file.conf
+            query_wait_timeout=5
+            default_pool_size=30
+
+            pool_mode = transaction
+            log_connections=1
+            log_disconnections=1
+            "#
+        )?;
+        let settings_ini = self.settings_ini();
+        f.persist(&settings_ini)?;
+
+        Ok(())
     }
 }
 
@@ -113,17 +202,19 @@ struct PgBouncer {
     child: Option<Child>,
     pid: Option<Pid>,
     started_at: Option<SystemTime>,
+    peer_offset: i32,
 }
 
 impl PgBouncer {
-    fn new(groupid: i32, index: i32) -> PgBouncer {
+    fn new(row: PgBouncerRow, index: i32) -> PgBouncer {
         PgBouncer {
-            groupid,
+            groupid: row.groupid,
             index,
             port: 6432,
             child: None,
             pid: None,
             started_at: None,
+            peer_offset: row.peer_offset,
         }
     }
 
@@ -207,23 +298,45 @@ impl PgBouncer {
         Ok(())
     }
 
+    fn shared_dir(&self) -> PathBuf {
+        Path::new(BASE_DIR).join(self.groupid.to_string())
+    }
+
     fn dir(&self) -> PathBuf {
-        Path::new(BASE_DIR)
-            .join(self.groupid.to_string())
-            .join(self.index.to_string())
+        self.shared_dir().join(self.index.to_string())
     }
 
     fn ini(&self) -> PathBuf {
         self.dir().join("config.ini")
     }
 
+    fn peers_ini(&self) -> PathBuf {
+        self.shared_dir().join("peers.ini")
+    }
+
+    fn databases_ini(&self) -> PathBuf {
+        self.shared_dir().join("databases.ini")
+    }
+
+    fn users_ini(&self) -> PathBuf {
+        self.shared_dir().join("users.ini")
+    }
+
+    fn settings_ini(&self) -> PathBuf {
+        self.shared_dir().join("settings.ini")
+    }
+
     fn pidfile(&self) -> PathBuf {
         self.dir().join("pgbouncer.pid")
     }
 
     fn socket(&self) -> PathBuf {
         self.dir().join(format!(".s.PGSQL.{}", self.port))
     }
+
+    fn peer_id(&self) -> i32 {
+        self.index + self.peer_offset
+    }
 }
 
 struct ManagerState {
@@ -301,14 +414,15 @@ fn cleanup_unneeded_processes(state: &mut ManagerState) -> Result<()> {
 }
 
 fn update_pgbouncer_group(state: &mut ManagerState, row: &PgBouncerRow) -> Result<()> {
-    let shared_dir = row.shared_dir();
-    create_dir_all(&shared_dir)?;
     let processes = state.processes.entry(row.groupid).or_insert_with(|| vec![]);
 
     // cleanup_unneeded_processes()
     for i in processes.len() as i32..row.num_processes {
-        processes.push(PgBouncer::new(row.groupid, i + 1))
+        processes.push(PgBouncer::new(row.clone(), i + 1))
     }
+
+    row.create_shared_configs(processes)?;
+
     for proc in processes.iter() {
         create_config(proc, row.peer_offset)?;
     }
@@ -329,34 +443,37 @@ fn create_config(proc: &PgBouncer, peer_offset: i32) -> Result<()> {
         proc.groupid,
         ini.to_string_lossy(),
     );
-    let peer_id = proc.index + peer_offset;
-
+    let peers_ini = proc.peers_ini();
+    let databases_ini = proc.databases_ini();
+    let settings_ini = proc.settings_ini();
+    let users_ini = proc.users_ini();
+
+    let dir_str = dir.to_str().context("no utf8")?;
+    let peers_ini_str = peers_ini.to_str().context("no utf8")?;
+    let databases_ini_str = databases_ini.to_str().context("no utf8")?;
+    let users_ini_str = users_ini.to_str().context("no utf8")?;
+    let settings_ini_str = settings_ini.to_str().context("no utf8")?;
+    let peer_id = proc.peer_id();
     // Atomically write to the file
     let mut f = NamedTempFile::new_in(TEMP_DIR)?;
-    let dir_str = dir.to_str().unwrap();
     writedoc!(
         f,
         r#"
+        [peers]
+        %include {peers_ini_str}
+
         [databases]
-        postgres = host=localhost
+        %include {databases_ini_str}
+
+        [users]
+        %include {users_ini_str}
 
         [pgbouncer]
-        so_reuseport = 1
+        %include {settings_ini_str}
+
         unix_socket_dir = {dir_str}
         peer_id = {peer_id}
 
-        listen_addr = 127.0.0.1
-        listen_port = 6432
-
-        auth_type = trust
-        admin_users = jelte
-        auth_file = /home/jelte/opensource/pgbouncer/auth_file.conf
-        query_wait_timeout=5
-        default_pool_size=30
-
-        pool_mode = transaction
-        log_connections=1
-        log_disconnections=1
         pidfile = {dir_str}/pgbouncer.pid
         "#,
     )?;