feat(tranquil-store): blockstore

Lewis: May this revision serve well! <lu5a@proton.me>
This commit is contained in:
Lewis
2026-04-10 13:54:41 +03:00
parent dbc3c23c95
commit a9e3ec63e1
84 changed files with 9771 additions and 406 deletions
Generated
+259 -4
View File
@@ -1040,6 +1040,21 @@ dependencies = [
"virtue",
]
[[package]]
name = "bit-set"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3"
dependencies = [
"bit-vec",
]
[[package]]
name = "bit-vec"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
[[package]]
name = "bitflags"
version = "2.11.0"
@@ -1274,6 +1289,12 @@ dependencies = [
"either",
]
[[package]]
name = "byteview"
version = "0.10.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1c53ba0f290bfc610084c05582d9c5d421662128fc69f4bf236707af6fd321b9"
[[package]]
name = "cbc"
version = "0.1.2"
@@ -1497,6 +1518,12 @@ dependencies = [
"tokio-util",
]
[[package]]
name = "compare"
version = "0.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ea0095f6103c2a8b44acd6fd15960c801dafebf02e21940360833e0673f48ba7"
[[package]]
name = "compression-codecs"
version = "0.4.37"
@@ -1715,6 +1742,16 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-skiplist"
version = "0.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df29de440c58ca2cc6e587ec3d22347551a32435fbde9d2bff64e78a9ffa151b"
dependencies = [
"crossbeam-epoch",
"crossbeam-utils",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.21"
@@ -2212,6 +2249,18 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "enum_dispatch"
version = "0.3.13"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aa18ce2bc66555b3218614519ac839ddb759a7d6720732f979ef8d13be147ecd"
dependencies = [
"once_cell",
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "equivalent"
version = "1.0.2"
@@ -2345,6 +2394,23 @@ version = "0.5.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
[[package]]
name = "fjall"
version = "3.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1a9530ff159bc3ad3a15da746da0f6e95375c2ac64708cbb85ec1ebd26761a84"
dependencies = [
"byteorder-lite",
"byteview",
"dashmap",
"flume 0.12.0",
"log",
"lsm-tree",
"lz4_flex",
"tempfile",
"xxhash-rust",
]
[[package]]
name = "flate2"
version = "1.1.9"
@@ -2364,6 +2430,16 @@ checksum = "da0e4dd2a88388a1f4ccc7c9ce104604dab68d9f408dc34cd45823d5a9069095"
dependencies = [
"futures-core",
"futures-sink",
"nanorand",
"spin 0.9.8",
]
[[package]]
name = "flume"
version = "0.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be"
dependencies = [
"spin 0.9.8",
]
@@ -3364,7 +3440,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "525e9ff3e1a4be2fbea1fdf0e98686a6d98b4d8f937e1bf7402245af1909e8c3"
dependencies = [
"byteorder-lite",
"quick-error",
"quick-error 2.0.1",
]
[[package]]
@@ -3418,6 +3494,15 @@ dependencies = [
"generic-array",
]
[[package]]
name = "interval-heap"
version = "0.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "11274e5e8e89b8607cfedc2910b6626e998779b48a019151c7604d0adcb86ac6"
dependencies = [
"compare",
]
[[package]]
name = "inventory"
version = "0.3.22"
@@ -3458,9 +3543,9 @@ checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
[[package]]
name = "iri-string"
version = "0.7.10"
version = "0.7.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c91338f0783edbd6195decb37bae672fd3b165faffb89bf7b9e6942f8b1a731a"
checksum = "d8e7418f59cc01c88316161279a7f665217ae316b388e58a0d10e29f54f1e5eb"
dependencies = [
"memchr",
"serde",
@@ -4052,6 +4137,37 @@ version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
[[package]]
name = "lsm-tree"
version = "3.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9d67f95fd716870329c30aaeedf87f23d426564e6ce46efa045a91444faf2a19"
dependencies = [
"byteorder-lite",
"byteview",
"crossbeam-skiplist",
"enum_dispatch",
"interval-heap",
"log",
"lz4_flex",
"quick_cache",
"rustc-hash",
"self_cell",
"sfa",
"tempfile",
"varint-rs",
"xxhash-rust",
]
[[package]]
name = "lz4_flex"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "db9a0d582c2874f68138a16ce1867e0ffde6c0bb0a0df85e1f36d04146db488a"
dependencies = [
"twox-hash",
]
[[package]]
name = "match-lookup"
version = "0.1.2"
@@ -4265,6 +4381,15 @@ dependencies = [
"web-time",
]
[[package]]
name = "nanorand"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6a51313c5820b0b02bd422f4b44776fbf47961755c74ce64afc73bfad10226c3"
dependencies = [
"getrandom 0.2.17",
]
[[package]]
name = "nom"
version = "7.1.3"
@@ -4987,6 +5112,25 @@ dependencies = [
"yansi",
]
[[package]]
name = "proptest"
version = "1.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37566cb3fdacef14c0737f9546df7cfeadbfbc9fef10991038bf5015d0c80532"
dependencies = [
"bit-set",
"bit-vec",
"bitflags",
"num-traits",
"rand 0.9.2",
"rand_chacha 0.9.0",
"rand_xorshift",
"regex-syntax 0.8.10",
"rusty-fork",
"tempfile",
"unarray",
]
[[package]]
name = "prost"
version = "0.13.5"
@@ -5179,6 +5323,12 @@ dependencies = [
"winapi",
]
[[package]]
name = "quick-error"
version = "1.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
[[package]]
name = "quick-error"
version = "2.0.1"
@@ -5194,6 +5344,16 @@ dependencies = [
"memchr",
]
[[package]]
name = "quick_cache"
version = "0.6.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a70b1b8b47e31d0498ecbc3c5470bb931399a8bfed1fd79d1717a61ce7f96e3"
dependencies = [
"equivalent",
"hashbrown 0.16.1",
]
[[package]]
name = "quinn"
version = "0.11.9"
@@ -5346,6 +5506,15 @@ version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba"
[[package]]
name = "rand_xorshift"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a"
dependencies = [
"rand_core 0.9.5",
]
[[package]]
name = "rand_xoshiro"
version = "0.7.0"
@@ -5771,6 +5940,18 @@ version = "1.0.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
[[package]]
name = "rusty-fork"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2"
dependencies = [
"fnv",
"quick-error 1.2.3",
"tempfile",
"wait-timeout",
]
[[package]]
name = "ryu"
version = "1.0.23"
@@ -5883,6 +6064,12 @@ dependencies = [
"libc",
]
[[package]]
name = "self_cell"
version = "1.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b12e76d157a900eb52e81bc6e9f3069344290341720e9178cde2407113ac8d89"
[[package]]
name = "semver"
version = "1.0.27"
@@ -6057,6 +6244,17 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "sfa"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1296838937cab56cd6c4eeeb8718ec777383700c33f060e2869867bd01d1175"
dependencies = [
"byteorder-lite",
"log",
"xxhash-rust",
]
[[package]]
name = "sha1"
version = "0.10.6"
@@ -6486,7 +6684,7 @@ checksum = "c2d12fe70b2c1b4401038055f90f151b78208de1f9f89a7dbfd41587a10c3eea"
dependencies = [
"atoi",
"chrono",
"flume",
"flume 0.11.1",
"futures-channel",
"futures-core",
"futures-executor",
@@ -7618,6 +7816,36 @@ dependencies = [
"uuid",
]
[[package]]
name = "tranquil-store"
version = "0.4.7"
dependencies = [
"async-trait",
"bytes",
"cid",
"fjall",
"flume 0.11.1",
"futures",
"jacquard-common",
"jacquard-repo",
"k256",
"memmap2",
"multihash",
"parking_lot",
"postcard",
"proptest",
"rand 0.8.5",
"serde",
"serde_ipld_dagcbor",
"sha2",
"sqlx",
"tempfile",
"tokio",
"tracing",
"tranquil-repo",
"xxhash-rust",
]
[[package]]
name = "tranquil-sync"
version = "0.4.7"
@@ -7696,6 +7924,12 @@ dependencies = [
"utf-8",
]
[[package]]
name = "twox-hash"
version = "2.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c"
[[package]]
name = "typed-path"
version = "0.12.3"
@@ -7708,6 +7942,12 @@ version = "1.19.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb"
[[package]]
name = "unarray"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94"
[[package]]
name = "unicase"
version = "2.9.0"
@@ -7903,6 +8143,12 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65"
[[package]]
name = "varint-rs"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8f54a172d0620933a27a4360d3db3e2ae0dd6cceae9730751a036bbf182c4b23"
[[package]]
name = "vcpkg"
version = "0.2.15"
@@ -7927,6 +8173,15 @@ version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64"
[[package]]
name = "wait-timeout"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11"
dependencies = [
"libc",
]
[[package]]
name = "want"
version = "0.3.1"
+2
View File
@@ -22,6 +22,7 @@ members = [
"crates/tranquil-api",
"crates/tranquil-lexicon",
"crates/tranquil-signal",
"crates/tranquil-store",
]
[workspace.package]
@@ -51,6 +52,7 @@ tranquil-sync = { path = "crates/tranquil-sync" }
tranquil-oauth-server = { path = "crates/tranquil-oauth-server" }
tranquil-api = { path = "crates/tranquil-api" }
tranquil-signal = { path = "crates/tranquil-signal" }
tranquil-store = { path = "crates/tranquil-store" }
presage = { git = "https://github.com/whisperfish/presage", rev = "fe3ed54c4844ae51c3a9fa49cf80a7816a31a425", default-features = false }
+2 -1
View File
@@ -188,7 +188,8 @@ pub async fn put_preferences(
.collect();
if state
.repos.infra
.repos
.infra
.replace_namespace_preferences(user_id, APP_BSKY_NAMESPACE, prefs_to_save)
.await
.is_err()
@@ -19,7 +19,8 @@ pub async fn delete_account(
) -> Result<Json<EmptyResponse>, ApiError> {
let did = &input.did;
let (user_id, handle) = state
.repos.user
.repos
.user
.get_id_and_handle_by_did(did)
.await
.log_db_err("in delete_account")?
@@ -27,7 +28,8 @@ pub async fn delete_account(
.map(|row| (row.id, row.handle))?;
state
.repos.user
.repos
.user
.admin_delete_account_complete(user_id, did)
.await
.log_db_err("deleting account")?;
@@ -31,7 +31,8 @@ pub async fn send_email(
return Err(ApiError::InvalidRequest("content is required".into()));
}
let user = state
.repos.user
.repos
.user
.get_by_did(&input.recipient_did)
.await
.log_db_err("in send_email")?
@@ -45,7 +46,8 @@ pub async fn send_email(
.clone()
.unwrap_or_else(|| format!("Message from {}", hostname));
let result = state
.repos.infra
.repos
.infra
.enqueue_comms(
Some(user_id),
tranquil_db_traits::CommsChannel::Email,
+18 -9
View File
@@ -69,7 +69,8 @@ pub async fn get_account_info(
Query(params): Query<GetAccountInfoParams>,
) -> Result<Json<AccountInfo>, ApiError> {
let account = state
.repos.infra
.repos
.infra
.get_admin_account_info_by_did(&params.did)
.await
.log_db_err("in get_account_info")?
@@ -98,7 +99,8 @@ pub async fn get_account_info(
async fn get_invited_by(state: &AppState, user_id: uuid::Uuid) -> Option<InviteCodeInfo> {
let code = state
.repos.infra
.repos
.infra
.get_invite_code_used_by_user(user_id)
.await
.ok()??;
@@ -111,7 +113,8 @@ async fn get_invites_for_user(
user_id: uuid::Uuid,
) -> Option<Vec<InviteCodeInfo>> {
let invite_codes = state
.repos.infra
.repos
.infra
.get_invites_created_by_user(user_id)
.await
.ok()?;
@@ -123,7 +126,8 @@ async fn get_invites_for_user(
let code_strings: Vec<String> = invite_codes.iter().map(|ic| ic.code.clone()).collect();
let uses = state
.repos.infra
.repos
.infra
.get_invite_code_uses_batch(&code_strings)
.await
.ok()?;
@@ -157,7 +161,8 @@ async fn get_invite_code_info(state: &AppState, code: &str) -> Option<InviteCode
let info = state.repos.infra.get_invite_code_info(code).await.ok()??;
let uses = state
.repos.infra
.repos
.infra
.get_invite_code_uses(code)
.await
.ok()
@@ -197,7 +202,8 @@ pub async fn get_account_infos(
let dids_typed: Vec<Did> = dids.iter().filter_map(|d| d.parse().ok()).collect();
let accounts = state
.repos.infra
.repos
.infra
.get_admin_account_infos_by_dids(&dids_typed)
.await
.log_db_err("fetching account infos")?;
@@ -205,7 +211,8 @@ pub async fn get_account_infos(
let user_ids: Vec<uuid::Uuid> = accounts.iter().map(|u| u.id).collect();
let all_invite_codes = state
.repos.infra
.repos
.infra
.get_invite_codes_by_users(&user_ids)
.await
.unwrap_or_default();
@@ -217,7 +224,8 @@ pub async fn get_account_infos(
let all_invite_uses = if !all_codes.is_empty() {
state
.repos.infra
.repos
.infra
.get_invite_code_uses_batch(&all_codes)
.await
.unwrap_or_default()
@@ -226,7 +234,8 @@ pub async fn get_account_infos(
};
let invited_by_map: HashMap<uuid::Uuid, String> = state
.repos.infra
.repos
.infra
.get_invite_code_uses_by_users(&user_ids)
.await
.unwrap_or_default()
@@ -55,7 +55,8 @@ pub async fn search_accounts(
let handle_filter = params.handle.as_deref().map(|h| format!("%{}%", h));
let cursor_did: Option<Did> = params.cursor.as_ref().and_then(|c| c.parse().ok());
let rows = state
.repos.user
.repos
.user
.search_accounts(
cursor_did.as_ref(),
email_filter.as_deref(),
@@ -30,7 +30,8 @@ pub async fn update_account_email(
.map_err(|_| ApiError::InvalidDid("Invalid DID format".into()))?;
match state
.repos.user
.repos
.user
.admin_update_email(&account_did, email)
.await
{
@@ -73,7 +74,8 @@ pub async fn update_account_handle(
};
let old_handle = state.repos.user.get_handle_by_did(did).await.ok().flatten();
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(did)
.await
.ok()
@@ -81,14 +83,16 @@ pub async fn update_account_handle(
.ok_or(ApiError::AccountNotFound)?;
let handle_for_check: Handle = handle.parse().map_err(|_| ApiError::InvalidHandle(None))?;
if let Ok(true) = state
.repos.user
.repos
.user
.check_handle_exists(&handle_for_check, user_id)
.await
{
return Err(ApiError::HandleTaken);
}
match state
.repos.user
.repos
.user
.admin_update_handle(did, &handle_for_check)
.await
{
@@ -149,7 +153,8 @@ pub async fn update_account_password(
let password_hash = crate::common::hash_or_internal_error(password)?;
match state
.repos.user
.repos
.user
.admin_update_password(did, &password_hash)
.await
{
+31 -15
View File
@@ -53,7 +53,8 @@ pub async fn get_server_config(
];
let rows = state
.repos.infra
.repos
.infra
.get_server_configs(keys)
.await
.log_db_err("fetching server config")?;
@@ -86,7 +87,8 @@ pub async fn update_server_config(
));
}
state
.repos.infra
.repos
.infra
.upsert_server_config("server_name", trimmed)
.await
.log_db_err("upserting server_name")?;
@@ -95,13 +97,15 @@ pub async fn update_server_config(
if let Some(ref color) = req.primary_color {
if color.is_empty() {
state
.repos.infra
.repos
.infra
.delete_server_config("primary_color")
.await
.log_db_err("deleting primary_color")?;
} else if is_valid_hex_color(color) {
state
.repos.infra
.repos
.infra
.upsert_server_config("primary_color", color)
.await
.log_db_err("upserting primary_color")?;
@@ -115,13 +119,15 @@ pub async fn update_server_config(
if let Some(ref color) = req.primary_color_dark {
if color.is_empty() {
state
.repos.infra
.repos
.infra
.delete_server_config("primary_color_dark")
.await
.log_db_err("deleting primary_color_dark")?;
} else if is_valid_hex_color(color) {
state
.repos.infra
.repos
.infra
.upsert_server_config("primary_color_dark", color)
.await
.log_db_err("upserting primary_color_dark")?;
@@ -135,13 +141,15 @@ pub async fn update_server_config(
if let Some(ref color) = req.secondary_color {
if color.is_empty() {
state
.repos.infra
.repos
.infra
.delete_server_config("secondary_color")
.await
.log_db_err("deleting secondary_color")?;
} else if is_valid_hex_color(color) {
state
.repos.infra
.repos
.infra
.upsert_server_config("secondary_color", color)
.await
.log_db_err("upserting secondary_color")?;
@@ -155,13 +163,15 @@ pub async fn update_server_config(
if let Some(ref color) = req.secondary_color_dark {
if color.is_empty() {
state
.repos.infra
.repos
.infra
.delete_server_config("secondary_color_dark")
.await
.log_db_err("deleting secondary_color_dark")?;
} else if is_valid_hex_color(color) {
state
.repos.infra
.repos
.infra
.upsert_server_config("secondary_color_dark", color)
.await
.log_db_err("upserting secondary_color_dark")?;
@@ -174,7 +184,8 @@ pub async fn update_server_config(
if let Some(ref logo_cid) = req.logo_cid {
let old_logo_cid = state
.repos.infra
.repos
.infra
.get_server_config("logo_cid")
.await
.ok()
@@ -189,8 +200,11 @@ pub async fn update_server_config(
if let Some(old_cid_str) = should_delete_old {
match CidLink::new(old_cid_str) {
Ok(old_cid) => {
if let Ok(Some(storage_key)) =
state.repos.infra.get_blob_storage_key_by_cid(&old_cid).await
if let Ok(Some(storage_key)) = state
.repos
.infra
.get_blob_storage_key_by_cid(&old_cid)
.await
{
if let Err(e) = state.blob_store.delete(&storage_key).await {
error!("Failed to delete old logo blob from storage: {:?}", e);
@@ -211,13 +225,15 @@ pub async fn update_server_config(
if logo_cid.is_empty() {
state
.repos.infra
.repos
.infra
.delete_server_config("logo_cid")
.await
.log_db_err("deleting logo_cid")?;
} else {
state
.repos.infra
.repos
.infra
.upsert_server_config("logo_cid", logo_cid)
.await
.log_db_err("upserting logo_cid")?;
+12 -6
View File
@@ -32,7 +32,8 @@ pub async fn disable_invite_codes(
let accounts_typed: Vec<tranquil_types::Did> =
accounts.iter().filter_map(|a| a.parse().ok()).collect();
if let Err(e) = state
.repos.infra
.repos
.infra
.disable_invite_codes_by_account(&accounts_typed)
.await
{
@@ -87,7 +88,8 @@ pub async fn get_invite_codes(
};
let codes_rows = state
.repos.infra
.repos
.infra
.list_invite_codes(params.cursor.as_deref(), limit, sort_order)
.await
.log_db_err("fetching invite codes")?;
@@ -96,7 +98,8 @@ pub async fn get_invite_codes(
let code_strings: Vec<String> = codes_rows.iter().map(|r| r.code.clone()).collect();
let creator_dids: std::collections::HashMap<uuid::Uuid, tranquil_types::Did> = state
.repos.infra
.repos
.infra
.get_user_dids_by_ids(&user_ids)
.await
.unwrap_or_default()
@@ -108,7 +111,8 @@ pub async fn get_invite_codes(
} else {
common::group_invite_uses_by_code(
state
.repos.infra
.repos
.infra
.get_invite_code_uses_batch(&code_strings)
.await
.unwrap_or_default(),
@@ -168,7 +172,8 @@ pub async fn disable_account_invites(
.map_err(|_| ApiError::InvalidDid("Invalid DID format".into()))?;
match state
.repos.user
.repos
.user
.set_invites_disabled(&account_did, true)
.await
{
@@ -200,7 +205,8 @@ pub async fn enable_account_invites(
.map_err(|_| ApiError::InvalidDid("Invalid DID format".into()))?;
match state
.repos.user
.repos
.user
.set_invites_disabled(&account_did, false)
.await
{
+6 -3
View File
@@ -172,7 +172,8 @@ pub async fn update_subject_status(
None
};
state
.repos.user
.repos
.user
.set_user_takedown(&did, takedown_ref)
.await
.map_err(|e| {
@@ -249,7 +250,8 @@ pub async fn update_subject_status(
None
};
state
.repos.repo
.repos
.repo
.set_record_takedown(&cid, takedown_ref)
.await
.map_err(|e| {
@@ -282,7 +284,8 @@ pub async fn update_subject_status(
None
};
state
.repos.blob
.repos
.blob
.update_blob_takedown(&cid, takedown_ref)
.await
.map_err(|e| {
+32 -16
View File
@@ -24,7 +24,8 @@ pub async fn list_controllers(
auth: Auth<Active>,
) -> Result<Json<ControllersOutput<Vec<tranquil_db_traits::ControllerInfo>>>, ApiError> {
let controllers = state
.repos.delegation
.repos
.delegation
.get_delegations_for_account(&auth.did)
.await
.map_err(|e| {
@@ -100,7 +101,8 @@ pub async fn add_controller(
if resolved.is_local
&& state
.repos.delegation
.repos
.delegation
.is_delegated_account(&input.controller_did)
.await
.unwrap_or(false)
@@ -111,7 +113,8 @@ pub async fn add_controller(
}
match state
.repos.delegation
.repos
.delegation
.create_delegation(
can_add.did(),
&input.controller_did,
@@ -122,7 +125,8 @@ pub async fn add_controller(
{
Ok(_) => {
let _ = state
.repos.delegation
.repos
.delegation
.log_delegation_action(
can_add.did(),
can_add.did(),
@@ -159,13 +163,15 @@ pub async fn remove_controller(
Json(input): Json<RemoveControllerInput>,
) -> Result<Json<SuccessResponse>, ApiError> {
match state
.repos.delegation
.repos
.delegation
.revoke_delegation(&auth.did, &input.controller_did, &auth.did)
.await
{
Ok(true) => {
let revoked_app_passwords = state
.repos.session
.repos
.session
.delete_app_passwords_by_controller(&auth.did, &input.controller_did)
.await
.unwrap_or(0)
@@ -173,13 +179,15 @@ pub async fn remove_controller(
.unwrap_or(0usize);
let revoked_oauth_tokens = state
.repos.oauth
.repos
.oauth
.revoke_tokens_for_controller(&auth.did, &input.controller_did)
.await
.unwrap_or(0);
let _ = state
.repos.delegation
.repos
.delegation
.log_delegation_action(
&auth.did,
&auth.did,
@@ -218,13 +226,15 @@ pub async fn update_controller_scopes(
Json(input): Json<UpdateControllerScopesInput>,
) -> Result<Json<SuccessResponse>, ApiError> {
match state
.repos.delegation
.repos
.delegation
.update_delegation_scopes(&auth.did, &input.controller_did, &input.granted_scopes)
.await
{
Ok(true) => {
let _ = state
.repos.delegation
.repos
.delegation
.log_delegation_action(
&auth.did,
&auth.did,
@@ -255,7 +265,8 @@ pub async fn list_controlled_accounts(
auth: Auth<Active>,
) -> Result<Json<AccountsOutput<Vec<tranquil_db_traits::DelegatedAccountInfo>>>, ApiError> {
let accounts = state
.repos.delegation
.repos
.delegation
.get_accounts_controlled_by(&auth.did)
.await
.map_err(|e| {
@@ -287,7 +298,8 @@ pub async fn get_audit_log(
let offset = params.offset.max(0);
let entries = state
.repos.delegation
.repos
.delegation
.get_audit_log_for_account(&auth.did, limit, offset)
.await
.map_err(|e| {
@@ -296,7 +308,8 @@ pub async fn get_audit_log(
})?;
let total = state
.repos.delegation
.repos
.delegation
.count_audit_log_entries(&auth.did)
.await
.unwrap_or_default();
@@ -388,7 +401,8 @@ pub async fn create_delegated_account(
};
let user_id = match state
.repos.user
.repos
.user
.create_delegated_account(&create_input)
.await
{
@@ -407,7 +421,8 @@ pub async fn create_delegated_account(
if let Some(validated) = validated_invite_code
&& let Err(e) = state
.repos.infra
.repos
.infra
.record_invite_code_use(&validated, user_id)
.await
{
@@ -424,7 +439,8 @@ pub async fn create_delegated_account(
.await;
let _ = state
.repos.delegation
.repos
.delegation
.log_delegation_action(
&did,
&auth.did,
+2 -1
View File
@@ -169,7 +169,8 @@ async fn handle_command(state: AppState, interaction: Interaction) -> Response {
);
match state
.repos.user
.repos
.user
.store_discord_user_id(&discord_username, &discord_user_id, handle.as_deref())
.await
{
+12 -4
View File
@@ -67,14 +67,16 @@ async fn try_reactivate_migration(
new_email: email.clone(),
};
match state
.repos.user
.repos
.user
.reactivate_migration_account(&reactivate_input)
.await
{
Ok(reactivated) => {
info!(did = %did, old_handle = %reactivated.old_handle, new_handle = %handle, "Preparing existing account for inbound migration");
let secret_key_bytes = match state
.repos.user
.repos
.user
.get_user_key_by_id(reactivated.user_id)
.await
{
@@ -399,7 +401,8 @@ pub async fn create_account(
Err(_) => return ApiError::InvalidHandle(None).into_response(),
};
let handle_available = match state
.repos.user
.repos
.user
.check_handle_available_for_new_account(&handle_typed)
.await
{
@@ -522,7 +525,12 @@ pub async fn create_account(
birthdate_pref,
};
let create_result = match state.repos.user.create_password_account(&create_input).await {
let create_result = match state
.repos
.user
.create_password_account(&create_input)
.await
{
Ok(r) => r,
Err(tranquil_db_traits::CreateAccountError::HandleTaken) => {
return ApiError::HandleNotAvailable(None).into_response();
+16 -8
View File
@@ -165,7 +165,8 @@ async fn serve_handle_did_doc(state: &AppState, handle: &str, hostname: &str) ->
Err(_) => return ApiError::InvalidRequest("Invalid DID format".into()).into_response(),
};
let user = match state
.repos.user
.repos
.user
.get_user_for_did_doc_build(&expected_did_typed)
.await
{
@@ -182,7 +183,8 @@ async fn serve_handle_did_doc(state: &AppState, handle: &str, hostname: &str) ->
let did = expected_did;
let overrides = state
.repos.user
.repos
.user
.get_did_web_overrides(user_id)
.await
.ok()
@@ -218,7 +220,8 @@ pub async fn user_did_doc(State(state): State<AppState>, Path(handle): Path<Stri
}
};
let user = match state
.repos.user
.repos
.user
.get_did_web_info_by_handle(&current_handle_typed)
.await
{
@@ -246,7 +249,8 @@ pub async fn user_did_doc(State(state): State<AppState>, Path(handle): Path<Stri
}
let overrides = state
.repos.user
.repos
.user
.get_did_web_overrides(user_id)
.await
.ok()
@@ -468,7 +472,8 @@ pub async fn get_recommended_did_credentials(
auth: Auth<NotTakendown>,
) -> Result<Json<GetRecommendedDidCredentialsOutput>, ApiError> {
let handle = state
.repos.user
.repos
.user
.get_handle_by_did(&auth.did)
.await
.log_db_err("fetching handle for DID credentials")?
@@ -539,7 +544,8 @@ pub async fn update_handle(
)
.await?;
let user_row = state
.repos.user
.repos
.user
.get_id_and_handle_by_did(&did)
.await
.log_db_err("fetching user for handle update")?
@@ -661,7 +667,8 @@ pub async fn update_handle(
.parse()
.map_err(|_| ApiError::InvalidHandle(Some("Invalid handle format".into())))?;
let handle_exists = state
.repos.user
.repos
.user
.check_handle_exists(&handle_typed, user_id)
.await
.log_db_err("checking handle existence")?;
@@ -669,7 +676,8 @@ pub async fn update_handle(
return Err(ApiError::HandleTaken);
}
state
.repos.user
.repos
.user
.update_handle(user_id, &handle_typed)
.await
.map_err(|e| {
@@ -20,7 +20,8 @@ pub async fn request_plc_operation_signature(
tranquil_pds::oauth::scopes::IdentityAttr::Wildcard,
)?;
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(&auth.did)
.await
.log_db_err("fetching user id")?
@@ -30,7 +31,8 @@ pub async fn request_plc_operation_signature(
let plc_token = generate_plc_token();
let expires_at = Utc::now() + Duration::minutes(10);
state
.repos.infra
.repos
.infra
.insert_plc_token(user_id, &plc_token, expires_at)
.await
.log_db_err("creating PLC token")?;
+6 -3
View File
@@ -55,14 +55,16 @@ pub async fn sign_plc_operation(
})?;
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(did)
.await
.log_db_err("fetching user id")?
.ok_or(ApiError::AccountNotFound)?;
let token_expiry = state
.repos.infra
.repos
.infra
.get_plc_token_expiry(user_id, token)
.await
.log_db_err("fetching PLC token expiry")?
@@ -73,7 +75,8 @@ pub async fn sign_plc_operation(
return Err(ApiError::ExpiredToken(Some("Token has expired".into())));
}
let key_row = state
.repos.user
.repos
.user
.get_user_key_by_id(user_id)
.await
.log_db_err("fetching user key")?
@@ -38,14 +38,16 @@ pub async fn submit_plc_operation(
let hostname = &tranquil_config::get().server.hostname;
let public_url = format!("https://{}", hostname);
let user = state
.repos.user
.repos
.user
.get_id_and_handle_by_did(did)
.await
.log_db_err("fetching user")?
.ok_or(ApiError::AccountNotFound)?;
let key_row = state
.repos.user
.repos
.user
.get_user_key_by_id(user.id)
.await
.log_db_err("fetching user key")?
@@ -128,7 +130,8 @@ pub async fn submit_plc_operation(
.map_err(ApiError::from)?;
match state
.repos.repo
.repos
.repo
.insert_identity_event(did, Some(&user.handle))
.await
{
@@ -136,7 +136,8 @@ pub async fn resolve_signing_key(
match signing_key_did {
Some(key_did) => {
let key = state
.repos.infra
.repos
.infra
.get_reserved_signing_key(key_did)
.await
.map_err(|e| {
@@ -295,7 +296,8 @@ pub async fn create_and_store_session(
app_password_name: None,
};
state
.repos.session
.repos
.session
.create_session(&session_data)
.await
.map_err(|e| {
+30 -15
View File
@@ -26,7 +26,8 @@ pub async fn get_notification_prefs(
auth: Auth<Active>,
) -> Result<Json<NotificationPrefsOutput>, ApiError> {
let prefs = state
.repos.user
.repos
.user
.get_notification_prefs(&auth.did)
.await
.log_db_err("get notification prefs")?
@@ -65,14 +66,16 @@ pub async fn get_notification_history(
auth: Auth<Active>,
) -> Result<Json<GetNotificationHistoryOutput>, ApiError> {
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(&auth.did)
.await
.log_db_err("get user id by did")?
.ok_or(ApiError::AccountNotFound)?;
let rows = state
.repos.infra
.repos
.infra
.get_notification_history(user_id, 50)
.await
.log_db_err("get notification history")?;
@@ -165,7 +168,8 @@ pub async fn request_channel_verification(
hostname, encoded_token, encoded_identifier
);
let prefs = state
.repos.user
.repos
.user
.get_comms_prefs(user_id)
.await
.ok()
@@ -185,7 +189,8 @@ pub async fn request_channel_verification(
);
let recipient = match channel {
CommsChannel::Telegram => state
.repos.user
.repos
.user
.get_telegram_chat_id(user_id)
.await
.ok()
@@ -195,7 +200,8 @@ pub async fn request_channel_verification(
_ => identifier.to_string(),
};
state
.repos.infra
.repos
.infra
.enqueue_comms(
Some(user_id),
channel,
@@ -238,17 +244,20 @@ async fn process_messaging_channel_update(
}
match channel {
CommsChannel::Discord => state
.repos.user
.repos
.user
.clear_discord(user_id)
.await
.log_db_err("clear discord")?,
CommsChannel::Telegram => state
.repos.user
.repos
.user
.clear_telegram(user_id)
.await
.log_db_err("clear telegram")?,
CommsChannel::Signal => state
.repos.user
.repos
.user
.clear_signal(user_id)
.await
.log_db_err("clear signal")?,
@@ -281,17 +290,20 @@ async fn process_messaging_channel_update(
match channel {
CommsChannel::Discord => state
.repos.user
.repos
.user
.set_unverified_discord(user_id, &clean)
.await
.log_db_err("set unverified discord")?,
CommsChannel::Telegram => state
.repos.user
.repos
.user
.set_unverified_telegram(user_id, &clean)
.await
.log_db_err("set unverified telegram")?,
CommsChannel::Signal => state
.repos.user
.repos
.user
.set_unverified_signal(user_id, &clean)
.await
.log_db_err("set unverified signal")?,
@@ -313,7 +325,8 @@ pub async fn update_notification_prefs(
Json(input): Json<UpdateNotificationPrefsInput>,
) -> Result<Json<UpdateNotificationPrefsOutput>, ApiError> {
let user_row = state
.repos.user
.repos
.user
.get_id_handle_email_by_did(&auth.did)
.await
.log_db_err("get user by did")?
@@ -324,7 +337,8 @@ pub async fn update_notification_prefs(
let current_email = user_row.email;
let current_prefs = state
.repos.user
.repos
.user
.get_notification_prefs(&auth.did)
.await
.log_db_err("get notification prefs for update")?
@@ -347,7 +361,8 @@ pub async fn update_notification_prefs(
if input.preferred_channel.is_some() {
state
.repos.user
.repos
.user
.update_preferred_comms_channel(&auth.did, effective_channel)
.await
.log_db_err("update preferred channel")?;
+12 -6
View File
@@ -66,7 +66,8 @@ pub async fn upload_blob(
};
if state
.repos.user
.repos
.user
.is_account_migrated(&did)
.await
.unwrap_or(false)
@@ -78,7 +79,8 @@ pub async fn upload_blob(
get_header_str(&headers, http::header::CONTENT_TYPE).unwrap_or("application/octet-stream");
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(&did)
.await
.log_db_err("fetching user id for blob upload")?
@@ -143,7 +145,8 @@ pub async fn upload_blob(
);
match state
.repos.blob
.repos
.blob
.insert_blob(
&cid_link,
&mime_type,
@@ -177,7 +180,8 @@ pub async fn upload_blob(
if let Some(ref controller) = controller_did
&& let Err(e) = state
.repos.delegation
.repos
.delegation
.log_delegation_action(
&did,
controller,
@@ -236,7 +240,8 @@ pub async fn list_missing_blobs(
) -> Result<Json<ListMissingBlobsOutput>, ApiError> {
let did = &auth.did;
let user = state
.repos.user
.repos
.user
.get_by_did(did)
.await
.log_db_err("fetching user")?
@@ -245,7 +250,8 @@ pub async fn list_missing_blobs(
let limit = params.limit.unwrap_or(500).clamp(1, 1000);
let cursor = params.cursor.as_deref();
let missing = state
.repos.blob
.repos
.blob
.list_missing_blobs(user.id, cursor, limit + 1)
.await
.log_db_err("fetching missing blobs")?;
+14 -7
View File
@@ -34,7 +34,8 @@ pub async fn import_repo(
}
let did = &auth.did;
let user = state
.repos.user
.repos
.user
.get_by_did(did)
.await
.log_db_err("fetching user")?
@@ -44,7 +45,8 @@ pub async fn import_repo(
}
let user_id = user.id;
let expected_root_cid = state
.repos.repo
.repos
.repo
.get_repo_root_cid_by_user_id(user_id)
.await
.map_err(|e| {
@@ -232,7 +234,8 @@ pub async fn import_repo(
blob_refs.into_iter().unzip();
match state
.repos.blob
.repos
.blob
.insert_record_blobs(user_id, &record_uris, &blob_cids)
.await
{
@@ -248,7 +251,8 @@ pub async fn import_repo(
}
}
let key_row = state
.repos.user
.repos
.user
.get_user_with_key_by_did(did)
.await
.map_err(|e| {
@@ -289,7 +293,8 @@ pub async fn import_repo(
})?;
let new_root_cid_link = CidLink::from(&new_root_cid);
state
.repos.repo
.repos
.repo
.update_repo_root(user_id, &new_root_cid_link, &new_rev_str)
.await
.map_err(|e| {
@@ -299,7 +304,8 @@ pub async fn import_repo(
let mut all_block_cids: Vec<Vec<u8>> = blocks.keys().map(|c| c.to_bytes()).collect();
all_block_cids.push(new_root_cid.to_bytes());
state
.repos.repo
.repos
.repo
.insert_user_blocks(user_id, &all_block_cids, &new_rev_str)
.await
.map_err(|e| {
@@ -322,7 +328,8 @@ pub async fn import_repo(
"birthDate": "1998-05-06T00:00:00.000Z"
});
if let Err(e) = state
.repos.infra
.repos
.infra
.insert_account_preference_if_not_exists(
user_id,
"app.bsky.actor.defs#personalDetailsPref",
+2 -1
View File
@@ -23,7 +23,8 @@ pub async fn describe_repo(
Err(e) => return e.into_response(),
};
let collections = state
.repos.repo
.repos
.repo
.list_collections(resolved.user_id)
.await
.unwrap_or_default();
+2 -1
View File
@@ -304,7 +304,8 @@ pub async fn apply_writes(
require_verified_or_delegated(&state, batch_proof.user()).await?;
let user_id: uuid::Uuid = state
.repos.user
.repos
.user
.get_id_by_did(&did)
.await
.log_db_err("fetching user for batch write")?
@@ -101,7 +101,8 @@ pub async fn delete_record(
let deleted_uri = AtUri::from_parts(&did, &input.collection, &input.rkey);
if let Err(e) = state
.repos.backlink
.repos
.backlink
.remove_backlinks_by_uri(&deleted_uri)
.await
{
@@ -130,7 +131,8 @@ pub async fn delete_record_internal(
let _write_lock = state.repo_write_locks.lock(user_id).await;
let root_cid_str = state
.repos.repo
.repos
.repo
.get_repo_root_cid_by_user_id(user_id)
.await
.map_err(|e| CommitError::DatabaseError(e.to_string()))?
+4 -2
View File
@@ -65,7 +65,8 @@ pub async fn get_record(
Err(e) => return e.into_response(),
};
let record_row = state
.repos.repo
.repos
.repo
.get_record_cid(user_id, &input.collection, &input.rkey)
.await;
let record_cid_link = match record_row {
@@ -139,7 +140,8 @@ pub async fn list_records(
.as_ref()
.and_then(|c| c.parse::<tranquil_pds::types::Rkey>().ok());
let rows = match state
.repos.repo
.repos
.repo
.list_records(
user_id,
&input.collection,
+9 -3
View File
@@ -46,7 +46,8 @@ pub async fn prepare_repo_write<A: RepoScopeAction>(
let _account_verified = require_verified_or_delegated(state, user).await?;
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(principal_did.as_did())
.await
.log_db_err("fetching user for repo write")?
@@ -128,7 +129,8 @@ pub async fn create_record(
if !backlinks.is_empty() {
let conflicts = state
.repos.backlink
.repos
.backlink
.get_backlink_conflicts(user_id, &input.collection, &backlinks)
.await
.log_db_err("checking backlink conflicts")?;
@@ -244,7 +246,11 @@ pub async fn create_record(
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
&& let Err(e) = state
.repos
.backlink
.add_backlinks(user_id, &backlinks)
.await
{
error!("Failed to add backlinks for {}: {}", created_uri, e);
}
@@ -41,13 +41,15 @@ pub async fn check_account_status(
) -> Result<Json<CheckAccountStatusOutput>, ApiError> {
let did = &auth.did;
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(did)
.await
.log_db_err("fetching user ID for account status")?
.ok_or(ApiError::InternalError(None))?;
let is_active = state
.repos.user
.repos
.user
.is_account_active_by_did(did)
.await
.ok()
@@ -58,7 +60,8 @@ pub async fn check_account_status(
.map(|r| (r.repo_root_cid.to_string(), r.repo_rev))
.unwrap_or_else(|| (String::new(), None));
let block_count: i64 = state
.repos.repo
.repos
.repo
.count_user_blocks(user_id)
.await
.unwrap_or(0);
@@ -82,12 +85,14 @@ pub async fn check_account_status(
};
let record_count: i64 = state.repos.repo.count_records(user_id).await.unwrap_or(0);
let imported_blobs: i64 = state
.repos.blob
.repos
.blob
.count_blobs_by_user(user_id)
.await
.unwrap_or(0);
let expected_blobs: i64 = state
.repos.blob
.repos
.blob
.count_distinct_record_blobs(user_id)
.await
.unwrap_or(0);
@@ -339,7 +344,13 @@ pub async fn activate_account(
did_validation_start.elapsed()
);
let handle = state.repos.user.get_handle_by_did(&did).await.ok().flatten();
let handle = state
.repos
.user
.get_handle_by_did(&did)
.await
.ok()
.flatten();
info!(
"[MIGRATION] activateAccount: Activating account did={} handle={:?}",
did, handle
@@ -406,7 +417,8 @@ pub async fn activate_account(
info!("[MIGRATION] activateAccount: Identity event sequenced successfully");
}
let repo_root = state
.repos.repo
.repos
.repo
.get_repo_root_by_did(&did)
.await
.ok()
@@ -480,9 +492,19 @@ pub async fn deactivate_account(
let did = auth.did.clone();
let handle = state.repos.user.get_handle_by_did(&did).await.ok().flatten();
let handle = state
.repos
.user
.get_handle_by_did(&did)
.await
.ok()
.flatten();
let result = state.repos.user.deactivate_account(&did, delete_after).await;
let result = state
.repos
.user
.deactivate_account(&did, delete_after)
.await;
match result {
Ok(true) => {
@@ -518,7 +540,8 @@ pub async fn request_account_delete(
let session_mfa = require_legacy_session_mfa(&state, &auth).await?;
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(session_mfa.did())
.await
.ok()
@@ -527,7 +550,8 @@ pub async fn request_account_delete(
let confirmation_token = Uuid::new_v4().to_string();
let expires_at = Utc::now() + Duration::minutes(15);
state
.repos.infra
.repos
.infra
.create_deletion_request(&confirmation_token, session_mfa.did(), expires_at)
.await
.log_db_err("creating deletion token")?;
@@ -572,7 +596,8 @@ pub async fn delete_account(
return Err(ApiError::InvalidToken(Some("token is required".into())));
}
let user = state
.repos.user
.repos
.user
.get_user_for_deletion(did)
.await
.map_err(|e| {
@@ -595,7 +620,8 @@ pub async fn delete_account(
)));
}
let deletion_request = state
.repos.infra
.repos
.infra
.get_deletion_request(token)
.await
.map_err(|e| {
@@ -615,7 +641,8 @@ pub async fn delete_account(
return Err(ApiError::ExpiredToken(None));
}
state
.repos.user
.repos
.user
.delete_account_complete(user_id, did)
.await
.map_err(|e| {
+22 -11
View File
@@ -31,14 +31,16 @@ pub async fn list_app_passwords(
auth: Auth<Permissive>,
) -> Result<Json<ListAppPasswordsOutput>, ApiError> {
let user = state
.repos.user
.repos
.user
.get_by_did(&auth.did)
.await
.log_db_err("getting user")?
.ok_or(ApiError::AccountNotFound)?;
let rows = state
.repos.session
.repos
.session
.list_app_passwords(user.id)
.await
.log_db_err("listing app passwords")?;
@@ -83,7 +85,8 @@ pub async fn create_app_password(
Json(input): Json<CreateAppPasswordInput>,
) -> Result<Json<CreateAppPasswordOutput>, ApiError> {
let user = state
.repos.user
.repos
.user
.get_by_did(&auth.did)
.await
.log_db_err("getting user")?
@@ -95,7 +98,8 @@ pub async fn create_app_password(
}
if state
.repos.session
.repos
.session
.get_app_password_by_name(user.id, name)
.await
.log_db_err("checking app password")?
@@ -106,7 +110,8 @@ pub async fn create_app_password(
let (final_scopes, controller_did) = if let Some(ref controller) = auth.controller_did {
let grant = state
.repos.delegation
.repos
.delegation
.get_delegation(&auth.did, controller)
.await
.ok()
@@ -149,14 +154,16 @@ pub async fn create_app_password(
};
state
.repos.session
.repos
.session
.create_app_password(&create_data)
.await
.log_db_err("creating app password")?;
if let Some(ref controller) = controller_did {
let _ = state
.repos.delegation
.repos
.delegation
.log_delegation_action(
&auth.did,
controller,
@@ -192,7 +199,8 @@ pub async fn revoke_app_password(
Json(input): Json<RevokeAppPasswordInput>,
) -> Result<Json<EmptyResponse>, ApiError> {
let user = state
.repos.user
.repos
.user
.get_by_did(&auth.did)
.await
.log_db_err("getting user")?
@@ -204,13 +212,15 @@ pub async fn revoke_app_password(
}
let sessions_to_invalidate = state
.repos.session
.repos
.session
.get_session_jtis_by_app_password(&auth.did, name)
.await
.unwrap_or_default();
state
.repos.session
.repos
.session
.delete_sessions_by_app_password(&auth.did, name)
.await
.log_db_err("revoking sessions for app password")?;
@@ -225,7 +235,8 @@ pub async fn revoke_app_password(
.await;
state
.repos.session
.repos
.session
.delete_app_password(user.id, name)
.await
.log_db_err("revoking app password")?;
+24 -11
View File
@@ -61,7 +61,8 @@ pub async fn request_email_update(
auth.check_account_scope(AccountAttr::Email, AccountAction::Manage)?;
let user = state
.repos.user
.repos
.user
.get_email_info_by_did(&auth.did)
.await
.log_db_err("getting email info")?
@@ -141,7 +142,8 @@ pub async fn confirm_email(
let did = &auth.did;
let user = state
.repos.user
.repos
.user
.get_email_info_by_did(did)
.await
.log_db_err("getting email info")?
@@ -185,7 +187,8 @@ pub async fn confirm_email(
}
state
.repos.user
.repos
.user
.set_email_verified(user.id, true)
.await
.log_db_err("confirming email")?;
@@ -212,7 +215,8 @@ pub async fn update_email(
let did = &auth.did;
let user = state
.repos.user
.repos
.user
.get_email_info_by_did(did)
.await
.log_db_err("getting email info")?
@@ -259,7 +263,8 @@ pub async fn update_email(
}
state
.repos.infra
.repos
.infra
.upsert_account_preference(user_id, "email_auth_factor", json!(email_auth_factor))
.await
.map_err(|e| {
@@ -342,7 +347,8 @@ pub async fn update_email(
}
state
.repos.user
.repos
.user
.update_email(user_id, &new_email)
.await
.log_db_err("updating email")?;
@@ -370,7 +376,8 @@ pub async fn update_email(
}
if let Err(e) = state
.repos.infra
.repos
.infra
.upsert_account_preference(
user_id,
"email_auth_factor",
@@ -396,7 +403,8 @@ pub async fn check_email_verified(
Json(input): Json<CheckEmailVerifiedInput>,
) -> Result<Json<VerifiedResponse>, ApiError> {
let verified = state
.repos.user
.repos
.user
.check_email_verified_by_identifier(&input.identifier)
.await
.map_err(|e| {
@@ -420,7 +428,8 @@ pub async fn check_channel_verified(
Json(input): Json<CheckChannelVerifiedInput>,
) -> Result<Json<VerifiedResponse>, ApiError> {
let verified = state
.repos.user
.repos
.user
.check_channel_verified_by_did(&input.did, input.channel)
.await
.map_err(|e| {
@@ -480,7 +489,10 @@ pub async fn authorize_email_update(
let mut pending = match get_pending_email_update(state.cache.as_ref(), &did).await {
Some(p) => p,
None => {
warn!("authorize_email_update: no pending email update in cache for did={}", did);
warn!(
"authorize_email_update: no pending email update in cache for did={}",
did
);
return ApiError::InvalidRequest("No pending email update found".into())
.into_response();
}
@@ -558,7 +570,8 @@ pub async fn check_email_in_use(
}
let count = state
.repos.user
.repos
.user
.count_accounts_by_email(&email)
.await
.map_err(|e| {
+6 -3
View File
@@ -40,7 +40,8 @@ pub async fn create_invite_code(
let code = gen_invite_code();
match state
.repos.infra
.repos
.infra
.create_invite_code(&code, input.use_count, Some(&for_account))
.await
{
@@ -97,7 +98,8 @@ pub async fn create_invite_codes(
};
let admin_user_id = state
.repos.user
.repos
.user
.get_any_admin_user_id()
.await
.log_db_err("looking up admin user")?
@@ -174,7 +176,8 @@ pub async fn get_account_invite_codes(
let include_used = params.include_used.unwrap_or(true);
let codes_info = state
.repos.infra
.repos
.infra
.get_invite_codes_for_account(&auth.did)
.await
.log_db_err("fetching invite codes")?;
+10 -5
View File
@@ -42,7 +42,8 @@ pub async fn update_did_document(
}
let user = state
.repos.user
.repos
.user
.get_user_for_did_doc(&auth.did)
.await
.log_db_err("getting user")?
@@ -97,7 +98,8 @@ pub async fn update_did_document(
let also_known_as: Option<Vec<String>> = input.also_known_as.clone();
state
.repos.user
.repos
.user
.upsert_did_web_overrides(user.id, verification_methods_json, also_known_as)
.await
.log_db_err("upserting did_web_overrides")?;
@@ -105,7 +107,8 @@ pub async fn update_did_document(
if let Some(ref endpoint) = input.service_endpoint {
let endpoint_clean = endpoint.trim().trim_end_matches('/');
state
.repos.user
.repos
.user
.update_migrated_to_pds(&auth.did, endpoint_clean)
.await
.log_db_err("updating service endpoint")?;
@@ -149,7 +152,8 @@ async fn build_did_document(state: &AppState, did: &tranquil_pds::types::Did) ->
};
let overrides = state
.repos.user
.repos
.user
.get_did_web_overrides(user.id)
.await
.ok()
@@ -193,7 +197,8 @@ async fn build_did_document(state: &AppState, did: &tranquil_pds::types::Did) ->
}
let key_info = state
.repos.user
.repos
.user
.get_user_key_by_id(user.id)
.await
.ok()
@@ -451,7 +451,12 @@ pub async fn complete_passkey_setup(
State(state): State<AppState>,
Json(input): Json<CompletePasskeySetupInput>,
) -> Result<Json<CompletePasskeySetupOutput>, ApiError> {
let user = match state.repos.user.get_user_for_passkey_setup(&input.did).await {
let user = match state
.repos
.user
.get_user_for_passkey_setup(&input.did)
.await
{
Ok(Some(u)) => u,
Ok(None) => {
return Err(ApiError::AccountNotFound);
@@ -484,7 +489,8 @@ pub async fn complete_passkey_setup(
let webauthn = &state.webauthn_config;
let reg_state = match state
.repos.user
.repos
.user
.load_webauthn_challenge(&input.did, WebauthnChallengeType::Registration)
.await
{
@@ -530,7 +536,8 @@ pub async fn complete_passkey_setup(
}
};
if let Err(e) = state
.repos.user
.repos
.user
.save_passkey(
&input.did,
&credential_id,
@@ -559,7 +566,8 @@ pub async fn complete_passkey_setup(
}
let _ = state
.repos.user
.repos
.user
.delete_webauthn_challenge(&input.did, WebauthnChallengeType::Registration)
.await;
@@ -577,7 +585,12 @@ pub async fn start_passkey_registration_for_setup(
State(state): State<AppState>,
Json(input): Json<StartPasskeyRegistrationInput>,
) -> Result<Json<OptionsResponse<serde_json::Value>>, ApiError> {
let user = match state.repos.user.get_user_for_passkey_setup(&input.did).await {
let user = match state
.repos
.user
.get_user_for_passkey_setup(&input.did)
.await
{
Ok(Some(u)) => u,
Ok(None) => {
return Err(ApiError::AccountNotFound);
@@ -610,7 +623,8 @@ pub async fn start_passkey_registration_for_setup(
let webauthn = &state.webauthn_config;
let existing_passkeys = state
.repos.user
.repos
.user
.get_passkeys_for_user(&input.did)
.await
.unwrap_or_default();
@@ -643,7 +657,8 @@ pub async fn start_passkey_registration_for_setup(
}
};
if let Err(e) = state
.repos.user
.repos
.user
.save_webauthn_challenge(&input.did, WebauthnChallengeType::Registration, &state_json)
.await
{
@@ -682,7 +697,8 @@ pub async fn request_passkey_recovery(
NormalizedLoginIdentifier::normalize(&input.email, hostname_for_handles);
let user = match state
.repos.user
.repos
.user
.get_user_for_passkey_recovery(identifier, normalized_handle.as_str())
.await
{
@@ -697,7 +713,8 @@ pub async fn request_passkey_recovery(
let expires_at = Utc::now() + Duration::hours(1);
if let Err(e) = state
.repos.user
.repos
.user
.set_recovery_token(&user.did, &recovery_token_hash, expires_at)
.await
{
@@ -771,7 +788,8 @@ pub async fn recover_passkey_account(
password_hash,
};
let result = match state
.repos.user
.repos
.user
.recover_passkey_account(&recover_input)
.await
{
@@ -789,7 +807,8 @@ pub async fn recover_passkey_account(
let actual_channel =
tranquil_pds::comms::resolve_delivery_channel(&prefs, user.preferred_comms_channel);
if let Err(e) = state
.repos.user
.repos
.user
.set_channel_verified(&input.did, actual_channel)
.await
{
+16 -8
View File
@@ -28,14 +28,16 @@ pub async fn start_passkey_registration(
let webauthn = &state.webauthn_config;
let handle = state
.repos.user
.repos
.user
.get_handle_by_did(&auth.did)
.await
.log_db_err("fetching user")?
.ok_or(ApiError::AccountNotFound)?;
let existing_passkeys = state
.repos.user
.repos
.user
.get_passkeys_for_user(&auth.did)
.await
.log_db_err("fetching existing passkeys")?;
@@ -60,7 +62,8 @@ pub async fn start_passkey_registration(
})?;
state
.repos.user
.repos
.user
.save_webauthn_challenge(&auth.did, WebauthnChallengeType::Registration, &state_json)
.await
.log_db_err("saving registration state")?;
@@ -94,7 +97,8 @@ pub async fn finish_passkey_registration(
let webauthn = &state.webauthn_config;
let reg_state_json = state
.repos.user
.repos
.user
.load_webauthn_challenge(&auth.did, WebauthnChallengeType::Registration)
.await
.log_db_err("loading registration state")?
@@ -125,7 +129,8 @@ pub async fn finish_passkey_registration(
})?;
let passkey_id = state
.repos.user
.repos
.user
.save_passkey(
&auth.did,
passkey.cred_id(),
@@ -136,7 +141,8 @@ pub async fn finish_passkey_registration(
.log_db_err("saving passkey")?;
if let Err(e) = state
.repos.user
.repos
.user
.delete_webauthn_challenge(&auth.did, WebauthnChallengeType::Registration)
.await
{
@@ -177,7 +183,8 @@ pub async fn list_passkeys(
auth: Auth<Active>,
) -> Result<Json<ListPasskeysOutput>, ApiError> {
let passkeys = state
.repos.user
.repos
.user
.get_passkeys_for_user(&auth.did)
.await
.log_db_err("fetching passkeys")?;
@@ -243,7 +250,8 @@ pub async fn update_passkey(
let id: uuid::Uuid = input.id.parse().map_err(|_| ApiError::InvalidId)?;
match state
.repos.user
.repos
.user
.update_passkey_name(id, &auth.did, &input.friendly_name)
.await
{
+24 -12
View File
@@ -50,7 +50,8 @@ pub async fn request_password_reset(
};
let user_id = match state
.repos.user
.repos
.user
.get_id_by_email_or_handle(normalized, normalized_handle.as_str())
.await
{
@@ -72,7 +73,8 @@ pub async fn request_password_reset(
let code = generate_reset_code();
let expires_at = Utc::now() + Duration::minutes(10);
if let Err(e) = state
.repos.user
.repos
.user
.set_password_reset_code(user_id, &code, expires_at)
.await
{
@@ -153,7 +155,8 @@ pub async fn reset_password(
}
let password_hash = crate::common::hash_password_async(&password).await?;
let result = match state
.repos.user
.repos
.user
.reset_password_with_sessions(user_id, &password_hash)
.await
{
@@ -180,7 +183,8 @@ pub async fn reset_password(
let actual_channel =
tranquil_pds::comms::resolve_delivery_channel(&prefs, user.preferred_comms_channel);
if let Err(e) = state
.repos.user
.repos
.user
.set_channel_verified(&user.did, actual_channel)
.await
{
@@ -225,7 +229,8 @@ pub async fn change_password(
let password_mfa = verify_password_mfa(&state, &auth, &input.current_password).await?;
let user = state
.repos.user
.repos
.user
.get_id_and_password_hash_by_did(password_mfa.did())
.await
.log_db_err("in change_password")?
@@ -234,7 +239,8 @@ pub async fn change_password(
let new_hash = crate::common::hash_password_async(&input.new_password).await?;
state
.repos.user
.repos
.user
.update_password_hash(user.id, &new_hash)
.await
.log_db_err("updating password")?;
@@ -248,7 +254,8 @@ pub async fn get_password_status(
auth: Auth<Active>,
) -> Result<Json<HasPasswordResponse>, ApiError> {
let has = state
.repos.user
.repos
.user
.has_password_by_did(&auth.did)
.await
.log_db_err("checking password status")?
@@ -265,7 +272,8 @@ pub async fn remove_password(
let reauth_mfa = require_reauth_window(&state, &auth).await?;
let has_passkeys = state
.repos.user
.repos
.user
.has_passkeys(reauth_mfa.did())
.await
.unwrap_or(false);
@@ -276,7 +284,8 @@ pub async fn remove_password(
}
let user = state
.repos.user
.repos
.user
.get_password_info_by_did(reauth_mfa.did())
.await
.log_db_err("getting password info")?
@@ -289,7 +298,8 @@ pub async fn remove_password(
}
state
.repos.user
.repos
.user
.remove_user_password(user.id)
.await
.log_db_err("removing password")?;
@@ -322,7 +332,8 @@ pub async fn set_password(
let did = reauth_mfa.as_ref().map(|m| m.did()).unwrap_or(&auth.did);
let user = state
.repos.user
.repos
.user
.get_password_info_by_did(did)
.await
.log_db_err("getting password info")?
@@ -337,7 +348,8 @@ pub async fn set_password(
let new_hash = crate::common::hash_password_async(&new_password).await?;
state
.repos.user
.repos
.user
.set_new_user_password(user.id, &new_hash)
.await
.log_db_err("setting password")?;
+18 -9
View File
@@ -33,7 +33,8 @@ pub async fn get_reauth_status(
auth: Auth<Active>,
) -> Result<Json<ReauthStatusOutput>, ApiError> {
let last_reauth_at = state
.repos.session
.repos
.session
.get_last_reauth_at(&auth.did)
.await
.log_db_err("getting last reauth")?;
@@ -66,7 +67,8 @@ pub async fn reauth_password(
Json(input): Json<PasswordReauthInput>,
) -> Result<Json<ReauthOutput>, ApiError> {
let password_hash = state
.repos.user
.repos
.user
.get_password_hash_by_did(&auth.did)
.await
.log_db_err("fetching password hash")?
@@ -76,7 +78,8 @@ pub async fn reauth_password(
if !password_valid {
let app_password_hashes = state
.repos.session
.repos
.session
.get_app_password_hashes_by_did(&auth.did)
.await
.unwrap_or_default();
@@ -146,7 +149,8 @@ pub async fn reauth_passkey_start(
auth: Auth<Active>,
) -> Result<Json<PasskeyReauthStartOutput>, ApiError> {
let stored_passkeys = state
.repos.user
.repos
.user
.get_passkeys_for_user(&auth.did)
.await
.log_db_err("getting passkeys")?;
@@ -179,7 +183,8 @@ pub async fn reauth_passkey_start(
})?;
state
.repos.user
.repos
.user
.save_webauthn_challenge(
&auth.did,
WebauthnChallengeType::Authentication,
@@ -204,7 +209,8 @@ pub async fn reauth_passkey_finish(
Json(input): Json<PasskeyReauthFinishInput>,
) -> Result<Json<ReauthOutput>, ApiError> {
let auth_state_json = state
.repos.user
.repos
.user
.load_webauthn_challenge(&auth.did, WebauthnChallengeType::Authentication)
.await
.log_db_err("loading authentication state")?
@@ -232,7 +238,8 @@ pub async fn reauth_passkey_finish(
let cred_id_bytes = auth_result.cred_id().as_ref();
match state
.repos.user
.repos
.user
.update_passkey_counter(
cred_id_bytes,
i32::try_from(auth_result.counter()).unwrap_or(i32::MAX),
@@ -242,7 +249,8 @@ pub async fn reauth_passkey_finish(
Ok(false) => {
warn!(did = %&auth.did, "Passkey counter anomaly detected - possible cloned key");
let _ = state
.repos.user
.repos
.user
.delete_webauthn_challenge(&auth.did, WebauthnChallengeType::Authentication)
.await;
return Err(ApiError::PasskeyCounterAnomaly);
@@ -254,7 +262,8 @@ pub async fn reauth_passkey_finish(
}
let _ = state
.repos.user
.repos
.user
.delete_webauthn_challenge(&auth.did, WebauthnChallengeType::Authentication)
.await;
@@ -44,7 +44,8 @@ pub async fn reserve_signing_key(
let expires_at = Utc::now() + Duration::hours(24);
let private_bytes: &[u8] = &private_key_bytes;
match state
.repos.infra
.repos
.infra
.reserve_signing_key(
input.did.as_ref(),
&public_key_did_key,
+18 -7
View File
@@ -41,7 +41,8 @@ pub async fn create_totp_secret(
let secret = generate_totp_secret();
let handle = state
.repos.user
.repos
.user
.get_handle_by_did(&auth.did)
.await
.log_db_err("fetching handle")?
@@ -61,7 +62,8 @@ pub async fn create_totp_secret(
})?;
state
.repos.user
.repos
.user
.upsert_totp_secret(&auth.did, &encrypted_secret, ENCRYPTION_VERSION)
.await
.log_db_err("storing TOTP secret")?;
@@ -139,7 +141,8 @@ pub async fn enable_totp(
})?;
state
.repos.user
.repos
.user
.enable_totp_with_backup_codes(&auth.did, &backup_hashes)
.await
.log_db_err("enabling TOTP")?;
@@ -173,7 +176,8 @@ pub async fn disable_totp(
let totp_mfa = verify_totp_mfa(&state, &auth, &input.code).await?;
state
.repos.user
.repos
.user
.delete_totp_and_backup_codes(totp_mfa.did())
.await
.log_db_err("deleting TOTP")?;
@@ -209,7 +213,8 @@ pub async fn get_totp_status(
};
let backup_count = state
.repos.user
.repos
.user
.count_unused_backup_codes(&auth.did)
.await
.log_db_err("counting backup codes")?;
@@ -259,7 +264,8 @@ pub async fn regenerate_backup_codes(
})?;
state
.repos.user
.repos
.user
.replace_backup_codes(totp_mfa.did(), &backup_hashes)
.await
.log_db_err("replacing backup codes")?;
@@ -332,5 +338,10 @@ pub async fn verify_totp_or_backup_for_user(
}
pub async fn has_totp_enabled(state: &AppState, did: &tranquil_pds::types::Did) -> bool {
state.repos.user.has_totp_enabled(did).await.unwrap_or(false)
state
.repos
.user
.has_totp_enabled(did)
.await
.unwrap_or(false)
}
@@ -72,7 +72,8 @@ pub async fn list_trusted_devices(
auth: Auth<Active>,
) -> Result<Json<ListTrustedDevicesOutput>, ApiError> {
let rows = state
.repos.oauth
.repos
.oauth
.list_trusted_devices(&auth.did)
.await
.log_db_err("listing trusted devices")?;
@@ -108,7 +109,8 @@ pub async fn revoke_trusted_device(
Json(input): Json<RevokeTrustedDeviceInput>,
) -> Result<Json<SuccessResponse>, ApiError> {
match state
.repos.oauth
.repos
.oauth
.device_belongs_to_user(&input.device_id, &auth.did)
.await
{
@@ -123,7 +125,8 @@ pub async fn revoke_trusted_device(
}
state
.repos.oauth
.repos
.oauth
.revoke_device_trust(&input.device_id)
.await
.log_db_err("revoking device trust")?;
@@ -145,7 +148,8 @@ pub async fn update_trusted_device(
Json(input): Json<UpdateTrustedDeviceInput>,
) -> Result<Json<SuccessResponse>, ApiError> {
match state
.repos.oauth
.repos
.oauth
.device_belongs_to_user(&input.device_id, &auth.did)
.await
{
@@ -160,7 +164,8 @@ pub async fn update_trusted_device(
}
state
.repos.oauth
.repos
.oauth
.update_device_friendly_name(&input.device_id, input.friendly_name.as_deref())
.await
.log_db_err("updating device friendly name")?;
+18 -9
View File
@@ -79,7 +79,8 @@ async fn handle_migration_verification(
identifier: &str,
) -> Result<Json<VerifyTokenOutput>, ApiError> {
let user = state
.repos.user
.repos
.user
.get_verification_info(did)
.await
.log_db_err("during migration verification")?
@@ -92,7 +93,8 @@ async fn handle_migration_verification(
}
if !user.channel_verification.email {
state
.repos.user
.repos
.user
.set_email_verified_flag(user.id)
.await
.log_db_err("updating email_verified status")?;
@@ -118,7 +120,8 @@ async fn handle_channel_update(
identifier: &str,
) -> Result<Json<VerifyTokenOutput>, ApiError> {
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(did)
.await
.log_db_err("fetching user id")?
@@ -127,7 +130,8 @@ async fn handle_channel_update(
match channel {
CommsChannel::Email => {
let success = state
.repos.user
.repos
.user
.verify_email_channel(user_id, identifier)
.await
.log_db_err("updating email channel")?;
@@ -137,21 +141,24 @@ async fn handle_channel_update(
}
CommsChannel::Discord => {
state
.repos.user
.repos
.user
.verify_discord_channel(user_id, identifier)
.await
.log_db_err("updating discord channel")?;
}
CommsChannel::Telegram => {
state
.repos.user
.repos
.user
.verify_telegram_channel(user_id, identifier)
.await
.log_db_err("updating telegram channel")?;
}
CommsChannel::Signal => {
state
.repos.user
.repos
.user
.verify_signal_channel(user_id, identifier)
.await
.log_db_err("updating signal channel")?;
@@ -178,7 +185,8 @@ async fn notify_channel_verified(
) {
let recipient = match channel {
CommsChannel::Telegram => state
.repos.user
.repos
.user
.get_telegram_chat_id(user_id)
.await
.ok()
@@ -208,7 +216,8 @@ async fn handle_signup_verification(
identifier: &str,
) -> Result<Json<VerifyTokenOutput>, ApiError> {
let user = state
.repos.user
.repos
.user
.get_verification_info(did)
.await
.log_db_err("during signup verification")?
+2 -1
View File
@@ -71,7 +71,8 @@ pub async fn handle_telegram_webhook(
"Received /start from Telegram user"
);
match state
.repos.user
.repos
.user
.store_telegram_chat_id(&username, from.id, handle.as_deref())
.await
{
@@ -50,7 +50,8 @@ pub async fn consent_get(
) -> Response {
let consent_request_id = RequestId::from(query.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&consent_request_id)
.await
{
@@ -105,7 +106,8 @@ pub async fn consent_get(
.and_then(|s| s.parse().ok());
let delegation_grant = if let Some(ref ctrl_did) = controller_did_parsed {
state
.repos.delegation
.repos
.delegation
.get_delegation(&did, ctrl_did)
.await
.ok()
@@ -136,7 +138,8 @@ pub async fn consent_get(
let requested_scopes: Vec<&str> = expanded_scope_str.split_whitespace().collect();
let consent_client_id = ClientId::from(request_data.parameters.client_id.clone());
let preferences = state
.repos.oauth
.repos
.oauth
.get_scope_preferences(&did, &consent_client_id)
.await
.unwrap_or_default();
@@ -206,7 +209,8 @@ pub async fn consent_get(
.collect();
let account_handle = state
.repos.user
.repos
.user
.get_handle_by_did(&did)
.await
.ok()
@@ -216,7 +220,8 @@ pub async fn consent_get(
let (is_delegation, controller_did_resp, controller_handle, delegation_level) =
if let Some(ref ctrl_did) = controller_did_parsed {
let ctrl_handle = state
.repos.user
.repos
.user
.get_handle_by_did(ctrl_did)
.await
.ok()
@@ -273,7 +278,8 @@ pub async fn consent_post(
);
let consent_post_request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&consent_post_request_id)
.await
{
@@ -302,7 +308,8 @@ pub async fn consent_post(
},
Err(_) => {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&consent_post_request_id)
.await;
return json_error(
@@ -327,7 +334,8 @@ pub async fn consent_post(
let delegation_grant = match controller_did_parsed.as_ref() {
Some(ctrl_did) => state
.repos.delegation
.repos
.delegation
.get_delegation(&did, ctrl_did)
.await
.ok()
@@ -397,12 +405,14 @@ pub async fn consent_post(
.collect();
let consent_post_client_id = ClientId::from(request_data.parameters.client_id.clone());
let _ = state
.repos.oauth
.repos
.oauth
.upsert_scope_preferences(&did, &consent_post_client_id, &preferences)
.await;
}
if let Err(e) = state
.repos.oauth
.repos
.oauth
.update_request_scope(&consent_post_request_id, &approved_scope_str)
.await
{
@@ -415,7 +425,8 @@ pub async fn consent_post(
.map(|d| DeviceIdType::new(d.0.clone()));
let consent_post_code = AuthorizationCode::from(code.0.clone());
if state
.repos.oauth
.repos
.oauth
.update_authorization_request(
&consent_post_request_id,
&did,
@@ -458,7 +469,8 @@ pub async fn authorize_renew(
) -> Response {
let request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&request_id)
.await
{
@@ -499,7 +511,8 @@ pub async fn authorize_renew(
let staleness = now - request_data.expires_at;
if staleness.num_seconds() > MAX_RENEWAL_STALENESS_SECONDS {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&request_id)
.await;
return json_error(
@@ -511,7 +524,8 @@ pub async fn authorize_renew(
let new_expires_at = now + chrono::Duration::seconds(RENEW_EXPIRY_SECONDS);
match state
.repos.oauth
.repos
.oauth
.extend_authorization_request_expiry(&request_id, new_expires_at)
.await
{
@@ -25,7 +25,8 @@ pub async fn authorize_get(
};
let request_id = RequestId::from(request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&request_id)
.await
{
@@ -61,7 +62,8 @@ pub async fn authorize_get(
};
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&request_id)
.await;
if wants_json(&headers) {
@@ -104,14 +106,16 @@ pub async fn authorize_get(
tracing::info!(normalized = %normalized, "Normalized login_hint");
match state
.repos.user
.repos
.user
.get_login_check_by_handle_or_email(normalized.as_str())
.await
{
Ok(Some(user)) => {
tracing::info!(did = %user.did, has_password = user.password_hash.is_some(), "Found user for login_hint");
let is_delegated = state
.repos.delegation
.repos
.delegation
.is_delegated_account(&user.did)
.await
.unwrap_or(false);
@@ -121,7 +125,8 @@ pub async fn authorize_get(
if is_delegated {
tracing::info!("Redirecting to delegation auth");
if let Err(e) = state
.repos.oauth
.repos
.oauth
.set_request_did(&request_id, &user.did)
.await
{
@@ -159,7 +164,8 @@ pub async fn authorize_get(
if !force_new_account
&& let Some(device_id) = extract_device_cookie(&headers)
&& let Ok(accounts) = state
.repos.oauth
.repos
.oauth
.get_device_accounts(&device_id.clone())
.await
&& !accounts.is_empty()
@@ -191,14 +197,16 @@ pub async fn authorize_get_json(
.ok_or_else(|| OAuthError::InvalidRequest("request_uri is required".to_string()))?;
let request_id_json = RequestId::from(request_uri.clone());
let request_data = state
.repos.oauth
.repos
.oauth
.get_authorization_request(&request_id_json)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?
.ok_or_else(|| OAuthError::InvalidRequest("Invalid or expired request_uri".to_string()))?;
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&request_id_json)
.await;
return Err(OAuthError::InvalidRequest(
@@ -307,7 +315,8 @@ pub async fn authorize_post(
let json_response = wants_json(&headers);
let form_request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&form_request_id)
.await
{
@@ -344,7 +353,8 @@ pub async fn authorize_post(
};
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&form_request_id)
.await;
if json_response {
@@ -389,7 +399,8 @@ pub async fn authorize_post(
"Normalized username for lookup"
);
let user = match state
.repos.user
.repos
.user
.get_login_info_by_handle_or_email(normalized_username.as_str())
.await
{
@@ -411,7 +422,8 @@ pub async fn authorize_post(
}
if user.account_type.is_delegated() {
if state
.repos.oauth
.repos
.oauth
.set_authorization_did(&form_request_id, &user.did, None)
.await
.is_err()
@@ -439,7 +451,8 @@ pub async fn authorize_post(
if !user.password_required {
if state
.repos.oauth
.repos
.oauth
.set_authorization_did(&form_request_id, &user.did, None)
.await
.is_err()
@@ -522,7 +535,8 @@ pub async fn authorize_post(
}
} else {
if state
.repos.oauth
.repos
.oauth
.set_authorization_did(&form_request_id, &user.did, None)
.await
.is_err()
@@ -543,11 +557,13 @@ pub async fn authorize_post(
}
if user.two_factor_enabled {
let _ = state
.repos.oauth
.repos
.oauth
.delete_2fa_challenge_by_request_uri(&form_request_id)
.await;
match state
.repos.oauth
.repos
.oauth
.create_2fa_challenge(&user.did, &form_request_id)
.await
{
@@ -602,7 +618,8 @@ pub async fn authorize_post(
last_seen_at: Utc::now(),
};
if state
.repos.oauth
.repos
.oauth
.create_device(&new_device_id_typed, &device_data)
.await
.is_ok()
@@ -613,13 +630,15 @@ pub async fn authorize_post(
new_device_id_typed
};
let _ = state
.repos.oauth
.repos
.oauth
.upsert_account_device(&user.did, &final_device_id)
.await;
}
let set_auth_device_id = device_id.clone();
if state
.repos.oauth
.repos
.oauth
.set_authorization_did(&form_request_id, &user.did, set_auth_device_id.as_ref())
.await
.is_err()
@@ -673,7 +692,8 @@ pub async fn authorize_post(
let auth_post_device_id = device_id.clone();
let auth_post_code = AuthorizationCode::from(code.0.clone());
if state
.repos.oauth
.repos
.oauth
.update_authorization_request(
&form_request_id,
&user.did,
@@ -738,7 +758,8 @@ pub async fn authorize_select(
};
let select_request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&select_request_id)
.await
{
@@ -760,7 +781,8 @@ pub async fn authorize_select(
};
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&select_request_id)
.await;
return json_error(
@@ -791,7 +813,8 @@ pub async fn authorize_select(
};
let verify_device_id = device_id.clone();
let account_valid = match state
.repos.oauth
.repos
.oauth
.verify_account_on_device(&verify_device_id, &did)
.await
{
@@ -851,7 +874,8 @@ pub async fn authorize_select(
.await;
if !device_is_trusted {
if state
.repos.oauth
.repos
.oauth
.set_authorization_did(&select_request_id, &did, Some(&select_early_device_typed))
.await
.is_err()
@@ -872,11 +896,13 @@ pub async fn authorize_select(
}
if user.two_factor_enabled {
let _ = state
.repos.oauth
.repos
.oauth
.delete_2fa_challenge_by_request_uri(&select_request_id)
.await;
match state
.repos.oauth
.repos
.oauth
.create_2fa_challenge(&did, &select_request_id)
.await
{
@@ -915,12 +941,14 @@ pub async fn authorize_select(
}
let select_device_typed = device_id.clone();
let _ = state
.repos.oauth
.repos
.oauth
.upsert_account_device(&did, &select_device_typed)
.await;
if state
.repos.oauth
.repos
.oauth
.set_authorization_did(&select_request_id, &did, Some(&select_device_typed))
.await
.is_err()
@@ -245,7 +245,8 @@ pub async fn authorize_deny(
) -> Response {
let deny_request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&deny_request_id)
.await
{
@@ -272,7 +273,8 @@ pub async fn authorize_deny(
}
};
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&deny_request_id)
.await;
let redirect_uri = &request_data.parameters.redirect_uri;
@@ -295,7 +297,6 @@ pub struct AuthorizeDenyForm {
pub request_uri: String,
}
mod consent;
mod login;
mod passkey;
@@ -20,7 +20,8 @@ pub async fn check_user_has_passkeys(
BareLoginIdentifier::from_identifier(&query.identifier, hostname_for_handles);
let user = state
.repos.user
.repos
.user
.get_login_check_by_handle_or_email(bare_identifier.as_str())
.await;
@@ -52,7 +53,8 @@ pub async fn check_user_security_status(
NormalizedLoginIdentifier::normalize(&query.identifier, hostname_for_handles);
let user = state
.repos.user
.repos
.user
.get_login_check_by_handle_or_email(normalized_identifier.as_str())
.await;
@@ -68,7 +70,8 @@ pub async fn check_user_security_status(
let totp = tranquil_api::server::has_totp_enabled(&state, &u.did).await;
let has_pw = u.password_hash.is_some();
let has_controllers = state
.repos.delegation
.repos
.delegation
.is_delegated_account(&u.did)
.await
.unwrap_or(false);
@@ -113,7 +116,8 @@ pub async fn passkey_start(
) -> Response {
let passkey_start_request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&passkey_start_request_id)
.await
{
@@ -142,7 +146,8 @@ pub async fn passkey_start(
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&passkey_start_request_id)
.await;
return (
@@ -160,7 +165,8 @@ pub async fn passkey_start(
NormalizedLoginIdentifier::normalize(&form.identifier, hostname_for_handles);
let user = match state
.repos.user
.repos
.user
.get_login_info_by_handle_or_email(normalized_username.as_str())
.await
{
@@ -299,7 +305,8 @@ pub async fn passkey_start(
};
if let Err(e) = state
.repos.user
.repos
.user
.save_webauthn_challenge(
&user.did,
WebauthnChallengeType::Authentication,
@@ -322,7 +329,8 @@ pub async fn passkey_start(
Some(delegated_did_str) => match delegated_did_str.parse::<tranquil_types::Did>() {
Ok(delegated_did) if delegated_did != user.did => {
match state
.repos.delegation
.repos
.delegation
.get_delegation(&delegated_did, &user.did)
.await
{
@@ -359,7 +367,8 @@ pub async fn passkey_start(
"Passkey auth with delegated_did param - setting delegation flow"
);
if state
.repos.oauth
.repos
.oauth
.set_authorization_did(&passkey_start_request_id, &delegated_did, None)
.await
.is_err()
@@ -367,7 +376,8 @@ pub async fn passkey_start(
return OAuthError::ServerError("An error occurred.".into()).into_response();
}
if state
.repos.oauth
.repos
.oauth
.set_controller_did(&passkey_start_request_id, &user.did)
.await
.is_err()
@@ -381,7 +391,8 @@ pub async fn passkey_start(
"Passkey auth in delegation flow - preserving delegated DID"
);
if state
.repos.oauth
.repos
.oauth
.set_controller_did(&passkey_start_request_id, &user.did)
.await
.is_err()
@@ -389,7 +400,8 @@ pub async fn passkey_start(
return OAuthError::ServerError("An error occurred.".into()).into_response();
}
} else if state
.repos.oauth
.repos
.oauth
.set_authorization_did(&passkey_start_request_id, &user.did, None)
.await
.is_err()
@@ -415,7 +427,8 @@ pub async fn passkey_finish(
) -> Response {
let passkey_finish_request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&passkey_finish_request_id)
.await
{
@@ -444,7 +457,8 @@ pub async fn passkey_finish(
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&passkey_finish_request_id)
.await;
return (
@@ -491,7 +505,8 @@ pub async fn passkey_finish(
let passkey_owner_did = controller_did.as_ref().unwrap_or(&did);
let auth_state_json = match state
.repos.user
.repos
.user
.load_webauthn_challenge(passkey_owner_did, WebauthnChallengeType::Authentication)
.await
{
@@ -570,7 +585,8 @@ pub async fn passkey_finish(
};
if let Err(e) = state
.repos.user
.repos
.user
.delete_webauthn_challenge(passkey_owner_did, WebauthnChallengeType::Authentication)
.await
{
@@ -580,7 +596,8 @@ pub async fn passkey_finish(
if auth_result.needs_update() {
let cred_id_bytes = auth_result.cred_id().as_slice();
match state
.repos.user
.repos
.user
.update_passkey_counter(
cred_id_bytes,
i32::try_from(auth_result.counter()).unwrap_or(i32::MAX),
@@ -640,7 +657,8 @@ pub async fn passkey_finish(
let passkey_final_device_id = device_id.clone();
let passkey_final_code = AuthorizationCode::from(code.0.clone());
if state
.repos.oauth
.repos
.oauth
.update_authorization_request(
&passkey_finish_request_id,
&did,
@@ -691,7 +709,8 @@ pub async fn authorize_passkey_start(
) -> Response {
let auth_passkey_start_request_id = RequestId::from(query.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&auth_passkey_start_request_id)
.await
{
@@ -720,7 +739,8 @@ pub async fn authorize_passkey_start(
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&auth_passkey_start_request_id)
.await;
return (
@@ -822,7 +842,8 @@ pub async fn authorize_passkey_start(
};
if let Err(e) = state
.repos.user
.repos
.user
.save_webauthn_challenge(&did, WebauthnChallengeType::Authentication, &state_json)
.await
{
@@ -858,7 +879,8 @@ pub async fn authorize_passkey_finish(
let passkey_finish_request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&passkey_finish_request_id)
.await
{
@@ -887,7 +909,8 @@ pub async fn authorize_passkey_finish(
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&passkey_finish_request_id)
.await;
return (
@@ -929,7 +952,8 @@ pub async fn authorize_passkey_finish(
};
let auth_state_json = match state
.repos.user
.repos
.user
.load_webauthn_challenge(&did, WebauthnChallengeType::Authentication)
.await
{
@@ -1003,12 +1027,14 @@ pub async fn authorize_passkey_finish(
};
let _ = state
.repos.user
.repos
.user
.delete_webauthn_challenge(&did, WebauthnChallengeType::Authentication)
.await;
match state
.repos.user
.repos
.user
.update_passkey_counter(
credential.id.as_ref(),
i32::try_from(auth_result.counter()).unwrap_or(i32::MAX),
@@ -1033,7 +1059,8 @@ pub async fn authorize_passkey_finish(
}
let has_totp = state
.repos.user
.repos
.user
.has_totp_enabled(&did)
.await
.unwrap_or(false);
@@ -1064,11 +1091,13 @@ pub async fn authorize_passkey_finish(
};
let _ = state
.repos.oauth
.repos
.oauth
.delete_2fa_challenge_by_request_uri(&passkey_finish_request_id)
.await;
match state
.repos.oauth
.repos
.oauth
.create_2fa_challenge(&did, &passkey_finish_request_id)
.await
{
@@ -16,7 +16,8 @@ pub async fn register_complete(
let request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&request_id)
.await
{
@@ -50,7 +51,8 @@ pub async fn register_complete(
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&request_id)
.await;
return (
@@ -114,7 +116,8 @@ pub async fn register_complete(
}
let password_hashes = match state
.repos.session
.repos
.session
.get_app_password_hashes_by_did(&did)
.await
{
@@ -202,7 +205,8 @@ pub async fn register_complete(
}
if let Err(e) = state
.repos.oauth
.repos
.oauth
.set_authorization_did(&request_id, &did, None)
.await
{
@@ -257,7 +261,8 @@ pub async fn register_complete(
let code = Code::generate();
let auth_code = AuthorizationCode::from(code.0.clone());
if let Err(e) = state
.repos.oauth
.repos
.oauth
.update_authorization_request(&request_id, &did, None, &auth_code)
.await
{
@@ -317,7 +322,8 @@ pub async fn establish_session(
};
if let Err(e) = state
.repos.oauth
.repos
.oauth
.create_device(&device_typed, &device_data)
.await
{
@@ -333,7 +339,8 @@ pub async fn establish_session(
}
if let Err(e) = state
.repos.oauth
.repos
.oauth
.upsert_account_device(did, &device_typed)
.await
{
@@ -44,7 +44,8 @@ pub async fn authorize_2fa_get(
);
}
let _request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&twofa_request_id)
.await
{
@@ -88,7 +89,8 @@ pub async fn authorize_2fa_post(
};
let twofa_post_request_id = RequestId::from(form.request_uri.clone());
let request_data = match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&twofa_post_request_id)
.await
{
@@ -110,7 +112,8 @@ pub async fn authorize_2fa_post(
};
if request_data.expires_at < Utc::now() {
let _ = state
.repos.oauth
.repos
.oauth
.delete_authorization_request(&twofa_post_request_id)
.await;
return json_error(
@@ -120,7 +123,8 @@ pub async fn authorize_2fa_post(
);
}
let challenge = state
.repos.oauth
.repos
.oauth
.get_2fa_challenge(&twofa_post_request_id)
.await
.ok()
@@ -162,7 +166,8 @@ pub async fn authorize_2fa_post(
let twofa_totp_device_id = device_id.clone();
let twofa_totp_code = AuthorizationCode::from(code.0.clone());
if state
.repos.oauth
.repos
.oauth
.update_authorization_request(
&twofa_post_request_id,
&challenge.did,
@@ -250,7 +255,8 @@ pub async fn authorize_2fa_post(
last_seen_at: Utc::now(),
};
if state
.repos.oauth
.repos
.oauth
.create_device(&new_device_id_typed, &device_data)
.await
.is_ok()
@@ -262,7 +268,8 @@ pub async fn authorize_2fa_post(
}
};
let _ = state
.repos.oauth
.repos
.oauth
.upsert_account_device(&did, &trust_device_id)
.await;
let _ =
@@ -305,7 +312,8 @@ pub async fn authorize_2fa_post(
let twofa_final_device_id = device_id.clone();
let twofa_final_code = AuthorizationCode::from(code.0.clone());
if state
.repos.oauth
.repos
.oauth
.update_authorization_request(
&twofa_post_request_id,
&did,
@@ -25,7 +25,8 @@ fn parse_did(s: &str, label: &str) -> Result<Did, Response> {
async fn get_auth_request(state: &AppState, request_uri: &str) -> Result<RequestData, Response> {
let request_id = RequestId::from(request_uri.to_string());
match state
.repos.oauth
.repos
.oauth
.get_authorization_request(&request_id)
.await
{
@@ -43,7 +44,8 @@ async fn get_delegation_grant(
controller_did: &Did,
) -> Result<tranquil_db_traits::DelegationGrant, Response> {
match state
.repos.delegation
.repos
.delegation
.get_delegation(delegated_did, controller_did)
.await
{
@@ -65,7 +67,8 @@ async fn finalize_delegation_auth(
user_agent: Option<&str>,
) -> Response {
let _ = state
.repos.delegation
.repos
.delegation
.log_delegation_action(
delegated_did,
controller_did,
@@ -87,12 +90,14 @@ async fn bind_delegation_to_request(
) -> Result<(), Response> {
let request_id = RequestId::from(request_uri.to_string());
state
.repos.oauth
.repos
.oauth
.set_request_did(&request_id, delegated_did)
.await
.map_err(|_| DelegationAuthResponse::err("Failed to update authorization request"))?;
state
.repos.oauth
.repos
.oauth
.set_controller_did(&request_id, controller_did)
.await
.map_err(|_| DelegationAuthResponse::err("Failed to update authorization request"))?;
@@ -211,7 +216,8 @@ pub async fn delegation_auth(
let is_cross_pds = form.auth_method.as_deref() == Some("cross_pds");
let controller_local = state
.repos.user
.repos
.user
.get_auth_info_by_did(&controller_did)
.await
.ok()
@@ -562,7 +568,8 @@ pub async fn delegation_callback(
}
let _ = state
.repos.delegation
.repos
.delegation
.log_delegation_action(
delegated_did,
controller_did,
@@ -115,7 +115,8 @@ pub async fn pushed_authorization_request(
};
let request_id_typed = RequestIdType::from(request_id.0.clone());
state
.repos.oauth
.repos
.oauth
.create_authorization_request(&request_id_typed, &request_data)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?;
@@ -47,7 +47,8 @@ pub async fn handle_authorization_code_grant(
};
let auth_code = AuthorizationCode::from(code);
let auth_request = state
.repos.oauth
.repos
.oauth
.consume_authorization_request_by_code(&auth_code)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?
@@ -104,7 +105,8 @@ pub async fn handle_authorization_code_grant(
let token_endpoint = format!("https://{}/oauth/token", pds_hostname);
let result = verifier.verify_proof(proof, Method::POST.as_str(), &token_endpoint, None)?;
if !state
.repos.oauth
.repos
.oauth
.check_and_record_dpop_jti(&result.jti)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?
@@ -140,7 +142,8 @@ pub async fn handle_authorization_code_grant(
.parse()
.map_err(|_| OAuthError::InvalidRequest("Invalid controller DID format".to_string()))?;
let grant = state
.repos.delegation
.repos
.delegation
.get_delegation(&did_parsed, &controller_parsed)
.await
.ok()
@@ -200,7 +203,8 @@ pub async fn handle_authorization_code_grant(
controller_did: controller_did.clone(),
};
state
.repos.oauth
.repos
.oauth
.create_token(&token_data)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?;
@@ -320,7 +324,8 @@ pub async fn handle_refresh_token_grant(
"Refresh token reuse detected, revoking token family"
);
state
.repos.oauth
.repos
.oauth
.delete_token_family(original_token_id)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?;
@@ -331,7 +336,8 @@ pub async fn handle_refresh_token_grant(
RefreshTokenLookup::Expired { db_id } => {
tracing::warn!(refresh_token_prefix = %token_prefix, "Refresh token has expired");
state
.repos.oauth
.repos
.oauth
.delete_token_family(db_id)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?;
@@ -353,7 +359,8 @@ pub async fn handle_refresh_token_grant(
let token_endpoint = format!("https://{}/oauth/token", pds_hostname);
let result = verifier.verify_proof(proof, Method::POST.as_str(), &token_endpoint, None)?;
if !state
.repos.oauth
.repos
.oauth
.check_and_record_dpop_jti(&result.jti)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?
@@ -386,7 +393,8 @@ pub async fn handle_refresh_token_grant(
let new_expires_at = Utc::now() + Duration::days(refresh_expiry_days);
let new_refresh_typed = RefreshTokenType::from(new_refresh_token.0.clone());
state
.repos.oauth
.repos
.oauth
.rotate_token(db_id, &new_refresh_typed, new_expires_at)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?;
@@ -24,20 +24,23 @@ pub async fn revoke_token(
if let Some(token) = &request.token {
let refresh_token = RefreshToken::from(token.clone());
if let Some((db_id, _)) = state
.repos.oauth
.repos
.oauth
.get_token_by_refresh_token(&refresh_token)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?
{
state
.repos.oauth
.repos
.oauth
.delete_token_family(db_id)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?;
} else {
let token_id = TokenId::from(token.clone());
state
.repos.oauth
.repos
.oauth
.delete_token(&token_id)
.await
.map_err(tranquil_pds::oauth::db_err_to_oauth)?;
@@ -112,7 +112,8 @@ pub async fn sso_initiate(
_ => {
let request_id = RequestId::new(request_uri.clone());
let _request_data = state
.repos.oauth
.repos
.oauth
.get_authorization_request(&request_id)
.await?
.ok_or(ApiError::InvalidRequest(
@@ -135,7 +136,8 @@ pub async fn sso_initiate(
})?;
state
.repos.sso
.repos
.sso
.create_sso_auth_state(
&sso_state,
&request_uri,
@@ -363,7 +365,8 @@ async fn handle_sso_login(
user_info: &tranquil_pds::sso::providers::SsoUserInfo,
) -> Response {
let identity = match state
.repos.sso
.repos
.sso
.get_external_identity_by_provider(provider, &user_info.provider_user_id)
.await
{
@@ -371,7 +374,8 @@ async fn handle_sso_login(
Ok(None) => {
let token = generate_registration_token();
if let Err(e) = state
.repos.sso
.repos
.sso
.create_pending_registration(
&token,
request_uri,
@@ -398,7 +402,12 @@ async fn handle_sso_login(
}
};
let is_verified = match state.repos.user.get_session_info_by_did(&identity.did).await {
let is_verified = match state
.repos
.user
.get_session_info_by_did(&identity.did)
.await
{
Ok(Some(info)) => info.channel_verification.has_any_verified(),
Ok(None) => {
tracing::error!("User not found for SSO login: {}", identity.did);
@@ -423,7 +432,8 @@ async fn handle_sso_login(
}
if let Err(e) = state
.repos.sso
.repos
.sso
.update_external_identity_login(
identity.id,
user_info.username.as_deref(),
@@ -436,7 +446,8 @@ async fn handle_sso_login(
let request_id = RequestId::new(request_uri.to_string());
if let Err(e) = state
.repos.oauth
.repos
.oauth
.set_authorization_did(&request_id, &identity.did, None)
.await
{
@@ -478,7 +489,8 @@ async fn handle_sso_link(
user_info: &tranquil_pds::sso::providers::SsoUserInfo,
) -> Response {
let existing = state
.repos.sso
.repos
.sso
.get_external_identity_by_provider(provider, &user_info.provider_user_id)
.await;
@@ -517,7 +529,8 @@ async fn handle_sso_link(
}
if let Err(e) = state
.repos.sso
.repos
.sso
.create_external_identity(
&did,
provider,
@@ -551,7 +564,8 @@ async fn handle_sso_register(
user_info: &tranquil_pds::sso::providers::SsoUserInfo,
) -> Response {
match state
.repos.sso
.repos
.sso
.get_external_identity_by_provider(provider, &user_info.provider_user_id)
.await
{
@@ -569,7 +583,8 @@ async fn handle_sso_register(
let token = generate_registration_token();
if let Err(e) = state
.repos.sso
.repos
.sso
.create_pending_registration(
&token,
request_uri,
@@ -612,7 +627,8 @@ pub async fn get_linked_accounts(
auth: tranquil_pds::auth::Auth<tranquil_pds::auth::Active>,
) -> Result<Json<LinkedAccountsResponse>, ApiError> {
let identities = state
.repos.sso
.repos
.sso
.get_external_identities_by_did(&auth.did)
.await?;
@@ -657,7 +673,8 @@ pub async fn unlink_account(
let id = uuid::Uuid::parse_str(&input.id).map_err(|_| ApiError::InvalidId)?;
let has_password = state
.repos.user
.repos
.user
.has_password_by_did(&auth.did)
.await?
.unwrap_or(false);
@@ -667,7 +684,8 @@ pub async fn unlink_account(
if !has_password && !has_passkeys {
let identities = state
.repos.sso
.repos
.sso
.get_external_identities_by_did(&auth.did)
.await?;
@@ -680,7 +698,8 @@ pub async fn unlink_account(
}
let deleted = state
.repos.sso
.repos
.sso
.delete_external_identity(id, &auth.did)
.await?;
@@ -718,7 +737,8 @@ pub async fn get_pending_registration(
}
let pending = state
.repos.sso
.repos
.sso
.get_pending_registration(&query.token)
.await?
.ok_or(ApiError::SsoSessionExpired)?;
@@ -780,7 +800,8 @@ pub async fn check_handle_available(
};
let db_available = state
.repos.user
.repos
.user
.check_handle_available_for_new_account(&handle_typed)
.await
.unwrap_or(false);
@@ -850,7 +871,8 @@ pub async fn complete_registration(
}
let pending_preview = state
.repos.sso
.repos
.sso
.get_pending_registration(&input.token)
.await?
.ok_or(ApiError::SsoSessionExpired)?;
@@ -977,7 +999,8 @@ pub async fn complete_registration(
let handle_typed: tranquil_pds::types::Handle =
handle.parse().map_err(|_| ApiError::InvalidHandle(None))?;
let reserved = state
.repos.user
.repos
.user
.reserve_handle(&handle_typed, client_ip)
.await
.unwrap_or(false);
@@ -1178,7 +1201,8 @@ pub async fn complete_registration(
};
let _ = state
.repos.user
.repos
.user
.release_handle_reservation(&handle_typed)
.await;
@@ -1216,8 +1240,7 @@ pub async fn complete_registration(
let app_password = generate_app_password();
let app_password_name = "bsky.app".to_string();
let app_password_hash =
tranquil_api::common::hash_or_internal_error(&app_password)?;
let app_password_hash = tranquil_api::common::hash_or_internal_error(&app_password)?;
let app_password_data = tranquil_db_traits::AppPasswordCreate {
user_id: create_result.user_id,
@@ -1228,7 +1251,8 @@ pub async fn complete_registration(
created_by_controller_did: None,
};
if let Err(e) = state
.repos.session
.repos
.session
.create_app_password(&app_password_data)
.await
{
@@ -1240,7 +1264,8 @@ pub async fn complete_registration(
if !is_standalone {
let request_id = RequestId::new(pending_preview.request_uri.clone());
if let Err(e) = state
.repos.oauth
.repos
.oauth
.set_authorization_did(&request_id, &did_typed, None)
.await
{
@@ -1259,7 +1284,8 @@ pub async fn complete_registration(
);
let user_id = state
.repos.user
.repos
.user
.get_id_by_did(&did_typed)
.await
.unwrap_or(None);
@@ -1270,7 +1296,8 @@ pub async fn complete_registration(
if channel_auto_verified {
let _ = state
.repos.user
.repos
.user
.set_channel_verified(&did_typed, tranquil_db_traits::CommsChannel::Email)
.await;
tracing::info!(did = %did, "Auto-verified email from SSO provider");
+3 -3
View File
@@ -8,7 +8,7 @@ pub use error::ApiError;
pub use proxy_client::{AtUriParts, proxy_client, validate_at_uri, validate_limit};
pub use responses::{
AccountsOutput, AuditLogOutput, ControllersOutput, DidResponse, EmailUpdateStatusOutput,
EmptyResponse, HasPasswordResponse, InUseOutput, OptionsResponse,
PasswordResetOutput, PreferredLocaleOutput, PresetsOutput, StatusResponse, SuccessResponse,
TokenRequiredResponse, VerifiedResponse,
EmptyResponse, HasPasswordResponse, InUseOutput, OptionsResponse, PasswordResetOutput,
PreferredLocaleOutput, PresetsOutput, StatusResponse, SuccessResponse, TokenRequiredResponse,
VerifiedResponse,
};
@@ -22,7 +22,8 @@ pub async fn require_verified_or_delegated<'a>(
user: &'a AuthenticatedUser,
) -> Result<AccountVerified<'a>, ApiError> {
let is_verified = state
.repos.user
.repos
.user
.has_verified_comms_channel(&user.did)
.await
.unwrap_or(false);
@@ -32,7 +33,8 @@ pub async fn require_verified_or_delegated<'a>(
}
let is_delegated = state
.repos.delegation
.repos
.delegation
.is_delegated_account(&user.did)
.await
.unwrap_or(false);
+4 -2
View File
@@ -242,7 +242,8 @@ async fn verify_oauth_token_and_build_user(
{
Ok(result) => {
let user_info = state
.repos.user
.repos
.user
.get_user_info_by_did(&result.did)
.await
.ok()
@@ -321,7 +322,8 @@ async fn extract_auth_internal(
.unwrap_or_else(|| parts.uri.path().to_string());
let uri = build_full_url(&original_uri);
match validate_bearer_token_for_service_auth(state.repos.user.as_ref(), &extracted.token).await {
match validate_bearer_token_for_service_auth(state.repos.user.as_ref(), &extracted.token).await
{
Ok(user) if !user.auth_source.is_oauth() => {
return Ok(ExtractedAuth::User(user));
}
+12 -6
View File
@@ -99,7 +99,8 @@ pub async fn require_reauth_window<'a>(
use chrono::Utc;
let status = state
.repos.session
.repos
.session
.get_session_mfa_status(&user.did)
.await
.ok()
@@ -144,19 +145,22 @@ pub async fn require_reauth_window_if_available<'a>(
use crate::auth::reauth::check_reauth_required_cached;
let has_password = state
.repos.user
.repos
.user
.has_password_by_did(&user.did)
.await
.ok()
.flatten()
.unwrap_or(false);
let has_passkeys = state
.repos.user
.repos
.user
.has_passkeys(&user.did)
.await
.unwrap_or(false);
let has_totp = state
.repos.user
.repos
.user
.has_totp_enabled(&user.did)
.await
.unwrap_or(false);
@@ -188,7 +192,8 @@ pub async fn verify_password_mfa<'a>(
password: &str,
) -> Result<MfaVerified<'a>, crate::api::error::ApiError> {
let hash = state
.repos.user
.repos
.user
.get_password_hash_by_did(&user.did)
.await
.ok()
@@ -220,7 +225,8 @@ pub async fn verify_totp_mfa<'a>(
if is_backup_code_format(code) {
let backup_codes = state
.repos.user
.repos
.user
.get_unused_backup_codes(&user.did)
.await
.ok()
+2 -4
View File
@@ -599,10 +599,8 @@ pub async fn validate_token_with_dpop(
if !allow_takendown && status.is_takendown() {
return Err(TokenValidationError::AccountTakedown);
}
let key_bytes = try_decrypt_user_key(
user_info.key_bytes.as_deref(),
user_info.encryption_version,
);
let key_bytes =
try_decrypt_user_key(user_info.key_bytes.as_deref(), user_info.encryption_version);
Ok(AuthenticatedUser {
did: result_did,
key_bytes,
+2 -1
View File
@@ -30,7 +30,8 @@ pub async fn resolve_identity(
did: &Did,
) -> Result<ResolvedIdentity, DidResolutionError> {
let is_local = state
.repos.user
.repos
.user
.get_by_did(did)
.await
.ok()
+22 -11
View File
@@ -85,7 +85,8 @@ impl From<CommitError> for ApiError {
pub async fn get_current_root_cid(state: &AppState, user_id: Uuid) -> Result<CommitCid, ApiError> {
let root_cid_str = state
.repos.repo
.repos
.repo
.get_repo_root_cid_by_user_id(user_id)
.await
.map_err(|e| {
@@ -168,7 +169,8 @@ pub async fn begin_repo_write(
let write_lock = state.repo_write_locks.lock(user_id).await;
let root_cid_str = state
.repos.repo
.repos
.repo
.get_repo_root_cid_by_user_id(user_id)
.await
.map_err(|e| {
@@ -254,7 +256,8 @@ pub async fn finalize_repo_write(
if let Some(controller_did) = params.controller_did
&& let Some(detail) = params.delegation_detail
&& let Err(e) = state
.repos.delegation
.repos
.delegation
.log_delegation_action(
params.did,
controller_did,
@@ -351,7 +354,8 @@ pub async fn commit_and_log(
obsolete_cids,
} = params;
let key_row = state
.repos.user
.repos
.user
.get_user_key_by_id(user_id)
.await
.map_err(|e| CommitError::DatabaseError(format!("Failed to fetch signing key: {}", e)))?
@@ -485,7 +489,8 @@ pub async fn commit_and_log(
};
let _result = state
.repos.repo
.repos
.repo
.apply_commit(input)
.await
.map_err(|e| match e {
@@ -507,7 +512,8 @@ pub async fn create_record_internal(
record: &serde_json::Value,
) -> Result<(String, Cid), CommitError> {
let user_id: Uuid = state
.repos.user
.repos
.user
.get_id_by_did(did)
.await
.map_err(|e| CommitError::DatabaseError(e.to_string()))?
@@ -516,7 +522,8 @@ pub async fn create_record_internal(
let _write_lock = state.repo_write_locks.lock(user_id).await;
let root_cid_link = state
.repos.repo
.repos
.repo
.get_repo_root_cid_by_user_id(user_id)
.await
.map_err(|e| CommitError::DatabaseError(e.to_string()))?
@@ -610,7 +617,8 @@ pub async fn sequence_identity_event(
handle: Option<&Handle>,
) -> Result<SequenceNumber, CommitError> {
state
.repos.repo
.repos
.repo
.insert_identity_event(did, handle)
.await
.map_err(|e| CommitError::DatabaseError(format!("identity event: {}", e)))
@@ -621,7 +629,8 @@ pub async fn sequence_account_event(
status: tranquil_db_traits::AccountStatus,
) -> Result<SequenceNumber, CommitError> {
state
.repos.repo
.repos
.repo
.insert_account_event(did, status)
.await
.map_err(|e| CommitError::DatabaseError(format!("account event: {}", e)))
@@ -636,7 +645,8 @@ pub async fn sequence_sync_event(
.parse()
.map_err(|_| CommitError::InvalidCid(commit_cid.to_string()))?;
state
.repos.repo
.repos
.repo
.insert_sync_event(did, &cid_link, rev)
.await
.map_err(|e| CommitError::DatabaseError(format!("sync event: {}", e)))
@@ -652,7 +662,8 @@ pub async fn sequence_genesis_commit(
let commit_cid_link = crate::types::CidLink::from(commit_cid);
let mst_root_cid_link = crate::types::CidLink::from(mst_root_cid);
state
.repos.repo
.repos
.repo
.insert_genesis_commit_event(did, &commit_cid_link, &mst_root_cid_link, rev)
.await
.map_err(|e| CommitError::DatabaseError(format!("genesis commit event: {}", e)))
+1 -2
View File
@@ -222,8 +222,7 @@ fn walk_mst_node(
if let Some(record_cid) = entry.value
&& let Ok(full_key) = String::from_utf8(current_key.clone())
&& let Some(record_block) = blocks.get(&record_cid)
&& let Ok(record_value) =
serde_ipld_dagcbor::from_slice::<Ipld>(record_block)
&& let Ok(record_value) = serde_ipld_dagcbor::from_slice::<Ipld>(record_block)
{
let blob_refs = find_blob_refs_ipld(&record_value, 0);
let parts: Vec<&str> = full_key.split('/').collect();
+8 -1
View File
@@ -221,7 +221,14 @@ fn serialize_event_frame<P: serde::Serialize>(
payload: &P,
capacity: usize,
) -> Result<Vec<u8>, SyncFrameError> {
serialize_cbor_pair(&FrameHeader { op: 1, t: frame_type }, payload, capacity)
serialize_cbor_pair(
&FrameHeader {
op: 1,
t: frame_type,
},
payload,
capacity,
)
}
fn format_identity_event(event: &SequencedEvent) -> Result<Vec<u8>, SyncFrameError> {
+44
View File
@@ -0,0 +1,44 @@
[package]
name = "tranquil-store"
description = "Embedded storage engine for tranquil-pds (experimental)"
version.workspace = true
edition.workspace = true
license.workspace = true
[dependencies]
async-trait = { workspace = true }
xxhash-rust = { version = "0.8", features = ["xxh3"] }
serde = { workspace = true }
postcard = { version = "1", features = ["alloc"] }
parking_lot = { workspace = true }
fjall = "3"
flume = "0.11"
tokio = { workspace = true, features = ["sync", "rt"] }
bytes = "1"
memmap2 = "0.9"
tracing = { workspace = true }
jacquard-repo = { workspace = true }
cid = { workspace = true }
multihash = { workspace = true }
sha2 = { workspace = true }
[features]
test-harness = []
[dev-dependencies]
proptest = "1"
tempfile = "3"
futures = { workspace = true }
tokio = { workspace = true, features = ["sync", "rt-multi-thread", "macros", "time"] }
jacquard-common = { workspace = true }
tranquil-repo = { workspace = true }
tranquil-db = { workspace = true }
sqlx = { workspace = true }
k256 = { workspace = true }
rand = { workspace = true }
serde_ipld_dagcbor = { workspace = true }
tikv-jemallocator = "0.6"
[[bench]]
name = "blockstore"
harness = false
+577
View File
@@ -0,0 +1,577 @@
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use bytes::Bytes;
use cid::Cid;
use futures::StreamExt;
use jacquard_repo::storage::BlockStore;
use multihash::Multihash;
use sha2::{Digest, Sha256};
use tranquil_store::blockstore::{
BlockStoreConfig, DEFAULT_MAX_FILE_SIZE, GroupCommitConfig, TranquilBlockStore,
};
const DAG_CBOR_CODEC: u64 = 0x71;
const SHA2_256_CODE: u64 = 0x12;
fn make_block(index: usize) -> Vec<u8> {
let size = if index.is_multiple_of(5) {
1024 + (index.wrapping_mul(997)) % (63 * 1024)
} else {
64 + (index.wrapping_mul(131)) % 960
};
(0..size)
.map(|i| (index.wrapping_mul(257).wrapping_add(i.wrapping_mul(131)) & 0xFF) as u8)
.collect()
}
fn make_cid(data: &[u8]) -> Cid {
let hash = Sha256::digest(data);
let mh = Multihash::wrap(SHA2_256_CODE, &hash).unwrap();
Cid::new_v1(DAG_CBOR_CODEC, mh)
}
struct LatencyStats {
p50: Duration,
p95: Duration,
p99: Duration,
max: Duration,
mean: Duration,
}
fn compute_stats(durations: &mut [Duration]) -> Option<LatencyStats> {
if durations.is_empty() {
return None;
}
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 open_store(dir: &Path) -> TranquilBlockStore {
TranquilBlockStore::open(BlockStoreConfig {
data_dir: dir.join("data"),
index_dir: dir.join("index"),
max_file_size: DEFAULT_MAX_FILE_SIZE,
group_commit: GroupCommitConfig::default(),
})
.unwrap()
}
fn format_latency(stats: Option<&LatencyStats>) -> String {
match stats {
Some(s) => format!(
" | p50={:?} p95={:?} p99={:?} max={:?} mean={:?}",
s.p50, s.p95, s.p99, s.max, s.mean
),
None => String::new(),
}
}
async fn bench_write_throughput(block_count: usize, concurrency: usize) {
let dir = tempfile::TempDir::new().unwrap();
let store = open_store(dir.path());
let blocks_per_task = block_count / concurrency;
let actual_count = blocks_per_task * concurrency;
let blocks: Vec<Vec<u8>> = (0..actual_count).map(make_block).collect();
let total_bytes: usize = blocks.iter().map(Vec::len).sum();
let first_error: Arc<std::sync::Once> = Arc::new(std::sync::Once::new());
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let store = store.clone();
let first_error = Arc::clone(&first_error);
let task_blocks: Vec<Vec<u8>> =
blocks[task_id * blocks_per_task..(task_id + 1) * blocks_per_task].to_vec();
tokio::spawn(async move {
let mut latencies = Vec::with_capacity(task_blocks.len());
let mut errors = 0u64;
futures::stream::iter(task_blocks)
.then(|block| {
let store = store.clone();
let first_error = Arc::clone(&first_error);
async move {
let t = Instant::now();
match store.put(&block).await {
Ok(_) => Ok(t.elapsed()),
Err(e) => {
first_error.call_once(|| {
eprintln!("first put error: {e:?}");
});
Err(())
}
}
}
})
.for_each(|result| {
match result {
Ok(d) => latencies.push(d),
Err(()) => errors += 1,
}
async {}
})
.await;
(latencies, errors)
})
})
.collect();
let results: Vec<_> = futures::future::join_all(handles)
.await
.into_iter()
.map(Result::unwrap)
.collect();
let elapsed = start.elapsed();
let total_errors: u64 = results.iter().map(|(_, e)| e).sum();
let mut all_latencies: Vec<Duration> = results.into_iter().flat_map(|(l, _)| l).collect();
let successful = all_latencies.len();
let stats = compute_stats(&mut all_latencies);
let lat = format_latency(stats.as_ref());
if total_errors > 0 {
println!(
"{successful} ok, {total_errors} errors, {:.0} blocks/sec, {:.1} MB/sec, {:.1}ms{lat}",
successful as f64 / elapsed.as_secs_f64(),
total_bytes as f64 / elapsed.as_secs_f64() / (1024.0 * 1024.0),
elapsed.as_secs_f64() * 1000.0,
);
} else {
println!(
"{:.0} blocks/sec, {:.1} MB/sec, {:.1}ms{lat}",
actual_count as f64 / elapsed.as_secs_f64(),
total_bytes as f64 / elapsed.as_secs_f64() / (1024.0 * 1024.0),
elapsed.as_secs_f64() * 1000.0,
);
}
}
async fn bench_read_throughput(block_count: usize, concurrency: usize) {
let dir = tempfile::TempDir::new().unwrap();
let store = open_store(dir.path());
let cids_per_task = block_count / concurrency;
let actual_count = cids_per_task * concurrency;
let blocks: Vec<Vec<u8>> = (0..actual_count).map(make_block).collect();
let cids: Vec<Cid> = {
let pairs: Vec<(Cid, Bytes)> = blocks
.iter()
.map(|b| (make_cid(b), Bytes::from(b.clone())))
.collect();
let cids: Vec<Cid> = pairs.iter().map(|(c, _)| *c).collect();
store.put_many(pairs).await.unwrap();
cids
};
let run_reads = |label: &'static str, store: TranquilBlockStore, cids: Vec<Cid>| async move {
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let store = store.clone();
let task_cids: Vec<Cid> =
cids[task_id * cids_per_task..(task_id + 1) * cids_per_task].to_vec();
tokio::spawn(async move {
futures::stream::iter(task_cids)
.then(|cid| {
let store = store.clone();
async move {
let t = Instant::now();
let result = store.get(&cid).await.unwrap();
assert!(result.is_some());
t.elapsed()
}
})
.collect::<Vec<Duration>>()
.await
})
})
.collect();
let mut all_latencies: Vec<Duration> = futures::future::join_all(handles)
.await
.into_iter()
.flat_map(Result::unwrap)
.collect();
let elapsed = start.elapsed();
let stats = compute_stats(&mut all_latencies);
let lat = format_latency(stats.as_ref());
println!(
"{label}: {:.0} blocks/sec, {:.1}ms{lat}",
actual_count as f64 / elapsed.as_secs_f64(),
elapsed.as_secs_f64() * 1000.0,
);
};
run_reads("hot", store.clone(), cids.clone()).await;
#[cfg(target_os = "linux")]
{
if std::fs::write("/proc/sys/vm/drop_caches", "3").is_ok() {
println!("dropped system page caches");
std::thread::sleep(Duration::from_millis(100));
run_reads("cold", store.clone(), cids).await;
} else {
println!("cold: skipped, no root");
}
}
#[cfg(not(target_os = "linux"))]
{
println!("cold: skipped, not linux");
}
}
async fn bench_mixed_workload(block_count: usize, concurrency: usize) {
let dir = tempfile::TempDir::new().unwrap();
let store = open_store(dir.path());
let ops_per_task = block_count / concurrency;
let actual_ops = ops_per_task * concurrency;
let pre_populate = actual_ops / 2;
let blocks: Vec<Vec<u8>> = (0..pre_populate).map(make_block).collect();
let cids: Arc<Vec<Cid>> = Arc::new({
let pairs: Vec<(Cid, Bytes)> = blocks
.iter()
.map(|b| (make_cid(b), Bytes::from(b.clone())))
.collect();
let cids: Vec<Cid> = pairs.iter().map(|(c, _)| *c).collect();
store.put_many(pairs).await.unwrap();
cids
});
let read_count = Arc::new(AtomicU64::new(0));
let write_count = Arc::new(AtomicU64::new(0));
let timer_jitters: Arc<parking_lot::Mutex<Vec<Duration>>> =
Arc::new(parking_lot::Mutex::new(Vec::new()));
let timer_jitters_ref = Arc::clone(&timer_jitters);
let timer_handle = tokio::spawn(async move {
let jitters: Vec<Duration> = futures::stream::iter(0..100)
.then(|_| async {
let expected = Duration::from_millis(1);
let t = Instant::now();
tokio::time::sleep(expected).await;
t.elapsed().saturating_sub(expected)
})
.collect()
.await;
*timer_jitters_ref.lock() = jitters;
});
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let store = store.clone();
let cids = Arc::clone(&cids);
let read_count = Arc::clone(&read_count);
let write_count = Arc::clone(&write_count);
tokio::spawn(async move {
futures::stream::iter(0..ops_per_task)
.then(|op_idx| {
let store = store.clone();
let cids = Arc::clone(&cids);
let read_count = Arc::clone(&read_count);
let write_count = Arc::clone(&write_count);
async move {
let is_read = op_idx % 5 != 0;
if is_read && !cids.is_empty() {
let global_idx = task_id * ops_per_task + op_idx;
let cid_idx = global_idx % cids.len();
if store.get(&cids[cid_idx]).await.is_ok() {
read_count.fetch_add(1, Ordering::Relaxed);
}
} else {
let global_idx = task_id * ops_per_task + op_idx;
let block = make_block(pre_populate + global_idx);
if store.put(&block).await.is_ok() {
write_count.fetch_add(1, Ordering::Relaxed);
}
}
}
})
.collect::<Vec<()>>()
.await;
})
})
.collect();
futures::future::join_all(handles).await;
timer_handle.await.unwrap();
let elapsed = start.elapsed();
let reads = read_count.load(Ordering::Relaxed);
let writes = write_count.load(Ordering::Relaxed);
let total = reads + writes;
let jitters = timer_jitters.lock();
let max_jitter = jitters.iter().max().copied().unwrap_or_default();
let mean_jitter: Duration = if jitters.is_empty() {
Duration::ZERO
} else {
let jitter_divisor = u32::try_from(jitters.len()).unwrap_or(u32::MAX);
jitters.iter().sum::<Duration>() / jitter_divisor
};
println!(
"{:.0} ops/sec, {} reads + {} writes, {:.1}ms total",
total as f64 / elapsed.as_secs_f64(),
reads,
writes,
elapsed.as_secs_f64() * 1000.0,
);
let starvation_warning = if max_jitter > Duration::from_millis(5) {
" worker starvation detected"
} else {
""
};
println!(
"timer jitter: mean={:?} max={:?}{starvation_warning}",
mean_jitter, max_jitter
);
}
async fn bench_group_commit_effectiveness(block_count: usize) {
println!("-- group commit effectiveness at {block_count} blocks --");
let baseline_cycle_time = {
let dir = tempfile::TempDir::new().unwrap();
let store = open_store(dir.path());
let start = Instant::now();
futures::stream::iter(0..20)
.then(|i| {
let store = store.clone();
async move {
let block = make_block(i);
store.put(&block).await.unwrap();
}
})
.collect::<Vec<()>>()
.await;
let elapsed = start.elapsed();
drop(store);
elapsed.checked_div(20).unwrap_or(Duration::from_micros(1))
};
futures::stream::iter([1usize, 10, 50, 100])
.then(|concurrency| {
let baseline = baseline_cycle_time;
async move {
if block_count < concurrency {
return;
}
let dir = tempfile::TempDir::new().unwrap();
let store = open_store(dir.path());
let blocks_per_task = block_count / concurrency;
let actual_count = blocks_per_task * concurrency;
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let store = store.clone();
tokio::spawn(async move {
futures::stream::iter(0..blocks_per_task)
.then(|i| {
let store = store.clone();
async move {
let block = make_block(task_id * blocks_per_task + i);
store.put(&block).await.unwrap();
}
})
.collect::<Vec<()>>()
.await;
})
})
.collect();
futures::future::join_all(handles).await;
let elapsed = start.elapsed();
let blocks_per_sec = actual_count as f64 / elapsed.as_secs_f64();
let est_fsyncs = elapsed.as_secs_f64() / baseline.as_secs_f64();
let blocks_per_cycle = actual_count as f64 / est_fsyncs;
println!(
"concurrency={concurrency}: {blocks_per_sec:.0} blocks/sec, {:.1}ms, ~{est_fsyncs:.0} commit cycles, {blocks_per_cycle:.1} blocks/cycle",
elapsed.as_secs_f64() * 1000.0,
);
drop(store);
drop(dir);
}
})
.collect::<Vec<()>>()
.await;
}
async fn bench_postgres_write_throughput(block_count: usize, concurrency: usize) {
let database_url = match std::env::var("DATABASE_URL") {
Ok(url) => url,
Err(_) => {
println!("skipped, set DATABASE_URL to enable");
return;
}
};
let max_conns = u32::try_from(concurrency).expect("concurrency exceeds u32") + 5;
let pool = sqlx::postgres::PgPoolOptions::new()
.max_connections(max_conns)
.connect(&database_url)
.await
.unwrap();
sqlx::query("CREATE TABLE IF NOT EXISTS blocks (cid bytea PRIMARY KEY, data bytea NOT NULL)")
.execute(&pool)
.await
.unwrap();
sqlx::query("TRUNCATE blocks").execute(&pool).await.unwrap();
let pg_store = tranquil_repo::PostgresBlockStore::new(pool.clone());
let blocks_per_task = block_count / concurrency;
let actual_count = blocks_per_task * concurrency;
let blocks: Vec<Vec<u8>> = (0..actual_count).map(make_block).collect();
let total_bytes: usize = blocks.iter().map(Vec::len).sum();
let start = Instant::now();
let handles: Vec<_> = (0..concurrency)
.map(|task_id| {
let store = pg_store.clone();
let task_blocks: Vec<Vec<u8>> =
blocks[task_id * blocks_per_task..(task_id + 1) * blocks_per_task].to_vec();
tokio::spawn(async move {
futures::stream::iter(task_blocks)
.then(|block| {
let store = store.clone();
async move {
let t = Instant::now();
store.put(&block).await.unwrap();
t.elapsed()
}
})
.collect::<Vec<Duration>>()
.await
})
})
.collect();
let mut all_latencies: Vec<Duration> = futures::future::join_all(handles)
.await
.into_iter()
.flat_map(Result::unwrap)
.collect();
let elapsed = start.elapsed();
let stats = compute_stats(&mut all_latencies);
let lat = format_latency(stats.as_ref());
println!(
"{:.0} blocks/sec, {:.1} MB/sec, {:.1}ms{lat}",
actual_count as f64 / elapsed.as_secs_f64(),
total_bytes as f64 / elapsed.as_secs_f64() / (1024.0 * 1024.0),
elapsed.as_secs_f64() * 1000.0,
);
sqlx::query("TRUNCATE blocks").execute(&pool).await.unwrap();
pool.close().await;
}
fn main() {
let worker_threads = std::env::var("BENCH_WORKER_THREADS")
.ok()
.and_then(|s| s.trim().parse::<usize>().ok())
.unwrap_or_else(|| {
std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(8)
});
let rt = tokio::runtime::Builder::new_multi_thread()
.worker_threads(worker_threads)
.enable_all()
.build()
.unwrap();
println!("tokio worker threads: {worker_threads}");
let parse_env_list = |var: &str, defaults: Vec<usize>| -> Vec<usize> {
std::env::var(var).map_or(defaults, |s| {
s.split(',')
.map(|n| {
let trimmed = n.trim();
trimmed
.replace('_', "")
.parse::<usize>()
.unwrap_or_else(|_| panic!("{var}: failed to parse {trimmed:?} as integer"))
})
.collect()
})
};
let block_counts = parse_env_list("BENCH_BLOCK_COUNTS", vec![1_000, 10_000]);
let concurrency_levels = parse_env_list("BENCH_CONCURRENCY", vec![1, 10, 50]);
println!(
"block counts: {:?}, concurrency: {:?}",
block_counts, concurrency_levels
);
block_counts.iter().for_each(|&block_count| {
concurrency_levels.iter().for_each(|&concurrency| {
if block_count < concurrency {
return;
}
println!(
"-- write throughput: {} blocks, {} writers --",
block_count, concurrency
);
rt.block_on(bench_write_throughput(block_count, concurrency));
println!(
"-- read throughput: {} blocks, {} readers --",
block_count, concurrency
);
rt.block_on(bench_read_throughput(block_count, concurrency));
println!(
"-- mixed workload 80/20 r/w: {} ops, {} workers --",
block_count, concurrency
);
rt.block_on(bench_mixed_workload(block_count, concurrency));
});
});
rt.block_on(bench_group_commit_effectiveness(1000));
if std::env::var("DATABASE_URL").is_ok() {
block_counts.iter().for_each(|&block_count| {
concurrency_levels.iter().for_each(|&concurrency| {
if block_count < concurrency {
return;
}
println!(
"-- postgres write: {} blocks, {} writers --",
block_count, concurrency
);
rt.block_on(bench_postgres_write_throughput(block_count, concurrency));
});
});
} else {
println!("set DATABASE_URL for postgres comparison");
}
}
@@ -0,0 +1,723 @@
use std::io;
use crate::io::{FileId, StorageIO};
use super::types::{BlockLength, BlockLocation, BlockOffset, DataFileId, MAX_BLOCK_SIZE};
pub const BLOCK_MAGIC: [u8; 4] = *b"TQBL";
pub const BLOCK_FORMAT_VERSION: u8 = 1;
pub const BLOCK_HEADER_SIZE: usize = 5;
pub const CID_SIZE: usize = 36;
pub const BLOCK_RECORD_OVERHEAD: usize = CID_SIZE + 4 + 4;
pub type ValidBlock = (BlockOffset, [u8; CID_SIZE], Vec<u8>);
fn block_record_checksum(cid_bytes: &[u8; CID_SIZE], length_bytes: &[u8; 4], data: &[u8]) -> u32 {
let mut hasher = xxhash_rust::xxh3::Xxh3::new();
hasher.update(cid_bytes);
hasher.update(length_bytes);
hasher.update(data);
hasher.digest() as u32
}
pub fn encode_block_record<S: StorageIO>(
io: &S,
fd: FileId,
offset: BlockOffset,
cid_bytes: &[u8; CID_SIZE],
data: &[u8],
) -> io::Result<u64> {
let length = u32::try_from(data.len())
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "block data exceeds u32::MAX"))?;
if length > MAX_BLOCK_SIZE {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"block data exceeds MAX_BLOCK_SIZE",
));
}
let length_bytes = length.to_le_bytes();
let checksum = block_record_checksum(cid_bytes, &length_bytes, data);
let mut cursor = offset.raw();
io.write_all_at(fd, cursor, cid_bytes)?;
cursor += CID_SIZE as u64;
io.write_all_at(fd, cursor, &length_bytes)?;
cursor += 4;
io.write_all_at(fd, cursor, data)?;
cursor += data.len() as u64;
io.write_all_at(fd, cursor, &checksum.to_le_bytes())?;
cursor += 4;
Ok(cursor - offset.raw())
}
pub fn decode_block_record<S: StorageIO>(
io: &S,
fd: FileId,
offset: BlockOffset,
file_size: u64,
) -> io::Result<Option<ReadBlockRecord>> {
let raw = offset.raw();
let remaining = match file_size.checked_sub(raw) {
Some(r) => r,
None => return Ok(None),
};
if remaining == 0 {
return Ok(None);
}
if remaining < (CID_SIZE + 4) as u64 {
return Ok(Some(ReadBlockRecord::Truncated { offset }));
}
let mut cid_bytes = [0u8; CID_SIZE];
io.read_exact_at(fd, raw, &mut cid_bytes)?;
let mut length_bytes = [0u8; 4];
io.read_exact_at(fd, raw + CID_SIZE as u64, &mut length_bytes)?;
let length = u32::from_le_bytes(length_bytes);
if length > MAX_BLOCK_SIZE {
return Ok(Some(ReadBlockRecord::Corrupted { offset }));
}
let record_size = BLOCK_RECORD_OVERHEAD as u64 + u64::from(length);
if record_size > remaining {
return Ok(Some(ReadBlockRecord::Truncated { offset }));
}
let data_offset = raw + CID_SIZE as u64 + 4;
let mut data = vec![0u8; length as usize];
io.read_exact_at(fd, data_offset, &mut data)?;
let mut checksum_bytes = [0u8; 4];
io.read_exact_at(fd, data_offset + u64::from(length), &mut checksum_bytes)?;
let stored_checksum = u32::from_le_bytes(checksum_bytes);
let computed_checksum = block_record_checksum(&cid_bytes, &length_bytes, &data);
if stored_checksum != computed_checksum {
return Ok(Some(ReadBlockRecord::Corrupted { offset }));
}
Ok(Some(ReadBlockRecord::Valid {
offset,
cid_bytes,
data,
}))
}
#[must_use]
#[derive(Debug)]
pub enum ReadBlockRecord {
Valid {
offset: BlockOffset,
cid_bytes: [u8; CID_SIZE],
data: Vec<u8>,
},
Corrupted {
offset: BlockOffset,
},
Truncated {
offset: BlockOffset,
},
}
pub struct DataFileWriter<'a, S: StorageIO> {
io: &'a S,
fd: FileId,
file_id: DataFileId,
position: BlockOffset,
}
impl<'a, S: StorageIO> DataFileWriter<'a, S> {
pub fn new(io: &'a S, fd: FileId, file_id: DataFileId) -> io::Result<Self> {
let mut header = [0u8; BLOCK_HEADER_SIZE];
header[..4].copy_from_slice(&BLOCK_MAGIC);
header[4] = BLOCK_FORMAT_VERSION;
io.write_all_at(fd, 0, &header)?;
Ok(Self {
io,
fd,
file_id,
position: BlockOffset::new(BLOCK_HEADER_SIZE as u64),
})
}
pub fn resume(io: &'a S, fd: FileId, file_id: DataFileId, position: BlockOffset) -> Self {
assert!(
position.raw() >= BLOCK_HEADER_SIZE as u64,
"resume position {position:?} is before header end"
);
Self {
io,
fd,
file_id,
position,
}
}
pub fn append_block(
&mut self,
cid_bytes: &[u8; CID_SIZE],
data: &[u8],
) -> io::Result<BlockLocation> {
let record_offset = self.position;
let bytes_written = encode_block_record(self.io, self.fd, record_offset, cid_bytes, data)?;
self.position = self.position.advance(bytes_written);
Ok(BlockLocation {
file_id: self.file_id,
offset: record_offset,
length: BlockLength::new(
u32::try_from(data.len()).expect("encode_block_record validated length"),
),
})
}
pub fn sync(&self) -> io::Result<()> {
self.io.sync(self.fd)
}
pub fn position(&self) -> BlockOffset {
self.position
}
pub fn fd(&self) -> FileId {
self.fd
}
pub fn file_id(&self) -> DataFileId {
self.file_id
}
}
pub struct DataFileReader<'a, S: StorageIO> {
io: &'a S,
fd: FileId,
position: BlockOffset,
file_size: u64,
}
impl<'a, S: StorageIO> DataFileReader<'a, S> {
pub fn open(io: &'a S, fd: FileId) -> io::Result<Self> {
let file_size = io.file_size(fd)?;
if file_size < BLOCK_HEADER_SIZE as u64 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"file too small for header",
));
}
let mut header = [0u8; BLOCK_HEADER_SIZE];
io.read_exact_at(fd, 0, &mut header)?;
if header[..4] != BLOCK_MAGIC {
return Err(io::Error::new(io::ErrorKind::InvalidData, "bad magic"));
}
if header[4] != BLOCK_FORMAT_VERSION {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"unsupported block format version",
));
}
Ok(Self {
io,
fd,
position: BlockOffset::new(BLOCK_HEADER_SIZE as u64),
file_size,
})
}
pub fn valid_blocks(self) -> io::Result<Vec<ValidBlock>> {
self.map_while(|r| match r {
Ok(ReadBlockRecord::Valid {
offset,
cid_bytes,
data,
}) => Some(Ok((offset, cid_bytes, data))),
Err(e) => Some(Err(e)),
_ => None,
})
.collect()
}
}
impl<S: StorageIO> Iterator for DataFileReader<'_, S> {
type Item = io::Result<ReadBlockRecord>;
fn next(&mut self) -> Option<Self::Item> {
match decode_block_record(self.io, self.fd, self.position, self.file_size) {
Err(e) => {
self.position = BlockOffset::new(self.file_size);
Some(Err(e))
}
Ok(None) => None,
Ok(Some(record)) => {
match &record {
ReadBlockRecord::Valid { data, .. } => {
self.position = self
.position
.advance(BLOCK_RECORD_OVERHEAD as u64 + data.len() as u64);
}
ReadBlockRecord::Corrupted { .. } | ReadBlockRecord::Truncated { .. } => {
self.position = BlockOffset::new(self.file_size);
}
}
Some(Ok(record))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::OpenOptions;
use crate::blockstore::test_cid;
use crate::sim::{FaultConfig, SimulatedIO};
use proptest::prelude::*;
use std::path::Path;
fn setup() -> (SimulatedIO, FileId) {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/test");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let fd = sim
.open(Path::new("/test/data.tqb"), OpenOptions::read_write())
.unwrap();
(sim, fd)
}
#[test]
fn write_and_read_single_block() {
let (sim, fd) = setup();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let cid = test_cid(1);
let data = b"hello blockstore";
let location = writer.append_block(&cid, data).unwrap();
writer.sync().unwrap();
assert_eq!(location.file_id, DataFileId::new(0));
assert_eq!(location.offset, BlockOffset::new(BLOCK_HEADER_SIZE as u64));
assert_eq!(location.length, BlockLength::new(data.len() as u32));
let reader = DataFileReader::open(&sim, fd).unwrap();
let blocks = reader.valid_blocks().unwrap();
assert_eq!(blocks.len(), 1);
assert_eq!(blocks[0].1, cid);
assert_eq!(blocks[0].2, data);
}
#[test]
fn write_and_read_multiple_blocks() {
let (sim, fd) = setup();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let payloads: Vec<(&[u8], u8)> = vec![(b"first", 1), (b"second", 2), (b"third", 3)];
let cids: Vec<[u8; CID_SIZE]> = payloads.iter().map(|(_, s)| test_cid(*s)).collect();
payloads
.iter()
.zip(cids.iter())
.for_each(|((data, _), cid)| {
let _ = writer.append_block(cid, data).unwrap();
});
writer.sync().unwrap();
let reader = DataFileReader::open(&sim, fd).unwrap();
let blocks = reader.valid_blocks().unwrap();
assert_eq!(blocks.len(), 3);
assert_eq!(blocks[0].2, b"first");
assert_eq!(blocks[1].2, b"second");
assert_eq!(blocks[2].2, b"third");
assert_eq!(blocks[0].1, cids[0]);
assert_eq!(blocks[1].1, cids[1]);
assert_eq!(blocks[2].1, cids[2]);
}
#[test]
fn empty_file_has_no_blocks() {
let (sim, fd) = setup();
DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let reader = DataFileReader::open(&sim, fd).unwrap();
let blocks = reader.valid_blocks().unwrap();
assert!(blocks.is_empty());
}
#[test]
fn detects_truncated_block() {
let (sim, fd) = setup();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let cid = test_cid(1);
let _ = writer.append_block(&cid, b"complete block").unwrap();
writer.sync().unwrap();
let partial_cid = test_cid(2);
sim.write_all_at(fd, writer.position().raw(), &partial_cid[..10])
.unwrap();
let mut reader = DataFileReader::open(&sim, fd).unwrap();
let first = reader.next().unwrap().unwrap();
assert!(matches!(first, ReadBlockRecord::Valid { .. }));
let second = reader.next().unwrap().unwrap();
assert!(matches!(second, ReadBlockRecord::Truncated { .. }));
}
#[test]
fn checksum_detects_corruption() {
let (sim, fd) = setup();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let cid = test_cid(1);
let data = vec![0xAA; 256];
let _ = writer.append_block(&cid, &data).unwrap();
writer.sync().unwrap();
let corrupt_offset = BLOCK_HEADER_SIZE as u64 + CID_SIZE as u64 + 4 + 128;
sim.write_all_at(fd, corrupt_offset, &[0x00]).unwrap();
let mut reader = DataFileReader::open(&sim, fd).unwrap();
let record = reader.next().unwrap().unwrap();
assert!(matches!(record, ReadBlockRecord::Corrupted { .. }));
}
#[test]
fn crash_before_sync_loses_blocks() {
let (sim, fd) = setup();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let cid1 = test_cid(1);
let _ = writer.append_block(&cid1, b"synced").unwrap();
writer.sync().unwrap();
sim.sync_dir(Path::new("/test")).unwrap();
let cid2 = test_cid(2);
let _ = writer.append_block(&cid2, b"not synced").unwrap();
sim.crash();
let fd = sim
.open(Path::new("/test/data.tqb"), OpenOptions::read())
.unwrap();
let reader = DataFileReader::open(&sim, fd).unwrap();
let blocks = reader.valid_blocks().unwrap();
assert_eq!(blocks.len(), 1);
assert_eq!(blocks[0].2, b"synced");
}
#[test]
fn rejects_oversized_block() {
let (sim, fd) = setup();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let cid = test_cid(1);
let oversized = vec![0u8; MAX_BLOCK_SIZE as usize + 1];
let result = writer.append_block(&cid, &oversized);
assert!(result.is_err());
}
#[test]
fn zero_length_block_round_trips() {
let (sim, fd) = setup();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let cid = test_cid(1);
let location = writer.append_block(&cid, &[]).unwrap();
writer.sync().unwrap();
assert_eq!(location.length, BlockLength::new(0));
let reader = DataFileReader::open(&sim, fd).unwrap();
let blocks = reader.valid_blocks().unwrap();
assert_eq!(blocks.len(), 1);
assert_eq!(blocks[0].1, cid);
assert!(blocks[0].2.is_empty());
}
#[test]
fn accepts_exact_max_block_size() {
let (sim, fd) = setup();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let cid = test_cid(1);
let max_data = vec![0xBB; MAX_BLOCK_SIZE as usize];
let location = writer.append_block(&cid, &max_data).unwrap();
assert_eq!(location.length, BlockLength::new(MAX_BLOCK_SIZE));
}
#[test]
fn bad_magic_rejected() {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/test");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let fd = sim
.open(Path::new("/test/bad.tqb"), OpenOptions::read_write())
.unwrap();
sim.write_all_at(fd, 0, b"NOPE\x01").unwrap();
let result = DataFileReader::open(&sim, fd);
assert!(result.is_err());
}
#[test]
fn encode_decode_round_trip_at_offset() {
let (sim, fd) = setup();
let cid = test_cid(42);
let data = b"round trip test data";
sim.write_all_at(fd, 0, &[0u8; 100]).unwrap();
let offset = BlockOffset::new(100);
let bytes_written = encode_block_record(&sim, fd, offset, &cid, data).unwrap();
let expected_size = BLOCK_RECORD_OVERHEAD as u64 + data.len() as u64;
assert_eq!(bytes_written, expected_size);
let file_size = sim.file_size(fd).unwrap();
let record = decode_block_record(&sim, fd, offset, file_size)
.unwrap()
.unwrap();
match record {
ReadBlockRecord::Valid {
cid_bytes,
data: decoded_data,
..
} => {
assert_eq!(cid_bytes, cid);
assert_eq!(decoded_data, data);
}
other => panic!("expected Valid, got {other:?}"),
}
}
#[test]
fn resume_writer_continues_at_position() {
let (sim, fd) = setup();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let cid1 = test_cid(1);
let _ = writer.append_block(&cid1, b"first").unwrap();
writer.sync().unwrap();
let resume_pos = writer.position();
let mut writer2 = DataFileWriter::resume(&sim, fd, DataFileId::new(0), resume_pos);
let cid2 = test_cid(2);
let _ = writer2.append_block(&cid2, b"second").unwrap();
writer2.sync().unwrap();
let reader = DataFileReader::open(&sim, fd).unwrap();
let blocks = reader.valid_blocks().unwrap();
assert_eq!(blocks.len(), 2);
assert_eq!(blocks[0].2, b"first");
assert_eq!(blocks[1].2, b"second");
}
fn run_crash_recovery_seed(seed: u64) {
let sim = SimulatedIO::new(seed, FaultConfig::aggressive());
let dir = Path::new("/data");
let _ = sim.mkdir(dir);
let _ = sim.sync_dir(dir);
let mut written_blocks: Vec<(u8, Vec<u8>)> = Vec::new();
if let Ok(fd) = sim.open(Path::new("/data/000000.tqb"), OpenOptions::read_write())
&& let Ok(mut writer) = DataFileWriter::new(&sim, fd, DataFileId::new(0))
{
(0u8..20).for_each(|i| {
let cid = test_cid(i);
let data = vec![i; (i as usize + 1) * 10];
if writer.append_block(&cid, &data).is_ok() {
written_blocks.push((i, data));
}
});
let _ = writer.sync();
}
let _ = sim.sync_dir(dir);
sim.crash();
if let Ok(fd) = sim.open(Path::new("/data/000000.tqb"), OpenOptions::read())
&& let Ok(reader) = DataFileReader::open(&sim, fd)
{
let recovered: Vec<_> = reader
.map_while(|r| match r {
Ok(ReadBlockRecord::Valid {
cid_bytes, data, ..
}) => Some((cid_bytes, data)),
_ => None,
})
.collect();
assert!(
recovered.len() <= written_blocks.len(),
"recovered {} blocks but only wrote {}",
recovered.len(),
written_blocks.len()
);
recovered
.iter()
.enumerate()
.for_each(|(i, (cid_bytes, data))| {
assert_eq!(cid_bytes[0], 0x01, "phantom block at index {i}");
assert_eq!(
*cid_bytes,
test_cid(written_blocks[i].0),
"block {i} cid does not match written order"
);
assert_eq!(
*data, written_blocks[i].1,
"block {i} data does not match what was written"
);
});
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(2000))]
#[test]
fn sim_crash_recovery_aggressive_faults(seed in 0u64..u64::MAX) {
run_crash_recovery_seed(seed);
}
}
#[test]
fn sim_partial_write_mid_block_reports_truncated() {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/data");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let fd = sim
.open(Path::new("/data/000000.tqb"), OpenOptions::read_write())
.unwrap();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
(0u8..5).for_each(|i| {
let _ = writer.append_block(&test_cid(i), &[i; 50]).unwrap();
});
writer.sync().unwrap();
let synced_pos = writer.position();
sim.write_all_at(fd, synced_pos.raw(), &test_cid(5)[..20])
.unwrap();
sim.sync(fd).unwrap();
sim.sync_dir(dir).unwrap();
sim.crash();
let fd = sim
.open(Path::new("/data/000000.tqb"), OpenOptions::read())
.unwrap();
let records: Vec<_> = DataFileReader::open(&sim, fd)
.unwrap()
.collect::<Result<Vec<_>, _>>()
.unwrap();
assert_eq!(records.len(), 6);
(0usize..5).for_each(|i| {
assert!(matches!(&records[i], ReadBlockRecord::Valid { .. }));
});
assert!(matches!(&records[5], ReadBlockRecord::Truncated { .. }));
match &records[5] {
ReadBlockRecord::Truncated { offset } => {
assert_eq!(offset.raw(), synced_pos.raw());
}
other => panic!("expected Truncated, got {other:?}"),
}
}
fn run_bit_flip_detection_seed(seed: u64) {
let sim = SimulatedIO::pristine(seed);
let dir = Path::new("/data");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let fd = sim
.open(Path::new("/data/000000.tqb"), OpenOptions::read_write())
.unwrap();
let mut writer = DataFileWriter::new(&sim, fd, DataFileId::new(0)).unwrap();
let data_len = ((seed % 256) as usize).max(1);
let cid = test_cid((seed % 256) as u8);
let data = vec![0xAA; data_len];
let _ = writer.append_block(&cid, &data).unwrap();
writer.sync().unwrap();
let data_start = BLOCK_HEADER_SIZE as u64 + CID_SIZE as u64 + 4;
let flip_pos = data_start + (seed.wrapping_mul(7) % data_len as u64);
let flip_bit = (seed.wrapping_mul(13) % 8) as u8;
let mut byte_buf = [0u8; 1];
sim.read_exact_at(fd, flip_pos, &mut byte_buf).unwrap();
byte_buf[0] ^= 1 << flip_bit;
sim.write_all_at(fd, flip_pos, &byte_buf).unwrap();
let mut reader = DataFileReader::open(&sim, fd).unwrap();
let record = reader.next().unwrap().unwrap();
assert!(matches!(record, ReadBlockRecord::Corrupted { .. }));
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(2000))]
#[test]
fn sim_bit_flip_detected_by_checksum(seed in 0u64..u64::MAX) {
run_bit_flip_detection_seed(seed);
}
}
#[test]
fn sim_rotation_without_dir_sync_loses_new_file() {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/data");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let fd0 = sim
.open(Path::new("/data/000000.tqb"), OpenOptions::read_write())
.unwrap();
let mut writer0 = DataFileWriter::new(&sim, fd0, DataFileId::new(0)).unwrap();
(0u8..3).for_each(|i| {
let _ = writer0.append_block(&test_cid(i), &[i; 50]).unwrap();
});
writer0.sync().unwrap();
sim.sync_dir(dir).unwrap();
let fd1 = sim
.open(Path::new("/data/000001.tqb"), OpenOptions::read_write())
.unwrap();
let mut writer1 = DataFileWriter::new(&sim, fd1, DataFileId::new(1)).unwrap();
let _ = writer1
.append_block(&test_cid(10), b"new file data")
.unwrap();
writer1.sync().unwrap();
sim.crash();
assert!(
sim.open(Path::new("/data/000001.tqb"), OpenOptions::read())
.is_err()
);
let fd0 = sim
.open(Path::new("/data/000000.tqb"), OpenOptions::read())
.unwrap();
let blocks = DataFileReader::open(&sim, fd0)
.unwrap()
.valid_blocks()
.unwrap();
assert_eq!(blocks.len(), 3);
}
}
@@ -0,0 +1,953 @@
use std::cell::Cell;
use std::collections::HashMap;
use std::io;
use std::sync::Arc;
use std::thread;
use crate::io::{FileId, OpenOptions, StorageIO};
use super::data_file::{CID_SIZE, DataFileWriter};
use super::hint::{HintFileWriter, hint_file_path};
use super::key_index::{KeyIndex, KeyIndexError};
use super::manager::DataFileManager;
use super::types::{BlockLocation, BlockOffset, DataFileId, HintOffset, WriteCursor};
#[derive(Debug, Clone)]
pub enum CommitError {
Io(Arc<io::Error>),
Index(Arc<KeyIndexError>),
ChannelClosed,
}
impl std::fmt::Display for CommitError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Io(e) => write!(f, "io: {}", e.as_ref()),
Self::Index(e) => write!(f, "index: {}", e.as_ref()),
Self::ChannelClosed => write!(f, "commit channel closed"),
}
}
}
impl std::error::Error for CommitError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(e) => Some(e.as_ref()),
Self::Index(e) => Some(e.as_ref()),
Self::ChannelClosed => None,
}
}
}
impl From<io::Error> for CommitError {
fn from(e: io::Error) -> Self {
Self::Io(Arc::new(e))
}
}
impl From<KeyIndexError> for CommitError {
fn from(e: KeyIndexError) -> Self {
Self::Index(Arc::new(e))
}
}
type PutResponse = tokio::sync::oneshot::Sender<Result<Vec<BlockLocation>, CommitError>>;
type ApplyResponse = tokio::sync::oneshot::Sender<Result<(), CommitError>>;
pub enum CommitRequest {
PutBlocks {
blocks: Vec<([u8; CID_SIZE], Vec<u8>)>,
response: PutResponse,
},
ApplyCommit {
blocks: Vec<([u8; CID_SIZE], Vec<u8>)>,
deleted_cids: Vec<[u8; CID_SIZE]>,
response: ApplyResponse,
},
Shutdown,
}
#[derive(Debug, Clone)]
pub struct GroupCommitConfig {
pub max_batch_size: usize,
pub channel_capacity: usize,
}
impl Default for GroupCommitConfig {
fn default() -> Self {
Self {
max_batch_size: 1024,
channel_capacity: 4096,
}
}
}
struct ActiveState {
file_id: DataFileId,
fd: FileId,
position: BlockOffset,
hint_fd: FileId,
hint_position: HintOffset,
}
fn log_thread_panic(payload: Box<dyn std::any::Any + Send>, context: &str) {
let msg = payload
.downcast_ref::<&str>()
.copied()
.or_else(|| payload.downcast_ref::<String>().map(|s| s.as_str()))
.unwrap_or("unknown panic");
tracing::error!(panic = msg, "{context}");
}
pub struct GroupCommitWriter {
sender: flume::Sender<CommitRequest>,
handle: Option<thread::JoinHandle<()>>,
}
impl GroupCommitWriter {
pub fn spawn<S: StorageIO + 'static>(
manager: DataFileManager<S>,
index: Arc<KeyIndex>,
config: GroupCommitConfig,
) -> Result<Self, CommitError> {
let cursor = index.read_write_cursor().map_err(CommitError::from)?;
let mut state = initialize_active_state(&manager, cursor)?;
let (sender, receiver) = flume::bounded(config.channel_capacity);
let handle = thread::Builder::new()
.name("blockstore-group-commit".into())
.spawn(move || {
commit_loop(&manager, &*index, &receiver, &config, &mut state);
})
.map_err(|e| CommitError::from(io::Error::other(e)))?;
Ok(Self {
sender,
handle: Some(handle),
})
}
pub fn sender(&self) -> &flume::Sender<CommitRequest> {
&self.sender
}
pub fn shutdown(mut self) {
let _ = self.sender.send(CommitRequest::Shutdown);
if let Some(handle) = self.handle.take()
&& let Err(payload) = handle.join()
{
log_thread_panic(payload, "group commit thread panicked");
}
}
}
impl Drop for GroupCommitWriter {
fn drop(&mut self) {
let _ = self.sender.try_send(CommitRequest::Shutdown);
if let Some(handle) = self.handle.take()
&& let Err(payload) = handle.join()
{
log_thread_panic(payload, "group commit thread panicked during drop");
}
}
}
fn initialize_active_state<S: StorageIO>(
manager: &DataFileManager<S>,
cursor: Option<WriteCursor>,
) -> Result<ActiveState, CommitError> {
let data_dir = manager.data_dir();
let existing_files = manager.list_files()?;
match cursor {
Some(wc) => {
let fd = manager.open_for_append(wc.file_id)?;
let file_size = manager.io().file_size(fd)?;
if file_size < wc.offset.raw() {
return Err(CommitError::from(io::Error::new(
io::ErrorKind::InvalidData,
"data file smaller than write cursor",
)));
}
let hint_path = hint_file_path(data_dir, wc.file_id);
let hint_fd = manager.io().open(&hint_path, OpenOptions::read_write())?;
let hint_size = manager.io().file_size(hint_fd)?;
Ok(ActiveState {
file_id: wc.file_id,
fd,
position: BlockOffset::new(file_size),
hint_fd,
hint_position: HintOffset::new(hint_size),
})
}
None => {
let file_id = existing_files
.last()
.copied()
.map(|id| id.next())
.unwrap_or_else(|| DataFileId::new(0));
let fd = manager.open_for_append(file_id)?;
let writer = DataFileWriter::new(manager.io(), fd, file_id)?;
writer.sync()?;
let position = writer.position();
let hint_path = hint_file_path(data_dir, file_id);
let hint_fd = manager.io().open(&hint_path, OpenOptions::read_write())?;
manager.io().sync_dir(data_dir)?;
Ok(ActiveState {
file_id,
fd,
position,
hint_fd,
hint_position: HintOffset::new(0),
})
}
}
}
enum BatchEntry {
Put {
blocks: Vec<([u8; CID_SIZE], Vec<u8>)>,
response: PutResponse,
},
Apply {
blocks: Vec<([u8; CID_SIZE], Vec<u8>)>,
deleted_cids: Vec<[u8; CID_SIZE]>,
response: ApplyResponse,
},
}
fn classify_request(req: CommitRequest) -> Result<BatchEntry, ()> {
match req {
CommitRequest::PutBlocks { blocks, response } => Ok(BatchEntry::Put { blocks, response }),
CommitRequest::ApplyCommit {
blocks,
deleted_cids,
response,
} => Ok(BatchEntry::Apply {
blocks,
deleted_cids,
response,
}),
CommitRequest::Shutdown => Err(()),
}
}
fn batch_entry_block_count(entry: &BatchEntry) -> usize {
match entry {
BatchEntry::Put { blocks, .. } | BatchEntry::Apply { blocks, .. } => blocks.len(),
}
}
fn drain_batch(
receiver: &flume::Receiver<CommitRequest>,
first: CommitRequest,
max_batch_size: usize,
) -> (Vec<BatchEntry>, bool) {
let first_entry = match classify_request(first) {
Err(()) => return (Vec::new(), true),
Ok(entry) => entry,
};
let block_count = Cell::new(batch_entry_block_count(&first_entry));
let mut entries = vec![first_entry];
let saw_shutdown = std::iter::from_fn(|| receiver.try_recv().ok())
.take_while(|_| block_count.get() < max_batch_size)
.try_for_each(|req| match classify_request(req) {
Err(()) => Err(()),
Ok(entry) => {
block_count.set(
block_count
.get()
.saturating_add(batch_entry_block_count(&entry)),
);
entries.push(entry);
Ok(())
}
})
.is_err();
(entries, saw_shutdown)
}
fn commit_loop<S: StorageIO>(
manager: &DataFileManager<S>,
index: &KeyIndex,
receiver: &flume::Receiver<CommitRequest>,
config: &GroupCommitConfig,
state: &mut ActiveState,
) {
loop {
let first = match receiver.recv() {
Ok(CommitRequest::Shutdown) => return,
Ok(msg) => msg,
Err(_) => return,
};
let (batch, shutdown_after) = drain_batch(receiver, first, config.max_batch_size);
tracing::debug!(
batch_size = batch.len(),
file_id = %state.file_id,
"processing commit batch"
);
let result = process_batch(manager, index, &batch, state);
if let Err(ref e) = result {
tracing::warn!(error = %e, "commit batch failed");
}
dispatch_responses(batch, result);
if shutdown_after {
drain_and_process_remaining(manager, index, receiver, state);
return;
}
}
}
fn drain_and_process_remaining<S: StorageIO>(
manager: &DataFileManager<S>,
index: &KeyIndex,
receiver: &flume::Receiver<CommitRequest>,
state: &mut ActiveState,
) {
let entries: Vec<BatchEntry> = std::iter::from_fn(|| receiver.try_recv().ok())
.filter_map(|req| classify_request(req).ok())
.collect();
if entries.is_empty() {
return;
}
let result = process_batch(manager, index, &entries, state);
dispatch_responses(entries, result);
}
struct RotationState {
file_id: DataFileId,
fd: FileId,
}
fn process_batch<S: StorageIO>(
manager: &DataFileManager<S>,
index: &KeyIndex,
batch: &[BatchEntry],
state: &mut ActiveState,
) -> Result<HashMap<[u8; CID_SIZE], BlockLocation>, CommitError> {
let mut dedup: HashMap<[u8; CID_SIZE], BlockLocation> = HashMap::new();
let mut index_entries: Vec<([u8; CID_SIZE], BlockLocation)> = Vec::new();
let mut all_decrements: Vec<[u8; CID_SIZE]> = Vec::new();
let mut current_hint_fd = state.hint_fd;
let mut rotation: Option<RotationState> = None;
let mut data_writer =
DataFileWriter::resume(manager.io(), state.fd, state.file_id, state.position);
let mut hint_writer =
HintFileWriter::resume(manager.io(), current_hint_fd, state.hint_position);
let write_result: Result<(), CommitError> = batch.iter().try_for_each(|entry| {
let (blocks, decrements) = match entry {
BatchEntry::Put { blocks, .. } => (blocks.as_slice(), None),
BatchEntry::Apply {
blocks,
deleted_cids,
..
} => (blocks.as_slice(), Some(deleted_cids.as_slice())),
};
blocks.iter().try_for_each(|(cid_bytes, data)| {
let location = match dedup.get(cid_bytes) {
Some(&loc) => loc,
None => {
if manager.should_rotate(data_writer.position()) {
data_writer.sync()?;
hint_writer.sync()?;
let (next_id, next_fd) = manager.prepare_rotation(data_writer.file_id())?;
tracing::info!(
from = %data_writer.file_id(),
to = %next_id,
"data file rotation"
);
data_writer = DataFileWriter::new(manager.io(), next_fd, next_id)?;
let new_hint_path = hint_file_path(manager.data_dir(), next_id);
let new_hint_fd = manager
.io()
.open(&new_hint_path, OpenOptions::read_write())?;
manager.io().sync_dir(manager.data_dir())?;
current_hint_fd = new_hint_fd;
hint_writer = HintFileWriter::new(manager.io(), new_hint_fd);
rotation = Some(RotationState {
file_id: next_id,
fd: next_fd,
});
}
let loc = data_writer.append_block(cid_bytes, data)?;
hint_writer.append_hint(cid_bytes, loc.file_id, loc.offset, loc.length)?;
dedup.insert(*cid_bytes, loc);
loc
}
};
index_entries.push((*cid_bytes, location));
Ok::<_, CommitError>(())
})?;
if let Some(decs) = decrements {
all_decrements.extend_from_slice(decs);
}
Ok::<_, CommitError>(())
});
if let Err(e) = write_result {
if let Some(rot) = rotation {
manager.rollback_rotation(rot.file_id, rot.fd);
}
return Err(e);
}
data_writer.sync()?;
hint_writer.sync()?;
if let Some(ref rot) = rotation {
manager.commit_rotation(rot.file_id, rot.fd);
}
state.file_id = data_writer.file_id();
state.fd = data_writer.fd();
state.position = data_writer.position();
state.hint_fd = current_hint_fd;
state.hint_position = hint_writer.position();
let cursor = WriteCursor {
file_id: state.file_id,
offset: state.position,
};
index
.batch_put(&index_entries, &all_decrements, cursor)
.map_err(CommitError::from)?;
Ok(dedup)
}
fn dispatch_responses(
batch: Vec<BatchEntry>,
result: Result<HashMap<[u8; CID_SIZE], BlockLocation>, CommitError>,
) {
match result {
Err(e) => {
batch.into_iter().for_each(|entry| {
let err = e.clone();
match entry {
BatchEntry::Put { response, .. } => {
let _ = response.send(Err(err));
}
BatchEntry::Apply { response, .. } => {
let _ = response.send(Err(err));
}
}
});
}
Ok(written) => {
batch.into_iter().for_each(|entry| match entry {
BatchEntry::Put {
blocks, response, ..
} => {
let result: Result<Vec<BlockLocation>, CommitError> = blocks
.iter()
.map(|(cid, _)| match written.get(cid) {
Some(&loc) => Ok(loc),
None => {
tracing::error!(
?cid,
"committed CID missing from dedup map, this is a bug"
);
Err(CommitError::from(io::Error::other(
"committed CID missing from dedup map",
)))
}
})
.collect();
let _ = response.send(result);
}
BatchEntry::Apply { response, .. } => {
let _ = response.send(Ok(()));
}
});
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::RealIO;
use crate::blockstore::data_file::DataFileReader;
use crate::blockstore::manager::DATA_FILE_EXTENSION;
use crate::blockstore::test_cid;
use futures::StreamExt;
fn setup_real(dir: &std::path::Path) -> (DataFileManager<RealIO>, Arc<KeyIndex>) {
let data_dir = dir.join("data");
std::fs::create_dir_all(&data_dir).unwrap();
let index_dir = dir.join("index");
let manager = DataFileManager::with_default_max_size(RealIO::new(), data_dir);
let index = Arc::new(KeyIndex::open(&index_dir).unwrap().into_inner());
(manager, index)
}
async fn put_blocks(
sender: &flume::Sender<CommitRequest>,
blocks: Vec<([u8; CID_SIZE], Vec<u8>)>,
) -> Result<Vec<BlockLocation>, CommitError> {
let (tx, rx) = tokio::sync::oneshot::channel();
sender
.send_async(CommitRequest::PutBlocks {
blocks,
response: tx,
})
.await
.map_err(|_| CommitError::ChannelClosed)?;
rx.await.map_err(|_| CommitError::ChannelClosed)?
}
async fn apply_commit_req(
sender: &flume::Sender<CommitRequest>,
blocks: Vec<([u8; CID_SIZE], Vec<u8>)>,
deleted_cids: Vec<[u8; CID_SIZE]>,
) -> Result<(), CommitError> {
let (tx, rx) = tokio::sync::oneshot::channel();
sender
.send_async(CommitRequest::ApplyCommit {
blocks,
deleted_cids,
response: tx,
})
.await
.map_err(|_| CommitError::ChannelClosed)?;
rx.await.map_err(|_| CommitError::ChannelClosed)?
}
fn count_data_file_blocks(data_dir: &std::path::Path) -> usize {
let io = RealIO::new();
let data_files =
super::super::list_files_by_extension(&io, data_dir, DATA_FILE_EXTENSION).unwrap();
data_files
.iter()
.map(|&fid| {
let path = data_dir.join(format!("{fid}.tqb"));
let fd = io.open(&path, OpenOptions::read_only_existing()).unwrap();
let count = DataFileReader::open(&io, fd)
.unwrap()
.valid_blocks()
.unwrap()
.len();
let _ = io.close(fd);
count
})
.sum()
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_100_writes_from_10_tasks() {
let dir = tempfile::TempDir::new().unwrap();
let (manager, index) = setup_real(dir.path());
let data_dir = manager.data_dir().to_path_buf();
let writer =
GroupCommitWriter::spawn(manager, index, GroupCommitConfig::default()).unwrap();
let sender = writer.sender().clone();
let handles: Vec<_> = (0u8..10)
.map(|task_id| {
let sender = sender.clone();
tokio::spawn(async move {
let blocks: Vec<_> = (0u8..10)
.map(|block_id| {
let idx = task_id * 10 + block_id;
(test_cid(idx), vec![idx; (idx as usize + 1) * 8])
})
.collect();
futures::stream::iter(blocks)
.fold(
Vec::<BlockLocation>::new(),
|mut acc, (cid, data): ([u8; CID_SIZE], Vec<u8>)| {
let sender = sender.clone();
async move {
let locs =
put_blocks(&sender, vec![(cid, data)]).await.unwrap();
acc.extend(locs);
acc
}
},
)
.await
})
})
.collect();
let all_locations: Vec<Vec<BlockLocation>> = futures::future::join_all(handles)
.await
.into_iter()
.map(|r| r.unwrap())
.collect();
let total: usize = all_locations.iter().map(|v| v.len()).sum();
assert_eq!(total, 100);
writer.shutdown();
let index_dir = dir.path().join("index");
let index = KeyIndex::open(&index_dir).unwrap().into_inner();
(0u8..100).for_each(|i| {
assert!(
index.has(&test_cid(i)).unwrap(),
"block {i} missing from index"
);
});
assert_eq!(count_data_file_blocks(&data_dir), 100);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn duplicate_cids_in_same_batch_write_once() {
let dir = tempfile::TempDir::new().unwrap();
let (manager, index) = setup_real(dir.path());
let data_dir = manager.data_dir().to_path_buf();
let writer =
GroupCommitWriter::spawn(manager, index, GroupCommitConfig::default()).unwrap();
let sender = writer.sender().clone();
let cid = test_cid(42);
let data = vec![0xAB; 128];
let blocks = vec![
(cid, data.clone()),
(cid, data.clone()),
(cid, data.clone()),
];
let locations = put_blocks(&sender, blocks).await.unwrap();
assert_eq!(locations.len(), 3);
assert_eq!(locations[0], locations[1]);
assert_eq!(locations[1], locations[2]);
writer.shutdown();
let index_dir = dir.path().join("index");
let index = KeyIndex::open(&index_dir).unwrap().into_inner();
let entry = index.get(&cid).unwrap().unwrap();
assert_eq!(entry.refcount.raw(), 3);
assert_eq!(
count_data_file_blocks(&data_dir),
1,
"duplicate CID should only be written once to data file"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn apply_commit_with_blocks_and_deletes() {
let dir = tempfile::TempDir::new().unwrap();
let (manager, index) = setup_real(dir.path());
let writer =
GroupCommitWriter::spawn(manager, index, GroupCommitConfig::default()).unwrap();
let sender = writer.sender().clone();
let cid_a = test_cid(1);
let cid_b = test_cid(2);
put_blocks(
&sender,
vec![(cid_a, vec![0x01; 64]), (cid_b, vec![0x02; 64])],
)
.await
.unwrap();
let cid_c = test_cid(3);
apply_commit_req(&sender, vec![(cid_c, vec![0x03; 64])], vec![cid_a])
.await
.unwrap();
writer.shutdown();
let index_dir = dir.path().join("index");
let index = KeyIndex::open(&index_dir).unwrap().into_inner();
assert_eq!(index.get(&cid_a).unwrap().unwrap().refcount.raw(), 0);
assert_eq!(index.get(&cid_b).unwrap().unwrap().refcount.raw(), 1);
assert_eq!(index.get(&cid_c).unwrap().unwrap().refcount.raw(), 1);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn graceful_shutdown_processes_remaining() {
let dir = tempfile::TempDir::new().unwrap();
let (manager, index) = setup_real(dir.path());
let writer =
GroupCommitWriter::spawn(manager, index, GroupCommitConfig::default()).unwrap();
let sender = writer.sender().clone();
let cid = test_cid(99);
let locations = put_blocks(&sender, vec![(cid, vec![0xFF; 32])])
.await
.unwrap();
assert_eq!(locations.len(), 1);
writer.shutdown();
let index_dir = dir.path().join("index");
let index = KeyIndex::open(&index_dir).unwrap().into_inner();
assert!(index.has(&cid).unwrap());
}
#[test]
fn sim_crash_between_write_and_fsync_loses_unsynced() {
use crate::SimulatedIO;
use std::path::Path;
use std::sync::Arc;
let sim = Arc::new(SimulatedIO::pristine(42));
let data_dir = Path::new("/data");
sim.mkdir(data_dir).unwrap();
sim.sync_dir(data_dir).unwrap();
let manager =
DataFileManager::with_default_max_size(Arc::clone(&sim), data_dir.to_path_buf());
let fd = manager.open_for_append(DataFileId::new(0)).unwrap();
let mut writer = DataFileWriter::new(&*sim, fd, DataFileId::new(0)).unwrap();
let synced_cids: Vec<_> = (0u8..5)
.map(|i| {
let cid = test_cid(i);
let _ = writer.append_block(&cid, &vec![i; 64]).unwrap();
cid
})
.collect();
writer.sync().unwrap();
sim.sync_dir(data_dir).unwrap();
(5u8..10).for_each(|i| {
let cid = test_cid(i);
let _ = writer.append_block(&cid, &vec![i; 64]).unwrap();
});
sim.crash();
let fd_after = sim
.open(Path::new("/data/000000.tqb"), OpenOptions::read())
.unwrap();
let recovered = DataFileReader::open(&*sim, fd_after)
.unwrap()
.valid_blocks()
.unwrap();
assert!(
recovered.len() <= 5,
"expected at most 5 synced blocks, got {}",
recovered.len()
);
recovered.iter().enumerate().for_each(|(i, (_, cid, _))| {
assert_eq!(*cid, synced_cids[i], "recovered block {i} CID mismatch");
});
}
#[test]
fn sim_crash_between_fsync_and_index_update_recovers_via_hints() {
use crate::SimulatedIO;
use crate::blockstore::data_file::{BLOCK_HEADER_SIZE, BLOCK_RECORD_OVERHEAD};
use crate::blockstore::hint::rebuild_index_from_hints;
use crate::blockstore::types::BlockLength;
use std::path::Path;
use std::sync::Arc;
let sim = Arc::new(SimulatedIO::pristine(42));
let data_dir = Path::new("/data");
sim.mkdir(data_dir).unwrap();
sim.sync_dir(data_dir).unwrap();
let manager =
DataFileManager::with_default_max_size(Arc::clone(&sim), data_dir.to_path_buf());
let fd = manager.open_for_append(DataFileId::new(0)).unwrap();
let mut writer = DataFileWriter::new(&*sim, fd, DataFileId::new(0)).unwrap();
let phase1_cids: Vec<_> = (0u8..3)
.map(|i| {
let cid = test_cid(i);
let _ = writer.append_block(&cid, &vec![i; 64]).unwrap();
cid
})
.collect();
writer.sync().unwrap();
let phase1_end = writer.position();
let real_dir = tempfile::TempDir::new().unwrap();
let index_path = real_dir.path().join("index");
let index = KeyIndex::open(&index_path).unwrap().into_inner();
let entries: Vec<_> = phase1_cids
.iter()
.enumerate()
.map(|(i, cid)| {
let offset = BlockOffset::new(
BLOCK_HEADER_SIZE as u64 + i as u64 * (BLOCK_RECORD_OVERHEAD as u64 + 64),
);
(
*cid,
BlockLocation {
file_id: DataFileId::new(0),
offset,
length: BlockLength::new(64),
},
)
})
.collect();
index
.batch_put(
&entries,
&[],
WriteCursor {
file_id: DataFileId::new(0),
offset: phase1_end,
},
)
.unwrap();
index.persist().unwrap();
let phase2_cids: Vec<_> = (10u8..15)
.map(|i| {
let cid = test_cid(i);
let _ = writer.append_block(&cid, &vec![i; 128]).unwrap();
cid
})
.collect();
writer.sync().unwrap();
sim.sync_dir(data_dir).unwrap();
let hint_path = hint_file_path(data_dir, DataFileId::new(0));
let hint_fd = sim.open(&hint_path, OpenOptions::read_write()).unwrap();
let mut hint_writer = HintFileWriter::new(&*sim, hint_fd);
let mut offset_tracker = BlockOffset::new(BLOCK_HEADER_SIZE as u64);
phase1_cids.iter().for_each(|cid| {
hint_writer
.append_hint(
cid,
DataFileId::new(0),
offset_tracker,
BlockLength::new(64),
)
.unwrap();
offset_tracker = offset_tracker.advance(BLOCK_RECORD_OVERHEAD as u64 + 64);
});
phase2_cids.iter().for_each(|cid| {
hint_writer
.append_hint(
cid,
DataFileId::new(0),
offset_tracker,
BlockLength::new(128),
)
.unwrap();
offset_tracker = offset_tracker.advance(BLOCK_RECORD_OVERHEAD as u64 + 128);
});
hint_writer.sync().unwrap();
sim.sync_dir(data_dir).unwrap();
sim.crash();
drop(index);
let rebuilt_index_path = real_dir.path().join("rebuilt_index");
let rebuilt_index = KeyIndex::open(&rebuilt_index_path).unwrap().into_inner();
rebuild_index_from_hints(&*sim, data_dir, &rebuilt_index).unwrap();
phase1_cids.iter().for_each(|cid| {
assert!(
rebuilt_index.has(cid).unwrap(),
"phase1 CID should be in rebuilt index"
);
});
phase2_cids.iter().for_each(|cid| {
assert!(
rebuilt_index.has(cid).unwrap(),
"phase2 CID should be in rebuilt index, was synced and hinted before crash"
);
});
let cursor = rebuilt_index.read_write_cursor().unwrap().unwrap();
assert!(
cursor.offset.raw() > phase1_end.raw(),
"cursor should be past phase1 after rebuild"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn rotation_during_batch() {
let dir = tempfile::TempDir::new().unwrap();
let data_dir = dir.path().join("data");
std::fs::create_dir_all(&data_dir).unwrap();
let index_dir = dir.path().join("index");
let small_max = 512u64;
let manager = DataFileManager::new(RealIO::new(), data_dir.clone(), small_max);
let index = Arc::new(KeyIndex::open(&index_dir).unwrap().into_inner());
let writer =
GroupCommitWriter::spawn(manager, index, GroupCommitConfig::default()).unwrap();
let sender = writer.sender().clone();
let all_cids: Vec<_> = (0u8..20).map(test_cid).collect();
let handles: Vec<_> = all_cids
.iter()
.map(|&cid| {
let sender = sender.clone();
tokio::spawn(async move {
put_blocks(&sender, vec![(cid, vec![cid[4]; 100])])
.await
.unwrap()
})
})
.collect();
let results: Vec<_> = futures::future::join_all(handles)
.await
.into_iter()
.map(|r| r.unwrap())
.collect();
assert_eq!(results.len(), 20);
writer.shutdown();
let io = RealIO::new();
let data_files =
super::super::list_files_by_extension(&io, &data_dir, DATA_FILE_EXTENSION).unwrap();
assert!(
data_files.len() > 1,
"expected rotation to create multiple files, got {}",
data_files.len()
);
let index = KeyIndex::open(&index_dir).unwrap().into_inner();
all_cids.iter().for_each(|cid| {
assert!(index.has(cid).unwrap());
});
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,539 @@
use std::collections::HashMap;
use std::path::Path;
use fjall::{
Database, Keyspace, KeyspaceCreateOptions, PersistMode,
config::{BloomConstructionPolicy, FilterPolicy, FilterPolicyEntry},
};
use super::data_file::CID_SIZE;
use super::types::{BlockLocation, IndexEntry, RefCount, WriteCursor};
const WRITE_CURSOR_KEY: &[u8] = b"\x00write_cursor";
const KEYSPACE_NAME: &str = "blocks";
fn bloom_options() -> KeyspaceCreateOptions {
KeyspaceCreateOptions::default().filter_policy(FilterPolicy::new([
FilterPolicyEntry::Bloom(BloomConstructionPolicy::FalsePositiveRate(0.01)),
FilterPolicyEntry::Bloom(BloomConstructionPolicy::FalsePositiveRate(0.01)),
]))
}
fn is_corruption_error(e: &fjall::Error) -> bool {
match e {
fjall::Error::Io(io_err) => matches!(
io_err.kind(),
std::io::ErrorKind::InvalidData | std::io::ErrorKind::UnexpectedEof
),
fjall::Error::Locked | fjall::Error::KeyspaceDeleted => false,
_ => true,
}
}
fn serialize_entry(entry: &IndexEntry) -> Vec<u8> {
postcard::to_allocvec(entry)
.expect("IndexEntry serialization is infallible for fixed-layout types")
}
fn deserialize_entry(bytes: &[u8]) -> Result<IndexEntry, KeyIndexError> {
postcard::from_bytes(bytes).map_err(KeyIndexError::Deserialize)
}
fn serialize_cursor(cursor: &WriteCursor) -> Vec<u8> {
postcard::to_allocvec(cursor)
.expect("WriteCursor serialization is infallible for fixed-layout types")
}
fn deserialize_cursor(bytes: &[u8]) -> Result<WriteCursor, KeyIndexError> {
postcard::from_bytes(bytes).map_err(KeyIndexError::Deserialize)
}
#[derive(Debug)]
pub enum KeyIndexError {
Fjall(fjall::Error),
Deserialize(postcard::Error),
MissingEntry,
}
impl std::fmt::Display for KeyIndexError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Fjall(e) => write!(f, "fjall: {e}"),
Self::Deserialize(e) => write!(f, "deserialize: {e}"),
Self::MissingEntry => write!(f, "entry not found"),
}
}
}
impl std::error::Error for KeyIndexError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Fjall(e) => Some(e),
Self::Deserialize(e) => Some(e),
Self::MissingEntry => None,
}
}
}
impl From<fjall::Error> for KeyIndexError {
fn from(e: fjall::Error) -> Self {
Self::Fjall(e)
}
}
pub enum KeyIndexOpenOutcome {
Opened(KeyIndex),
NeedsRebuild(KeyIndex),
}
impl KeyIndexOpenOutcome {
pub fn into_inner(self) -> KeyIndex {
match self {
Self::Opened(idx) | Self::NeedsRebuild(idx) => idx,
}
}
pub fn needs_rebuild(&self) -> bool {
matches!(self, Self::NeedsRebuild(_))
}
}
pub struct KeyIndex {
db: Database,
blocks: Keyspace,
}
impl KeyIndex {
pub fn open(path: &Path) -> Result<KeyIndexOpenOutcome, KeyIndexError> {
match Self::try_open(path) {
Ok(idx) => Ok(KeyIndexOpenOutcome::Opened(idx)),
Err(KeyIndexError::Fjall(ref e)) if is_corruption_error(e) => {
let _ = std::fs::remove_dir_all(path);
let idx = Self::try_open(path)?;
Ok(KeyIndexOpenOutcome::NeedsRebuild(idx))
}
Err(e) => Err(e),
}
}
fn try_open(path: &Path) -> Result<Self, KeyIndexError> {
let db = Database::builder(path).open()?;
let blocks = db.keyspace(KEYSPACE_NAME, bloom_options)?;
Ok(Self { db, blocks })
}
pub fn get(&self, cid_bytes: &[u8; CID_SIZE]) -> Result<Option<IndexEntry>, KeyIndexError> {
self.blocks
.get(cid_bytes)?
.map(|v| deserialize_entry(&v))
.transpose()
}
pub fn has(&self, cid_bytes: &[u8; CID_SIZE]) -> Result<bool, KeyIndexError> {
self.blocks.contains_key(cid_bytes).map_err(Into::into)
}
pub fn put(
&self,
cid_bytes: &[u8; CID_SIZE],
location: BlockLocation,
) -> Result<(), KeyIndexError> {
let entry = match self.get(cid_bytes)? {
Some(existing) => IndexEntry {
location: existing.location,
refcount: existing.refcount.increment(),
},
None => IndexEntry {
location,
refcount: RefCount::one(),
},
};
self.blocks
.insert(cid_bytes, serialize_entry(&entry))
.map_err(Into::into)
}
pub fn decrement_refcount(
&self,
cid_bytes: &[u8; CID_SIZE],
) -> Result<RefCount, KeyIndexError> {
let existing = self.get(cid_bytes)?.ok_or(KeyIndexError::MissingEntry)?;
let new_refcount = match existing.refcount.is_zero() {
true => {
tracing::warn!(?cid_bytes, "decrement on zero-refcount entry, skipping");
existing.refcount
}
false => existing.refcount.decrement(),
};
let updated = IndexEntry {
location: existing.location,
refcount: new_refcount,
};
self.blocks.insert(cid_bytes, serialize_entry(&updated))?;
Ok(new_refcount)
}
pub fn batch_put(
&self,
entries: &[([u8; CID_SIZE], BlockLocation)],
decrements: &[[u8; CID_SIZE]],
cursor: WriteCursor,
) -> Result<(), KeyIndexError> {
let mut batch = self.db.batch().durability(Some(PersistMode::SyncData));
let mut pending: HashMap<[u8; CID_SIZE], IndexEntry> = HashMap::new();
entries.iter().try_for_each(|(cid_bytes, location)| {
let entry = match pending.get(cid_bytes).copied().or(self.get(cid_bytes)?) {
Some(existing) => IndexEntry {
location: existing.location,
refcount: existing.refcount.increment(),
},
None => IndexEntry {
location: *location,
refcount: RefCount::one(),
},
};
pending.insert(*cid_bytes, entry);
batch.insert(&self.blocks, cid_bytes.as_slice(), serialize_entry(&entry));
Ok::<_, KeyIndexError>(())
})?;
decrements.iter().try_for_each(|cid_bytes| {
let existing = pending
.get(cid_bytes)
.copied()
.or(self.get(cid_bytes)?)
.ok_or(KeyIndexError::MissingEntry)?;
let new_refcount = match existing.refcount.is_zero() {
true => {
tracing::warn!(?cid_bytes, "decrement on zero-refcount entry, skipping");
existing.refcount
}
false => existing.refcount.decrement(),
};
let updated = IndexEntry {
location: existing.location,
refcount: new_refcount,
};
pending.insert(*cid_bytes, updated);
batch.insert(
&self.blocks,
cid_bytes.as_slice(),
serialize_entry(&updated),
);
Ok::<_, KeyIndexError>(())
})?;
batch.insert(&self.blocks, WRITE_CURSOR_KEY, serialize_cursor(&cursor));
batch.commit().map_err(Into::into)
}
pub fn read_write_cursor(&self) -> Result<Option<WriteCursor>, KeyIndexError> {
self.blocks
.get(WRITE_CURSOR_KEY)?
.map(|v| deserialize_cursor(&v))
.transpose()
}
pub fn persist(&self) -> Result<(), KeyIndexError> {
self.db.persist(PersistMode::SyncData).map_err(Into::into)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::blockstore::test_cid;
use crate::blockstore::types::{BlockLength, BlockOffset, DataFileId};
fn test_location(file_id: u32, offset: u64, length: u32) -> BlockLocation {
BlockLocation {
file_id: DataFileId::new(file_id),
offset: BlockOffset::new(offset),
length: BlockLength::new(length),
}
}
fn open_temp() -> (tempfile::TempDir, KeyIndex) {
let dir = tempfile::TempDir::new().unwrap();
let outcome = KeyIndex::open(dir.path()).unwrap();
assert!(!outcome.needs_rebuild());
(dir, outcome.into_inner())
}
#[test]
fn put_then_get_round_trips() {
let (_dir, idx) = open_temp();
let cid = test_cid(1);
let loc = test_location(0, 100, 256);
idx.put(&cid, loc).unwrap();
let entry = idx.get(&cid).unwrap().unwrap();
assert_eq!(entry.location, loc);
assert_eq!(entry.refcount, RefCount::one());
}
#[test]
fn get_missing_returns_none() {
let (_dir, idx) = open_temp();
assert!(idx.get(&test_cid(42)).unwrap().is_none());
}
#[test]
fn has_missing_returns_false() {
let (_dir, idx) = open_temp();
assert!(!idx.has(&test_cid(42)).unwrap());
}
#[test]
fn has_existing_returns_true() {
let (_dir, idx) = open_temp();
let cid = test_cid(1);
idx.put(&cid, test_location(0, 0, 10)).unwrap();
assert!(idx.has(&cid).unwrap());
}
#[test]
fn duplicate_put_increments_refcount() {
let (_dir, idx) = open_temp();
let cid = test_cid(1);
let loc = test_location(0, 100, 256);
idx.put(&cid, loc).unwrap();
idx.put(&cid, test_location(1, 200, 512)).unwrap();
let entry = idx.get(&cid).unwrap().unwrap();
assert_eq!(entry.refcount, RefCount::new(2));
assert_eq!(entry.location, loc);
}
#[test]
fn decrement_refcount_from_two_to_one() {
let (_dir, idx) = open_temp();
let cid = test_cid(1);
idx.put(&cid, test_location(0, 0, 10)).unwrap();
idx.put(&cid, test_location(0, 0, 10)).unwrap();
let rc = idx.decrement_refcount(&cid).unwrap();
assert_eq!(rc, RefCount::one());
let entry = idx.get(&cid).unwrap().unwrap();
assert_eq!(entry.refcount, RefCount::one());
}
#[test]
fn decrement_refcount_to_zero_keeps_entry() {
let (_dir, idx) = open_temp();
let cid = test_cid(1);
idx.put(&cid, test_location(0, 0, 10)).unwrap();
let rc = idx.decrement_refcount(&cid).unwrap();
assert!(rc.is_zero());
let entry = idx.get(&cid).unwrap().unwrap();
assert!(entry.refcount.is_zero());
}
#[test]
fn decrement_missing_entry_errors() {
let (_dir, idx) = open_temp();
let result = idx.decrement_refcount(&test_cid(99));
assert!(matches!(result, Err(KeyIndexError::MissingEntry)));
}
#[test]
fn batch_put_new_entries() {
let (_dir, idx) = open_temp();
let entries: Vec<_> = (0u8..3)
.map(|i| (test_cid(i), test_location(0, i as u64 * 100, 50)))
.collect();
let cursor = WriteCursor {
file_id: DataFileId::new(0),
offset: BlockOffset::new(300),
};
idx.batch_put(&entries, &[], cursor).unwrap();
entries.iter().for_each(|(cid, loc)| {
let entry = idx.get(cid).unwrap().unwrap();
assert_eq!(entry.location, *loc);
assert_eq!(entry.refcount, RefCount::one());
});
}
#[test]
fn batch_put_increments_existing() {
let (_dir, idx) = open_temp();
let cid = test_cid(1);
let original_loc = test_location(0, 100, 50);
idx.put(&cid, original_loc).unwrap();
let entries = vec![(cid, test_location(1, 200, 60))];
let cursor = WriteCursor {
file_id: DataFileId::new(1),
offset: BlockOffset::new(260),
};
idx.batch_put(&entries, &[], cursor).unwrap();
let entry = idx.get(&cid).unwrap().unwrap();
assert_eq!(entry.refcount, RefCount::new(2));
assert_eq!(entry.location, original_loc);
}
#[test]
fn batch_put_with_decrements() {
let (_dir, idx) = open_temp();
let cid_a = test_cid(1);
let cid_b = test_cid(2);
idx.put(&cid_b, test_location(0, 0, 10)).unwrap();
idx.put(&cid_b, test_location(0, 0, 10)).unwrap();
let entries = vec![(cid_a, test_location(0, 100, 50))];
let decrements = vec![cid_b];
let cursor = WriteCursor {
file_id: DataFileId::new(0),
offset: BlockOffset::new(150),
};
idx.batch_put(&entries, &decrements, cursor).unwrap();
let a = idx.get(&cid_a).unwrap().unwrap();
assert_eq!(a.refcount, RefCount::one());
let b = idx.get(&cid_b).unwrap().unwrap();
assert_eq!(b.refcount, RefCount::one());
}
#[test]
fn batch_put_mixed_new_and_duplicate() {
let (_dir, idx) = open_temp();
let existing_cid = test_cid(1);
let existing_loc = test_location(0, 0, 10);
idx.put(&existing_cid, existing_loc).unwrap();
let entries: Vec<_> = (1u8..=4)
.map(|i| (test_cid(i), test_location(0, i as u64 * 100, 50)))
.collect();
let cursor = WriteCursor {
file_id: DataFileId::new(0),
offset: BlockOffset::new(500),
};
idx.batch_put(&entries, &[], cursor).unwrap();
let existing = idx.get(&existing_cid).unwrap().unwrap();
assert_eq!(existing.refcount, RefCount::new(2));
assert_eq!(existing.location, existing_loc);
(2u8..=4).for_each(|i| {
let entry = idx.get(&test_cid(i)).unwrap().unwrap();
assert_eq!(entry.refcount, RefCount::one());
});
}
#[test]
fn batch_put_duplicate_cid_in_same_batch() {
let (_dir, idx) = open_temp();
let cid = test_cid(1);
let loc = test_location(0, 100, 50);
let entries = vec![(cid, loc), (cid, test_location(0, 200, 60))];
let cursor = WriteCursor {
file_id: DataFileId::new(0),
offset: BlockOffset::new(260),
};
idx.batch_put(&entries, &[], cursor).unwrap();
let entry = idx.get(&cid).unwrap().unwrap();
assert_eq!(entry.refcount, RefCount::new(2));
assert_eq!(entry.location, loc);
}
#[test]
fn batch_put_entry_then_decrement_same_cid() {
let (_dir, idx) = open_temp();
let cid = test_cid(1);
let loc = test_location(0, 100, 50);
let entries = vec![(cid, loc)];
let decrements = vec![cid];
let cursor = WriteCursor {
file_id: DataFileId::new(0),
offset: BlockOffset::new(150),
};
idx.batch_put(&entries, &decrements, cursor).unwrap();
let entry = idx.get(&cid).unwrap().unwrap();
assert!(entry.refcount.is_zero());
}
#[test]
fn write_cursor_round_trip() {
let (_dir, idx) = open_temp();
assert!(idx.read_write_cursor().unwrap().is_none());
let cursor = WriteCursor {
file_id: DataFileId::new(3),
offset: BlockOffset::new(65536),
};
let entries = vec![(test_cid(1), test_location(3, 0, 100))];
idx.batch_put(&entries, &[], cursor).unwrap();
let read_back = idx.read_write_cursor().unwrap().unwrap();
assert_eq!(read_back, cursor);
}
#[test]
fn write_cursor_persists_across_reopen() {
let dir = tempfile::TempDir::new().unwrap();
let cursor = WriteCursor {
file_id: DataFileId::new(7),
offset: BlockOffset::new(99999),
};
{
let idx = KeyIndex::open(dir.path()).unwrap().into_inner();
let entries = vec![(test_cid(1), test_location(7, 0, 100))];
idx.batch_put(&entries, &[], cursor).unwrap();
idx.persist().unwrap();
}
{
let idx = KeyIndex::open(dir.path()).unwrap().into_inner();
let read_back = idx.read_write_cursor().unwrap().unwrap();
assert_eq!(read_back, cursor);
let entry = idx.get(&test_cid(1)).unwrap().unwrap();
assert_eq!(entry.refcount, RefCount::one());
}
}
#[test]
fn corrupt_index_triggers_needs_rebuild() {
let dir = tempfile::TempDir::new().unwrap();
{
let idx = KeyIndex::open(dir.path()).unwrap().into_inner();
idx.put(&test_cid(1), test_location(0, 0, 10)).unwrap();
idx.persist().unwrap();
}
std::fs::read_dir(dir.path())
.unwrap()
.filter_map(|e| e.ok())
.for_each(|entry| {
let path = entry.path();
if path.is_file() {
std::fs::write(&path, b"corrupted").unwrap();
}
});
let outcome = KeyIndex::open(dir.path()).unwrap();
assert!(outcome.needs_rebuild());
let idx = outcome.into_inner();
assert!(idx.get(&test_cid(1)).unwrap().is_none());
assert!(idx.read_write_cursor().unwrap().is_none());
}
}
@@ -0,0 +1,344 @@
use std::collections::HashMap;
use std::io;
use std::path::{Path, PathBuf};
use parking_lot::RwLock;
use crate::io::{FileId, OpenOptions, StorageIO};
use super::list_files_by_extension;
use super::types::{BlockOffset, DataFileId};
pub const DEFAULT_MAX_FILE_SIZE: u64 = 256 * 1024 * 1024;
pub(crate) const DATA_FILE_EXTENSION: &str = "tqb";
struct CachedHandle {
fd: FileId,
writable: bool,
}
pub struct DataFileManager<S: StorageIO> {
io: S,
data_dir: PathBuf,
max_file_size: u64,
handles: RwLock<HashMap<DataFileId, CachedHandle>>,
}
impl<S: StorageIO> DataFileManager<S> {
pub fn new(io: S, data_dir: PathBuf, max_file_size: u64) -> Self {
Self {
io,
data_dir,
max_file_size,
handles: RwLock::new(HashMap::new()),
}
}
pub fn with_default_max_size(io: S, data_dir: PathBuf) -> Self {
Self::new(io, data_dir, DEFAULT_MAX_FILE_SIZE)
}
pub fn io(&self) -> &S {
&self.io
}
pub fn data_dir(&self) -> &Path {
&self.data_dir
}
pub fn max_file_size(&self) -> u64 {
self.max_file_size
}
pub fn data_file_path(&self, file_id: DataFileId) -> PathBuf {
self.data_dir
.join(format!("{file_id}.{DATA_FILE_EXTENSION}"))
}
pub fn open_for_append(&self, file_id: DataFileId) -> io::Result<FileId> {
{
let cache = self.handles.read();
if let Some(entry) = cache.get(&file_id)
&& entry.writable
{
return Ok(entry.fd);
}
}
let path = self.data_file_path(file_id);
let fd = self.io.open(&path, OpenOptions::read_write())?;
let mut cache = self.handles.write();
match cache.get(&file_id) {
Some(entry) if entry.writable => {
let _ = self.io.close(fd);
Ok(entry.fd)
}
Some(entry) => {
let old_fd = entry.fd;
cache.insert(file_id, CachedHandle { fd, writable: true });
let _ = self.io.close(old_fd);
Ok(fd)
}
None => {
cache.insert(file_id, CachedHandle { fd, writable: true });
Ok(fd)
}
}
}
pub fn open_for_read(&self, file_id: DataFileId) -> io::Result<FileId> {
if let Some(entry) = self.handles.read().get(&file_id) {
return Ok(entry.fd);
}
let path = self.data_file_path(file_id);
let fd = self.io.open(&path, OpenOptions::read_only_existing())?;
let mut cache = self.handles.write();
match cache.get(&file_id) {
Some(entry) => {
let _ = self.io.close(fd);
Ok(entry.fd)
}
None => {
cache.insert(
file_id,
CachedHandle {
fd,
writable: false,
},
);
Ok(fd)
}
}
}
pub fn prepare_rotation(&self, current: DataFileId) -> io::Result<(DataFileId, FileId)> {
let next = current.next();
let path = self.data_file_path(next);
let fd = self.io.open(&path, OpenOptions::read_write())?;
Ok((next, fd))
}
pub fn commit_rotation(&self, file_id: DataFileId, fd: FileId) {
self.handles
.write()
.insert(file_id, CachedHandle { fd, writable: true });
}
pub fn rollback_rotation(&self, file_id: DataFileId, fd: FileId) {
let _ = self.io.close(fd);
self.handles.write().remove(&file_id);
}
pub fn should_rotate(&self, position: BlockOffset) -> bool {
position.raw() >= self.max_file_size
}
pub fn list_files(&self) -> io::Result<Vec<DataFileId>> {
list_files_by_extension(&self.io, &self.data_dir, DATA_FILE_EXTENSION)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::blockstore::data_file::{BLOCK_HEADER_SIZE, DataFileReader, DataFileWriter};
use crate::blockstore::test_cid;
use crate::sim::SimulatedIO;
fn setup_manager(max_file_size: u64) -> DataFileManager<SimulatedIO> {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/data");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
DataFileManager::new(sim, dir.to_path_buf(), max_file_size)
}
#[test]
fn open_for_append_creates_file() {
let mgr = setup_manager(1024);
let fd = mgr.open_for_append(DataFileId::new(0)).unwrap();
assert_eq!(mgr.io().file_size(fd).unwrap(), 0);
}
#[test]
fn open_for_read_missing_file_errors() {
let mgr = setup_manager(1024);
assert!(mgr.open_for_read(DataFileId::new(99)).is_err());
}
#[test]
fn handle_cache_returns_same_fd() {
let mgr = setup_manager(1024);
let fd1 = mgr.open_for_append(DataFileId::new(0)).unwrap();
let fd2 = mgr.open_for_append(DataFileId::new(0)).unwrap();
assert_eq!(fd1, fd2);
}
#[test]
fn open_for_read_uses_cache_from_append() {
let mgr = setup_manager(1024);
let fd_write = mgr.open_for_append(DataFileId::new(0)).unwrap();
let fd_read = mgr.open_for_read(DataFileId::new(0)).unwrap();
assert_eq!(fd_write, fd_read);
}
#[test]
fn rotation_lifecycle_prepare_commit() {
let mgr = setup_manager(1024);
let _fd0 = mgr.open_for_append(DataFileId::new(0)).unwrap();
let (next_id, next_fd) = mgr.prepare_rotation(DataFileId::new(0)).unwrap();
assert_eq!(next_id, DataFileId::new(1));
assert_eq!(mgr.io().file_size(next_fd).unwrap(), 0);
mgr.io().sync_dir(mgr.data_dir()).unwrap();
mgr.commit_rotation(next_id, next_fd);
assert_eq!(mgr.open_for_read(next_id).unwrap(), next_fd);
}
#[test]
fn rotation_rollback_cleans_handle() {
let mgr = setup_manager(1024);
let _fd0 = mgr.open_for_append(DataFileId::new(0)).unwrap();
let (next_id, next_fd) = mgr.prepare_rotation(DataFileId::new(0)).unwrap();
mgr.commit_rotation(next_id, next_fd);
assert_eq!(mgr.open_for_read(next_id).unwrap(), next_fd);
mgr.rollback_rotation(next_id, next_fd);
let reopened_fd = mgr.open_for_read(next_id).unwrap();
assert_ne!(
reopened_fd, next_fd,
"rollback should have closed the cached fd"
);
}
#[test]
fn should_rotate_respects_threshold() {
let mgr = setup_manager(1024);
assert!(!mgr.should_rotate(BlockOffset::new(100)));
assert!(!mgr.should_rotate(BlockOffset::new(1023)));
assert!(mgr.should_rotate(BlockOffset::new(1024)));
assert!(mgr.should_rotate(BlockOffset::new(2000)));
}
#[test]
fn list_files_finds_data_files() {
let mgr = setup_manager(1024);
let _fd0 = mgr.open_for_append(DataFileId::new(0)).unwrap();
let _fd3 = mgr.open_for_append(DataFileId::new(3)).unwrap();
let files = mgr.list_files().unwrap();
assert_eq!(files, vec![DataFileId::new(0), DataFileId::new(3)]);
}
#[test]
fn list_files_ignores_non_data_files() {
let mgr = setup_manager(1024);
let _fd0 = mgr.open_for_append(DataFileId::new(0)).unwrap();
mgr.io()
.open(Path::new("/data/notes.txt"), OpenOptions::read_write())
.unwrap();
let files = mgr.list_files().unwrap();
assert_eq!(files, vec![DataFileId::new(0)]);
}
#[test]
fn data_file_path_format() {
let mgr = setup_manager(1024);
assert_eq!(
mgr.data_file_path(DataFileId::new(0)),
Path::new("/data/000000.tqb")
);
assert_eq!(
mgr.data_file_path(DataFileId::new(42)),
Path::new("/data/000042.tqb")
);
}
#[test]
fn rotate_and_write_across_files() {
let mgr = setup_manager(1024);
let fd0 = mgr.open_for_append(DataFileId::new(0)).unwrap();
let mut writer0 = DataFileWriter::new(mgr.io(), fd0, DataFileId::new(0)).unwrap();
let _ = writer0
.append_block(&test_cid(1), b"first file data")
.unwrap();
writer0.sync().unwrap();
let (id1, fd1) = mgr.prepare_rotation(DataFileId::new(0)).unwrap();
mgr.io().sync_dir(mgr.data_dir()).unwrap();
mgr.commit_rotation(id1, fd1);
let mut writer1 = DataFileWriter::new(mgr.io(), fd1, id1).unwrap();
let _ = writer1
.append_block(&test_cid(2), b"second file data")
.unwrap();
writer1.sync().unwrap();
let fd0_read = mgr.open_for_read(DataFileId::new(0)).unwrap();
let blocks0 = DataFileReader::open(mgr.io(), fd0_read)
.unwrap()
.valid_blocks()
.unwrap();
assert_eq!(blocks0.len(), 1);
assert_eq!(blocks0[0].2, b"first file data");
let fd1_read = mgr.open_for_read(id1).unwrap();
let blocks1 = DataFileReader::open(mgr.io(), fd1_read)
.unwrap()
.valid_blocks()
.unwrap();
assert_eq!(blocks1.len(), 1);
assert_eq!(blocks1[0].2, b"second file data");
}
#[test]
fn read_cache_hit_from_writable_entry() {
let mgr = setup_manager(1024);
let fd_write = mgr.open_for_append(DataFileId::new(0)).unwrap();
DataFileWriter::new(mgr.io(), fd_write, DataFileId::new(0)).unwrap();
let fd_read = mgr.open_for_read(DataFileId::new(0)).unwrap();
assert_eq!(fd_write, fd_read);
}
#[test]
fn read_only_cache_upgraded_on_append() {
let mgr = setup_manager(1024);
let raw_fd = mgr
.io()
.open(
&mgr.data_file_path(DataFileId::new(0)),
OpenOptions::read_write(),
)
.unwrap();
DataFileWriter::new(mgr.io(), raw_fd, DataFileId::new(0)).unwrap();
mgr.io().sync(raw_fd).unwrap();
mgr.io().sync_dir(mgr.data_dir()).unwrap();
mgr.io().close(raw_fd).unwrap();
let fd_read = mgr.open_for_read(DataFileId::new(0)).unwrap();
let _reader = DataFileReader::open(mgr.io(), fd_read).unwrap();
let fd_append = mgr.open_for_append(DataFileId::new(0)).unwrap();
assert_ne!(fd_read, fd_append);
let mut writer = DataFileWriter::resume(
mgr.io(),
fd_append,
DataFileId::new(0),
BlockOffset::new(BLOCK_HEADER_SIZE as u64),
);
let _ = writer
.append_block(&test_cid(1), b"written after upgrade")
.unwrap();
writer.sync().unwrap();
let blocks = DataFileReader::open(mgr.io(), fd_append)
.unwrap()
.valid_blocks()
.unwrap();
assert_eq!(blocks.len(), 1);
assert_eq!(blocks[0].2, b"written after upgrade");
}
}
@@ -0,0 +1,68 @@
mod data_file;
mod group_commit;
mod hint;
mod key_index;
mod manager;
mod reader;
mod store;
mod types;
pub use data_file::{
BLOCK_FORMAT_VERSION, BLOCK_HEADER_SIZE, BLOCK_MAGIC, BLOCK_RECORD_OVERHEAD, CID_SIZE,
DataFileReader, DataFileWriter, ReadBlockRecord, ValidBlock, decode_block_record,
encode_block_record,
};
pub use group_commit::{CommitError, CommitRequest, GroupCommitConfig, GroupCommitWriter};
pub use hint::{
HINT_FILE_EXTENSION, HINT_RECORD_SIZE, HintFileReader, HintFileWriter, ReadHintRecord,
RebuildError, decode_hint_record, hint_file_path, rebuild_index_from_data_files,
rebuild_index_from_hints,
};
pub use key_index::{KeyIndex, KeyIndexError, KeyIndexOpenOutcome};
pub use manager::{DEFAULT_MAX_FILE_SIZE, DataFileManager};
pub use reader::{BlockStoreReader, ReadError};
pub use store::{BlockStoreConfig, TranquilBlockStore};
pub use types::{
BlockLength, BlockLocation, BlockOffset, DataFileId, HintOffset, IndexEntry, MAX_BLOCK_SIZE,
RefCount, WriteCursor,
};
use std::io;
use std::path::Path;
use crate::io::StorageIO;
pub(crate) fn list_files_by_extension<S: StorageIO>(
io: &S,
dir: &Path,
extension: &str,
) -> io::Result<Vec<DataFileId>> {
let entries = io.list_dir(dir)?;
let mut ids: Vec<DataFileId> = entries
.iter()
.filter_map(|path| {
let stem = path.file_stem()?.to_str()?;
let ext = path.extension()?.to_str()?;
(ext == extension).then(|| stem.parse::<u32>().ok().map(DataFileId::new))?
})
.collect();
ids.sort();
Ok(ids)
}
#[cfg(test)]
pub(crate) fn test_cid(seed: u8) -> [u8; CID_SIZE] {
test_cid_u16(seed as u16)
}
#[cfg(test)]
pub(crate) fn test_cid_u16(seed: u16) -> [u8; CID_SIZE] {
let mut cid = [0u8; CID_SIZE];
cid[0] = 0x01;
cid[1] = 0x71;
cid[2] = 0x12;
cid[3] = 0x20;
cid[4..6].copy_from_slice(&seed.to_le_bytes());
(6..CID_SIZE).for_each(|i| cid[i] = (seed as u8).wrapping_add(i as u8));
cid
}
@@ -0,0 +1,611 @@
use std::collections::HashMap;
use std::io;
use std::sync::Arc;
use bytes::Bytes;
use crate::io::{FileId, StorageIO};
use super::data_file::{CID_SIZE, ReadBlockRecord, decode_block_record};
use super::key_index::{KeyIndex, KeyIndexError};
use super::manager::DataFileManager;
use super::types::{BlockLocation, BlockOffset, DataFileId};
#[derive(Debug, Clone)]
pub enum ReadError {
Io(Arc<io::Error>),
Index(Arc<KeyIndexError>),
Corrupted {
file_id: DataFileId,
offset: BlockOffset,
},
}
impl std::fmt::Display for ReadError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Io(e) => write!(f, "io: {e}"),
Self::Index(e) => write!(f, "index: {e}"),
Self::Corrupted { file_id, offset } => {
write!(f, "corrupted block at {file_id}:{}", offset.raw())
}
}
}
}
impl std::error::Error for ReadError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(e) => Some(e.as_ref()),
Self::Index(e) => Some(e.as_ref()),
Self::Corrupted { .. } => None,
}
}
}
impl From<io::Error> for ReadError {
fn from(e: io::Error) -> Self {
Self::Io(Arc::new(e))
}
}
impl From<KeyIndexError> for ReadError {
fn from(e: KeyIndexError) -> Self {
Self::Index(Arc::new(e))
}
}
pub struct BlockStoreReader<S: StorageIO> {
index: Arc<KeyIndex>,
manager: Arc<DataFileManager<S>>,
}
impl<S: StorageIO> Clone for BlockStoreReader<S> {
fn clone(&self) -> Self {
Self {
index: Arc::clone(&self.index),
manager: Arc::clone(&self.manager),
}
}
}
impl<S: StorageIO> BlockStoreReader<S> {
pub fn new(index: Arc<KeyIndex>, manager: Arc<DataFileManager<S>>) -> Self {
Self { index, manager }
}
pub fn get(&self, cid: &[u8; CID_SIZE]) -> Result<Option<Bytes>, ReadError> {
let entry = match self.index.get(cid)? {
Some(e) => e,
None => return Ok(None),
};
self.read_block_at(entry.location).map(Some)
}
pub fn has(&self, cid: &[u8; CID_SIZE]) -> Result<bool, ReadError> {
self.index.has(cid).map_err(ReadError::from)
}
pub fn get_many(&self, cids: &[[u8; CID_SIZE]]) -> Result<Vec<Option<Bytes>>, ReadError> {
let mut results: Vec<Option<Bytes>> = vec![None; cids.len()];
let lookups: Vec<(usize, BlockLocation)> = cids
.iter()
.enumerate()
.filter_map(|(i, cid)| match self.index.get(cid) {
Ok(Some(entry)) => Some(Ok((i, entry.location))),
Ok(None) => None,
Err(e) => Some(Err(ReadError::from(e))),
})
.collect::<Result<Vec<_>, _>>()?;
let mut by_file: HashMap<DataFileId, Vec<(usize, BlockLocation)>> = HashMap::new();
lookups.into_iter().for_each(|(idx, loc)| {
by_file.entry(loc.file_id).or_default().push((idx, loc));
});
by_file.into_iter().try_for_each(|(file_id, mut entries)| {
let fd = self.manager.open_for_read(file_id)?;
let file_size = self.manager.io().file_size(fd)?;
entries.sort_by_key(|(_, loc)| loc.offset);
entries.into_iter().try_for_each(|(orig_idx, loc)| {
let data = self.decode_and_validate(fd, file_size, loc)?;
results[orig_idx] = Some(data);
Ok::<_, ReadError>(())
})
})?;
Ok(results)
}
fn read_block_at(&self, location: BlockLocation) -> Result<Bytes, ReadError> {
let fd = self.manager.open_for_read(location.file_id)?;
let file_size = self.manager.io().file_size(fd)?;
self.decode_and_validate(fd, file_size, location)
}
fn decode_and_validate(
&self,
fd: FileId,
file_size: u64,
location: BlockLocation,
) -> Result<Bytes, ReadError> {
match decode_block_record(self.manager.io(), fd, location.offset, file_size)? {
Some(ReadBlockRecord::Valid { data, .. })
if data.len() == location.length.raw() as usize =>
{
Ok(Bytes::from(data))
}
Some(ReadBlockRecord::Valid { .. }) => Err(ReadError::Corrupted {
file_id: location.file_id,
offset: location.offset,
}),
Some(ReadBlockRecord::Corrupted { offset } | ReadBlockRecord::Truncated { offset }) => {
Err(ReadError::Corrupted {
file_id: location.file_id,
offset,
})
}
None => Err(ReadError::Corrupted {
file_id: location.file_id,
offset: location.offset,
}),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::RealIO;
use crate::blockstore::data_file::CID_SIZE;
use crate::blockstore::group_commit::{CommitRequest, GroupCommitConfig, GroupCommitWriter};
use crate::blockstore::key_index::KeyIndex;
use crate::blockstore::manager::DataFileManager;
use crate::blockstore::test_cid;
use futures::StreamExt;
struct TestHarness {
_dir: tempfile::TempDir,
index: Arc<KeyIndex>,
manager: Arc<DataFileManager<RealIO>>,
writer: Option<GroupCommitWriter>,
sender: flume::Sender<CommitRequest>,
}
impl TestHarness {
fn new() -> Self {
let dir = tempfile::TempDir::new().unwrap();
let data_dir = dir.path().join("data");
std::fs::create_dir_all(&data_dir).unwrap();
let index_dir = dir.path().join("index");
let manager = Arc::new(DataFileManager::with_default_max_size(
RealIO::new(),
data_dir,
));
let index = Arc::new(KeyIndex::open(&index_dir).unwrap().into_inner());
let writer = GroupCommitWriter::spawn(
DataFileManager::with_default_max_size(RealIO::new(), dir.path().join("data")),
Arc::clone(&index),
GroupCommitConfig::default(),
)
.unwrap();
let sender = writer.sender().clone();
Self {
_dir: dir,
index,
manager,
writer: Some(writer),
sender,
}
}
fn reader(&self) -> BlockStoreReader<RealIO> {
BlockStoreReader::new(Arc::clone(&self.index), Arc::clone(&self.manager))
}
async fn put_blocks(
&self,
blocks: Vec<([u8; CID_SIZE], Vec<u8>)>,
) -> Result<Vec<super::super::types::BlockLocation>, super::super::group_commit::CommitError>
{
let (tx, rx) = tokio::sync::oneshot::channel();
self.sender
.send_async(CommitRequest::PutBlocks {
blocks,
response: tx,
})
.await
.map_err(|_| super::super::group_commit::CommitError::ChannelClosed)?;
rx.await
.map_err(|_| super::super::group_commit::CommitError::ChannelClosed)?
}
fn shutdown(&mut self) {
if let Some(w) = self.writer.take() {
w.shutdown();
}
}
}
impl Drop for TestHarness {
fn drop(&mut self) {
self.shutdown();
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_existing_block() {
let mut harness = TestHarness::new();
let cid = test_cid(1);
let data = vec![0xAB; 256];
harness.put_blocks(vec![(cid, data.clone())]).await.unwrap();
harness.shutdown();
let reader = harness.reader();
let result = reader.get(&cid).unwrap().unwrap();
assert_eq!(&result[..], &data[..]);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_missing_block_returns_none() {
let mut harness = TestHarness::new();
harness.shutdown();
let reader = harness.reader();
assert!(reader.get(&test_cid(99)).unwrap().is_none());
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_many_mixed_hits_and_misses() {
let mut harness = TestHarness::new();
let blocks: Vec<_> = (0u8..5)
.map(|i| (test_cid(i), vec![i; (i as usize + 1) * 32]))
.collect();
harness.put_blocks(blocks.clone()).await.unwrap();
harness.shutdown();
let reader = harness.reader();
let query: Vec<[u8; CID_SIZE]> = vec![
test_cid(0),
test_cid(99),
test_cid(2),
test_cid(100),
test_cid(4),
];
let results = reader.get_many(&query).unwrap();
assert_eq!(results.len(), 5);
assert_eq!(&results[0].as_ref().unwrap()[..], &blocks[0].1[..]);
assert!(results[1].is_none());
assert_eq!(&results[2].as_ref().unwrap()[..], &blocks[2].1[..]);
assert!(results[3].is_none());
assert_eq!(&results[4].as_ref().unwrap()[..], &blocks[4].1[..]);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn has_returns_true_for_existing() {
let mut harness = TestHarness::new();
let cid = test_cid(1);
harness
.put_blocks(vec![(cid, vec![0xFF; 64])])
.await
.unwrap();
harness.shutdown();
let reader = harness.reader();
assert!(reader.has(&cid).unwrap());
assert!(!reader.has(&test_cid(99)).unwrap());
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn checksum_mismatch_returns_error() {
let mut harness = TestHarness::new();
let cid = test_cid(1);
let data = vec![0xAA; 256];
harness.put_blocks(vec![(cid, data)]).await.unwrap();
harness.shutdown();
let entry = harness.index.get(&cid).unwrap().unwrap();
let loc = entry.location;
let data_file_path = harness.manager.data_file_path(loc.file_id);
let corrupt_offset = loc.offset.raw() + super::super::data_file::CID_SIZE as u64 + 4 + 128;
let file_bytes = std::fs::read(&data_file_path).unwrap();
let mut corrupted = file_bytes;
corrupted[corrupt_offset as usize] ^= 0xFF;
std::fs::write(&data_file_path, &corrupted).unwrap();
let fresh_manager = Arc::new(DataFileManager::with_default_max_size(
RealIO::new(),
harness.manager.data_dir().to_path_buf(),
));
let reader = BlockStoreReader::new(Arc::clone(&harness.index), fresh_manager);
let result = reader.get(&cid);
assert!(
matches!(result, Err(ReadError::Corrupted { .. })),
"expected Corrupted error, got {result:?}"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn spawn_blocking_does_not_block_tokio_workers() {
let mut harness = TestHarness::new();
let blocks: Vec<_> = (0u8..200).map(|i| (test_cid(i), vec![i; 1024])).collect();
harness.put_blocks(blocks).await.unwrap();
harness.shutdown();
let reader = harness.reader();
let reader = Arc::new(reader);
let timer_handle = tokio::spawn(futures::stream::iter(0..100).fold(
std::time::Duration::ZERO,
|max_drift, _| async move {
let start = std::time::Instant::now();
tokio::time::sleep(std::time::Duration::from_millis(1)).await;
let drift = start
.elapsed()
.saturating_sub(std::time::Duration::from_millis(1));
max_drift.max(drift)
},
));
let read_handles: Vec<_> = (0..8)
.map(|_| {
let reader = Arc::clone(&reader);
tokio::spawn(futures::stream::iter(0u8..200).fold(
(0u64, 200u64),
move |(total_us, count), i| {
let reader = Arc::clone(&reader);
async move {
let cid = test_cid(i);
let start = std::time::Instant::now();
let result = tokio::task::spawn_blocking(move || reader.get(&cid))
.await
.unwrap();
let elapsed_us = start.elapsed().as_micros() as u64;
assert!(result.unwrap().is_some());
(total_us.saturating_add(elapsed_us), count)
}
},
))
})
.collect();
let timer_drift = timer_handle.await.unwrap();
assert!(
timer_drift < std::time::Duration::from_millis(5),
"timer drift {timer_drift:?} exceeds 5ms, reads may be blocking tokio workers"
);
let stats: Vec<(u64, u64)> = futures::future::join_all(read_handles)
.await
.into_iter()
.map(|r| r.unwrap())
.collect();
let total_us: u64 = stats.iter().map(|(us, _)| us).sum();
let total_count: u64 = stats.iter().map(|(_, c)| c).sum();
let avg_us = total_us / total_count.max(1);
eprintln!("avg read latency: {avg_us}us across {total_count} reads");
}
use crate::blockstore::test_cid_u16 as stress_cid;
#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
async fn stress_50_writers_20_readers() {
let dir = tempfile::TempDir::new().unwrap();
let data_dir = dir.path().join("data");
std::fs::create_dir_all(&data_dir).unwrap();
let index_dir = dir.path().join("index");
let index = Arc::new(KeyIndex::open(&index_dir).unwrap().into_inner());
let manager_for_writer =
DataFileManager::with_default_max_size(RealIO::new(), data_dir.clone());
let writer = GroupCommitWriter::spawn(
manager_for_writer,
Arc::clone(&index),
GroupCommitConfig::default(),
)
.unwrap();
let sender = writer.sender().clone();
let manager_for_reader = Arc::new(DataFileManager::with_default_max_size(
RealIO::new(),
data_dir,
));
let reader = BlockStoreReader::new(Arc::clone(&index), manager_for_reader);
let committed = Arc::new(std::sync::Mutex::new(Vec::<(u16, Vec<u8>)>::new()));
let writer_done = Arc::new(std::sync::atomic::AtomicBool::new(false));
let writer_handles: Vec<_> = (0u16..50)
.map(|writer_id| {
let sender = sender.clone();
let committed = Arc::clone(&committed);
tokio::spawn(async move {
futures::stream::iter(0u16..200)
.fold((), |(), block_id| {
let sender = sender.clone();
let committed = Arc::clone(&committed);
async move {
let seed = writer_id * 200 + block_id;
let cid = stress_cid(seed);
let size = ((seed as usize % 256) + 1) * 4;
let data = vec![seed as u8; size];
let (tx, rx) = tokio::sync::oneshot::channel();
sender
.send_async(CommitRequest::PutBlocks {
blocks: vec![(cid, data.clone())],
response: tx,
})
.await
.unwrap();
rx.await.unwrap().unwrap();
committed.lock().unwrap().push((seed, data));
}
})
.await;
})
})
.collect();
let reader_handles: Vec<_> = (0..20)
.map(|_| {
let reader = reader.clone();
let committed = Arc::clone(&committed);
let done = Arc::clone(&writer_done);
tokio::spawn(async move {
let reads = std::sync::atomic::AtomicU64::new(0);
(0..5000)
.take_while(|_| {
let is_done = done.load(std::sync::atomic::Ordering::Relaxed);
let has_reads = reads.load(std::sync::atomic::Ordering::Relaxed) > 100;
!(is_done && has_reads)
})
.for_each(|_| {
let snapshot = committed.lock().unwrap().clone();
if let Some((seed, expected)) = snapshot.last() {
let cid = stress_cid(*seed);
match reader.get(&cid) {
Ok(Some(actual)) => {
assert_eq!(&actual[..], &expected[..]);
reads.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
Ok(None) => {}
Err(e) => panic!("read error: {e}"),
}
}
std::thread::yield_now();
});
reads.load(std::sync::atomic::Ordering::Relaxed)
})
})
.collect();
futures::future::join_all(writer_handles)
.await
.into_iter()
.for_each(|r| r.unwrap());
writer_done.store(true, std::sync::atomic::Ordering::Relaxed);
let read_counts: Vec<u64> = futures::future::join_all(reader_handles)
.await
.into_iter()
.map(|r| r.unwrap())
.collect();
let total_reads: u64 = read_counts.iter().sum();
eprintln!("total reader reads: {total_reads}");
assert!(total_reads > 0);
writer.shutdown();
let final_committed = committed.lock().unwrap();
assert_eq!(final_committed.len(), 10_000);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_read_write() {
let mut harness = TestHarness::new();
let sender = harness.sender.clone();
let reader = harness.reader();
let written_cids = Arc::new(std::sync::Mutex::new(Vec::<(u8, Vec<u8>)>::new()));
let writer_done = Arc::new(std::sync::atomic::AtomicBool::new(false));
let writer_handle = {
let written = Arc::clone(&written_cids);
tokio::spawn(async move {
futures::stream::iter(0u8..50)
.fold((), |(), i| {
let sender = sender.clone();
let written = Arc::clone(&written);
async move {
let cid = test_cid(i);
let data = vec![i; (i as usize + 1) * 16];
let (tx, rx) = tokio::sync::oneshot::channel();
sender
.send_async(CommitRequest::PutBlocks {
blocks: vec![(cid, data.clone())],
response: tx,
})
.await
.unwrap();
rx.await.unwrap().unwrap();
written.lock().unwrap().push((i, data));
}
})
.await;
})
};
let reader_handles: Vec<_> = (0..4)
.map(|_| {
let reader = reader.clone();
let written = Arc::clone(&written_cids);
let done = Arc::clone(&writer_done);
tokio::spawn(async move {
let reads = std::sync::atomic::AtomicU64::new(0);
(0..2000)
.take_while(|_| {
let is_done = done.load(std::sync::atomic::Ordering::Relaxed);
let has_reads = reads.load(std::sync::atomic::Ordering::Relaxed) > 0;
!(is_done && has_reads)
})
.for_each(|_| {
let snapshot = written.lock().unwrap().clone();
snapshot.iter().for_each(|(seed, expected_data)| {
let cid = test_cid(*seed);
match reader.get(&cid) {
Ok(Some(actual)) => {
assert_eq!(
&actual[..],
&expected_data[..],
"data mismatch for block {seed}"
);
reads.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
}
Ok(None) => {}
Err(e) => panic!("read error for block {seed}: {e}"),
}
});
std::thread::yield_now();
});
reads.load(std::sync::atomic::Ordering::Relaxed)
})
})
.collect();
writer_handle.await.unwrap();
writer_done.store(true, std::sync::atomic::Ordering::Relaxed);
let read_counts: Vec<u64> = futures::future::join_all(reader_handles)
.await
.into_iter()
.map(|r| r.unwrap())
.collect();
let total_reads: u64 = read_counts.iter().sum();
assert!(
total_reads > 0,
"readers should have completed at least some reads"
);
let final_snapshot = written_cids.lock().unwrap().clone();
assert_eq!(final_snapshot.len(), 50);
final_snapshot.iter().for_each(|(seed, expected_data)| {
let cid = test_cid(*seed);
let actual = reader.get(&cid).unwrap().unwrap();
assert_eq!(
&actual[..],
&expected_data[..],
"final verification failed for block {seed}"
);
});
harness.shutdown();
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,206 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct DataFileId(u32);
impl DataFileId {
pub fn new(id: u32) -> Self {
Self(id)
}
pub fn raw(self) -> u32 {
self.0
}
pub fn next(self) -> Self {
Self(self.0.checked_add(1).expect("DataFileId overflow"))
}
}
impl std::fmt::Display for DataFileId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{:06}", self.0)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct BlockOffset(u64);
impl BlockOffset {
pub fn new(offset: u64) -> Self {
Self(offset)
}
pub fn raw(self) -> u64 {
self.0
}
pub fn advance(self, delta: u64) -> Self {
Self(self.0.checked_add(delta).expect("BlockOffset overflow"))
}
}
pub const MAX_BLOCK_SIZE: u32 = 4 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct BlockLength(u32);
impl BlockLength {
pub fn new(length: u32) -> Self {
assert!(
length <= MAX_BLOCK_SIZE,
"BlockLength {length} exceeds MAX_BLOCK_SIZE {MAX_BLOCK_SIZE}"
);
Self(length)
}
pub fn raw(self) -> u32 {
self.0
}
pub fn as_u64(self) -> u64 {
u64::from(self.0)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct RefCount(u32);
impl RefCount {
pub fn new(count: u32) -> Self {
Self(count)
}
pub fn raw(self) -> u32 {
self.0
}
pub fn one() -> Self {
Self(1)
}
pub fn is_zero(self) -> bool {
self.0 == 0
}
pub fn increment(self) -> Self {
Self(self.0.checked_add(1).expect("RefCount overflow"))
}
pub fn decrement(self) -> Self {
Self(self.0.saturating_sub(1))
}
}
#[must_use]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct BlockLocation {
pub file_id: DataFileId,
pub offset: BlockOffset,
pub length: BlockLength,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct IndexEntry {
pub location: BlockLocation,
pub refcount: RefCount,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct WriteCursor {
pub file_id: DataFileId,
pub offset: BlockOffset,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct HintOffset(u64);
impl HintOffset {
pub fn new(offset: u64) -> Self {
Self(offset)
}
pub fn raw(self) -> u64 {
self.0
}
pub fn advance(self, delta: u64) -> Self {
Self(self.0.checked_add(delta).expect("HintOffset overflow"))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn index_entry_postcard_round_trip() {
let entry = IndexEntry {
location: BlockLocation {
file_id: DataFileId::new(42),
offset: BlockOffset::new(1024),
length: BlockLength::new(256),
},
refcount: RefCount::one(),
};
let bytes = postcard::to_allocvec(&entry).unwrap();
let decoded: IndexEntry = postcard::from_bytes(&bytes).unwrap();
assert_eq!(entry, decoded);
}
#[test]
fn write_cursor_postcard_round_trip() {
let cursor = WriteCursor {
file_id: DataFileId::new(7),
offset: BlockOffset::new(65536),
};
let bytes = postcard::to_allocvec(&cursor).unwrap();
let decoded: WriteCursor = postcard::from_bytes(&bytes).unwrap();
assert_eq!(cursor, decoded);
}
#[test]
fn data_file_id_display_zero_padded() {
assert_eq!(DataFileId::new(0).to_string(), "000000");
assert_eq!(DataFileId::new(42).to_string(), "000042");
assert_eq!(DataFileId::new(999999).to_string(), "999999");
}
#[test]
fn data_file_id_next_increments() {
assert_eq!(DataFileId::new(0).next(), DataFileId::new(1));
assert_eq!(DataFileId::new(99).next(), DataFileId::new(100));
}
#[test]
#[should_panic(expected = "DataFileId overflow")]
fn data_file_id_overflow_panics() {
DataFileId::new(u32::MAX).next();
}
#[test]
fn block_offset_advance() {
let offset = BlockOffset::new(100);
assert_eq!(offset.advance(50), BlockOffset::new(150));
}
#[test]
fn refcount_lifecycle() {
let rc = RefCount::one();
assert!(!rc.is_zero());
assert_eq!(rc.raw(), 1);
let rc2 = rc.increment();
assert_eq!(rc2.raw(), 2);
let rc3 = rc2.decrement().decrement();
assert!(rc3.is_zero());
}
#[test]
fn refcount_underflow_saturates_at_zero() {
assert!(RefCount::new(0).decrement().is_zero());
}
}
+273
View File
@@ -0,0 +1,273 @@
use std::io;
use std::path::Path;
use crate::io::{OpenOptions, StorageIO};
use crate::record::{RecordReader, RecordWriter};
use crate::sim::{FaultConfig, SimulatedIO};
fn setup_sim_file(sim: &SimulatedIO, name: &str) -> io::Result<(crate::io::FileId, String)> {
let dir = Path::new("/harness");
sim.mkdir(dir)?;
sim.sync_dir(dir)?;
let path_str = format!("/harness/{name}");
let path = Path::new(&path_str);
let fd = sim.open(path, OpenOptions::read_write())?;
sim.sync_dir(dir)?;
Ok((fd, path_str))
}
fn reopen_after_crash(sim: &SimulatedIO, path: &str) -> io::Result<crate::io::FileId> {
sim.open(Path::new(path), OpenOptions::read())
}
pub struct CrashTestResult {
pub seed: u64,
pub records_written: usize,
pub records_synced: usize,
pub records_recovered: usize,
pub corrupted_detected: usize,
pub truncated_detected: usize,
}
pub fn run_crash_test(
seed: u64,
fault_config: FaultConfig,
payloads: &[Vec<u8>],
sync_after: usize,
) -> io::Result<CrashTestResult> {
let sim = SimulatedIO::new(seed, fault_config);
let (fd, path) = setup_sim_file(&sim, "crash_test.dat")?;
let mut writer = RecordWriter::new(&sim, fd)?;
let mut records_written = 0usize;
let mut records_synced = 0usize;
let _stop_reason = payloads
.iter()
.enumerate()
.try_fold((), |(), (i, payload)| {
writer.append(payload)?;
records_written += 1;
if sync_after > 0
&& (i + 1) % sync_after == 0
&& writer.sync().is_ok()
&& sim.last_sync_persisted()
{
records_synced = records_written;
}
Ok::<_, io::Error>(())
});
sim.crash();
let recovery_fd = match reopen_after_crash(&sim, &path) {
Ok(fd) => fd,
Err(_) => {
return Ok(CrashTestResult {
seed,
records_written,
records_synced,
records_recovered: 0,
corrupted_detected: 0,
truncated_detected: 0,
});
}
};
let reader = match RecordReader::open(&sim, recovery_fd) {
Ok(r) => r,
Err(_) => {
return Ok(CrashTestResult {
seed,
records_written,
records_synced,
records_recovered: 0,
corrupted_detected: 0,
truncated_detected: 0,
});
}
};
use crate::record::ReadRecord;
let collected: Vec<_> = reader
.scan(false, |stopped, record| {
if *stopped {
return None;
}
match record {
ReadRecord::Valid { .. } => Some(record),
other => {
*stopped = true;
Some(other)
}
}
})
.collect();
let records_recovered = collected
.iter()
.filter(|r| matches!(r, ReadRecord::Valid { .. }))
.count();
let corrupted_detected = collected.last().map_or(0, |r| {
usize::from(matches!(r, ReadRecord::Corrupted { .. }))
});
let truncated_detected = collected.last().map_or(0, |r| {
usize::from(matches!(r, ReadRecord::Truncated { .. }))
});
Ok(CrashTestResult {
seed,
records_written,
records_synced,
records_recovered,
corrupted_detected,
truncated_detected,
})
}
pub fn run_pristine_comparison(
seed: u64,
fault_config: FaultConfig,
payloads: &[Vec<u8>],
sync_after: usize,
) -> io::Result<PristineComparisonResult> {
let pristine = SimulatedIO::pristine(seed);
let (pristine_fd, _) = setup_sim_file(&pristine, "pristine.dat")?;
let mut pristine_writer = RecordWriter::new(&pristine, pristine_fd)?;
let synced_payloads = payloads
.iter()
.enumerate()
.fold(
(Vec::<Vec<u8>>::new(), Vec::<Vec<u8>>::new()),
|(mut synced, mut pending), (i, payload)| {
pristine_writer.append(payload).unwrap();
pending.push(payload.clone());
if sync_after > 0 && (i + 1) % sync_after == 0 {
pristine_writer.sync().unwrap();
synced.append(&mut pending);
}
(synced, pending)
},
)
.0;
let faulty = SimulatedIO::new(seed, fault_config);
let (faulty_fd, faulty_path) = setup_sim_file(&faulty, "faulty.dat")?;
let mut faulty_writer = RecordWriter::new(&faulty, faulty_fd)?;
payloads
.iter()
.enumerate()
.try_fold((), |(), (i, payload)| {
faulty_writer.append(payload)?;
if sync_after > 0 && (i + 1) % sync_after == 0 {
let _ = faulty_writer.sync();
}
Ok::<_, io::Error>(())
})
.ok();
faulty.crash();
let recovered = match reopen_after_crash(&faulty, &faulty_path)
.and_then(|fd| RecordReader::open(&faulty, fd))
{
Ok(reader) => reader.valid_records(),
Err(_) => Vec::new(),
};
let prefix_valid = recovered
.iter()
.zip(synced_payloads.iter())
.all(|(r, p)| r == p);
let recovery_within_bounds = recovered.len() <= payloads.len();
Ok(PristineComparisonResult {
seed,
synced_count: synced_payloads.len(),
recovered_count: recovered.len(),
prefix_matches_pristine: prefix_valid,
recovery_within_bounds,
})
}
pub struct PristineComparisonResult {
pub seed: u64,
pub synced_count: usize,
pub recovered_count: usize,
pub prefix_matches_pristine: bool,
pub recovery_within_bounds: bool,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn no_fault_recovers_all_synced() {
let payloads: Vec<Vec<u8>> = (0..10)
.map(|i| format!("record {i}").into_bytes())
.collect();
let result = run_crash_test(42, FaultConfig::none(), &payloads, 5).unwrap();
assert_eq!(result.records_synced, 10);
assert_eq!(result.records_recovered, 10);
assert_eq!(result.corrupted_detected, 0);
}
#[test]
fn no_fault_unsynced_records_lost() {
let payloads: Vec<Vec<u8>> = (0..10)
.map(|i| format!("record {i}").into_bytes())
.collect();
let result = run_crash_test(42, FaultConfig::none(), &payloads, 0).unwrap();
assert_eq!(result.records_recovered, 0);
}
#[test]
fn pristine_comparison_no_faults() {
let payloads: Vec<Vec<u8>> = (0..20)
.map(|i| format!("payload {i}").into_bytes())
.collect();
let result = run_pristine_comparison(42, FaultConfig::none(), &payloads, 5).unwrap();
assert!(result.prefix_matches_pristine);
assert!(result.recovery_within_bounds);
assert_eq!(result.synced_count, 20);
assert_eq!(result.recovered_count, 20);
}
#[test]
fn faulted_recovery_never_exceeds_written() {
(0..1000).for_each(|seed| {
let payloads: Vec<Vec<u8>> = (0..5).map(|i| format!("data-{i}").into_bytes()).collect();
let Ok(result) = run_crash_test(seed, FaultConfig::moderate(), &payloads, 2) else {
return;
};
assert!(
result.records_recovered <= result.records_written,
"seed {seed}: recovered {} > written {}",
result.records_recovered,
result.records_written,
);
});
}
#[test]
fn pristine_comparison_with_faults() {
(0..1000).for_each(|seed| {
let payloads: Vec<Vec<u8>> = (0..8).map(|i| format!("item-{i}").into_bytes()).collect();
let Ok(result) = run_pristine_comparison(seed, FaultConfig::moderate(), &payloads, 4)
else {
return;
};
assert!(
result.recovery_within_bounds,
"seed {seed}: recovered more than written"
);
});
}
}
+411
View File
@@ -0,0 +1,411 @@
use std::cell::Cell;
use std::collections::HashMap;
use std::fs;
use std::io;
use std::os::unix::fs::FileExt;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
pub enum MappedFile {
Mmap(memmap2::Mmap),
Buffer(Vec<u8>),
}
impl AsRef<[u8]> for MappedFile {
fn as_ref(&self) -> &[u8] {
match self {
MappedFile::Mmap(m) => m.as_ref(),
MappedFile::Buffer(b) => b.as_ref(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct FileId(u64);
impl FileId {
pub(crate) fn new(id: u64) -> Self {
Self(id)
}
pub fn raw(self) -> u64 {
self.0
}
}
#[derive(Debug, Clone, Copy)]
pub struct OpenOptions {
pub read: bool,
pub write: bool,
pub create: bool,
pub truncate: bool,
}
impl OpenOptions {
pub fn read() -> Self {
Self {
read: true,
write: false,
create: false,
truncate: false,
}
}
pub fn write() -> Self {
Self {
read: false,
write: true,
create: true,
truncate: false,
}
}
pub fn read_write() -> Self {
Self {
read: true,
write: true,
create: true,
truncate: false,
}
}
pub fn read_only_existing() -> Self {
Self {
read: true,
write: false,
create: false,
truncate: false,
}
}
pub fn read_write_existing() -> Self {
Self {
read: true,
write: true,
create: false,
truncate: false,
}
}
}
pub trait StorageIO: Send + Sync {
fn open(&self, path: &Path, opts: OpenOptions) -> io::Result<FileId>;
fn close(&self, fd: FileId) -> io::Result<()>;
fn read_at(&self, fd: FileId, offset: u64, buf: &mut [u8]) -> io::Result<usize>;
fn write_at(&self, fd: FileId, offset: u64, buf: &[u8]) -> io::Result<usize>;
fn sync(&self, fd: FileId) -> io::Result<()>;
fn file_size(&self, fd: FileId) -> io::Result<u64>;
fn truncate(&self, fd: FileId, size: u64) -> io::Result<()>;
fn rename(&self, from: &Path, to: &Path) -> io::Result<()>;
fn delete(&self, path: &Path) -> io::Result<()>;
fn mkdir(&self, path: &Path) -> io::Result<()>;
fn sync_dir(&self, path: &Path) -> io::Result<()>;
fn list_dir(&self, path: &Path) -> io::Result<Vec<PathBuf>>;
fn write_all_at(&self, fd: FileId, offset: u64, buf: &[u8]) -> io::Result<()> {
let written = Cell::new(0usize);
std::iter::from_fn(|| (written.get() < buf.len()).then_some(()))
.try_fold(offset, |pos, ()| {
let w = written.get();
let n = self.write_at(fd, pos, &buf[w..])?;
match n {
0 => Err(io::Error::new(
io::ErrorKind::WriteZero,
"write returned 0 bytes",
)),
n => {
written.set(w + n);
Ok(pos + n as u64)
}
}
})
.map(|_| ())
}
fn read_exact_at(&self, fd: FileId, offset: u64, buf: &mut [u8]) -> io::Result<()> {
let total = buf.len();
let progress = Cell::new(0usize);
std::iter::from_fn(|| (progress.get() < total).then_some(()))
.try_fold(offset, |pos, ()| {
let r = progress.get();
let n = self.read_at(fd, pos, &mut buf[r..])?;
match n {
0 => Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"unexpected eof",
)),
n => {
progress.set(r + n);
Ok(pos + n as u64)
}
}
})
.map(|_| ())
}
fn mmap_file(&self, fd: FileId) -> io::Result<MappedFile> {
let size = self.file_size(fd)?;
let mut buf = vec![0u8; size as usize];
self.read_exact_at(fd, 0, &mut buf)?;
Ok(MappedFile::Buffer(buf))
}
}
impl<S: StorageIO> StorageIO for Arc<S> {
fn open(&self, path: &Path, opts: OpenOptions) -> io::Result<FileId> {
(**self).open(path, opts)
}
fn close(&self, fd: FileId) -> io::Result<()> {
(**self).close(fd)
}
fn read_at(&self, fd: FileId, offset: u64, buf: &mut [u8]) -> io::Result<usize> {
(**self).read_at(fd, offset, buf)
}
fn write_at(&self, fd: FileId, offset: u64, buf: &[u8]) -> io::Result<usize> {
(**self).write_at(fd, offset, buf)
}
fn sync(&self, fd: FileId) -> io::Result<()> {
(**self).sync(fd)
}
fn file_size(&self, fd: FileId) -> io::Result<u64> {
(**self).file_size(fd)
}
fn truncate(&self, fd: FileId, size: u64) -> io::Result<()> {
(**self).truncate(fd, size)
}
fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
(**self).rename(from, to)
}
fn delete(&self, path: &Path) -> io::Result<()> {
(**self).delete(path)
}
fn mkdir(&self, path: &Path) -> io::Result<()> {
(**self).mkdir(path)
}
fn sync_dir(&self, path: &Path) -> io::Result<()> {
(**self).sync_dir(path)
}
fn list_dir(&self, path: &Path) -> io::Result<Vec<PathBuf>> {
(**self).list_dir(path)
}
fn mmap_file(&self, fd: FileId) -> io::Result<MappedFile> {
(**self).mmap_file(fd)
}
}
pub struct RealIO {
next_id: AtomicU64,
fds: Mutex<HashMap<FileId, Arc<fs::File>>>,
}
impl RealIO {
pub fn new() -> Self {
Self {
next_id: AtomicU64::new(1),
fds: Mutex::new(HashMap::new()),
}
}
fn lookup(&self, id: FileId) -> io::Result<Arc<fs::File>> {
self.fds
.lock()
.map_err(|_| io::Error::other("fd table lock poisoned"))?
.get(&id)
.cloned()
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "unknown file id"))
}
}
impl Default for RealIO {
fn default() -> Self {
Self::new()
}
}
impl StorageIO for RealIO {
fn open(&self, path: &Path, opts: OpenOptions) -> io::Result<FileId> {
let file = fs::OpenOptions::new()
.read(opts.read)
.write(opts.write)
.create(opts.create)
.truncate(opts.truncate)
.open(path)?;
let id = FileId(self.next_id.fetch_add(1, Ordering::Relaxed));
self.fds
.lock()
.map_err(|_| io::Error::other("fd table lock poisoned"))?
.insert(id, Arc::new(file));
Ok(id)
}
fn close(&self, id: FileId) -> io::Result<()> {
self.fds
.lock()
.map_err(|_| io::Error::other("fd table lock poisoned"))?
.remove(&id)
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "unknown file id"))?;
Ok(())
}
fn read_at(&self, id: FileId, offset: u64, buf: &mut [u8]) -> io::Result<usize> {
self.lookup(id)?.read_at(buf, offset)
}
fn write_at(&self, id: FileId, offset: u64, buf: &[u8]) -> io::Result<usize> {
self.lookup(id)?.write_at(buf, offset)
}
fn sync(&self, id: FileId) -> io::Result<()> {
self.lookup(id)?.sync_data()
}
fn file_size(&self, id: FileId) -> io::Result<u64> {
self.lookup(id)?.metadata().map(|m| m.len())
}
fn truncate(&self, id: FileId, size: u64) -> io::Result<()> {
self.lookup(id)?.set_len(size)
}
fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
fs::rename(from, to)
}
fn delete(&self, path: &Path) -> io::Result<()> {
fs::remove_file(path)
}
fn mkdir(&self, path: &Path) -> io::Result<()> {
fs::create_dir_all(path)
}
fn sync_dir(&self, path: &Path) -> io::Result<()> {
fs::File::open(path)?.sync_all()
}
fn list_dir(&self, path: &Path) -> io::Result<Vec<PathBuf>> {
fs::read_dir(path)?
.map(|entry| entry.map(|e| e.path()))
.collect()
}
fn mmap_file(&self, fd: FileId) -> io::Result<MappedFile> {
let file = self.lookup(fd)?;
let mmap = unsafe { memmap2::Mmap::map(&*file)? };
Ok(MappedFile::Mmap(mmap))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn real_io_round_trip() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("test.dat");
let io = RealIO::new();
let fd = io.open(&path, OpenOptions::read_write()).unwrap();
let data = b"hello tranquil-store";
let written = io.write_at(fd, 0, data).unwrap();
assert_eq!(written, data.len());
io.sync(fd).unwrap();
let mut buf = vec![0u8; data.len()];
let read = io.read_at(fd, 0, &mut buf).unwrap();
assert_eq!(read, data.len());
assert_eq!(&buf, data);
assert_eq!(io.file_size(fd).unwrap(), data.len() as u64);
io.truncate(fd, 5).unwrap();
assert_eq!(io.file_size(fd).unwrap(), 5);
io.close(fd).unwrap();
}
#[test]
fn real_io_write_all_at() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("writeall.dat");
let io = RealIO::new();
let fd = io.open(&path, OpenOptions::read_write()).unwrap();
let data = b"complete write via write_all_at";
io.write_all_at(fd, 0, data).unwrap();
io.sync(fd).unwrap();
let mut buf = vec![0u8; data.len()];
io.read_exact_at(fd, 0, &mut buf).unwrap();
assert_eq!(&buf, data);
io.close(fd).unwrap();
}
#[test]
fn real_io_rename_and_delete() {
let tmp = tempfile::TempDir::new().unwrap();
let path_a = tmp.path().join("a.dat");
let path_b = tmp.path().join("b.dat");
let io = RealIO::new();
let fd = io.open(&path_a, OpenOptions::read_write()).unwrap();
io.write_all_at(fd, 0, b"data").unwrap();
io.sync(fd).unwrap();
io.close(fd).unwrap();
io.rename(&path_a, &path_b).unwrap();
assert!(!path_a.exists());
assert!(path_b.exists());
io.delete(&path_b).unwrap();
assert!(!path_b.exists());
}
#[test]
fn real_io_mkdir_and_sync_dir() {
let tmp = tempfile::TempDir::new().unwrap();
let dir = tmp.path().join("subdir");
let io = RealIO::new();
io.mkdir(&dir).unwrap();
assert!(dir.is_dir());
io.sync_dir(&dir).unwrap();
}
#[test]
fn concurrent_read_write_no_deadlock() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("concurrent.dat");
let io = Arc::new(RealIO::new());
let fd = io.open(&path, OpenOptions::read_write()).unwrap();
io.write_all_at(fd, 0, &vec![0u8; 4096]).unwrap();
io.sync(fd).unwrap();
let handles: Vec<_> = (0..4)
.map(|i| {
let io = Arc::clone(&io);
std::thread::spawn(move || {
let offset = (i * 1024) as u64;
let data = vec![i as u8; 1024];
io.write_all_at(fd, offset, &data).unwrap();
let mut buf = vec![0u8; 1024];
io.read_exact_at(fd, offset, &mut buf).unwrap();
})
})
.collect();
handles.into_iter().for_each(|h| h.join().unwrap());
io.close(fd).unwrap();
}
}
+19
View File
@@ -0,0 +1,19 @@
pub mod blockstore;
mod harness;
mod io;
mod record;
#[cfg(any(test, feature = "test-harness"))]
mod sim;
pub use blockstore::BlocksSynced;
pub use fsync_order::PostBlockstoreHook;
#[cfg(any(test, feature = "test-harness"))]
pub use harness::{
CrashTestResult, PristineComparisonResult, run_crash_test, run_pristine_comparison,
};
pub use io::{FileId, MappedFile, OpenOptions, RealIO, StorageIO};
pub use record::{
FILE_MAGIC, FORMAT_VERSION, HEADER_SIZE, MAX_RECORD_PAYLOAD, RECORD_OVERHEAD, ReadRecord,
RecordReader, RecordWriter,
};
pub use sim::{FaultConfig, SimulatedIO};
+358
View File
@@ -0,0 +1,358 @@
use std::io;
use crate::io::{FileId, StorageIO};
pub const FILE_MAGIC: [u8; 4] = *b"TQST";
pub const FORMAT_VERSION: u8 = 2;
pub const HEADER_SIZE: usize = 5;
pub const RECORD_OVERHEAD: usize = 8;
pub const MAX_RECORD_PAYLOAD: usize = 16 * 1024 * 1024;
fn record_checksum(length_bytes: &[u8; 4], payload: &[u8]) -> u32 {
let mut hasher = xxhash_rust::xxh3::Xxh3::new();
hasher.update(length_bytes);
hasher.update(payload);
hasher.digest() as u32
}
pub struct RecordWriter<'a, S: StorageIO> {
io: &'a S,
fd: FileId,
position: u64,
}
impl<'a, S: StorageIO> RecordWriter<'a, S> {
pub fn new(io: &'a S, fd: FileId) -> io::Result<Self> {
let header = [
FILE_MAGIC[0],
FILE_MAGIC[1],
FILE_MAGIC[2],
FILE_MAGIC[3],
FORMAT_VERSION,
];
io.write_all_at(fd, 0, &header)?;
Ok(Self {
io,
fd,
position: HEADER_SIZE as u64,
})
}
pub fn resume(io: &'a S, fd: FileId, position: u64) -> Self {
Self { io, fd, position }
}
pub fn append(&mut self, payload: &[u8]) -> io::Result<u64> {
let length = u32::try_from(payload.len())
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "payload exceeds u32::MAX"))?;
if payload.len() > MAX_RECORD_PAYLOAD {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"payload exceeds MAX_RECORD_PAYLOAD",
));
}
let length_bytes = length.to_le_bytes();
let checksum = record_checksum(&length_bytes, payload);
let mut cursor = self.position;
self.io.write_all_at(self.fd, cursor, &length_bytes)?;
cursor += 4;
self.io.write_all_at(self.fd, cursor, payload)?;
cursor += payload.len() as u64;
let checksum_bytes = checksum.to_le_bytes();
self.io.write_all_at(self.fd, cursor, &checksum_bytes)?;
cursor += 4;
let record_start = self.position;
self.position = cursor;
Ok(record_start)
}
pub fn sync(&self) -> io::Result<()> {
self.io.sync(self.fd)
}
pub fn position(&self) -> u64 {
self.position
}
}
#[derive(Debug)]
pub enum ReadRecord {
Valid { offset: u64, payload: Vec<u8> },
Corrupted { offset: u64 },
Truncated { offset: u64 },
}
pub struct RecordReader<'a, S: StorageIO> {
io: &'a S,
fd: FileId,
position: u64,
file_size: u64,
}
impl<'a, S: StorageIO> RecordReader<'a, S> {
pub fn open(io: &'a S, fd: FileId) -> io::Result<Self> {
let file_size = io.file_size(fd)?;
if file_size < HEADER_SIZE as u64 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"file too small for header",
));
}
let mut header = [0u8; HEADER_SIZE];
io.read_exact_at(fd, 0, &mut header)?;
if header[..4] != FILE_MAGIC {
return Err(io::Error::new(io::ErrorKind::InvalidData, "bad magic"));
}
if header[4] != FORMAT_VERSION {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"unsupported format version",
));
}
Ok(Self {
io,
fd,
position: HEADER_SIZE as u64,
file_size,
})
}
pub fn valid_records(self) -> Vec<Vec<u8>> {
self.map_while(|r| match r {
ReadRecord::Valid { payload, .. } => Some(payload),
_ => None,
})
.collect()
}
fn advance_truncated(&mut self) -> ReadRecord {
let offset = self.position;
self.position = self.file_size;
ReadRecord::Truncated { offset }
}
}
impl<S: StorageIO> Iterator for RecordReader<'_, S> {
type Item = ReadRecord;
fn next(&mut self) -> Option<Self::Item> {
if self.position >= self.file_size {
return None;
}
let remaining = self.file_size - self.position;
if remaining < 4 {
return Some(self.advance_truncated());
}
let mut length_bytes = [0u8; 4];
if self
.io
.read_exact_at(self.fd, self.position, &mut length_bytes)
.is_err()
{
return Some(self.advance_truncated());
}
let length = u32::from_le_bytes(length_bytes) as u64;
if length as usize > MAX_RECORD_PAYLOAD {
let offset = self.position;
self.position = self.file_size;
return Some(ReadRecord::Corrupted { offset });
}
let record_size = 4 + length + 4;
if self.position + record_size > self.file_size {
return Some(self.advance_truncated());
}
let mut payload = vec![0u8; length as usize];
if self
.io
.read_exact_at(self.fd, self.position + 4, &mut payload)
.is_err()
{
return Some(self.advance_truncated());
}
let mut checksum_bytes = [0u8; 4];
if self
.io
.read_exact_at(self.fd, self.position + 4 + length, &mut checksum_bytes)
.is_err()
{
return Some(self.advance_truncated());
}
let stored_checksum = u32::from_le_bytes(checksum_bytes);
let computed_checksum = record_checksum(&length_bytes, &payload);
let offset = self.position;
if stored_checksum == computed_checksum {
self.position += record_size;
Some(ReadRecord::Valid { offset, payload })
} else {
self.position = self.file_size;
Some(ReadRecord::Corrupted { offset })
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::OpenOptions;
use crate::sim::SimulatedIO;
use std::path::Path;
fn setup() -> (SimulatedIO, FileId) {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/test");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let fd = sim
.open(Path::new("/test/records.dat"), OpenOptions::read_write())
.unwrap();
(sim, fd)
}
#[test]
fn write_and_read_single_record() {
let (sim, fd) = setup();
let mut writer = RecordWriter::new(&sim, fd).unwrap();
writer.append(b"hello world").unwrap();
writer.sync().unwrap();
let reader = RecordReader::open(&sim, fd).unwrap();
let records = reader.valid_records();
assert_eq!(records.len(), 1);
assert_eq!(records[0], b"hello world");
}
#[test]
fn write_and_read_multiple_records() {
let (sim, fd) = setup();
let mut writer = RecordWriter::new(&sim, fd).unwrap();
writer.append(b"first").unwrap();
writer.append(b"second").unwrap();
writer.append(b"third").unwrap();
writer.sync().unwrap();
let reader = RecordReader::open(&sim, fd).unwrap();
let records = reader.valid_records();
assert_eq!(records.len(), 3);
assert_eq!(records[0], b"first");
assert_eq!(records[1], b"second");
assert_eq!(records[2], b"third");
}
#[test]
fn empty_file_has_no_records() {
let (sim, fd) = setup();
RecordWriter::new(&sim, fd).unwrap();
let reader = RecordReader::open(&sim, fd).unwrap();
let records = reader.valid_records();
assert!(records.is_empty());
}
#[test]
fn detects_truncated_record() {
let (sim, fd) = setup();
let mut writer = RecordWriter::new(&sim, fd).unwrap();
writer.append(b"complete record").unwrap();
writer.sync().unwrap();
let length_bytes = 100u32.to_le_bytes();
sim.write_all_at(fd, writer.position(), &length_bytes)
.unwrap();
sim.write_all_at(fd, writer.position() + 4, b"short")
.unwrap();
let mut reader = RecordReader::open(&sim, fd).unwrap();
let first = reader.next().unwrap();
assert!(matches!(first, ReadRecord::Valid { .. }));
let second = reader.next().unwrap();
assert!(matches!(second, ReadRecord::Truncated { .. }));
}
#[test]
fn crash_before_sync_loses_records() {
let (sim, fd) = setup();
let mut writer = RecordWriter::new(&sim, fd).unwrap();
writer.append(b"synced").unwrap();
writer.sync().unwrap();
sim.sync_dir(Path::new("/test")).unwrap();
writer.append(b"not synced").unwrap();
sim.crash();
let fd = sim
.open(Path::new("/test/records.dat"), OpenOptions::read())
.unwrap();
let reader = RecordReader::open(&sim, fd).unwrap();
let records = reader.valid_records();
assert_eq!(records.len(), 1);
assert_eq!(records[0], b"synced");
}
#[test]
fn real_io_record_round_trip() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("records.dat");
let real = crate::RealIO::new();
let fd = real.open(&path, OpenOptions::read_write()).unwrap();
let mut writer = RecordWriter::new(&real, fd).unwrap();
writer.append(b"real record 1").unwrap();
writer.append(b"real record 2").unwrap();
writer.sync().unwrap();
let reader = RecordReader::open(&real, fd).unwrap();
let records = reader.valid_records();
assert_eq!(records.len(), 2);
assert_eq!(records[0], b"real record 1");
assert_eq!(records[1], b"real record 2");
real.close(fd).unwrap();
}
#[test]
fn checksum_detects_single_bit_flip() {
let (sim, fd) = setup();
let mut writer = RecordWriter::new(&sim, fd).unwrap();
let payload = vec![0xAA; 256];
writer.append(&payload).unwrap();
writer.sync().unwrap();
let mut contents = sim.buffered_contents(fd).unwrap();
let payload_start = HEADER_SIZE + 4;
contents[payload_start + 128] ^= 0x01;
let sim2 = SimulatedIO::pristine(99);
let dir2 = Path::new("/verify");
sim2.mkdir(dir2).unwrap();
sim2.sync_dir(dir2).unwrap();
let fd2 = sim2
.open(Path::new("/verify/check.dat"), OpenOptions::read_write())
.unwrap();
sim2.write_all_at(fd2, 0, &contents).unwrap();
let mut reader = RecordReader::open(&sim2, fd2).unwrap();
let record = reader.next().unwrap();
assert!(matches!(record, ReadRecord::Corrupted { .. }));
}
}
+838
View File
@@ -0,0 +1,838 @@
use std::collections::{HashMap, HashSet};
use std::io;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use crate::io::{FileId, OpenOptions, StorageIO};
#[derive(Debug, Clone, Copy)]
pub struct FaultConfig {
pub partial_write_probability: f64,
pub bit_flip_on_read_probability: f64,
pub sync_failure_probability: f64,
pub dir_sync_failure_probability: f64,
pub misdirected_write_probability: f64,
pub io_error_probability: f64,
}
impl FaultConfig {
pub fn none() -> Self {
Self {
partial_write_probability: 0.0,
bit_flip_on_read_probability: 0.0,
sync_failure_probability: 0.0,
dir_sync_failure_probability: 0.0,
misdirected_write_probability: 0.0,
io_error_probability: 0.0,
}
}
pub fn moderate() -> Self {
Self {
partial_write_probability: 0.05,
bit_flip_on_read_probability: 0.01,
sync_failure_probability: 0.03,
dir_sync_failure_probability: 0.02,
misdirected_write_probability: 0.01,
io_error_probability: 0.02,
}
}
pub fn aggressive() -> Self {
Self {
partial_write_probability: 0.15,
bit_flip_on_read_probability: 0.05,
sync_failure_probability: 0.10,
dir_sync_failure_probability: 0.05,
misdirected_write_probability: 0.05,
io_error_probability: 0.08,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct StorageId(u64);
struct SimStorage {
buffered: Vec<u8>,
durable: Vec<u8>,
dir_entry_durable: bool,
}
struct SimFd {
storage_id: StorageId,
readable: bool,
writable: bool,
}
#[derive(Debug, Clone)]
pub enum OpRecord {
Open {
fd: FileId,
path: PathBuf,
},
Close {
fd: FileId,
},
ReadAt {
fd: FileId,
offset: u64,
len: usize,
},
WriteAt {
fd: FileId,
offset: u64,
data: Vec<u8>,
actual_written: usize,
},
Sync {
fd: FileId,
succeeded: bool,
},
Truncate {
fd: FileId,
size: u64,
},
Rename {
from: PathBuf,
to: PathBuf,
},
Delete {
path: PathBuf,
},
Mkdir {
path: PathBuf,
},
SyncDir {
path: PathBuf,
},
}
struct SimState {
storage: HashMap<StorageId, SimStorage>,
paths: HashMap<PathBuf, StorageId>,
fds: HashMap<FileId, SimFd>,
dirs_durable: HashSet<PathBuf>,
op_log: Vec<OpRecord>,
rng_counter: u64,
next_fd_id: u64,
next_storage_id: u64,
}
impl SimState {
fn next_random(&mut self, seed: u64) -> f64 {
let counter = self.rng_counter;
self.rng_counter += 1;
let mixed = splitmix64(seed.wrapping_add(counter));
(mixed >> 11) as f64 / (1u64 << 53) as f64
}
fn next_random_usize(&mut self, seed: u64, max: usize) -> usize {
if max == 0 {
return 0;
}
let counter = self.rng_counter;
self.rng_counter += 1;
let mixed = splitmix64(seed.wrapping_add(counter));
(mixed as usize) % max
}
fn should_fault(&mut self, seed: u64, probability: f64) -> bool {
probability > 0.0 && self.next_random(seed) < probability
}
fn alloc_fd_id(&mut self) -> FileId {
let id = self.next_fd_id;
self.next_fd_id += 1;
FileId::new(id)
}
fn alloc_storage_id(&mut self) -> StorageId {
let id = self.next_storage_id;
self.next_storage_id += 1;
StorageId(id)
}
fn require_open(&self, id: FileId) -> io::Result<StorageId> {
let fd_info = self
.fds
.get(&id)
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "unknown file id"))?;
if !self.storage.contains_key(&fd_info.storage_id) {
return Err(io::Error::new(
io::ErrorKind::NotFound,
"underlying storage removed",
));
}
Ok(fd_info.storage_id)
}
fn require_readable(&self, id: FileId) -> io::Result<StorageId> {
let sid = self.require_open(id)?;
if !self.fds[&id].readable {
return Err(io::Error::new(
io::ErrorKind::PermissionDenied,
"file not opened for reading",
));
}
Ok(sid)
}
fn require_writable(&self, id: FileId) -> io::Result<StorageId> {
let sid = self.require_open(id)?;
if !self.fds[&id].writable {
return Err(io::Error::new(
io::ErrorKind::PermissionDenied,
"file not opened for writing",
));
}
Ok(sid)
}
}
pub struct SimulatedIO {
state: Mutex<SimState>,
fault_config: FaultConfig,
rng_seed: u64,
}
impl SimulatedIO {
pub fn new(seed: u64, fault_config: FaultConfig) -> Self {
Self {
state: Mutex::new(SimState {
storage: HashMap::new(),
paths: HashMap::new(),
fds: HashMap::new(),
dirs_durable: HashSet::new(),
op_log: Vec::new(),
rng_counter: 0,
next_fd_id: 1,
next_storage_id: 1,
}),
fault_config,
rng_seed: seed,
}
}
pub fn pristine(seed: u64) -> Self {
Self::new(seed, FaultConfig::none())
}
pub fn crash(&self) {
let mut state = self.state.lock().unwrap();
state.fds.clear();
let orphaned: Vec<StorageId> = state
.storage
.iter()
.filter(|(_, s)| !s.dir_entry_durable)
.map(|(sid, _)| *sid)
.collect();
orphaned.iter().for_each(|sid| {
state.storage.remove(sid);
});
let live_sids: HashSet<StorageId> = state.storage.keys().copied().collect();
state.paths.retain(|_, sid| live_sids.contains(sid));
state
.storage
.values_mut()
.for_each(|s| s.buffered = s.durable.clone());
}
pub fn op_log(&self) -> Vec<OpRecord> {
self.state.lock().unwrap().op_log.clone()
}
pub fn durable_contents(&self, fd: FileId) -> io::Result<Vec<u8>> {
let state = self.state.lock().unwrap();
let sid = state.require_open(fd)?;
Ok(state.storage.get(&sid).unwrap().durable.clone())
}
pub fn buffered_contents(&self, fd: FileId) -> io::Result<Vec<u8>> {
let state = self.state.lock().unwrap();
let sid = state.require_open(fd)?;
Ok(state.storage.get(&sid).unwrap().buffered.clone())
}
pub fn last_sync_persisted(&self) -> bool {
let state = self.state.lock().unwrap();
state
.op_log
.iter()
.rev()
.find_map(|op| match op {
OpRecord::Sync { succeeded, .. } => Some(*succeeded),
_ => None,
})
.unwrap_or(false)
}
}
impl StorageIO for SimulatedIO {
fn open(&self, path: &Path, opts: OpenOptions) -> io::Result<FileId> {
let mut state = self.state.lock().unwrap();
let seed = self.rng_seed;
if state.should_fault(seed, self.fault_config.io_error_probability) {
return Err(io::Error::other("simulated EIO on open"));
}
let path_buf = path.to_path_buf();
let fd_id = state.alloc_fd_id();
match state.paths.get(&path_buf).copied() {
Some(sid) => {
if opts.truncate {
state.storage.get_mut(&sid).unwrap().buffered.clear();
}
state.fds.insert(
fd_id,
SimFd {
storage_id: sid,
readable: opts.read,
writable: opts.write,
},
);
}
None => {
if !opts.create {
return Err(io::Error::new(
io::ErrorKind::NotFound,
"file not found and create not set",
));
}
let sid = state.alloc_storage_id();
state.storage.insert(
sid,
SimStorage {
buffered: Vec::new(),
durable: Vec::new(),
dir_entry_durable: false,
},
);
state.paths.insert(path_buf.clone(), sid);
state.fds.insert(
fd_id,
SimFd {
storage_id: sid,
readable: opts.read,
writable: opts.write,
},
);
}
};
state.op_log.push(OpRecord::Open {
fd: fd_id,
path: path_buf,
});
Ok(fd_id)
}
fn close(&self, id: FileId) -> io::Result<()> {
let mut state = self.state.lock().unwrap();
let fd_info = state
.fds
.remove(&id)
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "unknown file id"))?;
let sid = fd_info.storage_id;
let unlinked = !state.paths.values().any(|s| *s == sid);
let no_remaining_fds = !state.fds.values().any(|f| f.storage_id == sid);
if unlinked && no_remaining_fds {
state.storage.remove(&sid);
}
state.op_log.push(OpRecord::Close { fd: id });
Ok(())
}
fn read_at(&self, id: FileId, offset: u64, buf: &mut [u8]) -> io::Result<usize> {
let mut state = self.state.lock().unwrap();
let sid = state.require_readable(id)?;
let seed = self.rng_seed;
if state.should_fault(seed, self.fault_config.io_error_probability) {
return Err(io::Error::other("simulated EIO on read"));
}
let storage = state.storage.get(&sid).unwrap();
let off = usize::try_from(offset)
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "offset exceeds usize"))?;
if off >= storage.buffered.len() {
state.op_log.push(OpRecord::ReadAt {
fd: id,
offset,
len: 0,
});
return Ok(0);
}
let available = storage.buffered.len().saturating_sub(off);
let to_read = buf.len().min(available);
buf[..to_read].copy_from_slice(&storage.buffered[off..off + to_read]);
if state.should_fault(seed, self.fault_config.bit_flip_on_read_probability) && to_read > 0 {
let flip_pos = state.next_random_usize(seed, to_read);
let flip_bit = state.next_random_usize(seed, 8);
buf[flip_pos] ^= 1 << flip_bit;
}
state.op_log.push(OpRecord::ReadAt {
fd: id,
offset,
len: to_read,
});
Ok(to_read)
}
fn write_at(&self, id: FileId, offset: u64, buf: &[u8]) -> io::Result<usize> {
let mut state = self.state.lock().unwrap();
let sid = state.require_writable(id)?;
let seed = self.rng_seed;
if state.should_fault(seed, self.fault_config.io_error_probability) {
return Err(io::Error::other("simulated EIO on write"));
}
let actual_len = if buf.len() > 1
&& state.should_fault(seed, self.fault_config.partial_write_probability)
{
let partial = state.next_random_usize(seed, buf.len());
partial.max(1)
} else {
buf.len()
};
let misdirected = state.should_fault(seed, self.fault_config.misdirected_write_probability);
let write_offset = if misdirected {
let drift = state.next_random_usize(seed, 64) as u64;
if state.next_random(seed) < 0.5 {
offset.saturating_sub(drift)
} else {
offset.saturating_add(drift)
}
} else {
offset
};
let storage = state.storage.get_mut(&sid).unwrap();
let off = usize::try_from(write_offset)
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "offset exceeds usize"))?;
let end = off.saturating_add(actual_len);
if end > storage.buffered.len() {
storage.buffered.resize(end, 0);
}
storage.buffered[off..end].copy_from_slice(&buf[..actual_len]);
state.op_log.push(OpRecord::WriteAt {
fd: id,
offset,
data: buf[..actual_len].to_vec(),
actual_written: actual_len,
});
Ok(actual_len)
}
fn sync(&self, id: FileId) -> io::Result<()> {
let mut state = self.state.lock().unwrap();
let sid = state.require_open(id)?;
let seed = self.rng_seed;
if state.should_fault(seed, self.fault_config.io_error_probability) {
return Err(io::Error::other("simulated EIO on sync"));
}
let sync_succeeded = !state.should_fault(seed, self.fault_config.sync_failure_probability);
let storage = state.storage.get_mut(&sid).unwrap();
if sync_succeeded {
storage.durable = storage.buffered.clone();
}
state.op_log.push(OpRecord::Sync {
fd: id,
succeeded: sync_succeeded,
});
Ok(())
}
fn file_size(&self, id: FileId) -> io::Result<u64> {
let state = self.state.lock().unwrap();
let sid = state.require_open(id)?;
Ok(state.storage.get(&sid).unwrap().buffered.len() as u64)
}
fn truncate(&self, id: FileId, size: u64) -> io::Result<()> {
let mut state = self.state.lock().unwrap();
let sid = state.require_open(id)?;
let storage = state.storage.get_mut(&sid).unwrap();
let target = usize::try_from(size)
.map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "size exceeds usize"))?;
storage.buffered.resize(target, 0);
state.op_log.push(OpRecord::Truncate { fd: id, size });
Ok(())
}
fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
let mut state = self.state.lock().unwrap();
let from_buf = from.to_path_buf();
let to_buf = to.to_path_buf();
let sid = state
.paths
.remove(&from_buf)
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "source file not found"))?;
let storage = state.storage.get_mut(&sid).unwrap();
storage.dir_entry_durable = false;
state.paths.insert(to_buf.clone(), sid);
state.op_log.push(OpRecord::Rename {
from: from_buf,
to: to_buf,
});
Ok(())
}
fn delete(&self, path: &Path) -> io::Result<()> {
let mut state = self.state.lock().unwrap();
let path_buf = path.to_path_buf();
let sid = state
.paths
.remove(&path_buf)
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "file not found"))?;
let has_open_fds = state.fds.values().any(|fd_info| fd_info.storage_id == sid);
if !has_open_fds {
state.storage.remove(&sid);
}
state.op_log.push(OpRecord::Delete { path: path_buf });
Ok(())
}
fn mkdir(&self, path: &Path) -> io::Result<()> {
let mut state = self.state.lock().unwrap();
state.op_log.push(OpRecord::Mkdir {
path: path.to_path_buf(),
});
Ok(())
}
fn sync_dir(&self, path: &Path) -> io::Result<()> {
let mut state = self.state.lock().unwrap();
let seed = self.rng_seed;
if state.should_fault(seed, self.fault_config.io_error_probability) {
return Err(io::Error::other("simulated EIO on sync_dir"));
}
let dir_path = path.to_path_buf();
let actually_persisted =
!state.should_fault(seed, self.fault_config.dir_sync_failure_probability);
if actually_persisted {
state.dirs_durable.insert(dir_path.clone());
let sids_in_dir: Vec<StorageId> = state
.paths
.iter()
.filter(|(p, _)| p.parent().map(|parent| parent == path).unwrap_or(false))
.map(|(_, sid)| *sid)
.collect();
sids_in_dir.iter().for_each(|sid| {
if let Some(storage) = state.storage.get_mut(sid) {
storage.dir_entry_durable = true;
}
});
}
state.op_log.push(OpRecord::SyncDir { path: dir_path });
Ok(())
}
fn list_dir(&self, path: &Path) -> io::Result<Vec<PathBuf>> {
let state = self.state.lock().unwrap();
let entries: Vec<PathBuf> = state
.paths
.keys()
.filter(|p| p.parent() == Some(path))
.cloned()
.collect();
Ok(entries)
}
}
fn splitmix64(mut x: u64) -> u64 {
x = x.wrapping_add(0x9e3779b97f4a7c15);
x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb);
x ^ (x >> 31)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pristine_round_trip() {
let sim = SimulatedIO::pristine(42);
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
let data = b"hello simulation";
sim.write_at(fd, 0, data).unwrap();
let mut buf = vec![0u8; data.len()];
sim.read_at(fd, 0, &mut buf).unwrap();
assert_eq!(&buf, data);
}
#[test]
fn crash_resets_to_durable() {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/test");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"durable data").unwrap();
sim.sync(fd).unwrap();
sim.sync_dir(dir).unwrap();
sim.write_at(fd, 0, b"volatile!!!!").unwrap();
sim.crash();
let fd = sim.open(path, OpenOptions::read()).unwrap();
let mut buf = vec![0u8; 12];
sim.read_at(fd, 0, &mut buf).unwrap();
assert_eq!(&buf, b"durable data");
}
#[test]
fn crash_with_no_sync_loses_everything() {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/test");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"never synced").unwrap();
sim.sync_dir(dir).unwrap();
sim.crash();
let fd = sim.open(path, OpenOptions::read()).unwrap();
assert_eq!(sim.file_size(fd).unwrap(), 0);
}
#[test]
fn crash_without_dir_sync_loses_file() {
let sim = SimulatedIO::pristine(42);
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"data").unwrap();
sim.sync(fd).unwrap();
sim.crash();
let result = sim.open(path, OpenOptions::read());
assert!(result.is_err());
}
#[test]
fn dir_sync_makes_file_durable() {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/test");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"persistent").unwrap();
sim.sync(fd).unwrap();
sim.sync_dir(dir).unwrap();
sim.crash();
let fd = sim.open(path, OpenOptions::read()).unwrap();
let mut buf = vec![0u8; 10];
sim.read_at(fd, 0, &mut buf).unwrap();
assert_eq!(&buf, b"persistent");
}
#[test]
fn read_only_rejects_write() {
let sim = SimulatedIO::pristine(42);
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"data").unwrap();
let fd2 = sim.open(path, OpenOptions::read()).unwrap();
assert_ne!(fd, fd2);
let result = sim.write_at(fd2, 0, b"nope");
assert_eq!(result.unwrap_err().kind(), io::ErrorKind::PermissionDenied);
}
#[test]
fn write_only_rejects_read() {
let sim = SimulatedIO::pristine(42);
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::write()).unwrap();
sim.write_at(fd, 0, b"data").unwrap();
let mut buf = vec![0u8; 4];
let result = sim.read_at(fd, 0, &mut buf);
assert_eq!(result.unwrap_err().kind(), io::ErrorKind::PermissionDenied);
}
#[test]
fn open_without_create_fails_for_missing_file() {
let sim = SimulatedIO::pristine(42);
let result = sim.open(Path::new("/nonexistent"), OpenOptions::read());
assert_eq!(result.unwrap_err().kind(), io::ErrorKind::NotFound);
}
#[test]
fn truncate_on_open() {
let sim = SimulatedIO::pristine(42);
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"existing data").unwrap();
let opts = OpenOptions {
read: true,
write: true,
create: true,
truncate: true,
};
let fd2 = sim.open(path, opts).unwrap();
assert_eq!(sim.file_size(fd2).unwrap(), 0);
}
#[test]
fn rename_makes_entry_non_durable() {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/test");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let path_a = Path::new("/test/a.dat");
let fd = sim.open(path_a, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"data").unwrap();
sim.sync(fd).unwrap();
sim.sync_dir(dir).unwrap();
let path_b = Path::new("/test/b.dat");
sim.rename(path_a, path_b).unwrap();
sim.crash();
let result_a = sim.open(path_a, OpenOptions::read());
let result_b = sim.open(path_b, OpenOptions::read());
assert!(result_a.is_err());
assert!(result_b.is_err());
}
#[test]
fn durable_contents_accessible() {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/test");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"synced").unwrap();
sim.sync(fd).unwrap();
sim.sync_dir(dir).unwrap();
sim.write_at(fd, 6, b" unsynced").unwrap();
let durable = sim.durable_contents(fd).unwrap();
assert_eq!(&durable, b"synced");
let buffered = sim.buffered_contents(fd).unwrap();
assert_eq!(&buffered, b"synced unsynced");
}
#[test]
fn op_log_records_operations() {
let sim = SimulatedIO::pristine(42);
let dir = Path::new("/test");
sim.mkdir(dir).unwrap();
sim.sync_dir(dir).unwrap();
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"data").unwrap();
sim.sync(fd).unwrap();
sim.close(fd).unwrap();
let log = sim.op_log();
assert_eq!(log.len(), 6);
assert!(matches!(log[0], OpRecord::Mkdir { .. }));
assert!(matches!(log[1], OpRecord::SyncDir { .. }));
assert!(matches!(log[2], OpRecord::Open { .. }));
assert!(matches!(log[3], OpRecord::WriteAt { .. }));
assert!(matches!(
log[4],
OpRecord::Sync {
succeeded: true,
..
}
));
assert!(matches!(log[5], OpRecord::Close { .. }));
}
#[test]
fn multiple_fds_independent_permissions() {
let sim = SimulatedIO::pristine(42);
let path = Path::new("/test/file.dat");
let fd_rw = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd_rw, 0, b"shared data").unwrap();
let fd_ro = sim.open(path, OpenOptions::read()).unwrap();
assert_ne!(fd_rw, fd_ro);
let mut buf = vec![0u8; 11];
sim.read_at(fd_ro, 0, &mut buf).unwrap();
assert_eq!(&buf, b"shared data");
sim.write_at(fd_rw, 0, b"mutated!!!!").unwrap();
sim.read_at(fd_ro, 0, &mut buf).unwrap();
assert_eq!(&buf, b"mutated!!!!");
let result = sim.write_at(fd_ro, 0, b"nope");
assert_eq!(result.unwrap_err().kind(), io::ErrorKind::PermissionDenied);
}
#[test]
fn last_sync_persisted_tracks_truth() {
let sim = SimulatedIO::pristine(42);
let path = Path::new("/test/file.dat");
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.write_at(fd, 0, b"data").unwrap();
sim.sync(fd).unwrap();
assert!(sim.last_sync_persisted());
}
}
@@ -0,0 +1,304 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use bytes::Bytes;
use cid::Cid;
use futures::StreamExt;
use jacquard_common::types::string::Did;
use jacquard_common::types::tid::Ticker;
use jacquard_repo::car::{parse_car_bytes, write_car_bytes};
use jacquard_repo::commit::Commit;
use jacquard_repo::mst::Mst;
use jacquard_repo::repo::CommitData;
use jacquard_repo::storage::BlockStore;
use multihash::Multihash;
use sha2::{Digest, Sha256};
use tranquil_store::blockstore::{
BlockStoreConfig, DEFAULT_MAX_FILE_SIZE, GroupCommitConfig, TranquilBlockStore,
};
const DAG_CBOR_CODEC: u64 = 0x71;
const SHA2_256_CODE: u64 = 0x12;
fn test_config(dir: &std::path::Path) -> BlockStoreConfig {
BlockStoreConfig {
data_dir: dir.join("data"),
index_dir: dir.join("index"),
max_file_size: DEFAULT_MAX_FILE_SIZE,
group_commit: GroupCommitConfig::default(),
}
}
fn make_record(value: &str) -> Vec<u8> {
serde_ipld_dagcbor::to_vec(&BTreeMap::from([
("$type", "app.bsky.feed.post"),
("text", value),
]))
.unwrap()
}
fn compute_cid(data: &[u8]) -> Cid {
let hash = Sha256::digest(data);
let multihash = Multihash::wrap(SHA2_256_CODE, &hash).unwrap();
Cid::new_v1(DAG_CBOR_CODEC, multihash)
}
fn test_signing_key() -> k256::ecdsa::SigningKey {
k256::ecdsa::SigningKey::random(&mut rand::thread_rng())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn mst_insert_commit_and_car_round_trip() {
let dir = tempfile::TempDir::new().unwrap();
let store = TranquilBlockStore::open(test_config(dir.path())).unwrap();
let storage = Arc::new(store.clone());
let mut mst = Mst::new(storage.clone());
let records: Vec<(String, Cid, Vec<u8>)> = (0..100u32)
.map(|i| {
let data = make_record(&format!("post number {i}"));
let cid = compute_cid(&data);
(format!("app.bsky.feed.post/{i:010}"), cid, data)
})
.collect();
let mut record_blocks = BTreeMap::new();
for (key, cid, data) in &records {
record_blocks.insert(*cid, Bytes::from(data.clone()));
storage.put(data).await.unwrap();
mst = mst.add(key, *cid).await.unwrap();
}
let mst_root = mst.persist().await.unwrap();
futures::stream::iter(&records)
.for_each(|(key, expected_cid, _)| {
let mst = &mst;
async move {
let found = mst.get(key).await.unwrap();
assert_eq!(found, Some(*expected_cid), "record {key} missing from MST");
}
})
.await;
let signing_key = test_signing_key();
let did = Did::new("did:plc:testuser123").unwrap();
let mut ticker = Ticker::new();
let rev = ticker.next(None);
let commit = Commit::new_unsigned(did, mst_root, rev.clone(), None)
.sign(&signing_key)
.unwrap();
let commit_cbor = commit.to_cbor().unwrap();
let commit_cid = compute_cid(&commit_cbor);
let commit_bytes = Bytes::from(commit_cbor);
let empty_mst = Mst::new(storage.clone());
let diff = empty_mst.diff(&mst).await.unwrap();
let mut all_blocks = diff.new_mst_blocks.clone();
all_blocks.insert(commit_cid, commit_bytes);
all_blocks.extend(record_blocks);
let commit_data = CommitData {
cid: commit_cid,
rev,
since: None,
prev: None,
data: mst_root,
prev_data: None,
blocks: all_blocks.clone(),
relevant_blocks: BTreeMap::new(),
deleted_cids: Vec::new(),
};
store.apply_commit(commit_data).await.unwrap();
let car_bytes = write_car_bytes(commit_cid, all_blocks).await.unwrap();
let parsed = parse_car_bytes(&car_bytes).await.unwrap();
assert_eq!(parsed.root, commit_cid);
assert!(parsed.blocks.contains_key(&commit_cid));
assert!(parsed.blocks.contains_key(&mst_root));
records.iter().for_each(|(_, record_cid, _)| {
assert!(
parsed.blocks.contains_key(record_cid),
"record block {record_cid} missing from CAR"
);
});
let parsed_commit = Commit::from_cbor(parsed.blocks.get(&commit_cid).unwrap()).unwrap();
assert_eq!(*parsed_commit.data(), mst_root);
let loaded_mst = Mst::load(storage.clone(), mst_root, None);
futures::stream::iter(&records)
.for_each(|(key, expected_cid, _)| {
let loaded_mst = &loaded_mst;
async move {
let found = loaded_mst.get(key).await.unwrap();
assert_eq!(
found,
Some(*expected_cid),
"record {key} missing after reload"
);
}
})
.await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn mst_create_update_delete_with_refcounts() {
let dir = tempfile::TempDir::new().unwrap();
let store = TranquilBlockStore::open(test_config(dir.path())).unwrap();
let storage = Arc::new(store.clone());
let record_a_v1 = make_record("version 1 of record A");
let record_a_v2 = make_record("version 2 of record A");
let record_b = make_record("record B to be deleted");
let record_c = make_record("record C stays forever");
let record_shared = make_record("shared content");
let cid_a_v1 = storage.put(&record_a_v1).await.unwrap();
let cid_a_v2 = compute_cid(&record_a_v2);
let cid_b = storage.put(&record_b).await.unwrap();
let cid_c = storage.put(&record_c).await.unwrap();
let cid_shared = storage.put(&record_shared).await.unwrap();
let mut mst = Mst::new(storage.clone());
mst = mst.add("app.bsky.feed.post/aaaa", cid_a_v1).await.unwrap();
mst = mst.add("app.bsky.feed.post/bbbb", cid_b).await.unwrap();
mst = mst.add("app.bsky.feed.post/cccc", cid_c).await.unwrap();
mst = mst
.add("app.bsky.feed.post/dddd", cid_shared)
.await
.unwrap();
mst = mst
.add("app.bsky.feed.post/eeee", cid_shared)
.await
.unwrap();
let mst_root_v1 = mst.persist().await.unwrap();
let signing_key = test_signing_key();
let mut ticker = Ticker::new();
let rev1 = ticker.next(None);
let commit_v1 = Commit::new_unsigned(
Did::new("did:plc:testuser123").unwrap(),
mst_root_v1,
rev1.clone(),
None,
)
.sign(&signing_key)
.unwrap();
let commit_v1_cbor = commit_v1.to_cbor().unwrap();
let commit_v1_cid = compute_cid(&commit_v1_cbor);
let empty_mst = Mst::new(storage.clone());
let diff_v1 = empty_mst.diff(&mst).await.unwrap();
let mut blocks_v1 = diff_v1.new_mst_blocks.clone();
blocks_v1.insert(commit_v1_cid, Bytes::from(commit_v1_cbor));
store
.apply_commit(CommitData {
cid: commit_v1_cid,
rev: rev1.clone(),
since: None,
prev: None,
data: mst_root_v1,
prev_data: None,
blocks: blocks_v1,
relevant_blocks: BTreeMap::new(),
deleted_cids: Vec::new(),
})
.await
.unwrap();
let old_mst = mst.clone();
mst = mst.add("app.bsky.feed.post/aaaa", cid_a_v2).await.unwrap();
mst = mst.delete("app.bsky.feed.post/bbbb").await.unwrap();
mst = mst.delete("app.bsky.feed.post/eeee").await.unwrap();
let mst_root_v2 = mst.persist().await.unwrap();
let diff_v2 = old_mst.diff(&mst).await.unwrap();
let rev2 = ticker.next(Some(rev1.clone()));
let commit_v2 = Commit::new_unsigned(
Did::new("did:plc:testuser123").unwrap(),
mst_root_v2,
rev2.clone(),
Some(commit_v1_cid),
)
.sign(&signing_key)
.unwrap();
let commit_v2_cbor = commit_v2.to_cbor().unwrap();
let commit_v2_cid = compute_cid(&commit_v2_cbor);
let mut blocks_v2 = diff_v2.new_mst_blocks.clone();
blocks_v2.insert(commit_v2_cid, Bytes::from(commit_v2_cbor));
blocks_v2.insert(cid_a_v2, Bytes::from(record_a_v2.clone()));
let mut deleted: Vec<Cid> = diff_v2.removed_mst_blocks.clone();
deleted.extend(diff_v2.removed_cids.iter());
store
.apply_commit(CommitData {
cid: commit_v2_cid,
rev: rev2,
since: Some(rev1),
prev: Some(commit_v1_cid),
data: mst_root_v2,
prev_data: Some(mst_root_v1),
blocks: blocks_v2,
relevant_blocks: BTreeMap::new(),
deleted_cids: deleted,
})
.await
.unwrap();
let retrieved_a_v2 = store.get(&cid_a_v2).await.unwrap();
assert!(retrieved_a_v2.is_some(), "updated record A v2 should exist");
assert_eq!(&retrieved_a_v2.unwrap()[..], &record_a_v2);
assert!(
store.has(&cid_a_v1).await.unwrap(),
"cid_a_v1 should still exist, tombstoned but not GC'd"
);
assert!(
store.has(&cid_b).await.unwrap(),
"cid_b should still exist, tombstoned but not GC'd"
);
assert!(
store.has(&cid_c).await.unwrap(),
"untouched record C should still exist"
);
let retrieved_c = store.get(&cid_c).await.unwrap().unwrap();
assert_eq!(&retrieved_c[..], &record_c);
assert!(
store.has(&cid_shared).await.unwrap(),
"shared-content block should still exist, tombstoned but not GC'd"
);
let loaded_mst = Mst::load(storage.clone(), mst_root_v2, None);
let expected_entries: Vec<(&str, Option<Cid>)> = vec![
("app.bsky.feed.post/aaaa", Some(cid_a_v2)),
("app.bsky.feed.post/bbbb", None),
("app.bsky.feed.post/cccc", Some(cid_c)),
("app.bsky.feed.post/dddd", Some(cid_shared)),
("app.bsky.feed.post/eeee", None),
];
futures::stream::iter(expected_entries)
.for_each(|(key, expected)| {
let loaded_mst = &loaded_mst;
async move {
assert_eq!(loaded_mst.get(key).await.unwrap(), expected);
}
})
.await;
}
+194
View File
@@ -0,0 +1,194 @@
use proptest::prelude::*;
use std::path::Path;
use tranquil_store::{
FaultConfig, OpenOptions, ReadRecord, RecordReader, RecordWriter, SimulatedIO, StorageIO,
run_crash_test, run_pristine_comparison,
};
fn arb_payloads(max_count: usize, max_size: usize) -> BoxedStrategy<Vec<Vec<u8>>> {
proptest::collection::vec(
proptest::collection::vec(any::<u8>(), 0..max_size),
1..max_count,
)
.boxed()
}
fn sim_with_dir(seed: u64, config: FaultConfig) -> SimulatedIO {
let sim = SimulatedIO::new(seed, config);
sim.mkdir(Path::new("/test")).unwrap();
sim.sync_dir(Path::new("/test")).unwrap();
sim
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(2000))]
#[test]
fn synced_records_survive_crash(
seed in any::<u64>(),
payloads in arb_payloads(20, 256),
) {
let result = run_crash_test(
seed,
FaultConfig::none(),
&payloads,
5,
).unwrap();
prop_assert_eq!(result.records_recovered, result.records_synced);
prop_assert_eq!(result.corrupted_detected, 0);
}
#[test]
fn recovered_never_exceeds_written(
seed in any::<u64>(),
payloads in arb_payloads(15, 128),
) {
let Ok(result) = run_crash_test(
seed,
FaultConfig::moderate(),
&payloads,
3,
) else { return Ok(()); };
prop_assert!(
result.records_recovered <= result.records_written,
"seed {}: recovered {} > written {}",
seed, result.records_recovered, result.records_written,
);
}
#[test]
fn no_phantom_records_after_crash(
seed in any::<u64>(),
payloads in arb_payloads(10, 512),
) {
let Ok(result) = run_crash_test(
seed,
FaultConfig::aggressive(),
&payloads,
0,
) else { return Ok(()); };
prop_assert!(
result.records_recovered <= result.records_synced,
"seed {}: recovered {} records but only {} were synced",
seed, result.records_recovered, result.records_synced,
);
}
#[test]
fn pristine_prefix_holds_under_faults(
seed in any::<u64>(),
payloads in arb_payloads(12, 200),
) {
let Ok(result) = run_pristine_comparison(
seed,
FaultConfig::moderate(),
&payloads,
4,
) else { return Ok(()); };
prop_assert!(result.recovery_within_bounds);
}
#[test]
fn bit_flip_detected_by_u32_checksum(
seed in any::<u64>(),
payload in proptest::collection::vec(any::<u8>(), 8..1024),
flip_offset in any::<usize>(),
flip_bit in 0u8..8,
) {
let sim = sim_with_dir(seed, FaultConfig::none());
let fd = sim.open(Path::new("/test/bitflip.dat"), OpenOptions::read_write()).unwrap();
let mut writer = RecordWriter::new(&sim, fd).unwrap();
writer.append(&payload).unwrap();
writer.sync().unwrap();
let mut contents = sim.durable_contents(fd).unwrap();
let record_region = &mut contents[HEADER_SIZE..];
let idx = flip_offset % record_region.len();
record_region[idx] ^= 1 << flip_bit;
let sim2 = sim_with_dir(seed.wrapping_add(1), FaultConfig::none());
let fd2 = sim2.open(Path::new("/test/check.dat"), OpenOptions::read_write()).unwrap();
sim2.write_all_at(fd2, 0, &contents).unwrap();
sim2.sync(fd2).unwrap();
let mut reader = RecordReader::open(&sim2, fd2).unwrap();
match reader.next() {
Some(ReadRecord::Valid { payload: ref recovered, .. }) => {
prop_assert_eq!(recovered, &payload, "bit flip produced valid record with wrong data");
}
Some(ReadRecord::Corrupted { .. }) => {}
Some(ReadRecord::Truncated { .. }) => {}
None => {}
}
}
#[test]
fn aggressive_faults_many_seeds(
seed in 0u64..10_000,
) {
let payloads: Vec<Vec<u8>> = (0..5)
.map(|i| format!("aggressive-test-{i}-{seed}").into_bytes())
.collect();
let Ok(result) = run_crash_test(
seed,
FaultConfig::aggressive(),
&payloads,
2,
) else { return Ok(()); };
prop_assert!(result.records_recovered <= payloads.len());
}
#[test]
fn write_all_at_handles_partial_writes(
seed in any::<u64>(),
data in proptest::collection::vec(any::<u8>(), 64..4096),
) {
let config = FaultConfig {
partial_write_probability: 0.5,
..FaultConfig::none()
};
let dir = Path::new("/test");
let path = Path::new("/test/partial.dat");
let sim = sim_with_dir(seed, config);
let fd = sim.open(path, OpenOptions::read_write()).unwrap();
sim.sync_dir(dir).unwrap();
sim.write_all_at(fd, 0, &data).unwrap();
sim.sync(fd).unwrap();
sim.crash();
let fd = sim.open(path, OpenOptions::read()).unwrap();
let mut buf = vec![0u8; data.len()];
sim.read_exact_at(fd, 0, &mut buf).unwrap();
prop_assert_eq!(&buf, &data, "write_all_at must produce complete writes");
}
#[test]
fn dir_sync_required_for_file_survival(
seed in any::<u64>(),
data in proptest::collection::vec(any::<u8>(), 1..256),
) {
let sim = SimulatedIO::pristine(seed);
sim.mkdir(Path::new("/ephemeral")).unwrap();
let fd = sim.open(
Path::new("/ephemeral/file.dat"),
OpenOptions::read_write(),
).unwrap();
sim.write_all_at(fd, 0, &data).unwrap();
sim.sync(fd).unwrap();
sim.crash();
let result = sim.open(Path::new("/ephemeral/file.dat"), OpenOptions::read());
prop_assert!(result.is_err(), "file must vanish without dir sync");
}
}
+4 -2
View File
@@ -93,7 +93,8 @@ pub async fn list_blobs(
let cids_result: Result<Vec<String>, _> = if let Some(since) = &params.since {
state
.repos.blob
.repos
.blob
.list_blobs_since_rev(&did, since)
.await
.map(|cids| {
@@ -107,7 +108,8 @@ pub async fn list_blobs(
})
} else {
state
.repos.blob
.repos
.blob
.list_blobs_by_user(user_id, Some(cursor_cid), limit + 1)
.await
.map(|cids| cids.into_iter().map(|c| c.to_string()).collect())
+2 -1
View File
@@ -104,7 +104,8 @@ pub async fn list_repos(
let cursor_did: Option<Did> = params.cursor.as_ref().and_then(|s| s.parse().ok());
let cursor_ref = cursor_did.as_ref();
let result = state
.repos.repo
.repos
.repo
.list_repos_paginated(cursor_ref, limit + 1)
.await;
match result {
+10 -5
View File
@@ -8,7 +8,8 @@ static LAST_BROADCAST_SEQ: AtomicI64 = AtomicI64::new(0);
pub async fn start_sequencer_listener(state: AppState) {
let initial_seq = state
.repos.repo
.repos
.repo
.get_max_seq()
.await
.unwrap_or(SequenceNumber::ZERO);
@@ -30,14 +31,16 @@ pub async fn start_sequencer_listener(state: AppState) {
async fn listen_loop(state: AppState) -> anyhow::Result<()> {
let mut receiver = state
.repos.event_notifier
.repos
.event_notifier
.subscribe()
.await
.map_err(|e| anyhow::anyhow!("Failed to subscribe to events: {:?}", e))?;
info!("Connected to database and listening for repo updates");
let catchup_start = SequenceNumber::from_raw(LAST_BROADCAST_SEQ.load(Ordering::SeqCst));
let events = state
.repos.repo
.repos
.repo
.get_events_since_seq(catchup_start, None)
.await
.map_err(|e| anyhow::anyhow!("Failed to fetch catchup events: {:?}", e))?;
@@ -70,7 +73,8 @@ async fn listen_loop(state: AppState) -> anyhow::Result<()> {
}
if seq_id > last_seq + 1 {
let gap_events = state
.repos.repo
.repos
.repo
.get_events_in_seq_range(
SequenceNumber::from_raw(last_seq),
SequenceNumber::from_raw(seq_id),
@@ -88,7 +92,8 @@ async fn listen_loop(state: AppState) -> anyhow::Result<()> {
}
}
let event = state
.repos.repo
.repos
.repo
.get_event_by_seq(SequenceNumber::from_raw(seq_id))
.await
.ok()
+2 -1
View File
@@ -179,7 +179,8 @@ async fn get_repo_since(state: &AppState, did: &Did, head_cid: &Cid, since: &str
};
let block_cid_bytes = match state
.repos.repo
.repos
.repo
.get_user_block_cids_since_rev(user_id, since)
.await
{
+4
View File
@@ -19,6 +19,9 @@ fmt-check:
cargo fmt -- --check
lint: fmt-check clippy
test-store:
SQLX_OFFLINE=true cargo nextest run -p tranquil-store --features tranquil-store/test-harness
test-unit:
SQLX_OFFLINE=true cargo test --test dpop_unit --test validation_edge_cases --test scope_edge_cases
@@ -50,6 +53,7 @@ test-misc:
./scripts/run-tests.sh --test actor --test commit_signing --test image_processing --test lifecycle_social --test notifications --test server --test signing_key --test verify_live_commit
test *args:
@just test-store
@just test-unit
./scripts/run-tests.sh {{args}}