feat(tranquil-store): metastore

Lewis: May this revision serve well! <lu5a@proton.me>
This commit is contained in:
Lewis
2026-04-10 13:57:43 +03:00
parent 9ea33def13
commit a5c68a3506
60 changed files with 16813 additions and 92 deletions
@@ -0,0 +1,22 @@
{
"db_name": "PostgreSQL",
"query": "SELECT cid FROM blocks ORDER BY created_at ASC LIMIT $1",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "cid",
"type_info": "Bytea"
}
],
"parameters": {
"Left": [
"Int8"
]
},
"nullable": [
false
]
},
"hash": "12f5864ebff622fc52643de7151a40e984082851741b22f63a170728e734763b"
}
@@ -0,0 +1,22 @@
{
"db_name": "PostgreSQL",
"query": "\n SELECT t.cid FROM UNNEST($1::bytea[]) AS t(cid)\n WHERE NOT EXISTS (\n SELECT 1 FROM user_blocks WHERE block_cid = t.cid\n )\n ",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "cid",
"type_info": "Bytea"
}
],
"parameters": {
"Left": [
"ByteaArray"
]
},
"nullable": [
null
]
},
"hash": "18fa821e4bd00ccf5d1d8395ba728e4905d69f9fe527b4d4b49c69deff52cea8"
}
@@ -0,0 +1,17 @@
{
"db_name": "PostgreSQL",
"query": "\n INSERT INTO backlinks (uri, path, link_to, repo_id)\n SELECT unnest($1::text[]), unnest($2::text[]), unnest($3::text[]), $4\n ON CONFLICT (uri, path) DO NOTHING\n ",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"TextArray",
"TextArray",
"Uuid"
]
},
"nullable": []
},
"hash": "47149c0577ad9e9b9b089820b0c93417769a4a37affe0e3972e324ec27ec532f"
}
@@ -0,0 +1,14 @@
{
"db_name": "PostgreSQL",
"query": "DELETE FROM blocks WHERE cid = ANY($1)",
"describe": {
"columns": [],
"parameters": {
"Left": [
"ByteaArray"
]
},
"nullable": []
},
"hash": "8eecf8fef308716be88815eb59bb67ec7c534b3c821d55481b110e3e462ee366"
}
@@ -0,0 +1,14 @@
{
"db_name": "PostgreSQL",
"query": "DELETE FROM backlinks WHERE uri = ANY($1::text[])",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray"
]
},
"nullable": []
},
"hash": "cffe4c37fe949fbdc3d5cd83ccec5655aae248a0a69dc260d1da9cf1d9ed2c49"
}
Generated
+13
View File
@@ -6389,6 +6389,12 @@ dependencies = [
"time",
]
[[package]]
name = "siphasher"
version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e"
[[package]]
name = "sketches-ddsketch"
version = "0.3.1"
@@ -7705,6 +7711,7 @@ dependencies = [
"tranquil-scopes",
"tranquil-signal",
"tranquil-storage",
"tranquil-store",
"tranquil-sync",
"tranquil-types",
"urlencoding",
@@ -7833,12 +7840,14 @@ dependencies = [
"bytes",
"chrono",
"cid",
"dashmap",
"fjall",
"flume 0.11.1",
"futures",
"jacquard-common",
"jacquard-repo",
"k256",
"lsm-tree",
"memmap2",
"multihash",
"parking_lot",
@@ -7849,14 +7858,18 @@ dependencies = [
"serde_ipld_dagcbor",
"serde_json",
"sha2",
"siphasher",
"smallvec",
"sqlx",
"tempfile",
"thiserror 2.0.18",
"tokio",
"tracing",
"tranquil-db",
"tranquil-db-traits",
"tranquil-repo",
"tranquil-types",
"uuid",
"xxhash-rust",
]
+4
View File
@@ -144,3 +144,7 @@ lto = "fat"
strip = true
codegen-units = 1
panic = "abort"
[profile.bench]
debug = 1
strip = false
+18 -8
View File
@@ -195,11 +195,22 @@ pub async fn update_subject_status(
ApiError::InternalError(Some("Failed to update deactivation status".into()))
})?;
}
let takedown_update = input.takedown.as_ref().map(|t| t.applied);
let takedown_ref = input.takedown.as_ref().and_then(|t| t.r#ref.as_deref());
let deactivated_update = input.deactivated.as_ref().map(|d| d.applied);
if (takedown_update.is_some() || deactivated_update.is_some())
&& let Err(e) = state
.repos
.repo
.update_repo_status(&did, takedown_update, takedown_ref, deactivated_update)
.await
{
warn!("failed to sync status to repo backend: {e:?}");
}
if let Some(takedown) = &input.takedown {
let status = if takedown.applied {
tranquil_db_traits::AccountStatus::Takendown
} else {
tranquil_db_traits::AccountStatus::Active
let status = match takedown.applied {
true => tranquil_db_traits::AccountStatus::Takendown,
false => tranquil_db_traits::AccountStatus::Active,
};
if let Err(e) =
tranquil_pds::repo_ops::sequence_account_event(&state, &did, status).await
@@ -208,10 +219,9 @@ pub async fn update_subject_status(
}
}
if let Some(deactivated) = &input.deactivated {
let status = if deactivated.applied {
tranquil_db_traits::AccountStatus::Deactivated
} else {
tranquil_db_traits::AccountStatus::Active
let status = match deactivated.applied {
true => tranquil_db_traits::AccountStatus::Deactivated,
false => tranquil_db_traits::AccountStatus::Active,
};
if let Err(e) =
tranquil_pds::repo_ops::sequence_account_event(&state, &did, status).await
+17
View File
@@ -18,6 +18,7 @@ use tranquil_pds::delegation::{
use tranquil_pds::rate_limit::{AccountCreationLimit, RateLimited};
use tranquil_pds::state::AppState;
use tranquil_pds::types::{Did, Handle};
use tranquil_types::CidLink;
pub async fn list_controllers(
State(state): State<AppState>,
@@ -419,6 +420,22 @@ pub async fn create_delegated_account(
}
};
state
.repos
.repo
.create_repo(
user_id,
&did,
&handle,
&CidLink::from(&repo.commit_cid),
&repo.repo_rev,
)
.await
.map_err(|e| {
error!("failed to register repo in backend: {e:?}");
ApiError::InternalError(None)
})?;
if let Some(validated) = validated_invite_code
&& let Err(e) = state
.repos
+18 -2
View File
@@ -15,6 +15,7 @@ use tranquil_pds::rate_limit::{AccountCreationLimit, RateLimited};
use tranquil_pds::state::AppState;
use tranquil_pds::types::{Did, Handle, PlainPassword};
use tranquil_pds::validation::validate_password;
use tranquil_types::CidLink;
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -160,7 +161,7 @@ async fn try_reactivate_migration(
.fetch_did_document(did)
.await
.ok()
.and_then(|f| Some((*f).clone())),
.map(|f| (*f).clone()),
access_jwt: access_meta.token,
refresh_jwt: refresh_meta.token,
verification_required,
@@ -547,6 +548,21 @@ pub async fn create_account(
}
};
let user_id = create_result.user_id;
if let Err(e) = state
.repos
.repo
.create_repo(
user_id,
&did_for_commit,
&handle_typed,
&CidLink::from(&repo.commit_cid),
&repo.repo_rev,
)
.await
{
error!("failed to register repo in backend: {e:?}");
return ApiError::InternalError(None).into_response();
}
if !is_migration && !is_did_web_byod {
super::provision::sequence_new_account(
&state,
@@ -607,7 +623,7 @@ pub async fn create_account(
Json(CreateAccountOutput {
handle: handle.clone().into(),
did: did_for_commit,
did_doc: did_doc.and_then(|f| Some((*f).clone())),
did_doc: did_doc.map(|f| (*f).clone()),
access_jwt: session.access_jwt,
refresh_jwt: session.refresh_jwt,
verification_required: !is_migration,
+10
View File
@@ -195,6 +195,16 @@ pub async fn import_repo(
}
let max_blocks = tranquil_config::get().import.max_blocks as usize;
let _write_lock = state.repo_write_locks.lock(user_id).await;
state
.block_store
.put_many(blocks.clone())
.await
.map_err(|e| {
error!("Failed to store import blocks: {:?}", e);
ApiError::InternalError(None)
})?;
match apply_import(
&state.repos.repo,
user_id,
+23 -1
View File
@@ -6,6 +6,7 @@ use jacquard_repo::{mst::Mst, storage::BlockStore};
use serde::{Deserialize, Serialize};
use serde_json::json;
use tracing::info;
use tranquil_db_traits::Backlink;
use tranquil_pds::api::error::{ApiError, DbResultExt};
use tranquil_pds::auth::{
Active, Auth, WriteOpKind, require_not_migrated, require_verified_or_delegated,
@@ -13,7 +14,8 @@ use tranquil_pds::auth::{
};
use tranquil_pds::repo::TrackingBlockStore;
use tranquil_pds::repo_ops::{
FinalizeParams, RecordOp, begin_repo_write, extract_blob_cids, finalize_repo_write,
FinalizeParams, RecordOp, begin_repo_write, extract_backlinks, extract_blob_cids,
finalize_repo_write,
};
use tranquil_pds::state::AppState;
use tranquil_pds::types::{AtIdentifier, AtUri, Did, Nsid, Rkey};
@@ -27,6 +29,8 @@ struct WriteAccumulator {
ops: Vec<RecordOp>,
modified_keys: Vec<String>,
all_blob_cids: Vec<String>,
backlinks_to_add: Vec<Backlink>,
backlinks_to_remove: Vec<AtUri>,
}
async fn process_single_write(
@@ -42,6 +46,8 @@ async fn process_single_write(
mut ops,
mut modified_keys,
mut all_blob_cids,
mut backlinks_to_add,
mut backlinks_to_remove,
} = acc;
match write {
@@ -79,6 +85,7 @@ async fn process_single_write(
.await
.map_err(|_| ApiError::InternalError(Some("Failed to add to MST".into())))?;
let uri = AtUri::from_parts(did, collection, &rkey);
backlinks_to_add.extend(extract_backlinks(&uri, value));
results.push(WriteResult::CreateResult {
uri,
cid: record_cid.to_string(),
@@ -95,6 +102,8 @@ async fn process_single_write(
ops,
modified_keys,
all_blob_cids,
backlinks_to_add,
backlinks_to_remove,
})
}
WriteOp::Update {
@@ -131,6 +140,8 @@ async fn process_single_write(
.await
.map_err(|_| ApiError::InternalError(Some("Failed to update MST".into())))?;
let uri = AtUri::from_parts(did, collection, rkey);
backlinks_to_remove.push(uri.clone());
backlinks_to_add.extend(extract_backlinks(&uri, value));
results.push(WriteResult::UpdateResult {
uri,
cid: record_cid.to_string(),
@@ -148,6 +159,8 @@ async fn process_single_write(
ops,
modified_keys,
all_blob_cids,
backlinks_to_add,
backlinks_to_remove,
})
}
WriteOp::Delete { collection, rkey } => {
@@ -158,6 +171,7 @@ async fn process_single_write(
.delete(&key)
.await
.map_err(|_| ApiError::InternalError(Some("Failed to delete from MST".into())))?;
backlinks_to_remove.push(AtUri::from_parts(did, collection, rkey));
results.push(WriteResult::DeleteResult {});
ops.push(RecordOp::Delete {
collection: collection.clone(),
@@ -170,6 +184,8 @@ async fn process_single_write(
ops,
modified_keys,
all_blob_cids,
backlinks_to_add,
backlinks_to_remove,
})
}
}
@@ -189,6 +205,8 @@ async fn process_writes(
ops: Vec::new(),
modified_keys: Vec::new(),
all_blob_cids: Vec::new(),
backlinks_to_add: Vec::new(),
backlinks_to_remove: Vec::new(),
};
stream::iter(writes.iter().map(Ok::<_, ApiError>))
.try_fold(initial_acc, |acc, write| async move {
@@ -319,6 +337,8 @@ pub async fn apply_writes(
ops,
modified_keys,
all_blob_cids,
backlinks_to_add,
backlinks_to_remove,
} = process_writes(
&input.writes,
mst,
@@ -373,6 +393,8 @@ pub async fn apply_writes(
ops,
modified_keys: &modified_keys,
blob_cids: &all_blob_cids,
backlinks_to_add,
backlinks_to_remove,
},
)
.await?;
+6 -10
View File
@@ -76,6 +76,7 @@ pub async fn delete_record(
};
let modified_keys = [key];
let deleted_uri = AtUri::from_parts(&did, &input.collection, &input.rkey);
let commit_result = finalize_repo_write(
&state,
@@ -95,20 +96,12 @@ pub async fn delete_record(
ops: vec![op],
modified_keys: &modified_keys,
blob_cids: &[],
backlinks_to_add: vec![],
backlinks_to_remove: vec![deleted_uri],
},
)
.await?;
let deleted_uri = AtUri::from_parts(&did, &input.collection, &input.rkey);
if let Err(e) = state
.repos
.backlink
.remove_backlinks_by_uri(&deleted_uri)
.await
{
error!("Failed to remove backlinks for {}: {}", deleted_uri, e);
}
Ok(Json(DeleteRecordOutput {
commit: Some(CommitInfo {
cid: commit_result.commit_cid.to_string(),
@@ -216,6 +209,7 @@ pub async fn delete_record_internal(
let written_cids_str: Vec<String> = written_cids.iter().map(ToString::to_string).collect();
let deleted_uri = AtUri::from_parts(did.as_str(), collection.as_str(), rkey.as_str());
commit_and_log(
state,
CommitParams {
@@ -228,6 +222,8 @@ pub async fn delete_record_internal(
blocks_cids: &written_cids_str,
blobs: &[],
obsolete_cids,
backlinks_to_add: vec![],
backlinks_to_remove: vec![deleted_uri],
},
)
.await?;
+1 -1
View File
@@ -178,7 +178,7 @@ pub async fn list_records(
};
let records: Vec<Value> = parsed_rows
.iter()
.zip(blocks.into_iter())
.zip(blocks)
.filter_map(|((_, rkey, cid_str), block_opt)| {
block_opt.and_then(|block| {
serde_ipld_dagcbor::from_slice::<Ipld>(&block)
+30 -33
View File
@@ -147,20 +147,28 @@ pub async fn create_record(
let prev_cid = match mst.get(&conflict_key).await {
Ok(Some(cid)) => cid,
_ => continue,
};
mst = match mst.delete(&conflict_key).await {
Ok(m) => m,
Ok(None) => continue,
Err(e) => {
error!(
"Failed to delete conflict from MST {}: {:?}",
"Failed to read conflict record from MST {}: {:?}",
conflict_uri, e
);
continue;
return Err(ApiError::InternalError(Some(
"Failed to read conflicting record from MST".into(),
)));
}
};
mst = mst.delete(&conflict_key).await.map_err(|e| {
error!(
"Failed to delete conflict from MST {}: {:?}",
conflict_uri, e
);
ApiError::InternalError(Some(
"Failed to delete conflicting record from MST".into(),
))
})?;
ops.push(RecordOp::Delete {
collection: conflict_collection,
rkey: conflict_rkey,
@@ -208,6 +216,9 @@ pub async fn create_record(
.collect();
let blob_cids = extract_blob_cids(&input.record);
let created_uri = AtUri::from_parts(&did, &input.collection, &rkey);
let backlinks_to_add = extract_backlinks(&created_uri, &input.record);
let commit_result = finalize_repo_write(
&state,
ctx,
@@ -226,35 +237,12 @@ pub async fn create_record(
ops,
modified_keys: &modified_keys,
blob_cids: &blob_cids,
backlinks_to_add,
backlinks_to_remove: conflict_uris_to_cleanup,
},
)
.await?;
{
let backlink_repo = state.repos.backlink.clone();
futures::future::join_all(conflict_uris_to_cleanup.iter().map(|uri| {
let backlink_repo = backlink_repo.clone();
async move {
if let Err(e) = backlink_repo.remove_backlinks_by_uri(uri).await {
error!("Failed to remove backlinks for {}: {}", uri, e);
}
}
}))
.await;
}
let created_uri = AtUri::from_parts(&did, &input.collection, &rkey);
let backlinks = extract_backlinks(&created_uri, &input.record);
if !backlinks.is_empty()
&& let Err(e) = state
.repos
.backlink
.add_backlinks(user_id, &backlinks)
.await
{
error!("Failed to add backlinks for {}: {}", created_uri, e);
}
Ok(Json(CreateRecordOutput {
uri: created_uri,
cid: record_cid.to_string(),
@@ -379,6 +367,13 @@ pub async fn put_record(
let modified_keys = [key];
let blob_cids = extract_blob_cids(&input.record);
let record_uri = AtUri::from_parts(&did, &input.collection, &input.rkey);
let backlinks_to_add = extract_backlinks(&record_uri, &input.record);
let backlinks_to_remove = match is_update {
true => vec![record_uri.clone()],
false => vec![],
};
let commit_result = finalize_repo_write(
&state,
ctx,
@@ -397,12 +392,14 @@ pub async fn put_record(
ops: vec![op],
modified_keys: &modified_keys,
blob_cids: &blob_cids,
backlinks_to_add,
backlinks_to_remove,
},
)
.await?;
Ok(Json(PutRecordOutput {
uri: AtUri::from_parts(&did, &input.collection, &input.rkey),
uri: record_uri,
cid: record_cid.to_string(),
commit: Some(CommitInfo {
cid: commit_result.commit_cid.to_string(),
@@ -382,6 +382,14 @@ pub async fn activate_account(
did
);
}
if let Err(e) = state
.repos
.repo
.update_repo_status(&did, None, None, Some(false))
.await
{
warn!("failed to sync activation to repo backend: {e:?}");
}
info!(
"[MIGRATION] activateAccount: Sequencing account event (active=true) for did={}",
did
@@ -514,6 +522,14 @@ pub async fn deactivate_account(
.delete(&tranquil_pds::cache_keys::handle_key(h))
.await;
}
if let Err(e) = state
.repos
.repo
.update_repo_status(&did, None, None, Some(true))
.await
{
warn!("failed to sync deactivation to repo backend: {e:?}");
}
if let Err(e) = tranquil_pds::repo_ops::sequence_account_event(
&state,
&did,
@@ -15,6 +15,7 @@ use tranquil_pds::rate_limit::{AccountCreationLimit, PasswordResetLimit, RateLim
use tranquil_pds::state::AppState;
use tranquil_pds::types::{Did, Handle, PlainPassword};
use tranquil_pds::validation::validate_password;
use tranquil_types::CidLink;
fn generate_setup_token() -> String {
let mut rng = rand::thread_rng();
@@ -366,6 +367,22 @@ pub async fn create_passkey_account(
};
let user_id = create_result.user_id;
state
.repos
.repo
.create_repo(
user_id,
&did_typed,
&handle_typed,
&CidLink::from(&repo.commit_cid),
&repo.repo_rev,
)
.await
.map_err(|e| {
error!("failed to register repo in backend: {e:?}");
ApiError::InternalError(None)
})?;
if !is_byod_did_web {
crate::identity::provision::sequence_new_account(
&state,
+2 -2
View File
@@ -153,7 +153,7 @@ pub async fn reset_password(
}
return Err(ApiError::ExpiredToken(None));
}
let password_hash = crate::common::hash_password_async(&password).await?;
let password_hash = crate::common::hash_password_async(password).await?;
let result = match state
.repos
.user
@@ -345,7 +345,7 @@ pub async fn set_password(
));
}
let new_hash = crate::common::hash_password_async(&new_password).await?;
let new_hash = crate::common::hash_password_async(new_password).await?;
state
.repos
+3 -3
View File
@@ -313,7 +313,7 @@ pub async fn create_session(
refresh_jwt: refresh_meta.token,
handle,
did: row.did,
did_doc: did_doc.ok().and_then(|f| Some((*f).clone())),
did_doc: did_doc.ok().map(|f| (*f).clone()),
email: row.email,
email_confirmed: Some(row.channel_verification.email),
email_auth_factor: email_auth_factor_out,
@@ -406,7 +406,7 @@ pub async fn get_session(
status: account_state.status_for_session().map(String::from),
migrated_to_pds,
migrated_at,
did_doc: did_doc.ok().and_then(|f| Some((*f).clone())),
did_doc: did_doc.ok().map(|f| (*f).clone()),
}))
}
Ok(None) => Err(ApiError::AuthenticationFailed(None)),
@@ -604,7 +604,7 @@ pub async fn refresh_session(
preferred_locale: u.preferred_locale,
is_admin: u.is_admin,
active: account_state.is_active(),
did_doc: did_doc.ok().and_then(|f| Some((*f).clone())),
did_doc: did_doc.ok().map(|f| (*f).clone()),
status: account_state.status_for_session().map(String::from),
}))
}
+83
View File
@@ -143,6 +143,9 @@ pub struct TranquilConfig {
#[config(nested)]
pub scheduled: ScheduledConfig,
#[config(nested)]
pub tranquil_store: TranquilStoreConfig,
}
impl TranquilConfig {
@@ -250,6 +253,23 @@ impl TranquilConfig {
);
}
// -- repo backend -----------------------------------------------------
if let Err(e) = self.storage.repo_backend.parse::<RepoBackend>() {
errors.push(e);
}
// -- tranquil-store ---------------------------------------------------
if let Some(mb) = self.tranquil_store.memory_budget_mb
&& mb == 0
{
errors.push("tranquil_store.memory_budget_mb must be at least 1".to_string());
}
if let Some(threads) = self.tranquil_store.handler_threads
&& threads == 0
{
errors.push("tranquil_store.handler_threads must be at least 1".to_string());
}
// -- cache ------------------------------------------------------------
match self.cache.backend.as_str() {
"valkey" => {
@@ -561,6 +581,35 @@ impl SecretsConfig {
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RepoBackend {
Postgres,
TranquilStore,
}
impl std::str::FromStr for RepoBackend {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"postgres" => Ok(Self::Postgres),
"tranquil-store" => Ok(Self::TranquilStore),
other => Err(format!(
"unknown repo backend \"{other}\", expected \"postgres\" or \"tranquil-store\""
)),
}
}
}
impl fmt::Display for RepoBackend {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Postgres => f.write_str("postgres"),
Self::TranquilStore => f.write_str("tranquil-store"),
}
}
}
#[derive(Debug, Config)]
pub struct StorageConfig {
/// Storage backend: `filesystem` or `s3`.
@@ -578,6 +627,17 @@ pub struct StorageConfig {
/// Custom S3 endpoint URL (for MinIO, R2, etc.).
#[config(env = "S3_ENDPOINT")]
pub s3_endpoint: Option<String>,
#[config(env = "REPO_BACKEND", default = "postgres")]
pub repo_backend: String,
}
impl StorageConfig {
pub fn repo_backend(&self) -> RepoBackend {
self.repo_backend
.parse()
.expect("repo_backend must be validated before use")
}
}
#[derive(Debug, Config)]
@@ -996,6 +1056,29 @@ pub struct ScheduledConfig {
/// Interval in seconds between scheduled delete checks.
#[config(env = "SCHEDULED_DELETE_CHECK_INTERVAL_SECS", default = 3600)]
pub delete_check_interval_secs: u64,
/// Interval in seconds between block garbage collection cycles.
#[config(env = "BLOCK_GC_INTERVAL_SECS", default = 21600)]
pub block_gc_interval_secs: u64,
}
#[derive(Debug, Config)]
pub struct TranquilStoreConfig {
/// Directory for tranquil-store data (metastore, eventlog).
#[config(
env = "TRANQUIL_STORE_DATA_DIR",
default = "/var/lib/tranquil-pds/store"
)]
pub data_dir: String,
/// Fjall block cache size in megabytes. Defaults to 20% of system RAM
/// when unset.
#[config(env = "TRANQUIL_STORE_MEMORY_BUDGET_MB")]
pub memory_budget_mb: Option<u64>,
/// Number of handler threads. Defaults to available_parallelism / 2.
#[config(env = "TRANQUIL_STORE_HANDLER_THREADS")]
pub handler_threads: Option<usize>,
}
/// Generate a TOML configuration template with all available options,
+18
View File
@@ -5,6 +5,7 @@ use tranquil_types::{AtUri, CidLink, Did, Handle, Nsid, Rkey};
use uuid::Uuid;
use crate::DbError;
use crate::backlink::Backlink;
use crate::sequence::SequenceNumber;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, sqlx::Type)]
@@ -279,6 +280,8 @@ pub struct ApplyCommitInput {
pub obsolete_block_cids: Vec<Vec<u8>>,
pub record_upserts: Vec<RecordUpsert>,
pub record_deletes: Vec<RecordDelete>,
pub backlinks_to_add: Vec<Backlink>,
pub backlinks_to_remove: Vec<AtUri>,
pub commit_event: CommitEventData,
}
@@ -300,6 +303,8 @@ pub trait RepoRepository: Send + Sync {
async fn create_repo(
&self,
user_id: Uuid,
did: &Did,
handle: &Handle,
repo_root_cid: &CidLink,
repo_rev: &str,
) -> Result<(), DbError>;
@@ -313,6 +318,14 @@ pub trait RepoRepository: Send + Sync {
async fn update_repo_rev(&self, user_id: Uuid, repo_rev: &str) -> Result<(), DbError>;
async fn update_repo_status(
&self,
did: &Did,
takedown: Option<bool>,
takedown_ref: Option<&str>,
deactivated: Option<bool>,
) -> Result<(), DbError>;
async fn delete_repo(&self, user_id: Uuid) -> Result<(), DbError>;
async fn get_repo_root_for_update(&self, user_id: Uuid) -> Result<Option<CidLink>, DbError>;
@@ -400,6 +413,11 @@ pub trait RepoRepository: Send + Sync {
async fn count_user_blocks(&self, user_id: Uuid) -> Result<i64, DbError>;
async fn find_unreferenced_blocks(
&self,
candidate_cids: &[Vec<u8>],
) -> Result<Vec<Vec<u8>>, DbError>;
async fn insert_commit_event(&self, data: &CommitEventData) -> Result<SequenceNumber, DbError>;
async fn insert_identity_event(
@@ -23,6 +23,10 @@ impl SequenceNumber {
pub fn is_valid(&self) -> bool {
self.0 >= 0
}
pub fn as_u64(&self) -> Option<u64> {
u64::try_from(self.0).ok()
}
}
impl fmt::Display for SequenceNumber {
+83
View File
@@ -46,9 +46,21 @@ impl PostgresRepoRepository {
#[async_trait]
impl RepoRepository for PostgresRepoRepository {
async fn update_repo_status(
&self,
_did: &Did,
_takedown: Option<bool>,
_takedown_ref: Option<&str>,
_deactivated: Option<bool>,
) -> Result<(), DbError> {
Ok(())
}
async fn create_repo(
&self,
user_id: Uuid,
_did: &Did,
_handle: &Handle,
repo_root_cid: &CidLink,
repo_rev: &str,
) -> Result<(), DbError> {
@@ -606,6 +618,30 @@ impl RepoRepository for PostgresRepoRepository {
Ok(count)
}
async fn find_unreferenced_blocks(
&self,
candidate_cids: &[Vec<u8>],
) -> Result<Vec<Vec<u8>>, DbError> {
match candidate_cids.is_empty() {
true => Ok(Vec::new()),
false => {
let rows = sqlx::query!(
r#"
SELECT t.cid FROM UNNEST($1::bytea[]) AS t(cid)
WHERE NOT EXISTS (
SELECT 1 FROM user_blocks WHERE block_cid = t.cid
)
"#,
candidate_cids,
)
.fetch_all(&self.pool)
.await
.map_err(map_sqlx_error)?;
Ok(rows.into_iter().filter_map(|r| r.cid).collect())
}
}
}
async fn get_user_block_cids_since_rev(
&self,
user_id: Uuid,
@@ -1367,6 +1403,53 @@ impl RepoRepository for PostgresRepoRepository {
.map_err(|e| ApplyCommitError::Database(e.to_string()))?;
}
if !input.backlinks_to_remove.is_empty() {
let remove_uris: Vec<&str> = input
.backlinks_to_remove
.iter()
.map(|u| u.as_str())
.collect();
sqlx::query!(
"DELETE FROM backlinks WHERE uri = ANY($1::text[])",
&remove_uris as &[&str],
)
.execute(&mut *tx)
.await
.map_err(|e| ApplyCommitError::Database(e.to_string()))?;
}
if !input.backlinks_to_add.is_empty() {
let uris: Vec<&str> = input
.backlinks_to_add
.iter()
.map(|b| b.uri.as_str())
.collect();
let paths: Vec<&str> = input
.backlinks_to_add
.iter()
.map(|b| b.path.as_str())
.collect();
let link_tos: Vec<&str> = input
.backlinks_to_add
.iter()
.map(|b| b.link_to.as_str())
.collect();
sqlx::query!(
r#"
INSERT INTO backlinks (uri, path, link_to, repo_id)
SELECT unnest($1::text[]), unnest($2::text[]), unnest($3::text[]), $4
ON CONFLICT (uri, path) DO NOTHING
"#,
&uris as &[&str],
&paths as &[&str],
&link_tos as &[&str],
input.user_id,
)
.execute(&mut *tx)
.await
.map_err(|e| ApplyCommitError::Database(e.to_string()))?;
}
let event = &input.commit_event;
let seq: i64 = sqlx::query_scalar(
r#"
+1
View File
@@ -18,6 +18,7 @@ tranquil-comms = { workspace = true }
tranquil-signal = { workspace = true }
tranquil-db = { workspace = true }
tranquil-db-traits = { workspace = true }
tranquil-store = { workspace = true }
tranquil-lexicon = { workspace = true, features = ["resolve"] }
aes-gcm = { workspace = true }
+4 -5
View File
@@ -49,11 +49,10 @@ pub async fn resolve_identity(
None
}
});
let handle = did_doc.also_known_as.iter().find_map(|alias| {
alias
.strip_prefix("at://")
.and_then(|s| Some(s.to_string()))
});
let handle = did_doc
.also_known_as
.iter()
.find_map(|alias| alias.strip_prefix("at://").map(|s| s.to_string()));
Ok(ResolvedIdentity {
did: did.clone(),
+6 -4
View File
@@ -4,7 +4,7 @@ use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
use tracing::{debug, info, warn};
#[derive(Debug, thiserror::Error)]
pub enum DidResolutionError {
@@ -55,10 +55,12 @@ pub struct ResolvedService {
pub service_id: String,
}
type TimedCache<T> = RwLock<HashMap<Box<str>, (Instant, Arc<T>)>>;
pub struct DidResolver {
did_doc_cache: RwLock<HashMap<Box<str>, (Instant, Arc<serde_json::Value>)>>,
parsed_did_doc_cache: RwLock<HashMap<Box<str>, (Instant, Arc<DidDocument>)>>,
service_cache: RwLock<HashMap<Box<str>, (Instant, Arc<ResolvedService>)>>,
did_doc_cache: TimedCache<serde_json::Value>,
parsed_did_doc_cache: TimedCache<DidDocument>,
service_cache: TimedCache<ResolvedService>,
client: Client,
cache_ttl: Duration,
plc_directory_url: String,
+15
View File
@@ -159,6 +159,8 @@ pub struct FinalizeParams<'a> {
pub ops: Vec<RecordOp>,
pub modified_keys: &'a [String],
pub blob_cids: &'a [String],
pub backlinks_to_add: Vec<Backlink>,
pub backlinks_to_remove: Vec<AtUri>,
}
pub async fn begin_repo_write(
@@ -249,6 +251,8 @@ pub async fn finalize_repo_write(
blocks_cids: &written_cids_str,
blobs: params.blob_cids,
obsolete_cids: vec![ctx.current_root_cid],
backlinks_to_add: params.backlinks_to_add,
backlinks_to_remove: params.backlinks_to_remove,
},
)
.await?;
@@ -331,6 +335,8 @@ pub struct CommitParams<'a> {
pub blocks_cids: &'a [String],
pub blobs: &'a [String],
pub obsolete_cids: Vec<Cid>,
pub backlinks_to_add: Vec<Backlink>,
pub backlinks_to_remove: Vec<AtUri>,
}
pub async fn commit_and_log(
@@ -342,6 +348,8 @@ pub async fn commit_and_log(
RepoEventType,
};
let backlinks_to_add = params.backlinks_to_add;
let backlinks_to_remove = params.backlinks_to_remove;
let CommitParams {
did,
user_id,
@@ -352,6 +360,7 @@ pub async fn commit_and_log(
blocks_cids,
blobs,
obsolete_cids,
..
} = params;
let key_row = state
.repos
@@ -485,6 +494,8 @@ pub async fn commit_and_log(
obsolete_block_cids: obsolete_bytes,
record_upserts,
record_deletes,
backlinks_to_add,
backlinks_to_remove,
commit_event,
};
@@ -592,6 +603,8 @@ pub async fn create_record_internal(
.collect();
let written_cids_str: Vec<String> = written_cids.iter().map(|c| c.to_string()).collect();
let blob_cids = extract_blob_cids(record);
let record_uri = AtUri::from_parts(did.as_str(), collection.as_str(), rkey.as_str());
let backlinks = extract_backlinks(&record_uri, record);
let result = commit_and_log(
state,
CommitParams {
@@ -604,6 +617,8 @@ pub async fn create_record_internal(
blocks_cids: &written_cids_str,
blobs: &blob_cids,
obsolete_cids,
backlinks_to_add: backlinks,
backlinks_to_remove: vec![],
},
)
.await?;
+62 -2
View File
@@ -436,19 +436,26 @@ pub async fn start_scheduled_tasks(
blob_repo: Arc<dyn BlobRepository>,
blob_store: Arc<dyn BlobStorage>,
sso_repo: Arc<dyn SsoRepository>,
repo_repo: Arc<dyn RepoRepository>,
block_store: PostgresBlockStore,
shutdown: CancellationToken,
) {
let check_interval =
Duration::from_secs(tranquil_config::get().scheduled.delete_check_interval_secs);
let cfg = tranquil_config::get();
let check_interval = Duration::from_secs(cfg.scheduled.delete_check_interval_secs);
let gc_interval = Duration::from_secs(cfg.scheduled.block_gc_interval_secs);
info!(
check_interval_secs = check_interval.as_secs(),
gc_interval_secs = gc_interval.as_secs(),
"Starting scheduled tasks service"
);
let mut ticker = interval(check_interval);
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
let mut gc_ticker = interval(gc_interval);
gc_ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tokio::select! {
_ = shutdown.cancelled() => {
@@ -494,10 +501,63 @@ pub async fn start_scheduled_tasks(
}
}
}
_ = gc_ticker.tick() => {
if let Err(e) = run_block_gc(repo_repo.as_ref(), &block_store).await {
error!("Block GC error: {e}");
}
}
}
}
}
const BLOCK_GC_BATCH_SIZE: i64 = 1000;
async fn run_block_gc(
repo_repo: &dyn RepoRepository,
block_store: &PostgresBlockStore,
) -> anyhow::Result<()> {
let mut total_deleted: u64 = 0;
loop {
let candidates = block_store
.get_oldest_block_cids(BLOCK_GC_BATCH_SIZE)
.await
.context("failed to fetch candidate blocks")?;
match candidates.is_empty() {
true => break,
false => {
let batch_len = candidates.len();
let unreferenced = repo_repo
.find_unreferenced_blocks(&candidates)
.await
.context("failed to check block references")?;
let deleted = match unreferenced.is_empty() {
true => 0,
false => block_store
.delete_blocks(&unreferenced)
.await
.context("failed to delete unreferenced blocks")?,
};
total_deleted = total_deleted.saturating_add(deleted);
match unreferenced.len() == batch_len {
true => continue,
false => break,
}
}
}
}
match total_deleted > 0 {
true => info!(total_deleted, "Block GC cycle complete"),
false => debug!("Block GC cycle: no orphaned blocks found"),
}
Ok(())
}
async fn process_scheduled_deletions(
user_repo: &dyn UserRepository,
blob_repo: &dyn BlobRepository,
+75 -1
View File
@@ -12,6 +12,7 @@ use crate::sso::{SsoConfig, SsoManager};
use crate::storage::{BlobStorage, create_blob_storage};
use sqlx::PgPool;
use std::error::Error;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::sync::broadcast;
@@ -262,7 +263,14 @@ impl AppState {
AuthConfig::init();
init_rate_limit_override();
let repos = Arc::new(PostgresRepositories::new(db.clone()));
let mut repos = PostgresRepositories::new(db.clone());
let cfg = tranquil_config::get();
if cfg.storage.repo_backend() == tranquil_config::RepoBackend::TranquilStore {
wire_tranquil_store(&mut repos, &cfg.tranquil_store);
}
let repos = Arc::new(repos);
let block_store = PostgresBlockStore::new(db);
let blob_store = create_blob_storage().await;
@@ -377,3 +385,69 @@ impl AppState {
true
}
}
fn wire_tranquil_store(
repos: &mut PostgresRepositories,
store_cfg: &tranquil_config::TranquilStoreConfig,
) {
use tranquil_store::RealIO;
use tranquil_store::eventlog::{EventLog, EventLogBridge, EventLogConfig};
use tranquil_store::metastore::client::MetastoreClient;
use tranquil_store::metastore::handler::HandlerPool;
use tranquil_store::metastore::partitions::Partition;
use tranquil_store::metastore::{Metastore, MetastoreConfig};
let base_dir = PathBuf::from(&store_cfg.data_dir);
let data_dir = match std::env::var("TRANQUIL_PDS_TEST_INFRA_READY").as_deref() {
Ok("1") => base_dir.join(format!("pid-{}", std::process::id())),
_ => base_dir,
};
let metastore_dir = data_dir.join("metastore");
let segments_dir = data_dir.join("eventlog").join("segments");
std::fs::create_dir_all(&metastore_dir).expect("failed to create metastore directory");
std::fs::create_dir_all(&segments_dir).expect("failed to create eventlog segments directory");
let metastore_config = store_cfg
.memory_budget_mb
.map(|mb| MetastoreConfig {
cache_size_bytes: mb.saturating_mul(1024 * 1024),
})
.unwrap_or_default();
let metastore =
Metastore::open(&metastore_dir, metastore_config).expect("failed to open metastore");
let event_log = EventLog::open(
EventLogConfig {
segments_dir,
..EventLogConfig::default()
},
RealIO::new(),
)
.expect("failed to open eventlog");
let event_log = Arc::new(event_log);
let bridge = Arc::new(EventLogBridge::new(Arc::clone(&event_log)));
let indexes = metastore.partition(Partition::Indexes).clone();
let event_ops = metastore.event_ops(Arc::clone(&bridge));
let recovered = event_ops
.recover_metastore_mutations(&indexes)
.expect("metastore crash recovery failed");
if recovered > 0 {
tracing::info!(recovered, "replayed metastore mutations from eventlog");
}
let notifier = bridge.notifier();
let pool = HandlerPool::spawn::<RealIO>(metastore, bridge, None, store_cfg.handler_threads);
let client = MetastoreClient::<RealIO>::new(Arc::new(pool));
tracing::info!(data_dir = %store_cfg.data_dir, "tranquil-store data directory");
repos.repo = Arc::new(client.clone());
repos.backlink = Arc::new(client);
repos.event_notifier = Arc::new(notifier);
}
+44 -20
View File
@@ -302,30 +302,54 @@ pub async fn create_repost(
#[allow(dead_code)]
pub async fn set_account_takedown(did: &str, takedown_ref: Option<&str>) {
let pool = get_test_db_pool().await;
sqlx::query!(
"UPDATE users SET takedown_ref = $1 WHERE did = $2",
takedown_ref,
did
)
.execute(pool)
.await
.expect("Failed to update takedown_ref");
let client = client();
let (admin_jwt, _) = create_admin_account_and_login(&client).await;
let applied = takedown_ref.is_some();
let res = client
.post(format!(
"{}/xrpc/com.atproto.admin.updateSubjectStatus",
base_url().await,
))
.bearer_auth(&admin_jwt)
.json(&json!({
"subject": {
"$type": "com.atproto.admin.defs#repoRef",
"did": did
},
"takedown": {
"applied": applied,
"ref": takedown_ref
}
}))
.send()
.await
.expect("Failed to send takedown request");
assert_eq!(res.status(), StatusCode::OK, "Failed to set takedown");
}
#[allow(dead_code)]
pub async fn set_account_deactivated(did: &str, deactivated: bool) {
let pool = get_test_db_pool().await;
let deactivated_at: Option<chrono::DateTime<Utc>> =
if deactivated { Some(Utc::now()) } else { None };
sqlx::query!(
"UPDATE users SET deactivated_at = $1 WHERE did = $2",
deactivated_at,
did
)
.execute(pool)
.await
.expect("Failed to update deactivated_at");
let client = client();
let (admin_jwt, _) = create_admin_account_and_login(&client).await;
let res = client
.post(format!(
"{}/xrpc/com.atproto.admin.updateSubjectStatus",
base_url().await,
))
.bearer_auth(&admin_jwt)
.json(&json!({
"subject": {
"$type": "com.atproto.admin.defs#repoRef",
"did": did
},
"deactivated": {
"applied": deactivated
}
}))
.send()
.await
.expect("Failed to send deactivation request");
assert_eq!(res.status(), StatusCode::OK, "Failed to set deactivation");
}
#[allow(dead_code)]
+26
View File
@@ -24,6 +24,32 @@ impl PostgresBlockStore {
}
}
impl PostgresBlockStore {
pub async fn get_oldest_block_cids(&self, limit: i64) -> Result<Vec<Vec<u8>>, RepoError> {
let rows = sqlx::query!(
"SELECT cid FROM blocks ORDER BY created_at ASC LIMIT $1",
limit,
)
.fetch_all(&self.pool)
.await
.map_err(RepoError::storage)?;
Ok(rows.into_iter().map(|r| r.cid).collect())
}
pub async fn delete_blocks(&self, cids: &[Vec<u8>]) -> Result<u64, RepoError> {
match cids.is_empty() {
true => Ok(0),
false => {
let result = sqlx::query!("DELETE FROM blocks WHERE cid = ANY($1)", cids,)
.execute(&self.pool)
.await
.map_err(RepoError::storage)?;
Ok(result.rows_affected())
}
}
}
}
impl BlockStore for PostgresBlockStore {
async fn get(&self, cid: &Cid) -> Result<Option<Bytes>, RepoError> {
let cid_bytes = cid.to_bytes();
+2
View File
@@ -253,6 +253,8 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
state.repos.blob.clone(),
state.blob_store.clone(),
state.repos.sso.clone(),
state.repos.repo.clone(),
state.block_store.clone(),
shutdown.clone(),
));
+5
View File
@@ -12,6 +12,7 @@ serde = { workspace = true }
postcard = { version = "1", features = ["alloc"] }
parking_lot = { workspace = true }
fjall = "3"
lsm-tree = "3"
flume = "0.11"
tokio = { workspace = true, features = ["sync", "rt"] }
bytes = "1"
@@ -26,6 +27,10 @@ jacquard-repo = { workspace = true }
cid = { workspace = true }
multihash = { workspace = true }
sha2 = { workspace = true }
siphasher = "1"
dashmap = "6"
smallvec = "1"
uuid = { workspace = true }
[features]
test-harness = []
+812
View File
@@ -0,0 +1,812 @@
use std::sync::Arc;
use std::time::{Duration, Instant};
use futures::StreamExt;
use tokio::sync::oneshot;
use tranquil_db_traits::{
ApplyCommitInput, CommitEventData, RecordUpsert, RepoEventType, RepoRepository,
};
use tranquil_types::{CidLink, Did, Handle, Nsid, Rkey};
use uuid::Uuid;
use tranquil_store::RealIO;
use tranquil_store::eventlog::{EventLog, EventLogConfig};
use tranquil_store::metastore::handler::{
CommitRequest, HandlerPool, MetastoreRequest, RecordRequest, RepoRequest,
};
use tranquil_store::metastore::{Metastore, MetastoreConfig};
struct LatencyStats {
p50: Duration,
p95: Duration,
p99: Duration,
max: Duration,
mean: Duration,
}
fn compute_stats(durations: &mut [Duration]) -> Option<LatencyStats> {
match durations.is_empty() {
true => None,
false => {
durations.sort();
let len = durations.len();
let sum: Duration = durations.iter().sum();
let divisor = u32::try_from(len).unwrap_or(u32::MAX);
let last = len - 1;
Some(LatencyStats {
p50: durations[last * 50 / 100],
p95: durations[last * 95 / 100],
p99: durations[last * 99 / 100],
max: durations[last],
mean: sum / divisor,
})
}
}
}
fn print_result(ops: usize, elapsed: Duration, stats: Option<&LatencyStats>) {
let throughput = ops as f64 / elapsed.as_secs_f64();
match stats {
Some(s) => println!(
"{throughput:.0} ops/sec, {:.1}ms | p50={:?} p95={:?} p99={:?} max={:?} mean={:?}",
elapsed.as_secs_f64() * 1000.0,
s.p50,
s.p95,
s.p99,
s.max,
s.mean
),
None => println!(
"{throughput:.0} ops/sec, {:.1}ms",
elapsed.as_secs_f64() * 1000.0,
),
}
}
async fn collect_latencies(handles: Vec<tokio::task::JoinHandle<Vec<Duration>>>) -> Vec<Duration> {
futures::stream::iter(handles)
.fold(Vec::new(), |mut acc, h| async move {
acc.extend(h.await.unwrap());
acc
})
.await
}
fn test_cid(seed: u8) -> CidLink {
let digest: [u8; 32] = std::array::from_fn(|i| seed.wrapping_add(i as u8));
let mh = multihash::Multihash::<64>::wrap(0x12, &digest).unwrap();
let c = cid::Cid::new_v1(0x71, mh);
CidLink::from_cid(&c)
}
fn test_cid_bytes(seed: u8) -> Vec<u8> {
let digest: [u8; 32] = std::array::from_fn(|i| seed.wrapping_add(i as u8));
let mh = multihash::Multihash::<64>::wrap(0x12, &digest).unwrap();
cid::Cid::new_v1(0x71, mh).to_bytes()
}
fn make_rev(n: u64) -> String {
format!("rev{n:010}")
}
struct BenchHarness {
pool: Arc<HandlerPool>,
_metastore_dir: tempfile::TempDir,
_eventlog_dir: tempfile::TempDir,
}
fn setup(thread_count: usize) -> BenchHarness {
let metastore_dir = tempfile::TempDir::new().unwrap();
let eventlog_dir = tempfile::TempDir::new().unwrap();
let segments_dir = eventlog_dir.path().join("segments");
std::fs::create_dir_all(&segments_dir).unwrap();
let metastore = Metastore::open(
metastore_dir.path(),
MetastoreConfig {
cache_size_bytes: 256 * 1024 * 1024,
},
)
.unwrap();
let event_log = EventLog::open(
EventLogConfig {
segments_dir,
..EventLogConfig::default()
},
RealIO::new(),
)
.unwrap();
let bridge = Arc::new(tranquil_store::eventlog::EventLogBridge::new(Arc::new(
event_log,
)));
let pool = Arc::new(HandlerPool::spawn::<RealIO>(
metastore,
bridge,
None,
Some(thread_count),
));
BenchHarness {
pool,
_metastore_dir: metastore_dir,
_eventlog_dir: eventlog_dir,
}
}
async fn create_user(pool: &HandlerPool, user_id: Uuid, did: &Did, cid: &CidLink) {
let (tx, rx) = oneshot::channel();
pool.send(MetastoreRequest::Repo(RepoRequest::CreateRepoFull {
user_id,
did: did.clone(),
handle: Handle::from(format!("bench.{}.invalid", user_id.as_simple())),
repo_root_cid: cid.clone(),
repo_rev: "rev0000000000".to_string(),
tx,
}))
.unwrap();
rx.await.unwrap().unwrap();
}
fn make_commit_input(
user_id: Uuid,
did: &Did,
collection: &Nsid,
rev_n: u64,
cid_seed: u8,
) -> ApplyCommitInput {
ApplyCommitInput {
user_id,
did: did.clone(),
expected_root_cid: None,
new_root_cid: test_cid(cid_seed),
new_rev: make_rev(rev_n),
new_block_cids: vec![test_cid_bytes(cid_seed)],
obsolete_block_cids: vec![],
record_upserts: vec![RecordUpsert {
collection: collection.clone(),
rkey: Rkey::from(format!("r{rev_n:010}")),
cid: test_cid(cid_seed),
}],
record_deletes: vec![],
backlinks_to_add: vec![],
backlinks_to_remove: vec![],
commit_event: CommitEventData {
did: did.clone(),
event_type: RepoEventType::Commit,
commit_cid: Some(test_cid(cid_seed)),
prev_cid: None,
ops: None,
blobs: None,
blocks_cids: None,
prev_data_cid: None,
rev: Some(make_rev(rev_n)),
},
}
}
async fn seed_records(
pool: &HandlerPool,
user_id: Uuid,
did: &Did,
collection: &Nsid,
count: usize,
) {
let batches: Vec<(usize, usize, u64, u8)> = (0..)
.map(|i| {
let start = i * 50;
let end = (start + 50).min(count);
let rev_n = (i as u64) + 1;
let cid_seed = ((i + 10) & 0xFF) as u8;
(start, end, rev_n, cid_seed)
})
.take_while(|(start, _, _, _)| *start < count)
.collect();
futures::stream::iter(batches)
.fold((), |(), (batch_start, batch_end, rev_n, cid_seed)| {
let did = did.clone();
let collection = collection.clone();
async move {
let record_upserts: Vec<RecordUpsert> = (batch_start..batch_end)
.map(|i| RecordUpsert {
collection: collection.clone(),
rkey: Rkey::from(format!("rec{i:08}")),
cid: test_cid(((i * 7 + 3) & 0xFF) as u8),
})
.collect();
let new_block_cids: Vec<Vec<u8>> = (batch_start..batch_end)
.map(|i| test_cid_bytes(((i * 11 + 5) & 0xFF) as u8))
.collect();
let input = ApplyCommitInput {
user_id,
did: did.clone(),
expected_root_cid: None,
new_root_cid: test_cid(cid_seed),
new_rev: make_rev(rev_n),
new_block_cids,
obsolete_block_cids: vec![],
record_upserts,
record_deletes: vec![],
backlinks_to_add: vec![],
backlinks_to_remove: vec![],
commit_event: CommitEventData {
did,
event_type: RepoEventType::Commit,
commit_cid: Some(test_cid(cid_seed)),
prev_cid: None,
ops: None,
blobs: None,
blocks_cids: None,
prev_data_cid: None,
rev: Some(make_rev(rev_n)),
},
};
let (tx, rx) = oneshot::channel();
pool.send(MetastoreRequest::Commit(Box::new(
CommitRequest::ApplyCommit {
input: Box::new(input),
tx,
},
)))
.unwrap();
rx.await.unwrap().unwrap();
}
})
.await;
}
async fn bench_apply_commit(pool: &Arc<HandlerPool>, concurrency: usize, ops_per_task: usize) {
let user_ids: Vec<Uuid> = (0..concurrency).map(|_| Uuid::new_v4()).collect();
let dids: Vec<Did> = user_ids
.iter()
.map(|u| Did::from(format!("did:plc:bench{}", u.as_simple())))
.collect();
futures::stream::iter(user_ids.iter().zip(dids.iter()))
.fold((), |(), (uid, did)| async {
create_user(pool, *uid, did, &test_cid(1)).await;
})
.await;
let collection = Nsid::from("app.bsky.feed.post".to_string());
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let pool = Arc::clone(pool);
let user_id = user_ids[task_id];
let did = dids[task_id].clone();
let collection = collection.clone();
tokio::spawn(async move {
futures::stream::iter(0..ops_per_task)
.fold(Vec::with_capacity(ops_per_task), |mut latencies, i| {
let pool = &pool;
let did = &did;
let collection = &collection;
async move {
let rev_n = (task_id * ops_per_task + i + 1) as u64;
let cid_seed = ((task_id * 31 + i * 7) & 0xFF) as u8;
let input =
make_commit_input(user_id, did, collection, rev_n, cid_seed);
let t = Instant::now();
let (tx, rx) = oneshot::channel();
pool.send(MetastoreRequest::Commit(Box::new(
CommitRequest::ApplyCommit {
input: Box::new(input),
tx,
},
)))
.unwrap();
rx.await.unwrap().unwrap();
latencies.push(t.elapsed());
latencies
}
})
.await
})
})
.collect();
let mut all_latencies = collect_latencies(handles).await;
let elapsed = start.elapsed();
let total_ops = concurrency * ops_per_task;
let stats = compute_stats(&mut all_latencies);
print_result(total_ops, elapsed, stats.as_ref());
}
async fn bench_get_record_cid(pool: &Arc<HandlerPool>, concurrency: usize, ops_per_task: usize) {
let user_id = Uuid::new_v4();
let did = Did::from(format!("did:plc:getrecord{}", user_id.as_simple()));
let collection = Nsid::from("app.bsky.feed.post".to_string());
create_user(pool, user_id, &did, &test_cid(1)).await;
seed_records(pool, user_id, &did, &collection, 1000).await;
let total_records = 1000usize;
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let pool = Arc::clone(pool);
let collection = collection.clone();
tokio::spawn(async move {
futures::stream::iter(0..ops_per_task)
.fold(Vec::with_capacity(ops_per_task), |mut latencies, i| {
let pool = &pool;
let collection = &collection;
async move {
let rec_idx = (task_id * 7 + i * 13) % total_records;
let rkey = Rkey::from(format!("rec{rec_idx:08}"));
let t = Instant::now();
let (tx, rx) = oneshot::channel();
pool.send(MetastoreRequest::Record(RecordRequest::GetRecordCid {
repo_id: user_id,
collection: collection.clone(),
rkey,
tx,
}))
.unwrap();
let result = rx.await.unwrap().unwrap();
assert!(result.is_some());
latencies.push(t.elapsed());
latencies
}
})
.await
})
})
.collect();
let mut all_latencies = collect_latencies(handles).await;
let elapsed = start.elapsed();
let total_ops = concurrency * ops_per_task;
let stats = compute_stats(&mut all_latencies);
print_result(total_ops, elapsed, stats.as_ref());
}
async fn bench_list_records(pool: &Arc<HandlerPool>, concurrency: usize, ops_per_task: usize) {
let user_id = Uuid::new_v4();
let did = Did::from(format!("did:plc:listrecords{}", user_id.as_simple()));
let collection = Nsid::from("app.bsky.feed.post".to_string());
create_user(pool, user_id, &did, &test_cid(1)).await;
seed_records(pool, user_id, &did, &collection, 1000).await;
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|_| {
let pool = Arc::clone(pool);
let collection = collection.clone();
tokio::spawn(async move {
futures::stream::iter(0..ops_per_task)
.fold(Vec::with_capacity(ops_per_task), |mut latencies, _| {
let pool = &pool;
let collection = &collection;
async move {
let t = Instant::now();
let (tx, rx) = oneshot::channel();
pool.send(MetastoreRequest::Record(RecordRequest::ListRecords {
repo_id: user_id,
collection: collection.clone(),
cursor: None,
limit: 50,
reverse: false,
rkey_start: None,
rkey_end: None,
tx,
}))
.unwrap();
let result = rx.await.unwrap().unwrap();
assert!(!result.is_empty());
latencies.push(t.elapsed());
latencies
}
})
.await
})
})
.collect();
let mut all_latencies = collect_latencies(handles).await;
let elapsed = start.elapsed();
let total_ops = concurrency * ops_per_task;
let stats = compute_stats(&mut all_latencies);
print_result(total_ops, elapsed, stats.as_ref());
}
async fn setup_pg_bench_schema(pool: &sqlx::PgPool) {
sqlx::query(
"CREATE TABLE IF NOT EXISTS users (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
handle TEXT NOT NULL UNIQUE,
email TEXT,
did TEXT NOT NULL UNIQUE,
password_hash TEXT NOT NULL DEFAULT '',
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
deactivated_at TIMESTAMPTZ,
invites_disabled BOOLEAN DEFAULT FALSE,
takedown_ref TEXT,
preferred_comms_channel TEXT NOT NULL DEFAULT 'email',
password_reset_code TEXT,
password_reset_code_expires_at TIMESTAMPTZ,
email_verified BOOLEAN NOT NULL DEFAULT FALSE,
two_factor_enabled BOOLEAN NOT NULL DEFAULT FALSE,
discord_id TEXT,
discord_verified BOOLEAN NOT NULL DEFAULT FALSE,
telegram_username TEXT,
telegram_verified BOOLEAN NOT NULL DEFAULT FALSE,
signal_number TEXT,
signal_verified BOOLEAN NOT NULL DEFAULT FALSE,
is_admin BOOLEAN NOT NULL DEFAULT FALSE,
migrated_to_pds TEXT,
migrated_at TIMESTAMPTZ,
preferred_locale TEXT,
signal_uuid TEXT
)",
)
.execute(pool)
.await
.unwrap();
sqlx::query(
"CREATE TABLE IF NOT EXISTS repos (
user_id UUID PRIMARY KEY REFERENCES users(id) ON DELETE CASCADE,
repo_root_cid TEXT NOT NULL,
repo_rev TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
)",
)
.execute(pool)
.await
.unwrap();
sqlx::query(
"CREATE TABLE IF NOT EXISTS records (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
repo_id UUID NOT NULL REFERENCES repos(user_id) ON DELETE CASCADE,
collection TEXT NOT NULL,
rkey TEXT NOT NULL,
record_cid TEXT NOT NULL,
takedown_ref TEXT,
repo_rev TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
UNIQUE(repo_id, collection, rkey)
)",
)
.execute(pool)
.await
.unwrap();
}
async fn pg_create_user(pg: &sqlx::PgPool, user_id: Uuid, did: &str) {
sqlx::query("INSERT INTO users (id, handle, did) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING")
.bind(user_id)
.bind(format!("bench.{}.invalid", Uuid::new_v4().as_simple()))
.bind(did)
.execute(pg)
.await
.unwrap();
}
async fn bench_pg_upsert_records(
repo: &dyn RepoRepository,
pg: &sqlx::PgPool,
concurrency: usize,
ops_per_task: usize,
) {
let user_ids: Vec<Uuid> = (0..concurrency).map(|_| Uuid::new_v4()).collect();
let dids: Vec<String> = user_ids
.iter()
.map(|u| format!("did:plc:pgbench{}", u.as_simple()))
.collect();
futures::stream::iter(user_ids.iter().zip(dids.iter()))
.fold((), |(), (uid, did)| async {
pg_create_user(pg, *uid, did).await;
let did_typed = Did::from(did.clone());
let handle = Handle::from(format!("bench.{}.invalid", uid.as_simple()));
repo.create_repo(*uid, &did_typed, &handle, &test_cid(1), "rev0000000000")
.await
.unwrap();
})
.await;
let collection = Nsid::from("app.bsky.feed.post".to_string());
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let user_id = user_ids[task_id];
let collection = collection.clone();
let pg = pg.clone();
tokio::spawn(async move {
let repo = tranquil_db::postgres::PostgresRepoRepository::new(pg);
futures::stream::iter(0..ops_per_task)
.fold(Vec::with_capacity(ops_per_task), |mut latencies, i| {
let repo = &repo;
let collection = &collection;
async move {
let rev_n = (task_id * ops_per_task + i + 1) as u64;
let cid_seed = ((task_id * 31 + i * 7) & 0xFF) as u8;
let rkey = Rkey::from(format!("r{rev_n:010}"));
let t = Instant::now();
repo.upsert_records(
user_id,
std::slice::from_ref(collection),
&[rkey],
&[test_cid(cid_seed)],
&make_rev(rev_n),
)
.await
.unwrap();
latencies.push(t.elapsed());
latencies
}
})
.await
})
})
.collect();
let mut all_latencies = collect_latencies(handles).await;
let elapsed = start.elapsed();
let total_ops = concurrency * ops_per_task;
let stats = compute_stats(&mut all_latencies);
print_result(total_ops, elapsed, stats.as_ref());
}
async fn pg_seed_records(
repo: &dyn RepoRepository,
user_id: Uuid,
collection: &Nsid,
count: usize,
) {
let batches: Vec<(usize, usize)> = (0..)
.map(|i| {
let start = i * 50;
let end = (start + 50).min(count);
(start, end)
})
.take_while(|(start, _)| *start < count)
.collect();
futures::stream::iter(batches)
.fold((), |(), (batch_start, batch_end)| {
let collection = collection.clone();
async move {
let collections: Vec<Nsid> = (batch_start..batch_end)
.map(|_| collection.clone())
.collect();
let rkeys: Vec<Rkey> = (batch_start..batch_end)
.map(|i| Rkey::from(format!("rec{i:08}")))
.collect();
let cids: Vec<CidLink> = (batch_start..batch_end)
.map(|i| test_cid(((i * 7 + 3) & 0xFF) as u8))
.collect();
repo.upsert_records(
user_id,
&collections,
&rkeys,
&cids,
&make_rev(batch_start as u64),
)
.await
.unwrap();
}
})
.await;
}
async fn bench_pg_get_record_cid(pg: &sqlx::PgPool, concurrency: usize, ops_per_task: usize) {
let user_id = Uuid::new_v4();
let did = format!("did:plc:pgget{}", user_id.as_simple());
let collection = Nsid::from("app.bsky.feed.post".to_string());
pg_create_user(pg, user_id, &did).await;
let repo = tranquil_db::postgres::PostgresRepoRepository::new(pg.clone());
let did_typed = Did::from(did.clone());
let handle = Handle::from(format!("bench.{}.invalid", user_id.as_simple()));
repo.create_repo(user_id, &did_typed, &handle, &test_cid(1), "rev0000000000")
.await
.unwrap();
pg_seed_records(&repo, user_id, &collection, 1000).await;
let total_records = 1000usize;
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let pg = pg.clone();
let collection = collection.clone();
tokio::spawn(async move {
let repo = tranquil_db::postgres::PostgresRepoRepository::new(pg);
futures::stream::iter(0..ops_per_task)
.fold(Vec::with_capacity(ops_per_task), |mut latencies, i| {
let repo = &repo;
let collection = &collection;
async move {
let rec_idx = (task_id * 7 + i * 13) % total_records;
let rkey = Rkey::from(format!("rec{rec_idx:08}"));
let t = Instant::now();
let result = repo
.get_record_cid(user_id, collection, &rkey)
.await
.unwrap();
assert!(result.is_some());
latencies.push(t.elapsed());
latencies
}
})
.await
})
})
.collect();
let mut all_latencies = collect_latencies(handles).await;
let elapsed = start.elapsed();
let total_ops = concurrency * ops_per_task;
let stats = compute_stats(&mut all_latencies);
print_result(total_ops, elapsed, stats.as_ref());
}
async fn bench_pg_list_records(pg: &sqlx::PgPool, concurrency: usize, ops_per_task: usize) {
let user_id = Uuid::new_v4();
let did = format!("did:plc:pglist{}", user_id.as_simple());
let collection = Nsid::from("app.bsky.feed.post".to_string());
pg_create_user(pg, user_id, &did).await;
let repo = tranquil_db::postgres::PostgresRepoRepository::new(pg.clone());
let did_typed = Did::from(did.clone());
let handle = Handle::from(format!("bench.{}.invalid", user_id.as_simple()));
repo.create_repo(user_id, &did_typed, &handle, &test_cid(1), "rev0000000000")
.await
.unwrap();
pg_seed_records(&repo, user_id, &collection, 1000).await;
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|_| {
let pg = pg.clone();
let collection = collection.clone();
tokio::spawn(async move {
let repo = tranquil_db::postgres::PostgresRepoRepository::new(pg);
futures::stream::iter(0..ops_per_task)
.fold(Vec::with_capacity(ops_per_task), |mut latencies, _| {
let repo = &repo;
let collection = &collection;
async move {
let t = Instant::now();
let result = repo
.list_records(user_id, collection, None, 50, false, None, None)
.await
.unwrap();
assert!(!result.is_empty());
latencies.push(t.elapsed());
latencies
}
})
.await
})
})
.collect();
let mut all_latencies = collect_latencies(handles).await;
let elapsed = start.elapsed();
let total_ops = concurrency * ops_per_task;
let stats = compute_stats(&mut all_latencies);
print_result(total_ops, elapsed, stats.as_ref());
}
#[tokio::main]
async fn main() {
let handler_threads = std::thread::available_parallelism()
.map(|n| n.get().max(2) / 2)
.unwrap_or(2);
println!("handler threads: {handler_threads}");
let concurrency_levels = [1, 10, 100, 1000];
let ops_per_concurrency = |c: usize| match c {
1 => 5000,
10 => 1000,
100 => 200,
1000 => 50,
_ => 100,
};
futures::stream::iter(concurrency_levels.iter())
.fold((), |(), &c| async move {
let ops = ops_per_concurrency(c);
println!("-- apply_commit: {} ops, {} callers --", ops * c, c);
let h = setup(handler_threads);
bench_apply_commit(&h.pool, c, ops).await;
println!("-- get_record_cid: {} ops, {} callers --", ops * c, c);
let h = setup(handler_threads);
bench_get_record_cid(&h.pool, c, ops).await;
println!("-- list_records: {} ops, {} callers --", ops * c, c);
let h = setup(handler_threads);
bench_list_records(&h.pool, c, ops).await;
})
.await;
let pg_url = match std::env::var("DATABASE_URL") {
Ok(url) => url,
Err(_) => {
println!("set DATABASE_URL for postgres comparison");
return;
}
};
let pg_concurrency_levels: &[usize] = &[1, 10, 50];
let setup_pg = |max_conns: u32| {
let url = pg_url.clone();
async move {
sqlx::postgres::PgPoolOptions::new()
.max_connections(max_conns)
.connect(&url)
.await
.unwrap()
}
};
let pg = setup_pg(60).await;
setup_pg_bench_schema(&pg).await;
pg.close().await;
futures::stream::iter(pg_concurrency_levels.iter())
.fold((), |(), &c| {
let setup_pg = &setup_pg;
async move {
let ops = ops_per_concurrency(c);
let max_conns = u32::try_from(c).unwrap_or(50) + 10;
let pg = setup_pg(max_conns).await;
let repo = tranquil_db::postgres::PostgresRepoRepository::new(pg.clone());
println!(
"-- postgres upsert_records: {} ops, {} callers --",
ops * c,
c
);
bench_pg_upsert_records(&repo, &pg, c, ops).await;
println!(
"-- postgres get_record_cid: {} ops, {} callers --",
ops * c,
c
);
bench_pg_get_record_cid(&pg, c, ops).await;
println!(
"-- postgres list_records: {} ops, {} callers --",
ops * c,
c
);
bench_pg_list_records(&pg, c, ops).await;
sqlx::query("TRUNCATE records, repos, users CASCADE")
.execute(&pg)
.await
.unwrap();
pg.close().await;
}
})
.await;
let pg = setup_pg(5).await;
sqlx::query("DROP TABLE IF EXISTS records CASCADE")
.execute(&pg)
.await
.unwrap();
sqlx::query("DROP TABLE IF EXISTS repos CASCADE")
.execute(&pg)
.await
.unwrap();
sqlx::query("DROP TABLE IF EXISTS users CASCADE")
.execute(&pg)
.await
.unwrap();
pg.close().await;
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,358 @@
use std::sync::Arc;
use std::time::{Duration, Instant};
use futures::StreamExt;
use tokio::sync::oneshot;
use tranquil_db_traits::{ApplyCommitInput, CommitEventData, RecordUpsert, RepoEventType};
use tranquil_types::{CidLink, Did, Handle, Nsid, Rkey};
use uuid::Uuid;
use tranquil_store::RealIO;
use tranquil_store::eventlog::{EventLog, EventLogConfig};
use tranquil_store::metastore::handler::{
CommitRequest, HandlerPool, MetastoreRequest, RecordRequest, RepoRequest,
};
use tranquil_store::metastore::{Metastore, MetastoreConfig};
fn test_cid(seed: u8) -> CidLink {
let digest: [u8; 32] = std::array::from_fn(|i| seed.wrapping_add(i as u8));
let mh = multihash::Multihash::<64>::wrap(0x12, &digest).unwrap();
let c = cid::Cid::new_v1(0x71, mh);
CidLink::from_cid(&c)
}
fn test_cid_bytes(seed: u8) -> Vec<u8> {
let digest: [u8; 32] = std::array::from_fn(|i| seed.wrapping_add(i as u8));
let mh = multihash::Multihash::<64>::wrap(0x12, &digest).unwrap();
cid::Cid::new_v1(0x71, mh).to_bytes()
}
struct UserInfo {
user_id: Uuid,
did: Did,
}
async fn seed_users(pool: &HandlerPool, count: usize) -> Vec<UserInfo> {
let users: Vec<UserInfo> = (0..count)
.map(|i| {
let user_id = Uuid::new_v4();
UserInfo {
did: Did::from(format!("did:plc:prof{i:06x}{}", user_id.as_simple())),
user_id,
}
})
.collect();
let batch_size = 500;
let start = Instant::now();
let total_batches = count.div_ceil(batch_size);
futures::stream::iter(users.chunks(batch_size).enumerate())
.fold((), |(), (batch_idx, batch)| async move {
futures::stream::iter(batch.iter())
.fold((), |(), user| async {
let (tx, rx) = oneshot::channel();
pool.send(MetastoreRequest::Repo(RepoRequest::CreateRepoFull {
user_id: user.user_id,
did: user.did.clone(),
handle: Handle::from(format!("u{}.prof.invalid", user.user_id.as_simple())),
repo_root_cid: test_cid(1),
repo_rev: "rev0000000000".to_string(),
tx,
}))
.unwrap();
rx.await.unwrap().unwrap();
})
.await;
if (batch_idx + 1) % 100 == 0 || batch_idx + 1 == total_batches {
println!(
"seeded {}/{} users, {:.1}s",
((batch_idx + 1) * batch_size).min(count),
count,
start.elapsed().as_secs_f64()
);
}
})
.await;
println!(
"seeded {} users in {:.1}s, {:.0} users/sec",
count,
start.elapsed().as_secs_f64(),
count as f64 / start.elapsed().as_secs_f64()
);
users
}
async fn seed_records(pool: &Arc<HandlerPool>, users: &[UserInfo], records_per_user: usize) {
let start = Instant::now();
let total = users.len();
let batch_size = 500;
let total_batches = total.div_ceil(batch_size);
let collection = Nsid::from("app.bsky.feed.post".to_string());
futures::stream::iter(users.chunks(batch_size).enumerate())
.fold((), |(), (chunk_idx, chunk)| {
let pool = Arc::clone(pool);
let collection = collection.clone();
async move {
futures::stream::iter(chunk.iter())
.fold((), |(), user| {
let pool = &pool;
let collection = &collection;
async move {
let record_upserts: Vec<RecordUpsert> = (0..records_per_user)
.map(|i| RecordUpsert {
collection: collection.clone(),
rkey: Rkey::from(format!("rec{i:08}")),
cid: test_cid(((i * 7 + 3) & 0xFF) as u8),
})
.collect();
let new_block_cids: Vec<Vec<u8>> = (0..records_per_user)
.map(|i| test_cid_bytes(((i * 11 + 5) & 0xFF) as u8))
.collect();
let input = ApplyCommitInput {
user_id: user.user_id,
did: user.did.clone(),
expected_root_cid: None,
new_root_cid: test_cid(2),
new_rev: "rev0000000001".to_string(),
new_block_cids,
obsolete_block_cids: vec![],
record_upserts,
record_deletes: vec![],
backlinks_to_add: vec![],
backlinks_to_remove: vec![],
commit_event: CommitEventData {
did: user.did.clone(),
event_type: RepoEventType::Commit,
commit_cid: Some(test_cid(2)),
prev_cid: None,
ops: None,
blobs: None,
blocks_cids: None,
prev_data_cid: None,
rev: Some("rev0000000001".to_string()),
},
};
let (tx, rx) = oneshot::channel();
pool.send(MetastoreRequest::Commit(Box::new(
CommitRequest::ApplyCommit {
input: Box::new(input),
tx,
},
)))
.unwrap();
rx.await.unwrap().unwrap();
}
})
.await;
if (chunk_idx + 1) % 100 == 0 || chunk_idx + 1 == total_batches {
println!(
"seeded records for {}/{} users, {:.1}s",
((chunk_idx + 1) * batch_size).min(total),
total,
start.elapsed().as_secs_f64()
);
}
}
})
.await;
println!(
"seeded {} records across {} users in {:.1}s",
total * records_per_user,
total,
start.elapsed().as_secs_f64()
);
}
async fn profile_list_records(
pool: &Arc<HandlerPool>,
user_ids: &Arc<Vec<Uuid>>,
concurrency: usize,
seconds: u64,
) -> u64 {
let deadline = Instant::now() + Duration::from_secs(seconds);
let user_count = user_ids.len();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let pool = Arc::clone(pool);
let user_ids = Arc::clone(user_ids);
tokio::spawn(async move {
futures::stream::unfold(0usize, |i| {
let cont = Instant::now() < deadline;
async move { cont.then_some((i, i + 1)) }
})
.fold(0u64, |ops, i| {
let pool = &pool;
let user_ids = &user_ids;
async move {
let idx = (task_id * 997 + i * 31) % user_count;
let (tx, rx) = oneshot::channel();
pool.send(MetastoreRequest::Record(RecordRequest::ListRecords {
repo_id: user_ids[idx],
collection: Nsid::from("app.bsky.feed.post".to_string()),
cursor: None,
limit: 50,
reverse: false,
rkey_start: None,
rkey_end: None,
tx,
}))
.unwrap();
let _ = rx.await.unwrap().unwrap();
ops + 1
}
})
.await
})
})
.collect();
futures::stream::iter(handles)
.fold(0u64, |acc, h| async move { acc + h.await.unwrap() })
.await
}
async fn profile_get_record_cid(
pool: &Arc<HandlerPool>,
user_ids: &Arc<Vec<Uuid>>,
concurrency: usize,
seconds: u64,
records_per_user: usize,
) -> u64 {
let deadline = Instant::now() + Duration::from_secs(seconds);
let user_count = user_ids.len();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let pool = Arc::clone(pool);
let user_ids = Arc::clone(user_ids);
tokio::spawn(async move {
futures::stream::unfold(0usize, |i| {
let cont = Instant::now() < deadline;
async move { cont.then_some((i, i + 1)) }
})
.fold(0u64, |ops, i| {
let pool = &pool;
let user_ids = &user_ids;
async move {
let user_idx = (task_id * 997 + i * 31) % user_count;
let rec_idx = (task_id * 13 + i * 7) % records_per_user;
let rkey = Rkey::from(format!("rec{rec_idx:08}"));
let (tx, rx) = oneshot::channel();
pool.send(MetastoreRequest::Record(RecordRequest::GetRecordCid {
repo_id: user_ids[user_idx],
collection: Nsid::from("app.bsky.feed.post".to_string()),
rkey,
tx,
}))
.unwrap();
let _ = rx.await.unwrap().unwrap();
ops + 1
}
})
.await
})
})
.collect();
futures::stream::iter(handles)
.fold(0u64, |acc, h| async move { acc + h.await.unwrap() })
.await
}
#[tokio::main]
async fn main() {
let user_count = std::env::var("PROFILE_USERS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(300_000usize);
let records_per_user = 10;
let profile_seconds = std::env::var("PROFILE_SECONDS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(30u64);
let concurrency = 100usize;
let handler_threads = std::thread::available_parallelism()
.map(|n| n.get().max(2) / 2)
.unwrap_or(2);
println!("-- profile reads --");
println!("handler threads: {handler_threads}");
println!("{user_count} users, {records_per_user} records each, {profile_seconds}s per phase");
let metastore_dir = tempfile::TempDir::new().unwrap();
let eventlog_dir = tempfile::TempDir::new().unwrap();
let segments_dir = eventlog_dir.path().join("segments");
std::fs::create_dir_all(&segments_dir).unwrap();
let cache_bytes = (user_count as u64 * 10 * 300)
.saturating_mul(2)
.clamp(512 * 1024 * 1024, 8 * 1024 * 1024 * 1024);
println!("cache size: {} MB", cache_bytes / (1024 * 1024));
let metastore = Metastore::open(
metastore_dir.path(),
MetastoreConfig {
cache_size_bytes: cache_bytes,
},
)
.unwrap();
let event_log = EventLog::open(
EventLogConfig {
segments_dir,
..EventLogConfig::default()
},
RealIO::new(),
)
.unwrap();
let bridge = Arc::new(tranquil_store::eventlog::EventLogBridge::new(Arc::new(
event_log,
)));
let pool = Arc::new(HandlerPool::spawn::<RealIO>(
metastore.clone(),
bridge,
None,
Some(handler_threads),
));
let users = seed_users(&pool, user_count).await;
seed_records(&pool, &users, records_per_user).await;
println!("running major compaction...");
let t = Instant::now();
metastore.major_compact().unwrap();
println!(
"major compaction complete in {:.1}s",
t.elapsed().as_secs_f64()
);
let user_ids: Arc<Vec<Uuid>> = Arc::new(users.iter().map(|u| u.user_id).collect());
println!("-- listRecords, {concurrency} readers, {profile_seconds}s --");
let list_ops = profile_list_records(&pool, &user_ids, concurrency, profile_seconds).await;
println!(
"listRecords: {list_ops} ops, {:.0} ops/sec",
list_ops as f64 / profile_seconds as f64
);
println!("-- getRecordCid, {concurrency} readers, {profile_seconds}s --");
let get_ops = profile_get_record_cid(
&pool,
&user_ids,
concurrency,
profile_seconds,
records_per_user,
)
.await;
println!(
"getRecordCid: {get_ops} ops, {:.0} ops/sec",
get_ops as f64 / profile_seconds as f64
);
println!("-- profile reads complete :D --");
}
+1
View File
@@ -3,6 +3,7 @@ pub mod eventlog;
pub mod fsync_order;
mod harness;
mod io;
pub mod metastore;
mod record;
#[cfg(any(test, feature = "test-harness"))]
mod sim;
@@ -0,0 +1,638 @@
use std::sync::Arc;
use fjall::{Keyspace, OwnedWriteBatch};
use uuid::Uuid;
use super::MetastoreError;
use super::backlinks::{
BacklinkValue, backlink_by_user_key, backlink_by_user_prefix, backlink_by_user_record_prefix,
backlink_key, backlink_target_user_prefix, discriminant_to_path, path_to_discriminant,
};
use super::encoding::KeyReader;
use super::keys::{KeyTag, UserHash};
use super::user_hash::UserHashMap;
use tranquil_db_traits::Backlink;
use tranquil_types::{AtUri, Nsid};
pub(super) fn parse_backlink_by_user_fields(key_bytes: &[u8]) -> Option<(String, String, String)> {
let mut reader = KeyReader::new(key_bytes);
let tag = reader.tag()?;
match tag == KeyTag::BACKLINK_BY_USER.raw() {
true => {
let _user_hash = reader.u64()?;
let collection = reader.string()?;
let rkey = reader.string()?;
let link_target = reader.string()?;
Some((collection, rkey, link_target))
}
false => None,
}
}
pub(super) fn remove_backlinks_for_record(
batch: &mut OwnedWriteBatch,
indexes: &Keyspace,
user_hash: UserHash,
collection: &str,
rkey: &str,
) -> Result<(), MetastoreError> {
let record_prefix = backlink_by_user_record_prefix(user_hash, collection, rkey);
indexes
.prefix(record_prefix.as_slice())
.try_for_each(|guard| {
let (key_bytes, _) = guard.into_inner().map_err(MetastoreError::Fjall)?;
let (_, _, link_target) = parse_backlink_by_user_fields(&key_bytes).ok_or(
MetastoreError::CorruptData("unparseable BACKLINK_BY_USER key"),
)?;
let primary = backlink_key(&link_target, user_hash, collection, rkey);
batch.remove(indexes, primary.as_slice());
batch.remove(indexes, key_bytes.as_ref());
Ok::<_, MetastoreError>(())
})
}
pub struct BacklinkOps {
indexes: Keyspace,
user_hashes: Arc<UserHashMap>,
}
impl BacklinkOps {
pub fn new(indexes: Keyspace, user_hashes: Arc<UserHashMap>) -> Self {
Self {
indexes,
user_hashes,
}
}
pub fn add_backlinks(
&self,
batch: &mut OwnedWriteBatch,
user_hash: UserHash,
backlinks: &[Backlink],
) -> Result<(), MetastoreError> {
backlinks
.iter()
.try_for_each(|bl| self.add_single_backlink(batch, user_hash, bl))
}
fn add_single_backlink(
&self,
batch: &mut OwnedWriteBatch,
user_hash: UserHash,
bl: &Backlink,
) -> Result<(), MetastoreError> {
let collection = bl.uri.collection().ok_or(MetastoreError::InvalidInput(
"backlink uri missing collection",
))?;
let rkey = bl
.uri
.rkey()
.ok_or(MetastoreError::InvalidInput("backlink uri missing rkey"))?;
let primary = backlink_key(&bl.link_to, user_hash, collection, rkey);
let value = BacklinkValue {
source_uri: bl.uri.as_str().to_owned(),
path: path_to_discriminant(bl.path),
};
batch.insert(&self.indexes, primary.as_slice(), value.serialize());
let reverse = backlink_by_user_key(user_hash, collection, rkey, &bl.link_to);
batch.insert(&self.indexes, reverse.as_slice(), []);
Ok(())
}
pub fn remove_backlinks_by_uri(
&self,
batch: &mut OwnedWriteBatch,
user_hash: UserHash,
uri: &AtUri,
) -> Result<(), MetastoreError> {
let collection = uri.collection().ok_or(MetastoreError::InvalidInput(
"backlink uri missing collection",
))?;
let rkey = uri
.rkey()
.ok_or(MetastoreError::InvalidInput("backlink uri missing rkey"))?;
remove_backlinks_for_record(batch, &self.indexes, user_hash, collection, rkey)
}
pub fn remove_backlinks_by_repo(
&self,
batch: &mut OwnedWriteBatch,
user_hash: UserHash,
) -> Result<(), MetastoreError> {
let user_prefix = backlink_by_user_prefix(user_hash);
self.indexes
.prefix(user_prefix.as_slice())
.try_for_each(|guard| {
let (key_bytes, _) = guard.into_inner().map_err(MetastoreError::Fjall)?;
let (collection, rkey, link_target) = parse_backlink_by_user_fields(&key_bytes)
.ok_or(MetastoreError::CorruptData(
"unparseable BACKLINK_BY_USER key",
))?;
let primary = backlink_key(&link_target, user_hash, &collection, &rkey);
batch.remove(&self.indexes, primary.as_slice());
batch.remove(&self.indexes, key_bytes.as_ref());
Ok::<_, MetastoreError>(())
})
}
pub fn get_backlink_conflicts(
&self,
repo_id: Uuid,
collection: &Nsid,
backlinks: &[Backlink],
) -> Result<Vec<AtUri>, MetastoreError> {
if backlinks.is_empty() {
return Ok(Vec::new());
}
let user_hash = self
.user_hashes
.get(&repo_id)
.ok_or(MetastoreError::InvalidInput("unknown repo_id"))?;
let collection_str = collection.as_str();
let mut seen = std::collections::HashSet::new();
backlinks.iter().try_fold(Vec::new(), |mut conflicts, bl| {
let prefix = backlink_target_user_prefix(&bl.link_to, user_hash);
self.indexes
.prefix(prefix.as_slice())
.try_for_each(|guard| {
let (_, val_bytes) = guard.into_inner().map_err(MetastoreError::Fjall)?;
let val = BacklinkValue::deserialize(&val_bytes).ok_or(
MetastoreError::CorruptData("corrupt backlink value in indexes partition"),
)?;
let uri: AtUri = val.source_uri.into();
let is_self = uri.as_str() == bl.uri.as_str();
let matches_collection = uri.collection().is_some_and(|c| c == collection_str);
let matches_path = match discriminant_to_path(val.path) {
Some(p) => p == bl.path,
None => {
tracing::warn!(
discriminant = val.path,
uri = %uri,
"unknown backlink path discriminant in indexes partition"
);
false
}
};
if !is_self
&& matches_collection
&& matches_path
&& !seen.contains(uri.as_str())
{
seen.insert(uri.as_str().to_owned());
conflicts.push(uri);
}
Ok::<_, MetastoreError>(())
})?;
Ok(conflicts)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metastore::backlinks::{backlink_by_user_prefix, backlink_target_prefix};
use crate::metastore::partitions::Partition;
use crate::metastore::{Metastore, MetastoreConfig};
use tranquil_db_traits::{Backlink, BacklinkPath};
use tranquil_types::{Did, Handle, Nsid};
struct TestHarness {
_dir: tempfile::TempDir,
metastore: Metastore,
}
fn setup() -> TestHarness {
let dir = tempfile::TempDir::new().unwrap();
let metastore = Metastore::open(
dir.path(),
MetastoreConfig {
cache_size_bytes: 64 * 1024 * 1024,
},
)
.unwrap();
TestHarness {
_dir: dir,
metastore,
}
}
fn test_cid_link(seed: u8) -> tranquil_types::CidLink {
let digest: [u8; 32] = std::array::from_fn(|i| seed.wrapping_add(i as u8));
let mh = multihash::Multihash::<64>::wrap(0x12, &digest).unwrap();
let c = cid::Cid::new_v1(0x71, mh);
tranquil_types::CidLink::from_cid(&c)
}
fn create_repo(h: &TestHarness, name: &str, seed: u8) -> (Uuid, UserHash) {
let user_id = Uuid::new_v4();
let did = Did::from(format!("did:plc:{name}"));
let handle = Handle::from(format!("{name}.test.invalid"));
let cid = test_cid_link(seed);
h.metastore
.repo_ops()
.create_repo(h.metastore.database(), user_id, &did, &handle, &cid, "rev0")
.unwrap();
let user_hash = h.metastore.user_hashes().get(&user_id).unwrap();
(user_id, user_hash)
}
fn count_prefix(ks: &fjall::Keyspace, prefix: &[u8]) -> usize {
ks.prefix(prefix)
.map(|g| g.into_inner().expect("prefix scan must not fail in test"))
.fold(0, |acc, _| acc + 1)
}
#[test]
fn add_and_query_by_target() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (_user_id, user_hash) = create_repo(&h, "alice", 1);
let uri = AtUri::from_parts("did:plc:alice", "app.bsky.feed.like", "3k2abc");
let backlinks = vec![Backlink {
uri: uri.clone(),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:target/app.bsky.feed.post/3k2post".to_string(),
}];
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash, &backlinks)
.unwrap();
batch.commit().unwrap();
let indexes = h.metastore.partition(Partition::Indexes);
let target_prefix =
backlink_target_prefix("at://did:plc:target/app.bsky.feed.post/3k2post");
assert_eq!(count_prefix(indexes, target_prefix.as_slice()), 1);
let user_prefix = backlink_by_user_prefix(user_hash);
assert_eq!(count_prefix(indexes, user_prefix.as_slice()), 1);
}
#[test]
fn remove_by_uri_deletes_both_indexes() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (_user_id, user_hash) = create_repo(&h, "bob", 2);
let uri = AtUri::from_parts("did:plc:bob", "app.bsky.graph.follow", "3k2fol");
let backlinks = vec![Backlink {
uri: uri.clone(),
path: BacklinkPath::Subject,
link_to: "did:plc:target_user".to_string(),
}];
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash, &backlinks)
.unwrap();
batch.commit().unwrap();
let indexes = h.metastore.partition(Partition::Indexes);
assert_eq!(
count_prefix(
indexes,
backlink_target_prefix("did:plc:target_user").as_slice()
),
1
);
let mut batch = h.metastore.database().batch();
ops.remove_backlinks_by_uri(&mut batch, user_hash, &uri)
.unwrap();
batch.commit().unwrap();
assert_eq!(
count_prefix(
indexes,
backlink_target_prefix("did:plc:target_user").as_slice()
),
0
);
assert_eq!(
count_prefix(indexes, backlink_by_user_prefix(user_hash).as_slice()),
0
);
}
#[test]
fn remove_by_repo_deletes_all_user_backlinks() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (_user_id, user_hash) = create_repo(&h, "carol", 3);
let backlinks: Vec<Backlink> = (0..5)
.map(|i| Backlink {
uri: AtUri::from_parts("did:plc:carol", "app.bsky.feed.like", &format!("3k2r{i}")),
path: BacklinkPath::SubjectUri,
link_to: format!("at://did:plc:target{i}/app.bsky.feed.post/3k2p{i}"),
})
.collect();
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash, &backlinks)
.unwrap();
batch.commit().unwrap();
let indexes = h.metastore.partition(Partition::Indexes);
assert_eq!(
count_prefix(indexes, backlink_by_user_prefix(user_hash).as_slice()),
5
);
let mut batch = h.metastore.database().batch();
ops.remove_backlinks_by_repo(&mut batch, user_hash).unwrap();
batch.commit().unwrap();
assert_eq!(
count_prefix(indexes, backlink_by_user_prefix(user_hash).as_slice()),
0
);
(0..5).for_each(|i| {
let target = format!("at://did:plc:target{i}/app.bsky.feed.post/3k2p{i}");
let prefix = backlink_target_prefix(&target);
assert_eq!(count_prefix(indexes, prefix.as_slice()), 0);
});
}
#[test]
fn get_backlink_conflicts_finds_matching() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (user_id, user_hash) = create_repo(&h, "dave", 4);
let existing = Backlink {
uri: AtUri::from_parts("did:plc:dave", "app.bsky.feed.like", "3k2old"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:someone/app.bsky.feed.post/3k2p1".to_string(),
};
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash, &[existing])
.unwrap();
batch.commit().unwrap();
let proposed = vec![Backlink {
uri: AtUri::from_parts("did:plc:dave", "app.bsky.feed.like", "3k2new"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:someone/app.bsky.feed.post/3k2p1".to_string(),
}];
let collection = Nsid::from("app.bsky.feed.like".to_string());
let conflicts = ops
.get_backlink_conflicts(user_id, &collection, &proposed)
.unwrap();
assert_eq!(conflicts.len(), 1);
assert_eq!(
conflicts[0].as_str(),
"at://did:plc:dave/app.bsky.feed.like/3k2old"
);
}
#[test]
fn get_backlink_conflicts_ignores_different_collection() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (user_id, user_hash) = create_repo(&h, "eve", 5);
let existing = Backlink {
uri: AtUri::from_parts("did:plc:eve", "app.bsky.feed.like", "3k2old"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:someone/app.bsky.feed.post/3k2p1".to_string(),
};
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash, &[existing])
.unwrap();
batch.commit().unwrap();
let proposed = vec![Backlink {
uri: AtUri::from_parts("did:plc:eve", "app.bsky.feed.repost", "3k2new"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:someone/app.bsky.feed.post/3k2p1".to_string(),
}];
let collection = Nsid::from("app.bsky.feed.repost".to_string());
let conflicts = ops
.get_backlink_conflicts(user_id, &collection, &proposed)
.unwrap();
assert!(conflicts.is_empty());
}
#[test]
fn get_backlink_conflicts_ignores_different_path() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (user_id, user_hash) = create_repo(&h, "frank", 6);
let existing = Backlink {
uri: AtUri::from_parts("did:plc:frank", "app.bsky.graph.follow", "3k2old"),
path: BacklinkPath::Subject,
link_to: "did:plc:target".to_string(),
};
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash, &[existing])
.unwrap();
batch.commit().unwrap();
let proposed = vec![Backlink {
uri: AtUri::from_parts("did:plc:frank", "app.bsky.graph.follow", "3k2new"),
path: BacklinkPath::SubjectUri,
link_to: "did:plc:target".to_string(),
}];
let collection = Nsid::from("app.bsky.graph.follow".to_string());
let conflicts = ops
.get_backlink_conflicts(user_id, &collection, &proposed)
.unwrap();
assert!(conflicts.is_empty());
}
#[test]
fn get_backlink_conflicts_ignores_other_users() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (_user_id_a, user_hash_a) = create_repo(&h, "grace", 7);
let (user_id_b, _user_hash_b) = create_repo(&h, "henry", 8);
let existing = Backlink {
uri: AtUri::from_parts("did:plc:grace", "app.bsky.feed.like", "3k2old"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:target/app.bsky.feed.post/3k2p1".to_string(),
};
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash_a, &[existing])
.unwrap();
batch.commit().unwrap();
let proposed = vec![Backlink {
uri: AtUri::from_parts("did:plc:henry", "app.bsky.feed.like", "3k2new"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:target/app.bsky.feed.post/3k2p1".to_string(),
}];
let collection = Nsid::from("app.bsky.feed.like".to_string());
let conflicts = ops
.get_backlink_conflicts(user_id_b, &collection, &proposed)
.unwrap();
assert!(conflicts.is_empty());
}
#[test]
fn get_backlink_conflicts_excludes_self_match() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (user_id, user_hash) = create_repo(&h, "luna", 12);
let existing = Backlink {
uri: AtUri::from_parts("did:plc:luna", "app.bsky.feed.like", "3k2same"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:someone/app.bsky.feed.post/3k2p1".to_string(),
};
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash, &[existing])
.unwrap();
batch.commit().unwrap();
let proposed = vec![Backlink {
uri: AtUri::from_parts("did:plc:luna", "app.bsky.feed.like", "3k2same"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:someone/app.bsky.feed.post/3k2p1".to_string(),
}];
let collection = Nsid::from("app.bsky.feed.like".to_string());
let conflicts = ops
.get_backlink_conflicts(user_id, &collection, &proposed)
.unwrap();
assert!(conflicts.is_empty());
}
#[test]
fn empty_backlinks_returns_empty_conflicts() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (user_id, _user_hash) = create_repo(&h, "ivan", 9);
let collection = Nsid::from("app.bsky.feed.like".to_string());
let conflicts = ops
.get_backlink_conflicts(user_id, &collection, &[])
.unwrap();
assert!(conflicts.is_empty());
}
#[test]
fn remove_by_uri_only_removes_matching_rkey() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (_user_id, user_hash) = create_repo(&h, "julia", 10);
let bl1 = Backlink {
uri: AtUri::from_parts("did:plc:julia", "app.bsky.feed.like", "3k2aaa"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:t1/app.bsky.feed.post/p1".to_string(),
};
let bl2 = Backlink {
uri: AtUri::from_parts("did:plc:julia", "app.bsky.feed.like", "3k2bbb"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:t2/app.bsky.feed.post/p2".to_string(),
};
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash, &[bl1.clone(), bl2])
.unwrap();
batch.commit().unwrap();
let indexes = h.metastore.partition(Partition::Indexes);
assert_eq!(
count_prefix(indexes, backlink_by_user_prefix(user_hash).as_slice()),
2
);
let mut batch = h.metastore.database().batch();
ops.remove_backlinks_by_uri(&mut batch, user_hash, &bl1.uri)
.unwrap();
batch.commit().unwrap();
assert_eq!(
count_prefix(indexes, backlink_by_user_prefix(user_hash).as_slice()),
1
);
assert_eq!(
count_prefix(
indexes,
backlink_target_prefix("at://did:plc:t1/app.bsky.feed.post/p1").as_slice()
),
0
);
assert_eq!(
count_prefix(
indexes,
backlink_target_prefix("at://did:plc:t2/app.bsky.feed.post/p2").as_slice()
),
1
);
}
#[test]
fn conflicts_deduplicates_results() {
let h = setup();
let ops = h.metastore.backlink_ops();
let (user_id, user_hash) = create_repo(&h, "kate", 11);
let existing = Backlink {
uri: AtUri::from_parts("did:plc:kate", "app.bsky.feed.like", "3k2old"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:someone/app.bsky.feed.post/3k2p1".to_string(),
};
let mut batch = h.metastore.database().batch();
ops.add_backlinks(&mut batch, user_hash, &[existing])
.unwrap();
batch.commit().unwrap();
let proposed = vec![
Backlink {
uri: AtUri::from_parts("did:plc:kate", "app.bsky.feed.like", "3k2new1"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:someone/app.bsky.feed.post/3k2p1".to_string(),
},
Backlink {
uri: AtUri::from_parts("did:plc:kate", "app.bsky.feed.like", "3k2new2"),
path: BacklinkPath::SubjectUri,
link_to: "at://did:plc:someone/app.bsky.feed.post/3k2p1".to_string(),
},
];
let collection = Nsid::from("app.bsky.feed.like".to_string());
let conflicts = ops
.get_backlink_conflicts(user_id, &collection, &proposed)
.unwrap();
assert_eq!(conflicts.len(), 1);
}
}
@@ -0,0 +1,257 @@
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use super::encoding::KeyBuilder;
use super::keys::{KeyTag, UserHash};
use tranquil_db_traits::BacklinkPath;
const SCHEMA_VERSION: u8 = 1;
pub fn path_to_discriminant(path: BacklinkPath) -> u8 {
match path {
BacklinkPath::Subject => 0,
BacklinkPath::SubjectUri => 1,
}
}
pub fn discriminant_to_path(d: u8) -> Option<BacklinkPath> {
match d {
0 => Some(BacklinkPath::Subject),
1 => Some(BacklinkPath::SubjectUri),
_ => None,
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BacklinkValue {
pub source_uri: String,
pub path: u8,
}
impl BacklinkValue {
pub fn serialize(&self) -> Vec<u8> {
let payload = postcard::to_allocvec(self).expect("BacklinkValue serialization cannot fail");
let mut buf = Vec::with_capacity(1 + payload.len());
buf.push(SCHEMA_VERSION);
buf.extend_from_slice(&payload);
buf
}
pub fn deserialize(bytes: &[u8]) -> Option<Self> {
let (&version, payload) = bytes.split_first()?;
match version {
SCHEMA_VERSION => postcard::from_bytes(payload).ok(),
_ => None,
}
}
}
pub fn backlink_key(
link_target: &str,
user_hash: UserHash,
collection: &str,
rkey: &str,
) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::BACKLINKS)
.string(link_target)
.u64(user_hash.raw())
.string(collection)
.string(rkey)
.build()
}
pub fn backlink_target_prefix(link_target: &str) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::BACKLINKS)
.string(link_target)
.build()
}
pub fn backlink_target_user_prefix(link_target: &str, user_hash: UserHash) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::BACKLINKS)
.string(link_target)
.u64(user_hash.raw())
.build()
}
pub fn backlink_by_user_key(
user_hash: UserHash,
collection: &str,
rkey: &str,
link_target: &str,
) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::BACKLINK_BY_USER)
.u64(user_hash.raw())
.string(collection)
.string(rkey)
.string(link_target)
.build()
}
pub fn backlink_by_user_prefix(user_hash: UserHash) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::BACKLINK_BY_USER)
.u64(user_hash.raw())
.build()
}
pub fn backlink_by_user_record_prefix(
user_hash: UserHash,
collection: &str,
rkey: &str,
) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::BACKLINK_BY_USER)
.u64(user_hash.raw())
.string(collection)
.string(rkey)
.build()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metastore::encoding::KeyReader;
#[test]
fn discriminant_roundtrip() {
assert_eq!(
discriminant_to_path(path_to_discriminant(BacklinkPath::Subject)),
Some(BacklinkPath::Subject)
);
assert_eq!(
discriminant_to_path(path_to_discriminant(BacklinkPath::SubjectUri)),
Some(BacklinkPath::SubjectUri)
);
assert_eq!(discriminant_to_path(255), None);
assert_eq!(discriminant_to_path(2), None);
}
#[test]
fn backlink_value_roundtrip() {
let value = BacklinkValue {
source_uri: "at://did:plc:abc/app.bsky.feed.like/3k2xyz".to_string(),
path: path_to_discriminant(BacklinkPath::SubjectUri),
};
let bytes = value.serialize();
let decoded = BacklinkValue::deserialize(&bytes).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn schema_version_is_first_byte() {
let value = BacklinkValue {
source_uri: "at://x".to_string(),
path: path_to_discriminant(BacklinkPath::Subject),
};
let bytes = value.serialize();
assert_eq!(bytes[0], SCHEMA_VERSION);
}
#[test]
fn deserialize_rejects_unknown_version() {
let value = BacklinkValue {
source_uri: "at://x".to_string(),
path: path_to_discriminant(BacklinkPath::Subject),
};
let mut bytes = value.serialize();
bytes[0] = 99;
assert!(BacklinkValue::deserialize(&bytes).is_none());
}
#[test]
fn deserialize_rejects_empty() {
assert!(BacklinkValue::deserialize(&[]).is_none());
}
#[test]
fn backlink_key_roundtrip() {
let hash = UserHash::from_raw(0xCAFE_BABE_DEAD_BEEF);
let key = backlink_key(
"at://did:plc:target/app.bsky.feed.post/3k2abc",
hash,
"app.bsky.feed.like",
"3k2xyz",
);
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::BACKLINKS.raw()));
assert_eq!(
reader.string(),
Some("at://did:plc:target/app.bsky.feed.post/3k2abc".to_string())
);
assert_eq!(reader.u64(), Some(0xCAFE_BABE_DEAD_BEEF));
assert_eq!(reader.string(), Some("app.bsky.feed.like".to_string()));
assert_eq!(reader.string(), Some("3k2xyz".to_string()));
assert!(reader.is_empty());
}
#[test]
fn backlink_keys_sort_by_target_then_user_then_collection_then_rkey() {
let h1 = UserHash::from_raw(1);
let h2 = UserHash::from_raw(2);
let k1 = backlink_key("aaa", h1, "col_a", "r1");
let k2 = backlink_key("aaa", h1, "col_a", "r2");
let k3 = backlink_key("aaa", h1, "col_b", "r1");
let k4 = backlink_key("aaa", h2, "col_a", "r1");
let k5 = backlink_key("bbb", h1, "col_a", "r1");
assert!(k1.as_slice() < k2.as_slice());
assert!(k2.as_slice() < k3.as_slice());
assert!(k3.as_slice() < k4.as_slice());
assert!(k4.as_slice() < k5.as_slice());
}
#[test]
fn target_prefix_is_prefix_of_full_key() {
let hash = UserHash::from_raw(42);
let prefix = backlink_target_prefix("did:plc:target");
let full = backlink_key("did:plc:target", hash, "col", "rk");
assert!(full.as_slice().starts_with(prefix.as_slice()));
}
#[test]
fn by_user_key_roundtrip() {
let hash = UserHash::from_raw(0xDEAD_BEEF_1234_5678);
let key = backlink_by_user_key(hash, "app.bsky.feed.like", "3k2abc", "did:plc:target");
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::BACKLINK_BY_USER.raw()));
assert_eq!(reader.u64(), Some(0xDEAD_BEEF_1234_5678));
assert_eq!(reader.string(), Some("app.bsky.feed.like".to_string()));
assert_eq!(reader.string(), Some("3k2abc".to_string()));
assert_eq!(reader.string(), Some("did:plc:target".to_string()));
assert!(reader.is_empty());
}
#[test]
fn by_user_prefix_is_prefix_of_full_key() {
let hash = UserHash::from_raw(42);
let prefix = backlink_by_user_prefix(hash);
let full = backlink_by_user_key(hash, "col", "rk", "target");
assert!(full.as_slice().starts_with(prefix.as_slice()));
}
#[test]
fn by_user_record_prefix_is_prefix_of_full_key() {
let hash = UserHash::from_raw(42);
let prefix = backlink_by_user_record_prefix(hash, "col", "rk");
let full = backlink_by_user_key(hash, "col", "rk", "target");
assert!(full.as_slice().starts_with(prefix.as_slice()));
}
#[test]
fn same_rkey_different_collection_produces_distinct_keys() {
let hash = UserHash::from_raw(42);
let k1 = backlink_key("target", hash, "app.bsky.feed.like", "self");
let k2 = backlink_key("target", hash, "app.bsky.graph.follow", "self");
assert_ne!(k1.as_slice(), k2.as_slice());
let r1 = backlink_by_user_key(hash, "app.bsky.feed.like", "self", "target");
let r2 = backlink_by_user_key(hash, "app.bsky.graph.follow", "self", "target");
assert_ne!(r1.as_slice(), r2.as_slice());
}
}
@@ -0,0 +1,785 @@
use std::collections::{BTreeMap, BTreeSet};
use std::ops::Bound;
use std::sync::Arc;
use fjall::{Database, Keyspace};
use smallvec::SmallVec;
use uuid::Uuid;
use super::MetastoreError;
use super::blobs::{BlobMetaValue, blob_by_cid_key, blob_meta_key, blob_user_prefix, blobs_prefix};
use super::commit_ops::{RecordBlobsValue, record_blobs_user_prefix};
use super::encoding::{KeyReader, exclusive_upper_bound};
use super::keys::{KeyTag, UserHash};
use super::repo_ops::bytes_to_cid_link;
use super::scan::{count_prefix, point_lookup};
use super::user_hash::UserHashMap;
use tranquil_types::CidLink;
const DELETE_BATCH_SIZE: usize = 1024;
pub struct BlobOps {
db: Database,
repo_data: Keyspace,
user_hashes: Arc<UserHashMap>,
}
impl BlobOps {
pub fn new(db: Database, repo_data: Keyspace, user_hashes: Arc<UserHashMap>) -> Self {
Self {
db,
repo_data,
user_hashes,
}
}
fn resolve_user_hash(&self, user_id: Uuid) -> Result<UserHash, MetastoreError> {
self.user_hashes
.get(&user_id)
.ok_or(MetastoreError::InvalidInput("unknown user_id"))
}
pub fn insert_blob(
&self,
cid: &CidLink,
mime_type: &str,
size_bytes: i64,
created_by_user: Uuid,
storage_key: &str,
) -> Result<Option<CidLink>, MetastoreError> {
if size_bytes < 0 {
return Err(MetastoreError::InvalidInput(
"size_bytes must be non-negative",
));
}
let user_hash = self.resolve_user_hash(created_by_user)?;
let cid_str = cid.as_str();
let cid_index_key = blob_by_cid_key(cid_str);
let existing = self
.repo_data
.get(cid_index_key.as_slice())
.map_err(MetastoreError::Fjall)?;
if existing.is_some() {
return Ok(None);
}
let value = BlobMetaValue {
size_bytes,
mime_type: mime_type.to_owned(),
storage_key: storage_key.to_owned(),
takedown_ref: None,
created_at_ms: chrono::Utc::now().timestamp_millis(),
};
let primary_key = blob_meta_key(user_hash, cid_str);
let mut batch = self.db.batch();
batch.insert(&self.repo_data, primary_key.as_slice(), value.serialize());
batch.insert(
&self.repo_data,
cid_index_key.as_slice(),
user_hash.raw().to_be_bytes(),
);
batch.commit().map_err(MetastoreError::Fjall)?;
Ok(Some(cid.clone()))
}
fn lookup_user_hash_by_cid(&self, cid_str: &str) -> Result<Option<UserHash>, MetastoreError> {
let key = blob_by_cid_key(cid_str);
match self
.repo_data
.get(key.as_slice())
.map_err(MetastoreError::Fjall)?
{
Some(raw) => {
let arr: [u8; 8] = raw
.as_ref()
.try_into()
.map_err(|_| MetastoreError::CorruptData("blob_by_cid value not 8 bytes"))?;
Ok(Some(UserHash::from_raw(u64::from_be_bytes(arr))))
}
None => Ok(None),
}
}
fn get_blob_value(&self, cid: &CidLink) -> Result<Option<BlobMetaValue>, MetastoreError> {
let cid_str = cid.as_str();
let user_hash = match self.lookup_user_hash_by_cid(cid_str)? {
Some(h) => h,
None => return Ok(None),
};
let key = blob_meta_key(user_hash, cid_str);
point_lookup(
&self.repo_data,
key.as_slice(),
BlobMetaValue::deserialize,
"corrupt blob_meta value",
)
}
pub fn get_blob_metadata(
&self,
cid: &CidLink,
) -> Result<Option<tranquil_db_traits::BlobMetadata>, MetastoreError> {
Ok(self
.get_blob_value(cid)?
.map(|v| tranquil_db_traits::BlobMetadata {
storage_key: v.storage_key,
mime_type: v.mime_type,
size_bytes: v.size_bytes,
}))
}
pub fn get_blob_with_takedown(
&self,
cid: &CidLink,
) -> Result<Option<tranquil_db_traits::BlobWithTakedown>, MetastoreError> {
Ok(self
.get_blob_value(cid)?
.map(|v| tranquil_db_traits::BlobWithTakedown {
cid: cid.clone(),
takedown_ref: v.takedown_ref,
}))
}
pub fn get_blob_storage_key(&self, cid: &CidLink) -> Result<Option<String>, MetastoreError> {
Ok(self.get_blob_value(cid)?.map(|v| v.storage_key))
}
pub fn list_blobs_by_user(
&self,
user_id: Uuid,
cursor: Option<&str>,
limit: usize,
) -> Result<Vec<CidLink>, MetastoreError> {
let user_hash = self.resolve_user_hash(user_id)?;
let prefix = blob_user_prefix(user_hash);
let upper = exclusive_upper_bound(prefix.as_slice())
.expect("blob user prefix always contains non-0xFF bytes");
let range_start: SmallVec<[u8; 128]> = match cursor {
Some(c) => {
let mut cursor_key = blob_meta_key(user_hash, c);
cursor_key.push(0x00);
cursor_key
}
None => prefix,
};
self.repo_data
.range(range_start.as_slice()..upper.as_slice())
.map(|guard| {
let (key_bytes, _) = guard.into_inner().map_err(MetastoreError::Fjall)?;
parse_blob_cid_from_key(key_bytes.as_ref())
})
.take(limit)
.collect()
}
pub fn count_blobs_by_user(&self, user_id: Uuid) -> Result<i64, MetastoreError> {
let user_hash = self.resolve_user_hash(user_id)?;
let prefix = blob_user_prefix(user_hash);
count_prefix(&self.repo_data, prefix.as_slice())
}
pub fn sum_blob_storage(&self) -> Result<i64, MetastoreError> {
let prefix = blobs_prefix();
self.repo_data
.prefix(prefix.as_slice())
.try_fold(0i64, |acc, guard| {
let (_, val_bytes) = guard.into_inner().map_err(MetastoreError::Fjall)?;
let value = BlobMetaValue::deserialize(&val_bytes)
.ok_or(MetastoreError::CorruptData("corrupt blob_meta in sum"))?;
Ok::<_, MetastoreError>(acc.saturating_add(value.size_bytes))
})
}
pub fn update_blob_takedown(
&self,
cid: &CidLink,
takedown_ref: Option<&str>,
) -> Result<bool, MetastoreError> {
let cid_str = cid.as_str();
let user_hash = match self.lookup_user_hash_by_cid(cid_str)? {
Some(h) => h,
None => return Ok(false),
};
let key = blob_meta_key(user_hash, cid_str);
let mut value = match point_lookup(
&self.repo_data,
key.as_slice(),
BlobMetaValue::deserialize,
"corrupt blob_meta value",
)? {
Some(v) => v,
None => return Ok(false),
};
value.takedown_ref = takedown_ref.map(str::to_owned);
let mut batch = self.db.batch();
batch.insert(&self.repo_data, key.as_slice(), value.serialize());
batch.commit().map_err(MetastoreError::Fjall)?;
Ok(true)
}
pub fn delete_blob_by_cid(&self, cid: &CidLink) -> Result<bool, MetastoreError> {
let cid_str = cid.as_str();
let user_hash = match self.lookup_user_hash_by_cid(cid_str)? {
Some(h) => h,
None => return Ok(false),
};
let primary_key = blob_meta_key(user_hash, cid_str);
let exists = self
.repo_data
.get(primary_key.as_slice())
.map_err(MetastoreError::Fjall)?
.is_some();
if !exists {
return Ok(false);
}
let cid_index_key = blob_by_cid_key(cid_str);
let mut batch = self.db.batch();
batch.remove(&self.repo_data, primary_key.as_slice());
batch.remove(&self.repo_data, cid_index_key.as_slice());
batch.commit().map_err(MetastoreError::Fjall)?;
Ok(true)
}
pub fn delete_blobs_by_user(&self, user_id: Uuid) -> Result<u64, MetastoreError> {
let user_hash = self.resolve_user_hash(user_id)?;
let prefix = blob_user_prefix(user_hash);
let user_hash_bytes = user_hash.raw().to_be_bytes();
let (final_batch, remaining, total) = self
.repo_data
.prefix(prefix.as_slice())
.map(|guard| {
let (key_bytes, _) = guard.into_inner().map_err(MetastoreError::Fjall)?;
parse_blob_cid_from_key(key_bytes.as_ref()).map(|c| c.as_str().to_owned())
})
.try_fold(
(self.db.batch(), 0usize, 0u64),
|(mut batch, count, total), entry: Result<_, MetastoreError>| {
let cid_str = entry?;
batch.remove(
&self.repo_data,
blob_meta_key(user_hash, &cid_str).as_slice(),
);
let cid_index_key = blob_by_cid_key(&cid_str);
let owns_cid = self
.repo_data
.get(cid_index_key.as_slice())
.map_err(MetastoreError::Fjall)?
.is_some_and(|raw| raw.as_ref() == user_hash_bytes);
if owns_cid {
batch.remove(&self.repo_data, cid_index_key.as_slice());
}
let new_count = count + 1;
if new_count >= DELETE_BATCH_SIZE {
batch.commit().map_err(MetastoreError::Fjall)?;
let flushed = u64::try_from(new_count).unwrap_or(u64::MAX);
Ok::<_, MetastoreError>((self.db.batch(), 0, total.saturating_add(flushed)))
} else {
Ok((batch, new_count, total))
}
},
)?;
if remaining > 0 {
final_batch.commit().map_err(MetastoreError::Fjall)?;
let flushed = u64::try_from(remaining).unwrap_or(u64::MAX);
Ok(total.saturating_add(flushed))
} else {
Ok(total)
}
}
pub fn get_blob_storage_keys_by_user(
&self,
user_id: Uuid,
) -> Result<Vec<String>, MetastoreError> {
let user_hash = self.resolve_user_hash(user_id)?;
let prefix = blob_user_prefix(user_hash);
self.repo_data
.prefix(prefix.as_slice())
.map(|guard| {
let (_, val_bytes) = guard.into_inner().map_err(MetastoreError::Fjall)?;
let value = BlobMetaValue::deserialize(&val_bytes)
.ok_or(MetastoreError::CorruptData("corrupt blob_meta value"))?;
Ok(value.storage_key)
})
.collect()
}
pub fn list_missing_blobs(
&self,
repo_id: Uuid,
cursor: Option<&str>,
limit: usize,
) -> Result<Vec<tranquil_db_traits::MissingBlobInfo>, MetastoreError> {
let user_hash = self.resolve_user_hash(repo_id)?;
let rb_prefix = record_blobs_user_prefix(user_hash);
let missing: BTreeMap<String, String> = self
.repo_data
.prefix(rb_prefix.as_slice())
.try_fold(BTreeMap::new(), |mut acc, guard| {
let (key_bytes, val_bytes) = guard.into_inner().map_err(MetastoreError::Fjall)?;
let record_uri = parse_record_blobs_uri(&key_bytes)
.ok_or(MetastoreError::CorruptData("corrupt record_blobs key"))?;
let blob_cid_bytes = RecordBlobsValue::deserialize(&val_bytes)
.map(|v| v.blob_cid_bytes)
.ok_or(MetastoreError::CorruptData("corrupt record_blobs value"))?;
blob_cid_bytes.into_iter().try_for_each(
|cid_bytes| -> Result<(), MetastoreError> {
let cid_link = bytes_to_cid_link(&cid_bytes)?;
let cid_str = cid_link.as_str().to_owned();
if acc.contains_key(&cid_str) {
return Ok(());
}
let key = blob_meta_key(user_hash, &cid_str);
let exists = self
.repo_data
.get(key.as_slice())
.map_err(MetastoreError::Fjall)?
.is_some();
if !exists {
acc.insert(cid_str, record_uri.clone());
}
Ok(())
},
)?;
Ok::<_, MetastoreError>(acc)
})?;
let start = cursor.map_or(Bound::Unbounded, Bound::Excluded);
Ok(missing
.range::<str, _>((start, Bound::Unbounded))
.take(limit)
.map(|(cid_str, uri)| tranquil_db_traits::MissingBlobInfo {
blob_cid: CidLink::from(cid_str.clone()),
record_uri: tranquil_types::AtUri::from(uri.clone()),
})
.collect())
}
fn collect_referenced_cid_bytes(
&self,
user_hash: UserHash,
) -> Result<BTreeSet<Vec<u8>>, MetastoreError> {
let rb_prefix = record_blobs_user_prefix(user_hash);
self.repo_data
.prefix(rb_prefix.as_slice())
.try_fold(BTreeSet::new(), |mut acc, guard| {
let (_, val_bytes) = guard.into_inner().map_err(MetastoreError::Fjall)?;
let blob_cids = RecordBlobsValue::deserialize(&val_bytes)
.map(|v| v.blob_cid_bytes)
.ok_or(MetastoreError::CorruptData("corrupt record_blobs value"))?;
acc.extend(blob_cids);
Ok::<_, MetastoreError>(acc)
})
}
pub fn count_distinct_record_blobs(&self, repo_id: Uuid) -> Result<i64, MetastoreError> {
let user_hash = self.resolve_user_hash(repo_id)?;
let distinct = self.collect_referenced_cid_bytes(user_hash)?;
Ok(i64::try_from(distinct.len()).unwrap_or(i64::MAX))
}
pub fn get_blobs_for_export(
&self,
repo_id: Uuid,
) -> Result<Vec<tranquil_db_traits::BlobForExport>, MetastoreError> {
let user_hash = self.resolve_user_hash(repo_id)?;
let referenced_cids = self.collect_referenced_cid_bytes(user_hash)?;
referenced_cids
.into_iter()
.filter_map(|cid_bytes| {
let cid_link = match bytes_to_cid_link(&cid_bytes) {
Ok(c) => c,
Err(e) => return Some(Err(e)),
};
let key = blob_meta_key(user_hash, cid_link.as_str());
match point_lookup(
&self.repo_data,
key.as_slice(),
BlobMetaValue::deserialize,
"corrupt blob_meta value",
) {
Ok(Some(v)) => Some(Ok(tranquil_db_traits::BlobForExport {
cid: cid_link,
storage_key: v.storage_key,
mime_type: v.mime_type,
})),
Ok(None) => None,
Err(e) => Some(Err(e)),
}
})
.collect()
}
}
fn parse_blob_cid_from_key(key: &[u8]) -> Result<CidLink, MetastoreError> {
let mut reader = KeyReader::new(key);
let tag = reader
.tag()
.ok_or(MetastoreError::CorruptData("corrupt blob key: missing tag"))?;
if tag != KeyTag::BLOBS.raw() {
return Err(MetastoreError::CorruptData(
"corrupt blob key: unexpected tag",
));
}
reader.u64().ok_or(MetastoreError::CorruptData(
"corrupt blob key: missing user_hash",
))?;
reader
.string()
.map(CidLink::from)
.ok_or(MetastoreError::CorruptData("corrupt blob key: missing cid"))
}
fn parse_record_blobs_uri(key: &[u8]) -> Option<String> {
let mut reader = KeyReader::new(key);
let tag = reader.tag()?;
if tag != KeyTag::RECORD_BLOBS.raw() {
return None;
}
reader.u64()?;
reader.string()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metastore::{Metastore, MetastoreConfig};
fn open_fresh() -> (tempfile::TempDir, Metastore) {
let dir = tempfile::TempDir::new().unwrap();
let ms = Metastore::open(
dir.path(),
MetastoreConfig {
cache_size_bytes: 64 * 1024 * 1024,
},
)
.unwrap();
(dir, ms)
}
fn setup_user(ms: &Metastore) -> (Uuid, UserHash) {
let user_id = Uuid::new_v4();
let did = format!("did:plc:blob_test_{}", user_id);
let user_hash = UserHash::from_did(&did);
let mut batch = ms.database().batch();
ms.user_hashes()
.stage_insert(&mut batch, user_id, user_hash)
.unwrap();
batch.commit().unwrap();
(user_id, user_hash)
}
fn test_cid_link(seed: u8) -> CidLink {
let digest: [u8; 32] = std::array::from_fn(|i| seed.wrapping_add(i as u8));
let mh = multihash::Multihash::<64>::wrap(0x12, &digest).unwrap();
let c = cid::Cid::new_v1(0x71, mh);
CidLink::from_cid(&c)
}
#[test]
fn insert_and_get_metadata_roundtrip() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cid = test_cid_link(1);
let result = ops
.insert_blob(&cid, "image/png", 1024, user_id, "blobs/a/b")
.unwrap();
assert_eq!(result, Some(cid.clone()));
let meta = ops.get_blob_metadata(&cid).unwrap().unwrap();
assert_eq!(meta.storage_key, "blobs/a/b");
assert_eq!(meta.mime_type, "image/png");
assert_eq!(meta.size_bytes, 1024);
}
#[test]
fn insert_duplicate_returns_none() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cid = test_cid_link(2);
assert!(
ops.insert_blob(&cid, "image/png", 100, user_id, "k1")
.unwrap()
.is_some()
);
assert!(
ops.insert_blob(&cid, "image/png", 100, user_id, "k1")
.unwrap()
.is_none()
);
}
#[test]
fn insert_same_cid_different_user_returns_none() {
let (_dir, ms) = open_fresh();
let (user_a, _) = setup_user(&ms);
let (user_b, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cid = test_cid_link(80);
assert!(
ops.insert_blob(&cid, "image/png", 100, user_a, "ka")
.unwrap()
.is_some()
);
assert!(
ops.insert_blob(&cid, "image/png", 100, user_b, "kb")
.unwrap()
.is_none()
);
}
#[test]
fn get_blob_with_takedown_no_takedown() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cid = test_cid_link(3);
ops.insert_blob(&cid, "text/plain", 10, user_id, "k")
.unwrap();
let result = ops.get_blob_with_takedown(&cid).unwrap().unwrap();
assert_eq!(result.cid, cid);
assert!(result.takedown_ref.is_none());
}
#[test]
fn update_takedown_and_read_back() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cid = test_cid_link(4);
ops.insert_blob(&cid, "text/plain", 10, user_id, "k")
.unwrap();
assert!(ops.update_blob_takedown(&cid, Some("mod-42")).unwrap());
let result = ops.get_blob_with_takedown(&cid).unwrap().unwrap();
assert_eq!(result.takedown_ref.as_deref(), Some("mod-42"));
assert!(ops.update_blob_takedown(&cid, None).unwrap());
let result = ops.get_blob_with_takedown(&cid).unwrap().unwrap();
assert!(result.takedown_ref.is_none());
}
#[test]
fn update_takedown_nonexistent_returns_false() {
let (_dir, ms) = open_fresh();
let ops = ms.blob_ops();
let cid = test_cid_link(99);
assert!(!ops.update_blob_takedown(&cid, Some("x")).unwrap());
}
#[test]
fn get_blob_storage_key() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cid = test_cid_link(5);
ops.insert_blob(&cid, "image/jpeg", 500, user_id, "blobs/x/y")
.unwrap();
assert_eq!(
ops.get_blob_storage_key(&cid).unwrap().as_deref(),
Some("blobs/x/y")
);
}
#[test]
fn list_blobs_by_user_with_pagination() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cids: Vec<CidLink> = (0..5).map(|i| test_cid_link(10 + i)).collect();
cids.iter().enumerate().for_each(|(i, cid)| {
ops.insert_blob(cid, "image/png", i as i64, user_id, &format!("k{i}"))
.unwrap();
});
let page1 = ops.list_blobs_by_user(user_id, None, 3).unwrap();
assert_eq!(page1.len(), 3);
let cursor = page1.last().unwrap().as_str();
let page2 = ops.list_blobs_by_user(user_id, Some(cursor), 3).unwrap();
assert_eq!(page2.len(), 2);
let all = ops.list_blobs_by_user(user_id, None, 100).unwrap();
assert_eq!(all.len(), 5);
}
#[test]
fn count_blobs_by_user() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
assert_eq!(ops.count_blobs_by_user(user_id).unwrap(), 0);
(0..3).for_each(|i| {
ops.insert_blob(
&test_cid_link(20 + i),
"image/png",
100,
user_id,
&format!("k{i}"),
)
.unwrap();
});
assert_eq!(ops.count_blobs_by_user(user_id).unwrap(), 3);
}
#[test]
fn sum_blob_storage() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
assert_eq!(ops.sum_blob_storage().unwrap(), 0);
ops.insert_blob(&test_cid_link(30), "a/b", 100, user_id, "k0")
.unwrap();
ops.insert_blob(&test_cid_link(31), "a/b", 250, user_id, "k1")
.unwrap();
assert_eq!(ops.sum_blob_storage().unwrap(), 350);
}
#[test]
fn delete_blob_by_cid() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cid = test_cid_link(40);
ops.insert_blob(&cid, "image/png", 100, user_id, "k")
.unwrap();
assert!(ops.delete_blob_by_cid(&cid).unwrap());
assert!(ops.get_blob_metadata(&cid).unwrap().is_none());
assert!(!ops.delete_blob_by_cid(&cid).unwrap());
}
#[test]
fn delete_blob_cleans_up_indexes() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cid = test_cid_link(41);
ops.insert_blob(&cid, "image/png", 100, user_id, "storage/abc")
.unwrap();
assert!(ops.get_blob_storage_key(&cid).unwrap().is_some());
ops.delete_blob_by_cid(&cid).unwrap();
assert!(ops.lookup_user_hash_by_cid(cid.as_str()).unwrap().is_none());
}
#[test]
fn delete_blobs_by_user() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
(0..4).for_each(|i| {
ops.insert_blob(&test_cid_link(50 + i), "a/b", 10, user_id, &format!("k{i}"))
.unwrap();
});
let deleted = ops.delete_blobs_by_user(user_id).unwrap();
assert_eq!(deleted, 4);
assert_eq!(ops.count_blobs_by_user(user_id).unwrap(), 0);
}
#[test]
fn delete_blobs_by_user_cleans_all_indexes() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let cid = test_cid_link(81);
ops.insert_blob(&cid, "a/b", 10, user_id, "storage/del_test")
.unwrap();
ops.delete_blobs_by_user(user_id).unwrap();
assert!(ops.lookup_user_hash_by_cid(cid.as_str()).unwrap().is_none());
}
#[test]
fn insert_blob_rejects_negative_size() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
let result = ops.insert_blob(&test_cid_link(90), "a/b", -1, user_id, "k");
assert!(result.is_err());
}
#[test]
fn get_blob_storage_keys_by_user() {
let (_dir, ms) = open_fresh();
let (user_id, _) = setup_user(&ms);
let ops = ms.blob_ops();
ops.insert_blob(&test_cid_link(60), "a/b", 10, user_id, "alpha")
.unwrap();
ops.insert_blob(&test_cid_link(61), "a/b", 10, user_id, "beta")
.unwrap();
let mut keys = ops.get_blob_storage_keys_by_user(user_id).unwrap();
keys.sort();
assert_eq!(keys, vec!["alpha", "beta"]);
}
#[test]
fn get_metadata_for_nonexistent_returns_none() {
let (_dir, ms) = open_fresh();
let ops = ms.blob_ops();
assert!(ops.get_blob_metadata(&test_cid_link(99)).unwrap().is_none());
}
#[test]
fn blobs_isolated_between_users() {
let (_dir, ms) = open_fresh();
let (user_a, _) = setup_user(&ms);
let (user_b, _) = setup_user(&ms);
let ops = ms.blob_ops();
ops.insert_blob(&test_cid_link(70), "a/b", 10, user_a, "ka")
.unwrap();
ops.insert_blob(&test_cid_link(71), "a/b", 20, user_b, "kb")
.unwrap();
assert_eq!(ops.count_blobs_by_user(user_a).unwrap(), 1);
assert_eq!(ops.count_blobs_by_user(user_b).unwrap(), 1);
assert_eq!(ops.sum_blob_storage().unwrap(), 30);
}
}
@@ -0,0 +1,144 @@
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use super::encoding::KeyBuilder;
use super::keys::{KeyTag, UserHash};
const BLOB_META_SCHEMA_VERSION: u8 = 1;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BlobMetaValue {
pub size_bytes: i64,
pub mime_type: String,
pub storage_key: String,
pub takedown_ref: Option<String>,
pub created_at_ms: i64,
}
impl BlobMetaValue {
pub fn serialize(&self) -> Vec<u8> {
let payload = postcard::to_allocvec(self).expect("BlobMetaValue serialization cannot fail");
let mut buf = Vec::with_capacity(1 + payload.len());
buf.push(BLOB_META_SCHEMA_VERSION);
buf.extend_from_slice(&payload);
buf
}
pub fn deserialize(bytes: &[u8]) -> Option<Self> {
let (&version, payload) = bytes.split_first()?;
match version {
BLOB_META_SCHEMA_VERSION => postcard::from_bytes(payload).ok(),
_ => None,
}
}
}
pub fn blob_meta_key(user_hash: UserHash, cid_str: &str) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::BLOBS)
.u64(user_hash.raw())
.string(cid_str)
.build()
}
pub fn blob_user_prefix(user_hash: UserHash) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::BLOBS)
.u64(user_hash.raw())
.build()
}
pub fn blobs_prefix() -> SmallVec<[u8; 128]> {
KeyBuilder::new().tag(KeyTag::BLOBS).build()
}
pub fn blob_by_cid_key(cid_str: &str) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::BLOB_BY_CID)
.string(cid_str)
.build()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metastore::encoding::KeyReader;
#[test]
fn blob_meta_value_roundtrip() {
let val = BlobMetaValue {
size_bytes: 1024,
mime_type: "image/png".to_owned(),
storage_key: "blobs/abc/def".to_owned(),
takedown_ref: None,
created_at_ms: 1700000000000,
};
let bytes = val.serialize();
let decoded = BlobMetaValue::deserialize(&bytes).unwrap();
assert_eq!(val, decoded);
}
#[test]
fn blob_meta_value_with_takedown_roundtrip() {
let val = BlobMetaValue {
size_bytes: 42,
mime_type: "text/plain".to_owned(),
storage_key: "k".to_owned(),
takedown_ref: Some("mod-123".to_owned()),
created_at_ms: 0,
};
let bytes = val.serialize();
let decoded = BlobMetaValue::deserialize(&bytes).unwrap();
assert_eq!(val, decoded);
}
#[test]
fn blob_meta_key_roundtrip() {
let uh = UserHash::from_did("did:plc:test");
let key = blob_meta_key(uh, "bafyreiabc");
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::BLOBS.raw()));
assert_eq!(reader.u64(), Some(uh.raw()));
assert_eq!(reader.string(), Some("bafyreiabc".to_owned()));
assert!(reader.is_empty());
}
#[test]
fn blob_user_prefix_is_prefix_of_key() {
let uh = UserHash::from_did("did:plc:test");
let prefix = blob_user_prefix(uh);
let key = blob_meta_key(uh, "bafyreiabc");
assert!(key.starts_with(prefix.as_slice()));
}
#[test]
fn blob_by_cid_key_roundtrip() {
let key = blob_by_cid_key("bafyreiabc");
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::BLOB_BY_CID.raw()));
assert_eq!(reader.string(), Some("bafyreiabc".to_owned()));
assert!(reader.is_empty());
}
#[test]
fn blob_meta_key_ordering_by_cid() {
let uh = UserHash::from_did("did:plc:test");
let key_a = blob_meta_key(uh, "aaa");
let key_b = blob_meta_key(uh, "bbb");
assert!(key_a.as_slice() < key_b.as_slice());
}
#[test]
fn deserialize_unknown_version_returns_none() {
let val = BlobMetaValue {
size_bytes: 0,
mime_type: String::new(),
storage_key: String::new(),
takedown_ref: None,
created_at_ms: 0,
};
let mut bytes = val.serialize();
bytes[0] = 99;
assert!(BlobMetaValue::deserialize(&bytes).is_none());
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,599 @@
use smallvec::SmallVec;
const NULL_ESCAPE: u8 = 0x01;
const NULL_TERMINATOR: [u8; 2] = [0x00, 0x00];
pub fn encode_u64(buf: &mut SmallVec<[u8; 128]>, value: u64) {
buf.extend_from_slice(&value.to_be_bytes());
}
pub fn decode_u64(src: &[u8]) -> Option<(u64, &[u8])> {
let (bytes, rest) = src.split_first_chunk::<8>()?;
Some((u64::from_be_bytes(*bytes), rest))
}
pub fn encode_i64(buf: &mut SmallVec<[u8; 128]>, value: i64) {
let encoded = (value as u64) ^ (1u64 << 63);
buf.extend_from_slice(&encoded.to_be_bytes());
}
pub fn decode_i64(src: &[u8]) -> Option<(i64, &[u8])> {
let (bytes, rest) = src.split_first_chunk::<8>()?;
let raw = u64::from_be_bytes(*bytes) ^ (1u64 << 63);
Some((raw as i64, rest))
}
pub fn encode_u32(buf: &mut SmallVec<[u8; 128]>, value: u32) {
buf.extend_from_slice(&value.to_be_bytes());
}
pub fn decode_u32(src: &[u8]) -> Option<(u32, &[u8])> {
let (bytes, rest) = src.split_first_chunk::<4>()?;
Some((u32::from_be_bytes(*bytes), rest))
}
pub fn encode_u16(buf: &mut SmallVec<[u8; 128]>, value: u16) {
buf.extend_from_slice(&value.to_be_bytes());
}
pub fn decode_u16(src: &[u8]) -> Option<(u16, &[u8])> {
let (bytes, rest) = src.split_first_chunk::<2>()?;
Some((u16::from_be_bytes(*bytes), rest))
}
pub fn encode_bool(buf: &mut SmallVec<[u8; 128]>, value: bool) {
buf.push(u8::from(value));
}
pub fn decode_bool(src: &[u8]) -> Option<(bool, &[u8])> {
let (&byte, rest) = src.split_first()?;
match byte {
0 => Some((false, rest)),
1 => Some((true, rest)),
_ => None,
}
}
pub fn encode_bytes(buf: &mut SmallVec<[u8; 128]>, value: &[u8]) {
value.iter().for_each(|&b| match b {
0x00 => {
buf.push(0x00);
buf.push(NULL_ESCAPE);
}
other => buf.push(other),
});
buf.extend_from_slice(&NULL_TERMINATOR);
}
pub fn decode_bytes(src: &[u8]) -> Option<(Vec<u8>, &[u8])> {
let mut result = Vec::new();
let mut i = 0;
loop {
match src.get(i)? {
0x00 => match src.get(i + 1)? {
0x00 => return Some((result, &src[i + 2..])),
&NULL_ESCAPE => {
result.push(0x00);
i += 2;
}
_ => return None,
},
&b => {
result.push(b);
i += 1;
}
}
}
}
pub fn encode_string(buf: &mut SmallVec<[u8; 128]>, value: &str) {
encode_bytes(buf, value.as_bytes());
}
pub fn decode_string(src: &[u8]) -> Option<(String, &[u8])> {
let (bytes, rest) = decode_bytes(src)?;
String::from_utf8(bytes).ok().map(|s| (s, rest))
}
pub struct KeyBuilder(SmallVec<[u8; 128]>);
impl KeyBuilder {
pub fn new() -> Self {
Self(SmallVec::new())
}
pub fn with_capacity(cap: usize) -> Self {
Self(SmallVec::with_capacity(cap))
}
pub fn u64(mut self, value: u64) -> Self {
encode_u64(&mut self.0, value);
self
}
pub fn i64(mut self, value: i64) -> Self {
encode_i64(&mut self.0, value);
self
}
pub fn u32(mut self, value: u32) -> Self {
encode_u32(&mut self.0, value);
self
}
pub fn u16(mut self, value: u16) -> Self {
encode_u16(&mut self.0, value);
self
}
pub fn bool(mut self, value: bool) -> Self {
encode_bool(&mut self.0, value);
self
}
pub fn bytes(mut self, value: &[u8]) -> Self {
encode_bytes(&mut self.0, value);
self
}
pub fn string(mut self, value: &str) -> Self {
encode_string(&mut self.0, value);
self
}
pub fn tag(mut self, tag: super::keys::KeyTag) -> Self {
self.0.push(tag.raw());
self
}
pub fn fixed<const N: usize>(mut self, bytes: &[u8; N]) -> Self {
self.0.extend_from_slice(bytes);
self
}
pub fn raw(mut self, bytes: &[u8]) -> Self {
self.0.extend_from_slice(bytes);
self
}
pub fn build(self) -> SmallVec<[u8; 128]> {
self.0
}
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl Default for KeyBuilder {
fn default() -> Self {
Self::new()
}
}
pub struct KeyReader<'a>(&'a [u8]);
impl<'a> KeyReader<'a> {
pub fn new(src: &'a [u8]) -> Self {
Self(src)
}
pub fn u64(&mut self) -> Option<u64> {
let (val, rest) = decode_u64(self.0)?;
self.0 = rest;
Some(val)
}
pub fn i64(&mut self) -> Option<i64> {
let (val, rest) = decode_i64(self.0)?;
self.0 = rest;
Some(val)
}
pub fn u32(&mut self) -> Option<u32> {
let (val, rest) = decode_u32(self.0)?;
self.0 = rest;
Some(val)
}
pub fn u16(&mut self) -> Option<u16> {
let (val, rest) = decode_u16(self.0)?;
self.0 = rest;
Some(val)
}
pub fn bool(&mut self) -> Option<bool> {
let (val, rest) = decode_bool(self.0)?;
self.0 = rest;
Some(val)
}
pub fn bytes(&mut self) -> Option<Vec<u8>> {
let (val, rest) = decode_bytes(self.0)?;
self.0 = rest;
Some(val)
}
pub fn string(&mut self) -> Option<String> {
let (val, rest) = decode_string(self.0)?;
self.0 = rest;
Some(val)
}
pub fn tag(&mut self) -> Option<u8> {
let (&tag, rest) = self.0.split_first()?;
self.0 = rest;
Some(tag)
}
pub fn remaining(&self) -> &'a [u8] {
self.0
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
}
pub fn exclusive_upper_bound(prefix: &[u8]) -> Option<SmallVec<[u8; 128]>> {
prefix.iter().rposition(|&b| b != 0xFF).map(|pos| {
let mut result = SmallVec::from_slice(&prefix[..=pos]);
result[pos] = prefix[pos].wrapping_add(1);
result
})
}
#[cfg(test)]
mod tests {
use super::*;
use proptest::prelude::*;
#[test]
fn u64_roundtrip_boundaries() {
[0u64, 1, u64::MAX / 2, u64::MAX - 1, u64::MAX]
.iter()
.for_each(|&v| {
let mut buf = SmallVec::new();
encode_u64(&mut buf, v);
let (decoded, rest) = decode_u64(&buf).unwrap();
assert_eq!(decoded, v);
assert!(rest.is_empty());
});
}
#[test]
fn i64_roundtrip_boundaries() {
[i64::MIN, -1, 0, 1, i64::MAX].iter().for_each(|&v| {
let mut buf = SmallVec::new();
encode_i64(&mut buf, v);
let (decoded, rest) = decode_i64(&buf).unwrap();
assert_eq!(decoded, v);
assert!(rest.is_empty());
});
}
#[test]
fn bool_roundtrip() {
[false, true].iter().for_each(|&v| {
let mut buf = SmallVec::new();
encode_bool(&mut buf, v);
let (decoded, rest) = decode_bool(&buf).unwrap();
assert_eq!(decoded, v);
assert!(rest.is_empty());
});
}
#[test]
fn bytes_with_nulls() {
let input = &[0x00, 0x01, 0x00, 0xFF, 0x00];
let mut buf = SmallVec::new();
encode_bytes(&mut buf, input);
let (decoded, rest) = decode_bytes(&buf).unwrap();
assert_eq!(decoded, input);
assert!(rest.is_empty());
}
#[test]
fn empty_bytes_roundtrip() {
let mut buf = SmallVec::new();
encode_bytes(&mut buf, &[]);
let (decoded, rest) = decode_bytes(&buf).unwrap();
assert!(decoded.is_empty());
assert!(rest.is_empty());
}
#[test]
fn empty_string_roundtrip() {
let mut buf = SmallVec::new();
encode_string(&mut buf, "");
let (decoded, rest) = decode_string(&buf).unwrap();
assert_eq!(decoded, "");
assert!(rest.is_empty());
}
#[test]
fn string_with_null_bytes() {
let input = "hello\x00world";
let mut buf = SmallVec::new();
encode_string(&mut buf, input);
let (decoded, rest) = decode_string(&buf).unwrap();
assert_eq!(decoded, input);
assert!(rest.is_empty());
}
#[test]
fn key_builder_composite_roundtrip() {
let key = KeyBuilder::new()
.tag(super::super::keys::KeyTag::RECORDS)
.u64(42)
.string("app.bsky.feed.post")
.string("3k2a")
.build();
let mut reader = KeyReader::new(&key);
assert_eq!(
reader.tag(),
Some(super::super::keys::KeyTag::RECORDS.raw())
);
assert_eq!(reader.u64(), Some(42));
assert_eq!(reader.string(), Some("app.bsky.feed.post".to_string()));
assert_eq!(reader.string(), Some("3k2a".to_string()));
assert!(reader.is_empty());
}
#[test]
fn key_builder_ordering_preserves_field_order() {
let key_a = KeyBuilder::new().u64(1).string("aaa").build();
let key_b = KeyBuilder::new().u64(1).string("bbb").build();
let key_c = KeyBuilder::new().u64(2).string("aaa").build();
assert!(key_a.as_slice() < key_b.as_slice());
assert!(key_b.as_slice() < key_c.as_slice());
}
#[test]
fn decode_bytes_rejects_invalid_escape() {
assert!(decode_bytes(&[0x00, 0x02]).is_none());
assert!(decode_bytes(&[0x00, 0xFF]).is_none());
assert!(decode_bytes(&[0x41, 0x00, 0x03]).is_none());
}
#[test]
fn decode_bytes_rejects_truncated_input() {
assert!(decode_bytes(&[]).is_none());
assert!(decode_bytes(&[0x00]).is_none());
assert!(decode_bytes(&[0x41]).is_none());
assert!(decode_bytes(&[0x41, 0x00]).is_none());
assert!(decode_bytes(&[0x00, 0x01]).is_none());
}
#[test]
fn decode_bool_rejects_invalid_byte() {
assert!(decode_bool(&[0x02]).is_none());
assert!(decode_bool(&[0xFF]).is_none());
assert!(decode_bool(&[]).is_none());
}
#[test]
fn decode_string_rejects_invalid_utf8() {
let mut buf = SmallVec::new();
encode_bytes(&mut buf, &[0xFF, 0xFE]);
assert!(decode_string(&buf).is_none());
}
#[test]
fn decode_u64_rejects_short_input() {
assert!(decode_u64(&[]).is_none());
assert!(decode_u64(&[0x00; 7]).is_none());
}
#[test]
fn decode_u32_rejects_short_input() {
assert!(decode_u32(&[]).is_none());
assert!(decode_u32(&[0x00; 3]).is_none());
}
#[test]
fn decode_u16_rejects_short_input() {
assert!(decode_u16(&[]).is_none());
assert!(decode_u16(&[0x00]).is_none());
}
#[test]
fn decode_i64_rejects_short_input() {
assert!(decode_i64(&[]).is_none());
assert!(decode_i64(&[0x00; 7]).is_none());
}
#[test]
fn fixed_key_roundtrip() {
let data: [u8; 4] = [0xDE, 0xAD, 0xBE, 0xEF];
let key = KeyBuilder::new()
.tag(super::super::keys::KeyTag::RECORDS)
.fixed(&data)
.build();
let mut reader = KeyReader::new(&key);
assert_eq!(
reader.tag(),
Some(super::super::keys::KeyTag::RECORDS.raw())
);
assert_eq!(reader.remaining(), &data);
}
proptest! {
#[test]
fn prop_u64_roundtrip(v: u64) {
let mut buf = SmallVec::new();
encode_u64(&mut buf, v);
let (decoded, rest) = decode_u64(&buf).unwrap();
prop_assert_eq!(decoded, v);
prop_assert!(rest.is_empty());
}
#[test]
fn prop_u64_ordering(a: u64, b: u64) {
let mut buf_a = SmallVec::new();
let mut buf_b = SmallVec::new();
encode_u64(&mut buf_a, a);
encode_u64(&mut buf_b, b);
prop_assert_eq!(buf_a.as_slice().cmp(buf_b.as_slice()), a.cmp(&b));
}
#[test]
fn prop_i64_roundtrip(v: i64) {
let mut buf = SmallVec::new();
encode_i64(&mut buf, v);
let (decoded, rest) = decode_i64(&buf).unwrap();
prop_assert_eq!(decoded, v);
prop_assert!(rest.is_empty());
}
#[test]
fn prop_i64_ordering(a: i64, b: i64) {
let mut buf_a = SmallVec::new();
let mut buf_b = SmallVec::new();
encode_i64(&mut buf_a, a);
encode_i64(&mut buf_b, b);
prop_assert_eq!(buf_a.as_slice().cmp(buf_b.as_slice()), a.cmp(&b));
}
#[test]
fn prop_u32_roundtrip(v: u32) {
let mut buf = SmallVec::new();
encode_u32(&mut buf, v);
let (decoded, rest) = decode_u32(&buf).unwrap();
prop_assert_eq!(decoded, v);
prop_assert!(rest.is_empty());
}
#[test]
fn prop_u32_ordering(a: u32, b: u32) {
let mut buf_a = SmallVec::new();
let mut buf_b = SmallVec::new();
encode_u32(&mut buf_a, a);
encode_u32(&mut buf_b, b);
prop_assert_eq!(buf_a.as_slice().cmp(buf_b.as_slice()), a.cmp(&b));
}
#[test]
fn prop_u16_roundtrip(v: u16) {
let mut buf = SmallVec::new();
encode_u16(&mut buf, v);
let (decoded, rest) = decode_u16(&buf).unwrap();
prop_assert_eq!(decoded, v);
prop_assert!(rest.is_empty());
}
#[test]
fn prop_u16_ordering(a: u16, b: u16) {
let mut buf_a = SmallVec::new();
let mut buf_b = SmallVec::new();
encode_u16(&mut buf_a, a);
encode_u16(&mut buf_b, b);
prop_assert_eq!(buf_a.as_slice().cmp(buf_b.as_slice()), a.cmp(&b));
}
#[test]
fn prop_bool_roundtrip(v: bool) {
let mut buf = SmallVec::new();
encode_bool(&mut buf, v);
let (decoded, rest) = decode_bool(&buf).unwrap();
prop_assert_eq!(decoded, v);
prop_assert!(rest.is_empty());
}
#[test]
fn prop_bool_ordering(a: bool, b: bool) {
let mut buf_a = SmallVec::new();
let mut buf_b = SmallVec::new();
encode_bool(&mut buf_a, a);
encode_bool(&mut buf_b, b);
prop_assert_eq!(buf_a.as_slice().cmp(buf_b.as_slice()), a.cmp(&b));
}
#[test]
fn prop_bytes_roundtrip(v in proptest::collection::vec(any::<u8>(), 0..256)) {
let mut buf = SmallVec::new();
encode_bytes(&mut buf, &v);
let (decoded, rest) = decode_bytes(&buf).unwrap();
prop_assert_eq!(decoded, v);
prop_assert!(rest.is_empty());
}
#[test]
fn prop_bytes_ordering(
a in proptest::collection::vec(any::<u8>(), 0..64),
b in proptest::collection::vec(any::<u8>(), 0..64),
) {
let mut buf_a = SmallVec::new();
let mut buf_b = SmallVec::new();
encode_bytes(&mut buf_a, &a);
encode_bytes(&mut buf_b, &b);
prop_assert_eq!(buf_a.as_slice().cmp(buf_b.as_slice()), a.cmp(&b));
}
#[test]
fn prop_string_roundtrip(v in "\\PC{0,128}") {
let mut buf = SmallVec::new();
encode_string(&mut buf, &v);
let (decoded, rest) = decode_string(&buf).unwrap();
prop_assert_eq!(decoded, v);
prop_assert!(rest.is_empty());
}
#[test]
fn prop_string_ordering(
a in "[\\x00-\\xff]{0,32}",
b in "[\\x00-\\xff]{0,32}",
) {
let mut buf_a = SmallVec::new();
let mut buf_b = SmallVec::new();
encode_string(&mut buf_a, &a);
encode_string(&mut buf_b, &b);
prop_assert_eq!(
buf_a.as_slice().cmp(buf_b.as_slice()),
a.as_bytes().cmp(b.as_bytes())
);
}
#[test]
fn prop_composite_roundtrip(
tag_raw in 0u8..=255,
num in any::<u64>(),
s1 in "\\PC{0,32}",
s2 in "\\PC{0,32}",
) {
let tag = super::super::keys::KeyTag::from_raw_unchecked(tag_raw);
let key = KeyBuilder::new()
.tag(tag)
.u64(num)
.string(&s1)
.string(&s2)
.build();
let mut reader = KeyReader::new(&key);
prop_assert_eq!(reader.tag(), Some(tag_raw));
prop_assert_eq!(reader.u64(), Some(num));
prop_assert_eq!(reader.string(), Some(s1));
prop_assert_eq!(reader.string(), Some(s2));
prop_assert!(reader.is_empty());
}
#[test]
fn prop_composite_ordering(
tag_raw in 0u8..=10,
a_num in any::<u64>(),
b_num in any::<u64>(),
a_str in "[a-z]{0,8}",
b_str in "[a-z]{0,8}",
) {
let tag = super::super::keys::KeyTag::from_raw_unchecked(tag_raw);
let key_a = KeyBuilder::new().tag(tag).u64(a_num).string(&a_str).build();
let key_b = KeyBuilder::new().tag(tag).u64(b_num).string(&b_str).build();
let expected = a_num.cmp(&b_num).then_with(|| a_str.as_bytes().cmp(b_str.as_bytes()));
prop_assert_eq!(key_a.as_slice().cmp(key_b.as_slice()), expected);
}
}
}
@@ -0,0 +1,230 @@
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use super::encoding::KeyBuilder;
use super::keys::{KeyTag, UserHash};
const SEQ_META_SCHEMA_VERSION: u8 = 1;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SeqMetaValue {
pub blocks_cids: Vec<String>,
}
impl SeqMetaValue {
pub fn serialize(&self) -> Vec<u8> {
let payload = postcard::to_allocvec(self).expect("SeqMetaValue serialization cannot fail");
let mut buf = Vec::with_capacity(1 + payload.len());
buf.push(SEQ_META_SCHEMA_VERSION);
buf.extend_from_slice(&payload);
buf
}
pub fn deserialize(bytes: &[u8]) -> Option<Self> {
let (&version, payload) = bytes.split_first()?;
match version {
SEQ_META_SCHEMA_VERSION => postcard::from_bytes(payload).ok(),
_ => None,
}
}
}
pub fn rev_to_seq_key(user_hash: UserHash, rev: &str) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::REV_TO_SEQ)
.u64(user_hash.raw())
.string(rev)
.build()
}
pub fn rev_to_seq_user_prefix(user_hash: UserHash) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::REV_TO_SEQ)
.u64(user_hash.raw())
.build()
}
pub fn seq_meta_key(seq: u64) -> SmallVec<[u8; 128]> {
KeyBuilder::new().tag(KeyTag::SEQ_META).u64(seq).build()
}
pub fn seq_tombstone_key(seq: u64) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::SEQ_TOMBSTONE)
.u64(seq)
.build()
}
pub fn did_events_key(user_hash: UserHash, seq: u64) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::DID_EVENTS)
.u64(user_hash.raw())
.u64(seq)
.build()
}
pub fn did_events_prefix(user_hash: UserHash) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::DID_EVENTS)
.u64(user_hash.raw())
.build()
}
pub fn metastore_cursor_key() -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::METASTORE_CURSOR)
.raw(&[0x00])
.build()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metastore::encoding::KeyReader;
#[test]
fn seq_meta_value_roundtrip() {
let value = SeqMetaValue {
blocks_cids: vec!["bafyreiblock1".to_owned(), "bafyreiblock2".to_owned()],
};
let bytes = value.serialize();
let decoded = SeqMetaValue::deserialize(&bytes).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn seq_meta_value_empty_blocks() {
let value = SeqMetaValue {
blocks_cids: vec![],
};
let bytes = value.serialize();
let decoded = SeqMetaValue::deserialize(&bytes).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn seq_meta_schema_version_first_byte() {
let value = SeqMetaValue {
blocks_cids: vec![],
};
let bytes = value.serialize();
assert_eq!(bytes[0], SEQ_META_SCHEMA_VERSION);
}
#[test]
fn seq_meta_rejects_unknown_version() {
let value = SeqMetaValue {
blocks_cids: vec![],
};
let mut bytes = value.serialize();
bytes[0] = 99;
assert!(SeqMetaValue::deserialize(&bytes).is_none());
}
#[test]
fn seq_meta_rejects_empty_input() {
assert!(SeqMetaValue::deserialize(&[]).is_none());
}
#[test]
fn rev_to_seq_key_roundtrip() {
let hash = UserHash::from_raw(0xDEAD_BEEF_CAFE_BABE);
let key = rev_to_seq_key(hash, "3k2abcde");
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::REV_TO_SEQ.raw()));
assert_eq!(reader.u64(), Some(0xDEAD_BEEF_CAFE_BABE));
assert_eq!(reader.string(), Some("3k2abcde".to_owned()));
assert!(reader.is_empty());
}
#[test]
fn rev_to_seq_keys_sort_by_user_then_rev() {
let h1 = UserHash::from_raw(1);
let h2 = UserHash::from_raw(2);
let k1 = rev_to_seq_key(h1, "abc");
let k2 = rev_to_seq_key(h1, "def");
let k3 = rev_to_seq_key(h2, "abc");
assert!(k1.as_slice() < k2.as_slice());
assert!(k2.as_slice() < k3.as_slice());
}
#[test]
fn rev_to_seq_user_prefix_is_prefix_of_full_key() {
let hash = UserHash::from_raw(42);
let prefix = rev_to_seq_user_prefix(hash);
let full = rev_to_seq_key(hash, "some_rev");
assert!(full.as_slice().starts_with(prefix.as_slice()));
}
#[test]
fn seq_meta_key_roundtrip() {
let key = seq_meta_key(12345);
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::SEQ_META.raw()));
assert_eq!(reader.u64(), Some(12345));
assert!(reader.is_empty());
}
#[test]
fn seq_meta_keys_sort_by_seq() {
let k1 = seq_meta_key(1);
let k2 = seq_meta_key(2);
let k3 = seq_meta_key(100);
assert!(k1.as_slice() < k2.as_slice());
assert!(k2.as_slice() < k3.as_slice());
}
#[test]
fn seq_tombstone_key_roundtrip() {
let key = seq_tombstone_key(999);
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::SEQ_TOMBSTONE.raw()));
assert_eq!(reader.u64(), Some(999));
assert!(reader.is_empty());
}
#[test]
fn did_events_key_roundtrip() {
let hash = UserHash::from_raw(0xCAFE_BABE_DEAD_BEEF);
let key = did_events_key(hash, 42);
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::DID_EVENTS.raw()));
assert_eq!(reader.u64(), Some(0xCAFE_BABE_DEAD_BEEF));
assert_eq!(reader.u64(), Some(42));
assert!(reader.is_empty());
}
#[test]
fn did_events_keys_sort_by_user_then_seq() {
let h1 = UserHash::from_raw(1);
let h2 = UserHash::from_raw(2);
let k1 = did_events_key(h1, 10);
let k2 = did_events_key(h1, 20);
let k3 = did_events_key(h2, 5);
assert!(k1.as_slice() < k2.as_slice());
assert!(k2.as_slice() < k3.as_slice());
}
#[test]
fn did_events_prefix_is_prefix_of_full_key() {
let hash = UserHash::from_raw(99);
let prefix = did_events_prefix(hash);
let full = did_events_key(hash, 1);
assert!(full.as_slice().starts_with(prefix.as_slice()));
}
#[test]
fn metastore_cursor_key_roundtrip() {
let key = metastore_cursor_key();
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::METASTORE_CURSOR.raw()));
assert_eq!(reader.remaining(), &[0x00]);
}
#[test]
fn metastore_cursor_key_is_stable() {
let k1 = metastore_cursor_key();
let k2 = metastore_cursor_key();
assert_eq!(k1.as_slice(), k2.as_slice());
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+149
View File
@@ -0,0 +1,149 @@
use serde::{Deserialize, Serialize};
use siphasher::sip::SipHasher24;
use std::hash::Hasher;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct UserHash(u64);
const SIPHASH_KEY0: u64 = 0x7472_616e_7175_696c;
const SIPHASH_KEY1: u64 = 0x7064_735f_7573_6572;
impl UserHash {
pub fn from_did(did: &str) -> Self {
let mut hasher = SipHasher24::new_with_keys(SIPHASH_KEY0, SIPHASH_KEY1);
hasher.write(did.as_bytes());
Self(hasher.finish())
}
pub fn from_raw(raw: u64) -> Self {
Self(raw)
}
pub fn raw(self) -> u64 {
self.0
}
}
impl std::fmt::Display for UserHash {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{:016x}", self.0)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct KeyTag(u8);
impl KeyTag {
pub const REPO_META: Self = Self(0x01);
pub const RECORDS: Self = Self(0x02);
pub const USER_BLOCKS: Self = Self(0x03);
pub const HANDLES: Self = Self(0x04);
pub const BLOBS: Self = Self(0x05);
pub const BACKLINKS: Self = Self(0x06);
pub const BLOB_BY_CID: Self = Self(0x07);
pub const USER_MAP: Self = Self(0x10);
pub const USER_MAP_REVERSE: Self = Self(0x11);
pub const REV_TO_SEQ: Self = Self(0x20);
pub const SEQ_META: Self = Self(0x21);
pub const SEQ_TOMBSTONE: Self = Self(0x22);
pub const METASTORE_CURSOR: Self = Self(0x23);
pub const DID_EVENTS: Self = Self(0x24);
pub const RECORD_BLOBS: Self = Self(0x30);
pub const BACKLINK_BY_USER: Self = Self(0x31);
pub const FORMAT_VERSION: Self = Self(0xFF);
pub const fn raw(self) -> u8 {
self.0
}
pub fn exclusive_prefix_bound(self) -> [u8; 1] {
match self.0.checked_add(1) {
Some(next) => [next],
None => panic!("cannot compute exclusive upper bound for tag 0xFF"),
}
}
#[cfg(test)]
pub fn from_raw_unchecked(raw: u8) -> Self {
Self(raw)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn user_hash_deterministic() {
let a = UserHash::from_did("did:plc:abc123");
let b = UserHash::from_did("did:plc:abc123");
assert_eq!(a, b);
}
#[test]
fn user_hash_different_dids_differ() {
let a = UserHash::from_did("did:plc:abc123");
let b = UserHash::from_did("did:plc:xyz789");
assert_ne!(a, b);
}
#[test]
fn user_hash_display_is_hex() {
let h = UserHash::from_raw(0xDEAD_BEEF_CAFE_BABE);
assert_eq!(h.to_string(), "deadbeefcafebabe");
}
#[test]
fn key_tags_are_distinct() {
let tags = [
KeyTag::REPO_META,
KeyTag::RECORDS,
KeyTag::USER_BLOCKS,
KeyTag::HANDLES,
KeyTag::BLOBS,
KeyTag::BACKLINKS,
KeyTag::BLOB_BY_CID,
KeyTag::USER_MAP,
KeyTag::USER_MAP_REVERSE,
KeyTag::REV_TO_SEQ,
KeyTag::SEQ_META,
KeyTag::SEQ_TOMBSTONE,
KeyTag::METASTORE_CURSOR,
KeyTag::DID_EVENTS,
KeyTag::RECORD_BLOBS,
KeyTag::BACKLINK_BY_USER,
KeyTag::FORMAT_VERSION,
];
let mut raw: Vec<u8> = tags.iter().map(|t| t.raw()).collect();
let original_len = raw.len();
raw.sort();
raw.dedup();
assert_eq!(raw.len(), original_len);
}
#[test]
fn key_tag_ordering() {
assert!(KeyTag::REPO_META < KeyTag::RECORDS);
assert!(KeyTag::RECORDS < KeyTag::USER_BLOCKS);
}
#[test]
fn exclusive_prefix_bound_is_tag_plus_one() {
assert_eq!(
KeyTag::REPO_META.exclusive_prefix_bound(),
[KeyTag::REPO_META.raw() + 1]
);
assert_eq!(KeyTag::HANDLES.exclusive_prefix_bound(), [0x05]);
}
#[test]
#[should_panic(expected = "cannot compute exclusive upper bound for tag 0xFF")]
fn exclusive_prefix_bound_panics_for_0xff() {
KeyTag::FORMAT_VERSION.exclusive_prefix_bound();
}
}
+417
View File
@@ -0,0 +1,417 @@
pub mod backlink_ops;
pub mod backlinks;
pub mod blob_ops;
pub mod blobs;
pub mod commit_ops;
pub mod encoding;
pub mod event_keys;
pub mod event_ops;
pub mod keys;
pub mod partitions;
pub mod record_ops;
pub mod records;
pub mod recovery;
pub mod repo_meta;
pub mod repo_ops;
pub mod scan;
pub mod user_block_ops;
pub mod user_blocks;
pub mod user_hash;
use std::path::Path;
use std::sync::Arc;
use fjall::{Database, Keyspace};
use self::keys::KeyTag;
use self::partitions::Partition;
use self::user_hash::UserHashMap;
const CURRENT_FORMAT_VERSION: u64 = 1;
#[derive(Debug, Clone)]
pub struct MetastoreConfig {
pub cache_size_bytes: u64,
}
impl Default for MetastoreConfig {
fn default() -> Self {
let total_ram = total_system_ram_bytes();
let twenty_percent = total_ram / 5;
Self {
cache_size_bytes: twenty_percent,
}
}
}
fn total_system_ram_bytes() -> u64 {
#[cfg(target_os = "linux")]
{
std::fs::read_to_string("/proc/meminfo")
.ok()
.and_then(|contents| {
contents
.lines()
.find(|line| line.starts_with("MemTotal:"))
.and_then(|line| {
line.split_whitespace()
.nth(1)
.and_then(|kb| kb.parse::<u64>().ok())
.map(|kb| kb.saturating_mul(1024))
})
})
.unwrap_or(4 * 1024 * 1024 * 1024)
}
#[cfg(not(target_os = "linux"))]
{
tracing::warn!("cannot detect system RAM on this platform, defaulting to 4GB");
4 * 1024 * 1024 * 1024
}
}
#[derive(Debug)]
pub enum MetastoreError {
Fjall(fjall::Error),
Lsm(lsm_tree::Error),
VersionMismatch {
expected: u64,
found: u64,
},
CorruptData(&'static str),
InvalidInput(&'static str),
UserHashCollision {
hash: keys::UserHash,
existing_uuid: uuid::Uuid,
new_uuid: uuid::Uuid,
},
}
impl std::fmt::Display for MetastoreError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Fjall(e) => write!(f, "fjall: {e}"),
Self::Lsm(e) => write!(f, "lsm: {e}"),
Self::VersionMismatch { expected, found } => {
write!(
f,
"format version mismatch: expected {expected}, found {found}"
)
}
Self::CorruptData(msg) => write!(f, "corrupt data: {msg}"),
Self::InvalidInput(msg) => write!(f, "invalid input: {msg}"),
Self::UserHashCollision {
hash,
existing_uuid,
new_uuid,
} => write!(
f,
"user hash collision: hash {hash} maps to both {existing_uuid} and {new_uuid}"
),
}
}
}
impl std::error::Error for MetastoreError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Fjall(e) => Some(e),
Self::Lsm(e) => Some(e),
_ => None,
}
}
}
impl From<fjall::Error> for MetastoreError {
fn from(e: fjall::Error) -> Self {
Self::Fjall(e)
}
}
impl From<lsm_tree::Error> for MetastoreError {
fn from(e: lsm_tree::Error) -> Self {
Self::Lsm(e)
}
}
type CompactionFilterFn =
Arc<dyn Fn(&str) -> Option<Arc<dyn fjall::compaction::filter::Factory>> + Send + Sync>;
pub mod client;
pub mod handler;
#[derive(Clone)]
pub struct Metastore {
db: Database,
partitions: [Keyspace; Partition::ALL.len()],
user_hashes: Arc<UserHashMap>,
}
impl Metastore {
pub fn open(path: &Path, config: MetastoreConfig) -> Result<Self, MetastoreError> {
let auth_name = Partition::Auth.name();
let filter_factory: CompactionFilterFn =
Arc::new(move |name: &str| match name == auth_name {
true => Some(Arc::new(partitions::TtlFilterFactory)),
false => None,
});
let db = Database::builder(path)
.cache_size(config.cache_size_bytes)
.with_compaction_filter_factories(filter_factory)
.open()?;
let opened: Vec<Keyspace> = Partition::ALL
.iter()
.map(|&p| {
let opts = p.create_options();
db.keyspace(p.name(), || opts)
})
.collect::<Result<_, fjall::Error>>()?;
let partitions: [Keyspace; Partition::ALL.len()] = opened
.try_into()
.ok()
.expect("opened exactly Partition::ALL.len() keyspaces");
let repo_data = partitions[Partition::RepoData.index()].clone();
Self::check_or_write_version(&db, &repo_data)?;
let user_hashes = Arc::new(UserHashMap::new(repo_data));
let loaded = user_hashes.load_all()?;
tracing::info!(count = loaded, "loaded user hash mappings");
Ok(Self {
db,
partitions,
user_hashes,
})
}
fn check_or_write_version(db: &Database, repo_data: &Keyspace) -> Result<(), MetastoreError> {
let version_key = [KeyTag::FORMAT_VERSION.raw()];
let version_bytes = CURRENT_FORMAT_VERSION.to_be_bytes();
match repo_data.get(version_key)? {
Some(existing) => {
let found_bytes: [u8; 8] = existing
.as_ref()
.try_into()
.map_err(|_| MetastoreError::CorruptData("format version not 8 bytes"))?;
let found = u64::from_be_bytes(found_bytes);
match found == CURRENT_FORMAT_VERSION {
true => Ok(()),
false => Err(MetastoreError::VersionMismatch {
expected: CURRENT_FORMAT_VERSION,
found,
}),
}
}
None => {
repo_data.insert(version_key, version_bytes)?;
db.persist(fjall::PersistMode::SyncData)?;
Ok(())
}
}
}
pub fn partition(&self, p: Partition) -> &Keyspace {
&self.partitions[p.index()]
}
pub fn user_hashes(&self) -> &Arc<UserHashMap> {
&self.user_hashes
}
pub fn database(&self) -> &Database {
&self.db
}
pub fn repo_ops(&self) -> repo_ops::RepoOps {
repo_ops::RepoOps::new(
self.partitions[Partition::RepoData.index()].clone(),
Arc::clone(&self.user_hashes),
)
}
pub fn record_ops(&self) -> record_ops::RecordOps {
record_ops::RecordOps::new(
self.partitions[Partition::RepoData.index()].clone(),
Arc::clone(&self.user_hashes),
)
}
pub fn user_block_ops(&self) -> user_block_ops::UserBlockOps {
user_block_ops::UserBlockOps::new(
self.partitions[Partition::RepoData.index()].clone(),
Arc::clone(&self.user_hashes),
)
}
pub fn event_ops<S: crate::io::StorageIO>(
&self,
bridge: Arc<crate::eventlog::EventLogBridge<S>>,
) -> event_ops::EventOps<S> {
event_ops::EventOps::new(
self.db.clone(),
self.partitions[Partition::RepoData.index()].clone(),
bridge,
)
}
pub fn blob_ops(&self) -> blob_ops::BlobOps {
blob_ops::BlobOps::new(
self.db.clone(),
self.partitions[Partition::RepoData.index()].clone(),
Arc::clone(&self.user_hashes),
)
}
pub fn backlink_ops(&self) -> backlink_ops::BacklinkOps {
backlink_ops::BacklinkOps::new(
self.partitions[Partition::Indexes.index()].clone(),
Arc::clone(&self.user_hashes),
)
}
pub fn commit_ops<S: crate::io::StorageIO>(
&self,
bridge: Arc<crate::eventlog::EventLogBridge<S>>,
) -> commit_ops::CommitOps<S> {
commit_ops::CommitOps::new(
self.db.clone(),
self.partitions[Partition::RepoData.index()].clone(),
self.partitions[Partition::Indexes.index()].clone(),
Arc::clone(&self.user_hashes),
bridge,
)
}
pub fn persist(&self) -> Result<(), MetastoreError> {
self.db
.persist(fjall::PersistMode::SyncData)
.map_err(MetastoreError::Fjall)
}
pub fn major_compact(&self) -> Result<(), MetastoreError> {
Partition::ALL.iter().try_for_each(|&p| {
tracing::info!(partition = p.name(), "starting major compaction");
self.partitions[p.index()]
.major_compact()
.map_err(MetastoreError::Fjall)?;
tracing::info!(partition = p.name(), "major compaction complete");
Ok::<(), MetastoreError>(())
})
}
}
#[cfg(test)]
mod tests {
use super::*;
fn open_fresh() -> (tempfile::TempDir, Metastore) {
let dir = tempfile::TempDir::new().unwrap();
let ms = Metastore::open(
dir.path(),
MetastoreConfig {
cache_size_bytes: 64 * 1024 * 1024,
},
)
.unwrap();
(dir, ms)
}
fn test_config() -> MetastoreConfig {
MetastoreConfig {
cache_size_bytes: 64 * 1024 * 1024,
}
}
#[test]
fn open_fresh_directory_succeeds() {
let (_dir, ms) = open_fresh();
assert_eq!(ms.user_hashes().len(), 0);
}
#[test]
fn all_partitions_accessible() {
let (_dir, ms) = open_fresh();
Partition::ALL.iter().for_each(|&p| {
let _ = ms.partition(p);
});
}
#[test]
fn reopen_preserves_partitions() {
let dir = tempfile::TempDir::new().unwrap();
{
let ms = Metastore::open(dir.path(), test_config()).unwrap();
let repo_data = ms.partition(Partition::RepoData);
repo_data.insert(b"test_key", b"test_value").unwrap();
ms.persist().unwrap();
}
{
let ms = Metastore::open(dir.path(), test_config()).unwrap();
let repo_data = ms.partition(Partition::RepoData);
let val = repo_data.get(b"test_key").unwrap().unwrap();
assert_eq!(val.as_ref(), b"test_value");
}
}
#[test]
fn version_mismatch_returns_error() {
let dir = tempfile::TempDir::new().unwrap();
{
let ms = Metastore::open(dir.path(), test_config()).unwrap();
let repo_data = ms.partition(Partition::RepoData);
let version_key = [KeyTag::FORMAT_VERSION.raw()];
repo_data.insert(version_key, 999u64.to_be_bytes()).unwrap();
ms.persist().unwrap();
}
{
let result = Metastore::open(dir.path(), test_config());
assert!(matches!(
result,
Err(MetastoreError::VersionMismatch {
expected: 1,
found: 999
})
));
}
}
#[test]
fn user_hash_mappings_survive_reopen() {
let dir = tempfile::TempDir::new().unwrap();
let uuid = uuid::Uuid::new_v4();
let hash = keys::UserHash::from_did("did:plc:survivor");
{
let ms = Metastore::open(dir.path(), test_config()).unwrap();
let mut batch = ms.database().batch();
ms.user_hashes()
.stage_insert(&mut batch, uuid, hash)
.unwrap();
batch.commit().unwrap();
ms.persist().unwrap();
}
{
let ms = Metastore::open(dir.path(), test_config()).unwrap();
assert_eq!(ms.user_hashes().len(), 1);
assert_eq!(ms.user_hashes().get(&uuid), Some(hash));
assert_eq!(ms.user_hashes().get_uuid(&hash), Some(uuid));
}
}
#[test]
fn default_config_has_reasonable_cache_size() {
let config = MetastoreConfig::default();
assert!(config.cache_size_bytes > 0);
assert!(config.cache_size_bytes <= 4 * 1024 * 1024 * 1024);
}
}
@@ -0,0 +1,129 @@
use std::time::{SystemTime, UNIX_EPOCH};
use fjall::KeyspaceCreateOptions;
use fjall::compaction::filter::{CompactionFilter, Context, Factory, ItemAccessor, Verdict};
use fjall::config::{BloomConstructionPolicy, FilterPolicy, FilterPolicyEntry};
pub const EXPIRES_AT_MS_SIZE: usize = 8;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Partition {
RepoData,
Auth,
Users,
Infra,
Indexes,
}
impl Partition {
pub const ALL: [Partition; 5] = [
Partition::RepoData,
Partition::Auth,
Partition::Users,
Partition::Infra,
Partition::Indexes,
];
pub const fn index(self) -> usize {
match self {
Self::RepoData => 0,
Self::Auth => 1,
Self::Users => 2,
Self::Infra => 3,
Self::Indexes => 4,
}
}
pub fn name(self) -> &'static str {
match self {
Self::RepoData => "repo_data",
Self::Auth => "auth",
Self::Users => "users",
Self::Infra => "infra",
Self::Indexes => "indexes",
}
}
pub fn create_options(self) -> KeyspaceCreateOptions {
match self {
Self::RepoData | Self::Indexes => {
KeyspaceCreateOptions::default().filter_policy(FilterPolicy::new([
FilterPolicyEntry::Bloom(BloomConstructionPolicy::FalsePositiveRate(0.01)),
FilterPolicyEntry::Bloom(BloomConstructionPolicy::BitsPerKey(10.0)),
]))
}
Self::Auth | Self::Users | Self::Infra => KeyspaceCreateOptions::default(),
}
}
}
pub(crate) struct TtlFilterFactory;
impl Factory for TtlFilterFactory {
fn name(&self) -> &str {
"ttl_expiry"
}
fn make_filter(&self, _ctx: &Context) -> Box<dyn CompactionFilter> {
let now_ms = u64::try_from(
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("system clock before unix epoch")
.as_millis(),
)
.unwrap_or(u64::MAX);
Box::new(TtlFilter { now_ms })
}
}
struct TtlFilter {
now_ms: u64,
}
impl CompactionFilter for TtlFilter {
fn filter_item(&mut self, item: ItemAccessor<'_>, _ctx: &Context) -> lsm_tree::Result<Verdict> {
let value = item.value()?;
match value.get(..EXPIRES_AT_MS_SIZE) {
Some(bytes) => {
let expires_at_ms =
u64::from_be_bytes(bytes.try_into().expect("slice is exactly 8 bytes"));
match expires_at_ms > 0 && expires_at_ms < self.now_ms {
true => Ok(Verdict::Remove),
false => Ok(Verdict::Keep),
}
}
None => Ok(Verdict::Keep),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn partition_names_are_distinct() {
let names: Vec<_> = Partition::ALL.iter().map(|p| p.name()).collect();
let mut deduped = names.clone();
deduped.sort();
deduped.dedup();
assert_eq!(names.len(), deduped.len());
}
#[test]
fn all_partitions_covered() {
assert_eq!(Partition::ALL.len(), 5);
}
#[test]
fn auth_partition_has_filter() {
assert_eq!(Partition::Auth.name(), "auth");
}
#[test]
fn index_matches_all_array_position() {
Partition::ALL.iter().enumerate().for_each(|(i, &p)| {
assert_eq!(p.index(), i, "Partition::{:?} index mismatch", p);
});
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,162 @@
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use super::encoding::KeyBuilder;
use super::keys::{KeyTag, UserHash};
const SCHEMA_VERSION: u8 = 1;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RecordValue {
pub record_cid: Vec<u8>,
pub takedown_ref: Option<String>,
}
impl RecordValue {
pub fn serialize(&self) -> Vec<u8> {
let payload = postcard::to_allocvec(self).expect("RecordValue serialization cannot fail");
let mut buf = Vec::with_capacity(1 + payload.len());
buf.push(SCHEMA_VERSION);
buf.extend_from_slice(&payload);
buf
}
pub fn deserialize(bytes: &[u8]) -> Option<Self> {
let (&version, payload) = bytes.split_first()?;
match version {
SCHEMA_VERSION => postcard::from_bytes(payload).ok(),
_ => None,
}
}
}
pub fn record_key(user_hash: UserHash, collection: &str, rkey: &str) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::RECORDS)
.u64(user_hash.raw())
.string(collection)
.string(rkey)
.build()
}
pub fn record_collection_prefix(user_hash: UserHash, collection: &str) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::RECORDS)
.u64(user_hash.raw())
.string(collection)
.build()
}
pub fn record_user_prefix(user_hash: UserHash) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::RECORDS)
.u64(user_hash.raw())
.build()
}
pub fn records_prefix() -> SmallVec<[u8; 128]> {
KeyBuilder::new().tag(KeyTag::RECORDS).build()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metastore::encoding::KeyReader;
#[test]
fn record_value_roundtrip() {
let value = RecordValue {
record_cid: vec![0x01, 0x71, 0x12, 0x20, 0xAB],
takedown_ref: None,
};
let bytes = value.serialize();
let decoded = RecordValue::deserialize(&bytes).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn record_value_with_takedown() {
let value = RecordValue {
record_cid: vec![0x01],
takedown_ref: Some("DMCA-456".to_string()),
};
let bytes = value.serialize();
let decoded = RecordValue::deserialize(&bytes).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn schema_version_is_first_byte() {
let value = RecordValue {
record_cid: vec![0x01],
takedown_ref: None,
};
let bytes = value.serialize();
assert_eq!(bytes[0], SCHEMA_VERSION);
}
#[test]
fn deserialize_rejects_unknown_schema_version() {
let value = RecordValue {
record_cid: vec![0x01],
takedown_ref: None,
};
let mut bytes = value.serialize();
bytes[0] = 99;
assert!(RecordValue::deserialize(&bytes).is_none());
}
#[test]
fn deserialize_rejects_empty_input() {
assert!(RecordValue::deserialize(&[]).is_none());
}
#[test]
fn record_key_roundtrip() {
let hash = UserHash::from_raw(0xDEAD_BEEF_CAFE_BABE);
let key = record_key(hash, "app.bsky.feed.post", "3k2abcd");
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::RECORDS.raw()));
assert_eq!(reader.u64(), Some(0xDEAD_BEEF_CAFE_BABE));
assert_eq!(reader.string(), Some("app.bsky.feed.post".to_string()));
assert_eq!(reader.string(), Some("3k2abcd".to_string()));
assert!(reader.is_empty());
}
#[test]
fn record_keys_sort_by_user_then_collection_then_rkey() {
let h1 = UserHash::from_raw(1);
let h2 = UserHash::from_raw(2);
let k1 = record_key(h1, "app.bsky.feed.like", "aaa");
let k2 = record_key(h1, "app.bsky.feed.post", "aaa");
let k3 = record_key(h1, "app.bsky.feed.post", "bbb");
let k4 = record_key(h2, "app.bsky.feed.like", "aaa");
assert!(k1.as_slice() < k2.as_slice());
assert!(k2.as_slice() < k3.as_slice());
assert!(k3.as_slice() < k4.as_slice());
}
#[test]
fn collection_prefix_is_prefix_of_full_key() {
let hash = UserHash::from_raw(42);
let prefix = record_collection_prefix(hash, "app.bsky.feed.post");
let full = record_key(hash, "app.bsky.feed.post", "some_rkey");
assert!(full.as_slice().starts_with(prefix.as_slice()));
}
#[test]
fn user_prefix_is_prefix_of_collection_prefix() {
let hash = UserHash::from_raw(42);
let user_pfx = record_user_prefix(hash);
let coll_pfx = record_collection_prefix(hash, "app.bsky.feed.post");
assert!(coll_pfx.as_slice().starts_with(user_pfx.as_slice()));
}
#[test]
fn records_prefix_is_just_tag() {
let pfx = records_prefix();
assert_eq!(pfx.as_slice(), &[KeyTag::RECORDS.raw()]);
}
}
@@ -0,0 +1,307 @@
use std::collections::HashSet;
use serde::{Deserialize, Serialize};
use super::backlink_ops::remove_backlinks_for_record;
use super::backlinks::{BacklinkValue, backlink_by_user_key, backlink_key, discriminant_to_path};
use super::encoding::KeyReader;
use super::keys::{KeyTag, UserHash};
use super::records::{RecordValue, record_key};
use super::repo_meta::{RepoMetaValue, repo_meta_key};
use super::user_blocks::{user_block_key, user_block_user_prefix};
use crate::metastore::MetastoreError;
const MUTATION_SET_VERSION: u8 = 1;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CommitMutationSet {
pub new_root_cid: Vec<u8>,
pub new_rev: String,
pub record_upserts: Vec<RecordMutationUpsert>,
pub record_deletes: Vec<RecordMutationDelete>,
pub block_inserts: Vec<Vec<u8>>,
pub block_deletes: Vec<Vec<u8>>,
pub backlink_adds: Vec<BacklinkMutation>,
pub backlink_remove_uris: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RecordMutationUpsert {
pub collection: String,
pub rkey: String,
pub cid_bytes: Vec<u8>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RecordMutationDelete {
pub collection: String,
pub rkey: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BacklinkMutation {
pub uri: String,
pub path: u8,
pub link_to: String,
}
const MAX_MUTATION_SET_ENTRIES: usize = 50_000;
impl CommitMutationSet {
pub fn serialize(&self) -> Result<Vec<u8>, MetastoreError> {
self.validate_size()?;
let payload = postcard::to_allocvec(self)
.map_err(|_| MetastoreError::CorruptData("CommitMutationSet serialization failed"))?;
let mut buf = Vec::with_capacity(1 + payload.len());
buf.push(MUTATION_SET_VERSION);
buf.extend_from_slice(&payload);
Ok(buf)
}
fn validate_size(&self) -> Result<(), MetastoreError> {
let total = self.record_upserts.len()
+ self.record_deletes.len()
+ self.block_inserts.len()
+ self.block_deletes.len()
+ self.backlink_adds.len()
+ self.backlink_remove_uris.len();
match total <= MAX_MUTATION_SET_ENTRIES {
true => Ok(()),
false => {
tracing::warn!(
total_entries = total,
max = MAX_MUTATION_SET_ENTRIES,
"CommitMutationSet exceeds entry limit"
);
Err(MetastoreError::InvalidInput(
"CommitMutationSet exceeds maximum entry count",
))
}
}
}
pub fn deserialize(bytes: &[u8]) -> Option<Self> {
let (&version, payload) = bytes.split_first()?;
match version {
MUTATION_SET_VERSION => match postcard::from_bytes(payload) {
Ok(v) => Some(v),
Err(e) => {
tracing::warn!(%e, "failed to deserialize CommitMutationSet payload");
None
}
},
_ => {
tracing::warn!(version, "unknown CommitMutationSet version");
None
}
}
}
}
pub fn replay_mutation_set(
batch: &mut fjall::OwnedWriteBatch,
repo_data: &fjall::Keyspace,
indexes: &fjall::Keyspace,
user_hash: UserHash,
current_meta: &RepoMetaValue,
mutation_set: &CommitMutationSet,
) -> Result<(), MetastoreError> {
mutation_set.validate_size()?;
let updated_meta = RepoMetaValue {
repo_root_cid: mutation_set.new_root_cid.clone(),
repo_rev: mutation_set.new_rev.clone(),
..current_meta.clone()
};
let meta_key = repo_meta_key(user_hash);
batch.insert(repo_data, meta_key.as_slice(), updated_meta.serialize());
mutation_set.record_upserts.iter().for_each(|u| {
let key = record_key(user_hash, &u.collection, &u.rkey);
let value = RecordValue {
record_cid: u.cid_bytes.clone(),
takedown_ref: None,
};
batch.insert(repo_data, key.as_slice(), value.serialize());
});
mutation_set.record_deletes.iter().for_each(|d| {
let key = record_key(user_hash, &d.collection, &d.rkey);
batch.remove(repo_data, key.as_slice());
});
mutation_set.block_inserts.iter().for_each(|cid_bytes| {
let key = user_block_key(user_hash, &mutation_set.new_rev, cid_bytes);
batch.insert(repo_data, key.as_slice(), []);
});
delete_user_blocks_by_cid_scan(batch, repo_data, user_hash, &mutation_set.block_deletes)?;
mutation_set
.backlink_remove_uris
.iter()
.try_for_each(|uri_str| {
let uri = tranquil_types::AtUri::from(uri_str.clone());
let collection = uri.collection().ok_or(MetastoreError::CorruptData(
"backlink URI missing collection",
))?;
let rkey = uri
.rkey()
.ok_or(MetastoreError::CorruptData("backlink URI missing rkey"))?;
remove_backlinks_for_record(batch, indexes, user_hash, collection, rkey)
})?;
mutation_set.backlink_adds.iter().try_for_each(|bl| {
let uri = tranquil_types::AtUri::from(bl.uri.clone());
let collection = uri.collection().ok_or(MetastoreError::CorruptData(
"backlink URI missing collection",
))?;
let rkey = uri
.rkey()
.ok_or(MetastoreError::CorruptData("backlink URI missing rkey"))?;
match discriminant_to_path(bl.path) {
None => {
tracing::warn!(
path = bl.path,
uri = %bl.uri,
"skipping backlink with unknown path discriminant during recovery"
);
}
Some(_) => {
let primary = backlink_key(&bl.link_to, user_hash, collection, rkey);
let value = BacklinkValue {
source_uri: bl.uri.clone(),
path: bl.path,
};
batch.insert(indexes, primary.as_slice(), value.serialize());
let reverse = backlink_by_user_key(user_hash, collection, rkey, &bl.link_to);
batch.insert(indexes, reverse.as_slice(), []);
}
}
Ok::<_, MetastoreError>(())
})
}
fn delete_user_blocks_by_cid_scan(
batch: &mut fjall::OwnedWriteBatch,
repo_data: &fjall::Keyspace,
user_hash: UserHash,
block_cids: &[Vec<u8>],
) -> Result<(), MetastoreError> {
match block_cids.is_empty() {
true => Ok(()),
false => {
let cid_set: HashSet<&[u8]> = block_cids.iter().map(|c| c.as_slice()).collect();
let prefix = user_block_user_prefix(user_hash);
repo_data.prefix(prefix.as_slice()).try_for_each(|guard| {
let (key_bytes, _) = guard.into_inner().map_err(MetastoreError::Fjall)?;
match extract_cid_from_user_block_key(&key_bytes) {
Some(cid) if cid_set.contains(cid) => {
batch.remove(repo_data, key_bytes.as_ref());
Ok(())
}
_ => Ok(()),
}
})
}
}
}
fn extract_cid_from_user_block_key(key_bytes: &[u8]) -> Option<&[u8]> {
let mut reader = KeyReader::new(key_bytes);
let tag = reader.tag()?;
if tag != KeyTag::USER_BLOCKS.raw() {
tracing::warn!(
tag,
"unexpected key tag in user_block prefix scan during recovery"
);
return None;
}
if reader.u64().and_then(|_| reader.string()).is_none() {
tracing::warn!("user_block key has corrupt user_hash or rev during recovery");
return None;
}
let remaining = reader.remaining();
match remaining.is_empty() {
true => {
tracing::warn!("user_block key has no CID suffix during recovery");
None
}
false => Some(remaining),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn mutation_set_roundtrip() {
let ms = CommitMutationSet {
new_root_cid: vec![0x01, 0x71, 0x12, 0x20],
new_rev: "rev1".to_owned(),
record_upserts: vec![RecordMutationUpsert {
collection: "app.bsky.feed.post".to_owned(),
rkey: "3k2abc".to_owned(),
cid_bytes: vec![0xDE, 0xAD],
}],
record_deletes: vec![RecordMutationDelete {
collection: "app.bsky.feed.like".to_owned(),
rkey: "3k2del".to_owned(),
}],
block_inserts: vec![vec![0x01, 0x02]],
block_deletes: vec![vec![0x03, 0x04]],
backlink_adds: vec![BacklinkMutation {
uri: "at://did:plc:alice/app.bsky.feed.like/3k2abc".to_owned(),
path: 1,
link_to: "at://did:plc:bob/app.bsky.feed.post/3k2xyz".to_owned(),
}],
backlink_remove_uris: vec!["at://did:plc:alice/app.bsky.feed.like/3k2old".to_owned()],
};
let bytes = ms.serialize().unwrap();
assert_eq!(bytes[0], MUTATION_SET_VERSION);
let recovered = CommitMutationSet::deserialize(&bytes).unwrap();
assert_eq!(recovered, ms);
}
#[test]
fn mutation_set_empty_roundtrip() {
let ms = CommitMutationSet {
new_root_cid: vec![],
new_rev: String::new(),
record_upserts: vec![],
record_deletes: vec![],
block_inserts: vec![],
block_deletes: vec![],
backlink_adds: vec![],
backlink_remove_uris: vec![],
};
let recovered = CommitMutationSet::deserialize(&ms.serialize().unwrap()).unwrap();
assert_eq!(recovered, ms);
}
#[test]
fn unknown_version_returns_none() {
let ms = CommitMutationSet {
new_root_cid: vec![],
new_rev: String::new(),
record_upserts: vec![],
record_deletes: vec![],
block_inserts: vec![],
block_deletes: vec![],
backlink_adds: vec![],
backlink_remove_uris: vec![],
};
let mut bytes = ms.serialize().unwrap();
bytes[0] = 99;
assert!(CommitMutationSet::deserialize(&bytes).is_none());
}
}
@@ -0,0 +1,226 @@
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use super::encoding::KeyBuilder;
use super::keys::{KeyTag, UserHash};
const SCHEMA_VERSION: u8 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[repr(u8)]
pub enum RepoStatus {
Active = 0,
Takendown = 1,
Suspended = 2,
Deactivated = 3,
Deleted = 4,
}
impl RepoStatus {
pub fn is_active(self) -> bool {
matches!(self, Self::Active)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RepoMetaValue {
pub repo_root_cid: Vec<u8>,
pub repo_rev: String,
pub handle: String,
pub status: RepoStatus,
pub deactivated_at_ms: Option<u64>,
pub takedown_ref: Option<String>,
#[serde(default)]
pub did: Option<String>,
}
impl RepoMetaValue {
pub fn serialize(&self) -> Vec<u8> {
let payload = postcard::to_allocvec(self).expect("RepoMetaValue serialization cannot fail");
let mut buf = Vec::with_capacity(1 + payload.len());
buf.push(SCHEMA_VERSION);
buf.extend_from_slice(&payload);
buf
}
pub fn deserialize(bytes: &[u8]) -> Option<Self> {
let (&version, payload) = bytes.split_first()?;
match version {
SCHEMA_VERSION => postcard::from_bytes(payload).ok(),
_ => None,
}
}
}
pub fn repo_meta_key(user_hash: UserHash) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::REPO_META)
.u64(user_hash.raw())
.build()
}
pub fn repo_meta_prefix() -> SmallVec<[u8; 128]> {
KeyBuilder::new().tag(KeyTag::REPO_META).build()
}
pub fn handle_key(handle_lower: &str) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::HANDLES)
.string(handle_lower)
.build()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metastore::encoding::KeyReader;
#[test]
fn repo_meta_value_roundtrip() {
let value = RepoMetaValue {
repo_root_cid: vec![0x01, 0x71, 0x12, 0x20, 0xAB],
repo_rev: "3k2a7bcd".to_string(),
handle: "alice.bsky.social".to_string(),
status: RepoStatus::Active,
deactivated_at_ms: None,
takedown_ref: None,
did: None,
};
let bytes = value.serialize();
let decoded = RepoMetaValue::deserialize(&bytes).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn repo_meta_value_with_optional_fields() {
let value = RepoMetaValue {
repo_root_cid: vec![0x01],
repo_rev: "rev1".to_string(),
handle: "bob.example.com".to_string(),
status: RepoStatus::Deactivated,
deactivated_at_ms: Some(1700000000000),
takedown_ref: Some("DMCA-123".to_string()),
did: Some("did:plc:bob".to_string()),
};
let bytes = value.serialize();
let decoded = RepoMetaValue::deserialize(&bytes).unwrap();
assert_eq!(decoded, value);
}
#[test]
fn repo_meta_key_roundtrip() {
let hash = UserHash::from_raw(0xDEAD_BEEF_CAFE_BABE);
let key = repo_meta_key(hash);
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::REPO_META.raw()));
assert_eq!(reader.u64(), Some(0xDEAD_BEEF_CAFE_BABE));
assert!(reader.is_empty());
}
#[test]
fn repo_meta_keys_sort_by_user_hash() {
let k1 = repo_meta_key(UserHash::from_raw(1));
let k2 = repo_meta_key(UserHash::from_raw(2));
let k3 = repo_meta_key(UserHash::from_raw(u64::MAX));
assert!(k1.as_slice() < k2.as_slice());
assert!(k2.as_slice() < k3.as_slice());
}
#[test]
fn handle_key_roundtrip() {
let key = handle_key("alice.bsky.social");
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::HANDLES.raw()));
assert_eq!(reader.string(), Some("alice.bsky.social".to_string()));
assert!(reader.is_empty());
}
#[test]
fn handle_keys_sort_lexicographically() {
let k1 = handle_key("alice.example.com");
let k2 = handle_key("bob.example.com");
assert!(k1.as_slice() < k2.as_slice());
}
#[test]
fn all_repo_statuses_roundtrip() {
[
RepoStatus::Active,
RepoStatus::Takendown,
RepoStatus::Suspended,
RepoStatus::Deactivated,
RepoStatus::Deleted,
]
.iter()
.for_each(|&status| {
let value = RepoMetaValue {
repo_root_cid: vec![0x01],
repo_rev: "r".to_string(),
handle: "h.test".to_string(),
status,
deactivated_at_ms: None,
takedown_ref: None,
did: None,
};
let decoded = RepoMetaValue::deserialize(&value.serialize()).unwrap();
assert_eq!(decoded.status, status);
});
}
#[test]
fn repo_status_serialization_stability() {
[
(RepoStatus::Active, 0u8),
(RepoStatus::Takendown, 1),
(RepoStatus::Suspended, 2),
(RepoStatus::Deactivated, 3),
(RepoStatus::Deleted, 4),
]
.iter()
.for_each(|&(status, expected_byte)| {
let bytes = postcard::to_allocvec(&status).unwrap();
assert_eq!(
bytes,
[expected_byte],
"{status:?} serialized to {bytes:?}, expected [{expected_byte}]"
);
});
}
#[test]
fn schema_version_is_first_byte() {
let value = RepoMetaValue {
repo_root_cid: vec![0x01],
repo_rev: "r".to_string(),
handle: "h.test".to_string(),
status: RepoStatus::Active,
deactivated_at_ms: None,
takedown_ref: None,
did: None,
};
let bytes = value.serialize();
assert_eq!(bytes[0], SCHEMA_VERSION);
assert_eq!(bytes[0], 1, "schema version must remain 1 for this format");
}
#[test]
fn deserialize_rejects_unknown_schema_version() {
let value = RepoMetaValue {
repo_root_cid: vec![0x01],
repo_rev: "r".to_string(),
handle: "h.test".to_string(),
status: RepoStatus::Active,
deactivated_at_ms: None,
takedown_ref: None,
did: None,
};
let mut bytes = value.serialize();
bytes[0] = 99;
assert!(RepoMetaValue::deserialize(&bytes).is_none());
}
#[test]
fn deserialize_rejects_empty_input() {
assert!(RepoMetaValue::deserialize(&[]).is_none());
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,36 @@
use fjall::Keyspace;
use super::MetastoreError;
pub fn count_prefix(keyspace: &Keyspace, prefix: &[u8]) -> Result<i64, MetastoreError> {
keyspace.prefix(prefix).try_fold(0i64, |acc, guard| {
guard.into_inner().map_err(MetastoreError::Fjall)?;
Ok::<_, MetastoreError>(acc.saturating_add(1))
})
}
pub fn delete_all_by_prefix(
keyspace: &Keyspace,
batch: &mut fjall::OwnedWriteBatch,
prefix: &[u8],
) -> Result<(), MetastoreError> {
keyspace.prefix(prefix).try_for_each(|guard| {
let (key_bytes, _) = guard.into_inner().map_err(MetastoreError::Fjall)?;
batch.remove(keyspace, key_bytes.as_ref());
Ok::<(), MetastoreError>(())
})
}
pub fn point_lookup<T>(
keyspace: &Keyspace,
key: &[u8],
deserialize: impl FnOnce(&[u8]) -> Option<T>,
corrupt_msg: &'static str,
) -> Result<Option<T>, MetastoreError> {
match keyspace.get(key).map_err(MetastoreError::Fjall)? {
Some(raw) => deserialize(&raw)
.ok_or(MetastoreError::CorruptData(corrupt_msg))
.map(Some),
None => Ok(None),
}
}
@@ -0,0 +1,560 @@
use std::collections::HashSet;
use std::sync::Arc;
use fjall::Keyspace;
use uuid::Uuid;
use super::MetastoreError;
use super::encoding::{KeyReader, exclusive_upper_bound};
use super::keys::UserHash;
use super::scan::{count_prefix, delete_all_by_prefix};
use super::user_blocks::{user_block_key, user_block_rev_prefix, user_block_user_prefix};
use super::user_hash::UserHashMap;
pub struct UserBlockOps {
repo_data: Keyspace,
user_hashes: Arc<UserHashMap>,
}
impl UserBlockOps {
pub fn new(repo_data: Keyspace, user_hashes: Arc<UserHashMap>) -> Self {
Self {
repo_data,
user_hashes,
}
}
pub fn insert_user_blocks<C: AsRef<[u8]>>(
&self,
batch: &mut fjall::OwnedWriteBatch,
user_hash: UserHash,
block_cids: &[C],
repo_rev: &str,
) -> Result<(), MetastoreError> {
let existing: HashSet<Vec<u8>> = match block_cids.is_empty() {
true => HashSet::new(),
false => {
let prefix = user_block_user_prefix(user_hash);
self.repo_data
.prefix(prefix.as_slice())
.filter_map(|guard| {
let (key_bytes, _) = guard.into_inner().ok()?;
extract_cid_from_key(&key_bytes).map(|c| c.to_vec())
})
.collect()
}
};
block_cids.iter().try_for_each(|cid| {
let cid = cid.as_ref();
match cid.is_empty() {
true => Err(MetastoreError::InvalidInput("block CID must not be empty")),
false => {
if !existing.contains(cid) {
let key = user_block_key(user_hash, repo_rev, cid);
batch.insert(&self.repo_data, key.as_slice(), []);
}
Ok(())
}
}
})
}
pub fn delete_user_blocks<C: AsRef<[u8]>>(
&self,
batch: &mut fjall::OwnedWriteBatch,
user_hash: UserHash,
block_cids: &[C],
rev: &str,
) -> Result<(), MetastoreError> {
block_cids.iter().try_for_each(|cid| {
let cid = cid.as_ref();
match cid.is_empty() {
true => Err(MetastoreError::InvalidInput("block CID must not be empty")),
false => {
let key = user_block_key(user_hash, rev, cid);
batch.remove(&self.repo_data, key.as_slice());
Ok(())
}
}
})
}
pub fn delete_user_blocks_by_cid<C: AsRef<[u8]>>(
&self,
batch: &mut fjall::OwnedWriteBatch,
user_hash: UserHash,
block_cids: &[C],
) -> Result<(), MetastoreError> {
match block_cids.is_empty() {
true => Ok(()),
false => {
let cid_set: HashSet<&[u8]> = block_cids.iter().map(|c| c.as_ref()).collect();
let prefix = user_block_user_prefix(user_hash);
self.repo_data
.prefix(prefix.as_slice())
.try_for_each(|guard| {
let (key_bytes, _) = guard.into_inner().map_err(MetastoreError::Fjall)?;
match extract_cid_from_key(&key_bytes) {
Some(cid) if cid_set.contains(cid.as_slice()) => {
batch.remove(&self.repo_data, key_bytes.as_ref());
Ok(())
}
_ => Ok(()),
}
})
}
}
}
pub fn delete_all_user_blocks(
&self,
batch: &mut fjall::OwnedWriteBatch,
user_hash: UserHash,
) -> Result<(), MetastoreError> {
let prefix = user_block_user_prefix(user_hash);
delete_all_by_prefix(&self.repo_data, batch, prefix.as_slice())
}
pub fn get_user_block_cids_since_rev(
&self,
user_id: Uuid,
since_rev: &str,
) -> Result<Vec<Vec<u8>>, MetastoreError> {
let user_hash = match self.user_hashes.get(&user_id) {
Some(h) => h,
None => return Ok(Vec::new()),
};
let since_prefix = user_block_rev_prefix(user_hash, since_rev);
let since_upper = exclusive_upper_bound(since_prefix.as_slice())
.expect("user block rev prefix always contains non-0xFF bytes");
let user_prefix = user_block_user_prefix(user_hash);
let user_upper = exclusive_upper_bound(user_prefix.as_slice())
.expect("user block user prefix always contains non-0xFF bytes");
self.repo_data
.range(since_upper.as_slice()..user_upper.as_slice())
.map(|guard| {
let (key_bytes, _) = guard.into_inner().map_err(MetastoreError::Fjall)?;
extract_cid_from_key(&key_bytes)
.ok_or(MetastoreError::CorruptData("invalid user_blocks key"))
})
.collect()
}
pub fn find_unreferenced(&self, candidate_cids: &[Vec<u8>]) -> Vec<Vec<u8>> {
match candidate_cids.is_empty() {
true => Vec::new(),
false => {
let mut remaining: HashSet<Vec<u8>> = candidate_cids.iter().cloned().collect();
let tag_prefix = super::keys::KeyTag::USER_BLOCKS.raw();
let mut iter = self.repo_data.prefix([tag_prefix]);
loop {
match remaining.is_empty() {
true => break,
false => match iter.next() {
None => break,
Some(guard) => {
if let Ok((key_bytes, _)) = guard.into_inner()
&& let Some(cid) = extract_cid_from_key(&key_bytes)
{
remaining.remove(&cid);
}
}
},
}
}
remaining.into_iter().collect()
}
}
}
pub fn count_user_blocks(&self, user_id: Uuid) -> Result<i64, MetastoreError> {
let user_hash = match self.user_hashes.get(&user_id) {
Some(h) => h,
None => return Ok(0),
};
let prefix = user_block_user_prefix(user_hash);
count_prefix(&self.repo_data, prefix.as_slice())
}
}
fn extract_cid_from_key(key_bytes: &[u8]) -> Option<Vec<u8>> {
let mut reader = KeyReader::new(key_bytes);
reader.tag()?;
reader.u64()?;
reader.string()?;
match reader.remaining().is_empty() {
true => None,
false => Some(reader.remaining().to_vec()),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metastore::{Metastore, MetastoreConfig};
fn open_fresh() -> (tempfile::TempDir, Metastore) {
let dir = tempfile::TempDir::new().unwrap();
let ms = Metastore::open(
dir.path(),
MetastoreConfig {
cache_size_bytes: 64 * 1024 * 1024,
},
)
.unwrap();
(dir, ms)
}
fn setup_user(ms: &Metastore) -> (Uuid, UserHash) {
let uuid = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:testuser1");
let mut batch = ms.database().batch();
ms.user_hashes()
.stage_insert(&mut batch, uuid, hash)
.unwrap();
batch.commit().unwrap();
(uuid, hash)
}
#[test]
fn insert_and_count() {
let (_dir, ms) = open_fresh();
let (uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let cids = vec![vec![0x01, 0x71], vec![0x02, 0x72], vec![0x03, 0x73]];
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash, &cids, "rev1")
.unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid).unwrap(), 3);
}
#[test]
fn get_since_rev_returns_later_revisions() {
let (_dir, ms) = open_fresh();
let (uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let cids_abc = vec![vec![0x01], vec![0x02]];
let cids_def = vec![vec![0x03]];
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash, &cids_abc, "abc")
.unwrap();
ops.insert_user_blocks(&mut batch, hash, &cids_def, "def")
.unwrap();
batch.commit().unwrap();
let since_abc = ops.get_user_block_cids_since_rev(uuid, "abc").unwrap();
assert_eq!(since_abc.len(), 1);
assert_eq!(since_abc[0], vec![0x03]);
let since_def = ops.get_user_block_cids_since_rev(uuid, "def").unwrap();
assert!(since_def.is_empty());
}
#[test]
fn get_since_rev_with_both_revisions() {
let (_dir, ms) = open_fresh();
let (uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let cids_r1 = vec![vec![0x01]];
let cids_r2 = vec![vec![0x02]];
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash, &cids_r1, "aaa")
.unwrap();
ops.insert_user_blocks(&mut batch, hash, &cids_r2, "bbb")
.unwrap();
batch.commit().unwrap();
let since_before = ops.get_user_block_cids_since_rev(uuid, "aaa").unwrap();
assert_eq!(since_before.len(), 1);
assert_eq!(since_before[0], vec![0x02]);
let all = ops.get_user_block_cids_since_rev(uuid, "").unwrap();
assert_eq!(all.len(), 2);
}
#[test]
fn delete_blocks_at_rev() {
let (_dir, ms) = open_fresh();
let (uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let cids = vec![vec![0x01], vec![0x02], vec![0x03]];
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash, &cids, "rev1")
.unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid).unwrap(), 3);
let to_delete = vec![vec![0x01], vec![0x03]];
let mut batch = ms.database().batch();
ops.delete_user_blocks(&mut batch, hash, &to_delete, "rev1")
.unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid).unwrap(), 1);
}
#[test]
fn delete_wrong_rev_is_noop() {
let (_dir, ms) = open_fresh();
let (uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let cids = vec![vec![0x01], vec![0x02]];
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash, &cids, "rev1")
.unwrap();
batch.commit().unwrap();
let mut batch = ms.database().batch();
ops.delete_user_blocks(&mut batch, hash, &cids, "wrong_rev")
.unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid).unwrap(), 2);
}
#[test]
fn delete_all_user_blocks_clears_all_revisions() {
let (_dir, ms) = open_fresh();
let (uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash, &[vec![0x01], vec![0x02]], "rev1")
.unwrap();
ops.insert_user_blocks(&mut batch, hash, &[vec![0x03]], "rev2")
.unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid).unwrap(), 3);
let mut batch = ms.database().batch();
ops.delete_all_user_blocks(&mut batch, hash).unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid).unwrap(), 0);
}
#[test]
fn empty_rev_scan_returns_empty() {
let (_dir, ms) = open_fresh();
let (uuid, _hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let result = ops.get_user_block_cids_since_rev(uuid, "anything").unwrap();
assert!(result.is_empty());
}
#[test]
fn unknown_user_returns_zero_count() {
let (_dir, ms) = open_fresh();
let ops = ms.user_block_ops();
let unknown = Uuid::new_v4();
assert_eq!(ops.count_user_blocks(unknown).unwrap(), 0);
}
#[test]
fn blocks_survive_reopen() {
let dir = tempfile::TempDir::new().unwrap();
let uuid = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:persist");
{
let ms = Metastore::open(
dir.path(),
MetastoreConfig {
cache_size_bytes: 64 * 1024 * 1024,
},
)
.unwrap();
let mut batch = ms.database().batch();
ms.user_hashes()
.stage_insert(&mut batch, uuid, hash)
.unwrap();
batch.commit().unwrap();
let ops = ms.user_block_ops();
let cids = vec![vec![0xAA, 0xBB], vec![0xCC, 0xDD]];
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash, &cids, "rev1")
.unwrap();
batch.commit().unwrap();
ms.persist().unwrap();
}
{
let ms = Metastore::open(
dir.path(),
MetastoreConfig {
cache_size_bytes: 64 * 1024 * 1024,
},
)
.unwrap();
let ops = ms.user_block_ops();
assert_eq!(ops.count_user_blocks(uuid).unwrap(), 2);
}
}
#[test]
fn multiple_users_isolated() {
let (_dir, ms) = open_fresh();
let uuid1 = Uuid::new_v4();
let hash1 = UserHash::from_did("did:plc:user1");
let uuid2 = Uuid::new_v4();
let hash2 = UserHash::from_did("did:plc:user2");
let mut batch = ms.database().batch();
ms.user_hashes()
.stage_insert(&mut batch, uuid1, hash1)
.unwrap();
ms.user_hashes()
.stage_insert(&mut batch, uuid2, hash2)
.unwrap();
batch.commit().unwrap();
let ops = ms.user_block_ops();
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash1, &[vec![0x01], vec![0x02]], "rev1")
.unwrap();
ops.insert_user_blocks(&mut batch, hash2, &[vec![0x03]], "rev1")
.unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid1).unwrap(), 2);
assert_eq!(ops.count_user_blocks(uuid2).unwrap(), 1);
let mut batch = ms.database().batch();
ops.delete_user_blocks(&mut batch, hash1, &[vec![0x01]], "rev1")
.unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid1).unwrap(), 1);
assert_eq!(ops.count_user_blocks(uuid2).unwrap(), 1);
}
#[test]
fn delete_all_does_not_affect_other_users() {
let (_dir, ms) = open_fresh();
let uuid1 = Uuid::new_v4();
let hash1 = UserHash::from_did("did:plc:user1");
let uuid2 = Uuid::new_v4();
let hash2 = UserHash::from_did("did:plc:user2");
let mut batch = ms.database().batch();
ms.user_hashes()
.stage_insert(&mut batch, uuid1, hash1)
.unwrap();
ms.user_hashes()
.stage_insert(&mut batch, uuid2, hash2)
.unwrap();
batch.commit().unwrap();
let ops = ms.user_block_ops();
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash1, &[vec![0x01]], "rev1")
.unwrap();
ops.insert_user_blocks(&mut batch, hash2, &[vec![0x02]], "rev1")
.unwrap();
batch.commit().unwrap();
let mut batch = ms.database().batch();
ops.delete_all_user_blocks(&mut batch, hash1).unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid1).unwrap(), 0);
assert_eq!(ops.count_user_blocks(uuid2).unwrap(), 1);
}
#[test]
fn cids_with_null_bytes_roundtrip_through_storage() {
let (_dir, ms) = open_fresh();
let (uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let cids = vec![
vec![0x00, 0x00, 0x01],
vec![0x00, 0x01, 0x00, 0x00],
vec![0x00, 0x00],
];
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash, &cids, "rev1")
.unwrap();
batch.commit().unwrap();
assert_eq!(ops.count_user_blocks(uuid).unwrap(), 3);
let retrieved = ops.get_user_block_cids_since_rev(uuid, "").unwrap();
assert_eq!(retrieved.len(), 3);
let mut expected = cids.clone();
expected.sort();
assert_eq!(retrieved, expected);
}
#[test]
fn insert_empty_cid_is_rejected() {
let (_dir, ms) = open_fresh();
let (_uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let cids = vec![vec![]];
let mut batch = ms.database().batch();
let result = ops.insert_user_blocks(&mut batch, hash, &cids, "rev1");
assert!(matches!(result, Err(MetastoreError::InvalidInput(_))));
}
#[test]
fn delete_empty_cid_is_rejected() {
let (_dir, ms) = open_fresh();
let (_uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let cids = vec![vec![]];
let mut batch = ms.database().batch();
let result = ops.delete_user_blocks(&mut batch, hash, &cids, "rev1");
assert!(matches!(result, Err(MetastoreError::InvalidInput(_))));
}
#[test]
fn since_rev_nonexistent_returns_later_revisions() {
let (_dir, ms) = open_fresh();
let (uuid, hash) = setup_user(&ms);
let ops = ms.user_block_ops();
let mut batch = ms.database().batch();
ops.insert_user_blocks(&mut batch, hash, &[vec![0x01]], "aaa")
.unwrap();
ops.insert_user_blocks(&mut batch, hash, &[vec![0x02]], "bbb")
.unwrap();
ops.insert_user_blocks(&mut batch, hash, &[vec![0x03]], "ddd")
.unwrap();
batch.commit().unwrap();
let result = ops.get_user_block_cids_since_rev(uuid, "aab").unwrap();
assert_eq!(result.len(), 2);
assert_eq!(result[0], vec![0x02]);
assert_eq!(result[1], vec![0x03]);
let result = ops.get_user_block_cids_since_rev(uuid, "ccc").unwrap();
assert_eq!(result.len(), 1);
assert_eq!(result[0], vec![0x03]);
}
}
@@ -0,0 +1,112 @@
use smallvec::SmallVec;
use super::encoding::KeyBuilder;
use super::keys::{KeyTag, UserHash};
pub fn user_block_key(user_hash: UserHash, rev: &str, cid_bytes: &[u8]) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::USER_BLOCKS)
.u64(user_hash.raw())
.string(rev)
.raw(cid_bytes)
.build()
}
pub fn user_block_user_prefix(user_hash: UserHash) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::USER_BLOCKS)
.u64(user_hash.raw())
.build()
}
pub fn user_block_rev_prefix(user_hash: UserHash, rev: &str) -> SmallVec<[u8; 128]> {
KeyBuilder::new()
.tag(KeyTag::USER_BLOCKS)
.u64(user_hash.raw())
.string(rev)
.build()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::metastore::encoding::KeyReader;
#[test]
fn user_block_key_roundtrip() {
let hash = UserHash::from_raw(0xDEAD_BEEF_CAFE_BABE);
let cid = [0x01, 0x71, 0x12, 0x20, 0xAB];
let key = user_block_key(hash, "3k2abcde", &cid);
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::USER_BLOCKS.raw()));
assert_eq!(reader.u64(), Some(0xDEAD_BEEF_CAFE_BABE));
assert_eq!(reader.string(), Some("3k2abcde".to_string()));
assert_eq!(reader.remaining(), &cid);
}
#[test]
fn keys_sort_by_user_then_rev_then_cid() {
let h1 = UserHash::from_raw(1);
let h2 = UserHash::from_raw(2);
let k1 = user_block_key(h1, "abc", &[0x01]);
let k2 = user_block_key(h1, "abc", &[0x02]);
let k3 = user_block_key(h1, "def", &[0x01]);
let k4 = user_block_key(h2, "abc", &[0x01]);
assert!(k1.as_slice() < k2.as_slice());
assert!(k2.as_slice() < k3.as_slice());
assert!(k3.as_slice() < k4.as_slice());
}
#[test]
fn user_prefix_is_prefix_of_rev_prefix() {
let hash = UserHash::from_raw(42);
let user_pfx = user_block_user_prefix(hash);
let rev_pfx = user_block_rev_prefix(hash, "some_rev");
assert!(rev_pfx.as_slice().starts_with(user_pfx.as_slice()));
}
#[test]
fn rev_prefix_is_prefix_of_full_key() {
let hash = UserHash::from_raw(42);
let rev_pfx = user_block_rev_prefix(hash, "some_rev");
let full = user_block_key(hash, "some_rev", &[0x01, 0x02]);
assert!(full.as_slice().starts_with(rev_pfx.as_slice()));
}
#[test]
fn cid_with_null_bytes_roundtrips() {
let hash = UserHash::from_raw(99);
let cid = [0x00, 0x01, 0x00, 0xFF];
let key = user_block_key(hash, "rev1", &cid);
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::USER_BLOCKS.raw()));
assert_eq!(reader.u64(), Some(99));
assert_eq!(reader.string(), Some("rev1".to_string()));
assert_eq!(reader.remaining(), &cid);
}
#[test]
fn cid_with_double_null_bytes_roundtrips() {
let hash = UserHash::from_raw(99);
let cid = [0x00, 0x00, 0x01, 0x00, 0x00];
let key = user_block_key(hash, "rev1", &cid);
let mut reader = KeyReader::new(&key);
assert_eq!(reader.tag(), Some(KeyTag::USER_BLOCKS.raw()));
assert_eq!(reader.u64(), Some(99));
assert_eq!(reader.string(), Some("rev1".to_string()));
assert_eq!(reader.remaining(), &cid);
}
#[test]
fn empty_cid_produces_key_equal_to_rev_prefix() {
let hash = UserHash::from_raw(42);
let rev_pfx = user_block_rev_prefix(hash, "rev1");
let full = user_block_key(hash, "rev1", &[]);
assert_eq!(full.as_slice(), rev_pfx.as_slice());
}
}
@@ -0,0 +1,437 @@
use dashmap::DashMap;
use fjall::Keyspace;
use parking_lot::Mutex;
use uuid::Uuid;
use super::MetastoreError;
use super::encoding::KeyBuilder;
use super::keys::{KeyTag, UserHash};
pub struct UserHashMap {
cache: DashMap<Uuid, UserHash>,
reverse: DashMap<UserHash, Uuid>,
repo_data: Keyspace,
write_guard: Mutex<()>,
}
impl UserHashMap {
pub fn new(repo_data: Keyspace) -> Self {
Self {
cache: DashMap::new(),
reverse: DashMap::new(),
repo_data,
write_guard: Mutex::new(()),
}
}
pub fn load_all(&self) -> Result<usize, MetastoreError> {
let prefix = [KeyTag::USER_MAP.raw()];
let mut count = 0usize;
self.repo_data.prefix(prefix).try_for_each(|guard| {
let (key_bytes, value_bytes) = guard.into_inner().map_err(MetastoreError::Fjall)?;
let uuid_bytes: [u8; 16] = key_bytes
.get(1..17)
.and_then(|s| <[u8; 16]>::try_from(s).ok())
.ok_or(MetastoreError::CorruptData(
"user_map key too short for UUID",
))?;
let hash_bytes: [u8; 8] = <[u8; 8]>::try_from(value_bytes.as_ref())
.map_err(|_| MetastoreError::CorruptData("user_map value not 8 bytes"))?;
let uuid = Uuid::from_bytes(uuid_bytes);
let user_hash = UserHash::from_raw(u64::from_be_bytes(hash_bytes));
let existing = self.reverse.get(&user_hash).map(|r| *r);
if let Some(existing_uuid) = existing
&& existing_uuid != uuid
{
tracing::error!(
existing_uuid = %existing_uuid,
new_uuid = %uuid,
user_hash = %user_hash,
"user hash collision in persisted data"
);
return Err(MetastoreError::UserHashCollision {
hash: user_hash,
existing_uuid,
new_uuid: uuid,
});
}
self.cache.insert(uuid, user_hash);
self.reverse.insert(user_hash, uuid);
count = count.saturating_add(1);
if count.is_multiple_of(100_000) {
tracing::info!(count, "loading user hash mappings");
}
Ok::<_, MetastoreError>(())
})?;
Ok(count)
}
pub fn stage_insert(
&self,
batch: &mut fjall::OwnedWriteBatch,
uuid: Uuid,
user_hash: UserHash,
) -> Result<(), MetastoreError> {
let _guard = self.write_guard.lock();
let existing = self.reverse.get(&user_hash).map(|r| *r);
if let Some(existing_uuid) = existing
&& existing_uuid != uuid
{
tracing::error!(
existing_uuid = %existing_uuid,
new_uuid = %uuid,
user_hash = %user_hash,
"user hash collision detected"
);
return Err(MetastoreError::UserHashCollision {
hash: user_hash,
existing_uuid,
new_uuid: uuid,
});
}
let forward_key = KeyBuilder::new()
.tag(KeyTag::USER_MAP)
.fixed(uuid.as_bytes())
.build();
let reverse_key = KeyBuilder::new()
.tag(KeyTag::USER_MAP_REVERSE)
.u64(user_hash.raw())
.build();
batch.insert(
&self.repo_data,
forward_key.as_slice(),
user_hash.raw().to_be_bytes(),
);
batch.insert(
&self.repo_data,
reverse_key.as_slice(),
uuid.as_bytes().as_slice(),
);
self.cache.insert(uuid, user_hash);
self.reverse.insert(user_hash, uuid);
Ok(())
}
pub fn rollback_insert(&self, uuid: &Uuid, user_hash: &UserHash) {
self.cache.remove(uuid);
self.reverse.remove(user_hash);
}
pub fn stage_remove(
&self,
batch: &mut fjall::OwnedWriteBatch,
uuid: &Uuid,
) -> Option<UserHash> {
let _guard = self.write_guard.lock();
let (_, user_hash) = self.cache.remove(uuid)?;
self.reverse.remove(&user_hash);
let forward_key = KeyBuilder::new()
.tag(KeyTag::USER_MAP)
.fixed(uuid.as_bytes())
.build();
let reverse_key = KeyBuilder::new()
.tag(KeyTag::USER_MAP_REVERSE)
.u64(user_hash.raw())
.build();
batch.remove(&self.repo_data, forward_key.as_slice());
batch.remove(&self.repo_data, reverse_key.as_slice());
Some(user_hash)
}
pub fn rollback_remove(&self, uuid: Uuid, user_hash: UserHash) {
self.cache.insert(uuid, user_hash);
self.reverse.insert(user_hash, uuid);
}
pub fn get(&self, uuid: &Uuid) -> Option<UserHash> {
self.cache.get(uuid).map(|r| *r)
}
pub fn get_uuid(&self, user_hash: &UserHash) -> Option<Uuid> {
self.reverse.get(user_hash).map(|r| *r)
}
pub fn len(&self) -> usize {
self.cache.len()
}
pub fn is_empty(&self) -> bool {
self.cache.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn open_temp() -> (tempfile::TempDir, fjall::Database, Keyspace) {
let dir = tempfile::TempDir::new().unwrap();
let db = fjall::Database::builder(dir.path()).open().unwrap();
let ks = db
.keyspace("repo_data", fjall::KeyspaceCreateOptions::default)
.unwrap();
(dir, db, ks)
}
#[test]
fn insert_and_lookup() {
let (_dir, db, ks) = open_temp();
let map = UserHashMap::new(ks);
let uuid = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:test123");
let mut batch = db.batch();
map.stage_insert(&mut batch, uuid, hash).unwrap();
batch.commit().unwrap();
assert_eq!(map.get(&uuid), Some(hash));
assert_eq!(map.get_uuid(&hash), Some(uuid));
assert_eq!(map.len(), 1);
}
#[test]
fn cache_populated_at_stage_time() {
let (_dir, db, ks) = open_temp();
let map = UserHashMap::new(ks);
let uuid = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:staged_only");
let mut batch = db.batch();
map.stage_insert(&mut batch, uuid, hash).unwrap();
assert_eq!(map.get(&uuid), Some(hash));
assert_eq!(map.get_uuid(&hash), Some(uuid));
assert_eq!(map.len(), 1);
}
#[test]
fn rollback_removes_from_cache() {
let (_dir, db, ks) = open_temp();
let map = UserHashMap::new(ks);
let uuid = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:rollback");
let mut batch = db.batch();
map.stage_insert(&mut batch, uuid, hash).unwrap();
assert_eq!(map.len(), 1);
map.rollback_insert(&uuid, &hash);
assert!(map.is_empty());
assert_eq!(map.get(&uuid), None);
assert_eq!(map.get_uuid(&hash), None);
drop(batch);
}
#[test]
fn load_all_after_reopen() {
let dir = tempfile::TempDir::new().unwrap();
let uuid = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:persist");
{
let db = fjall::Database::builder(dir.path()).open().unwrap();
let ks = db
.keyspace("repo_data", fjall::KeyspaceCreateOptions::default)
.unwrap();
let map = UserHashMap::new(ks);
let mut batch = db.batch();
map.stage_insert(&mut batch, uuid, hash).unwrap();
batch.commit().unwrap();
db.persist(fjall::PersistMode::SyncData).unwrap();
}
{
let db = fjall::Database::builder(dir.path()).open().unwrap();
let ks = db
.keyspace("repo_data", fjall::KeyspaceCreateOptions::default)
.unwrap();
let map = UserHashMap::new(ks);
let count = map.load_all().unwrap();
assert_eq!(count, 1);
assert_eq!(map.get(&uuid), Some(hash));
assert_eq!(map.get_uuid(&hash), Some(uuid));
}
}
#[test]
fn multiple_users() {
let (_dir, db, ks) = open_temp();
let map = UserHashMap::new(ks);
let pairs: Vec<_> = (0..10)
.map(|i| {
let uuid = Uuid::new_v4();
let hash = UserHash::from_did(&format!("did:plc:user{i}"));
(uuid, hash)
})
.collect();
let mut batch = db.batch();
pairs.iter().for_each(|(uuid, hash)| {
map.stage_insert(&mut batch, *uuid, *hash).unwrap();
});
batch.commit().unwrap();
assert_eq!(map.len(), 10);
pairs.iter().for_each(|(uuid, hash)| {
assert_eq!(map.get(uuid), Some(*hash));
assert_eq!(map.get_uuid(hash), Some(*uuid));
});
}
#[test]
fn stage_insert_idempotent_for_same_uuid() {
let (_dir, db, ks) = open_temp();
let map = UserHashMap::new(ks);
let uuid = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:same");
let mut batch = db.batch();
map.stage_insert(&mut batch, uuid, hash).unwrap();
batch.commit().unwrap();
let mut batch2 = db.batch();
map.stage_insert(&mut batch2, uuid, hash).unwrap();
}
#[test]
fn stage_insert_rejects_collision() {
let (_dir, db, ks) = open_temp();
let map = UserHashMap::new(ks);
let uuid_a = Uuid::new_v4();
let uuid_b = Uuid::new_v4();
let same_hash = UserHash::from_raw(0xDEAD_BEEF);
let mut batch = db.batch();
map.stage_insert(&mut batch, uuid_a, same_hash).unwrap();
batch.commit().unwrap();
let mut batch2 = db.batch();
let result = map.stage_insert(&mut batch2, uuid_b, same_hash);
assert!(matches!(
result,
Err(MetastoreError::UserHashCollision { .. })
));
}
#[test]
fn stage_remove_clears_cache_and_persists() {
let dir = tempfile::TempDir::new().unwrap();
let uuid = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:removable");
let db = fjall::Database::builder(dir.path()).open().unwrap();
let ks = db
.keyspace("repo_data", fjall::KeyspaceCreateOptions::default)
.unwrap();
let map = UserHashMap::new(ks);
let mut batch = db.batch();
map.stage_insert(&mut batch, uuid, hash).unwrap();
batch.commit().unwrap();
assert_eq!(map.len(), 1);
let mut remove_batch = db.batch();
let removed = map.stage_remove(&mut remove_batch, &uuid);
assert_eq!(removed, Some(hash));
remove_batch.commit().unwrap();
assert!(map.is_empty());
assert_eq!(map.get(&uuid), None);
assert_eq!(map.get_uuid(&hash), None);
db.persist(fjall::PersistMode::SyncData).unwrap();
drop(map);
drop(db);
let db2 = fjall::Database::builder(dir.path()).open().unwrap();
let ks2 = db2
.keyspace("repo_data", fjall::KeyspaceCreateOptions::default)
.unwrap();
let map2 = UserHashMap::new(ks2);
assert_eq!(map2.load_all().unwrap(), 0);
}
#[test]
fn stage_remove_returns_none_for_unknown() {
let (_dir, db, ks) = open_temp();
let map = UserHashMap::new(ks);
let mut batch = db.batch();
assert_eq!(map.stage_remove(&mut batch, &Uuid::new_v4()), None);
}
#[test]
fn rollback_remove_restores_cache() {
let (_dir, db, ks) = open_temp();
let map = UserHashMap::new(ks);
let uuid = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:rollback_remove");
let mut batch = db.batch();
map.stage_insert(&mut batch, uuid, hash).unwrap();
batch.commit().unwrap();
let mut remove_batch = db.batch();
map.stage_remove(&mut remove_batch, &uuid);
assert!(map.is_empty());
map.rollback_remove(uuid, hash);
assert_eq!(map.get(&uuid), Some(hash));
assert_eq!(map.get_uuid(&hash), Some(uuid));
assert_eq!(map.len(), 1);
drop(remove_batch);
}
#[test]
fn stage_remove_then_reinsert_same_did() {
let (_dir, db, ks) = open_temp();
let map = UserHashMap::new(ks);
let uuid_a = Uuid::new_v4();
let hash = UserHash::from_did("did:plc:reinsert");
let mut batch = db.batch();
map.stage_insert(&mut batch, uuid_a, hash).unwrap();
batch.commit().unwrap();
let mut remove_batch = db.batch();
map.stage_remove(&mut remove_batch, &uuid_a);
remove_batch.commit().unwrap();
let uuid_b = Uuid::new_v4();
let mut batch2 = db.batch();
map.stage_insert(&mut batch2, uuid_b, hash).unwrap();
batch2.commit().unwrap();
assert_eq!(map.get(&uuid_b), Some(hash));
assert_eq!(map.get_uuid(&hash), Some(uuid_b));
assert_eq!(map.get(&uuid_a), None);
}
}
+34
View File
@@ -180,6 +180,32 @@
# Can also be specified via environment variable `S3_ENDPOINT`.
#s3_endpoint =
# Repository backend: `postgres` by default, or `tranquil-store`, our embedded db.
# tranquil-store is EXPERIMENTAL!!!! RISK OF TOTAL DATA LOSS.
#
# Can also be specified via environment variable `REPO_BACKEND`.
#
# Default value: "postgres"
#repo_backend = "postgres"
[tranquil_store]
# Directory for tranquil-store data: the metastore, eventlog, and blockstore.
#
# Can also be specified via environment variable `TRANQUIL_STORE_DATA_DIR`.
#
# Default value: "/var/lib/tranquil-pds/store"
#data_dir = "/var/lib/tranquil-pds/store"
# Fjall block cache size in megabytes. Defaults to 20% of system RAM when unset.
#
# Can also be specified via environment variable `TRANQUIL_STORE_MEMORY_BUDGET_MB`.
#memory_budget_mb =
# Number of handler threads. Defaults to available_parallelism / 2.
#
# Can also be specified via environment variable `TRANQUIL_STORE_HANDLER_THREADS`.
#handler_threads =
[cache]
# Cache backend: `ripple` (default, built-in gossip) or `valkey`.
#
@@ -482,3 +508,11 @@
#
# Default value: 3600
#delete_check_interval_secs = 3600
# Interval in seconds between block garbage collection cycles.
# Reclaims orphaned ipld blocks that were stored but never committed.
#
# Can also be specified via environment variable `BLOCK_GC_INTERVAL_SECS`.
#
# Default value: 21600
#block_gc_interval_secs = 21600