Added resource mocking pass to simplify termination of tasks

.gitignore

@@ -2,6 +2,7 @@ __pycache__
 build
 venv
 .vscode
+.idea
 .devcontainer/**
 .cache/**
 third_party/**

codegen/resmockpass.cpp

@@ -0,0 +1,142 @@
+#include <llvm/Demangle/Demangle.h>
+#include <llvm/IR/Verifier.h>
+
+#include <functional>
+#include <map>
+#include <utility>
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Passes/PassPlugin.h"
+
+using namespace llvm;
+
+// ignore handling by exceptions memory allocation failure. This is very
+// unlikely case and shouldn't be in tested program.
+// Supporting this case would require -fexcept
+constexpr std::string_view ltest_allocation_fun_name = "LtestMemAlloc";
+constexpr std::string_view ltest_deallocation_fun_name = "LtestMemDealloc";
+
+SmallVector<Value*> collectArgs(CallBase& call) {
+  SmallVector<Value*> res_args;
+  for (auto& arg : call.args()) {
+    res_args.push_back(arg);
+  }
+  return res_args;
+}
+
+class ResMockInserter {
+ public:
+  explicit ResMockInserter(Module& m) : m(m) {
+    auto& context = m.getContext();
+    auto ptr_type = PointerType::get(context, 0);
+
+    ltest_allocation_fun = m.getOrInsertFunction(
+        ltest_allocation_fun_name,
+        FunctionType::get(ptr_type, {IntegerType::get(context, 64)}, false));
+    ltest_deallocation_fun = m.getOrInsertFunction(
+        ltest_deallocation_fun_name,
+        FunctionType::get(Type::getVoidTy(context), {ptr_type}, false));
+  }
+  void Run() {
+    IRBuilder<> builder(m.getContext());
+    for (auto& f : m) {
+      // dirty hack, but otherwise boost contextes allocation is alos mocked,
+      // which produces warning by asan i couldn't debug double free here
+      std::string demangled_parent = demangle(f.getName());
+      if (demangled_parent.find("boost::context") != std::string::npos) {
+        continue;
+      }
+      for (auto& bb : f) {
+        for (auto& in : bb) {
+          if (!isa<CallBase>(&in)) {
+            continue;
+          }
+          CallBase& call = *dyn_cast<CallBase>(&in);
+          auto called_fun = call.getCalledFunction();
+          if (!called_fun) {
+            continue;
+          }
+          std::string demangled = demangle(called_fun->getName());
+          auto it = replaces.find(demangled);
+          if (it == replaces.end()) {
+            continue;
+          }
+          auto [rep_callee, rep_args] = (it->second)(call);
+          builder.SetInsertPoint(&in);
+          // todo - need we look at invoke?
+          auto rep = builder.CreateCall(rep_callee, rep_args);
+          call.replaceAllUsesWith(rep);
+          to_delete.push_back(&call);
+        }
+      }
+    }
+
+    for (auto& el : to_delete) {
+      el->eraseFromParent();
+    }
+  }
+
+ private:
+  FunctionCallee ltest_allocation_fun, ltest_deallocation_fun;
+  Module& m;
+  // to avoid problems with correct iteration through instrcution remove
+  // instructions only at the end
+  std::vector<CallBase*> to_delete;
+  std::map<
+      std::string_view,
+      std::function<std::pair<FunctionCallee, SmallVector<Value*>>(CallBase&)>>
+      replaces = {{"malloc",
+                   [this](CallBase& a) -> auto {
+                     assert(a.arg_size() == 1 && "args count is 1");
+                     return std::pair{ltest_allocation_fun, collectArgs(a)};
+                   }},
+                  // {"operator new(unsigned long)",
+                  //  [this](CallBase& a) -> auto {
+                  //    assert(a.arg_size() == 1 && "args count is 1");
+                  //    return std::pair{ltest_allocation_fun, collectArgs(a)};
+                  //  }},
+                  {"free", [this](CallBase& a) -> auto {
+                     assert(a.arg_size() == 1 && "args count is 1");
+                     return std::pair{ltest_deallocation_fun, collectArgs(a)};
+                   }}};
+};
+namespace {
+
+struct ResMockPass final : public PassInfoMixin<ResMockPass> {
+  PreservedAnalyses run(Module& M, ModuleAnalysisManager& AM) {
+    ResMockInserter inserter(M);
+    inserter.Run();
+    if (verifyModule(M, &errs())) {
+      report_fatal_error("module verification failed", false);
+    }
+    return PreservedAnalyses::none();
+  };
+};
+
+}  // namespace
+
+extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
+llvmGetPassPluginInfo() {
+  return {.APIVersion = LLVM_PLUGIN_API_VERSION,
+          .PluginName = "resmockpass",
+          .PluginVersion = "v0.1",
+          .RegisterPassBuilderCallbacks = [](PassBuilder& pb) {
+            // This parsing we need for testing with opt
+            pb.registerPipelineParsingCallback(
+                [](StringRef Name, ModulePassManager& mpm,
+                   ArrayRef<PassBuilder::PipelineElement>) {
+                  if (Name == "resmock") {
+                    mpm.addPass(ResMockPass());
+                    return true;
+                  }
+                  return false;
+                });
+            pb.registerPipelineStartEPCallback(
+                [](ModulePassManager& mpm, OptimizationLevel level) {
+                  mpm.addPass(ResMockPass());
+                });
+          }};
+}

runtime/CMakeLists.txt

@@ -8,6 +8,8 @@ set (SOURCE_FILES
         minimization.cpp
         minimization_smart.cpp
         coro_ctx_guard.cpp
+        os_simulator.cpp
+        mock_res.cpp
 )
 
 add_library(runtime SHARED ${SOURCE_FILES})

runtime/include/lib.h

@@ -11,6 +11,7 @@
 #include <vector>
 
 #include "block_manager.h"
+#include "mock_res.h"
 #include "value_wrapper.h"
 
 #define panic() assert(false)
@@ -42,7 +43,7 @@ struct CoroBase : public std::enable_shared_from_this<CoroBase> {
 
   // Restart the coroutine from the beginning passing this_ptr as this.
   // Returns restarted coroutine.
-  virtual std::shared_ptr<CoroBase> Restart(void* this_ptr) = 0;
+  virtual void Restart(void* this_ptr) = 0;
 
   // Resume the coroutine to the next yield.
   void Resume();
@@ -122,16 +123,18 @@ struct Coro final : public CoroBase {
       std::function<std::vector<std::string>(std::shared_ptr<void>)>;
 
   // unsafe: caller must ensure that this_ptr points to Target.
-  std::shared_ptr<CoroBase> Restart(void* this_ptr) override {
-    /**
-     *  The task must be returned if we want to restart it.
-     *   We can't just Terminate() it because it is the runtime responsibility
-     * to decide, in which order the tasks should be terminated.
-     *
-     */
-    assert(IsReturned());
-    auto coro = New(func, this_ptr, args, args_to_strings, name, id);
-    return coro;
+  void Restart(void* this_ptr) override {
+    is_returned = false;
+    this->this_ptr = this_ptr;
+    ctx = boost::context::fiber_context(
+        [this, this_ptr](boost::context::fiber_context&& ctx) {
+          auto real_args = reinterpret_cast<std::tuple<Args...>*>(args.get());
+          auto this_arg =
+              std::tuple<Target*>{reinterpret_cast<Target*>(this_ptr)};
+          ret = std::apply(func, std::tuple_cat(this_arg, *real_args));
+          is_returned = true;
+          return std::move(ctx);
+        });
   }
 
   // unsafe: caller must ensure that this_ptr points to Target.

runtime/include/mock_res.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#include <cstddef>
+#include <deque>
+
+class MemoryHandler {
+  std::deque<void*> memory;
+
+ public:
+  void* Allocate(size_t);
+  void Deallocate(void*);
+  void FreeAllMemory();
+};
+
+extern MemoryHandler* memory_handler;
+
+extern "C" void* LtestMemAlloc(std::size_t size);
+
+extern "C" void LtestMemDealloc(void* ptr);

runtime/include/os_simulator.h

@@ -0,0 +1,15 @@
+#pragma once
+#include <cstddef>
+
+#include "mock_res.h"
+
+class OSSimulator {
+  MemoryHandler os_memory;
+
+ public:
+  OSSimulator() { memory_handler = &os_memory; }
+  bool CanThreadContinue(std::size_t number);
+  void UpdateState();
+  void ResetState();
+  ~OSSimulator() { ResetState(); }
+};

runtime/include/scheduler.h

@@ -15,6 +15,7 @@
 #include "logger.h"
 #include "minimization.h"
 #include "minimization_smart.h"
+#include "os_simulator.h"
 #include "pretty_print.h"
 #include "scheduler_fwd.h"
 #include "stable_vector.h"
@@ -191,7 +192,7 @@ struct BaseStrategyWithThreads : public Strategy {
       size_t tasks_in_thread = thread.size();
       for (size_t i = 0; i < tasks_in_thread; ++i) {
         if (!IsTaskRemoved(thread[i]->GetId())) {
-          thread[i] = thread[i]->Restart(state.get());
+          thread[i]->Restart(state.get());
         }
       }
     }
@@ -282,58 +283,11 @@ struct BaseStrategyWithThreads : public Strategy {
   }
 
  protected:
-  // Terminates all running tasks.
-  // We do it in a dangerous way: in random order.
-  // Actually, we assume obstruction free here.
   void TerminateTasks() {
-    auto& round_schedule = this->round_schedule;
-    assert(round_schedule.size() == this->threads.size() &&
-           "sizes expected to be the same");
-    round_schedule.assign(round_schedule.size(), -1);
-
-    std::vector<size_t> task_indexes(this->threads.size(), 0);
-    bool has_nonterminated_threads = true;
-    while (has_nonterminated_threads) {
-      has_nonterminated_threads = false;
-
-      for (size_t thread_index = 0; thread_index < this->threads.size();
-           ++thread_index) {
-        auto& thread = this->threads[thread_index];
-        auto& task_index = task_indexes[thread_index];
-
-        // find first non-finished task in the thread
-        while (task_index < thread.size() && thread[task_index]->IsReturned()) {
-          task_index++;
-        }
-
-        if (task_index == thread.size()) {
-          std::optional<std::string> releaseTask =
-              this->sched_checker.ReleaseTask(thread_index);
-          // Check if we should schedule release task to unblock other tasks
-          if (releaseTask) {
-            auto constructor =
-                *std::find_if(constructors.begin(), constructors.end(),
-                              [=](const TaskBuilder& b) {
-                                return b.GetName() == *releaseTask;
-                              });
-            auto task =
-                constructor.Build(this->state.get(), thread_index, task_index);
-            auto verified = this->sched_checker.Verify(
-                std::string(task->GetName()), thread_index);
-            thread.emplace_back(task);
-          }
-        }
-
-        if (task_index < thread.size() && !thread[task_index]->IsBlocked()) {
-          auto& task = thread[task_index];
-          has_nonterminated_threads = true;
-          // do a single step in this task
-          task->Resume();
-          if (task->IsReturned()) {
-            OnVerifierTaskFinish(task, thread_index);
-            debug(stderr, "Terminated: %ld\n", thread_index);
-          }
-        }
+    simulator.ResetState();
+    for (auto& thread : threads) {
+      if (!thread.empty() && !thread.back()->IsReturned()) {
+        thread.back()->Terminate();
       }
     }
     state.reset(new TargetObj{});
@@ -368,6 +322,7 @@ struct BaseStrategyWithThreads : public Strategy {
   std::uniform_int_distribution<std::mt19937::result_type>
       constructors_distribution;
   std::mt19937 rng;
+  OSSimulator simulator;
 };
 
 // StrategyScheduler generates different sequential histories (using Strategy)
@@ -653,15 +608,14 @@ struct TLAScheduler : Scheduler {
   TLAScheduler(size_t max_tasks, size_t max_rounds, size_t threads_count,
                size_t max_switches, size_t max_depth,
                std::vector<TaskBuilder> constructors, ModelChecker& checker,
-               PrettyPrinter& pretty_printer, std::function<void()> cancel_func)
+               PrettyPrinter& pretty_printer)
       : max_tasks{max_tasks},
         max_rounds{max_rounds},
         max_switches{max_switches},
         constructors{std::move(constructors)},
         checker{checker},
         pretty_printer{pretty_printer},
-        max_depth(max_depth),
-        cancel(cancel_func) {
+        max_depth(max_depth) {
     for (size_t i = 0; i < threads_count; ++i) {
       threads.emplace_back(Thread{
           .id = i,
@@ -708,7 +662,7 @@ struct TLAScheduler : Scheduler {
   // Actually, we assume obstruction free here.
   // cancel() func takes care for graceful shutdown
   void TerminateTasks() {
-    cancel();
+    simulator.ResetState();
     for (size_t i = 0; i < threads.size(); ++i) {
       for (size_t j = 0; j < threads[i].tasks.size(); ++j) {
         auto& task = threads[i].tasks[j];
@@ -732,7 +686,7 @@ struct TLAScheduler : Scheduler {
       if (frame.is_new) {
         // It was a new task.
         // So restart it from the beginning with the same args.
-        *task = (*task)->Restart(state.get());
+        (*task)->Restart(state.get());
       } else {
         // It was a not new task, hence, we recreated in early.
       }
@@ -942,5 +896,5 @@ struct TLAScheduler : Scheduler {
   StableVector<Thread> threads;
   StableVector<Frame> frames;
   Verifier verifier{};
-  std::function<void()> cancel;
+  OSSimulator simulator;
 };