From 2db572cff78c0c9df603030290c8a94d36fcadca Mon Sep 17 00:00:00 2001 From: yyjeqhc <1772413353@qq.com> Date: Thu, 10 Jul 2025 16:37:41 +0800 Subject: [PATCH] feat: optimize the fetch_code function for concurrency to improve code fetching performance. --- ceres/src/api_service/mono_api_service.rs | 4 +- ceres/src/model/mod.rs | 2 +- ceres/src/pack/import_repo.rs | 4 +- context/src/lib.rs | 15 +- extensions/rag/chat/src/main.rs | 4 +- extensions/rag/index/src/lib.rs | 1 - extensions/rag/index/src/main.rs | 4 +- gemini/src/lfs/mod.rs | 4 +- jupiter/callisto/src/entity_ext/label.rs | 3 +- jupiter/callisto/src/entity_ext/mod.rs | 2 +- jupiter/callisto/src/mod.rs | 2 +- jupiter/src/storage/git_db_storage.rs | 4 +- jupiter/src/storage/issue_storage.rs | 5 +- jupiter/src/storage/mono_storage.rs | 5 +- jupiter/src/storage/mr_storage.rs | 5 +- libra/src/command/branch.rs | 8 +- libra/src/command/diff.rs | 4 +- libra/src/command/fetch.rs | 4 +- libra/src/command/log.rs | 4 +- libra/src/command/push.rs | 8 +- libra/src/command/switch.rs | 4 +- libra/src/utils/lfs.rs | 4 +- libra/tests/command/log_test.rs | 4 +- mercury/src/internal/object/mod.rs | 6 +- mono/src/api/oauth/model.rs | 2 +- orion-server/src/api.rs | 19 +- orion-server/src/model/builds.rs | 2 +- orion-server/src/server.rs | 4 +- scorpio/README.md | 4 + scorpio/doc/perf_test.md | 113 +++++ scorpio/scorpio.toml | 7 +- scorpio/script/run_1000_files.sh | 111 +++++ scorpio/src/dicfuse/store.rs | 5 +- scorpio/src/manager/fetch.rs | 500 ++++++++++++++++++++-- scorpio/src/manager/store.rs | 4 +- scorpio/src/util/config.rs | 23 +- vault/src/integration/vault_core.rs | 18 +- 37 files changed, 777 insertions(+), 145 deletions(-) create mode 100755 scorpio/script/run_1000_files.sh diff --git a/ceres/src/api_service/mono_api_service.rs b/ceres/src/api_service/mono_api_service.rs index 68710cf71..c3d744358 100644 --- a/ceres/src/api_service/mono_api_service.rs +++ b/ceres/src/api_service/mono_api_service.rs @@ -328,9 +328,7 @@ impl MonoApiService { }) .collect(); - storage.batch_save_model(save_trees) - .await - .unwrap(); + storage.batch_save_model(save_trees).await.unwrap(); Ok(p_commit_id) } diff --git a/ceres/src/model/mod.rs b/ceres/src/model/mod.rs index 62f977232..54f088e42 100644 --- a/ceres/src/model/mod.rs +++ b/ceres/src/model/mod.rs @@ -1,2 +1,2 @@ pub mod git; -pub mod mr; \ No newline at end of file +pub mod mr; diff --git a/ceres/src/pack/import_repo.rs b/ceres/src/pack/import_repo.rs index 9dce67269..5749de99c 100644 --- a/ceres/src/pack/import_repo.rs +++ b/ceres/src/pack/import_repo.rs @@ -395,9 +395,7 @@ impl ImportRepo { }) .collect(); - storage.batch_save_model( save_trees) - .await - .unwrap(); + storage.batch_save_model(save_trees).await.unwrap(); root_ref.ref_commit_hash = new_commit.id.to_string(); root_ref.ref_tree_hash = new_commit.tree_id.to_string(); diff --git a/context/src/lib.rs b/context/src/lib.rs index a22e49117..1fd225db8 100644 --- a/context/src/lib.rs +++ b/context/src/lib.rs @@ -22,17 +22,17 @@ pub struct AppContext { impl AppContext { /// Creates a new application context with the given configuration. pub async fn new(config: common::config::Config) -> Self { - let config = Arc::new(config); - + let storage = jupiter::storage::Storage::new(config.clone()).await; - + let storage_for_vault = storage.clone(); let vault = tokio::task::spawn_blocking(move || { vault::integration::vault_core::VaultCore::new(storage_for_vault) - }).await.expect("VaultCore::new panicked"); - - + }) + .await + .expect("VaultCore::new panicked"); + #[cfg(feature = "p2p")] let client = gemini::p2p::client::P2PClient::new(storage.clone(), vault.clone()); @@ -49,12 +49,9 @@ impl AppContext { #[cfg(feature = "p2p")] client, } - - } pub fn wrapped_context(&self) -> Arc { Arc::new(self.clone()) } } - diff --git a/extensions/rag/chat/src/main.rs b/extensions/rag/chat/src/main.rs index e120762d1..b7b02d37f 100644 --- a/extensions/rag/chat/src/main.rs +++ b/extensions/rag/chat/src/main.rs @@ -1,7 +1,7 @@ use chat::command::{Cli, Commands}; use chat::generation::GenerationNode; use chat::search::SearchNode; -use chat::{GENERATION_NODE, SEARCH_NODE, vect_url, qdrant_url, llm_url}; +use chat::{llm_url, qdrant_url, vect_url, GENERATION_NODE, SEARCH_NODE}; use clap::Parser; use dagrs::utils::env::EnvVar; use dagrs::{DefaultNode, Graph, Node, NodeTable}; @@ -10,7 +10,7 @@ use std::env; use std::thread; fn main() -> Result<(), Box> { - // Initialize logger + // Initialize logger env::set_var("RUST_LOG", "info"); env_logger::init(); diff --git a/extensions/rag/index/src/lib.rs b/extensions/rag/index/src/lib.rs index a9bae16bf..2d45547c6 100644 --- a/extensions/rag/index/src/lib.rs +++ b/extensions/rag/index/src/lib.rs @@ -24,7 +24,6 @@ pub fn llm_url() -> String { env::var("LLM_URL").unwrap_or_else(|_| "http://ollama:11434/api/chat".to_string()) } - pub fn consumer_group() -> String { std::env::var("CONSUMER_GROUP").unwrap_or_else(|_| "test-group".to_string()) } diff --git a/extensions/rag/index/src/main.rs b/extensions/rag/index/src/main.rs index 73105a949..130bb8f15 100644 --- a/extensions/rag/index/src/main.rs +++ b/extensions/rag/index/src/main.rs @@ -6,7 +6,7 @@ use index::indexer::WalkDirAction; use index::qdrant::QdrantNode; use index::vectorization::VectClient; use index::{broker, consumer_group, crates_path, topic}; -use index::{PROCESS_ITEMS_NODE, QDRANT_NODE, VECT_CLIENT_NODE, qdrant_url, vect_url}; +use index::{qdrant_url, vect_url, PROCESS_ITEMS_NODE, QDRANT_NODE, VECT_CLIENT_NODE}; use observatory::facilities::Telescope; use observatory::model::crates::CrateMessage; use std::env; @@ -41,7 +41,7 @@ fn main() { let rt = Runtime::new().unwrap(); rt.block_on(async { let telescope = Telescope::new(&broker(), &consumer_group(), &topic()); - + // 1. Clone the Arc *before* the loop's closure. // This clone is moved into the closure, allowing the original to remain. let id_counter_for_loop = Arc::clone(&shared_id_counter); diff --git a/gemini/src/lfs/mod.rs b/gemini/src/lfs/mod.rs index 674cf5a34..44ec95973 100644 --- a/gemini/src/lfs/mod.rs +++ b/gemini/src/lfs/mod.rs @@ -78,9 +78,7 @@ pub async fn share_lfs( /// This method will send a GET request to the relay to get lfs chunks info /// pub async fn get_lfs_chunks_info(bootstrap_node: String, file_hash: String) -> Option { - let url = format!( - "{bootstrap_node}/api/v1/lfs_chunk?file_hash={file_hash}" - ); + let url = format!("{bootstrap_node}/api/v1/lfs_chunk?file_hash={file_hash}"); let lfs_info: LFSInfoRes = match get(url.clone()).await { Ok(response) => { if !response.status().is_success() { diff --git a/jupiter/callisto/src/entity_ext/label.rs b/jupiter/callisto/src/entity_ext/label.rs index e791ec9df..6a0fd4aa5 100644 --- a/jupiter/callisto/src/entity_ext/label.rs +++ b/jupiter/callisto/src/entity_ext/label.rs @@ -2,7 +2,6 @@ use sea_orm::entity::prelude::*; use crate::label::Entity; - #[derive(Copy, Clone, Debug, EnumIter)] pub enum Relation { ItemLabels, @@ -14,4 +13,4 @@ impl RelationTrait for Relation { Self::ItemLabels => Entity::has_many(crate::item_labels::Entity).into(), } } -} \ No newline at end of file +} diff --git a/jupiter/callisto/src/entity_ext/mod.rs b/jupiter/callisto/src/entity_ext/mod.rs index 5845e5d38..95b4ac5bd 100644 --- a/jupiter/callisto/src/entity_ext/mod.rs +++ b/jupiter/callisto/src/entity_ext/mod.rs @@ -1,6 +1,6 @@ pub mod item_assignees; pub mod item_labels; pub mod label; +pub mod mega_conversation; pub mod mega_issue; pub mod mega_mr; -pub mod mega_conversation; \ No newline at end of file diff --git a/jupiter/callisto/src/mod.rs b/jupiter/callisto/src/mod.rs index 43b08847a..66ad72715 100644 --- a/jupiter/callisto/src/mod.rs +++ b/jupiter/callisto/src/mod.rs @@ -1,7 +1,7 @@ //! `SeaORM` Entity, @generated by sea-orm-codegen 1.1.12 -pub mod prelude; pub mod entity_ext; +pub mod prelude; pub mod access_token; pub mod builds; diff --git a/jupiter/src/storage/git_db_storage.rs b/jupiter/src/storage/git_db_storage.rs index b8f32d3a1..fdfa117c1 100644 --- a/jupiter/src/storage/git_db_storage.rs +++ b/jupiter/src/storage/git_db_storage.rs @@ -4,8 +4,8 @@ use std::sync::Arc; use futures::{stream, Stream, StreamExt}; use sea_orm::sea_query::Expr; use sea_orm::{ - ActiveModelTrait, ColumnTrait, DbBackend, DbErr, EntityTrait, - IntoActiveModel, QueryFilter, QueryTrait, Set, + ActiveModelTrait, ColumnTrait, DbBackend, DbErr, EntityTrait, IntoActiveModel, QueryFilter, + QueryTrait, Set, }; use sea_orm::{PaginatorTrait, QueryOrder}; use tokio::sync::Mutex; diff --git a/jupiter/src/storage/issue_storage.rs b/jupiter/src/storage/issue_storage.rs index 81e6fe282..97414f6d6 100644 --- a/jupiter/src/storage/issue_storage.rs +++ b/jupiter/src/storage/issue_storage.rs @@ -2,8 +2,8 @@ use std::collections::HashMap; use std::ops::Deref; use sea_orm::{ - ActiveModelTrait, ColumnTrait, Condition, EntityTrait, IntoActiveModel, - JoinType, PaginatorTrait, QueryFilter, QuerySelect, RelationTrait, Set, TransactionTrait, + ActiveModelTrait, ColumnTrait, Condition, EntityTrait, IntoActiveModel, JoinType, + PaginatorTrait, QueryFilter, QuerySelect, RelationTrait, Set, TransactionTrait, }; use callisto::sea_orm_active_enums::ConvTypeEnum; @@ -30,7 +30,6 @@ impl Deref for IssueStorage { } impl IssueStorage { - pub async fn get_issue_list( &self, params: ListParams, diff --git a/jupiter/src/storage/mono_storage.rs b/jupiter/src/storage/mono_storage.rs index 355b1fc26..80c261126 100644 --- a/jupiter/src/storage/mono_storage.rs +++ b/jupiter/src/storage/mono_storage.rs @@ -3,8 +3,8 @@ use std::sync::{Arc, Mutex}; use futures::{stream, StreamExt}; use sea_orm::{ - ActiveModelTrait, ColumnTrait, EntityTrait, IntoActiveModel, QueryFilter, - QueryOrder, QuerySelect, + ActiveModelTrait, ColumnTrait, EntityTrait, IntoActiveModel, QueryFilter, QueryOrder, + QuerySelect, }; use callisto::{mega_blob, mega_commit, mega_refs, mega_tag, mega_tree, raw_blob}; @@ -39,7 +39,6 @@ struct GitObjects { } impl MonoStorage { - pub async fn save_ref( &self, path: &str, diff --git a/jupiter/src/storage/mr_storage.rs b/jupiter/src/storage/mr_storage.rs index dbe26fcf4..e10cf6bf2 100644 --- a/jupiter/src/storage/mr_storage.rs +++ b/jupiter/src/storage/mr_storage.rs @@ -3,8 +3,8 @@ use std::ops::Deref; use common::model::Pagination; use sea_orm::{ - ActiveModelTrait, ColumnTrait, Condition, EntityTrait, IntoActiveModel, - JoinType, PaginatorTrait, QueryFilter, QuerySelect, RelationTrait, Set, + ActiveModelTrait, ColumnTrait, Condition, EntityTrait, IntoActiveModel, JoinType, + PaginatorTrait, QueryFilter, QuerySelect, RelationTrait, Set, }; use callisto::sea_orm_active_enums::MergeStatusEnum; @@ -30,7 +30,6 @@ impl Deref for MrStorage { } impl MrStorage { - pub async fn get_open_mr_by_path( &self, path: &str, diff --git a/libra/src/command/branch.rs b/libra/src/command/branch.rs index 3ad0b9d58..b309cc7de 100644 --- a/libra/src/command/branch.rs +++ b/libra/src/command/branch.rs @@ -81,9 +81,7 @@ pub async fn set_upstream(branch: &str, upstream: &str) { ) .await; } - println!( - "Branch '{branch}' set up to track remote branch '{upstream}'" - ); + println!("Branch '{branch}' set up to track remote branch '{upstream}'"); } pub async fn create_branch(new_branch: String, branch_or_commit: Option) { @@ -131,9 +129,7 @@ async fn delete_branch(branch_name: String) { if let Head::Branch(name) = head { if name == branch_name { - panic!( - "fatal: Cannot delete the branch '{branch_name}' which you are currently on" - ); + panic!("fatal: Cannot delete the branch '{branch_name}' which you are currently on"); } } diff --git a/libra/src/command/diff.rs b/libra/src/command/diff.rs index 2be20af32..4da778f1e 100644 --- a/libra/src/command/diff.rs +++ b/libra/src/command/diff.rs @@ -71,9 +71,7 @@ pub async fn execute(args: DiffArgs) { Some(ref path) => { let file = std::fs::File::create(path) .map_err(|e| { - eprintln!( - "fatal: could not open to file '{path}' for writing: {e}" - ); + eprintln!("fatal: could not open to file '{path}' for writing: {e}"); }) .unwrap(); Some(file) diff --git a/libra/src/command/fetch.rs b/libra/src/command/fetch.rs index 5461a8dce..915d378e5 100644 --- a/libra/src/command/fetch.rs +++ b/libra/src/command/fetch.rs @@ -68,9 +68,7 @@ pub async fn execute(args: FetchArgs) { Some(remote_config) => fetch_repository(&remote_config, args.refspec).await, None => { tracing::error!("remote config '{}' not found", remote); - eprintln!( - "fatal: '{remote}' does not appear to be a libra repository" - ); + eprintln!("fatal: '{remote}' does not appear to be a libra repository"); } } } diff --git a/libra/src/command/log.rs b/libra/src/command/log.rs index 5323c73b4..5f750ccbb 100644 --- a/libra/src/command/log.rs +++ b/libra/src/command/log.rs @@ -66,9 +66,7 @@ pub async fn execute(args: LogArgs) { if let Head::Branch(branch_name) = head.to_owned() { let branch = Branch::find_branch(&branch_name, None).await; if branch.is_none() { - panic!( - "fatal: your current branch '{branch_name}' does not have any commits yet " - ); + panic!("fatal: your current branch '{branch_name}' does not have any commits yet "); } } diff --git a/libra/src/command/push.rs b/libra/src/command/push.rs index fa95f9dd1..01948ab25 100644 --- a/libra/src/command/push.rs +++ b/libra/src/command/push.rs @@ -75,9 +75,7 @@ pub async fn execute(args: PushArgs) { .commit .to_string(); - println!( - "pushing {branch}({commit_hash}) to {repository}({repo_url})" - ); + println!("pushing {branch}({commit_hash}) to {repository}({repo_url})"); let url = Url::parse(&repo_url).unwrap(); let client = HttpsClient::from_url(&url); @@ -106,9 +104,7 @@ pub async fn execute(args: PushArgs) { let mut data = BytesMut::new(); add_pkt_line_string( &mut data, - format!( - "{remote_hash} {commit_hash} {tracked_branch}\0report-status\n" - ), + format!("{remote_hash} {commit_hash} {tracked_branch}\0report-status\n"), ); data.extend_from_slice(b"0000"); tracing::debug!("{:?}", data); diff --git a/libra/src/command/switch.rs b/libra/src/command/switch.rs index 9193cb297..3974819a9 100644 --- a/libra/src/command/switch.rs +++ b/libra/src/command/switch.rs @@ -89,9 +89,7 @@ async fn switch_to_branch(branch_name: String) { let target_branch = Branch::find_branch(&branch_name, None).await; if target_branch.is_none() { if !Branch::search_branch(&branch_name).await.is_empty() { - eprintln!( - "fatal: a branch is expected, got remote branch {branch_name}" - ); + eprintln!("fatal: a branch is expected, got remote branch {branch_name}"); } else { eprintln!("fatal: branch '{}' not found", &branch_name); } diff --git a/libra/src/utils/lfs.rs b/libra/src/utils/lfs.rs index 05380bc62..ce9b7bb5f 100644 --- a/libra/src/utils/lfs.rs +++ b/libra/src/utils/lfs.rs @@ -74,9 +74,7 @@ pub fn generate_pointer_file(path: impl AsRef) -> (String, String) { } pub fn format_pointer_string(oid: &str, size: u64) -> String { - format!( - "version {LFS_VERSION}\noid {LFS_HASH_ALGO}:{oid}\nsize {size}\n" - ) + format!("version {LFS_VERSION}\noid {LFS_HASH_ALGO}:{oid}\nsize {size}\n") } /// Generate LFS Server Url from repo Url. diff --git a/libra/tests/command/log_test.rs b/libra/tests/command/log_test.rs index f675c7287..41498402f 100644 --- a/libra/tests/command/log_test.rs +++ b/libra/tests/command/log_test.rs @@ -30,9 +30,7 @@ async fn test_execute_log() { if let Head::Branch(branch_name) = head.to_owned() { let branch = Branch::find_branch(&branch_name, None).await; if branch.is_none() { - panic!( - "fatal: your current branch '{branch_name}' does not have any commits yet " - ); + panic!("fatal: your current branch '{branch_name}' does not have any commits yet "); } } diff --git a/mercury/src/internal/object/mod.rs b/mercury/src/internal/object/mod.rs index f6b708eb6..4eac9a06b 100644 --- a/mercury/src/internal/object/mod.rs +++ b/mercury/src/internal/object/mod.rs @@ -40,11 +40,7 @@ pub trait ObjectTrait: Send + Sync + Display { read.read_to_end(&mut content).unwrap(); let h = read.hash.clone(); let hash_str = h.finalize(); - Self::from_bytes( - &content, - SHA1::from_str(&format!("{hash_str:x}")).unwrap(), - ) - .unwrap() + Self::from_bytes(&content, SHA1::from_str(&format!("{hash_str:x}")).unwrap()).unwrap() } /// Returns the type of the object. diff --git a/mono/src/api/oauth/model.rs b/mono/src/api/oauth/model.rs index a933ca0e6..c573daecb 100644 --- a/mono/src/api/oauth/model.rs +++ b/mono/src/api/oauth/model.rs @@ -77,7 +77,7 @@ impl From for LoginUser { email: value.email, created_at: value.created_at, campsite_user_id: String::new(), - username: String::new() + username: String::new(), } } } diff --git a/orion-server/src/api.rs b/orion-server/src/api.rs index 6f4e295c4..e3476842f 100644 --- a/orion-server/src/api.rs +++ b/orion-server/src/api.rs @@ -1,4 +1,5 @@ use crate::model::builds; +use axum::Json; use axum::body::Bytes; use axum::extract::ws::{Message, Utf8Bytes, WebSocket}; use axum::extract::{ConnectInfo, Path, State, WebSocketUpgrade}; @@ -6,7 +7,6 @@ use axum::http::StatusCode; use axum::response::sse::{Event, KeepAlive}; use axum::response::{IntoResponse, Sse}; use axum::routing::{any, get}; -use axum::{Json}; use axum_extra::json; use dashmap::DashMap; use futures_util::{SinkExt, Stream, StreamExt, stream}; @@ -17,7 +17,7 @@ use scopeguard::defer; use sea_orm::ActiveValue::Set; use sea_orm::prelude::DateTimeUtc; use sea_orm::sqlx::types::chrono; -use sea_orm::{ActiveModelTrait, DatabaseConnection, EntityTrait,ColumnTrait,QueryFilter as _}; +use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter as _}; use serde::{Deserialize, Serialize}; use std::convert::Infallible; use std::io::Write; @@ -28,19 +28,19 @@ use std::time::Duration; use tokio::io::AsyncReadExt; use tokio::sync::mpsc; use tokio::sync::mpsc::UnboundedSender; -use uuid::Uuid; use utoipa::ToSchema; use utoipa_axum::{router::OpenApiRouter, routes}; +use uuid::Uuid; static BUILD_LOG_DIR: Lazy = Lazy::new(|| std::env::var("BUILD_LOG_DIR").expect("BUILD_LOG_DIR must be set")); -#[derive(Debug, Deserialize,ToSchema)] +#[derive(Debug, Deserialize, ToSchema)] pub struct BuildRequest { repo: String, target: String, args: Option>, - mr: Option + mr: Option, } pub struct BuildInfo { @@ -51,7 +51,7 @@ pub struct BuildInfo { mr: Option, } -#[derive(Debug, Serialize, Default,ToSchema)] +#[derive(Debug, Serialize, Default, ToSchema)] pub enum TaskStatusEnum { Building, Interrupted, // exit code is None @@ -61,7 +61,7 @@ pub enum TaskStatusEnum { NotFound, } -#[derive(Debug, Serialize, Default,ToSchema)] +#[derive(Debug, Serialize, Default, ToSchema)] pub struct TaskStatus { status: TaskStatusEnum, #[serde(skip_serializing_if = "Option::is_none")] @@ -83,8 +83,7 @@ pub fn routers() -> OpenApiRouter { .routes(routes!(task_handler)) .routes(routes!(task_status_handler)) .route("/task-output/{id}", get(task_output_handler)) - .routes(routes!(task_query_by_mr)) - + .routes(routes!(task_query_by_mr)) } #[utoipa::path( @@ -472,7 +471,7 @@ async fn task_query_by_mr( Ok(models) if !models.is_empty() => { let dtos = models.into_iter().map(BuildDTO::from_model).collect(); Ok(Json(dtos)) - }, + } Ok(_) => Err(( StatusCode::NOT_FOUND, Json(serde_json::json!({ "message": "No builds found for the given MR" })), diff --git a/orion-server/src/model/builds.rs b/orion-server/src/model/builds.rs index 1ed2668c1..9afe99b2f 100644 --- a/orion-server/src/model/builds.rs +++ b/orion-server/src/model/builds.rs @@ -1,7 +1,7 @@ use sea_orm::entity::prelude::*; use serde::Serialize; -#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq,Serialize)] +#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize)] #[sea_orm(table_name = "builds")] pub struct Model { #[sea_orm(primary_key, auto_increment = false)] diff --git a/orion-server/src/server.rs b/orion-server/src/server.rs index a4ebe1726..951d3d449 100644 --- a/orion-server/src/server.rs +++ b/orion-server/src/server.rs @@ -1,15 +1,15 @@ use crate::api; use crate::api::AppState; use crate::model::builds; +use api::{BuildDTO, BuildRequest, TaskStatus, TaskStatusEnum}; use axum::Router; use axum::routing::get; use dashmap::DashMap; use sea_orm::{ConnectionTrait, Database, DatabaseConnection, DbErr, Schema, TransactionTrait}; -use utoipa::OpenApi; use std::net::SocketAddr; use std::sync::Arc; use tower_http::trace::TraceLayer; -use api::{BuildRequest, BuildDTO, TaskStatus, TaskStatusEnum}; +use utoipa::OpenApi; use utoipa_swagger_ui::SwaggerUi; #[derive(OpenApi)] #[openapi( diff --git a/scorpio/README.md b/scorpio/README.md index 9b7b29833..6fdbea052 100644 --- a/scorpio/README.md +++ b/scorpio/README.md @@ -80,6 +80,7 @@ git_author = "MEGA" git_email = "admin@mega.org" dicfuse_readable = "true" load_dir_depth = "3" +fetch_file_thread = "10" ``` ### `scorpio.toml` Configuration Guide: @@ -107,6 +108,9 @@ load_dir_depth = "3" - **`load_dir_depth`** Specifies how deep the file system should load and preload directories during initialization. +- **`fetch_file_thread`** + Sets the number of threads used for downloading files concurrently. + ### How to Contribute? diff --git a/scorpio/doc/perf_test.md b/scorpio/doc/perf_test.md index cdf7b52e4..e4644d295 100644 --- a/scorpio/doc/perf_test.md +++ b/scorpio/doc/perf_test.md @@ -174,3 +174,116 @@ Stat 操作速率: 195071.20 文件/秒 总结：cd切换文件夹路径，如果文件夹内容没有变化，也至少有一次网络开销，但实际上，目录变化不会那么频繁，所以不用每次都要进行同步 ``` + + +优化fetch_code之前，1000个文件拉取测试 + +文件创建脚本：script/run_1000_files.sh + +```sh +server running...ID:96cc1dd3-d9a8-453c-8976-4796e579a1d7|[init] REQRequest { unique: 2, uid: 0, gid: 0, pid: 0 }- - Call: +ID:96cc1dd3-d9a8-453c-8976-4796e579a1d7 [init] - Success: ReplyInit { max_write: 131072 } +new tree:903752ee634b1708c4d046a40cc62e572bb87a32 +new tree:68dfb43f330a0990660fbe5951be4cbbe4fda88e +new tree:29e9768f1ae1595cfc8229815a252adca1b75465 +new tree:03d5ccb9b5b209ba7d69fb5b3cc440c477847451 +new tree:15d966af87d8be73d17d0275e1e144ee1fd41e89 +new tree:352fa9971717d83086b387df815e1290ee1722a4 +new tree:49c68383b97e3ea515fd97c1badf4ddcc2d44851 +new tree:98218ce618733c6b21ae092cd2eb9c490a6f51fa +new tree:4bf75ab6553e47a5a13beb6af0168382626e008f +new tree:4078d08857c555366433a97f0492139751192b67 +new tree:370ce8cc1c2a1a70ed8605d16f3b2ebb2070c98d +new tree:ef3ab426ae85fc791cc7bc73f46c2847bcbeca90 +new tree:79648eba285af8fe672abc2447591665d37ac49f +finish store.... +finish code for third-party/t1...fetch 1000 files finished,use time: 68.8379431s + +new tree:903752ee634b1708c4d046a40cc62e572bb87a32 +new tree:68dfb43f330a0990660fbe5951be4cbbe4fda88e +new tree:29e9768f1ae1595cfc8229815a252adca1b75465 +new tree:03d5ccb9b5b209ba7d69fb5b3cc440c477847451 +new tree:15d966af87d8be73d17d0275e1e144ee1fd41e89 +new tree:98218ce618733c6b21ae092cd2eb9c490a6f51fa +new tree:352fa9971717d83086b387df815e1290ee1722a4 +new tree:4bf75ab6553e47a5a13beb6af0168382626e008f +new tree:49c68383b97e3ea515fd97c1badf4ddcc2d44851 +new tree:370ce8cc1c2a1a70ed8605d16f3b2ebb2070c98d +new tree:4078d08857c555366433a97f0492139751192b67 +new tree:ef3ab426ae85fc791cc7bc73f46c2847bcbeca90 +new tree:79648eba285af8fe672abc2447591665d37ac49f +finish store.... +finish code for third-party/t1...fetch 1000 files finished,use time: 65.558764133s + +new tree:903752ee634b1708c4d046a40cc62e572bb87a32 +new tree:68dfb43f330a0990660fbe5951be4cbbe4fda88e +new tree:29e9768f1ae1595cfc8229815a252adca1b75465 +new tree:03d5ccb9b5b209ba7d69fb5b3cc440c477847451 +new tree:15d966af87d8be73d17d0275e1e144ee1fd41e89 +new tree:98218ce618733c6b21ae092cd2eb9c490a6f51fa +new tree:352fa9971717d83086b387df815e1290ee1722a4 +new tree:4bf75ab6553e47a5a13beb6af0168382626e008f +new tree:49c68383b97e3ea515fd97c1badf4ddcc2d44851 +new tree:4078d08857c555366433a97f0492139751192b67 +new tree:370ce8cc1c2a1a70ed8605d16f3b2ebb2070c98d +new tree:79648eba285af8fe672abc2447591665d37ac49f +new tree:ef3ab426ae85fc791cc7bc73f46c2847bcbeca90 +finish store.... +finish code for third-party/t1...fetch 1000 files finished,use time: 66.745361716s +``` + +优化思路： + +添加一个全局的文件下载管理器，维护一个下载队列，在遍历目录的过程中，将需要下载的内容加载到队列中，此队列有10个固定协程进行下载 + +优化结果： + +```sh +Worker 4 shutting down +new tree:79648eba285af8fe672abc2447591665d37ac49f +Worker 0 shutting down +Worker 2 shutting down +Worker 3 shutting down +Worker 1 shutting down +finish store.... +Finished downloading code for third-party/t1 +fetch 1000 files finished,use time: 42.854706696s + +Worker 4 shutting down +finish store.... +Finished downloading code for third-party/t1 +fetch 1000 files finished,use time: 41.484116077s + +Worker 2 shutting down +Worker 0 shutting down +Worker 4 shutting down +new tree:79648eba285af8fe672abc2447591665d37ac49f +Worker 1 shutting down +Worker 3 shutting down +finish store.... +Finished downloading code for third-party/t1 +fetch 1000 files finished,use time: 44.768321595s + + +new tree:ef3ab426ae85fc791cc7bc73f46c2847bcbeca90 +Worker 0 shutting down +Worker 1 shutting down +Worker 4 shutting down +Directory processing completed for third-party/t1 +finish store.... +Finished downloading code for third-party/t1 +fetch 1000 files finished,use time: 45.409070254s + +Worker 0 shutting down +Worker 3 shutting down +Worker 1 shutting down +Worker 4 shutting down +new tree:ef3ab426ae85fc791cc7bc73f46c2847bcbeca90 +Worker 2 shutting down +Directory processing completed for third-party/t1 +finish store.... +Finished downloading code for third-party/t1 +fetch 1000 files finished,use time: 52.19500406s +``` + +由于系统本身性能存在波动，综合来看，大概可以实现50%的下载提速 diff --git a/scorpio/scorpio.toml b/scorpio/scorpio.toml index 0b9ffd8d7..340a02c7d 100644 --- a/scorpio/scorpio.toml +++ b/scorpio/scorpio.toml @@ -1,9 +1,10 @@ lfs_url = "http://localhost:8000" -store_path = "/home/luxian/megadir/store" +store_path = "/tmp/megadir/store" config_file = "config.toml" git_author = "MEGA" git_email = "admin@mega.org" -workspace = "/home/luxian/megadir/mount" +workspace = "/tmp/megadir/mount" base_url = "http://localhost:8000" dicfuse_readable = "true" -load_dir_depth = "3" \ No newline at end of file +load_dir_depth = "3" +fetch_file_thread = "10" \ No newline at end of file diff --git a/scorpio/script/run_1000_files.sh b/scorpio/script/run_1000_files.sh new file mode 100755 index 000000000..d02a8fc9e --- /dev/null +++ b/scorpio/script/run_1000_files.sh @@ -0,0 +1,111 @@ +#!/bin/bash + +set -e +# 参数 +TEST_DIR_BASE="/tmp" # 测试目录基础路径 +TOTAL_FILES=1000 # 总文件数 +MAX_DEPTH=6 # 最大目录深度 + +# 1. 创建一个时间相关的临时目录 +timestamp=$(date +%Y%m%d%H%M) +git_dir="test_$timestamp" +base_dir="$TEST_DIR_BASE/$git_dir" +echo "创建目录：$base_dir" +mkdir -p "$base_dir" + +# 2. 创建深层目录结构 (类似原始run.sh但支持更多文件) +generate_deep_structure() { + local base_path=$1 + local prefix=$2 + local depth_2_3_files=$3 # 深度2和3的文件数量 + local depth_4_files=$4 # 深度4的文件数量 + local depth_5_6_files=$5 # 深度5和6的文件数量 + + # 创建深层目录路径：prefix/2/3/4/5/6 + local current_path="$base_path/$prefix" + for level in {2..6}; do + current_path="$current_path/$level" + mkdir -p "$current_path" + + # 根据深度决定文件大小 (缩小10倍，总大小控制在100MB) + local file_size_kb + local files_count + case $level in + 2) file_size_kb=100; files_count=$depth_2_3_files ;; # 原来1MB -> 100KB + 3) file_size_kb=80; files_count=$depth_2_3_files ;; # 原来800KB -> 80KB + 4) file_size_kb=60; files_count=$depth_4_files ;; # 原来600KB -> 60KB + 5) file_size_kb=40; files_count=$depth_5_6_files ;; # 原来400KB -> 40KB + 6) file_size_kb=20; files_count=$depth_5_6_files ;; # 原来200KB -> 20KB + esac + + # 在每个层级创建多个文件 + for ((i=1; i<=files_count; i++)); do + local filename="${file_size_kb}K_${prefix}_${level}_${i}.bin" + echo "创建文件：$current_path/$filename" + head -c ${file_size_kb}K "$current_path/$filename" + done + done +} + +# 3. 计算文件分布以达到1000个文件 +calculate_file_distribution() { + echo "文件分布规划（基于原始run.sh结构，但扩展到1000个文件）：" + echo " 两个主要分支：1/ 和 2/" + echo " 每个分支各有5个深度层级 (2-6)" + echo " 每个层级的文件数量：" + echo " 深度2: 每分支50个文件，100KB each (100个文件，10MB)" + echo " 深度3: 每分支50个文件，80KB each (100个文件，8MB)" + echo " 深度4: 每分支100个文件，60KB each (200个文件，12MB)" + echo " 深度5: 每分支150个文件，40KB each (300个文件，12MB)" + echo " 深度6: 每分支150个文件，20KB each (300个文件，6MB)" + echo " 总计: 1000个文件，约48MB" +} + +# 4. 主要执行逻辑 +main() { + echo "开始创建1000个文件的测试结构..." + calculate_file_distribution + + local created_files=0 + + # 创建两个深层目录结构，每个结构500个文件 + echo "创建结构1..." + # 深度2: 50个文件，深度3: 50个文件，深度4: 100个文件，深度5: 150个文件，深度6: 150个文件 + generate_deep_structure "$base_dir" "1" 50 100 150 + created_files=$((created_files + 500)) + + echo "创建结构2..." + generate_deep_structure "$base_dir" "2" 50 100 150 + created_files=$((created_files + 500)) + + echo "" + echo "文件创建完成！" + echo " 深层目录结构1：1/2/3/4/5/6 (深度2,3各50个文件，深度4各100个文件，深度5,6各150个文件)" + echo " 深层目录结构2：2/2/3/4/5/6 (深度2,3各50个文件，深度4各100个文件，深度5,6各150个文件)" + echo " 总层数：6层" + echo " 总文件数：$created_files" + echo " 文件大小：20KB-100KB（随深度递减）" + echo " 预估总大小：约48MB" + echo " 基础目录: $base_dir" + + # 统计信息 + echo "" + echo "统计信息：" + echo " 总目录数: $(find "$base_dir" -type d | wc -l)" + echo " 总文件数: $(find "$base_dir" -type f | wc -l)" + echo " 总大小: $(du -sh "$base_dir" | cut -f1)" + + # 显示目录结构示例 + echo "" + echo "目录结构示例：" + tree "$base_dir" -L 3 -d 2>/dev/null || echo " (tree命令未安装，无法显示结构)" + + echo "" + echo "使用方法：" + echo " cd $base_dir" + echo " ls -la" + echo " find . -name '*.bin' | head -10" +} + +# 执行主函数 +main diff --git a/scorpio/src/dicfuse/store.rs b/scorpio/src/dicfuse/store.rs index 407682cab..2cb418f6c 100644 --- a/scorpio/src/dicfuse/store.rs +++ b/scorpio/src/dicfuse/store.rs @@ -827,10 +827,7 @@ pub async fn load_dir_depth(store: Arc, parent_path: String, ma let producers = Arc::clone(&active_producers); workers.push(tokio::spawn(async move { - while { - // If there are active producers or the queue is not empty, continue - producers.load(Ordering::Acquire) > 0 || !queue.is_empty() - } { + while producers.load(Ordering::Acquire) > 0 || !queue.is_empty() { if let Some(inode) = queue.pop() { //get the whole path. let path = diff --git a/scorpio/src/manager/fetch.rs b/scorpio/src/manager/fetch.rs index ce8c09502..d088cd947 100644 --- a/scorpio/src/manager/fetch.rs +++ b/scorpio/src/manager/fetch.rs @@ -1,28 +1,332 @@ +use super::{ScorpioManager, WorkDir}; +use crate::manager::store::store_trees; +use crate::scolfs; +use crate::util::config; +use crate::util::GPath; +use async_recursion::async_recursion; use axum::async_trait; use ceres::model::git::LatestCommitInfo; +use crossbeam::queue::SegQueue; +use futures::future::join_all; use mercury::hash::SHA1; use mercury::internal::object::tree::{Tree, TreeItemMode}; -use reqwest::Client; -use std::path::{Path, PathBuf}; -use std::str::FromStr; -use std::{collections::VecDeque, sync::Arc, time::Duration}; - -use async_recursion::async_recursion; use mercury::internal::object::{ commit::Commit, signature::{Signature, SignatureType}, }; +use reqwest::Client; +use std::collections::VecDeque; use std::io::Write; +use std::path::{Path, PathBuf}; +use std::str::FromStr; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::sync::OnceLock; +use tokio::sync::mpsc; use tokio::sync::mpsc::Sender; +use tokio::sync::watch; use tokio::sync::Mutex; +use tokio::sync::Notify; use tokio::time; +use tokio::time::Duration; + +///Download a file needs it's blob_id and save_path. +#[derive(Debug, Clone)] +pub struct DownloadTask { + file_id: SHA1, + save_path: PathBuf, + retry_count: u32, +} -use crate::manager::store::store_trees; -use crate::scolfs; -use crate::util::config; -use crate::util::GPath; +impl DownloadTask { + pub fn new(file_id: SHA1, save_path: PathBuf) -> Self { + Self { + file_id, + save_path, + retry_count: 0, + } + } -use super::{ScorpioManager, WorkDir}; + /// Create a retry task with incremented retry count + pub fn retry(&self) -> Self { + Self { + file_id: self.file_id, + save_path: self.save_path.clone(), + retry_count: self.retry_count + 1, + } + } + + /// Check if the task has exceeded maximum retry attempts + pub fn is_max_retries_exceeded(&self) -> bool { + self.retry_count >= 3 + } +} + +/// DownloadManager is responsible for managing file download operations in a concurrent manner. +/// +/// ## File Download Flow: +/// 1. **Directory Processing Phase**: +/// - Directory workers traverse the file tree using BFS (Breadth-First Search) +/// - They create directories inline and enqueue file download tasks +/// - The `directory_processing_sender` tracks whether directory traversal is still ongoing +/// +/// 2. **File Download Phase**: +/// - File download tasks are processed by a fixed pool of worker threads +/// - Each task downloads a file from the server using its SHA1 hash +/// - Workers update the `pending_tasks` counter as they complete downloads +/// +/// 3. **Completion Coordination**: +/// - The system only considers the entire operation complete when: +/// a) Directory processing is finished (directory_processing_sender = false) +/// b) All file downloads are complete (pending_tasks = 0) +/// - A completion coordinator monitors both conditions and notifies waiters +/// +/// ## Key Components: +/// - `sender`: Channel for enqueuing new download tasks +/// - `pending_tasks`: Atomic counter tracking active download operations +/// - `completion_notify`: Notifies when all operations are complete +/// - `directory_processing_sender/receiver`: Tracks directory traversal state +/// +/// This design ensures that we don't prematurely consider downloads complete +/// while directories are still being processed and potentially creating new download tasks. +pub struct DownloadManager { + sender: mpsc::UnboundedSender, // used to add new download tasks + #[allow(unused)] + worker_handles: Vec>, + pending_tasks: Arc, //the number of pending download tasks + completion_notify: Arc, // used to notify when all tasks are done + directory_processing_sender: watch::Sender, // watch the dir processing state + directory_processing_receiver: watch::Receiver, +} + +static DOWNLOAD_MANAGER: OnceLock = OnceLock::new(); + +impl DownloadManager { + /// Creates a new DownloadManager with the specified number of worker threads. + /// + /// The manager starts with directory processing enabled (true) and spawns + /// worker threads to handle file download tasks concurrently. + pub fn new(worker_count: usize) -> Self { + let (sender, receiver) = mpsc::unbounded_channel(); + let receiver = Arc::new(Mutex::new(receiver)); + + let pending_tasks = Arc::new(AtomicUsize::new(0)); + let completion_notify = Arc::new(Notify::new()); + + let (directory_processing_sender, directory_processing_receiver) = watch::channel(true); + + let worker_handles = (0..worker_count) + .map(|worker_id| { + let receiver = receiver.clone(); + let sender = sender.clone(); + let pending_tasks = pending_tasks.clone(); + let completion_notify = completion_notify.clone(); + let directory_receiver = directory_processing_receiver.clone(); + tokio::spawn(async move { + Self::worker_loop( + worker_id, + receiver, + sender, + pending_tasks, + completion_notify, + directory_receiver, + ) + .await; + }) + }) + .collect(); + + Self { + sender, + worker_handles, + pending_tasks, + completion_notify, + directory_processing_sender, + directory_processing_receiver, + } + } + + /// Starts the completion coordinator to handle task completion notifications. + /// + /// The coordinator monitors directory processing state changes and automatically + /// notifies waiters when both directory processing is complete AND no files are pending. + pub fn start_completion_coordinator(&self) { + let completion_notify = self.completion_notify.clone(); + let pending_tasks = self.pending_tasks.clone(); + let mut directory_receiver = self.directory_processing_receiver.clone(); + + tokio::spawn(async move { + loop { + if directory_receiver.changed().await.is_err() { + break; + } + + let directory_processing = *directory_receiver.borrow(); + let has_pending = pending_tasks.load(Ordering::Relaxed) > 0; + + if !directory_processing && !has_pending { + completion_notify.notify_waiters(); + break; + } + } + }); + } + + /// Worker loop that processes download tasks from the queue. + /// + /// Each worker continuously: + /// 1. Waits for download tasks from the shared queue + /// 2. Downloads files using fetch_and_save_file() + /// 3. If download fails and retries are available, re-enqueues the task + /// 4. Updates the pending task counter only on success or max retries exceeded + /// 5. Notifies completion when it's the last task AND directory processing is done + async fn worker_loop( + worker_id: usize, + receiver: Arc>>, + sender: mpsc::UnboundedSender, + pending_tasks: Arc, + completion_notify: Arc, + directory_receiver: watch::Receiver, + ) { + loop { + let task = { + let mut rx = receiver.lock().await; + rx.recv().await + }; + + match task { + Some(task) => { + match fetch_and_save_file(&task.file_id, &task.save_path).await { + Ok(_) => { + // Download successful, proceed to decrement counter + } + Err(e) => { + if task.is_max_retries_exceeded() { + eprintln!( + "Worker {}: Failed to download file {} (path: {}) after {} retries, giving up: {}", + worker_id, task.file_id, task.save_path.display(), task.retry_count, e + ); + // Max retries exceeded, proceed to decrement counter + } else { + eprintln!( + "Worker {}: Failed to download file {} (path: {}) on attempt {}, retrying: {}", + worker_id, task.file_id, task.save_path.display(), task.retry_count + 1, e + ); + + // Create retry task and re-enqueue + let retry_task = task.retry(); + if let Err(retry_err) = sender.send(retry_task) { + eprintln!( + "Worker {}: Failed to re-enqueue retry task for file {} (path: {}): {}", + worker_id, task.file_id, task.save_path.display(), retry_err + ); + // If we can't re-enqueue, we still need to decrement the counter + } else { + // Successfully re-enqueued, don't decrement counter + continue; + } + } + } + } + + let remaining = pending_tasks.fetch_sub(1, Ordering::Relaxed); + + if remaining == 1 { + let directory_processing = *directory_receiver.borrow(); + if !directory_processing { + completion_notify.notify_waiters(); + } + } + } + None => { + break; + } + } + } + } + + /// Enqueue a file download task to be processed by worker threads. + /// + /// This increments the pending task counter and sends the task to workers. + /// Returns an error if the channel is closed. + pub fn enqueue_download(&self, task: DownloadTask) -> Result<(), String> { + self.pending_tasks.fetch_add(1, Ordering::Relaxed); + self.sender + .send(task) + .map_err(|_| "Failed to enqueue download task".to_string()) + } + + pub fn has_pending_tasks(&self) -> bool { + self.pending_tasks.load(Ordering::Relaxed) > 0 + } + + /// Notifies that directory processing has completed. + /// + /// This is called when all directory traversal workers have finished. + /// If no files are pending, it immediately notifies completion. + pub fn notify_directory_processing_complete(&self) { + let _ = self.directory_processing_sender.send(false); + + if !self.has_pending_tasks() { + self.completion_notify.notify_waiters(); + } + } + + /// Returns true if directory processing is still ongoing. + pub fn is_directory_processing(&self) -> bool { + *self.directory_processing_receiver.borrow() + } + + /// Waits for all operations to complete using an efficient notification system. + /// + /// This method waits until BOTH conditions are met: + /// 1. Directory processing is complete (no more directories being traversed) + /// 2. All file downloads are complete (pending_tasks == 0) + /// + /// Uses tokio::select! to efficiently wait for state changes rather than busy polling. + pub async fn wait_for_completion(&self) { + loop { + let directory_processing = self.is_directory_processing(); + let has_pending = self.has_pending_tasks(); + + if !directory_processing && !has_pending { + return; + } + + let mut directory_receiver = self.directory_processing_receiver.clone(); + + tokio::select! { + _ = self.completion_notify.notified() => { + continue; + } + _ = directory_receiver.changed() => { + continue; + } + // _ = tokio::time::sleep(tokio::time::Duration::from_millis(100)) => { + // continue; + // } + } + } + } + + pub fn get_global() -> &'static DownloadManager { + DOWNLOAD_MANAGER.get_or_init(|| { + let worker_count = config::fetch_file_thread(); + println!("Initializing global download manager with {worker_count} workers"); + DownloadManager::new(worker_count) + }) + } +} + +/// Enqueue a file download task,the common entry point for downloading files. +pub fn enqueue_file_download(file_id: SHA1, save_path: PathBuf) { + let download_manager = DownloadManager::get_global(); + let task = DownloadTask::new(file_id, save_path); + + if let Err(e) = download_manager.enqueue_download(task) { + eprintln!("Failed to enqueue download task for file {file_id}: {e}"); + } +} #[allow(unused)] #[async_trait] @@ -65,6 +369,7 @@ impl CheckHash for ScorpioManager { } } + #[allow(unused)] async fn fetch + std::marker::Send>( &mut self, inode: u64, @@ -209,6 +514,7 @@ async fn worker_thread( /// Network operations, recursively read the main Tree from the pipe, and download Blobs objects #[async_recursion] +#[allow(unused)] async fn worker_ro_thread( root_path: GPath, target_path: Arc, @@ -251,39 +557,154 @@ async fn worker_ro_thread( /// /// Download remote data to local and store it in Overlay format async fn fetch_code(path: &GPath, save_path: impl AsRef) -> std::io::Result<()> { - let target_path: Arc = Arc::new(save_path.as_ref().to_path_buf()); + let target_path = save_path.as_ref().to_path_buf(); + // println!("fetch_code starting for path: {path}"); + + let download_manager = DownloadManager::get_global(); + download_manager.start_completion_coordinator(); + let _ = download_manager.directory_processing_sender.send(true); // Create the save_path directory if it doesn't exist tokio::fs::create_dir_all(&save_path).await?; - let rece; - let handle; - { - // Set up a pipeline with a capacity of 100 - let (send, _rece) = tokio::sync::mpsc::channel::<(GPath, Tree)>(100); - rece = _rece; - - let p: GPath = path.clone(); - let sc = send.clone(); - let ps = target_path.clone(); - // Use child threads to operate mounted directory - handle = tokio::spawn(async move { - worker_ro_thread(p.clone(), ps, p, sc).await; - }); + + // Setup tree storage channel + let (tree_sender, tree_receiver) = mpsc::channel::<(GPath, Tree)>(1000); + + // Use simple queue for directory processing, similar to load_dir_depth + let queue = Arc::new(SegQueue::new()); + + // Fetch and process the initial/root directory + let initial_tree = fetch_tree(path).await.map_err(std::io::Error::other)?; + + // Send tree to storage + if let Err(e) = tree_sender.send((path.clone(), initial_tree.clone())).await { + eprintln!("Failed to send initial tree: {e}"); } + let dir_count = initial_tree + .tree_items + .iter() + .filter(|item| item.mode == TreeItemMode::Tree) + .count(); + + let active_producers = Arc::new(AtomicUsize::new(dir_count)); + + for item in initial_tree.tree_items { + let item_name = item.name.clone(); + let real_path = target_path.join(&item_name); + + if item.mode == TreeItemMode::Tree { + // Create directory and add to queue for processing + if let Err(e) = tokio::fs::create_dir_all(&real_path).await { + eprintln!("Failed to create directory {real_path:?}: {e}"); + // If we failed to create directory, reduce the producer count + active_producers.fetch_sub(1, Ordering::Release); + continue; + } + + let mut subpath = path.clone(); + subpath.push(item_name); + queue.push((subpath, real_path)); + } else { + // Enqueue file for download + enqueue_file_download(item.id, real_path); + } + } + + let worker_count = 5; + let mut workers = Vec::with_capacity(worker_count); + + for worker_id in 0..worker_count { + let queue = Arc::clone(&queue); + let tree_sender = tree_sender.clone(); + let producers = Arc::clone(&active_producers); + + workers.push(tokio::spawn(async move { + while producers.load(Ordering::Acquire) > 0 || !queue.is_empty() { + // println!("worker_id {:?} left {:?} {:?}", worker_id, current_count, queue_len); + if let Some((current_path, current_target)) = queue.pop() { + // println!("Worker {} processing path: {}", worker_id, current_path); + + // Fetch tree for this directory + match fetch_tree(¤t_path).await { + Ok(tree) => { + // Send tree to storage + if let Err(e) = + tree_sender.send((current_path.clone(), tree.clone())).await + { + eprintln!("Worker {worker_id}: Failed to send tree: {e}"); + } + + // Process each item in the tree + for item in tree.tree_items { + let item_name = item.name.clone(); + let item_real_path = current_target.join(&item_name); + + if item.mode == TreeItemMode::Tree { + // Create directory + if let Err(_e) = + tokio::fs::create_dir_all(&item_real_path).await + { + continue; + } + + // Add subdirectory to queue and increment producer count + let mut subpath = current_path.clone(); + subpath.push(item_name); + producers.fetch_add(1, Ordering::Release); + queue.push((subpath, item_real_path)); + } else { + // Enqueue file for download + enqueue_file_download(item.id, item_real_path); + } + } + } + Err(e) => { + eprintln!( + "Worker {worker_id}: Failed to fetch tree for {current_path}: {e}", + ); + } + } + + producers.fetch_sub(1, Ordering::Release); + } else { + if producers.load(Ordering::Acquire) == 0 { + return; + } + tokio::task::yield_now().await; + } + } + // println!("Worker {worker_id} shutting down"); + })); + } + + // Drop tree sender to allow storage to finish + drop(tree_sender); + let storepath = save_path.as_ref().parent().unwrap().join("tree.db"); - store_trees(storepath.to_str().unwrap(), rece).await?; + let store_handle = tokio::spawn(async move { + if let Err(e) = store_trees(storepath.to_str().unwrap(), tree_receiver).await { + eprintln!("Failed to store trees: {e}"); + } + }); + + // Wait for all workers to complete + join_all(workers).await; + + DownloadManager::get_global().notify_directory_processing_complete(); + // println!("Directory processing completed for {path}"); + + let _ = store_handle.await; - // Clean up workers (depends on how you implement worker_thread termination) - let _ = handle.await; + // Wait for all downloads to complete + DownloadManager::get_global().wait_for_completion().await; - //get lfs file + // Get LFS files scolfs::lfs::lfs_restore(&path.to_string(), save_path.as_ref().to_str().unwrap()) .await .unwrap(); - print!("finish code for {path}..."); - + println!("Finished downloading code for {path}"); Ok(()) } @@ -379,16 +800,21 @@ async fn fetch_and_save_file( /// Network operations, extracting Tree objects from HTTP byte streams #[allow(unused)] -pub async fn fetch_tree(path: &GPath) -> Result> { +pub async fn fetch_tree(path: &GPath) -> Result { let url = format!("{}{}", config::tree_file_endpoint(), path); - let response = reqwest::get(&url).await?; + let response = reqwest::get(&url) + .await + .map_err(|e| format!("Request failed: {e}"))?; if response.status().is_success() { - let bytes = response.bytes().await?; - let tree = Tree::try_from(&bytes[..])?; + let bytes = response + .bytes() + .await + .map_err(|e| format!("Failed to read response: {e}"))?; + let tree = Tree::try_from(&bytes[..]).map_err(|e| format!("Failed to parse tree: {e}"))?; Ok(tree) } else { - Err(format!("Failed to fetch tree: {}", response.status()).into()) + Err(format!("Failed to fetch tree: {}", response.status())) } } diff --git a/scorpio/src/manager/store.rs b/scorpio/src/manager/store.rs index 1af625394..c031fbed4 100644 --- a/scorpio/src/manager/store.rs +++ b/scorpio/src/manager/store.rs @@ -84,11 +84,11 @@ impl CommitStore for sled::Db { pub async fn store_trees(storepath: &str, mut tree_channel: Receiver<(GPath, Tree)>) -> Result<()> { let db = sled::open(storepath)?; while let Some((path, tree)) = tree_channel.recv().await { - println!("new tree:{}", tree.id); + // println!("new tree:{}", tree.id); db.insert_tree(path.into(), tree); } - println!("finish store...."); + // println!("finish store...."); Ok(()) } diff --git a/scorpio/src/util/config.rs b/scorpio/src/util/config.rs index 0b7ace083..134d21ef7 100644 --- a/scorpio/src/util/config.rs +++ b/scorpio/src/util/config.rs @@ -15,7 +15,10 @@ pub type ConfigResult = Result; pub struct ScorpioConfig { config: HashMap, } + const DEFAULT_LOAD_DIR_DEPTH: usize = 3; +const DEFAULT_FETCH_FILE_THREAD: usize = 10; + // Global configuration management static SCORPIO_CONFIG: OnceLock = OnceLock::new(); @@ -27,8 +30,8 @@ static SCORPIO_CONFIG: OnceLock = OnceLock::new(); /// # Returns /// `ConfigResult<()>` - Success or error pub fn init_config(path: &str) -> ConfigResult<()> { - let content = fs::read_to_string(path) - .map_err(|e| format!("Config file not found at '{path}': {e}"))?; + let content = + fs::read_to_string(path).map_err(|e| format!("Config file not found at '{path}': {e}"))?; let mut config: HashMap = toml::from_str(&content).map_err(|e| format!("Invalid config format: {e}"))?; @@ -156,6 +159,10 @@ fn get_config() -> &'static ScorpioConfig { "http://localhost:8000/lfs".to_string(), ); config.insert("dicfuse_readable".to_string(), "true".to_string()); + config.insert( + "fetch_file_thread".to_string(), + DEFAULT_FETCH_FILE_THREAD.to_string(), + ); // Create required directories for path in [config["workspace"].as_str(), config["store_path"].as_str()] { @@ -196,6 +203,7 @@ fn validate(config: &mut HashMap) -> ConfigResult<()> { "lfs_url", "dicfuse_readable", "load_dir_depth", + "fetch_file_thread", ]; for key in required_keys { @@ -254,6 +262,15 @@ pub fn load_dir_depth() -> usize { .and_then(|s| s.parse().ok()) .unwrap_or(DEFAULT_LOAD_DIR_DEPTH) } + +///Get the number of file download threads +pub fn fetch_file_thread() -> usize { + get_config() + .config + .get("fetch_file_thread") + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_FETCH_FILE_THREAD) +} #[cfg(test)] mod tests { use super::*; @@ -269,6 +286,7 @@ mod tests { lfs_url = "http://localhost:8000/lfs" dicfuse_readable = "true" load_dir_depth = "3" + fetch_file_thread = "10" "#; let config_path = "/tmp/scorpio.toml"; std::fs::write(config_path, config_content).expect("Failed to write test config file"); @@ -289,6 +307,7 @@ mod tests { "http://localhost:8000/api/v1/file/blob" ); assert_eq!(load_dir_depth(), 3); + assert_eq!(fetch_file_thread(), 10); assert_eq!(config_file(), "config.toml"); } } diff --git a/vault/src/integration/vault_core.rs b/vault/src/integration/vault_core.rs index 199aa05ab..c2f8b27ed 100644 --- a/vault/src/integration/vault_core.rs +++ b/vault/src/integration/vault_core.rs @@ -1,7 +1,10 @@ -use std::{path::PathBuf, sync::{Arc, RwLock}}; use crate::integration::jupiter_backend::JupiterBackend; use common::errors::MegaError; use jupiter::storage::Storage; +use std::{ + path::PathBuf, + sync::{Arc, RwLock}, +}; use rusty_vault::{ core::Core, @@ -14,8 +17,6 @@ use tracing::log; const CORE_KEY_FILE: &str = "core_key.json"; // where the core key is stored, like `root_token` - - #[derive(Debug, Clone, Serialize, Deserialize)] struct CoreKey { secret_shares: Vec>, @@ -69,7 +70,6 @@ impl VaultCore { }; let core = Arc::new(RwLock::new(core)); - let key = { let mut managed_core = core.write().unwrap(); managed_core @@ -77,8 +77,6 @@ impl VaultCore { .expect("Failed to configure vault core"); let core_key = if !key_path.exists() { - - let result = managed_core .init(&seal_config) .expect("Failed to initialize vault"); @@ -86,7 +84,10 @@ impl VaultCore { secret_shares: Vec::from(&result.secret_shares[..]), root_token: result.root_token, }; - println!("[vault] Creating new core_key.json at: {}", key_path.display()); + println!( + "[vault] Creating new core_key.json at: {}", + key_path.display() + ); let file = std::fs::File::create(&key_path).unwrap(); serde_json::to_writer_pretty(file, &core_key).unwrap(); core_key @@ -200,8 +201,7 @@ mod tests { "Vault core should be initialized" ); } - - + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn test_vault_api() { let temp_dir = tempfile::tempdir().expect("Failed to create temporary directory");