state_exec: max_memory counts the deep-data work a process waits on

docs/PYCVC_STATE_EXEC_PARITY.md

@@ -118,7 +118,7 @@ are listed too. Status after PR 4:
 | | 38 `builtin_*` functions | none | exception: DSL builtins, called from DSL source |
 | | `percent_encode`, `percent_decode` | `pycvc.percent_encode`, `pycvc.percent_decode` | wrapped (PR 2) |
 | stackless_evaluator.h, evaluator.h | `stackless_evaluator`, `evaluator`, `evaluator_state`, `eval_frame`, `eval_phase`, `evaluation_stats`, `step_kind` (#658), `evaluation_interrupted`, `evaluation_timeout` | the same names under `pycvc.` | wrapped (PR 3). `eval_phase` and `step_kind` are `IntEnum`s, and `evaluation_stats` has #658's counts by kind (`steps`, `step_time`, `data_nodes`, ...); the two exceptions are `RuntimeError` subclasses. `on_complete` is a Python callable. `evaluator.user_macros()` returns a copy, a dict of name -> (params, body). `restrict_special_forms` takes any iterable of str, or None |
-| | `evaluator_state.data_work`, `evaluator_state.backlog`, `pending_data_work` (#658) | none | exception: deep-data work in flight between steps, internal to the evaluator; `data_step_nodes` is wrapped |
+| | `evaluator_state.data_work`, `evaluator_state.backlog`, `pending_data_work` (#658), `waits_on_data_work` | none | exception: deep-data work in flight between steps, internal to the evaluator; `data_step_nodes` is wrapped |
 | | `return_exception` | none | exception: the internal unwinding of `(return ...)`, never seen outside an evaluation |
 | async_evaluator.h, async_stackless_evaluator.h | the async evaluators | none | exception: coroutine API |
 | resource_policy.h | `resource_policy`, `process_limits`, `validate_limits` | `pycvc.resource_policy` (its `mode` is `pycvc.resource_policy.mode`, an `IntEnum`), `pycvc.process_limits`, `pycvc.validate_limits` | wrapped (PR 3) |

docs/STATE_EXEC_DEVELOPER_GUIDE.md

@@ -884,10 +884,16 @@ these holds:
 
 - natives have reported allocating 64 KiB, or half the last measurement if that is more, since
   the last one (the functions with an allocation check, §11, plus copies like `cons` and
-  `slice`);
+  `slice`, and the jobs of calls left waiting on the evaluator);
 - what they reported could cross the remaining budget;
 - `limits.heap_check_steps` steps have passed (default 4096; `0` = only the two above).
 
+A measurement can start while the process waits on other deep-data work, such as a call the
+evaluator makes for a native function: it runs first, and the work goes on after it. It counts
+that work too, with what the work's jobs hold (a waiting call's function and arguments, the
+values `collect` has gathered). So a process that never stops waiting, like a class whose `init`
+constructs again, is still held to `max_memory`.
+
 The kill decision always uses an actual measurement, never the reported bytes, so temporaries a
 program drops don't count. Only a process with `max_memory > 0` is measured.
 `process_info::heap_memory` and `heap_measurements` report the last measurement. The per-step

docs/STATE_EXEC_GRAMMAR.md

@@ -293,7 +293,8 @@ the evaluator's: it makes the call at its own level, in a data step (`call_funct
   `init`), hands its place to the call it makes, so a chain of them runs in constant space. A
   call with more to do after its callee (a constructor returns the instance once `init` is done)
   waits, a link of pending work per level, like a function that calls itself before doing
-  anything else. Pending work is not yet part of a process's measured heap (`max_memory`).
+  anything else. Pending work and what its jobs hold count toward the process's measured heap,
+  so `max_memory` stops such a chain too.
 - **Cost.** A call the evaluator makes costs about 300 ns more than one made directly (MSVC
   Release, `state_exec_bench_test` `NativeCallThroughput`). The one direct level keeps `apply`
   and `send` at their old cost, and only a native call two levels down pays it. Constructing a

inc/cvc/state/state_exec/data_job.h

@@ -78,6 +78,17 @@ class data_job {
   /// work take the job's place, so a chain of such calls runs in constant space.
   virtual bool forwards_call() const { return false; }
 
+  /// For measuring the heap of a process whose work waits on the job (heap.h), which happens
+  /// between its slices: `visit` is called with each value the job owns a copy of (what it was
+  /// given, what it has built so far). A job that only borrows its caller's arguments, or whose
+  /// results are held to the allocation limit as they grow, need not report them.
+  virtual void held_values(const std::function<void(const value_t &)> &visit) const { (void)visit; }
+
+  /// The heap bytes the job takes itself, its object and its own buffers, besides what the
+  /// values it reports hold. A job left to run after its call reports these as allocated
+  /// (eval_alloc_note), so a heap measurement comes due as such jobs pile up.
+  virtual std::size_t held_bytes() const { return 0; }
+
 protected:
   std::size_t nodes_ = 0;
   std::optional<call_request> call_;

inc/cvc/state/state_exec/heap.h

@@ -64,8 +64,9 @@ std::size_t heap_remeasure_bytes(std::size_t last_measured);
 /// The heap accounting a scheduler does after each step of `proc`: take a measurement that just
 /// finished (proc.heap_bytes), add `allocated` (the step guard's eval_alloc_guard::allocated()),
 /// and start the next measurement as the process's data work when one is due
-/// (process::heap_measure_due). `state_bytes` is the memory tracker's count for the process.
-/// Nothing is measured for a process without a memory budget.
+/// (process::heap_measure_due), on top of any deep-data work in flight. The measurement counts
+/// that work too, with what its jobs hold (data_job::held_values). `state_bytes` is the memory
+/// tracker's count for the process. Nothing is measured for a process without a memory budget.
 void heap_after_step(process &proc, std::size_t allocated, uint64_t state_bytes);
 
 // ---------------------------------------------------------------------------

inc/cvc/state/state_exec/stackless_evaluator.h

@@ -127,6 +127,10 @@ struct evaluator_state {
   std::optional<uint64_t> heap_measured;
 };
 
+/// Whether a frame of `st` waits on deep-data work: any work but a heap measurement, which waits
+/// on no frame. A process cannot fork while one does (the child could not finish the call).
+bool waits_on_data_work(const evaluator_state &st);
+
 /// Synchronous stackless evaluator (Tier 2 — local contexts only).
 ///
 /// Uses an explicit stack of eval_frame objects instead of the C++ call stack.

src/cvc/state/state_exec/async_scheduler.cpp

@@ -739,9 +739,8 @@ int async_scheduler::fork(int pid) {
     return -1;
   // Deep-data work in flight (a function collections.map calls, say) is not copied: a child
   // could not finish the call its stack waits on. A heap measurement waits on no frame, so a
-  // child without it is whole.
-  if (parent.state.data_work &&
-      parent.state.data_work->what != pending_data_work::kind::heap_measure)
+  // child without it is whole, unless the measurement runs on top of such work.
+  if (waits_on_data_work(parent.state))
     return -1;
 
   int child_pid = next_pid_++;

src/cvc/state/state_exec/data_job.cpp

@@ -503,6 +503,17 @@ class call_walk final : public data_job {
 
   bool forwards_call() const override { return true; }
 
+  void held_values(const std::function<void(const value_t &)> &visit) const override {
+    visit(fn_);
+    for (const value_t &a : args_)
+      visit(a);
+    visit(result_);
+  }
+
+  std::size_t held_bytes() const override {
+    return sizeof(*this) + args_.capacity() * sizeof(value_t);
+  }
+
 private:
   value_t fn_;
   std::vector<value_t> args_;
@@ -520,6 +531,16 @@ class generator_walk final : public data_job {
 
   value_t result() override { return make_list(std::move(values_)); }
 
+  // The generator itself is the waiting call's argument, which its frame holds.
+  void held_values(const std::function<void(const value_t &)> &visit) const override {
+    for (const value_t &v : values_)
+      visit(v);
+  }
+
+  std::size_t held_bytes() const override {
+    return sizeof(*this) + values_.capacity() * sizeof(value_t);
+  }
+
 private:
   // Keep a value it yielded. True when the walk goes on (collect).
   bool keep(value_t v) {
@@ -760,6 +781,7 @@ value_t run_data_job(std::unique_ptr<data_job> job, std::function<value_t(value_
     job->call_returned(std::move(value));
   }
   job->own_inputs();
+  eval_alloc_note(job->held_bytes()); // it outlives the call (heap.h)
   scope->deferred_ = std::make_unique<detail::deferred_data_job>(
       detail::deferred_data_job{std::move(job), std::move(then), std::move(waits_on)});
   return nil_value;

src/cvc/state/state_exec/heap.cpp

@@ -266,6 +266,18 @@ class heap_share_job final : public data_job {
     // The roots, whole: the references the process holds directly.
     auto root = [this](node_kind k, const auto &sp) { add(k, sp, /*from_root=*/true); };
     process_edges(p, root);
+    // And the deep-data work it waits on, with what that work's jobs hold: a call that waits on
+    // a call like it (a constructor whose init constructs again) piles them up.
+    for (const pending_data_work *w = p.state.data_work.get(); w; w = w->outer.get()) {
+      own_ += static_cast<double>(sizeof(pending_data_work) + kControlBlock);
+      if (!w->job)
+        continue;
+      own_ += static_cast<double>(w->job->held_bytes());
+      w->job->held_values([&](const value_t &v) {
+        own_ += static_cast<double>(inline_heap(v));
+        value_edges(v, root);
+      });
+    }
   }
 
   bool step(std::size_t budget) override {
@@ -517,21 +529,26 @@ void heap_after_step(process &proc, std::size_t allocated, uint64_t state_bytes_
   proc.heap_alloc_bytes += allocated;
   ++proc.heap_steps;
   if (proc.heap_measuring) {
-    if (!st.data_work)
-      proc.heap_measuring = false; // dropped with the state it ran in (a signal, a restore)
+    // Until it finishes, a measurement is the innermost work: if not, it was dropped with the
+    // state it ran in (a signal, a restore).
+    if (!st.data_work || st.data_work->what != pending_data_work::kind::heap_measure)
+      proc.heap_measuring = false;
     return;
   }
-  if (st.done || st.reading || st.data_work || st.stack.empty())
+  if (st.done || st.reading || st.stack.empty())
     return;
   if (!proc.heap_measure_due(state_bytes_used))
     return;
   // It waits on no frame: evaluation resumes at the top of the stack, as it was, so `base` and
   // `frame_args` are the whole stack and the top frame's values (a snapshot keeps all of them).
+  // Started in the middle of other deep-data work (a call in flight, which a process can wait on
+  // for as long as its budget lasts), it runs first, and that work goes on after it.
   auto work = std::make_shared<pending_data_work>();
   work->job = make_heap_share_job(proc);
   work->what = pending_data_work::kind::heap_measure;
   work->base = st.stack.size();
   work->frame_args = st.stack.back().results.size();
+  work->outer = std::move(st.data_work);
   st.data_work = std::move(work);
   proc.heap_measuring = true;
 }

src/cvc/state/state_exec/scheduler.cpp

@@ -753,9 +753,8 @@ int scheduler::fork(int pid) {
     return -1;
   // Deep-data work in flight (a function collections.map calls, say) is not copied: a child
   // could not finish the call its stack waits on. A heap measurement waits on no frame, so a
-  // child without it is whole.
-  if (parent.state.data_work &&
-      parent.state.data_work->what != pending_data_work::kind::heap_measure)
+  // child without it is whole, unless the measurement runs on top of such work.
+  if (waits_on_data_work(parent.state))
     return -1;
 
   int child_pid = next_pid_++;

src/cvc/state/state_exec/stackless_evaluator.cpp

@@ -18,6 +18,13 @@ pending_data_work::~pending_data_work() {
     next = std::move(next->outer);
 }
 
+bool waits_on_data_work(const evaluator_state &st) {
+  for (const pending_data_work *w = st.data_work.get(); w; w = w->outer.get())
+    if (w->what != pending_data_work::kind::heap_measure)
+      return true;
+  return false;
+}
+
 // ---------------------------------------------------------------------------
 // ctor
 // ---------------------------------------------------------------------------

src/cvc/tests/state_exec_heap_test.cpp

@@ -454,3 +454,92 @@ TEST(HeapBudgetTest, AProcessForksDuringAMeasurement) {
   EXPECT_GT(sched.get_process_info(pid)->heap_measurements, 0u);
   EXPECT_EQ(std::get<int64_t>(sched.get_result(pid)->v), 400);
 }
+
+// A measurement due in the middle of a call in flight (here collections.map calling a DSL
+// function) runs on top of it, waiting on no frame, and the call goes on after it. The process
+// cannot fork meanwhile: a frame still waits on the call.
+TEST(HeapBudgetTest, AMeasurementRunsOnTopOfACallInFlight) {
+  cvc::app app;
+  auto env = builtins::make_default_environment();
+  env->set("keep", value_t{int_list(200)});
+  stackless_evaluator ev(env);
+  process p;
+  p.max_memory = std::size_t{64} << 20;
+  p.state = ev.create_state(parse("(length (collections.map keep (lambda (x) (+ x 1))))"));
+  p.state.global_env = env;
+  p.state.data_step_nodes = 8;
+  p.state.stats.start();
+  while (!(p.state.data_work && p.state.data_work->calling))
+    ASSERT_FALSE(ev.step(p.state));
+  const std::shared_ptr<pending_data_work> call = p.state.data_work;
+  const std::size_t frames = p.state.stack.size();
+  heap_after_step(p, std::size_t{1} << 20, 0); // allocations call for a measurement
+  ASSERT_TRUE(p.state.data_work);
+  ASSERT_EQ(p.state.data_work->what, pending_data_work::kind::heap_measure);
+  EXPECT_EQ(p.state.data_work->outer, call);
+  EXPECT_EQ(p.state.data_work->base, frames);
+  EXPECT_TRUE(waits_on_data_work(p.state));
+  while (!p.state.heap_measured) {
+    ASSERT_FALSE(ev.step(p.state));
+    ASSERT_EQ(p.state.stack.size(), frames); // nothing else runs meanwhile
+  }
+  EXPECT_EQ(p.state.data_work, call);
+  heap_after_step(p, 0, 0);
+  EXPECT_GT(p.heap_bytes, 0u);
+  EXPECT_FALSE(p.heap_measuring);
+  while (!ev.step(p.state)) {
+  }
+  EXPECT_EQ(std::get<int64_t>(p.state.result.v), 200);
+}
+
+// --- work that never finishes is still held to max_memory ------------------------------------
+//
+// A call the evaluator makes can wait on another like it for as long as the budget lasts: a class
+// whose init is its own constructor, or a closure init that constructs again. A measurement counts
+// that work, with what its jobs hold (each waiting call's function and arguments), so max_memory
+// stops the process; max_steps is only a backstop here.
+
+const char *const kSelfInit = "(begin (defclass P) (set-attr P.__cls_dict__ \"init\" P) (P))";
+const char *const kInitConstructs = "(begin (defclass P (init (self) (P))) (P))";
+
+TEST(HeapBudgetTest, TheWorkAProcessWaitsOnCounts) {
+  // After a hundred steps of kSelfInit, ~180 constructions wait, the k-th on a call with k
+  // arguments: some 16000 values held by jobs alone.
+  stackless_evaluator ev(builtins::make_default_environment());
+  process p;
+  p.state = ev.create_state(parse(kSelfInit));
+  for (int i = 0; i < 100; ++i)
+    ASSERT_FALSE(ev.step(p.state));
+  EXPECT_GT(measure_heap_share(p), 10000 * sizeof(value_t));
+}
+
+TEST(HeapBudgetTest, ACallChainThatNeverEndsRunsOutOfMemory) {
+  // kSelfInit's arguments grow quadratically: 2 MiB within a few hundred steps.
+  for (const char *script : {kSelfInit, kInitConstructs}) {
+    SCOPED_TRACE(script);
+    const uint64_t backstop = script == kSelfInit ? 1000 : 1000000;
+    {
+      async_scheduler sched;
+      execute_options opts;
+      opts.max_memory = std::size_t{2} << 20;
+      opts.max_steps = backstop;
+      const int pid = sched.execute(std::string(script), opts);
+      sched.sync_run();
+      EXPECT_EQ(sched.get_exit_error(pid).value_or(""), "memory_limit_exceeded");
+      EXPECT_GT(sched.get_process_info(pid)->heap_measurements, 0u);
+    }
+    {
+      scheduler sched;
+      execute_options opts;
+      opts.max_memory = std::size_t{2} << 20;
+      opts.max_steps = backstop;
+      const int pid = sched.execute(std::string(script), opts);
+      sched.run();
+      const auto info = sched.get_process_info(pid);
+      ASSERT_TRUE(info);
+      EXPECT_EQ(info->status, process_status::killed);
+      EXPECT_LT(info->step_count, backstop); // memory, not the step budget
+      EXPECT_GE(info->heap_memory, opts.max_memory);
+    }
+  }
+}