From 3868c1134d67519efcd928a25cad3828bd92bc27 Mon Sep 17 00:00:00 2001 From: "benjamin.747" Date: Fri, 28 Aug 2026 16:37:58 +0800 Subject: [PATCH 1/3] perf(import): cut receive-pack filepath writes and root-lock wait Batch blob file_path updates, hold the monorepo root lock only for the CAS attach (10ms Redis retry), and replace git_repo LIKE seq scans with index range / parent-walk lookups. crates-sync ETA now uses push rate instead of skip-inflated throughput. --- Cargo.lock | 104 ++--- .../src/application/api_service/mono/sync.rs | 21 +- .../code_edit/post_receive/import.rs | 32 +- ceres/src/transport/pack/import_repo.rs | 214 +++++++-- ceres/src/transport/pack/monorepo.rs | 134 +++--- jupiter/src/redis/lock.rs | 6 +- jupiter/src/storage/git_db_storage.rs | 419 ++++++++++++++++-- jupiter/src/storage/mono_storage.rs | 119 ++++- scripts/crates-sync/crates-sync.py | 32 +- 9 files changed, 879 insertions(+), 202 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8dc41e619..44db20ab6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -42,13 +42,13 @@ dependencies = [ [[package]] name = "aes" -version = "0.9.2" +version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8eb277bec05f56a0e0591f155a484cbd0f4f07ff2905051a48c72f004f7ed58" +checksum = "35f0f96ce78e38c3dc6d8948aa8163d06385be74000f3c7a95bf1eef35d3ea32" dependencies = [ "cipher 0.5.2", "cpubits", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "zeroize", ] @@ -73,7 +73,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f2b8006a0c83f52b62ba44a97b58bf76fe2f70a329e588f67f89691d93d498f" dependencies = [ "aead 0.6.1", - "aes 0.9.2", + "aes 0.9.3", "cipher 0.5.2", "ctr 0.10.1", "ctutils", @@ -292,13 +292,13 @@ dependencies = [ [[package]] name = "argon2" -version = "0.6.0-rc.8" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7af50940b73bf4e16c15c448a2b121c63f2d68e3e54b6a8731673cb4aa0cdff5" +checksum = "134c52ddac6d63c576bef8168db10c83c49c26444ecbc68060fef078925a901c" dependencies = [ "base64ct", "blake2 0.11.0", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "password-hash 0.6.1", ] @@ -1059,7 +1059,7 @@ dependencies = [ "cc", "cfg-if", "constant_time_eq", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", ] [[package]] @@ -1491,13 +1491,13 @@ checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" [[package]] name = "chacha20" -version = "0.10.1" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", "cipher 0.5.2", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "rand_core 0.10.1", "zeroize", ] @@ -1844,9 +1844,9 @@ dependencies = [ [[package]] name = "cpufeatures" -version = "0.3.0" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" dependencies = [ "libc", ] @@ -2076,7 +2076,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b5eed333089e2e1c1ac8c6c0398e5e2497b4c9926ca6d0365ed1e099afa5bc23" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "curve25519-dalek-derive", "digest 0.11.3", "fiat-crypto 0.3.0", @@ -2977,9 +2977,9 @@ checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" [[package]] name = "flate2" -version = "1.1.9" +version = "1.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" dependencies = [ "crc32fast", "miniz_oxide", @@ -3375,7 +3375,7 @@ dependencies = [ "futures-core", "futures-sink", "http", - "indexmap 2.14.0", + "indexmap 2.14.1", "slab", "tokio", "tokio-util", @@ -3483,7 +3483,7 @@ checksum = "48af7144c49a8db969e8a9d00cd470e1a446a3a73f6fa5eafc1eeb3d44d61ff4" dependencies = [ "hcl-edit", "hcl-primitives", - "indexmap 2.14.0", + "indexmap 2.14.1", "itoa", "serde", "vecmap-rs", @@ -3671,9 +3671,9 @@ dependencies = [ [[package]] name = "hyper" -version = "1.11.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d22053281f852e11534f5198498373cbb59295120a20771d90f7ed1897490a72" +checksum = "27b501faa50e7a26c3d3560ca625132f4078a17771f4810baf70475ae48cbe43" dependencies = [ "atomic-waker", "bytes", @@ -3929,9 +3929,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.14.0" +version = "2.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +checksum = "07aa2048142242915a31d35844fb311e0e53fcca590c3a0a40dcf1b841fa09eb" dependencies = [ "equivalent", "hashbrown 0.17.1", @@ -4253,7 +4253,7 @@ dependencies = [ "hex", "hmac 0.13.0", "idgenerator", - "indexmap 2.14.0", + "indexmap 2.14.1", "io-orbit", "jupiter-migrate", "pgp 0.20.0", @@ -4322,7 +4322,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8f198d1db720e4940b5a493201d199d9f24f568f8f746bd13706243a2f71598" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", ] [[package]] @@ -4561,9 +4561,9 @@ dependencies = [ [[package]] name = "libredox" -version = "0.1.20" +version = "0.1.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28d0a00925a9f930d679b6789b721e3a7f9ed110f41b86d2497caa780c3a070a" +checksum = "d7955dfc218a8afb29dfeffd540e3a6e96baeb94fe7138228dd7cc6937fbbf96" dependencies = [ "bitflags 2.13.1", "libc", @@ -4898,9 +4898,9 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" [[package]] name = "miniz_oxide" -version = "0.8.9" +version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" dependencies = [ "adler2", "simd-adler32", @@ -5141,7 +5141,7 @@ version = "0.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e493c05128df7a83b9676b709d590e0ebc285c7ed3152bc679668e8c1e506af5" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "memchr", "rustc-hash", "serde", @@ -5984,7 +5984,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" dependencies = [ "fixedbitset", - "indexmap 2.14.0", + "indexmap 2.14.1", ] [[package]] @@ -6213,7 +6213,7 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63d440a804ec8d6fafbb6b84471e013286658d373248927692ab3366686220ca" dependencies = [ - "aes 0.9.2", + "aes 0.9.3", "aes-gcm 0.11.1", "cbc", "der 0.8.1", @@ -6274,7 +6274,7 @@ version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e2d0073b297041425c7c3df6eb4792d598a15323fe63346852b092eca02904c" dependencies = [ - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "universal-hash 0.6.1", "zeroize", ] @@ -6298,7 +6298,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0fa31d631f2b2cb2a544d0aa321ce847a94764d701ca2becc411138b93d49cd" dependencies = [ "cpubits", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "universal-hash 0.6.1", "zeroize", ] @@ -6402,7 +6402,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "93980406f12d9f8140ed5abe7155acb10bb1e69ea55c88960b9c2f117445ef96" dependencies = [ "equivalent", - "indexmap 2.14.0", + "indexmap 2.14.1", "serde", ] @@ -7134,7 +7134,7 @@ dependencies = [ "bytecheck 0.8.3", "bytes", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap 2.14.1", "munge", "ptr_meta 0.3.2", "rancor", @@ -7226,7 +7226,7 @@ version = "0.62.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9decb68e4e44e1079700e54f17c8f23806ec53d7e0db73ab1c71d9dabc666812" dependencies = [ - "aes 0.9.2", + "aes 0.9.3", "aws-lc-rs", "bitflags 2.13.1", "block-padding 0.4.2", @@ -7298,7 +7298,7 @@ version = "0.63.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35bab1b87d915817d5d9cc352637cd40d5f0b298a48c6309af9156a4addc3031" dependencies = [ - "aes 0.9.2", + "aes 0.9.3", "aws-lc-rs", "bitflags 2.13.1", "block-padding 0.4.2", @@ -8064,7 +8064,7 @@ version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "itoa", "memchr", "serde", @@ -8124,7 +8124,7 @@ dependencies = [ "chrono", "hex", "indexmap 1.9.3", - "indexmap 2.14.0", + "indexmap 2.14.1", "jiff", "schemars 0.9.0", "schemars 1.2.2", @@ -8152,7 +8152,7 @@ version = "0.9.34+deprecated" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "itoa", "ryu", "serde", @@ -8242,7 +8242,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "digest 0.11.3", ] @@ -8281,7 +8281,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" dependencies = [ "cfg-if", - "cpufeatures 0.3.0", + "cpufeatures 0.3.1", "digest 0.11.3", ] @@ -8577,7 +8577,7 @@ dependencies = [ "futures-util", "hashbrown 0.16.1", "hashlink", - "indexmap 2.14.0", + "indexmap 2.14.1", "log", "memchr", "percent-encoding", @@ -8748,7 +8748,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d801accda99469cde6d73da741422610fdf6508a72d9a69d1b55cb241c720597" dependencies = [ "aead 0.6.1", - "aes 0.9.2", + "aes 0.9.3", "aes-gcm 0.11.1", "chacha20", "cipher 0.5.2", @@ -8813,7 +8813,7 @@ version = "0.7.0-rc.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9a32fae177b74a22aa9c5b01bf7e68b33545be32d9e381e248058d2adc15ce3" dependencies = [ - "argon2 0.6.0-rc.8", + "argon2 0.6.0", "bcrypt-pbkdf", "ctutils", "ed25519-dalek 3.0.0", @@ -9429,7 +9429,7 @@ version = "0.9.12+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "serde_core", "serde_spanned 1.1.1", "toml_datetime 0.7.5+spec-1.1.0", @@ -9444,7 +9444,7 @@ version = "1.1.4+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3aace63f4bbcdfc2c965b059de67119c89c4017a70d633be6c104910f67056f5" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "serde_core", "serde_spanned 1.1.1", "toml_datetime 1.1.1+spec-1.1.0", @@ -9486,7 +9486,7 @@ version = "0.22.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "serde", "serde_spanned 0.6.9", "toml_datetime 0.6.11", @@ -9500,7 +9500,7 @@ version = "0.25.13+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "toml_datetime 1.1.1+spec-1.1.0", "toml_parser", "winnow 1.0.4", @@ -9564,7 +9564,7 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", - "indexmap 2.14.0", + "indexmap 2.14.1", "pin-project-lite", "slab", "sync_wrapper", @@ -10058,7 +10058,7 @@ version = "5.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8bde15df68e80b16c7d16b9616e80770ad158988daa56a27dccd1e55558b0160" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "serde", "serde_json", "utoipa-gen", @@ -11001,7 +11001,7 @@ dependencies = [ "arbitrary", "crc32fast", "flate2", - "indexmap 2.14.0", + "indexmap 2.14.1", "memchr", "zopfli", ] diff --git a/ceres/src/application/api_service/mono/sync.rs b/ceres/src/application/api_service/mono/sync.rs index 9dcd64283..c7257b495 100644 --- a/ceres/src/application/api_service/mono/sync.rs +++ b/ceres/src/application/api_service/mono/sync.rs @@ -37,7 +37,8 @@ impl MonoApiService { )); for attempt in 0..MAX_ATTACH_ATTEMPTS { - let guard = redlock.clone().lock().await?; + // Same split as import receive-pack attach: walk/upload off the + // global root lock; CAS on mega_refs still serializes the root update. let root_ref = storage .get_main_ref("/") .await? @@ -68,6 +69,24 @@ impl MonoApiService { .save_blobs(&new_commit.id.to_string(), vec![gitkeep_blob]) .await?; + let guard = redlock.clone().lock().await?; + let current_root = storage + .get_main_ref("/") + .await? + .ok_or_else(|| MegaError::Other("root ref not found".to_string()))?; + if current_root.ref_commit_hash != expected_commit + || current_root.ref_tree_hash != expected_tree + || current_root.id != root_ref_id + { + let _ = guard.unlock().await; + tracing::warn!( + attempt, + repo_path = %path, + "attach_project_path_to_monorepo_root: root ref moved before txn, retrying" + ); + continue; + } + let txn = self.storage().begin_db_transaction().await?; match storage .attach_to_monorepo_parent_in_txn( diff --git a/ceres/src/application/code_edit/post_receive/import.rs b/ceres/src/application/code_edit/post_receive/import.rs index 6203e5b0a..1b250f0d5 100644 --- a/ceres/src/application/code_edit/post_receive/import.rs +++ b/ceres/src/application/code_edit/post_receive/import.rs @@ -67,12 +67,9 @@ pub async fn dispatch_import_receive_pack_finalized( let mut root_lock_wait_sum_ms: u128 = 0; for attempt in 0..MAX_ATTACH_ATTEMPTS { - let t_lock = Instant::now(); - let guard = unpack_redlock.clone().lock().await?; - let lock_wait_ms = t_lock.elapsed().as_millis(); - root_lock_wait_max_ms = root_lock_wait_max_ms.max(lock_wait_ms); - root_lock_wait_sum_ms += lock_wait_ms; - + // Tree walk + .gitkeep upload are the expensive part of attach (~20ms of + // the former 25ms hold). Do them without the global root lock; the CAS + // on mega_refs still rejects a stale snapshot. let root_ref = mono_storage .get_main_ref("/") .await? @@ -99,6 +96,29 @@ pub async fn dispatch_import_receive_pack_finalized( .save_blobs(&new_commit.id.to_string(), vec![gitkeep_blob]) .await?; + let t_lock = Instant::now(); + let guard = unpack_redlock.clone().lock().await?; + let lock_wait_ms = t_lock.elapsed().as_millis(); + root_lock_wait_max_ms = root_lock_wait_max_ms.max(lock_wait_ms); + root_lock_wait_sum_ms += lock_wait_ms; + + let current_root = mono_storage + .get_main_ref("/") + .await? + .ok_or_else(|| MegaError::Other("root ref not found".to_string()))?; + if current_root.ref_commit_hash != expected_commit + || current_root.ref_tree_hash != expected_tree + || current_root.id != root_ref_id + { + let _ = guard.unlock().await; + tracing::warn!( + attempt = attempt, + repo_path = %repo_path.display(), + "attach_to_monorepo_parent: root ref moved before txn, retrying" + ); + continue; + } + let txn = storage.begin_db_transaction().await?; let git_db = storage.git_db_storage(); for &cmd in &branch_cmds { diff --git a/ceres/src/transport/pack/import_repo.rs b/ceres/src/transport/pack/import_repo.rs index ee714effb..cfcaa3291 100644 --- a/ceres/src/transport/pack/import_repo.rs +++ b/ceres/src/transport/pack/import_repo.rs @@ -472,44 +472,50 @@ impl RepoHandler for ImportRepo { .unwrap() .clone(), ); - self.traverses_and_update_filepath(root_tree, PathBuf::new()) + let pairs = collect_git_blob_filepaths( + self.storage.git_db_storage(), + self.repo.repo_id, + root_tree, + PathBuf::new(), + ) + .await?; + self.storage + .git_db_storage() + .update_git_blob_filepaths(self.repo.repo_id, pairs) .await?; Ok(()) } } -impl ImportRepo { - #[async_recursion] - async fn traverses_and_update_filepath( - &self, - tree: Tree, - path: PathBuf, - ) -> Result<(), MegaError> { - for item in tree.tree_items { - if item.is_tree() { - let tree = Tree::from_git_model( - self.storage - .git_db_storage() - .get_tree_by_hash(self.repo.repo_id, &item.id.to_string()) - .await? - .unwrap() - .clone(), - ); - - // 递归调用 - self.traverses_and_update_filepath(tree, path.join(item.name)) - .await?; - } else { - let id = item.id.to_string(); - self.storage - .git_db_storage() - .update_git_blob_filepath(&id, path.join(item.name).to_str().unwrap()) - .await?; - } +#[async_recursion] +pub(crate) async fn collect_git_blob_filepaths( + storage: GitDbStorage, + repo_id: i64, + tree: Tree, + path: PathBuf, +) -> Result, MegaError> { + let mut pairs = Vec::new(); + for item in tree.tree_items { + if item.is_tree() { + let child = Tree::from_git_model( + storage + .get_tree_by_hash(repo_id, &item.id.to_string()) + .await? + .unwrap() + .clone(), + ); + pairs.extend( + collect_git_blob_filepaths(storage.clone(), repo_id, child, path.join(item.name)) + .await?, + ); + } else { + pairs.push(( + item.id.to_string(), + path.join(item.name).to_str().unwrap().to_string(), + )); } - - Ok(()) } + Ok(pairs) } async fn process_objects( @@ -604,6 +610,23 @@ async fn process_objects( #[cfg(test)] mod test { use std::path::PathBuf; + + use callisto::{git_blob, git_tree}; + use git_internal::internal::object::{ + ObjectTrait, + blob::Blob, + tree::{Tree, TreeItem, TreeItemMode}, + }; + use jupiter::{ + sea_orm::{ColumnTrait, EntityTrait, IntoActiveModel, QueryFilter}, + storage::base_storage::StorageConnector, + tests::test_storage, + utils::converter::FromGitModel, + }; + use tempfile::TempDir; + + use super::collect_git_blob_filepaths; + #[test] pub fn test_recurse_tree() { let path = PathBuf::from("/third-party/crates/tokio/tokio-console"); @@ -612,4 +635,133 @@ mod test { println!("{path:?}"); } } + + #[tokio::test] + async fn collect_and_batch_update_nested_crate_tree() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + let repo_id = 11i64; + + let cargo = Blob::from_content("[package]\nname = \"demo\"\n"); + let lib = Blob::from_content("pub fn f() {}\n"); + let src_tree = Tree::from_tree_items(vec![TreeItem { + mode: TreeItemMode::Blob, + id: lib.id, + name: "lib.rs".into(), + }]) + .unwrap(); + let root_tree = Tree::from_tree_items(vec![ + TreeItem { + mode: TreeItemMode::Blob, + id: cargo.id, + name: "Cargo.toml".into(), + }, + TreeItem { + mode: TreeItemMode::Tree, + id: src_tree.id, + name: "src".into(), + }, + ]) + .unwrap(); + + let now = chrono::Utc::now().naive_utc(); + git_blob::Entity::insert_many([ + git_blob::Model { + id: 1, + repo_id, + blob_id: cargo.id.to_string(), + name: None, + size: 0, + created_at: now, + pack_id: String::new(), + file_path: String::new(), + pack_offset: 0, + is_delta_in_pack: false, + } + .into_active_model(), + git_blob::Model { + id: 2, + repo_id, + blob_id: lib.id.to_string(), + name: None, + size: 0, + created_at: now, + pack_id: String::new(), + file_path: String::new(), + pack_offset: 0, + is_delta_in_pack: false, + } + .into_active_model(), + ]) + .exec(stg.get_connection()) + .await + .unwrap(); + + git_tree::Entity::insert_many([ + git_tree::Model { + id: 3, + repo_id, + tree_id: src_tree.id.to_string(), + sub_trees: src_tree.to_data().unwrap(), + size: 0, + created_at: now, + pack_id: String::new(), + pack_offset: 0, + } + .into_active_model(), + git_tree::Model { + id: 4, + repo_id, + tree_id: root_tree.id.to_string(), + sub_trees: root_tree.to_data().unwrap(), + size: 0, + created_at: now, + pack_id: String::new(), + pack_offset: 0, + } + .into_active_model(), + ]) + .exec(stg.get_connection()) + .await + .unwrap(); + + let loaded_root = Tree::from_git_model( + stg.get_tree_by_hash(repo_id, &root_tree.id.to_string()) + .await + .unwrap() + .unwrap(), + ); + let mut pairs = + collect_git_blob_filepaths(stg.clone(), repo_id, loaded_root, PathBuf::new()) + .await + .unwrap(); + pairs.sort_by(|a, b| a.1.cmp(&b.1)); + assert_eq!( + pairs, + vec![ + (cargo.id.to_string(), "Cargo.toml".into()), + (lib.id.to_string(), "src/lib.rs".into()), + ] + ); + + stg.update_git_blob_filepaths(repo_id, pairs).await.unwrap(); + + let cargo_row = git_blob::Entity::find() + .filter(git_blob::Column::RepoId.eq(repo_id)) + .filter(git_blob::Column::BlobId.eq(cargo.id.to_string())) + .one(stg.get_connection()) + .await + .unwrap() + .unwrap(); + let lib_row = git_blob::Entity::find() + .filter(git_blob::Column::RepoId.eq(repo_id)) + .filter(git_blob::Column::BlobId.eq(lib.id.to_string())) + .one(stg.get_connection()) + .await + .unwrap() + .unwrap(); + assert_eq!(cargo_row.file_path, "Cargo.toml"); + assert_eq!(lib_row.file_path, "src/lib.rs"); + } } diff --git a/ceres/src/transport/pack/monorepo.rs b/ceres/src/transport/pack/monorepo.rs index 1778f686e..4f305dc92 100644 --- a/ceres/src/transport/pack/monorepo.rs +++ b/ceres/src/transport/pack/monorepo.rs @@ -30,7 +30,11 @@ use git_internal::{ }, }; use io_orbit::object_storage::MultiObjectByteStream; -use jupiter::{sea_orm::DatabaseTransaction, storage::Storage, utils::converter::FromMegaModel}; +use jupiter::{ + sea_orm::DatabaseTransaction, + storage::{Storage, mono_storage::MonoStorage}, + utils::converter::FromMegaModel, +}; use tokio::sync::{RwLock, mpsc}; use tokio_stream::wrappers::ReceiverStream; @@ -490,7 +494,18 @@ impl RepoHandler for MonoRepo { commit_opt.tree_id ); - self.traverses_and_update_filepath(root_tree, PathBuf::new()) + let pairs = + collect_mega_blob_filepaths(self.storage.mono_storage(), root_tree, PathBuf::new()) + .await + .map_err(|e| { + MegaError::Other(format!( + "Failed to update file paths for commit {}: {}", + commit_opt.id, e + )) + })?; + self.storage + .mono_storage() + .update_blob_filepaths(pairs) .await .map_err(|e| { MegaError::Other(format!( @@ -571,43 +586,41 @@ impl MonoRepo { } } -impl MonoRepo { - #[async_recursion] - async fn traverses_and_update_filepath( - &self, - tree: Tree, - path: PathBuf, - ) -> Result<(), MegaError> { - for item in tree.tree_items { - let item_path = path.join(&item.name); - - if item.is_tree() { - let tree_hash = item.id.to_string(); - let trees = self - .storage - .mono_storage() - .get_trees_by_hashes(vec![tree_hash.clone()]) - .await - .map_err(|e| { - MegaError::Other(format!( - "Failed to retrieve tree {} at path '{}': {}", - tree_hash, - item_path.display(), - e - )) - })?; - - if trees.is_empty() { - return Err(MegaError::Other(format!( - "Tree {} not found at path '{}'", +#[async_recursion] +async fn collect_mega_blob_filepaths( + storage: MonoStorage, + tree: Tree, + path: PathBuf, +) -> Result, MegaError> { + let mut pairs = Vec::new(); + for item in tree.tree_items { + let item_path = path.join(&item.name); + + if item.is_tree() { + let tree_hash = item.id.to_string(); + let trees = storage + .get_trees_by_hashes(vec![tree_hash.clone()]) + .await + .map_err(|e| { + MegaError::Other(format!( + "Failed to retrieve tree {} at path '{}': {}", tree_hash, - item_path.display() - ))); - } + item_path.display(), + e + )) + })?; - let child_tree = Tree::from_mega_model(trees[0].clone()); + if trees.is_empty() { + return Err(MegaError::Other(format!( + "Tree {} not found at path '{}'", + tree_hash, + item_path.display() + ))); + } - self.traverses_and_update_filepath(child_tree, item_path.clone()) + let child_tree = Tree::from_mega_model(trees[0].clone()); + pairs.extend( + collect_mega_blob_filepaths(storage.clone(), child_tree, item_path.clone()) .await .map_err(|e| { MegaError::Other(format!( @@ -616,38 +629,29 @@ impl MonoRepo { item_path.display(), e )) - })?; - } else { - let blob_id = item.id.to_string(); - let file_path_str = item_path.to_str().ok_or_else(|| { - MegaError::Other(format!( - "Invalid UTF-8 path for blob {}: '{}'", - blob_id, - item_path.display() - )) - })?; - - self.storage - .mono_storage() - .update_blob_filepath(&blob_id, file_path_str) - .await - .map_err(|e| { - MegaError::Other(format!( - "Failed to update file path for blob {} at '{}': {}", - blob_id, file_path_str, e - )) - })?; - - tracing::debug!( - "Updated file path for blob {} to '{}'", + })?, + ); + } else { + let blob_id = item.id.to_string(); + let file_path_str = item_path.to_str().ok_or_else(|| { + MegaError::Other(format!( + "Invalid UTF-8 path for blob {}: '{}'", blob_id, - file_path_str - ); - } + item_path.display() + )) + })?; + tracing::debug!( + "Queued file path for blob {} to '{}'", + blob_id, + file_path_str + ); + pairs.push((blob_id, file_path_str.to_string())); } - - Ok(()) } + Ok(pairs) +} + +impl MonoRepo { async fn fetch_or_new_cl_link(&self) -> Result { let storage = self.storage.cl_storage(); let path_str = self.path.to_str().unwrap(); diff --git a/jupiter/src/redis/lock.rs b/jupiter/src/redis/lock.rs index 67f2c284d..ddeea0ca4 100644 --- a/jupiter/src/redis/lock.rs +++ b/jupiter/src/redis/lock.rs @@ -50,11 +50,15 @@ impl RedLock { Ok(result.is_some()) } + /// Retry interval while waiting for SET NX. A 200ms sleep turned a ~25ms + /// hold into 200/400/600ms waits for anyone who missed the unlock. + const LOCK_RETRY_SLEEP: Duration = Duration::from_millis(10); + /// Lock with retry pub async fn lock(self: Arc) -> Result { let t0 = Instant::now(); while !self.try_lock().await? { - sleep(Duration::from_millis(200)).await; + sleep(Self::LOCK_RETRY_SLEEP).await; } self.spawn_auto_renew(); diff --git a/jupiter/src/storage/git_db_storage.rs b/jupiter/src/storage/git_db_storage.rs index f0580f71c..c7cdfdcb2 100644 --- a/jupiter/src/storage/git_db_storage.rs +++ b/jupiter/src/storage/git_db_storage.rs @@ -1,4 +1,4 @@ -use std::ops::Deref; +use std::{collections::HashMap, ops::Deref}; use api_model::common::Pagination; use callisto::{ @@ -11,9 +11,9 @@ use common::{ }; use futures::Stream; use sea_orm::{ - ActiveModelTrait, ColumnTrait, ConnectionTrait, DatabaseTransaction, DbBackend, DbErr, - EntityTrait, IntoActiveModel, PaginatorTrait, QueryFilter, QueryOrder, QueryTrait, Set, - TransactionTrait, sea_query::Expr, + ActiveModelTrait, ColumnTrait, ConnectionTrait, DatabaseTransaction, DbErr, EntityTrait, + IntoActiveModel, PaginatorTrait, QueryFilter, QueryOrder, Set, TransactionTrait, + sea_query::{CaseStatement, Expr, ExprTrait}, }; use crate::storage::base_storage::{BaseStorage, StorageConnector}; @@ -282,21 +282,46 @@ impl GitDbStorage { pub async fn update_git_blob_filepath( &self, - blob_id: &String, + repo_id: i64, + blob_id: &str, file_path: &str, ) -> Result<(), MegaError> { - if let Some(model) = git_blob::Entity::find() - .filter(git_blob::Column::BlobId.eq(blob_id)) - .one(self.get_connection()) - .await? - { - let mut active: git_blob::ActiveModel = model.into(); - - active.file_path = Set(file_path.to_string()); + self.update_git_blob_filepaths(repo_id, vec![(blob_id.to_string(), file_path.to_string())]) + .await + } - active.update(self.get_connection()).await?; + /// Batch-assign `file_path` for blobs in one repo. + /// + /// Duplicate `blob_id`s keep the last path (same as sequential UPDATE). Missing ids are + /// skipped. Empty input is a no-op. + pub async fn update_git_blob_filepaths( + &self, + repo_id: i64, + pairs: Vec<(String, String)>, + ) -> Result<(), MegaError> { + if pairs.is_empty() { + return Ok(()); } + let collapsed = last_wins_filepaths(pairs); + for chunk in collapsed.chunks(::BATCH_CHUNK_SIZE) { + let blob_ids: Vec = chunk.iter().map(|(id, _)| id.clone()).collect(); + let mut case = CaseStatement::new(); + for (blob_id, file_path) in chunk { + case = case.case( + Expr::col(git_blob::Column::BlobId).eq(blob_id.clone()), + file_path.clone(), + ); + } + case = case.finally(Expr::col(git_blob::Column::FilePath)); + + git_blob::Entity::update_many() + .col_expr(git_blob::Column::FilePath, case.into()) + .filter(git_blob::Column::RepoId.eq(repo_id)) + .filter(git_blob::Column::BlobId.is_in(blob_ids)) + .exec(self.get_connection()) + .await?; + } Ok(()) } @@ -335,8 +360,12 @@ impl GitDbStorage { } // Descendant of the new path (new path would become a parent import repo). + // Btree range `[path/, path0)` uses idx_ir_repo_path; `LIKE 'path/%'` does not + // (needs varchar_pattern_ops / C collation) and seq-scans ~1M rows. + let (lo, hi) = git_repo_descendant_bounds(path); if let Some(descendant) = git_repo::Entity::find() - .filter(git_repo::Column::RepoPath.like(format!("{path}/%"))) + .filter(git_repo::Column::RepoPath.gte(lo)) + .filter(git_repo::Column::RepoPath.lt(hi)) .one(self.get_connection()) .await? { @@ -363,26 +392,36 @@ impl GitDbStorage { Ok(None) } - /// Finds a Git repository with a path that matches the beginning of the provided repository path using a LIKE query. - /// - /// # Arguments - /// - /// * `repo_path` - A string slice that holds the beginning of the path of the repository to search for. - /// - /// # Returns + /// Longest import repo whose path is a segment-prefix of `repo_path`. /// - /// A `Result` containing an `Option` with the Git repository model if found, or `None` if not found. - /// Returns a `MegaError` if an error occurs during the search. + /// Walks parents with the unique `repo_path` index instead of + /// `'{path}' LIKE repo_path || '%'` (seq scan, and a false match on + /// `/third-party/rust` vs `/third-party/rust_v1`). pub async fn find_git_repo_like_path( &self, repo_path: &str, ) -> Result, MegaError> { - let query = git_repo::Entity::find() - .filter(Expr::cust(format!("'{repo_path}' LIKE repo_path || '%'"))) - .order_by_desc(Expr::cust("LENGTH(repo_path)")); - tracing::debug!("{}", query.build(DbBackend::Postgres).to_string()); - let result = query.one(self.get_connection()).await?; - Ok(result) + let path = repo_path.trim_end_matches('/'); + if path.is_empty() { + return Ok(None); + } + + let mut current = std::path::PathBuf::from(path); + loop { + let candidate = current.to_string_lossy(); + let candidate = if candidate.is_empty() { + "/".to_string() + } else { + candidate.to_string() + }; + if let Some(repo) = self.find_git_repo_exact_match(&candidate).await? { + return Ok(Some(repo)); + } + if candidate == "/" || !current.pop() { + break; + } + } + Ok(None) } pub async fn save_git_repo(&self, repo: git_repo::Model) -> Result<(), MegaError> { @@ -618,3 +657,323 @@ impl GitDbStorage { .unwrap() } } + +fn last_wins_filepaths(pairs: Vec<(String, String)>) -> Vec<(String, String)> { + let mut map = HashMap::with_capacity(pairs.len()); + for (blob_id, file_path) in pairs { + map.insert(blob_id, file_path); + } + map.into_iter().collect() +} + +/// Inclusive lower / exclusive upper bound for `repo_path` values that are +/// strict descendants of `path` (`path/…`). +/// +/// `'/'` is ASCII 47 and `'0'` is 48, so `[path/, path0)` is the btree range +/// of keys that start with `path/`. +fn git_repo_descendant_bounds(path: &str) -> (String, String) { + (format!("{path}/"), format!("{path}0")) +} + +#[cfg(test)] +mod tests { + use callisto::{git_blob, git_repo}; + use sea_orm::{ColumnTrait, EntityTrait, IntoActiveModel, QueryFilter}; + use tempfile::TempDir; + + use super::*; + use crate::tests::test_storage; + + fn blob_row(id: i64, repo_id: i64, blob_id: &str, file_path: &str) -> git_blob::Model { + git_blob::Model { + id, + repo_id, + blob_id: blob_id.to_string(), + name: None, + size: 0, + created_at: chrono::Utc::now().naive_utc(), + pack_id: String::new(), + file_path: file_path.to_string(), + pack_offset: 0, + is_delta_in_pack: false, + } + } + + async fn insert_blob(stg: &GitDbStorage, model: git_blob::Model) { + git_blob::Entity::insert(model.into_active_model()) + .exec(stg.get_connection()) + .await + .expect("insert git_blob"); + } + + async fn filepath_of(stg: &GitDbStorage, repo_id: i64, blob_id: &str) -> Option { + git_blob::Entity::find() + .filter(git_blob::Column::RepoId.eq(repo_id)) + .filter(git_blob::Column::BlobId.eq(blob_id)) + .one(stg.get_connection()) + .await + .unwrap() + .map(|m| m.file_path) + } + + #[tokio::test] + async fn update_git_blob_filepaths_empty_is_noop() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + stg.update_git_blob_filepaths(1, vec![]) + .await + .expect("empty batch"); + } + + #[tokio::test] + async fn update_git_blob_filepaths_sets_paths_in_one_repo() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + let repo_id = 7; + insert_blob(&stg, blob_row(1, repo_id, "blob-a", "")).await; + insert_blob(&stg, blob_row(2, repo_id, "blob-b", "")).await; + + stg.update_git_blob_filepaths( + repo_id, + vec![ + ("blob-a".into(), "src/lib.rs".into()), + ("blob-b".into(), "Cargo.toml".into()), + ], + ) + .await + .unwrap(); + + assert_eq!( + filepath_of(&stg, repo_id, "blob-a").await.as_deref(), + Some("src/lib.rs") + ); + assert_eq!( + filepath_of(&stg, repo_id, "blob-b").await.as_deref(), + Some("Cargo.toml") + ); + } + + #[tokio::test] + async fn update_git_blob_filepaths_is_scoped_to_repo() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + insert_blob(&stg, blob_row(1, 1, "shared-blob", "old-a")).await; + insert_blob(&stg, blob_row(2, 2, "shared-blob", "old-b")).await; + + stg.update_git_blob_filepaths(1, vec![("shared-blob".into(), "src/lib.rs".into())]) + .await + .unwrap(); + + assert_eq!( + filepath_of(&stg, 1, "shared-blob").await.as_deref(), + Some("src/lib.rs") + ); + assert_eq!( + filepath_of(&stg, 2, "shared-blob").await.as_deref(), + Some("old-b") + ); + } + + #[tokio::test] + async fn update_git_blob_filepaths_duplicate_blob_keeps_last_path() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + insert_blob(&stg, blob_row(1, 1, "dup", "")).await; + + stg.update_git_blob_filepaths( + 1, + vec![ + ("dup".into(), "first.rs".into()), + ("dup".into(), "last.rs".into()), + ], + ) + .await + .unwrap(); + + assert_eq!( + filepath_of(&stg, 1, "dup").await.as_deref(), + Some("last.rs") + ); + } + + #[tokio::test] + async fn update_git_blob_filepaths_ignores_missing_blob_id() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + stg.update_git_blob_filepaths(1, vec![("missing".into(), "nope.rs".into())]) + .await + .expect("missing blob is not an error"); + } + + #[tokio::test] + async fn update_git_blob_filepaths_chunks_above_batch_size() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + let repo_id = 3; + let n = ::BATCH_CHUNK_SIZE + 1; + let mut pairs = Vec::with_capacity(n); + for i in 0..n { + let blob_id = format!("b{i:04}"); + insert_blob(&stg, blob_row(i as i64 + 1, repo_id, &blob_id, "")).await; + pairs.push((blob_id, format!("f{i}.rs"))); + } + + stg.update_git_blob_filepaths(repo_id, pairs.clone()) + .await + .unwrap(); + + let last_id = format!("b{:04}", n - 1); + assert_eq!( + filepath_of(&stg, repo_id, &last_id).await.as_deref(), + Some(format!("f{}.rs", n - 1).as_str()) + ); + assert_eq!( + filepath_of(&stg, repo_id, "b0000").await.as_deref(), + Some("f0.rs") + ); + } + + #[tokio::test] + async fn update_git_blob_filepaths_binds_quotes_in_path() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + insert_blob(&stg, blob_row(1, 1, "q", "")).await; + stg.update_git_blob_filepaths(1, vec![("q".into(), "foo's/bar.rs".into())]) + .await + .unwrap(); + assert_eq!( + filepath_of(&stg, 1, "q").await.as_deref(), + Some("foo's/bar.rs") + ); + } + + async fn insert_repo(stg: &GitDbStorage, id: i64, path: &str) { + stg.save_git_repo(git_repo::Model { + id, + repo_path: path.to_string(), + repo_name: path.rsplit('/').next().unwrap_or(path).to_string(), + created_at: chrono::Utc::now().naive_utc(), + updated_at: chrono::Utc::now().naive_utc(), + }) + .await + .expect("insert git_repo"); + } + + #[test] + fn git_repo_descendant_bounds_are_prefix_range() { + assert_eq!( + git_repo_descendant_bounds("/third-party/rust"), + ( + "/third-party/rust/".to_string(), + "/third-party/rust0".to_string() + ) + ); + let (lo, hi) = git_repo_descendant_bounds("/third-party/rust"); + let crate_path = "/third-party/rust/crates/sw/ay/swayws/1.3.0"; + assert!(crate_path >= lo.as_str() && crate_path < hi.as_str()); + assert!(!("/third-party/rust_v1" >= lo.as_str() && "/third-party/rust_v1" < hi.as_str())); + assert!(!("/third-party/rust" >= lo.as_str() && "/third-party/rust" < hi.as_str())); + } + + #[tokio::test] + async fn nested_conflict_finds_descendant_not_sibling_prefix() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + insert_repo(&stg, 1, "/third-party/rust/crates/sw/ay/swayws/1.3.0").await; + insert_repo(&stg, 2, "/third-party/rust_v1").await; + + let hit = stg + .find_nested_import_repo_conflict("/third-party/rust") + .await + .unwrap() + .expect("descendant is a conflict"); + assert_eq!(hit.repo_path, "/third-party/rust/crates/sw/ay/swayws/1.3.0"); + + assert!( + stg.find_nested_import_repo_conflict("/third-party/rust_v1") + .await + .unwrap() + .is_none() + ); + assert!( + stg.find_nested_import_repo_conflict("/third-party/foo") + .await + .unwrap() + .is_none() + ); + assert!( + stg.find_nested_import_repo_conflict("/third-party/rust/crates/sw/ay/swayws/1.3.0") + .await + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn nested_conflict_finds_ancestor() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + insert_repo(&stg, 1, "/third-party/rust").await; + + let hit = stg + .find_nested_import_repo_conflict("/third-party/rust/crates/to/ki/tokio/1.0.0") + .await + .unwrap() + .expect("ancestor is a conflict"); + assert_eq!(hit.repo_path, "/third-party/rust"); + } + + #[tokio::test] + async fn like_path_walks_to_longest_segment_prefix() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + insert_repo(&stg, 1, "/third-party/rust").await; + insert_repo(&stg, 2, "/third-party/rust/crates/foo/1.0.0").await; + + let exact = stg + .find_git_repo_like_path("/third-party/rust/crates/foo/1.0.0") + .await + .unwrap() + .unwrap(); + assert_eq!(exact.repo_path, "/third-party/rust/crates/foo/1.0.0"); + + let under = stg + .find_git_repo_like_path("/third-party/rust/crates/foo/1.0.0/src/lib.rs") + .await + .unwrap() + .unwrap(); + assert_eq!(under.repo_path, "/third-party/rust/crates/foo/1.0.0"); + + let parent = stg + .find_git_repo_like_path("/third-party/rust/crates/bar/2.0.0") + .await + .unwrap() + .unwrap(); + assert_eq!(parent.repo_path, "/third-party/rust"); + } + + #[tokio::test] + async fn like_path_does_not_match_string_prefix_sibling() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + insert_repo(&stg, 1, "/third-party/rust").await; + + assert!( + stg.find_git_repo_like_path("/third-party/rust_v1") + .await + .unwrap() + .is_none() + ); + } +} diff --git a/jupiter/src/storage/mono_storage.rs b/jupiter/src/storage/mono_storage.rs index 6ea4627a8..687849210 100644 --- a/jupiter/src/storage/mono_storage.rs +++ b/jupiter/src/storage/mono_storage.rs @@ -16,7 +16,7 @@ use sea_orm::{ ActiveValue::Set, ColumnTrait, Condition, ConnectionTrait, DatabaseTransaction, DbErr, EntityTrait, IntoActiveModel, PaginatorTrait, QueryFilter, QueryOrder, QuerySelect, TransactionTrait, - sea_query::{Expr, OnConflict}, + sea_query::{CaseStatement, Expr, ExprTrait, OnConflict}, }; use crate::{ @@ -455,18 +455,38 @@ impl MonoStorage { blob_id: &str, file_path: &str, ) -> Result<(), MegaError> { - if let Some(model) = mega_blob::Entity::find() - .filter(mega_blob::Column::BlobId.eq(blob_id)) - .one(self.get_connection()) - .await? - { - let mut active: mega_blob::ActiveModel = model.into(); - - active.file_path = Set(file_path.to_string()); + self.update_blob_filepaths(vec![(blob_id.to_string(), file_path.to_string())]) + .await + } - active.update(self.get_connection()).await?; + /// Batch-assign `file_path` on `mega_blob`. Duplicate `blob_id`s keep the last path. + /// Missing ids are skipped. Empty input is a no-op. + pub async fn update_blob_filepaths( + &self, + pairs: Vec<(String, String)>, + ) -> Result<(), MegaError> { + if pairs.is_empty() { + return Ok(()); } + let collapsed = last_wins_mega_filepaths(pairs); + for chunk in collapsed.chunks(::BATCH_CHUNK_SIZE) { + let blob_ids: Vec = chunk.iter().map(|(id, _)| id.clone()).collect(); + let mut case = CaseStatement::new(); + for (blob_id, file_path) in chunk { + case = case.case( + Expr::col(mega_blob::Column::BlobId).eq(blob_id.clone()), + file_path.clone(), + ); + } + case = case.finally(Expr::col(mega_blob::Column::FilePath)); + + mega_blob::Entity::update_many() + .col_expr(mega_blob::Column::FilePath, case.into()) + .filter(mega_blob::Column::BlobId.is_in(blob_ids)) + .exec(self.get_connection()) + .await?; + } Ok(()) } @@ -888,10 +908,13 @@ impl MonoStorage { #[cfg(test)] mod tests { + use callisto::mega_blob; use common::utils; + use sea_orm::{ColumnTrait, EntityTrait, IntoActiveModel, QueryFilter}; use tempfile::TempDir; - use crate::tests::test_storage; + use super::*; + use crate::{storage::base_storage::StorageConnector, tests::test_storage}; #[tokio::test] async fn ensure_cl_ref_creates_missing_canonical_ref() { @@ -1027,4 +1050,78 @@ mod tests { assert_eq!(updated.ref_tree_hash, "4".repeat(40)); assert!(updated.is_cl); } + + fn mega_blob_row(id: i64, blob_id: &str, file_path: &str) -> mega_blob::Model { + mega_blob::Model { + id, + blob_id: blob_id.to_string(), + name: String::new(), + size: 0, + created_at: chrono::Utc::now().naive_utc(), + pack_id: String::new(), + file_path: file_path.to_string(), + pack_offset: 0, + is_delta_in_pack: false, + commit_id: String::new(), + } + } + + async fn insert_mega_blob(stg: &MonoStorage, model: mega_blob::Model) { + mega_blob::Entity::insert(model.into_active_model()) + .exec(stg.get_connection()) + .await + .expect("insert mega_blob"); + } + + async fn mega_filepath(stg: &MonoStorage, blob_id: &str) -> Option { + mega_blob::Entity::find() + .filter(mega_blob::Column::BlobId.eq(blob_id)) + .one(stg.get_connection()) + .await + .unwrap() + .map(|m| m.file_path) + } + + #[tokio::test] + async fn update_blob_filepaths_sets_multiple_paths() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.mono_storage(); + insert_mega_blob(&stg, mega_blob_row(1, "blob-a", "")).await; + insert_mega_blob(&stg, mega_blob_row(2, "blob-b", "")).await; + + stg.update_blob_filepaths(vec![ + ("blob-a".into(), "src/lib.rs".into()), + ("blob-b".into(), "Cargo.toml".into()), + ]) + .await + .unwrap(); + + assert_eq!( + mega_filepath(&stg, "blob-a").await.as_deref(), + Some("src/lib.rs") + ); + assert_eq!( + mega_filepath(&stg, "blob-b").await.as_deref(), + Some("Cargo.toml") + ); + } + + #[tokio::test] + async fn update_blob_filepaths_ignores_missing_id() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.mono_storage(); + stg.update_blob_filepaths(vec![("missing".into(), "nope.rs".into())]) + .await + .expect("missing blob is not an error"); + } +} + +fn last_wins_mega_filepaths(pairs: Vec<(String, String)>) -> Vec<(String, String)> { + let mut map = HashMap::with_capacity(pairs.len()); + for (blob_id, file_path) in pairs { + map.insert(blob_id, file_path); + } + map.into_iter().collect() } diff --git a/scripts/crates-sync/crates-sync.py b/scripts/crates-sync/crates-sync.py index 00e668a6c..75d75eec5 100644 --- a/scripts/crates-sync/crates-sync.py +++ b/scripts/crates-sync/crates-sync.py @@ -177,17 +177,39 @@ def _progress_set_total(total: int) -> None: _progress_total = total _progress_versions_queued = total +def _eta_push_rate_per_s() -> float: + """Successful pushes/s. Never use skip rate — remaining work is extract+push.""" + window = _pushes_per_sec_last_60s() + if _push_ok_last_60s() >= 5 and window > 0: + return window + if _run_start_mono is None: + return 0.0 + ok_total, _ = _push_totals() + if ok_total < 5: + return 0.0 + elapsed = max(1e-6, time.monotonic() - _run_start_mono) + return ok_total / elapsed + + def _format_eta(done: int, total: int | None) -> str: - if _run_start_mono is None or done <= 0 or not total or total <= done: + """ETA = remaining queue / successful push rate. + + Skip (ls-remote) is orders of magnitude faster than push. Counting skip + toward the rate makes a multi-day import look like ~1–2h. + """ + if not total or total <= done: return "eta=?" - elapsed = max(1e-6, time.monotonic() - _run_start_mono) - rate = done / elapsed - remain = (total - done) / max(1e-9, rate) + rate = _eta_push_rate_per_s() + if rate <= 0: + return "eta=?" + remain = (total - done) / rate if remain < 60: return f"eta={remain:.0f}s" if remain < 3600: return f"eta={remain / 60:.1f}m" - return f"eta={remain / 3600:.1f}h" + if remain < 86400: + return f"eta={remain / 3600:.1f}h" + return f"eta={remain / 86400:.1f}d" def _format_progress_bar(done: int, total: int | None, width: int = 30) -> str: """ASCII progress bar, e.g. [############--------------] 40.0%""" From aa668b747fbf5e6cfb8df6c34ad8a5e5c47c7b30 Mon Sep 17 00:00:00 2001 From: "benjamin.747" Date: Sat, 29 Aug 2026 11:14:22 +0800 Subject: [PATCH 2/3] fix(import): retry git_blob insert deadlocks and upsert git_repo by path Serialize 1000-row chunks, insert hashes in sorted order, and retry Postgres 40P01/40001 so concurrent receive-packs stop deadlocking on git_blob unique indexes. Reuse the persisted git_repo row on repo_path conflict instead of the caller's unused snowflake id. --- ceres/src/transport/protocol/mod.rs | 10 ++- common/src/errors.rs | 55 +++++++++++++++ jupiter/src/service/import_service.rs | 10 ++- jupiter/src/service/mono_service.rs | 12 +++- jupiter/src/storage/base_storage.rs | 98 ++++++++++++++++++--------- jupiter/src/storage/git_db_storage.rs | 85 +++++++++++++++++++---- jupiter/src/utils/converter/mod.rs | 5 ++ 7 files changed, 221 insertions(+), 54 deletions(-) diff --git a/ceres/src/transport/protocol/mod.rs b/ceres/src/transport/protocol/mod.rs index 353da5c83..9b4b590d1 100644 --- a/ceres/src/transport/protocol/mod.rs +++ b/ceres/src/transport/protocol/mod.rs @@ -196,16 +196,14 @@ impl SmartSession { nested_import_repo_conflict_message(path_str, &conflict.repo_path), )); } - let repo = Repo::new(self.repo_path.clone(), false); - storage - .save_git_repo(repo.clone().into()) - .await - .map_err(|e| { + let created = Repo::new(self.repo_path.clone(), false); + let persisted = + storage.save_git_repo(created.into()).await.map_err(|e| { ProtocolError::InvalidInput(format!( "failed to create import repo: {e}" )) })?; - repo + Repo::from(persisted) } } }; diff --git a/common/src/errors.rs b/common/src/errors.rs index 3059be656..ff1f606c2 100644 --- a/common/src/errors.rs +++ b/common/src/errors.rs @@ -130,6 +130,29 @@ impl MegaError { } Self::Other(msg) } + + /// True for Postgres deadlock (`40P01`) or serialization failure (`40001`). + pub fn is_retryable_db_serialization(&self) -> bool { + match self { + MegaError::Db(err) => db_err_is_retryable_serialization(err), + _ => is_retryable_pg_conflict(&self.to_string()), + } + } +} + +/// True when a SeaORM error is a Postgres deadlock (`40P01`) or +/// serialization failure (`40001`) that is safe to retry. +pub fn db_err_is_retryable_serialization(err: &sea_orm::DbErr) -> bool { + is_retryable_pg_conflict(&err.to_string()) || is_retryable_pg_conflict(&format!("{err:?}")) +} + +fn is_retryable_pg_conflict(text: &str) -> bool { + let lower = text.to_ascii_lowercase(); + lower.contains("deadlock detected") + || lower.contains("40p01") + || lower.contains("40001") + || lower.contains("serialization failure") + || lower.contains("could not serialize access") } impl From for MegaError { @@ -450,4 +473,36 @@ mod tests { assert!(matches!(err, ProtocolError::NotFound(_))); assert_eq!(protocol_error_http_status(&err), 404); } + + #[test] + fn db_err_detects_deadlock_message() { + let err = sea_orm::DbErr::Custom( + "deadlock detected\nCONTEXT: while inserting index tuple in relation \"git_blob\"" + .into(), + ); + assert!(db_err_is_retryable_serialization(&err)); + assert!(MegaError::Db(err).is_retryable_db_serialization()); + } + + #[test] + fn db_err_detects_sqlstate_40p01() { + let err = sea_orm::DbErr::Custom("ERROR: 40P01 deadlock detected".into()); + assert!(db_err_is_retryable_serialization(&err)); + } + + #[test] + fn db_err_detects_sqlstate_40001() { + let err = sea_orm::DbErr::Custom("ERROR: 40001 could not serialize access".into()); + assert!(db_err_is_retryable_serialization(&err)); + assert!(MegaError::Db(err).is_retryable_db_serialization()); + } + + #[test] + fn db_err_ignores_unrelated_unique_violations() { + let err = sea_orm::DbErr::Custom( + "duplicate key value violates unique constraint \"git_repo_pkey\"".into(), + ); + assert!(!db_err_is_retryable_serialization(&err)); + assert!(!MegaError::Db(err).is_retryable_db_serialization()); + } } diff --git a/jupiter/src/service/import_service.rs b/jupiter/src/service/import_service.rs index e544e6c01..f8f0992a9 100644 --- a/jupiter/src/service/import_service.rs +++ b/jupiter/src/service/import_service.rs @@ -16,7 +16,7 @@ use crate::{ base_storage::{BaseStorage, StorageConnector}, git_db_storage::GitDbStorage, }, - utils::converter::{GitObjectModel, process_entry}, + utils::converter::{GitObjectModel, active_hash, process_entry}, }; #[derive(Clone)] @@ -103,10 +103,16 @@ impl ImportService { return Err(err); } - let git_objects = Arc::try_unwrap(git_objects) + let mut git_objects = Arc::try_unwrap(git_objects) .expect("Failed to unwrap Arc") .into_inner(); + git_objects.blobs.sort_by_key(|a| active_hash(&a.blob_id)); + git_objects.trees.sort_by_key(|a| active_hash(&a.tree_id)); + git_objects + .commits + .sort_by_key(|a| active_hash(&a.commit_id)); + self.git_db_storage .batch_save_model(git_objects.commits) .await?; diff --git a/jupiter/src/service/mono_service.rs b/jupiter/src/service/mono_service.rs index b6f7cbf40..b6dccd806 100644 --- a/jupiter/src/service/mono_service.rs +++ b/jupiter/src/service/mono_service.rs @@ -17,7 +17,9 @@ use crate::{ base_storage::{BaseStorage, StorageConnector}, mono_storage::MonoStorage, }, - utils::converter::{IntoMegaModel, MegaModelConverter, MegaObjectModel, process_entry}, + utils::converter::{ + IntoMegaModel, MegaModelConverter, MegaObjectModel, active_hash, process_entry, + }, }; #[derive(Clone)] @@ -134,10 +136,16 @@ impl MonoService { return Err(err); } - let git_objects = Arc::try_unwrap(git_objects) + let mut git_objects = Arc::try_unwrap(git_objects) .expect("Failed to unwrap Arc") .into_inner(); + git_objects.blobs.sort_by_key(|a| active_hash(&a.blob_id)); + git_objects.trees.sort_by_key(|a| active_hash(&a.tree_id)); + git_objects + .commits + .sort_by_key(|a| active_hash(&a.commit_id)); + self.mono_storage .batch_save_model(git_objects.commits.clone()) .await?; diff --git a/jupiter/src/storage/base_storage.rs b/jupiter/src/storage/base_storage.rs index df03357f9..fc610be37 100644 --- a/jupiter/src/storage/base_storage.rs +++ b/jupiter/src/storage/base_storage.rs @@ -1,12 +1,51 @@ -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; use async_trait::async_trait; -use common::errors::MegaError; +use common::errors::{MegaError, db_err_is_retryable_serialization}; use sea_orm::{ ActiveModelTrait, DatabaseConnection, DatabaseTransaction, DbErr, EntityTrait, sea_query::OnConflict, }; +const INSERT_RETRY_ATTEMPTS: u32 = 5; +const INSERT_RETRY_BASE_MS: u64 = 10; + +async fn insert_many_with_deadlock_retry( + conn: &DatabaseConnection, + txn: Option<&DatabaseTransaction>, + models: Vec, + onconflict: &OnConflict, +) -> Result<(), MegaError> +where + E: EntityTrait, + A: ActiveModelTrait + From<::Model> + Send + Clone, +{ + let mut attempt = 0u32; + loop { + attempt += 1; + let insert = E::insert_many(models.clone()).on_conflict(onconflict.clone()); + let result = if let Some(txn) = txn { + insert.exec(txn).await + } else { + insert.exec(conn).await + }; + match result { + Ok(_) | Err(DbErr::RecordNotInserted) => return Ok(()), + Err(e) if db_err_is_retryable_serialization(&e) && attempt < INSERT_RETRY_ATTEMPTS => { + let backoff_ms = INSERT_RETRY_BASE_MS << (attempt - 1); + tracing::warn!( + attempt, + backoff_ms, + error = %e, + "retrying batch insert after deadlock or serialization failure" + ); + tokio::time::sleep(Duration::from_millis(backoff_ms)).await; + } + Err(e) => return Err(e.into()), + } + } +} + #[async_trait] pub trait StorageConnector { const BATCH_CHUNK_SIZE: usize = 1000; @@ -21,7 +60,7 @@ pub trait StorageConnector { async fn batch_save_model(&self, save_models: Vec) -> Result<(), MegaError> where E: EntityTrait, - A: ActiveModelTrait + From<::Model> + Send, + A: ActiveModelTrait + From<::Model> + Send + Clone, { let onconflict = OnConflict::new().do_nothing().to_owned(); Self::batch_save_model_with_conflict(self, save_models, onconflict).await @@ -34,7 +73,7 @@ pub trait StorageConnector { ) -> Result<(), MegaError> where E: EntityTrait, - A: ActiveModelTrait + From<::Model> + Send, + A: ActiveModelTrait + From<::Model> + Send + Clone, { let onconflict = OnConflict::new().do_nothing().to_owned(); Self::batch_save_model_with_conflict_and_txn(self, save_models, onconflict, txn).await @@ -48,24 +87,20 @@ pub trait StorageConnector { ) -> Result<(), MegaError> where E: EntityTrait, - A: ActiveModelTrait + From<::Model> + Send, + A: ActiveModelTrait + From<::Model> + Send + Clone, { let mut i = 0; let len = save_models.len(); while i < len { let end = (i + Self::BATCH_CHUNK_SIZE).min(len); - let models = save_models[i..end].to_vec(); - let insert = E::insert_many(models).on_conflict(onconflict.clone()); - let _ = match if let Some(txn) = txn { - insert.exec(txn).await - } else { - insert.exec(self.get_connection()).await - } { - Ok(_) => Ok(()), - Err(DbErr::RecordNotInserted) => Ok(()), - Err(e) => Err(e), - }; + insert_many_with_deadlock_retry::( + self.get_connection(), + txn, + save_models[i..end].to_vec(), + &onconflict, + ) + .await?; i = end; } Ok(()) @@ -98,23 +133,22 @@ pub trait StorageConnector { ) -> Result<(), MegaError> where E: EntityTrait, - A: ActiveModelTrait + From<::Model> + Send, + A: ActiveModelTrait + From<::Model> + Send + Clone, { - let futures = save_models.chunks(Self::BATCH_CHUNK_SIZE).map(|chunk| { - let insert = E::insert_many(chunk.iter().cloned()).on_conflict(onconflict.clone()); - - async move { - match insert.exec(self.get_connection()).await { - Ok(_) => Ok(()), - Err(DbErr::RecordNotInserted) => { - // ignore not inserted err - Ok(()) - } - Err(e) => Err(e), - } - } - }); - futures::future::try_join_all(futures).await?; + let mut i = 0; + let len = save_models.len(); + + while i < len { + let end = (i + Self::BATCH_CHUNK_SIZE).min(len); + insert_many_with_deadlock_retry::( + self.get_connection(), + None, + save_models[i..end].to_vec(), + &onconflict, + ) + .await?; + i = end; + } Ok(()) } } diff --git a/jupiter/src/storage/git_db_storage.rs b/jupiter/src/storage/git_db_storage.rs index c7cdfdcb2..fcf963b6e 100644 --- a/jupiter/src/storage/git_db_storage.rs +++ b/jupiter/src/storage/git_db_storage.rs @@ -13,7 +13,7 @@ use futures::Stream; use sea_orm::{ ActiveModelTrait, ColumnTrait, ConnectionTrait, DatabaseTransaction, DbErr, EntityTrait, IntoActiveModel, PaginatorTrait, QueryFilter, QueryOrder, Set, TransactionTrait, - sea_query::{CaseStatement, Expr, ExprTrait}, + sea_query::{CaseStatement, Expr, ExprTrait, OnConflict}, }; use crate::storage::base_storage::{BaseStorage, StorageConnector}; @@ -50,16 +50,14 @@ impl GitDbStorage { &conflict.repo_path, ))); } - let repo_id = generate_id(); let repo = git_repo::Model { - id: repo_id, + id: generate_id(), repo_path: repo_path.to_string(), repo_name: repo_name.to_string(), created_at: chrono::Utc::now().naive_utc(), updated_at: chrono::Utc::now().naive_utc(), }; - self.save_git_repo(repo).await?; - repo_id + self.save_git_repo(repo).await?.id }; let refs = import_refs::Model { @@ -424,13 +422,22 @@ impl GitDbStorage { Ok(None) } - pub async fn save_git_repo(&self, repo: git_repo::Model) -> Result<(), MegaError> { - let a_model = repo.into_active_model(); - git_repo::Entity::insert(a_model) - .exec(self.get_connection()) - .await - .map_err(|e| MegaError::Other(format!("Failed to insert git_repo: {e}")))?; - Ok(()) + pub async fn save_git_repo(&self, repo: git_repo::Model) -> Result { + let repo_path = repo.repo_path.clone(); + let insert = git_repo::Entity::insert(repo.into_active_model()).on_conflict( + OnConflict::column(git_repo::Column::RepoPath) + .do_nothing() + .to_owned(), + ); + match insert.exec(self.get_connection()).await { + Ok(_) | Err(DbErr::RecordNotInserted) => {} + Err(e) => { + return Err(MegaError::Other(format!("Failed to insert git_repo: {e}"))); + } + } + self.find_git_repo_exact_match(&repo_path) + .await? + .ok_or_else(|| MegaError::Other(format!("git_repo missing after save: {repo_path}"))) } pub async fn get_commit_by_hash( @@ -866,6 +873,60 @@ mod tests { .expect("insert git_repo"); } + fn git_repo_row(id: i64, path: &str, name: &str) -> git_repo::Model { + git_repo::Model { + id, + repo_path: path.to_string(), + repo_name: name.to_string(), + created_at: chrono::Utc::now().naive_utc(), + updated_at: chrono::Utc::now().naive_utc(), + } + } + + #[tokio::test] + async fn save_git_repo_same_path_returns_existing_id() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + let path = "/third-party/rust/crates/foo/1.0.0"; + let first = stg + .save_git_repo(git_repo_row(11, path, "1.0.0")) + .await + .unwrap(); + let second = stg + .save_git_repo(git_repo_row(22, path, "other")) + .await + .unwrap(); + assert_eq!(first.id, 11); + assert_eq!(second.id, first.id); + assert_eq!(second.repo_name, first.repo_name); + } + + #[tokio::test] + async fn create_repo_and_save_ref_reuses_repo_id() { + let dir = TempDir::new().unwrap(); + let storage = test_storage(dir.path()).await; + let stg = storage.git_db_storage(); + let path = "/third-party/foo"; + stg.create_repo_and_save_ref(path, "foo", "refs/heads/master", "aaa") + .await + .unwrap(); + let first = stg + .find_git_repo_exact_match(path) + .await + .unwrap() + .expect("repo after first create"); + stg.create_repo_and_save_ref(path, "foo", "refs/heads/master", "bbb") + .await + .unwrap(); + let second = stg + .find_git_repo_exact_match(path) + .await + .unwrap() + .expect("repo after second create"); + assert_eq!(first.id, second.id); + } + #[test] fn git_repo_descendant_bounds_are_prefix_range() { assert_eq!( diff --git a/jupiter/src/utils/converter/mod.rs b/jupiter/src/utils/converter/mod.rs index 5396c7b2f..619a52b5d 100644 --- a/jupiter/src/utils/converter/mod.rs +++ b/jupiter/src/utils/converter/mod.rs @@ -6,7 +6,12 @@ mod traits; pub use init_monorepo::*; pub use pack::*; +use sea_orm::ActiveValue; pub use traits::*; +pub(crate) fn active_hash(value: &ActiveValue) -> String { + value.clone().take().unwrap_or_default() +} + #[cfg(test)] mod test; From 9cf471de5a92691fc1b9990976e8a560be61b1a4 Mon Sep 17 00:00:00 2001 From: "benjamin.747" Date: Sat, 29 Aug 2026 11:46:57 +0800 Subject: [PATCH 3/3] fix(idgen): claim a unique snowflake worker_id per pod Hardcoded worker_id(1) made every replica share one partition, so concurrent generate_id() hit git_repo_pkey. Claim mega:snowflake:worker:N in Redis at boot (8+8 bits), with env/hash fallback when Redis is down. --- jupiter/src/redis/mod.rs | 2 + jupiter/src/redis/snowflake_worker.rs | 90 ++++++++++++++++ jupiter/src/utils/id_generator.rs | 150 ++++++++++++++++++++++++-- mono/src/bootstrap/mod.rs | 23 +++- 4 files changed, 255 insertions(+), 10 deletions(-) create mode 100644 jupiter/src/redis/snowflake_worker.rs diff --git a/jupiter/src/redis/mod.rs b/jupiter/src/redis/mod.rs index d0999b340..5801f3109 100644 --- a/jupiter/src/redis/mod.rs +++ b/jupiter/src/redis/mod.rs @@ -1,7 +1,9 @@ pub mod lock; +pub mod snowflake_worker; pub use ::redis::{AsyncCommands, aio::ConnectionManager}; use common::config::RedisConfig; +pub use snowflake_worker::claim_snowflake_worker; /// Initializes a Redis multiplexed asynchronous connection from the given configuration. /// diff --git a/jupiter/src/redis/snowflake_worker.rs b/jupiter/src/redis/snowflake_worker.rs new file mode 100644 index 000000000..604069b4d --- /dev/null +++ b/jupiter/src/redis/snowflake_worker.rs @@ -0,0 +1,90 @@ +use std::time::Duration; + +use redis::{Script, aio::ConnectionManager}; +use tokio::time::sleep; + +use crate::utils::id_generator::{self, MAX_WORKER_ID}; + +const SLOT_KEY_PREFIX: &str = "mega:snowflake:worker:"; +const SLOT_TTL_MS: u64 = 30_000; + +/// Try to exclusive-claim a snowflake worker slot in Redis (`SET NX PX`). +/// +/// Returns `None` if Redis errors or every slot in `0..=MAX_WORKER_ID` is taken. +/// On success, spawns a background PEXPIRE refresh so a live process keeps the slot. +pub async fn claim_snowflake_worker(connection: &ConnectionManager) -> Option { + let identity = id_generator::process_identity(); + let mut conn = connection.clone(); + for id in 0..=MAX_WORKER_ID { + let key = format!("{SLOT_KEY_PREFIX}{id}"); + let result: Result, _> = redis::cmd("SET") + .arg(&key) + .arg(&identity) + .arg("NX") + .arg("PX") + .arg(SLOT_TTL_MS) + .query_async(&mut conn) + .await; + match result { + Ok(Some(_)) => { + tracing::info!( + worker_id = id, + slot_key = %key, + identity = %identity, + "claimed snowflake worker slot" + ); + spawn_slot_refresh(connection.clone(), key, identity); + return Some(id); + } + Ok(None) => continue, + Err(e) => { + tracing::warn!( + error = %e, + "snowflake worker slot claim failed; falling back to hash" + ); + return None; + } + } + } + tracing::warn!( + max = MAX_WORKER_ID, + "all snowflake worker slots taken; falling back to hash" + ); + None +} + +fn spawn_slot_refresh(connection: ConnectionManager, key: String, value: String) { + tokio::spawn(async move { + let half = SLOT_TTL_MS / 2; + let mut conn = connection; + let script = Script::new( + r#" + if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("PEXPIRE", KEYS[1], ARGV[2]) + else + return 0 + end + "#, + ); + loop { + sleep(Duration::from_millis(half)).await; + let ok: Result = script + .key(&key) + .arg(&value) + .arg(SLOT_TTL_MS) + .invoke_async(&mut conn) + .await; + match ok { + Ok(1) => {} + Ok(_) => { + tracing::warn!(key = %key, "snowflake worker slot refresh lost"); + break; + } + Err(e) => { + tracing::warn!(key = %key, error = %e, "snowflake worker slot refresh failed"); + break; + } + } + } + }); +} diff --git a/jupiter/src/utils/id_generator.rs b/jupiter/src/utils/id_generator.rs index decd0e91b..b5f4b3d9d 100644 --- a/jupiter/src/utils/id_generator.rs +++ b/jupiter/src/utils/id_generator.rs @@ -1,8 +1,80 @@ -use std::sync::Once; +use std::sync::{Once, OnceLock}; use idgenerator::*; +/// Matches `idgenerator` defaults. `worker_id_bit_len + seq_bit_len` must be +/// `<= 22`. 8+8 gives 256 concurrent workers and 256 ids/ms per worker. +pub const WORKER_ID_BIT_LEN: u8 = 8; +pub const SEQ_BIT_LEN: u8 = 8; +pub const MAX_WORKER_ID: u32 = (1 << WORKER_ID_BIT_LEN) - 1; +pub const ENV_WORKER_ID: &str = "MEGA_ID_GENERATOR_WORKER_ID"; + static ID_GENERATOR_INIT: Once = Once::new(); +static CLAIMED_WORKER_ID: OnceLock = OnceLock::new(); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WorkerIdSource { + Env, + Redis, + Hash, +} + +/// Record a Redis-claimed worker id. Must run before [`set_up_options`]. +/// Returns false if a claim was already stored. +pub fn claim_worker_id(id: u32) -> bool { + CLAIMED_WORKER_ID.set(id.min(MAX_WORKER_ID)).is_ok() +} + +pub fn process_identity() -> String { + std::env::var("POD_UID") + .or_else(|_| std::env::var("HOSTNAME")) + .unwrap_or_else(|_| "mega-local".to_string()) +} + +/// FNV-1a 32-bit. Stable across rustc versions (unlike `DefaultHasher`). +pub fn fnv1a_32(bytes: &[u8]) -> u32 { + let mut hash = 0x811c9dc5u32; + for b in bytes { + hash ^= u32::from(*b); + hash = hash.wrapping_mul(0x0100_0193); + } + hash +} + +pub fn hash_worker_id(identity: &str) -> u32 { + fnv1a_32(identity.as_bytes()) % (MAX_WORKER_ID + 1) +} + +pub fn resolve_worker_id() -> (u32, WorkerIdSource) { + resolve_worker_id_from( + std::env::var(ENV_WORKER_ID).ok().as_deref(), + CLAIMED_WORKER_ID.get().copied(), + &process_identity(), + ) +} + +pub fn resolve_worker_id_from( + env_val: Option<&str>, + claimed: Option, + identity: &str, +) -> (u32, WorkerIdSource) { + if let Some(raw) = env_val { + match raw.parse::() { + Ok(id) if id <= MAX_WORKER_ID => return (id, WorkerIdSource::Env), + _ => { + tracing::warn!( + raw, + max = MAX_WORKER_ID, + "ignoring out-of-range or invalid MEGA_ID_GENERATOR_WORKER_ID" + ); + } + } + } + if let Some(id) = claimed { + return (id.min(MAX_WORKER_ID), WorkerIdSource::Redis); + } + (hash_worker_id(identity), WorkerIdSource::Hash) +} /// Ensures [`IdInstance`] is configured (idempotent; safe if [`set_up_options`] already ran, e.g. via [`crate::storage::init::database_connection`]). pub fn ensure_initialized() { @@ -17,15 +89,77 @@ pub fn ensure_initialized() { } pub fn set_up_options() -> Result<(), OptionError> { - // Setup the option for the id generator instance. - let options = IdGeneratorOptions::new().worker_id(1).worker_id_bit_len(6); + let (worker_id, source) = resolve_worker_id(); + let options = IdGeneratorOptions::new() + .worker_id(worker_id) + .worker_id_bit_len(WORKER_ID_BIT_LEN) + .seq_bit_len(SEQ_BIT_LEN); - // Initialize the id generator instance with the option. - // Other options not set will be given the default value. IdInstance::init(options)?; - // Get the option from the id generator instance. - let options = IdInstance::get_options(); - tracing::info!("First setting: {:?}", options); + tracing::info!( + worker_id, + worker_id_bit_len = WORKER_ID_BIT_LEN, + seq_bit_len = SEQ_BIT_LEN, + ?source, + identity = %process_identity(), + "snowflake id generator initialized" + ); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn env_in_range_wins() { + let (id, source) = resolve_worker_id_from(Some("7"), Some(3), "pod-a"); + assert_eq!(id, 7); + assert_eq!(source, WorkerIdSource::Env); + } + + #[test] + fn env_out_of_range_falls_through_to_claimed() { + let too_big = (MAX_WORKER_ID + 1).to_string(); + let (id, source) = resolve_worker_id_from(Some(&too_big), Some(3), "pod-a"); + assert_eq!(id, 3); + assert_eq!(source, WorkerIdSource::Redis); + } + + #[test] + fn env_invalid_falls_through_to_hash() { + let expected = hash_worker_id("fixture-pod"); + let (id, source) = resolve_worker_id_from(Some("abc"), None, "fixture-pod"); + assert_eq!(id, expected); + assert_eq!(source, WorkerIdSource::Hash); + } + + #[test] + fn missing_env_without_claim_hashes_identity() { + let (id, source) = resolve_worker_id_from(None, None, "fixture-pod"); + assert_eq!(id, hash_worker_id("fixture-pod")); + assert_eq!(source, WorkerIdSource::Hash); + assert!(id <= MAX_WORKER_ID); + } + + #[test] + fn worker_id_space_matches_crate_default() { + assert_eq!(WORKER_ID_BIT_LEN, 8); + assert_eq!(SEQ_BIT_LEN, 8); + assert_eq!(MAX_WORKER_ID, 255); + } + + #[test] + fn hash_is_stable() { + assert_eq!(hash_worker_id("fixture-pod"), hash_worker_id("fixture-pod")); + } + + #[test] + fn two_identities_hash_differently() { + assert_ne!( + hash_worker_id("mono-engine-5f6d8d7cc9-45tx8"), + hash_worker_id("mono-engine-5f6d8d7cc9-rknxw") + ); + } +} diff --git a/mono/src/bootstrap/mod.rs b/mono/src/bootstrap/mod.rs index 9ed2dd475..77928c5a1 100644 --- a/mono/src/bootstrap/mod.rs +++ b/mono/src/bootstrap/mod.rs @@ -1,6 +1,9 @@ use std::sync::Arc; -use jupiter::redis::{ConnectionManager, init_connection}; +use jupiter::{ + redis::{ConnectionManager, claim_snowflake_worker, init_connection}, + utils::id_generator, +}; /// Main application context for the Mono application. #[derive(Clone)] @@ -15,10 +18,16 @@ impl AppContext { pub async fn new(config: common::config::Config) -> Self { let config = Arc::new(config); + let connection = init_connection(&config.redis).await; + if env_worker_id_is_set() { + tracing::info!("MEGA_ID_GENERATOR_WORKER_ID set; skipping Redis snowflake slot claim"); + } else if let Some(id) = claim_snowflake_worker(&connection).await { + id_generator::claim_worker_id(id); + } + let storage = jupiter::storage::Storage::new(config.clone()) .await .expect("init monorepo storage err"); - let connection = init_connection(&config.redis).await; let storage_for_vault = storage.clone(); let vault = vault::integration::vault_core::VaultCore::new(storage_for_vault).await; @@ -41,3 +50,13 @@ impl AppContext { Arc::new(self.clone()) } } + +fn env_worker_id_is_set() -> bool { + match std::env::var(id_generator::ENV_WORKER_ID) { + Ok(raw) => raw + .parse::() + .ok() + .is_some_and(|id| id <= id_generator::MAX_WORKER_ID), + Err(_) => false, + } +}