compression example works again

Author: rklaehn

examples/compression.rs_ renamed to examples/compression.rs

@@ -9,7 +9,7 @@
 /// grade code you might nevertheless put the tasks into a [tokio::task::JoinSet] or
 /// [n0_future::FuturesUnordered].
 mod common;
-use std::{path::PathBuf, time::Instant};
+use std::{io, path::PathBuf, time::Instant};
 
 use anyhow::Result;
 use async_compression::tokio::{bufread::Lz4Decoder, write::Lz4Encoder};
@@ -22,7 +22,8 @@ use iroh_blobs::{
     get::fsm::{AtConnected, ConnectedNext, EndBlobNext},
     protocol::{ChunkRangesSeq, GetRequest, Request},
     provider::{
-        events::{ClientConnected, EventSender, HasErrorCode}, handle_get, AsyncReadRecvStream, AsyncWriteSendStream, ErrorHandler, StreamPair
+        events::{ClientConnected, EventSender, HasErrorCode},
+        handle_get, AsyncReadRecvStream, AsyncWriteSendStream, StreamPair,
     },
     store::mem::MemStore,
     ticket::BlobTicket,
@@ -50,11 +51,91 @@ pub enum Args {
     },
 }
 
-type CompressedWriter =
-    AsyncWriteSendStream<async_compression::tokio::write::Lz4Encoder<iroh::endpoint::SendStream>>;
-type CompressedReader = AsyncReadRecvStream<
+struct CompressedWriter(async_compression::tokio::write::Lz4Encoder<iroh::endpoint::SendStream>);
+struct CompressedReader(
     async_compression::tokio::bufread::Lz4Decoder<BufReader<iroh::endpoint::RecvStream>>,
->;
+);
+
+impl iroh_blobs::provider::SendStream for CompressedWriter {
+    async fn send_bytes(&mut self, bytes: bytes::Bytes) -> io::Result<()> {
+        AsyncWriteSendStream::new(self).send_bytes(bytes).await
+    }
+
+    async fn send<const L: usize>(&mut self, buf: &[u8; L]) -> io::Result<()> {
+        AsyncWriteSendStream::new(self).send(buf).await
+    }
+
+    async fn sync(&mut self) -> io::Result<()> {
+        AsyncWriteSendStream::new(self).sync().await
+    }
+}
+
+impl iroh_blobs::provider::RecvStream for CompressedReader {
+    async fn recv_bytes(&mut self, len: usize) -> io::Result<bytes::Bytes> {
+        AsyncReadRecvStream::new(self).recv_bytes(len).await
+    }
+
+    async fn recv_bytes_exact(&mut self, len: usize) -> io::Result<bytes::Bytes> {
+        AsyncReadRecvStream::new(self).recv_bytes_exact(len).await
+    }
+
+    async fn recv<const L: usize>(&mut self) -> io::Result<[u8; L]> {
+        AsyncReadRecvStream::new(self).recv::<L>().await
+    }
+}
+
+impl tokio::io::AsyncRead for CompressedReader {
+    fn poll_read(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &mut tokio::io::ReadBuf<'_>,
+    ) -> std::task::Poll<io::Result<()>> {
+        std::pin::Pin::new(&mut self.0).poll_read(cx, buf)
+    }
+}
+
+impl tokio::io::AsyncWrite for CompressedWriter {
+    fn poll_write(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> std::task::Poll<io::Result<usize>> {
+        std::pin::Pin::new(&mut self.0).poll_write(cx, buf)
+    }
+
+    fn poll_flush(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<io::Result<()>> {
+        std::pin::Pin::new(&mut self.0).poll_flush(cx)
+    }
+
+    fn poll_shutdown(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<io::Result<()>> {
+        std::pin::Pin::new(&mut self.0).poll_shutdown(cx)
+    }
+}
+
+impl iroh_blobs::provider::SendStreamSpecific for CompressedWriter {
+    fn reset(&mut self, code: quinn::VarInt) -> io::Result<()> {
+        self.0.get_mut().reset(code)?;
+        Ok(())
+    }
+
+    async fn stopped(&mut self) -> io::Result<Option<quinn::VarInt>> {
+        let res = self.0.get_mut().stopped().await?;
+        Ok(res)
+    }
+}
+
+impl iroh_blobs::provider::RecvStreamSpecific for CompressedReader {
+    fn stop(&mut self, code: quinn::VarInt) -> io::Result<()> {
+        self.0.get_mut().get_mut().stop(code)?;
+        Ok(())
+    }
+}
 
 #[derive(Debug, Clone)]
 struct CompressedBlobsProtocol {
@@ -71,22 +152,6 @@ impl CompressedBlobsProtocol {
     }
 }
 
-struct CompressedErrorHandler;
-
-impl ErrorHandler for CompressedErrorHandler {
-    type W = CompressedWriter;
-
-    type R = CompressedReader;
-
-    async fn stop(reader: &mut Self::R, code: quinn::VarInt) {
-        reader.0.get_mut().get_mut().stop(code).ok();
-    }
-
-    async fn reset(writer: &mut Self::W, code: quinn::VarInt) {
-        writer.0.get_mut().reset(code).ok();
-    }
-}
-
 impl ProtocolHandler for CompressedBlobsProtocol {
     async fn accept(
         &self,
@@ -108,15 +173,15 @@ impl ProtocolHandler for CompressedBlobsProtocol {
         }
         while let Ok((send, recv)) = connection.accept_bi().await {
             let stream_id = send.id().index();
-            let send = TokioStreamWriter(Lz4Encoder::new(send));
-            let recv = TokioStreamReader(Lz4Decoder::new(BufReader::new(recv)));
+            let send = CompressedWriter(Lz4Encoder::new(send));
+            let recv = CompressedReader(Lz4Decoder::new(BufReader::new(recv)));
             let store = self.store.clone();
             let mut pair =
                 StreamPair::new(connection_id, stream_id, recv, send, self.events.clone());
             tokio::spawn(async move {
                 let request = pair.read_request().await?;
                 if let Request::Get(request) = request {
-                    handle_get::<CompressedErrorHandler>(pair, store, request).await?;
+                    handle_get(pair, store, request).await?;
                 }
                 anyhow::Ok(())
             });
@@ -156,8 +221,8 @@ async fn main() -> Result<()> {
         Args::Get { ticket, target } => {
             let conn = endpoint.connect(ticket.node_addr().clone(), ALPN).await?;
             let (send, recv) = conn.open_bi().await?;
-            let send = AsyncWriteSendStream(Lz4Encoder::new(send));
-            let recv = AsyncReadRecvStream::new(Lz4Decoder::new(BufReader::new(recv)));
+            let send = CompressedWriter(Lz4Encoder::new(send));
+            let recv = CompressedReader(Lz4Decoder::new(BufReader::new(recv)));
             let request = GetRequest {
                 hash: ticket.hash(),
                 ranges: ChunkRangesSeq::root(),

src/api/blobs.rs

@@ -23,7 +23,7 @@ use bao_tree::{
 };
 use bytes::Bytes;
 use genawaiter::sync::Gen;
-use iroh_io::{AsyncStreamReader, AsyncStreamWriter};
+use iroh_io::AsyncStreamWriter;
 use irpc::channel::{mpsc, oneshot};
 use n0_future::{future, stream, Stream, StreamExt};
 use range_collections::{range_set::RangeSetRange, RangeSet2};
@@ -55,7 +55,7 @@ use super::{
 };
 use crate::{
     api::proto::{BatchRequest, ImportByteStreamUpdate},
-    provider::events::ClientResult,
+    provider::{events::ClientResult, RecvStreamAsyncStreamReader},
     store::IROH_BLOCK_SIZE,
     util::temp_tag::TempTag,
     BlobFormat, Hash, HashAndFormat,
@@ -429,13 +429,13 @@ impl Blobs {
     }
 
     #[cfg_attr(feature = "hide-proto-docs", doc(hidden))]
-    pub async fn import_bao_reader<R: AsyncStreamReader>(
+    pub async fn import_bao_reader<R: crate::provider::RecvStream>(
         &self,
         hash: Hash,
         ranges: ChunkRanges,
         mut reader: R,
     ) -> RequestResult<R> {
-        let size = u64::from_le_bytes(reader.read::<8>().await.map_err(super::Error::other)?);
+        let size = u64::from_le_bytes(reader.recv::<8>().await.map_err(super::Error::other)?);
         let Some(size) = NonZeroU64::new(size) else {
             return if hash == Hash::EMPTY {
                 Ok(reader)
@@ -444,7 +444,12 @@ impl Blobs {
             };
         };
         let tree = BaoTree::new(size.get(), IROH_BLOCK_SIZE);
-        let mut decoder = ResponseDecoder::new(hash.into(), ranges, tree, reader);
+        let mut decoder = ResponseDecoder::new(
+            hash.into(),
+            ranges,
+            tree,
+            RecvStreamAsyncStreamReader::new(reader),
+        );
         let options = ImportBaoOptions { hash, size };
         let handle = self.import_bao_with_opts(options, 32).await?;
         let driver = async move {
@@ -463,7 +468,7 @@ impl Blobs {
         let fut = async move { handle.rx.await.map_err(io::Error::other)? };
         let (reader, res) = tokio::join!(driver, fut);
         res?;
-        Ok(reader?)
+        Ok(reader?.into_inner())
     }
 
     #[cfg_attr(feature = "hide-proto-docs", doc(hidden))]
@@ -1068,7 +1073,7 @@ impl ExportBaoProgress {
     }
 
     /// Write quinn variant that also feeds a progress writer.
-    pub(crate) async fn write_with_progress<W: AsyncStreamWriter>(
+    pub(crate) async fn write_with_progress<W: crate::provider::SendStream>(
         self,
         writer: &mut W,
         progress: &mut impl WriteProgress,
@@ -1080,19 +1085,19 @@ impl ExportBaoProgress {
             match item {
                 EncodedItem::Size(size) => {
                     progress.send_transfer_started(index, hash, size).await;
-                    writer.write(&size.to_le_bytes()).await?;
+                    writer.send(&size.to_le_bytes()).await?;
                     progress.log_other_write(8);
                 }
                 EncodedItem::Parent(parent) => {
-                    let mut data = vec![0u8; 64];
+                    let mut data = [0u8; 64];
                     data[..32].copy_from_slice(parent.pair.0.as_bytes());
                     data[32..].copy_from_slice(parent.pair.1.as_bytes());
-                    writer.write(&data).await?;
+                    writer.send(&data).await?;
                     progress.log_other_write(64);
                 }
                 EncodedItem::Leaf(leaf) => {
                     let len = leaf.data.len();
-                    writer.write_bytes(leaf.data).await?;
+                    writer.send_bytes(leaf.data).await?;
                     progress
                         .notify_payload_write(index, leaf.offset, len)
                         .await?;

src/api/remote.rs

@@ -16,7 +16,7 @@ use crate::{
     get::{
         fsm::DecodeError,
         get_error::{BadRequestSnafu, LocalFailureSnafu},
-        GetError, GetResult, IrohStreamWriter, Stats,
+        GetError, GetResult, Stats,
     },
     protocol::{
         GetManyRequest, ObserveItem, ObserveRequest, PushRequest, Request, RequestType,
@@ -593,7 +593,6 @@ impl Remote {
         let mut request_ranges = request.ranges.iter_infinite();
         let root = request.hash;
         let root_ranges = request_ranges.next().expect("infinite iterator");
-        let mut send = IrohStreamWriter(send);
         if !root_ranges.is_empty() {
             self.store()
                 .export_bao(root, root_ranges.clone())
@@ -602,7 +601,7 @@ impl Remote {
         }
         if request.ranges.is_blob() {
             // we are done
-            send.0.finish()?;
+            send.finish()?;
             return Ok(Default::default());
         }
         let hash_seq = self.store().get_bytes(root).await?;
@@ -617,7 +616,7 @@ impl Remote {
                     .await?;
             }
         }
-        send.0.finish()?;
+        send.finish()?;
         Ok(Default::default())
     }