streamdal/go-sdk

package streamdal

import (
    "context"
    "fmt"
    "io"
    "sync"

    "github.com/pkg/errors"
    "github.com/tetratelabs/wazero"
    "github.com/tetratelabs/wazero/api"
    "github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"

    "github.com/streamdal/streamdal/libs/protos/build/go/protos"
)

type function struct {
    ID      string
    Inst    api.Module
    entry   api.Function
    alloc   api.Function
    dealloc api.Function
    mtx     *sync.Mutex
}

func (f *function) Exec(ctx context.Context, req []byte) ([]byte, error) {
    ptrLen := uint64(len(req))

    inputPtr, err := f.alloc.Call(ctx, ptrLen)
    if err != nil {
        return nil, errors.Wrap(err, "unable to allocate memory")
    }

    if len(inputPtr) == 0 {
        return nil, errors.New("unable to allocate memory")
    }

    ptrVal := inputPtr[0]

    if !f.Inst.Memory().Write(uint32(ptrVal), req) {
        return nil, fmt.Errorf("Memory.Write(%d, %d) out of range of memory size %d",
            ptrVal, len(req), f.Inst.Memory().Size())
    }

    result, err := f.entry.Call(ctx, ptrVal, ptrLen)
    if err != nil {
        // Clear mem on error
        if _, err := f.dealloc.Call(ctx, ptrVal, ptrLen); err != nil {
            return nil, errors.Wrap(err, "unable to deallocate memory")
        }
        return nil, errors.Wrap(err, "error during func call")
    }

    resultPtr := uint32(result[0] >> 32)
    resultSize := uint32(result[0])

    // Dealloc request memory
    if _, err := f.dealloc.Call(ctx, ptrVal, ptrLen); err != nil {
        return nil, errors.Wrap(err, "unable to deallocate memory")
    }

    // Read memory starting from result ptr
    resBytes, err := f.readMemory(resultPtr, resultSize)
    if err != nil {
        // Dealloc response memory
        if _, err := f.dealloc.Call(ctx, uint64(resultPtr), uint64(resultSize)); err != nil {
            return nil, errors.Wrap(err, "unable to deallocate memory")
        }
        return nil, errors.Wrap(err, "unable to read memory")
    }

    // Dealloc response memory
    if _, err := f.dealloc.Call(ctx, uint64(resultPtr), uint64(resultSize)); err != nil {
        return nil, errors.Wrap(err, "unable to deallocate memory")
    }

    return resBytes, nil
}

func (s *Streamdal) setFunctionCache(wasmID string, f *function) {
    s.functionsMtx.Lock()
    defer s.functionsMtx.Unlock()

    s.functions[wasmID] = f
}

func (s *Streamdal) getFunction(_ context.Context, step *protos.PipelineStep) (*function, error) {
    // check cache
    fc, ok := s.getFunctionFromCache(step.GetXWasmId())
    if ok {
        return fc, nil
    }

    fi, err := s.createFunction(step)
    if err != nil {
        return nil, errors.Wrap(err, "failed to create function")
    }

    // Cache function
    s.setFunctionCache(step.GetXWasmId(), fi)

    return fi, nil
}

func (s *Streamdal) getFunctionFromCache(wasmID string) (*function, bool) {
    s.functionsMtx.RLock()
    defer s.functionsMtx.RUnlock()

    f, ok := s.functions[wasmID]
    return f, ok
}

func (s *Streamdal) createFunction(step *protos.PipelineStep) (*function, error) {
    inst, err := s.createWASMInstance(step.GetXWasmBytes())
    if err != nil {
        return nil, errors.Wrap(err, "unable to create WASM instance")
    }

    // This is the actual function we'll be executing
    f := inst.ExportedFunction(step.GetXWasmFunction())
    if f == nil {
        return nil, fmt.Errorf("unable to get exported function '%s'", step.GetXWasmFunction())
    }

    // alloc allows us to pre-allocate memory in order to pass data to the WASM module
    alloc := inst.ExportedFunction("alloc")
    if alloc == nil {
        return nil, errors.New("unable to get alloc func")
    }

    // dealloc allows us to free memory passed to the wasm module after we're done with it
    dealloc := inst.ExportedFunction("dealloc")
    if dealloc == nil {
        return nil, errors.New("unable to get dealloc func")
    }

    return &function{
        ID:      step.GetXWasmId(),
        Inst:    inst,
        entry:   f,
        alloc:   alloc,
        dealloc: dealloc,
        mtx:     &sync.Mutex{},
    }, nil
}

func (s *Streamdal) createWASMInstance(wasmBytes []byte) (api.Module, error) {
    if len(wasmBytes) == 0 {
        return nil, errors.New("wasm data is empty")
    }

    hostFuncs := map[string]func(_ context.Context, module api.Module, ptr, length int32) uint64{
        "kvExists":    s.hf.KVExists,
        "httpRequest": s.hf.HTTPRequest,
    }

    rCfg := wazero.NewRuntimeConfig().
        WithMemoryLimitPages(1000) // 64MB (default is 1MB)

    ctx := context.Background()
    r := wazero.NewRuntimeWithConfig(ctx, rCfg)

    wasi_snapshot_preview1.MustInstantiate(ctx, r)

    cfg := wazero.NewModuleConfig().
        WithStderr(io.Discard).
        WithStdout(io.Discard).
        WithSysNanotime().
        WithSysNanosleep().
        WithSysWalltime().
        WithStartFunctions("") // We don't need _start() to be called for our purposes

    builder := r.NewHostModuleBuilder("env")

    // This is how multiple host funcs are exported:
    // https://github.com/tetratelabs/wazero/blob/b7e8191cceb83c7335d6b8922b40b957475beecf/examples/import-go/age-calculator.go#L41
    for name, fn := range hostFuncs {
        builder = builder.NewFunctionBuilder().
            WithFunc(fn).
            Export(name)
    }

    if _, err := builder.Instantiate(ctx); err != nil {
        return nil, errors.Wrap(err, "failed to instantiate module")
    }

    mod, err := r.InstantiateWithConfig(ctx, wasmBytes, cfg)
    if err != nil {
        return nil, errors.Wrap(err, "failed to instantiate wasm module")
    }

    return mod, nil
}

func (f *function) readMemory(ptr, length uint32) ([]byte, error) {
    mem, ok := f.Inst.Memory().Read(ptr, length)
    if !ok {
        return nil, fmt.Errorf("unable to read memory at '%d' with length '%d'", ptr, length)
    }

    return mem, nil

}