Integration example

demo/app/app.go

@@ -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 {
@@ -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(),
 	)
 
@@ -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(
@@ -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

x/meshsecurity/abci.go

@@ -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)
+		}
 	}
 }
 

x/meshsecurity/contract/out_message.go

@@ -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"`
+	}
+)

x/meshsecurity/keeper/adapter.go

@@ -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

x/meshsecurity/keeper/common_test.go

@@ -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)
 	}
@@ -257,6 +257,7 @@ func CreateDefaultTestInput(t testing.TB) (sdk.Context, TestKeepers) {
 	msKeeper := NewKeeper(
 		appCodec,
 		keys[types.StoreKey],
+		memKeys[types.MemStoreKey],
 		bankKeeper,
 		stakingKeeper,
 		wasmKeeper,

x/meshsecurity/keeper/evidence.go

@@ -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).

x/meshsecurity/keeper/keeper.go

@@ -18,6 +18,7 @@ import (
 
 type Keeper struct {
 	storeKey storetypes.StoreKey
+	memKey   storetypes.StoreKey
 	cdc      codec.Codec
 	bank     types.XBankKeeper
 	staking  types.XStakingKeeper
@@ -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,
@@ -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)

x/meshsecurity/keeper/msg_server.go

@@ -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

x/meshsecurity/keeper/params.go

@@ -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)
 }
 

x/meshsecurity/keeper/scheduler.go

@@ -4,8 +4,6 @@ import (
 	"bytes"
 	"math"
 
-	"github.com/golang/protobuf/proto"
-
 	errorsmod "cosmossdk.io/errors"
 
 	"github.com/cosmos/cosmos-sdk/store/prefix"
@@ -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
@@ -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)
 	}
@@ -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
@@ -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)
@@ -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
 			}
 		}