Tiered Storage POC

api/dependencies.go

@@ -11,7 +11,7 @@
 	"github.com/ava-labs/avalanchego/trace"
 	"github.com/ava-labs/avalanchego/utils/logging"
 
 	"github.com/ava-labs/hypersdk/chain"
 	"github.com/ava-labs/hypersdk/codec"
 	"github.com/ava-labs/hypersdk/fees"
 	"github.com/ava-labs/hypersdk/genesis"
@@ -41,7 +41,13 @@
 		context.Context,
 	) (map[ids.NodeID]*validators.GetValidatorOutput, map[string]struct{})
 	GetVerifyAuth() bool
-	ReadState(ctx context.Context, keys [][]byte) ([][]byte, []error)
+	// Access to Raw State
+	ReadState(ctx context.Context, keys [][]byte) (state.Immutable, error)
 	ImmutableState(ctx context.Context) (state.Immutable, error)
 	BalanceHandler() chain.BalanceHandler
 }
+
+type Chain interface {
+	// Access to executable state
+	Get()
+}

api/jsonrpc/server.go

@@ -209,17 +209,17 @@ func (j *JSONRPCServer) ExecuteActions(
 			storageKeysToRead = append(storageKeysToRead, []byte(key))
 		}
 
-		values, errs := j.vm.ReadState(ctx, storageKeysToRead)
-		for _, err := range errs {
+		im, err := j.vm.ReadState(ctx, storageKeysToRead)
+		if err != nil {
+			return fmt.Errorf("failed to read state: %w", err)
+		}
+
+		for _, k := range storageKeysToRead {
+			v, err := im.GetValue(ctx, k)
 			if err != nil && !errors.Is(err, database.ErrNotFound) {
 				return fmt.Errorf("failed to read state: %w", err)
 			}
-		}
-		for i, value := range values {
-			if value == nil {
-				continue
-			}
-			storage[string(storageKeysToRead[i])] = value
+			storage[string(k)] = v
 		}
 
 		tsv := ts.NewView(

api/state/server.go

@@ -51,9 +51,13 @@ func (s *JSONRPCStateServer) ReadState(req *http.Request, args *ReadStateRequest
 	ctx, span := s.stateReader.Tracer().Start(req.Context(), "Server.ReadState")
 	defer span.End()
 
-	var errs []error
-	res.Values, errs = s.stateReader.ReadState(ctx, args.Keys)
-	for _, err := range errs {
+	im, err := s.stateReader.ReadState(ctx, args.Keys)
+	if err != nil {
+		return err
+	}
+	for _, k := range args.Keys {
+		v, err := im.GetValue(ctx, k)
+		res.Values = append(res.Values, v)
 		res.Errors = append(res.Errors, err.Error())
 	}
 	return nil

chain/builder.go

@@ -23,6 +23,7 @@ import (
 	"github.com/ava-labs/hypersdk/internal/fees"
 	"github.com/ava-labs/hypersdk/keys"
 	"github.com/ava-labs/hypersdk/state"
+	"github.com/ava-labs/hypersdk/state/shim"
 	"github.com/ava-labs/hypersdk/state/tstate"
 )
 
@@ -61,15 +62,19 @@ func HandlePreExecute(log logging.Logger, err error) bool {
 }
 
 type Builder struct {
-	tracer          trace.Tracer
-	ruleFactory     RuleFactory
-	log             logging.Logger
-	metadataManager MetadataManager
-	balanceHandler  BalanceHandler
-	mempool         Mempool
-	validityWindow  ValidityWindow
-	metrics         *chainMetrics
-	config          Config
+	tracer             trace.Tracer
+	ruleFactory        RuleFactory
+	log                logging.Logger
+	metadataManager    MetadataManager
+	balanceHandler     BalanceHandler
+	mempool            Mempool
+	validityWindow     ValidityWindow
+	metrics            *chainMetrics
+	executionShim      shim.Execution
+	exportStateDiff    ExportStateDiffFunc
+	resultModifierFunc ResultModifierFunc
+	refundFunc         RefundFunc
+	config             Config
 }
 
 func NewBuilder(
@@ -81,18 +86,26 @@ func NewBuilder(
 	mempool Mempool,
 	validityWindow ValidityWindow,
 	metrics *chainMetrics,
+	executionShim shim.Execution,
+	afterBlock ExportStateDiffFunc,
+	resultModifierFunc ResultModifierFunc,
+	refundFunc RefundFunc,
 	config Config,
 ) *Builder {
 	return &Builder{
-		tracer:          tracer,
-		ruleFactory:     ruleFactory,
-		log:             log,
-		metadataManager: metadataManager,
-		balanceHandler:  balanceHandler,
-		mempool:         mempool,
-		validityWindow:  validityWindow,
-		metrics:         metrics,
-		config:          config,
+		tracer:             tracer,
+		ruleFactory:        ruleFactory,
+		log:                log,
+		metadataManager:    metadataManager,
+		balanceHandler:     balanceHandler,
+		mempool:            mempool,
+		validityWindow:     validityWindow,
+		metrics:            metrics,
+		executionShim:      executionShim,
+		exportStateDiff:    afterBlock,
+		resultModifierFunc: resultModifierFunc,
+		refundFunc:         refundFunc,
+		config:             config,
 	}
 }
 
@@ -289,10 +302,15 @@ func (c *Builder) BuildBlock(ctx context.Context, parentView state.View, parent
 					}()
 				}
 
+				state, err := c.executionShim.ImmutableView(ctx, stateKeys, state.ImmutableStorage(storage), height)
+				if err != nil {
+					return err
+				}
+
 				// Execute block
 				tsv := ts.NewView(
 					stateKeys,
-					state.ImmutableStorage(storage),
+					state,
 					len(stateKeys),
 				)
 				if err := tx.PreExecute(ctx, feeManager, c.balanceHandler, r, tsv, nextTime); err != nil {
@@ -322,6 +340,15 @@ func (c *Builder) BuildBlock(ctx context.Context, parentView state.View, parent
 				blockLock.Lock()
 				defer blockLock.Unlock()
 
+				resultChanges, err := c.resultModifierFunc(state, result, feeManager)
+				if err != nil {
+					return err
+				}
+
+				if err := c.refundFunc(ctx, resultChanges, c.balanceHandler, tx.Auth.Sponsor(), tsv); err != nil {
+					return err
+				}
+
 				// Ensure block isn't too big
 				if ok, dimension := feeManager.Consume(result.Units, maxUnits); !ok {
 					c.log.Debug(
@@ -440,7 +467,7 @@ func (c *Builder) BuildBlock(ctx context.Context, parentView state.View, parent
 	}
 
 	// Get view from [tstate] after writing all changed keys
-	view, err := ts.ExportMerkleDBView(ctx, c.tracer, parentView)
+	view, err := c.exportStateDiff(ctx, ts, parentView, c.metadataManager, height)
 	if err != nil {
 		return nil, nil, nil, err
 	}

chain/chain.go

@@ -36,11 +36,16 @@ func NewChain(
 	authVM AuthVM,
 	validityWindow ValidityWindow,
 	config Config,
+	ops []Option,
 ) (*Chain, error) {
 	metrics, err := newMetrics(registerer)
 	if err != nil {
 		return nil, err
 	}
+
+	options := NewDefaultOptions()
+	applyOptions(options, ops)
+
 	return &Chain{
 		builder: NewBuilder(
 			tracer,
@@ -51,6 +56,10 @@ func NewChain(
 			mempool,
 			validityWindow,
 			metrics,
+			options.executionShim,
+			options.exportStateDiffFunc,
+			options.resultModifierFunc,
+			options.refundFunc,
 			config,
 		),
 		processor: NewProcessor(
@@ -63,6 +72,11 @@ func NewChain(
 			balanceHandler,
 			validityWindow,
 			metrics,
+			options.executionShim,
+			options.exportStateDiffFunc,
+			options.dimsModifierFunc,
+			options.resultModifierFunc,
+			options.refundFunc,
 			config,
 		),
 		accepter: NewAccepter(
@@ -74,6 +88,7 @@ func NewChain(
 			ruleFactory,
 			validityWindow,
 			metadataManager,
+			options.executionShim,
 			balanceHandler,
 		),
 		blockParser: NewBlockParser(tracer, parser),

chain/default.go

@@ -0,0 +1,34 @@
+// Copyright (C) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package chain
+
+import (
+	"context"
+
+	"github.com/ava-labs/avalanchego/x/merkledb"
+
+	"github.com/ava-labs/hypersdk/codec"
+	"github.com/ava-labs/hypersdk/state"
+	"github.com/ava-labs/hypersdk/state/tstate"
+
+	internalfees "github.com/ava-labs/hypersdk/internal/fees"
+)
+
+func DefaultExportStateDiff(ctx context.Context, ts *tstate.TState, view state.View, _ MetadataManager, _ uint64) (merkledb.View, error) {
+	return view.NewView(ctx, merkledb.ViewChanges{MapOps: ts.ChangedKeys(), ConsumeBytes: true})
+}
+
+// This modifies the result passed in
+// Futhermore, any changes that were made in Result should be documented in ResultChanges
+func DefaultResultModifier(state.Immutable, *Result, *internalfees.Manager) (*ResultChanges, error) {
+	return nil, nil
+}
+
+func DefaultRefundFunc(context.Context, *ResultChanges, BalanceHandler, codec.Address, state.Mutable) error {
+	return nil
+}
+
+func FeeManagerModifier(*internalfees.Manager, *ResultChanges) error {
+	return nil
+}

chain/dependencies.go

@@ -8,11 +8,15 @@ import (
 
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/utils/set"
+	"github.com/ava-labs/avalanchego/x/merkledb"
 
 	"github.com/ava-labs/hypersdk/codec"
 	"github.com/ava-labs/hypersdk/fees"
 	"github.com/ava-labs/hypersdk/internal/validitywindow"
 	"github.com/ava-labs/hypersdk/state"
+	"github.com/ava-labs/hypersdk/state/tstate"
+
+	internalfees "github.com/ava-labs/hypersdk/internal/fees"
 )
 
 type Parser interface {
@@ -221,3 +225,17 @@ type ValidityWindow interface {
 		oldestAllowed int64,
 	) (set.Bits, error)
 }
+
+// ExportStateDiffFunc finalizes the state diff
+type ExportStateDiffFunc func(context.Context, *tstate.TState, state.View, MetadataManager, uint64) (merkledb.View, error)
+
+// ResultModifierFunc modifies the result of a transaction
+// Any changes should be reflected in ResultChanges
+type ResultModifierFunc func(state.Immutable, *Result, *internalfees.Manager) (*ResultChanges, error)
+
+// RefundFunc refunds any fees to the transaction sponsor
+type RefundFunc func(context.Context, *ResultChanges, BalanceHandler, codec.Address, state.Mutable) error
+
+// FeeManagerModifierFunc can be used to modify the fee manager
+// Any changes to a result's units should be reflected here
+type FeeManagerModifierFunc func(*internalfees.Manager, *ResultChanges) error

chain/option.go

@@ -0,0 +1,63 @@
+// Copyright (C) 2024, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package chain
+
+import "github.com/ava-labs/hypersdk/state/shim"
+
+type Options struct {
+	executionShim shim.Execution
+	// exportStateDiffFunc allows users to override the state diff at the end of block execution
+	exportStateDiffFunc ExportStateDiffFunc
+	refundFunc          RefundFunc
+	dimsModifierFunc    FeeManagerModifierFunc
+	resultModifierFunc  ResultModifierFunc
+}
+
+type Option func(*Options)
+
+func NewDefaultOptions() *Options {
+	return &Options{
+		executionShim:       &shim.ExecutionNoOp{},
+		exportStateDiffFunc: DefaultExportStateDiff,
+		refundFunc:          DefaultRefundFunc,
+		dimsModifierFunc:    FeeManagerModifier,
+		resultModifierFunc:  DefaultResultModifier,
+	}
+}
+
+func WithExecutionShim(shim shim.Execution) Option {
+	return func(opts *Options) {
+		opts.executionShim = shim
+	}
+}
+
+func WithExportStateDiffFunc(exportStateDiff ExportStateDiffFunc) Option {
+	return func(opts *Options) {
+		opts.exportStateDiffFunc = exportStateDiff
+	}
+}
+
+func WithRefundFunc(refund RefundFunc) Option {
+	return func(opts *Options) {
+		opts.refundFunc = refund
+	}
+}
+
+func WithDimsModifierFunc(dimsModifier FeeManagerModifierFunc) Option {
+	return func(opts *Options) {
+		opts.dimsModifierFunc = dimsModifier
+	}
+}
+
+func WithResultModifierFunc(resultModifier ResultModifierFunc) Option {
+	return func(opts *Options) {
+		opts.resultModifierFunc = resultModifier
+	}
+}
+
+func applyOptions(option *Options, opts []Option) {
+	for _, o := range opts {
+		o(option)
+	}
+}