heimoshuiyu/rust-synapse-compress-state
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Split up main (#39)

Browse Source 
Moved code from inside the main function into function calls
This commit is contained in:
Azrenbeth
2021-07-27 15:31:48 +01:00
committed by GitHub
parent f3b1dd0b8a
commit 31e37996ea
2 changed files with 90 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -2,3 +2,4 @@
**/*.rs.bk
*.data
*.old
out.sql

132
src/main.rs
View File

@@ -92,9 +92,19 @@ impl FromStr for LevelSizes {
}
}
fn main() {
#[allow(deprecated)]
let matches = App::new(crate_name!())
struct Config {
db_url: String,
output_file: Option<File>,
room_id: String,
max_state_group: Option<i64>,
min_saved_rows: Option<i32>,
transactions: bool,
level_sizes: LevelSizes,
}
impl Config {
fn parse_arguments() -> Config {
let matches = App::new(crate_name!())
.version(crate_version!())
.author(crate_authors!("\n"))
.about(crate_description!())
@@ -156,33 +166,50 @@ fn main() {
.takes_value(true),
).get_matches();
let db_url = matches
.value_of("postgres-url")
.expect("db url should be required");
let db_url = matches
.value_of("postgres-url")
.expect("db url should be required");
let mut output_file = matches
.value_of("output_file")
.map(|path| File::create(path).unwrap());
let output_file = matches
.value_of("output_file")
.map(|path| File::create(path).unwrap());
let room_id = matches
.value_of("room_id")
.expect("room_id should be required since no file");
let room_id = matches
.value_of("room_id")
.expect("room_id should be required since no file");
let max_state_group = matches
.value_of("max_state_group")
.map(|s| s.parse().expect("max_state_group must be an integer"));
let max_state_group = matches
.value_of("max_state_group")
.map(|s| s.parse().expect("max_state_group must be an integer"));
let min_saved_rows = matches
.value_of("min_saved_rows")
.map(|v| v.parse().expect("COUNT must be an integer"));
let min_saved_rows = matches
.value_of("min_saved_rows")
.map(|v| v.parse().expect("COUNT must be an integer"));
let transactions = matches.is_present("transactions");
let transactions = matches.is_present("transactions");
let level_sizes = value_t_or_exit!(matches, "level_sizes", LevelSizes);
let level_sizes = value_t_or_exit!(matches, "level_sizes", LevelSizes);
Config {
db_url: String::from(db_url),
output_file,
room_id: String::from(room_id),
max_state_group,
min_saved_rows,
transactions,
level_sizes,
}
}
}
fn main() {
let mut config = Config::parse_arguments();
// First we need to get the current state groups
println!("Fetching state from DB for room '{}'...", room_id);
let state_group_map = database::get_data_from_db(db_url, room_id, max_state_group);
println!("Fetching state from DB for room '{}'...", config.room_id);
let state_group_map =
database::get_data_from_db(&config.db_url, &config.room_id, config.max_state_group);
println!("Number of state groups: {}", state_group_map.len());
@@ -196,7 +223,7 @@ fn main() {
println!("Compressing state...");
let compressor = Compressor::compress(&state_group_map, &level_sizes.0);
let compressor = Compressor::compress(&state_group_map, &config.level_sizes.0);
let new_state_group_map = compressor.new_state_group_map;
@@ -228,7 +255,7 @@ fn main() {
compressor.stats.state_groups_changed
);
if let Some(min) = min_saved_rows {
if let Some(min) = config.min_saved_rows {
let saving = (original_summed_size - compressed_summed_size) as i32;
if saving < min {
println!(
@@ -239,25 +266,39 @@ fn main() {
}
}
check_that_maps_match(&state_group_map, &new_state_group_map);
// If we are given an output file, we output the changes as SQL. If the
// `transactions` argument is set we wrap each change to a state group in a
// transaction.
if let Some(output) = &mut output_file {
println!("Writing changes...");
output_sql(&mut config, &state_group_map, &new_state_group_map);
}
let pb = ProgressBar::new(state_group_map.len() as u64);
pb.set_style(
ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"),
);
pb.set_message("state groups");
pb.enable_steady_tick(100);
fn output_sql(
config: &mut Config,
old_map: &BTreeMap<i64, StateGroupEntry>,
new_map: &BTreeMap<i64, StateGroupEntry>,
) {
if config.output_file.is_none() {
return;
}
for (sg, old_entry) in &state_group_map {
let new_entry = &new_state_group_map[sg];
println!("Writing changes...");
let pb = ProgressBar::new(old_map.len() as u64);
pb.set_style(
ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"),
);
pb.set_message("state groups");
pb.enable_steady_tick(100);
if let Some(output) = &mut config.output_file {
for (sg, old_entry) in old_map {
let new_entry = &new_map[sg];
if old_entry != new_entry {
if transactions {
if config.transactions {
writeln!(output, "BEGIN;").unwrap();
}
@@ -292,7 +333,7 @@ fn main() {
output,
"({}, {}, {}, {}, {})",
sg,
PGEscape(room_id),
PGEscape(&config.room_id),
PGEscape(t),
PGEscape(s),
PGEscape(e)
@@ -302,7 +343,7 @@ fn main() {
writeln!(output, ";").unwrap();
}
if transactions {
if config.transactions {
writeln!(output, "COMMIT;").unwrap();
}
writeln!(output).unwrap();
@@ -310,13 +351,18 @@ fn main() {
pb.inc(1);
}
pb.finish();
}
pb.finish();
}
fn check_that_maps_match(
old_map: &BTreeMap<i64, StateGroupEntry>,
new_map: &BTreeMap<i64, StateGroupEntry>,
) {
println!("Checking that state maps match...");
let pb = ProgressBar::new(state_group_map.len() as u64);
let pb = ProgressBar::new(old_map.len() as u64);
pb.set_style(
ProgressStyle::default_bar().template("[{elapsed_precise}] {bar} {pos}/{len} {msg}"),
);
@@ -325,11 +371,11 @@ fn main() {
// Now let's iterate through and assert that the state for each group
// matches between the two versions.
state_group_map
old_map
.par_iter() // This uses rayon to run the checks in parallel
.try_for_each(|(sg, _)| {
let expected = collapse_state_maps(&state_group_map, *sg);
let actual = collapse_state_maps(&new_state_group_map, *sg);
let expected = collapse_state_maps(&old_map, *sg);
let actual = collapse_state_maps(&new_map, *sg);
pb.inc(1);