Merge pull request #312 from ComposableFi/iterate-delegation-storage

iterate delegation storage

custom/staking/keeper/msg_server.go

@@ -5,6 +5,7 @@ import (
 
 	"cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 )
@@ -47,6 +48,12 @@ func (k msgServer) Delegate(goCtx context.Context, msg *types.MsgDelegate) (*typ
 	k.stakingmiddleware.SetLastTotalPower(ctx, math.Int{})
 
 	k.stakingmiddleware.SetDelegation(ctx, msg.DelegatorAddress, msg.ValidatorAddress, msg.Amount.Denom, msg.Amount.Amount)
+	delegations := k.stakingmiddleware.DequeueAllDelegation(ctx)
+	if len(delegations) >= 1 {
+		return nil, sdkerrors.Wrapf(
+			sdkerrors.ErrInvalidRequest, "should always be less then 2 : got %s, expected %s", len(delegations),
+		)
+	}
 
 	return &types.MsgDelegateResponse{}, nil
 	// return nil, fmt.Errorf("My custom error: Nikita")

x/stakingmiddleware/keeper/keeper.go

@@ -124,6 +124,28 @@ func (k Keeper) IterateDelegations(ctx sdk.Context, fn func(index int64, ubd typ
 	}
 }
 
+// DequeueAllMatureUBDQueue returns a concatenated list of all the timeslices inclusively previous to
+// currTime, and deletes the timeslices from the queue.
+func (k Keeper) DequeueAllDelegation(ctx sdk.Context) (delegations []types.Delegation) {
+	store := ctx.KVStore(k.storeKey)
+
+	// gets an iterator for all timeslices from time 0 until the current Blockheader time
+	delegationIterator := sdk.KVStorePrefixIterator(store, types.DelegationKey)
+	defer delegationIterator.Close()
+
+	for ; delegationIterator.Valid(); delegationIterator.Next() {
+		delegation := types.Delegation{}
+		value := delegationIterator.Value()
+		k.cdc.MustUnmarshal(value, &delegation)
+
+		delegations = append(delegations, delegation)
+
+		store.Delete(delegationIterator.Key())
+	}
+
+	return delegations
+}
+
 func GetValidatorAddr(d types.Delegation) sdk.ValAddress {
 	addr, err := sdk.ValAddressFromBech32(d.ValidatorAddress)
 	if err != nil {