From 550250fa833e6f5378e11489f78e9a16a760b61a Mon Sep 17 00:00:00 2001
From: chn <chn@chn.moe>
Date: Sun, 2 Feb 2025 23:20:14 +0800
Subject: [PATCH] packages.info: finish

---
 packages/info/src/main.cpp | 58 ++++++++++++++++++++++++++------------
 1 file changed, 40 insertions(+), 18 deletions(-)

diff --git a/packages/info/src/main.cpp b/packages/info/src/main.cpp
index 98dc631b..d118937d 100644
--- a/packages/info/src/main.cpp
+++ b/packages/info/src/main.cpp
@@ -7,36 +7,41 @@ extern "C"
 # include <slurm/spank.h>
 # include <slurm/slurm.h>
 // ac argv: configuration count and flags in plugstack.conf
-[[gnu::visibility("default")]] int slurm_spank_exit(spank_t spank, int ac, char** argv);
+[[gnu::visibility("default")]] int slurm_spank_job_epilog(spank_t spank, int ac, char** argv);
 }
 
 struct switch_user
 {
+  std::uint32_t uid, gid;
+  switch_user(std::uint32_t uid, std::uint32_t gid) : uid(uid), gid(gid) {}
   boost::system::error_code on_exec_setup(auto&&...)
   {
     // first set gid then set uid, otherwise failed
-    if (setegid(1000) != 0 || seteuid(1000) != 0)
+    if (setegid(gid) != 0 || seteuid(uid) != 0)
       return boost::system::error_code{errno, boost::system::system_category()};
     else return {};
-    return {};
   }
 };
 
-int slurm_spank_exit(spank_t spank, int ac, char** argv)
+int slurm_spank_job_epilog(spank_t spank, int ac, char** argv)
 {
   using namespace biu::literals;
-  if (spank_context() == S_CTX_REMOTE)
+  auto [info, outfile, uid, gid] = [&]
   {
     std::stringstream ss;
+    std::optional<std::string> outfile;
     ss << "------------------------------------------------------------\n";
-    std::uint32_t jid;
+    std::uint32_t jid, uid = -1, gid = -1;
     auto result = spank_get_item(spank, S_JOB_ID, &jid);
     if (result != ESPANK_SUCCESS) ss << "error getting job id: {}\n"_f(int(result));
     else
     {
       ss << "info for job {}:\n"_f(jid);
+      YAML::Node info;
+
+      // gather info from slurmctld
       job_info_msg_t* job_info;
-      // slurm_init(nullptr);
+      slurm_init(nullptr);
       auto result = slurm_load_job(&job_info, jid, 0);
       if (result != SLURM_SUCCESS) ss << "error loading job info: {}\n"_f(slurm_strerror(result));
       else if (job_info->record_count != 1) ss << "record_count {} != 1\n"_f(job_info->record_count);
@@ -47,11 +52,13 @@ int slurm_spank_exit(spank_t spank, int ac, char** argv)
           { return "{:%Y-%m-%d %H:%M:%S}"_f(*std::localtime(&time)); };
         auto timespan = [](time_t time)
           { return "{:%H:%M:%S}"_f(std::chrono::seconds(time)); };
-        YAML::Node info;
+        auto get_status = [](int code)
+          { return std::vector{ "{}"_f(job_states(code & 0xff)), "{:#x}"_f(code) }; };
         info["Job Id"] = job_info->job_array->job_id;
         info["Job Name"] = null_to_empty(job_info->job_array->name);
         info["User Id"] = job_info->job_array->user_id;
         info["Work Directory"] = null_to_empty(job_info->job_array->work_dir);
+        info["Output File"] = null_to_empty(job_info->job_array->std_out);
         info["Partition"] = null_to_empty(job_info->job_array->partition);
         info["Submit Time"] = timepoint(job_info->job_array->submit_time);
         info["Start Time"] = timepoint(job_info->job_array->start_time);
@@ -59,23 +66,38 @@ int slurm_spank_exit(spank_t spank, int ac, char** argv)
         info["Nodes"] = null_to_empty(job_info->job_array->nodes);
         info["TREs Allocated"] = null_to_empty(job_info->job_array->tres_alloc_str);
         info["GREs Allocated"] = null_to_empty(job_info->job_array->gres_total);
-        info["Status"] = job_info->job_array->job_state;
-        ss << "------------------------------------------------------------\n" << info;
+        info["Status"] = get_status(job_info->job_array->job_state);
+        if (job_info->job_array->std_out != nullptr) outfile = job_info->job_array->std_out;
+        uid = job_info->job_array->user_id;
+        gid = job_info->job_array->group_id;
       }
       slurm_free_job_info_msg(job_info);
+      slurm_fini();
 
+      ss << "------------------------------------------------------------\n" << info << '\n';
+    }
+    return std::tuple(ss.str(), outfile, uid, gid);
+  }();
+  slurm_spank_log("%s", info.c_str());
+  if (outfile)
+  {
+    try
+    {
       boost::asio::io_context context;
       boost::system::error_code ec;
-      boost::asio::readable_pipe rp{context};
-      boost::process::v2::process proc(context, "/run/current-system/sw/bin/capsh", { "--print" }, boost::process::v2::process_stdio{nullptr, rp, nullptr}, switch_user{});
-
-      std::string output;
-      boost::asio::read(rp, boost::asio::dynamic_buffer(output), ec);
-      if (ec != boost::asio::error::eof) ss << "error reading whoami: {}\n"_f(ec.message());
-      ss << "\nusername: {}"_f(output);
+      boost::asio::writable_pipe wp{context};
+      boost::process::v2::process proc
+      (
+        context, "/run/current-system/sw/bin/tee", { "-a", outfile->c_str() },
+        boost::process::v2::process_stdio{wp, nullptr, nullptr}, switch_user(uid, gid)
+      );
+      boost::asio::write(wp, boost::asio::buffer(info));
+      wp.close();
       proc.wait();
     }
-    slurm_spank_log("%s", ss.str().c_str());
+    catch (boost::system::system_error& e) { slurm_spank_log("boost error writing to output file: %s", e.what()); }
+    catch (std::exception& e) { slurm_spank_log("error writing to output file: %s", e.what()); }
+    catch (...) { slurm_spank_log("error writing to output file"); }
   }
   return 0;
 }