Add snapshot ignore region for wasm stack and thread stacks

include/wasm/WasmModule.h

@@ -157,6 +157,8 @@ class WasmModule
 
     int awaitPthreadCall(const faabric::Message* msg, int pthreadPtr);
 
+    std::vector<uint32_t> getThreadStacks();
+
     // ----- Debugging -----
     virtual void printDebugInfo();
 
@@ -207,6 +209,8 @@ class WasmModule
     // Snapshots
     void snapshotWithKey(const std::string& snapKey, bool locallyRestorable);
 
+    void ignoreAllStacksInSnapshot(const std::string& snapshotKey);
+
     // Threads
     void createThreadStacks();
 };

src/wamr/WAMRWasmModule.cpp

@@ -76,7 +76,6 @@ WAMRWasmModule* getExecutingWAMRModule()
 // ----- Module lifecycle -----
 void WAMRWasmModule::doBindToFunction(faabric::Message& msg, bool cache)
 {
-
     // Prepare the filesystem
     filesystem.prepareFilesystem();
 

src/wasm/WasmModule.cpp

@@ -1,8 +1,7 @@
-#include "wasm/WasmModule.h"
-
 #include <conf/FaasmConfig.h>
 #include <threads/ThreadState.h>
 #include <wasm/WasmExecutionContext.h>
+#include <wasm/WasmModule.h>
 
 #include <faabric/scheduler/Scheduler.h>
 #include <faabric/snapshot/SnapshotRegistry.h>
@@ -14,6 +13,7 @@
 #include <faabric/util/locks.h>
 #include <faabric/util/logging.h>
 #include <faabric/util/memory.h>
+#include <faabric/util/snapshot.h>
 #include <faabric/util/timing.h>
 
 #include <boost/filesystem.hpp>
@@ -301,6 +301,40 @@ void WasmModule::bindToFunction(faabric::Message& msg, bool cache)
     doBindToFunction(msg, cache);
 }
 
+void WasmModule::ignoreAllStacksInSnapshot(const std::string& snapshotKey)
+{
+    faabric::util::SnapshotData& snapData =
+      faabric::snapshot::getSnapshotRegistry().getSnapshot(snapshotKey);
+
+    // First ignore the main wasm stack
+    SPDLOG_TRACE("Ignoring snapshot diffs for {} for wasm stack: 0-{}",
+                 snapshotKey,
+                 STACK_SIZE);
+
+    snapData.addMergeRegion(0,
+                            STACK_SIZE,
+                            faabric::util::SnapshotDataType::Raw,
+                            faabric::util::SnapshotMergeOperation::Ignore);
+
+    uint32_t threadStackRegionStart =
+      threadStacks.at(0) + 1 - THREAD_STACK_SIZE - GUARD_REGION_SIZE;
+    uint32_t threadStackRegionSize =
+      threadPoolSize * (THREAD_STACK_SIZE + (2 * GUARD_REGION_SIZE));
+
+    SPDLOG_TRACE("Ignoring snapshot diffs for {} for thread stacks: {}-{}",
+                 snapshotKey,
+                 threadStackRegionStart,
+                 threadStackRegionStart + threadStackRegionSize);
+
+    // Note - the merge regions for a snapshot are keyed on the offset, so
+    // we will just overwrite the same region if another module has already
+    // set it
+    snapData.addMergeRegion(threadStackRegionStart,
+                            threadStackRegionSize,
+                            faabric::util::SnapshotDataType::Raw,
+                            faabric::util::SnapshotMergeOperation::Ignore);
+}
+
 void WasmModule::prepareArgcArgv(const faabric::Message& msg)
 {
     // Here we set up the arguments to main(), i.e. argc and argv
@@ -398,6 +432,11 @@ int32_t WasmModule::executeTask(
     assert(!threadStacks.empty());
     uint32_t stackTop = threadStacks.at(threadPoolIdx);
 
+    // Ensure we ignore all stacks in a snapshot if it exists
+    if (!msg.snapshotkey().empty()) {
+        ignoreAllStacksInSnapshot(msg.snapshotkey());
+    }
+
     // Perform the appropriate type of execution
     int returnValue;
     if (req->type() == faabric::BatchExecuteRequest::THREADS) {
@@ -551,7 +590,7 @@ void WasmModule::createThreadStacks()
     SPDLOG_DEBUG("Creating {} thread stacks", threadPoolSize);
 
     for (int i = 0; i < threadPoolSize; i++) {
-        // Allocate thread and guard pages
+        // Allocate thread stack and guard pages
         uint32_t memSize = THREAD_STACK_SIZE + (2 * GUARD_REGION_SIZE);
         uint32_t memBase = growMemory(memSize);
 
@@ -566,6 +605,11 @@ void WasmModule::createThreadStacks()
     }
 }
 
+std::vector<uint32_t> WasmModule::getThreadStacks()
+{
+    return threadStacks;
+}
+
 threads::MutexManager& WasmModule::getMutexes()
 {
     return mutexes;

tests/test/wasm/test_snapshots.cpp

@@ -6,9 +6,14 @@
 #include <boost/filesystem.hpp>
 
 #include <faabric/proto/faabric.pb.h>
+#include <faabric/runner/FaabricMain.h>
+#include <faabric/snapshot/SnapshotClient.h>
+#include <faabric/util/config.h>
+#include <faabric/util/environment.h>
 #include <faabric/util/func.h>
 #include <faabric/util/memory.h>
 #include <faabric/util/snapshot.h>
+#include <faabric/util/testing.h>
 
 #include <faaslet/Faaslet.h>
 #include <wavm/WAVMWasmModule.h>
@@ -21,11 +26,12 @@ class WasmSnapTestFixture
   : public RedisTestFixture
   , public SchedulerTestFixture
   , public SnapshotTestFixture
+  , public ConfTestFixture
 {
   public:
-    WasmSnapTestFixture() {}
+    WasmSnapTestFixture() { wasm::getWAVMModuleCache().clear(); }
 
-    ~WasmSnapTestFixture() {}
+    ~WasmSnapTestFixture() { wasm::getWAVMModuleCache().clear(); }
 };
 
 TEST_CASE_METHOD(WasmSnapTestFixture,
@@ -247,7 +253,7 @@ TEST_CASE_METHOD(WasmSnapTestFixture,
     std::string snapshotKey = "foobar-snap";
 
     size_t snapSize = 64 * faabric::util::HOST_PAGE_SIZE;
-    uint8_t* snapMemory = (uint8_t*)mmap(
+    uint8_t* snapMemory = (uint8_t*)::mmap(
       nullptr, snapSize, PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
 
     // Write dummy data to the snapshot
@@ -328,5 +334,109 @@ TEST_CASE_METHOD(WasmSnapTestFixture,
         REQUIRE(dataAfter == dummyDataB);
         REQUIRE(f.module->getCurrentBrk() == snapSize);
     }
+
+    ::munmap(snapMemory, snapSize);
+}
+
+TEST_CASE_METHOD(WasmSnapTestFixture,
+                 "Test ignoring stacks",
+                 "[wasm][snapshot]")
+{
+    int nCores = 10;
+    conf.overrideCpuCount = nCores;
+
+    std::string user = "demo";
+    std::string function = "echo";
+
+    // Set up a snapshot for this function
+    faabric::Message msgA = faabric::util::messageFactory(user, function);
+    wasm::WAVMWasmModule moduleA;
+    moduleA.bindToFunction(msgA);
+    std::string originalSnapshotKey = moduleA.snapshot(true);
+
+    // Set up another module from this snapshot
+    std::shared_ptr<faabric::BatchExecuteRequest> req =
+      faabric::util::batchExecFactory(user, function, 1);
+    faabric::Message& msgB = req->mutable_messages()->at(0);
+    msgB.set_snapshotkey(originalSnapshotKey);
+
+    wasm::WAVMWasmModule moduleB;
+    moduleB.bindToFunction(msgB);
+    std::vector<uint32_t> threadStacksB = moduleB.getThreadStacks();
+
+    // Check thread stacks are same size
+    std::vector<uint32_t> threadStacksA = moduleA.getThreadStacks();
+    REQUIRE(threadStacksA.size() == nCores);
+    REQUIRE(threadStacksB == threadStacksA);
+
+    // Execute the function to make sure the ignores are set up
+    int32_t returnValue = moduleB.executeTask(0, 0, req);
+    REQUIRE(returnValue == 0);
+
+    // Reset dirty tracking otherwise we'll have all wasm pages marked as dirty
+    // (after being restored from the snapshot)
+    faabric::util::resetDirtyTracking();
+
+    // Modify a couple of places in the wasm stack
+    auto* wasmStackBottom = (int*)(moduleB.wasmPointerToNative(0));
+    auto* wasmStackTop = (int*)(moduleB.wasmPointerToNative(STACK_SIZE - 1));
+
+    *wasmStackBottom = 345;
+    *wasmStackTop = 123;
+
+    // Modify the top and bottom of each thread stack
+    // Note that each stack grows downwards
+    for (auto t : threadStacksB) {
+        auto* stackTop =
+          (int*)(moduleB.wasmPointerToNative((t - 1) - sizeof(int)));
+        auto* stackBottom =
+          (int*)(moduleB.wasmPointerToNative(t - THREAD_STACK_SIZE));
+
+        *stackTop = 123;
+        *stackBottom = 345;
+    }
+
+    // Modify some other places in the heap
+    auto* heapA =
+      (int*)(moduleB.wasmPointerToNative(threadStacksB.back() + 100));
+    auto* heapB =
+      (int*)(moduleB.wasmPointerToNative(threadStacksB.back() + 300));
+    *heapA = 123;
+    *heapB = 345;
+
+    // Get post execution snapshot
+    std::string snapKeyPostExecution = moduleB.snapshot();
+    faabric::util::SnapshotData snapshotPostExecution =
+      reg.getSnapshot(snapKeyPostExecution);
+
+    // Diff with original snapshot
+    faabric::util::SnapshotData snapshotPreExecution =
+      reg.getSnapshot(originalSnapshotKey);
+    std::vector<faabric::util::SnapshotDiff> diffs =
+      snapshotPreExecution.getChangeDiffs(snapshotPostExecution.data,
+                                          snapshotPostExecution.size);
+
+    // Check that we have some diffs, but that none of them are in the thread
+    // stacks region
+    REQUIRE(!diffs.empty());
+    uint32_t stacksMin =
+      threadStacksB.at(0) + 1 - THREAD_STACK_SIZE - GUARD_REGION_SIZE;
+    uint32_t stacksMax = threadStacksB.back();
+
+    for (auto d : diffs) {
+        bool isBelow = (d.offset + d.size) < stacksMin;
+        bool isAbove = d.offset > stacksMax;
+        bool isNotInRegion = isBelow || isAbove;
+
+        if (!isNotInRegion) {
+            SPDLOG_ERROR("Stack not ignored, {} + {} between {} and {}",
+                         d.offset,
+                         d.size,
+                         stacksMin,
+                         stacksMax);
+        }
+
+        REQUIRE(isNotInRegion);
+    }
 }
 }