Snapshots overhaul

include/faabric/scheduler/Scheduler.h

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

include/faabric/util/snapshot.h

@@ -40,13 +40,13 @@ class SnapshotDiff
                  uint32_t offsetIn,
                  std::span<const uint8_t> dataIn);
 
-    SnapshotDataType getDataType() const;
+    SnapshotDataType getDataType() const { return dataType; }
 
-    SnapshotMergeOperation getOperation() const;
+    SnapshotMergeOperation getOperation() const { return operation; }
 
-    uint32_t getOffset() const;
+    uint32_t getOffset() const { return offset; }
 
-    std::span<const uint8_t> getData() const;
+    std::span<const uint8_t> getData() const { return data; }
 
     std::vector<uint8_t> getDataCopy() const;
 
@@ -114,13 +114,13 @@ class SnapshotData
 
     size_t getQueuedDiffsCount();
 
-    void queueDiffs(const std::vector<SnapshotDiff>& diffs);
+    void queueDiffs(std::span<SnapshotDiff> diffs);
 
     void writeQueuedDiffs();
 
-    size_t getSize();
+    size_t getSize() const { return size; }
 
-    size_t getMaxSize();
+    size_t getMaxSize() const { return maxSize; }
 
   private:
     size_t size = 0;
@@ -150,18 +150,16 @@ class MemoryView
   public:
     // Note - this object is just a view of a section of memory, and does not
     // own the underlying data
-    explicit MemoryView(std::span<const uint8_t> dataIn);
-
-    MemoryView(const MemoryView&) = delete;
+    MemoryView() = default;
 
-    MemoryView& operator=(const MemoryView&) = delete;
+    explicit MemoryView(std::span<const uint8_t> dataIn);
 
     std::vector<SnapshotDiff> getDirtyRegions();
 
     std::vector<SnapshotDiff> diffWithSnapshot(
       std::shared_ptr<SnapshotData> snap);
 
-    std::span<const uint8_t> getData();
+    std::span<const uint8_t> getData() { return data; }
 
   private:
     std::span<const uint8_t> data;

src/scheduler/Executor.cpp

@@ -289,8 +289,9 @@ void Executor::threadPoolThread(int threadPoolIdx)
                      oldTaskCount - 1);
 
         // Handle snapshot diffs _before_ we reset the executor
-        std::shared_ptr<faabric::util::MemoryView> funcMemory = getMemoryView();
-        if (funcMemory != nullptr && isLastInBatch && task.needsSnapshotSync) {
+        faabric::util::MemoryView funcMemory = getMemoryView();
+        if (!funcMemory.getData().empty() && isLastInBatch &&
+            task.needsSnapshotSync) {
             auto snap = faabric::snapshot::getSnapshotRegistry().getSnapshot(
               msg.snapshotkey());
 
@@ -300,7 +301,7 @@ void Executor::threadPoolThread(int threadPoolIdx)
             // If we're on master, we write the diffs straight to the snapshot
             // otherwise we push them to the master.
             std::vector<faabric::util::SnapshotDiff> diffs =
-              funcMemory->diffWithSnapshot(snap);
+              funcMemory.diffWithSnapshot(snap);
 
             if (isMaster) {
                 SPDLOG_DEBUG("Queueing {} diffs for {} to snapshot {} on "
@@ -424,11 +425,11 @@ void Executor::postFinish() {}
 
 void Executor::reset(faabric::Message& msg) {}
 
-std::shared_ptr<faabric::util::MemoryView> Executor::getMemoryView()
+faabric::util::MemoryView Executor::getMemoryView()
 {
     SPDLOG_WARN("Executor for {} has not implemented memory view method",
                 faabric::util::funcToString(boundMessage, false));
-    return nullptr;
+    return faabric::util::MemoryView();
 }
 
 void Executor::restore(faabric::Message& msg)

src/snapshot/SnapshotRegistry.cpp

@@ -42,14 +42,14 @@ void SnapshotRegistry::registerSnapshotIfNotExists(
   const std::string& key,
   std::shared_ptr<faabric::util::SnapshotData> data)
 {
-    doRegisterSnapshot(key, data, false);
+    doRegisterSnapshot(key, std::move(data), false);
 }
 
 void SnapshotRegistry::registerSnapshot(
   const std::string& key,
   std::shared_ptr<faabric::util::SnapshotData> data)
 {
-    doRegisterSnapshot(key, data, true);
+    doRegisterSnapshot(key, std::move(data), true);
 }
 
 void SnapshotRegistry::doRegisterSnapshot(

src/snapshot/SnapshotServer.cpp

@@ -123,6 +123,7 @@ SnapshotServer::recvPushSnapshotDiffs(const uint8_t* buffer, size_t bufferSize)
 
     // Convert chunks to snapshot diff objects
     std::vector<SnapshotDiff> diffs;
+    diffs.reserve(r->chunks()->size());
     for (const auto* chunk : *r->chunks()) {
         diffs.emplace_back(
           static_cast<SnapshotDataType>(chunk->dataType()),

src/util/snapshot.cpp

@@ -20,26 +20,6 @@ SnapshotDiff::SnapshotDiff(SnapshotDataType dataTypeIn,
   , data(dataIn.begin(), dataIn.end())
 {}
 
-SnapshotDataType SnapshotDiff::getDataType() const
-{
-    return dataType;
-}
-
-SnapshotMergeOperation SnapshotDiff::getOperation() const
-{
-    return operation;
-}
-
-uint32_t SnapshotDiff::getOffset() const
-{
-    return offset;
-}
-
-std::span<const uint8_t> SnapshotDiff::getData() const
-{
-    return data;
-}
-
 std::vector<uint8_t> SnapshotDiff::getDataCopy() const
 {
     return std::vector<uint8_t>(data.begin(), data.end());
@@ -238,7 +218,7 @@ size_t SnapshotData::getQueuedDiffsCount()
     return queuedDiffs.size();
 }
 
-void SnapshotData::queueDiffs(const std::vector<SnapshotDiff>& diffs)
+void SnapshotData::queueDiffs(const std::span<SnapshotDiff> diffs)
 {
     faabric::util::FullLock lock(snapMx);
     for (const auto& diff : diffs) {
@@ -351,25 +331,14 @@ void SnapshotData::writeQueuedDiffs()
     queuedDiffs.clear();
 }
 
-size_t SnapshotData::getSize()
-{
-    return size;
-}
-
-size_t SnapshotData::getMaxSize()
-{
-    return maxSize;
-}
-
 MemoryView::MemoryView(std::span<const uint8_t> dataIn)
   : data(dataIn)
 {}
 
 std::vector<SnapshotDiff> MemoryView::getDirtyRegions()
 {
     if (data.empty()) {
-        std::vector<SnapshotDiff> empty;
-        return empty;
+        return {};
     }
 
     // Get dirty regions
@@ -382,13 +351,12 @@ std::vector<SnapshotDiff> MemoryView::getDirtyRegions()
 
     // Convert to snapshot diffs
     std::vector<SnapshotDiff> diffs;
-    for (auto& p : regions) {
-        size_t regionSize = p.second - p.first;
-        const uint8_t* regionData = data.data() + p.first;
+    diffs.reserve(regions.size());
+    for (auto [regionBegin, regionEnd] : regions) {
         diffs.emplace_back(SnapshotDataType::Raw,
                            SnapshotMergeOperation::Overwrite,
-                           p.first,
-                           std::span<const uint8_t>(regionData, regionSize));
+                           regionBegin,
+                           data.subspan(regionBegin, regionEnd - regionBegin));
     }
 
     SPDLOG_DEBUG("Memory view has {}/{} dirty pages", diffs.size(), nPages);
@@ -441,11 +409,6 @@ std::vector<SnapshotDiff> MemoryView::diffWithSnapshot(
     return diffs;
 }
 
-std::span<const uint8_t> MemoryView::getData()
-{
-    return data;
-}
-
 std::string snapshotDataTypeStr(SnapshotDataType dt)
 {
     switch (dt) {

tests/dist/DistTestExecutor.cpp

@@ -75,10 +75,9 @@ void DistTestExecutor::restore(faabric::Message& msg)
     reg.mapSnapshot(msg.snapshotkey(), dummyMemory.get());
 }
 
-std::shared_ptr<faabric::util::MemoryView> DistTestExecutor::getMemoryView()
+faabric::util::MemoryView DistTestExecutor::getMemoryView()
 {
-    return std::make_shared<faabric::util::MemoryView>(
-      std::span<uint8_t>(dummyMemory.get(), dummyMemorySize));
+    return faabric::util::MemoryView({ dummyMemory.get(), dummyMemorySize });
 }
 
 std::span<uint8_t> DistTestExecutor::getDummyMemory()

tests/dist/DistTestExecutor.h

@@ -23,7 +23,7 @@ class DistTestExecutor final : public faabric::scheduler::Executor
 
     void restore(faabric::Message& msg) override;
 
-    std::shared_ptr<faabric::util::MemoryView> getMemoryView() override;
+    faabric::util::MemoryView getMemoryView() override;
 
     std::span<uint8_t> getDummyMemory();
 

tests/dist/scheduler/functions.cpp

@@ -71,8 +71,6 @@ int handleFakeDiffsFunction(tests::DistTestExecutor* exec,
       faabric::snapshot::getSnapshotRegistry();
 
     auto originalSnap = reg.getSnapshot(snapshotKey);
-    std::shared_ptr<faabric::util::MemoryView> funcMemory =
-      exec->getMemoryView();
 
     // Add a single merge region to catch both diffs
     int offsetA = 10;

tests/test/scheduler/test_executor.cpp

@@ -53,10 +53,9 @@ void TestExecutor::restore(faabric::Message& msg)
     reg.mapSnapshot(msg.snapshotkey(), dummyMemory.get());
 }
 
-std::shared_ptr<faabric::util::MemoryView> TestExecutor::getMemoryView()
+faabric::util::MemoryView TestExecutor::getMemoryView()
 {
-    return std::make_shared<faabric::util::MemoryView>(
-      std::span<uint8_t>(dummyMemory.get(), dummyMemorySize));
+    return faabric::util::MemoryView({ dummyMemory.get(), dummyMemorySize });
 }
 
 int32_t TestExecutor::executeTask(

tests/utils/fixtures.h

@@ -305,7 +305,7 @@ class TestExecutor final : public faabric::scheduler::Executor
 
     void restore(faabric::Message& msg) override;
 
-    std::shared_ptr<faabric::util::MemoryView> getMemoryView() override;
+    faabric::util::MemoryView getMemoryView() override;
 
     int32_t executeTask(
       int threadPoolIdx,