Reset slice progress after a potentially successful write

Author: tubignat

service/history/queuev2/queue_base.go

@@ -278,8 +278,8 @@ func (q *queueBase) insertSingleTask(task task.Task) bool {
 	return q.virtualQueueManager.InsertSingleTaskToRootQueue(task)
 }
 
-func (q *queueBase) removeScheduledTasksAfter(t time.Time) {
-	q.virtualQueueManager.RemoveScheduledTasksAfter(t)
+func (q *queueBase) removeScheduledTasksAfter(key persistence.HistoryTaskKey) {
+	q.virtualQueueManager.RemoveScheduledTasksAfter(key)
 }
 
 func (q *queueBase) updateQueueState(ctx context.Context) {

service/history/queuev2/queue_scheduled.go

@@ -175,7 +175,7 @@ func (q *scheduledQueue) NotifyNewTask(clusterName string, info *hcommon.NotifyT
 	}
 
 	if !nextReadTime.IsZero() {
-		q.base.removeScheduledTasksAfter(nextReadTime)
+		q.base.removeScheduledTasksAfter(persistence.NewHistoryTaskKey(nextReadTime, 0))
 		q.notify(nextReadTime)
 	}
 

service/history/queuev2/virtual_queue.go

@@ -70,7 +70,7 @@ type (
 		// InsertSingleTask inserts a single task to the virtual queue. Return false if the task's timestamp is out of range of the current queue slice..
 		InsertSingleTask(task task.Task) bool
 		// RemoveScheduledTasksAfter removes the scheduled tasks after the given time
-		RemoveScheduledTasksAfter(time.Time)
+		RemoveScheduledTasksAfter(persistence.HistoryTaskKey)
 	}
 
 	VirtualQueueOptions struct {
@@ -452,31 +452,22 @@ func (q *virtualQueueImpl) InsertSingleTask(task task.Task) bool {
 	return true
 }
 
-func (q *virtualQueueImpl) RemoveScheduledTasksAfter(t time.Time) {
+func (q *virtualQueueImpl) RemoveScheduledTasksAfter(key persistence.HistoryTaskKey) {
 	q.Lock()
 	defer q.Unlock()
 
 	for e := q.virtualSlices.Front(); e != nil; e = e.Next() {
-		s := e.Value.(VirtualSlice)
-		r := s.GetState().Range
-
-		if t.Before(r.InclusiveMinTaskKey.GetScheduledTime()) {
-			continue
-		}
-
-		s.CancelTasks(func(task task.Task) bool {
-			taskTime := task.GetTaskKey().GetScheduledTime()
-			return taskTime.After(t) || taskTime.Equal(t)
-		})
-
-		q.monitor.SetSlicePendingTaskCount(s, s.GetPendingTaskCount())
+		slice := e.Value.(VirtualSlice)
+		slice.CancelTasksAfter(key)
+		q.monitor.SetSlicePendingTaskCount(slice, slice.GetPendingTaskCount())
 	}
 }
 
 func (q *virtualQueueImpl) submitTask(now time.Time, task task.Task) error {
 	if persistence.IsTaskCorrupted(task) {
 		q.logger.Error("Virtual queue encountered a corrupted task", tag.Dynamic("task", task))
 		q.metricsScope.IncCounter(metrics.CorruptedHistoryTaskCounter)
+
 		task.Ack()
 		return nil
 	}

service/history/queuev2/virtual_queue_manager.go

@@ -34,6 +34,7 @@ import (
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/log/tag"
 	"github.com/uber/cadence/common/metrics"
+	"github.com/uber/cadence/common/persistence"
 	"github.com/uber/cadence/common/quotas"
 	"github.com/uber/cadence/service/history/task"
 )
@@ -58,7 +59,7 @@ type (
 		AddNewVirtualSliceToRootQueue(VirtualSlice)
 		// Insert a single task to the current slice. Return false if the task's timestamp is out of range of the current slice.
 		InsertSingleTaskToRootQueue(task.Task) bool
-		RemoveScheduledTasksAfter(time.Time)
+		RemoveScheduledTasksAfter(persistence.HistoryTaskKey)
 	}
 
 	virtualQueueManagerImpl struct {
@@ -232,13 +233,14 @@ func (m *virtualQueueManagerImpl) InsertSingleTaskToRootQueue(t task.Task) bool
 	return false
 }
 
-func (m *virtualQueueManagerImpl) RemoveScheduledTasksAfter(t time.Time) {
+func (m *virtualQueueManagerImpl) RemoveScheduledTasksAfter(key persistence.HistoryTaskKey) {
 	m.Lock()
 	defer m.Unlock()
 	for _, vq := range m.virtualQueues {
-		vq.RemoveScheduledTasksAfter(t)
+		vq.RemoveScheduledTasksAfter(key)
 	}
 }
+
 func (m *virtualQueueManagerImpl) appendOrMergeSlice(vq VirtualQueue, s VirtualSlice) {
 	now := m.timeSource.Now()
 	newVirtualSliceState := s.GetState()

service/history/queuev2/virtual_queue_manager_mock.go

@@ -43,7 +43,7 @@ type (
 		GetPendingTaskCount() int
 		Clear()
 		PendingTaskStats() PendingTaskStats
-		CancelTasks(predicate func(task.Task) bool)
+		CancelTasksAfter(key persistence.HistoryTaskKey)
 
 		TrySplitByTaskKey(persistence.HistoryTaskKey) (VirtualSlice, VirtualSlice, bool)
 		TrySplitByPredicate(Predicate) (VirtualSlice, VirtualSlice, bool)
@@ -197,13 +197,40 @@ func (s *virtualSliceImpl) UpdateAndGetState() VirtualSliceState {
 	return s.state
 }
 
-func (s *virtualSliceImpl) CancelTasks(predicate func(task.Task) bool) {
+func (s *virtualSliceImpl) CancelTasksAfter(key persistence.HistoryTaskKey) {
 	taskMap := s.pendingTaskTracker.GetTasks()
 	for _, task := range taskMap {
-		if predicate(task) {
+		if task.GetTaskKey().Compare(key) >= 0 {
 			task.Cancel()
 		}
 	}
+
+	if len(s.progress) == 0 {
+		s.progress = []*GetTaskProgress{
+			{
+				Range: Range{
+					InclusiveMinTaskKey: key,
+					ExclusiveMaxTaskKey: s.state.Range.ExclusiveMaxTaskKey,
+				},
+				NextPageToken: nil,
+				NextTaskKey:   key,
+			},
+		}
+		return
+	}
+
+	for i, progress := range s.progress {
+		if progress.NextTaskKey.Compare(key) > 0 {
+			s.progress[i] = &GetTaskProgress{
+				Range: Range{
+					InclusiveMinTaskKey: key,
+					ExclusiveMaxTaskKey: progress.Range.ExclusiveMaxTaskKey,
+				},
+				NextPageToken: nil,
+				NextTaskKey:   key,
+			}
+		}
+	}
 }
 
 func (s *virtualSliceImpl) TrySplitByTaskKey(taskKey persistence.HistoryTaskKey) (VirtualSlice, VirtualSlice, bool) {