feat: wazero fork upgrade to use uint64 for memory.Size() (#3800)

go.mod

@@ -202,6 +202,6 @@ require (
 
 go 1.21
 
-replace github.com/tetratelabs/wazero => github.com/ChainSafe/wazero v0.0.0-20231114190045-1d874d099362
+replace github.com/tetratelabs/wazero => github.com/ChainSafe/wazero v0.0.0-20240319130522-78b21a59bd5f
 
 replace github.com/centrifuge/go-substrate-rpc-client/v4 => github.com/timwu20/go-substrate-rpc-client/v4 v4.0.0-20231110032757-3d8e441b7303

go.sum

@@ -10,8 +10,8 @@ git.apache.org/thrift.git v0.0.0-20180902110319-2566ecd5d999/go.mod h1:fPE2ZNJGy
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/ChainSafe/go-schnorrkel v1.1.0 h1:rZ6EU+CZFCjB4sHUE1jIu8VDoB/wRKZxoe1tkcO71Wk=
 github.com/ChainSafe/go-schnorrkel v1.1.0/go.mod h1:ABkENxiP+cvjFiByMIZ9LYbRoNNLeBLiakC1XeTFxfE=
-github.com/ChainSafe/wazero v0.0.0-20231114190045-1d874d099362 h1:hbvvSSB436JJalwq/2fRZwJpptvq9HMOLYVZX9oVHKM=
-github.com/ChainSafe/wazero v0.0.0-20231114190045-1d874d099362/go.mod h1:0U0G41+ochRKoPKCJlh0jMg1CHkyfK8kDqiirMmKY8A=
+github.com/ChainSafe/wazero v0.0.0-20240319130522-78b21a59bd5f h1:TUrrP3YSwSv2TB3Q02dWJD7A3VQoCzEQ2LNhX9sb0Jo=
+github.com/ChainSafe/wazero v0.0.0-20240319130522-78b21a59bd5f/go.mod h1:ytl6Zuh20R/eROuyDaGPkp82O9C/DJfXAwJfQ3X6/7Y=
 github.com/DataDog/zstd v1.4.1/go.mod h1:1jcaCB/ufaK+sKp1NBhlGmpz41jOoPQ35bpF36t7BBo=
 github.com/DataDog/zstd v1.4.5 h1:EndNeuB0l9syBZhut0wns3gV1hL8zX8LIu6ZiVHWLIQ=
 github.com/DataDog/zstd v1.4.5/go.mod h1:1jcaCB/ufaK+sKp1NBhlGmpz41jOoPQ35bpF36t7BBo=

lib/runtime/allocator/freeing_bump.go

@@ -10,6 +10,7 @@ import (
 	"math/big"
 	"math/bits"
 
+	"github.com/ChainSafe/gossamer/internal/log"
 	"github.com/ChainSafe/gossamer/lib/runtime"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -37,7 +38,7 @@ const (
 	MaxPossibleAllocations uint32 = 33554432
 
 	PageSize     = 65536
-	MaxWasmPages = (4 * 1024 * 1024 * 1024 / PageSize) - 1
+	MaxWasmPages = 4 * 1024 * 1024 * 1024 / PageSize
 )
 
 var (
@@ -58,6 +59,10 @@ var (
 		Help: "the amount of address space (in bytes) used by the allocator this is calculated as " +
 			"the difference between the allocator's bumper and the heap base.",
 	})
+
+	logger = log.NewFromGlobal(
+		log.AddContext("pkg", "runtime-allocator"),
+	)
 )
 
 var (
@@ -332,7 +337,7 @@ type FreeingBumpHeapAllocator struct {
 	bumper                 uint32
 	freeLists              *FreeLists
 	poisoned               bool
-	lastObservedMemorySize uint32
+	lastObservedMemorySize uint64
 	stats                  AllocationStats
 }
 
@@ -393,7 +398,7 @@ func (f *FreeingBumpHeapAllocator) Allocate(mem runtime.Memory, size uint32) (pt
 	link := f.freeLists.heads[order]
 	switch value := link.(type) {
 	case Ptr:
-		if uint64(value.headerPtr)+uint64(order.size())+uint64(HeaderSize) > uint64(mem.Size()) {
+		if uint64(value.headerPtr)+uint64(order.size())+uint64(HeaderSize) > mem.Size() {
 			return 0, fmt.Errorf("%w: pointer: %d, order size: %d",
 				ErrInvalidHeaderPointerDetected, value.headerPtr, order.size())
 		}
@@ -505,16 +510,15 @@ func (f *FreeingBumpHeapAllocator) Deallocate(mem runtime.Memory, ptr uint32) (e
 func bump(bumper *uint32, size uint32, mem runtime.Memory) (uint32, error) {
 	requiredSize := uint64(*bumper) + uint64(size)
 
-	if requiredSize > uint64(mem.Size()) {
+	if requiredSize > mem.Size() {
 		requiredPages, ok := pagesFromSize(requiredSize)
 		if !ok {
-			return 0, fmt.Errorf("%w: required size %d dont fit uint32",
-				ErrAllocatorOutOfSpace, requiredSize)
+			return 0, fmt.Errorf("%w: cannot calculate number of pages from size %d", ErrAllocatorOutOfSpace, requiredSize)
 		}
 
-		currentPages := mem.Size() / PageSize
-		if currentPages >= requiredPages {
-			panic(fmt.Sprintf("current pages %d >= required pages %d", currentPages, requiredPages))
+		currentPages, ok := pagesFromSize(mem.Size())
+		if !ok {
+			panic(fmt.Sprintf("page size cannot fit into uint32, current memory size: %d", mem.Size()))
 		}
 
 		if currentPages >= MaxWasmPages {
@@ -539,9 +543,9 @@ func bump(bumper *uint32, size uint32, mem runtime.Memory) (uint32, error) {
 				ErrCannotGrowLinearMemory, currentPages, nextPages)
 		}
 
-		pagesIncrease := (mem.Size() / PageSize) == nextPages
+		pagesIncrease := (mem.Size() / PageSize) == uint64(nextPages)
 		if !pagesIncrease {
-			panic(fmt.Sprintf("number of pages should have increased! previous: %d, desired: %d", currentPages, nextPages))
+			logger.Errorf("number of pages should have increased! previous: %d, desired: %d", currentPages, nextPages)
 		}
 	}
 
@@ -553,7 +557,7 @@ func bump(bumper *uint32, size uint32, mem runtime.Memory) (uint32, error) {
 // pagesFromSize convert the given `size` in bytes into the number of pages.
 // The returned number of pages is ensured to be big enough to hold memory
 // with the given `size`.
-// Returns false if the number of pages do not fit into `uint32`
+// Returns false if the number of pages does not fit into `uint32`
 func pagesFromSize(size uint64) (uint32, bool) {
 	value := (size + uint64(PageSize) - 1) / uint64(PageSize)
 

lib/runtime/allocator/freeing_bump_test.go

@@ -248,7 +248,7 @@ func TestShouldReturnErrorWhenBumperGreaterThanHeapSize(t *testing.T) {
 	// further allocation which would increment the bumper must fail.
 	// we try to allocate 8 bytes here, which will increment the
 	// bumper since no 8 byte item has been freed before.
-	require.Equal(t, heap.bumper, mem.Size())
+	require.Equal(t, uint64(heap.bumper), mem.Size())
 
 	ptr, err := heap.Allocate(mem, 8)
 	require.Zero(t, ptr)

lib/runtime/allocator/memory_test.go

@@ -26,8 +26,8 @@ func (m *MemoryInstance) pages() uint32 {
 	return pages
 }
 
-func (m *MemoryInstance) Size() uint32 {
-	return m.pages() * PageSize
+func (m *MemoryInstance) Size() uint64 {
+	return uint64(m.pages() * PageSize)
 }
 
 func (m *MemoryInstance) Grow(pages uint32) (uint32, bool) {
@@ -54,7 +54,7 @@ func (m *MemoryInstance) WriteUint64Le(offset uint32, v uint64) bool {
 	copy(m.data[offset:offset+8], encoded)
 	return true
 }
-func (*MemoryInstance) Read(_, _ uint32) ([]byte, bool) {
+func (*MemoryInstance) Read(_ uint32, _ uint64) ([]byte, bool) {
 	return nil, false
 }
 

lib/runtime/memory.go

@@ -9,7 +9,7 @@ type Memory interface {
 	// has 1 page: 65536
 	//
 	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#-hrefsyntax-instr-memorymathsfmemorysize%E2%91%A0
-	Size() uint32
+	Size() uint64
 
 	// Grow increases memory by the delta in pages (65536 bytes per page).
 	// The return val is the previous memory size in pages, or false if the
@@ -66,7 +66,7 @@ type Memory interface {
 	// shared. Those who need a stable view must set Wasm memory min=max, or
 	// use wazero.RuntimeConfig WithMemoryCapacityPages to ensure max is always
 	// allocated.
-	Read(offset, byteCount uint32) ([]byte, bool)
+	Read(offset uint32, byteCount uint64) ([]byte, bool)
 
 	// WriteByte writes a single byte to the underlying buffer at the offset in or returns false if out of range.
 	WriteByte(offset uint32, v byte) bool //nolint:govet

lib/runtime/wazero/imports.go

@@ -51,8 +51,8 @@ func newPointerSize(ptr, size uint32) (pointerSize uint64) {
 
 // splitPointerSize converts a 64bit pointer size to an
 // uint32 pointer and a uint32 size.
-func splitPointerSize(pointerSize uint64) (ptr, size uint32) {
-	return uint32(pointerSize), uint32(pointerSize >> 32)
+func splitPointerSize(pointerSize uint64) (ptr uint32, size uint64) {
+	return uint32(pointerSize), pointerSize >> 32
 }
 
 // read will read from 64 bit pointer size and return a byte slice
@@ -2265,7 +2265,7 @@ func ext_storage_read_version_1(ctx context.Context, m api.Module, keySpan, valu
 
 	var written uint
 	valueOutPtr, valueOutSize := splitPointerSize(valueOut)
-	if uint32(len(data)) <= valueOutSize {
+	if uint64(len(data)) <= valueOutSize {
 		written = uint(len(data))
 	} else {
 		written = uint(valueOutSize)

lib/runtime/wazero/imports_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/crypto/sr25519"
 	"github.com/ChainSafe/gossamer/lib/keystore"
 	"github.com/ChainSafe/gossamer/lib/runtime"
+	"github.com/ChainSafe/gossamer/lib/runtime/allocator"
 	"github.com/ChainSafe/gossamer/lib/runtime/storage"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 	"github.com/ChainSafe/gossamer/pkg/trie"
@@ -873,8 +874,10 @@ func Test_ext_misc_runtime_version_version_1(t *testing.T) {
 		}
 	}
 
-	data := bytes
+	allocator := allocator.NewFreeingBumpHeapAllocator(0)
+	inst.Context.Allocator = allocator
 
+	data := bytes
 	dataLength := uint32(len(data))
 	inputPtr, err := inst.Context.Allocator.Allocate(inst.Module.Memory(), dataLength)
 	if err != nil {

lib/runtime/wazero/instance.go

@@ -42,6 +42,7 @@ type Instance struct {
 	Module   api.Module
 	Context  *runtime.Context
 	codeHash common.Hash
+	heapBase uint32
 	sync.Mutex
 }
 
@@ -408,24 +409,23 @@ func NewInstance(code []byte, cfg Config) (instance *Instance, err error) {
 		return nil, err
 	}
 
-	global := mod.ExportedGlobal("__heap_base")
-	if global == nil {
+	encodedHeapBase := mod.ExportedGlobal("__heap_base")
+	if encodedHeapBase == nil {
 		return nil, fmt.Errorf("wazero error: nil global for __heap_base")
 	}
 
-	hb := api.DecodeU32(global.Get())
+	heapBase := api.DecodeU32(encodedHeapBase.Get())
 	// hb = runtime.DefaultHeapBase
 
 	mem := mod.Memory()
 	if mem == nil {
 		return nil, fmt.Errorf("wazero error: nil memory for module")
 	}
 
-	allocator := allocator.NewFreeingBumpHeapAllocator(hb)
 	instance = &Instance{
-		Runtime: rt,
+		heapBase: heapBase,
+		Runtime:  rt,
 		Context: &runtime.Context{
-			Allocator:       allocator,
 			Keystore:        cfg.Keystore,
 			Validator:       cfg.Role == common.AuthorityRole,
 			NodeStorage:     cfg.NodeStorage,
@@ -458,7 +458,13 @@ var ErrExportFunctionNotFound = errors.New("export function not found")
 
 func (i *Instance) Exec(function string, data []byte) (result []byte, err error) {
 	i.Lock()
-	defer i.Unlock()
+	i.Context.Allocator = allocator.NewFreeingBumpHeapAllocator(i.heapBase)
+
+	defer func() {
+		i.Context.Allocator = nil
+		i.Unlock()
+	}()
+	// instantiate a new allocator on every execution func
 
 	dataLength := uint32(len(data))
 	inputPtr, err := i.Context.Allocator.Allocate(i.Module.Memory(), dataLength)
@@ -471,6 +477,7 @@ func (i *Instance) Exec(function string, data []byte) (result []byte, err error)
 	if mem == nil {
 		panic("nil memory")
 	}
+
 	ok := mem.Write(inputPtr, data)
 	if !ok {
 		panic("write overflow")
@@ -482,7 +489,6 @@ func (i *Instance) Exec(function string, data []byte) (result []byte, err error)
 	}
 
 	ctx := context.WithValue(context.Background(), runtimeContextKey, i.Context)
-
 	values, err := runtimeFunc.Call(ctx, api.EncodeU32(inputPtr), api.EncodeU32(dataLength))
 	if err != nil {
 		return nil, fmt.Errorf("running runtime function: %w", err)
@@ -497,6 +503,7 @@ func (i *Instance) Exec(function string, data []byte) (result []byte, err error)
 	if !ok {
 		panic("write overflow")
 	}
+
 	return result, nil
 }