Something went wrong. Try again.
A public mirror for the whole atmosphere hubble.microcosm.blue
Something went wrong. Try again.
3.6 kB · 133 lines
Rust
at backup
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134use clap::Parser;use rocksdb::{CompactOptions, DB, DBCompressionType, Options};
use std::path::{Path, PathBuf};use std::sync::{Arc, atomic::Ordering};
use space_efficiency_check::work::{ProcessError, run_workers};
/// import a directory of car files into rocks to check space efficiency#[derive(Parser)]struct Args { /// directory of car files car_dir: PathBuf,
/// where to kick rocks #[arg(long, default_value = "./local/space-efficiency-check-db")] db_dir: PathBuf,
/// concurrent car workers #[arg(long, default_value_t = 10)] workers: usize,
/// repo-stream memory limit: we don't disk spill; set higher than largest car #[arg(long, default_value_t = 1024)] mem_limit_mb: usize,
/// skip import and just do manual compact #[arg(long, action)] just_compact: bool,
/// sample fraction of DIDs to import (0.0–1.0, default: all) #[arg(long)] sample: Option<f64>,
/// store sequential u64s in place of dids #[arg(long, action)] intern: bool,}
fn open_db(path: &Path) -> Result<DB, rocksdb::Error> { let mut opts = Options::default(); opts.create_if_missing(true);
opts.set_compression_type(DBCompressionType::Lz4);
opts.set_max_subcompactions(4);
let mut block_opts = rocksdb::BlockBasedOptions::default(); block_opts.set_block_size(32_768); opts.set_block_based_table_factory(&block_opts);
DB::open(&opts, path)}
#[tokio::main]async fn main() -> Result<(), ProcessError> { tracing_subscriber::fmt::init();
let args = Args::parse();
let db = Arc::new(open_db(&args.db_dir)?);
if args.just_compact { manually_compact(db)?; return Ok(()); }
let stats = run_workers( &args.car_dir, db.clone(), args.workers, args.mem_limit_mb, args.sample, args.intern, ) .await?;
let repos = stats.repos.load(Ordering::Relaxed); let empty = stats.empty_repos.load(Ordering::Relaxed); let records = stats.records.load(Ordering::Relaxed); let failed = stats.failed_repos.load(Ordering::Relaxed); let car_bytes = stats.car_bytes.load(Ordering::Relaxed);
tracing::info!(repos, empty, records, failed, "import complete"); tracing::info!( car_bytes, car_mib = car_bytes / 2_u64.pow(20), "total input size" );
let size_before = db_size(&db); tracing::info!( bytes = size_before, mib = size_before / 2_u64.pow(20), "db size before compaction" );
let size_after = manually_compact(db)?;
if car_bytes > 0 { let ratio = size_after as f64 / car_bytes as f64; tracing::info!(ratio = format!("{ratio:.3}"), "db size / car bytes"); }
Ok(())}
fn manually_compact(db: Arc<DB>) -> Result<u64, ProcessError> { let mut compact_opts = CompactOptions::default(); compact_opts.set_bottommost_level_compaction(rocksdb::BottommostLevelCompaction::Force);
tracing::info!("running manual compaction..."); let compact_start = std::time::Instant::now(); db.compact_range_opt(None::<&[u8]>, None::<&[u8]>, &compact_opts); let compact_elapsed = compact_start.elapsed();
let size_after = db_size(&db); tracing::info!( bytes = size_after, mib = size_after / 2_u64.pow(20), ?compact_elapsed, "db size after compaction", ); Ok(size_after)}
fn db_size(db: &DB) -> u64 { db.property_value("rocksdb.total-sst-files-size") .ok() .flatten() .and_then(|v| v.parse::<u64>().ok()) .unwrap_or(0)}