From 038541f94e63984ad84bef97c0a1483a85996171 Mon Sep 17 00:00:00 2001
From: Marcus Holland-Moritz <github@mhxnet.de>
Date: Wed, 21 May 2025 23:05:11 +0200
Subject: [PATCH] feat: show memory usage in `mkdwarfs`

---
 include/dwarfs/util.h                  |  3 ++
 include/dwarfs/writer/console_writer.h |  4 +++
 src/util.cpp                           | 40 ++++++++++++++++++++++++++
 src/writer/console_writer.cpp          | 10 +++++--
 tools/src/mkdwarfs_main.cpp            |  5 ++++
 5 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/include/dwarfs/util.h b/include/dwarfs/util.h
index 89072ecb..ad944cb3 100644
--- a/include/dwarfs/util.h
+++ b/include/dwarfs/util.h
@@ -33,6 +33,7 @@
 #include <exception>
 #include <filesystem>
 #include <iosfwd>
+#include <optional>
 #include <span>
 #include <string>
 #include <string_view>
@@ -85,4 +86,6 @@ void install_signal_handlers();
 
 std::tm safe_localtime(std::time_t t);
 
+std::optional<size_t> get_self_memory_usage();
+
 } // namespace dwarfs
diff --git a/include/dwarfs/writer/console_writer.h b/include/dwarfs/writer/console_writer.h
index 3e2d6468..275c9a88 100644
--- a/include/dwarfs/writer/console_writer.h
+++ b/include/dwarfs/writer/console_writer.h
@@ -44,6 +44,7 @@ class console_writer : public stream_logger {
  public:
   enum display_mode { NORMAL, REWRITE };
   enum progress_mode { NONE, SIMPLE, ASCII, UNICODE };
+  using mem_usage_fn = std::function<size_t()>;
 
   console_writer(std::shared_ptr<terminal const> term, std::ostream& os,
                  progress_mode pg_mode, display_mode mode = NORMAL,
@@ -51,6 +52,8 @@ class console_writer : public stream_logger {
 
   void update(writer_progress& prog, bool last);
 
+  void set_memory_usage_function(mem_usage_fn func) { mem_usage_ = func; }
+
  private:
   void preamble(std::ostream& os) override;
   void postamble(std::ostream& os) override;
@@ -63,6 +66,7 @@ class console_writer : public stream_logger {
   std::atomic<size_t> counter_{0};
   progress_mode const pg_mode_;
   display_mode const mode_;
+  mem_usage_fn mem_usage_;
 };
 
 } // namespace writer
diff --git a/src/util.cpp b/src/util.cpp
index 3a350b78..3b414087 100644
--- a/src/util.cpp
+++ b/src/util.cpp
@@ -32,6 +32,7 @@
 #include <clocale>
 #include <cstdio>
 #include <cstdlib>
+#include <fstream>
 #include <iostream>
 #include <mutex>
 #include <optional>
@@ -63,6 +64,10 @@
 #endif
 #endif
 
+#ifdef __APPLE__
+#include <mach/mach.h>
+#endif
+
 #include <dwarfs/conv.h>
 #include <dwarfs/error.h>
 #include <dwarfs/util.h>
@@ -584,4 +589,39 @@ std::tm safe_localtime(std::time_t t) {
   return buf;
 }
 
+std::optional<size_t> get_self_memory_usage() {
+#if defined(_WIN32)
+  // TODO
+#elif defined(__APPLE__)
+  task_vm_info info{};
+  mach_msg_type_number_t count = TASK_VM_INFO_COUNT;
+
+  if (task_info(mach_task_self(), TASK_VM_INFO,
+                reinterpret_cast<task_info_t>(&info), &count) == KERN_SUCCESS) {
+    return info.phys_footprint;
+  }
+#elif defined(__linux__)
+  static constexpr auto kSmapsPath{"/proc/self/smaps_rollup"};
+  std::ifstream smaps(kSmapsPath);
+
+  if (smaps) {
+    std::string line;
+
+    while (std::getline(smaps, line)) {
+      if (line.starts_with("Pss:")) {
+        std::string_view size_str(line);
+        size_str.remove_prefix(size_str.find_first_not_of(" \t", 4));
+        size_t size;
+        auto [endp, ec] = std::from_chars(
+            size_str.data(), size_str.data() + size_str.size(), size);
+        if (ec == std::errc() && std::string_view(endp).starts_with(" kB")) {
+          return size * 1024;
+        }
+      }
+    }
+  }
+#endif
+  return std::nullopt; // Not implemented
+}
+
 } // namespace dwarfs
diff --git a/src/writer/console_writer.cpp b/src/writer/console_writer.cpp
index 91aeaefa..12ad1747 100644
--- a/src/writer/console_writer.cpp
+++ b/src/writer/console_writer.cpp
@@ -285,7 +285,7 @@ void console_writer::update(writer_progress& prog, bool last) {
           << newline
 
           << "compressed filesystem: " << p.blocks_written << " blocks/"
-          << size_with_unit(p.compressed_size) << " written" << newline;
+          << size_with_unit(p.compressed_size) << " written";
       break;
 
     case REWRITE:
@@ -295,9 +295,15 @@ void console_writer::update(writer_progress& prog, bool last) {
 
           << "compressed filesystem: " << p.blocks_written << "/"
           << p.block_count << " blocks/" << size_with_unit(p.compressed_size)
-          << " written" << newline;
+          << " written";
       break;
     }
+
+    if (mem_usage_) {
+      oss << ", using " << size_with_unit(mem_usage_()) << " of RAM";
+    }
+
+    oss << newline;
   }
 
   if (pg_mode_ == NONE) {
diff --git a/tools/src/mkdwarfs_main.cpp b/tools/src/mkdwarfs_main.cpp
index 0fad9ca6..ff41d418 100644
--- a/tools/src/mkdwarfs_main.cpp
+++ b/tools/src/mkdwarfs_main.cpp
@@ -996,6 +996,11 @@ int mkdwarfs_main(int argc, sys_char** argv, iolayer const& iol) {
                                         : writer::console_writer::NORMAL,
                              logopts);
 
+  if (get_self_memory_usage()) {
+    lgr.set_memory_usage_function(
+        [] { return get_self_memory_usage().value_or(0); });
+  }
+
   std::unique_ptr<writer::rule_based_entry_filter> rule_filter;
 
   if (!filter.empty()) {