From c60272268874d34d6103c1a9ed4cda5f8a15dc48 Mon Sep 17 00:00:00 2001 From: Qihang Cai Date: Tue, 28 May 2024 17:55:10 +0800 Subject: [PATCH] `Pack`: add `decode_stream` for `futures_core::stream::Stream`, using `Channel` as adapter Signed-off-by: Qihang Cai --- mercury/Cargo.toml | 4 ++ mercury/src/internal/pack/channel_reader.rs | 50 ++++++++++++++ mercury/src/internal/pack/decode.rs | 76 +++++++++++++++++++-- mercury/src/internal/pack/mod.rs | 1 + 4 files changed, 126 insertions(+), 5 deletions(-) create mode 100644 mercury/src/internal/pack/channel_reader.rs diff --git a/mercury/Cargo.toml b/mercury/Cargo.toml index 21da31155..d62805dad 100644 --- a/mercury/Cargo.toml +++ b/mercury/Cargo.toml @@ -30,6 +30,10 @@ tokio.workspace = true lru-mem = "0.3.0" bincode = "1.3.3" byteorder = "1.5.0" +tokio-util = { version = "0.7.11", features = ["io"] } +futures-util = "0.3.30" +bytes = { workspace = true } +axum = { workspace = true } [target.'cfg(windows)'.dependencies] # only on Windows mimalloc = "0.1.39" # avoid sticking on dropping on Windows diff --git a/mercury/src/internal/pack/channel_reader.rs b/mercury/src/internal/pack/channel_reader.rs new file mode 100644 index 000000000..138a06c39 --- /dev/null +++ b/mercury/src/internal/pack/channel_reader.rs @@ -0,0 +1,50 @@ +use std::io; +use std::io::{BufRead, Read}; +use std::sync::mpsc::Receiver; + +/// Custom BufRead implementation that reads from the channel +pub(crate) struct ChannelReader { + receiver: Receiver>, + buffer: io::Cursor>, +} + +impl ChannelReader { + pub(crate) fn new(receiver: Receiver>) -> Self { + ChannelReader { + receiver, + buffer: io::Cursor::new(Vec::new()), + } + } +} + +impl Read for ChannelReader { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + if self.buffer.position() as usize == self.buffer.get_ref().len() { // buffer has been read completely + match self.receiver.recv() { + Ok(data) => { + self.buffer = io::Cursor::new(data); + } + Err(_) => return Ok(0), // Channel is closed + } + } + self.buffer.read(buf) + } +} + +impl BufRead for ChannelReader { + fn fill_buf(&mut self) -> io::Result<&[u8]> { + if self.buffer.position() as usize == self.buffer.get_ref().len() { + match self.receiver.recv() { + Ok(data) => { + self.buffer = io::Cursor::new(data); + } + Err(_) => return Ok(&[]), // Channel is closed + } + } + self.buffer.fill_buf() + } + + fn consume(&mut self, amt: usize) { + self.buffer.consume(amt); + } +} \ No newline at end of file diff --git a/mercury/src/internal/pack/decode.rs b/mercury/src/internal/pack/decode.rs index 7c1608b20..edfe2a3e9 100644 --- a/mercury/src/internal/pack/decode.rs +++ b/mercury/src/internal/pack/decode.rs @@ -1,12 +1,15 @@ -use std::io::{self, BufRead, Cursor, ErrorKind, Read, Seek}; +use std::io::{self, BufRead, Cursor, ErrorKind, Read}; use std::path::PathBuf; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::Arc; -use std::sync::mpsc::Sender; +use std::sync::{Arc, mpsc}; +use std::sync::mpsc::{Receiver, Sender}; use std::thread::{self, JoinHandle}; use std::time::Instant; +use axum::Error; +use bytes::Bytes; use flate2::bufread::ZlibDecoder; +use futures_util::{Stream, StreamExt}; use threadpool::ThreadPool; use crate::errors::GitError; @@ -20,6 +23,7 @@ use crate::internal::pack::waitlist::Waitlist; use crate::internal::pack::wrapper::Wrapper; use crate::internal::pack::{utils, Pack, DEFAULT_TMP_DIR}; use uuid::Uuid; +use crate::internal::pack::channel_reader::ChannelReader; use crate::internal::pack::entry::Entry; /// For Convenient to pass Params @@ -320,7 +324,7 @@ impl Pack { /// Decodes a pack file from a given Read and BufRead source and get a vec of objects. /// /// - pub fn decode(&mut self, pack: &mut (impl BufRead + Seek + Send), callback: F) -> Result<(), GitError> + pub fn decode(&mut self, pack: &mut (impl BufRead + Send), callback: F) -> Result<(), GitError> where F: Fn(Entry, usize) + Sync + Send + 'static { @@ -460,7 +464,7 @@ impl Pack { /// Decode Pack in a new thread and send the CacheObjects while decoding. ///
Attention: It will consume the `pack` and return in JoinHandle - pub fn decode_async(mut self, mut pack: (impl BufRead + Seek + Send + 'static), sender: Sender) -> JoinHandle { + pub fn decode_async(mut self, mut pack: (impl BufRead + Send + 'static), sender: Sender) -> JoinHandle { thread::spawn(move || { self.decode(&mut pack, move |entry, _| { sender.send(entry).unwrap(); @@ -469,6 +473,31 @@ impl Pack { }) } + /// Decode `Pack` with inputting a `Stream` of `Bytes`, and send the `Entry` while decoding. + pub async fn decode_stream(mut self, + mut stream: impl Stream> + Unpin + Send + 'static, + sender: Sender) + -> Self + { + let (tx, rx): (Sender>, Receiver>) = mpsc::channel(); + let mut reader = ChannelReader::new(rx); + tokio::spawn(async move { + // use Channel to connect `async` & `sync` + while let Some(chunk) = stream.next().await { + let data = chunk.unwrap().to_vec(); + tx.send(data).unwrap(); + } + }); + // CPU-bound task, so use spawn_blocking + // DO NOT use thread::spawn, because it will block tokio runtime (if single-threaded runtime, like in tests) + tokio::task::spawn_blocking(move || { + self.decode(&mut reader, move |entry, _| { + sender.send(entry).unwrap(); + }).unwrap(); + self + }).await.unwrap() + } + /// CacheObjects + Index size of Caches fn memory_used(&self) -> usize { self.cache_objs_mem_used() + self.caches.memory_used_index() @@ -598,12 +627,16 @@ mod tests { use std::io::BufReader; use std::io::Cursor; use std::{env, path::PathBuf}; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use flate2::write::ZlibEncoder; use flate2::Compression; + use tokio_util::io::ReaderStream; use tracing_subscriber::util::SubscriberInitExt; use crate::internal::pack::Pack; + use futures_util::TryStreamExt; fn init_logger() { let _ = tracing_subscriber::fmt::Subscriber::builder() @@ -707,6 +740,39 @@ mod tests { } } // it will be stuck on dropping `Pack` on Windows if `mem_size` is None, so we need `mimalloc` + #[tokio::test] + async fn test_decode_large_file_stream() { + init_logger(); + let mut source = PathBuf::from(env::current_dir().unwrap().parent().unwrap()); + source.push("tests/data/packs/git-2d187177923cd618a75da6c6db45bb89d92bd504.pack"); + + let tmp = PathBuf::from("/tmp/.cache_temp"); + let f = tokio::fs::File::open(source).await.unwrap(); + let stream = ReaderStream::new(f).map_err(|e| { + axum::Error::new(e) + }); + let p = Pack::new(Some(20), Some(1024*1024*1024*2), Some(tmp.clone()), true); + + let (tx, rx) = std::sync::mpsc::channel(); + let handle = tokio::spawn(async move { + p.decode_stream(stream, tx).await + }); + + let count = Arc::new(AtomicUsize::new(0)); + let count_c = count.clone(); + // in tests, RUNTIME is single-threaded, so `sync code` will block the tokio runtime + tokio::task::spawn_blocking(move || { + let mut cnt = 0; + for _entry in rx { + cnt += 1; //use entry here + } + tracing::info!("Received: {}", cnt); + count_c.store(cnt, Ordering::Relaxed); + }).await.unwrap(); + let p = handle.await.unwrap(); + assert_eq!(count.load(Ordering::Relaxed), p.number); + } + #[test] fn test_decode_large_file_async() { let mut source = PathBuf::from(env::current_dir().unwrap().parent().unwrap()); diff --git a/mercury/src/internal/pack/mod.rs b/mercury/src/internal/pack/mod.rs index a3a54475a..8b11af0e5 100644 --- a/mercury/src/internal/pack/mod.rs +++ b/mercury/src/internal/pack/mod.rs @@ -10,6 +10,7 @@ pub mod cache; pub mod waitlist; pub mod cache_object; pub mod entry; +pub mod channel_reader; use crate::hash::SHA1; use threadpool::ThreadPool;