From 9e50e44fca9bbbe3dab09051b69bcc55014d2e1b Mon Sep 17 00:00:00 2001
From: Jonathan Lifflander <jliffla@sandia.gov>
Date: Wed, 30 Jun 2021 09:47:06 -0700
Subject: [PATCH] #1490: vrt_coll: implement new restoreToFileInPlace

---
 src/vt/vrt/collection/manager.h      | 23 +++++++++
 src/vt/vrt/collection/manager.impl.h | 73 ++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/src/vt/vrt/collection/manager.h b/src/vt/vrt/collection/manager.h
index 3e091de46c..9422522f6d 100644
--- a/src/vt/vrt/collection/manager.h
+++ b/src/vt/vrt/collection/manager.h
@@ -1861,6 +1861,29 @@ struct CollectionManager
     typename ColT::IndexType range, std::string const& file_base
   );
 
+  template <typename ColT, typename IdxT = typename ColT::IndexType>
+  struct RestoreMigrateMsg : vt::Message {
+    RestoreMigrateMsg() = default;
+    RestoreMigrateMsg(NodeType in_to_node, IdxT in_idx, CollectionProxyWrapType<ColT> in_proxy)
+      : to_node_(in_to_node),
+        idx_(in_idx),
+        proxy_(in_proxy)
+    { }
+
+    NodeType to_node_ = uninitialized_destination;
+    IdxT idx_;
+    CollectionProxyWrapType<ColT> proxy_;
+  };
+
+  template <typename ColT>
+  static void restoreHandler(RestoreMigrateMsg<ColT>* msg);
+
+  template <typename ColT>
+  void restoreFromFileInPlace(
+    CollectionProxyWrapType<ColT> proxy, typename ColT::IndexType range,
+    std::string const& file_base
+  );
+
   template <typename SerializerT>
   void serialize(SerializerT& s) {
     s | buffers_
diff --git a/src/vt/vrt/collection/manager.impl.h b/src/vt/vrt/collection/manager.impl.h
index f3ae86164b..b9100017ce 100644
--- a/src/vt/vrt/collection/manager.impl.h
+++ b/src/vt/vrt/collection/manager.impl.h
@@ -3148,6 +3148,79 @@ void CollectionManager::checkpointToFile(
   checkpoint::serializeToFile(directory, directory_name);
 }
 
+template <typename ColT>
+/*static*/ void CollectionManager::restoreHandler(RestoreMigrateMsg<ColT>* msg) {
+  auto idx = msg->idx_;
+  auto node = msg->to_node_;
+  auto proxy = msg->proxy_;
+  theCollection()->migrate(proxy(idx), node);
+}
+
+template <typename ColT>
+void CollectionManager::restoreFromFileInPlace(
+  CollectionProxyWrapType<ColT> proxy, typename ColT::IndexType range,
+  std::string const& file_base
+) {
+  using IndexType = typename ColT::IndexType;
+  using DirectoryType = CollectionDirectory<IndexType>;
+
+  auto proxy_bits = proxy.getProxy();
+
+  auto metadata_file_name = makeMetaFilename<IndexType>(file_base, false);
+
+  if (access(metadata_file_name.c_str(), F_OK) == -1) {
+    // file doesn't exist, try looking in sub-directory
+    metadata_file_name = makeMetaFilename<IndexType>(file_base, true);
+  }
+
+  if (access(metadata_file_name.c_str(), F_OK) == -1) {
+    throw std::runtime_error("Collection directory file cannot be found");
+  }
+
+  auto directory = checkpoint::deserializeFromFile<DirectoryType>(
+    metadata_file_name
+  );
+
+  runInEpochCollective([&]{
+    for (auto&& elm : directory->elements_) {
+      auto idx = elm.idx_;
+      auto file_name = elm.file_name_;
+
+      if (proxy(idx).tryGetLocalPtr() == nullptr) {
+        auto mapped_node = getMappedNode<ColT, IndexType>(proxy, idx);
+        vtAssertExpr(mapped_node != uninitialized_destination);
+        auto this_node = theContext()->getNode();
+
+        auto msg = makeMessage<RestoreMigrateMsg<IndexType>>(this_node, idx, proxy);
+        theMsg()->sendMsg<
+          RestoreMigrateMsg<IndexType>, restoreHandler<IndexType>
+        >(mapped_node, msg);
+      }
+    }
+  });
+
+  for (auto&& elm : directory->elements_) {
+    auto idx = elm.idx_;
+    auto file_name = elm.file_name_;
+    vtAssertExpr(proxy(idx).tryGetLocalPtr() != nullptr);
+
+    auto col_ptr = checkpoint::deserializeFromFile<ColT>(file_name);
+    col_ptr->stats_.resetPhase();
+
+    auto holder = findColHolder<ColT, IndexType>(proxy_bits);
+    vtAssertExpr(holder != nullptr);
+
+    auto elm_holder = findElmHolder<ColT,IndexType>(proxy_bits);
+    auto const elm_exists = elm_holder->exists(idx);
+    vtAssertExpr(elm_exists);
+
+    auto map_han = UniversalIndexHolder<>::getMap(proxy_bits);
+    elm_holder->insert(idx, typename Holder<ColT, IndexType>::InnerHolder{
+      std::move(col_ptr), map_han, range
+    });
+  }
+}
+
 template <typename ColT>
 CollectionManager::CollectionProxyWrapType<ColT>
 CollectionManager::restoreFromFile(