mirror of
https://tangled.org/tranquil.farm/tranquil-pds
synced 2026-09-03 08:46:55 +00:00
feat(tranquil-store): metastore
Lewis: May this revision serve well! <lu5a@proton.me>
This commit is contained in:
+22
@@ -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"
|
||||
}
|
||||
+22
@@ -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"
|
||||
}
|
||||
+17
@@ -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"
|
||||
}
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"db_name": "PostgreSQL",
|
||||
"query": "DELETE FROM blocks WHERE cid = ANY($1)",
|
||||
"describe": {
|
||||
"columns": [],
|
||||
"parameters": {
|
||||
"Left": [
|
||||
"ByteaArray"
|
||||
]
|
||||
},
|
||||
"nullable": []
|
||||
},
|
||||
"hash": "8eecf8fef308716be88815eb59bb67ec7c534b3c821d55481b110e3e462ee366"
|
||||
}
|
||||
+14
@@ -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
@@ -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",
|
||||
]
|
||||
|
||||
|
||||
@@ -144,3 +144,7 @@ lto = "fat"
|
||||
strip = true
|
||||
codegen-units = 1
|
||||
panic = "abort"
|
||||
|
||||
[profile.bench]
|
||||
debug = 1
|
||||
strip = false
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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#"
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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(),
|
||||
));
|
||||
|
||||
|
||||
@@ -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 = []
|
||||
|
||||
@@ -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 --");
|
||||
}
|
||||
@@ -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
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user