feat(provider): emit events about outgoing transfers

Author: dignifiedquire

src/lib.rs

@@ -28,8 +28,8 @@ mod tests {
     use anyhow::{anyhow, Context, Result};
     use rand::RngCore;
     use testdir::testdir;
-    use tokio::fs;
     use tokio::io::{self, AsyncReadExt, AsyncWriteExt};
+    use tokio::{fs, sync::broadcast};
     use tracing_subscriber::{prelude::*, EnvFilter};
 
     use crate::protocol::AuthToken;
@@ -224,20 +224,31 @@ mod tests {
         let mut provider_events = provider.subscribe();
         let events_task = tokio::task::spawn(async move {
             let mut events = Vec::new();
-            while let Ok(event) = provider_events.recv().await {
-                match event {
-                    Event::TransferCompleted { .. } | Event::TransferAborted { .. } => {
-                        events.push(event);
-                        break;
-                    }
-                    _ => events.push(event),
+            loop {
+                match provider_events.recv().await {
+                    Ok(event) => match event {
+                        Event::TransferCollectionCompleted { .. }
+                        | Event::TransferAborted { .. } => {
+                            events.push(event);
+                            break;
+                        }
+                        _ => events.push(event),
+                    },
+                    Err(e) => match e {
+                        broadcast::error::RecvError::Closed => {
+                            break;
+                        }
+                        broadcast::error::RecvError::Lagged(num) => {
+                            panic!("unable to keep up, skipped {num} messages");
+                        }
+                    },
                 }
             }
             events
         });
 
         let opts = get::Options {
-            addr: provider.listen_addr(),
+            addr: dbg!(provider.listen_addr()),
             peer_id: Some(provider.peer_id()),
             keylog: true,
         };
@@ -281,16 +292,35 @@ mod tests {
         provider.shutdown();
         provider.await?;
 
-        assert_events(events);
+        assert_events(events, num_blobs);
 
         Ok(())
     }
 
-    fn assert_events(events: Vec<Event>) {
-        assert_eq!(events.len(), 3);
+    fn assert_events(events: Vec<Event>, num_blobs: usize) {
+        let num_basic_events = 4;
+        let num_total_events = num_basic_events + num_blobs;
+        assert_eq!(
+            events.len(),
+            num_total_events,
+            "missing events, only got {:#?}",
+            events
+        );
         assert!(matches!(events[0], Event::ClientConnected { .. }));
         assert!(matches!(events[1], Event::RequestReceived { .. }));
-        assert!(matches!(events[2], Event::TransferCompleted { .. }));
+        assert!(matches!(events[2], Event::TransferCollectionStarted { .. }));
+        for (i, event) in events[3..num_total_events - 1].iter().enumerate() {
+            match event {
+                Event::TransferBlobCompleted { index, .. } => {
+                    assert_eq!(*index, i as u64);
+                }
+                _ => panic!("unexpected event {:?}", event),
+            }
+        }
+        assert!(matches!(
+            events.last().unwrap(),
+            Event::TransferCollectionCompleted { .. }
+        ));
     }
 
     fn setup_logging() {
@@ -328,7 +358,7 @@ mod tests {
                         match maybe_event {
                             Ok(event) => {
                                 match event {
-                                    Event::TransferCompleted { .. } => provider.shutdown(),
+                                    Event::TransferCollectionCompleted { .. } => provider.shutdown(),
                                     Event::TransferAborted { .. } => {
                                         break Err(anyhow!("transfer aborted"));
                                     }

src/provider/mod.rs

@@ -327,12 +327,36 @@ pub enum Event {
         /// The hash for which the client wants to receive data.
         hash: Hash,
     },
-    /// A request was completed and the data was sent to the client.
-    TransferCompleted {
+    /// A collection has been found and is being transferred.
+    TransferCollectionStarted {
         /// An unique connection id.
         connection_id: u64,
         /// An identifier uniquely identifying this transfer request.
         request_id: u64,
+        /// The number of blobs in the collection.
+        num_blobs: u64,
+        /// The total blob size of the data.
+        total_blobs_size: u64,
+    },
+    /// A collection request was completed and the data was sent to the client.
+    TransferCollectionCompleted {
+        /// An unique connection id.
+        connection_id: u64,
+        /// An identifier uniquely identifying this transfer request.
+        request_id: u64,
+    },
+    /// A blob in a collection was transferred.
+    TransferBlobCompleted {
+        /// An unique connection id.
+        connection_id: u64,
+        /// An identifier uniquely identifying this transfer request.
+        request_id: u64,
+        /// The hash of the blob
+        hash: Hash,
+        /// The index of the blob in the collection.
+        index: u64,
+        /// The size of the blob transferred.
+        size: u64,
     },
     /// A request was aborted because the client disconnected.
     TransferAborted {
@@ -625,6 +649,7 @@ async fn read_request(mut reader: quinn::RecvStream, buffer: &mut BytesMut) -> R
 /// close the writer, and return with `Ok(SentStatus::NotFound)`.
 ///
 /// If the transfer does _not_ end in error, the buffer will be empty and the writer is gracefully closed.
+#[allow(clippy::too_many_arguments)]
 async fn transfer_collection(
     // Database from which to fetch blobs.
     db: &Database,
@@ -636,6 +661,9 @@ async fn transfer_collection(
     outboard: &Bytes,
     // The actual blob data.
     data: &Bytes,
+    events: broadcast::Sender<Event>,
+    connection_id: u64,
+    request_id: u64,
 ) -> Result<SentStatus> {
     // We only respond to requests for collections, not individual blobs
     let mut extractor = SliceExtractor::new_outboard(
@@ -652,6 +680,13 @@ async fn transfer_collection(
 
     let c: Collection = postcard::from_bytes(data)?;
 
+    let _ = events.send(Event::TransferCollectionStarted {
+        connection_id,
+        request_id,
+        num_blobs: c.blobs.len() as u64,
+        total_blobs_size: c.total_blobs_size,
+    });
+
     // TODO: we should check if the blobs referenced in this container
     // actually exist in this provider before returning `FoundCollection`
     write_response(
@@ -667,12 +702,21 @@ async fn transfer_collection(
     writer.write_buf(&mut data).await?;
     for (i, blob) in c.blobs.iter().enumerate() {
         debug!("writing blob {}/{}", i, c.blobs.len());
-        let (status, writer1) = send_blob(db.clone(), blob.hash, writer, buffer).await?;
+        tokio::task::yield_now().await;
+        let (status, writer1, size) = send_blob(db.clone(), blob.hash, writer, buffer).await?;
         writer = writer1;
         if SentStatus::NotFound == status {
             writer.finish().await?;
             return Ok(status);
         }
+
+        let _ = events.send(Event::TransferBlobCompleted {
+            connection_id,
+            request_id,
+            hash: blob.hash,
+            index: i as u64,
+            size,
+        });
     }
 
     writer.finish().await?;
@@ -740,9 +784,20 @@ async fn handle_stream(
     };
 
     // 5. Transfer data!
-    match transfer_collection(&db, writer, &mut out_buffer, &outboard, &data).await {
+    match transfer_collection(
+        &db,
+        writer,
+        &mut out_buffer,
+        &outboard,
+        &data,
+        events.clone(),
+        connection_id,
+        request_id,
+    )
+    .await
+    {
         Ok(SentStatus::Sent) => {
-            let _ = events.send(Event::TransferCompleted {
+            let _ = events.send(Event::TransferCollectionCompleted {
                 connection_id,
                 request_id,
             });
@@ -771,7 +826,7 @@ async fn send_blob<W: AsyncWrite + Unpin + Send + 'static>(
     name: Hash,
     mut writer: W,
     buffer: &mut BytesMut,
-) -> Result<(SentStatus, W)> {
+) -> Result<(SentStatus, W, u64)> {
     match db.get(&name) {
         Some(BlobOrCollection::Blob(Data {
             outboard,
@@ -796,11 +851,12 @@ async fn send_blob<W: AsyncWrite + Unpin + Send + 'static>(
                 std::io::Result::Ok(writer)
             })
             .await??;
-            Ok((SentStatus::Sent, writer))
+
+            Ok((SentStatus::Sent, writer, size))
         }
         _ => {
             write_response(&mut writer, buffer, Res::NotFound).await?;
-            Ok((SentStatus::NotFound, writer))
+            Ok((SentStatus::NotFound, writer, 0))
         }
     }
 }