From 7f463fa0b71f06fc62fc0382f0ba71dbd59dfd53 Mon Sep 17 00:00:00 2001
From: Sean McGovern <st.mcgovern@gmail.com>
Date: Tue, 15 Nov 2022 01:35:32 +0000
Subject: [PATCH] #1938: have scheduler store current time and lb context use
 it for start time

---
 src/vt/context/runnable_context/lb_data.cc | 3 ++-
 src/vt/elm/elm_lb_data.cc                  | 4 ++--
 src/vt/elm/elm_lb_data.h                   | 2 +-
 src/vt/scheduler/scheduler.cc              | 1 +
 src/vt/scheduler/scheduler.h               | 8 ++++++++
 5 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/vt/context/runnable_context/lb_data.cc b/src/vt/context/runnable_context/lb_data.cc
index 66b6e47ab2..250435aab9 100644
--- a/src/vt/context/runnable_context/lb_data.cc
+++ b/src/vt/context/runnable_context/lb_data.cc
@@ -49,7 +49,8 @@ namespace vt { namespace ctx {
 void LBData::begin() {
   // record start time
   if (should_instrument_) {
-    lb_data_->startTime();
+    TimeType current_time = theSched()->getCurrentTimeFromScheduler();
+    lb_data_->startTime(current_time);
   }
 }
 
diff --git a/src/vt/elm/elm_lb_data.cc b/src/vt/elm/elm_lb_data.cc
index d5dc5d5584..f61350e795 100644
--- a/src/vt/elm/elm_lb_data.cc
+++ b/src/vt/elm/elm_lb_data.cc
@@ -50,8 +50,8 @@
 
 namespace vt { namespace elm {
 
-void ElementLBData::startTime() {
-  TimeTypeWrapper const start_time = timing::getCurrentTime();
+void ElementLBData::startTime(TimeType current_time) {
+  TimeTypeWrapper const start_time = current_time;
   cur_time_ = start_time.seconds();
   cur_time_started_ = true;
 
diff --git a/src/vt/elm/elm_lb_data.h b/src/vt/elm/elm_lb_data.h
index 554856f8d1..cdd6a841a5 100644
--- a/src/vt/elm/elm_lb_data.h
+++ b/src/vt/elm/elm_lb_data.h
@@ -61,7 +61,7 @@ struct ElementLBData {
   ElementLBData(ElementLBData const&) = default;
   ElementLBData(ElementLBData&&) = default;
 
-  void startTime();
+  void startTime(TimeType current_time);
   void stopTime();
   void addTime(TimeTypeWrapper const& time);
 
diff --git a/src/vt/scheduler/scheduler.cc b/src/vt/scheduler/scheduler.cc
index 58890baeda..854cc80e20 100644
--- a/src/vt/scheduler/scheduler.cc
+++ b/src/vt/scheduler/scheduler.cc
@@ -251,6 +251,7 @@ void Scheduler::runProgress(bool msg_only, TimeType current_time) {
 void Scheduler::runSchedulerOnceImpl(bool msg_only) {
   if (special_progress_) {
     auto current_time = timing::getCurrentTime();
+    current_sched_time_ = current_time;
     auto time_since_last_progress = current_time - last_progress_time_;
     if (shouldCallProgress(processed_after_last_progress_, time_since_last_progress)) {
       runProgress(msg_only, current_time);
diff --git a/src/vt/scheduler/scheduler.h b/src/vt/scheduler/scheduler.h
index 4cc2e96d95..89ea63f649 100644
--- a/src/vt/scheduler/scheduler.h
+++ b/src/vt/scheduler/scheduler.h
@@ -332,6 +332,12 @@ struct Scheduler : runtime::component::Component<Scheduler> {
    */
   ThreadManager* getThreadManager();
 #endif
+  /**
+   * \brief return the time set at last scheduler call to getCurrentTime
+   *
+   * \return scheduler's current time
+   */
+  TimeType getCurrentTimeFromScheduler(){ return current_sched_time_; }
 
   template <typename SerializerT>
   void serialize(SerializerT& s) {
@@ -348,6 +354,7 @@ struct Scheduler : runtime::component::Component<Scheduler> {
       | event_triggers
       | event_triggers_once
       | last_progress_time_
+      | current_sched_time_
       | progress_time_enabled_
       | processed_after_last_progress_
       | last_threshold_memory_usage_
@@ -417,6 +424,7 @@ struct Scheduler : runtime::component::Component<Scheduler> {
   EventTriggerContType event_triggers_once;
 
   TimeType last_progress_time_ = 0.0;
+  TimeType current_sched_time_ = 0.0;
   bool progress_time_enabled_ = false;
   int32_t processed_after_last_progress_ = 0;