rust: add Wasmtime bindings

The C++ bindings provided by wasmtime are lacking a crucial
capability: asynchronous execution of the wasm functions.
This forces us to stop the execution of the function after
a short time to prevent increasing the latency. Fortunately,
this feature is implemented in the native language
of Wasmtime - Rust. Support for Rust was recently added to
scylla, so we can implement the async bindings ourselves,
which is done in this patch.

The bindings expose all the objects necessary for creating
and calling wasm functions. The majority of code implemented
in Rust is a translation of code that was previously present
in C++.

Types exported from Rust are currently required to be defined
by the  same crate that contains the bridge using them, so
wasmtime types can't be exported directly. Instead, for each
class that was supposed to be exported, a wrapper type is
created, where its first member is the wasmtime class. Note
that the members are not visible from C++ anyway, the
difference only applies to Rust code.

Aside from wasmtime types and methods, two additional types
are exported with some associated methods.
- The first one is ValVec, which is a wrapper for a rust Vec
of wasmtime Vals. The underlying vector is required by
wasmtime methods for calling wasm functions. By having it
exported we avoid multiple conversions from a Val wrapper
to a wasmtime Val, as would be required if we exported a
rust Vec of Val wrappers (the rust Vec itself does not
require wrappers if the type it contains is already wrapped)
- The second one is Fut. This class represents an computation
tha may or may not be ready. We're currently using it
to control the execution of wasm functions from C++. This
class exposes one method: resume(), which returns a bool
that signals whether the computation is finished or not.

Signed-off-by: Wojciech Mitros <wojciech.mitros@scylladb.com>

rust/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 
 [dependencies]
 inc = { path = "inc", version = "0.1.0" }
+wasmtime_bindings = { path = "wasmtime_bindings", version = "0.1.0" }
 
 [lib]
 crate-type = ["staticlib"]

rust/src/lib.rs

@@ -7,3 +7,4 @@
  */
 
 extern crate inc;
+extern crate wasmtime_bindings;

rust/wasmtime_bindings/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "wasmtime_bindings"
+version = "0.1.0"
+edition = "2021"
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+cxx = { version = "1.0.73", features = ["c++20"] }
+wasmtime-wasi = "1.0.0"
+futures = "0.3.23"
+anyhow = "1.0.62"
+
+[dependencies.wasmtime]
+version = "1.0.0"
+default-features = false
+features = ["async", "wat", "cranelift"]

rust/wasmtime_bindings/src/lib.rs

@@ -0,0 +1,342 @@
+/*
+ * Copyright (C) 2022-present ScyllaDB
+ */
+
+/*
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+use anyhow::Context;
+use anyhow::{anyhow, Result};
+use core::task::Poll;
+use futures::future::{BoxFuture, Future};
+mod memory_creator;
+use memory_creator::ScyllaMemoryCreator;
+
+#[cxx::bridge(namespace = "wasmtime")]
+mod ffi {
+    enum ValKind {
+        I32,
+        I64,
+        F32,
+        F64,
+        V128,
+        FuncRef,
+        ExternRef,
+    }
+    extern "Rust" {
+        // We export opaque types directly correlated to wasmtime types with the same names.
+        // Each of these types is a wrapper that for each of its methods calls the corresponding
+        // wasmtime method on the underlying struct.
+        type Instance;
+        fn create_instance(
+            engine: &Engine,
+            module: &Module,
+            store: &mut Store,
+        ) -> Result<Box<Instance>>;
+
+        type Module;
+        fn create_module(engine: &mut Engine, script: &str) -> Result<Box<Module>>;
+
+        type Store;
+        fn create_store(engine: &mut Engine) -> Result<Box<Store>>;
+
+        type Memory;
+        fn data(self: &Memory, store: &Store) -> *mut u8;
+        fn size(self: &Memory, store: &Store) -> u64;
+        fn grow(self: &mut Memory, store: &mut Store, delta: u64) -> Result<u64>;
+        fn get_memory(instance: &Instance, store: &mut Store) -> Result<Box<Memory>>;
+        fn get_abi(instance: &Instance, store: &mut Store, memory: &Memory) -> Result<u32>;
+
+        type Engine;
+        fn create_engine() -> Result<Box<Engine>>;
+
+        type Func;
+        fn create_func(
+            instance: &Instance,
+            store: &mut Store,
+            function_name: &str,
+        ) -> Result<Box<Func>>;
+
+        type Val;
+        fn kind(self: &Val) -> Result<ValKind>;
+        fn i32(self: &Val) -> Result<i32>;
+        fn i64(self: &Val) -> Result<i64>;
+        fn f32(self: &Val) -> Result<f32>;
+        fn f64(self: &Val) -> Result<f64>;
+
+        // vector of wasmtime values, used for passing arguments to a function and returning results from a function
+        type ValVec;
+        fn push_i32(self: &mut ValVec, val: i32) -> Result<()>;
+        fn push_i64(self: &mut ValVec, val: i64) -> Result<()>;
+        fn push_f32(self: &mut ValVec, val: f32) -> Result<()>;
+        fn push_f64(self: &mut ValVec, val: f64) -> Result<()>;
+        fn pop_val(self: &mut ValVec) -> Result<Box<Val>>;
+        fn get_val_vec() -> Result<Box<ValVec>>;
+
+        // type and methods for returning and executing a WebAssembly function asynchronously
+        type Fut<'a>;
+        unsafe fn resume(self: &mut Fut<'_>) -> Result<bool>;
+        unsafe fn get_func_future<'a>(
+            store: &'a mut Store,
+            func: &'a Func,
+            args: &'a ValVec,
+            rets: &'a mut ValVec,
+        ) -> Result<Box<Fut<'a>>>;
+    }
+}
+
+pub struct Instance {
+    wasmtime_instance: wasmtime::Instance,
+}
+
+fn create_instance(engine: &Engine, module: &Module, store: &mut Store) -> Result<Box<Instance>> {
+    let mut linker = wasmtime::Linker::new(&engine.wasmtime_engine);
+    wasmtime_wasi::add_to_linker(&mut linker, |s| s).context("Failed to add wasi to linker")?;
+    let mut inst_fut =
+        Box::pin(linker.instantiate_async(&mut store.wasmtime_store, &module.wasmtime_module));
+    let mut ctx = core::task::Context::from_waker(futures::task::noop_waker_ref());
+
+    loop {
+        // If the instantiation uses async imports in the future, we should return the future here.
+        // For now, we just poll it to completion.
+        match inst_fut.as_mut().poll(&mut ctx) {
+            Poll::Pending => {}
+            Poll::Ready(Ok(inst)) => {
+                return Ok(Box::new(Instance {
+                    wasmtime_instance: inst,
+                }));
+            }
+            Poll::Ready(Err(e)) => {
+                return Err(anyhow!("Failed to instantiate module: {:?}", e));
+            }
+        }
+    }
+}
+
+pub struct Module {
+    wasmtime_module: wasmtime::Module,
+}
+
+fn create_module(engine: &mut Engine, script: &str) -> Result<Box<Module>> {
+    let module = wasmtime::Module::new(&engine.wasmtime_engine, script)
+        .map_err(|e| anyhow!("Compilation failed: {:?}", e))?;
+    Ok(Box::new(Module {
+        wasmtime_module: module,
+    }))
+}
+
+pub struct Store {
+    wasmtime_store: wasmtime::Store<wasmtime_wasi::WasiCtx>,
+}
+
+fn create_store(engine: &mut Engine) -> Result<Box<Store>> {
+    let wasi = wasmtime_wasi::WasiCtxBuilder::new().build();
+    let mut store = wasmtime::Store::new(&engine.wasmtime_engine, wasi);
+    // TODO: make this configurable
+    store.out_of_fuel_async_yield(10000000, 1000);
+    Ok(Box::new(Store {
+        wasmtime_store: store,
+    }))
+}
+
+pub struct Memory {
+    wasmtime_memory: wasmtime::Memory,
+}
+
+impl Memory {
+    fn data(&self, store: &Store) -> *mut u8 {
+        self.wasmtime_memory.data_ptr(&store.wasmtime_store)
+    }
+    fn size(&self, store: &Store) -> u64 {
+        self.wasmtime_memory.size(&store.wasmtime_store)
+    }
+    fn grow(&self, store: &mut Store, delta: u64) -> Result<u64> {
+        self.wasmtime_memory.grow(&mut store.wasmtime_store, delta)
+    }
+}
+
+fn get_memory(instance: &Instance, store: &mut Store) -> Result<Box<Memory>> {
+    let export = instance
+        .wasmtime_instance
+        .get_export(&mut store.wasmtime_store, "memory")
+        .ok_or_else(|| {
+            anyhow!("Memory export not found - please export `memory` in the wasm module")
+        })?;
+    let memory = export
+        .into_memory()
+        .ok_or_else(|| anyhow!("Exported object memory is not a WebAssembly memory"))?;
+    Ok(Box::new(Memory {
+        wasmtime_memory: memory,
+    }))
+}
+
+fn get_abi(instance: &Instance, store: &mut Store, memory: &Memory) -> Result<u32> {
+    let export = instance
+        .wasmtime_instance
+        .get_export(&mut store.wasmtime_store, "_scylla_abi")
+        .ok_or_else(|| {
+            anyhow!("ABI export not found - please export `_scylla_abi` in the wasm module")
+        })?;
+    let global = export
+        .into_global()
+        .ok_or_else(|| anyhow!("Exported object _scylla_abi is not a WebAssembly global"))?;
+    if let wasmtime::Val::I32(x) = global.get(&mut store.wasmtime_store) {
+        let mut bytes = [0u8; 4];
+        memory
+            .wasmtime_memory
+            .read(&store.wasmtime_store, x as usize, &mut bytes)?;
+        Ok(u32::from_le_bytes(bytes))
+    } else {
+        Err(anyhow!("Exported global _scylla_abi is not an i32"))
+    }
+}
+
+pub struct Engine {
+    wasmtime_engine: wasmtime::Engine,
+}
+
+fn create_engine() -> Result<Box<Engine>> {
+    let mut config = wasmtime::Config::new();
+    config.async_support(true);
+    config.consume_fuel(true);
+    // ScyllaMemoryCreator uses malloc (from seastar) to allocate linear memory
+    config.with_host_memory(std::sync::Arc::new(ScyllaMemoryCreator {}));
+    // The following configuration settings make wasmtime allocate only as much memory as it needs
+    config.static_memory_maximum_size(0);
+    config.dynamic_memory_reserved_for_growth(0);
+    config.dynamic_memory_guard_size(0);
+
+    let engine =
+        wasmtime::Engine::new(&config).map_err(|e| anyhow!("Failed to create engine: {:?}", e))?;
+    Ok(Box::new(Engine {
+        wasmtime_engine: engine,
+    }))
+}
+
+pub struct Func {
+    wasmtime_func: wasmtime::Func,
+}
+
+fn create_func(instance: &Instance, store: &mut Store, name: &str) -> Result<Box<Func>> {
+    let export = instance
+        .wasmtime_instance
+        .get_export(&mut store.wasmtime_store, name)
+        .ok_or_else(|| anyhow!("Function {name} was not found in given wasm source code"))?;
+    let func = export
+        .into_func()
+        .ok_or_else(|| anyhow!("Exported object {name} is not a function"))?;
+    Ok(Box::new(Func {
+        wasmtime_func: func,
+    }))
+}
+
+pub struct Val {
+    wasmtime_val: wasmtime::Val,
+}
+
+impl Val {
+    fn kind(&self) -> Result<ffi::ValKind> {
+        match self.wasmtime_val {
+            wasmtime::Val::I32(_) => Ok(ffi::ValKind::I32),
+            wasmtime::Val::I64(_) => Ok(ffi::ValKind::I64),
+            wasmtime::Val::F32(_) => Ok(ffi::ValKind::F32),
+            wasmtime::Val::F64(_) => Ok(ffi::ValKind::F64),
+            wasmtime::Val::V128(_) => Ok(ffi::ValKind::V128),
+            wasmtime::Val::FuncRef(_) => Ok(ffi::ValKind::FuncRef),
+            wasmtime::Val::ExternRef(_) => Ok(ffi::ValKind::ExternRef),
+        }
+    }
+    fn i32(&self) -> Result<i32> {
+        match self.wasmtime_val {
+            wasmtime::Val::I32(val) => Ok(val),
+            _ => Err(anyhow!("Expected i32")),
+        }
+    }
+    fn i64(&self) -> Result<i64> {
+        match self.wasmtime_val {
+            wasmtime::Val::I64(val) => Ok(val),
+            _ => Err(anyhow!("Expected i64")),
+        }
+    }
+    fn f32(&self) -> Result<f32> {
+        match self.wasmtime_val {
+            wasmtime::Val::F32(val) => Ok(f32::from_bits(val)),
+            _ => Err(anyhow!("Expected f32")),
+        }
+    }
+    fn f64(&self) -> Result<f64> {
+        match self.wasmtime_val {
+            wasmtime::Val::F64(val) => Ok(f64::from_bits(val)),
+            _ => Err(anyhow!("Expected f64")),
+        }
+    }
+}
+
+pub struct ValVec {
+    wasmtime_val_vec: Vec<wasmtime::Val>,
+}
+
+impl ValVec {
+    fn push_i32(&mut self, val: i32) -> Result<()> {
+        self.wasmtime_val_vec.push(wasmtime::Val::I32(val));
+        Ok(())
+    }
+    fn push_i64(&mut self, val: i64) -> Result<()> {
+        self.wasmtime_val_vec.push(wasmtime::Val::I64(val));
+        Ok(())
+    }
+    fn push_f32(&mut self, val: f32) -> Result<()> {
+        self.wasmtime_val_vec
+            .push(wasmtime::Val::F32(val.to_bits()));
+        Ok(())
+    }
+    fn push_f64(&mut self, val: f64) -> Result<()> {
+        self.wasmtime_val_vec
+            .push(wasmtime::Val::F64(val.to_bits()));
+        Ok(())
+    }
+    fn pop_val(&mut self) -> Result<Box<Val>> {
+        match self.wasmtime_val_vec.pop() {
+            Some(val) => Ok(Box::new(Val { wasmtime_val: val })),
+            None => Err(anyhow!("Popping from an empty value vector")),
+        }
+    }
+}
+
+fn get_val_vec() -> Result<Box<ValVec>> {
+    Ok(Box::new(ValVec {
+        wasmtime_val_vec: Vec::<wasmtime::Val>::new(),
+    }))
+}
+
+pub struct Fut<'a> {
+    fut: BoxFuture<'a, Result<()>>,
+}
+
+impl<'a> Fut<'a> {
+    fn resume(&mut self) -> Result<bool> {
+        match self.fut.as_mut().poll(&mut core::task::Context::from_waker(
+            futures::task::noop_waker_ref(),
+        )) {
+            Poll::Pending => Ok(false),
+            Poll::Ready(Ok(())) => Ok(true),
+            Poll::Ready(Err(e)) => Err(e),
+        }
+    }
+}
+
+fn get_func_future<'a>(
+    store: &'a mut Store,
+    func: &'a Func,
+    args: &'a ValVec,
+    rets: &'a mut ValVec,
+) -> Result<Box<Fut<'a>>> {
+    Ok(Box::new(Fut {
+        fut: Box::pin(func.wasmtime_func.call_async(
+            &mut store.wasmtime_store,
+            args.wasmtime_val_vec.as_slice(),
+            rets.wasmtime_val_vec.as_mut_slice(),
+        )),
+    }))
+}

rust/wasmtime_bindings/src/memory_creator.rs

@@ -0,0 +1,115 @@
+/*
+ * Copyright (C) 2022-present ScyllaDB
+ */
+
+/*
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+use anyhow::{anyhow, Result};
+use std::{cmp, ptr, u32};
+use wasmtime::LinearMemory;
+
+const WASM_PAGE_SIZE: usize = 64 * 1024;
+
+extern "C" {
+    fn aligned_alloc(align: usize, size: usize) -> *mut u8;
+    fn free(ptr: *mut u8);
+}
+
+pub struct ScyllaLinearMemory {
+    ptr: *mut u8,
+    size: usize,
+    maximum_size: Option<usize>,
+}
+
+// The entire ScyllaLinearMemory is used only in a single thread,
+// because we're not sharing it between seastar shards
+unsafe impl Send for ScyllaLinearMemory {}
+unsafe impl Sync for ScyllaLinearMemory {}
+
+impl Drop for ScyllaLinearMemory {
+    fn drop(&mut self) {
+        // previously allocated or reset to nullptr in grow_to()
+        unsafe { free(self.ptr) };
+    }
+}
+
+unsafe impl LinearMemory for ScyllaLinearMemory {
+    fn byte_size(&self) -> usize {
+        self.size
+    }
+    fn maximum_byte_size(&self) -> Option<usize> {
+        self.maximum_size
+    }
+    fn grow_to(&mut self, new_size: usize) -> Result<()> {
+        let new_size_maybe_extra_page = new_size.checked_add(WASM_PAGE_SIZE - 1).ok_or(anyhow!(
+            "memory grow failed: new size {} is too large",
+            new_size
+        ))?;
+        let new_size_aligned = new_size_maybe_extra_page & !(WASM_PAGE_SIZE - 1);
+        if new_size_aligned == self.size {
+            return Ok(());
+        }
+        let max_size = self.maximum_size.unwrap_or(u32::MAX as usize + 1);
+        if new_size_aligned > max_size {
+            return Err(anyhow!(
+                "memory grow failed: new size {} exceeds maximum size {}",
+                new_size_aligned,
+                max_size
+            ));
+        }
+        let new_ptr: *mut u8;
+        if new_size_aligned == 0 {
+            new_ptr = ptr::null_mut()
+        } else {
+            new_ptr = unsafe { aligned_alloc(WASM_PAGE_SIZE, new_size_aligned) };
+            if new_ptr.is_null() {
+                return Err(anyhow!("Failed to grow WASM memory: allocation error"));
+            }
+        }
+        let copy_size = cmp::min(self.size, new_size_aligned);
+        if copy_size > 0 {
+            // new_ptr is not null, because new_size_aligned > 0
+            // self.ptr is not null, because self.size > 0
+            unsafe {
+                std::ptr::copy_nonoverlapping(self.ptr, new_ptr, copy_size);
+            };
+        }
+        unsafe { free(self.ptr) };
+        self.size = new_size_aligned;
+        self.ptr = new_ptr;
+        Ok(())
+    }
+    fn as_ptr(&self) -> *mut u8 {
+        self.ptr
+    }
+}
+
+// In order to use the Seastar memory allocator instead of mmap,
+// create our own MemoryCreator which directly calls aligned_alloc
+// and free, both of which came from Seastar
+pub struct ScyllaMemoryCreator;
+
+unsafe impl wasmtime::MemoryCreator for ScyllaMemoryCreator {
+    fn new_memory(
+        &self,
+        ty: wasmtime::MemoryType,
+        minimum: usize,
+        maximum: Option<usize>,
+        reserved_size_in_bytes: Option<usize>,
+        guard_size_in_bytes: usize,
+    ) -> Result<Box<dyn wasmtime::LinearMemory>, String> {
+        // assert that this is a memory that only allocates as much as it needs
+        assert_eq!(guard_size_in_bytes, 0);
+        assert!(reserved_size_in_bytes.is_none());
+        assert!(!ty.is_64());
+        let mut mem = ScyllaLinearMemory {
+            ptr: ptr::null_mut(),
+            size: 0,
+            maximum_size: maximum,
+        };
+        mem.grow_to(minimum).map_err(|e| e.to_string())?;
+        Ok(Box::new(mem))
+    }
+}