From ec8ab0d85bd992dcde81e9f9610e7f2c15598eff Mon Sep 17 00:00:00 2001
From: Jakub Strzebonski <jakub.strzebonski@ng-analytics.com>
Date: Mon, 21 Sep 2020 23:11:09 +0200
Subject: [PATCH] #1014 dump node stats even if LB was not run

---
 .../vrt/collection/balance/lb_invoke/lb_manager.cc   |  4 ++++
 src/vt/vrt/collection/balance/lb_invoke/lb_manager.h | 11 +++++++++++
 src/vt/vrt/collection/balance/node_stats.cc          | 12 +++++++-----
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
index 440a839f33..05eeb291d5 100644
--- a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
+++ b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
@@ -353,4 +353,8 @@ void LBManager::triggerListeners(PhaseType phase) {
   }
 }
 
+void LBManager::flushNodeStats(PhaseType phase) {
+  theNodeStats()->outputStatsForPhase(phase);
+}
+
 }}}} /* end namespace vt::vrt::collection::balance */
diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.h b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.h
index 2d080e8b95..43617551cf 100644
--- a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.h
+++ b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.h
@@ -203,6 +203,7 @@ struct LBManager : runtime::component::Component<LBManager> {
   void sysReleaseLB(MsgT* msg) {
     vt_debug_print(lb, node, "sysReleaseLB\n");
     printMemoryUsage(msg->phase_);
+    flushNodeStats(msg->phase_);
     flushTraceNextPhase();
     setTraceEnabledNextPhase(msg->phase_);
     return releaseImpl(msg->phase_, msg->num_collections_);
@@ -263,6 +264,16 @@ struct LBManager : runtime::component::Component<LBManager> {
   template <typename LB>
   void makeLB(MsgSharedPtr<StartLBMsg> msg);
 
+private:
+  /**
+   * \internal \brief Communicate to the NodeStats component that a new phase
+   * occurred without running LB so flushing of previous node stats
+   * can occur if required
+   *
+   * \param[in] phase the phase
+   */
+  void flushNodeStats(PhaseType phase);
+
 private:
   std::size_t num_invocations_             = 0;
   std::size_t num_release_                 = 0;
diff --git a/src/vt/vrt/collection/balance/node_stats.cc b/src/vt/vrt/collection/balance/node_stats.cc
index 5e1700f6a4..4eb2449a9e 100644
--- a/src/vt/vrt/collection/balance/node_stats.cc
+++ b/src/vt/vrt/collection/balance/node_stats.cc
@@ -142,10 +142,7 @@ void NodeStats::startIterCleanup(PhaseType phase, int look_back) {
   }
   node_data_[phase] = std::move(new_data);
 
-  // Statistics output when LB is enabled and appropriate flag is enabled
-  if (theConfig()->vt_lb_stats) {
-    outputStatsForPhase(phase);
-  }
+  outputStatsForPhase(phase);
 
   if (phase - look_back >= 0) {
     node_data_.erase(phase - look_back);
@@ -192,7 +189,7 @@ void NodeStats::createStatsFile() {
 
   vt_debug_print(
     lb, node,
-    "NodeStats: createStatsFile file={}\n", file_name
+    "NodeStats::createStatsFile file={}\n", file_name
   );
 
   // Node 0 creates the directory
@@ -232,6 +229,11 @@ void NodeStats::closeStatsFile() {
 }
 
 void NodeStats::outputStatsForPhase(PhaseType phase) {
+  // Statistics output when LB is enabled and appropriate flag is enabled
+  if (!theConfig()->vt_lb_stats) {
+    return;
+  }
+
   vtAssertExpr(stats_file_ != nullptr);
 
   vt_print(lb, "NodeStats::outputStatsForPhase: file={}, phase={}\n", print_ptr(stats_file_), phase);