feat: added execution stage for schedules #NTRN-339

proto/neutron/cron/query.proto

@@ -30,25 +30,31 @@ service Query {
   // this line is used by starport scaffolding # 2
 }
 
+// QueryParamsRequest is the request type for the Query/Params RPC method.
 message QueryParamsRequest {}
 
+// QueryParamsResponse is the response type for the Query/Params RPC method.
 message QueryParamsResponse {
   // params holds all the parameters of this module.
   Params params = 1 [(gogoproto.nullable) = false];
 }
 
+// QueryGetScheduleRequest is the request type for the Query/Schedule RPC method.
 message QueryGetScheduleRequest {
   string name = 1;
 }
 
+// QueryGetScheduleResponse is the response type for the Query/Params RPC method.
 message QueryGetScheduleResponse {
   Schedule schedule = 1 [(gogoproto.nullable) = false];
 }
 
+// QuerySchedulesRequest is the request type for the Query/Schedules RPC method.
 message QuerySchedulesRequest {
   cosmos.base.query.v1beta1.PageRequest pagination = 1;
 }
 
+// QuerySchedulesResponse is the response type for the Query/Params RPC method.
 message QuerySchedulesResponse {
   repeated Schedule schedules = 1 [(gogoproto.nullable) = false];
   cosmos.base.query.v1beta1.PageResponse pagination = 2;

proto/neutron/cron/schedule.proto

@@ -5,24 +5,37 @@ import "gogoproto/gogo.proto";
 
 option go_package = "github.com/neutron-org/neutron/v4/x/cron/types";
 
+// ExecutionStage defines when messages will be executed in the block
+enum ExecutionStage {
+  // Execution at the end of the block
+  EXECUTION_STAGE_END_BLOCKER = 0;
+  // Execution at the beginning of the block
+  EXECUTION_STAGE_BEGIN_BLOCKER = 1;
+}
+
+// Schedule defines the schedule for execution
 message Schedule {
   // Name of schedule
   string name = 1;
   // Period in blocks
   uint64 period = 2;
-  // Msgs that will be executed every period amount of time
+  // Msgs that will be executed every certain number of blocks, specified in the `period` field
   repeated MsgExecuteContract msgs = 3 [(gogoproto.nullable) = false];
   // Last execution's block height
   uint64 last_execute_height = 4;
+  // Execution stages when messages will be executed
+  repeated ExecutionStage execution_stages = 5 [(gogoproto.nullable) = false];
 }
 
+// MsgExecuteContract defines the contract and the message to pass
 message MsgExecuteContract {
   // Contract is the address of the smart contract
   string contract = 1;
   // Msg is json encoded message to be passed to the contract
   string msg = 2;
 }
 
+// ScheduleCount defines the number of current schedules
 message ScheduleCount {
   // Count is the number of current schedules
   int32 count = 1;

proto/neutron/cron/tx.proto

@@ -6,6 +6,7 @@ import "cosmos/msg/v1/msg.proto";
 import "cosmos_proto/cosmos.proto";
 import "gogoproto/gogo.proto";
 import "neutron/cron/params.proto";
+import "neutron/cron/schedule.proto";
 
 // this line is used by starport scaffolding # proto/tx/import
 
@@ -15,10 +16,51 @@ option go_package = "github.com/neutron-org/neutron/v4/x/cron/types";
 service Msg {
   option (cosmos.msg.v1.service) = true;
 
+  // Adds new schedule.
+  rpc AddSchedule(MsgAddSchedule) returns (MsgAddScheduleResponse);
+  // Removes schedule.
+  rpc RemoveSchedule(MsgRemoveSchedule) returns (MsgRemoveScheduleResponse);
+  // Updates the module parameters.
   rpc UpdateParams(MsgUpdateParams) returns (MsgUpdateParamsResponse);
   // this line is used by starport scaffolding # proto/tx/rpc
 }
 
+// MsgAddSchedule is the MsgAddSchedule request type.
+message MsgAddSchedule {
+  option (amino.name) = "cron/MsgAddSchedule";
+  option (cosmos.msg.v1.signer) = "authority";
+
+  // Authority is the address of the governance account.
+  string authority = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+  // Name of the schedule
+  string name = 2;
+  // Period in blocks
+  uint64 period = 3;
+  // Msgs that will be executed every certain number of blocks, specified in the `period` field
+  repeated MsgExecuteContract msgs = 4 [(gogoproto.nullable) = false];
+  // Execution stages when messages will be executed
+  repeated ExecutionStage execution_stages = 5 [(gogoproto.nullable) = false];
+}
+
+// MsgAddScheduleResponse defines the response structure for executing a
+// MsgAddSchedule message.
+message MsgAddScheduleResponse {}
+
+// MsgRemoveSchedule is the MsgRemoveSchedule request type.
+message MsgRemoveSchedule {
+  option (amino.name) = "cron/MsgRemoveSchedule";
+  option (cosmos.msg.v1.signer) = "authority";
+
+  // Authority is the address of the governance account.
+  string authority = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
+  // Name of the schedule
+  string name = 2;
+}
+
+// MsgRemoveScheduleResponse defines the response structure for executing a
+// MsgRemoveSchedule message.
+message MsgRemoveScheduleResponse {}
+
 // this line is used by starport scaffolding # proto/tx/message
 
 // MsgUpdateParams is the MsgUpdateParams request type.

proto/neutron/cron/v1/schedule.proto

@@ -0,0 +1,32 @@
+syntax = "proto3";
+package neutron.cron.v1;
+
+import "gogoproto/gogo.proto";
+
+option go_package = "github.com/neutron-org/neutron/v4/x/cron/types/v1";
+
+// Schedule defines the schedule for execution
+message Schedule {
+  // Name of schedule
+  string name = 1;
+  // Period in blocks
+  uint64 period = 2;
+  // Msgs that will be executed every certain number of blocks, specified in the `period` field
+  repeated MsgExecuteContract msgs = 3 [(gogoproto.nullable) = false];
+  // Last execution's block height
+  uint64 last_execute_height = 4;
+}
+
+// MsgExecuteContract defines the contract and the message to pass
+message MsgExecuteContract {
+  // Contract is the address of the smart contract
+  string contract = 1;
+  // Msg is json encoded message to be passed to the contract
+  string msg = 2;
+}
+
+// ScheduleCount defines the number of current schedules
+message ScheduleCount {
+  // Count is the number of current schedules
+  int32 count = 1;
+}

wasmbinding/bindings/msg.go

@@ -195,9 +195,10 @@ type ForceTransfer struct {
 
 // AddSchedule adds new schedule to the cron module
 type AddSchedule struct {
-	Name   string               `json:"name"`
-	Period uint64               `json:"period"`
-	Msgs   []MsgExecuteContract `json:"msgs"`
+	Name            string               `json:"name"`
+	Period          uint64               `json:"period"`
+	Msgs            []MsgExecuteContract `json:"msgs"`
+	ExecutionStages []string             `json:"execution_stages"`
 }
 
 // AddScheduleResponse holds response AddSchedule

wasmbinding/message_plugin.go

@@ -985,7 +985,12 @@ func (m *CustomMessenger) addSchedule(ctx sdk.Context, contractAddr sdk.AccAddre
 		})
 	}
 
-	err := m.CronKeeper.AddSchedule(ctx, addSchedule.Name, addSchedule.Period, msgs)
+	executionStages := make([]crontypes.ExecutionStage, 0, len(addSchedule.ExecutionStages))
+	for _, executionStage := range addSchedule.ExecutionStages {
+		executionStages = append(executionStages, crontypes.ExecutionStage(crontypes.ExecutionStage_value[executionStage]))
+	}
+
+	err := m.CronKeeper.AddSchedule(ctx, addSchedule.Name, addSchedule.Period, msgs, executionStages)
 	if err != nil {
 		ctx.Logger().Error("failed to addSchedule",
 			"from_address", contractAddr.String(),

x/cron/genesis.go

@@ -11,7 +11,7 @@ import (
 func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState) {
 	// Set all the schedules
 	for _, elem := range genState.ScheduleList {
-		err := k.AddSchedule(ctx, elem.Name, elem.Period, elem.Msgs)
+		err := k.AddSchedule(ctx, elem.Name, elem.Period, elem.Msgs, elem.ExecutionStages)
 		if err != nil {
 			panic(err)
 		}

x/cron/keeper/grpc_query_schedule_test.go

@@ -133,8 +133,9 @@ func createNSchedule(t *testing.T, ctx sdk.Context, k *cronkeeper.Keeper, n int3
 		item.Period = 1000
 		item.Msgs = nil
 		item.LastExecuteHeight = uint64(ctx.BlockHeight())
+		item.ExecutionStages = []types.ExecutionStage{types.ExecutionStage_EXECUTION_STAGE_END_BLOCKER}
 
-		err := k.AddSchedule(ctx, item.Name, item.Period, item.Msgs)
+		err := k.AddSchedule(ctx, item.Name, item.Period, item.Msgs, item.ExecutionStages)
 		require.NoError(t, err)
 
 		res[idx] = item

x/cron/keeper/keeper.go

@@ -27,6 +27,8 @@ var (
 
 	MetricLabelSuccess      = "success"
 	MetricLabelScheduleName = "schedule_name"
+
+	schedulesExecutionStages map[string]map[types.ExecutionStage]struct{}
 )
 
 type (
@@ -47,6 +49,8 @@ func NewKeeper(
 	accountKeeper types.AccountKeeper,
 	authority string,
 ) *Keeper {
+	schedulesExecutionStages = make(map[string]map[types.ExecutionStage]struct{})
+
 	return &Keeper{
 		cdc:           cdc,
 		storeKey:      storeKey,
@@ -66,30 +70,52 @@ func (k *Keeper) Logger(ctx sdk.Context) log.Logger {
 
 // ExecuteReadySchedules gets all schedules that are due for execution (with limit that is equal to Params.Limit)
 // and executes messages in each one
-// NOTE that errors in contract calls rollback all already executed messages
-func (k *Keeper) ExecuteReadySchedules(ctx sdk.Context) {
+func (k *Keeper) ExecuteReadySchedules(ctx sdk.Context, executionStage types.ExecutionStage) {
 	telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), LabelExecuteReadySchedules)
 	schedules := k.getSchedulesReadyForExecution(ctx)
 
 	for _, schedule := range schedules {
-		err := k.executeSchedule(ctx, schedule)
-		recordExecutedSchedule(err, schedule)
+		if _, ok := schedulesExecutionStages[schedule.Name][executionStage]; ok {
+			err := k.executeSchedule(ctx, schedule)
+			recordExecutedSchedule(err, schedule)
+		}
 	}
 }
 
-// AddSchedule adds new schedule to execution for every block `period`.
+// AddSchedule adds a new schedule to be executed every certain number of blocks, specified in the `period`.
 // First schedule execution is supposed to be on `now + period` block.
-func (k *Keeper) AddSchedule(ctx sdk.Context, name string, period uint64, msgs []types.MsgExecuteContract) error {
+func (k *Keeper) AddSchedule(
+	ctx sdk.Context,
+	name string,
+	period uint64,
+	msgs []types.MsgExecuteContract,
+	executionStages []types.ExecutionStage,
+) error {
 	if k.scheduleExists(ctx, name) {
 		return fmt.Errorf("schedule already exists with name=%v", name)
 	}
 
+	schedulesExecutionStages[name] = make(map[types.ExecutionStage]struct{})
+	execStages := make([]types.ExecutionStage, 0)
+	for _, executionStage := range executionStages {
+		if _, ok := types.ExecutionStage_name[int32(executionStage)]; !ok {
+			executionStage = types.ExecutionStage_EXECUTION_STAGE_END_BLOCKER
+		}
+
+		if _, ok := schedulesExecutionStages[name][executionStage]; !ok {
+			schedulesExecutionStages[name][executionStage] = struct{}{}
+			execStages = append(execStages, executionStage)
+		}
+	}
+
 	schedule := types.Schedule{
 		Name:              name,
 		Period:            period,
 		Msgs:              msgs,
 		LastExecuteHeight: uint64(ctx.BlockHeight()), // let's execute newly added schedule on `now + period` block
+		ExecutionStages:   execStages,
 	}
+
 	k.storeSchedule(ctx, schedule)
 	k.changeTotalCount(ctx, 1)
 
@@ -102,6 +128,8 @@ func (k *Keeper) RemoveSchedule(ctx sdk.Context, name string) {
 		return
 	}
 
+	delete(schedulesExecutionStages, name)
+
 	k.changeTotalCount(ctx, -1)
 	k.removeSchedule(ctx, name)
 }