Something went wrong. Try again.
A public mirror for the whole atmosphere hubble.microcosm.blue
Something went wrong. Try again.
4.7 kB · 150 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151use clap::Parser;use rocksdb::{DB, Options, WriteBatch, WriteOptions};use std::path::PathBuf;use std::time::Instant;
/// Create a smaller test DB by copying contiguous key ranges from a compacted source DB./// Picks every Nth SST file's key range to preserve DID-level key locality.#[derive(Parser)]struct Args { /// source DB directory (should be fully compacted to L6) source: PathBuf,
/// destination DB directory (will be created) dest: PathBuf,
/// fraction of SST files to sample (0.0–1.0) #[arg(long, default_value_t = 0.1)] fraction: f64,}
fn main() { let args = Args::parse(); assert!( args.fraction > 0.0 && args.fraction <= 1.0, "fraction must be in (0.0, 1.0]" );
// open source read-only to get file list let source_opts = Options::default(); let source = DB::open_for_read_only(&source_opts, &args.source, false) .expect("failed to open source db");
let mut files = source.live_files().expect("failed to get live files"); files.sort_by(|a, b| a.start_key.cmp(&b.start_key));
let total_size: usize = files.iter().map(|f| f.size).sum(); let total_entries: u64 = files.iter().map(|f| f.num_entries).sum(); eprintln!( "source: {} SST files, {} MiB, {} entries", files.len(), total_size / (1 << 20), total_entries, );
// pick every Nth file's key range let step = (1.0 / args.fraction).round().max(1.0) as usize; let selected: Vec<_> = files.iter().step_by(step).collect();
let sel_entries: u64 = selected.iter().map(|f| f.num_entries).sum(); eprintln!( "sampling {} file ranges (every {}th), ~{} entries", selected.len(), step, sel_entries, );
// collect key ranges to copy let ranges: Vec<_> = selected .iter() .filter_map(|f| { let start = f.start_key.as_ref()?; let end = f.end_key.as_ref()?; Some((start.clone(), end.clone())) }) .collect();
// create dest DB let mut dest_opts = Options::default(); dest_opts.create_if_missing(true); let dest = DB::open(&dest_opts, &args.dest).expect("failed to create dest db");
let mut write_opts = WriteOptions::default(); write_opts.disable_wal(true);
// iterate source key ranges and write to dest let start = Instant::now(); let mut total_written: u64 = 0;
for (i, (range_start, range_end)) in ranges.iter().enumerate() { let mut iter = source.raw_iterator(); iter.seek(range_start);
let mut batch = WriteBatch::default(); let mut batch_count = 0;
while iter.valid() { let key = iter.key().unwrap(); if key > range_end.as_slice() { break; } batch.put(key, iter.value().unwrap()); batch_count += 1; total_written += 1;
if batch_count >= 32_768 { dest.write_opt( std::mem::replace(&mut batch, WriteBatch::default()), &write_opts, ) .expect("failed to write batch"); batch_count = 0; }
iter.next(); }
if batch_count > 0 { dest.write_opt(batch, &write_opts) .expect("failed to write batch"); }
if (i + 1) % 10 == 0 || i + 1 == ranges.len() { eprintln!( " copied {}/{} ranges, {} entries so far ({:.1}s)", i + 1, ranges.len(), total_written, start.elapsed().as_secs_f64(), ); } }
eprintln!("flushing..."); dest.flush().expect("failed to flush");
// compact so the dest DB is in a clean state eprintln!("compacting..."); let compact_start = Instant::now(); let mut compact_opts = rocksdb::CompactOptions::default(); compact_opts.set_bottommost_level_compaction(rocksdb::BottommostLevelCompaction::Force); dest.compact_range_opt(None::<&[u8]>, None::<&[u8]>, &compact_opts); eprintln!("compacted in {:.1}s", compact_start.elapsed().as_secs_f64());
// report let dest_size: u64 = dest .property_value("rocksdb.total-sst-files-size") .ok() .flatten() .and_then(|v| v.parse().ok()) .unwrap_or(0); let dest_files = dest.live_files().map(|f| f.len()).unwrap_or(0);
eprintln!( "done: {} entries → {} SST files, {} MiB ({:.1}s total)", total_written, dest_files, dest_size / (1 << 20), start.elapsed().as_secs_f64(), );}