add logic to EndBlock custom staking wrapper to dequeue all request.

custom/staking/keeper/msg_server.go

@@ -29,11 +29,12 @@
 }
 
 func (k msgServer) Delegate(goCtx context.Context, msg *types.MsgDelegate) (*types.MsgDelegateResponse, error) {
+	//todo add validation. check that this account really has some coins
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
 	bondDenom := k.BondDenom(ctx)
 	if msg.Amount.Denom != bondDenom {
 		return nil, sdkerrors.Wrapf(
 			sdkerrors.ErrInvalidRequest, "invalid coin denomination: got %s, expected %s", msg.Amount.Denom, bondDenom,
 		)
 	}
@@ -56,15 +57,27 @@
 }
 
 func (k msgServer) BeginRedelegate(goCtx context.Context, msg *types.MsgBeginRedelegate) (*types.MsgBeginRedelegateResponse, error) {
-	return k.msgServer.BeginRedelegate(goCtx, msg)
+	//todo add validation. check that this account really has some coins staked with this validator
+	ctx := sdk.UnwrapSDKContext(goCtx)
+	k.Stakingmiddleware.SetBeginRedelegation(ctx, msg.DelegatorAddress, msg.ValidatorSrcAddress, msg.ValidatorDstAddress, msg.Amount.Denom, msg.Amount.Amount)
+	return &types.MsgBeginRedelegateResponse{}, nil
+	// return k.msgServer.BeginRedelegate(goCtx, msg)
 }
 
 func (k msgServer) Undelegate(goCtx context.Context, msg *types.MsgUndelegate) (*types.MsgUndelegateResponse, error) {
-	return k.msgServer.Undelegate(goCtx, msg)
+	//todo add validation. check that this account really has some coins staked with this validator
+	ctx := sdk.UnwrapSDKContext(goCtx)
+	k.Stakingmiddleware.SetUndelegation(ctx, msg.DelegatorAddress, msg.ValidatorAddress, msg.Amount.Denom, msg.Amount.Amount)
+	return &types.MsgUndelegateResponse{}, nil
+	// return k.msgServer.Undelegate(goCtx, msg)
 }
 
 func (k msgServer) CancelUnbondingDelegation(goCtx context.Context, msg *types.MsgCancelUnbondingDelegation) (*types.MsgCancelUnbondingDelegationResponse, error) {
-	return k.msgServer.CancelUnbondingDelegation(goCtx, msg)
+	//todo add validation. check that this account really has some coins staked with this validator
+	ctx := sdk.UnwrapSDKContext(goCtx)
+	k.Stakingmiddleware.SetCancelUndelegation(ctx, msg.DelegatorAddress, msg.ValidatorAddress, msg.Amount.Denom, msg.Amount.Amount, msg.CreationHeight)
+	return &types.MsgCancelUnbondingDelegationResponse{}, nil
+	// return k.msgServer.CancelUnbondingDelegation(goCtx, msg)
 }
 
 func (ms msgServer) UpdateParams(goCtx context.Context, msg *types.MsgUpdateParams) (*types.MsgUpdateParamsResponse, error) {

custom/staking/module.go

@@ -58,21 +58,14 @@
 	}
 }
 
-// func (am AppModule) BeginBlock(ctx sdk.Context, _abc abcitype.RequestBeginBlock) {
-// 	//Define the logic around the batching.
-// 	am.AppModule.BeginBlock(ctx, _abc)
-// }
-
 func (am AppModule) EndBlock(ctx sdk.Context, _abc abcitype.RequestEndBlock) []abcitype.ValidatorUpdate {
-	//Define the logic around the batching.
-	//TODO!!!
 
 	println("EndBlock Custom Staking Module")
 
 	delegations := am.keeper.Stakingmiddleware.DequeueAllDelegation(ctx)
 	println("Delegations: ", delegations)
 	println("Delegations len: ", len(delegations))
 	//for delegations print the delegator address and the validator address
 	for _, delegation := range delegations {
 		println("Delegator Address: ", delegation.DelegatorAddress)
 		println("Validator Address: ", delegation.ValidatorAddress)
@@ -83,17 +76,62 @@
 		if err != nil {
 			println("Error for Delegator Address: ", delegation.DelegatorAddress)
 		}
+	}
+
+	beginredelegations := am.keeper.Stakingmiddleware.DequeueAllRedelegation(ctx)
+	println("BeginRedelegations: ", beginredelegations)
+	println("BeginRedelegations len: ", len(beginredelegations))
+	for _, beginredelegation := range beginredelegations {
+		println("Delegator Address: ", beginredelegation.DelegatorAddress)
+		println("Validator Address: ", beginredelegation.ValidatorSrcAddress)
+		// fmt.Println("Amount", delegation.Amount.Amount)
+
+		// msgDelegate := stakingtypes.MsgDelegate{DelegatorAddress: delegation.DelegatorAddress, ValidatorAddress: delegation.ValidatorAddress, Amount: delegation.Amount}
+		// _, err := am.msgServer.Delegate(ctx, &msgDelegate)
+		// if err != nil {
+		// 	println("Error for Delegator Address: ", delegation.DelegatorAddress)
+		// }
+	}
+
+	undelegations := am.keeper.Stakingmiddleware.DequeueAllUndelegation(ctx)
+	println("Undelegation: ", beginredelegations)
+	println("Undelegation len: ", len(beginredelegations))
+	for _, undelegation := range undelegations {
+		println("Undelegation Delegator Address: ", undelegation.DelegatorAddress)
+		println("Undelegation Validator Address: ", undelegation.ValidatorAddress)
+		// fmt.Println("Amount", delegation.Amount.Amount)
+
+		// msgDelegate := stakingtypes.MsgDelegate{DelegatorAddress: delegation.DelegatorAddress, ValidatorAddress: delegation.ValidatorAddress, Amount: delegation.Amount}
+		// _, err := am.msgServer.Delegate(ctx, &msgDelegate)
+		// if err != nil {
+		// 	println("Error for Delegator Address: ", delegation.DelegatorAddress)
+		// }
+	}
 
+	//DequeueAllCancelUnbondingDelegation
+	cancel_unbonding_delegations := am.keeper.Stakingmiddleware.DequeueAllCancelUnbondingDelegation(ctx)
+	println("Cancel Unbonding Delegations: ", cancel_unbonding_delegations)
+	println("Cancel Ubonding Delegations len: ", len(cancel_unbonding_delegations))
+	for _, undelegation := range cancel_unbonding_delegations {
+		println("Cancel Unbonding Delegation  Delegator Address: ", undelegation.DelegatorAddress)
+		println("Cancel Unbonding Delegations Validator Address: ", undelegation.ValidatorAddress)
+		// fmt.Println("Amount", delegation.Amount.Amount)
+
+		// msgDelegate := stakingtypes.MsgDelegate{DelegatorAddress: delegation.DelegatorAddress, ValidatorAddress: delegation.ValidatorAddress, Amount: delegation.Amount}
+		// _, err := am.msgServer.Delegate(ctx, &msgDelegate)
+		// if err != nil {
+		// 	println("Error for Delegator Address: ", delegation.DelegatorAddress)
+		// }
 	}
 
-	ctx.EventManager().EmitEvent(
-		sdk.NewEvent(
-			"DequeueAllDelegation",
-			// sdk.NewAttribute(sdk.AttributeKeyAmount, balances.String()),
-			// sdk.NewAttribute(types.AttributeKeyValidator, dvPair.ValidatorAddress),
-			// sdk.NewAttribute(types.AttributeKeyDelegator, dvPair.DelegatorAddress),
-		),
-	)
+	// ctx.EventManager().EmitEvent(
+	// 	sdk.NewEvent(
+	// 		"DequeueAllDelegation",
+	// 		// sdk.NewAttribute(sdk.AttributeKeyAmount, balances.String()),
+	// 		// sdk.NewAttribute(types.AttributeKeyValidator, dvPair.ValidatorAddress),
+	// 		// sdk.NewAttribute(types.AttributeKeyDelegator, dvPair.DelegatorAddress),
+	// 	),
+	// )
 
 	return am.AppModule.EndBlock(ctx, _abc)
 }

x/stakingmiddleware/keeper/keeper.go

@@ -83,12 +83,12 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 func (k Keeper) SetLastTotalPower(ctx sdk.Context, power sdkmath.Int) {
 	store := ctx.KVStore(k.storeKey)
 	bz := k.cdc.MustMarshal(&sdk.IntProto{Int: power})
-	store.Set(types.DelegationKey, bz)
+	store.Set(types.DelegateKey, bz)
 }
 
 func (k Keeper) GetLastTotalPower(ctx sdk.Context) sdkmath.Int {
 	store := ctx.KVStore(k.storeKey)
-	bz := store.Get(types.DelegationKey)
+	bz := store.Get(types.DelegateKey)
 
 	if bz == nil {
 		return sdkmath.ZeroInt()
@@ -106,31 +106,40 @@ func (k Keeper) SetDelegation(ctx sdk.Context, sourceDelegatorAddress, validator
 
 	store := ctx.KVStore(k.storeKey)
 	b := k.cdc.MustMarshal(&delegation)
-	store.Set(types.GetDelegationKey(delegatorAddress, GetValidatorAddr(delegation)), b)
+	store.Set(types.GetDelegateKey(delegatorAddress, GetValidatorAddr(delegation)), b)
 }
 
-func (k Keeper) IterateDelegations(ctx sdk.Context, fn func(index int64, ubd types.Delegation) (stop bool)) {
+func (k Keeper) SetBeginRedelegation(ctx sdk.Context, sourceDelegatorAddress, validatorSrcAddress, validatorDstAddress, denom string, amount sdkmath.Int) {
+	begindelegation := types.BeginRedelegate{DelegatorAddress: sourceDelegatorAddress, ValidatorSrcAddress: validatorSrcAddress, ValidatorDstAddress: validatorDstAddress, Amount: sdk.NewCoin(denom, amount)}
+	delegatorAddress := sdk.MustAccAddressFromBech32(begindelegation.DelegatorAddress)
+
 	store := ctx.KVStore(k.storeKey)
+	b := k.cdc.MustMarshal(&begindelegation)
+	store.Set(types.GetBeginRedelegateKey(delegatorAddress, GetValidatorAddrFromStr(validatorSrcAddress)), b)
+}
 
-	iterator := sdk.KVStorePrefixIterator(store, types.DelegationKey)
-	defer iterator.Close()
+func (k Keeper) SetUndelegation(ctx sdk.Context, sourceDelegatorAddress, validatorAddress, denom string, amount sdkmath.Int) {
+	undelegation := types.Undelegate{DelegatorAddress: sourceDelegatorAddress, ValidatorAddress: validatorAddress, Amount: sdk.NewCoin(denom, amount)}
+	delegatorAddress := sdk.MustAccAddressFromBech32(undelegation.DelegatorAddress)
 
-	for i := int64(0); iterator.Valid(); iterator.Next() {
-		ubd := MustUnmarshalUBD(k.cdc, iterator.Value())
-		if stop := fn(i, ubd); stop {
-			break
-		}
-		i++
-	}
+	store := ctx.KVStore(k.storeKey)
+	b := k.cdc.MustMarshal(&undelegation)
+	store.Set(types.GetUndelegateKey(delegatorAddress, GetValidatorAddrFromStr(validatorAddress)), b)
+}
+
+func (k Keeper) SetCancelUndelegation(ctx sdk.Context, sourceDelegatorAddress, validatorAddress, denom string, amount sdkmath.Int, height int64) {
+	undelegation := types.CancelUnbondingDelegation{DelegatorAddress: sourceDelegatorAddress, ValidatorAddress: validatorAddress, Amount: sdk.NewCoin(denom, amount), CreationHeight: height}
+	delegatorAddress := sdk.MustAccAddressFromBech32(undelegation.DelegatorAddress)
+
+	store := ctx.KVStore(k.storeKey)
+	b := k.cdc.MustMarshal(&undelegation)
+	store.Set(types.GetCancelUnbondingDelegateKey(delegatorAddress, GetValidatorAddrFromStr(validatorAddress)), b)
 }
 
-// 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)
+	delegationIterator := sdk.KVStorePrefixIterator(store, types.DelegateKey)
 	defer delegationIterator.Close()
 
 	for ; delegationIterator.Valid(); delegationIterator.Next() {
@@ -146,6 +155,70 @@ func (k Keeper) DequeueAllDelegation(ctx sdk.Context) (delegations []types.Deleg
 	return delegations
 }
 
+func (k Keeper) DequeueAllRedelegation(ctx sdk.Context) (redelegations []types.BeginRedelegate) {
+	store := ctx.KVStore(k.storeKey)
+
+	redelegationIterator := sdk.KVStorePrefixIterator(store, types.BeginRedelegateKey)
+	defer redelegationIterator.Close()
+
+	for ; redelegationIterator.Valid(); redelegationIterator.Next() {
+		redelegation := types.BeginRedelegate{}
+		value := redelegationIterator.Value()
+		k.cdc.MustUnmarshal(value, &redelegation)
+
+		redelegations = append(redelegations, redelegation)
+
+		store.Delete(redelegationIterator.Key())
+	}
+
+	return redelegations
+}
+
+func (k Keeper) DequeueAllUndelegation(ctx sdk.Context) (undelegations []types.Undelegate) {
+	store := ctx.KVStore(k.storeKey)
+
+	undelegationIterator := sdk.KVStorePrefixIterator(store, types.UndelegateKey)
+	defer undelegationIterator.Close()
+
+	for ; undelegationIterator.Valid(); undelegationIterator.Next() {
+		undelegation := types.Undelegate{}
+		value := undelegationIterator.Value()
+		k.cdc.MustUnmarshal(value, &undelegation)
+
+		undelegations = append(undelegations, undelegation)
+
+		store.Delete(undelegationIterator.Key())
+	}
+
+	return undelegations
+}
+
+func (k Keeper) DequeueAllCancelUnbondingDelegation(ctx sdk.Context) (undelegations []types.CancelUnbondingDelegation) {
+	store := ctx.KVStore(k.storeKey)
+
+	cancelunbondingundelegationIterator := sdk.KVStorePrefixIterator(store, types.CancelUnbondingDelegationKey)
+	defer cancelunbondingundelegationIterator.Close()
+
+	for ; cancelunbondingundelegationIterator.Valid(); cancelunbondingundelegationIterator.Next() {
+		cancelunbondingdelegation := types.CancelUnbondingDelegation{}
+		value := cancelunbondingundelegationIterator.Value()
+		k.cdc.MustUnmarshal(value, &cancelunbondingdelegation)
+
+		undelegations = append(undelegations, cancelunbondingdelegation)
+
+		store.Delete(cancelunbondingundelegationIterator.Key())
+	}
+	return undelegations
+}
+
+func GetValidatorAddrFromStr(d string) sdk.ValAddress {
+	addr, err := sdk.ValAddressFromBech32(d)
+	if err != nil {
+		panic(err)
+	}
+	return addr
+}
+
 func GetValidatorAddr(d types.Delegation) sdk.ValAddress {
 	addr, err := sdk.ValAddressFromBech32(d.ValidatorAddress)
 	if err != nil {

x/stakingmiddleware/types/keys.go

@@ -7,7 +7,10 @@ import (
 
 // MinterKey is the key to use for the keeper store.
 var (
-	DelegationKey = []byte{0x01} // key for a delegation
+	DelegateKey                  = []byte{0x01} // key for a delegation
+	BeginRedelegateKey           = []byte{0x02} // key for a delegation
+	UndelegateKey                = []byte{0x03} // key for a delegation
+	CancelUnbondingDelegationKey = []byte{0x04} // key for a delegation
 )
 
 const (
@@ -19,12 +22,36 @@ const (
 	StoreKey = "customstmiddleware" // not using the module name because of collisions with key "ibc"
 )
 
-// GetDelegationKey creates the key for delegator bond with validator
+// GetDelegateKey creates the key for delegator bond with validator
 // VALUE: staking/Delegation
-func GetDelegationKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
-	return append(GetDelegationsKey(delAddr), address.MustLengthPrefix(valAddr)...)
+func GetDelegateKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
+	return append(DelegationKey(delAddr), address.MustLengthPrefix(valAddr)...)
 }
 
-func GetDelegationsKey(delAddr sdk.AccAddress) []byte {
-	return append(DelegationKey, address.MustLengthPrefix(delAddr)...)
+func DelegationKey(delAddr sdk.AccAddress) []byte {
+	return append(DelegateKey, address.MustLengthPrefix(delAddr)...)
+}
+
+func GetBeginRedelegateKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
+	return append(BeginRedelegationsKey(delAddr), address.MustLengthPrefix(valAddr)...)
+}
+
+func BeginRedelegationsKey(delAddr sdk.AccAddress) []byte {
+	return append(BeginRedelegateKey, address.MustLengthPrefix(delAddr)...)
+}
+
+func GetUndelegateKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
+	return append(UndelegateionKey(delAddr), address.MustLengthPrefix(valAddr)...)
+}
+
+func UndelegateionKey(delAddr sdk.AccAddress) []byte {
+	return append(UndelegateKey, address.MustLengthPrefix(delAddr)...)
+}
+
+func GetCancelUnbondingDelegateKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
+	return append(CancelUnbondingDelegateKey(delAddr), address.MustLengthPrefix(valAddr)...)
+}
+
+func CancelUnbondingDelegateKey(delAddr sdk.AccAddress) []byte {
+	return append(CancelUnbondingDelegationKey, address.MustLengthPrefix(delAddr)...)
 }