Update Rust example and convert C to C++ storage example

Author: fwsGonzo

guest/rust_storage/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+emballoc = "0.3.0"
 libc = "0.2"
 
 [profile.release]

guest/rust_storage/src/main.rs

@@ -1,6 +1,12 @@
+#[global_allocator]
+static ALLOCATOR: emballoc::Allocator<4096> = emballoc::Allocator::new();
+extern crate alloc;
+
 use std::arch::asm;
-extern "C" { fn remote_function(arg: fn(i32) -> i32, value: i32) -> i32; }
 use std::process::ExitCode;
+extern "C" { fn remote_function(arg: fn(i32) -> i32, value: i32) -> i32; }
+// Perform a remote allocation with a local allocator
+extern "C" { fn remote_allocation(fsbase: u64, alloc: fn() -> Vec<i32>) -> Vec<i32>; }
 
 fn double_int(input: i32) -> i32 {
     return input * 2;
@@ -30,11 +36,34 @@ fn set_callback(name: &str, cb: extern "C" fn(i32) -> i32) {
 	}
 }
 
+fn get_current_fsbase() -> u64 {
+	let fsbase: u64;
+	unsafe {
+		asm!(
+			"rdfsbase {}",
+			out(reg) fsbase,
+		);
+	}
+	return fsbase;
+}
+fn alloc_vec() -> Vec<i32> {
+	let mut v = Vec::with_capacity(10);
+	for i in 0..10 {
+		v.push(i * 10);
+	}
+	return v;
+}
+
 fn main() -> ExitCode
 {
 	println!("Hello, world!");
 	let result = unsafe { remote_function(double_int, 21) };
 	println!("Result from remote function: {}", result);
+
+	let vec = unsafe { remote_allocation(get_current_fsbase(), alloc_vec) };
+	println!("Received vector from remote allocation: {:?}", vec);
+
+	// Register callbacks
 	set_callback("do_calculation", do_calculation);
 	set_callback("do_nothing", do_nothing);
 	return ExitCode::from(result as u8);

guest/rust_storage/storage_program/src/main.rs

@@ -1,5 +1,6 @@
 use std::process::ExitCode;
 use std::sync::{LazyLock, Mutex};
+use std::arch::asm;
 static ARRAY: LazyLock<Mutex<Vec<i32>>> = LazyLock::new(|| Mutex::new(vec![]));
 
 #[no_mangle]
@@ -10,6 +11,24 @@ extern "C" fn remote_function(arg: fn(i32) -> i32, value: i32) -> i32 {
 	return arg(value);
 }
 
+fn set_fsbase(fsbase: u64) {
+	unsafe {
+		asm!(
+			"wrfsbase {}",
+			in(reg) fsbase,
+		);
+	}
+}
+
+#[no_mangle]
+extern "C" fn remote_allocation(fsbase: u64, alloc: fn() -> Vec<i32>) -> Vec<i32> {
+	// Set the FSBASE to the provided value
+	set_fsbase(fsbase);
+	// Call the allocator function
+	return alloc();
+	//return vec![42, 43, 44];
+}
+
 fn main() -> ExitCode
 {
 	println!("Hello, Storage World!");

guest/storage/build.sh

@@ -1,5 +1,5 @@
 set -v
-gcc-12 -static -O2 -Wl,-Ttext-segment=0x44000000 storage.c -o storage
+clang++-20 -static -O2 -std=c++20 -Wl,-Ttext-segment=0x44000000 storage.cpp -o storage
 
 objcopy -w --extract-symbol --strip-symbol=!remote* --strip-symbol=* storage storage.syms
-gcc-12 -static -O2 -Wl,--just-symbols=storage.syms main.c -o main
+clang++-20 -static -O2 -std=c++20 -Wl,--just-symbols=storage.syms main.cpp -o main

guest/storage/main.c guest/storage/main.cppguest/storage/main.c renamed to guest/storage/main.cpp

@@ -1,5 +1,5 @@
 #include <stdio.h>
-extern int remote_function(int(*callback)(int), int value);
+extern "C" int remote_function(int(*callback)(int), int value);
 
 static int double_int(int value)
 {
@@ -8,7 +8,7 @@ static int double_int(int value)
 
 int main()
 {
-	printf("Jumping to %p\n", &remote_function);
+	printf("Hello from Main VM!\n");
 	fflush(stdout);
 	for (int i = 0; i < 10; i++) {
 		const int val = remote_function(double_int, 21);
@@ -17,12 +17,8 @@ int main()
 	return 0;
 }
 
-int do_calculation(int value)
+extern "C" int do_calculation(int value)
 {
 	return remote_function(double_int, value);
 }
-
-int simple_calculation(int value)
-{
-	return value;
-}
+extern "C" void do_nothing(int) { }

guest/storage/storage.c guest/storage/storage.cppguest/storage/storage.c renamed to guest/storage/storage.cpp

@@ -1,9 +1,9 @@
 #include <stdio.h>
 #include <unistd.h>
 
-extern int remote_function(int (*arg)(int), int value)
+extern "C" int remote_function(int (*arg)(int), int value)
 {
-	write(1, "In remote_function\n", 20);
+	//write(1, "In remote_function\n", 20);
 	return arg(value);
 }
 

src/storage.cpp

@@ -215,35 +215,46 @@ int main(int argc, char** argv)
 	assert(vm.has_remote());
 
 	/* Measure call overhead */
+	uint64_t do_nothing_addr = 0x0;
+	uint64_t do_calculation_addr = 0x0;
 	auto do_it = callback_address.find("do_nothing");
 	if (do_it == callback_address.end()) {
-		fprintf(stderr, "Error: no do_nothing() in guest\n");
-		exit(1);
+		do_nothing_addr = master_vm.address_of("do_nothing");
+		if (do_nothing_addr == 0x0) {
+			fprintf(stderr, "Error: no do_nothing() in guest\n");
+			exit(1);
+		}
+	} else {
+		do_nothing_addr = do_it->second;
 	}
 	auto calc_it = callback_address.find("do_calculation");
 	if (calc_it == callback_address.end()) {
-		fprintf(stderr, "Error: no do_calculation() in guest\n");
-		exit(1);
+		do_calculation_addr = master_vm.address_of("do_calculation");
+		if (do_calculation_addr == 0x0) {
+			fprintf(stderr, "Error: no do_calculation() in guest\n");
+			exit(1);
+		}
+	} else {
+		do_calculation_addr = calc_it->second;
 	}
-	auto call_overhead = timed_action([&] {
-		for (int i = 0; i < 100; i++)
-			vm.vmcall(do_it->second, 5.0f, 21);
-	}) / 100.0;
-	printf("Call overhead: %.2fus\n", call_overhead * 1e6);
+	TimedResult call_overhead = timed_action(1000, [&] {
+		vm.vmcall(do_nothing_addr, 5.0f, 21);
+	});
+	printf("Call overhead: %.2fus\n", call_overhead.median * 1e6);
 
 	/* Call 'do_calculation' with 21 as argument */
-	printf("Calling do_calculation() @ 0x%lX\n", calc_it->second);
+	printf("Calling do_calculation() @ 0x%lX\n", do_calculation_addr);
 	for (int i = 0; i < 50; i++)
-		vm.vmcall(calc_it->second, 21);
+		vm.vmcall(do_calculation_addr, 21);
 	TimedResult fork_tdiff = timed_action(15000, [&] {
-		vm.vmcall(calc_it->second, 21);
+		vm.vmcall(do_calculation_addr, 21);
 	});
 	if (vm.remote_connection_count() < 15000) {
 		fprintf(stderr, "Error: only %u remote connections were made, expected 500\n",
 			vm.remote_connection_count());
 		exit(1);
 	}
-	fork_tdiff -= call_overhead;
+	fork_tdiff -= call_overhead.median;
 	printf("* Remote call time: avg %.2fus  med %.2fus  p50 %.2fus  p90 %.2fus  p99 %.2fus  min %.2fus  max %.2fus\n",
 		fork_tdiff.average * 1e6, fork_tdiff.median * 1e6,
 		fork_tdiff.p50 * 1e6, fork_tdiff.p90 * 1e6,

storage.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+set -e
+pushd build
+make -j8
+popd
+# We can't pass a symlink to the main VM, since we allow dynamic linker loading
+# there, and TinyKVM refuses to use symlinks for security reasons.
+
+# --rust:
+# ./build/storagekvm guest/rust_storage/target/release/demo guest/rust_storage/storage
+# --cpp
+# ./build/storagekvm guest/storage/main guest/storage/storage
+
+# Parse args
+if [ "$1" = "--rust" ]; then
+	MAIN=guest/rust_storage/target/release/demo
+	SECONDARY=guest/rust_storage/storage
+elif [ "$1" = "--cpp" ]; then
+	MAIN=guest/storage/main
+	SECONDARY=guest/storage/storage
+else
+	echo "Usage: $0 [--rust|--cpp]"
+	exit 1
+fi
+
+./build/storagekvm $MAIN $SECONDARY