[antlir2][sendstream_parser] convert to more idiomatic stream interface

Summary:
The old callback style was pretty much required prior to this diff stack, but
it's not very idiomatic in Rust. This diff converts it to a `futures::Stream`
which is more standard and less confusing

Test Plan:
```
❯ buck2 test fbcode//antlir/antlir2/sendstream_parser/...
Buck UI: https://www.internalfb.com/buck2/31c8c888-6847-4828-9261-32f3b78b7845
Tests finished: Pass 4. Fail 0. Fatal 0. Skip 0. Build failure 0
```
https://www.internalfb.com/intern/testinfra/testrun/4785074928638786

Differential Revision: D86802448

fbshipit-source-id: 6d4240286d3b85ef473bcd8a5ef4d592b820cc2a

Author: vmagro

antlir/antlir2/sendstream_parser/src/lib.rs

@@ -719,6 +719,8 @@ mod tests {
     use std::fmt::Write;
     use std::io::Cursor;
 
+    use futures::StreamExt;
+    use futures::future;
     use similar_asserts::SimpleDiff;
 
     use super::*;
@@ -746,12 +748,13 @@ mod tests {
         let cursor = Cursor::new(data);
         let mut parsed_txt = String::new();
         let mut sendstream_index = 0;
-        let num_cmds_parsed = wire::parse(cursor, |cmd| {
-            serialize_cmd(&mut sendstream_index, &mut parsed_txt, cmd);
-            wire::ParserControl::KeepGoing
-        })
-        .await
-        .expect("while parsing");
+        let mut num_cmds_parsed = 0;
+        let mut parser = wire::parse(cursor);
+        while let Some(cmd_res) = parser.next().await {
+            let cmd = cmd_res.expect("failed to parse");
+            serialize_cmd(&mut sendstream_index, &mut parsed_txt, &cmd);
+            num_cmds_parsed += 1;
+        }
         if let Some(dst) = std::env::var_os("UPDATE_DEMO_TXT") {
             std::fs::write(dst, parsed_txt).expect("while writing to {dst}");
         } else {
@@ -766,32 +769,6 @@ mod tests {
         assert_eq!(num_cmds_parsed, 106);
     }
 
-    /// Demonstrate how we might eagerly abort parsing after collecting information embedded in an
-    /// early command.
-    #[tokio::test]
-    async fn partial_parse() {
-        let data = include_bytes!("../testdata/demo.sendstream");
-        let cursor = Cursor::new(data);
-        let mut uuid: Option<Uuid> = None;
-        let num_cmds_parsed = wire::parse(cursor, |cmd| {
-            if let Command::Subvol(sv) = cmd {
-                uuid = Some(sv.uuid());
-                return wire::ParserControl::Enough;
-            }
-            wire::ParserControl::KeepGoing
-        })
-        .await
-        .expect("while parsing");
-        assert_eq!(
-            uuid,
-            Some(
-                Uuid::parse_str("0fbf2b5f-ff82-a748-8b41-e35aec190b49")
-                    .expect("while parsing uuid")
-            )
-        );
-        assert_eq!(num_cmds_parsed, 1);
-    }
-
     #[tokio::test]
     async fn sendstream_covers_all_commands() {
         let all_cmds: BTreeSet<_> = wire::cmd::PARSED_SUBTYPES
@@ -803,13 +780,15 @@ mod tests {
             .collect();
         let data = include_bytes!("../testdata/demo.sendstream");
         let cursor = Cursor::new(data);
-        let mut seen_cmds: BTreeSet<wire::cmd::CommandType> = BTreeSet::new();
-        wire::parse(cursor, |cmd| {
-            seen_cmds.insert(cmd.command_type());
-            wire::ParserControl::KeepGoing
-        })
-        .await
-        .expect("while parsing");
+        let seen_cmds: BTreeSet<_> = wire::parse(cursor)
+            .filter_map(|i| {
+                future::ready(match i {
+                    Ok(cmd) => Some(cmd.command_type()),
+                    Err(e) => panic!("failed to parse: {e}"),
+                })
+            })
+            .collect()
+            .await;
         assert_eq!(seen_cmds, all_cmds);
     }
 }

antlir/antlir2/sendstream_parser/src/wire/mod.rs

@@ -5,7 +5,9 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use futures::Stream;
 use futures::StreamExt;
+use futures::future;
 use tokio::io::AsyncRead;
 use tokio_util::codec::FramedRead;
 
@@ -15,45 +17,83 @@ mod nombytes;
 mod tlv;
 pub use nombytes::NomBytes;
 
-#[derive(Debug)]
-pub enum ParserControl {
-    KeepGoing,
-    Enough,
-}
-
 /// Parse an async source of bytes, expecting to find it to contain one or more sendstreams.
 /// Because the parsed commands reference data owned by the source, we do not collect the commands.
 /// Instead, we allow the caller to process them via `f`, which can instruct the processing to
 /// continue or shut down gracefully via the returned `ParserControl`.
 ///
 /// Each sendstream is expected to (1) start with a header, followed by (2) either a Subvol or
 /// Snapshot command, followed by (3) 0 or more additional commands, terminated by (4) an End
-/// command. Note that we don't validate #2 here, but we do expect #1 and #4.
-///
-/// Returns number of commands parsed.
+/// command. Note that only (1) is actually enforced, afterward every command
+/// will be emitted into the stream as long as it could be read and parsed.
 ///
 /// See https://btrfs.readthedocs.io/en/latest/dev/dev-send-stream.html for reference.
-pub async fn parse<R, F>(reader: R, mut f: F) -> crate::Result<u128>
+pub fn parse<R>(reader: R) -> impl Stream<Item = crate::Result<crate::Command>>
 where
-    R: AsyncRead + Unpin + Send,
-    F: FnMut(&crate::Command) -> ParserControl + Send,
+    R: AsyncRead,
 {
-    let mut reader = FramedRead::new(reader, framed::SendstreamDecoder::new());
-    let mut command_count = 0;
-    while let Some(item_res) = reader.next().await {
-        match item_res {
-            Ok(framed::Item::Command(command)) => {
-                command_count += 1;
-                if let ParserControl::Enough = f(&command) {
-                    // caller got what they needed, no need to continue parsing
-                    break;
+    let reader = FramedRead::new(reader, framed::SendstreamDecoder::new());
+    reader.filter_map(|item_res| {
+        future::ready(match item_res {
+            Ok(framed::Item::Command(command)) => Some(Ok(command)),
+            Ok(framed::Item::SendstreamStart(_)) => None,
+            Err(e) => Some(Err(e)),
+        })
+    })
+}
+
+#[cfg(test)]
+#[allow(clippy::expect_used)]
+mod tests {
+    use std::io::Cursor;
+    use std::time::Duration;
+
+    use futures::StreamExt;
+    use tokio::io::AsyncWriteExt;
+    use tokio::time::sleep;
+
+    use super::*;
+
+    /// Historically, we couldn't stream commands as they were parsed very well,
+    /// so the early exit was implemented with a not-very-Rusty callback
+    /// interface. This test proves that the parser stops operating as soon as
+    /// the caller stops asking for commands (via a reader that will only
+    /// provide enough bytes for a few commands before then stalling forever)
+    #[tokio::test]
+    async fn early_exit() {
+        let make_parser = |truncate: Option<usize>| async move {
+            let src = include_bytes!("../../testdata/demo.sendstream");
+            // only use simplex stream if we're going to truncate it
+            // prematurely, otherwise use a Cursor on top of the static byte
+            // slice so that the framed codec actually sees the EOF
+            let reader: Box<dyn AsyncRead + Unpin> = match truncate {
+                Some(size) => {
+                    let (receiver, mut sender) = tokio::io::simplex(size);
+                    sender
+                        .write_all(&src[..size])
+                        .await
+                        .expect("failed to write input data");
+                    Box::new(receiver)
                 }
-            }
-            Ok(framed::Item::SendstreamStart(_)) => {}
-            Err(e) => {
-                return Err(e);
-            }
-        }
+                None => Box::new(Cursor::new(src)),
+            };
+            parse(reader)
+        };
+
+        // empirically determined that 46 commands fit in the first 4k
+        // first, prove that we can read the expected number of commands, before
+        // dropping the stream which while then stop consuming its upstream input
+        let parser = make_parser(Some(4 * 1024)).await;
+        let count = parser
+            .take_until(sleep(Duration::from_millis(100)))
+            .count()
+            .await;
+        assert_eq!(count, 46);
+
+        // the entire stream should be consumable (but obviously only when given
+        // the entire input)
+        let parser = make_parser(None).await;
+        let count = parser.count().await;
+        assert_eq!(count, 106);
     }
-    Ok(command_count)
 }