Add store tracer (#579)

## Describe your changes and provide context
Add a tracer that can dump all relevant prestate in the db for a traced
tx, which can then be used to construct past-scenario test cases.

## Testing performed to validate your change
unit test

Author: codchen

store/gaskv/store.go

@@ -8,22 +8,35 @@ import (
 	"github.com/cosmos/cosmos-sdk/telemetry"
 )
 
+type IStoreTracer interface {
+	Get([]byte, []byte, string)
+	Has([]byte, string)
+	Set([]byte, []byte, string)
+	Delete([]byte, string)
+	DerivePrestateToJson() []byte
+	Clear()
+}
+
 var _ types.KVStore = &Store{}
 
 // Store applies gas tracking to an underlying KVStore. It implements the
 // KVStore interface.
 type Store struct {
-	gasMeter  types.GasMeter
-	gasConfig types.GasConfig
-	parent    types.KVStore
+	gasMeter   types.GasMeter
+	gasConfig  types.GasConfig
+	parent     types.KVStore
+	moduleName string
+	tracer     IStoreTracer
 }
 
 // NewStore returns a reference to a new GasKVStore.
-func NewStore(parent types.KVStore, gasMeter types.GasMeter, gasConfig types.GasConfig) *Store {
+func NewStore(parent types.KVStore, gasMeter types.GasMeter, gasConfig types.GasConfig, moduleName string, tracer IStoreTracer) *Store {
 	kvs := &Store{
-		gasMeter:  gasMeter,
-		gasConfig: gasConfig,
-		parent:    parent,
+		gasMeter:   gasMeter,
+		gasConfig:  gasConfig,
+		parent:     parent,
+		moduleName: moduleName,
+		tracer:     tracer,
 	}
 	return kvs
 }
@@ -45,6 +58,9 @@ func (gs *Store) Get(key []byte) (value []byte) {
 	// TODO overflow-safe math?
 	gs.gasMeter.ConsumeGas(gs.gasConfig.ReadCostPerByte*types.Gas(len(key)), types.GasReadPerByteDesc)
 	gs.gasMeter.ConsumeGas(gs.gasConfig.ReadCostPerByte*types.Gas(len(value)), types.GasReadPerByteDesc)
+	if gs.tracer != nil {
+		gs.tracer.Get(key, value, gs.moduleName)
+	}
 
 	return value
 }
@@ -64,7 +80,11 @@ func (gs *Store) Set(key []byte, value []byte) {
 func (gs *Store) Has(key []byte) bool {
 	defer telemetry.MeasureSince(time.Now(), "store", "gaskv", "has")
 	gs.gasMeter.ConsumeGas(gs.gasConfig.HasCost, types.GasHasDesc)
-	return gs.parent.Has(key)
+	res := gs.parent.Has(key)
+	if gs.tracer != nil && res {
+		gs.tracer.Has(key, gs.moduleName)
+	}
+	return res
 }
 
 // Implements KVStore.
@@ -113,7 +133,7 @@ func (gs *Store) iterator(start, end []byte, ascending bool) types.Iterator {
 		parent = gs.parent.ReverseIterator(start, end)
 	}
 
-	gi := newGasIterator(gs.gasMeter, gs.gasConfig, parent)
+	gi := newGasIterator(gs.gasMeter, gs.gasConfig, parent, gs.moduleName, gs.tracer)
 	defer func() {
 		if err := recover(); err != nil {
 			// if there is a panic, we close the iterator then reraise
@@ -139,16 +159,20 @@ func (gs *Store) GetAllKeyStrsInRange(start, end []byte) (res []string) {
 }
 
 type gasIterator struct {
-	gasMeter  types.GasMeter
-	gasConfig types.GasConfig
-	parent    types.Iterator
+	gasMeter   types.GasMeter
+	gasConfig  types.GasConfig
+	parent     types.Iterator
+	moduleName string
+	tracer     IStoreTracer
 }
 
-func newGasIterator(gasMeter types.GasMeter, gasConfig types.GasConfig, parent types.Iterator) types.Iterator {
+func newGasIterator(gasMeter types.GasMeter, gasConfig types.GasConfig, parent types.Iterator, moduleName string, tracer IStoreTracer) types.Iterator {
 	return &gasIterator{
-		gasMeter:  gasMeter,
-		gasConfig: gasConfig,
-		parent:    parent,
+		gasMeter:   gasMeter,
+		gasConfig:  gasConfig,
+		parent:     parent,
+		moduleName: moduleName,
+		tracer:     tracer,
 	}
 }
 
@@ -174,13 +198,19 @@ func (gi *gasIterator) Next() {
 // not incur any gas cost.
 func (gi *gasIterator) Key() (key []byte) {
 	key = gi.parent.Key()
+	if gi.tracer != nil {
+		gi.tracer.Has(key, gi.moduleName)
+	}
 	return key
 }
 
 // Value implements the Iterator interface. It returns the current value and it
 // does not incur any gas cost.
 func (gi *gasIterator) Value() (value []byte) {
 	value = gi.parent.Value()
+	if gi.tracer != nil {
+		gi.tracer.Get(gi.Key(), value, gi.moduleName)
+	}
 	return value
 }
 

store/gaskv/store_test.go

@@ -20,7 +20,7 @@ func valFmt(i int) []byte { return bz(fmt.Sprintf("value%0.8d", i)) }
 func TestGasKVStoreBasic(t *testing.T) {
 	mem := dbadapter.Store{DB: dbm.NewMemDB()}
 	meter := types.NewMultiplierGasMeter(10000, 1, 1)
-	st := gaskv.NewStore(mem, meter, types.KVGasConfig())
+	st := gaskv.NewStore(mem, meter, types.KVGasConfig(), "", nil)
 
 	require.Equal(t, types.StoreTypeDB, st.GetStoreType())
 	require.Panics(t, func() { st.CacheWrap(nil) })
@@ -41,7 +41,7 @@ func TestGasKVStoreBasic(t *testing.T) {
 func TestGasKVStoreIterator(t *testing.T) {
 	mem := dbadapter.Store{DB: dbm.NewMemDB()}
 	meter := types.NewMultiplierGasMeter(100000, 1, 1)
-	st := gaskv.NewStore(mem, meter, types.KVGasConfig())
+	st := gaskv.NewStore(mem, meter, types.KVGasConfig(), "", nil)
 	require.False(t, st.Has(keyFmt(1)))
 	require.Empty(t, st.Get(keyFmt(1)), "Expected `key1` to be empty")
 	require.Empty(t, st.Get(keyFmt(2)), "Expected `key2` to be empty")
@@ -109,14 +109,14 @@ func TestGasKVStoreIterator(t *testing.T) {
 func TestGasKVStoreOutOfGasSet(t *testing.T) {
 	mem := dbadapter.Store{DB: dbm.NewMemDB()}
 	meter := types.NewMultiplierGasMeter(0, 1, 1)
-	st := gaskv.NewStore(mem, meter, types.KVGasConfig())
+	st := gaskv.NewStore(mem, meter, types.KVGasConfig(), "", nil)
 	require.Panics(t, func() { st.Set(keyFmt(1), valFmt(1)) }, "Expected out-of-gas")
 }
 
 func TestGasKVStoreOutOfGasIterator(t *testing.T) {
 	mem := dbadapter.Store{DB: dbm.NewMemDB()}
 	meter := types.NewMultiplierGasMeter(20000, 1, 1)
-	st := gaskv.NewStore(mem, meter, types.KVGasConfig())
+	st := gaskv.NewStore(mem, meter, types.KVGasConfig(), "", nil)
 	st.Set(keyFmt(1), valFmt(1))
 	iterator := st.Iterator(nil, nil)
 	iterator.Next()

store/prefix/store_test.go

@@ -100,7 +100,7 @@ func TestIAVLStorePrefix(t *testing.T) {
 func TestPrefixKVStoreNoNilSet(t *testing.T) {
 	meter := types.NewMultiplierGasMeter(100000000, 1, 1)
 	mem := dbadapter.Store{DB: dbm.NewMemDB()}
-	gasStore := gaskv.NewStore(mem, meter, types.KVGasConfig())
+	gasStore := gaskv.NewStore(mem, meter, types.KVGasConfig(), "", nil)
 	require.Panics(t, func() { gasStore.Set([]byte("key"), nil) }, "setting a nil value should panic")
 }
 

types/context.go

@@ -70,7 +70,8 @@ type Context struct {
 
 	traceSpanContext context.Context
 
-	isTracing bool
+	isTracing   bool
+	storeTracer gaskv.IStoreTracer
 }
 
 // Proposed rename, not done to avoid API breakage
@@ -231,6 +232,13 @@ func (c Context) IsTracing() bool {
 	return c.isTracing
 }
 
+func (c Context) StoreTracer() gaskv.IStoreTracer {
+	if c.storeTracer == nil {
+		return nil
+	}
+	return c.storeTracer
+}
+
 // WithEventManager returns a Context with an updated tx priority
 func (c Context) WithPriority(p int64) Context {
 	c.priority = p
@@ -504,6 +512,9 @@ func (c Context) WithExpireTxHandler(expireTxHandler func()) Context {
 
 func (c Context) WithIsTracing(it bool) Context {
 	c.isTracing = it
+	if it {
+		c.storeTracer = NewStoreTracer()
+	}
 	return c
 }
 
@@ -554,20 +565,20 @@ func (c Context) Value(key interface{}) interface{} {
 func (c Context) KVStore(key StoreKey) KVStore {
 	if c.isTracing {
 		if _, ok := c.nextStoreKeys[key.Name()]; ok {
-			return gaskv.NewStore(c.nextMs.GetKVStore(key), c.GasMeter(), stypes.KVGasConfig())
+			return gaskv.NewStore(c.nextMs.GetKVStore(key), c.GasMeter(), stypes.KVGasConfig(), key.Name(), c.StoreTracer())
 		}
 	}
-	return gaskv.NewStore(c.MultiStore().GetKVStore(key), c.GasMeter(), stypes.KVGasConfig())
+	return gaskv.NewStore(c.MultiStore().GetKVStore(key), c.GasMeter(), stypes.KVGasConfig(), key.Name(), c.StoreTracer())
 }
 
 // TransientStore fetches a TransientStore from the MultiStore.
 func (c Context) TransientStore(key StoreKey) KVStore {
 	if c.isTracing {
 		if _, ok := c.nextStoreKeys[key.Name()]; ok {
-			return gaskv.NewStore(c.nextMs.GetKVStore(key), c.GasMeter(), stypes.TransientGasConfig())
+			return gaskv.NewStore(c.nextMs.GetKVStore(key), c.GasMeter(), stypes.TransientGasConfig(), key.Name(), c.StoreTracer())
 		}
 	}
-	return gaskv.NewStore(c.MultiStore().GetKVStore(key), c.GasMeter(), stypes.TransientGasConfig())
+	return gaskv.NewStore(c.MultiStore().GetKVStore(key), c.GasMeter(), stypes.TransientGasConfig(), key.Name(), c.StoreTracer())
 }
 
 // CacheContext returns a new Context with the multi-store cached and a new

types/tracer.go

@@ -0,0 +1,152 @@
+package types
+
+import (
+	"encoding/hex"
+	"encoding/json"
+	"sync"
+)
+
+type StoreTracer struct {
+	Modules map[string]*ModuleTrace
+	mu      *sync.Mutex
+}
+
+type ModuleTrace struct {
+	Accesses []Access
+}
+
+type Access struct {
+	Op    OpType
+	Key   []byte
+	Value []byte
+}
+
+type OpType int
+
+const (
+	Get OpType = iota
+	Has
+	Set
+	Delete
+)
+
+func NewStoreTracer() *StoreTracer {
+	return &StoreTracer{
+		Modules: map[string]*ModuleTrace{},
+		mu:      &sync.Mutex{},
+	}
+}
+
+func (st *StoreTracer) Get(key []byte, value []byte, module string) {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+	accesses := st.getOrSetModuleTrace(module)
+	accesses.Accesses = append(accesses.Accesses, Access{
+		Op:    Get,
+		Key:   key,
+		Value: value,
+	})
+}
+
+func (st *StoreTracer) Set(key []byte, value []byte, module string) {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+	accesses := st.getOrSetModuleTrace(module)
+	accesses.Accesses = append(accesses.Accesses, Access{
+		Op:    Set,
+		Key:   key,
+		Value: value,
+	})
+}
+
+func (st *StoreTracer) Has(key []byte, module string) {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+	accesses := st.getOrSetModuleTrace(module)
+	accesses.Accesses = append(accesses.Accesses, Access{
+		Op:  Has,
+		Key: key,
+	})
+}
+
+func (st *StoreTracer) Delete(key []byte, module string) {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+	accesses := st.getOrSetModuleTrace(module)
+	accesses.Accesses = append(accesses.Accesses, Access{
+		Op:  Delete,
+		Key: key,
+	})
+}
+
+func (st *StoreTracer) getOrSetModuleTrace(module string) (mt *ModuleTrace) {
+	if _, ok := st.Modules[module]; !ok {
+		mt = &ModuleTrace{
+			Accesses: []Access{},
+		}
+		st.Modules[module] = mt
+	} else {
+		mt = st.Modules[module]
+	}
+	return
+}
+
+func (st *StoreTracer) Clear() {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+	st.Modules = map[string]*ModuleTrace{}
+}
+
+type StoreTraceDump struct {
+	Modules map[string]ModuleTraceDump `json:"modules"`
+}
+
+type ModuleTraceDump struct {
+	Reads map[string]string `json:"reads"`
+	Has   []string          `json:"has"`
+}
+
+func (st *StoreTracer) DerivePrestateToJson() []byte {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+	d := StoreTraceDump{
+		Modules: make(map[string]ModuleTraceDump, len(st.Modules)),
+	}
+	for name, module := range st.Modules {
+		mtd := ModuleTraceDump{
+			Reads: make(map[string]string),
+			Has:   []string{},
+		}
+		// any read for key XYZ after a Set/Delete to XYZ is discarded
+		// because the result doesn't represent prestate.
+		writtenKey := map[string]struct{}{}
+		hasMap := map[string]struct{}{}
+		for _, a := range module.Accesses {
+			switch a.Op {
+			case Get:
+				if _, ok := writtenKey[string(a.Key)]; ok {
+					continue
+				}
+				// no need to check if it's already in dump because the value
+				// must be the same if there is no preceding write
+				mtd.Reads[hex.EncodeToString(a.Key)] = hex.EncodeToString(a.Value)
+			case Has:
+				if _, ok := writtenKey[string(a.Key)]; ok {
+					continue
+				}
+				hasMap[hex.EncodeToString(a.Key)] = struct{}{}
+			default:
+				writtenKey[string(a.Key)] = struct{}{}
+			}
+		}
+		for k := range hasMap {
+			mtd.Has = append(mtd.Has, k)
+		}
+		d.Modules[name] = mtd
+	}
+	bz, err := json.Marshal(&d)
+	if err != nil {
+		panic(err)
+	}
+	return bz
+}

x/bank/keeper/keeper.go

@@ -53,6 +53,7 @@ type Keeper interface {
 	UndelegateCoins(ctx sdk.Context, moduleAccAddr, delegatorAddr sdk.AccAddress, amt sdk.Coins) error
 
 	GetStoreKey() sdk.StoreKey
+	GetCdc() codec.BinaryCodec
 
 	types.QueryServer
 }
@@ -684,6 +685,11 @@ func (k BaseKeeper) GetStoreKey() sdk.StoreKey {
 	return k.storeKey
 }
 
+// for testing
+func (k BaseKeeper) GetCdc() codec.BinaryCodec {
+	return k.cdc
+}
+
 // IterateTotalSupply iterates over the total supply calling the given cb (callback) function
 // with the balance of each coin.
 // The iteration stops if the callback returns true.