Make PackagesWriter streaming - you don't need to give count when creating it

4 files changed, 30 insertions, 22 deletions

diff --git a/src/monitor-gui.cc b/src/monitor-gui.cc
index b02c4de..67b9b63 100644
--- a/src/monitor-gui.cc
+++ b/src/monitor-gui.cc
@@ -1633,13 +1633,12 @@ private:
       show_error("Unable to open " + file + " for writing");
       return false;
     }
-    auto writer = std::unique_ptr<PackagesWriter>(
-        PackagesWriter::create(packages->rows(), &out));
+    auto writer = std::unique_ptr<PackagesWriter>(PackagesWriter::create(&out));
     for (size_t i = 0; i < packages->rows(); ++i) {
       auto const& pkg = packages->package(i);
       writer->write(pkg.pkg, pkg.data);
     }
-    writer.reset();
+    writer->flush();
     if (!out.good()) {
       show_error("Error writing " + file);
       return false;
diff --git a/src/packages.cc b/src/packages.cc
index 0309df1..f2c4ba6 100644
--- a/src/packages.cc
+++ b/src/packages.cc
@@ -14,26 +14,23 @@ namespace {
 
 class PackagesWriterImpl : public PackagesWriter {
 public:
-  PackagesWriterImpl(size_t count, std::ostream* out)
-    : count_(count), out_(out) {
+  PackagesWriterImpl(std::ostream* out)
+    : open_(true), count_(0), out_(out) {
     uint8_t header[8];
     memcpy(header, "TPP", 3);
     header[3] = 0x1;  // Version
-    write_u32(header + 4, count_);  // Count
+    write_u32(header + 4, 0);  // Count
     out->write(reinterpret_cast<char*>(header), 8);
-
-    if (count == 0) {
-      write_u64(header, 0);  // EOF
-      out_->write(reinterpret_cast<char*>(header), 8);
-    }
   }
 
   ~PackagesWriterImpl() {
-    assert(count_ == 0);
+    if (open_) {
+      flush();
+    }
   }
 
   void write(Package const& package, std::string const& data) override {
-    if (count_ == 0) {
+    if (!open_) {
       assert(false);
       return;
     }
@@ -52,14 +49,25 @@ public:
     out_->write(reinterpret_cast<char*>(ptr), need);
     backup.reset();
     out_->write(data.data(), data.size());
+    count_++;
+  }
 
-    if (--count_ == 0) {
-      write_u64(size, 0);  // EOF
-      out_->write(reinterpret_cast<char*>(size), 8);
+  void flush() override {
+    if (!open_) return;
+    open_ = false;
+    uint8_t footer[8];
+    write_u64(footer, 0);  // EOF
+    out_->write(reinterpret_cast<char*>(footer), 8);
+    if (count_ > 0) {
+      out_->seekp(4);
+      write_u32(footer, count_);
+      out_->write(reinterpret_cast<char*>(footer), 4);
+      count_ = 0;
     }
   }
 
 private:
+  bool open_;
   size_t count_;
   std::ostream* const out_;
 };
@@ -67,8 +75,8 @@ private:
 }  // namespace
 
 // static
-PackagesWriter* PackagesWriter::create(size_t count, std::ostream* out) {
-  return new PackagesWriterImpl(count, out);
+PackagesWriter* PackagesWriter::create(std::ostream* out) {
+  return new PackagesWriterImpl(out);
 }
 
 // static
diff --git a/src/packages.hh b/src/packages.hh
index e646f92..f696212 100644
--- a/src/packages.hh
+++ b/src/packages.hh
@@ -9,9 +9,10 @@ class PackagesWriter {
 public:
   virtual ~PackagesWriter() {}
 
-  static PackagesWriter* create(size_t count, std::ostream* out);
+  static PackagesWriter* create(std::ostream* out);
 
   virtual void write(Package const& package, std::string const& data) = 0;
+  virtual void flush() = 0;
 
 protected:
   PackagesWriter() {}
diff --git a/test/test-packages.cc b/test/test-packages.cc
index d0c526c..fe56077 100644
--- a/test/test-packages.cc
+++ b/test/test-packages.cc
@@ -77,9 +77,9 @@ bool test_sanity() {
 
   std::stringstream stream;
   auto writer = std::unique_ptr<PackagesWriter>(
-      PackagesWriter::create(1, &stream));
+      PackagesWriter::create(&stream));
   writer->write(pkg1, "Hello World!");
-  writer.reset();
+  writer->flush();
 
   Delegate delegate;
   ASSERT_EQ(PackagesReader::GOOD, PackagesReader::read(stream, &delegate));
@@ -93,7 +93,7 @@ bool test_sanity() {
 
 bool test_empty() {
   std::stringstream stream;
-  delete PackagesWriter::create(0, &stream);
+  delete PackagesWriter::create(&stream);
   ASSERT_TRUE(!stream.str().empty());
   Delegate delegate;
   ASSERT_EQ(PackagesReader::GOOD, PackagesReader::read(stream, &delegate));