change: improve TX tracking

x/gastracker/ante/ante.go

@@ -8,7 +8,7 @@ import (
 
 type GasTrackingKeeper interface {
 	GetParams(ctx sdk.Context) gastracker.Params
-	TrackNewTx(ctx sdk.Context, rewardCoins []sdk.DecCoin, gas uint64) error
+	TrackNewTx(ctx sdk.Context, rewardCoins []sdk.DecCoin, gas uint64)
 }
 
 type TxGasTrackingDecorator struct {
@@ -37,10 +37,7 @@ func (t TxGasTrackingDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simulate
 		rewardCoins[i] = reward
 	}
 
-	err = t.gasTrackingKeeper.TrackNewTx(ctx, rewardCoins, feeTx.GetGas())
-	if err != nil {
-		return ctx, err
-	}
+	t.gasTrackingKeeper.TrackNewTx(ctx, rewardCoins, feeTx.GetGas())
 
 	return next(ctx, tx, simulate)
 }

x/gastracker/ante/ante_test.go

@@ -1,31 +1,11 @@
 package ante
 
 import (
-	"os"
-	"testing"
-	"time"
-
-	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
-	wasmTypes "github.com/CosmWasm/wasmd/x/wasm/types"
-	"github.com/cosmos/cosmos-sdk/simapp"
-	"github.com/cosmos/cosmos-sdk/store"
+	"github.com/archway-network/archway/x/gastracker/testutil"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
-	paramskeeper "github.com/cosmos/cosmos-sdk/x/params/keeper"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-	tmLog "github.com/tendermint/tendermint/libs/log"
-	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
-	db "github.com/tendermint/tm-db"
-
-	"github.com/archway-network/archway/x/gastracker"
-	gstTypes "github.com/archway-network/archway/x/gastracker"
-	"github.com/archway-network/archway/x/gastracker/keeper"
-)
-
-// NOTE: this is needed to allow the keeper to set BlockGasTracking
-var (
-	storeKey = sdk.NewKVStoreKey(gastracker.StoreKey)
+	"testing"
 )
 
 type dummyTx struct {
@@ -73,7 +53,7 @@ func dummyNextAnteHandler(_ sdk.Context, _ sdk.Tx, _ bool) (newCtx sdk.Context,
 }
 
 func TestGasTrackingAnteHandler(t *testing.T) {
-	ctx, keeper := CreateTestKeeperAndContext(t, sdk.AccAddress{})
+	ctx, keeper := testutil.CreateTestKeeperAndContext(t, sdk.AccAddress{})
 
 	testTxGasTrackingDecorator := NewTxGasTrackingDecorator(keeper)
 
@@ -98,16 +78,7 @@ func TestGasTrackingAnteHandler(t *testing.T) {
 		expectedDecCoins[i] = sdk.NewDecCoinFromCoin(sdk.NewCoin(coin.Denom, coin.Amount.QuoRaw(2)))
 	}
 
-	_, err = testTxGasTrackingDecorator.AnteHandle(ctx, testTx, false, dummyNextAnteHandler)
-	assert.EqualError(
-		t,
-		err,
-		gstTypes.ErrBlockTrackingDataNotFound.Error(),
-		"Gastracking ante handler should return expected error",
-	)
-
-	err = keeper.TrackNewBlock(ctx)
-	assert.NoError(t, err, "New block gas tracking should succeed")
+	keeper.TrackNewBlock(ctx)
 
 	_, err = testTxGasTrackingDecorator.AnteHandle(ctx, testTx, false, dummyNextAnteHandler)
 	assert.NoError(
@@ -116,7 +87,7 @@ func TestGasTrackingAnteHandler(t *testing.T) {
 		"Gastracking ante handler should not return an error",
 	)
 
-	currentBlockTrackingInfo, err := keeper.GetCurrentBlockTracking(ctx)
+	currentBlockTrackingInfo := keeper.GetCurrentBlockTracking(ctx)
 	assert.NoError(t, err, "Current block tracking info should exists")
 
 	assert.Equal(t, 1, len(currentBlockTrackingInfo.TxTrackingInfos), "Only 1 txtracking info should be there")
@@ -140,84 +111,9 @@ func TestGasTrackingAnteHandler(t *testing.T) {
 		"Gastracking ante handler should not return an error",
 	)
 
-	currentBlockTrackingInfo, err = keeper.GetCurrentBlockTracking(ctx)
-	assert.NoError(t, err, "Current block tracking info should exists")
+	currentBlockTrackingInfo = keeper.GetCurrentBlockTracking(ctx)
 
 	assert.Equal(t, 2, len(currentBlockTrackingInfo.TxTrackingInfos), "Only 1 txtracking info should be there")
 	assert.Equal(t, testTx.Gas, currentBlockTrackingInfo.TxTrackingInfos[1].MaxGasAllowed, "MaxGasAllowed must match the Gas field of tx")
 	assert.Equal(t, expectedDecCoins, currentBlockTrackingInfo.TxTrackingInfos[1].MaxContractRewards, "MaxContractReward must be half of the tx fees")
 }
-
-// TODO: this is shared test util, that is copied
-// from /keeper/keeper_test, refactor
-func createTestBaseKeeperAndContext(t *testing.T, contractAdmin sdk.AccAddress) (sdk.Context, keeper.Keeper) {
-	encodingConfig := simapp.MakeTestEncodingConfig()
-	appCodec := encodingConfig.Marshaler
-
-	memDB := db.NewMemDB()
-	ms := store.NewCommitMultiStore(memDB)
-
-	mkey := sdk.NewKVStoreKey("test")
-	tkey := sdk.NewTransientStoreKey("transient_test")
-	tstoreKey := sdk.NewTransientStoreKey(gastracker.TStoreKey)
-
-	ms.MountStoreWithDB(mkey, sdk.StoreTypeIAVL, memDB)
-	ms.MountStoreWithDB(tkey, sdk.StoreTypeIAVL, memDB)
-	ms.MountStoreWithDB(storeKey, sdk.StoreTypeIAVL, memDB)
-	ms.MountStoreWithDB(tstoreKey, sdk.StoreTypeTransient, memDB)
-
-	err := ms.LoadLatestVersion()
-	require.NoError(t, err, "Loading latest version should not fail")
-
-	pkeeper := paramskeeper.NewKeeper(appCodec, encodingConfig.Amino, mkey, tkey)
-	subspace := pkeeper.Subspace(gstTypes.ModuleName)
-
-	keeper := keeper.NewGasTrackingKeeper(
-		storeKey,
-		appCodec,
-		subspace,
-		NewTestContractInfoView(contractAdmin.String()),
-		wasmkeeper.NewDefaultWasmGasRegister(),
-		nil,
-	)
-
-	ctx := sdk.NewContext(ms, tmproto.Header{
-		Height: 10000,
-		Time:   time.Date(2020, time.April, 22, 12, 0, 0, 0, time.UTC),
-	}, false, tmLog.NewTMLogger(os.Stdout))
-
-	params := gstTypes.DefaultParams()
-	subspace.SetParamSet(ctx, &params)
-	return ctx, keeper
-}
-
-func CreateTestKeeperAndContext(t *testing.T, contractAdmin sdk.AccAddress) (sdk.Context, keeper.Keeper) {
-	return createTestBaseKeeperAndContext(t, contractAdmin)
-}
-
-type TestContractInfoView struct {
-	keeper.ContractInfoView
-	adminMap     map[string]string
-	defaultAdmin string
-}
-
-func NewTestContractInfoView(defaultAdmin string) *TestContractInfoView {
-	return &TestContractInfoView{
-		adminMap:     make(map[string]string),
-		defaultAdmin: defaultAdmin,
-	}
-}
-
-func (t *TestContractInfoView) GetContractInfo(_ sdk.Context, contractAddress sdk.AccAddress) *wasmTypes.ContractInfo {
-	if admin, ok := t.adminMap[contractAddress.String()]; ok {
-		return &wasmTypes.ContractInfo{Admin: admin}
-	} else {
-		return &wasmTypes.ContractInfo{Admin: t.defaultAdmin}
-	}
-}
-
-func (t *TestContractInfoView) AddContractToAdminMapping(contractAddress string, admin string) {
-	t.adminMap[contractAddress] = admin
-}
-
-var _ keeper.ContractInfoView = &TestContractInfoView{}

x/gastracker/keeper/grpc_query.go

@@ -47,10 +47,7 @@ func (g *queryServer) BlockGasTracking(ctx context.Context, request *gstTypes.Qu
 		return nil, status.Error(codes.InvalidArgument, "empty request")
 	}
 
-	blockGasTracking, err := g.keeper.GetCurrentBlockTracking(sdk.UnwrapSDKContext(ctx))
-	if err != nil {
-		return nil, err
-	}
+	blockGasTracking := g.keeper.GetCurrentBlockTracking(sdk.UnwrapSDKContext(ctx))
 
 	return &gstTypes.QueryBlockGasTrackingResponse{
 		BlockGasTracking: &blockGasTracking,

x/gastracker/keeper/keeper.go

@@ -1,13 +1,13 @@
 package keeper
 
 import (
+	"fmt"
 	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
 	wasmTypes "github.com/CosmWasm/wasmd/x/wasm/types"
-	"github.com/cosmos/cosmos-sdk/store/prefix"
-	minttypes "github.com/cosmos/cosmos-sdk/x/mint/types"
-
 	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	minttypes "github.com/cosmos/cosmos-sdk/x/mint/types"
 	paramsTypes "github.com/cosmos/cosmos-sdk/x/params/types"
 
 	gastracker "github.com/archway-network/archway/x/gastracker"
@@ -23,13 +23,13 @@ type MintKeeper interface {
 }
 
 type Keeper struct {
+	WasmGasRegister wasmkeeper.GasRegister // can safely be exported since it's readonly.
+
 	key              sdk.StoreKey
 	cdc              codec.Codec
 	paramSpace       gastracker.Subspace
 	contractInfoView ContractInfoView
-	wasmGasRegister  wasmkeeper.GasRegister
-
-	mintKeeper MintKeeper
+	mintKeeper       MintKeeper
 }
 
 func NewGasTrackingKeeper(
@@ -48,7 +48,7 @@ func NewGasTrackingKeeper(
 		cdc:              appCodec,
 		paramSpace:       paramSpace,
 		contractInfoView: contractInfoView,
-		wasmGasRegister:  gasRegister,
+		WasmGasRegister:  gasRegister,
 		mintKeeper:       mintKeeper,
 	}
 }
@@ -159,84 +159,92 @@ func (k Keeper) CommitPendingContractMetadata(ctx sdk.Context) (int, error) {
 	return len(keysToBeDeleted), nil
 }
 
-func (k Keeper) TrackNewBlock(ctx sdk.Context) error {
-	gstKvStore := ctx.KVStore(k.key)
-	bz, err := k.cdc.Marshal(&gastracker.BlockGasTracking{})
-	if err != nil {
-		return err
+func (k Keeper) TrackNewBlock(ctx sdk.Context) {
+	// reset tx identifier
+	k.ResetTxIdentifier(ctx)
+	// delete tx tracking information
+	store := prefix.NewStore(ctx.KVStore(k.key), append(gastracker.PrefixGasTrackingTxTracking, sdk.Uint64ToBigEndian(uint64(ctx.BlockHeight()))...))
+	iter := store.Iterator(nil, nil)
+	var keys [][]byte
+	for ; iter.Valid(); iter.Next() {
+		keys = append(keys, iter.Key())
 	}
-	gstKvStore.Set([]byte(gastracker.CurrentBlockTrackingKey), bz)
-	return nil
-}
 
-func (k Keeper) GetCurrentBlockTracking(ctx sdk.Context) (gastracker.BlockGasTracking, error) {
-	gstKvStore := ctx.KVStore(k.key)
+	for _, key := range keys {
+		store.Delete(key)
+	}
+}
 
+func (k Keeper) GetCurrentBlockTracking(ctx sdk.Context) gastracker.BlockGasTracking {
+	store := prefix.NewStore(ctx.KVStore(k.key), gastracker.PrefixGasTrackingTxTracking)
+	// we prefix over current block height
+	iter := prefix.NewStore(store, sdk.Uint64ToBigEndian(uint64(ctx.BlockHeight()))).Iterator(nil, nil)
+	defer iter.Close()
 	var currentBlockTracking gastracker.BlockGasTracking
-	bz := gstKvStore.Get([]byte(gastracker.CurrentBlockTrackingKey))
-	if bz == nil {
-		return currentBlockTracking, gastracker.ErrBlockTrackingDataNotFound
+
+	for ; iter.Valid(); iter.Next() {
+		v := gastracker.TransactionTracking{}
+		k.cdc.MustUnmarshal(iter.Value(), &v)
+		currentBlockTracking.TxTrackingInfos = append(currentBlockTracking.TxTrackingInfos, v)
 	}
-	err := k.cdc.Unmarshal(bz, &currentBlockTracking)
-	return currentBlockTracking, err
+
+	return currentBlockTracking
 }
 
-func (k Keeper) TrackNewTx(ctx sdk.Context, fee []sdk.DecCoin, gasLimit uint64) error {
-	gstKvStore := ctx.KVStore(k.key)
+func (k Keeper) TrackNewTx(ctx sdk.Context, fee []sdk.DecCoin, gasLimit uint64) {
+	txIdentifier := k.GetAndIncreaseTxIdentifier(ctx)
+	store := prefix.NewStore(ctx.KVStore(k.key), gastracker.PrefixGasTrackingTxTracking)
 
 	var currentTxGasTracking gastracker.TransactionTracking
 	currentTxGasTracking.MaxContractRewards = fee
 	currentTxGasTracking.MaxGasAllowed = gasLimit
 
-	bz := gstKvStore.Get([]byte(gastracker.CurrentBlockTrackingKey))
-	if bz == nil {
-		return gastracker.ErrBlockTrackingDataNotFound
-	}
-	var currentBlockTracking gastracker.BlockGasTracking
-	err := k.cdc.Unmarshal(bz, &currentBlockTracking)
-	if err != nil {
-		return err
-	}
-	currentBlockTracking.TxTrackingInfos = append(currentBlockTracking.TxTrackingInfos, currentTxGasTracking)
-	bz, err = k.cdc.Marshal(&currentBlockTracking)
-	if err != nil {
-		return err
-	}
-	gstKvStore.Set([]byte(gastracker.CurrentBlockTrackingKey), bz)
-	return nil
+	store.Set(
+		append(sdk.Uint64ToBigEndian(uint64(ctx.BlockHeight())), sdk.Uint64ToBigEndian(txIdentifier)...),
+		k.cdc.MustMarshal(&currentTxGasTracking),
+	)
 }
 
-func (k Keeper) TrackContractGasUsage(ctx sdk.Context, contractAddress sdk.AccAddress, originalGas wasmTypes.GasConsumptionInfo, operation gastracker.ContractOperation) error {
+func (k Keeper) TrackContractGasUsage(ctx sdk.Context, contractAddress sdk.AccAddress, originalGas wasmTypes.GasConsumptionInfo, operation gastracker.ContractOperation) {
+	store := prefix.NewStore(ctx.KVStore(k.key), gastracker.PrefixGasTrackingTxTracking)
+	key := append(sdk.Uint64ToBigEndian(uint64(ctx.BlockHeight())), sdk.Uint64ToBigEndian(k.GetCurrentTxIdentifier(ctx))...)
 
-	gstKvStore := ctx.KVStore(k.key)
-	bz := gstKvStore.Get([]byte(gastracker.CurrentBlockTrackingKey))
-	if bz == nil {
-		return gastracker.ErrBlockTrackingDataNotFound
-	}
-	var currentBlockGasTracking gastracker.BlockGasTracking
-	err := k.cdc.Unmarshal(bz, &currentBlockGasTracking)
-	if err != nil {
-		return err
+	bytes := store.Get(key)
+	if bytes == nil {
+		panic(fmt.Errorf("no gas tracking found for tx"))
 	}
+	var transactionTracking gastracker.TransactionTracking
+	k.cdc.MustUnmarshal(bytes, &transactionTracking)
 
-	txsLen := len(currentBlockGasTracking.TxTrackingInfos)
-	if txsLen == 0 {
-		return gastracker.ErrTxTrackingDataNotFound
-	}
-	currentTxGasTracking := currentBlockGasTracking.TxTrackingInfos[txsLen-1]
-	currentBlockGasTracking.TxTrackingInfos[txsLen-1].ContractTrackingInfos = append(currentTxGasTracking.ContractTrackingInfos, gastracker.ContractGasTracking{
+	transactionTracking.ContractTrackingInfos = append(transactionTracking.ContractTrackingInfos, gastracker.ContractGasTracking{
 		Address:        contractAddress.String(),
 		OriginalVmGas:  originalGas.VMGas,
 		OriginalSdkGas: originalGas.SDKGas,
 		Operation:      operation,
 	})
-	bz, err = k.cdc.Marshal(&currentBlockGasTracking)
-	if err != nil {
-		return err
-	}
+	store.Set(key, k.cdc.MustMarshal(&transactionTracking))
+}
 
-	gstKvStore.Set([]byte(gastracker.CurrentBlockTrackingKey), bz)
-	return nil
+// GetAndIncreaseTxIdentifier gets the current Tx identifier.
+// Then increases the current tx identifier by 1.
+// Assumes there is already a valid value saved in the store.
+func (k Keeper) GetAndIncreaseTxIdentifier(ctx sdk.Context) uint64 {
+	store := prefix.NewStore(ctx.KVStore(k.key), gastracker.PrefixGasTrackingTxIdentifier)
+	value := sdk.BigEndianToUint64(store.Get(gastracker.KeyTxIdentifier))
+	store.Set(gastracker.KeyTxIdentifier, sdk.Uint64ToBigEndian(value+1))
+	return value
+}
+
+// GetCurrentTxIdentifier gets the current Tx identifier.
+// Contract: assumes GetAndIncreaseTxIdentifier was called
+// at least once in this block.
+func (k Keeper) GetCurrentTxIdentifier(ctx sdk.Context) uint64 {
+	return sdk.BigEndianToUint64(prefix.NewStore(ctx.KVStore(k.key), gastracker.PrefixGasTrackingTxIdentifier).Get(gastracker.KeyTxIdentifier)) - 1
+}
+
+// ResetTxIdentifier resets the current Tx identifier to 0.
+func (k Keeper) ResetTxIdentifier(ctx sdk.Context) {
+	prefix.NewStore(ctx.KVStore(k.key), gastracker.PrefixGasTrackingTxIdentifier).Set(gastracker.KeyTxIdentifier, sdk.Uint64ToBigEndian(0))
 }
 
 // UpdateDappInflationaryRewards sets the current block inflationary rewards.

x/gastracker/keeper/keeper_contract_gas_processor.go

@@ -64,12 +64,10 @@ func (k Keeper) IngestGasRecord(ctx types.Context, records []wasmTypes.ContractG
 			operation = gastracker.ContractOperation_CONTRACT_OPERATION_UNSPECIFIED
 		}
 
-		if err := k.TrackContractGasUsage(ctx, contractAddress, wasmTypes.GasConsumptionInfo{
+		k.TrackContractGasUsage(ctx, contractAddress, wasmTypes.GasConsumptionInfo{
 			SDKGas: record.OriginalGas.SDKGas,
-			VMGas:  k.wasmGasRegister.FromWasmVMGas(record.OriginalGas.VMGas),
-		}, operation); err != nil {
-			return err
-		}
+			VMGas:  k.WasmGasRegister.FromWasmVMGas(record.OriginalGas.VMGas),
+		}, operation)
 	}
 
 	return nil