fix(dot/state): track runtime per-block, fix runtime upgrades differing between forks (ChainSafe#1638)

* fix runtime upgrade

* Add runtime in blocktree.

* Remove runtime instance from babe service.

* Remove runtime instance from sync service.

* Self review.

* Fix chain reorg test.

* Fix failing test.

* Self review.

* Address comments.

* Remove unused functions from interface.

* Fix failing test.

* fix TestService_HandleSubmittedExtrinsic

* update HandleTransactionMessage to set runtime storage before validating

* Fix failing test.

* address comment

* lint

* cleanup

Co-authored-by: noot <[email protected]>
Co-authored-by: noot <[email protected]>

Author: 2 people

chain/dev/genesis-spec.json

@@ -7,6 +7,9 @@
 	"protocolId": "/gossamer/gssmr/0",
 	"genesis": {
 		"runtime": {
+		    "Sudo": {
+ 		        "Key": "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY"
+ 		    },
 			"Babe": {
 				"Authorities": [
 					[

chain/dev/genesis.json

@@ -22,6 +22,7 @@ import (
 	"github.com/ChainSafe/gossamer/dot/network"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/ChainSafe/gossamer/lib/runtime"
 	rtstorage "github.com/ChainSafe/gossamer/lib/runtime/storage"
 	"github.com/ChainSafe/gossamer/lib/transaction"
 )
@@ -48,6 +49,9 @@ type BlockState interface {
 	HighestCommonAncestor(a, b common.Hash) (common.Hash, error)
 	SubChain(start, end common.Hash) ([]common.Hash, error)
 	GetBlockBody(hash common.Hash) (*types.Body, error)
+	HandleRuntimeChanges(newState *rtstorage.TrieState, in runtime.Instance, bHash common.Hash) error
+	GetRuntime(*common.Hash) (runtime.Instance, error)
+	StoreRuntime(common.Hash, runtime.Instance)
 }
 
 // StorageState interface for storage state methods

dot/core/interface.go

@@ -31,10 +31,23 @@ func (s *Service) HandleTransactionMessage(msg *network.TransactionMessage) (boo
 	// get transactions from message extrinsics
 	txs := msg.Extrinsics
 	var toPropagate []types.Extrinsic
+
+	rt, err := s.blockState.GetRuntime(nil)
+	if err != nil {
+		return false, err
+	}
+
 	for _, tx := range txs {
+		ts, err := s.storageState.TrieState(nil)
+		if err != nil {
+			return false, err
+		}
+
+		rt.SetContextStorage(ts)
+
 		// validate each transaction
 		externalExt := types.Extrinsic(append([]byte{byte(types.TxnExternal)}, tx...))
-		val, err := s.rt.ValidateTransaction(externalExt)
+		val, err := rt.ValidateTransaction(externalExt)
 		if err != nil {
 			logger.Debug("failed to validate transaction", "err", err)
 			continue

dot/core/messages.go

@@ -24,6 +24,7 @@ import (
 	. "github.com/ChainSafe/gossamer/dot/core/mocks" // nolint
 	"github.com/ChainSafe/gossamer/dot/network"
 	"github.com/ChainSafe/gossamer/dot/state"
+	"github.com/ChainSafe/gossamer/dot/sync"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/crypto/sr25519"
@@ -140,14 +141,22 @@ func TestService_HandleTransactionMessage(t *testing.T) {
 
 	s := NewTestService(t, cfg)
 	genHash := s.blockState.GenesisHash()
-	header, err := types.NewHeader(genHash, common.Hash{}, common.Hash{}, big.NewInt(1), types.NewEmptyDigest())
+	genHeader, err := s.blockState.BestBlockHeader()
 	require.NoError(t, err)
 
-	// initialise block header
-	err = s.rt.InitializeBlock(header)
+	rt, err := s.blockState.GetRuntime(nil)
 	require.NoError(t, err)
 
-	extBytes := createExtrinsic(t, s.rt, genHash, 0)
+	ts, err := s.storageState.TrieState(nil)
+	require.NoError(t, err)
+	rt.SetContextStorage(ts)
+
+	block := sync.BuildBlock(t, rt, genHeader, nil)
+
+	err = s.handleBlock(block, ts)
+	require.NoError(t, err)
+
+	extBytes := createExtrinsic(t, rt, genHash, 0)
 	msg := &network.TransactionMessage{Extrinsics: []types.Extrinsic{extBytes}}
 	b, err := s.HandleTransactionMessage(msg)
 	require.NoError(t, err)

dot/core/messages_test.go

@@ -16,9 +16,7 @@
 package core
 
 import (
-	"bytes"
 	"context"
-	"fmt"
 	"math/big"
 	"os"
 	"sync"
@@ -59,10 +57,6 @@ type Service struct {
 	net              Network
 	digestHandler    DigestHandler
 
-	// Current runtime and hash of the current runtime code
-	rt       runtime.Instance
-	codeHash common.Hash
-
 	// map of code substitutions keyed by block hash
 	codeSubstitute       map[common.Hash]string
 	codeSubstitutedState CodeSubstitutedState
@@ -103,10 +97,6 @@ func NewService(cfg *Config) (*Service, error) {
 		return nil, ErrNilStorageState
 	}
 
-	if cfg.Runtime == nil {
-		return nil, ErrNilRuntime
-	}
-
 	if cfg.Network == nil {
 		return nil, ErrNilNetwork
 	}
@@ -123,24 +113,12 @@ func NewService(cfg *Config) (*Service, error) {
 	h = log.CallerFileHandler(h)
 	logger.SetHandler(log.LvlFilterHandler(cfg.LogLvl, h))
 
-	sr, err := cfg.BlockState.BestBlockStateRoot()
-	if err != nil {
-		return nil, err
-	}
-
-	codeHash, err := cfg.StorageState.LoadCodeHash(&sr)
-	if err != nil {
-		return nil, err
-	}
-
 	blockAddCh := make(chan *types.Block, 256)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	srv := &Service{
 		ctx:                  ctx,
 		cancel:               cancel,
-		rt:                   cfg.Runtime,
-		codeHash:             codeHash,
 		keys:                 cfg.Keystore,
 		blockState:           cfg.BlockState,
 		epochState:           cfg.EpochState,
@@ -221,7 +199,7 @@ func (s *Service) handleBlock(block *types.Block, state *rtstorage.TrieState) er
 	}
 
 	// store block in database
-	if err := s.blockState.AddBlock(block); err != nil {
+	if err = s.blockState.AddBlock(block); err != nil {
 		if err == blocktree.ErrParentNotFound && block.Header.Number.Cmp(big.NewInt(0)) != 0 {
 			return err
 		} else if err == blocktree.ErrBlockExists || block.Header.Number.Cmp(big.NewInt(0)) == 0 {
@@ -236,8 +214,13 @@ func (s *Service) handleBlock(block *types.Block, state *rtstorage.TrieState) er
 	// handle consensus digests
 	s.digestHandler.HandleDigests(block.Header)
 
+	rt, err := s.blockState.GetRuntime(&block.Header.ParentHash)
+	if err != nil {
+		return err
+	}
+
 	// check for runtime changes
-	if err := s.handleRuntimeChanges(state); err != nil {
+	if err := s.blockState.HandleRuntimeChanges(state, rt, block.Header.Hash()); err != nil {
 		logger.Crit("failed to update runtime code", "error", err)
 		return err
 	}
@@ -267,57 +250,6 @@ func (s *Service) handleBlock(block *types.Block, state *rtstorage.TrieState) er
 	return nil
 }
 
-func (s *Service) handleRuntimeChanges(newState *rtstorage.TrieState) error {
-	currCodeHash, err := newState.LoadCodeHash()
-	if err != nil {
-		return err
-	}
-
-	if bytes.Equal(s.codeHash[:], currCodeHash[:]) {
-		return nil
-	}
-
-	logger.Info("🔄 detected runtime code change, upgrading...", "block", s.blockState.BestBlockHash(), "previous code hash", s.codeHash, "new code hash", currCodeHash)
-	code := newState.LoadCode()
-	if len(code) == 0 {
-		return ErrEmptyRuntimeCode
-	}
-
-	codeSubBlockHash := s.codeSubstitutedState.LoadCodeSubstitutedBlockHash()
-
-	if !codeSubBlockHash.Equal(common.Hash{}) {
-		// don't do runtime change if using code substitution and runtime change spec version are equal
-		//  (do a runtime change if code substituted and runtime spec versions are different, or code not substituted)
-		newVersion, err := s.rt.CheckRuntimeVersion(code) //nolint
-		if err != nil {
-			return err
-		}
-
-		previousVersion, _ := s.rt.Version()
-		if previousVersion.SpecVersion() == newVersion.SpecVersion() {
-			return nil
-		}
-
-		logger.Info("🔄 detected runtime code change, upgrading...", "block", s.blockState.BestBlockHash(),
-			"previous code hash", s.codeHash, "new code hash", currCodeHash,
-			"previous spec version", previousVersion.SpecVersion(), "new spec version", newVersion.SpecVersion())
-	}
-
-	err = s.rt.UpdateRuntimeCode(code)
-	if err != nil {
-		return err
-	}
-
-	s.codeHash = currCodeHash
-
-	err = s.codeSubstitutedState.StoreCodeSubstitutedBlockHash(common.Hash{})
-	if err != nil {
-		return fmt.Errorf("failed to update code substituted block hash: %w", err)
-	}
-
-	return nil
-}
-
 func (s *Service) handleCodeSubstitution(hash common.Hash) error {
 	value := s.codeSubstitute[hash]
 	if value == "" {
@@ -330,7 +262,12 @@ func (s *Service) handleCodeSubstitution(hash common.Hash) error {
 		return ErrEmptyRuntimeCode
 	}
 
-	err := s.rt.UpdateRuntimeCode(code)
+	rt, err := s.blockState.GetRuntime(&hash)
+	if err != nil {
+		return err
+	}
+
+	err = rt.UpdateRuntimeCode(code)
 	if err != nil {
 		return err
 	}
@@ -415,6 +352,12 @@ func (s *Service) handleChainReorg(prev, curr common.Hash) error {
 		subchain = subchain[1:]
 	}
 
+	// Check transaction validation on the best block.
+	rt, err := s.blockState.GetRuntime(nil)
+	if err != nil {
+		return err
+	}
+
 	// for each block in the previous chain, re-add its extrinsics back into the pool
 	for _, hash := range subchain {
 		body, err := s.blockState.GetBlockBody(hash)
@@ -448,7 +391,7 @@ func (s *Service) handleChainReorg(prev, curr common.Hash) error {
 			}
 
 			externalExt := types.Extrinsic(append([]byte{byte(types.TxnExternal)}, encExt...))
-			txv, err := s.rt.ValidateTransaction(externalExt)
+			txv, err := rt.ValidateTransaction(externalExt)
 			if err != nil {
 				logger.Debug("failed to validate transaction", "error", err, "extrinsic", ext)
 				continue
@@ -519,6 +462,7 @@ func (s *Service) HasKey(pubKeyStr, keyType string) (bool, error) {
 // GetRuntimeVersion gets the current RuntimeVersion
 func (s *Service) GetRuntimeVersion(bhash *common.Hash) (runtime.Version, error) {
 	var stateRootHash *common.Hash
+
 	// If block hash is not nil then fetch the state root corresponding to the block.
 	if bhash != nil {
 		var err error
@@ -533,16 +477,36 @@ func (s *Service) GetRuntimeVersion(bhash *common.Hash) (runtime.Version, error)
 		return nil, err
 	}
 
-	s.rt.SetContextStorage(ts)
-	return s.rt.Version()
+	rt, err := s.blockState.GetRuntime(bhash)
+	if err != nil {
+		return nil, err
+	}
+
+	rt.SetContextStorage(ts)
+	return rt.Version()
 }
 
 // HandleSubmittedExtrinsic is used to send a Transaction message containing a Extrinsic @ext
 func (s *Service) HandleSubmittedExtrinsic(ext types.Extrinsic) error {
+	if s.net == nil {
+		return nil
+	}
+
+	ts, err := s.storageState.TrieState(nil)
+	if err != nil {
+		return err
+	}
+
+	rt, err := s.blockState.GetRuntime(nil)
+	if err != nil {
+		logger.Crit("failed to get runtime")
+		return err
+	}
+
+	rt.SetContextStorage(ts)
 	// the transaction source is External
 	externalExt := types.Extrinsic(append([]byte{byte(types.TxnExternal)}, ext...))
-
-	txv, err := s.rt.ValidateTransaction(externalExt)
+	txv, err := rt.ValidateTransaction(externalExt)
 	if err != nil {
 		return err
 	}
@@ -576,6 +540,11 @@ func (s *Service) GetMetadata(bhash *common.Hash) ([]byte, error) {
 		return nil, err
 	}
 
-	s.rt.SetContextStorage(ts)
-	return s.rt.Metadata()
+	rt, err := s.blockState.GetRuntime(bhash)
+	if err != nil {
+		return nil, err
+	}
+
+	rt.SetContextStorage(ts)
+	return rt.Metadata()
 }