use 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, but the compaction keeps running in the background",
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 keyspace's compression is saved when it's created, so a dictionary only gets used
/// by keyspaces created after it's there, a restart doesn't apply it to existing ones.
/// 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, but the training keeps running in the background",
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(())
}
}