Add snapshot if exists function to snapshot registry

include/faabric/snapshot/SnapshotRegistry.h

@@ -24,6 +24,10 @@ class SnapshotRegistry
                       faabric::util::SnapshotData data,
                       bool locallyRestorable = true);
 
+    void takeSnapshotIfNotExists(const std::string& key,
+                                 faabric::util::SnapshotData data,
+                                 bool locallyRestorable = true);
+
     void deleteSnapshot(const std::string& key);
 
     size_t getSnapshotCount();
@@ -36,6 +40,11 @@ class SnapshotRegistry
     std::mutex snapshotsMx;
 
     int writeSnapshotToFd(const std::string& key);
+
+    void doTakeSnapshot(const std::string& key,
+                        faabric::util::SnapshotData data,
+                        bool locallyRestorable,
+                        bool overwrite);
 };
 
 SnapshotRegistry& getSnapshotRegistry();

src/snapshot/SnapshotRegistry.cpp

@@ -56,23 +56,44 @@ void SnapshotRegistry::mapSnapshot(const std::string& key, uint8_t* target)
     }
 }
 
+void SnapshotRegistry::takeSnapshotIfNotExists(const std::string& key,
+                                               faabric::util::SnapshotData data,
+                                               bool locallyRestorable)
+{
+    doTakeSnapshot(key, data, locallyRestorable, false);
+}
+
 void SnapshotRegistry::takeSnapshot(const std::string& key,
                                     faabric::util::SnapshotData data,
                                     bool locallyRestorable)
+{
+    doTakeSnapshot(key, data, locallyRestorable, true);
+}
+
+void SnapshotRegistry::doTakeSnapshot(const std::string& key,
+                                      faabric::util::SnapshotData data,
+                                      bool locallyRestorable,
+                                      bool overwrite)
 {
     if (data.size == 0) {
         SPDLOG_ERROR("Cannot take snapshot {} of size zero", key);
         throw std::runtime_error("Taking snapshot size zero");
     }
 
+    faabric::util::UniqueLock lock(snapshotsMx);
+
+    if (snapshotExists(key) && !overwrite) {
+        SPDLOG_TRACE("Skipping already existing snapshot {}", key);
+        return;
+    }
+
     SPDLOG_TRACE("Registering snapshot {} size {} (restorable={})",
                  key,
                  data.size,
                  locallyRestorable);
 
     // Note - we only preserve the snapshot in the in-memory file, and do not
     // take ownership for the original data referenced in SnapshotData
-    faabric::util::UniqueLock lock(snapshotsMx);
     snapshotMap[key] = data;
 
     // Write to fd to be locally restorable

tests/test/snapshot/test_snapshot_registry.cpp

@@ -113,6 +113,55 @@ TEST_CASE_METHOD(SnapshotTestFixture,
     deallocatePages(actualDataC, 3);
 }
 
+TEST_CASE_METHOD(SnapshotTestFixture,
+                 "Test set snapshot if not exists",
+                 "[snapshot]")
+{
+    REQUIRE(reg.getSnapshotCount() == 0);
+
+    std::string keyA = "snapA";
+    std::string keyB = "snapB";
+
+    REQUIRE(!reg.snapshotExists(keyA));
+    REQUIRE(!reg.snapshotExists(keyB));
+
+    // Take one of the snapshots
+    SnapshotData snapBefore = takeSnapshot(keyA, 1, true);
+
+    REQUIRE(reg.snapshotExists(keyA));
+    REQUIRE(!reg.snapshotExists(keyB));
+    REQUIRE(reg.getSnapshotCount() == 1);
+
+    // Set up some different data
+    std::vector<uint8_t> otherDataA(snapBefore.size + 10, 1);
+    std::vector<uint8_t> otherDataB(snapBefore.size + 5, 2);
+
+    SnapshotData snapUpdateA;
+    snapUpdateA.data = otherDataA.data();
+    snapUpdateA.size = otherDataA.size();
+
+    SnapshotData snapUpdateB;
+    snapUpdateB.data = otherDataB.data();
+    snapUpdateB.size = otherDataB.size();
+
+    // Check existing snapshot is not overwritten
+    reg.takeSnapshotIfNotExists(keyA, snapUpdateA, true);
+    SnapshotData snapAfterA = reg.getSnapshot(keyA);
+    REQUIRE(snapAfterA.data == snapBefore.data);
+    REQUIRE(snapAfterA.size == snapBefore.size);
+
+    // Check new snapshot is still created
+    reg.takeSnapshotIfNotExists(keyB, snapUpdateB, true);
+
+    REQUIRE(reg.snapshotExists(keyA));
+    REQUIRE(reg.snapshotExists(keyB));
+    REQUIRE(reg.getSnapshotCount() == 2);
+
+    SnapshotData snapAfterB = reg.getSnapshot(keyB);
+    REQUIRE(snapAfterB.data == otherDataB.data());
+    REQUIRE(snapAfterB.size == otherDataB.size());
+}
+
 TEST_CASE_METHOD(SnapshotTestFixture,
                  "Test can't get snapshot with empty key",
                  "[snapshot]")