#2171: Allow for multiple key/value pairs to be added via addUserDefi…

…nedData
DARMA-tasking · Nov 6, 2023 · 242e4ed · 242e4ed
1 parent 1fefe4f
commit 242e4ed
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/src/vt/phase/phase_manager.h b/src/vt/phase/phase_manager.h
@@ -188,10 +188,16 @@ void addUserDefinedData(PhaseType phase, const KeyT& key, const ValueT& value) {
     "PhaseManager::addUserDefinedData: ValueT type is not jsonable"
   );
 
-  auto j = std::make_shared<nlohmann::json>();
-  j->emplace(key, value);
+  auto perPhase = theNodeLBData()->getLBData()->user_per_phase_json_;
 
-  theNodeLBData()->getLBData()->user_per_phase_json_[phase] = j;
+  if (perPhase.find(phase) == perPhase.end()) {
+    auto j = std::make_shared<nlohmann::json>();
+    j->emplace(key, value);
+
+    theNodeLBData()->getLBData()->user_per_phase_json_[phase] = j;
+  } else {
+    perPhase[phase]->emplace(key, value);
+  }
 }
 
   /**

diff --git a/tests/unit/collection/test_lb.extended.cc b/tests/unit/collection/test_lb.extended.cc
@@ -411,6 +411,9 @@ TEST_P(TestNodeLBDataDumper, test_node_lb_data_dumping_with_interval) {
     vt::thePhase()->addUserDefinedData(
       phase, std::string{"time"}, static_cast<double>(phase));
 
+    vt::thePhase()->addUserDefinedData(
+      phase, std::string{"new_time"}, static_cast<double>(phase));
+
     // Go to the next phase
     vt::thePhase()->nextPhaseCollective();
   }
@@ -431,6 +434,8 @@ TEST_P(TestNodeLBDataDumper, test_node_lb_data_dumping_with_interval) {
       EXPECT_TRUE(phase.find("user_defined") != phase.end());
       EXPECT_TRUE(phase["user_defined"].contains("time"));
       EXPECT_EQ(phase["user_defined"]["time"], static_cast<double>(phase["id"]));
+      EXPECT_TRUE(phase["user_defined"].contains("new_time"));
+      EXPECT_EQ(phase["user_defined"]["new_time"], static_cast<double>(phase["id"]));
     }
 
   });