Add snapshotExists function to registry (#136)

* Small fixes for snapshot handling * Remove app id change * Add tests
faasm · Sep 8, 2021 · 68bdb05 · 68bdb05
1 parent 0ff6027
commit 68bdb05
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 3 deletions.
diff --git a/include/faabric/snapshot/SnapshotRegistry.h b/include/faabric/snapshot/SnapshotRegistry.h
@@ -16,6 +16,8 @@ class SnapshotRegistry
 
     faabric::util::SnapshotData& getSnapshot(const std::string& key);
 
+    bool snapshotExists(const std::string& key);
+
     void mapSnapshot(const std::string& key, uint8_t* target);
 
     void takeSnapshot(const std::string& key,

diff --git a/src/snapshot/SnapshotRegistry.cpp b/src/snapshot/SnapshotRegistry.cpp
@@ -25,6 +25,11 @@ faabric::util::SnapshotData& SnapshotRegistry::getSnapshot(
     return snapshotMap[key];
 }
 
+bool SnapshotRegistry::snapshotExists(const std::string& key)
+{
+    return snapshotMap.find(key) != snapshotMap.end();
+}
+
 void SnapshotRegistry::mapSnapshot(const std::string& key, uint8_t* target)
 {
     faabric::util::SnapshotData d = getSnapshot(key);

diff --git a/tests/test/snapshot/test_snapshot_registry.cpp b/tests/test/snapshot/test_snapshot_registry.cpp
@@ -21,9 +21,24 @@ TEST_CASE_METHOD(SnapshotTestFixture,
     std::string keyB = "snapB";
     std::string keyC = "snapC";
 
+    REQUIRE(!reg.snapshotExists(keyA));
+    REQUIRE(!reg.snapshotExists(keyB));
+    REQUIRE(!reg.snapshotExists(keyC));
+
     SnapshotData snapA = takeSnapshot(keyA, 1, true);
     SnapshotData snapB = takeSnapshot(keyB, 2, false);
-    SnapshotData snapC = takeSnapshot(keyA, 3, true);
+
+    REQUIRE(reg.snapshotExists(keyA));
+    REQUIRE(reg.snapshotExists(keyB));
+    REQUIRE(!reg.snapshotExists(keyC));
+    REQUIRE(reg.getSnapshotCount() == 2);
+
+    SnapshotData snapC = takeSnapshot(keyC, 3, true);
+
+    REQUIRE(reg.snapshotExists(keyA));
+    REQUIRE(reg.snapshotExists(keyB));
+    REQUIRE(reg.snapshotExists(keyC));
+    REQUIRE(reg.getSnapshotCount() == 3);
 
     // Add some random bits of data to the vectors
     for (int i = 0; i < HOST_PAGE_SIZE - 10; i += 50) {
@@ -32,12 +47,11 @@ TEST_CASE_METHOD(SnapshotTestFixture,
         snapC.data[i + 2] = i;
     }
 
+    // Take snapshots again with updated data
     reg.takeSnapshot(keyA, snapA);
     reg.takeSnapshot(keyB, snapB, false);
     reg.takeSnapshot(keyC, snapC);
 
-    REQUIRE(reg.getSnapshotCount() == 3);
-
     SnapshotData actualA = reg.getSnapshot(keyA);
     SnapshotData actualB = reg.getSnapshot(keyB);
     SnapshotData actualC = reg.getSnapshot(keyC);
@@ -81,7 +95,19 @@ TEST_CASE_METHOD(SnapshotTestFixture,
 
     removeSnapshot(keyA, 1);
     removeSnapshot(keyB, 2);
+
+    REQUIRE(!reg.snapshotExists(keyA));
+    REQUIRE(!reg.snapshotExists(keyB));
+    REQUIRE(reg.snapshotExists(keyC));
+    REQUIRE(reg.getSnapshotCount() == 1);
+
     removeSnapshot(keyC, 3);
+
+    REQUIRE(!reg.snapshotExists(keyA));
+    REQUIRE(!reg.snapshotExists(keyB));
+    REQUIRE(!reg.snapshotExists(keyC));
+    REQUIRE(reg.getSnapshotCount() == 0);
+
     deallocatePages(actualDataA, 1);
     deallocatePages(actualDataB, 2);
     deallocatePages(actualDataC, 3);