[x/programs]: improve metering and add tests

go.mod

@@ -143,4 +143,4 @@ require (
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
-replace github.com/tetratelabs/wazero => github.com/ava-labs/wazero v0.0.1-hypersdk
+replace github.com/tetratelabs/wazero => github.com/ava-labs/wazero v0.0.2-hypersdk

go.sum

@@ -65,8 +65,8 @@ github.com/ava-labs/avalanchego v1.10.8 h1:fUudA4J37y8wyNG3iiX0kpoZXunsWpCgvsGDg
 github.com/ava-labs/avalanchego v1.10.8/go.mod h1:2zuce+beHe25wQ5RlrEeDpa+SqY/sjEOjDky+Q1NxfU=
 github.com/ava-labs/coreth v0.12.5-rc.1 h1:rLJ6mT/44jgc+vPpKaspGpxDIdQGLkSK6wDSmKgYPxY=
 github.com/ava-labs/coreth v0.12.5-rc.1/go.mod h1:K7Xm2jqx90wxKZXfLvkLEL+zlM5843gGq9XkqVDwKds=
-github.com/ava-labs/wazero v0.0.1-hypersdk h1:l5lvgGYTSDgDqPFm7QawNmA5rDUb45Fhz7gCUCqO7lc=
-github.com/ava-labs/wazero v0.0.1-hypersdk/go.mod h1:0U0G41+ochRKoPKCJlh0jMg1CHkyfK8kDqiirMmKY8A=
+github.com/ava-labs/wazero v0.0.2-hypersdk h1:zsl0saXFEU/dnZ7/cFVDXoSKeavzjrqv1PoR/vKaKmY=
+github.com/ava-labs/wazero v0.0.2-hypersdk/go.mod h1:0U0G41+ochRKoPKCJlh0jMg1CHkyfK8kDqiirMmKY8A=
 github.com/aymerick/raymond v2.0.3-0.20180322193309-b565731e1464+incompatible/go.mod h1:osfaiScAUVup+UC9Nfq76eWqDhXlp+4UYaA8uhTBO6g=
 github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=

scripts/mocks.mockgen.txt

@@ -3,3 +3,4 @@ github.com/ava-labs/hypersdk/chain=Auth=chain/mock_auth.go
 github.com/ava-labs/hypersdk/chain=AuthFactory=chain/mock_auth_factory.go
 github.com/ava-labs/hypersdk/chain=Rules=chain/mock_rules.go
 github.com/ava-labs/hypersdk/vm=Controller=vm/mock_controller.go
+github.com/ava-labs/hypersdk/x/programs/runtime=Storage=x/programs/runtime/mock_storage.go

x/programs/examples/examples.go

@@ -14,9 +14,9 @@ import (
 	"github.com/ava-labs/hypersdk/x/programs/runtime"
 )
 
-var _ runtime.ProgramStorage = (*programStorage)(nil)
+var _ runtime.Storage = (*programStorage)(nil)
 
-// newProgramStorage returns an instance of program storage used for examples
+// newProgramStorage returns an instance of runtime storage used for examples
 // and backed by memDb.
 func newProgramStorage(db chain.Database) *programStorage {
 	return &programStorage{

x/programs/examples/token.go

@@ -9,7 +9,6 @@ import (
 	"github.com/ava-labs/avalanchego/utils/logging"
 
 	"github.com/ava-labs/hypersdk/crypto/ed25519"
-	"github.com/ava-labs/hypersdk/x/programs/meter"
 	"github.com/ava-labs/hypersdk/x/programs/runtime"
 	"github.com/ava-labs/hypersdk/x/programs/utils"
 
@@ -46,7 +45,7 @@ func (t *Token) Run(ctx context.Context) error {
 		"init_program",
 	}
 
-	meter := meter.New(t.maxFee, t.costMap)
+	meter := runtime.NewMeter(t.maxFee, t.costMap)
 	db := utils.NewTestDB()
 	store := newProgramStorage(db)
 

x/programs/runtime/dependencies.go

@@ -22,7 +22,7 @@ type Runtime interface {
 	Stop(context.Context) error
 }
 
-type ProgramStorage interface {
+type Storage interface {
 	Get(context.Context, uint32) ([]byte, bool, error)
 	Set(context.Context, uint32, uint32, []byte) error
 }

x/programs/runtime/map.go

@@ -10,7 +10,6 @@ import (
 	"github.com/tetratelabs/wazero/api"
 
 	"github.com/ava-labs/avalanchego/utils/logging"
-	"github.com/ava-labs/hypersdk/x/programs/meter"
 	"github.com/ava-labs/hypersdk/x/programs/utils"
 )
 
@@ -31,15 +30,15 @@ type storage struct {
 }
 
 type MapModule struct {
-	meter meter.Meter
+	meter Meter
 	log   logging.Logger
 	store storage
 }
 
 // NewMapModule returns a new map host module which can manage in memory state.
 // This is a placeholder storage system intended to show how a wasm program
 // would access/modify persistent storage.
-func NewMapModule(log logging.Logger, meter meter.Meter) *MapModule {
+func NewMapModule(log logging.Logger, meter Meter) *MapModule {
 	return &MapModule{
 		meter: meter,
 		log:   log,

x/programs/runtime/meter.go

@@ -0,0 +1,67 @@
+// Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package runtime
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+
+	"github.com/tetratelabs/wazero/api"
+)
+
+var (
+	_                      Meter = (*meter)(nil)
+	ErrMeterInvalidBalance       = errors.New("operation cost greater than balance")
+)
+
+type Meter interface {
+	api.Meter
+	Run(context.Context)
+}
+
+func NewMeter(maxFee uint64, costMap map[string]uint64) *meter {
+	return &meter{
+		costMap: costMap,
+		balance: maxFee,
+		waitCh:  make(chan struct{}),
+	}
+}
+
+type meter struct {
+	lock    sync.RWMutex
+	costMap map[string]uint64
+	balance uint64
+
+	waitCh chan struct{}
+}
+
+func (m *meter) AddCost(_ context.Context, op string) error {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	cost := m.costMap[op]
+	if cost > m.balance {
+		defer close(m.waitCh)
+		m.waitCh <- struct{}{}
+		return fmt.Errorf("%w: %d: %d", ErrMeterInvalidBalance, cost, m.balance)
+	}
+
+	m.balance -= cost
+	return nil
+}
+
+func (m *meter) GetBalance() uint64 {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+	return m.balance
+}
+
+func (m *meter) Run(ctx context.Context) {
+	select {
+	case <-m.waitCh:
+	case <-ctx.Done():
+	}
+}

x/programs/runtime/meter_test.go

@@ -0,0 +1,59 @@
+// Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package runtime
+
+import (
+	"context"
+	_ "embed"
+	"os"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/ava-labs/avalanchego/utils/logging"
+)
+
+var (
+	//go:embed testdata/get_guest.wasm
+	tokenProgramBytes []byte
+
+	// example cost map
+	costMap = map[string]uint64{
+		"ConstI32 0x0": 1, // initialize i32
+		"ConstI32 0x1": 1, // set i32 to value 1
+	}
+	maxFee uint64 = 3
+	log           = logging.NewLogger(
+		"",
+		logging.NewWrappedCore(
+			logging.Debug,
+			os.Stderr,
+			logging.Plain.ConsoleEncoder(),
+		))
+)
+
+// go test -v -run ^TestMeter$ github.com/ava-labs/hypersdk/x/programs/runtime
+func TestMeter(t *testing.T) {
+	require := require.New(t)
+	ctrl := gomock.NewController(t)
+	storage := NewMockStorage(ctrl)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	meter := NewMeter(maxFee, costMap)
+	runtime := New(log, meter, storage)
+	err := runtime.Initialize(ctx, tokenProgramBytes, []string{"get"})
+	require.NoError(err)
+
+	// first call should pass
+	resp, err := runtime.Call(ctx, "get")
+	require.NoError(err)
+	require.Equal(uint64(1), resp[0])
+
+	// second call should fail invalid
+	_, err = runtime.Call(ctx, "get")
+	require.ErrorIs(err, ErrMeterInvalidBalance)
+}