From 413f98c442ff79f0c0ab939937b252424403183b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cezary=20Skrzy=C5=84ski?=
 <cezary.skrzynski@ng-analytics.com>
Date: Mon, 20 Jun 2022 10:32:09 +0200
Subject: [PATCH] #1830: extract more helper methods

---
 .../vrt/collection/balance/temperedlb/temperedlb.cc   | 11 ++++++++---
 src/vt/vrt/collection/balance/temperedlb/temperedlb.h |  3 +++
 .../collection/balance/temperedwmin/temperedwmin.h    |  2 ++
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc b/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc
index 20b187ce8f..e55b08794c 100644
--- a/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc
+++ b/src/vt/vrt/collection/balance/temperedlb/temperedlb.cc
@@ -681,7 +681,7 @@ void TemperedLB::informAsync() {
   auto propagate_epoch = theTerm()->makeEpochCollective("TemperedLB: informAsync");
 
   // Underloaded start the round
-  if (is_underloaded_) {
+  if (canPropagate()) {
     uint8_t k_cur_async = 0;
     propagateRound(k_cur_async, false, propagate_epoch);
   }
@@ -721,7 +721,7 @@ void TemperedLB::informSync() {
     underloaded_.insert(this_node);
   }
 
-  auto propagate_this_round = is_underloaded_;
+  auto propagate_this_round = canPropagate();
   propagate_next_round_ = false;
   new_underloaded_ = underloaded_;
   new_load_info_ = load_info_;
@@ -792,6 +792,9 @@ void TemperedLB::propagateRound(uint8_t k_cur, bool sync, EpochType epoch) {
     gen_propagate_.seed(seed_());
   }
 
+  // remove 1 line
+  // selected = getSelectedNodessss
+  // selected = all?
   auto& selected = selected_;
   selected = underloaded_;
   if (selected.find(this_node) == selected.end()) {
@@ -813,6 +816,8 @@ void TemperedLB::propagateRound(uint8_t k_cur, bool sync, EpochType epoch) {
       return;
     }
 
+    // extract generateRandomNode
+    // auto random_node = generateRandomNode(dist);
     // First, randomly select a node
     NodeType random_node = uninitialized_destination;
 
@@ -1202,7 +1207,7 @@ void TemperedLB::decide() {
 
   int n_transfers = 0, n_rejected = 0;
 
-  if (is_overloaded_) {
+  if (canMigrate()) {
     std::vector<NodeType> under = makeUnderloaded();
     std::unordered_map<NodeType, ObjsType> migrate_objs;
 
diff --git a/src/vt/vrt/collection/balance/temperedlb/temperedlb.h b/src/vt/vrt/collection/balance/temperedlb/temperedlb.h
index dd7849ee57..8086e08f13 100644
--- a/src/vt/vrt/collection/balance/temperedlb/temperedlb.h
+++ b/src/vt/vrt/collection/balance/temperedlb/temperedlb.h
@@ -96,6 +96,9 @@ struct TemperedLB : BaseLB {
   void migrate();
   void clearDataStructures();
 
+  virtual bool canMigrate() const {  return is_overloaded_; }
+  virtual bool canPropagate() const {  return is_underloaded_; }
+
   void propagateRound(uint8_t k_cur_async, bool sync, EpochType epoch = no_epoch);
   void propagateIncomingAsync(LoadMsgAsync* msg);
   void propagateIncomingSync(LoadMsgSync* msg);
diff --git a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h
index c008741b4a..df6cbe39af 100644
--- a/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h
+++ b/src/vt/vrt/collection/balance/temperedwmin/temperedwmin.h
@@ -61,6 +61,8 @@ struct TemperedWMin : TemperedLB {
   void inputParams(balance::SpecEntry* spec) override;
 
 protected:
+  bool canPropagate() const override { return true; }
+
   TimeType getModeledWork(const elm::ElementIDStruct& obj) override;
 
 private: