Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193use super::Hydrant;use crate::state::AppState;use futures::FutureExt;use miette::{IntoDiagnostic, Result};use std::sync::Arc;
/// control over database maintenance operations.////// all methods pause the crawler, firehose, and backfill worker for the duration/// of the operation and restore their prior state on completion, whether or not/// the operation succeeds.#[derive(Clone)]pub struct DbControl(pub(crate) Arc<AppState>);
impl DbControl { /// trigger a major compaction of all keyspaces in parallel. /// /// compaction reclaims disk space from deleted/updated keys and improves /// read performance. can take several minutes on large datasets. pub async fn compact(&self) -> Result<()> { tracing::info!("compaction request received"); let start = std::time::Instant::now(); let state = self.0.clone();
let timeout_duration = state.maintenance_timeout; let fut = state.with_ingestion_paused(async || { tracing::info!("compaction starting (ingestion paused)"); state.db.compact().await });
let res = if timeout_duration.is_zero() { fut.await } else { match tokio::time::timeout(timeout_duration, fut).await { Ok(inner_res) => inner_res, Err(_) => { tracing::error!( "compaction timed out after {:?}: ingestion has been resumed and the operation was aborted", start.elapsed() ); miette::bail!("compaction timed out"); } } };
match &res { Ok(_) => { tracing::info!("compaction completed in {:?}", start.elapsed()); } Err(e) => { tracing::error!(err = ?e, "compaction failed after {:?}", start.elapsed()); } } res }
/// train zstd compression dictionaries for every trainable keyspace /// (see `db::schema`: `repos`, `records`, `events`, and mode-dependent ones). /// /// dictionaries are written to `dict_{name}.bin` files inside the database folder. /// a restart is required to apply them. training samples data blocks from the /// existing database, so the database must have a reasonable amount of data first. pub async fn train_dicts(&self) -> Result<()> { tracing::info!("dictionary training request received"); let start = std::time::Instant::now(); let state = self.0.clone();
let timeout_duration = state.maintenance_timeout; let fut = state.with_ingestion_paused(async || { tracing::info!("dictionary training starting (ingestion paused)"); let train = |name: &'static str| state.db.run(move |db| db.train_dict(name)); futures::future::try_join_all( crate::db::registry::trainable() .into_iter() .map(|(name, _)| train(name)), ) .await .map(|_| ()) });
let res = if timeout_duration.is_zero() { fut.await } else { match tokio::time::timeout(timeout_duration, fut).await { Ok(inner_res) => inner_res, Err(_) => { tracing::error!( "dictionary training timed out after {:?}: ingestion has been resumed and the operation was aborted", start.elapsed() ); miette::bail!("dictionary training timed out"); } } };
match &res { Ok(_) => { tracing::info!("dictionary training completed in {:?}", start.elapsed()); } Err(e) => { tracing::error!(err = ?e, "dictionary training failed after {:?}", start.elapsed()); } } res }
/// open a user-defined keyspace. /// /// the keyspace name is prefixed to prevent collision with hydrant's internal keyspaces. #[cfg(feature = "user-keyspace")] pub fn open_keyspace<F>(&self, name: &str, setup: F) -> Result<fjall::Keyspace> where F: FnOnce() -> fjall::KeyspaceCreateOptions, { let name = format!("user:{name}"); self.0.db.inner.keyspace(&name, setup).into_diagnostic() }
/// create an owned write batch for batch operations. #[cfg(feature = "user-keyspace")] pub fn write_batch(&self) -> fjall::OwnedWriteBatch { self.0.db.inner.batch() }}
#[cfg(all(test, feature = "user-keyspace"))]mod tests { use super::*; use crate::config::Config;
fn test_config(path: &std::path::Path) -> Config { Config { database_path: path.to_path_buf(), ..Default::default() } }
#[tokio::test] async fn test_db_control_user_keyspace() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let cfg = test_config(tmp.path()); let hydrant = Hydrant::new(cfg).await?;
// 1. Open keyspace with custom options let ks = hydrant .db .open_keyspace("custom", fjall::KeyspaceCreateOptions::default)?;
// 2. Put a key/value ks.insert("key1", "value1").into_diagnostic()?; assert_eq!( ks.get("key1").into_diagnostic()?.as_deref(), Some(b"value1" as &[u8]) );
// 3. Verify it is prefixed (maps to user_custom) let prefixed_ks = hydrant .state .db .inner .keyspace("user:custom", fjall::KeyspaceCreateOptions::default) .into_diagnostic()?; assert_eq!( prefixed_ks.get("key1").into_diagnostic()?.as_deref(), Some(b"value1" as &[u8]) );
// 4. Verify it does not map to unprefixed "custom" let unprefixed_ks = hydrant .state .db .inner .keyspace("custom", fjall::KeyspaceCreateOptions::default) .into_diagnostic()?; assert_eq!( unprefixed_ks.get("key1").into_diagnostic()?.as_deref(), None );
// 5. Test batch functionality let mut batch = hydrant.db.write_batch(); batch.insert(&ks, "key2", "value2"); batch.commit().into_diagnostic()?;
assert_eq!( ks.get("key2").into_diagnostic()?.as_deref(), Some(b"value2" as &[u8]) );
Ok(()) }}