Skip to content

Commit

Permalink
attempt parallel pg_dump
Browse files Browse the repository at this point in the history
  • Loading branch information
dullbananas committed May 20, 2024
1 parent ced9bb5 commit 592a127
Show file tree
Hide file tree
Showing 4 changed files with 126 additions and 37 deletions.
5 changes: 3 additions & 2 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/db_schema/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ moka.workspace = true
serial_test = { workspace = true }
pretty_assertions = { workspace = true }
diff = "0.1.13"
crossbeam-channel = "0.5.13"

[package.metadata.cargo-machete]
ignored = ["strum"]
23 changes: 17 additions & 6 deletions crates/db_schema/src/schema_setup.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ const REPLACEABLE_SCHEMA: &[&str] = &[
struct MigrationHarnessWrapper<'a, 'b> {
conn: &'a mut PgConnection,
options: &'b Options,
#[cfg(test)]
diff_checker:diff_check::DiffChecker,
}

impl<'a, 'b> MigrationHarness<Pg> for MigrationHarnessWrapper<'a, 'b> {
Expand All @@ -59,13 +61,12 @@ impl<'a, 'b> MigrationHarness<Pg> for MigrationHarnessWrapper<'a, 'b> {

#[cfg(test)]
if self.options.enable_diff_check {
let before = diff_check::get_dump(&mut self.conn);
let before = self.diff_checker.get_dump();
self.conn.run_migration(migration)?;
self.conn.revert_migration(migration)?;
diff_check::check_dump_diff(
&mut self.conn,
self.diff_checker.check_dump_diff(
before,
&format!("migrations/{name}/down.sql"),
format!("migrations/{name}/down.sql"),
);
}

Expand All @@ -75,6 +76,10 @@ impl<'a, 'b> MigrationHarness<Pg> for MigrationHarnessWrapper<'a, 'b> {

let duration = start_time.elapsed().as_millis();
info!("{duration}ms run {name}");
#[cfg(test)]
if result.is_err() {
self.diff_checker.finish();
}

result
}
Expand Down Expand Up @@ -153,6 +158,8 @@ impl Options {
// TODO return struct with field `ran_replaceable_schema`
pub fn run(options: Options) -> LemmyResult<()> {
let db_url = SETTINGS.get_database_url();
#[cfg(test)]
let mut diff_checker = diff_check::DiffChecker::new(&db_url)?;

// Migrations don't support async connection, and this function doesn't need to be async
let mut conn =
Expand Down Expand Up @@ -198,6 +205,8 @@ pub fn run(options: Options) -> LemmyResult<()> {
let mut wrapper = MigrationHarnessWrapper {
conn,
options: &options,
#[cfg(test)]
diff_checker,
};

// * Prevent other lemmy_server processes from running this transaction simultaneously by repurposing
Expand Down Expand Up @@ -247,7 +256,7 @@ pub fn run(options: Options) -> LemmyResult<()> {
if !(options.revert && !options.redo_after_revert) {
#[cfg(test)]
if options.enable_diff_check {
let before = diff_check::get_dump(&mut wrapper.conn);
let before = wrapper.diff_checker.get_dump();
// todo move replaceable_schema dir path to let/const?
wrapper
.conn
Expand All @@ -258,7 +267,7 @@ pub fn run(options: Options) -> LemmyResult<()> {
.conn
.batch_execute("DROP SCHEMA IF EXISTS r CASCADE;")?;
// todo use different first output line in this case
diff_check::check_dump_diff(&mut wrapper.conn, before, "replaceable_schema");
wrapper.diff_checker.check_dump_diff(before, "replaceable_schema".to_owned());
}

wrapper
Expand All @@ -272,6 +281,8 @@ pub fn run(options: Options) -> LemmyResult<()> {

debug_assert_eq!(num_rows_updated, 1);
}
#[cfg(test)]
wrapper.diff_checker.finish();

Ok(())
};
Expand Down
134 changes: 105 additions & 29 deletions crates/db_schema/src/schema_setup/diff_check.rs
Original file line number Diff line number Diff line change
@@ -1,46 +1,122 @@
use diesel::{PgConnection, RunQueryDsl};
use diesel::{PgConnection, RunQueryDsl,connection::SimpleConnection,Connection};
use lemmy_utils::settings::SETTINGS;
use std::{
borrow::Cow,
collections::BTreeSet,
fmt::Write,
process::{Command, Stdio},
process::{Command, Stdio},thread,cell::OnceCell,sync::{Arc,Mutex},collections::HashMap,any::Any
};
use crossbeam_channel::{Sender, Receiver};

enum DumpAction {
Send(Sender<String>),
Compare(Receiver<String>, String),
}

pub struct DiffChecker {
snapshot_conn: PgConnection,
handles: Vec<thread::JoinHandle<()>>,
snapshot_sender: Option<Sender<(String, DumpAction)>>,
error: Receiver<Box<dyn Any + Send + 'static>>,
// todo rename to channels
//dump_receivers: Arc<Mutex<HashMap<String, Receiver<String>>>>,
}

diesel::sql_function! {
fn pg_export_snapshot() -> diesel::sql_types::Text;
}

pub fn get_dump(conn: &mut PgConnection) -> String {
/*// Required for pg_dump to see uncommitted changes from a different database connection
// The pg_dump command runs independently from `conn`, which means it can't see changes from
// an uncommitted transaction. NASA made each migration run in a separate transaction. When
// it was discovered that
let snapshot = diesel::select(pg_export_snapshot())
.get_result::<String>(conn)
.expect("pg_export_snapshot failed");
let snapshot_arg = format!("--snapshot={snapshot}");*/
let output = Command::new("pg_dump")
.args(["--schema-only"])
.env("DATABASE_URL", SETTINGS.get_database_url())
.stderr(Stdio::inherit())
.output()
.expect("failed to start pg_dump process");

// TODO: use exit_ok method when it's stable
assert!(output.status.success());

String::from_utf8(output.stdout).expect("pg_dump output is not valid UTF-8 text")
impl DiffChecker {
pub fn new(db_url: &str) -> diesel::result::QueryResult<Self> {
// todo use settings
let mut snapshot_conn = PgConnection::establish(db_url).expect("conn");
snapshot_conn.batch_execute("BEGIN;")?;
let (tx, rx) = crossbeam_channel::unbounded();
let (error_t, error_r) = crossbeam_channel::unbounded();
//let dump_receivers = Arc::new(Mutex::new(HashMap::new()));

let mut handles = Vec::new();
let n = usize::from(thread::available_parallelism().expect("parallelism"));
// todo remove
assert_eq!(16,n);
for _ in 0..(n){
let rx2 = rx.clone();
let error_t = error_t.clone();
handles.push(thread::spawn(move || if let Err(e) = std::panic::catch_unwind(move || {
while let Ok((snapshot, action)) = rx2.recv() {
let snapshot_arg = format!("--snapshot={snapshot}");
let output = Command::new("pg_dump")
.args(["--schema-only", &snapshot_arg])
.env("DATABASE_URL", SETTINGS.get_database_url())
.output()
.expect("failed to start pg_dump process");

if !output.status.success() {
panic!("{}", String::from_utf8(output.stderr).expect(""));
}

let output_string = String::from_utf8(output.stdout).expect("pg_dump output is not valid UTF-8 text");
match action {
DumpAction::Send(x) => {x.send(output_string).ok();},
DumpAction::Compare(x, name) => {
if let Ok(before) = x.recv() {
if let Some(e) = check_dump_diff(before, output_string, &name) {
panic!("{e}");
}
}
}
}
}
}){
error_t.send(e).ok();
}));
}

Ok(DiffChecker {snapshot_conn,handles,snapshot_sender:Some(tx),error:error_r})
}

fn check_err(&mut self) {
if let Ok(e) = self.error.try_recv() {
std::panic::resume_unwind(e);
}
}

pub fn finish(&mut self) {
self.snapshot_sender.take(); // stop threads from waiting
for handle in self.handles.drain(..) {
handle.join().expect("");
}
self.check_err();
}

fn get_snapshot(&mut self) -> String {
diesel::select(pg_export_snapshot())
.get_result::<String>(&mut self.snapshot_conn)
.expect("pg_export_snapshot failed")
}

pub fn get_dump(&mut self) -> Receiver<String> {
self.check_err();
let snapshot = self.get_snapshot();
let (tx, rx) = crossbeam_channel::unbounded(); // ::bounded(1);
self.snapshot_sender.as_mut().expect("").send((snapshot, DumpAction::Send(tx))).expect("send msg");
rx
}

pub fn check_dump_diff(&mut self, before: Receiver<String>, name: String) {
self.check_err();
let snapshot = self.get_snapshot();
self.snapshot_sender.as_mut().expect("").send((snapshot, DumpAction::Compare(before, name))).expect("compare msg");
}
}


const PATTERN_LEN: usize = 19;

// TODO add unit test for output
pub fn check_dump_diff(conn: &mut PgConnection, mut before: String, name: &str) {
let mut after = get_dump(conn);
pub fn check_dump_diff(mut before: String, mut after: String, name: &str) -> Option<String> {
if after == before {
return;
return None;
}
// Ignore timestamp differences by removing timestamps
for dump in [&mut before, &mut after] {
Expand Down Expand Up @@ -97,7 +173,7 @@ pub fn check_dump_diff(conn: &mut PgConnection, mut before: String, name: &str)
.collect::<Vec<_>>()
});
if only_in_before.is_empty() && only_in_after.is_empty() {
return;
return None;
}
let after_has_more =
only_in_before.len() < only_in_after.len();
Expand Down Expand Up @@ -142,13 +218,13 @@ pub fn check_dump_diff(conn: &mut PgConnection, mut before: String, name: &str)
}
.expect("failed to build string");
}
write!(&mut output, "\n{most_similar_chunk_filtered}");
write!(&mut output, "\n{most_similar_chunk_filtered}").expect("");
if !chunks_gt.is_empty() {
chunks_gt.swap_remove(most_similar_chunk_index);}
}
// should have all been removed
assert_eq!(chunks_gt.len(), 0);
panic!("{output}");
Some(output)
}

// todo inline?
Expand Down

0 comments on commit 592a127

Please sign in to comment.