Skip to content

Commit

Permalink
Integration example
Browse files Browse the repository at this point in the history
  • Loading branch information
alpe committed Sep 28, 2023
1 parent fbea852 commit 96afc6d
Show file tree
Hide file tree
Showing 17 changed files with 366 additions and 98 deletions.
31 changes: 16 additions & 15 deletions demo/app/app.go
Original file line number Diff line number Diff line change
Expand Up @@ -343,7 +343,7 @@ func NewMeshApp(
)

tkeys := sdk.NewTransientStoreKeys(paramstypes.TStoreKey)
memKeys := sdk.NewMemoryStoreKeys(capabilitytypes.MemStoreKey)
memKeys := sdk.NewMemoryStoreKeys(capabilitytypes.MemStoreKey, meshsectypes.MemStoreKey)

// load state streaming if enabled
if _, _, err := streaming.LoadStreamingServices(bApp, appOpts, appCodec, logger, keys); err != nil {
Expand Down Expand Up @@ -430,11 +430,24 @@ func NewMeshApp(
authtypes.NewModuleAddress(govtypes.ModuleName).String(),
)

// setup mesh-security keeper with vanilla Cosmos-SDK
// see also NewKeeperX constructor for integration with Osmosis SDK fork
// should be initialized before wasm keeper for custom query/msg handlers
app.MeshSecKeeper = meshseckeeper.NewKeeper(
app.appCodec,
keys[meshsectypes.StoreKey],
memKeys[meshsectypes.MemStoreKey],
app.BankKeeper,
app.StakingKeeper,
&app.WasmKeeper, // ensure this is a pointer as we instantiate the keeper a bit later
authtypes.NewModuleAddress(govtypes.ModuleName).String(),
)

app.SlashingKeeper = slashingkeeper.NewKeeper(
appCodec,
legacyAmino,
keys[slashingtypes.StoreKey],
app.StakingKeeper,
meshseckeeper.NewStakingDecorator(app.StakingKeeper, app.MeshSecKeeper),
authtypes.NewModuleAddress(govtypes.ModuleName).String(),
)

Expand All @@ -450,18 +463,6 @@ func NewMeshApp(

app.FeeGrantKeeper = feegrantkeeper.NewKeeper(appCodec, keys[feegrant.StoreKey], app.AccountKeeper)

// setup mesh-security keeper with vanilla Cosmos-SDK
// see also NewKeeperX constructor for integration with Osmosis SDK fork
// should be initialized before wasm keeper for custom query/msg handlers
app.MeshSecKeeper = meshseckeeper.NewKeeper(
app.appCodec,
keys[meshsectypes.StoreKey],
app.BankKeeper,
app.StakingKeeper,
&app.WasmKeeper, // ensure this is a pointer as we instantiate the keeper a bit later
authtypes.NewModuleAddress(govtypes.ModuleName).String(),
)

// register the staking hooks
// NOTE: stakingKeeper above is passed by reference, so that it will contain these hooks
app.StakingKeeper.SetHooks(
Expand Down Expand Up @@ -751,7 +752,7 @@ func NewMeshApp(
icatypes.ModuleName,
ibcfeetypes.ModuleName,
wasmtypes.ModuleName,
meshsectypes.ModuleName, // after evidence
meshsectypes.ModuleName, // last to capture all chain events
)

// NOTE: The genutils module must occur after staking so that pools are
Expand Down
33 changes: 26 additions & 7 deletions x/meshsecurity/abci.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,34 @@ type TaskExecutionResponseHandler interface {
// EndBlocker is called after every block
func EndBlocker(ctx sdk.Context, k *keeper.Keeper, h TaskExecutionResponseHandler) {
defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyEndBlocker)

do := rspHandler(ctx, h)
epochLength := k.GetRebalanceEpochLength(ctx)
results, err := k.ExecScheduledTasks(ctx, types.SchedulerTaskRebalance, epochLength, func(ctx sdk.Context, contract sdk.AccAddress) error {
return k.Rebalance(ctx, contract)
})
if err != nil {
panic(fmt.Sprintf("task scheduler: %s", err)) // todo: log or fail?
var valsetUpdated bool
do(k.ExecScheduledTasks(ctx, types.SchedulerTaskValsetUpdate, epochLength, func(ctx sdk.Context, contract sdk.AccAddress) error {
valsetUpdated = true
report, err := k.ValsetUpdateReport(ctx)
if err != nil {
return err
}
return k.SendValsetUpdate(ctx, contract, report)
}))
if valsetUpdated {
k.ClearPipedValsetOperations(ctx)
}
for _, r := range results {
h.Handle(ctx, r)
do(k.ExecScheduledTasks(ctx, types.SchedulerTaskRebalance, epochLength, func(ctx sdk.Context, contract sdk.AccAddress) error {
return k.SendRebalance(ctx, contract)
}))
}

func rspHandler(ctx sdk.Context, h TaskExecutionResponseHandler) func(results []keeper.ExecResult, err error) {
return func(results []keeper.ExecResult, err error) {
if err != nil {
panic(fmt.Sprintf("task scheduler: %s", err)) // todo: log or fail?
}
for _, r := range results {
h.Handle(ctx, r)
}
}
}

Expand Down
28 changes: 25 additions & 3 deletions x/meshsecurity/contract/out_message.go
Original file line number Diff line number Diff line change
@@ -1,5 +1,27 @@
package contract

type SudoMsg struct {
Rebalance *struct{} `json:"rebalance"`
}
import (
wasmvmtypes "github.com/CosmWasm/wasmvm/types"
)

type (
SudoMsg struct {
Rebalance *struct{} `json:"rebalance,omitempty"`
ValsetUpdate *ValsetUpdate `json:"valset_update,omitempty"`
}

// Validator alias to wasmVM type
Validator = wasmvmtypes.Validator
// ValidatorAddr alias for the Bech32 address string of sdk.ValAddress
ValidatorAddr = string

// ValsetUpdate updates to the active validator set
ValsetUpdate struct {
Additions []Validator `json:"additions,omitempty"`
Removals []ValidatorAddr `json:"removals,omitempty"`
Updated []Validator `json:"updated,omitempty"`
Jailed []ValidatorAddr `json:"slashed,omitempty"`
Unjailed []ValidatorAddr `json:"unjailed,omitempty"`
Tombstoned []ValidatorAddr `json:"tombstoned,omitempty"`
}
)
1 change: 1 addition & 0 deletions x/meshsecurity/keeper/adapter.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ func (b BankKeeperAdapter) AddSupplyOffset(ctx sdk.Context, denom string, offset
var _ types.XStakingKeeper = &StakingKeeperAdapter{}

// StakingKeeperAdapter is an adapter to enhance the vanilla sdk staking keeper with additional functionality
// required for MS. The methods match Osmosis SDK fork.
type StakingKeeperAdapter struct {
types.SDKStakingKeeper
bank types.SDKBankKeeper
Expand Down
3 changes: 2 additions & 1 deletion x/meshsecurity/keeper/common_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ func CreateDefaultTestInput(t testing.TB) (sdk.Context, TestKeepers) {
for _, v := range keys {
ms.MountStoreWithDB(v, storetypes.StoreTypeIAVL, db)
}
memKeys := sdk.NewMemoryStoreKeys(capabilitytypes.MemStoreKey)
memKeys := sdk.NewMemoryStoreKeys(capabilitytypes.MemStoreKey, types.MemStoreKey)
for _, v := range memKeys {
ms.MountStoreWithDB(v, storetypes.StoreTypeMemory, db)
}
Expand Down Expand Up @@ -257,6 +257,7 @@ func CreateDefaultTestInput(t testing.TB) (sdk.Context, TestKeepers) {
msKeeper := NewKeeper(
appCodec,
keys[types.StoreKey],
memKeys[types.MemStoreKey],
bankKeeper,
stakingKeeper,
wasmKeeper,
Expand Down
12 changes: 6 additions & 6 deletions x/meshsecurity/keeper/evidence.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,22 +7,22 @@ import (
"github.com/osmosis-labs/mesh-security-sdk/x/meshsecurity/types"
)

var _ evidencetypes.SlashingKeeper = EvidenceDecorator{}
var _ evidencetypes.SlashingKeeper = SlashingKeeperDecorator{}

// EvidenceDecorator to capture TompStone events
type EvidenceDecorator struct {
// SlashingKeeperDecorator to capture TompStone events
type SlashingKeeperDecorator struct {
evidencetypes.SlashingKeeper
stakingKeeper types.SDKStakingKeeper
k *Keeper
}

// CaptureTombstoneDecorator constructor
func CaptureTombstoneDecorator(k *Keeper, slashingKeeper evidencetypes.SlashingKeeper, stakingKeeper types.SDKStakingKeeper) *EvidenceDecorator {
return &EvidenceDecorator{SlashingKeeper: slashingKeeper, stakingKeeper: stakingKeeper, k: k}
func CaptureTombstoneDecorator(k *Keeper, slashingKeeper evidencetypes.SlashingKeeper, stakingKeeper types.SDKStakingKeeper) *SlashingKeeperDecorator {
return &SlashingKeeperDecorator{SlashingKeeper: slashingKeeper, stakingKeeper: stakingKeeper, k: k}
}

// Tombstone is executed in endblocker by evidence module
func (e EvidenceDecorator) Tombstone(ctx sdk.Context, address sdk.ConsAddress) {
func (e SlashingKeeperDecorator) Tombstone(ctx sdk.Context, address sdk.ConsAddress) {
v, ok := e.stakingKeeper.GetValidatorByConsAddr(ctx, address)
if !ok {
ModuleLogger(ctx).
Expand Down
8 changes: 6 additions & 2 deletions x/meshsecurity/keeper/keeper.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (

type Keeper struct {
storeKey storetypes.StoreKey
memKey storetypes.StoreKey
cdc codec.Codec
bank types.XBankKeeper
staking types.XStakingKeeper
Expand All @@ -31,25 +32,28 @@ type Keeper struct {
func NewKeeper(
cdc codec.Codec,
storeKey storetypes.StoreKey,
memoryStoreKey storetypes.StoreKey,
bank types.SDKBankKeeper,
staking types.SDKStakingKeeper,
wasm types.WasmKeeper,
authority string,
) *Keeper {
return NewKeeperX(cdc, storeKey, NewBankKeeperAdapter(bank), NewStakingKeeperAdapter(staking, bank), wasm, authority)
return NewKeeperX(cdc, storeKey, memoryStoreKey, NewBankKeeperAdapter(bank), NewStakingKeeperAdapter(staking, bank), wasm, authority)
}

// NewKeeperX constructor with extended Osmosis SDK keepers
func NewKeeperX(
cdc codec.Codec,
storeKey storetypes.StoreKey,
memoryStoreKey storetypes.StoreKey,
bank types.XBankKeeper,
staking types.XStakingKeeper,
wasm types.WasmKeeper,
authority string,
) *Keeper {
return &Keeper{
storeKey: storeKey,
memKey: memoryStoreKey,
cdc: cdc,
bank: bank,
staking: staking,
Expand Down Expand Up @@ -145,7 +149,7 @@ func (k Keeper) mustLoadInt(ctx sdk.Context, storeKey storetypes.StoreKey, key [
return r
}

// IterateMaxCapLimit iterate over contract addresses with max cap limt set
// IterateMaxCapLimit iterate over contract addresses with max cap limit set
// Callback can return true to stop early
func (k Keeper) IterateMaxCapLimit(ctx sdk.Context, cb func(sdk.AccAddress, math.Int) bool) {
prefixStore := prefix.NewStore(ctx.KVStore(k.storeKey), types.MaxCapLimitKeyPrefix)
Expand Down
2 changes: 1 addition & 1 deletion x/meshsecurity/keeper/msg_server.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ func (m msgServer) SetVirtualStakingMaxCap(goCtx context.Context, req *types.Msg
}

// schedule last rebalance callback to let the contract do undelegates and housekeeping
if err := m.k.ScheduleTask(ctx, types.SchedulerTaskRebalance, acc, uint64(ctx.BlockHeight()), false, nil); err != nil {
if err := m.k.ScheduleOneShotTask(ctx, types.SchedulerTaskRebalance, acc, uint64(ctx.BlockHeight())); err != nil {
return nil, errorsmod.Wrap(err, "schedule one shot rebalance task")
}
return &types.MsgSetVirtualStakingMaxCapResponse{}, nil
Expand Down
2 changes: 1 addition & 1 deletion x/meshsecurity/keeper/params.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ func (k Keeper) GetParams(clientCtx sdk.Context) (params types.Params) {
return params
}

func (k Keeper) GetRebalanceGasLimit(ctx sdk.Context) sdk.Gas {
func (k Keeper) GetMaxSudoGas(ctx sdk.Context) sdk.Gas {
return sdk.Gas(k.GetParams(ctx).MaxGasEndBlocker)
}

Expand Down
36 changes: 25 additions & 11 deletions x/meshsecurity/keeper/scheduler.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,6 @@ import (
"bytes"
"math"

"github.com/golang/protobuf/proto"

errorsmod "cosmossdk.io/errors"

"github.com/cosmos/cosmos-sdk/store/prefix"
Expand All @@ -22,7 +20,7 @@ func (k Keeper) ScheduleRegularRebalanceTask(ctx sdk.Context, contract sdk.AccAd
}
epochLength := k.GetRebalanceEpochLength(ctx)
nextExecBlock := uint64(ctx.BlockHeight()) + epochLength
return k.ScheduleTask(ctx, types.SchedulerTaskRebalance, contract, nextExecBlock, true, nil)
return k.ScheduleRepeatingTask(ctx, types.SchedulerTaskRebalance, contract, nextExecBlock)
}

// HasScheduledTask returns true if the contract has a task scheduled of the given type and repeat setting
Expand Down Expand Up @@ -55,8 +53,9 @@ func (k Keeper) getScheduledTaskAt(ctx sdk.Context, tp types.SchedulerTaskType,
return isRepeat(bz), bz != nil
}

// ScheduleTask register a new task to be executed at given block height
func (k Keeper) ScheduleTask(ctx sdk.Context, tp types.SchedulerTaskType, contract sdk.AccAddress, execBlockHeight uint64, repeat bool, payload proto.Message) error {
// ScheduleOneShotTask register a new task to be executed at given block height.
// The task is not repeating and registered only once for a given contract and height. Duplicates are silently ignored
func (k Keeper) ScheduleOneShotTask(ctx sdk.Context, tp types.SchedulerTaskType, contract sdk.AccAddress, execBlockHeight uint64) error {
if execBlockHeight < uint64(ctx.BlockHeight()) { // sanity check
return types.ErrInvalid.Wrapf("can not schedule for past block: %d", execBlockHeight)
}
Expand All @@ -65,11 +64,26 @@ func (k Keeper) ScheduleTask(ctx sdk.Context, tp types.SchedulerTaskType, contra
return err
}
store := ctx.KVStore(k.storeKey)
if !repeat { // ensure that we do not overwrite a repeating scheduled event
if bz := store.Get(storeKey); bz != nil {
repeat = isRepeat(bz)
}
if store.Has(storeKey) {
return nil
}
store.Set(storeKey, []byte{toByte(false)})
types.EmitSchedulerRegisteredEvent(ctx, contract, execBlockHeight, false)
return nil
}

// ScheduleRepeatingTask register a recurring task to be executed at given block height.
// Duplicates are overwritten
func (k Keeper) ScheduleRepeatingTask(ctx sdk.Context, tp types.SchedulerTaskType, contract sdk.AccAddress, execBlockHeight uint64) error {
if execBlockHeight < uint64(ctx.BlockHeight()) { // sanity check
return types.ErrInvalid.Wrapf("can not schedule for past block: %d", execBlockHeight)
}
storeKey, err := types.BuildSchedulerContractKey(tp, execBlockHeight, contract)
if err != nil {
return err
}
store := ctx.KVStore(k.storeKey)
const repeat = true
store.Set(storeKey, []byte{toByte(repeat)})
types.EmitSchedulerRegisteredEvent(ctx, contract, execBlockHeight, repeat)
return nil
Expand Down Expand Up @@ -158,7 +172,7 @@ func (k Keeper) ExecScheduledTasks(pCtx sdk.Context, tp types.SchedulerTaskType,
currentHeight := uint64(pCtx.BlockHeight())
// iterator is most gas cost-efficient currently
err := k.IterateScheduledTasks(pCtx, tp, currentHeight, func(contract sdk.AccAddress, scheduledHeight uint64, repeat bool) bool {
gasLimit := k.GetRebalanceGasLimit(pCtx)
gasLimit := k.GetMaxSudoGas(pCtx)
cachedCtx, done := pCtx.CacheContext()
gasMeter := sdk.NewGasMeter(gasLimit)
cachedCtx = cachedCtx.WithGasMeter(gasMeter)
Expand All @@ -179,7 +193,7 @@ func (k Keeper) ExecScheduledTasks(pCtx sdk.Context, tp types.SchedulerTaskType,
// re-schedule
nextExecBlock := uint64(pCtx.BlockHeight()) + epochLength
result.NextRunHeight = nextExecBlock
if err := k.ScheduleTask(pCtx, tp, contract, nextExecBlock, repeat, nil); err != nil {
if err := k.ScheduleRepeatingTask(pCtx, tp, contract, nextExecBlock); err != nil {
result.RescheduleErr = err
}
}
Expand Down
Loading

0 comments on commit 96afc6d

Please sign in to comment.