cc-lb-runtime-wasmtime 0.1.1

Wasmtime-based plugin runtime for cc-lb. Host-side wasm plugin admission + dispatch.
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

use wasmtime::{Instance, Memory, Store, TypedFunc};

use crate::cell::PluginCell;
use crate::engine::HostState;
use crate::error::WasmtimeRuntimeError;

use super::HookFn;

pub(super) struct WorkerInstance {
    pub(super) store: Store<HostState>,
    pub(super) memory: Memory,
    pub(super) alloc_fn: TypedFunc<(u32, u32), u32>,
    pub(super) free_fn: TypedFunc<(u32, u32, u32), ()>,
    pub(super) filter_fn: Option<TypedFunc<(u32, u32), u64>>,
    pub(super) shape_fn: Option<TypedFunc<(u32, u32), u64>>,
    pub(super) observe_fn: Option<TypedFunc<(u32, u32), u64>>,
}

pub(super) fn build_worker_instance(
    cell: &PluginCell,
    hook: HookFn,
) -> Result<WorkerInstance, WasmtimeRuntimeError> {
    let engine = cell.instance_pre.module().engine();
    let mut store = Store::new(engine, HostState::new(cell.memory_max_pages));
    store.limiter(|state| state.limits());

    let instance: Instance = cell.instance_pre.instantiate(&mut store).map_err(|e| {
        if e.downcast_ref::<wasmtime::PoolConcurrencyLimitError>()
            .is_some()
        {
            WasmtimeRuntimeError::PoolSaturated { resource: "pool" }
        } else {
            WasmtimeRuntimeError::InstantiateFailed(anyhow::Error::from(e))
        }
    })?;

    let memory = instance.get_memory(&mut store, "memory").ok_or_else(|| {
        WasmtimeRuntimeError::ModuleRejected {
            reason: "module does not export `memory`".into(),
        }
    })?;

    let alloc_fn = instance
        .get_typed_func::<(u32, u32), u32>(&mut store, "cc_lb_alloc")
        .map_err(|e| WasmtimeRuntimeError::ModuleRejected {
            reason: format!("missing or mistyped `cc_lb_alloc` export: {e}"),
        })?;

    let free_fn = instance
        .get_typed_func::<(u32, u32, u32), ()>(&mut store, "cc_lb_free")
        .map_err(|e| WasmtimeRuntimeError::ModuleRejected {
            reason: format!("missing or mistyped `cc_lb_free` export: {e}"),
        })?;

    let (filter_fn, shape_fn, observe_fn) = match hook {
        HookFn::Filter => (
            Some(
                instance
                    .get_typed_func::<(u32, u32), u64>(&mut store, "cc_lb_filter")
                    .map_err(|e| WasmtimeRuntimeError::ModuleRejected {
                        reason: format!("missing or mistyped `cc_lb_filter` export: {e}"),
                    })?,
            ),
            None,
            None,
        ),
        HookFn::Shape => (
            None,
            Some(
                instance
                    .get_typed_func::<(u32, u32), u64>(&mut store, "cc_lb_shape")
                    .map_err(|e| WasmtimeRuntimeError::ModuleRejected {
                        reason: format!("missing or mistyped `cc_lb_shape` export: {e}"),
                    })?,
            ),
            None,
        ),
        HookFn::Observe => (
            None,
            None,
            Some(
                instance
                    .get_typed_func::<(u32, u32), u64>(&mut store, "cc_lb_observe")
                    .map_err(|e| WasmtimeRuntimeError::ModuleRejected {
                        reason: format!("missing or mistyped `cc_lb_observe` export: {e}"),
                    })?,
            ),
        ),
    };

    Ok(WorkerInstance {
        store,
        memory,
        alloc_fn,
        free_fn,
        filter_fn,
        shape_fn,
        observe_fn,
    })
}