Fix issues with remote resume, add forking test

Author: fwsGonzo

lib/tinykvm/memory_bank.cpp

@@ -28,6 +28,11 @@ void MemoryBanks::init_from(const MemoryBanks& other)
 {
 	this->m_arena_begin = other.m_arena_begin;
 	this->m_arena_next = other.m_arena_next;
+	// Verify that there aren't any existing banks allocated
+	if (!this->m_mem.empty()) {
+		throw MemoryException("Cannot init_from() when banks are already allocated (arena_begin will be wrong)",
+			this->m_arena_begin, m_mem[0].addr);
+	}
 }
 void MemoryBanks::set_max_pages(size_t new_max, size_t new_hugepages)
 {

lib/tinykvm/remote.cpp

@@ -40,6 +40,7 @@ void Machine::remote_connect(Machine& remote, bool connect_now)
 			this->delete_memory(1);
 			this->memory.delete_foreign_mmap_ranges();
 			this->memory.delete_foreign_banks();
+			this->m_remote_cow_counter = ~0u;
 		}
 		// Install the remote memory in this machine
 		this->install_memory(1, remote_vmem, false);
@@ -115,6 +116,29 @@ void Machine::ipre_remote_resume_now(bool save_all, std::function<void(Machine&)
 	// 2. Connect to remote now
 	const auto remote_fsbase = this->remote_activate_now();
 
+	const auto remote_cow_counter = remote().memory.cow_written_pages.size();
+	bool do_prepare = false;
+	if (this->m_remote_cow_counter != remote_cow_counter) {
+		this->m_remote_cow_counter = remote_cow_counter;
+		// New working memory pages have been created in the remote,
+		// so we need to make sure we see the latest changes.
+		for (auto& bank : remote().memory.banks)
+		{
+			const VirtualMem vmem = bank.to_vmem();
+			if (this->memory.remote_bank_break < bank.addr) {
+				this->memory.remote_bank_break = bank.addr;
+				if constexpr (VERBOSE_REMOTE) {
+					fprintf(stderr, "IPRE remote: mapped bank %u at 0x%lX-0x%lX\n",
+						bank.idx, bank.addr, bank.addr + bank.size());
+				}
+				const unsigned new_idx = memory.allocate_region_idx();
+				this->install_memory(new_idx, vmem, false);
+				memory.foreign_banks.push_back(new_idx);
+				do_prepare = true;
+			}
+		}
+	}
+
 	// 3. Copy remote registers into current state
 	tinykvm::Machine& remote_vm = remote();
 	copy_callee_saved_registers(save_all, this->registers(), remote_vm.registers());
@@ -130,6 +154,8 @@ void Machine::ipre_remote_resume_now(bool save_all, std::function<void(Machine&)
 		// 4. Resume execution
 		// Set RDI to our FSBASE for the remote VM
 		this->registers().rdi = remote_fsbase;
+		if (do_prepare)
+			this->prepare_vmresume(remote_fsbase, true);
 		this->run(0.0f);
 	} catch (const std::exception& e) {
 		// If an exception occurred, disconnect and restore FSBASE

lib/tinykvm/vcpu_run.cpp

@@ -283,13 +283,18 @@ long vCPU::run_once()
 				} else if (machine().is_foreign_address(addr)) {
 					/* Check that the error code is instruction fetch failed */
 					const uint32_t errcode = regs.rax;
+					if constexpr (VERBOSE_REMOTE) {
+						printf("Remote VM page fault at 0x%lX, errcode=0x%X\n", addr, errcode);
+					}
 					if ((errcode & 0x10) == 0) {
 						if (machine().remote().is_remote_connected()) {
 							// Not an instruction fetch, but a memory read or write
 							// Since it's foreign memory, we try to handle it in the remote VM
 							WritablePageOptions zero_opts;
 							zero_opts.zeroes = false;
 							(void)writable_page_at(machine().remote().memory, addr, PDE64_USER | PDE64_RW, zero_opts);
+							regs.rax = 0; /* Indicate that it was local */
+							this->set_registers(regs);
 							return KVM_EXIT_IO;
 						} else {
 							// Not connected, so we can't handle it
@@ -326,7 +331,17 @@ long vCPU::run_once()
 
 				WritablePageOptions zero_opts;
 				zero_opts.zeroes = false;
-				(void)writable_page_at(memory, addr, PDE64_USER | PDE64_RW, zero_opts);
+				auto result = writable_page_at(memory, addr, PDE64_USER | PDE64_RW, zero_opts);
+				if constexpr (false) {
+					char buffer[256];
+					PRINTER(machine().m_printer, buffer,
+						"Page fault on 0x%lX handled, mapped to host page %p\n", addr, result.page);
+					PRINTER(machine().m_printer, buffer,
+						"  Entry value 0x%lX\n", result.entry);
+					PRINTER(machine().m_printer, buffer,
+						"  Our bank arena: begin=0x%lX\n",
+						memory.banks.arena_begin());
+				}
 				return KVM_EXIT_IO;
 			}
 			else if (intr == 1) /* Debug trap */

tests/unit/remote.cpp

@@ -342,6 +342,7 @@ int main() {
 	while (1) {
 		void* p = NULL;
 		size_t len = storage_wait_paused(&p);
+		memset(p, 0, len);
 		strcpy((char*)p, "Data from storage");
 	}
 	return 1234;
@@ -369,7 +370,7 @@ extern long remote_resume(void* data, size_t len);
 long test_remote() {
 	for (int i = 0; i < 100; i++) {
 		char buffer[8192];
-		remote_resume(buffer, sizeof(buffer));
+		remote_resume(buffer, 8192);
 		if (strcmp(buffer, "Data from storage") == 0) {
 			continue;
 		}
@@ -380,8 +381,7 @@ long test_remote() {
 int main() {
 	return test_remote();
 }
-)M",
-											"-Wl,--just-symbols=storage.syms");
+)M", "-Wl,--just-symbols=storage.syms");
 
 	static bool is_waiting = false;
 	tinykvm::Machine::install_unhandled_syscall_handler(
@@ -438,8 +438,32 @@ int main() {
 	REQUIRE(fork.has_remote());
 
 	// Test remote resume
-	fork.vmcall("test_remote");
-	REQUIRE(fork.return_value() == 2345);
-	REQUIRE(!fork.is_remote_connected());
-	REQUIRE(fork.remote_connection_count() == 100);
+	for (int i = 0; i < 100; i++) {
+		is_waiting = false;
+		fork.vmcall("test_remote");
+		REQUIRE(fork.return_value() == 2345);
+		REQUIRE(!fork.is_remote_connected());
+		REQUIRE(fork.remote_connection_count() == (i + 1) * 100);
+		REQUIRE(is_waiting);
+	}
+
+	// Test remote resume against a forked storage VM
+	storage.prepare_copy_on_write();
+	tinykvm::Machine storage_fork(storage, {
+		.max_mem = 16ULL << 20, // MB
+		.max_cow_mem = MAX_COWMEM,
+		.split_hugepages = true
+	});
+	fork.remote_connect(storage_fork);
+
+	// Test remote resume against the forked storage VM
+	printf("Testing remote resume against forked storage VM\n");
+	for (int i = 0; i < 100; i++) {
+		is_waiting = false;
+		fork.vmcall("test_remote");
+		REQUIRE(fork.return_value() == 2345);
+		REQUIRE(!fork.is_remote_connected());
+		REQUIRE(fork.remote_connection_count() == 10000 + (i + 1) * 100);
+		REQUIRE(is_waiting);
+	}
 }