chore (optimisation): optimise pkg/extension

This commit is contained in:
MickaelK
2026-09-02 14:42:41 +10:00
parent ff977b6240
commit 280e4e089a
7 changed files with 88 additions and 55 deletions
@@ -22,9 +22,9 @@ func stateAuthorisation(ctx context.Context) *authData {
func exportAuthorisation(b *runtime.HostModuleBuilder) {
b.Export("ffi_authorisation_pull_path", func(ctx context.Context, mem runtime.IMemory, outPtr, outCap uint32) uint32 {
return mem.Write(outPtr, outCap, []byte(stateAuthorisation(ctx).path))
return mem.WriteString(outPtr, outCap, stateAuthorisation(ctx).path)
}).Export("ffi_authorisation_pull_target", func(ctx context.Context, mem runtime.IMemory, outPtr, outCap uint32) uint32 {
return mem.Write(outPtr, outCap, []byte(stateAuthorisation(ctx).target))
return mem.WriteString(outPtr, outCap, stateAuthorisation(ctx).target)
}).Export("ffi_authorisation_push_allow", func(ctx context.Context) {
stateAuthorisation(ctx).allow = true
})
+4 -3
View File
@@ -38,11 +38,12 @@ type route struct {
}
func exportHttp(b *runtime.HostModuleBuilder) {
b.Export("ffi_http_push_route", func(ctx context.Context, mem runtime.IMemory, mPtr, mLen, pPtr, pLen, wPtr, wLen uint32) {
b.Export("ffi_http_push_route", func(ctx context.Context, mem runtime.IMemory, rPtr, rLen, wPtr, wLen uint32) {
routes, _ := ctx.Value(routesKey{}).(*[]route)
line := strings.SplitN(string(mem.Read(rPtr, rLen)), " ", 2)
r := route{
method: string(mem.Read(mPtr, mLen)),
path: string(mem.Read(pPtr, pLen)),
method: line[0],
path: line[1],
}
if names := string(mem.Read(wPtr, wLen)); names != "" {
r.middleware = strings.Split(names, ",")
+4 -3
View File
@@ -15,6 +15,7 @@ type Instance struct {
func NewInstance(wasm []byte) (*Instance, error) {
rt, err := runtime.New(wasm, runtime.WithExports(func(b *runtime.HostModuleBuilder) {
exportShared(b)
exportAuthentication(b)
exportAuthorisation(b)
exportMiddleware(b)
exportHttp(b)
@@ -70,10 +71,10 @@ func exportShared(b *runtime.HostModuleBuilder) {
}).Export("resp_write", func(ctx context.Context, mem runtime.IMemory, ptr, length uint32) {
stateHttp(ctx).w.Write(mem.Read(ptr, length))
}).Export("req_method", func(ctx context.Context, mem runtime.IMemory, outPtr, outCap uint32) uint32 {
return mem.Write(outPtr, outCap, []byte(stateHttp(ctx).r.Method))
return mem.WriteString(outPtr, outCap, stateHttp(ctx).r.Method)
}).Export("req_path", func(ctx context.Context, mem runtime.IMemory, outPtr, outCap uint32) uint32 {
return mem.Write(outPtr, outCap, []byte(stateHttp(ctx).r.URL.Path))
return mem.WriteString(outPtr, outCap, stateHttp(ctx).r.URL.Path)
}).Export("req_header", func(ctx context.Context, mem runtime.IMemory, nPtr, nLen, outPtr, outCap uint32) uint32 {
return mem.Write(outPtr, outCap, []byte(stateHttp(ctx).r.Header.Get(string(mem.Read(nPtr, nLen)))))
return mem.WriteString(outPtr, outCap, stateHttp(ctx).r.Header.Get(string(mem.Read(nPtr, nLen))))
})
}
+8 -34
View File
@@ -1,26 +1,20 @@
package runtime
import (
"reflect"
"github.com/tetratelabs/wazero/api"
)
type IMemory interface {
Read(ptr, length uint32) []byte
Write(ptr, cap uint32, data []byte) uint32
WriteString(ptr, cap uint32, data string) uint32
}
type wazeroMem struct{ m api.Memory }
func (w wazeroMem) Read(ptr, length uint32) []byte {
b, ok := w.m.Read(ptr, length)
if !ok {
return nil
}
cp := make([]byte, len(b))
copy(cp, b)
return cp
b, _ := w.m.Read(ptr, length)
return b
}
func (w wazeroMem) Write(ptr, cap uint32, data []byte) uint32 {
@@ -31,30 +25,10 @@ func (w wazeroMem) Write(ptr, cap uint32, data []byte) uint32 {
return n
}
var (
memType = reflect.TypeOf((*IMemory)(nil)).Elem()
modType = reflect.TypeOf((*api.Module)(nil)).Elem()
)
// convert IMemory to wasmMemory so adapters don't need to know about wazero
func wazeroMemoryAdapter(fn any) any {
fv := reflect.ValueOf(fn)
ft := fv.Type()
if ft.NumIn() > 1 && ft.In(1) == memType {
in := make([]reflect.Type, ft.NumIn())
for i := range in {
in[i] = ft.In(i)
}
in[1] = modType
out := make([]reflect.Type, ft.NumOut())
for i := range out {
out[i] = ft.Out(i)
}
return reflect.MakeFunc(reflect.FuncOf(in, out, false), func(args []reflect.Value) []reflect.Value {
mod := args[1].Interface().(api.Module)
args[1] = reflect.ValueOf(IMemory(wazeroMem{mod.Memory()}))
return fv.Call(args)
}).Interface()
func (w wazeroMem) WriteString(ptr, cap uint32, data string) uint32 {
n := uint32(len(data))
if n > cap || (n > 0 && !w.m.WriteString(ptr, data)) {
return 0
}
return fn
return n
}
+56 -1
View File
@@ -4,6 +4,7 @@ import (
"context"
"github.com/tetratelabs/wazero"
"github.com/tetratelabs/wazero/api"
)
type HostModuleBuilder struct {
@@ -14,8 +15,62 @@ func NewHostModuleBuilder(wrt wazero.Runtime, name string) *HostModuleBuilder {
return &HostModuleBuilder{inner: wrt.NewHostModuleBuilder(name)}
}
var i32 = api.ValueTypeI32
func (b *HostModuleBuilder) Export(name string, fn any) *HostModuleBuilder {
b.inner.NewFunctionBuilder().WithFunc(wazeroMemoryAdapter(fn)).Export(name)
var (
gofn api.GoModuleFunction
params []api.ValueType
results []api.ValueType
cached IMemory
)
mem := func(mod api.Module) IMemory {
if cached == nil {
cached = wazeroMem{mod.Memory()}
}
return cached
}
switch f := fn.(type) {
case func(context.Context):
gofn = api.GoModuleFunc(func(ctx context.Context, mod api.Module, stack []uint64) {
f(ctx)
})
case func(context.Context, uint32):
params = []api.ValueType{i32}
gofn = api.GoModuleFunc(func(ctx context.Context, mod api.Module, stack []uint64) {
f(ctx, uint32(stack[0]))
})
case func(context.Context, IMemory, uint32, uint32):
params = []api.ValueType{i32, i32}
gofn = api.GoModuleFunc(func(ctx context.Context, mod api.Module, stack []uint64) {
f(ctx, mem(mod), uint32(stack[0]), uint32(stack[1]))
})
case func(context.Context, IMemory, uint32, uint32) uint32:
params = []api.ValueType{i32, i32}
results = []api.ValueType{i32}
gofn = api.GoModuleFunc(func(ctx context.Context, mod api.Module, stack []uint64) {
stack[0] = uint64(f(ctx, mem(mod), uint32(stack[0]), uint32(stack[1])))
})
case func(context.Context, IMemory, uint32, uint32, uint32):
params = []api.ValueType{i32, i32, i32}
gofn = api.GoModuleFunc(func(ctx context.Context, mod api.Module, stack []uint64) {
f(ctx, mem(mod), uint32(stack[0]), uint32(stack[1]), uint32(stack[2]))
})
case func(context.Context, IMemory, uint32, uint32, uint32, uint32):
params = []api.ValueType{i32, i32, i32, i32}
gofn = api.GoModuleFunc(func(ctx context.Context, mod api.Module, stack []uint64) {
f(ctx, mem(mod), uint32(stack[0]), uint32(stack[1]), uint32(stack[2]), uint32(stack[3]))
})
case func(context.Context, IMemory, uint32, uint32, uint32, uint32) uint32:
params = []api.ValueType{i32, i32, i32, i32}
results = []api.ValueType{i32}
gofn = api.GoModuleFunc(func(ctx context.Context, mod api.Module, stack []uint64) {
stack[0] = uint64(f(ctx, mem(mod), uint32(stack[0]), uint32(stack[1]), uint32(stack[2]), uint32(stack[3])))
})
default:
panic("runtime: unsupported export signature for " + name)
}
b.inner.NewFunctionBuilder().WithGoModuleFunction(gofn, params, results).Export(name)
return b
}
@@ -14,8 +14,9 @@ type Runtime struct {
wrt wazero.Runtime
ctx context.Context
mu sync.Mutex
mod api.Module
mu sync.Mutex
mod api.Module
_cache map[string]api.Function
}
func New(wasm []byte, opts ...Option) (*Runtime, error) {
@@ -40,7 +41,7 @@ func New(wasm []byte, opts ...Option) (*Runtime, error) {
wrt.Close(ctx)
return nil, err
}
return &Runtime{ctx: ctx, wrt: wrt, mod: mod}, nil
return &Runtime{ctx: ctx, wrt: wrt, mod: mod, _cache: map[string]api.Function{}}, nil
}
func (r *Runtime) HasExport(name string) bool {
@@ -50,7 +51,11 @@ func (r *Runtime) HasExport(name string) bool {
func (r *Runtime) Call(ctx context.Context, fnName string, key, val any) error {
r.mu.Lock()
defer r.mu.Unlock()
fn := r.mod.ExportedFunction(fnName)
fn, ok := r._cache[fnName]
if !ok {
fn = r.mod.ExportedFunction(fnName)
r._cache[fnName] = fn
}
if fn == nil {
return fmt.Errorf("%w: %s", ErrNoExport, fnName)
}
+5 -8
View File
@@ -3,10 +3,8 @@ use crate::{Request, Response};
#[link(wasm_import_module = "env")]
extern "C" {
fn ffi_http_push_route(
method: *const u8,
method_len: u32,
path: *const u8,
path_len: u32,
route: *const u8,
route_len: u32,
middleware: *const u8,
middleware_len: u32,
);
@@ -50,13 +48,12 @@ impl<A> Router<A> {
pub fn describe(&self) {
for route in &self.routes {
let line = format!("{} {}", route.method, route.path);
let middleware = route.middleware.join(",");
unsafe {
ffi_http_push_route(
route.method.as_ptr(),
route.method.len() as u32,
route.path.as_ptr(),
route.path.len() as u32,
line.as_ptr(),
line.len() as u32,
middleware.as_ptr(),
middleware.len() as u32,
);