Fix scheduled task rescheduling on failover (#5377)

temporalio · Feb 1, 2024 · 2647b36 · 2647b36
1 parent fb61704
commit 2647b36
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 6 deletions.
diff --git a/service/history/queues/reader.go b/service/history/queues/reader.go
@@ -527,9 +527,10 @@ func (r *ReaderImpl) submit(
 	executable Executable,
 ) {
 	now := r.timeSource.Now()
-	// Persistence layer may lose precision when persisting the task, which essentially move
-	// task fire time forward. Need to account for that when submitting the task.
-	if fireTime := executable.GetKey().FireTime.Add(persistence.ScheduledTaskMinPrecision); now.Before(fireTime) {
+	// Persistence layer may lose precision when persisting the task, which essentially moves
+	// task fire time backward. Need to account for that when submitting the task.
+	fireTime := executable.GetKey().FireTime.Add(persistence.ScheduledTaskMinPrecision)
+	if now.Before(fireTime) {
 		r.rescheduler.Add(executable, fireTime)
 		return
 	}

diff --git a/service/history/queues/rescheduler.go b/service/history/queues/rescheduler.go
@@ -38,8 +38,10 @@ import (
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/log/tag"
 	"go.temporal.io/server/common/metrics"
+	"go.temporal.io/server/common/persistence"
 	ctasks "go.temporal.io/server/common/tasks"
 	"go.temporal.io/server/common/timer"
+	"go.temporal.io/server/common/util"
 )
 
 const (
@@ -176,7 +178,12 @@ func (r *reschedulerImpl) Reschedule(
 		items := make([]rescheduledExecuable, 0, pq.Len())
 		for !pq.IsEmpty() {
 			rescheduled := pq.Remove()
-			rescheduled.rescheduleTime = now
+			// scheduled queue pre-fetches tasks,
+			// so we need to make sure the reschedule time is not before the task scheduled time
+			rescheduled.rescheduleTime = util.MaxTime(
+				rescheduled.executable.GetKey().FireTime.Add(persistence.ScheduledTaskMinPrecision),
+				now,
+			)
 			items = append(items, rescheduled)
 		}
 		r.pqMap[key] = r.newPriorityQueue(items)

diff --git a/service/history/queues/rescheduler_test.go b/service/history/queues/rescheduler_test.go
@@ -38,6 +38,7 @@ import (
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/metrics"
 	ctasks "go.temporal.io/server/common/tasks"
+	"go.temporal.io/server/service/history/tasks"
 )
 
 type (
@@ -221,7 +222,7 @@ func (s *rescheudulerSuite) TestReschedule_DropCancelled() {
 	s.Equal(0, s.rescheduler.Len())
 }
 
-func (s *rescheudulerSuite) TestImmdiateReschedule() {
+func (s *rescheudulerSuite) TestForceReschedule_ImmediateTask() {
 	now := time.Now()
 	s.timeSource.Update(now)
 	namespaceID := s.mockScheduler.TaskChannelKeyFn()(nil).NamespaceID
@@ -234,8 +235,9 @@ func (s *rescheudulerSuite) TestImmdiateReschedule() {
 	taskWG.Add(numTask)
 	for i := 0; i != numTask; i++ {
 		mockTask := NewMockExecutable(s.controller)
-		mockTask.EXPECT().State().Return(ctasks.TaskStatePending).Times(1)
+		mockTask.EXPECT().State().Return(ctasks.TaskStatePending).AnyTimes()
 		mockTask.EXPECT().SetScheduledTime(gomock.Any()).AnyTimes()
+		mockTask.EXPECT().GetKey().Return(tasks.NewImmediateKey(int64(i))).AnyTimes()
 		s.rescheduler.Add(
 			mockTask,
 			now.Add(time.Minute+time.Duration(rand.Int63n(time.Minute.Nanoseconds()))),
@@ -251,3 +253,44 @@ func (s *rescheudulerSuite) TestImmdiateReschedule() {
 	taskWG.Wait()
 	s.Equal(0, s.rescheduler.Len())
 }
+
+func (s *rescheudulerSuite) TestForceReschedule_ScheduledTask() {
+	now := time.Now()
+	s.timeSource.Update(now)
+	namespaceID := s.mockScheduler.TaskChannelKeyFn()(nil).NamespaceID
+
+	s.rescheduler.Start()
+	defer s.rescheduler.Stop()
+
+	taskWG := &sync.WaitGroup{}
+	taskWG.Add(1)
+
+	retryingTask := NewMockExecutable(s.controller)
+	retryingTask.EXPECT().State().Return(ctasks.TaskStatePending).AnyTimes()
+	retryingTask.EXPECT().SetScheduledTime(gomock.Any()).AnyTimes()
+	retryingTask.EXPECT().GetKey().Return(tasks.NewKey(now.Add(-time.Minute), int64(1))).AnyTimes()
+	s.rescheduler.Add(
+		retryingTask,
+		now.Add(time.Minute),
+	)
+
+	// schedule queue pre-fetches tasks
+	futureTaskTimestamp := now.Add(time.Second)
+	futureTask := NewMockExecutable(s.controller)
+	futureTask.EXPECT().State().Return(ctasks.TaskStatePending).AnyTimes()
+	futureTask.EXPECT().SetScheduledTime(gomock.Any()).AnyTimes()
+	futureTask.EXPECT().GetKey().Return(tasks.NewKey(futureTaskTimestamp, int64(2))).AnyTimes()
+	s.rescheduler.Add(
+		futureTask,
+		futureTaskTimestamp,
+	)
+
+	s.mockScheduler.EXPECT().TrySubmit(gomock.Any()).DoAndReturn(func(_ Executable) bool {
+		taskWG.Done()
+		return true
+	}).Times(1)
+
+	s.rescheduler.Reschedule(namespaceID)
+	taskWG.Wait()
+	s.Equal(1, s.rescheduler.Len())
+}