From 3138a72b267ffcd8347e251d17dd40c27c78d944 Mon Sep 17 00:00:00 2001
From: Jakub Domagala <domagala.jacob@gmail.com>
Date: Fri, 25 Sep 2020 00:20:39 +0200
Subject: [PATCH] #1024: Schedule the PendingSend call from collective bcast
 instead of invoking it

---
 src/vt/vrt/collection/manager.h      |  8 +++---
 src/vt/vrt/collection/manager.impl.h | 40 +++++++++++++---------------
 2 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/src/vt/vrt/collection/manager.h b/src/vt/vrt/collection/manager.h
index 240742e7fe..3fec762b35 100644
--- a/src/vt/vrt/collection/manager.h
+++ b/src/vt/vrt/collection/manager.h
@@ -964,7 +964,7 @@ struct CollectionManager
   >
   messaging::PendingSend broadcastMsgCollective(
     CollectionProxyWrapType<typename MsgT::CollectionType> const& proxy,
-    MsgT* msg, bool instrument = true);
+    messaging::MsgPtrThief<MsgT> msg, bool instrument = true);
 
   /**
    * \brief Broadcast collective a message with action member handler
@@ -982,7 +982,7 @@ struct CollectionManager
   >
   messaging::PendingSend broadcastMsgCollective(
     CollectionProxyWrapType<typename MsgT::CollectionType> const& proxy,
-    MsgT* msg, bool instrument = true);
+    messaging::MsgPtrThief<MsgT> msg, bool instrument = true);
 
   /**
    * \internal \brief Broadcast collective a message
@@ -996,8 +996,8 @@ struct CollectionManager
    */
   template <typename MsgT, typename ColT>
   messaging::PendingSend broadcastMsgCollectiveImpl(
-    CollectionProxyWrapType<ColT> const& proxy, MsgT* msg,
-    bool instrument = true);
+    CollectionProxyWrapType<ColT> const& proxy, MsgPtr<MsgT>& msg,
+    bool instrument);
 
   /**
    * \brief Broadcast a message with action function handler
diff --git a/src/vt/vrt/collection/manager.impl.h b/src/vt/vrt/collection/manager.impl.h
index 1ceb648572..76baa9aae4 100644
--- a/src/vt/vrt/collection/manager.impl.h
+++ b/src/vt/vrt/collection/manager.impl.h
@@ -854,14 +854,15 @@ template <
 >
 messaging::PendingSend CollectionManager::broadcastMsgCollective(
   CollectionProxyWrapType<typename MsgT::CollectionType> const& proxy,
-  MsgT* msg, bool instrument
+  messaging::MsgPtrThief<MsgT> msg, bool instrument
 ) {
   using ColT = typename MsgT::CollectionType;
 
-  msg->setVrtHandler(auto_registry::makeAutoHandlerCollection<ColT, MsgT, f>());
-  msg->setMember(false);
+  auto& msgPtr = msg.msg_;
+  msgPtr->setVrtHandler(auto_registry::makeAutoHandlerCollection<ColT, MsgT, f>());
+  msgPtr->setMember(false);
 
-  return broadcastMsgCollectiveImpl<MsgT, ColT>(proxy, msg, instrument);
+  return broadcastMsgCollectiveImpl<MsgT, ColT>(proxy, msgPtr, instrument);
 }
 
 template <
@@ -870,24 +871,24 @@ template <
 >
 messaging::PendingSend CollectionManager::broadcastMsgCollective(
   CollectionProxyWrapType<typename MsgT::CollectionType> const& proxy,
-  MsgT* msg, bool instrument
+  messaging::MsgPtrThief<MsgT> msg, bool instrument
 ) {
   using ColT = typename MsgT::CollectionType;
 
-  msg->setVrtHandler(
+  auto& msgPtr = msg.msg_;
+  msgPtr->setVrtHandler(
     auto_registry::makeAutoHandlerCollectionMem<ColT, MsgT, f>()
   );
-  msg->setMember(true);
+  msgPtr->setMember(true);
 
-  return broadcastMsgCollectiveImpl<MsgT, ColT>(proxy, msg, instrument);
+  return broadcastMsgCollectiveImpl<MsgT, ColT>(proxy, msgPtr, instrument);
 }
 
 template <typename MsgT, typename ColT>
 messaging::PendingSend CollectionManager::broadcastMsgCollectiveImpl(
-  CollectionProxyWrapType<ColT> const& proxy, MsgT* raw_msg, bool instrument
+  CollectionProxyWrapType<ColT> const& proxy, MsgPtr<MsgT>& msg, bool instrument
 ) {
   using IndexT = typename ColT::IndexType;
-  auto msg = promoteMsg(raw_msg);
 
   msg->setFromNode(theContext()->getNode());
   msg->setBcastProxy(proxy.getProxy());
@@ -910,20 +911,17 @@ messaging::PendingSend CollectionManager::broadcastMsgCollectiveImpl(
   msg->setCat(balance::CommCategory::NodeToCollection);
 #endif
 
-  theMsg()->setupEpochMsg(msg);
-
-  return messaging::PendingSend(
-    msg, [proxy](MsgSharedPtr<BaseMsgType>& msgIn){
-      auto col_msg = reinterpret_cast<MsgT*>(msgIn.get());
+  auto const cur_epoch = theMsg()->setupEpochMsg(msg);
 
-      auto elm_holder = theCollection()->findElmHolder<ColT, IndexT>(proxy);
-      auto const bcast_epoch = elm_holder->cur_bcast_epoch_++;
-      col_msg->setBcastEpoch(bcast_epoch);
+  return schedule(msg, false, cur_epoch, [proxy, msg] {
+    auto elm_holder = theCollection()->findElmHolder<ColT, IndexT>(proxy);
+    auto const bcast_epoch = elm_holder->cur_bcast_epoch_++;
+    msg->setBcastEpoch(bcast_epoch);
 
-      theMsg()->markAsCollectionMessage(col_msg);
+    theMsg()->markAsCollectionMessage(msg);
 
-      collectionBcastHandler<ColT, IndexT, MsgT>(col_msg);
-    }
+    collectionBcastHandler<ColT, IndexT, MsgT>(msg.get());
+  }
   );
 }