Snapshots overhaul

include/faabric/scheduler/Scheduler.h

@@ -69,7 +69,7 @@ class Executor
 
     void releaseClaim();
 
-    virtual faabric::util::SnapshotData snapshot();
+    virtual faabric::util::SnapshotData& snapshot();
 
   protected:
     virtual void restore(faabric::Message& msg);
@@ -80,6 +80,8 @@ class Executor
 
     uint32_t threadPoolSize = 0;
 
+    faabric::util::SnapshotData _snapshot;
+
   private:
     std::atomic<bool> claimed = false;
 

include/faabric/snapshot/SnapshotClient.h

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

include/faabric/snapshot/SnapshotRegistry.h

@@ -23,15 +23,17 @@ class SnapshotRegistry
     void mapSnapshot(const std::string& key, uint8_t* target);
 
     void takeSnapshot(const std::string& key,
-                      faabric::util::SnapshotData data,
+                      faabric::util::SnapshotData& data,
                       bool locallyRestorable = true);
 
     void takeSnapshotIfNotExists(const std::string& key,
-                                 faabric::util::SnapshotData data,
+                                 faabric::util::SnapshotData& data,
                                  bool locallyRestorable = true);
 
     void deleteSnapshot(const std::string& key);
 
+    void changeSnapshotSize(const std::string& key, size_t newSize);
+
     size_t getSnapshotCount();
 
     void clear();
@@ -47,7 +49,7 @@ class SnapshotRegistry
                           faabric::util::SnapshotData& data);
 
     void doTakeSnapshot(const std::string& key,
-                        faabric::util::SnapshotData data,
+                        const faabric::util::SnapshotData& data,
                         bool locallyRestorable,
                         bool overwrite);
 };

include/faabric/util/memory.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstdint>
+#include <string>
 #include <unistd.h>
 #include <vector>
 
@@ -41,4 +42,25 @@ 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
+// -------------------------
+
+void deallocateMemory(uint8_t* memory, size_t size);
+
+void deallocatePages(uint8_t* memory, int nPages);
+
+uint8_t* allocateSharedMemory(size_t size);
+
+uint8_t* allocatePrivateMemory(size_t size);
+
+uint8_t* allocateVirtualMemory(size_t size);
+
+void claimVirtualMemory(uint8_t* start, size_t size);
+
+void mapMemory(uint8_t* target, size_t size, int fd);
+
+int writeMemoryToFd(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

@@ -9,6 +9,8 @@
 #include <faabric/util/logging.h>
 #include <faabric/util/macros.h>
 
+#define FOUR_GB (size_t)(1024L * 1024L * 1024L * 4L)
+
 namespace faabric::util {
 
 enum SnapshotDataType
@@ -73,13 +75,20 @@ class SnapshotMergeRegion
 class SnapshotData
 {
   public:
-    size_t size = 0;
     uint8_t* data = nullptr;
-    int fd = 0;
+    size_t size = 0;
 
     SnapshotData() = default;
 
-    std::vector<SnapshotDiff> getDirtyPages();
+    SnapshotData(const SnapshotData&);
+
+    SnapshotData& operator=(const SnapshotData&) = delete;
+
+    ~SnapshotData();
+
+    bool isRestorable();
+
+    std::vector<SnapshotDiff> getDirtyRegions();
 
     std::vector<SnapshotDiff> getChangeDiffs(const uint8_t* updated,
                                              size_t updatedSize);
@@ -90,7 +99,22 @@ class SnapshotData
                         SnapshotMergeOperation operation,
                         bool overwrite = false);
 
+    void clearMergeRegions();
+
+    void setSnapshotSize(size_t newSize);
+
+    void updateFd();
+
+    void writeToFd(const std::string& fdLabel);
+
+    void mapToMemory(uint8_t* target);
+
+    std::map<uint32_t, SnapshotMergeRegion> getMergeRegions();
+
   private:
+    int fd = 0;
+    size_t fdSize = 0;
+
     // 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,8 @@ 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;
+    // 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
@@ -299,6 +300,10 @@ void Executor::threadPoolThread(int threadPoolIdx)
 
             // 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 +408,10 @@ void Executor::postFinish() {}
 
 void Executor::reset(faabric::Message& msg) {}
 
-faabric::util::SnapshotData Executor::snapshot()
+faabric::util::SnapshotData& Executor::snapshot()
 {
     SPDLOG_WARN("Executor has not implemented snapshot method");
-    faabric::util::SnapshotData d;
-    return d;
+    return _snapshot;
 }
 
 void Executor::restore(faabric::Message& msg)

src/scheduler/Scheduler.cpp

@@ -504,7 +504,7 @@ 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 {
@@ -515,11 +515,11 @@ faabric::util::SchedulingDecision Scheduler::doCallFunctions(
     }
 
     // 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

src/snapshot/SnapshotClient.cpp

@@ -26,7 +26,7 @@ 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, faabric::util::SnapshotData>>&
 getSnapshotPushes()
 {
     faabric::util::UniqueLock lock(mockMutex);
@@ -73,7 +73,8 @@ SnapshotClient::SnapshotClient(const std::string& hostIn)
 
 void SnapshotClient::pushSnapshot(const std::string& key,
                                   int groupId,
-                                  const faabric::util::SnapshotData& data)
+                                  const faabric::util::SnapshotData& data,
+                                  size_t maxSize)
 {
     if (data.size == 0) {
         SPDLOG_ERROR("Cannot push snapshot {} with size zero to {}", key, host);
@@ -84,15 +85,16 @@ void SnapshotClient::pushSnapshot(const std::string& key,
 
     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 requestOffset = CreateSnapshotPushRequest(
+          mb, keyOffset, groupId, maxSize, dataOffset);
         mb.Finish(requestOffset);
 
         // Send it

src/snapshot/SnapshotRegistry.cpp

@@ -33,45 +33,26 @@ bool SnapshotRegistry::snapshotExists(const std::string& key)
 void SnapshotRegistry::mapSnapshot(const std::string& key, uint8_t* target)
 {
     auto d = getSnapshot(key);
-
-    if (!faabric::util::isPageAligned((void*)target)) {
-        SPDLOG_ERROR(
-          "Mapping snapshot {} to non page-aligned address {}", key, target);
-        throw std::runtime_error(
-          "Mapping snapshot to non page-aligned address");
-    }
-
-    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);
-
-    if (mmapRes == MAP_FAILED) {
-        SPDLOG_ERROR(
-          "mmapping snapshot failed: {} ({})", errno, ::strerror(errno));
-        throw std::runtime_error("mmapping snapshot failed");
-    }
+    d->mapToMemory(target);
 }
 
-void SnapshotRegistry::takeSnapshotIfNotExists(const std::string& key,
-                                               faabric::util::SnapshotData data,
-                                               bool locallyRestorable)
+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,
+                                    faabric::util::SnapshotData& data,
                                     bool locallyRestorable)
 {
     doTakeSnapshot(key, data, locallyRestorable, true);
 }
 
 void SnapshotRegistry::doTakeSnapshot(const std::string& key,
-                                      faabric::util::SnapshotData data,
+                                      const faabric::util::SnapshotData& data,
                                       bool locallyRestorable,
                                       bool overwrite)
 {
@@ -94,13 +75,13 @@ 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
-    auto shared_data =
+    auto sharedData =
       std::make_shared<faabric::util::SnapshotData>(std::move(data));
-    snapshotMap[key] = shared_data;
+    snapshotMap[key] = sharedData;
 
     // Write to fd to be locally restorable
     if (locallyRestorable) {
-        writeSnapshotToFd(key, *shared_data);
+        sharedData->writeToFd(key);
     }
 }
 
@@ -112,17 +93,16 @@ void SnapshotRegistry::deleteSnapshot(const std::string& key)
         return;
     }
 
-    auto d = snapshotMap[key];
+    snapshotMap.erase(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);
-        d->fd = 0;
-    }
+void SnapshotRegistry::changeSnapshotSize(const std::string& key,
+                                          size_t newSize)
+{
+    faabric::util::FullLock lock(snapshotsMx);
 
-    snapshotMap.erase(key);
+    auto d = getSnapshot(key);
+    d->setSnapshotSize(newSize);
 }
 
 size_t SnapshotRegistry::getSnapshotCount()
@@ -140,38 +120,6 @@ SnapshotRegistry& getSnapshotRegistry()
 void SnapshotRegistry::clear()
 {
     faabric::util::FullLock lock(snapshotsMx);
-    for (auto p : snapshotMap) {
-        if (p.second->fd > 0) {
-            ::close(p.second->fd);
-        }
-    }
-
     snapshotMap.clear();
 }
-
-int SnapshotRegistry::writeSnapshotToFd(const std::string& key,
-                                        faabric::util::SnapshotData& data)
-{
-    int fd = ::memfd_create(key.c_str(), 0);
-
-    // Make the fd big enough
-    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, 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
-    data.fd = fd;
-
-    SPDLOG_DEBUG("Wrote snapshot {} to fd {}", key, fd);
-    return fd;
-}
 }