More race condition fixes

.gitignore

@@ -1,4 +1,5 @@
 build/
+work/
 conan-cache/
 # Clang
 .clangd

include/faabric/snapshot/SnapshotRegistry.h

@@ -1,10 +1,11 @@
 #pragma once
 
-#include <mutex>
+#include <shared_mutex>
 #include <string>
 #include <unordered_map>
 
 #include <faabric/proto/faabric.pb.h>
+#include <faabric/util/locks.h>
 #include <faabric/util/snapshot.h>
 
 namespace faabric::snapshot {
@@ -14,7 +15,8 @@ class SnapshotRegistry
   public:
     SnapshotRegistry() = default;
 
-    faabric::util::SnapshotData& getSnapshot(const std::string& key);
+    std::shared_ptr<faabric::util::SnapshotData> getSnapshot(
+      const std::string& key);
 
     bool snapshotExists(const std::string& key);
 
@@ -35,11 +37,14 @@ class SnapshotRegistry
     void clear();
 
   private:
-    std::unordered_map<std::string, faabric::util::SnapshotData> snapshotMap;
+    std::unordered_map<std::string,
+                       std::shared_ptr<faabric::util::SnapshotData>>
+      snapshotMap;
 
-    std::mutex snapshotsMx;
+    std::shared_mutex snapshotsMx;
 
-    int writeSnapshotToFd(const std::string& key);
+    int writeSnapshotToFd(const std::string& key,
+                          faabric::util::SnapshotData& data);
 
     void doTakeSnapshot(const std::string& key,
                         faabric::util::SnapshotData data,

include/faabric/util/bytes.h

@@ -25,6 +25,21 @@ int safeCopyToBuffer(const uint8_t* dataIn,
                      uint8_t* buffer,
                      int bufferLen);
 
+template<class T>
+T unalignedRead(const std::byte* bytes)
+{
+    T value;
+    std::copy_n(bytes, sizeof(T), reinterpret_cast<std::byte*>(&value));
+    return value;
+}
+
+template<class T>
+void unalignedWrite(const T& value, std::byte* destination)
+{
+    std::copy_n(
+      reinterpret_cast<const std::byte*>(&value), sizeof(T), destination);
+}
+
 template<class T>
 void appendBytesOf(std::vector<uint8_t>& container, T value)
 {

src/scheduler/Executor.cpp

@@ -282,18 +282,18 @@ void Executor::threadPoolThread(int threadPoolIdx)
         // Handle snapshot diffs _before_ we reset the executor
         if (isLastInBatch && task.needsSnapshotPush) {
             // Get diffs between original snapshot and after execution
-            faabric::util::SnapshotData snapshotPostExecution = snapshot();
+            auto snapshotPostExecution = snapshot();
 
-            faabric::util::SnapshotData snapshotPreExecution =
+            auto snapshotPreExecution =
               faabric::snapshot::getSnapshotRegistry().getSnapshot(
                 msg.snapshotkey());
 
             SPDLOG_TRACE("Diffing pre and post execution snapshots for {}",
                          msg.snapshotkey());
 
             std::vector<faabric::util::SnapshotDiff> diffs =
-              snapshotPreExecution.getChangeDiffs(snapshotPostExecution.data,
-                                                  snapshotPostExecution.size);
+              snapshotPreExecution->getChangeDiffs(snapshotPostExecution.data,
+                                                   snapshotPostExecution.size);
 
             sch.pushSnapshotDiffs(msg, diffs);
 

src/scheduler/Scheduler.cpp

@@ -497,18 +497,18 @@ faabric::util::SchedulingDecision Scheduler::doCallFunctions(
     if (!snapshotKey.empty()) {
         for (const auto& host : getFunctionRegisteredHosts(firstMsg, false)) {
             SnapshotClient& c = getSnapshotClient(host);
-            faabric::util::SnapshotData snapshotData =
+            auto snapshotData =
               faabric::snapshot::getSnapshotRegistry().getSnapshot(snapshotKey);
 
             // See if we've already pushed this snapshot to the given host,
             // if so, just push the diffs
             if (pushedSnapshotsMap[snapshotKey].contains(host)) {
                 std::vector<faabric::util::SnapshotDiff> snapshotDiffs =
-                  snapshotData.getDirtyPages();
+                  snapshotData->getDirtyPages();
                 c.pushSnapshotDiffs(
                   snapshotKey, firstMsg.groupid(), snapshotDiffs);
             } else {
-                c.pushSnapshot(snapshotKey, firstMsg.groupid(), snapshotData);
+                c.pushSnapshot(snapshotKey, firstMsg.groupid(), *snapshotData);
                 pushedSnapshotsMap[snapshotKey].insert(host);
             }
         }
@@ -909,18 +909,21 @@ void Scheduler::pushSnapshotDiffs(
 
 void Scheduler::setThreadResultLocally(uint32_t msgId, int32_t returnValue)
 {
+    faabric::util::SharedLock lock(mx);
     SPDLOG_DEBUG("Setting result for thread {} to {}", msgId, returnValue);
-    threadResults[msgId].set_value(returnValue);
+    threadResults.at(msgId).set_value(returnValue);
 }
 
 int32_t Scheduler::awaitThreadResult(uint32_t messageId)
 {
-    if (threadResults.count(messageId) == 0) {
+    faabric::util::SharedLock lock(mx);
+    auto it = threadResults.find(messageId);
+    if (it == threadResults.end()) {
         SPDLOG_ERROR("Thread {} not registered on this host", messageId);
         throw std::runtime_error("Awaiting unregistered thread");
     }
-
-    return threadResults[messageId].get_future().get();
+    lock.unlock();
+    return it->second.get_future().get();
 }
 
 faabric::Message Scheduler::getFunctionResult(unsigned int messageId,

src/snapshot/SnapshotRegistry.cpp

@@ -7,9 +7,11 @@
 #include <sys/mman.h>
 
 namespace faabric::snapshot {
-faabric::util::SnapshotData& SnapshotRegistry::getSnapshot(
+std::shared_ptr<faabric::util::SnapshotData> SnapshotRegistry::getSnapshot(
   const std::string& key)
 {
+    faabric::util::SharedLock lock(snapshotsMx);
+
     if (key.empty()) {
         SPDLOG_ERROR("Attempting to get snapshot with empty key");
         throw std::runtime_error("Getting snapshot with empty key");
@@ -30,7 +32,7 @@ bool SnapshotRegistry::snapshotExists(const std::string& key)
 
 void SnapshotRegistry::mapSnapshot(const std::string& key, uint8_t* target)
 {
-    faabric::util::SnapshotData d = getSnapshot(key);
+    auto d = getSnapshot(key);
 
     if (!faabric::util::isPageAligned((void*)target)) {
         SPDLOG_ERROR(
@@ -39,13 +41,13 @@ void SnapshotRegistry::mapSnapshot(const std::string& key, uint8_t* target)
           "Mapping snapshot to non page-aligned address");
     }
 
-    if (d.fd == 0) {
+    if (d->fd == 0) {
         SPDLOG_ERROR("Attempting to map non-restorable snapshot");
         throw std::runtime_error("Mapping non-restorable snapshot");
     }
 
     void* mmapRes =
-      mmap(target, d.size, PROT_WRITE, MAP_PRIVATE | MAP_FIXED, d.fd, 0);
+      mmap(target, d->size, PROT_WRITE, MAP_PRIVATE | MAP_FIXED, d->fd, 0);
 
     if (mmapRes == MAP_FAILED) {
         SPDLOG_ERROR(
@@ -78,7 +80,7 @@ void SnapshotRegistry::doTakeSnapshot(const std::string& key,
         throw std::runtime_error("Taking snapshot size zero");
     }
 
-    faabric::util::UniqueLock lock(snapshotsMx);
+    faabric::util::FullLock lock(snapshotsMx);
 
     if (snapshotExists(key) && !overwrite) {
         SPDLOG_TRACE("Skipping already existing snapshot {}", key);
@@ -92,37 +94,40 @@ void SnapshotRegistry::doTakeSnapshot(const std::string& key,
 
     // Note - we only preserve the snapshot in the in-memory file, and do not
     // take ownership for the original data referenced in SnapshotData
-    snapshotMap[key] = data;
+    auto shared_data =
+      std::make_shared<faabric::util::SnapshotData>(std::move(data));
+    snapshotMap[key] = shared_data;
 
     // Write to fd to be locally restorable
     if (locallyRestorable) {
-        writeSnapshotToFd(key);
+        writeSnapshotToFd(key, *shared_data);
     }
 }
 
 void SnapshotRegistry::deleteSnapshot(const std::string& key)
 {
-    faabric::util::UniqueLock lock(snapshotsMx);
+    faabric::util::FullLock lock(snapshotsMx);
 
     if (snapshotMap.count(key) == 0) {
         return;
     }
 
-    faabric::util::SnapshotData d = snapshotMap[key];
+    auto d = snapshotMap[key];
 
     // Note - the data referenced by the SnapshotData object is not owned by the
     // snapshot registry so we don't delete it here. We only remove the file
     // descriptor used for mapping memory
-    if (d.fd > 0) {
-        ::close(d.fd);
+    if (d->fd > 0) {
+        ::close(d->fd);
+        d->fd = 0;
     }
 
     snapshotMap.erase(key);
 }
 
 size_t SnapshotRegistry::getSnapshotCount()
 {
-    faabric::util::UniqueLock lock(snapshotsMx);
+    faabric::util::FullLock lock(snapshotsMx);
     return snapshotMap.size();
 }
 
@@ -134,36 +139,37 @@ SnapshotRegistry& getSnapshotRegistry()
 
 void SnapshotRegistry::clear()
 {
+    faabric::util::FullLock lock(snapshotsMx);
     for (auto p : snapshotMap) {
-        if (p.second.fd > 0) {
-            ::close(p.second.fd);
+        if (p.second->fd > 0) {
+            ::close(p.second->fd);
         }
     }
 
     snapshotMap.clear();
 }
 
-int SnapshotRegistry::writeSnapshotToFd(const std::string& key)
+int SnapshotRegistry::writeSnapshotToFd(const std::string& key,
+                                        faabric::util::SnapshotData& data)
 {
     int fd = ::memfd_create(key.c_str(), 0);
-    faabric::util::SnapshotData snapData = getSnapshot(key);
 
     // Make the fd big enough
-    int ferror = ::ftruncate(fd, snapData.size);
+    int ferror = ::ftruncate(fd, data.size);
     if (ferror) {
         SPDLOG_ERROR("ferror call failed with error {}", ferror);
         throw std::runtime_error("Failed writing memory to fd (ftruncate)");
     }
 
     // Write the data
-    ssize_t werror = ::write(fd, snapData.data, snapData.size);
+    ssize_t werror = ::write(fd, data.data, data.size);
     if (werror == -1) {
         SPDLOG_ERROR("Write call failed with error {}", werror);
         throw std::runtime_error("Failed writing memory to fd (write)");
     }
 
     // Record the fd
-    getSnapshot(key).fd = fd;
+    data.fd = fd;
 
     SPDLOG_DEBUG("Wrote snapshot {} to fd {}", key, fd);
     return fd;

src/snapshot/SnapshotServer.cpp

@@ -127,7 +127,7 @@ SnapshotServer::recvPushSnapshotDiffs(const uint8_t* buffer, size_t bufferSize)
     // Get the snapshot
     faabric::snapshot::SnapshotRegistry& reg =
       faabric::snapshot::getSnapshotRegistry();
-    faabric::util::SnapshotData& snap = reg.getSnapshot(r->key()->str());
+    auto snap = reg.getSnapshot(r->key()->str());
 
     // Lock the function group if it exists
     if (groupId > 0 &&
@@ -138,7 +138,7 @@ SnapshotServer::recvPushSnapshotDiffs(const uint8_t* buffer, size_t bufferSize)
 
     // Iterate through the chunks passed in the request
     for (const auto* chunk : *r->chunks()) {
-        uint8_t* dest = snap.data + chunk->offset();
+        uint8_t* dest = snap->data + chunk->offset();
 
         SPDLOG_TRACE("Applying snapshot diff to {} at {}-{}",
                      r->key()->str(),

tests/dist/DistTestExecutor.cpp

@@ -66,12 +66,12 @@ void DistTestExecutor::restore(faabric::Message& msg)
     // Initialise the dummy memory and map to snapshot
     faabric::snapshot::SnapshotRegistry& reg =
       faabric::snapshot::getSnapshotRegistry();
-    faabric::util::SnapshotData& snap = reg.getSnapshot(msg.snapshotkey());
+    auto snap = reg.getSnapshot(msg.snapshotkey());
 
     // Note this has to be mmapped to be page-aligned
     snapshotMemory = (uint8_t*)mmap(
-      nullptr, snap.size, PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
-    snapshotSize = snap.size;
+      nullptr, snap->size, PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    snapshotSize = snap->size;
 
     reg.mapSnapshot(msg.snapshotkey(), snapshotMemory);
 }

tests/dist/scheduler/functions.cpp

@@ -67,15 +67,15 @@ int handleFakeDiffsFunction(faabric::scheduler::Executor* exec,
     faabric::snapshot::SnapshotRegistry& reg =
       faabric::snapshot::getSnapshotRegistry();
 
-    faabric::util::SnapshotData& originalSnap = reg.getSnapshot(snapshotKey);
+    auto originalSnap = reg.getSnapshot(snapshotKey);
     faabric::util::SnapshotData updatedSnap = exec->snapshot();
 
     // Add a single merge region to catch both diffs
     int offsetA = 10;
     int offsetB = 100;
     std::vector<uint8_t> inputBytes = faabric::util::stringToBytes(msgInput);
 
-    originalSnap.addMergeRegion(
+    originalSnap->addMergeRegion(
       0,
       offsetB + inputBytes.size() + 10,
       faabric::util::SnapshotDataType::Raw,
@@ -209,13 +209,12 @@ int handleFakeDiffsThreadedFunction(
         std::vector<uint8_t> inputBytes =
           faabric::util::stringToBytes(msgInput);
 
-        faabric::util::SnapshotData& originalSnap =
-          reg.getSnapshot(snapshotKey);
+        auto originalSnap = reg.getSnapshot(snapshotKey);
         faabric::util::SnapshotData updatedSnap = exec->snapshot();
 
         // Make sure it's captured by the region
         int regionLength = 20 + inputBytes.size();
-        originalSnap.addMergeRegion(
+        originalSnap->addMergeRegion(
           regionOffset,
           regionLength,
           faabric::util::SnapshotDataType::Raw,