From 949c05b6f9ba1fa98bbf89455a108642656f5331 Mon Sep 17 00:00:00 2001
From: Bogdan Nicolae <bogdan.nicolae@acm.org>
Date: Fri, 1 Mar 2019 19:15:41 -0600
Subject: [PATCH] support for partial variable recovery on restart

---
 include/veloc.h    |  2 ++
 src/lib/client.cpp | 30 +++++++++++++++++-------------
 src/lib/veloc.cpp  |  7 +++++++
 3 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/include/veloc.h b/include/veloc.h
index 5817022..ad7fa8e 100644
--- a/include/veloc.h
+++ b/include/veloc.h
@@ -21,6 +21,8 @@
 #define VELOC_MAX_NAME (1024)
 
 #define VELOC_RECOVER_ALL (0)
+#define VELOC_RECOVER_SOME (1)
+#define VELOC_RECOVER_REST (2)
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/lib/client.cpp b/src/lib/client.cpp
index 5c4d06a..7d1c5b3 100644
--- a/src/lib/client.cpp
+++ b/src/lib/client.cpp
@@ -177,12 +177,9 @@ bool veloc_client_t::restart_begin(const char *name, int version) {
 }
 
 bool veloc_client_t::recover_mem(int mode, std::set<int> &ids) {
-    if (mode != VELOC_RECOVER_ALL) {
-	ERROR("only VELOC_RECOVER_ALL mode currently supported");
-	return false;
-    }
-
     std::ifstream f;
+    std::map<int, size_t> region_info;
+
     f.exceptions(std::ifstream::failbit | std::ifstream::badbit);
     try {
 	f.open(current_ckpt.filename(cfg.get("scratch")), std::ifstream::in | std::ifstream::binary);
@@ -192,19 +189,26 @@ bool veloc_client_t::recover_mem(int mode, std::set<int> &ids) {
 	for (unsigned int i = 0; i < no_regions; i++) {
 	    f.read((char *)&id, sizeof(int));
 	    f.read((char *)&region_size, sizeof(size_t));
-	    if (mem_regions.find(id) == mem_regions.end()) {
-		ERROR("protected memory region " << id << " does not exist");
+	    region_info.insert(std::make_pair(id, region_size));
+	}
+	for (auto &e : region_info) {
+	    bool found = ids.find(e.first) != ids.end();
+	    if ((mode == VELOC_RECOVER_SOME && !found) || (mode == VELOC_RECOVER_REST && found)) {
+		f.seekg(e.second, std::ifstream::cur);
+		continue;
+	    }
+	    if (mem_regions.find(e.first) == mem_regions.end()) {
+		ERROR("no protected memory region defined for id " << e.first);
 		return false;
 	    }
-	    if (mem_regions[id].second < region_size) {
-		ERROR("protected memory region " << id << " is too small ("
-		      << mem_regions[id].second << ") to hold required size ("
-		      << region_size << ")");
+	    if (mem_regions[e.first].second < e.second) {
+		ERROR("protected memory region " << e.first << " is too small ("
+		      << mem_regions[e.first].second << ") to hold required size ("
+		      << e.second << ")");
 		return false;
 	    }
+	    f.read((char *)mem_regions[e.first].first, e.second);
 	}
-	for (auto &e : mem_regions)
-	    f.read((char *)e.second.first, e.second.second);
     } catch (std::ifstream::failure &e) {
 	ERROR("cannot read checkpoint file " << current_ckpt << ", reason: " << e.what());
 	return false;
diff --git a/src/lib/veloc.cpp b/src/lib/veloc.cpp
index 340ab22..fe11677 100644
--- a/src/lib/veloc.cpp
+++ b/src/lib/veloc.cpp
@@ -71,6 +71,13 @@ extern "C" int VELOC_Recover_mem() {
     return CLIENT_CALL(veloc_client->recover_mem(VELOC_RECOVER_ALL, ids));
 }
 
+extern "C" int VELOC_Recover_selective(int mode, int *ids, int no_ids) {
+    std::set<int> id_set = {};
+    for (int i = 0; i < no_ids; i++)
+	id_set.insert(ids[i]);
+    return CLIENT_CALL(veloc_client->recover_mem(mode, id_set));
+}
+
 extern "C" int VELOC_Restart_end(int success) {
     return CLIENT_CALL(veloc_client->restart_end(success));
 }