From c3647caacceb69fb10bc344a16761c304ed0f4e0 Mon Sep 17 00:00:00 2001
From: chn <chn@chn.moe>
Date: Tue, 4 Jun 2024 21:29:28 +0800
Subject: [PATCH] =?UTF-8?q?localPackages.hpcstat:=20=E7=BB=9F=E8=AE=A1?=
 =?UTF-8?q?=E7=A3=81=E7=9B=98=E4=BD=BF=E7=94=A8=E6=83=85=E5=86=B5=E6=97=B6?=
 =?UTF-8?q?=E5=8F=AA=E6=89=AB=E6=8F=8F=E4=B8=80=E6=AC=A1?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 local/pkgs/hpcstat/src/disk.cpp | 59 ++++++++++++++++++++-------------
 1 file changed, 36 insertions(+), 23 deletions(-)

diff --git a/local/pkgs/hpcstat/src/disk.cpp b/local/pkgs/hpcstat/src/disk.cpp
index 437c3f35..9b787b63 100644
--- a/local/pkgs/hpcstat/src/disk.cpp
+++ b/local/pkgs/hpcstat/src/disk.cpp
@@ -25,39 +25,52 @@ bool hpcstat::disk::stat(boost::interprocess::file_lock &lock)
     { std::cerr << "HOME not set\n"; return false; }
   else
   {
-    auto get_size  = [](std::string path) -> std::optional<double>
+    auto all_size  = [&] -> std::optional<std::map<std::string, double>>
     {
-      if (auto result = exec("/usr/bin/du", { "-s", path }); !result)
-        { std::cerr << fmt::format("failed to stat {}\n", path); return std::nullopt; }
+      if
+      (
+        auto result =
+          exec("/usr/bin/du", { "-a", "--max-depth=2", *homedir });
+        !result
+      )
+        { std::cerr << fmt::format("failed to stat home.\n"); return std::nullopt; }
       else
       {
-        std::smatch match;
-        if (!std::regex_search(*result, match, std::regex(R"((\d+))")))
-          { std::cerr << fmt::format("failed to parse {}\n", *result); return std::nullopt; }
-        return std::stod(match[1]) / 1024 / 1024;
+        std::map<std::string, double> size;
+        std::regex re(R"((\d+)\s+(.*)\n)");
+        for (auto i = std::sregex_iterator(result->begin(), result->end(), re);
+          i != std::sregex_iterator(); ++i)
+        {
+          if (i->str(2).find(*homedir) != 0)
+            { std::cerr << fmt::format("invalid path: {}\n", i->str(2)); return std::nullopt; }
+          else size[i->str(2).substr(homedir->size())] = std::stoul(i->str(1)) / 1024. / 1024;
+        }
+        return size;
       }
-    };
-    auto get_subdir = [](std::string path) -> std::vector<std::string>
-    {
-      std::filesystem::directory_iterator it(path);
-      std::vector<std::string> result;
-      for (const auto& entry : it)
-        if (entry.is_directory()) result.push_back(entry.path().filename().string());
-      return result;
-    };
+    }();
+    if (!all_size) return false;
     Usage usage;
     usage.Time = now();
-    if (auto size = get_size(*homedir); size) usage.Total = *size; else return false;
+    if (!all_size->contains("")) { std::cerr << "Failed to get size of home\n"; return false; }
+    usage.Total = (*all_size)[""];
     for (const auto& [dir, recursive] : Directories)
     {
-      if (auto size = get_size(*homedir + "/" + dir); size)
-        usage.Teacher.push_back({ dir, *size });
-      else return false;
+      if (!all_size->contains(dir))
+        { std::cerr << fmt::format("Failed to get size of {}\n", dir); return false; }
+      else usage.Teacher.push_back({ dir, (*all_size)[dir] });
+      auto get_subdir = [](std::string path) -> std::vector<std::string>
+      {
+        std::filesystem::directory_iterator it(path);
+        std::vector<std::string> result;
+        for (const auto& entry : it)
+          if (entry.is_directory()) result.push_back(entry.path().filename().string());
+        return result;
+      };
       if (recursive) for (const auto& subdir : get_subdir(*homedir + "/" + dir))
       {
-        if (auto size = get_size(*homedir + "/" + dir + "/" + subdir); size)
-          usage.Student.push_back({ dir + "/" + subdir, *size });
-        else return false;
+        if (!all_size->contains(dir + "/" + subdir))
+          { std::cerr << fmt::format("Failed to get size of {}/{}\n", dir, subdir); return false; }
+        else usage.Student.push_back({ dir + "/" + subdir, (*all_size)[dir + "/" + subdir] });
       }
     }
     std::sort(usage.Teacher.begin(), usage.Teacher.end(),