diff --git a/src/dwarfs/filesystem_writer.cpp b/src/dwarfs/filesystem_writer.cpp index 9701dbec..5f08b80d 100644 --- a/src/dwarfs/filesystem_writer.cpp +++ b/src/dwarfs/filesystem_writer.cpp @@ -84,7 +84,8 @@ class fsblock { folly::Function set_block_cb = nullptr); fsblock(section_type type, compression_type compression, - std::span data); + std::span data, + std::shared_ptr pctx = nullptr); fsblock(section_type type, block_compressor const& bc, std::span data, compression_type data_comp_type, @@ -254,10 +255,12 @@ class raw_fsblock : public fsblock::impl { class compressed_fsblock : public fsblock::impl { public: compressed_fsblock(section_type type, compression_type compression, - std::span range) + std::span range, + std::shared_ptr pctx) : type_{type} , compression_{compression} - , range_{range} {} + , range_{range} + , pctx_{std::move(pctx)} {} void compress(worker_group& wg, std::optional /* meta */) override { @@ -266,6 +269,10 @@ class compressed_fsblock : public fsblock::impl { wg.add_job([this, prom = std::move(prom)]() mutable { fsblock::build_section_header(header_, *this); + if (pctx_) { + pctx_->bytes_in += size(); + pctx_->bytes_out += size(); + } prom.set_value(); }); } @@ -295,6 +302,7 @@ class compressed_fsblock : public fsblock::impl { std::future future_; std::optional number_; section_header_v2 header_; + std::shared_ptr pctx_; }; class rewritten_fsblock : public fsblock::impl { @@ -412,8 +420,10 @@ fsblock::fsblock(section_type type, block_compressor const& bc, std::move(set_block_cb))) {} fsblock::fsblock(section_type type, compression_type compression, - std::span data) - : impl_(std::make_unique(type, compression, data)) {} + std::span data, + std::shared_ptr pctx) + : impl_(std::make_unique(type, compression, data, + std::move(pctx))) {} fsblock::fsblock(section_type type, block_compressor const& bc, std::span data, compression_type data_comp_type, @@ -856,7 +866,11 @@ void filesystem_writer_::write_compressed_section( { std::lock_guard lock(mx_); - auto fsb = std::make_unique(type, compression, data); + if (!pctx_) { + pctx_ = prog_.create_context(); + } + + auto fsb = std::make_unique(type, compression, data, pctx_); fsb->set_block_no(section_number_++); fsb->compress(wg_);