use fjall::{CompressionType, Database};
use lsm_tree::compaction::Factory;
use miette::{Context, IntoDiagnostic, Result};
use scc::HashMap;
use smol_str::SmolStr;
use std::collections::{BTreeSet, HashMap as StdHashMap};
#[cfg(test)]
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use crate::config::{Compression, Config};
use super::compaction::{CountsGcFilterFactory, HistoryRetentionFilterFactory};
use super::counts::{load_count_delta_watermark, read_u64_counter, replay_count_deltas};
use super::keys;
use super::keyspaces::OpenCx;
use super::schema::{Ks, mb};
use super::{Db, migration, registry, schema};
impl Db {
pub fn open(cfg: &Config) -> Result<Self> {
cfg.validate()?;
tracing::info!("opening database at {}...", cfg.database_path.display());
let start = std::time::Instant::now();
let (db, count_delta_gc_watermark, history_ttl) = Self::open_database(cfg)?;
let dicts = Self::load_dicts(cfg);
let get_compression = |name: &str, level: i32| match cfg.data_compression {
Compression::Lz4 => CompressionType::Lz4,
Compression::Zstd => dicts
.get(name)
.map(|dict| CompressionType::ZstdDict {
level,
dict: dict.clone(),
})
.unwrap_or_else(|| CompressionType::Zstd { level }),
Compression::None => CompressionType::None,
};
let cx = OpenCx {
db: &db,
cfg,
compression: &get_compression,
opened: std::cell::RefCell::new(Vec::new()),
poison: super::Poison::default(),
};
let this = Self::assemble_keyspaces_and_verify(&cx, count_delta_gc_watermark)?;
// a marked database can reject a mode mismatch before migrations or
// compaction filters mutate anything. pre-marker databases are
// adopted after migration, when their record layout is inferable.
#[cfg(feature = "indexer")]
let storage_mode_marked = super::storage_mode::preflight(&this, cfg)?;
migration::run(&this)?;
#[cfg(feature = "indexer")]
if !storage_mode_marked {
super::storage_mode::adopt(&this, cfg)?;
}
// history compaction stays fail-safe (keep everything) until storage
// mode validation and migrations have both succeeded.
history_ttl
.set(cfg.history_ttl)
.map_err(|_| miette::miette!("history retention was initialized twice"))?;
this.init_modes(cfg)?;
this.load_persisted_counts()?;
this.restore_count_deltas_and_next_id()?;
tracing::info!("database ready (recovered in {:?})", start.elapsed());
Ok(this)
}
fn open_database(
cfg: &Config,
) -> Result<(
Arc<Database>,
Arc<AtomicU64>,
Arc<OnceLock<Option<std::time::Duration>>>,
)> {
let count_delta_gc_watermark = Arc::new(AtomicU64::new(0));
let history_ttl = Arc::new(OnceLock::new());
let db = Database::builder(&cfg.database_path)
.cache_size(mb(cfg.cache_size))
.manual_journal_persist(true)
.journal_compression(match cfg.journal_compression {
Compression::Lz4 => CompressionType::Lz4,
Compression::Zstd => CompressionType::Zstd { level: 3 },
Compression::None => CompressionType::None,
})
.worker_threads(cfg.db_worker_threads)
.max_journaling_size(mb(cfg.db_max_journaling_size_mb))
.with_compaction_filter_factories({
let history_retention_enabled = cfg.history_ttl.is_some();
let history_ttl = history_ttl.clone();
let count_delta_gc_watermark = count_delta_gc_watermark.clone();
let f = move |ks: &str| match ks {
"counts" => Some(Arc::new(CountsGcFilterFactory {
delta_gc_watermark: count_delta_gc_watermark.clone(),
}) as Arc<dyn Factory>),
"history" if history_retention_enabled => {
Some(Arc::new(HistoryRetentionFilterFactory {
ttl: history_ttl.clone(),
}) as Arc<dyn Factory>)
}
_ => None,
};
Arc::new(f)
})
.open()
.into_diagnostic()?;
let db = Arc::new(db);
Ok((db, count_delta_gc_watermark, history_ttl))
}
fn load_dicts(cfg: &Config) -> StdHashMap<&'static str, Arc<[u8]>> {
let load_dict = |name: &str| -> Option<Arc<[u8]>> {
let path = cfg.database_path.join(format!("dict_{name}.bin"));
if path.exists()
&& let Ok(bytes) = std::fs::read(&path)
{
tracing::debug!(
"loaded zstd dictionary for keyspace {name} ({} bytes)",
bytes.len()
);
return Some(bytes.into());
}
None
};
registry::trainable()
.into_iter()
.filter_map(|(name, _)| load_dict(name).map(|dict| (name, dict)))
.collect()
}
fn assemble_keyspaces_and_verify(
cx: &OpenCx,
count_delta_gc_watermark: Arc<AtomicU64>,
) -> Result<Self> {
let repos = Ks::<schema::Repos>::open(cx)?;
let repo_metadata = Ks::<schema::RepoMetadata>::open(cx)?;
let cursors = Ks::<schema::Cursors>::open(cx)?;
let counts = Ks::<schema::Counts>::open(cx)?;
let filter = Ks::<schema::Filter>::open(cx)?;
let crawler = Ks::<schema::Crawler>::open(cx)?;
#[cfg(feature = "backlinks")]
let backlinks = Ks::<schema::Backlinks>::open(cx)?;
#[cfg(feature = "indexer")]
let indexer = super::keyspaces::IndexerDb::open(cx)?;
#[cfg(feature = "indexer_stream")]
let stream = super::keyspaces::StreamDb::open(cx)?;
#[cfg(feature = "jetstream")]
let jetstream = super::keyspaces::JetstreamDb::open(cx)?;
#[cfg(feature = "relay")]
let relay = super::keyspaces::RelayDb::open(cx)?;
#[cfg(any(feature = "indexer_stream", feature = "relay"))]
let sequencer = Arc::new(super::sequencer::Sequencer::new(cursors.clone()));
// every opened keyspace must have a registry row and vice versa, so the
// by-name, /stats, /debug, and training tables cannot silently drift.
let opened: BTreeSet<&'static str> = cx.opened.borrow().iter().copied().collect();
let registered: BTreeSet<&'static str> = registry::names().into_iter().collect();
if opened != registered {
let missing: Vec<_> = registered.difference(&opened).collect();
let unregistered: Vec<_> = opened.difference(®istered).collect();
miette::bail!(
"keyspace registry drift: registered but not opened: {missing:?}, opened but not registered: {unregistered:?}"
);
}
Ok(Self {
inner: Arc::clone(cx.db),
path: cx.cfg.database_path.clone(),
repos,
repo_metadata,
cursors,
counts,
filter,
crawler,
#[cfg(feature = "indexer")]
indexer,
#[cfg(feature = "indexer_stream")]
stream,
#[cfg(feature = "jetstream")]
jetstream,
#[cfg(feature = "relay")]
relay,
#[cfg(feature = "backlinks")]
backlinks,
counts_map: Arc::new(HashMap::new()),
next_count_delta_id: Arc::new(AtomicU64::new(0)),
count_delta_checkpoint_watermark: Arc::new(AtomicU64::new(0)),
count_delta_gc_watermark,
count_delta_in_flight: Arc::new(Mutex::new(BTreeSet::new())),
compaction_running: Arc::new(std::sync::atomic::AtomicBool::new(false)),
poison: cx.poison.clone(),
#[cfg(test)]
persist_failures: Arc::new(AtomicUsize::new(0)),
#[cfg(test)]
before_persist: Arc::default(),
#[cfg(test)]
compaction_hold: Arc::new(Mutex::new(())),
#[cfg(any(feature = "indexer_stream", feature = "relay"))]
sequencer,
})
}
#[cfg_attr(not(feature = "indexer_stream"), allow(unused_variables))]
fn init_modes(&self, cfg: &Config) -> Result<()> {
#[cfg(feature = "relay")]
self.relay.init(&self.cursors)?;
#[cfg(feature = "indexer_stream")]
self.stream.init(&self.cursors, cfg)?;
#[cfg(feature = "jetstream")]
self.jetstream.init(&self.cursors)?;
Ok(())
}
fn load_persisted_counts(&self) -> Result<()> {
// load counts into memory
for guard in self.counts.prefix(keys::COUNT_KS_PREFIX) {
let (k, v) = guard.into_inner().into_diagnostic()?;
let name = std::str::from_utf8(&k[keys::COUNT_KS_PREFIX.len()..])
.into_diagnostic()
.wrap_err("expected valid utf8 for ks count key")?;
let _ = self
.counts_map
.insert_sync(SmolStr::new(name), read_u64_counter(&v)?);
}
Ok(())
}
fn restore_count_deltas_and_next_id(&self) -> Result<()> {
let durable_watermark = load_count_delta_watermark(self)?;
replay_count_deltas(self, durable_watermark)?;
self.count_delta_checkpoint_watermark
.store(durable_watermark, Ordering::Relaxed);
self.count_delta_gc_watermark
.store(durable_watermark, Ordering::Relaxed);
// always stay strictly above the durable watermark so that after a migration
// deletes delta keys, new deltas are not assigned ids that checkpoint/replay
// would silently skip (finding 5f309024bc588191aa1a79eb449e3630).
let next_count_delta_id = self
.counts
.prefix(keys::COUNT_DELTA_PREFIX)
.next_back()
.map(|guard| -> Result<u64> {
let key = guard.key().into_diagnostic()?;
let (id, _) = keys::parse_count_delta_key(&key)?;
Ok(id + 1)
})
.transpose()?
.unwrap_or(0)
.max(durable_watermark.saturating_add(1));
self.next_count_delta_id
.store(next_count_delta_id, Ordering::Relaxed);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_block_cache_is_cache_size_megabytes() {
let tmp = tempfile::tempdir().unwrap();
let db = Db::open(&Config {
database_path: tmp.path().to_path_buf(),
cache_size: 64,
..Default::default()
})
.unwrap();
assert_eq!(db.inner.cache_capacity(), 64 * 1024 * 1024);
}
}