planner_multi interface: add span partial cancel

resource/planner/c/planner_multi.h

@@ -12,6 +12,7 @@
 #define PLANNER_MULTI_H
 
 #include "planner.h"
+#include <unordered_map>
 
 #ifdef __cplusplus
 extern "C" {
@@ -278,6 +279,35 @@ int64_t planner_multi_add_span (planner_multi_t *ctx, int64_t start_time,
  */
 int planner_multi_rem_span (planner_multi_t *ctx, int64_t span_id);
 
+/*! Reduce the existing span's resources from the planner.
+ *  This function will be called for a partial release/cancel.
+ *  If the number of resources to be removed is equal to those 
+ *  allocated to the span, completely remove the span.
+ *
+ *  \param ctx          opaque multi-planner context returned
+ *                      from planner_multi_new.
+ *  \param span_id      span_id returned from planner_add_span.
+ *  \param reduced_totals
+ *                      64-bit unsigned integer array of size len where each
+ *                      element contains the total count of available resources
+ *                      of a single resource type.
+ *  \param resource_types
+ *                      string array of size len where each element contains
+ *                      the resource type corresponding to each corresponding
+ *                      element in the resource_totals array.
+ *  \param len          length of the resource_totals and resource_types arrays.
+ *  \param removed      bool indicating if the entire span was removed.
+ *  \return             0 on success; -1 on error with errno set as follows:
+ *                          EINVAL: invalid argument.
+ *                          EKEYREJECTED: span could not be removed from
+ *                                        the planner's internal data structures.
+ *                          ERANGE: a resource state became out of a valid range.
+ */
+int planner_multi_reduce_span (planner_multi_t *ctx, int64_t span_id,
+                               const uint64_t *reduced_totals,
+                               const char **resource_types, size_t len,
+                               bool &removed);
+
 //! Span iterators -- there is no specific iteration order
 //  return -1 when you no longer can iterate: EINVAL when ctx is NULL.
 //  ENOENT when you reached the end of the spans

resource/planner/c/planner_multi_c_interface.cpp

@@ -15,6 +15,7 @@
 #include <vector>
 #include <map>
 #include <unordered_set>
+#include <numeric>
 
 #include "planner_multi.h"
 #include "resource/planner/c++/planner_multi.hpp"
@@ -414,6 +415,122 @@ extern "C" int planner_multi_rem_span (planner_multi_t *ctx, int64_t span_id)
     return rc;
 }
 
+extern "C" int planner_multi_reduce_span (planner_multi_t *ctx,
+                                          int64_t span_id,
+                                          const uint64_t *reduced_totals,
+                                          const char **resource_types,
+                                          size_t len,
+                                          bool &removed)
+{
+    size_t i;
+    int rc = -1;
+    bool tmp_removed;
+    size_t mspan_idx;
+    int64_t mspan_sum = 0;
+    std::unordered_set<std::string> ext_res_types;
+
+    if (!ctx || span_id < 0 || !reduced_totals || !resource_types
+                    || len > ctx->plan_multi->get_planners_size ()) {
+        errno = EINVAL;
+        return -1;
+    }
+    for (i = 0; i < len; ++i) {
+        if (reduced_totals[i] >
+            static_cast<uint64_t> (std::numeric_limits<int64_t>::max ())) {
+            errno = ERANGE;
+            return -1;
+        }
+    }
+    auto span_it = ctx->plan_multi->get_span_lookup ().find (span_id);
+    if (span_it == ctx->plan_multi->get_span_lookup ().end ()) {
+        errno = ENOENT;
+        goto error;
+    }
+
+    for (i = 0; i < len; ++i) {
+        // Index could be different than the span_lookup due to order of
+        // iteration in the reader differing from the graph initialization
+        // order.
+        mspan_idx = ctx->plan_multi->get_resource_type_idx (resource_types[i]);
+        tmp_removed = false;
+        if ( (rc = planner_reduce_span (ctx->plan_multi->get_planner_at (
+                                                resource_types[i]),
+                                                span_it->second.at (mspan_idx),
+                                                reduced_totals[i],
+                                                tmp_removed)) == -1) {
+            // Could return -1 if the span with 0 resource request had been removed
+            // by a previous cancellation, so need to check if the span exists.
+            if (ctx->plan_multi->get_planner_at (
+                resource_types[i])->plan->get_span_lookup ().find (
+                                                span_it->second.at (mspan_idx))
+                != ctx->plan_multi->get_planner_at (
+                    resource_types[i])->plan->get_span_lookup ().end ()) {
+                // We know the span is valid, so planner_reduce_span
+                // encountered another error.
+                errno = EINVAL;
+                goto error;
+            }
+        }
+        ext_res_types.insert (resource_types[i]);
+        // Enter invalid span ID in the span_lookup to indicate the resource
+        // removal.
+        if (tmp_removed)
+            span_it->second[mspan_idx] = -1;
+    }
+    // Iterate over planner_multi resources since resource_types may not cover
+    // all planner_multi resources. If resource_types contains fewer types
+    // than the total planner_multi resources, this means the reader partial
+    // cancel didn't encounter those resource types. This can happen since
+    // agfilter requests for 0 resources are entered for resource types
+    // tracked by the agfilter that the job didn't request. Ex: job requests
+    // cores, but the agfilter tracks cores and memory. A span will be created
+    // for cores and memory, but the memory request will be 0. We need to
+    // remove these spans.
+    for (i = 0; i < ctx->plan_multi->get_planners_size (); ++i) {
+        const char *tmp_type = ctx->plan_multi->get_resource_type_at (i);
+        tmp_removed = false;
+        // Check if the resource type was already processed in the previous
+        // for loop.
+        if (ext_res_types.find (tmp_type) == ext_res_types.end ()) {
+            if ( (rc = planner_reduce_span (ctx->plan_multi->get_planner_at (
+                                                    tmp_type),
+                                                    span_it->second.at (i),
+                                                    0,
+                                                    tmp_removed)) == -1) {
+                // Could return -1 if the span with 0 resource request had been
+                // removed by a previous cancellation, so need to check if the
+                // span exists.
+                if (ctx->plan_multi->get_planner_at (
+                    tmp_type)->plan->get_span_lookup ().find (mspan_idx)
+                    != ctx->plan_multi->get_planner_at (
+                        tmp_type)->plan->get_span_lookup ().end ()) {
+                    // We know the span is valid, so planner_reduce_span
+                    // encountered another error.
+                    errno = EINVAL;
+                    goto error;
+                }
+            }
+            // Enter invalid span ID in the span_lookup to indicate the
+            // resource removal.
+            if (tmp_removed)
+                span_it->second[i] = -1;
+        }
+    }
+
+    mspan_sum = std::accumulate (span_it->second.begin (),
+                                 span_it->second.end (),
+                                 0, std::plus<int64_t> ());
+    // Delete if all entries are -1
+    if (mspan_sum == (-1 * span_it->second.size ())) {
+        ctx->plan_multi->get_span_lookup ().erase (span_it);
+        removed = true;
+    }
+
+    rc  = 0;
+error:
+    return rc;
+}
+
 int64_t planner_multi_span_first (planner_multi_t *ctx)
 {
     int64_t rc = -1;