FIX: Fix Zstd decoding and file reading

Author: threecgreen

.gitignore

@@ -1 +1,2 @@
 build
+!*.zst

include/databento/detail/file_stream.hpp

@@ -2,7 +2,8 @@
 
 #include <cstddef>  // size_t
 #include <cstdint>  // uint8_t
-#include <fstream>
+#include <fstream>  // ifstream
+#include <string>
 
 #include "databento/ireadable.hpp"
 

include/databento/detail/zstd_stream.hpp

@@ -13,8 +13,9 @@ namespace databento {
 namespace detail {
 class ZstdStream : public IReadable {
  public:
-  explicit ZstdStream(std::unique_ptr<IReadable> input,
-                      std::vector<std::uint8_t>&& in_buffer);
+  explicit ZstdStream(std::unique_ptr<IReadable> input);
+  ZstdStream(std::unique_ptr<IReadable> input,
+             std::vector<std::uint8_t>&& in_buffer);
 
   // Read exactly `length` bytes into `buffer`.
   void ReadExact(std::uint8_t* buffer, std::size_t length) override;
@@ -28,8 +29,6 @@ class ZstdStream : public IReadable {
   std::size_t read_suggestion_;
   std::vector<std::uint8_t> in_buffer_;
   ZSTD_inBuffer z_in_buffer_;
-  // std::vector<std::uint8_t> out_buffer_;
-  // ZSTD_outBuffer z_out_buffer_;
 };
 }  // namespace detail
 }  // namespace databento

src/detail/file_stream.cpp

@@ -1,6 +1,7 @@
 #include "databento/detail/file_stream.hpp"
 
 #include <ios>  // streamsize
+#include <sstream>
 
 #include "databento/exceptions.hpp"
 
@@ -14,12 +15,17 @@ FileStream::FileStream(const std::string& file_path) : stream_{file_path} {
 }
 
 void FileStream::ReadExact(std::uint8_t* buffer, std::size_t length) {
-  stream_.read(reinterpret_cast<char*>(buffer),
-               static_cast<std::streamsize>(length));
+  const auto size = ReadSome(buffer, length);
+  if (size != length) {
+    std::ostringstream err_msg;
+    err_msg << "Unexpected end of file, expected " << length << " bytes, got "
+            << size;
+    throw DbnResponseError{err_msg.str()};
+  }
 }
 
 std::size_t FileStream::ReadSome(std::uint8_t* buffer, std::size_t max_length) {
-  return static_cast<std::size_t>(
-      stream_.readsome(reinterpret_cast<char*>(buffer),
-                       static_cast<std::streamsize>(max_length)));
+  stream_.read(reinterpret_cast<char*>(buffer),
+               static_cast<std::streamsize>(max_length));
+  return static_cast<std::size_t>(stream_.gcount());
 }

src/detail/zstd_stream.cpp

@@ -7,6 +7,9 @@
 
 using databento::detail::ZstdStream;
 
+ZstdStream::ZstdStream(std::unique_ptr<IReadable> input)
+    : ZstdStream{std::move(input), {}} {}
+
 ZstdStream::ZstdStream(std::unique_ptr<IReadable> input,
                        std::vector<std::uint8_t>&& in_buffer)
     : input_{std::move(input)},
@@ -31,26 +34,33 @@ void ZstdStream::ReadExact(std::uint8_t* buffer, std::size_t length) {
 
 size_t ZstdStream::ReadSome(std::uint8_t* buffer, std::size_t max_length) {
   ZSTD_outBuffer z_out_buffer{buffer, max_length, 0};
-  const auto unread_input = z_in_buffer_.size - z_in_buffer_.pos;
-  if (unread_input > 0) {
-    std::copy(
-        in_buffer_.cbegin() + static_cast<std::ptrdiff_t>(z_in_buffer_.pos),
-        in_buffer_.cend(), in_buffer_.begin());
-  }
-  const auto new_size = unread_input + read_suggestion_;
-  if (new_size != in_buffer_.size()) {
-    in_buffer_.resize(new_size);
-    z_in_buffer_.src = in_buffer_.data();
-  }
-  z_in_buffer_.size = unread_input + input_->ReadSome(&in_buffer_[unread_input],
-                                                      read_suggestion_);
-  z_in_buffer_.pos = 0;
-
-  read_suggestion_ =
-      ::ZSTD_decompressStream(z_dstream_.get(), &z_out_buffer, &z_in_buffer_);
-  if (::ZSTD_isError(read_suggestion_)) {
-    throw DbnResponseError{std::string{"Zstd error decompressing record: "} +
-                           ::ZSTD_getErrorName(read_suggestion_)};
-  }
+  std::size_t read_size = 0;
+  do {
+    const auto unread_input = z_in_buffer_.size - z_in_buffer_.pos;
+    if (unread_input > 0) {
+      std::copy(
+          in_buffer_.cbegin() + static_cast<std::ptrdiff_t>(z_in_buffer_.pos),
+          in_buffer_.cend(), in_buffer_.begin());
+    }
+    if (read_suggestion_ == 0) {
+      // next frame
+      read_suggestion_ = ::ZSTD_initDStream(z_dstream_.get());
+    }
+    const auto new_size = unread_input + read_suggestion_;
+    if (new_size != in_buffer_.size()) {
+      in_buffer_.resize(new_size);
+      z_in_buffer_.src = in_buffer_.data();
+    }
+    read_size = input_->ReadSome(&in_buffer_[unread_input], read_suggestion_);
+    z_in_buffer_.size = unread_input + read_size;
+    z_in_buffer_.pos = 0;
+
+    read_suggestion_ =
+        ::ZSTD_decompressStream(z_dstream_.get(), &z_out_buffer, &z_in_buffer_);
+    if (::ZSTD_isError(read_suggestion_)) {
+      throw DbnResponseError{std::string{"Zstd error decompressing: "} +
+                             ::ZSTD_getErrorName(read_suggestion_)};
+    }
+  } while (z_out_buffer.pos == 0 && read_size > 0);
   return z_out_buffer.pos;
 }

test/CMakeLists.txt

@@ -29,6 +29,7 @@ set(
   src/datetime_tests.cpp
   src/dbn_tests.cpp
   src/dbn_decoder_tests.cpp
+  src/file_stream_tests.cpp
   src/flag_set_tests.cpp
   src/historical_tests.cpp
   src/live_tests.cpp
@@ -44,6 +45,7 @@ set(
   src/stream_op_helper_tests.cpp
   src/symbology_tests.cpp
   src/tcp_client_tests.cpp
+  src/zstd_stream_tests.cpp
 )
 add_executable(${PROJECT_NAME} ${test_headers} ${test_sources})
 find_package(Threads REQUIRED)

test/data/multi-frame.definition.zst

@@ -0,0 +1,33 @@
+#include <gtest/gtest.h>
+
+#include "databento/detail/file_stream.hpp"
+#include "databento/exceptions.hpp"
+
+namespace databento {
+namespace detail {
+namespace test {
+TEST(FileStreamTests, TestReadExactInsufficient) {
+  const std::string file_path = TEST_BUILD_DIR "/data/test_data.ohlcv-1d.dbn";
+  databento::detail::FileStream target{file_path};
+  std::vector<std::uint8_t> buffer(1024);  // File is less than 1KiB
+  try {
+    target.ReadExact(buffer.data(), buffer.size());
+    FAIL() << "Expected throw";
+  } catch (const databento::Exception& exc) {
+    ASSERT_STREQ(exc.what(),
+                 "Unexpected end of file, expected 1024 bytes, got 206");
+  }
+}
+
+TEST(FileStreamTests, TestReadSomeLessThanMax) {
+  const std::string file_path = TEST_BUILD_DIR "/data/test_data.ohlcv-1d.dbn";
+  databento::detail::FileStream target{file_path};
+  std::vector<std::uint8_t> buffer(1024);  // File is less than 1KiB
+  const auto read_size = target.ReadSome(buffer.data(), buffer.size());
+  ASSERT_GT(read_size, 0);
+  ASSERT_TRUE(std::any_of(buffer.cbegin(), buffer.cend(),
+                          [](std::uint8_t byte) { return byte != 0; }));
+}
+}  // namespace test
+}  // namespace detail
+}  // namespace databento

test/src/file_stream_tests.cpp

@@ -0,0 +1,32 @@
+#include <gtest/gtest.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include "databento/detail/file_stream.hpp"
+#include "databento/detail/zstd_stream.hpp"
+#include "databento/enums.hpp"
+#include "databento/ireadable.hpp"
+#include "databento/record.hpp"
+
+namespace databento {
+namespace detail {
+namespace test {
+TEST(ZstdStreamTests, TestMultiFrameFiles) {
+  constexpr auto kRecordCount = 8;
+  const std::string file_path =
+      TEST_BUILD_DIR "/data/multi-frame.definition.zst";
+
+  databento::detail::ZstdStream target{std::unique_ptr<databento::IReadable>{
+      new databento::detail::FileStream{file_path}}};
+  for (std::size_t i = 0; i < kRecordCount; ++i) {
+    databento::InstrumentDefMsg def_msg;
+    target.ReadExact(reinterpret_cast<std::uint8_t*>(&def_msg),
+                     sizeof(def_msg));
+    EXPECT_EQ(def_msg.hd.rtype, databento::rtype::InstrumentDef);
+  }
+}
+}  // namespace test
+}  // namespace detail
+}  // namespace databento