Update schedule logic for zero max cap

osmosis-labs · Sep 5, 2024 · 082fbe5 · 082fbe5
1 parent 616de21
commit 082fbe5
Show file tree

Hide file tree

Showing 15 changed files with 34 additions and 62 deletions.
diff --git a/tests/testdata/mesh_converter.wasm.gz b/tests/testdata/mesh_converter.wasm.gz
diff --git a/tests/testdata/mesh_external_staking.wasm.gz b/tests/testdata/mesh_external_staking.wasm.gz
diff --git a/tests/testdata/mesh_native_staking.wasm.gz b/tests/testdata/mesh_native_staking.wasm.gz
diff --git a/tests/testdata/mesh_native_staking_proxy.wasm.gz b/tests/testdata/mesh_native_staking_proxy.wasm.gz
diff --git a/tests/testdata/mesh_osmosis_price_provider.wasm.gz b/tests/testdata/mesh_osmosis_price_provider.wasm.gz
diff --git a/tests/testdata/mesh_remote_price_feed.wasm.gz b/tests/testdata/mesh_remote_price_feed.wasm.gz
diff --git a/tests/testdata/mesh_simple_price_feed.wasm.gz b/tests/testdata/mesh_simple_price_feed.wasm.gz
diff --git a/tests/testdata/mesh_vault.wasm.gz b/tests/testdata/mesh_vault.wasm.gz
diff --git a/tests/testdata/mesh_virtual_staking.wasm.gz b/tests/testdata/mesh_virtual_staking.wasm.gz
diff --git a/tests/testdata/version.txt b/tests/testdata/version.txt
@@ -1 +1 @@
-a0bba9b5f7eed6768d7a84ac766afb9d89dea4c2
+ef0e3840b092ed66cd0968e2b9b253ae243ffe52
diff --git a/x/meshsecurity/contract/in_message.go b/x/meshsecurity/contract/in_message.go
@@ -7,9 +7,10 @@ type (
 		VirtualStake *VirtualStakeMsg `json:"virtual_stake,omitempty"`
 	}
 	VirtualStakeMsg struct {
-		Bond             *BondMsg             `json:"bond,omitempty"`
-		Unbond           *UnbondMsg           `json:"unbond,omitempty"`
-		UpdateDelegation *UpdateDelegationMsg `json:"update_delegation,omitempty"`
+		Bond                    *BondMsg                    `json:"bond,omitempty"`
+		Unbond                  *UnbondMsg                  `json:"unbond,omitempty"`
+		UpdateDelegation        *UpdateDelegationMsg        `json:"update_delegation,omitempty"`
+		DeleteAllScheduledTasks *DeleteAllScheduledTasksMsg `json:"delete_all_scheduled_tasks,omitempty"`
 	}
 	BondMsg struct {
 		Amount    wasmvmtypes.Coin `json:"amount"`
@@ -25,4 +26,6 @@ type (
 		Delegator string           `json:"delegator"`
 		Validator string           `json:"validator"`
 	}
+
+	DeleteAllScheduledTasksMsg struct{}
 )
diff --git a/x/meshsecurity/keeper/handler_plugin.go b/x/meshsecurity/keeper/handler_plugin.go
@@ -26,6 +26,7 @@ type msKeeper interface {
 	Delegate(ctx sdk.Context, actor sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin) (sdk.Dec, error)
 	Undelegate(ctx sdk.Context, actor sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin) error
 	UpdateDelegation(ctx sdk.Context, actor, delAddr sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin, isDeduct bool)
+	DeleteAllScheduledTasks(ctx sdk.Context, tp types.SchedulerTaskType, contract sdk.AccAddress) error
 }
 
 type CustomMsgHandler struct {
@@ -78,6 +79,8 @@ func (h CustomMsgHandler) DispatchMsg(ctx sdk.Context, contractAddr sdk.AccAddre
 		return h.handleUnbondMsg(ctx, contractAddr, customMsg.VirtualStake.Unbond)
 	case customMsg.VirtualStake.UpdateDelegation != nil:
 		return h.handleUpdateDelegationMsg(ctx, contractAddr, customMsg.VirtualStake.UpdateDelegation)
+	case customMsg.VirtualStake.DeleteAllScheduledTasks != nil:
+		return nil, nil, h.k.DeleteAllScheduledTasks(ctx, types.SchedulerTaskHandleEpoch, contractAddr)
 	}
 	return nil, nil, wasmtypes.ErrUnknownMsg
 }

diff --git a/x/meshsecurity/keeper/handler_plugin_test.go b/x/meshsecurity/keeper/handler_plugin_test.go
@@ -36,6 +36,7 @@ func TestCustomMeshSecDispatchMsg(t *testing.T) {
 	validUnbondMsg := []byte(fmt.Sprintf(
 		`{"virtual_stake":{"unbond":{"amount":{"denom":"ALX", "amount":"1234"},"delegator":%q,"validator":%q}}}`,
 		myDelegatorAddr.String(), myValidatorAddr.String()))
+	validDeleteScheduledTasks := []byte(`{"virtual_stake":{"delete_all_scheduled_tasks":{}}}`)
 
 	specs := map[string]struct {
 		src       wasmvmtypes.CosmosMsg
@@ -95,6 +96,17 @@ func TestCustomMeshSecDispatchMsg(t *testing.T) {
 			},
 			expErr: myErr,
 		},
+		"handle delete tasks": {
+			src:  wasmvmtypes.CosmosMsg{Custom: validDeleteScheduledTasks},
+			auth: allAuthZ,
+			setup: func(t *testing.T) (msKeeper, func()) {
+				m := msKeeperMock{DeleteAllScheduledTasksFn: func(_ sdk.Context, tp types.SchedulerTaskType, contract sdk.AccAddress) error {
+					return myErr
+				}}
+				return &m, t.FailNow
+			},
+			expErr: myErr,
+		},
 		"non custom msg- skip": {
 			src:  wasmvmtypes.CosmosMsg{},
 			auth: panicAuthZ,
@@ -182,9 +194,10 @@ func captureCall(t *testing.T, myContractAddr sdk.AccAddress, myValidatorAddr sd
 var _ msKeeper = msKeeperMock{}
 
 type msKeeperMock struct {
-	DelegateFn         func(ctx sdk.Context, actor sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin) (sdk.Dec, error)
-	UndelegateFn       func(ctx sdk.Context, actor sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin) error
-	UpdateDelegationFn func(ctx sdk.Context, actor, delAddr sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin, isDeduct bool)
+	DelegateFn                func(ctx sdk.Context, actor sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin) (sdk.Dec, error)
+	UndelegateFn              func(ctx sdk.Context, actor sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin) error
+	UpdateDelegationFn        func(ctx sdk.Context, actor, delAddr sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin, isDeduct bool)
+	DeleteAllScheduledTasksFn func(ctx sdk.Context, tp types.SchedulerTaskType, contract sdk.AccAddress) error
 }
 
 func (m msKeeperMock) Delegate(ctx sdk.Context, actor sdk.AccAddress, valAddr sdk.ValAddress, coin sdk.Coin) (sdk.Dec, error) {
@@ -207,6 +220,14 @@ func (m msKeeperMock) UpdateDelegation(ctx sdk.Context, actor, delAddr sdk.AccAd
 	}
 	m.UpdateDelegationFn(ctx, actor, delAddr, valAddr, coin, isDeduct)
 }
+
+func (m msKeeperMock) DeleteAllScheduledTasks(ctx sdk.Context, tp types.SchedulerTaskType, contract sdk.AccAddress) error {
+	if m.DeleteAllScheduledTasksFn == nil {
+		panic("not expected to be called")
+	}
+	return m.DeleteAllScheduledTasksFn(ctx, tp, contract)
+}
+
 func TestIntegrityHandler(t *testing.T) {
 	myContractAddr := sdk.AccAddress(rand.Bytes(32))
 	specs := map[string]struct {

diff --git a/x/meshsecurity/keeper/msg_server.go b/x/meshsecurity/keeper/msg_server.go
@@ -47,16 +47,5 @@ func (m msgServer) SetVirtualStakingMaxCap(goCtx context.Context, req *types.Msg
 		}
 		return &types.MsgSetVirtualStakingMaxCapResponse{}, nil
 	}
-	if req.MaxCap.IsZero() {
-		// no need to run regular rebalances with a new limit of 0
-		if err := m.k.DeleteAllScheduledTasks(ctx, types.SchedulerTaskHandleEpoch, acc); err != nil {
-			return nil, err
-		}
-	}
-
-	// schedule last rebalance callback to let the contract do undelegates and housekeeping
-	if err := m.k.ScheduleOneShotTask(ctx, types.SchedulerTaskHandleEpoch, acc, uint64(ctx.BlockHeight())); err != nil {
-		return nil, errorsmod.Wrap(err, "schedule one shot rebalance task")
-	}
 	return &types.MsgSetVirtualStakingMaxCapResponse{}, nil
 }
diff --git a/x/meshsecurity/keeper/msg_server_test.go b/x/meshsecurity/keeper/msg_server_test.go
@@ -43,50 +43,6 @@ func TestSetVirtualStakingMaxCap(t *testing.T) {
 				assert.True(t, k.HasScheduledTask(ctx, types.SchedulerTaskHandleEpoch, myContract, true))
 			},
 		},
-		"existing limit updated": {
-			setup: func(ctx sdk.Context) {
-				_, err := m.SetVirtualStakingMaxCap(sdk.WrapSDKContext(ctx), &types.MsgSetVirtualStakingMaxCap{
-					Authority: k.GetAuthority(),
-					Contract:  myContract.String(),
-					MaxCap:    sdk.NewInt64Coin(denom, 456),
-				})
-				require.NoError(t, err)
-			},
-			src: types.MsgSetVirtualStakingMaxCap{
-				Authority: k.GetAuthority(),
-				Contract:  myContract.String(),
-				MaxCap:    myAmount,
-			},
-			expLimit: myAmount,
-			expSchedule: func(t *testing.T, ctx sdk.Context) {
-				repeat, exists := k.getScheduledTaskAt(ctx, types.SchedulerTaskHandleEpoch, myContract, uint64(ctx.BlockHeight()))
-				require.True(t, exists)
-				assert.False(t, repeat)
-				assert.True(t, k.HasScheduledTask(ctx, types.SchedulerTaskHandleEpoch, myContract, true))
-			},
-		},
-		"existing limit set to empty value": {
-			setup: func(ctx sdk.Context) {
-				_, err := m.SetVirtualStakingMaxCap(sdk.WrapSDKContext(ctx), &types.MsgSetVirtualStakingMaxCap{
-					Authority: k.GetAuthority(),
-					Contract:  myContract.String(),
-					MaxCap:    myAmount,
-				})
-				require.NoError(t, err)
-			},
-			src: types.MsgSetVirtualStakingMaxCap{
-				Authority: k.GetAuthority(),
-				Contract:  myContract.String(),
-				MaxCap:    sdk.NewInt64Coin(denom, 0),
-			},
-			expLimit: sdk.NewInt64Coin(denom, 0),
-			expSchedule: func(t *testing.T, ctx sdk.Context) {
-				repeat, exists := k.getScheduledTaskAt(ctx, types.SchedulerTaskHandleEpoch, myContract, uint64(ctx.BlockHeight()))
-				require.True(t, exists)
-				assert.False(t, repeat)
-				assert.False(t, k.HasScheduledTask(ctx, types.SchedulerTaskHandleEpoch, myContract, true))
-			},
-		},
 		"fails for non existing contract": {
 			setup: func(ctx sdk.Context) {},
 			src: types.MsgSetVirtualStakingMaxCap{