Allow attaching accessed pages to snapshots

Author: fwsGonzo

lib/tinykvm/amd64/paging.cpp

@@ -524,9 +524,24 @@ void foreach_page_makecow(vMemory& mem, uint64_t kernel_end, uint64_t shared_mem
 				entry &= ~PDE64_RW;
 				entry |= PDE64_CLONEABLE | PDE64_G; // Global bit for read-only pages
 			}
+			// Clear accessed bit for *all* pages
+			// Doing this makes it possible to send a dummy request to estimate which
+			// pages are needed after a fork, for use with MAP_POPULATE-like optimizations.
+			entry &= ~PDE64_ACCESSED;
 		}
 	});
 }
+std::vector<uint64_t> get_accessed_pages(const vMemory& memory)
+{
+	std::vector<uint64_t> accessed_pages;
+	foreach_page(memory,
+	[&accessed_pages] (uint64_t addr, uint64_t& entry, size_t /*size*/) {
+		if ((entry & (PDE64_ACCESSED | PDE64_PRESENT)) == (PDE64_ACCESSED | PDE64_PRESENT)) {
+			accessed_pages.push_back(addr & PDE64_ADDR_MASK);
+		}
+	}, false);
+	return accessed_pages;
+}
 
 void page_at(vMemory& memory, uint64_t addr, foreach_page_t callback, bool ignore_missing)
 {

lib/tinykvm/amd64/paging.hpp

@@ -14,6 +14,7 @@ using foreach_page_t = std::function<void(uint64_t, uint64_t&, size_t)>;
 extern void foreach_page(vMemory&, foreach_page_t callback, bool skip_oob_addresses = true);
 extern void foreach_page(const vMemory&, foreach_page_t callback, bool skip_oob_addresses = true);
 extern void foreach_page_makecow(vMemory&, uint64_t kernel_end, uint64_t shared_memory_boundary);
+extern std::vector<uint64_t> get_accessed_pages(const vMemory& memory);
 
 extern void page_at(vMemory&, uint64_t addr, foreach_page_t, bool ignore_missing = false);
 struct WritablePage {

lib/tinykvm/machine.hpp

@@ -245,6 +245,7 @@ struct Machine
 	void set_main_memory_writable(bool v) { memory.main_memory_writes = v; }
 	bool is_forked() const noexcept { return m_forked; }
 	bool uses_cow_memory() const noexcept { return m_forked || m_prepped; }
+	std::vector<uint64_t> get_accessed_pages() const;
 
 	/* Remote VM through address space merging */
 	void remote_connect(Machine& other, bool connect_now = false);
@@ -298,7 +299,7 @@ struct Machine
 	/* Store non-memory VM state to the already existing cold
 	   start state area in memory. Any failure will throw an
 	   exception. The memory must have been pre-allocated. */
-	void save_snapshot_state_now() const;
+	void save_snapshot_state_now(const std::vector<uint64_t>& populate_pages = {}) const;
 	/* Check if the VM was loaded from a snapshot state. */
 	bool has_snapshot_state() const noexcept { return m_loaded_from_snapshot; }
 	/* Get pointer to user area in snapshot state memory, or nullptr

lib/tinykvm/machine_state.cpp

@@ -1,19 +1,27 @@
 #include "machine.hpp"
 
+#include <cstring>
 #include <fcntl.h>
 #include <linux/kvm.h>
 #include <stdexcept>
+#include <span>
 #include <sys/mman.h>
 #include <sys/stat.h>
 #include <unistd.h>
 #ifdef TINYKVM_ARCH_AMD64
 #include "amd64/amd64.hpp"
+#include "amd64/paging.hpp"
 #endif
 #include "linux/fds.hpp"
 #include "linux/threads.hpp"
 
 namespace tinykvm {
 
+struct ColdStartAccessedRange {
+	uint64_t start;
+	uint64_t end;
+};
+
 struct ColdStartThreadState {
 	int tid;
 	tinykvm_x86regs    regs;
@@ -73,10 +81,13 @@ struct SnapshotState {
 	//MMapCache m_mmap_cache;
 	Machine::address_t mmap_current;
 
-	bool main_memory_writes;
 	Machine::address_t m_page_tables;
+	bool main_memory_writes;
+	uint32_t num_access_ranges;
 
-	static constexpr size_t Size() noexcept { return 4096ul; }
+	char current[0];
+
+	static constexpr size_t Size() noexcept { return vMemory::ColdStartStateSize(); }
 
 	template <typename T>
 	T* next(void*& current) {
@@ -106,6 +117,7 @@ bool Machine::load_snapshot_state()
 		throw std::runtime_error("No valid snapshot state found");
 	}
 	if (state.size < sizeof(SnapshotState) || state.size > SnapshotState::Size()) {
+		fprintf(stderr, "Invalid snapshot state size: %u\n", state.size);
 		throw std::runtime_error("Invalid snapshot state size");
 	}
 
@@ -129,7 +141,22 @@ bool Machine::load_snapshot_state()
 		this->memory.main_memory_writes = state.main_memory_writes;
 		this->memory.page_tables = state.m_page_tables;
 
-		void* current = reinterpret_cast<char*>(&state) + sizeof(SnapshotState);
+		void* current = state.current;
+		// Load populate pages
+		for (unsigned i = 0; i < state.num_access_ranges; i++) {
+			ColdStartAccessedRange* range = state.next<ColdStartAccessedRange>(current);
+			if (range->start >= MemoryBanks::ARENA_BASE_ADDRESS)
+				continue;
+			try {
+				printf("Populating pages from 0x%lX -> 0x%lX\n", range->start, range->end);
+				char* page = this->memory.get_userpage_at(range->start);
+				madvise(page, range->end - range->start, MADV_WILLNEED);
+			} catch (const std::exception& e) {
+				fprintf(stderr, "Failed to access page at 0x%lX: %s\n", range->start, e.what());
+				continue;
+			}
+		}
+
 		// Load the thread states
 		ColdStartThreads* threads = state.next<ColdStartThreads>(current);
 		if (threads->count > 0) {
@@ -181,7 +208,7 @@ bool Machine::load_snapshot_state()
 	}
 	return true;
 }
-void Machine::save_snapshot_state_now() const
+void Machine::save_snapshot_state_now(const std::vector<uint64_t>& populate_pages) const
 {
 	if (this->is_forked()) {
 		throw std::runtime_error("Cannot save snapshot state of a forked VM");
@@ -209,7 +236,40 @@ void Machine::save_snapshot_state_now() const
 		state.main_memory_writes = this->memory.main_memory_writes;
 		state.m_page_tables = this->memory.page_tables;
 
-		void* current = reinterpret_cast<char*>(&state) + sizeof(SnapshotState);
+		void* current = state.current;
+		// Save populate pages
+		state.num_access_ranges = 0;
+		if (!populate_pages.empty()) {
+			uint64_t current_begin = 0;
+			uint64_t current_end = 0;
+			for (uint64_t page_addr : populate_pages) {
+				if (page_addr >= MemoryBanks::ARENA_BASE_ADDRESS)
+					continue;
+				// Merge contiguous ranges
+				if (current_end == page_addr) {
+					current_end += vMemory::PageSize();
+					continue;
+				}
+				// Store previous range
+				if (current_end != current_begin) {
+					ColdStartAccessedRange* range = state.next<ColdStartAccessedRange>(current);
+					range->start = current_begin;
+					range->end = current_end;
+					state.num_access_ranges++;
+				}
+				// Start new range
+				current_begin = page_addr;
+				current_end = page_addr;
+			}
+			// Store last range
+			if (current_end != current_begin) {
+				ColdStartAccessedRange* range = state.next<ColdStartAccessedRange>(current);
+				range->start = current_begin;
+				range->end = current_end;
+				state.num_access_ranges++;
+			}
+		}
+
 		// Save the multi-threading state
 		ColdStartThreads* threads = state.next<ColdStartThreads>(current);
 		if (this->has_threads()) {
@@ -274,7 +334,14 @@ void Machine::save_snapshot_state_now() const
 		}
 
 	} catch (const MachineException& me) {
+		fprintf(stderr, "Failed to get snapshot state: %s Data: 0x%#lX\n",
+			me.what(), me.data());
+		state.magic = 0; // Invalidate
 		throw std::runtime_error(std::string("Failed to get snapshot state: ") + me.what());
+	} catch (const std::exception& e) {
+		fprintf(stderr, "Failed to get snapshot state: %s\n", e.what());
+		state.magic = 0; // Invalidate
+		throw;
 	}
 }
 

lib/tinykvm/memory.cpp

@@ -425,6 +425,9 @@ vMemory::AllocationResult
 		// Map an existing file, which should not be modified on disk
 		ptr = (char*) mmap(NULL, size, PROT_READ | PROT_WRITE,
 			MAP_PRIVATE | MAP_NORESERVE, fd, 0);
+		// Advise the kernel that we will be immediately accessing this memory's
+		// state and user region
+		madvise(ptr + size - ColdStartStateSize(), vMemory::PageSize(), MADV_WILLNEED);
 	}
 	close(fd);
 	if (ptr == MAP_FAILED) {
@@ -509,6 +512,10 @@ char* vMemory::get_userpage_at(uint64_t addr) const
 #endif
 }
 
+std::vector<uint64_t> Machine::get_accessed_pages() const
+{
+	return tinykvm::get_accessed_pages(this->main_memory());
+}
 size_t Machine::banked_memory_pages() const noexcept
 {
 	size_t count = 0;

lib/tinykvm/memory.hpp

@@ -110,7 +110,7 @@ struct vMemory {
 		static constexpr uint64_t ALIGN = 1ULL << 21;
 		return (size + (ALIGN - 1)) & ~(ALIGN - 1);
 	}
-	static size_t ColdStartStateSize() {
+	static constexpr size_t ColdStartStateSize() {
 		return 2UL << 20; // 2MB
 	}
 	bool has_loadable_snapshot_state() const noexcept;