From a8f59ecf32a10d61fcf5e6f6c8f7c47ff252baf9 Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Sat, 11 Jul 2026 21:51:42 +0300 Subject: [PATCH] [db] migrate async boundaries to Db::run --- src/api/debug.rs | 313 +++++++++++++++-------------- src/api/repos.rs | 24 +-- src/api/xrpc/list_repos.rs | 76 +++---- src/api/xrpc/request_crawl.rs | 15 +- src/backfill/manager.rs | 48 ++--- src/backfill/sparse.rs | 15 +- src/backfill/worker/process.rs | 119 +++++------ src/backfill/worker/task.rs | 109 ++++------ src/backlinks/mod.rs | 12 +- src/control/crawler.rs | 58 +++--- src/control/db.rs | 4 +- src/control/filter.rs | 65 +++--- src/control/firehose.rs | 109 +++++----- src/control/hosts.rs | 69 +++---- src/control/hydrant.rs | 41 ++-- src/control/hydrant/run.rs | 21 +- src/control/pds.rs | 18 +- src/control/repos/indexer.rs | 132 +++++------- src/control/repos/mod.rs | 36 ++-- src/control/stats.rs | 10 +- src/control/stream/relay.rs | 6 +- src/crawler/by_collection.rs | 5 +- src/crawler/list_repos/producer.rs | 17 +- src/crawler/list_repos/retry.rs | 4 +- src/crawler/worker.rs | 57 +++--- src/db/counts.rs | 4 +- src/db/ephemeral.rs | 18 +- src/db/keyspaces.rs | 4 + src/db/lifecycle_counts.rs | 43 ++-- src/db/migration/v8.rs | 6 +- src/db/mod.rs | 38 ++-- src/db/open.rs | 130 +++++++----- src/db/txn.rs | 138 +++++++++++-- src/ingest/firehose_stats.rs | 4 +- src/ingest/indexer/shard.rs | 29 +-- src/ingest/relay/context.rs | 5 +- src/ingest/relay/handlers.rs | 10 +- src/ingest/relay/sink/indexer.rs | 78 +++---- src/ingest/relay/sink/none.rs | 12 +- src/ingest/relay/sink/relay.rs | 43 ++-- src/ingest/relay/worker.rs | 37 +--- src/ingest/validation.rs | 2 +- src/jetstream.rs | 3 +- src/types.rs | 1 - 44 files changed, 1014 insertions(+), 974 deletions(-) diff --git a/src/api/debug.rs b/src/api/debug.rs index e06d6a0..80b23f0 100644 --- a/src/api/debug.rs +++ b/src/api/debug.rs @@ -10,6 +10,7 @@ use axum::{ }; #[cfg(feature = "indexer")] use jacquard_common::types::ident::AtIdentifier; +use miette::IntoDiagnostic; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::sync::Arc; @@ -59,25 +60,25 @@ pub async fn handle_debug_count( .await .map_err(|_| StatusCode::BAD_REQUEST)?; - let db = &state.db; - let ks = db - .keyspace_by_name("records") - .expect("records keyspace exists in indexer mode"); - - // {TrimmedDid}|{collection}| let prefix = keys::record_prefix_collection(&did, &req.collection); - let count = tokio::task::spawn_blocking(move || { - let start_key = prefix.clone(); - let mut end_key = prefix.clone(); - if let Some(msg) = end_key.last_mut() { - *msg += 1; - } + let count = state + .db + .run(move |db| { + let ks = db + .keyspace_by_name("records") + .expect("records keyspace exists in indexer mode"); + + let start_key = prefix.clone(); + let mut end_key = prefix.clone(); + if let Some(msg) = end_key.last_mut() { + *msg += 1; + } - ks.range(start_key..end_key).count() - }) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + Ok::<_, miette::Report>(ks.range(start_key..end_key).count()) + }) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok(Json(DebugCountResponse { count })) } @@ -102,7 +103,13 @@ pub async fn handle_debug_get( let key = registry::debug_parse_key(&req.partition, &req.key).ok_or(StatusCode::BAD_REQUEST)?; let partition = req.partition.clone(); - let value = crate::db::Db::get(ks, key) + let value = state + .db + .run(move |_| { + ks.get(key) + .inspect_err(crate::db::check_poisoned) + .into_diagnostic() + }) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? .map(|v| registry::debug_value(&partition, &v)); @@ -128,7 +135,6 @@ pub async fn handle_debug_iter( State(state): State>, Query(req): Query, ) -> Result, StatusCode> { - let ks = get_keyspace_by_name(&state.db, &req.partition)?; let partition = req.partition.clone(); let parse_bound = |s: Option| -> Result>, StatusCode> { @@ -139,55 +145,59 @@ pub async fn handle_debug_iter( let start = parse_bound(req.start)?; let end = parse_bound(req.end)?; - let items = tokio::task::spawn_blocking(move || { - let limit = req.limit.unwrap_or(50); + let items = state + .db + .run(move |db| { + let ks = get_keyspace_by_name(db, &req.partition) + .map_err(|_| miette::miette!("bad request"))?; + let limit = req.limit.unwrap_or(50); - let collect = |iter: &mut dyn Iterator| { - let mut items = Vec::new(); - for guard in iter.take(limit) { - let (k, v) = guard - .into_inner() - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + let collect = |iter: &mut dyn Iterator| { + let mut items = Vec::new(); + for guard in iter.take(limit) { + let (k, v) = guard + .into_inner() + .map_err(|_| miette::miette!("internal error"))?; - let key_str = registry::debug_render_key(&partition, &k); + let key_str = registry::debug_render_key(&partition, &k); - items.push((key_str, registry::debug_value(&partition, &v))); - } - Ok::<_, StatusCode>(items) - }; - - let start_bound = if let Some(ref s) = start { - std::ops::Bound::Included(s.as_slice()) - } else { - std::ops::Bound::Unbounded - }; - - let end_bound = if let Some(ref e) = end { - std::ops::Bound::Included(e.as_slice()) - } else { - std::ops::Bound::Unbounded - }; - - if req.reverse == Some(true) { - collect( - &mut ks - .range::<&[u8], (std::ops::Bound<&[u8]>, std::ops::Bound<&[u8]>)>(( + items.push((key_str, registry::debug_value(&partition, &v))); + } + Ok::<_, miette::Report>(items) + }; + + let start_bound = if let Some(s) = &start { + std::ops::Bound::Included(s.as_slice()) + } else { + std::ops::Bound::Unbounded + }; + + let end_bound = if let Some(e) = &end { + std::ops::Bound::Included(e.as_slice()) + } else { + std::ops::Bound::Unbounded + }; + + if req.reverse == Some(true) { + collect( + &mut ks + .range::<&[u8], (std::ops::Bound<&[u8]>, std::ops::Bound<&[u8]>)>(( + start_bound, + end_bound, + )) + .rev(), + ) + } else { + collect( + &mut ks.range::<&[u8], (std::ops::Bound<&[u8]>, std::ops::Bound<&[u8]>)>(( start_bound, end_bound, - )) - .rev(), - ) - } else { - collect( - &mut ks.range::<&[u8], (std::ops::Bound<&[u8]>, std::ops::Bound<&[u8]>)>(( - start_bound, - end_bound, - )), - ) - } - }) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)??; + )), + ) + } + }) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok(Json(DebugIterResponse { items })) } @@ -205,18 +215,21 @@ 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 || { - ks.remove(b"dummy_tombstone123")?; - state_clone.db.inner.persist(fjall::PersistMode::Buffer)?; - ks.rotate_memtable_and_wait()?; - ks.major_compact() - }) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + state + .db + .run(move |db| { + let ks = get_keyspace_by_name(db, &req.partition) + .map_err(|_| miette::miette!("bad request"))?; + ks.remove(b"dummy_tombstone123").into_diagnostic()?; + db.inner + .persist(fjall::PersistMode::Buffer) + .into_diagnostic()?; + ks.rotate_memtable_and_wait().into_diagnostic()?; + ks.major_compact().into_diagnostic()?; + Ok::<_, miette::Report>(()) + }) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok(StatusCode::OK) } @@ -225,18 +238,20 @@ pub async fn handle_debug_compact( pub async fn handle_debug_ephemeral_ttl_tick( State(state): State>, ) -> Result { - tokio::task::spawn_blocking(move || -> miette::Result<()> { - #[cfg(feature = "indexer_stream")] - crate::db::ephemeral::ephemeral_ttl_tick(&state.db, &state.ephemeral_ttl)?; - #[cfg(feature = "relay")] - crate::db::ephemeral::relay_events_ttl_tick(&state.db, &state.ephemeral_ttl)?; - #[cfg(feature = "jetstream")] - crate::db::ephemeral::jetstream_events_ttl_tick(&state.db, &state.ephemeral_ttl)?; - Ok(()) - }) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + let ephemeral_ttl = state.ephemeral_ttl.clone(); + state + .db + .run(move |db| { + #[cfg(feature = "indexer_stream")] + crate::db::ephemeral::ephemeral_ttl_tick(db, &ephemeral_ttl)?; + #[cfg(feature = "relay")] + crate::db::ephemeral::relay_events_ttl_tick(db, &ephemeral_ttl)?; + #[cfg(feature = "jetstream")] + crate::db::ephemeral::jetstream_events_ttl_tick(db, &ephemeral_ttl)?; + Ok::<_, miette::Report>(()) + }) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok(StatusCode::OK) } @@ -258,29 +273,27 @@ pub async fn handle_debug_seed_watermark( State(state): State>, Query(req): Query, ) -> Result { - tokio::task::spawn_blocking(move || -> Result<(), StatusCode> { - #[cfg(feature = "indexer_stream")] - state - .db - .cursors - .insert( - crate::db::keys::event_watermark_key(req.ts), - req.event_id.to_be_bytes(), - ) - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - #[cfg(feature = "relay")] - state - .db - .cursors - .insert( - crate::db::keys::relay_event_watermark_key(req.ts), - req.event_id.to_be_bytes(), - ) - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - Ok(()) - }) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)??; + state + .db + .run(move |db| { + #[cfg(feature = "indexer_stream")] + db.cursors + .insert( + crate::db::keys::event_watermark_key(req.ts), + req.event_id.to_be_bytes(), + ) + .into_diagnostic()?; + #[cfg(feature = "relay")] + db.cursors + .insert( + crate::db::keys::relay_event_watermark_key(req.ts), + req.event_id.to_be_bytes(), + ) + .into_diagnostic()?; + Ok::<_, miette::Report>(()) + }) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok(StatusCode::OK) } @@ -297,50 +310,50 @@ pub async fn handle_debug_seed_events( State(state): State>, Query(req): Query, ) -> Result { - tokio::task::spawn_blocking(move || -> Result<(), StatusCode> { - let mut batch = state.db.inner.batch(); - if req.partition == "events" { - #[cfg(feature = "indexer_stream")] - { - for _ in 0..req.count { - let seq = state - .db - .stream - .next_event_id - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - state.db.stream.stage_event( - &mut batch, - crate::db::keys::event_key(seq), - b"dummy", - ); + if req.partition != "events" && req.partition != "relay_events" { + return Err(StatusCode::BAD_REQUEST); + } + + state + .db + .run(move |db| { + let mut batch = db.inner.batch(); + if req.partition == "events" { + #[cfg(feature = "indexer_stream")] + { + for _ in 0..req.count { + let seq = db + .stream + .next_event_id + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + db.stream.stage_event( + &mut batch, + crate::db::keys::event_key(seq), + b"dummy", + ); + } } - } - } else if req.partition == "relay_events" { - #[cfg(feature = "relay")] - { - for _ in 0..req.count { - let seq = state - .db - .relay - .next_seq - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - batch.insert( - &state.db.relay.events, - crate::db::keys::relay_event_key(seq), - b"dummy", - ); + } else if req.partition == "relay_events" { + #[cfg(feature = "relay")] + { + for _ in 0..req.count { + let seq = db + .relay + .next_seq + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + batch.insert( + &db.relay.events, + crate::db::keys::relay_event_key(seq), + b"dummy", + ); + } } } - } else { - return Err(StatusCode::BAD_REQUEST); - } - batch - .commit() - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - Ok(()) - }) - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)??; + batch.commit().into_diagnostic()?; + Ok::<_, miette::Report>(()) + }) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok(StatusCode::OK) } diff --git a/src/api/repos.rs b/src/api/repos.rs index b86c142..8662e8d 100644 --- a/src/api/repos.rs +++ b/src/api/repos.rs @@ -12,7 +12,6 @@ use axum::{ routing::get, }; use jacquard_common::types::did::Did; -use miette::IntoDiagnostic; use serde::Deserialize; pub fn router() -> Router { @@ -54,17 +53,18 @@ pub async fn handle_get_repos( .transpose() .map_err(bad_request)?; - let items = tokio::task::spawn_blocking(move || { - hydrant - .repos - .iter(cursor.as_ref()) - .take(limit) - .collect::>>() - }) - .await - .into_diagnostic() - .flatten() - .map_err(internal)?; + let items = hydrant + .state + .db + .run(move |_db| { + hydrant + .repos + .iter(cursor.as_ref()) + .take(limit) + .collect::>>() + }) + .await + .map_err(internal)?; if prefers_json(&headers) { return Ok(Json(items).into_response()); diff --git a/src/api/xrpc/list_repos.rs b/src/api/xrpc/list_repos.rs index bc84b04..3642cbc 100644 --- a/src/api/xrpc/list_repos.rs +++ b/src/api/xrpc/list_repos.rs @@ -26,51 +26,53 @@ pub async fn handle( .transpose() .map_err(|e| bad_request(nsid, e))?; - let (repos, next_cursor) = tokio::task::spawn_blocking(move || { - let mut repos: Vec> = Vec::new(); - let mut next_cursor: Option> = None; + let (repos, next_cursor) = hydrant + .state + .db + .run(move |_db| { + let mut repos: Vec> = Vec::new(); + let mut next_cursor: Option> = None; - for item in hydrant.repos.iter_states(cursor.as_ref()) { - let (did, state) = item?; + for item in hydrant.repos.iter_states(cursor.as_ref()) { + let (did, state) = item?; - // skip repos that haven't been synced at least once - let Some(commit) = state.root else { - continue; - }; + // skip repos that haven't been synced at least once + let Some(commit) = state.root else { + continue; + }; - let Some(atp_commit) = commit.into_atp_commit(did.clone()) else { - tracing::warn!(did = %did, "repo needs migration"); - continue; - }; + let Some(atp_commit) = commit.into_atp_commit(did.clone()) else { + tracing::warn!(did = %did, "repo needs migration"); + continue; + }; - let Ok(commit_cid) = atp_commit.to_cid() else { - tracing::warn!(did = %did, "failed to compute commit CID"); - continue; - }; + let Ok(commit_cid) = atp_commit.to_cid() else { + tracing::warn!(did = %did, "failed to compute commit CID"); + continue; + }; - let status = repo_status_to_api(state.status); - let repo = Repo { - active: Some(state.active), - did: did.clone(), - head: Cid::Str(CowStr::Owned(commit_cid.to_smolstr())), - rev: atp_commit.rev, - status, - extra_data: None, - }; + let status = repo_status_to_api(state.status); + let repo = Repo { + active: Some(state.active), + did: did.clone(), + head: Cid::Str(CowStr::Owned(commit_cid.to_smolstr())), + rev: atp_commit.rev, + status, + extra_data: None, + }; - if repos.len() < limit { - repos.push(repo); - } else { - next_cursor = repos.last().map(|r| r.did.clone()); - break; + if repos.len() < limit { + repos.push(repo); + } else { + next_cursor = repos.last().map(|r| r.did.clone()); + break; + } } - } - Ok::<_, miette::Report>((repos, next_cursor)) - }) - .await - .map_err(|e| internal_error(nsid, e))? - .map_err(|e| internal_error(nsid, e))?; + Ok::<_, miette::Report>((repos, next_cursor)) + }) + .await + .map_err(|e| internal_error(nsid, e))?; Ok(Json(ListReposOutput { cursor: next_cursor.map(|d| CowStr::Owned(d.as_str().to_smolstr())), diff --git a/src/api/xrpc/request_crawl.rs b/src/api/xrpc/request_crawl.rs index 644f576..767d191 100644 --- a/src/api/xrpc/request_crawl.rs +++ b/src/api/xrpc/request_crawl.rs @@ -1,7 +1,6 @@ use jacquard_api::com_atproto::sync::request_crawl::{ RequestCrawlError, RequestCrawlRequest, RequestCrawlResponse, }; -use miette::IntoDiagnostic; use url::Url; use super::*; @@ -43,14 +42,12 @@ pub async fn handle( // persist the new count before returning so a crash cannot reset the counter // and allow the budget to be replayed. if let Some((day, count)) = to_persist { - let state = hydrant.state.clone(); - tokio::task::spawn_blocking(move || { - crate::db::save_pds_daily_adds(&state.db, day, count) - }) - .await - .into_diagnostic() - .flatten() - .map_err(|e| internal_error(nsid, e))?; + hydrant + .state + .db + .run(move |db| crate::db::save_pds_daily_adds(db, day, count)) + .await + .map_err(|e| internal_error(nsid, e))?; } } diff --git a/src/backfill/manager.rs b/src/backfill/manager.rs index c1f9f82..b16baa4 100644 --- a/src/backfill/manager.rs +++ b/src/backfill/manager.rs @@ -12,8 +12,7 @@ pub fn queue_gone_backfills(state: &Arc) -> Result<()> { debug!("scanning for deactivated/takendown repos to retry..."); let mut transitions = 0usize; - let mut batch = state.db.inner.batch(); - let mut lifecycle_counts = state.db.lifecycle_counts(); + let mut txn = db::Txn::new(&state.db); for guard in state.db.indexer.resync.iter() { let (key, val) = guard.into_inner().into_diagnostic()?; @@ -48,22 +47,22 @@ pub fn queue_gone_backfills(state: &Arc) -> Result<()> { let mut metadata = crate::db::deser_repo_meta(&metadata_bytes)?; // move from resync back into pending - batch.remove(&state.db.indexer.resync, key.clone()); + txn.batch.remove(&state.db.indexer.resync, key.clone()); let old_pending = keys::pending_key(metadata.index_id); - batch.remove(&state.db.indexer.pending, old_pending); + txn.batch.remove(&state.db.indexer.pending, old_pending); metadata.index_id = rand::random::(); - batch.insert( + txn.batch.insert( &state.db.indexer.pending, keys::pending_key(metadata.index_id), key.clone(), ); - batch.insert( + txn.batch.insert( &state.db.repo_metadata, &metadata_key, crate::db::ser_repo_meta(&metadata)?, ); - lifecycle_counts.transition(&mut batch, &did, GaugeState::Pending)?; + txn.transition_lifecycle(&did, GaugeState::Pending)?; transitions += 1; } } @@ -72,9 +71,7 @@ pub fn queue_gone_backfills(state: &Arc) -> Result<()> { return Ok(()); } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - batch.commit().into_diagnostic()?; - state.db.apply_lifecycle_counts(lifecycle_reservation); + txn.commit()?; state.notify_backfill(); @@ -92,8 +89,7 @@ pub fn retry_worker(state: Arc) { let now = chrono::Utc::now().timestamp(); let mut transitions = 0usize; - let mut batch = state.db.inner.batch(); - let mut lifecycle_counts = state.db.lifecycle_counts(); + let mut txn = db::Txn::new(&state.db); for guard in db.indexer.resync.iter() { let (key, value) = match guard.into_inner() { @@ -121,7 +117,7 @@ pub fn retry_worker(state: Arc) { if let Ok(repo_state) = rmp_serde::from_slice::(&state_bytes) { - if let Some(ref pds_str) = repo_state.pds { + if let Some(pds_str) = &repo_state.pds { if let Ok(pds_url) = url::Url::parse(pds_str.as_ref()) { let now_ts = chrono::Utc::now().timestamp(); state.throttler.snapshot(&pds_url).is_throttled(now_ts) @@ -178,18 +174,21 @@ pub fn retry_worker(state: Arc) { continue; } }; - if let Err(e) = - lifecycle_counts.transition(&mut batch, &did, GaugeState::Pending) - { - error!(did = %did, err = %e, "failed to stage lifecycle transition"); + if let Err(e) = txn.transition_lifecycle(&did, GaugeState::Pending) { + error!(did = %did, err = %e, "failed to transition lifecycle"); continue; } // move from resync back into pending - batch.remove(&state.db.indexer.resync, key.clone()); - batch.remove(&state.db.indexer.pending, old_pending); - batch.insert(&state.db.indexer.pending, new_pending, key.clone()); - batch.insert(&state.db.repo_metadata, &metadata_key, serialized_metadata); + txn.batch.remove(&state.db.indexer.resync, key.clone()); + txn.batch.remove(&state.db.indexer.pending, old_pending); + txn.batch + .insert(&state.db.indexer.pending, new_pending, key.clone()); + txn.batch.insert( + &state.db.repo_metadata, + &metadata_key, + serialized_metadata, + ); transitions += 1; } } @@ -207,14 +206,11 @@ pub fn retry_worker(state: Arc) { continue; } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - if let Err(e) = batch.commit() { + if let Err(e) = txn.commit() { error!(err = %e, "failed to commit batch"); - db::check_poisoned(&e); - drop(lifecycle_reservation); + db::check_poisoned_report(&e); continue; } - state.db.apply_lifecycle_counts(lifecycle_reservation); state.notify_backfill(); info!(count = transitions, "queued retries"); } diff --git a/src/backfill/sparse.rs b/src/backfill/sparse.rs index 014e436..9c9f753 100644 --- a/src/backfill/sparse.rs +++ b/src/backfill/sparse.rs @@ -411,7 +411,8 @@ async fn persist_sparse_backfill( ) -> Result)>, BackfillError> { let app_state = app_state.clone(); let did = did.clone(); - tokio::task::spawn_blocking(move || { + let db = app_state.db.clone(); + db.run(move |db| { let filter = app_state.filter.load(); let ephemeral = app_state.ephemeral; let mut count = 0; @@ -421,7 +422,7 @@ async fn persist_sparse_backfill( let mut existing_cids: HashMap<(SmolStr, DbRkey), SmolStr> = HashMap::new(); if !ephemeral { - for guard in app_state.db.indexer.record_prefix(&prefix) { + for guard in db.indexer.record_prefix(&prefix) { let (key, cid_bytes) = guard.into_inner().into_diagnostic()?; let mut remaining = key[prefix.len()..].splitn(2, |b| keys::SEP.eq(b)); let collection_raw = remaining @@ -448,7 +449,7 @@ async fn persist_sparse_backfill( } let mut signal_seen = filter.mode == FilterMode::Full || filter.signals.is_empty(); - let mut txn = DbTxn::new(&app_state.db); + let mut txn = DbTxn::new(db); let mut record_txn = txn.backfill_records(&app_state, &root_commit.rev, &did); for (key, cid) in leaves { @@ -506,8 +507,7 @@ async fn persist_sparse_backfill( let _events = record_txn.finish()?; let metadata_key = keys::repo_metadata_key(&did); - let metadata_bytes = app_state - .db + let metadata_bytes = db .repo_metadata .get(&metadata_key) .into_diagnostic()? @@ -515,7 +515,7 @@ async fn persist_sparse_backfill( let mut metadata = crate::db::deser_repo_meta(&metadata_bytes)?; metadata.tracked = true; txn.batch.insert( - &app_state.db.repo_metadata, + &db.repo_metadata, &metadata_key, crate::db::ser_repo_meta(&metadata)?, ); @@ -523,7 +523,7 @@ async fn persist_sparse_backfill( if !ephemeral { db::replace_record_counts_matching( &mut txn.batch, - &app_state.db, + db, &did, |collection| filter.matches_collection(collection), collection_counts.iter().map(|(col, cnt)| (col.as_str(), *cnt)), @@ -535,7 +535,6 @@ async fn persist_sparse_backfill( Ok::<_, miette::Report>(Some((count, state))) }) .await - .into_diagnostic()? .map_err(BackfillError::from) } diff --git a/src/backfill/worker/process.rs b/src/backfill/worker/process.rs index 165fdb6..4314e1e 100644 --- a/src/backfill/worker/process.rs +++ b/src/backfill/worker/process.rs @@ -20,7 +20,7 @@ use crate::backfill::error::BackfillError; use crate::backfill::sparse::{SparseBackfillResult, process_did_sparse}; use crate::config::BackfillStrategy; use crate::db::types::{DbAction, DbRkey}; -use crate::db::{self, CountDeltas, Db, Txn as DbTxn, keys}; +use crate::db::{self, Txn as DbTxn, keys}; use crate::filter::FilterMode; use crate::ops; use crate::sparse_mst::sparse_probe_collection; @@ -48,7 +48,18 @@ pub(crate) async fn process_did( let db = &app_state.db; let did_key = keys::repo_key(did); - let Some(state_bytes) = Db::get(db.repos.keyspace(), did_key).await? else { + let Some(state_bytes) = db + .run({ + let did_key = did_key.clone(); + move |db| { + db.repos + .get(did_key) + .inspect_err(crate::db::check_poisoned) + .into_diagnostic() + } + }) + .await? + else { return Err(BackfillError::Deleted); }; let mut state: RepoState<'static> = rmp_serde::from_slice::(&state_bytes) @@ -141,36 +152,24 @@ pub(crate) async fn process_did( let app_state_clone = app_state.clone(); let did = did.clone(); let pending_key = pending_key.clone(); - tokio::task::spawn_blocking(move || { - let mut batch = app_state_clone.db.inner.batch(); - let mut count_deltas = CountDeltas::default(); - let mut lifecycle_counts = app_state_clone.db.lifecycle_counts(); - let applied = lifecycle_counts.transition_pending_key( - &mut batch, + app_state_clone.db.run(move |db| { + let mut txn = DbTxn::new(db); + let applied = txn.transition_pending_key( &did, pending_key.as_ref(), GaugeState::Synced, )?; - batch.remove(&app_state_clone.db.indexer.pending, pending_key.clone()); + txn.batch + .remove(&db.indexer.pending, pending_key); if applied { - batch.remove(&app_state_clone.db.repos, &did_key); - batch.remove(&app_state_clone.db.repo_metadata, &metadata_key); - count_deltas.add_repos(-1); + txn.batch.remove(&db.repos, &did_key); + txn.batch + .remove(&db.repo_metadata, &metadata_key); + txn.counts.add_repos(-1); } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - let reservation = app_state_clone - .db - .stage_count_deltas(&mut batch, &count_deltas); - batch.commit().into_diagnostic().inspect(|_| { - app_state_clone - .db - .apply_lifecycle_counts(lifecycle_reservation); - app_state_clone.db.apply_count_deltas(&count_deltas); - drop(reservation); - }) + txn.commit() }) - .await - .into_diagnostic()??; + .await?; return Ok(None); } @@ -226,23 +225,16 @@ pub(crate) async fn process_did( Err(XrpcError::Xrpc(e)) => { if matches!(e, GetRepoError::RepoNotFound(_)) { warn!("repo not found, deleting"); - let mut batch = db.inner.batch(); - let mut lifecycle_counts = db.lifecycle_counts(); - let applied = lifecycle_counts.transition_pending_key( - &mut batch, - did, - pending_key.as_ref(), - GaugeState::Synced, - )?; - batch.remove(&db.indexer.pending, pending_key.clone()); + let mut txn = DbTxn::new(db); + let applied = + txn.transition_pending_key(did, pending_key.as_ref(), GaugeState::Synced)?; + txn.batch.remove(&db.indexer.pending, pending_key.clone()); if applied { - if let Err(e) = crate::ops::delete_repo(&mut batch, db, did, &state) { + if let Err(e) = crate::ops::delete_repo(&mut txn.batch, db, did, &state) { error!(err = %e, "failed to wipe repo during backfill"); } } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - batch.commit().into_diagnostic()?; - db.apply_lifecycle_counts(lifecycle_reservation); + txn.commit()?; // return None so did_task skips sending BackfillFinished (nothing to drain for a deleted repo) return Ok(None); } @@ -268,20 +260,17 @@ pub(crate) async fn process_did( let app_state_clone = app_state.clone(); let did = did.clone(); let pending_key = pending_key.clone(); - tokio::task::spawn_blocking(move || { - let db = &app_state_clone.db; - let mut batch = db.inner.batch(); - let mut lifecycle_counts = db.lifecycle_counts(); - let applied = lifecycle_counts.transition_pending_key( - &mut batch, + app_state_clone.db.run(move |db| { + let mut txn = DbTxn::new(db); + let applied = txn.transition_pending_key( &did, pending_key.as_ref(), GaugeState::Resync(None), )?; - batch.remove(&db.indexer.pending, pending_key.clone()); + txn.batch.remove(&db.indexer.pending, pending_key.clone()); if applied { - Db::update_repo_state( - &mut batch, + crate::db::Db::update_repo_state( + &mut txn.batch, &db.repos, &did, move |state, (key, batch)| { @@ -292,13 +281,10 @@ pub(crate) async fn process_did( }, )?; } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - batch.commit().into_diagnostic()?; - db.apply_lifecycle_counts(lifecycle_reservation); + txn.commit()?; Ok::<_, miette::Report>(()) }) - .await - .into_diagnostic()??; + .await?; return Ok(None); } @@ -544,32 +530,23 @@ pub(crate) async fn process_did( let backfill_pending_key = pending_key.clone(); let app_state = app_state.clone(); let did = did.clone(); - tokio::task::spawn_blocking(move || { - let mut batch = app_state.db.inner.batch(); - let mut count_deltas = CountDeltas::default(); - let mut lifecycle_counts = app_state.db.lifecycle_counts(); - let applied = lifecycle_counts.transition_pending_key( - &mut batch, + app_state.db.run(move |db| { + let mut txn = DbTxn::new(db); + let applied = txn.transition_pending_key( &did, backfill_pending_key.as_ref(), GaugeState::Synced, )?; - batch.remove(&app_state.db.indexer.pending, backfill_pending_key.clone()); + txn.batch + .remove(&db.indexer.pending, backfill_pending_key); if applied { - batch.remove(&app_state.db.repos, &did_key); - batch.remove(&app_state.db.repo_metadata, &metadata_key); - count_deltas.add_repos(-1); + txn.batch.remove(&db.repos, &did_key); + txn.batch.remove(&db.repo_metadata, &metadata_key); + txn.counts.add_repos(-1); } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - let reservation = app_state.db.stage_count_deltas(&mut batch, &count_deltas); - batch.commit().into_diagnostic().inspect(|_| { - app_state.db.apply_lifecycle_counts(lifecycle_reservation); - app_state.db.apply_count_deltas(&count_deltas); - drop(reservation); - }) + txn.commit() }) - .await - .into_diagnostic()??; + .await?; return Ok(None); }; diff --git a/src/backfill/worker/task.rs b/src/backfill/worker/task.rs index c95e395..7db0926 100644 --- a/src/backfill/worker/task.rs +++ b/src/backfill/worker/task.rs @@ -8,7 +8,7 @@ use tracing::{debug, error, warn}; use crate::backfill::client::ThrottledHttpClient; use crate::backfill::error::BackfillError; use crate::config::BackfillStrategy; -use crate::db::{Db, keys}; +use crate::db::{Txn as DbTxn, keys}; use crate::ingest::indexer::{IndexerMessage, IndexerTx}; use crate::state::AppState; use crate::types::{GaugeState, RepoState, RepoStatus, ResyncErrorKind, ResyncState}; @@ -38,48 +38,35 @@ pub(crate) async fn did_task( .await { Ok(Some(_repo_state)) => { - let applied = tokio::task::spawn_blocking({ - let state = state.clone(); + let applied = state.db.run({ let did = did.clone(); let pending_key = pending_key.clone(); - move || { - let db = &state.db; + move |db| { let did_key = keys::repo_key(&did); - let mut batch = db.inner.batch(); - let mut lifecycle_counts = db.lifecycle_counts(); - let applied = lifecycle_counts.transition_pending_key( - &mut batch, - &did, - pending_key.as_ref(), - GaugeState::Synced, - )?; - batch.remove(&db.indexer.pending, pending_key.clone()); + let mut txn = DbTxn::new(db); + let applied = + txn.transition_pending_key(&did, pending_key.as_ref(), GaugeState::Synced)?; + txn.batch.remove(&db.indexer.pending, pending_key.clone()); if applied { - batch.remove(&db.indexer.resync, &did_key); + txn.batch.remove(&db.indexer.resync, &did_key); } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - batch.commit().into_diagnostic()?; - db.apply_lifecycle_counts(lifecycle_reservation); + txn.commit()?; Ok::<_, miette::Report>(applied) } }) - .await - .into_diagnostic()??; + .await?; if !applied { return Ok(()); } let state = state.clone(); - tokio::task::spawn_blocking(move || { - state - .db - .inner + state.db.run(move |db| { + db.inner .persist(fjall::PersistMode::Buffer) .into_diagnostic() }) - .await - .into_diagnostic()??; + .await?; if let Err(e) = buffer_tx .send(IndexerMessage::BackfillFinished(did.clone())) @@ -92,29 +79,18 @@ pub(crate) async fn did_task( Ok(None) => Ok(()), Err(BackfillError::Deleted) => { warn!("orphaned pending entry, cleaning up"); - tokio::task::spawn_blocking({ - let state = state.clone(); + state.db.run({ let did = did.clone(); let pending_key = pending_key.clone(); - move || { - let db = &state.db; - let mut batch = db.inner.batch(); - let mut lifecycle_counts = db.lifecycle_counts(); - lifecycle_counts.transition_pending_key( - &mut batch, - &did, - pending_key.as_ref(), - GaugeState::Synced, - )?; - batch.remove(&db.indexer.pending, pending_key); - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - batch.commit().into_diagnostic()?; - db.apply_lifecycle_counts(lifecycle_reservation); + move |db| { + let mut txn = DbTxn::new(db); + txn.transition_pending_key(&did, pending_key.as_ref(), GaugeState::Synced)?; + txn.batch.remove(&db.indexer.pending, pending_key); + txn.commit()?; Ok::<_, miette::Report>(()) } }) - .await - .into_diagnostic()??; + .await?; Ok(()) } Err(e) => { @@ -143,10 +119,19 @@ pub(crate) async fn did_task( let did_key = keys::repo_key(did); // 1. get current retry count - let existing_state = Db::get(db.indexer.resync.keyspace(), &did_key).await.and_then(|b| { - b.map(|b| rmp_serde::from_slice::(&b).into_diagnostic()) - .transpose() - })?; + let did_key_clone = did_key.clone(); + let existing_state = db + .run(move |db| { + db.indexer + .resync + .get(&did_key_clone) + .into_diagnostic() + .and_then(|b| { + b.map(|b| rmp_serde::from_slice::(&b).into_diagnostic()) + .transpose() + }) + }) + .await?; let mut retry_count = match existing_state { Some(ResyncState::Error { retry_count, .. }) => retry_count, @@ -169,19 +154,18 @@ pub(crate) async fn did_task( }; let error_string = e.to_string(); - tokio::task::spawn_blocking({ - let state = state.clone(); + state.db.run({ let did_key = did_key.into_static(); let did = did.clone(); let pending_key = pending_key.clone(); - move || { + move |db| { // 3. save to resync let serialized_resync_state = rmp_serde::to_vec(&resync_state).into_diagnostic()?; // 4. and update the main repo state let serialized_repo_state = if let Some(state_bytes) = - state.db.repos.get(&did_key).into_diagnostic()? + db.repos.get(&did_key).into_diagnostic()? { let mut state: RepoState = rmp_serde::from_slice(&state_bytes).into_diagnostic()?; @@ -191,30 +175,25 @@ pub(crate) async fn did_task( } else { None }; - - let mut batch = state.db.inner.batch(); - let mut lifecycle_counts = state.db.lifecycle_counts(); - let applied = lifecycle_counts.transition_pending_key( - &mut batch, + let mut txn = DbTxn::new(db); + let applied = txn.transition_pending_key( &did, pending_key.as_ref(), GaugeState::Resync(Some(error_kind)), )?; - batch.remove(&state.db.indexer.pending, pending_key.clone()); + txn.batch.remove(&db.indexer.pending, pending_key.clone()); if applied { - batch.insert(&state.db.indexer.resync, &did_key, serialized_resync_state); + txn.batch + .insert(&db.indexer.resync, &did_key, serialized_resync_state); if let Some(state_bytes) = serialized_repo_state { - batch.insert(&state.db.repos, &did_key, state_bytes); + txn.batch.insert(&db.repos, &did_key, state_bytes); } } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - batch.commit().into_diagnostic()?; - state.db.apply_lifecycle_counts(lifecycle_reservation); + txn.commit()?; Ok::<_, miette::Report>(()) } }) - .await - .into_diagnostic()??; + .await?; Err(e) } diff --git a/src/backlinks/mod.rs b/src/backlinks/mod.rs index 7b188ab..68bd365 100644 --- a/src/backlinks/mod.rs +++ b/src/backlinks/mod.rs @@ -117,8 +117,7 @@ impl BacklinksFetch { /// execute and return a page of backlink entries. pub async fn run(self) -> Result { - tokio::task::spawn_blocking(move || { - let db = &self.state.db; + self.state.db.run(move |db| { let scan_prefix = store::reverse_scan_prefix( &self.subject, self.collection.as_deref(), @@ -126,7 +125,7 @@ impl BacklinksFetch { ); let iter: Box> = if !self.reverse { - if let Some(ref cursor_bytes) = self.cursor { + if let Some(cursor_bytes) = &self.cursor { Box::new(db.backlinks.range::<&[u8], _>(( Bound::Excluded(cursor_bytes.as_slice()), Bound::Unbounded, @@ -137,7 +136,7 @@ impl BacklinksFetch { } else { // for reverse scans, bound the end at the cursor (exclusive) or the // prefix end (increment last byte to get an exclusive upper bound) - let end: Vec = if let Some(ref cursor_bytes) = self.cursor { + let end: Vec = if let Some(cursor_bytes) = &self.cursor { cursor_bytes.clone() } else { let mut end = scan_prefix.clone(); @@ -192,7 +191,6 @@ impl BacklinksFetch { }) }) .await - .into_diagnostic()? } } @@ -238,8 +236,7 @@ impl BacklinksCount { /// execute and return the total count of matching entries. pub async fn run(self) -> Result { - tokio::task::spawn_blocking(move || { - let db = &self.state.db; + self.state.db.run(move |db| { let scan_prefix = store::reverse_scan_prefix( &self.subject, self.collection.as_deref(), @@ -264,6 +261,5 @@ impl BacklinksCount { } }) .await - .into_diagnostic()? } } diff --git a/src/control/crawler.rs b/src/control/crawler.rs index 170b7a2..c5fa29d 100644 --- a/src/control/crawler.rs +++ b/src/control/crawler.rs @@ -147,23 +147,23 @@ impl CrawlerHandle { /// delete all cursor entries associated with the given URL. pub async fn reset_cursor(&self, url: &str) -> Result<()> { - let state = self.state.clone(); let point_keys = [keys::crawler_cursor_key(url)]; let by_collection_prefix = keys::by_collection_cursor_prefix(url); - tokio::task::spawn_blocking(move || { - let mut batch = state.db.inner.batch(); - for k in point_keys { - batch.remove(&state.db.cursors, k); - } - for entry in state.db.cursors.prefix(&by_collection_prefix) { - let k = entry.key().into_diagnostic()?; - batch.remove(&state.db.cursors, k); - } - batch.commit().into_diagnostic()?; - state.db.persist() - }) - .await - .into_diagnostic()??; + self.state + .db + .run(move |db| { + let mut batch = db.inner.batch(); + for k in point_keys { + batch.remove(&db.cursors, k); + } + for entry in db.cursors.prefix(&by_collection_prefix) { + let k = entry.key().into_diagnostic()?; + batch.remove(&db.cursors, k); + } + batch.commit().into_diagnostic()?; + db.persist() + }) + .await?; Ok(()) } @@ -196,15 +196,15 @@ impl CrawlerHandle { miette::bail!("crawler not yet started: call Hydrant::run() first"); }; - let state = self.state.clone(); let key = keys::crawler_source_key(source.url.as_str()); let val = rmp_serde::to_vec(&source.mode).into_diagnostic()?; - tokio::task::spawn_blocking(move || { - state.db.crawler.insert(key, val).into_diagnostic()?; - state.db.persist() - }) - .await - .into_diagnostic()??; + self.state + .db + .run(move |db| { + db.crawler.insert(key, val).into_diagnostic()?; + db.persist() + }) + .await?; let enabled_rx = self.state.crawler_enabled.subscribe(); let handle = spawn_crawler_producer( @@ -250,14 +250,14 @@ impl CrawlerHandle { // remove from DB if it was a persisted source if self.persisted.remove_async(url).await.is_some() { - let state = self.state.clone(); let key = keys::crawler_source_key(url.as_str()); - tokio::task::spawn_blocking(move || { - state.db.crawler.remove(key).into_diagnostic()?; - state.db.persist() - }) - .await - .into_diagnostic()??; + self.state + .db + .run(move |db| { + db.crawler.remove(key).into_diagnostic()?; + db.persist() + }) + .await?; } Ok(true) diff --git a/src/control/db.rs b/src/control/db.rs index f783291..6fedf9e 100644 --- a/src/control/db.rs +++ b/src/control/db.rs @@ -35,9 +35,7 @@ impl DbControl { state .with_ingestion_paused(async || { let train = |name: &'static str| { - let state = state.clone(); - tokio::task::spawn_blocking(move || state.db.train_dict(name)) - .map(|res: Result<_, _>| res.into_diagnostic().flatten()) + state.db.run(move |db| db.train_dict(name)) }; futures::future::try_join_all( crate::db::registry::trainable() diff --git a/src/control/filter.rs b/src/control/filter.rs index 32cd6c0..329474e 100644 --- a/src/control/filter.rs +++ b/src/control/filter.rs @@ -44,19 +44,19 @@ pub struct FilterControl(pub(super) Arc); impl FilterControl { /// return the current filter configuration from the database. pub async fn get(&self) -> Result { - let filter_ks = self.0.db.filter.clone(); - tokio::task::spawn_blocking(move || { - let hot = db_filter::load(&filter_ks)?; - let excludes = db_filter::read_set(&filter_ks, db_filter::EXCLUDE_PREFIX)?; - Ok(FilterSnapshot { - mode: hot.mode, - signals: hot.signals.iter().map(|s| s.to_string()).collect(), - collections: hot.collections.iter().map(|s| s.to_string()).collect(), - excludes, + self.0 + .db + .run(move |db| { + let hot = db_filter::load(&db.filter)?; + let excludes = db_filter::read_set(&db.filter, db_filter::EXCLUDE_PREFIX)?; + Ok(FilterSnapshot { + mode: hot.mode, + signals: hot.signals.iter().map(|s| s.to_string()).collect(), + collections: hot.collections.iter().map(|s| s.to_string()).collect(), + excludes, + }) }) - }) - .await - .into_diagnostic()? + .await } /// set the indexing mode. see [`FilterControl`] for mode semantics. @@ -268,8 +268,6 @@ impl FilterPatch { /// commit the patch atomically to the database and update the in-memory filter. /// returns the updated [`FilterSnapshot`]. pub async fn apply(self) -> Result { - let filter_ks = self.state.db.filter.clone(); - let inner = self.state.db.inner.clone(); let filter_handle = self.state.filter.clone(); let state = self.state.clone(); let mode = self.mode; @@ -277,28 +275,27 @@ impl FilterPatch { let collections = self.collections; let excludes = self.excludes; - let new_filter = tokio::task::spawn_blocking(move || { - let mut batch = inner.batch(); - db_filter::apply_patch(&mut batch, &filter_ks, mode, signals, collections, excludes)?; - batch.commit().into_diagnostic()?; - state.db.persist()?; - db_filter::load(&filter_ks) - }) - .await - .into_diagnostic()? - .map_err(|e| { - error!(err = %e, "failed to apply filter patch"); - e - })?; - - let exclude_list = { - let filter_ks = self.state.db.filter.clone(); - tokio::task::spawn_blocking(move || { - db_filter::read_set(&filter_ks, db_filter::EXCLUDE_PREFIX) + let new_filter = state + .db + .run(move |db| { + let mut batch = db.inner.batch(); + db_filter::apply_patch(&mut batch, &db.filter, mode, signals, collections, excludes)?; + batch.commit().into_diagnostic()?; + db.persist()?; + db_filter::load(&db.filter) }) .await - .into_diagnostic()?? - }; + .map_err(|e| { + error!(err = %e, "failed to apply filter patch"); + e + })?; + + let exclude_list = state + .db + .run(move |db| { + db_filter::read_set(&db.filter, db_filter::EXCLUDE_PREFIX) + }) + .await?; let snapshot = FilterSnapshot { mode: new_filter.mode, diff --git a/src/control/firehose.rs b/src/control/firehose.rs index 97adfcc..7688773 100644 --- a/src/control/firehose.rs +++ b/src/control/firehose.rs @@ -307,20 +307,21 @@ impl FirehoseHandle { // persist to db first let key = keys::firehose_source_key(url.as_str()); - tokio::task::spawn_blocking({ - let state = self.state.clone(); - move || { - let mut batch = state.db.inner.batch(); - let value = rmp_serde::to_vec(&db::FirehoseSourceMeta { is_pds }).map_err(|e| { - miette::miette!("failed to serialize firehose source meta: {e}") - })?; - batch.insert(&state.db.crawler, key, &value); - batch.commit().into_diagnostic()?; - state.db.persist() - } - }) - .await - .into_diagnostic()??; + self.state + .db + .run({ + let is_pds = is_pds; + move |db| { + let mut batch = db.inner.batch(); + let value = rmp_serde::to_vec(&db::FirehoseSourceMeta { is_pds }).map_err(|e| { + miette::miette!("failed to serialize firehose source meta: {e}") + })?; + batch.insert(&db.crawler, key, &value); + batch.commit().into_diagnostic()?; + db.persist() + } + }) + .await?; let _ = self.known_sources.upsert_async(url.clone(), is_pds).await; @@ -359,30 +360,30 @@ impl FirehoseHandle { return Ok(0); } - tokio::task::spawn_blocking({ - let state = self.state.clone(); - let sources = sources.clone(); - move || { - let mut batch = state.db.inner.batch(); - for source in &sources { - let value = rmp_serde::to_vec(&db::FirehoseSourceMeta { - is_pds: source.is_pds, - }) - .map_err(|e| { - miette::miette!("failed to serialize firehose source meta: {e}") - })?; - batch.insert( - &state.db.crawler, - keys::firehose_source_key(source.url.as_str()), - &value, - ); + self.state + .db + .run({ + let sources = sources.clone(); + move |db| { + let mut batch = db.inner.batch(); + for source in &sources { + let value = rmp_serde::to_vec(&db::FirehoseSourceMeta { + is_pds: source.is_pds, + }) + .map_err(|e| { + miette::miette!("failed to serialize firehose source meta: {e}") + })?; + batch.insert( + &db.crawler, + keys::firehose_source_key(source.url.as_str()), + &value, + ); + } + batch.commit().into_diagnostic()?; + db.persist() } - batch.commit().into_diagnostic()?; - state.db.persist() - } - }) - .await - .into_diagnostic()??; + }) + .await?; let mut added = 0usize; for source in sources { @@ -408,19 +409,15 @@ impl FirehoseHandle { pub async fn remove_source(&self, url: &Url) -> Result { if self.known_sources.contains_async(url).await { let url_str = url.to_string(); - tokio::task::spawn_blocking({ - let state = self.state.clone(); - move || { - state - .db - .crawler + self.state + .db + .run(move |db| { + db.crawler .remove(keys::firehose_source_key(&url_str)) .into_diagnostic()?; - state.db.persist() - } - }) - .await - .into_diagnostic()??; + db.persist() + }) + .await?; self.known_sources.remove_async(url).await; } @@ -446,15 +443,13 @@ impl FirehoseHandle { pub async fn reset_cursor(&self, url: &str) -> Result<()> { let url = Url::parse(url).into_diagnostic()?; let key = keys::firehose_cursor_key_from_url(&url); - tokio::task::spawn_blocking({ - let state = self.state.clone(); - move || { - state.db.cursors.remove(key).into_diagnostic()?; - state.db.persist() - } - }) - .await - .into_diagnostic()??; + self.state + .db + .run(move |db| { + db.cursors.remove(key).into_diagnostic()?; + db.persist() + }) + .await?; self.state.firehose_cursors.remove_async(&url).await; diff --git a/src/control/hosts.rs b/src/control/hosts.rs index 39c878b..a900099 100644 --- a/src/control/hosts.rs +++ b/src/control/hosts.rs @@ -56,20 +56,21 @@ impl Hydrant { let state = self.state.clone(); let hostname = hostname.to_smolstr(); - tokio::task::spawn_blocking(move || { + let state_closure = state.clone(); + state.db.run(move |db| { let key = keys::firehose_cursor_key(&hostname); let mut seq = 0; - if let Some(cursor_bytes) = state.db.cursors.get(&key).into_diagnostic()? { + if let Some(cursor_bytes) = db.cursors.get(&key).into_diagnostic()? { seq = i64::from_be_bytes(cursor_bytes.as_ref().try_into().into_diagnostic()?); } else { // if it has no cursor, check if it's explicitly tracked in hosts map // or firehose tasks (recently added via API but no messages yet) - let meta = state.pds_meta.load(); + let meta = state_closure.pds_meta.load(); if !meta.hosts.contains_key(hostname.as_str()) { // we should also allow it if it's an active firehose ingestor let mut found_in_cursors = false; - state.firehose_cursors.iter_sync(|u, _| { + state_closure.firehose_cursors.iter_sync(|u, _| { if u.host_str() == Some(hostname.as_str()) { found_in_cursors = true; } @@ -82,11 +83,9 @@ impl Hydrant { } } - let account_count = state - .db + let account_count = db .get_count_sync(&keys::pds_account_count_key(&hostname)); - let status = state.pds_meta.load().status(&hostname); - + let status = state_closure.pds_meta.load().status(&hostname); Ok(Some(Host { name: hostname, seq, @@ -95,7 +94,6 @@ impl Hydrant { })) }) .await - .into_diagnostic()? } /// enumerates all hosts hydrant is consuming from. @@ -108,8 +106,9 @@ impl Hydrant { ) -> Result<(Vec, Option)> { let state = self.state.clone(); let cursor = cursor.map(str::to_string); + let state_closure = state.clone(); - tokio::task::spawn_blocking(move || { + state.db.run(move |db| { let start_bound = match &cursor { Some(after) => std::ops::Bound::Included(keys::firehose_cursor_key(after)), None => std::ops::Bound::Included(keys::FIREHOSE_CURSOR_PREFIX.to_vec()), @@ -123,7 +122,7 @@ impl Hydrant { let end_bound = std::ops::Bound::Excluded(prefix_end); let mut db_hosts = Vec::new(); - for item in state.db.cursors.range((start_bound, end_bound)) { + for item in db.cursors.range((start_bound, end_bound)) { let (k, _) = item.into_inner().into_diagnostic()?; let hostname = std::str::from_utf8(&k[keys::FIREHOSE_CURSOR_PREFIX.len()..]) .into_diagnostic() @@ -143,7 +142,7 @@ impl Hydrant { let mut meta_hosts = Vec::new(); { - let meta = state.pds_meta.load(); + let meta = state_closure.pds_meta.load(); for hostname in meta.hosts.keys() { if let Some(after) = &cursor { if hostname.as_str() <= after.as_str() { @@ -168,8 +167,7 @@ impl Hydrant { let mut hosts: Vec = Vec::with_capacity(selected.len().min(limit)); for hostname in selected.iter().take(limit) { - let seq = state - .db + let seq = db .cursors .get(keys::firehose_cursor_key(hostname)) .into_diagnostic()? @@ -182,10 +180,9 @@ impl Hydrant { }) .transpose()? .unwrap_or(0); - let account_count = state - .db + let account_count = db .get_count_sync(&keys::pds_account_count_key(hostname)); - let status = state.pds_meta.load().status(hostname); + let status = state_closure.pds_meta.load().status(hostname); hosts.push(Host { name: hostname.clone(), seq, @@ -203,10 +200,8 @@ impl Hydrant { Ok((hosts, next_cursor)) }) .await - .into_diagnostic()? } } - #[cfg(test)] mod host_listing_tests { use super::*; @@ -228,34 +223,33 @@ mod host_listing_tests { { let state = hydrant.state.clone(); - tokio::task::spawn_blocking(move || -> Result<()> { - let mut batch = state.db.inner.batch(); + state.db.run(move |db| -> Result<()> { + let mut batch = db.inner.batch(); crate::db::pds_meta::set_status( &mut batch, - &state.db.filter, + &db.filter, "offline.example", HostStatus::Offline, )?; crate::db::pds_meta::set_status( &mut batch, - &state.db.filter, + &db.filter, "active.example", HostStatus::Active, )?; set_ks_count( &mut batch, - &state.db, + db, &keys::pds_account_count_key("offline.example"), 7, ); set_ks_count( &mut batch, - &state.db, + db, &keys::pds_account_count_key("active.example"), 42, ); - state - .db + db .cursors .insert( keys::firehose_cursor_key("active.example"), @@ -263,10 +257,9 @@ mod host_listing_tests { ) .into_diagnostic()?; batch.commit().into_diagnostic()?; - state.db.persist() + db.persist() }) - .await - .into_diagnostic()??; + .await?; crate::pds_meta::PdsMeta::update_host( &hydrant.state.pds_meta, @@ -313,26 +306,22 @@ mod host_listing_tests { // Seed some in DB: host2, host4, host6 { let state = hydrant.state.clone(); - tokio::task::spawn_blocking(move || -> Result<()> { - state - .db + state.db.run(move |db| -> Result<()> { + db .cursors .insert(keys::firehose_cursor_key("host2"), 2_i64.to_be_bytes()) .into_diagnostic()?; - state - .db + db .cursors .insert(keys::firehose_cursor_key("host4"), 4_i64.to_be_bytes()) .into_diagnostic()?; - state - .db + db .cursors .insert(keys::firehose_cursor_key("host6"), 6_i64.to_be_bytes()) .into_diagnostic()?; - state.db.persist() + db.persist() }) - .await - .into_diagnostic()??; + .await?; } // Seed some in memory: host1, host3, host5 diff --git a/src/control/hydrant.rs b/src/control/hydrant.rs index 2206b04..8ee9b10 100644 --- a/src/control/hydrant.rs +++ b/src/control/hydrant.rs @@ -88,8 +88,6 @@ impl Hydrant { || config.filter_collections.is_some() || config.filter_excludes.is_some() { - let filter_ks = state.db.filter.clone(); - let inner = state.db.inner.clone(); let mode = config.full_network.then_some(FilterMode::Full); let signals = config .filter_signals @@ -104,28 +102,27 @@ impl Hydrant { .clone() .map(crate::patch::SetUpdate::Set); - tokio::task::spawn_blocking(move || { - let mut batch = inner.batch(); - db_filter::apply_patch( - &mut batch, - &filter_ks, - mode, - signals, - collections, - excludes, - )?; - batch.commit().into_diagnostic() - }) - .await - .into_diagnostic()??; + state + .db + .run(move |db| { + let mut batch = db.inner.batch(); + db_filter::apply_patch( + &mut batch, + &db.filter, + mode, + signals, + collections, + excludes, + )?; + batch.commit().into_diagnostic() + }) + .await?; // 3. reload the live filter into the hot-path arc-swap - let new_filter = tokio::task::spawn_blocking({ - let filter_ks = state.db.filter.clone(); - move || db_filter::load(&filter_ks) - }) - .await - .into_diagnostic()??; + let new_filter = state + .db + .run(move |db| db_filter::load(&db.filter)) + .await?; state.filter.store(Arc::new(new_filter)); } diff --git a/src/control/hydrant/run.rs b/src/control/hydrant/run.rs index ae73849..fdfb9b0 100644 --- a/src/control/hydrant/run.rs +++ b/src/control/hydrant/run.rs @@ -90,12 +90,11 @@ impl Hydrant { // 6. re-queue any repos that lost their backfill state, then start the retry worker #[cfg(feature = "indexer")] { - if let Err(e) = tokio::task::spawn_blocking({ + if let Err(e) = state.db.run({ let state = state.clone(); - move || crate::backfill::manager::queue_gone_backfills(&state) + move |_db| crate::backfill::manager::queue_gone_backfills(&state) }) .await - .into_diagnostic()? { error!(err = %e, "failed to queue gone backfills"); db::check_poisoned_report(&e); @@ -262,12 +261,10 @@ impl Hydrant { } // add persisted hosts - let persisted_sources = tokio::task::spawn_blocking({ - let state = state.clone(); - move || load_persisted_firehose_sources(&state.db) + let persisted_sources = state.db.run({ + move |db| load_persisted_firehose_sources(db) }) - .await - .into_diagnostic()??; + .await?; for source in &persisted_sources { let _ = firehose .known_sources @@ -442,12 +439,10 @@ impl Hydrant { let _ = crawler.tasks.insert_async(source.url.clone(), handle).await; } - let persisted_sources = tokio::task::spawn_blocking({ - let state = state.clone(); - move || load_persisted_crawler_sources(&state.db) + let persisted_sources = state.db.run({ + move |db| load_persisted_crawler_sources(db) }) - .await - .into_diagnostic()??; + .await?; for source in &persisted_sources { let _ = crawler.persisted.insert_async(source.url.clone()).await; diff --git a/src/control/pds.rs b/src/control/pds.rs index d9802e4..ca80bbc 100644 --- a/src/control/pds.rs +++ b/src/control/pds.rs @@ -51,15 +51,15 @@ impl PdsControl { F: FnOnce(&mut fjall::OwnedWriteBatch, &fjall::Keyspace) + Send + 'static, G: FnOnce(&mut PdsMeta), { - let state = self.0.clone(); - tokio::task::spawn_blocking(move || -> Result<()> { - let mut batch = state.db.inner.batch(); - db_op(&mut batch, &state.db.filter); - batch.commit().into_diagnostic()?; - state.db.persist() - }) - .await - .into_diagnostic()??; + self.0 + .db + .run(move |db| { + let mut batch = db.inner.batch(); + db_op(&mut batch, &db.filter); + batch.commit().into_diagnostic()?; + db.persist() + }) + .await?; let mut snapshot = (**self.0.pds_meta.load()).clone(); mem_op(&mut snapshot); diff --git a/src/control/repos/indexer.rs b/src/control/repos/indexer.rs index 52555e4..2a307ae 100644 --- a/src/control/repos/indexer.rs +++ b/src/control/repos/indexer.rs @@ -2,7 +2,6 @@ use futures::{FutureExt, TryFutureExt}; use rand::Rng; use super::*; -use crate::db::LifecycleCountBatch; impl ReposControl { /// iterates through pending repositories, returning their state. @@ -101,12 +100,7 @@ impl ReposControl { .filter_map(|b| b.transpose()) } - pub(crate) fn _resync( - db: &Db, - did: &Did<'_>, - batch: &mut fjall::OwnedWriteBatch, - lifecycle_counts: &mut LifecycleCountBatch<'_>, - ) -> Result { + pub(crate) fn _resync(db: &Db, did: &Did<'_>, txn: &mut crate::db::Txn<'_>) -> Result { let did_key = keys::repo_key(did); let metadata_key = keys::repo_metadata_key(did); @@ -135,20 +129,20 @@ impl ReposControl { metadata.tracked = true; // insert into pending with new index_id let old_pending = keys::pending_key(metadata.index_id); - batch.remove(&db.indexer.pending, old_pending); + txn.batch.remove(&db.indexer.pending, old_pending); metadata.index_id = rand::Rng::next_u64(&mut rand::rng()); - batch.insert( + txn.batch.insert( &db.indexer.pending, keys::pending_key(metadata.index_id), &did_key, ); - batch.remove(&db.indexer.resync, &did_key); - batch.insert( + txn.batch.remove(&db.indexer.resync, &did_key); + txn.batch.insert( &db.repo_metadata, &metadata_key, crate::db::ser_repo_meta(&metadata)?, ); - lifecycle_counts.transition(batch, did, GaugeState::Pending)?; + txn.transition_lifecycle(did, GaugeState::Pending)?; return Ok(true); } } @@ -169,28 +163,23 @@ impl ReposControl { dids: impl IntoIterator>, ) -> Result>> { let dids: Vec> = dids.into_iter().map(|d| d.into_static()).collect(); - let state = self.0.clone(); - let queued = tokio::task::spawn_blocking(move || { - let db = &state.db; - let mut batch = db.inner.batch(); + let queued = self.0.db.run(move |db| { + let mut txn = crate::db::Txn::new(db); let mut queued: Vec> = Vec::new(); - let mut lifecycle_counts = db.lifecycle_counts(); for did in dids { - if Self::_resync(db, &did, &mut batch, &mut lifecycle_counts)? { + if Self::_resync(db, &did, &mut txn)? { queued.push(did); } } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - batch.commit().into_diagnostic()?; - db.apply_lifecycle_counts(lifecycle_reservation); - state.db.persist()?; - Ok::<_, miette::Report>(queued) + txn.commit()?; + db.persist()?; + Ok(queued) }) - .await - .into_diagnostic()??; + .await?; + if !queued.is_empty() { self.0.notify_backfill(); } @@ -208,14 +197,10 @@ impl ReposControl { dids: impl IntoIterator>, ) -> Result>> { let dids: Vec> = dids.into_iter().map(|d| d.into_static()).collect(); - let state = self.0.clone(); - let queued = tokio::task::spawn_blocking(move || { - let db = &state.db; - let mut batch = db.inner.batch(); + let queued = self.0.db.run(move |db| { + let mut txn = crate::db::Txn::new(db); let mut queued: Vec> = Vec::new(); - let mut count_deltas = crate::db::CountDeltas::default(); - let mut lifecycle_counts = db.lifecycle_counts(); for did in dids { let did_key = keys::repo_key(&did); @@ -227,42 +212,35 @@ impl ReposControl { .transpose()?; if let Some(metadata) = existing_metadata { - if !metadata.tracked - && Self::_resync(db, &did, &mut batch, &mut lifecycle_counts)? - { + if !metadata.tracked && Self::_resync(db, &did, &mut txn)? { queued.push(did); } } else { let repo_state = RepoState::backfilling(); let metadata = RepoMetadata::backfilling(rand::random()); - batch.insert(&db.repos, &did_key, crate::db::ser_repo_state(&repo_state)?); - batch.insert( + txn.batch + .insert(&db.repos, &did_key, crate::db::ser_repo_state(&repo_state)?); + txn.batch.insert( &db.repo_metadata, &metadata_key, crate::db::ser_repo_meta(&metadata)?, ); - batch.insert( + txn.batch.insert( &db.indexer.pending, keys::pending_key(metadata.index_id), &did_key, ); - count_deltas.add_repos(1); - lifecycle_counts.transition(&mut batch, &did, GaugeState::Pending)?; + txn.counts.add_repos(1); + txn.transition_lifecycle(&did, GaugeState::Pending)?; queued.push(did); } } - - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - let reservation = db.stage_count_deltas(&mut batch, &count_deltas); - batch.commit().into_diagnostic()?; - db.apply_lifecycle_counts(lifecycle_reservation); - db.apply_count_deltas(&count_deltas); - drop(reservation); - state.db.persist()?; - Ok::<_, miette::Report>(queued) + txn.commit()?; + db.persist()?; + Ok(queued) }) - .await - .into_diagnostic()??; + .await?; + self.0.notify_backfill(); Ok(queued) } @@ -275,13 +253,10 @@ impl ReposControl { dids: impl IntoIterator>, ) -> Result>> { let dids: Vec> = dids.into_iter().map(|d| d.into_static()).collect(); - let state = self.0.clone(); - let untracked = tokio::task::spawn_blocking(move || { - let db = &state.db; - let mut batch = db.inner.batch(); + let untracked = self.0.db.run(move |db| { + let mut txn = crate::db::Txn::new(db); let mut untracked: Vec> = Vec::new(); - let mut lifecycle_counts = db.lifecycle_counts(); for did in dids { let did_key = keys::repo_key(&did); @@ -303,27 +278,26 @@ impl ReposControl { && metadata.tracked { metadata.tracked = false; - batch.insert( + txn.batch.insert( &db.repo_metadata, &metadata_key, crate::db::ser_repo_meta(&metadata)?, ); - batch.remove(&db.indexer.pending, keys::pending_key(metadata.index_id)); - batch.remove(&db.indexer.resync, &did_key); - lifecycle_counts.transition(&mut batch, &did, GaugeState::Synced)?; + txn.batch + .remove(&db.indexer.pending, keys::pending_key(metadata.index_id)); + txn.batch.remove(&db.indexer.resync, &did_key); + txn.transition_lifecycle(&did, GaugeState::Synced)?; untracked.push(did); } } } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - batch.commit().into_diagnostic()?; - db.apply_lifecycle_counts(lifecycle_reservation); - state.db.persist()?; - Ok::<_, miette::Report>(untracked) + txn.commit()?; + db.persist()?; + Ok(untracked) }) - .await - .into_diagnostic()??; + .await?; + Ok(untracked) } } @@ -339,18 +313,17 @@ impl<'i> RepoHandle<'i> { let db_key = keys::record_key(&did, collection, &DbRkey::new(rkey)); let collection = collection.to_smolstr(); - let state = self.state.clone(); - tokio::task::spawn_blocking(move || { + self.state.db.run(move |db| { use miette::WrapErr; - let cid_bytes = state.db.indexer.record(db_key).into_diagnostic()?; + let cid_bytes = db.indexer.record(db_key).into_diagnostic()?; let Some(cid_bytes) = cid_bytes else { return Ok(None); }; // lookup block using col|cid key let block_key = keys::block_key(&collection, &cid_bytes); - let Some(block_bytes) = state.db.indexer.block(block_key).into_diagnostic()? else { + let Some(block_bytes) = db.indexer.block(block_key).into_diagnostic()? else { miette::bail!("block {cid_bytes:?} not found, this is a bug!!"); }; @@ -366,7 +339,6 @@ impl<'i> RepoHandle<'i> { Ok(Some(Record { did, cid, value })) }) .await - .into_diagnostic()? } /// lists records from this repository. @@ -382,12 +354,11 @@ impl<'i> RepoHandle<'i> { } let did = self.did.clone().into_static(); - let state = self.state.clone(); let prefix = keys::record_prefix_collection(&did, collection); let collection = collection.to_smolstr(); let cursor = cursor.map(|c| c.to_smolstr()); - tokio::task::spawn_blocking(move || { + self.state.db.run(move |db| { let mut results = Vec::new(); let mut next_cursor = None; @@ -407,8 +378,7 @@ impl<'i> RepoHandle<'i> { }; Box::new( - state - .db + db .indexer .record_range(prefix.as_slice()..end_key.as_slice()) .rev(), @@ -424,7 +394,7 @@ impl<'i> RepoHandle<'i> { prefix.clone() }; - Box::new(state.db.indexer.record_range(start_key.as_slice()..)) + Box::new(db.indexer.record_range(start_key.as_slice()..)) }; for item in iter { @@ -441,8 +411,7 @@ impl<'i> RepoHandle<'i> { } // look up using col|cid key built from collection and binary cid bytes - if let Ok(Some(block_bytes)) = state - .db + if let Ok(Some(block_bytes)) = db .indexer .block(keys::block_key(collection.as_str(), &cid_bytes)) { @@ -458,10 +427,9 @@ impl<'i> RepoHandle<'i> { }); } } - Result::<_, miette::Report>::Ok((results, next_cursor)) + Ok((results, next_cursor)) }) .await - .into_diagnostic()? .map(|(records, next_cursor)| RecordList { records, cursor: next_cursor.map(|rkey| { @@ -614,10 +582,8 @@ impl<'i> RepoHandle<'i> { /// gets how many records of a collection this repository has. pub async fn count_records(&self, collection: &str) -> Result { let did = self.did.clone().into_static(); - let state = self.state.clone(); let collection = collection.to_string(); - tokio::task::spawn_blocking(move || db::get_record_count(&state.db, &did, &collection)) + self.state.db.run(move |db| db::get_record_count(db, &did, &collection)) .await - .into_diagnostic()? } } diff --git a/src/control/repos/mod.rs b/src/control/repos/mod.rs index 9c68268..ea946b7 100644 --- a/src/control/repos/mod.rs +++ b/src/control/repos/mod.rs @@ -254,10 +254,9 @@ pub struct RepoHandle<'i> { impl<'i> RepoHandle<'i> { pub(crate) async fn state(&self) -> Result>> { let did_key = keys::repo_key(&self.did); - let app_state = self.state.clone(); - tokio::task::spawn_blocking(move || { - let bytes = app_state.db.repos.get(&did_key).into_diagnostic()?; + self.state.db.run(move |db| { + let bytes = db.repos.get(&did_key).into_diagnostic()?; bytes .as_deref() .map(db::deser_repo_state) @@ -265,7 +264,6 @@ impl<'i> RepoHandle<'i> { .map(|opt| opt.map(IntoStatic::into_static)) }) .await - .into_diagnostic()? } /// fetch the current state of this repository. @@ -275,18 +273,16 @@ impl<'i> RepoHandle<'i> { let did_key = keys::repo_key(&did); #[cfg(feature = "indexer")] let metadata_key = keys::repo_metadata_key(&did); - let app_state = self.state.clone(); - tokio::task::spawn_blocking(move || { - let state_bytes = app_state.db.repos.get(&did_key).into_diagnostic()?; + self.state.db.run(move |db| { + let state_bytes = db.repos.get(&did_key).into_diagnostic()?; let Some(state_bytes) = state_bytes else { return Ok(None); }; let repo_state = crate::db::deser_repo_state(&state_bytes)?; #[cfg(feature = "indexer")] - let tracked = app_state - .db + let tracked = db .repo_metadata .get(&metadata_key) .into_diagnostic()? @@ -299,18 +295,16 @@ impl<'i> RepoHandle<'i> { Ok(Some(repo_state_to_info(did, repo_state, tracked))) }) .await - .into_diagnostic()? } /// returns the collections of this repository and the number of records it has in each. pub async fn collections(&self) -> Result, u64>> { let did = self.did.clone().into_static(); - let state = self.state.clone(); - tokio::task::spawn_blocking(move || { + self.state.db.run(move |db| { let prefix = keys::did_collection_prefix(&did); let mut res = HashMap::new(); - for item in state.db.counts.prefix(&prefix) { + for item in db.counts.prefix(&prefix) { if res.len() >= 1000 { break; } @@ -332,7 +326,6 @@ impl<'i> RepoHandle<'i> { Ok(res) }) .await - .into_diagnostic()? } /// returns a bi-directionally validated mini doc. @@ -345,27 +338,24 @@ impl<'i> RepoHandle<'i> { #[cfg(feature = "indexer")] let is_pending = { let metadata_key = keys::repo_metadata_key(&self.did); - let app_state = self.state.clone(); - tokio::task::spawn_blocking(move || { - let metadata_bytes = app_state - .db + self.state.db.run(move |db| { + let metadata_bytes = db .repo_metadata .get(&metadata_key) .into_diagnostic()?; let Some(metadata_bytes) = metadata_bytes else { - return Ok::<_, miette::Report>(false); + return Ok(false); }; let metadata = crate::db::deser_repo_meta(metadata_bytes.as_ref())?; - Ok(app_state - .db - .indexer.pending + Ok(db + .indexer + .pending .get(crate::db::keys::pending_key(metadata.index_id)) .into_diagnostic()? .is_some()) }) .await - .map_err(|e| MiniDocError::Other(miette::miette!(e)))? .map_err(MiniDocError::Other)? }; #[cfg(not(feature = "indexer"))] diff --git a/src/control/stats.rs b/src/control/stats.rs index 56978e3..84ce371 100644 --- a/src/control/stats.rs +++ b/src/control/stats.rs @@ -64,16 +64,14 @@ impl Hydrant { state.db.jetstream.events.approximate_len() as u64, ); - let sizes = tokio::task::spawn_blocking(move || { - state - .db + let sizes = state.db.run(move |db| { + Ok(db .all_keyspaces() .into_iter() .map(|(name, ks)| (name, ks.disk_space())) - .collect::>() + .collect::>()) }) - .await - .into_diagnostic()?; + .await?; Ok(StatsResponse { counts, sizes }) } diff --git a/src/control/stream/relay.rs b/src/control/stream/relay.rs index 639c040..55e4c03 100644 --- a/src/control/stream/relay.rs +++ b/src/control/stream/relay.rs @@ -23,7 +23,8 @@ pub(crate) fn relay_stream_thread( None => Some( state .db - .relay.next_seq + .relay + .next_seq .load(Ordering::SeqCst) .saturating_sub(1), ), @@ -32,7 +33,8 @@ pub(crate) fn relay_stream_thread( .and_then(|_| { state .db - .relay.next_seq + .relay + .next_seq .load(Ordering::SeqCst) .checked_sub(1) }) diff --git a/src/crawler/by_collection.rs b/src/crawler/by_collection.rs index ef74fd7..fb36cb4 100644 --- a/src/crawler/by_collection.rs +++ b/src/crawler/by_collection.rs @@ -73,7 +73,9 @@ impl ByCollectionProducer { let cursor_key = by_collection_cursor_key(self.index_url.as_str(), collection); // resume from any persisted cursor, so a restart mid-pass doesn't rescan from scratch. - let mut cursor: Option = Db::get(db.cursors.keyspace(), &cursor_key) + let cursor_lookup_key = cursor_key.clone(); + let mut cursor: Option = db + .run(move |db| db.cursors.get(cursor_lookup_key).into_diagnostic()) .await .ok() .flatten() @@ -142,6 +144,7 @@ impl ByCollectionProducer { tokio::time::timeout( BLOCKING_TASK_TIMEOUT, + // CPU-bound JSON parsing + read-only DB scan, stays on spawn_blocking tokio::task::spawn_blocking(move || -> miette::Result> { let output = match serde_json::from_slice::(&bytes) { diff --git a/src/crawler/list_repos/producer.rs b/src/crawler/list_repos/producer.rs index 02f8105..4c26d1b 100644 --- a/src/crawler/list_repos/producer.rs +++ b/src/crawler/list_repos/producer.rs @@ -42,7 +42,12 @@ impl ListReposProducer { async fn get_cursor(&self) -> Result> { let key = crawler_cursor_key(self.url.as_str()); - let cursor_bytes = Db::get(self.checker.state.db.cursors.keyspace(), &key).await?; + let cursor_bytes = self + .checker + .state + .db + .run(move |db| db.cursors.get(key).into_diagnostic()) + .await?; Ok(cursor_bytes .as_deref() .and_then(|b| rmp_serde::from_slice::(b).ok())) @@ -115,6 +120,7 @@ impl ListReposProducer { match tokio::time::timeout( BLOCKING_TASK_TIMEOUT, + // CPU-bound JSON parsing + read-only DB scan, stays on spawn_blocking tokio::task::spawn_blocking(move || -> Result> { let output = serde_json::from_slice::(&bytes) .into_diagnostic() @@ -216,10 +222,9 @@ impl ListReposProducer { .await?; tokio::time::timeout( BLOCKING_TASK_TIMEOUT, - tokio::task::spawn_blocking(move || retry_batch.commit().into_diagnostic()), + db.run(move |_db| retry_batch.commit().into_diagnostic()), ) .await - .into_diagnostic()? .map_err(|_| miette::miette!("retry state commit timed out"))??; confirmed } else { @@ -251,7 +256,11 @@ impl ListReposProducer { pub(crate) async fn cursor_display(state: &AppState, relay_host: &Url) -> SmolStr { let key = crawler_cursor_key(relay_host.as_str()); - let cursor_bytes = match Db::get(state.db.cursors.keyspace(), &key).await { + let cursor_bytes = match state + .db + .run(move |db| db.cursors.get(key).into_diagnostic()) + .await + { Ok(b) => b, Err(e) => return e.to_smolstr(), }; diff --git a/src/crawler/list_repos/retry.rs b/src/crawler/list_repos/retry.rs index 211ccec..7ca24b2 100644 --- a/src/crawler/list_repos/retry.rs +++ b/src/crawler/list_repos/retry.rs @@ -50,6 +50,7 @@ impl RetryProducer { had_more: bool, } + // CPU-bound JSON parsing + read-only DB scan, stays on spawn_blocking let ScanResult { ready, existing, @@ -118,9 +119,8 @@ impl RetryProducer { .check_signals_batch(in_flight, &filter, &mut retry_batch, &existing) .await?; - tokio::task::spawn_blocking(move || retry_batch.commit().into_diagnostic()) + self.checker.state.db.run(move |_db| retry_batch.commit().into_diagnostic()) .await - .into_diagnostic()? .inspect_err(|e| error!(err = ?e, "retry state commit failed")) .ok(); diff --git a/src/crawler/worker.rs b/src/crawler/worker.rs index 37d1dd8..540bca1 100644 --- a/src/crawler/worker.rs +++ b/src/crawler/worker.rs @@ -1,4 +1,4 @@ -use crate::db::{CountDeltas, keys, ser_repo_state}; +use crate::db::{keys, ser_repo_state}; use crate::state::AppState; use crate::types::{RepoMetadata, RepoState}; use miette::{IntoDiagnostic, Result}; @@ -140,17 +140,14 @@ impl CrawlerWorker { let app_state = self.state.clone(); let surviving = tokio::time::timeout( BLOCKING_TASK_TIMEOUT, - tokio::task::spawn_blocking(move || -> Result> { + app_state.db.run(move |db| { let mut rng: SmallRng = rand::make_rng(); - let mut batch = app_state.db.inner.batch(); + let mut txn = crate::db::Txn::new(db); let mut surviving = Vec::new(); - let mut count_deltas = CountDeltas::default(); - let mut lifecycle_counts = app_state.db.lifecycle_counts(); for guard in guards { let did_key = keys::repo_key(&guard); let metadata_key = keys::repo_metadata_key(&guard); - if app_state - .db + if db .repos .contains_key(&did_key) .into_diagnostic()? @@ -159,45 +156,38 @@ impl CrawlerWorker { } let state = RepoState::backfilling(); let metadata = RepoMetadata::backfilling(rng.next_u64()); - batch.insert(&app_state.db.repos, &did_key, ser_repo_state(&state)?); - batch.insert( - &app_state.db.repo_metadata, + txn.batch + .insert(&db.repos, &did_key, ser_repo_state(&state)?); + txn.batch.insert( + &db.repo_metadata, &metadata_key, crate::db::ser_repo_meta(&metadata)?, ); #[cfg(feature = "indexer")] - batch.insert( - &app_state.db.indexer.pending, + txn.batch.insert( + &db.indexer.pending, keys::pending_key(metadata.index_id), &did_key, ); // clear any stale retry entry, this DID is confirmed and being enqueued - batch.remove(&app_state.db.crawler, keys::crawler_retry_key(&guard)); + txn.batch + .remove(&db.crawler, keys::crawler_retry_key(&guard)); trace!(did = %*guard, "enqueuing repo"); - count_deltas.add_repos(1); + txn.counts.add_repos(1); #[cfg(feature = "indexer")] - lifecycle_counts.transition( - &mut batch, - &guard, - crate::types::GaugeState::Pending, - )?; + txn.transition_lifecycle(&guard, crate::types::GaugeState::Pending)?; surviving.push(guard); } if let Some(cursor) = cursor_update { - batch.insert(&app_state.db.cursors, cursor.key, cursor.value); + txn.batch + .insert(&db.cursors, cursor.key, cursor.value); } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - let reservation = app_state.db.stage_count_deltas(&mut batch, &count_deltas); // todo: repo state overwrites here are acceptable? - batch.commit().into_diagnostic()?; - app_state.db.apply_lifecycle_counts(lifecycle_reservation); - app_state.db.apply_count_deltas(&count_deltas); - drop(reservation); + txn.commit()?; Ok(surviving) }), ) .await - .into_diagnostic()? .map_err(|_| { error!("enqueue batch timed out after {BLOCKING_TASK_TIMEOUT:?}"); miette::miette!("enqueue batch timed out") @@ -222,14 +212,15 @@ impl CrawlerWorker { let state = self.state.clone(); tokio::time::timeout( BLOCKING_TASK_TIMEOUT, - tokio::task::spawn_blocking(move || { - let mut batch = state.db.inner.batch(); - batch.insert(&state.db.cursors, cursor.key, cursor.value); - batch.commit().into_diagnostic() + state.db.run(move |db| { + let mut txn = crate::db::Txn::new(db); + txn.batch + .insert(&db.cursors, cursor.key, cursor.value); + txn.commit() }), ) .await - .into_diagnostic()? - .map_err(|_| miette::miette!("cursor-only commit timed out"))? + .map_err(|_| miette::miette!("cursor-only commit timed out"))??; + Ok(()) } } diff --git a/src/db/counts.rs b/src/db/counts.rs index ca6c135..4b57f2b 100644 --- a/src/db/counts.rs +++ b/src/db/counts.rs @@ -217,7 +217,7 @@ fn get_persisted_ks_count(db: &Db, name: &str) -> Result { } impl Db { - pub(crate) fn stage_count_deltas( + pub(super) fn stage_count_deltas( &self, batch: &mut OwnedWriteBatch, deltas: &CountDeltas, @@ -248,7 +248,7 @@ impl Db { }) } - pub(crate) fn apply_count_deltas(&self, deltas: &CountDeltas) { + pub(super) fn apply_count_deltas(&self, deltas: &CountDeltas) { for (key, delta) in deltas.iter() { self.update_count(key, *delta); } diff --git a/src/db/ephemeral.rs b/src/db/ephemeral.rs index c7d9951..a05fde1 100644 --- a/src/db/ephemeral.rs +++ b/src/db/ephemeral.rs @@ -71,14 +71,16 @@ pub fn jetstream_events_ttl_tick(db: &Db, ttl: &Duration) -> miette::Result<()> let cutoff_ts = now.saturating_sub(ttl.as_secs()); let cutoff_us = cutoff_ts.saturating_mul(1_000_000); - db.jetstream.events + db.jetstream + .events .rotate_memtable_and_wait() .into_diagnostic() .wrap_err("failed to rotate memtable before Jetstream TTL range drop")?; let before_space = db.jetstream.events.disk_space(); let before_tables = db.jetstream.events.table_count(); - db.jetstream.events + db.jetstream + .events .drop_range(..keys::jetstream_event_key(cutoff_us, 0)) .into_diagnostic() .wrap_err("failed Jetstream TTL range drop for old events")?; @@ -276,7 +278,8 @@ mod tests { payload: &[u8], ) -> miette::Result<()> { insert_relay_events(db, start_seq, count, payload)?; - db.relay.events + db.relay + .events .rotate_memtable_and_wait() .into_diagnostic()?; Ok(()) @@ -303,7 +306,8 @@ mod tests { } fn compact_relay_events_once(db: &crate::db::Db) -> miette::Result<()> { - db.relay.events + db.relay + .events .compact(Arc::new(fjall::compaction::Leveled::default())) .into_diagnostic() } @@ -331,7 +335,8 @@ mod tests { batch.remove(&db.relay.events, keys::relay_event_key(seq)); } batch.commit().into_diagnostic()?; - db.relay.events + db.relay + .events .rotate_memtable_and_wait() .into_diagnostic()?; Ok(()) @@ -485,7 +490,8 @@ mod tests { let after_delete = db.relay.events.disk_space(); for _ in 0..16 { - db.relay.events + db.relay + .events .compact(Arc::new(fjall::compaction::Leveled::default())) .into_diagnostic()?; } diff --git a/src/db/keyspaces.rs b/src/db/keyspaces.rs index b9ce42f..f275b6d 100644 --- a/src/db/keyspaces.rs +++ b/src/db/keyspaces.rs @@ -49,6 +49,7 @@ impl OpenCx<'_> { } #[cfg(feature = "indexer")] +#[derive(Clone)] pub struct IndexerDb { /// maps `{DID}|{COL}|{RKey}` -> record CID pub(super) records: Ks, @@ -98,6 +99,7 @@ impl IndexerDb { } #[cfg(feature = "indexer_stream")] +#[derive(Clone)] pub struct StreamDb { /// maps `{ID}` (u64 BE) -> `StoredEvent`, the source for the json stream api pub(super) events: Ks, @@ -162,6 +164,7 @@ impl StreamDb { } #[cfg(feature = "jetstream")] +#[derive(Clone)] pub(crate) struct JetstreamDb { /// maps `{time_us}|{ID}` (16 bytes) -> jetstream event data pub(crate) events: Ks, @@ -207,6 +210,7 @@ impl JetstreamDb { } #[cfg(feature = "relay")] +#[derive(Clone)] pub(crate) struct RelayDb { /// maps `{SEQ}` (u64 BE) -> re-encoded relay frame pub(crate) events: Ks, diff --git a/src/db/lifecycle_counts.rs b/src/db/lifecycle_counts.rs index 3fc8ff3..b7453f4 100644 --- a/src/db/lifecycle_counts.rs +++ b/src/db/lifecycle_counts.rs @@ -9,25 +9,26 @@ use crate::types::{GaugeState, ResyncErrorKind, ResyncState}; use super::{CountDeltaReservation, CountDeltas, Db, deser_repo_meta, keys}; -pub(crate) struct LifecycleCountBatch<'a> { +pub(super) struct LifecycleCountBatch<'a> { db: &'a Db, _lock: MutexGuard<'a, ()>, deltas: CountDeltas, staged: BTreeMap, GaugeState>, } -pub(crate) struct LifecycleCountReservation<'a> { +pub(super) struct LifecycleCountReservation<'a> { _lock: MutexGuard<'a, ()>, _count_reservation: Option, deltas: CountDeltas, } impl Db { - pub(crate) fn lifecycle_counts(&self) -> LifecycleCountBatch<'_> { + pub(super) fn lifecycle_counts(&self) -> LifecycleCountBatch<'_> { LifecycleCountBatch { db: self, _lock: self - .indexer.lifecycle_count_lock + .indexer + .lifecycle_count_lock .lock() .expect("lifecycle count lock poisoned"), deltas: CountDeltas::default(), @@ -35,14 +36,14 @@ impl Db { } } - pub(crate) fn apply_lifecycle_counts(&self, reservation: LifecycleCountReservation<'_>) { + pub(super) fn apply_lifecycle_counts(&self, reservation: LifecycleCountReservation<'_>) { self.apply_count_deltas(&reservation.deltas); drop(reservation); } } impl<'a> LifecycleCountBatch<'a> { - pub(crate) fn transition( + pub(super) fn transition( &mut self, batch: &mut OwnedWriteBatch, did: &Did<'_>, @@ -64,7 +65,7 @@ impl<'a> LifecycleCountBatch<'a> { Ok(false) } - pub(crate) fn transition_pending_key( + pub(super) fn transition_pending_key( &mut self, batch: &mut OwnedWriteBatch, did: &Did<'_>, @@ -79,7 +80,7 @@ impl<'a> LifecycleCountBatch<'a> { Ok(true) } - pub(crate) fn stage(self, batch: &mut OwnedWriteBatch) -> LifecycleCountReservation<'a> { + pub(super) fn stage(self, batch: &mut OwnedWriteBatch) -> LifecycleCountReservation<'a> { let count_reservation = self.db.stage_count_deltas(batch, &self.deltas); LifecycleCountReservation { _lock: self._lock, @@ -112,7 +113,8 @@ impl<'a> LifecycleCountBatch<'a> { if let Some(index_id) = index_id { if self .db - .indexer.pending + .indexer + .pending .get(keys::pending_key(index_id)) .into_diagnostic()? .is_some() @@ -123,7 +125,9 @@ impl<'a> LifecycleCountBatch<'a> { let gauge = if is_pending { GaugeState::Pending - } else if let Some(resync_bytes) = self.db.indexer.resync.get(did_key).into_diagnostic()? { + } else if let Some(resync_bytes) = + self.db.indexer.resync.get(did_key).into_diagnostic()? + { gauge_from_resync(&resync_bytes) } else { GaugeState::Synced @@ -141,7 +145,13 @@ impl<'a> LifecycleCountBatch<'a> { }; Ok(current.as_slice() == pending_key - && self.db.indexer.pending.get(current).into_diagnostic()?.is_some()) + && self + .db + .indexer + .pending + .get(current) + .into_diagnostic()? + .is_some()) } fn stage_membership( @@ -301,7 +311,13 @@ mod tests { assert_eq!(db.get_count_sync("pending"), 1); assert!(!complete_pending(&db, &did, stale_pending)?); assert_eq!(db.get_count_sync("pending"), 1); - assert!(db.indexer.pending.get(current_pending).into_diagnostic()?.is_some()); + assert!( + db.indexer + .pending + .get(current_pending) + .into_diagnostic()? + .is_some() + ); Ok(()) } @@ -337,7 +353,8 @@ mod tests { let db = Db::open(&cfg)?; assert_eq!(db.get_count_sync("pending"), 0); assert!( - db.indexer.pending + db.indexer + .pending .get(keys::pending_key(1)) .into_diagnostic()? .is_none() diff --git a/src/db/migration/v8.rs b/src/db/migration/v8.rs index 018765d..ef4ed0f 100644 --- a/src/db/migration/v8.rs +++ b/src/db/migration/v8.rs @@ -80,7 +80,8 @@ fn primary_lifecycle_gauge(db: &Db, repo_key: &[u8]) -> Result { let metadata = deser_repo_meta(metadata_bytes.as_ref()) .wrap_err("invalid repo metadata during lifecycle count rebuild")?; if db - .indexer.pending + .indexer + .pending .get(keys::pending_key(metadata.index_id)) .into_diagnostic()? .is_some() @@ -89,7 +90,8 @@ fn primary_lifecycle_gauge(db: &Db, repo_key: &[u8]) -> Result { } } - db.indexer.resync + db.indexer + .resync .get(repo_key) .into_diagnostic()? .map(|bytes| crate::db::lifecycle_counts::gauge_from_resync(bytes.as_ref())) diff --git a/src/db/mod.rs b/src/db/mod.rs index cff00aa..fc91bf0 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -1,6 +1,6 @@ use crate::types::{RepoMetadata, RepoState}; -use fjall::{Database, Keyspace, PersistMode, Slice}; +use fjall::{Database, Keyspace, PersistMode}; use miette::{Context, IntoDiagnostic, Result}; use scc::HashMap; use smol_str::SmolStr; @@ -14,7 +14,7 @@ use url::Url; pub mod compaction; pub mod counts; #[cfg(feature = "indexer")] -pub(crate) use counts::CountDeltaReservation; +use counts::CountDeltaReservation; pub use counts::{CountDeltas, load_count_delta_watermark, set_ks_count}; pub mod ephemeral; pub mod filter; @@ -30,7 +30,6 @@ mod open; pub mod registry; pub mod schema; mod train; -#[cfg(feature = "indexer")] mod txn; #[cfg(feature = "indexer")] @@ -45,12 +44,12 @@ pub use keyspaces::StreamDb; pub use schema::Ks; #[cfg(feature = "indexer")] -pub(crate) use lifecycle_counts::LifecycleCountBatch; -#[cfg(feature = "indexer")] -pub(crate) use txn::Txn; +use lifecycle_counts::LifecycleCountBatch; +pub(crate) use txn::{Txn, TxnCommitTimings}; use tracing::error; +#[derive(Clone)] pub struct Db { pub inner: Arc, pub path: std::path::PathBuf, @@ -70,7 +69,7 @@ pub struct Db { pub(crate) relay: RelayDb, #[cfg(feature = "backlinks")] pub(crate) backlinks: Ks, - pub counts_map: HashMap, + pub counts_map: Arc>, next_count_delta_id: Arc, count_delta_checkpoint_watermark: Arc, count_delta_gc_watermark: Arc, @@ -79,6 +78,18 @@ pub struct Db { } impl Db { + /// runs synchronous database work without blocking the async runtime. + pub async fn run(&self, f: F) -> Result + where + T: Send + 'static, + F: FnOnce(&Db) -> Result + Send + 'static, + { + let db = self.clone(); + tokio::task::spawn_blocking(move || f(&db)) + .await + .into_diagnostic()? + } + pub fn persist(&self) -> Result<()> { #[cfg(not(feature = "__persist_sync_all"))] const MODE: PersistMode = PersistMode::Buffer; @@ -121,15 +132,6 @@ impl Db { Ok(()) } - - pub async fn get(ks: Keyspace, key: impl Into) -> Result> { - let key = key.into(); - tokio::task::spawn_blocking(move || { - ks.get(key).inspect_err(check_poisoned).into_diagnostic() - }) - .await - .into_diagnostic()? - } } #[cfg(feature = "indexer")] @@ -155,9 +157,9 @@ pub fn set_firehose_cursor(db: &Db, relay: &Url, cursor: i64) -> Result<()> { pub async fn get_firehose_cursor(db: &Db, relay: &Url) -> Result> { let key = keys::firehose_cursor_key_from_url(relay); - Db::get(db.cursors.keyspace(), key) + db.run(move |db| db.cursors.get(key).into_diagnostic()) .await? - .map(|v: Slice| { + .map(|v| { Ok(i64::from_be_bytes( v.as_ref() .try_into() diff --git a/src/db/open.rs b/src/db/open.rs index bd4f9d7..65d14fd 100644 --- a/src/db/open.rs +++ b/src/db/open.rs @@ -18,6 +18,41 @@ use super::{Db, migration, registry, schema}; impl Db { pub fn open(cfg: &Config) -> Result { + let (db, count_delta_gc_watermark) = 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()), + }; + + let this = Self::assemble_keyspaces_and_verify(&cx, count_delta_gc_watermark)?; + + migration::run(&this)?; + + this.init_modes()?; + + this.load_persisted_counts()?; + + this.restore_count_deltas_and_next_id()?; + + Ok(this) + } + + fn open_database(cfg: &Config) -> Result<(Arc, Arc)> { let count_delta_gc_watermark = Arc::new(AtomicU64::new(0)); let db = Database::builder(&cfg.database_path) .cache_size(cfg.cache_size * 2_u64.pow(20) / 2) @@ -44,7 +79,10 @@ impl Db { .open() .into_diagnostic()?; let db = Arc::new(db); + Ok((db, count_delta_gc_watermark)) + } + fn load_dicts(cfg: &Config) -> StdHashMap<&'static str, Arc<[u8]>> { let load_dict = |name: &str| -> Option> { let path = cfg.database_path.join(format!("dict_{name}.bin")); if path.exists() @@ -58,7 +96,7 @@ impl Db { } None }; - let dicts = registry::trainable() + registry::trainable() .into_iter() .fold(StdHashMap::new(), |mut acc, (name, _)| { let Some(dict) = load_dict(name) else { @@ -66,41 +104,29 @@ impl Db { }; acc.insert(name, dict); acc - }); - 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()), - }; + }) + } - let repos = Ks::::open(&cx)?; - let repo_metadata = Ks::::open(&cx)?; - let cursors = Ks::::open(&cx)?; - let counts = Ks::::open(&cx)?; - let filter = Ks::::open(&cx)?; - let crawler = Ks::::open(&cx)?; + fn assemble_keyspaces_and_verify( + cx: &OpenCx, + count_delta_gc_watermark: Arc, + ) -> Result { + let repos = Ks::::open(cx)?; + let repo_metadata = Ks::::open(cx)?; + let cursors = Ks::::open(cx)?; + let counts = Ks::::open(cx)?; + let filter = Ks::::open(cx)?; + let crawler = Ks::::open(cx)?; #[cfg(feature = "backlinks")] - let backlinks = Ks::::open(&cx)?; + let backlinks = Ks::::open(cx)?; #[cfg(feature = "indexer")] - let indexer = super::keyspaces::IndexerDb::open(&cx)?; + let indexer = super::keyspaces::IndexerDb::open(cx)?; #[cfg(feature = "indexer_stream")] - let stream = super::keyspaces::StreamDb::open(&cx)?; + let stream = super::keyspaces::StreamDb::open(cx)?; #[cfg(feature = "jetstream")] - let jetstream = super::keyspaces::JetstreamDb::open(&cx)?; + let jetstream = super::keyspaces::JetstreamDb::open(cx)?; #[cfg(feature = "relay")] - let relay = super::keyspaces::RelayDb::open(&cx)?; + let relay = super::keyspaces::RelayDb::open(cx)?; // every opened keyspace must have a registry row and vice versa, so the // by-name, /stats, /debug, and training tables cannot silently drift. @@ -114,9 +140,9 @@ impl Db { ); } - let this = Self { - inner: db, - path: cfg.database_path.clone(), + Ok(Self { + inner: Arc::clone(cx.db), + path: cx.cfg.database_path.clone(), repos, repo_metadata, cursors, @@ -133,45 +159,51 @@ impl Db { relay, #[cfg(feature = "backlinks")] backlinks, - counts_map: HashMap::new(), + 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)), - }; - - migration::run(&this)?; + }) + } + fn init_modes(&self) -> Result<()> { #[cfg(feature = "relay")] - this.relay.init()?; + self.relay.init()?; #[cfg(feature = "indexer_stream")] - this.stream.init()?; + self.stream.init()?; #[cfg(feature = "jetstream")] - this.jetstream.init()?; + self.jetstream.init()?; + Ok(()) + } + fn load_persisted_counts(&self) -> Result<()> { // load counts into memory - for guard in this.counts.prefix(keys::COUNT_KS_PREFIX) { + 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 _ = this + let _ = self .counts_map .insert_sync(SmolStr::new(name), read_u64_counter(&v)?); } + Ok(()) + } - let durable_watermark = load_count_delta_watermark(&this)?; - replay_count_deltas(&this, durable_watermark)?; - this.count_delta_checkpoint_watermark + 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); - this.count_delta_gc_watermark + 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 = this + let next_count_delta_id = self .counts .prefix(keys::COUNT_DELTA_PREFIX) .next_back() @@ -183,9 +215,9 @@ impl Db { .transpose()? .unwrap_or(0) .max(durable_watermark.saturating_add(1)); - this.next_count_delta_id + self.next_count_delta_id .store(next_count_delta_id, Ordering::Relaxed); - Ok(this) + Ok(()) } } diff --git a/src/db/txn.rs b/src/db/txn.rs index 3a2c45f..aded08f 100644 --- a/src/db/txn.rs +++ b/src/db/txn.rs @@ -1,18 +1,52 @@ +#[cfg(feature = "indexer")] use std::collections::HashMap; +use std::time::Duration; +#[cfg(feature = "indexer")] use bytes::Bytes; -use jacquard_common::IntoStatic; +#[cfg(feature = "indexer")] use jacquard_common::types::cid::IpldCid; +#[cfg(feature = "indexer")] use jacquard_common::types::did::Did; use miette::{IntoDiagnostic, Result}; +#[cfg(feature = "indexer")] use crate::db::types::{DbAction, DbRkey, DbTid}; -use crate::db::{CountDeltas, Db, keys}; +use crate::db::{CountDeltas, Db}; +#[cfg(feature = "indexer")] +use crate::db::{LifecycleCountBatch, keys}; +#[cfg(feature = "indexer")] use crate::ops::record_events::{EmitOp, RecordEmitter, RecordEventOrigin, RecordEvents}; +#[cfg(feature = "indexer")] use crate::state::AppState; #[cfg(feature = "indexer")] -use crate::types::GaugeState; -use crate::types::RepoState; +use crate::types::{GaugeState, RepoState}; + +#[cfg(feature = "firehose-diagnostics")] +type TxnInstant = std::time::Instant; + +#[cfg(not(feature = "firehose-diagnostics"))] +#[derive(Clone, Copy)] +struct TxnInstant; + +#[cfg(not(feature = "firehose-diagnostics"))] +impl TxnInstant { + #[inline(always)] + fn now() -> Self { + Self + } + + #[inline(always)] + fn elapsed(&self) -> Duration { + Duration::ZERO + } +} + +pub(crate) struct TxnCommitTimings { + pub(crate) stage_counts: Duration, + pub(crate) stage_and_commit: Duration, + pub(crate) apply_counts: Duration, +} /// one atomic database write, including its in-memory count projections. /// @@ -23,20 +57,36 @@ pub(crate) struct Txn<'db> { pub(crate) db: &'db Db, pub(crate) counts: CountDeltas, #[cfg(feature = "indexer")] - lifecycle_transitions: Vec<(Did<'static>, GaugeState)>, + lifecycle: Option>, } impl<'db> Txn<'db> { + #[allow(dead_code)] pub(crate) fn new(db: &'db Db) -> Self { Self { batch: db.inner.batch(), db, counts: CountDeltas::default(), #[cfg(feature = "indexer")] - lifecycle_transitions: Vec::new(), + lifecycle: None, } } + pub(crate) fn from_parts( + db: &'db Db, + batch: fjall::OwnedWriteBatch, + counts: CountDeltas, + ) -> Self { + Self { + batch, + db, + counts, + #[cfg(feature = "indexer")] + lifecycle: None, + } + } + + #[cfg(feature = "indexer")] pub(crate) fn records<'txn, 'did, 'repo>( &'txn mut self, state: &AppState, @@ -46,6 +96,7 @@ impl<'db> Txn<'db> { self.record_scope(state, commit_rev, did, RecordEventOrigin::Live, false) } + #[cfg(feature = "indexer")] pub(crate) fn backfill_records<'txn, 'did, 'repo>( &'txn mut self, state: &AppState, @@ -55,6 +106,7 @@ impl<'db> Txn<'db> { self.record_scope(state, commit_rev, did, RecordEventOrigin::Backfill, true) } + #[cfg(feature = "indexer")] fn record_scope<'txn, 'did, 'repo>( &'txn mut self, state: &AppState, @@ -77,33 +129,77 @@ impl<'db> Txn<'db> { } #[cfg(feature = "indexer")] - pub(crate) fn transition_lifecycle(&mut self, did: &Did<'_>, gauge: GaugeState) { - self.lifecycle_transitions - .push((did.clone().into_static(), gauge)); + pub(crate) fn transition_lifecycle( + &mut self, + did: &Did<'_>, + gauge: GaugeState, + ) -> Result { + let lifecycle = self + .lifecycle + .get_or_insert_with(|| self.db.lifecycle_counts()); + lifecycle.transition(&mut self.batch, did, gauge) + } + + #[cfg(feature = "indexer")] + pub(crate) fn transition_pending_key( + &mut self, + did: &Did<'_>, + pending_key: &[u8], + gauge: GaugeState, + ) -> Result { + let lifecycle = self + .lifecycle + .get_or_insert_with(|| self.db.lifecycle_counts()); + lifecycle.transition_pending_key(&mut self.batch, did, pending_key, gauge) + } + + #[allow(dead_code)] + pub(crate) fn commit(self) -> Result<()> { + self.commit_with(|_| Ok(())).map(|_| ()) } - pub(crate) fn commit(mut self) -> Result<()> { + /// stages count reservations, lets the caller add commit-coupled writes, + /// commits once, then applies the in-memory projections. + #[allow(dead_code)] + pub(crate) fn commit_with(mut self, stage: F) -> Result<(T, TxnCommitTimings)> + where + F: FnOnce(&mut fjall::OwnedWriteBatch) -> Result, + { + let stage_counts_started = TxnInstant::now(); #[cfg(feature = "indexer")] - let lifecycle_reservation = { - let mut lifecycle = self.db.lifecycle_counts(); - for (did, gauge) in self.lifecycle_transitions { - lifecycle.transition(&mut self.batch, &did, gauge)?; - } - lifecycle.stage(&mut self.batch) - }; + let lifecycle_reservation = self + .lifecycle + .map(|lifecycle| lifecycle.stage(&mut self.batch)); let count_reservation = self.db.stage_count_deltas(&mut self.batch, &self.counts); + let stage_counts = stage_counts_started.elapsed(); + let stage_and_commit_started = TxnInstant::now(); + let staged = stage(&mut self.batch)?; self.batch.commit().into_diagnostic()?; + let stage_and_commit = stage_and_commit_started.elapsed(); + let apply_counts_started = TxnInstant::now(); self.db.apply_count_deltas(&self.counts); drop(count_reservation); #[cfg(feature = "indexer")] - self.db.apply_lifecycle_counts(lifecycle_reservation); - Ok(()) + if let Some(reservation) = lifecycle_reservation { + self.db.apply_lifecycle_counts(reservation); + } + let apply_counts = apply_counts_started.elapsed(); + + Ok(( + staged, + TxnCommitTimings { + stage_counts, + stage_and_commit, + apply_counts, + }, + )) } } /// one commit's record mutations within a larger atomic transaction. +#[cfg(feature = "indexer")] pub(crate) struct RecordTxn<'txn, 'db, 'did, 'repo> { txn: &'txn mut Txn<'db>, emitter: RecordEmitter, @@ -116,6 +212,7 @@ pub(crate) struct RecordTxn<'txn, 'db, 'did, 'repo> { collection_deltas: HashMap, } +#[cfg(feature = "indexer")] impl RecordTxn<'_, '_, '_, '_> { pub(crate) fn put_record( &mut self, @@ -130,7 +227,6 @@ impl RecordTxn<'_, '_, '_, '_> { self.records_delta += 1; } - #[cfg(feature = "indexer")] if !self.ephemeral { let cid_bytes = cid.to_bytes(); if !self.only_index_links { @@ -180,7 +276,6 @@ impl RecordTxn<'_, '_, '_, '_> { self.records_delta -= 1; } - #[cfg(feature = "indexer")] if !self.ephemeral { self.txn.batch.remove( &self.txn.db.indexer.records, @@ -223,7 +318,6 @@ impl RecordTxn<'_, '_, '_, '_> { } pub(crate) fn finish(self) -> Result { - #[cfg(feature = "indexer")] if !self.ephemeral { for (collection, delta) in &self.collection_deltas { crate::db::update_record_count( diff --git a/src/ingest/firehose_stats.rs b/src/ingest/firehose_stats.rs index 68dea85..1f8bd7b 100644 --- a/src/ingest/firehose_stats.rs +++ b/src/ingest/firehose_stats.rs @@ -10,9 +10,7 @@ pub use kinds::*; #[cfg(not(feature = "firehose-diagnostics"))] pub use noop::*; #[cfg(feature = "firehose-diagnostics")] -pub use relay::{ - RelayShardStats, RelayWorkerStats, RelayWorkerStatsSnapshot, -}; +pub use relay::{RelayShardStats, RelayWorkerStats, RelayWorkerStatsSnapshot}; #[cfg(feature = "firehose-diagnostics")] pub use source::{FirehoseSourceStats, FirehoseStats, FirehoseStatsSnapshot}; diff --git a/src/ingest/indexer/shard.rs b/src/ingest/indexer/shard.rs index f30311e..4120e1f 100644 --- a/src/ingest/indexer/shard.rs +++ b/src/ingest/indexer/shard.rs @@ -244,19 +244,24 @@ impl FirehoseWorker { } } - for (did, gauge) in ctx.lifecycle_transitions.drain(..) { - ctx.txn.transition_lifecycle(&did, gauge); + let lifecycle_result = ctx + .lifecycle_transitions + .drain(..) + .try_for_each(|(did, gauge)| ctx.txn.transition_lifecycle(&did, gauge).map(|_| ())); + if let Err(e) = lifecycle_result { + error!(shard = id, err = %e, "failed to stage lifecycle transitions"); + continue; } if let Err(e) = ctx.txn.commit() { error!(shard = id, err = %e, "failed to commit transaction"); continue; } #[cfg(feature = "indexer_stream")] - for evt in broadcast_events.drain(..) { + for evt in ctx.broadcast_events.drain(..) { let _ = state.db.stream.event_tx.send(evt); } #[cfg(feature = "jetstream")] - for evt in jetstream_events.drain(..) { + for evt in ctx.jetstream_events.drain(..) { let _ = state.db.jetstream.tx.send(evt); } @@ -520,8 +525,7 @@ impl FirehoseWorker { repo_state: RepoState<'s>, ) -> Result, IngestError> { let db = &ctx.state.db; - let mut batch = db.inner.batch(); - let mut lifecycle_counts = db.lifecycle_counts(); + let mut txn = Txn::new(db); let repo_key = keys::repo_key(did); let meta_key = keys::repo_metadata_key(did); @@ -548,22 +552,21 @@ impl FirehoseWorker { // remove old pending entry and insert new one with fresh index_id if had_metadata { // only remove if we had one so we dont delete a random entry - batch.remove(&db.indexer.pending, old_pkey); + txn.batch.remove(&db.indexer.pending, old_pkey); } metadata.index_id = rand::random::(); - batch.insert( + txn.batch.insert( &db.indexer.pending, keys::pending_key(metadata.index_id), &repo_key, ); - batch.insert(&db.repo_metadata, &meta_key, ser_repo_meta(&metadata)?); + txn.batch + .insert(&db.repo_metadata, &meta_key, ser_repo_meta(&metadata)?); if !was_pending { - lifecycle_counts.transition(&mut batch, did, GaugeState::Pending)?; + txn.transition_lifecycle(did, GaugeState::Pending)?; } - let lifecycle_reservation = lifecycle_counts.stage(&mut batch); - batch.commit().into_diagnostic()?; - db.apply_lifecycle_counts(lifecycle_reservation); + txn.commit()?; if !was_pending { ctx.state.notify_backfill(); diff --git a/src/ingest/relay/context.rs b/src/ingest/relay/context.rs index f381c36..a30c122 100644 --- a/src/ingest/relay/context.rs +++ b/src/ingest/relay/context.rs @@ -456,9 +456,8 @@ impl WorkerContext<'_> { self.count_deltas.add_repos(1); - self.stats.record_repo_state_outcome( - crate::ingest::firehose_stats::RepoStateLoadOutcome::Miss, - ); + self.stats + .record_repo_state_outcome(crate::ingest::firehose_stats::RepoStateLoadOutcome::Miss); self.stats.record_new_account(new_account_started.elapsed()); Ok(Some(repo_state)) diff --git a/src/ingest/relay/handlers.rs b/src/ingest/relay/handlers.rs index c66128d..df34a43 100644 --- a/src/ingest/relay/handlers.rs +++ b/src/ingest/relay/handlers.rs @@ -142,8 +142,14 @@ impl RelayWorker { let repo_key = keys::repo_key(&identity.did); - ctx.sink - .identity(ctx.state, &mut ctx.batch, firehose, identity, repo_state, snapshot)?; + ctx.sink.identity( + ctx.state, + &mut ctx.batch, + firehose, + identity, + repo_state, + snapshot, + )?; ctx.batch.insert( &ctx.state.db.repos, diff --git a/src/ingest/relay/sink/indexer.rs b/src/ingest/relay/sink/indexer.rs index 5a6123f..f943fe2 100644 --- a/src/ingest/relay/sink/indexer.rs +++ b/src/ingest/relay/sink/indexer.rs @@ -1,16 +1,16 @@ use fjall::OwnedWriteBatch; use jacquard_common::IntoStatic; use jacquard_common::types::string::Handle; -use miette::{IntoDiagnostic, Result}; +use miette::Result; use url::Url; +use crate::db::types::DidKey; use crate::ingest::indexer::{ IndexerAccountData, IndexerCommitData, IndexerEvent, IndexerEventData, IndexerIdentityData, IndexerMessage, IndexerTx, }; use crate::ingest::stream::{Account, Commit, Identity, Sync}; use crate::state::AppState; -use crate::db::types::DidKey; use crate::types::{RepoState, RepoStatus}; /// per-worker seed from which each shard builds its sink. @@ -79,15 +79,16 @@ impl EventSink { chain_break: bool, parsed_blocks: jacquard_repo::car::reader::ParsedCar, ) -> Result<()> { - self.pending.push(IndexerMessage::Event(Box::new(IndexerEvent { - seq: commit.seq, - firehose: firehose.clone(), - data: IndexerEventData::Commit(IndexerCommitData { - commit, - chain_break, - parsed_blocks, - }), - }))); + self.pending + .push(IndexerMessage::Event(Box::new(IndexerEvent { + seq: commit.seq, + firehose: firehose.clone(), + data: IndexerEventData::Commit(IndexerCommitData { + commit, + chain_break, + parsed_blocks, + }), + }))); Ok(()) } @@ -98,11 +99,12 @@ impl EventSink { firehose: &Url, sync: Sync<'static>, ) -> Result<()> { - self.pending.push(IndexerMessage::Event(Box::new(IndexerEvent { - seq: sync.seq, - firehose: firehose.clone(), - data: IndexerEventData::Sync(sync.did.into_static()), - }))); + self.pending + .push(IndexerMessage::Event(Box::new(IndexerEvent { + seq: sync.seq, + firehose: firehose.clone(), + data: IndexerEventData::Sync(sync.did.into_static()), + }))); Ok(()) } @@ -115,13 +117,14 @@ impl EventSink { repo_state: &RepoState, snapshot: IdentitySnapshot, ) -> Result<()> { - let changed = repo_state.handle != snapshot.handle - || repo_state.signing_key != snapshot.signing_key; - self.pending.push(IndexerMessage::Event(Box::new(IndexerEvent { - seq: identity.seq, - firehose: firehose.clone(), - data: IndexerEventData::Identity(IndexerIdentityData { identity, changed }), - }))); + let changed = + repo_state.handle != snapshot.handle || repo_state.signing_key != snapshot.signing_key; + self.pending + .push(IndexerMessage::Event(Box::new(IndexerEvent { + seq: identity.seq, + firehose: firehose.clone(), + data: IndexerEventData::Identity(IndexerIdentityData { identity, changed }), + }))); Ok(()) } @@ -137,25 +140,26 @@ impl EventSink { was_active: bool, ) -> Result<()> { let changed = repo_state.active != was_active || repo_state.status != snapshot.status; - self.pending.push(IndexerMessage::Event(Box::new(IndexerEvent { - seq: account.seq, - firehose: firehose.clone(), - data: IndexerEventData::Account(IndexerAccountData { - account, - was_active, - changed, - }), - }))); + self.pending + .push(IndexerMessage::Event(Box::new(IndexerEvent { + seq: account.seq, + firehose: firehose.clone(), + data: IndexerEventData::Account(IndexerAccountData { + account, + was_active, + changed, + }), + }))); Ok(()) } - pub(crate) fn commit_batch( + pub(crate) fn commit_txn( &mut self, _state: &AppState, - batch: OwnedWriteBatch, - ) -> Result { - batch.commit().into_diagnostic()?; - Ok(Staged) + txn: crate::db::Txn<'_>, + ) -> Result<(Staged, crate::db::TxnCommitTimings)> { + let (_, timings) = txn.commit_with(|_| Ok(()))?; + Ok((Staged, timings)) } pub(crate) fn flush(&mut self, _state: &AppState, _staged: Staged) { diff --git a/src/ingest/relay/sink/none.rs b/src/ingest/relay/sink/none.rs index 73507a9..0fb22ec 100644 --- a/src/ingest/relay/sink/none.rs +++ b/src/ingest/relay/sink/none.rs @@ -2,7 +2,7 @@ //! repo state but are not forwarded anywhere. use fjall::OwnedWriteBatch; -use miette::{IntoDiagnostic, Result}; +use miette::Result; use url::Url; use crate::ingest::stream::{Account, Commit, Identity, Sync}; @@ -87,13 +87,13 @@ impl EventSink { Ok(()) } - pub(crate) fn commit_batch( + pub(crate) fn commit_txn( &mut self, _state: &AppState, - batch: OwnedWriteBatch, - ) -> Result { - batch.commit().into_diagnostic()?; - Ok(Staged) + txn: crate::db::Txn<'_>, + ) -> Result<(Staged, crate::db::TxnCommitTimings)> { + let (_, timings) = txn.commit_with(|_| Ok(()))?; + Ok((Staged, timings)) } pub(crate) fn flush(&mut self, _state: &AppState, _staged: Staged) {} diff --git a/src/ingest/relay/sink/relay.rs b/src/ingest/relay/sink/relay.rs index 4a1a263..990f09b 100644 --- a/src/ingest/relay/sink/relay.rs +++ b/src/ingest/relay/sink/relay.rs @@ -1,13 +1,13 @@ use fjall::OwnedWriteBatch; -use miette::{IntoDiagnostic, Result}; +use miette::Result; #[cfg(feature = "jetstream")] use smol_str::ToSmolStr; use std::sync::atomic::Ordering; use url::Url; +use crate::db::keys; #[cfg(feature = "jetstream")] use crate::db::types::TrimmedDid; -use crate::db::keys; use crate::ingest::stream::{Account, Commit, Identity, Sync, encode_frame}; use crate::state::AppState; #[cfg(feature = "jetstream")] @@ -121,7 +121,9 @@ impl EventSink { let has_subscribers = state.db.jetstream.tx.receiver_count() > 0; for (op_index, collection) in &jetstream_ops { let ephemeral = has_subscribers - .then(|| build_relay_commit_ephemeral(&commit, *op_index, collection, &parsed_blocks)) + .then(|| { + build_relay_commit_ephemeral(&commit, *op_index, collection, &parsed_blocks) + }) .flatten(); self.jetstream_events.push(( StoredJetstreamEvent::RelayCommit { @@ -201,31 +203,34 @@ impl EventSink { Ok(()) } - pub(crate) fn commit_batch( + pub(crate) fn commit_txn( &mut self, state: &AppState, - batch: OwnedWriteBatch, - ) -> Result { + txn: crate::db::Txn<'_>, + ) -> Result<(Staged, crate::db::TxnCommitTimings)> { #[cfg(feature = "jetstream")] { - let mut batch = batch; - let mut jetstream_broadcasts = Vec::new(); let _lock = state.db.jetstream.lock.lock(); - for (event, ephemeral) in self.jetstream_events.drain(..) { - jetstream_broadcasts.push(crate::jetstream::stage_event( - &mut batch, &state.db, event, ephemeral, - )?); - } - batch.commit().into_diagnostic()?; - Ok(Staged { - jetstream_broadcasts, - }) + let (jetstream_broadcasts, timings) = txn.commit_with(|batch| { + self.jetstream_events + .drain(..) + .map(|(event, ephemeral)| { + crate::jetstream::stage_event(batch, &state.db, event, ephemeral) + }) + .collect::>>() + })?; + Ok(( + Staged { + jetstream_broadcasts, + }, + timings, + )) } #[cfg(not(feature = "jetstream"))] { let _ = state; - batch.commit().into_diagnostic()?; - Ok(Staged {}) + let (_, timings) = txn.commit_with(|_| Ok(()))?; + Ok((Staged {}, timings)) } } diff --git a/src/ingest/relay/worker.rs b/src/ingest/relay/worker.rs index e2c93e0..957c003 100644 --- a/src/ingest/relay/worker.rs +++ b/src/ingest/relay/worker.rs @@ -71,16 +71,7 @@ impl RelayWorker { .name(format!("relay-shard-{i}")) .spawn(move || { let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - Self::shard( - i, - rx, - state, - seed, - verify, - h, - opts, - http, - ); + Self::shard(i, rx, state, seed, verify, h, opts, http); })); let _ = exit_tx.send((i, res)); }) @@ -172,29 +163,17 @@ impl RelayWorker { } let process_message = process_started.elapsed(); - let mut batch = std::mem::replace(&mut ctx.batch, ctx.state.db.inner.batch()); - let stage_counts_started = StatsInstant::now(); - let reservation = ctx - .state - .db - .stage_count_deltas(&mut batch, &ctx.count_deltas); - let stage_counts = stage_counts_started.elapsed(); - - let stage_and_commit_started = StatsInstant::now(); - let staged = match ctx.sink.commit_batch(&state, batch) { + let batch = std::mem::replace(&mut ctx.batch, ctx.state.db.inner.batch()); + let count_deltas = std::mem::take(&mut ctx.count_deltas); + let txn = crate::db::Txn::from_parts(&ctx.state.db, batch, count_deltas); + let (staged, commit_timings) = match ctx.sink.commit_txn(&state, txn) { Ok(staged) => staged, Err(e) => { shard_stats.record_commit_error(); error!(shard = id, err = %e, "relay shard: failed to commit batch"); - drop(reservation); continue; } }; - let stage_and_commit = stage_and_commit_started.elapsed(); - let apply_counts_started = StatsInstant::now(); - ctx.state.db.apply_count_deltas(&ctx.count_deltas); - drop(reservation); - let apply_counts = apply_counts_started.elapsed(); let broadcast_started = StatsInstant::now(); ctx.sink.flush(&state, staged); @@ -206,9 +185,9 @@ impl RelayWorker { ctx.sink.advance_cursor(&state, &firehose, seq); shard_stats.record_processed(crate::ingest::firehose_stats::RelayShardTimings { process_message, - stage_counts, - stage_and_commit, - apply_counts, + stage_counts: commit_timings.stage_counts, + stage_and_commit: commit_timings.stage_and_commit, + apply_counts: commit_timings.apply_counts, broadcast, cursor: cursor_started.elapsed(), total: message_started.elapsed(), diff --git a/src/ingest/validation.rs b/src/ingest/validation.rs index 8fb2b2a..d737f16 100644 --- a/src/ingest/validation.rs +++ b/src/ingest/validation.rs @@ -1,10 +1,10 @@ use jacquard_common::IntoStatic; use jacquard_common::types::crypto::PublicKey; +use jacquard_repo::MemoryBlockStore; use jacquard_repo::Mst; use jacquard_repo::car::reader::{ParsedCar, parse_car_bytes}; use jacquard_repo::commit::Commit as AtpCommit; use jacquard_repo::mst::VerifiedWriteOp; -use jacquard_repo::MemoryBlockStore; use miette::IntoDiagnostic; use smol_str::ToSmolStr; use std::sync::Arc; diff --git a/src/jetstream.rs b/src/jetstream.rs index fb331e9..ed4ffed 100644 --- a/src/jetstream.rs +++ b/src/jetstream.rs @@ -72,7 +72,8 @@ fn next_time_us(db: &Db) -> i64 { let now = chrono::Utc::now().timestamp_micros(); let next = now.max(last.saturating_add(1)); if db - .jetstream.last_time_us + .jetstream + .last_time_us .compare_exchange(last, next, Ordering::SeqCst, Ordering::SeqCst) .is_ok() { diff --git a/src/types.rs b/src/types.rs index d6eec38..c53189a 100644 --- a/src/types.rs +++ b/src/types.rs @@ -64,7 +64,6 @@ impl Display for RepoStatus { } } - #[cfg(feature = "indexer")] mod indexer { use super::*; -- 2.51.2