keep the original context for GetCommittedState api

x/evm/keeper/grpc_query.go

@@ -444,15 +444,14 @@ func (k Keeper) EstimateGas(c context.Context, req *types.EthCallRequest) (*type
  args.Gas = (*hexutil.Uint64)(&gas)
 
  // Execute the call in an isolated context
- sandboxCtx, _ := ctx.CacheContext()
- k.WithContext(sandboxCtx)
+ k.BeginCachedContext()
 
  msg := args.ToMessage(req.GasCap)
  evm := k.NewEVM(msg, ethCfg, params, coinbase)
  // pass true means execute in query mode, which don't do actual gas refund.
  rsp, err := k.ApplyMessage(evm, msg, ethCfg, true)
 
- k.WithContext(ctx)
+ k.EndCachedContext()
 
  if err != nil {
  if errors.Is(stacktrace.RootCause(err), core.ErrIntrinsicGas) {

x/evm/keeper/keeper.go

@@ -38,6 +38,8 @@ type Keeper struct {
  stakingKeeper types.StakingKeeper
 
  ctx sdk.Context
+ // set in `BeginCachedContext`, used by `GetCommittedState` api.
+ committedCtx sdk.Context
 
  // chain ID number obtained from the context's chain id
  eip155ChainID *big.Int
@@ -75,6 +77,11 @@ func NewKeeper(
  }
 }
 
+// CommittedCtx returns the committed context
+func (k Keeper) CommittedCtx() sdk.Context {
+ return k.committedCtx
+}
+
 // Logger returns a module-specific logger.
 func (k Keeper) Logger(ctx sdk.Context) log.Logger {
  return ctx.Logger().With("module", types.ModuleName)
@@ -83,6 +90,7 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 // WithContext sets an updated SDK context to the keeper
 func (k *Keeper) WithContext(ctx sdk.Context) {
  k.ctx = ctx
+ k.committedCtx = ctx
 }
 
 // WithChainID sets the chain id to the local variable in the keeper

x/evm/keeper/state_transition.go

@@ -137,18 +137,16 @@ func (k *Keeper) ApplyTransaction(tx *ethtypes.Transaction) (*types.MsgEthereumT
  return nil, stacktrace.Propagate(err, "failed to return ethereum transaction as core message")
  }
 
- // create an ethereum StateTransition instance and run TransitionDb
- // we use a ctx context to avoid modifying to state in case EVM msg is reverted
- originalCtx := k.ctx
- cacheCtx, commit := k.ctx.CacheContext()
- k.ctx = cacheCtx
+ // we use a cached context to avoid modifying to state in case EVM msg is reverted
+ commit := k.BeginCachedContext()
 
  // get the coinbase address from the block proposer
  coinbase, err := k.GetCoinbaseAddress()
  if err != nil {
  return nil, stacktrace.Propagate(err, "failed to obtain coinbase address")
  }
 
+ // create an ethereum EVM instance and run the message
  evm := k.NewEVM(msg, ethCfg, params, coinbase)
 
  k.SetTxHashTransient(tx.Hash())
@@ -163,11 +161,12 @@ func (k *Keeper) ApplyTransaction(tx *ethtypes.Transaction) (*types.MsgEthereumT
  res.Hash = txHash.Hex()
  logs := k.GetTxLogs(txHash)
 
- // Commit and switch to original context
+ // Commit and switch to committed context
  if !res.Failed() {
  commit()
  }
- k.ctx = originalCtx
+
+ k.EndCachedContext()
 
  // Logs needs to be ignored when tx is reverted
  // Set the log and bloom filter only when the tx is NOT REVERTED
@@ -373,3 +372,15 @@ func (k Keeper) GetCoinbaseAddress() (common.Address, error) {
  coinbase := common.BytesToAddress(validator.GetOperator())
  return coinbase, nil
 }
+
+// BeginCachedContext create the cached context
+func (k *Keeper) BeginCachedContext() (commit func()) {
+ k.committedCtx = k.ctx
+ k.ctx, commit = k.ctx.CacheContext()
+ return
+}
+
+// EndCachedContext recover the committed context
+func (k *Keeper) EndCachedContext() {
+ k.ctx = k.committedCtx
+}

x/evm/keeper/statedb.go

@@ -352,10 +352,8 @@ func (k *Keeper) GetRefund() uint64 {
 // State
 // ----------------------------------------------------------------------------
 
-// GetCommittedState returns the value set in store for the given key hash. If the key is not registered
-// this function returns the empty hash.
-func (k *Keeper) GetCommittedState(addr common.Address, hash common.Hash) common.Hash {
- store := prefix.NewStore(k.ctx.KVStore(k.storeKey), types.AddressStoragePrefix(addr))
+func doGetState(ctx sdk.Context, storeKey sdk.StoreKey, addr common.Address, hash common.Hash) common.Hash {
+ store := prefix.NewStore(ctx.KVStore(storeKey), types.AddressStoragePrefix(addr))
 
  key := types.KeyAddressStorage(addr, hash)
  value := store.Get(key.Bytes())
@@ -366,10 +364,16 @@ func (k *Keeper) GetCommittedState(addr common.Address, hash common.Hash) common
  return common.BytesToHash(value)
 }
 
+// GetCommittedState returns the value set in store for the given key hash. If the key is not registered
+// this function returns the empty hash.
+func (k *Keeper) GetCommittedState(addr common.Address, hash common.Hash) common.Hash {
+ return doGetState(k.committedCtx, k.storeKey, addr, hash)
+}
+
 // GetState returns the committed state for the given key hash, as all changes are committed directly
 // to the KVStore.
 func (k *Keeper) GetState(addr common.Address, hash common.Hash) common.Hash {
- return k.GetCommittedState(addr, hash)
+ return doGetState(k.ctx, k.storeKey, addr, hash)
 }
 
 // SetState sets the given hashes (key, value) to the KVStore. If the value hash is empty, this

x/evm/keeper/statedb_test.go

@@ -383,6 +383,26 @@ func (suite *KeeperTestSuite) TestState() {
  }
 }
 
+func (suite *KeeperTestSuite) TestCommittedState() {
+ suite.SetupTest()
+
+ var key = common.BytesToHash([]byte("key"))
+ var value1 = common.BytesToHash([]byte("value1"))
+ var value2 = common.BytesToHash([]byte("value2"))
+
+ suite.app.EvmKeeper.SetState(suite.address, key, value1)
+
+ suite.app.EvmKeeper.BeginCachedContext()
+
+ suite.app.EvmKeeper.SetState(suite.address, key, value2)
+ tmp := suite.app.EvmKeeper.GetState(suite.address, key)
+ suite.Require().Equal(value2, tmp)
+ tmp = suite.app.EvmKeeper.GetCommittedState(suite.address, key)
+ suite.Require().Equal(value1, tmp)
+
+ suite.app.EvmKeeper.EndCachedContext()
+}
+
 func (suite *KeeperTestSuite) TestSuicide() {
  testCases := []struct {
  name string