use anyhow::{Context, Result};
use backbeat::{
record::RecordView,
ring::walk,
wire::{DumpReader, OwnedSchema},
};
use bytes::Bytes;
use rayon::prelude::*;
use std::{
collections::HashMap,
fs,
path::{Path, PathBuf},
};
pub struct Rec {
pub ts_nanos: u64,
pub shard_id: u32,
pub local_seq: u64,
pub schema_idx: usize,
pub fields: Bytes,
}
pub struct Loaded {
pub path: PathBuf,
pub instance_id: u64,
pub host: String,
pub schemas: Vec<OwnedSchema>,
pub intern: HashMap<u32, String>,
pub records: Vec<Rec>,
}
pub fn load(path: &Path, bytes: Bytes) -> Result<Loaded> {
let reader = DumpReader::new(bytes).map_err(|e| anyhow::anyhow!("{e}"))?;
let mut schemas = reader.schemas().map_err(|e| anyhow::anyhow!("{e}"))?;
let intern_pairs = reader.intern_table().map_err(|e| anyhow::anyhow!("{e}"))?;
let meta = reader.meta().map_err(|e| anyhow::anyhow!("{e}"))?;
let shards = reader.shards().map_err(|e| anyhow::anyhow!("{e}"))?;
schemas.sort_by(|a, b| a.qualified_name.cmp(&b.qualified_name));
let intern: HashMap<u32, String> = intern_pairs
.into_iter()
.map(|(id, bytes)| (id, String::from_utf8_lossy(&bytes).into_owned()))
.collect();
let by_id: HashMap<u64, usize> = schemas
.iter()
.enumerate()
.map(|(i, s)| (s.id.get(), i))
.collect();
let mut records: Vec<Rec> = shards
.par_iter()
.flat_map(|shard| {
let mut shard_recs: Vec<Rec> = Vec::new();
let mut seq = u64::MAX;
walk(
&shard.region,
shard.head as usize,
shard.capacity as usize,
|payload| {
let Some(rec) = RecordView::parse(&payload[..]) else {
return false;
};
match by_id.get(&rec.event_id.get()) {
Some(&idx) if rec.fields.len() == schemas[idx].record_size as usize => {
let fields_off = payload.len() - rec.fields.len();
shard_recs.push(Rec {
ts_nanos: rec.ts_nanos,
shard_id: shard.shard_id,
local_seq: seq,
schema_idx: idx,
fields: payload.slice(fields_off..),
});
seq -= 1;
true
}
_ => false,
}
},
);
shard_recs
})
.collect();
records.sort_by_key(|r| (r.ts_nanos, r.shard_id, r.local_seq));
let (instance_id, host) = meta
.map(|m| (m.instance_id, m.host))
.unwrap_or((0, String::new()));
Ok(Loaded {
path: path.to_path_buf(),
instance_id,
host,
schemas,
intern,
records,
})
}
pub fn load_many(paths: &[PathBuf]) -> Result<Vec<Loaded>> {
paths
.par_iter()
.map(|p| {
let bytes = fs::read(p).with_context(|| format!("reading dump {}", p.display()))?;
load(p, bytes.into()).with_context(|| format!("decoding dump {}", p.display()))
})
.collect()
}