Add indexing store for ValidatorDelegations query

app/app.go

@@ -95,6 +95,8 @@ import (
 	iscnkeeper "github.com/likecoin/likechain/x/iscn/keeper"
 	iscntypes "github.com/likecoin/likechain/x/iscn/types"
 
+	stakingwithindex "github.com/likecoin/likechain/x/staking"
+
 	bech32authmigration "github.com/likecoin/likechain/bech32-migration/auth"
 	bech32govmigration "github.com/likecoin/likechain/bech32-migration/gov"
 	bech32slashingmigration "github.com/likecoin/likechain/bech32-migration/slashing"
@@ -313,12 +315,14 @@ func NewLikeApp(
 	app.registerUpgradeHandlers()
 	app.IscnKeeper = iscnkeeper.NewKeeper(appCodec, keys[iscntypes.StoreKey], app.AccountKeeper, app.BankKeeper, iscnSubspace)
 
-	// register the staking hooks
-	// NOTE: stakingKeeper above is passed by reference, so that it will contain these hooks
-	app.StakingKeeper = *stakingKeeper.SetHooks(
-		stakingtypes.NewMultiStakingHooks(app.DistrKeeper.Hooks(), app.SlashingKeeper.Hooks()),
+	hookedStakingKeeper, stakingAppModule := stakingwithindex.SetupStakingModule(
+		homePath, appCodec,
+		&stakingKeeper, app.AccountKeeper, app.BankKeeper, app.DistrKeeper, app.SlashingKeeper,
+		appOpts,
 	)
 
+	app.StakingKeeper = *hookedStakingKeeper
+
 	// Create IBC Keeper
 	app.IBCKeeper = ibckeeper.NewKeeper(
 		appCodec, keys[ibchost.StoreKey], ibcHostSubspace, app.StakingKeeper, app.UpgradeKeeper, scopedIBCKeeper,
@@ -381,7 +385,7 @@ func NewLikeApp(
 		distr.NewAppModule(
 			appCodec, app.DistrKeeper, app.AccountKeeper, app.BankKeeper, app.StakingKeeper,
 		),
-		staking.NewAppModule(appCodec, app.StakingKeeper, app.AccountKeeper, app.BankKeeper),
+		stakingAppModule,
 		upgrade.NewAppModule(app.UpgradeKeeper),
 		evidence.NewAppModule(app.EvidenceKeeper),
 		ibc.NewAppModule(app.IBCKeeper),

cmd/liked/cmd/cmd.go

@@ -45,6 +45,7 @@ import (
 	simappcli "github.com/cosmos/cosmos-sdk/simapp/simd/cmd"
 
 	"github.com/likecoin/likechain/ip"
+	stakingwithindex "github.com/likecoin/likechain/x/staking"
 
 	serverconfig "github.com/cosmos/cosmos-sdk/server/config"
 )
@@ -88,9 +89,14 @@ func addCrisisFlag(startCmd *cobra.Command) {
 	crisis.AddModuleInitFlags(startCmd)
 }
 
+func addEnableCustomIndexFlag(startCmd *cobra.Command) {
+	stakingwithindex.AddEnableCustomIndexFlag(startCmd)
+}
+
 func addStartFlags(startCmd *cobra.Command) {
 	addCrisisFlag(startCmd)
 	addGetIpFlag(startCmd)
+	addEnableCustomIndexFlag(startCmd)
 }
 
 func queryCommand() *cobra.Command {

x/staking/cmd.go

@@ -0,0 +1,78 @@
+package staking
+
+import (
+	"fmt"
+	"path"
+
+	"github.com/spf13/cobra"
+
+	dbm "github.com/tendermint/tm-db"
+
+	"github.com/cosmos/cosmos-sdk/codec"
+	servertypes "github.com/cosmos/cosmos-sdk/server/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+	distrkeeper "github.com/cosmos/cosmos-sdk/x/distribution/keeper"
+	slashingkeeper "github.com/cosmos/cosmos-sdk/x/slashing/keeper"
+	originalstaking "github.com/cosmos/cosmos-sdk/x/staking"
+	"github.com/cosmos/cosmos-sdk/x/staking/keeper"
+	"github.com/cosmos/cosmos-sdk/x/staking/types"
+)
+
+const FlagEnableCustomIndex = "enable-custom-index"
+
+func AddEnableCustomIndexFlag(startCmd *cobra.Command) {
+	startCmd.Flags().Bool(FlagEnableCustomIndex, false, "Enable custom index (currently for ValidatorDelegations only)")
+}
+
+func setupOriginalStakingModule(
+	homePath string, appCodec codec.Codec,
+	stakingKeeper *keeper.Keeper, ak types.AccountKeeper, bk types.BankKeeper,
+	distrKeeper distrkeeper.Keeper, slashingKeeper slashingkeeper.Keeper,
+) (*keeper.Keeper, module.AppModule) {
+	stakingKeeper = stakingKeeper.SetHooks(
+		types.NewMultiStakingHooks(
+			distrKeeper.Hooks(),
+			slashingKeeper.Hooks(),
+		),
+	)
+	am := originalstaking.NewAppModule(appCodec, *stakingKeeper, ak, bk)
+	return stakingKeeper, am
+}
+
+func setupStakingModuleWithCustomIndex(
+	homePath string, appCodec codec.Codec,
+	stakingKeeper *keeper.Keeper, ak types.AccountKeeper, bk types.BankKeeper,
+	distrKeeper distrkeeper.Keeper, slashingKeeper slashingkeeper.Keeper,
+) (*keeper.Keeper, module.AppModule) {
+	stakingIndexDB, err := dbm.NewGoLevelDB("index", path.Join(homePath, "data"))
+	if err != nil {
+		panic(fmt.Errorf("failed to create indexing DB for staking module: %s", err))
+	}
+	// TODO: When to close this DB?
+	// This should be handled by the finalizer set in runtime so should not be a problem.
+	// But we still want to do something in code
+	stakingQuerier := NewQuerier(stakingKeeper, appCodec, stakingIndexDB)
+	stakingKeeper = stakingKeeper.SetHooks(
+		types.NewMultiStakingHooks(
+			distrKeeper.Hooks(),
+			slashingKeeper.Hooks(),
+			NewHooks(stakingQuerier),
+		),
+	)
+	am := NewAppModule(appCodec, *stakingKeeper, ak, bk, stakingQuerier)
+	return stakingKeeper, am
+}
+
+func SetupStakingModule(
+	homePath string, appCodec codec.Codec,
+	stakingKeeper *keeper.Keeper, ak types.AccountKeeper, bk types.BankKeeper,
+	distrKeeper distrkeeper.Keeper, slashingKeeper slashingkeeper.Keeper,
+	appOpts servertypes.AppOptions,
+) (*keeper.Keeper, module.AppModule) {
+	opt := appOpts.Get(FlagEnableCustomIndex)
+	if shouldEnableCustomIndex, ok := opt.(bool); ok && shouldEnableCustomIndex {
+		return setupStakingModuleWithCustomIndex(homePath, appCodec, stakingKeeper, ak, bk, distrKeeper, slashingKeeper)
+	} else {
+		return setupOriginalStakingModule(homePath, appCodec, stakingKeeper, ak, bk, distrKeeper, slashingKeeper)
+	}
+}

x/staking/debug/db.go

@@ -0,0 +1,85 @@
+package debug
+
+import (
+	"github.com/tendermint/tendermint/libs/log"
+	dbm "github.com/tendermint/tm-db"
+)
+
+type DBDebugWrap struct {
+	dbm.DB
+	Logger log.Logger
+}
+
+func (d *DBDebugWrap) Set(key []byte, value []byte) error {
+	return d.DB.Set(key, value)
+}
+
+func (d *DBDebugWrap) Delete(key []byte) error {
+	return d.DB.Delete(key)
+}
+
+func (d *DBDebugWrap) Close() error {
+	d.Logger.Debug("DB Close")
+	return d.DB.Close()
+}
+
+func (d *DBDebugWrap) NewBatch() dbm.Batch {
+	return &DBDebugBatch{
+		Batch:  d.DB.NewBatch(),
+		Logger: d.Logger,
+	}
+}
+
+func (d *DBDebugWrap) Iterator(start, end []byte) (dbm.Iterator, error) {
+	// d.Logger.Debug("DB Iterator", "start", start, "end", end)
+	it, err := d.DB.Iterator(start, end)
+	if err != nil {
+		return nil, err
+	}
+	return &DBDebugIterator{
+		Iterator: it,
+		Logger:   d.Logger,
+	}, nil
+}
+
+type DBDebugBatch struct {
+	dbm.Batch
+	Logger log.Logger
+}
+
+func (b *DBDebugBatch) Set(key, value []byte) error {
+	// b.Logger.Debug("DB Batch Set", "key", key, "value", value)
+	return b.Batch.Set(key, value)
+}
+
+func (b *DBDebugBatch) Delete(key []byte) error {
+	// b.Logger.Debug("DB Batch Set", "key", key)
+	return b.Batch.Delete(key)
+}
+
+func (b *DBDebugBatch) Write() error {
+	// b.Logger.Debug("DB Batch Write")
+	return b.Batch.Write()
+}
+
+func (b *DBDebugBatch) WriteSync() error {
+	// b.Logger.Debug("DB Batch WriteSync")
+	return b.Batch.WriteSync()
+}
+
+type DBDebugIterator struct {
+	dbm.Iterator
+	Logger log.Logger
+}
+
+func (it *DBDebugIterator) Key() []byte {
+	key := it.Iterator.Key()
+	// it.Logger.Debug("DB Iterator Key", "key", key)
+	return key
+}
+
+func (it *DBDebugIterator) Value() []byte {
+	value := it.Iterator.Value()
+	// it.Logger.Debug("DB Iterator value", "value", value)
+	return value
+}

x/staking/grpc_query.go

@@ -0,0 +1,109 @@
+package staking
+
+import (
+	"context"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	prefixstore "github.com/cosmos/cosmos-sdk/store/prefix"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/query"
+
+	"github.com/cosmos/cosmos-sdk/x/staking/keeper"
+	"github.com/cosmos/cosmos-sdk/x/staking/types"
+)
+
+// keeper.DelegationsToDelegationResponses is so inefficient that it does one query for bonded token + one query for
+// validator for each delegation, which are both unchanged
+func DelegationsToDelegationResponses(
+	ctx sdk.Context, k keeper.Keeper, delegations types.Delegations, valAddr sdk.ValAddress,
+) (types.DelegationResponses, error) {
+	val, found := k.GetValidator(ctx, valAddr)
+	if !found {
+		return nil, types.ErrNoValidatorFound
+	}
+	bondDenom := k.BondDenom(ctx)
+	resp := make(types.DelegationResponses, len(delegations))
+	for i, del := range delegations {
+		resp[i] = types.NewDelegationResp(
+			del.GetDelegatorAddr(),
+			valAddr,
+			del.Shares,
+			sdk.NewCoin(bondDenom, val.TokensFromShares(del.Shares).TruncateInt()),
+		)
+	}
+
+	return resp, nil
+}
+
+func (q *Querier) ValidatorDelegations(c context.Context, req *types.QueryValidatorDelegationsRequest) (*types.QueryValidatorDelegationsResponse, error) {
+	ctx := sdk.UnwrapSDKContext(c)
+	blockHeight := ctx.BlockHeight()
+	height := q.GetHeight()
+	ctx.Logger().Debug(
+		"Querying ValidatorDelegations from indexer",
+		"req", req,
+		"block_height", blockHeight,
+		"indexer_height", height,
+	)
+	if blockHeight != int64(height) {
+		// Indexing store does not store old heights, so fallback to the slow path
+		// Maybe also because the indexing is not done yet
+		ctx.Logger().Info(
+			"Querying height not current height in ValidatorDelegations query (could be due to index not yet built), fallback to slow path",
+			"block_height", blockHeight,
+			"indexer_height", height,
+		)
+		return q.QueryServer.ValidatorDelegations(c, req)
+	}
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "empty request")
+	}
+	if req.ValidatorAddr == "" {
+		return nil, status.Error(codes.InvalidArgument, "validator address cannot be empty")
+	}
+	valAddr, err := sdk.ValAddressFromBech32(req.ValidatorAddr)
+	if err != nil {
+		return nil, err
+	}
+	store := prefixstore.NewStore(q.readStore, getIndexValidatorPrefix(valAddr))
+	delegations := []types.Delegation{}
+	pageRes, err := query.FilteredPaginate(store, req.Pagination, func(key, value []byte, accumulate bool) (bool, error) {
+		if accumulate {
+			delegation := types.MustUnmarshalDelegation(q.cdc, value)
+			delegations = append(delegations, delegation)
+		}
+		return true, nil
+	})
+	heightAfterQuery := q.GetHeight()
+	if heightAfterQuery != height {
+		// Write occurred during query, queried data could be invalid.
+		// Since height changed, the query must be querying an old height (relative to the new height) which we don't
+		// support, so fallback to the slow path.
+
+		// Note:
+		// In most of the cases, the user actually wants to query the newest height (i.e. height = 0 in query),
+		// but the SDK will automatically set the height to current height before passing to the QueryServer,
+		// so maybe in this case we actually should not fallback to the slow path?
+		// But the SDK will set the result's height to the old height... So it seems not possible to get consistent result.
+		ctx.Logger().Debug(
+			"Write occurred during ValidatorDelegations query, fallback to slow path",
+			"height", height,
+			"height_after_query", heightAfterQuery,
+		)
+		return q.QueryServer.ValidatorDelegations(c, req)
+	}
+	delResponses := types.DelegationResponses{}
+	if len(delegations) > 0 {
+		delResponses, err = DelegationsToDelegationResponses(
+			ctx, *q.keeper, delegations, delegations[0].GetValidatorAddr(),
+		)
+		if err != nil {
+			return nil, status.Error(codes.Internal, err.Error())
+		}
+	}
+	return &types.QueryValidatorDelegationsResponse{
+		DelegationResponses: delResponses, Pagination: pageRes,
+	}, nil
+}

x/staking/hooks.go

@@ -0,0 +1,61 @@
+package staking
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/staking/types"
+)
+
+type IndexingHooks struct {
+	querier *Querier
+}
+
+var _ types.StakingHooks = IndexingHooks{}
+
+func NewHooks(querier *Querier) IndexingHooks {
+	return IndexingHooks{
+		querier: querier,
+	}
+}
+
+func (hooks IndexingHooks) BeforeDelegationRemoved(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) {
+	hooks.querier.RemoveIndex(ctx, delAddr, valAddr)
+}
+
+func (hooks IndexingHooks) AfterDelegationModified(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) {
+	hooks.querier.UpdateIndex(ctx, delAddr, valAddr)
+}
+
+func (IndexingHooks) BeforeDelegationSharesModified(sdk.Context, sdk.AccAddress, sdk.ValAddress) {}
+func (IndexingHooks) BeforeDelegationCreated(sdk.Context, sdk.AccAddress, sdk.ValAddress)        {}
+func (IndexingHooks) AfterValidatorCreated(sdk.Context, sdk.ValAddress)                          {}
+func (IndexingHooks) BeforeValidatorModified(sdk.Context, sdk.ValAddress)                        {}
+func (IndexingHooks) AfterValidatorRemoved(sdk.Context, sdk.ConsAddress, sdk.ValAddress)         {}
+func (IndexingHooks) AfterValidatorBonded(sdk.Context, sdk.ConsAddress, sdk.ValAddress)          {}
+func (IndexingHooks) AfterValidatorBeginUnbonding(sdk.Context, sdk.ConsAddress, sdk.ValAddress)  {}
+func (IndexingHooks) BeforeValidatorSlashed(sdk.Context, sdk.ValAddress, sdk.Dec)                {}
+
+type DebugHooks struct{}
+
+func (DebugHooks) BeforeDelegationRemoved(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) {
+	ctx.Logger().Debug(
+		"ValidatorDelegations indexer: BeforeDelegationRemoved",
+		"delegator_addr", delAddr.String(),
+		"validator_addr", valAddr.String(),
+	)
+}
+
+func (DebugHooks) AfterDelegationModified(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) {
+	ctx.Logger().Debug(
+		"ValidatorDelegations indexer: AfterDelegationModified",
+		"delegator_addr", delAddr.String(),
+		"validator_addr", valAddr.String(),
+	)
+}
+func (DebugHooks) BeforeDelegationSharesModified(sdk.Context, sdk.AccAddress, sdk.ValAddress) {}
+func (DebugHooks) BeforeDelegationCreated(sdk.Context, sdk.AccAddress, sdk.ValAddress)        {}
+func (DebugHooks) AfterValidatorCreated(sdk.Context, sdk.ValAddress)                          {}
+func (DebugHooks) BeforeValidatorModified(sdk.Context, sdk.ValAddress)                        {}
+func (DebugHooks) AfterValidatorRemoved(sdk.Context, sdk.ConsAddress, sdk.ValAddress)         {}
+func (DebugHooks) AfterValidatorBonded(sdk.Context, sdk.ConsAddress, sdk.ValAddress)          {}
+func (DebugHooks) AfterValidatorBeginUnbonding(sdk.Context, sdk.ConsAddress, sdk.ValAddress)  {}
+func (DebugHooks) BeforeValidatorSlashed(sdk.Context, sdk.ValAddress, sdk.Dec)                {}