epoching: keeper functions, queries, and testing infra (#19)

testutil/keeper/epoching.go

@@ -0,0 +1,52 @@
+package keeper
+
+import (
+	"testing"
+
+	"github.com/babylonchain/babylon/x/epoching/keeper"
+	"github.com/babylonchain/babylon/x/epoching/types"
+	"github.com/cosmos/cosmos-sdk/codec"
+	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+	"github.com/cosmos/cosmos-sdk/store"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	typesparams "github.com/cosmos/cosmos-sdk/x/params/types"
+	"github.com/stretchr/testify/require"
+	"github.com/tendermint/tendermint/libs/log"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+	tmdb "github.com/tendermint/tm-db"
+)
+
+func EpochingKeeper(t testing.TB) (*keeper.Keeper, sdk.Context) {
+	storeKey := sdk.NewKVStoreKey(types.StoreKey)
+	memStoreKey := storetypes.NewMemoryStoreKey(types.MemStoreKey)
+
+	db := tmdb.NewMemDB()
+	stateStore := store.NewCommitMultiStore(db)
+	stateStore.MountStoreWithDB(storeKey, sdk.StoreTypeIAVL, db)
+	stateStore.MountStoreWithDB(memStoreKey, sdk.StoreTypeMemory, nil)
+	require.NoError(t, stateStore.LoadLatestVersion())
+
+	registry := codectypes.NewInterfaceRegistry()
+	cdc := codec.NewProtoCodec(registry)
+
+	paramsSubspace := typesparams.NewSubspace(cdc,
+		types.Amino,
+		storeKey,
+		memStoreKey,
+		"EpochingParams",
+	)
+	k := keeper.NewKeeper(
+		cdc,
+		storeKey,
+		memStoreKey,
+		paramsSubspace,
+	)
+
+	ctx := sdk.NewContext(stateStore, tmproto.Header{}, false, log.NewNopLogger())
+
+	// Initialize params
+	k.SetParams(ctx, types.DefaultParams())
+
+	return &k, ctx
+}

x/epoching/keeper/grpc_query.go

@@ -4,7 +4,9 @@ import (
 	"context"
 
 	"github.com/babylonchain/babylon/x/epoching/types"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/query"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
@@ -25,10 +27,52 @@ func (k Keeper) Params(c context.Context, req *types.QueryParamsRequest) (*types
 	return &types.QueryParamsResponse{Params: k.GetParams(ctx)}, nil
 }
 
+// CurrentEpoch handles the QueryCurrentEpochRequest query
 func (k Keeper) CurrentEpoch(c context.Context, req *types.QueryCurrentEpochRequest) (*types.QueryCurrentEpochResponse, error) {
-	panic("TODO: unimplemented")
+	ctx := sdk.UnwrapSDKContext(c)
+	epochNumber, err := k.GetEpochNumber(ctx)
+	if err != nil {
+		return nil, err
+	}
+	epochBoundary, err := k.GetEpochBoundary(ctx)
+	if err != nil {
+		return nil, err
+	}
+	resp := &types.QueryCurrentEpochResponse{
+		CurrentEpoch:  epochNumber.BigInt().Uint64(),
+		EpochBoundary: epochBoundary.BigInt().Uint64(),
+	}
+	return resp, nil
 }
 
+// EpochMsgs handles the QueryEpochMsgsRequest query
 func (k Keeper) EpochMsgs(c context.Context, req *types.QueryEpochMsgsRequest) (*types.QueryEpochMsgsResponse, error) {
-	panic("TODO: unimplemented")
+	ctx := sdk.UnwrapSDKContext(c)
+	var msgs []*types.QueuedMessage
+	store := ctx.KVStore(k.storeKey)
+	epochMsgsStore := prefix.NewStore(store, types.QueuedMsgKey)
+
+	// handle pagination
+	// TODO: the epoch might end between pagination requests, leading inconsistent results by the time someone gets to the end. Possible fixes:
+	// - We could add the epoch number to the query, and return nothing if the current epoch number is different. But it's a bit of pain to have to set it and not know why there's no result.
+	// - We could not reset the key to 0 when the queue is cleared, and just keep incrementing the ID forever. That way when the next query comes, it might skip some records that have been deleted, then resume from the next available record which has a higher key than the value in the pagination data structure.
+	// - We can do nothing, in which case some records that have been inserted after the delete might be skipped because their keys are lower than the pagionation state.
+	pageRes, err := query.Paginate(epochMsgsStore, req.Pagination, func(key, value []byte) error {
+		// unmarshal to queuedMsg
+		var queuedMsg types.QueuedMessage
+		if err := k.cdc.Unmarshal(value, &queuedMsg); err != nil {
+			return err
+		}
+		// append to msgs
+		msgs = append(msgs, &queuedMsg)
+		return nil
+	})
+	if err != nil {
+		return nil, status.Error(codes.Internal, err.Error())
+	}
+	resp := &types.QueryEpochMsgsResponse{
+		Msgs:       msgs,
+		Pagination: pageRes,
+	}
+	return resp, nil
 }

x/epoching/keeper/grpc_query_test.go

@@ -1 +1,22 @@
 package keeper_test
+
+import (
+	"testing"
+
+	testkeeper "github.com/babylonchain/babylon/testutil/keeper"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/babylonchain/babylon/x/epoching/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParamsQuery(t *testing.T) {
+	keeper, ctx := testkeeper.EpochingKeeper(t)
+	wctx := sdk.WrapSDKContext(ctx)
+	params := types.DefaultParams()
+	keeper.SetParams(ctx, params)
+
+	response, err := keeper.Params(wctx, &types.QueryParamsRequest{})
+	require.NoError(t, err)
+	require.Equal(t, &types.QueryParamsResponse{Params: params}, response)
+}

x/epoching/keeper/keeper.go

@@ -3,12 +3,15 @@ package keeper
 import (
 	"fmt"
 
-	"github.com/tendermint/tendermint/libs/log"
-
 	"github.com/babylonchain/babylon/x/epoching/types"
 	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
+	"github.com/tendermint/tendermint/libs/log"
+)
+
+const (
+	DefaultEpochNumber = 0
 )
 
 type (
@@ -46,27 +49,160 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 }
 
 // Set the validator hooks
-func (k *Keeper) SetHooks(sh types.EpochingHooks) *Keeper {
+func (k *Keeper) SetHooks(eh types.EpochingHooks) *Keeper {
 	if k.hooks != nil {
 		panic("cannot set validator hooks twice")
 	}
 
-	k.hooks = sh
+	k.hooks = eh
 
 	return k
 }
 
-// GetCurrentEpoch returns the current epoch number
-func (k Keeper) GetCurrentEpoch(ctx sdk.Context) sdk.Uint {
-	panic("TODO: unimplemented")
+// GetEpochNumber fetches epoch number
+func (k Keeper) GetEpochNumber(ctx sdk.Context) (sdk.Uint, error) {
+	store := ctx.KVStore(k.storeKey)
+
+	bz := store.Get(types.EpochNumberKey)
+	if bz == nil {
+		return sdk.NewUint(uint64(DefaultEpochNumber)), nil
+	}
+	var epochNumber sdk.Uint
+	err := epochNumber.Unmarshal(bz)
+
+	return epochNumber, err
+}
+
+// setEpochNumber sets epoch number
+func (k Keeper) setEpochNumber(ctx sdk.Context, epochNumber sdk.Uint) error {
+	store := ctx.KVStore(k.storeKey)
+
+	epochNumberBytes, err := epochNumber.Marshal()
+	if err != nil {
+		return err
+	}
+
+	store.Set(types.EpochNumberKey, epochNumberBytes)
+
+	return nil
+}
+
+// IncEpochNumber adds epoch number by 1
+func (k Keeper) IncEpochNumber(ctx sdk.Context) error {
+	epochNumber, err := k.GetEpochNumber(ctx)
+	if err != nil {
+		return err
+	}
+	incrementedEpochNumber := epochNumber.AddUint64(1)
+	return k.setEpochNumber(ctx, incrementedEpochNumber)
+}
+
+// GetEpochBoundary gets the epoch boundary, i.e., the height of the block that ends this epoch
+func (k Keeper) GetEpochBoundary(ctx sdk.Context) (sdk.Uint, error) {
+	epochNumber, err := k.GetEpochNumber(ctx)
+	if err != nil {
+		return sdk.NewUint(0), err
+	}
+	epochInterval := sdk.NewUint(k.GetParams(ctx).EpochInterval)
+	// example: in epoch 0, epoch interval is 5 blocks, boundary will be (0+1)*5=5
+	return epochNumber.AddUint64(1).Mul(epochInterval), nil
 }
 
-// GetEpochMsgs returns the set of messages queued of the current epoch
-func (k Keeper) GetEpochMsgs(ctx sdk.Context) []types.QueuedMessage {
-	panic("TODO: unimplemented")
+// GetQueueLength fetches the number of queued messages
+func (k Keeper) GetQueueLength(ctx sdk.Context) (sdk.Uint, error) {
+	store := ctx.KVStore(k.storeKey)
+
+	bz := store.Get(types.QueueLengthKey)
+	if bz == nil {
+		return sdk.NewUint(0), nil
+	}
+	var queueLen sdk.Uint
+	err := queueLen.Unmarshal(bz)
+
+	return queueLen, err
+}
+
+// setQueueLength sets the msg queue length
+func (k Keeper) setQueueLength(ctx sdk.Context, queueLen sdk.Uint) error {
+	store := ctx.KVStore(k.storeKey)
+
+	queueLenBytes, err := queueLen.Marshal()
+	if err != nil {
+		return err
+	}
+
+	store.Set(types.QueueLengthKey, queueLenBytes)
+
+	return nil
+}
+
+// incQueueLength adds the queue length by 1
+func (k Keeper) incQueueLength(ctx sdk.Context) error {
+	queueLen, err := k.GetQueueLength(ctx)
+	if err != nil {
+		return err
+	}
+	incrementedQueueLen := queueLen.AddUint64(1)
+	return k.setQueueLength(ctx, incrementedQueueLen)
 }
 
 // EnqueueMsg enqueues a message to the queue of the current epoch
 func (k Keeper) EnqueueMsg(ctx sdk.Context, msg types.QueuedMessage) error {
-	panic("TODO: unimplemented")
+	store := ctx.KVStore(k.storeKey)
+
+	// insert KV pair, where
+	// - key: QueuedMsgKey || queueLenBytes
+	// - value: msgBytes
+	queueLen, err := k.GetQueueLength(ctx)
+	if err != nil {
+		return err
+	}
+	queueLenBytes, err := queueLen.Marshal()
+	if err != nil {
+		return err
+	}
+	msgBytes, err := k.cdc.Marshal(&msg)
+	if err != nil {
+		return err
+	}
+	store.Set(append(types.QueuedMsgKey, queueLenBytes...), msgBytes)
+
+	// increment queue length
+	return k.incQueueLength(ctx)
+}
+
+// GetEpochMsgs returns the set of messages queued in the current epoch
+func (k Keeper) GetEpochMsgs(ctx sdk.Context) ([]*types.QueuedMessage, error) {
+	queuedMsgs := []*types.QueuedMessage{}
+	store := ctx.KVStore(k.storeKey)
+
+	// add each queued msg to queuedMsgs
+	iterator := sdk.KVStorePrefixIterator(store, types.QueuedMsgKey)
+	defer iterator.Close()
+	for ; iterator.Valid(); iterator.Next() {
+		queuedMsgBytes := iterator.Value()
+		var queuedMsg types.QueuedMessage
+		if err := k.cdc.Unmarshal(queuedMsgBytes, &queuedMsg); err != nil {
+			return nil, err
+		}
+		queuedMsgs = append(queuedMsgs, &queuedMsg)
+	}
+
+	return queuedMsgs, nil
+}
+
+// ClearEpochMsgs removes all messages in the queue
+func (k Keeper) ClearEpochMsgs(ctx sdk.Context) error {
+	store := ctx.KVStore(k.storeKey)
+
+	// remove all epoch msgs
+	iterator := sdk.KVStorePrefixIterator(ctx.KVStore(k.storeKey), types.QueuedMsgKey)
+	defer iterator.Close()
+	for ; iterator.Valid(); iterator.Next() {
+		key := iterator.Key()
+		store.Delete(key)
+	}
+
+	// set queue len to zero
+	return k.setQueueLength(ctx, sdk.NewUint(0))
 }

x/epoching/keeper/keeper_test.go

@@ -0,0 +1 @@
+package keeper_test

x/epoching/keeper/msg_server_test.go

@@ -1 +1,16 @@
 package keeper_test
+
+import (
+	"context"
+	"testing"
+
+	keepertest "github.com/babylonchain/babylon/testutil/keeper"
+	"github.com/babylonchain/babylon/x/epoching/keeper"
+	"github.com/babylonchain/babylon/x/epoching/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+func setupMsgServer(t testing.TB) (types.MsgServer, context.Context) {
+	k, ctx := keepertest.EpochingKeeper(t)
+	return keeper.NewMsgServerImpl(*k), sdk.WrapSDKContext(ctx)
+}

x/epoching/keeper/params_test.go

@@ -1 +1,18 @@
 package keeper_test
+
+import (
+	"testing"
+
+	testkeeper "github.com/babylonchain/babylon/testutil/keeper"
+	"github.com/babylonchain/babylon/x/epoching/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetParams(t *testing.T) {
+	k, ctx := testkeeper.EpochingKeeper(t)
+	params := types.DefaultParams()
+
+	k.SetParams(ctx, params)
+
+	require.EqualValues(t, params, k.GetParams(ctx))
+}

x/epoching/module.go

@@ -168,10 +168,13 @@ func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) json.Raw
 func (AppModule) ConsensusVersion() uint64 { return 2 }
 
 // BeginBlock executes all ABCI BeginBlock logic respective to the capability module.
-func (am AppModule) BeginBlock(_ sdk.Context, _ abci.RequestBeginBlock) {}
+func (am AppModule) BeginBlock(_ sdk.Context, _ abci.RequestBeginBlock) {
+	// TODO: trigger BeginBlock stuff upon the first block of an epoch
+}
 
 // EndBlock executes all ABCI EndBlock logic respective to the capability module. It
 // returns no validator updates.
 func (am AppModule) EndBlock(_ sdk.Context, _ abci.RequestEndBlock) []abci.ValidatorUpdate {
+	// TODO: trigger EndBlock stuff upon the last block of an epoch
 	return []abci.ValidatorUpdate{}
 }