148 lines
4.9 KiB
Python
148 lines
4.9 KiB
Python
from typing import Any, Dict, List, Union
|
|
|
|
import ctypes
|
|
import functools
|
|
import struct
|
|
from queue import Empty, Queue
|
|
|
|
import wasmtime
|
|
|
|
from phasmplatform.common import valuetype
|
|
from phasmplatform.common.container import Container
|
|
from phasmplatform.common.method import Method
|
|
from phasmplatform.common.methodcall import MethodCall, MethodCallError, MethodNotFoundError
|
|
from phasmplatform.common.router import MethodCallRouterInterface
|
|
from phasmplatform.common.value import Value
|
|
from phasmplatform.common.valuetype import ValueType
|
|
|
|
from .base import BaseRunner, WasmValue
|
|
|
|
|
|
class WasmTimeRunner(BaseRunner):
|
|
__slots__ = ('store', 'module', 'instance', 'exports')
|
|
|
|
def __init__(self, method_call_router: MethodCallRouterInterface, container: Container) -> None:
|
|
super().__init__(method_call_router, container)
|
|
|
|
with open(f'./{container.image.path}', 'rb') as fil: # TODO: ImageLoader?
|
|
wasm_bin = fil.read()
|
|
|
|
self.store = wasmtime.Store()
|
|
self.module = wasmtime.Module(self.store.engine, wasm_bin)
|
|
|
|
import_map: Dict[str, wasmtime.Func] = {
|
|
f'{imprt_service}.{imprt_method.name}': wasmtime.Func(
|
|
self.store,
|
|
build_func_type(imprt_method),
|
|
functools.partial(self.send_service_call, imprt_service, imprt_method)
|
|
)
|
|
for (imprt_service, imprt_method, ) in container.image.imports
|
|
}
|
|
|
|
# Make sure the given import lists order matches the one given by wasmtime
|
|
# Otherwise, wasmtime can't match them up.
|
|
imports: List[wasmtime.Func] = [
|
|
import_map[f'{imprt.module}.{imprt.name}']
|
|
for imprt in self.module.imports
|
|
]
|
|
|
|
self.instance = wasmtime.Instance(self.store, self.module, imports)
|
|
|
|
self.exports = self.instance.exports(self.store)
|
|
|
|
def alloc_bytes(self, data: bytes) -> int:
|
|
memory = self.exports['memory']
|
|
assert isinstance(memory, wasmtime.Memory) # type hint
|
|
|
|
data_ptr = memory.data_ptr(self.store)
|
|
data_len = memory.data_len(self.store)
|
|
|
|
alloc_bytes = self.exports['stdlib.types.__alloc_bytes__']
|
|
assert isinstance(alloc_bytes, wasmtime.Func)
|
|
|
|
ptr = alloc_bytes(self.store, len(data))
|
|
assert isinstance(ptr, int) # type hint
|
|
|
|
idx = ptr + 4 # Skip the header from header from __alloc_bytes__
|
|
for byt in data:
|
|
assert idx < data_len
|
|
data_ptr[idx] = ctypes.c_ubyte(byt)
|
|
idx += 1
|
|
|
|
return ptr
|
|
|
|
def read_bytes(self, ptr: int) -> bytes:
|
|
memory = self.exports['memory']
|
|
assert isinstance(memory, wasmtime.Memory) # type hint
|
|
|
|
data_ptr = memory.data_ptr(self.store)
|
|
data_len = memory.data_len(self.store)
|
|
|
|
raw = ctypes.string_at(data_ptr, data_len)
|
|
|
|
length, = struct.unpack('<I', raw[ptr:ptr + 4]) # Header prefixed by __alloc_bytes__
|
|
|
|
return raw[ptr + 4:ptr + 4 + length]
|
|
|
|
def do_call(self, call: MethodCall) -> Union[Value, MethodCallError]:
|
|
try:
|
|
wasm_method = self.exports[call.method.name]
|
|
except KeyError:
|
|
return MethodNotFoundError()
|
|
|
|
assert isinstance(wasm_method, wasmtime.Func)
|
|
|
|
act_args = [self.value_to_wasm(x) for x in call.args]
|
|
result = wasm_method(self.store, *act_args)
|
|
assert result is None or isinstance(result, (int, float, )) # type hint
|
|
|
|
return self.value_from_wasm(call.method.return_type, result)
|
|
|
|
def send_service_call(self, service: str, method: Method, *args: Any) -> WasmValue:
|
|
assert len(method.arg_types) == len(args) # type hint
|
|
|
|
call_args = [
|
|
self.value_from_wasm(x, y)
|
|
for x, y in zip(method.arg_types, args)
|
|
]
|
|
|
|
queue: Queue[Value] = Queue(maxsize=1)
|
|
|
|
def on_result(res: Union[Value, MethodCallError]) -> None:
|
|
if isinstance(res, Value):
|
|
queue.put(res)
|
|
else:
|
|
raise Exception(res)
|
|
|
|
call = MethodCall(method, call_args)
|
|
|
|
self.method_call_router.route_call(service, call, self.container, on_result)
|
|
|
|
try:
|
|
value = queue.get(block=True, timeout=10)
|
|
except Empty:
|
|
raise Exception('Did not receive value from remote call') # TODO
|
|
|
|
return self.value_to_wasm(value)
|
|
|
|
|
|
def build_func_type(method: Method) -> wasmtime.FuncType:
|
|
if method.return_type is valuetype.none:
|
|
returns = []
|
|
else:
|
|
returns = [build_wasm_type(method.return_type)]
|
|
|
|
args = []
|
|
for arg_type in method.arg_types:
|
|
assert arg_type is not valuetype.none # type hint
|
|
args.append(build_wasm_type(arg_type))
|
|
|
|
return wasmtime.FuncType(args, returns)
|
|
|
|
|
|
def build_wasm_type(value_type: ValueType) -> wasmtime.ValType:
|
|
if value_type is valuetype.bytes:
|
|
return wasmtime.ValType.i32() # Bytes are passed as pointer
|
|
|
|
raise NotImplementedError
|