Snapshots overhaul

.github/workflows/sanitisers.yml

@@ -99,7 +99,7 @@ jobs:
             ${{ runner.os }}-
       # --- Build with thread sanitising
       - name: "Build tests with address sanitising"
-        run: inv dev.sanitise Address --noclean
+        run: inv dev.sanitise Address faabric_tests --noclean
       - name: "Run tests with address sanitising on and print to stderr to file"
         run: ./bin/faabric_tests
         working-directory: /build/faabric/static
@@ -149,7 +149,7 @@ jobs:
             ${{ runner.os }}-
       # --- Build with thread sanitising
       - name: "Build tests with thread sanitising"
-        run: inv dev.sanitise Thread --noclean
+        run: inv dev.sanitise Thread faabric_tests --noclean
       # --- Tests ---
       - name: "Run tests with thread sanitising on and print stderr to file"
         run: ./bin/faabric_tests
@@ -200,7 +200,7 @@ jobs:
             ${{ runner.os }}-
       # --- Build with thread sanitising
       - name: "Build tests with undefined sanitising"
-        run: inv dev.sanitise Undefined --noclean
+        run: inv dev.sanitise Undefined faabric_tests --noclean
       # --- Tests ---
       - name: "Run tests with undefined sanitising on and print stderr to file"
         run: ./bin/faabric_tests
@@ -251,7 +251,7 @@ jobs:
             ${{ runner.os }}-
       # --- Build with thread sanitising
       - name: "Build tests with leak sanitising"
-        run: inv dev.sanitise Leak --noclean
+        run: inv dev.sanitise Leak faabric_tests --noclean
       # --- Tests ---
       - name: "Run tests with leak sanitising on and print stderr to file"
         run: ./bin/faabric_tests

include/faabric/scheduler/Scheduler.h

@@ -69,7 +69,7 @@ class Executor
 
     void releaseClaim();
 
-    virtual faabric::util::SnapshotData snapshot();
+    virtual std::shared_ptr<faabric::util::SnapshotData> snapshot();
 
   protected:
     virtual void restore(faabric::Message& msg);

include/faabric/snapshot/SnapshotClient.h

@@ -12,7 +12,8 @@ namespace faabric::snapshot {
 // Mocking
 // -----------------------------------
 
-std::vector<std::pair<std::string, faabric::util::SnapshotData>>
+std::vector<
+  std::pair<std::string, std::shared_ptr<faabric::util::SnapshotData>>>
 getSnapshotPushes();
 
 std::vector<std::pair<std::string, std::vector<faabric::util::SnapshotDiff>>>
@@ -36,7 +37,7 @@ class SnapshotClient final : public faabric::transport::MessageEndpointClient
 
     void pushSnapshot(const std::string& key,
                       int32_t groupId,
-                      const faabric::util::SnapshotData& data);
+                      std::shared_ptr<faabric::util::SnapshotData> data);
 
     void pushSnapshotDiffs(std::string snapshotKey,
                            int32_t groupId,

include/faabric/snapshot/SnapshotRegistry.h

@@ -22,13 +22,12 @@ class SnapshotRegistry
 
     void mapSnapshot(const std::string& key, uint8_t* target);
 
-    void takeSnapshot(const std::string& key,
-                      faabric::util::SnapshotData data,
-                      bool locallyRestorable = true);
+    void registerSnapshot(const std::string& key,
+                          std::shared_ptr<faabric::util::SnapshotData> data);
 
-    void takeSnapshotIfNotExists(const std::string& key,
-                                 faabric::util::SnapshotData data,
-                                 bool locallyRestorable = true);
+    void registerSnapshotIfNotExists(
+      const std::string& key,
+      std::shared_ptr<faabric::util::SnapshotData> data);
 
     void deleteSnapshot(const std::string& key);
 
@@ -46,10 +45,9 @@ class SnapshotRegistry
     int writeSnapshotToFd(const std::string& key,
                           faabric::util::SnapshotData& data);
 
-    void doTakeSnapshot(const std::string& key,
-                        faabric::util::SnapshotData data,
-                        bool locallyRestorable,
-                        bool overwrite);
+    void doRegisterSnapshot(const std::string& key,
+                            std::shared_ptr<faabric::util::SnapshotData> data,
+                            bool overwrite);
 };
 
 SnapshotRegistry& getSnapshotRegistry();

include/faabric/util/memory.h

@@ -1,6 +1,9 @@
 #pragma once
 
 #include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
 #include <unistd.h>
 #include <vector>
 
@@ -41,4 +44,26 @@ std::vector<int> getDirtyPageNumbers(const uint8_t* ptr, int nPages);
 
 std::vector<std::pair<uint32_t, uint32_t>> getDirtyRegions(const uint8_t* ptr,
                                                            int nPages);
+// -------------------------
+// Allocation
+// -------------------------
+
+typedef std::unique_ptr<uint8_t[], std::function<void(uint8_t*)>>
+  OwnedMmapRegion;
+
+OwnedMmapRegion allocateSharedMemory(size_t size);
+
+OwnedMmapRegion allocatePrivateMemory(size_t size);
+
+OwnedMmapRegion allocateVirtualMemory(size_t size);
+
+void claimVirtualMemory(uint8_t* start, size_t size);
+
+void mapMemory(uint8_t* target, size_t size, int fd);
+
+int writeMemoryToFd(const uint8_t* source,
+                    size_t size,
+                    const std::string& fdLabel);
+
+void appendDataToFd(int fd, size_t oldSize, size_t newSize, uint8_t* newData);
 }

include/faabric/util/snapshot.h

@@ -3,11 +3,13 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <shared_mutex>
 #include <string>
 #include <vector>
 
 #include <faabric/util/logging.h>
 #include <faabric/util/macros.h>
+#include <faabric/util/memory.h>
 
 namespace faabric::util {
 
@@ -74,12 +76,45 @@ class SnapshotData
 {
   public:
     size_t size = 0;
-    uint8_t* data = nullptr;
-    int fd = 0;
+    size_t maxSize = 0;
 
     SnapshotData() = default;
 
-    std::vector<SnapshotDiff> getDirtyPages();
+    explicit SnapshotData(size_t sizeIn);
+
+    explicit SnapshotData(std::vector<uint8_t> dataIn);
+
+    SnapshotData(size_t sizeIn, size_t maxSizeIn);
+
+    SnapshotData(uint8_t* dataIn, size_t sizeIn);
+
+    SnapshotData(uint8_t* dataIn, size_t sizeIn, size_t maxSizeIn);
+
+    SnapshotData(const SnapshotData&) = delete;
+
+    SnapshotData& operator=(const SnapshotData&) = delete;
+
+    ~SnapshotData();
+
+    bool isRestorable();
+
+    void makeRestorable(const std::string& fdLabel);
+
+    void setSnapshotSize(size_t newSize);
+
+    void copyInData(std::vector<uint8_t> buffer, uint32_t offset = 0);
+
+    void copyInData(uint8_t* buffer, size_t bufferSize, uint32_t offset = 0);
+
+    const uint8_t* getDataPtr(uint32_t offset = 0);
+
+    uint8_t* getMutableDataPtr(uint32_t offset = 0);
+
+    std::vector<uint8_t> getDataCopy();
+
+    std::vector<uint8_t> getDataCopy(uint32_t offset, size_t dataSize);
+
+    std::vector<SnapshotDiff> getDirtyRegions();
 
     std::vector<SnapshotDiff> getChangeDiffs(const uint8_t* updated,
                                              size_t updatedSize);
@@ -90,7 +125,22 @@ class SnapshotData
                         SnapshotMergeOperation operation,
                         bool overwrite = false);
 
+    void clearMergeRegions();
+
+    void mapToMemory(uint8_t* target);
+
+    std::map<uint32_t, SnapshotMergeRegion> getMergeRegions();
+
   private:
+    int fd = 0;
+    size_t fdSize = 0;
+
+    std::shared_mutex snapMx;
+
+    bool owner = false;
+    uint8_t* rawData = nullptr;
+    OwnedMmapRegion ownedData = nullptr;
+
     // Note - we care about the order of this map, as we iterate through it
     // in order of offsets
     std::map<uint32_t, SnapshotMergeRegion> mergeRegions;

src/flat/faabric.fbs

@@ -1,6 +1,7 @@
 table SnapshotPushRequest {
   key:string;
   groupid:int;
+  maxSize:ulong;
   contents:[ubyte];
 }
 

src/scheduler/Executor.cpp

@@ -144,7 +144,10 @@ void Executor::executeTasks(std::vector<int> msgIdxs,
     // Work out if we should skip the reset after this batch. This only needs to
     // happen when we're executing threads on the master host, in which case the
     // original function call will cause a reset
-    bool skipReset = isMaster && isThreads;
+
+    // TODO - should this be on master or all hosts?
+    // bool skipReset = isMaster && isThreads;
+    bool skipReset = isThreads;
 
     // Iterate through and invoke tasks. By default, we allocate tasks
     // one-to-one with thread pool threads. Only once the pool is exhausted do
@@ -292,13 +295,18 @@ void Executor::threadPoolThread(int threadPoolIdx)
                          msg.snapshotkey());
 
             std::vector<faabric::util::SnapshotDiff> diffs =
-              snapshotPreExecution->getChangeDiffs(snapshotPostExecution.data,
-                                                   snapshotPostExecution.size);
+              snapshotPreExecution->getChangeDiffs(
+                snapshotPostExecution->getDataPtr(),
+                snapshotPostExecution->size);
 
             sch.pushSnapshotDiffs(msg, diffs);
 
             // Reset dirty page tracking now that we've pushed the diffs
             faabric::util::resetDirtyTracking();
+
+            // Clear any merge regions
+            SPDLOG_DEBUG("Clearing merge regions for {}", msg.snapshotkey());
+            snapshotPreExecution->clearMergeRegions();
         }
 
         // If this batch is finished, reset the executor and release its claim.
@@ -403,11 +411,10 @@ void Executor::postFinish() {}
 
 void Executor::reset(faabric::Message& msg) {}
 
-faabric::util::SnapshotData Executor::snapshot()
+std::shared_ptr<faabric::util::SnapshotData> Executor::snapshot()
 {
     SPDLOG_WARN("Executor has not implemented snapshot method");
-    faabric::util::SnapshotData d;
-    return d;
+    return nullptr;
 }
 
 void Executor::restore(faabric::Message& msg)

src/scheduler/Scheduler.cpp

@@ -83,8 +83,7 @@ void Scheduler::addHostToGlobalSet()
 
 void Scheduler::resetThreadLocalCache()
 {
-    auto tid = (pid_t)syscall(SYS_gettid);
-    SPDLOG_DEBUG("Resetting scheduler thread-local cache for thread {}", tid);
+    SPDLOG_DEBUG("Resetting scheduler thread-local cache");
 
     functionCallClients.clear();
     snapshotClients.clear();
@@ -504,22 +503,22 @@ faabric::util::SchedulingDecision Scheduler::doCallFunctions(
             // if so, just push the diffs
             if (pushedSnapshotsMap[snapshotKey].contains(host)) {
                 std::vector<faabric::util::SnapshotDiff> snapshotDiffs =
-                  snapshotData->getDirtyPages();
+                  snapshotData->getDirtyRegions();
                 c.pushSnapshotDiffs(
                   snapshotKey, firstMsg.groupid(), snapshotDiffs);
             } else {
-                c.pushSnapshot(snapshotKey, firstMsg.groupid(), *snapshotData);
+                c.pushSnapshot(snapshotKey, firstMsg.groupid(), snapshotData);
                 pushedSnapshotsMap[snapshotKey].insert(host);
             }
         }
     }
 
     // Now reset the dirty page tracking just before we start executing
-    SPDLOG_DEBUG("Resetting dirty tracking after pushing diffs {}", funcStr);
+    SPDLOG_DEBUG("Resetting dirty tracking before executing {}", funcStr);
     faabric::util::resetDirtyTracking();
 
     // -------------------------------------------
-    // EXECTUION
+    // EXECUTION
     // -------------------------------------------
 
     // Records for tests - copy messages before execution to avoid racing on msg
@@ -785,9 +784,11 @@ std::shared_ptr<Executor> Scheduler::claimExecutor(
 
     // We have no warm executors available, so scale up
     if (claimed == nullptr) {
-        int nExecutors = thisExecutors.size();
-        SPDLOG_DEBUG(
-          "Scaling {} from {} -> {}", funcStr, nExecutors, nExecutors + 1);
+        SPDLOG_DEBUG("Scaling {} from {} -> {}",
+                     funcStr,
+                     thisExecutors.size(),
+                     thisExecutors.size() + 1);
+
         // Spinning up a new executor can be lengthy, allow other things
         // to run in parallel
         schedulerLock.unlock();

src/snapshot/SnapshotClient.cpp

@@ -14,7 +14,8 @@ namespace faabric::snapshot {
 
 static std::mutex mockMutex;
 
-static std::vector<std::pair<std::string, faabric::util::SnapshotData>>
+static std::vector<
+  std::pair<std::string, std::shared_ptr<faabric::util::SnapshotData>>>
   snapshotPushes;
 
 static std::vector<
@@ -26,7 +27,8 @@ static std::vector<std::pair<std::string, std::string>> snapshotDeletes;
 static std::vector<std::pair<std::string, std::pair<uint32_t, int>>>
   threadResults;
 
-std::vector<std::pair<std::string, faabric::util::SnapshotData>>
+std::vector<
+  std::pair<std::string, std::shared_ptr<faabric::util::SnapshotData>>>
 getSnapshotPushes()
 {
     faabric::util::UniqueLock lock(mockMutex);
@@ -71,28 +73,31 @@ SnapshotClient::SnapshotClient(const std::string& hostIn)
                                               SNAPSHOT_SYNC_PORT)
 {}
 
-void SnapshotClient::pushSnapshot(const std::string& key,
-                                  int groupId,
-                                  const faabric::util::SnapshotData& data)
+void SnapshotClient::pushSnapshot(
+  const std::string& key,
+  int groupId,
+  std::shared_ptr<faabric::util::SnapshotData> data)
 {
-    if (data.size == 0) {
+    if (data->size == 0) {
         SPDLOG_ERROR("Cannot push snapshot {} with size zero to {}", key, host);
         throw std::runtime_error("Pushing snapshot with zero size");
     }
 
-    SPDLOG_DEBUG("Pushing snapshot {} to {} ({} bytes)", key, host, data.size);
+    SPDLOG_DEBUG("Pushing snapshot {} to {} ({} bytes)", key, host, data->size);
 
     if (faabric::util::isMockMode()) {
         faabric::util::UniqueLock lock(mockMutex);
+
         snapshotPushes.emplace_back(host, data);
     } else {
         // Set up the main request
         // TODO - avoid copying data here
         flatbuffers::FlatBufferBuilder mb;
         auto keyOffset = mb.CreateString(key);
-        auto dataOffset = mb.CreateVector<uint8_t>(data.data, data.size);
-        auto requestOffset =
-          CreateSnapshotPushRequest(mb, keyOffset, groupId, dataOffset);
+        auto dataOffset =
+          mb.CreateVector<uint8_t>(data->getDataPtr(), data->size);
+        auto requestOffset = CreateSnapshotPushRequest(
+          mb, keyOffset, groupId, data->maxSize, dataOffset);
         mb.Finish(requestOffset);
 
         // Send it