Skip to content

Commit

Permalink
feat(lib/runtime/wazero): implement compiler cache (#3879)
Browse files Browse the repository at this point in the history
Co-authored-by: EclesioMeloJunior <[email protected]>
Co-authored-by: Eric Tu <[email protected]>
  • Loading branch information
3 people authored Apr 24, 2024
1 parent 0e6d32b commit 3fd6ed7
Show file tree
Hide file tree
Showing 3 changed files with 102 additions and 48 deletions.
14 changes: 13 additions & 1 deletion chain/paseo/chain-spec-raw.json
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,19 @@
"/dns/paseo.bootnode.amforc.com/tcp/30333/wss/p2p/12D3KooWFD81HC9memUwuGMLvhDDEfmXjn6jC4n7zyNs3vToXapS",
"/dns/paseo.bootnode.amforc.com/tcp/30344/p2p/12D3KooWFD81HC9memUwuGMLvhDDEfmXjn6jC4n7zyNs3vToXapS",
"/dns/boot.stake.plus/tcp/43334/wss/p2p/12D3KooWNhgAC3hjZHxaT52EpPFZohkCL1AHFAijqcN8xB9Rwud2",
"/dns/boot.stake.plus/tcp/43333/p2p/12D3KooWNhgAC3hjZHxaT52EpPFZohkCL1AHFAijqcN8xB9Rwud2"
"/dns/boot.stake.plus/tcp/43333/p2p/12D3KooWNhgAC3hjZHxaT52EpPFZohkCL1AHFAijqcN8xB9Rwud2",
"/dns/boot.metaspan.io/tcp/36017/wss/p2p/12D3KooWSW6nDfM3SS8rUtjMyjdszivK31bu4a1sRngGa2hFETz7",
"/dns/boot.metaspan.io/tcp/36018/p2p/12D3KooWSW6nDfM3SS8rUtjMyjdszivK31bu4a1sRngGa2hFETz7",
"/dns/paseo.bootnodes.polkadotters.com/tcp/30538/p2p/12D3KooWPbbFy4TefEGTRF5eTYhq8LEzc4VAHdNUVCbY4nAnhqPP",
"/dns/paseo.bootnodes.polkadotters.com/tcp/30540/wss/p2p/12D3KooWPbbFy4TefEGTRF5eTYhq8LEzc4VAHdNUVCbY4nAnhqPP",
"/dns/boot-node.helikon.io/tcp/10020/p2p/12D3KooWBetfzZpf6tGihKrqCo5z854Ub4ZNAUUTRT6eYHNh7FYi",
"/dns/boot-node.helikon.io/tcp/10022/wss/p2p/12D3KooWBetfzZpf6tGihKrqCo5z854Ub4ZNAUUTRT6eYHNh7FYi",
"/dns/boot.gatotech.network/tcp/33400/p2p/12D3KooWEvz5Ygv3MhCUNTVQbUTVhzhvf4KKcNoe5M5YbVLPBeeW",
"/dns/boot.gatotech.network/tcp/35400/wss/p2p/12D3KooWEvz5Ygv3MhCUNTVQbUTVhzhvf4KKcNoe5M5YbVLPBeeW",
"/dns/paseo-bootnode.turboflakes.io/tcp/30630/p2p/12D3KooWMjCN2CrnN71hAdehn6M2iYKeGdGbZ1A3SKhf4hxrgG9e",
"/dns/paseo-bootnode.turboflakes.io/tcp/30730/wss/p2p/12D3KooWMjCN2CrnN71hAdehn6M2iYKeGdGbZ1A3SKhf4hxrgG9e",
"/dns/pso16.rotko.net/tcp/33246/p2p/12D3KooWRH8eBMhw8c7bucy6pJfy94q4dKpLkF3pmeGohHmemdRu",
"/dns/pso16.rotko.net/tcp/35246/wss/p2p/12D3KooWRH8eBMhw8c7bucy6pJfy94q4dKpLkF3pmeGohHmemdRu"
],
"telemetryEndpoints": null,
"protocolId": "pas",
Expand Down
133 changes: 88 additions & 45 deletions lib/runtime/wazero/instance.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,20 @@ const runtimeContextKey = contextKey("runtime.Context")

var _ runtime.Instance = &Instance{}

type wazeroMeta struct {
config wazero.RuntimeConfig
cache wazero.CompilationCache
guestModule wazero.CompiledModule
}

// Instance backed by wazero.Runtime
type Instance struct {
Runtime wazero.Runtime
Module api.Module
Context *runtime.Context
codeHash common.Hash
heapBase uint32
Runtime wazero.Runtime
Module api.Module
Context *runtime.Context
wasmByteCode []byte
codeHash common.Hash
metadata wazeroMeta
sync.Mutex
}

Expand Down Expand Up @@ -100,15 +107,13 @@ func NewInstanceFromTrie(t trie.Trie, cfg Config) (*Instance, error) {
return NewInstance(code, cfg)
}

// NewInstance instantiates a runtime from raw wasm bytecode
func NewInstance(code []byte, cfg Config) (instance *Instance, err error) {
logger.Info("instantiating a runtime!")
logger.Patch(log.SetLevel(cfg.LogLvl), log.SetCallerFunc(true))

ctx := context.Background()
rt := wazero.NewRuntime(ctx)
func newRuntime(ctx context.Context,
code []byte,
config wazero.RuntimeConfig,
) (api.Module, wazero.Runtime, wazero.CompiledModule, error) {
rt := wazero.NewRuntimeWithConfig(ctx, config)

_, err = rt.NewHostModuleBuilder("env").
hostCompiledModule, err := rt.NewHostModuleBuilder("env").
// values from newer kusama/polkadot runtimes
ExportMemory("memory", 23).
NewFunctionBuilder().
Expand Down Expand Up @@ -393,38 +398,51 @@ func NewInstance(code []byte, cfg Config) (instance *Instance, err error) {
NewFunctionBuilder().
WithFunc(ext_crypto_ecdsa_generate_version_1).
Export("ext_crypto_ecdsa_generate_version_1").
Instantiate(ctx)
Compile(ctx)

if err != nil {
return nil, err
return nil, nil, nil, err
}

_, err = rt.InstantiateModule(ctx, hostCompiledModule, wazero.NewModuleConfig())
if err != nil {
return nil, nil, nil, err
}

code, err = decompressWasm(code)
if err != nil {
return nil, err
return nil, nil, nil, err
}

guestCompiledModule, err := rt.CompileModule(ctx, code)
if err != nil {
return nil, nil, nil, err
}
mod, err := rt.Instantiate(ctx, code)
if err != nil {
return nil, err
return nil, nil, nil, err
}

encodedHeapBase := mod.ExportedGlobal("__heap_base")
if encodedHeapBase == nil {
return nil, fmt.Errorf("wazero error: nil global for __heap_base")
}
return mod, rt, guestCompiledModule, nil
}

heapBase := api.DecodeU32(encodedHeapBase.Get())
// hb = runtime.DefaultHeapBase
// NewInstance instantiates a runtime from raw wasm bytecode
func NewInstance(code []byte, cfg Config) (instance *Instance, err error) {
logger.Debug("instantiating a runtime!")
logger.Patch(log.SetLevel(cfg.LogLvl), log.SetCallerFunc(true))

mem := mod.Memory()
if mem == nil {
return nil, fmt.Errorf("wazero error: nil memory for module")
// Prepare a cache directory.
ctx := context.Background()
cache := wazero.NewCompilationCache()
config := wazero.NewRuntimeConfig().WithCompilationCache(cache)
mod, rt, guestCompiledModule, err := newRuntime(ctx, code, config)
if err != nil {
return nil, fmt.Errorf("creating runtime instance: %w", err)
}

instance = &Instance{
heapBase: heapBase,
Runtime: rt,
wasmByteCode: code,
Runtime: rt,
Context: &runtime.Context{
Keystore: cfg.Keystore,
Validator: cfg.Role == common.AuthorityRole,
Expand All @@ -436,6 +454,11 @@ func NewInstance(code []byte, cfg Config) (instance *Instance, err error) {
},
Module: mod,
codeHash: cfg.CodeHash,
metadata: wazeroMeta{
config: config,
cache: cache,
guestModule: guestCompiledModule,
},
}

if cfg.DefaultVersion == nil {
Expand All @@ -456,34 +479,50 @@ func NewInstance(code []byte, cfg Config) (instance *Instance, err error) {

var ErrExportFunctionNotFound = errors.New("export function not found")

func (i *Instance) Exec(function string, data []byte) (result []byte, err error) {
func (i *Instance) Exec(function string, data []byte) ([]byte, error) {
i.Lock()
i.Context.Allocator = allocator.NewFreeingBumpHeapAllocator(i.heapBase)
defer i.Unlock()

mod, err := i.Runtime.InstantiateModule(context.Background(), i.metadata.guestModule, wazero.NewModuleConfig())
if mod == nil {
return nil, fmt.Errorf("instantiate guest module: nil")
}
if err != nil {
return nil, fmt.Errorf("instantiate guest module: %w", err)
}

defer func() {
i.Context.Allocator = nil
i.Unlock()
err = mod.Close(context.Background())
if err != nil {
logger.Criticalf("guest module not closed: %w", err)
}
}()
// instantiate a new allocator on every execution func

dataLength := uint32(len(data))
inputPtr, err := i.Context.Allocator.Allocate(i.Module.Memory(), dataLength)
if err != nil {
return nil, fmt.Errorf("allocating input memory: %w", err)
encodedHeapBase := mod.ExportedGlobal("__heap_base")
if encodedHeapBase == nil {
return nil, fmt.Errorf("wazero error: nil global for __heap_base")
}

// Store the data into memory
mem := i.Module.Memory()
if mem == nil {
heapBase := api.DecodeU32(encodedHeapBase.Get())
i.Context.Allocator = allocator.NewFreeingBumpHeapAllocator(heapBase)

memory := mod.Memory()
if memory == nil {
panic("nil memory")
}

ok := mem.Write(inputPtr, data)
dataLength := uint32(len(data))
inputPtr, err := i.Context.Allocator.Allocate(memory, dataLength)
if err != nil {
return nil, fmt.Errorf("allocating input memory: %w", err)
}

ok := memory.Write(inputPtr, data)
if !ok {
panic("write overflow")
}

runtimeFunc := i.Module.ExportedFunction(function)
runtimeFunc := mod.ExportedFunction(function)
if runtimeFunc == nil {
return nil, fmt.Errorf("%w: %s", ErrExportFunctionNotFound, function)
}
Expand All @@ -497,11 +536,10 @@ func (i *Instance) Exec(function string, data []byte) (result []byte, err error)
return nil, fmt.Errorf("no returned values from runtime function: %s", function)
}
wasmValue := values[0]

outputPtr, outputLength := splitPointerSize(wasmValue)
result, ok = mem.Read(outputPtr, outputLength)
result, ok := memory.Read(outputPtr, outputLength)
if !ok {
panic("write overflow")
panic("read overflow")
}

return result, nil
Expand Down Expand Up @@ -899,4 +937,9 @@ func (in *Instance) Stop() {
if err != nil {
log.Errorf("runtime failed to close: %v", err)
}

err = in.metadata.cache.Close(context.Background())
if err != nil {
log.Errorf("closing the wazero compilation cache: %v", err)
}
}
3 changes: 1 addition & 2 deletions lib/runtime/wazero/instance_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ import (
"github.com/ChainSafe/gossamer/lib/utils"
"github.com/ChainSafe/gossamer/pkg/scale"
"github.com/ChainSafe/gossamer/pkg/trie"
"github.com/ChainSafe/gossamer/pkg/trie/inmemory"
inmemory_trie "github.com/ChainSafe/gossamer/pkg/trie/inmemory"
"github.com/centrifuge/go-substrate-rpc-client/v4/signature"

Expand Down Expand Up @@ -1096,7 +1095,7 @@ func newTrieFromPairs(t *testing.T, filename string) trie.Trie {
entries[pairArr[0].(string)] = pairArr[1].(string)
}

tr, err := inmemory.LoadFromMap(entries, trie.V0)
tr, err := inmemory_trie.LoadFromMap(entries, trie.V0)
require.NoError(t, err)
return tr
}
Expand Down

0 comments on commit 3fd6ed7

Please sign in to comment.