lunatic_process/runtimes/
mod.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
//! WebAssembly runtimes powering lunatic.
//!
//! Currently only Wasmtime is supported, but it should be "easy" to add any runtime that has a
//! `Linker` abstraction and supports `async` host functions.
//!
//! NOTE: This traits are not used at all. Until rust supports async-traits all functions working
//!       with a runtime will directly take `wasmtime::WasmtimeRuntime` instead of a generic.

use std::sync::Arc;

use anyhow::Result;
use dashmap::DashMap;
use tokio::task::JoinHandle;

use crate::state::ProcessState;

use self::wasmtime::{WasmtimeCompiledModule, WasmtimeRuntime};

pub mod wasmtime;

pub struct RawWasm {
    // Id returned by control and used when spawning modules on other nodes
    pub id: Option<u64>,
    pub bytes: Vec<u8>,
}

impl RawWasm {
    pub fn new(id: Option<u64>, bytes: Vec<u8>) -> Self {
        Self { id, bytes }
    }

    pub fn as_slice(&self) -> &[u8] {
        self.bytes.as_slice()
    }
}

impl From<Vec<u8>> for RawWasm {
    fn from(bytes: Vec<u8>) -> Self {
        Self::new(None, bytes)
    }
}

/// A `WasmRuntime` is a compiler that can generate runnable code from raw .wasm files.
///
/// It also provides a mechanism to register host functions that are accessible to the wasm guest
/// code through the generic type `T`. The type `T` must implement the [`ProcessState`] trait and
/// expose a `register` function for host functions.
pub trait WasmRuntime<T>: Clone
where
    T: crate::state::ProcessState + Default + Send,
{
    type WasmInstance: WasmInstance;

    /// Takes a raw binary WebAssembly module and returns the index of a compiled module.
    fn compile_module(&mut self, data: RawWasm) -> anyhow::Result<usize>;

    /// Returns a reference to the raw binary WebAssembly module if the index exists.
    fn wasm_module(&self, index: usize) -> Option<&RawWasm>;

    // Creates a wasm instance from compiled module if the index exists.
    /* async fn instantiate(
        &self,
        index: usize,
        state: T,
        config: ProcessConfig,
    ) -> Result<WasmtimeInstance<T>>; */
}

pub trait WasmInstance {
    type Param;

    // Calls a wasm function by name with the specified arguments. Ignores the returned values.
    /* async fn call(&mut self, function: &str, params: Vec<Self::Param>) -> Result<()>; */
}

pub struct Modules<T> {
    modules: Arc<DashMap<u64, Arc<WasmtimeCompiledModule<T>>>>,
}

impl<T> Clone for Modules<T> {
    fn clone(&self) -> Self {
        Self {
            modules: self.modules.clone(),
        }
    }
}

impl<T> Default for Modules<T> {
    fn default() -> Self {
        Self {
            modules: Arc::new(DashMap::new()),
        }
    }
}

impl<T: ProcessState + 'static> Modules<T> {
    pub fn get(&self, module_id: u64) -> Option<Arc<WasmtimeCompiledModule<T>>> {
        self.modules.get(&module_id).map(|m| m.clone())
    }

    pub fn compile(
        &self,
        runtime: WasmtimeRuntime,
        wasm: RawWasm,
    ) -> JoinHandle<Result<Arc<WasmtimeCompiledModule<T>>>> {
        let modules = self.modules.clone();
        tokio::task::spawn_blocking(move || {
            let id = wasm.id;
            match runtime.compile_module(wasm) {
                Ok(m) => {
                    let module = Arc::new(m);
                    if let Some(id) = id {
                        modules.insert(id, Arc::clone(&module));
                    }
                    Ok(module)
                }
                Err(e) => Err(e),
            }
        })
    }
}