Snapshots overhaul (#176)

* Start on changing snapshot sizes

* Resizing

* Up to failing test

* Tests for memory utils

* Fixing up tests

* Test for expanding and remapping snapshots

* Formatting

* Checking for uninitialised fds

* Test for pushing snapshots

* Test fixes, formatting

* Add target to dev.sanitise

* Half-way through refactor

* Small typos etc.

* Continued refactor

* Compiling

* Fixing up tests

* Fix tests

* Formatting

* Fixing up dist tests

* Fix failing asan test

* Small tidy-up

* Typos

* PR comments WIP

* Remove vector/ span duplicates

* More tests

* Adding private and shared snapshot mappings

* Compilation fixes

* Sketching out memory view class

* Follow through with MemoryView refactor

* Fixing tests

* Fixing tests

* Tidying up

* Move merging logic into SnapshotData

* Master writing back snapshot changes locally

* Failing distributed reduction test

* Formatting and tests

* Adding repeats into dist test

* Fix up distributed test

* Tidy-up

* Override CPU count in dist tests

* Fix race condition in master snapshot

* Added locking log statement

* Fix scheduler test

* Fix straggler error in tests

* Immutable snapshot diffs with ownership

* Fix up broken tests

* Fix compile error in dist tests

* Test fix-up

* Formatting

* Fix dist tests

* Rename _data param

* PR comments

Author: Shillaker

dist-test/dev_server.sh

@@ -5,6 +5,8 @@ THIS_DIR=$(dirname $(readlink -f $0))
 PROJ_ROOT=${THIS_DIR}/..
 pushd ${PROJ_ROOT} > /dev/null
 
+export OVERRIDE_CPU_COUNT=4
+
 if [[ -z "$1" ]]; then
     docker-compose up -d dist-test-server
 elif [[ "$1" == "restart" ]]; then

dist-test/run.sh

@@ -6,6 +6,8 @@ pushd ${PROJ_ROOT} >> /dev/null
 export CONAN_CACHE_MOUNT_SOURCE=$HOME/.conan/
 RETURN_VAL=0
 
+export OVERRIDE_CPU_COUNT=4
+
 # Run the test server in the background
 docker-compose \
     up \

docker-compose.yml

@@ -20,6 +20,10 @@ services:
       - LOG_LEVEL=debug
       - REDIS_STATE_HOST=redis
       - REDIS_QUEUE_HOST=redis
+      - OVERRIDE_CPU_COUNT=${OVERRIDE_CPU_COUNT:-0}
+      - ASAN_OPTIONS=verbosity=1:halt_on_error=1
+      - TSAN_OPTIONS=halt_on_error=1:suppressions=/code/faabric/thread-sanitizer-ignorelist.txt:history_size=7:second_deadlock_stack=1
+      - UBSAN_OPTIONS="print_stacktrace=1:halt_on_error=1
     depends_on:
       - redis
 
@@ -34,6 +38,7 @@ services:
       - LOG_LEVEL=debug
       - REDIS_STATE_HOST=redis
       - REDIS_QUEUE_HOST=redis
+      - OVERRIDE_CPU_COUNT=${OVERRIDE_CPU_COUNT:-0}
     command: ./bin/faabric_dist_test_server
     depends_on:
       - redis

include/faabric/scheduler/Scheduler.h

@@ -32,13 +32,13 @@ class ExecutorTask
     ExecutorTask(int messageIndexIn,
                  std::shared_ptr<faabric::BatchExecuteRequest> reqIn,
                  std::shared_ptr<std::atomic<int>> batchCounterIn,
-                 bool needsSnapshotPushIn,
+                 bool needsSnapshotSyncIn,
                  bool skipResetIn);
 
     std::shared_ptr<faabric::BatchExecuteRequest> req;
     std::shared_ptr<std::atomic<int>> batchCounter;
     int messageIndex = 0;
-    bool needsSnapshotPush = false;
+    bool needsSnapshotSync = false;
     bool skipReset = false;
 };
 
@@ -69,7 +69,7 @@ class Executor
 
     void releaseClaim();
 
-    virtual faabric::util::SnapshotData snapshot();
+    virtual faabric::util::MemoryView getMemoryView();
 
   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>>>
@@ -35,12 +36,12 @@ class SnapshotClient final : public faabric::transport::MessageEndpointClient
     explicit SnapshotClient(const std::string& hostIn);
 
     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,
-                           std::vector<faabric::util::SnapshotDiff> diffs);
+    void pushSnapshotDiffs(
+      std::string snapshotKey,
+      bool force,
+      const std::vector<faabric::util::SnapshotDiff>& diffs);
 
     void deleteSnapshot(const std::string& key);
 

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/snapshot/SnapshotServer.h

@@ -14,10 +14,6 @@ class SnapshotServer final : public faabric::transport::MessageEndpointServer
   public:
     SnapshotServer();
 
-    // Returns how many diffs have been applied since started, useful for
-    // testing
-    size_t diffsApplied() const;
-
   protected:
     void doAsyncRecv(int header,
                      const uint8_t* buffer,
@@ -40,6 +36,5 @@ class SnapshotServer final : public faabric::transport::MessageEndpointServer
 
   private:
     faabric::transport::PointToPointBroker& broker;
-    std::atomic_size_t diffsAppliedCounter = 0;
 };
 }

include/faabric/util/bytes.h

@@ -26,18 +26,18 @@ int safeCopyToBuffer(const uint8_t* dataIn,
                      int bufferLen);
 
 template<class T>
-T unalignedRead(const std::byte* bytes)
+T unalignedRead(const uint8_t* bytes)
 {
     T value;
-    std::copy_n(bytes, sizeof(T), reinterpret_cast<std::byte*>(&value));
+    std::copy_n(bytes, sizeof(T), reinterpret_cast<uint8_t*>(&value));
     return value;
 }
 
 template<class T>
-void unalignedWrite(const T& value, std::byte* destination)
+void unalignedWrite(const T& value, uint8_t* destination)
 {
     std::copy_n(
-      reinterpret_cast<const std::byte*>(&value), sizeof(T), destination);
+      reinterpret_cast<const uint8_t*>(&value), sizeof(T), destination);
 }
 
 template<class T>

include/faabric/util/memory.h

@@ -1,6 +1,10 @@
 #pragma once
 
 #include <cstdint>
+#include <functional>
+#include <memory>
+#include <span>
+#include <string>
 #include <unistd.h>
 #include <vector>
 
@@ -22,7 +26,7 @@ struct AlignedChunk
 
 static const long HOST_PAGE_SIZE = sysconf(_SC_PAGESIZE);
 
-bool isPageAligned(void* ptr);
+bool isPageAligned(const void* ptr);
 
 size_t getRequiredHostPages(size_t nBytes);
 
@@ -41,4 +45,27 @@ 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*)>> MemoryRegion;
+
+MemoryRegion allocateSharedMemory(size_t size);
+
+MemoryRegion allocateVirtualMemory(size_t size);
+
+void claimVirtualMemory(std::span<uint8_t> region);
+
+void mapMemoryPrivate(std::span<uint8_t> target, int fd);
+
+void mapMemoryShared(std::span<uint8_t> target, int fd);
+
+void resizeFd(int fd, size_t size);
+
+void writeToFd(int fd, off_t offset, std::span<const uint8_t> data);
+
+int createFd(size_t size, const std::string& fdLabel);
+
+void appendDataToFd(int fd, std::span<uint8_t> data);
 }

include/faabric/util/snapshot.h

@@ -3,11 +3,14 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <shared_mutex>
+#include <span>
 #include <string>
 #include <vector>
 
 #include <faabric/util/logging.h>
 #include <faabric/util/macros.h>
+#include <faabric/util/memory.h>
 
 namespace faabric::util {
 
@@ -30,28 +33,28 @@ enum SnapshotMergeOperation
 class SnapshotDiff
 {
   public:
-    const uint8_t* data = nullptr;
-    size_t size = 0;
-    SnapshotDataType dataType = SnapshotDataType::Raw;
-    SnapshotMergeOperation operation = SnapshotMergeOperation::Overwrite;
-    uint32_t offset = 0;
-
-    bool noChange = false;
-
     SnapshotDiff() = default;
 
     SnapshotDiff(SnapshotDataType dataTypeIn,
                  SnapshotMergeOperation operationIn,
                  uint32_t offsetIn,
-                 const uint8_t* dataIn,
-                 size_t sizeIn)
-    {
-        dataType = dataTypeIn;
-        operation = operationIn;
-        offset = offsetIn;
-        data = dataIn;
-        size = sizeIn;
-    }
+                 std::span<const uint8_t> dataIn);
+
+    SnapshotDataType getDataType() const { return dataType; }
+
+    SnapshotMergeOperation getOperation() const { return operation; }
+
+    uint32_t getOffset() const { return offset; }
+
+    std::span<const uint8_t> getData() const { return data; }
+
+    std::vector<uint8_t> getDataCopy() const;
+
+  private:
+    SnapshotDataType dataType = SnapshotDataType::Raw;
+    SnapshotMergeOperation operation = SnapshotMergeOperation::Overwrite;
+    uint32_t offset = 0;
+    std::vector<uint8_t> data;
 };
 
 class SnapshotMergeRegion
@@ -73,27 +76,93 @@ class SnapshotMergeRegion
 class SnapshotData
 {
   public:
-    size_t size = 0;
-    uint8_t* data = nullptr;
-    int fd = 0;
-
     SnapshotData() = default;
 
-    std::vector<SnapshotDiff> getDirtyPages();
+    explicit SnapshotData(size_t sizeIn);
+
+    SnapshotData(size_t sizeIn, size_t maxSizeIn);
+
+    explicit SnapshotData(std::span<const uint8_t> dataIn);
+
+    SnapshotData(std::span<const uint8_t> dataIn, size_t maxSizeIn);
+
+    SnapshotData(const SnapshotData&) = delete;
+
+    SnapshotData& operator=(const SnapshotData&) = delete;
+
+    ~SnapshotData();
+
+    void copyInData(std::span<const uint8_t> buffer, uint32_t offset = 0);
+
+    const uint8_t* getDataPtr(uint32_t offset = 0);
+
+    std::vector<uint8_t> getDataCopy();
 
-    std::vector<SnapshotDiff> getChangeDiffs(const uint8_t* updated,
-                                             size_t updatedSize);
+    std::vector<uint8_t> getDataCopy(uint32_t offset, size_t dataSize);
+
+    void mapToMemory(uint8_t* target);
 
     void addMergeRegion(uint32_t offset,
                         size_t length,
                         SnapshotDataType dataType,
                         SnapshotMergeOperation operation,
                         bool overwrite = false);
 
+    void clearMergeRegions();
+
+    std::map<uint32_t, SnapshotMergeRegion> getMergeRegions();
+
+    size_t getQueuedDiffsCount();
+
+    void queueDiffs(std::span<SnapshotDiff> diffs);
+
+    void writeQueuedDiffs();
+
+    size_t getSize() const { return size; }
+
+    size_t getMaxSize() const { return maxSize; }
+
   private:
+    size_t size = 0;
+    size_t maxSize = 0;
+
+    int fd = -1;
+
+    std::shared_mutex snapMx;
+
+    MemoryRegion data = nullptr;
+
+    std::vector<SnapshotDiff> queuedDiffs;
+
     // Note - we care about the order of this map, as we iterate through it
     // in order of offsets
     std::map<uint32_t, SnapshotMergeRegion> mergeRegions;
+
+    uint8_t* validatedOffsetPtr(uint32_t offset);
+
+    void mapToMemory(uint8_t* target, bool shared);
+
+    void writeData(std::span<const uint8_t> buffer, uint32_t offset = 0);
+};
+
+class MemoryView
+{
+  public:
+    // Note - this object is just a view of a section of memory, and does not
+    // own the underlying data
+    MemoryView() = default;
+
+    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() { return data; }
+
+  private:
+    std::span<const uint8_t> data;
 };
 
 std::string snapshotDataTypeStr(SnapshotDataType dt);