rpc: multipart reader needs no completion

Author: nicoburniske

rpc/src/rpc/parts.rs

@@ -6,6 +6,8 @@ use ql_common::ResetCode;
 use crate::{codec, read_bytes, ChunkQueue, RpcCodec, RpcError, RpcRead};
 
 /// reads a stream of framed byte parts
+///
+/// dropping the reader stops receiving further parts
 pub struct MultipartReader<H, R>
 where
     H: RpcCodec,
@@ -66,22 +68,6 @@ where
         Err(crate::Error::UnexpectedFrameKind(kind.tag()).into())
     }
 
-    /// rejects any remaining part
-    pub async fn complete(mut self) -> Result<(), RpcError<H::Error, R::Error>> {
-        let Some(frame) = self.read_frame().await? else {
-            return Ok(());
-        };
-
-        let kind = match frame {
-            PartFrame::PartHeader(_) => FrameKind::PartHeader,
-            PartFrame::BodyBytes(_) => FrameKind::BodyChunk,
-            PartFrame::EndPart => FrameKind::EndPart,
-        };
-
-        self.reset_inner(ResetCode::PROTOCOL);
-        Err(crate::Error::UnexpectedFrameKind(kind.tag()).into())
-    }
-
     pub fn reset(mut self, code: ResetCode) {
         self.reset_inner(code);
     }
@@ -117,6 +103,18 @@ where
     }
 }
 
+impl<H, R> Drop for MultipartReader<H, R>
+where
+    H: RpcCodec,
+    R: RpcRead,
+{
+    fn drop(&mut self) {
+        if !self.finished {
+            self.stream.reset(ResetCode::CANCELLED);
+        }
+    }
+}
+
 impl<H, R> MultipartPart<'_, H, R>
 where
     H: RpcCodec,

rpc/src/rpc/tests.rs

@@ -440,6 +440,7 @@ fn download() {
     for chunk_size in chunk_sizes() {
         let (client, server) = stream_pair(chunk_size);
         let client = async move {
+            let pipe = client.reader.pipe.clone();
             let download = download::start::<Rpc, _>(client, &b"file".to_vec())
                 .await
                 .unwrap();
@@ -457,7 +458,8 @@ fn download() {
             }
             assert_eq!(body, b"abcdef".to_vec());
             drop(part);
-            reader.complete().await.unwrap();
+            drop(reader);
+            assert_eq!(pipe.borrow().reset, Some(ResetCode::CANCELLED));
         };
         let server = download::handle_download::<(), Rpc, _, _, _, _>(
             (),
@@ -524,7 +526,7 @@ fn upload() {
                 }
                 assert_eq!(body, b"abcdef".to_vec());
                 drop(part);
-                upload.complete().await.unwrap();
+                assert!(upload.next_part().await.unwrap().is_none());
                 responder.respond(b"stored".to_vec()).await.unwrap();
             },
             |_, error| panic!("upload failed: {error:?}"),