From 1277524172dca03e7ab7e190ec30e7c3e180376f Mon Sep 17 00:00:00 2001
From: Daniel Milroy <daniel.milroy@gmail.com>
Date: Fri, 21 Jun 2024 01:57:13 -0700
Subject: [PATCH] resource-query: add support for partial-cancel

---
 resource/utilities/command.cpp | 84 ++++++++++++++++++++++++++++++++++
 resource/utilities/command.hpp |  2 +
 2 files changed, 86 insertions(+)

diff --git a/resource/utilities/command.cpp b/resource/utilities/command.cpp
index f87236a06..6a73f3a33 100644
--- a/resource/utilities/command.cpp
+++ b/resource/utilities/command.cpp
@@ -50,6 +50,8 @@ command_t commands[] = {
 "resource-query> find status=down and sched-now=allocated" },
     { "cancel", "c", cmd_cancel, "Cancel an allocation or reservation: "
 "resource-query> cancel jobid" },
+    { "partial-cancel", "pc", cmd_partial_cancel, "Partially release an allocation: "
+"resource-query> partial-cancel jobid (file format: jgf | rv1exec) R_to_cancel.file" },
     { "set-property", "p", cmd_set_property, "Add a property to a resource: "
 "resource-query> set-property resource PROPERTY=VALUE" },
 { "get-property", "g", cmd_get_property, "Get all properties of a resource: "
@@ -84,6 +86,26 @@ static int do_remove (std::shared_ptr<resource_context_t> &ctx, int64_t jobid)
     return rc;
 }
 
+static int do_partial_remove (std::shared_ptr<resource_context_t> &ctx,
+                              std::shared_ptr<resource_reader_base_t> &reader,
+                              int64_t jobid, const std::string &R_cancel)
+{
+    int rc = -1;
+    bool full_cancel = false;
+
+    if ( (rc = ctx->traverser->remove (R_cancel, reader, (int64_t)jobid,
+                                       full_cancel)) == 0) {
+        if (full_cancel && (ctx->jobs.find (jobid) != ctx->jobs.end ())) {
+           std::shared_ptr<job_info_t> info = ctx->jobs[jobid];
+           info->state = job_lifecycle_t::CANCELED;
+        }
+    } else {
+        std::cout << ctx->traverser->err_message ();
+        ctx->traverser->clear_err_message ();
+    }
+    return rc;
+}
+
 static void print_sat_info (std::shared_ptr<resource_context_t> &ctx,
                             std::ostream &out, bool sat, double elapse,
                             unsigned int pre, unsigned int post)
@@ -600,6 +622,68 @@ int cmd_cancel (std::shared_ptr<resource_context_t> &ctx,
     return 0;
 }
 
+int cmd_partial_cancel (std::shared_ptr<resource_context_t> &ctx,
+                        std::vector<std::string> &args)
+{
+    int rc = -1;
+    std::stringstream buffer{};
+    std::string jobid_str = args[1];
+    std::string reader = args[2];
+    std::ifstream cancel_file (args[3]);
+    std::shared_ptr<resource_reader_base_t> rd;
+
+    if (args.size () != 4) {
+        std::cerr << "ERROR: malformed command" << std::endl;
+        return 0;
+    }
+
+    uint64_t jobid = (uint64_t)std::strtoll (jobid_str.c_str (), NULL, 10);
+
+    if (!(reader == "jgf" || reader == "rv1exec")) {
+        std::cerr << "ERROR: unsupported reader " << args[2] << std::endl;
+        goto done;
+    }
+
+    if (!cancel_file) {
+        std::cerr << "ERROR: can't open " << args[3] << std::endl;
+        goto done;
+    }
+    buffer << cancel_file.rdbuf ();
+    cancel_file.close ();
+
+    if (reader == "rv1exec") {
+        if ( (rd = create_resource_reader ("rv1exec")) == nullptr) {
+            std::cerr << "ERROR: can't create rv1exec reader " << std::endl;
+            goto done;
+        }
+    } else { // must be JGF
+        if ( (rd = create_resource_reader ("jgf")) == nullptr) {
+            std::cerr << "ERROR: can't create rv1exec reader " << std::endl;
+            goto done;
+        }
+    }
+
+    if (ctx->allocations.find (jobid) != ctx->allocations.end ()) {
+        if ( (rc = do_partial_remove (ctx, rd, jobid, buffer.str ())) == 0)
+            ctx->allocations.erase (jobid);
+    } else if (ctx->reservations.find (jobid) != ctx->reservations.end ()) {
+        std::cerr << "ERROR: reservations not currently supported by partial cancel"
+                  << std::endl;
+        goto done;
+    } else {
+        std::cerr << "ERROR: nonexistent job " << jobid << std::endl;
+        goto done;
+    }
+
+    if (rc != 0) {
+        std::cerr << "ERROR: error encountered while removing job "
+                  << jobid << std::endl;
+    }
+
+done:
+    return 0;
+}
+
 int cmd_set_property (std::shared_ptr<resource_context_t> &ctx,
                       std::vector<std::string> &args)
 {
diff --git a/resource/utilities/command.hpp b/resource/utilities/command.hpp
index 54e139122..a257e4b2d 100644
--- a/resource/utilities/command.hpp
+++ b/resource/utilities/command.hpp
@@ -81,6 +81,8 @@ int cmd_find (std::shared_ptr<resource_context_t> &ctx,
                 std::vector<std::string> &args);
 int cmd_cancel (std::shared_ptr<resource_context_t> &ctx,
                 std::vector<std::string> &args);
+int cmd_partial_cancel (std::shared_ptr<resource_context_t> &ctx,
+                std::vector<std::string> &args);
 int cmd_set_property (std::shared_ptr<resource_context_t> &ctx,
                       std::vector<std::string> &args);
 int cmd_get_property (std::shared_ptr<resource_context_t> &ctx,