fix(mono): Unified the behavior during unpacking for both monorepo an…

ceres/src/pack/import_repo.rs

@@ -2,18 +2,12 @@ use std::{
     collections::{HashMap, HashSet},
     path::PathBuf,
     str::FromStr,
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc,
-    },
+    sync::atomic::{AtomicUsize, Ordering},
 };
 
 use async_trait::async_trait;
-use futures::{future::join_all, StreamExt};
-use tokio::sync::{
-    mpsc::{self, UnboundedReceiver},
-    Semaphore,
-};
+use futures::StreamExt;
+use tokio::sync::mpsc::{self};
 use tokio_stream::wrappers::ReceiverStream;
 
 use callisto::{mega_tree, raw_blob, sea_orm_active_enums::RefTypeEnum};
@@ -57,46 +51,17 @@ impl PackHandler for ImportRepo {
         self.find_head_hash(refs)
     }
 
-    async fn handle_receiver(
-        &self,
-        mut receiver: UnboundedReceiver<Entry>,
-    ) -> Result<Option<Commit>, GitError> {
+    async fn save_entry(&self, entry_list: Vec<Entry>) -> Result<(), MegaError> {
         let storage = self.storage.git_db_storage();
-        let mut entry_list = vec![];
-        let semaphore = Arc::new(Semaphore::new(8));
-        let mut join_tasks = vec![];
-        let repo_id = self.repo.repo_id;
+        storage.save_entry(self.repo.repo_id, entry_list).await
+    }
 
-        while let Some(entry) = receiver.recv().await {
-            entry_list.push(entry);
-            if entry_list.len() >= 1000 {
-                let stg_clone = storage.clone();
-                let acquired = semaphore.clone().acquire_owned().await.unwrap();
-
-                let handle = tokio::spawn(async move {
-                    let _acquired = acquired;
-                    stg_clone.save_entry(repo_id, entry_list).await
-                });
-                join_tasks.push(handle);
-                entry_list = vec![];
-            }
-        }
-        let results = join_all(join_tasks).await;
-        for (i, res) in results.into_iter().enumerate() {
-            match res {
-                Ok(Ok(())) => {}
-                Ok(Err(e)) => {
-                    tracing::error!("Task {} save_entry Err: {:?}", i, e);
-                }
-                Err(join_err) => {
-                    tracing::error!("Task {} panic or cancle: {:?}", i, join_err);
-                }
-            }
-        }
+    async fn post_receiver_handler(&self) -> Result<(), GitError> {
+        self.attach_to_monorepo_parent().await
+    }
 
-        storage.save_entry(repo_id, entry_list).await.unwrap();
-        self.attach_to_monorepo_parent().await.unwrap();
-        Ok(None)
+    async fn check_entry(&self, _: &Entry) -> Result<(), GitError> {
+        Ok(())
     }
 
     async fn full_pack(&self, _: Vec<String>) -> Result<ReceiverStream<Vec<u8>>, GitError> {
@@ -304,7 +269,7 @@ impl PackHandler for ImportRepo {
             .await
     }
 
-    async fn update_refs(&self, _: Option<Commit>, refs: &RefCommand) -> Result<(), GitError> {
+    async fn update_refs(&self, refs: &RefCommand) -> Result<(), GitError> {
         let storage = self.storage.git_db_storage();
         match refs.command_type {
             CommandType::Create => {

ceres/src/pack/mod.rs

@@ -1,14 +1,17 @@
 use std::{
     collections::HashSet,
     pin::Pin,
-    sync::atomic::{AtomicUsize, Ordering},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
 };
 
 use async_trait::async_trait;
 use bytes::Bytes;
-use futures::Stream;
+use futures::{future::join_all, Stream};
 use sysinfo::System;
-use tokio::sync::mpsc::UnboundedReceiver;
+use tokio::sync::{mpsc::UnboundedReceiver, Semaphore};
 use tokio_stream::wrappers::ReceiverStream;
 
 use crate::protocol::import_refs::{RefCommand, Refs};
@@ -18,7 +21,7 @@ use common::{
     errors::{MegaError, ProtocolError},
     utils::ZERO_ID,
 };
-use mercury::internal::{object::commit::Commit, pack::Pack};
+use mercury::internal::pack::Pack;
 use mercury::{
     errors::GitError,
     internal::{
@@ -34,13 +37,59 @@ pub mod import_repo;
 pub mod monorepo;
 
 #[async_trait]
-pub trait PackHandler: Send + Sync {
+pub trait PackHandler: Send + Sync + 'static {
     async fn head_hash(&self) -> (String, Vec<Refs>);
 
-    async fn handle_receiver(
-        &self,
+    async fn receiver_handler(
+        self: Arc<Self>,
         mut rx: UnboundedReceiver<Entry>,
-    ) -> Result<Option<Commit>, GitError>;
+    ) -> Result<(), GitError> {
+        let mut entry_list = vec![];
+        let semaphore = Arc::new(Semaphore::new(4));
+        let mut join_tasks = vec![];
+
+        while let Some(entry) = rx.recv().await {
+            self.check_entry(&entry).await?;
+            entry_list.push(entry);
+            if entry_list.len() >= 1000 {
+                let acquired = semaphore.clone().acquire_owned().await.unwrap();
+                let entries = std::mem::take(&mut entry_list);
+                let shared = self.clone();
+                let handle = tokio::spawn(async move {
+                    let _acquired = acquired;
+                    shared.save_entry(entries).await
+                });
+                join_tasks.push(handle);
+            }
+        }
+        // process left entries
+        if !entry_list.is_empty() {
+            let handler = self.clone();
+            let entries = std::mem::take(&mut entry_list);
+            let handle = tokio::spawn(async move { handler.save_entry(entries).await });
+            join_tasks.push(handle);
+        }
+
+        let results = join_all(join_tasks).await;
+        for (i, res) in results.into_iter().enumerate() {
+            match res {
+                Ok(Ok(())) => {}
+                Ok(Err(e)) => {
+                    tracing::error!("Task {} save_entry Err: {:?}", i, e);
+                }
+                Err(join_err) => {
+                    tracing::error!("Task {} panic or cancle: {:?}", i, join_err);
+                }
+            }
+        }
+        self.post_receiver_handler().await
+    }
+
+    async fn save_entry(&self, entry_list: Vec<Entry>) -> Result<(), MegaError>;
+
+    async fn post_receiver_handler(&self) -> Result<(), GitError>;
+
+    async fn check_entry(&self, entry: &Entry) -> Result<(), GitError>;
 
     /// Asynchronously retrieves the full pack data for the specified repository path.
     /// This function collects commits and nodes from the storage and packs them into
@@ -65,7 +114,7 @@ pub trait PackHandler: Send + Sync {
         hashes: Vec<String>,
     ) -> Result<Vec<raw_blob::Model>, MegaError>;
 
-    async fn update_refs(&self, commit: Option<Commit>, refs: &RefCommand) -> Result<(), GitError>;
+    async fn update_refs(&self, refs: &RefCommand) -> Result<(), GitError>;
 
     async fn check_commit_exist(&self, hash: &str) -> bool;
 

ceres/src/pack/monorepo.rs

@@ -2,13 +2,18 @@ use std::{
     collections::{HashMap, HashSet},
     path::{Component, PathBuf},
     str::FromStr,
-    sync::atomic::{AtomicUsize, Ordering},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
     vec,
 };
 
 use async_trait::async_trait;
-use futures::future::join_all;
-use tokio::sync::mpsc::{self, UnboundedReceiver};
+use tokio::sync::{
+    mpsc::{self},
+    RwLock,
+};
 use tokio_stream::wrappers::ReceiverStream;
 
 use callisto::{mega_mr, raw_blob, sea_orm_active_enums::ConvTypeEnum};
@@ -37,6 +42,7 @@ pub struct MonoRepo {
     pub path: PathBuf,
     pub from_hash: String,
     pub to_hash: String,
+    pub current_commit: Arc<RwLock<Option<Commit>>>,
 }
 
 #[async_trait]
@@ -122,46 +128,34 @@ impl PackHandler for MonoRepo {
         self.find_head_hash(refs)
     }
 
-    async fn handle_receiver(
-        &self,
-        mut receiver: UnboundedReceiver<Entry>,
-    ) -> Result<Option<Commit>, GitError> {
+    async fn save_entry(&self, entry_list: Vec<Entry>) -> Result<(), MegaError> {
         let storage = self.storage.mono_storage();
-        let mut entry_list = Vec::new();
-        let mut join_tasks = vec![];
-        let mut current_commit_id = String::new();
-        let mut current_commit = None;
-        while let Some(entry) = receiver.recv().await {
-            if current_commit.is_none() {
-                if entry.obj_type == ObjectType::Commit {
-                    current_commit_id = entry.hash.to_string();
-                    let commit = Commit::from_bytes(&entry.data, entry.hash).unwrap();
-                    current_commit = Some(commit);
-                }
-            } else {
-                if entry.obj_type == ObjectType::Commit {
-                    return Err(GitError::CustomError(
-                        "only single commit support in each push".to_string(),
-                    ));
-                }
-                if entry_list.len() >= 1000 {
-                    let stg_clone = storage.clone();
-                    let commit_id = current_commit_id.clone();
-                    let handle = tokio::spawn(async move {
-                        stg_clone.save_entry(&commit_id, entry_list).await.unwrap();
-                    });
-                    join_tasks.push(handle);
-                    entry_list = vec![];
-                }
+        let current_commit = self.current_commit.read().await;
+        let commit_id = if let Some(commit) = &*current_commit {
+            commit.id.to_string()
+        } else {
+            String::new()
+        };
+        storage.save_entry(&commit_id, entry_list).await
+    }
+
+    async fn post_receiver_handler(&self) -> Result<(), GitError> {
+        Ok(())
+    }
+
+    async fn check_entry(&self, entry: &Entry) -> Result<(), GitError> {
+        if self.current_commit.read().await.is_none() {
+            if entry.obj_type == ObjectType::Commit {
+                let commit = Commit::from_bytes(&entry.data, entry.hash).unwrap();
+                let mut current = self.current_commit.write().await;
+                *current = Some(commit);
             }
-            entry_list.push(entry);
+        } else if entry.obj_type == ObjectType::Commit {
+            return Err(GitError::CustomError(
+                "only single commit support in each push".to_string(),
+            ));
         }
-        join_all(join_tasks).await;
-        storage
-            .save_entry(&current_commit_id, entry_list)
-            .await
-            .unwrap();
-        Ok(current_commit)
+        Ok(())
     }
 
     // monorepo full pack should follow the shallow clone command 'git clone --depth=1'
@@ -281,10 +275,10 @@ impl PackHandler for MonoRepo {
             .await
     }
 
-    async fn update_refs(&self, commit: Option<Commit>, refs: &RefCommand) -> Result<(), GitError> {
+    async fn update_refs(&self, refs: &RefCommand) -> Result<(), GitError> {
         let storage = self.storage.mono_storage();
-
-        if let Some(c) = commit {
+        let current_commit = self.current_commit.read().await;
+        if let Some(c) = &*current_commit {
             let mr_link = self.handle_mr(&c.format_message()).await.unwrap();
             let ref_name = utils::mr_ref_name(&mr_link);
             if let Some(mut mr_ref) = storage.get_mr_ref(&ref_name).await.unwrap() {

ceres/src/protocol/mod.rs

@@ -9,6 +9,7 @@ use common::{
 use import_refs::RefCommand;
 use jupiter::storage::Storage;
 use repo::Repo;
+use tokio::sync::RwLock;
 
 use crate::pack::{import_repo::ImportRepo, monorepo::MonoRepo, PackHandler};
 
@@ -178,6 +179,7 @@ impl SmartProtocol {
                 path: self.path.clone(),
                 from_hash: String::new(),
                 to_hash: String::new(),
+                current_commit: Arc::new(RwLock::new(None)),
             };
             if let Some(command) = self
                 .command_list

ceres/src/protocol/smart.rs

@@ -216,14 +216,7 @@ impl SmartProtocol {
             .unpack_stream(&self.storage.config().pack, data_stream)
             .await?;
 
-        // do not block main thread here.
-        let handler_clone = pack_handler.clone();
-        let unpack_result = tokio::task::spawn_blocking(move || {
-            let handle = tokio::runtime::Handle::current();
-            handle.block_on(async { handler_clone.handle_receiver(receiver).await })
-        })
-        .await
-        .unwrap();
+        let unpack_result = pack_handler.clone().receiver_handler(receiver).await;
 
         // write "unpack ok\n to report"
         add_pkt_line_string(&mut report_status, "unpack ok\n".to_owned());
@@ -234,19 +227,19 @@ impl SmartProtocol {
         for command in &mut self.command_list {
             if command.ref_type == RefTypeEnum::Tag {
                 // just update if refs type is tag
-                pack_handler.update_refs(None, command).await.unwrap();
+                pack_handler.update_refs(command).await.unwrap();
             } else {
                 // Updates can be unsuccessful for a number of reasons.
                 // a.The reference can have changed since the reference discovery phase was originally sent, meaning someone pushed in the meantime.
                 // b.The reference being pushed could be a non-fast-forward reference and the update hooks or configuration could be set to not allow that, etc.
                 // c.Also, some references can be updated while others can be rejected.
                 match unpack_result {
-                    Ok(ref commit) => {
+                    Ok(_) => {
                         if !default_exist {
                             command.default_branch = true;
                             default_exist = true;
                         }
-                        if let Err(e) = pack_handler.update_refs(commit.clone(), command).await {
+                        if let Err(e) = pack_handler.update_refs(command).await {
                             command.failed(e.to_string());
                         }
                     }