Cleaned up rust codebase

Author: SuperB3333

src/checkpoint.rs

@@ -4,7 +4,8 @@ use std::sync::atomic::Ordering::Relaxed;
 use std::sync::LazyLock;
 use std::io::Write;
 use std::ops::Range;
-use crate::{CHECKPOINT_FILE, PROGRESS_WORKERS};
+use crate::{CHECKPOINT_FILE, END_SEED, PROGRESS_WORKERS, START_SEED, WORKER_THREADS};
+use crate::misc::split_chunks;
 
 static MAX_BUFFER: LazyLock<i32> = LazyLock::new(|| {
     2
@@ -24,6 +25,7 @@ pub fn write_checkpoints() -> Result<()> {
         });
     Ok(())
 }
-pub fn load_checkpoints() -> Result<Vec<Range<i32>>> {
-    todo!()
+pub fn load_workloads() -> Result<Vec<Range<i32>>> {
+    return Ok(split_chunks(*START_SEED..*END_SEED, *WORKER_THREADS));
+    todo!("load workloads from checkpoints file")
 }

src/macros.rs

@@ -1,4 +1,3 @@
-mod macros;
 mod algorithm;
 mod generate_csv;
 mod misc;
@@ -11,20 +10,17 @@ use crossbeam_channel::{bounded, Receiver, Sender};
 use lazy_static::lazy_static;
 use std::{
     time::{Duration, Instant},
-    io::{Write, stdout},
+    io::stdout,
     thread,
-    sync::atomic::{
-        AtomicI32
-    }
+    sync::atomic::AtomicI32
 };
 use crossterm::ExecutableCommand;
 use crossterm::terminal::{EnterAlternateScreen, LeaveAlternateScreen};
 use crate::{
-    misc::split_chunks,
     metrics::write_metrics,
     threads::*
 };
-use crate::checkpoint::write_checkpoints;
+use crate::checkpoint::{load_workloads, write_checkpoints};
 
 const STAR_COUNT: usize = 64;
 const REC_MULTIPLIER: f32 = 1.0;
@@ -39,8 +35,8 @@ lazy_static! {
     static ref END_SEED: i32 = env_int!("END_SEED", 10_000);
     static ref WORKER_THREADS: i32 = env_int!("WORKER_THREADS", 8);
     static ref WRITER_THREADS: i32 = env_int!("WRITER_THREADS", 4);
-    static ref COMMIT_COUNT: i32 = env_int!("COMMIT_COUNT", 1000);
-    static ref CHANNEL_SIZE: i32 = env_int!("CHANNEL_SIZE", 1000);
+    static ref COMMIT_COUNT: usize = env_int!("COMMIT_COUNT", 1000) as usize;
+    static ref CHANNEL_SIZE: usize = env_int!("CHANNEL_SIZE", 1000) as usize;
     static ref CHECKPOINT_FILE: String = env_str!("CHECKPOINT_FILE", "checkpoints.txt");
     static ref BENCHMARK: bool = env_int!("BENCHMARK", 0) == 1;
 
@@ -53,20 +49,19 @@ lazy_static! {
         format!("postgres://{user}:{pass}@{netloc}:{port}/{db_name}?sslmode=disable")
     };
 
-    static ref MAX_BUFFER: i32 = *CHANNEL_SIZE + *COMMIT_COUNT * *WORKER_THREADS;
+    static ref MAX_BUFFER: usize = *CHANNEL_SIZE + *COMMIT_COUNT * *WORKER_THREADS as usize;
 }
 fn main() {
     assert!(*START_SEED < *END_SEED);
     assert!(*WORKER_THREADS < *END_SEED);
     assert!(*WORKER_THREADS < MAX_WORKERS as i32);
 
-
+    // capture start time for performance evaluation
     let start = Instant::now();
 
     // Prepare thread resources
-    let all_seeds = *START_SEED..*END_SEED;
-    let workloads = split_chunks(all_seeds, *WORKER_THREADS);
-    let (entry_sender, entry_reciever): (Sender<(String, String)>, Receiver<(String, String)>) = bounded(*CHANNEL_SIZE as usize);
+    let workloads = load_workloads().unwrap();
+    let (entry_sender, entry_reciever): (Sender<(String, String)>, Receiver<(String, String)>) = bounded(*CHANNEL_SIZE);
 
     let mut work_handles = vec![];
     let mut commit_handles = vec![];
@@ -79,18 +74,17 @@ fn main() {
         }))
     }
     if !*BENCHMARK {
-        let conf = ((*DB_STR).clone(), *COMMIT_COUNT);
         // Launch database threads
         for _ in 0..*WRITER_THREADS {
             let thread_receiver = entry_reciever.clone();
-            let thread_conf = conf.clone();
             commit_handles.push(thread::spawn(move || {
-                commit_thread(thread_receiver, thread_conf)
+                commit_thread(thread_receiver)
             }))
         }
     }
     else {
         for _ in 0..*WRITER_THREADS {
+            // launch dummy db threads that will void all results
             let thread_receiver = entry_reciever.clone();
             commit_handles.push(thread::spawn(move || {
                 writer_sink(thread_receiver)

src/main.rs

@@ -18,3 +18,24 @@ pub fn split_chunks(r: Range<i32>, chunks: i32) -> Vec<Range<i32>> {
     }
     out
 }
+#[macro_export]
+macro_rules! env_int {
+    ($var:expr) => {
+        match std::env::var($var) {
+            Ok(val) => val.parse(),
+            Err(err) => Err(err)
+        }
+    };
+    ($var:expr, $default:expr) => {
+        std::env::var($var).map(|e| e.parse::<i32>().unwrap_or($default)).unwrap_or($default)
+    };
+}
+#[macro_export]
+macro_rules! env_str {
+    ($var:expr) => {
+        std::env::var($var)
+    };
+    ($var:expr, $default:expr) => {
+        std::env::var($var).unwrap_or($default.to_string())
+    }
+}

src/misc.rs

@@ -6,7 +6,7 @@ use crossbeam_channel::{Receiver, RecvTimeoutError, Sender};
 use postgres::{Client, NoTls};
 use anyhow::Result;
 use crate::generate_csv::gen_formatted;
-use crate::{COMMITTED_SEEDS, PROGRESS_WORKERS, REC_MULTIPLIER, STAR_COUNT};
+use crate::{COMMITTED_SEEDS, COMMIT_COUNT, DB_STR, PROGRESS_WORKERS, REC_MULTIPLIER, STAR_COUNT};
 use crate::misc::{COPY_PLANET, COPY_STAR};
 
 pub fn worker_thread(seeds: Range<i32>, send: Sender<(String, String)>, id: usize) -> Result<()> {
@@ -18,24 +18,19 @@ pub fn worker_thread(seeds: Range<i32>, send: Sender<(String, String)>, id: usiz
     }
     Ok(())
 }
-pub fn commit_thread(rec: Receiver<(String, String)>, config: (String, i32)) -> Result<()> {
-    let mut client = Client::connect(config.0.as_str(), NoTls)?;
-    let commit_size = config.1 as usize;
+pub fn commit_thread(rec: Receiver<(String, String)>) -> Result<()> {
+    let mut client = Client::connect(&*DB_STR.as_str(), NoTls)?;
 
-    loop {
-        let mut batch: Vec<(String, String)> = Vec::with_capacity(commit_size);
-        for _ in 0..commit_size {
+    'outer: loop {
+        let mut batch: Vec<(String, String)> = Vec::with_capacity(*COMMIT_COUNT);
+        'inner: for i in 0..*COMMIT_COUNT {
             match rec.recv_timeout(Duration::new(1, 0)) {
                 Ok(msg) => batch.push(msg),
                 Err(RecvTimeoutError::Timeout) => panic!("commit_thread: recv_timeout reached - channel stall detected (>1s lull)"),
-                Err(RecvTimeoutError::Disconnected) => break,
+                Err(RecvTimeoutError::Disconnected) => if i == 0 {break 'outer} else {break 'inner},
             }
         }
 
-        if batch.is_empty() {
-            break;
-        }
-
         let mut txn = client.transaction()?;
 
         {