From b4f37099a30434a64ff0b0c5f53171c093a05c17 Mon Sep 17 00:00:00 2001
From: Sean McGovern <st.mcgovern@gmail.com>
Date: Wed, 14 Sep 2022 01:39:43 +0000
Subject: [PATCH] #1943: add vt_checks in scheduler

---
 src/vt/scheduler/base_unit.cc       | 12 ++++++++----
 src/vt/scheduler/scheduler.cc       |  6 ++++++
 src/vt/scheduler/suspended_units.cc |  4 ++++
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/vt/scheduler/base_unit.cc b/src/vt/scheduler/base_unit.cc
index 344f7deed1..ed2e257358 100644
--- a/src/vt/scheduler/base_unit.cc
+++ b/src/vt/scheduler/base_unit.cc
@@ -50,10 +50,14 @@ namespace vt { namespace sched {
 void BaseUnit::execute() {
   if (r_) {
     r_->run();
-    if (not r_->isDone()) {
-      auto tid = r_->getThreadID();
-      theSched()->suspend(tid, std::move(r_));
-    }
+    #if vt_check_enabled(fcontext)
+      if (not r_->isDone()) {
+        auto tid = r_->getThreadID();
+        theSched()->suspend(tid, std::move(r_));
+      }
+    #endif
+    delete r_;
+    r_ = nullptr;
   } else if (work_) {
     work_();
   }
diff --git a/src/vt/scheduler/scheduler.cc b/src/vt/scheduler/scheduler.cc
index e5894ffefe..fe574166c2 100644
--- a/src/vt/scheduler/scheduler.cc
+++ b/src/vt/scheduler/scheduler.cc
@@ -128,9 +128,15 @@ void Scheduler::runWorkUnit(UnitType& work) {
 
   workUnitCount.increment(1);
 
+#if vt_check_enabled(mpi_access_guards)
   ++action_depth_;
+#endif
+
   work();
+
+#if vt_check_enabled(mpi_access_guards)
   --action_depth_;
+#endif
 
 #if vt_check_enabled(mpi_access_guards)
   if (action_depth_ == 0) {
diff --git a/src/vt/scheduler/suspended_units.cc b/src/vt/scheduler/suspended_units.cc
index 2a700f72c8..6150bd263c 100644
--- a/src/vt/scheduler/suspended_units.cc
+++ b/src/vt/scheduler/suspended_units.cc
@@ -50,6 +50,7 @@ namespace vt { namespace sched {
 void SuspendedUnits::addSuspended(
   ThreadIDType tid, RunnablePtrType runnable, PriorityType p
 ) {
+#if vt_check_enabled(fcontext)
   vtAssert(runnable->isSuspended(), "Runnable must be suspended to add");
 
   units_.emplace(
@@ -59,15 +60,18 @@ void SuspendedUnits::addSuspended(
       detail::SuspendedRunnable{std::move(runnable), p}
     )
   );
+#endif
 }
 
 void SuspendedUnits::resumeRunnable(ThreadIDType tid) {
+#if vt_check_enabled(fcontext)
   auto iter = units_.find(tid);
   vtAbortIf(iter == units_.end(), "Must have valid thread ID to resume");
   auto r = std::move(iter->second.runnable_);
   auto p = iter->second.priority_;
   theSched()->enqueue(p, std::move(r));
   units_.erase(iter);
+#endif
 }
 
 }} /* end namespace vt::sched */