Genericize provider side a bit

Author: rklaehn

src/api/blobs.rs

@@ -23,14 +23,12 @@ use bao_tree::{
 };
 use bytes::Bytes;
 use genawaiter::sync::Gen;
-use iroh_io::{AsyncStreamReader, TokioStreamReader};
+use iroh_io::{AsyncStreamReader, AsyncStreamWriter};
 use irpc::channel::{mpsc, oneshot};
 use n0_future::{future, stream, Stream, StreamExt};
-use quinn::SendStream;
 use range_collections::{range_set::RangeSetRange, RangeSet2};
 use ref_cast::RefCast;
 use serde::{Deserialize, Serialize};
-use tokio::io::AsyncWriteExt;
 use tracing::trace;
 mod reader;
 pub use reader::BlobReader;
@@ -431,7 +429,7 @@ impl Blobs {
     }
 
     #[cfg_attr(feature = "hide-proto-docs", doc(hidden))]
-    async fn import_bao_reader<R: AsyncStreamReader>(
+    pub async fn import_bao_reader<R: AsyncStreamReader>(
         &self,
         hash: Hash,
         ranges: ChunkRanges,
@@ -468,18 +466,6 @@ impl Blobs {
         Ok(reader?)
     }
 
-    #[cfg_attr(feature = "hide-proto-docs", doc(hidden))]
-    pub async fn import_bao_quinn(
-        &self,
-        hash: Hash,
-        ranges: ChunkRanges,
-        stream: &mut iroh::endpoint::RecvStream,
-    ) -> RequestResult<()> {
-        let reader = TokioStreamReader::new(stream);
-        self.import_bao_reader(hash, ranges, reader).await?;
-        Ok(())
-    }
-
     #[cfg_attr(feature = "hide-proto-docs", doc(hidden))]
     pub async fn import_bao_bytes(
         &self,
@@ -1058,24 +1044,21 @@ impl ExportBaoProgress {
         Ok(data)
     }
 
-    pub async fn write_quinn(self, target: &mut quinn::SendStream) -> super::ExportBaoResult<()> {
+    pub async fn write<W: AsyncStreamWriter>(self, target: &mut W) -> super::ExportBaoResult<()> {
         let mut rx = self.inner.await?;
         while let Some(item) = rx.recv().await? {
             match item {
                 EncodedItem::Size(size) => {
-                    target.write_u64_le(size).await?;
+                    target.write(&size.to_le_bytes()).await?;
                 }
                 EncodedItem::Parent(parent) => {
                     let mut data = vec![0u8; 64];
                     data[..32].copy_from_slice(parent.pair.0.as_bytes());
                     data[32..].copy_from_slice(parent.pair.1.as_bytes());
-                    target.write_all(&data).await.map_err(io::Error::from)?;
+                    target.write(&data).await?;
                 }
                 EncodedItem::Leaf(leaf) => {
-                    target
-                        .write_chunk(leaf.data)
-                        .await
-                        .map_err(io::Error::from)?;
+                    target.write_bytes(leaf.data).await?;
                 }
                 EncodedItem::Done => break,
                 EncodedItem::Error(cause) => return Err(cause.into()),
@@ -1085,9 +1068,9 @@ impl ExportBaoProgress {
     }
 
     /// Write quinn variant that also feeds a progress writer.
-    pub(crate) async fn write_quinn_with_progress(
+    pub(crate) async fn write_with_progress<W: AsyncStreamWriter>(
         self,
-        writer: &mut SendStream,
+        writer: &mut W,
         progress: &mut impl WriteProgress,
         hash: &Hash,
         index: u64,
@@ -1097,22 +1080,19 @@ impl ExportBaoProgress {
             match item {
                 EncodedItem::Size(size) => {
                     progress.send_transfer_started(index, hash, size).await;
-                    writer.write_u64_le(size).await?;
+                    writer.write(&size.to_le_bytes()).await?;
                     progress.log_other_write(8);
                 }
                 EncodedItem::Parent(parent) => {
                     let mut data = vec![0u8; 64];
                     data[..32].copy_from_slice(parent.pair.0.as_bytes());
                     data[32..].copy_from_slice(parent.pair.1.as_bytes());
-                    writer.write_all(&data).await.map_err(io::Error::from)?;
+                    writer.write(&data).await?;
                     progress.log_other_write(64);
                 }
                 EncodedItem::Leaf(leaf) => {
                     let len = leaf.data.len();
-                    writer
-                        .write_chunk(leaf.data)
-                        .await
-                        .map_err(io::Error::from)?;
+                    writer.write_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, Stats,
+        GetError, GetResult, IrohStreamWriter, Stats,
     },
     protocol::{
         GetManyRequest, ObserveItem, ObserveRequest, PushRequest, Request, RequestType,
@@ -594,15 +594,16 @@ 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())
-                .write_quinn_with_progress(&mut send, &mut context, &root, 0)
+                .write_with_progress(&mut send, &mut context, &root, 0)
                 .await?;
         }
         if request.ranges.is_blob() {
             // we are done
-            send.finish()?;
+            send.0.finish()?;
             return Ok(Default::default());
         }
         let hash_seq = self.store().get_bytes(root).await?;
@@ -613,16 +614,11 @@ impl Remote {
             if !child_ranges.is_empty() {
                 self.store()
                     .export_bao(child_hash, child_ranges.clone())
-                    .write_quinn_with_progress(
-                        &mut send,
-                        &mut context,
-                        &child_hash,
-                        (child + 1) as u64,
-                    )
+                    .write_with_progress(&mut send, &mut context, &child_hash, (child + 1) as u64)
                     .await?;
             }
         }
-        send.finish()?;
+        send.0.finish()?;
         Ok(Default::default())
     }
 

src/get.rs

@@ -41,7 +41,7 @@ pub mod request;
 pub(crate) use error::get_error;
 pub use error::{GetError, GetResult};
 
-pub struct IrohStreamWriter(iroh::endpoint::SendStream);
+pub struct IrohStreamWriter(pub iroh::endpoint::SendStream);
 
 impl AsyncStreamWriter for IrohStreamWriter {
     async fn write(&mut self, data: &[u8]) -> io::Result<()> {
@@ -57,7 +57,7 @@ impl AsyncStreamWriter for IrohStreamWriter {
     }
 }
 
-pub struct IrohStreamReader(iroh::endpoint::RecvStream);
+pub struct IrohStreamReader(pub iroh::endpoint::RecvStream);
 
 impl AsyncStreamReader for IrohStreamReader {
     async fn read<const N: usize>(&mut self) -> io::Result<[u8; N]> {

src/provider.rs

@@ -12,6 +12,7 @@ use std::{
 use anyhow::{Context, Result};
 use bao_tree::ChunkRanges;
 use iroh::endpoint::{self, RecvStream, SendStream};
+use iroh_io::{AsyncStreamReader, AsyncStreamWriter};
 use n0_future::StreamExt;
 use quinn::{ClosedStream, ConnectionError, ReadToEndError};
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
@@ -23,6 +24,7 @@ use crate::{
         blobs::{Bitfield, WriteProgress},
         ExportBaoResult, Store,
     },
+    get::{IrohStreamReader, IrohStreamWriter},
     hashseq::HashSeq,
     protocol::{GetManyRequest, GetRequest, ObserveItem, ObserveRequest, PushRequest, Request},
     provider::events::{ClientConnected, ClientResult, ConnectionClosed, RequestTracker},
@@ -31,6 +33,9 @@ use crate::{
 pub mod events;
 use events::EventSender;
 
+type DefaultWriter = IrohStreamWriter;
+type DefaultReader = IrohStreamReader;
+
 /// Statistics about a successful or failed transfer.
 #[derive(Debug, Serialize, Deserialize)]
 pub struct TransferStats {
@@ -106,7 +111,7 @@ impl StreamPair {
             return Err(e);
         };
         Ok(ProgressWriter::new(
-            self.writer,
+            IrohStreamWriter(self.writer),
             WriterContext {
                 t0: self.t0,
                 other_bytes_read: self.other_bytes_read,
@@ -130,7 +135,7 @@ impl StreamPair {
             return Err(e);
         };
         Ok(ProgressReader {
-            inner: self.reader,
+            inner: IrohStreamReader(self.reader),
             context: ReaderContext {
                 t0: self.t0,
                 other_bytes_read: self.other_bytes_read,
@@ -282,14 +287,14 @@ impl WriteProgress for WriterContext {
 
 /// Wrapper for a [`quinn::SendStream`] with additional per request information.
 #[derive(Debug)]
-pub struct ProgressWriter {
+pub struct ProgressWriter<W: AsyncStreamWriter = DefaultWriter> {
     /// The quinn::SendStream to write to
-    pub inner: SendStream,
+    pub inner: W,
     pub(crate) context: WriterContext,
 }
 
-impl ProgressWriter {
-    fn new(inner: SendStream, context: WriterContext) -> Self {
+impl<W: AsyncStreamWriter> ProgressWriter<W> {
+    fn new(inner: W, context: WriterContext) -> Self {
         Self { inner, context }
     }
 
@@ -465,7 +470,7 @@ pub async fn handle_push(
     if !root_ranges.is_empty() {
         // todo: send progress from import_bao_quinn or rename to import_bao_quinn_with_progress
         store
-            .import_bao_quinn(hash, root_ranges.clone(), &mut reader.inner)
+            .import_bao_reader(hash, root_ranges.clone(), &mut reader.inner)
             .await?;
     }
     if request.ranges.is_blob() {
@@ -480,7 +485,7 @@ pub async fn handle_push(
             continue;
         }
         store
-            .import_bao_quinn(child_hash, child_ranges.clone(), &mut reader.inner)
+            .import_bao_reader(child_hash, child_ranges.clone(), &mut reader.inner)
             .await?;
     }
     Ok(())
@@ -496,7 +501,7 @@ pub(crate) async fn send_blob(
 ) -> ExportBaoResult<()> {
     store
         .export_bao(hash, ranges)
-        .write_quinn_with_progress(&mut writer.inner, &mut writer.context, &hash, index)
+        .write_with_progress(&mut writer.inner, &mut writer.context, &hash, index)
         .await
 }
 
@@ -527,7 +532,7 @@ pub async fn handle_observe(
                 send_observe_item(writer, &diff).await?;
                 old = new;
             }
-            _ = writer.inner.stopped() => {
+            _ = writer.inner.0.stopped() => {
                 debug!("observer closed");
                 break;
             }
@@ -539,13 +544,13 @@ pub async fn handle_observe(
 async fn send_observe_item(writer: &mut ProgressWriter, item: &Bitfield) -> Result<()> {
     use irpc::util::AsyncWriteVarintExt;
     let item = ObserveItem::from(item);
-    let len = writer.inner.write_length_prefixed(item).await?;
+    let len = writer.inner.0.write_length_prefixed(item).await?;
     writer.context.log_other_write(len);
     Ok(())
 }
 
-pub struct ProgressReader {
-    inner: RecvStream,
+pub struct ProgressReader<R: AsyncStreamReader = DefaultReader> {
+    inner: R,
     context: ReaderContext,
 }