diff --git a/ceres/src/pack/import_repo.rs b/ceres/src/pack/import_repo.rs index 11e17d20d..6127f7a1d 100644 --- a/ceres/src/pack/import_repo.rs +++ b/ceres/src/pack/import_repo.rs @@ -1,19 +1,20 @@ use std::{ collections::{HashMap, HashSet}, str::FromStr, - sync::{atomic::{AtomicUsize, Ordering}, mpsc::Receiver}, + sync::{ + atomic::{AtomicUsize, Ordering}, + mpsc::Receiver, + }, }; use async_trait::async_trait; +use futures::{future::join_all, StreamExt}; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use callisto::raw_blob; use common::errors::MegaError; -use jupiter::{ - context::Context, - storage::{batch_query_by_columns, GitStorageProvider}, -}; +use jupiter::{context::Context, storage::GitStorageProvider}; use mercury::{ errors::GitError, internal::{ @@ -50,14 +51,21 @@ impl PackHandler for ImportRepo { async fn save_entry(&self, receiver: Receiver) -> Result<(), GitError> { let storage = self.context.services.git_db_storage.clone(); - let mut entry_list = Vec::new(); + let mut entry_list = vec![]; + let mut join_tasks = vec![]; for entry in receiver { entry_list.push(entry); if entry_list.len() >= 1000 { - storage.save_entry(&self.repo, entry_list).await.unwrap(); - entry_list = Vec::new(); + let stg_clone = storage.clone(); + let repo_clone = self.repo.clone(); + let handle = tokio::spawn(async move { + stg_clone.save_entry(&repo_clone, entry_list).await.unwrap(); + }); + join_tasks.push(handle); + entry_list = vec![]; } } + join_all(join_tasks).await; storage.save_entry(&self.repo, entry_list).await.unwrap(); Ok(()) } @@ -74,43 +82,65 @@ impl PackHandler for ImportRepo { let repo = self.repo.clone(); tokio::spawn(async move { - let commits = storage.get_commits_by_repo_id(&repo).await.unwrap(); - for m in commits.into_iter() { - let c: Commit = m.into(); - let entry: Entry = c.into(); - entry_tx.send(entry).await.unwrap(); + let mut commit_stream = storage.get_commits_by_repo_id(&repo).await.unwrap(); + + while let Some(model) = commit_stream.next().await { + match model { + Ok(m) => { + let c: Commit = m.into(); + let entry = c.into(); + entry_tx.send(entry).await.unwrap(); + } + Err(err) => eprintln!("Error: {:?}", err), + } } + tracing::info!("send commits end"); - let trees: Vec = - storage.get_trees_by_repo_id(&repo).await.unwrap(); - for m in trees.into_iter() { - let c: Tree = m.into(); - let entry: Entry = c.into(); - entry_tx.send(entry).await.unwrap(); + let mut tree_stream = storage.get_trees_by_repo_id(&repo).await.unwrap(); + while let Some(model) = tree_stream.next().await { + match model { + Ok(m) => { + let t: Tree = m.into(); + let entry = t.into(); + entry_tx.send(entry).await.unwrap(); + } + Err(err) => eprintln!("Error: {:?}", err), + } } + tracing::info!("send trees end"); - let bids: Vec = storage - .get_blobs_by_repo_id(&repo) - .await - .unwrap() - .into_iter() - .map(|b| b.blob_id) - .collect(); - let raw_blobs = batch_query_by_columns::( - storage.get_connection(), - raw_blob::Column::Sha1, - bids, - None, - None, - ) - .await - .unwrap(); - for m in raw_blobs { - // todo handle storage type - let c: Blob = m.into(); - let entry: Entry = c.into(); - entry_tx.send(entry).await.unwrap(); + let mut bid_stream = storage.get_blobs_by_repo_id(&repo).await.unwrap(); + let mut bids = vec![]; + while let Some(model) = bid_stream.next().await { + match model { + Ok(m) => bids.push(m.blob_id), + Err(err) => eprintln!("Error: {:?}", err), + } + } + + let mut blob_handler = vec![]; + for chunk in bids.chunks(10000) { + let stg_clone = storage.clone(); + let sender_clone = entry_tx.clone(); + let chunk_clone = chunk.to_vec(); + let handler = tokio::spawn(async move { + let mut blob_stream = stg_clone.get_raw_blobs(chunk_clone).await.unwrap(); + while let Some(model) = blob_stream.next().await { + match model { + Ok(m) => { + // todo handle storage type + let b: Blob = m.into(); + let entry: Entry = b.into(); + sender_clone.send(entry).await.unwrap(); + } + Err(err) => eprintln!("Error: {:?}", err), + } + } + }); + blob_handler.push(handler); } + join_all(blob_handler).await; + tracing::info!("send blobs end"); let tags = storage.get_tags_by_repo_id(&repo).await.unwrap(); for m in tags.into_iter() { @@ -119,6 +149,7 @@ impl PackHandler for ImportRepo { entry_tx.send(entry).await.unwrap(); } drop(entry_tx); + tracing::info!("sending all object end..."); }); Ok(ReceiverStream::new(stream_rx)) diff --git a/jupiter/src/storage/git_db_storage.rs b/jupiter/src/storage/git_db_storage.rs index 56cac3828..9431f6991 100644 --- a/jupiter/src/storage/git_db_storage.rs +++ b/jupiter/src/storage/git_db_storage.rs @@ -1,14 +1,15 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; +use futures::Stream; use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; use sea_orm::{ - ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, IntoActiveModel, QueryFilter, - Set, + ActiveModelTrait, ColumnTrait, DatabaseConnection, DbErr, EntityTrait, IntoActiveModel, + QueryFilter, Set, }; use sea_orm::{PaginatorTrait, QueryOrder}; -use callisto::{git_blob, git_commit, git_repo, git_tag, git_tree, import_refs}; +use callisto::{git_blob, git_commit, git_repo, git_tag, git_tree, import_refs, raw_blob}; use common::config::StorageConfig; use common::errors::MegaError; use mercury::internal::object::GitObjectModel; @@ -22,7 +23,7 @@ use crate::{ storage::GitStorageProvider, }; -use super::batch_save_model; +use crate::storage::batch_save_model; #[derive(Clone)] pub struct GitDbStorage { @@ -232,24 +233,25 @@ impl GitDbStorage { .unwrap()) } - pub async fn get_commits_by_repo_id( - &self, - repo: &Repo, - ) -> Result, MegaError> { - Ok(git_commit::Entity::find() + pub async fn get_commits_by_repo_id<'a>( + &'a self, + repo: &'a Repo, + ) -> Result> + 'a + Send, MegaError> { + let stream = git_commit::Entity::find() .filter(git_commit::Column::RepoId.eq(repo.repo_id)) - .all(self.get_connection()) + .stream(self.get_connection()) .await - .unwrap()) + .unwrap(); + Ok(stream) } - pub async fn get_trees_by_repo_id( - &self, - repo: &Repo, - ) -> Result, MegaError> { + pub async fn get_trees_by_repo_id<'a>( + &'a self, + repo: &'a Repo, + ) -> Result> + 'a + Send, MegaError> { Ok(git_tree::Entity::find() .filter(git_tree::Column::RepoId.eq(repo.repo_id)) - .all(self.get_connection()) + .stream(self.get_connection()) .await .unwrap()) } @@ -280,13 +282,24 @@ impl GitDbStorage { .unwrap()) } - pub async fn get_blobs_by_repo_id( - &self, - repo: &Repo, - ) -> Result, MegaError> { + pub async fn get_blobs_by_repo_id<'a>( + &'a self, + repo: &'a Repo, + ) -> Result> + 'a + Send, MegaError> { Ok(git_blob::Entity::find() .filter(git_blob::Column::RepoId.eq(repo.repo_id)) - .all(self.get_connection()) + .stream(self.get_connection()) + .await + .unwrap()) + } + + pub async fn get_raw_blobs( + &self, + hashes: Vec, + ) -> Result> + '_ + Send, MegaError> { + Ok(raw_blob::Entity::find() + .filter(raw_blob::Column::Sha1.is_in(hashes)) + .stream(self.get_connection()) .await .unwrap()) } @@ -312,13 +325,11 @@ impl GitDbStorage { .await .unwrap(); - let bids: Vec = self - .get_blobs_by_repo_id(repo) + let b_count = git_blob::Entity::find() + .filter(git_blob::Column::RepoId.eq(repo.repo_id)) + .count(self.get_connection()) .await - .unwrap() - .into_iter() - .map(|b| b.blob_id) - .collect(); + .unwrap(); let tag_count = git_tag::Entity::find() .filter(git_tag::Column::RepoId.eq(repo.repo_id)) @@ -326,7 +337,7 @@ impl GitDbStorage { .await .unwrap(); - (c_count + t_count + bids.len() as u64 + tag_count) + (c_count + t_count + b_count + tag_count) .try_into() .unwrap() } diff --git a/jupiter/src/storage/mega_storage.rs b/jupiter/src/storage/mega_storage.rs index 2532f01aa..a2d7c2364 100644 --- a/jupiter/src/storage/mega_storage.rs +++ b/jupiter/src/storage/mega_storage.rs @@ -290,13 +290,6 @@ impl MegaStorage { .unwrap()) } - pub async fn get_commits(&self) -> Result, MegaError> { - Ok(mega_commit::Entity::find() - .all(self.get_connection()) - .await - .unwrap()) - } - pub async fn get_tree_by_hash( &self, hash: &str, diff --git a/jupiter/src/storage/mod.rs b/jupiter/src/storage/mod.rs index 97a6477de..88e89b94c 100644 --- a/jupiter/src/storage/mod.rs +++ b/jupiter/src/storage/mod.rs @@ -7,10 +7,7 @@ pub mod mega_storage; use async_trait::async_trait; use common::errors::MegaError; -use sea_orm::{ - sea_query::OnConflict, ActiveModelTrait, ColumnTrait, ConnectionTrait, DatabaseConnection, - EntityTrait, QueryFilter, -}; +use sea_orm::{sea_query::OnConflict, ActiveModelTrait, ConnectionTrait, EntityTrait}; use venus::{ import_repo::import_refs::{RefCommand, Refs}, import_repo::repo::Repo, @@ -96,30 +93,3 @@ where futures::future::join_all(results).await; Ok(()) } - -#[allow(unused)] -pub async fn batch_query_by_columns( - connection: &DatabaseConnection, - column: C, - ids: Vec, - filter_column: Option, - value: Option, -) -> Result, MegaError> -where - T: EntityTrait, - C: ColumnTrait, -{ - let mut result = Vec::::new(); - for chunk in ids.chunks(1000) { - let query_builder = T::find().filter(column.is_in(chunk)); - - // Conditionally add the filter based on the value parameter - let query_builder = match value { - Some(ref v) => query_builder.filter(filter_column.unwrap().eq(v)), - None => query_builder, - }; - - result.extend(query_builder.all(connection).await?); - } - Ok(result) -} diff --git a/mercury/src/internal/object/commit.rs b/mercury/src/internal/object/commit.rs index 011bb4843..16acb5b48 100644 --- a/mercury/src/internal/object/commit.rs +++ b/mercury/src/internal/object/commit.rs @@ -14,6 +14,8 @@ use std::fmt::Display; use std::str::FromStr; +use bstr::ByteSlice; + use crate::errors::GitError; use crate::hash::SHA1; use crate::internal::object::signature::Signature; @@ -108,64 +110,48 @@ impl ObjectTrait for Commit { where Self: Sized, { - let commit = data; - - // Find the tree id - let tree_end = commit - .iter() - .position(|&b| b == 0x0a) - .ok_or(GitError::InvalidCommitObject)?; - let tree_id = SHA1::from_str( - std::str::from_utf8(&commit[5..tree_end]).map_err(|_| GitError::InvalidCommitObject)?, + let mut commit = data; + // Find the tree id and remove it from the data + let tree_end = commit.find_byte(0x0a).unwrap(); + let tree_id: SHA1 = SHA1::from_str( + String::from_utf8(commit[5..tree_end].to_owned()) + .unwrap() + .as_str(), ) .unwrap(); - - // Find the parent commit ids - let author_begin = commit - .windows(6) - .position(|window| window == b"author") - .ok_or(GitError::InvalidCommitObject)?; - let parent_commit_ids = commit[..author_begin] - .windows(6) - .enumerate() - .filter_map(|(i, window)| { - if window == b"parent" { - let parent_end = commit[i + 7..].iter().position(|&b| b == 0x0a)?; - Some( - SHA1::from_str( - std::str::from_utf8(&commit[i + 7..i + 7 + parent_end]).ok()?, - ) - .ok(), - ) - } else { - None - } + let binding = commit[tree_end + 1..].to_vec(); + commit = &binding; + + // Find the parent commit ids and remove them from the data + let author_begin = commit.find("author").unwrap(); + let parent_commit_ids: Vec = commit[..author_begin] + .find_iter("parent") + .map(|parent| { + let parent_end = commit[parent..].find_byte(0x0a).unwrap(); + SHA1::from_str( + String::from_utf8(commit[parent + 7..parent + parent_end].to_owned()) + .unwrap() + .as_str(), + ) + .unwrap() }) - .collect::>>() - .ok_or(GitError::InvalidCommitObject)?; - - // Extract author - let author_end = commit[author_begin..] - .iter() - .position(|&b| b == 0x0a) - .ok_or(GitError::InvalidCommitObject)? - + author_begin; - let author = Signature::from_data(commit[author_begin..author_end].to_vec())?; - - // Extract committer - let committer_begin = author_end + 1; - let committer_end = commit[committer_begin..] - .iter() - .position(|&b| b == 0x0a) - .ok_or(GitError::InvalidCommitObject)? - + committer_begin; - let committer = Signature::from_data(commit[committer_begin..committer_end].to_vec())?; + .collect(); + let binding = commit[author_begin..].to_vec(); + commit = &binding; + + // Find the author and committer and remove them from the data + let author = + Signature::from_data(commit[..commit.find_byte(0x0a).unwrap()].to_vec()).unwrap(); + let binding = commit[commit.find_byte(0x0a).unwrap() + 1..].to_vec(); + commit = &binding; + let committer = + Signature::from_data(commit[..commit.find_byte(0x0a).unwrap()].to_vec()).unwrap(); // The rest is the message - let message = std::str::from_utf8(&commit[committer_end + 1..]) - .map_err(|_| GitError::InvalidCommitObject)? - .to_string(); - + let message = unsafe { + String::from_utf8_unchecked(commit[commit.find_byte(0x0a).unwrap() + 1..].to_vec()) + }; + Ok(Commit { id: hash, tree_id,