From 41096cfe4e3f8440daa488787d344eed9ad2a6b9 Mon Sep 17 00:00:00 2001 From: dawn <90008@gaze.systems> Date: Tue, 10 Mar 2026 10:13:13 +0300 Subject: [PATCH] [all] implement garbage collection for blocks, expose /stream/ack --- .agent/rules/read-agents.md | 5 - AGENTS.md | 2 +- Cargo.toml | 3 + README.md | 3 + src/api/debug.rs | 85 ++++++++++- src/api/repos.rs | 13 +- src/api/stream.rs | 33 ++-- src/backfill/mod.rs | 125 ++++++++++----- src/db/compaction.rs | 25 --- src/db/gc.rs | 294 ++++++++++++++++++++++++++++++++++++ src/db/keys.rs | 17 ++- src/db/mod.rs | 83 +++++++--- src/db/refcount.rs | 61 ++++++++ src/ingest/worker.rs | 54 ++++--- src/main.rs | 24 +++ src/ops.rs | 101 +++++++++---- tests/common.nu | 4 +- tests/gc_test.nu | 159 +++++++++++++++++++ tests/repos_api_test.nu | 4 +- 19 files changed, 935 insertions(+), 160 deletions(-) delete mode 100644 .agent/rules/read-agents.md create mode 100644 src/db/gc.rs create mode 100644 src/db/refcount.rs create mode 100755 tests/gc_test.nu diff --git a/.agent/rules/read-agents.md b/.agent/rules/read-agents.md deleted file mode 100644 index 728f4b7..0000000 --- a/.agent/rules/read-agents.md +++ /dev/null @@ -1,5 +0,0 @@ ---- -trigger: always_on ---- - -ALWAYS read AGENTS.md at the start of a conversation, and keep it in mind throught. \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md index aae1c6e..d3da4de 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -117,7 +117,7 @@ Examples: # GitNexus — Code Intelligence -This project is indexed by GitNexus as **hydrant** (599 symbols, 1624 relationships, 51 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely. +This project is indexed by GitNexus as **hydrant** (655 symbols, 1810 relationships, 55 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely. > If any GitNexus tool warns the index is stale, run `npx gitnexus analyze` in terminal first. diff --git a/Cargo.toml b/Cargo.toml index 6618cfb..9fe22a2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,3 +50,6 @@ lsm-tree = "3.1.0" [dev-dependencies] tempfile = "3.26.0" + +[profile.dev] +opt-level = 2 \ No newline at end of file diff --git a/README.md b/README.md index 4cf6d42..66e59ff 100644 --- a/README.md +++ b/README.md @@ -131,6 +131,9 @@ returns `{ count }`. - `GET /stream`: subscribe to the event stream. - query parameters: - `cursor` (optional): start streaming from a specific event ID. +- `POST /stream/ack`: ack events. + - body: + - `ids`: list of event IDs to acknowledge. ### stats diff --git a/src/api/debug.rs b/src/api/debug.rs index 3ff06e2..fb23c7e 100644 --- a/src/api/debug.rs +++ b/src/api/debug.rs @@ -1,6 +1,7 @@ use crate::api::AppState; use crate::db::keys; use crate::types::{RepoState, ResyncState, StoredEvent}; +use axum::routing::{get, post}; use axum::{ Json, extract::{Query, State}, @@ -26,9 +27,12 @@ pub struct DebugCountResponse { pub fn router() -> axum::Router> { axum::Router::new() - .route("/debug/count", axum::routing::get(handle_debug_count)) - .route("/debug/get", axum::routing::get(handle_debug_get)) - .route("/debug/iter", axum::routing::get(handle_debug_iter)) + .route("/debug/count", get(handle_debug_count)) + .route("/debug/get", get(handle_debug_get)) + .route("/debug/iter", get(handle_debug_iter)) + .route("/debug/refcount", get(handle_debug_refcount)) + .route("/debug/refcount", post(handle_set_debug_refcount)) + .route("/debug/compact", post(handle_debug_compact)) } pub async fn handle_debug_count( @@ -182,7 +186,7 @@ pub async fn handle_debug_iter( let end = parse_bound(req.end)?; let items = tokio::task::spawn_blocking(move || { - let limit = req.limit.unwrap_or(50).min(1000); + let limit = req.limit.unwrap_or(50); let collect = |iter: &mut dyn Iterator| { let mut items = Vec::new(); @@ -197,6 +201,11 @@ pub async fn handle_debug_iter( } else { "invalid_u64".to_string() } + } else if partition == "blocks" { + match cid::Cid::read_bytes(k.as_ref()) { + Ok(cid) => cid.to_string(), + Err(_) => String::from_utf8_lossy(&k).into_owned(), + } } else { String::from_utf8_lossy(&k).into_owned() }; @@ -255,3 +264,71 @@ fn get_keyspace_by_name(db: &crate::db::Db, name: &str) -> Result Err(StatusCode::BAD_REQUEST), } } + +#[derive(Deserialize)] +pub struct DebugCompactRequest { + pub partition: String, +} + +pub async fn handle_debug_compact( + State(state): State>, + Query(req): Query, +) -> Result { + let ks = get_keyspace_by_name(&state.db, &req.partition)?; + let state_clone = state.clone(); + + tokio::task::spawn_blocking(move || { + let _ = ks.remove(b"dummy_tombstone123"); + let _ = state_clone.db.persist(); + let _ = ks.rotate_memtable_and_wait(); + ks.major_compact() + }) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + Ok(StatusCode::OK) +} + +#[derive(Deserialize)] +pub struct DebugRefcountRequest { + pub cid: String, +} + +#[derive(Serialize)] +pub struct DebugRefcountResponse { + pub count: Option, +} + +pub async fn handle_debug_refcount( + State(state): State>, + Query(req): Query, +) -> Result, StatusCode> { + let cid = cid::Cid::from_str(&req.cid).map_err(|_| StatusCode::BAD_REQUEST)?; + let cid_bytes = fjall::Slice::from(cid.to_bytes()); + + let count = state + .db + .block_refcounts + .read_sync(cid_bytes.as_ref(), |_, v| *v); + + Ok(Json(DebugRefcountResponse { count })) +} + +#[derive(Deserialize)] +pub struct DebugSetRefcountRequest { + pub cid: String, + pub count: i64, +} + +pub async fn handle_set_debug_refcount( + State(state): State>, + axum::extract::Json(req): axum::extract::Json, +) -> Result { + let cid = cid::Cid::from_str(&req.cid).map_err(|_| StatusCode::BAD_REQUEST)?; + let cid_bytes = fjall::Slice::from(cid.to_bytes()); + + let _ = state.db.block_refcounts.insert_sync(cid_bytes, req.count); + + Ok(StatusCode::OK) +} diff --git a/src/api/repos.rs b/src/api/repos.rs index de88a5c..88f506e 100644 --- a/src/api/repos.rs +++ b/src/api/repos.rs @@ -14,6 +14,7 @@ use rand::Rng; use serde::{Deserialize, Serialize}; use crate::api::AppState; +use crate::db::refcount::RefcountedBatch; use crate::db::{keys, ser_repo_state}; use crate::types::{GaugeState, RepoState}; @@ -28,7 +29,7 @@ pub fn router() -> Router> { #[derive(Deserialize, Debug)] pub struct RepoRequest { pub did: String, - #[serde(skip_serializing_if = "Option::is_none", rename = "deleteData")] + #[serde(skip_serializing_if = "Option::is_none")] pub delete_data: Option, } @@ -276,7 +277,7 @@ pub async fn handle_delete_repos( let state_task = state.clone(); let (deleted_count, gauge_decrements) = tokio::task::spawn_blocking(move || { let db = &state_task.db; - let mut batch = db.inner.batch(); + let mut batch = RefcountedBatch::new(db); let mut deleted_count = 0i64; let mut gauge_decrements = Vec::new(); @@ -306,13 +307,15 @@ pub async fn handle_delete_repos( } else if repo_state.tracked { let mut repo_state = repo_state.into_static(); repo_state.tracked = false; - batch.insert( + batch.batch_mut().insert( &db.repos, &did_key, ser_repo_state(&repo_state).map_err(internal)?, ); - batch.remove(&db.pending, keys::pending_key(repo_state.index_id)); - batch.remove(&db.resync, &did_key); + batch + .batch_mut() + .remove(&db.pending, keys::pending_key(repo_state.index_id)); + batch.batch_mut().remove(&db.resync, &did_key); if old_gauge != GaugeState::Synced { gauge_decrements.push(old_gauge); } diff --git a/src/api/stream.rs b/src/api/stream.rs index 597e2f9..4df11ba 100644 --- a/src/api/stream.rs +++ b/src/api/stream.rs @@ -1,9 +1,10 @@ use crate::api::AppState; use crate::db::keys; +use crate::db::refcount::RefcountedBatch; use crate::types::{BroadcastEvent, MarshallableEvt, RecordEvt, StoredEvent}; use axum::Router; use axum::http::StatusCode; -use axum::routing::get; +use axum::routing::{get, post}; use axum::{ extract::{ Query, State, @@ -19,17 +20,16 @@ use tokio::sync::{broadcast, mpsc, oneshot}; use tracing::{error, info_span, warn}; pub fn router() -> Router> { - Router::new().route("/", get(handle_stream)) - // .route("/ack", post(handle_ack)) + Router::new() + .route("/", get(handle_stream)) + .route("/ack", post(handle_ack)) } -#[allow(dead_code)] #[derive(Deserialize)] pub struct AckBody { pub ids: Vec, } -#[allow(dead_code)] pub async fn handle_ack( State(state): State>, axum::Json(body): axum::Json, @@ -41,21 +41,34 @@ pub async fn handle_ack( let state = state.clone(); let ids = body.ids; tokio::task::spawn_blocking(move || { - let mut batch = state.db.inner.batch(); + let mut batch = RefcountedBatch::new(&state.db); for &id in &ids { - batch.remove(&state.db.events, keys::event_key(id)); + let _entered = tracing::info_span!("ack", id).entered(); + let key = keys::event_key(id); + let Some(event_bytes) = state.db.events.get(&key).into_diagnostic()? else { + tracing::warn!("event bytes not found"); + continue; + }; + let evt = rmp_serde::from_slice::(&event_bytes).into_diagnostic()?; + if let Some(cid) = evt.cid { + tracing::debug!(cid = %cid, "acking event"); + batch.update_block_refcount(fjall::Slice::from(cid.to_bytes()), -1)?; + } else { + tracing::debug!("acking event with NO cid"); + } + batch.batch_mut().remove(&state.db.events, key); } batch .commit() .into_diagnostic() - .wrap_err("failed to delete events") - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + .wrap_err("failed to delete events")?; Ok(StatusCode::OK) }) .await .into_diagnostic() .wrap_err("panicked while deleting events") - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .flatten() + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) } #[derive(Deserialize)] diff --git a/src/backfill/mod.rs b/src/backfill/mod.rs index 3bb11eb..bace142 100644 --- a/src/backfill/mod.rs +++ b/src/backfill/mod.rs @@ -22,6 +22,7 @@ use miette::{Diagnostic, IntoDiagnostic, Result}; use reqwest::StatusCode; use smol_str::{SmolStr, ToSmolStr}; use std::collections::HashMap; +use std::str::FromStr; use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::{Duration, Instant}; @@ -39,6 +40,7 @@ pub struct BackfillWorker { http: reqwest::Client, semaphore: Arc, verify_signatures: bool, + ephemeral: bool, in_flight: Arc>>, } @@ -49,6 +51,7 @@ impl BackfillWorker { timeout: Duration, concurrency_limit: usize, verify_signatures: bool, + ephemeral: bool, ) -> Self { Self { state, @@ -62,6 +65,7 @@ impl BackfillWorker { .expect("failed to build http client"), semaphore: Arc::new(Semaphore::new(concurrency_limit)), verify_signatures, + ephemeral, in_flight: Arc::new(scc::HashSet::new()), } } @@ -123,10 +127,14 @@ impl BackfillWorker { let did = did.clone(); let buffer_tx = self.buffer_tx.clone(); let verify = self.verify_signatures; + let ephemeral = self.ephemeral; tokio::spawn(async move { let _guard = guard; - let res = did_task(&state, http, buffer_tx, &did, key, permit, verify).await; + let res = did_task( + &state, http, buffer_tx, &did, key, permit, verify, ephemeral, + ) + .await; if let Err(e) = res { error!(did = %did, err = %e, "backfill process failed"); @@ -161,10 +169,11 @@ async fn did_task( pending_key: Slice, _permit: tokio::sync::OwnedSemaphorePermit, verify_signatures: bool, + ephemeral: bool, ) -> Result<(), BackfillError> { let db = &state.db; - match process_did(&state, &http, &did, verify_signatures).await { + match process_did(&state, &http, &did, verify_signatures, ephemeral).await { Ok(Some(repo_state)) => { let did_key = keys::repo_key(&did); @@ -367,6 +376,7 @@ async fn process_did<'i>( http: &reqwest::Client, did: &Did<'static>, verify_signatures: bool, + ephemeral: bool, ) -> Result>, BackfillError> { debug!(did = %did, "backfilling"); @@ -426,7 +436,7 @@ async fn process_did<'i>( Err(XrpcError::Xrpc(e)) => { if matches!(e, GetRepoError::RepoNotFound(_)) { warn!(did = %did, "repo not found, deleting"); - let mut batch = db.inner.batch(); + let mut batch = db::refcount::RefcountedBatch::new(db); if let Err(e) = crate::ops::delete_repo(&mut batch, db, did, &state) { tracing::error!(err = %e, "failed to wipe repo during backfill"); } @@ -539,35 +549,37 @@ async fn process_did<'i>( let mut delta = 0; let mut added_blocks = 0; let mut collection_counts: HashMap = HashMap::new(); - let mut batch = app_state.db.inner.batch(); + let mut batch = db::refcount::RefcountedBatch::new(&app_state.db); let store = mst.storage(); let prefix = keys::record_prefix_did(&did); let mut existing_cids: HashMap<(SmolStr, DbRkey), SmolStr> = HashMap::new(); - for guard in app_state.db.records.prefix(&prefix) { - let (key, cid_bytes) = guard.into_inner().into_diagnostic()?; - // key is did|collection|rkey - // skip did| - let mut remaining = key[prefix.len()..].splitn(2, |b| keys::SEP.eq(b)); - let collection_raw = remaining - .next() - .ok_or_else(|| miette::miette!("invalid record key format: {key:?}"))?; - let rkey_raw = remaining - .next() - .ok_or_else(|| miette::miette!("invalid record key format: {key:?}"))?; - - let collection = std::str::from_utf8(collection_raw) - .map_err(|e| miette::miette!("invalid collection utf8: {e}"))?; - - let rkey = keys::parse_rkey(rkey_raw) - .map_err(|e| miette::miette!("invalid rkey '{key:?}' for {did}: {e}"))?; - - let cid = cid::Cid::read_bytes(cid_bytes.as_ref()) - .map_err(|e| miette::miette!("invalid cid '{cid_bytes:?}' for {did}: {e}"))? - .to_smolstr(); - - existing_cids.insert((collection.into(), rkey), cid); + if !ephemeral { + for guard in app_state.db.records.prefix(&prefix) { + let (key, cid_bytes) = guard.into_inner().into_diagnostic()?; + // key is did|collection|rkey + // skip did| + let mut remaining = key[prefix.len()..].splitn(2, |b| keys::SEP.eq(b)); + let collection_raw = remaining + .next() + .ok_or_else(|| miette::miette!("invalid record key format: {key:?}"))?; + let rkey_raw = remaining + .next() + .ok_or_else(|| miette::miette!("invalid record key format: {key:?}"))?; + + let collection = std::str::from_utf8(collection_raw) + .map_err(|e| miette::miette!("invalid collection utf8: {e}"))?; + + let rkey = keys::parse_rkey(rkey_raw) + .map_err(|e| miette::miette!("invalid rkey '{key:?}' for {did}: {e}"))?; + + let cid = cid::Cid::read_bytes(cid_bytes.as_ref()) + .map_err(|e| miette::miette!("invalid cid '{cid_bytes:?}' for {did}: {e}"))? + .to_smolstr(); + + existing_cids.insert((collection.into(), rkey), cid); + } } let mut signal_seen = filter.mode == FilterMode::Full || filter.signals.is_empty(); @@ -602,25 +614,28 @@ async fn process_did<'i>( let cid_obj = Cid::ipld(cid); // check if this record already exists with same CID - let (action, is_new) = if let Some(existing_cid) = existing_cids.remove(&path) { + let existing_cid = existing_cids.remove(&path); + let action = if let Some(existing_cid) = &existing_cid { if existing_cid == cid_obj.as_str() { trace!(did = %did, collection = %collection, rkey = %rkey, cid = %cid, "skip unchanged record"); continue; // skip unchanged record } - (DbAction::Update, false) + DbAction::Update } else { - (DbAction::Create, true) + DbAction::Create }; trace!(did = %did, collection = %collection, rkey = %rkey, cid = %cid, ?action, "action record"); // key is did|collection|rkey let db_key = keys::record_key(&did, collection, &rkey); - batch.insert(&app_state.db.blocks, cid.to_bytes(), val.as_ref()); - batch.insert(&app_state.db.records, db_key, cid.to_bytes()); + batch.batch_mut().insert(&app_state.db.blocks, cid.to_bytes(), val.as_ref()); + if !ephemeral { + batch.batch_mut().insert(&app_state.db.records, db_key, cid.to_bytes()); + } added_blocks += 1; - if is_new { + if action == DbAction::Create { delta += 1; *collection_counts.entry(path.0.clone()).or_default() += 1; } @@ -636,7 +651,27 @@ async fn process_did<'i>( cid: Some(cid_obj.to_ipld().expect("valid cid")), }; let bytes = rmp_serde::to_vec(&evt).into_diagnostic()?; - batch.insert(&app_state.db.events, keys::event_key(event_id), bytes); + batch.batch_mut().insert(&app_state.db.events, keys::event_key(event_id), bytes); + + // update block refcount + let cid_bytes = Slice::from(cid.to_bytes()); + // if ephemeral, we only care about events, so its 1 + // if not, then its 2 since we also insert to records + batch.update_block_refcount(cid_bytes, ephemeral.then_some(1).unwrap_or(2))?; + // for Update, also decrement old CID refcount + // event will still be there, so we only decrement for records + // which means only if not ephemeral + if !ephemeral && action == DbAction::Update { + let existing_cid = existing_cid.expect("that cid exists since this is Update"); + if existing_cid != cid_obj.as_str() { + let old_cid_bytes = Slice::from( + cid::Cid::from_str(&existing_cid) + .expect("valid cid from existing_cids") + .to_bytes() + ); + batch.update_block_refcount(old_cid_bytes, -1)?; + } + } count += 1; } @@ -646,11 +681,21 @@ async fn process_did<'i>( for ((collection, rkey), cid) in existing_cids { trace!(did = %did, collection = %collection, rkey = %rkey, cid = %cid, "remove existing record"); - batch.remove( + // we dont have to put if ephemeral around here since + // existing_cids will be empty anyyway + batch.batch_mut().remove( &app_state.db.records, keys::record_key(&did, &collection, &rkey), ); + // decrement block refcount + let cid_bytes = Slice::from( + cid::Cid::from_str(&cid) + .expect("valid cid from existing_cids") + .to_bytes() + ); + batch.update_block_refcount(cid_bytes, -1)?; + let event_id = app_state.db.next_event_id.fetch_add(1, Ordering::SeqCst); let evt = StoredEvent { live: false, @@ -662,7 +707,7 @@ async fn process_did<'i>( cid: None, }; let bytes = rmp_serde::to_vec(&evt).into_diagnostic()?; - batch.insert(&app_state.db.events, keys::event_key(event_id), bytes); + batch.batch_mut().insert(&app_state.db.events, keys::event_key(event_id), bytes); delta -= 1; count += 1; @@ -679,15 +724,17 @@ async fn process_did<'i>( state.data = Some(root_commit.data); state.last_updated_at = chrono::Utc::now().timestamp(); - batch.insert( + batch.batch_mut().insert( &app_state.db.repos, keys::repo_key(&did), ser_repo_state(&state)?, ); // add the counts - for (col, cnt) in collection_counts { - db::set_record_count(&mut batch, &app_state.db, &did, &col, cnt); + if !ephemeral { + for (col, cnt) in collection_counts { + db::set_record_count(batch.batch_mut(), &app_state.db, &did, &col, cnt); + } } batch.commit().into_diagnostic()?; diff --git a/src/db/compaction.rs b/src/db/compaction.rs index 8bcdbbe..da67c18 100644 --- a/src/db/compaction.rs +++ b/src/db/compaction.rs @@ -2,31 +2,6 @@ use fjall::compaction::filter::Context; use lsm_tree::compaction::{CompactionFilter, Factory}; use lsm_tree::compaction::{ItemAccessor, Verdict}; -mod drop_all { - use super::*; - - pub struct DropAllFilter; - - impl CompactionFilter for DropAllFilter { - fn filter_item(&mut self, _: ItemAccessor<'_>, _: &Context) -> lsm_tree::Result { - Ok(Verdict::Destroy) - } - } - - pub struct DropAllFilterFactory; - - impl Factory for DropAllFilterFactory { - fn name(&self) -> &str { - "drop_all" - } - - fn make_filter(&self, _: &Context) -> Box { - Box::new(DropAllFilter) - } - } -} -pub use drop_all::*; - mod drop_prefix { use super::*; diff --git a/src/db/gc.rs b/src/db/gc.rs new file mode 100644 index 0000000..211546c --- /dev/null +++ b/src/db/gc.rs @@ -0,0 +1,294 @@ +use fjall::Slice; +use fjall::compaction::filter::{CompactionFilter, Context, Factory, ItemAccessor, Verdict}; +use miette::{IntoDiagnostic, WrapErr}; +use scc::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; +use tracing::{error, info}; + +const EVENT_TTL_SECS: u64 = 3600; + +use crate::db::{Db, keys}; +use crate::types::StoredEvent; + +pub struct BlocksGcFilterFactory { + pub gc_ready: Arc, + pub refcounts: Arc>, +} + +struct BlocksGcFilter { + gc_ready: Arc, + refcounts: Arc>, +} + +impl Factory for BlocksGcFilterFactory { + fn make_filter(&self, _ctx: &Context) -> Box { + Box::new(BlocksGcFilter { + gc_ready: self.gc_ready.clone(), + refcounts: self.refcounts.clone(), + }) + } + + fn name(&self) -> &str { + "blocks_gc" + } +} + +impl CompactionFilter for BlocksGcFilter { + fn filter_item(&mut self, item: ItemAccessor<'_>, _ctx: &Context) -> lsm_tree::Result { + if !self.gc_ready.load(Ordering::SeqCst) { + return Ok(Verdict::Keep); + } + + let count = self + .refcounts + .read_sync(item.key().as_ref(), |_, v| *v) + .unwrap_or(0); + + #[cfg(debug_assertions)] + if let Ok(cid) = cid::Cid::read_bytes(item.key().as_ref()) { + tracing::debug!(cid = %cid, count, "BlocksGcFilter checking block"); + } + + Ok((count <= 0) + .then_some(Verdict::Destroy) + .unwrap_or(Verdict::Keep)) + } +} + +pub fn startup_load_refcounts(db: &Db) -> miette::Result<()> { + let checkpoint_seq = db + .cursors + .get(keys::BLOCK_REFS_CHECKPOINT_SEQ_KEY) + .into_diagnostic()? + .map(|v| { + u64::from_be_bytes( + v.as_ref() + .try_into() + .expect("checkpoint seq should be 8 bytes"), + ) + }) + .unwrap_or(0); + + // check if we need to run the one-time migration + let needs_migration = db + .counts + .get(keys::count_keyspace_key("gc_schema_version")) + .into_diagnostic()? + .is_none(); + + if needs_migration { + migrate_build_refcounts(db)?; + } + + // load snapshot + for guard in db.block_refs.iter() { + let (k, v) = guard.into_inner().into_diagnostic()?; + let count = i64::from_be_bytes( + v.as_ref() + .try_into() + .into_diagnostic() + .wrap_err("invalid block_refs count bytes")?, + ); + let _ = db.block_refcounts.insert_sync(k, count); + } + + // replay WAL since checkpoint + let start_key = keys::reflog_key(checkpoint_seq.saturating_add(1)); + for guard in db.block_reflog.range(start_key..) { + let (_, v) = guard.into_inner().into_diagnostic()?; + let (cid, delta): (Vec, i8) = rmp_serde::from_slice(&v).into_diagnostic()?; + let cid = Slice::from(cid); + let mut entry = db.block_refcounts.entry_sync(cid).or_insert(0); + *entry += delta as i64; + } + + db.gc_ready.store(true, Ordering::SeqCst); + info!("block refcounts loaded, gc ready"); + Ok(()) +} + +fn migrate_build_refcounts(db: &Db) -> miette::Result<()> { + info!("building initial block refcounts from existing records (one-time migration)"); + let mut batch = db.inner.batch(); + + // scan records + for guard in db.records.iter() { + let cid_bytes = guard.value().into_diagnostic()?; + let mut entry = db.block_refcounts.entry_sync(cid_bytes).or_insert(0); + *entry += 1i64; + } + + // events with cids + for guard in db.events.iter() { + let v = guard.value().into_diagnostic()?; + let evt = rmp_serde::from_slice::(&v).into_diagnostic()?; + let Some(cid) = evt.cid else { + continue; + }; + let cid_bytes = Slice::from(cid.to_bytes()); + let mut entry = db.block_refcounts.entry_sync(cid_bytes).or_insert(0); + *entry += 1i64; + } + + // persist as initial checkpoint + db.block_refcounts.iter_sync(|k, v| { + batch.insert(&db.block_refs, k.as_ref(), v.to_be_bytes()); + true + }); + + let seq = db.next_reflog_seq.load(Ordering::SeqCst); + batch.insert( + &db.cursors, + keys::BLOCK_REFS_CHECKPOINT_SEQ_KEY, + seq.to_be_bytes(), + ); + // mark migration done + let one: u64 = 1; + batch.insert( + &db.counts, + keys::count_keyspace_key("gc_schema_version"), + one.to_be_bytes(), + ); + batch.commit().into_diagnostic()?; + + info!("block refcount migration complete"); + Ok(()) +} + +pub fn checkpoint_worker(state: Arc) { + info!("block refs checkpoint worker started"); + loop { + std::thread::sleep(Duration::from_secs(300)); + if let Err(e) = checkpoint(&state.db) { + error!(err = %e, "block refs checkpoint failed"); + } + } +} + +fn checkpoint(db: &Db) -> miette::Result<()> { + let checkpoint_seq = db.next_reflog_seq.load(Ordering::SeqCst).saturating_sub(1); + + let mut batch = db.inner.batch(); + + db.block_refcounts.iter_sync(|k, v| { + batch.insert(&db.block_refs, k.as_ref(), v.to_be_bytes()); + true + }); + + batch.insert( + &db.cursors, + keys::BLOCK_REFS_CHECKPOINT_SEQ_KEY, + checkpoint_seq.to_be_bytes(), + ); + + // truncate reflog up to and including checkpoint_seq + for guard in db.block_reflog.range(..=keys::reflog_key(checkpoint_seq)) { + let k = guard.key().into_diagnostic()?; + batch.remove(&db.block_reflog, k); + } + + batch.commit().into_diagnostic()?; + info!(seq = checkpoint_seq, "block refs checkpoint complete"); + Ok(()) +} + +pub fn ephemeral_startup_load_refcounts(db: &Db) -> miette::Result<()> { + info!("rebuilding block refcounts from events (ephemeral mode)"); + + for guard in db.events.iter() { + let v = guard.value().into_diagnostic()?; + let evt = rmp_serde::from_slice::(&v).into_diagnostic()?; + let Some(cid) = evt.cid else { + continue; + }; + let cid_bytes = Slice::from(cid.to_bytes()); + let mut entry = db.block_refcounts.entry_sync(cid_bytes).or_insert(0); + *entry += 1; + } + + db.gc_ready.store(true, Ordering::SeqCst); + info!("ephemeral block refcounts ready"); + Ok(()) +} + +pub fn ephemeral_ttl_worker(state: Arc) { + info!("ephemeral TTL worker started"); + loop { + std::thread::sleep(Duration::from_secs(60)); + if let Err(e) = ephemeral_ttl_tick(&state.db) { + error!(err = %e, "ephemeral TTL tick failed"); + } + } +} + +fn ephemeral_ttl_tick(db: &Db) -> miette::Result<()> { + let now = chrono::Utc::now().timestamp() as u64; + let cutoff_ts = now.saturating_sub(EVENT_TTL_SECS); + + // write current watermark + let current_event_id = db.next_event_id.load(Ordering::SeqCst); + db.cursors + .insert( + keys::event_watermark_key(now), + current_event_id.to_be_bytes(), + ) + .into_diagnostic()?; + + // find the watermark entry closest to and <= cutoff_ts + let cutoff_key = keys::event_watermark_key(cutoff_ts); + let cutoff_event_id = db + .cursors + .range(..=cutoff_key.clone()) + .next_back() + .map(|g| g.into_inner().into_diagnostic()) + .transpose()? + .filter(|(k, _)| k.starts_with(keys::EVENT_WATERMARK_PREFIX)) + .map(|(_, v)| { + v.as_ref() + .try_into() + .into_diagnostic() + .wrap_err("expected cutoff event id to be u64") + }) + .transpose()? + .map(u64::from_be_bytes); + + let Some(cutoff_event_id) = cutoff_event_id else { + // no watermark old enough yet, nothing to prune + return Ok(()); + }; + + let cutoff_key_events = keys::event_key(cutoff_event_id); + let mut batch = db.inner.batch(); + let mut pruned = 0usize; + + for guard in db.events.range(..cutoff_key_events) { + let (k, v) = guard.into_inner().into_diagnostic()?; + let evt = rmp_serde::from_slice::(&v).into_diagnostic()?; + let Some(cid) = evt.cid else { + continue; + }; + let cid_bytes = Slice::from(cid.to_bytes()); + let mut entry = db.block_refcounts.entry_sync(cid_bytes).or_insert(0); + *entry -= 1; + batch.remove(&db.events, k); + pruned += 1; + } + + // clean up consumed watermark entries (everything up to and including cutoff_ts) + for guard in db.cursors.range(..=cutoff_key) { + let k = guard.key().into_diagnostic()?; + if k.starts_with(keys::EVENT_WATERMARK_PREFIX) { + batch.remove(&db.cursors, k); + } + } + + batch.commit().into_diagnostic()?; + + if pruned > 0 { + info!(pruned, "pruned old events"); + } + + Ok(()) +} diff --git a/src/db/keys.rs b/src/db/keys.rs index 386ce2c..b623671 100644 --- a/src/db/keys.rs +++ b/src/db/keys.rs @@ -8,7 +8,11 @@ pub const SEP: u8 = b'|'; pub const CURSOR_KEY: &[u8] = b"firehose_cursor"; -// Key format: {DID} +pub const BLOCK_REFS_CHECKPOINT_SEQ_KEY: &[u8] = b"block_refs_checkpoint_seq"; + +pub const EVENT_WATERMARK_PREFIX: &[u8] = b"ewm|"; + +// key format: {DID} pub fn repo_key<'a>(did: &'a Did) -> Vec { let mut vec = Vec::with_capacity(32); TrimmedDid::from(did).write_to_vec(&mut vec); @@ -19,6 +23,17 @@ pub fn pending_key(id: u64) -> [u8; 8] { id.to_be_bytes() } +pub fn reflog_key(seq: u64) -> [u8; 8] { + seq.to_be_bytes() +} + +pub fn event_watermark_key(timestamp_secs: u64) -> Vec { + let mut key = Vec::with_capacity(EVENT_WATERMARK_PREFIX.len() + 8); + key.extend_from_slice(EVENT_WATERMARK_PREFIX); + key.extend_from_slice(×tamp_secs.to_be_bytes()); + key +} + // prefix format: {DID}| (DID trimmed) pub fn record_prefix_did(did: &Did) -> Vec { let repo = TrimmedDid::from(did); diff --git a/src/db/mod.rs b/src/db/mod.rs index 3c93ac9..9f4a1e8 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -1,6 +1,6 @@ -use crate::db::compaction::{DropAllFilterFactory, DropPrefixFilterFactory}; +use crate::db::compaction::DropPrefixFilterFactory; use crate::types::{BroadcastEvent, RepoState}; -use fjall::compaction::{Fifo, Levelled}; + use fjall::config::BlockSizePolicy; use fjall::{Database, Keyspace, KeyspaceCreateOptions, OwnedWriteBatch, PersistMode, Slice}; use jacquard_common::IntoStatic; @@ -11,13 +11,15 @@ use scc::HashMap; use smol_str::SmolStr; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64}; pub mod compaction; pub mod filter; +pub mod gc; pub mod keys; +pub mod refcount; pub mod types; -use std::sync::atomic::AtomicU64; use tokio::sync::broadcast; use tracing::error; @@ -38,6 +40,11 @@ pub struct Db { pub counts: Keyspace, pub filter: Keyspace, pub crawler: Keyspace, + pub block_refs: Keyspace, + pub block_reflog: Keyspace, + pub block_refcounts: Arc>, + pub gc_ready: Arc, + pub next_reflog_seq: Arc, pub event_tx: broadcast::Sender, pub next_event_id: Arc, pub counts_map: HashMap, @@ -98,6 +105,12 @@ impl Db { v * 1024 * 1024 } + let gc_ready = Arc::new(AtomicBool::new(false)); + let block_refcounts: Arc> = Arc::new(HashMap::new()); + + let gc_ready_factory = gc_ready.clone(); + let refcounts_factory = block_refcounts.clone(); + let db = Database::builder(&cfg.database_path) .cache_size(cfg.cache_size * 2_u64.pow(20) / 2) .manual_journal_persist(true) @@ -110,16 +123,20 @@ impl Db { .max_journaling_size(mb(cfg.db_max_journaling_size_mb)) .with_compaction_filter_factories({ let ephemeral = cfg.ephemeral; - let f = move |ks: &str| match ks { - "records" => { - ephemeral.then(|| -> Arc { Arc::new(DropAllFilterFactory) }) + let f = move |ks: &str| { + tracing::info!("with_compaction_filter_factories queried for keyspace: {ks}",); + match ks { + "counts" => ephemeral.then(|| -> Arc { + Arc::new(DropPrefixFilterFactory { + prefix: keys::COUNT_COLLECTION_PREFIX, + }) + }), + "blocks" => Some(Arc::new(gc::BlocksGcFilterFactory { + gc_ready: gc_ready_factory.clone(), + refcounts: refcounts_factory.clone(), + }) as Arc), + _ => None, } - "counts" => ephemeral.then(|| -> Arc { - Arc::new(DropPrefixFilterFactory { - prefix: keys::COUNT_COLLECTION_PREFIX, - }) - }), - _ => None, }; Arc::new(f) }) @@ -197,12 +214,7 @@ impl Db { // only iterators are used here, no point reads .expect_point_read_hits(true) .max_memtable_size(mb(cfg.db_events_memtable_size_mb)) - .data_block_size_policy(BlockSizePolicy::new([kb(16), kb(32)])) - .compaction_strategy(if cfg.ephemeral { - Arc::new(Fifo::new(mb(512), Some(60 * 60))) - } else { - Arc::new(Levelled::default()) - }), + .data_block_size_policy(BlockSizePolicy::new([kb(16), kb(32)])), )?; let counts = open_ks( "counts", @@ -231,6 +243,24 @@ impl Db { .data_block_size_policy(BlockSizePolicy::all(kb(1))), )?; + let block_refs = open_ks( + "block_refs", + opts() + // only ever iterated on + .expect_point_read_hits(true) + .max_memtable_size(mb(cfg.db_records_memtable_size_mb / 2)) + .data_block_size_policy(BlockSizePolicy::all(kb(64))), + )?; + + let block_reflog = open_ks( + "block_reflog", + opts() + // only ever iterated on + .expect_point_read_hits(true) + .max_memtable_size(mb(4)) + .data_block_size_policy(BlockSizePolicy::all(kb(64))), + )?; + let mut last_id = 0; if let Some(guard) = events.iter().next_back() { let k = guard.key().into_diagnostic()?; @@ -242,6 +272,18 @@ impl Db { ); } + let mut last_reflog_seq = 0u64; + if let Some(guard) = block_reflog.iter().next_back() { + let k = guard.key().into_diagnostic()?; + last_reflog_seq = u64::from_be_bytes( + k.as_ref() + .try_into() + .into_diagnostic() + .wrap_err("expected to be reflog seq (8 bytes)")?, + ); + } + let next_reflog_seq = Arc::new(AtomicU64::new(last_reflog_seq.saturating_add(1))); + // load counts into memory let counts_map = HashMap::new(); for guard in counts.prefix(keys::COUNT_KS_PREFIX) { @@ -284,6 +326,11 @@ impl Db { counts, filter, crawler, + block_refs, + block_reflog, + block_refcounts, + gc_ready, + next_reflog_seq, event_tx, counts_map, next_event_id: Arc::new(AtomicU64::new(last_id + 1)), diff --git a/src/db/refcount.rs b/src/db/refcount.rs new file mode 100644 index 0000000..83406ba --- /dev/null +++ b/src/db/refcount.rs @@ -0,0 +1,61 @@ +use crate::db::{Db, keys}; + +use fjall::{OwnedWriteBatch, Slice}; +use miette::{IntoDiagnostic, Result}; +use std::sync::atomic::Ordering; + +/// a write batch that tracks block refcount deltas and applies them +/// automatically on commit. this prevents the bug class where callers +/// forget to call `apply_block_refcount_deltas` after committing. +pub struct RefcountedBatch<'db> { + batch: OwnedWriteBatch, + db: &'db Db, + pending_deltas: Vec<(Slice, i8)>, +} + +impl<'db> RefcountedBatch<'db> { + pub fn new(db: &'db Db) -> Self { + Self { + batch: db.inner.batch(), + db, + pending_deltas: Vec::new(), + } + } + + /// records a refcount delta for the given CID. writes to the reflog + /// in the batch and accumulates the delta for in-memory application on commit. + pub fn update_block_refcount(&mut self, cid_bytes: Slice, delta: i8) -> Result<()> { + #[cfg(debug_assertions)] + if let Ok(cid) = cid::Cid::read_bytes(cid_bytes.as_ref()) { + tracing::debug!(delta, %cid, "update_block_refcount"); + } + + let value = rmp_serde::to_vec(&(cid_bytes.as_ref(), delta)).into_diagnostic()?; + let seq = self.db.next_reflog_seq.fetch_add(1, Ordering::SeqCst); + self.batch + .insert(&self.db.block_reflog, keys::reflog_key(seq), value); + self.pending_deltas.push((cid_bytes, delta)); + Ok(()) + } + + /// commits the batch and applies all accumulated refcount deltas to the in-memory map. + pub fn commit(self) -> Result<(), fjall::Error> { + self.batch.commit()?; + apply_deltas(self.db, &self.pending_deltas); + Ok(()) + } + + pub fn batch_mut(&mut self) -> &mut OwnedWriteBatch { + &mut self.batch + } +} + +fn apply_deltas(db: &Db, deltas: &[(Slice, i8)]) { + for (cid_bytes, delta) in deltas { + let mut entry = db + .block_refcounts + .entry_sync(cid_bytes.clone()) + .or_insert(0); + *entry += *delta as i64; + } +} diff --git a/src/ingest/worker.rs b/src/ingest/worker.rs index 30a9fbc..9c2b8bb 100644 --- a/src/ingest/worker.rs +++ b/src/ingest/worker.rs @@ -7,7 +7,7 @@ use crate::resolver::{NoSigningKeyError, ResolverError}; use crate::state::AppState; use crate::types::{AccountEvt, BroadcastEvent, GaugeState, IdentityEvt, RepoState, RepoStatus}; -use fjall::OwnedWriteBatch; +use crate::db::refcount::RefcountedBatch; use jacquard_common::IntoStatic; use jacquard_common::cowstr::ToCowStr; @@ -59,13 +59,15 @@ pub struct FirehoseWorker { state: Arc, rx: BufferRx, verify_signatures: bool, + ephemeral: bool, num_shards: usize, } struct WorkerContext<'a> { verify_signatures: bool, + ephemeral: bool, state: &'a AppState, - batch: &'a mut OwnedWriteBatch, + batch: RefcountedBatch<'a>, added_blocks: &'a mut i64, records_delta: &'a mut i64, broadcast_events: &'a mut Vec, @@ -77,12 +79,14 @@ impl FirehoseWorker { state: Arc, rx: BufferRx, verify_signatures: bool, + ephemeral: bool, num_shards: usize, ) -> Self { Self { state, rx, verify_signatures, + ephemeral, num_shards, } } @@ -99,12 +103,13 @@ impl FirehoseWorker { let state = self.state.clone(); let verify = self.verify_signatures; + let ephemeral = self.ephemeral; let handle = handle.clone(); std::thread::Builder::new() .name(format!("ingest-shard-{i}")) .spawn(move || { - Self::shard(i, rx, state, verify, handle); + Self::shard(i, rx, state, verify, ephemeral, handle); }) .into_diagnostic()?; } @@ -149,6 +154,7 @@ impl FirehoseWorker { mut rx: mpsc::UnboundedReceiver, state: Arc, verify_signatures: bool, + ephemeral: bool, handle: tokio::runtime::Handle, ) { let _guard = handle.enter(); @@ -157,7 +163,7 @@ impl FirehoseWorker { let mut broadcast_events = Vec::new(); while let Some(msg) = rx.blocking_recv() { - let mut batch = state.db.inner.batch(); + let batch = RefcountedBatch::new(&state.db); broadcast_events.clear(); let mut added_blocks = 0; @@ -165,12 +171,13 @@ impl FirehoseWorker { let mut ctx = WorkerContext { state: &state, - batch: &mut batch, + batch, added_blocks: &mut added_blocks, records_delta: &mut records_delta, broadcast_events: &mut broadcast_events, handle: &handle, verify_signatures, + ephemeral, }; match msg { @@ -191,7 +198,7 @@ impl FirehoseWorker { // while the resync buffer is being drained, we should handle that probably // but also it should still be fine since we'll sync eventually anyway let res = ops::update_repo_status( - &mut batch, + ctx.batch.batch_mut(), &state.db, &did, s, @@ -260,7 +267,7 @@ impl FirehoseWorker { } } - if let Err(e) = batch.commit() { + if let Err(e) = ctx.batch.commit() { error!(shard = id, err = %e, "failed to commit batch"); } @@ -366,7 +373,7 @@ impl FirehoseWorker { let handle = identity.handle.as_ref().map(|h| h.clone()); repo_state.handle = handle.or(repo_state.handle); - ctx.batch.insert( + ctx.batch.batch_mut().insert( &ctx.state.db.repos, keys::repo_key(did), crate::db::ser_repo_state(&repo_state)?, @@ -394,7 +401,12 @@ impl FirehoseWorker { match &account.status { Some(AccountStatus::Deleted) => { debug!(did = %did, "account deleted, wiping data"); - crate::ops::delete_repo(ctx.batch, &ctx.state.db, did, &repo_state)?; + crate::ops::delete_repo( + &mut ctx.batch, + &ctx.state.db, + did, + &repo_state, + )?; return Ok(RepoProcessResult::Deleted); } status => { @@ -432,7 +444,7 @@ impl FirehoseWorker { } repo_state = ops::update_repo_status( - ctx.batch, + ctx.batch.batch_mut(), &ctx.state.db, did, repo_state, @@ -509,13 +521,15 @@ impl FirehoseWorker { return Ok(RepoProcessResult::Syncing(Some(commit))); } + let signing_key = Self::fetch_key(ctx, did)?; let res = ops::apply_commit( - ctx.batch, + &mut ctx.batch, &ctx.state.db, repo_state, &commit, - Self::fetch_key(ctx, did)?.as_ref(), + signing_key.as_ref(), &ctx.state.filter.load(), + ctx.ephemeral, )?; let repo_state = res.repo_state; *ctx.added_blocks += res.blocks_count; @@ -629,7 +643,7 @@ impl FirehoseWorker { } } repo_state = ops::update_repo_status( - ctx.batch, + ctx.batch.batch_mut(), &ctx.state.db, did, repo_state, @@ -660,21 +674,27 @@ impl FirehoseWorker { Ok(r) => r, Err(e) => { if !Self::check_if_retriable_failure(&e) { - ctx.batch.remove(&ctx.state.db.resync_buffer, key); + ctx.batch + .batch_mut() + .remove(&ctx.state.db.resync_buffer, key); } return Err(e); } }; match res { RepoProcessResult::Ok(rs) => { - ctx.batch.remove(&ctx.state.db.resync_buffer, key); + ctx.batch + .batch_mut() + .remove(&ctx.state.db.resync_buffer, key); repo_state = rs; } RepoProcessResult::Syncing(_) => { return Ok(RepoProcessResult::Syncing(None)); } RepoProcessResult::Deleted => { - ctx.batch.remove(&ctx.state.db.resync_buffer, key); + ctx.batch + .batch_mut() + .remove(&ctx.state.db.resync_buffer, key); return Ok(RepoProcessResult::Deleted); } } @@ -692,7 +712,7 @@ impl FirehoseWorker { ctx.state.resolver.invalidate_sync(did); let doc = ctx.handle.block_on(ctx.state.resolver.resolve_doc(did))?; repo_state.update_from_doc(doc); - ctx.batch.insert( + ctx.batch.batch_mut().insert( &ctx.state.db.repos, keys::repo_key(did), crate::db::ser_repo_state(&repo_state)?, diff --git a/src/main.rs b/src/main.rs index 4b62ed1..486e953 100644 --- a/src/main.rs +++ b/src/main.rs @@ -29,6 +29,13 @@ async fn main() -> miette::Result<()> { let state = AppState::new(&cfg)?; + // load block refcounts for GC - must complete before any ingest workers start + if cfg.ephemeral { + db::gc::ephemeral_startup_load_refcounts(&state.db)?; + } else { + db::gc::startup_load_refcounts(&state.db)?; + } + if cfg.full_network || cfg.filter_signals.is_some() || cfg.filter_collections.is_some() @@ -76,6 +83,21 @@ async fn main() -> miette::Result<()> { let (buffer_tx, buffer_rx) = mpsc::unbounded_channel(); let state = Arc::new(state); + // spawn GC workers + if cfg.ephemeral { + let state_ttl = state.clone(); + std::thread::Builder::new() + .name("ephemeral-ttl".into()) + .spawn(move || db::gc::ephemeral_ttl_worker(state_ttl)) + .into_diagnostic()?; + } else { + let state_gc = state.clone(); + std::thread::Builder::new() + .name("gc-checkpoint".into()) + .spawn(move || db::gc::checkpoint_worker(state_gc)) + .into_diagnostic()?; + } + if cfg.enable_backfill { tokio::spawn({ let state = state.clone(); @@ -89,6 +111,7 @@ async fn main() -> miette::Result<()> { cfg.verify_signatures, SignatureVerification::Full | SignatureVerification::BackfillOnly ), + cfg.ephemeral, ) .run() }); @@ -213,6 +236,7 @@ async fn main() -> miette::Result<()> { state, buffer_rx, matches!(cfg.verify_signatures, SignatureVerification::Full), + cfg.ephemeral, cfg.firehose_workers, ) .run(handle) diff --git a/src/ops.rs b/src/ops.rs index 6f4ba75..f6303e8 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -1,4 +1,7 @@ use fjall::OwnedWriteBatch; +use fjall::Slice; + +use crate::db::refcount::RefcountedBatch; use jacquard_common::CowStr; use jacquard_common::IntoStatic; use jacquard_common::types::cid::Cid; @@ -64,8 +67,8 @@ pub fn make_account_event(db: &Db, evt: AccountEvt<'static>) -> BroadcastEvent { BroadcastEvent::Ephemeral(Box::new(marshallable)) } -pub fn delete_repo<'batch>( - batch: &'batch mut OwnedWriteBatch, +pub fn delete_repo( + batch: &mut RefcountedBatch<'_>, db: &Db, did: &Did, repo_state: &RepoState, @@ -76,14 +79,14 @@ pub fn delete_repo<'batch>( let pending_key = keys::pending_key(repo_state.index_id); // 1. delete from repos, pending, resync - batch.remove(&db.repos, &repo_key); + batch.batch_mut().remove(&db.repos, &repo_key); match repo_state.status { RepoStatus::Synced => {} RepoStatus::Backfilling => { - batch.remove(&db.pending, &pending_key); + batch.batch_mut().remove(&db.pending, &pending_key); } _ => { - batch.remove(&db.resync, &repo_key); + batch.batch_mut().remove(&db.resync, &repo_key); } } @@ -91,14 +94,15 @@ pub fn delete_repo<'batch>( let resync_prefix = keys::resync_buffer_prefix(did); for guard in db.resync_buffer.prefix(&resync_prefix) { let k = guard.key().into_diagnostic()?; - batch.remove(&db.resync_buffer, k); + batch.batch_mut().remove(&db.resync_buffer, k); } // 3. delete from records let records_prefix = keys::record_prefix_did(did); for guard in db.records.prefix(&records_prefix) { - let k = guard.key().into_diagnostic()?; - batch.remove(&db.records, k); + let (k, cid_bytes) = guard.into_inner().into_diagnostic()?; + batch.update_block_refcount(cid_bytes, -1); + batch.batch_mut().remove(&db.records, k); } // 4. reset collection counts @@ -110,7 +114,7 @@ pub fn delete_repo<'batch>( for guard in db.counts.prefix(&count_prefix) { let k = guard.key().into_diagnostic()?; - batch.remove(&db.counts, k); + batch.batch_mut().remove(&db.counts, k); } Ok(()) @@ -218,13 +222,14 @@ pub struct ApplyCommitResults<'s> { pub blocks_count: i64, } -pub fn apply_commit<'batch, 'db, 'commit, 's>( - batch: &'batch mut OwnedWriteBatch, +pub fn apply_commit<'db, 'commit, 's>( + batch: &mut RefcountedBatch<'db>, db: &'db Db, mut repo_state: RepoState<'s>, commit: &'commit Commit<'commit>, signing_key: Option<&PublicKey>, filter: &FilterConfig, + ephemeral: bool, ) -> Result> { let did = &commit.repo; debug!(did = %did, commit = %commit.commit, "applying commit"); @@ -257,7 +262,9 @@ pub fn apply_commit<'batch, 'db, 'commit, 's>( repo_state.data = Some(repo_commit.data); repo_state.last_updated_at = chrono::Utc::now().timestamp(); - batch.insert(&db.repos, keys::repo_key(did), ser_repo_state(&repo_state)?); + batch + .batch_mut() + .insert(&db.repos, keys::repo_key(did), ser_repo_state(&repo_state)?); // 2. iterate ops and update records index let mut records_delta = 0; @@ -287,25 +294,53 @@ pub fn apply_commit<'batch, 'db, 'commit, 's>( .into_diagnostic() .wrap_err("expected valid cid from relay")?; - if let Some(bytes) = parsed.blocks.get(&cid_ipld) { - batch.insert(&db.blocks, cid_ipld.to_bytes(), bytes.to_vec()); - blocks_count += 1; - } - - batch.insert(&db.records, db_key.clone(), cid_ipld.to_bytes()); - - // accumulate counts - if action == DbAction::Create { - records_delta += 1; - *collection_deltas.entry(collection).or_default() += 1; + let Some(bytes) = parsed.blocks.get(&cid_ipld) else { + return Err(miette::miette!( + "block {cid} not found in CAR for record {did}/{collection}/{rkey}" + )); + }; + let cid_bytes = Slice::from(cid_ipld.to_bytes()); + batch + .batch_mut() + .insert(&db.blocks, cid_bytes.clone(), bytes.to_vec()); + blocks_count += 1; + batch.update_block_refcount(cid_bytes.clone(), ephemeral.then_some(1).unwrap_or(2)); + + if !ephemeral { + batch + .batch_mut() + .insert(&db.records, db_key.clone(), cid_ipld.to_bytes()); + // for Update, also decrement old CID refcount + if action == DbAction::Update { + let Some(old_cid_bytes) = db.records.get(&db_key).into_diagnostic()? else { + return Err(miette::miette!( + "!!! THIS IS A BUG !!! expected previous cid to be there for record being updated ({did}/{collection}/{rkey}). how did we get here?" + )); + }; + if old_cid_bytes != cid_bytes { + batch.update_block_refcount(old_cid_bytes, -1); + } + } + // accumulate counts + if action == DbAction::Create { + records_delta += 1; + *collection_deltas.entry(collection).or_default() += 1; + } } } DbAction::Delete => { - batch.remove(&db.records, db_key); - - // accumulate counts - records_delta -= 1; - *collection_deltas.entry(collection).or_default() -= 1; + if !ephemeral { + // decrement block refcount + let old_cid_bytes = db.records.get(&db_key).into_diagnostic()?; + if let Some(cid_bytes) = old_cid_bytes { + batch.update_block_refcount(cid_bytes, -1); + } + batch.batch_mut().remove(&db.records, db_key); + + // accumulate counts + records_delta -= 1; + *collection_deltas.entry(collection).or_default() -= 1; + } } } @@ -320,12 +355,16 @@ pub fn apply_commit<'batch, 'db, 'commit, 's>( }; let bytes = rmp_serde::to_vec(&evt).into_diagnostic()?; - batch.insert(&db.events, keys::event_key(event_id), bytes); + batch + .batch_mut() + .insert(&db.events, keys::event_key(event_id), bytes); } // update counts - for (col, delta) in collection_deltas { - db::update_record_count(batch, db, did, col, delta)?; + if !ephemeral { + for (col, delta) in collection_deltas { + db::update_record_count(batch.batch_mut(), db, did, col, delta)?; + } } Ok(ApplyCommitResults { diff --git a/tests/common.nu b/tests/common.nu index f96a22e..6c585da 100644 --- a/tests/common.nu +++ b/tests/common.nu @@ -52,8 +52,8 @@ export def activate-account [pds_url: string, jwt: string] { # build the hydrant binary export def build-hydrant [] { print "building hydrant..." - cargo build --release - "target/release/hydrant" + cargo build + "target/debug/hydrant" } # start hydrant in the background diff --git a/tests/gc_test.nu b/tests/gc_test.nu new file mode 100755 index 0000000..7c36fcc --- /dev/null +++ b/tests/gc_test.nu @@ -0,0 +1,159 @@ +#!/usr/bin/env nu +use common.nu * + +def run-test-instance [name: string, scenario_closure: closure] { + let port = 3004 + let debug_port = $port + 1 + let url = $"http://localhost:($port)" + let debug_url = $"http://localhost:($debug_port)" + let db_path = (mktemp -d -t hydrant_gc_test.XXXXXX) + + print $"--- running scenario: ($name) ---" + print $"database path: ($db_path)" + + let binary = build-hydrant + let instance = start-hydrant $binary $db_path $port + + try { + if not (wait-for-api $url) { + error make {msg: "api failed to start"} + } + + do $scenario_closure $url $debug_url + + print $"PASSED: ($name)\n" + } catch { |e| + print $"test failed: ($e.msg)" + try { kill $instance.pid } + exit 1 + } + + try { kill $instance.pid } +} + +def wait-for-blocks [debug_url: string] { + print "waiting for blocks to appear..." + mut blocks = {} + mut count = 0 + for i in 1..30 { + $blocks = (http get $"($debug_url)/debug/iter?partition=blocks&limit=1000") + $count = ($blocks.items | length) + if $count > 0 { + break + } + sleep 2sec + } + if $count == 0 { + error make {msg: "FAILED: no blocks found after backfill"} + } + $count +} + +def compact-and-check-blocks [debug_url: string, expected_count: int] { + print "triggering major compaction on blocks partition..." + http post -H [Content-Length 0] $"($debug_url)/debug/compact?partition=blocks" "" + + let blocks_after = http get $"($debug_url)/debug/iter?partition=blocks&limit=1000" + let after_count = ($blocks_after.items | length) + + if $after_count != $expected_count { + error make {msg: $"FAILED: expected ($expected_count) blocks after compaction, found ($after_count)"} + } +} + +def ack-all-events [debug_url: string, url: string] { + print "acking all events..." + mut items = [] + for i in 1..30 { + let events = http get $"($debug_url)/debug/iter?partition=events&limit=1000" + $items = $events.items + if ($items | length) > 0 { + break + } + sleep 2sec + } + + if ($items | length) == 0 { + error make {msg: "FAILED: no events to ack"} + } + + let event_ids = ($items | each { |x| ($x | first | into int) }) + + http post -t application/json $"($url)/stream/ack" { ids: $event_ids } + print $"acked ($event_ids | length) events" +} + +def main [] { + let repo1 = "did:web:guestbook.gaze.systems" + let repo2 = "did:plc:dfl62fgb7wtjj3fcbb72naae" + + run-test-instance "delete repo only" { |url, debug_url| + print $"adding repo ($repo1) to tracking..." + http put -t application/json $"($url)/repos" [ { did: ($repo1) } ] + + let before_count = (wait-for-blocks $debug_url) + print $"found ($before_count) blocks before GC" + + print "deleting repo..." + http delete -t application/json $"($url)/repos" --data [ { did: ($repo1), delete_data: true } ] + sleep 1sec + + compact-and-check-blocks $debug_url $before_count + } + + run-test-instance "ack events only" { |url, debug_url| + print $"adding repo ($repo1) to tracking..." + http put -t application/json $"($url)/repos" [ { did: ($repo1) } ] + + let before_count = (wait-for-blocks $debug_url) + print $"found ($before_count) blocks before GC" + + ack-all-events $debug_url $url + sleep 1sec + + compact-and-check-blocks $debug_url $before_count + } + + run-test-instance "delete repo, ack events" { |url, debug_url| + print $"adding repo ($repo1) to tracking..." + http put -t application/json $"($url)/repos" [ { did: ($repo1) } ] + + let before_count = (wait-for-blocks $debug_url) + print $"found ($before_count) blocks before GC" + + print "deleting repo..." + http delete -t application/json $"($url)/repos" --data [ { did: ($repo1), delete_data: true } ] + + ack-all-events $debug_url $url + sleep 1sec + + compact-and-check-blocks $debug_url 0 + } + + run-test-instance "multiple repos" { |url, debug_url| + print $"adding repo ($repo2) to tracking..." + http put -t application/json $"($url)/repos" [ { did: ($repo2) } ] + let repo2_blocks = (wait-for-blocks $debug_url) + print $"found ($repo2_blocks) blocks for repo2" + + print $"adding repo ($repo1) to tracking..." + http put -t application/json $"($url)/repos" [ { did: ($repo1) } ] + + # wait a bit more for repo1 blocks to finish + sleep 5sec + let total_blocks = (http get $"($debug_url)/debug/iter?partition=blocks&limit=1000000" | get items | length) + print $"found ($total_blocks) total blocks before GC" + + print $"deleting repo ($repo1)..." + http delete -t application/json $"($url)/repos" --data [ { did: ($repo1), delete_data: true } ] + + # ack events specifically for repo1? Actually wait, the events endpoint contains all events. + # we will ack all events to be safe. Since repo2 is NOT deleted, its refcount should be fine even if events are acked. + ack-all-events $debug_url $url + sleep 1sec + + compact-and-check-blocks $debug_url $repo2_blocks + } + + print "all gc tests passed!" +} diff --git a/tests/repos_api_test.nu b/tests/repos_api_test.nu index ce10ca3..42eec51 100644 --- a/tests/repos_api_test.nu +++ b/tests/repos_api_test.nu @@ -110,8 +110,8 @@ def main [] { # Helper to build hydrant def build-hydrant [] { - cargo build --release --quiet - "./target/release/hydrant" + cargo build --quiet + "./target/debug/hydrant" } # Helper to start hydrant -- 2.51.2