[WIP] feat: compress blobs on transfer

Cargo.toml

@@ -14,6 +14,7 @@ rust-version = "1.63"
 [dependencies]
 abao = { version = "0.2.0", features = ["group_size_16k", "tokio_io"], default-features = false }
 anyhow = { version = "1", features = ["backtrace"] }
+async-compression = { version = "0.3.15", features = ["tokio", "brotli"] }
 base64 = "0.21.0"
 blake3 = "1.3.3"
 bytes = "1"

src/get.rs

@@ -19,7 +19,7 @@ use anyhow::{anyhow, bail, Context, Result};
 use bytes::BytesMut;
 use futures::Future;
 use postcard::experimental::max_size::MaxSize;
-use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};
+use tokio::io::{AsyncRead, AsyncReadExt, BufReader, ReadBuf};
 use tracing::{debug, error};
 
 pub use crate::util::Hash;
@@ -103,18 +103,21 @@ impl Stats {
 /// We guarantee that the data is correct by incrementally verifying a hash
 #[repr(transparent)]
 #[derive(Debug)]
-pub struct DataStream(AsyncSliceDecoder<quinn::RecvStream>);
+pub struct DataStream(AsyncSliceDecoder<RecvStream>);
+
+type RecvStream =
+    async_compression::tokio::bufread::BrotliDecoder<tokio::io::BufReader<quinn::RecvStream>>;
 
 impl DataStream {
-    fn new(inner: quinn::RecvStream, hash: Hash) -> Self {
+    fn new(inner: RecvStream, hash: Hash) -> Self {
         DataStream(AsyncSliceDecoder::new(inner, &hash.into(), 0, u64::MAX))
     }
 
     async fn read_size(&mut self) -> io::Result<u64> {
         self.0.read_size().await
     }
 
-    fn into_inner(self) -> quinn::RecvStream {
+    fn into_inner(self) -> RecvStream {
         self.0.into_inner()
     }
 }
@@ -149,7 +152,7 @@ where
     let now = Instant::now();
     let connection = setup(opts).await?;
 
-    let (mut writer, mut reader) = connection.open_bi().await?;
+    let (mut writer, reader) = connection.open_bi().await?;
 
     on_connected().await?;
 
@@ -181,6 +184,7 @@ where
     {
         debug!("reading response");
         let mut in_buffer = BytesMut::with_capacity(1024);
+        let mut reader = BufReader::new(reader);
 
         // track total amount of blob data transferred
         let mut data_len = 0;
@@ -218,7 +222,7 @@ where
                             if blob_reader.read_exact(&mut [0u8; 1]).await.is_ok() {
                                 bail!("`on_blob` callback did not fully read the blob content")
                             }
-                            reader = blob_reader.into_inner();
+                            reader = blob_reader.into_inner().into_inner();
                         }
                     }
 
@@ -236,11 +240,12 @@ where
                 }
 
                 // Shut down the stream
-                if let Some(chunk) = reader.read_chunk(8, false).await? {
-                    reader.stop(0u8.into()).ok();
-                    error!("Received unexpected data from the provider: {chunk:?}");
+                if let Ok(bytes) = reader.read_u8().await {
+                    reader.into_inner().stop(0u8.into()).ok();
+                    error!("Received unexpected data from the provider: {bytes:?}");
+                } else {
+                    drop(reader);
                 }
-                drop(reader);
 
                 let elapsed = now.elapsed();
 
@@ -261,7 +266,7 @@ where
 /// The `AsyncReader` can be used to read the content.
 async fn handle_blob_response(
     hash: Hash,
-    mut reader: quinn::RecvStream,
+    mut reader: BufReader<quinn::RecvStream>,
     buffer: &mut BytesMut,
 ) -> Result<DataStream> {
     match read_lp(&mut reader, buffer).await? {
@@ -277,7 +282,10 @@ async fn handle_blob_response(
                 // next blob in collection will be sent over
                 Res::Found => {
                     assert!(buffer.is_empty());
-                    let decoder = DataStream::new(reader, hash);
+                    // Decompress data
+                    let decompress_reader =
+                        async_compression::tokio::bufread::BrotliDecoder::new(reader);
+                    let decoder = DataStream::new(decompress_reader, hash);
                     Ok(decoder)
                 }
             }

src/lib.rs

@@ -41,6 +41,7 @@ mod tests {
 
     #[tokio::test]
     async fn basics() -> Result<()> {
+        setup_logging();
         transfer_data(vec![("hello_world", "hello world!".as_bytes().to_vec())]).await
     }
 

src/provider/mod.rs

@@ -821,6 +821,38 @@ enum SentStatus {
     NotFound,
 }
 
+struct ShutdownCatcher<W>(W);
+
+impl<W> ShutdownCatcher<W> {
+    fn into_inner(self) -> W {
+        self.0
+    }
+}
+
+impl<W: AsyncWrite + Unpin> AsyncWrite for ShutdownCatcher<W> {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> Poll<std::result::Result<usize, std::io::Error>> {
+        Pin::new(&mut self.0).poll_write(cx, buf)
+    }
+
+    fn poll_flush(
+        mut self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<std::result::Result<(), std::io::Error>> {
+        Pin::new(&mut self.0).poll_flush(cx)
+    }
+
+    fn poll_shutdown(
+        mut self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<std::result::Result<(), std::io::Error>> {
+        Pin::new(&mut self.0).poll_flush(cx)
+    }
+}
+
 async fn send_blob<W: AsyncWrite + Unpin + Send + 'static>(
     db: Database,
     name: Hash,
@@ -837,21 +869,29 @@ async fn send_blob<W: AsyncWrite + Unpin + Send + 'static>(
             // need to thread the writer though the spawn_blocking, since
             // taking a reference does not work. spawn_blocking requires
             // 'static lifetime.
-            writer = tokio::task::spawn_blocking(move || {
+
+            // Compress data
+            let mut compressed_writer =
+                async_compression::tokio::write::BrotliEncoder::new(ShutdownCatcher(writer));
+
+            compressed_writer = tokio::task::spawn_blocking(move || {
                 let file_reader = std::fs::File::open(&path)?;
                 let outboard_reader = std::io::Cursor::new(outboard);
-                let mut wrapper = SyncIoBridge::new(&mut writer);
+                let mut wrapper = SyncIoBridge::new(&mut compressed_writer);
                 let mut slice_extractor = abao::encode::SliceExtractor::new_outboard(
                     file_reader,
                     outboard_reader,
                     0,
                     size,
                 );
                 let _copied = std::io::copy(&mut slice_extractor, &mut wrapper)?;
-                std::io::Result::Ok(writer)
+                std::io::Result::Ok(compressed_writer)
             })
             .await??;
 
+            compressed_writer.shutdown().await?;
+            let writer = compressed_writer.into_inner().into_inner();
+
             Ok((SentStatus::Sent, writer, size))
         }
         _ => {
@@ -1069,6 +1109,7 @@ async fn write_response<W: AsyncWrite + Unpin>(
     }
     let used = postcard::to_slice(&response, buffer)?;
 
+    // Write lp
     write_lp(&mut writer, used).await?;
 
     debug!("written response of length {}", used.len());