Add support for VM snapshots

CMakeLists.txt

@@ -13,6 +13,7 @@ add_executable(kvmserver
 	src/file.cpp
 	src/warmup.cpp
 	src/vm.cpp
+	src/vm_state.cpp
 )
 target_compile_features(kvmserver PUBLIC cxx_std_20)
 target_link_libraries(kvmserver

src/config.cpp

@@ -315,7 +315,7 @@ Configuration Configuration::FromArgs(int argc, char* argv[])
 	app.add_option("-t,--threads", config.concurrency, "Number of request VMs (0 to use cpu count)")->capture_default_str();
 	app.add_flag("-e,--ephemeral", config.ephemeral, "Use ephemeral VMs");
 	app.add_option("-w,--warmup", config.warmup_connect_requests, "Number of warmup requests")->capture_default_str();
-	app.add_option("--cold-start-file", config.coldstart_filename, "Cold start snapshot filename");
+	app.add_option("--snapshot-file", config.snapshot_filename, "Snapshot filename");
 
 	app.add_flag("-v,--verbose", config.verbose, "Enable verbose output")->group("Verbose");
 	app.add_flag("--verbose-syscalls", config.verbose_syscalls, "Enable verbose syscall output")->group("Verbose");

src/config.hpp

@@ -10,7 +10,7 @@ struct Configuration
 {
 	std::string main_filename;
 	std::string storage_filename;
-	std::string coldstart_filename;
+	std::string snapshot_filename;
 	uint16_t concurrency = 1; /* Request VMs */
 	uint16_t warmup_connect_requests = 0; /* Warmup requests, individual connections */
 	uint16_t warmup_intra_connect_requests = 1; /* Send N requests while connected */

src/vm.cpp

@@ -166,8 +166,8 @@ VirtualMachine::VirtualMachine(std::string_view binary, const Configuration& con
 		.master_direct_memory_writes = true,
 		.split_hugepages = false,
 		.executable_heap = config.executable_heap,
-		.mmap_backed_files = config.mmap_backed_files,
-		.fast_cold_start_file = storage ? "" : config.coldstart_filename,
+		.mmap_backed_files = config.mmap_backed_files && config.snapshot_filename.empty(),
+		.snapshot_file = storage ? "" : config.snapshot_filename,
 		.hugepages_arena_size = config.hugepage_arena_size,
 	}),
 	m_config(config),
@@ -414,9 +414,28 @@ void VirtualMachine::reset_to(const VirtualMachine& other)
 	this->m_blocking_connections = false;
 }
 
+VirtualMachine::InitResult VirtualMachine::initialize_from_file()
+{
+	InitResult result;
+	auto start = std::chrono::high_resolution_clock::now();
+	this->set_waiting_for_requests(true);
+	this->machine().prepare_copy_on_write();
+	if (this->machine().has_snapshot_state()) {
+		this->load_state();
+	}
+	auto end = std::chrono::high_resolution_clock::now();
+	result.initialization_time += std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
+	result.warmup_time = std::chrono::milliseconds(0);
+	return result;
+}
+
 VirtualMachine::InitResult VirtualMachine::initialize(std::function<void()> warmup_callback, bool just_one_vm)
 {
 	InitResult result;
+	if (machine().has_snapshot_state()) {
+		return this->initialize_from_file();
+	}
+
 	auto start = std::chrono::high_resolution_clock::now();
 	try {
 		// Use constrained working memory
@@ -620,6 +639,10 @@ VirtualMachine::InitResult VirtualMachine::initialize(std::function<void()> warm
 			machine().set_registers(regs);
 		}
 
+		if (!m_is_storage && machine().main_memory().has_snapshot_area()) {
+			this->save_state();
+		}
+
 		// Finish measuring initialization time
 		end = std::chrono::high_resolution_clock::now();
 		result.initialization_time += std::chrono::duration_cast<std::chrono::milliseconds>(end - start);

src/vm.hpp

@@ -60,6 +60,9 @@ struct VirtualMachine
 	void stop_warmup_client();
 	bool connect_and_send_requests(const sockaddr* serv_addr, socklen_t serv_addr_len);
 	bool validate_listener(int fd);
+	InitResult initialize_from_file();
+	void save_state();
+	void load_state();
 
 	tinykvm::Machine m_machine;
 	const Configuration& m_config;

src/vm_state.cpp

@@ -0,0 +1,124 @@
+#include "vm.hpp"
+#include <cstring>
+#include <fcntl.h>
+#include <stdexcept>
+#include <sys/socket.h>
+static constexpr bool VERBOSE_SNAPSHOT = false;
+
+struct AppSnapshotState {
+	VirtualMachine::PollMethod poll_method;
+	int tracked_client_vfd;
+	int backlog;
+	int domain;
+	int type;
+	int protocol;
+	int flags;
+	int reuseaddr;
+	socklen_t addr_len;
+	struct sockaddr_storage addr;
+};
+
+void VirtualMachine::save_state()
+{
+	machine().save_snapshot_state_now();
+	void* map = machine().get_snapshot_state_user_area();
+	if (map == nullptr) {
+		throw std::runtime_error("snapshot user area is null");
+	}
+	AppSnapshotState& state = *reinterpret_cast<AppSnapshotState*>(map);
+	state.poll_method = this->m_poll_method;
+	state.tracked_client_vfd = this->m_tracked_client_vfd;
+	state.backlog = 128; // XXX
+
+	const auto fd = this->m_tracked_client_fd;
+	int len = sizeof(state.domain);
+	if(getsockopt(fd, SOL_SOCKET, SO_DOMAIN, &(state.domain), (socklen_t*) &len) < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	len = sizeof(state.type);
+	if(getsockopt(fd, SOL_SOCKET, SO_TYPE, &(state.type), (socklen_t*) &len) < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	len = sizeof(state.protocol);
+	if(getsockopt(fd, SOL_SOCKET, SO_PROTOCOL, &(state.protocol), (socklen_t*) &len) < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	len = sizeof(state.reuseaddr);
+	if (getsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &(state.reuseaddr), (socklen_t*) &len) < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	state.flags = fcntl(fd, F_GETFL, 0);
+	if (state.flags < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	state.addr_len = sizeof(state.addr);
+	if (getsockname(fd, (struct sockaddr*)&(state.addr), &(state.addr_len)) < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+}
+
+void VirtualMachine::load_state()
+{
+	auto map = machine().get_snapshot_state_user_area();
+	if (map == NULL) {
+		throw std::runtime_error("snapshot user area is null");
+	}
+	AppSnapshotState& state = *reinterpret_cast<AppSnapshotState*>(map);
+	const auto fdm = machine().fds();
+	this->m_poll_method = state.poll_method;
+	int fd = socket(state.domain, state.type, state.protocol);
+	if (fd < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &(state.reuseaddr), sizeof(state.reuseaddr)) < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	if (fcntl(fd, F_SETFL, state.flags) < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	if (bind(fd, (struct sockaddr*) &(state.addr), state.addr_len) < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	if (listen(fd, state.backlog) < 0) {
+		throw std::runtime_error(strerror(errno));
+	}
+	this->m_tracked_client_vfd = state.tracked_client_vfd;
+	this->m_tracked_client_fd = fd;
+	this->machine().fds().manage_as(state.tracked_client_vfd, fd, true, true);
+
+	// Look through epoll systems
+	for (auto& [vfd, epoll_entry] : fdm.get_epoll_entries())
+	{
+		// Find the tracked client vfd in the epoll entries
+		auto it = epoll_entry->epoll_fds.find(this->m_tracked_client_vfd);
+		if (it != epoll_entry->epoll_fds.end()) {
+			const int entry_vfd = it->first;
+			auto& event = it->second;
+			const int epoll_fd = this->machine().fds().translate(vfd);
+			// Remove old entry (with old fd)
+			int old_fd = this->machine().fds().translate(entry_vfd);
+			if (old_fd >= 0) {
+				epoll_ctl(epoll_fd, EPOLL_CTL_DEL, old_fd, &event);
+			}
+			// Add new entry (with new fd)
+			epoll_ctl(epoll_fd, EPOLL_CTL_ADD, fd, &event);
+			if constexpr (VERBOSE_SNAPSHOT) {
+				printf("TinyKVM: Restored epoll entry for vfd %d to new fd %d\n", entry_vfd, fd);
+			}
+		}
+		// Remove entries where we don't have the tracked fd
+		for (auto it = epoll_entry->epoll_fds.begin(); it != epoll_entry->epoll_fds.end(); ) {
+			const int entry_vfd = it->first;
+			const int entry_fd = this->machine().fds().translate(entry_vfd);
+			if (entry_fd < 0) {
+				// Remove the fd from the epoll entry since we can't use it anymore
+				it = epoll_entry->epoll_fds.erase(it);
+				if constexpr (VERBOSE_SNAPSHOT) {
+					printf("TinyKVM: Removed stale epoll entry for vfd %d\n", entry_vfd);
+				}
+			} else {
+				++it;
+			}
+		}
+	}
+}