diff --git a/libra/Cargo.toml b/libra/Cargo.toml index 8b206138e..1911f6731 100644 --- a/libra/Cargo.toml +++ b/libra/Cargo.toml @@ -43,11 +43,11 @@ ring = "0.17.8" hex = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -async_static = "0.1.3" once_cell = "1.19.0" byte-unit = "5.1.4" scopeguard = "1.2.0" lru-mem = "0.3.0" +anyhow = { workspace = true } [target.'cfg(unix)'.dependencies] # only on Unix pager = "0.16.0" diff --git a/libra/src/command/lfs.rs b/libra/src/command/lfs.rs index e0ad914b3..808fc5100 100644 --- a/libra/src/command/lfs.rs +++ b/libra/src/command/lfs.rs @@ -8,7 +8,7 @@ use ceres::lfs::lfs_structs::LockListQuery; use mercury::internal::index::Index; use crate::command::{ask_basic_auth, status}; use crate::internal::head::Head; -use crate::internal::protocol::lfs_client::LFS_CLIENT; +use crate::internal::protocol::lfs_client::LFSClient; use crate::utils::{lfs, path, util}; use crate::utils::path_ext::PathExt; @@ -94,7 +94,7 @@ pub async fn execute(cmd: LfsCmds) { cursor: "".to_string(), refspec, }; - let locks = LFS_CLIENT.await.get_locks(query).await.locks; + let locks = LFSClient::get().await.get_locks(query).await.locks; if !locks.is_empty() { let max_path_len = locks.iter().map(|l| l.path.len()).max().unwrap(); for lock in locks { @@ -112,7 +112,7 @@ pub async fn execute(cmd: LfsCmds) { let refspec = current_refspec().await.unwrap(); let mut auth = None; loop { - let code = LFS_CLIENT.await.lock(path.clone(), refspec.clone(), auth.clone()).await; + let code = LFSClient::get().await.lock(path.clone(), refspec.clone(), auth.clone()).await; if code.is_success() { println!("Locked {}", path); } else if code == StatusCode::FORBIDDEN { @@ -140,7 +140,7 @@ pub async fn execute(cmd: LfsCmds) { let id = match id { None => { // get id by path - let locks = LFS_CLIENT.await.get_locks(LockListQuery { + let locks = LFSClient::get().await.get_locks(LockListQuery { refspec: refspec.clone(), path: path.clone(), id: "".to_string(), @@ -157,7 +157,7 @@ pub async fn execute(cmd: LfsCmds) { }; let mut auth = None; loop { - let code = LFS_CLIENT.await.unlock(id.clone(), refspec.clone(), force, auth.clone()).await; + let code = LFSClient::get().await.unlock(id.clone(), refspec.clone(), force, auth.clone()).await; if code.is_success() { println!("Unlocked {}", path); } else if code == StatusCode::FORBIDDEN { diff --git a/libra/src/command/restore.rs b/libra/src/command/restore.rs index 8cf359a9f..0074404a1 100644 --- a/libra/src/command/restore.rs +++ b/libra/src/command/restore.rs @@ -8,13 +8,13 @@ use clap::Parser; use std::collections::{HashMap, HashSet}; use std::{fs, io}; use std::path::PathBuf; +use crate::internal::protocol::lfs_client::LFSClient; use mercury::hash::SHA1; use mercury::internal::object::blob::Blob; use mercury::internal::object::commit::Commit; use mercury::internal::object::tree::Tree; use mercury::internal::object::types::ObjectType; use crate::command::calc_file_blob_hash; -use crate::internal::protocol::lfs_client::LFS_CLIENT; #[derive(Parser, Debug)] pub struct RestoreArgs { @@ -154,7 +154,9 @@ async fn restore_to_file(hash: &SHA1, path: &PathBuf) -> io::Result<()> { fs::copy(&lfs_obj_path, &path_abs)?; } else { // not exist, download from server - LFS_CLIENT.await.download_object(&oid, size, &path_abs).await; + if let Err(e) = LFSClient::get().await.download_object(&oid, size, &path_abs, None).await { + eprintln!("fatal: {}", e); + } } } None => { diff --git a/libra/src/internal/protocol/lfs_client.rs b/libra/src/internal/protocol/lfs_client.rs index ae3eb395b..de38ba84b 100644 --- a/libra/src/internal/protocol/lfs_client.rs +++ b/libra/src/internal/protocol/lfs_client.rs @@ -3,8 +3,7 @@ use crate::internal::config::Config; use crate::internal::protocol::https_client::BasicAuth; use crate::internal::protocol::ProtocolClient; use crate::utils::lfs; -use async_static::async_static; -use ceres::lfs::lfs_structs::{BatchRequest, FetchchunkResponse, Link, LockList, LockListQuery, LockRequest, Ref, Representation, RequestVars, UnlockRequest, VerifiableLockList, VerifiableLockRequest}; +use ceres::lfs::lfs_structs::{BatchRequest, ChunkRepresentation, FetchchunkResponse, LockList, LockListQuery, LockRequest, Ref, Representation, RequestVars, UnlockRequest, VerifiableLockList, VerifiableLockRequest}; use futures_util::StreamExt; use mercury::internal::object::types::ObjectType; use mercury::internal::pack::entry::Entry; @@ -13,18 +12,27 @@ use ring::digest::{Context, SHA256}; use serde::{Deserialize, Serialize}; use std::collections::HashSet; use std::path::Path; -use tokio::io::AsyncWriteExt; +use anyhow::anyhow; +use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; +use tokio::sync::OnceCell; use url::Url; -async_static! { - pub static ref LFS_CLIENT: LFSClient = LFSClient::new().await; -} - +#[derive(Debug)] pub struct LFSClient { pub batch_url: Url, pub lfs_url: Url, pub client: Client, } +static LFS_CLIENT: OnceCell = OnceCell::const_new(); +impl LFSClient { + /// Get LFSClient instance + /// - DO NOT use `async_static!`: No IDE Code Completion & lagging + pub async fn get() -> &'static LFSClient { + LFS_CLIENT.get_or_init(|| async { + LFSClient::new().await + }).await + } +} /// see [successful-responses](https://github.com/git-lfs/git-lfs/blob/main/docs/api/batch.md#successful-responses) #[derive(Serialize, Deserialize)] @@ -208,8 +216,31 @@ impl LFSClient { Ok(()) } + /// Just for resume download + async fn update_file_checksum(file: &mut tokio::fs::File, checksum: &mut Context) { + file.seek(tokio::io::SeekFrom::Start(0)).await.unwrap(); + let mut buf = [0u8; 8192]; + loop { + let n = file.read(&mut buf).await.unwrap(); + if n == 0 { + break; + } + checksum.update(&buf[..n]); + } + } + + #[allow(clippy::type_complexity)] /// download (GET) one LFS file from remote server - pub async fn download_object(&self, oid: &str, size: u64, path: impl AsRef) { + pub async fn download_object( + &self, + oid: &str, + size: u64, + path: impl AsRef, + mut reporter: Option<( + &mut (dyn FnMut(f64) -> anyhow::Result<()> + Send), // progress callback + f64 // step + )>) -> anyhow::Result<()> + { let batch_request = BatchRequest { operation: "download".to_string(), transfers: vec![lfs::LFS_TRANSFER_API.to_string()], @@ -225,34 +256,66 @@ impl LFSClient { .post(self.batch_url.clone()) .json(&batch_request) .headers(lfs::LFS_HEADERS.clone()); - let response = request.send().await.unwrap(); + let response = request.send().await?; - let text = response.text().await.unwrap(); - tracing::debug!("LFS download response:\n {:#?}", serde_json::from_str::(&text).unwrap()); - let resp = serde_json::from_str::(&text).unwrap(); + let text = response.text().await?; + tracing::debug!("LFS download response:\n {:#?}", serde_json::from_str::(&text)?); + let resp = serde_json::from_str::(&text)?; let link = resp.objects[0].actions.as_ref().unwrap().get("download").unwrap(); let mut is_chunked = false; // Chunk API - let links = match self.fetch_chunk_links(&link.href).await { + let mut chunk_size = None; // infer that all chunks are the same size! + let links = match self.fetch_chunks(&link.href).await { Ok(chunks) => { is_chunked = true; + chunk_size = chunks.first().map(|c| c.size); tracing::info!("LFS Chunk API supported."); - chunks + chunks.into_iter().map(|c| c.link).collect() }, Err(_) => vec![link.clone()], }; - let mut file = tokio::fs::File::create(path).await.unwrap(); let mut checksum = Context::new(&SHA256); + let mut got_parts = 0; + let mut file = if links.len() <= 1 || lfs::parse_pointer_file(&path).is_ok() { + // pointer file or Not Chunks, truncate + tokio::fs::File::create(path).await? + } else { + // for Chunks, calc offset to resume download + let mut file = tokio::fs::File::options() + .write(true) + .read(true) + .create(true) + .truncate(false) + .open(&path).await?; + let file_len = file.metadata().await?.len(); + if file_len > size { + println!("Local file size is larger than remote, truncate to 0."); + file.set_len(0).await?; // clear + file.seek(tokio::io::SeekFrom::Start(0)).await?; + } else if file_len > 0 { + let chunk_size = chunk_size.unwrap() as u64; + got_parts = file_len / chunk_size; + let file_offset = got_parts * chunk_size; + println!("Resume download from offset: {}, part: {}", file_offset, got_parts + 1); + file.set_len(file_offset).await?; // truncate + Self::update_file_checksum(&mut file, &mut checksum).await; // resume checksum + file.seek(tokio::io::SeekFrom::End(0)).await?; + } + file + }; + println!("Downloading LFS file: {}", oid); - let mut cnt = 0; - let total = links.len(); - for link in links { - cnt += 1; + let parts = links.len(); + let mut downloaded: u64 = file.metadata().await?.len(); + let mut last_progress = 0.0; + let start_part = got_parts as usize; + for link in links.iter().skip(start_part) { + got_parts += 1; if is_chunked { - println!("- part: {}/{}", cnt, total); + println!("- part: {}/{}", got_parts, parts); } let mut request = self.client.get(&link.href); @@ -260,33 +323,46 @@ impl LFSClient { request = request.header(k, v); } - let response = request.send().await.unwrap(); + let response = request.send().await?; if !response.status().is_success() { - eprintln!("fatal: LFS download failed. Status: {}, Message: {}", response.status(), response.text().await.unwrap()); - return; + eprintln!("fatal: LFS download failed. Status: {}, Message: {}", response.status(), response.text().await?); + return Err(anyhow!("LFS download failed.")); } let mut stream = response.bytes_stream(); - while let Some(chunk) = stream.next().await { // TODO: progress bar - let chunk = chunk.unwrap(); - file.write_all(&chunk).await.unwrap(); + while let Some(chunk) = stream.next().await { // TODO: progress bar TODO: multi-thread or async + let chunk = chunk?; + file.write_all(&chunk).await?; checksum.update(&chunk); + + // report progress + if let Some((ref mut report_fn, step)) = reporter { + downloaded += chunk.len() as u64; + let progress = (downloaded as f64 / size as f64) * 100.0; + if progress >= last_progress + step { + last_progress = progress; + report_fn(progress)?; + } + } } } let checksum = hex::encode(checksum.finish().as_ref()); if checksum == oid { println!("Downloaded."); + Ok(()) } else { eprintln!("fatal: LFS download failed. Checksum mismatch: {} != {}. Fallback to pointer file.", checksum, oid); let pointer = lfs::format_pointer_string(oid, size); - file.set_len(0).await.unwrap(); // clear - file.write_all(pointer.as_bytes()).await.unwrap(); + file.set_len(0).await?; // clear + file.seek(tokio::io::SeekFrom::Start(0)).await?; // ensure + file.write_all(pointer.as_bytes()).await?; + Err(anyhow!("Checksum mismatch, fallback to pointer file.")) } } /// Only for MonoRepo (mega) - async fn fetch_chunk_links(&self, obj_link: &str) -> Result, ()> { + async fn fetch_chunks(&self, obj_link: &str) -> Result, ()> { let mut url = Url::parse(obj_link).unwrap(); let path = url.path().trim_end_matches('/'); url.set_path(&(path.to_owned() + "/chunks")); // reserve query params (for GitHub link) @@ -304,7 +380,7 @@ impl LFSClient { let mut res = resp.json::().await.unwrap(); // sort by offset res.chunks.sort_by(|a, b| a.offset.cmp(&b.offset)); - Ok(res.chunks.into_iter().map(|c| c.link).collect()) + Ok(res.chunks) } } diff --git a/libra/src/lib.rs b/libra/src/lib.rs index e54f61ef1..7e06d37fb 100644 --- a/libra/src/lib.rs +++ b/libra/src/lib.rs @@ -1,7 +1,7 @@ use mercury::errors::GitError; mod command; -mod internal; +pub mod internal; mod utils; pub mod cli; @@ -34,4 +34,27 @@ mod tests { std::env::set_current_dir(tmp_dir.path()).unwrap(); exec(vec!["init"]).unwrap(); } + + #[tokio::test] + async fn test_lfs_client() { + use url::Url; + use crate::internal::protocol::lfs_client::LFSClient; + use crate::internal::protocol::ProtocolClient; + + let client = LFSClient::from_url(&Url::parse("https://git.gitmono.org").unwrap()); + println!("{:?}", client); + let mut report_fn = |progress: f64| { + println!("progress: {:.2}%", progress); + Ok(()) + }; + client.download_object( + "a744b4beab939d899e22c8a070b7041a275582fb942483c9436d455173c7e23d", + 338607424, + "/home/bean/projects/tmp/Qwen2.5-0.5B-Instruct-Q2_K.gguf", + Some(( + &mut report_fn, + 0.1 + )) + ).await.expect("Failed to download object"); + } } \ No newline at end of file diff --git a/mega/tests/lfs_test.rs b/mega/tests/lfs_test.rs index 8c0350826..b3d128b58 100644 --- a/mega/tests/lfs_test.rs +++ b/mega/tests/lfs_test.rs @@ -223,6 +223,7 @@ fn lfs_split_with_git() { git_lfs_clone(url).expect("Failed to clone large file from mega server"); mega.kill().expect("Failed to kill mega server"); + thread::sleep(Duration::from_secs(1)); // wait for server to stop, avoiding affecting other tests } #[test] @@ -241,4 +242,5 @@ fn lfs_split_with_libra() { libra_lfs_clone(url).expect("(libra)Failed to clone large file from mega server"); mega.kill().expect("Failed to kill mega server"); + thread::sleep(Duration::from_secs(1)); // wait for server to stop, avoiding affecting other tests } \ No newline at end of file