NBSNEBIUS-89: [Tasks] add 'crontab' parameters to TaskSchedule (not s…

…upported yet); refactor tasks/storage.CreateRegularTasks
ydb-platform · Jan 26, 2024 · c2ac4bf · c2ac4bf
1 parent 01ea147
commit c2ac4bf
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 13 deletions.
diff --git a/cloud/tasks/scheduler.go b/cloud/tasks/scheduler.go
@@ -13,6 +13,12 @@ import (
 type TaskSchedule struct {
 	ScheduleInterval time.Duration
 	MaxTasksInflight int
+
+	// Crontab params.
+	// Schedules task every day - only 'hour' and 'min' are supported.
+	UseCrontab bool // If set, ScheduleInterval is ignored.
+	Hour       int  // (0 - 23)
+	Min        int  // (0 - 59)
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/cloud/tasks/scheduler_impl.go b/cloud/tasks/scheduler_impl.go
@@ -155,6 +155,9 @@ func (s *scheduler) ScheduleRegularTasks(
 			schedule := tasks_storage.TaskSchedule{
 				ScheduleInterval: schedule.ScheduleInterval,
 				MaxTasksInflight: schedule.MaxTasksInflight,
+
+				Hour: schedule.Hour,
+				Min:  schedule.Min,
 			}
 
 			err = s.storage.CreateRegularTasks(ctx, tasks_storage.TaskState{

diff --git a/cloud/tasks/storage/storage.go b/cloud/tasks/storage/storage.go
@@ -258,6 +258,12 @@ type TaskInfo struct {
 type TaskSchedule struct {
 	ScheduleInterval time.Duration
 	MaxTasksInflight int
+
+	// Crontab params.
+	// Schedules task every day - only 'hour' and 'min' are supported.
+	UseCrontab bool // If set, ScheduleInterval is ignored.
+	Hour       int  // (0 - 23)
+	Min        int  // (0 - 59)
 }
 
 ////////////////////////////////////////////////////////////////////////////////

diff --git a/cloud/tasks/storage/storage_ydb_impl.go b/cloud/tasks/storage/storage_ydb_impl.go
@@ -702,32 +702,26 @@ func (s *storageYDB) createRegularTasks(
 		}
 	}
 
-	scheduled := false
+	shouldSchedule := false
 
 	if found {
 		schedulingTime := schState.scheduledAt.Add(schedule.ScheduleInterval)
 
 		if schState.tasksInflight == 0 && state.CreatedAt.After(schedulingTime) {
-			err := s.addRegularTasks(ctx, tx, state, schedule.MaxTasksInflight)
-			if err != nil {
-				return err
-			}
-
-			schState.tasksInflight = uint64(schedule.MaxTasksInflight)
-			scheduled = true
+			shouldSchedule = true
 		}
 	} else {
+		shouldSchedule = true
+	}
+
+	if shouldSchedule {
 		err := s.addRegularTasks(ctx, tx, state, schedule.MaxTasksInflight)
 		if err != nil {
 			return err
 		}
 
 		schState.taskType = state.TaskType
 		schState.tasksInflight = uint64(schedule.MaxTasksInflight)
-		scheduled = true
-	}
-
-	if scheduled {
 		schState.scheduledAt = state.CreatedAt
 
 		_, err = tx.Execute(ctx, fmt.Sprintf(`
@@ -754,7 +748,7 @@ func (s *storageYDB) createRegularTasks(
 		return err
 	}
 
-	if scheduled {
+	if shouldSchedule {
 		s.metrics.OnTaskCreated(state, schedule.MaxTasksInflight)
 	}