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

3 files changed, 27 insertions, 19 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() {}