FIX: Fix support for multi-frame in AsyncDynRead

Author: threecgreen

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## 0.19.2 - TBD
+
+### Bug fixes
+- Fixed issue where `AsyncDynReader` would only decode the first frame of multi-frame
+  Zstandard files
+
 ## 0.19.1 - 2024-07-16
 
 ### Bug fixes

rust/dbn/src/decode.rs

@@ -570,6 +570,8 @@ mod r#async {
 
     use crate::enums::Compression;
 
+    use super::zstd::zstd_decoder;
+
     /// A type for runtime polymorphism on compressed and uncompressed input.
     /// The async version of [`DynReader`](super::DynReader).
     pub struct DynReader<R>(DynReaderImpl<R>)
@@ -629,7 +631,7 @@ mod r#async {
                 .await
                 .map_err(|e| crate::Error::io(e, "creating buffer to infer encoding"))?;
             Ok(if super::zstd::starts_with_prefix(first_bytes) {
-                Self(DynReaderImpl::ZStd(ZstdDecoder::new(reader)))
+                Self(DynReaderImpl::ZStd(zstd_decoder(reader)))
             } else {
                 Self(DynReaderImpl::Uncompressed(reader))
             })
@@ -693,4 +695,35 @@ mod r#async {
             }
         }
     }
+
+    #[cfg(test)]
+    mod tests {
+        use crate::{
+            compat::InstrumentDefMsgV1,
+            decode::{tests::TEST_DATA_PATH, AsyncDbnRecordDecoder},
+            VersionUpgradePolicy,
+        };
+
+        use super::*;
+
+        #[tokio::test]
+        async fn test_decode_multiframe_zst() {
+            let mut decoder = AsyncDbnRecordDecoder::with_version(
+                DynReader::from_file(&format!(
+                    "{TEST_DATA_PATH}/multi-frame.definition.v1.dbn.frag.zst"
+                ))
+                .await
+                .unwrap(),
+                1,
+                VersionUpgradePolicy::AsIs,
+                false,
+            )
+            .unwrap();
+            let mut count = 0;
+            while let Some(_rec) = decoder.decode::<InstrumentDefMsgV1>().await.unwrap() {
+                count += 1;
+            }
+            assert_eq!(count, 8);
+        }
+    }
 }

rust/dbn/src/decode/dbn/async.rs

@@ -8,21 +8,13 @@ use tokio::{
 
 use crate::{
     compat,
-    decode::{r#async::ZSTD_FILE_BUFFER_CAPACITY, FromLittleEndianSlice, VersionUpgradePolicy},
+    decode::{
+        r#async::ZSTD_FILE_BUFFER_CAPACITY, zstd::zstd_decoder, FromLittleEndianSlice,
+        VersionUpgradePolicy,
+    },
     HasRType, Metadata, Record, RecordHeader, RecordRef, Result, DBN_VERSION, METADATA_FIXED_LEN,
 };
 
-/// Helper to always set multiple members.
-fn zstd_decoder<R>(reader: R) -> ZstdDecoder<R>
-where
-    R: io::AsyncBufReadExt + Unpin,
-{
-    let mut zstd_decoder = ZstdDecoder::new(reader);
-    // explicitly enable decoding multiple frames
-    zstd_decoder.multiple_members(true);
-    zstd_decoder
-}
-
 /// An async decoder for Databento Binary Encoding (DBN), both metadata and records.
 pub struct Decoder<R>
 where

rust/dbn/src/decode/zstd.rs

@@ -15,6 +15,18 @@ pub fn starts_with_prefix(bytes: &[u8]) -> bool {
     ZSTD_MAGIC_NUMBER == magic
 }
 
+/// Helper to always set multiple members.
+#[cfg(feature = "async")]
+pub(crate) fn zstd_decoder<R>(reader: R) -> async_compression::tokio::bufread::ZstdDecoder<R>
+where
+    R: tokio::io::AsyncBufReadExt + Unpin,
+{
+    let mut zstd_decoder = async_compression::tokio::bufread::ZstdDecoder::new(reader);
+    // explicitly enable decoding multiple frames
+    zstd_decoder.multiple_members(true);
+    zstd_decoder
+}
+
 #[cfg(test)]
 mod tests {
     use std::{fs::File, io::Read};