pass reader id to persistence as a hint

common/constants.go

@@ -97,3 +97,8 @@ const (
 	// Limit for schedule notes field
 	ScheduleNotesSizeLimit = 1000
 )
+
+const (
+	// DefaultQueueReaderID is the default readerID when loading history tasks
+	DefaultQueueReaderID = 0
+)

common/persistence/tests/execution_mutable_state_task.go

@@ -40,6 +40,7 @@ import (
 	"go.temporal.io/api/serviceerror"
 
 	persistencespb "go.temporal.io/server/api/persistence/v1"
+	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/debug"
 	"go.temporal.io/server/common/definition"
 	"go.temporal.io/server/common/dynamicconfig"
@@ -518,6 +519,7 @@ func (s *ExecutionMutableStateTaskSuite) TestGetScheduledTasksOrdered() {
 	response, err := s.ExecutionManager.GetHistoryTasks(s.Ctx, &p.GetHistoryTasksRequest{
 		ShardID:             s.ShardID,
 		TaskCategory:        fakeScheduledTaskCategory,
+		ReaderID:            common.DefaultQueueReaderID,
 		InclusiveMinTaskKey: tasks.NewKey(now, 0),
 		ExclusiveMaxTaskKey: tasks.NewKey(now.Add(time.Second), 0),
 		BatchSize:           10,
@@ -565,6 +567,7 @@ func (s *ExecutionMutableStateTaskSuite) PaginateTasks(
 	request := &p.GetHistoryTasksRequest{
 		ShardID:             s.ShardID,
 		TaskCategory:        category,
+		ReaderID:            common.DefaultQueueReaderID,
 		InclusiveMinTaskKey: inclusiveMinTaskKey,
 		ExclusiveMaxTaskKey: exclusiveMaxTaskKey,
 		BatchSize:           batchSize,

service/frontend/adminHandler.go

@@ -747,6 +747,7 @@ func (adh *AdminHandler) ListHistoryTasks(
 	resp, err := adh.persistenceExecutionManager.GetHistoryTasks(ctx, &persistence.GetHistoryTasksRequest{
 		ShardID:             request.ShardId,
 		TaskCategory:        taskCategory,
+		ReaderID:            common.DefaultQueueReaderID, // TODO: 1. does this work? need to move to history?
 		InclusiveMinTaskKey: minTaskKey,
 		ExclusiveMaxTaskKey: maxTaskKey,
 		BatchSize:           int(request.BatchSize),

service/history/queues/iterator.go

@@ -33,7 +33,8 @@ import (
 
 type (
 	Iterator interface {
-		collection.Iterator[tasks.Task]
+		HasNext(readerID int32) bool
+		Next(readerID int32) (tasks.Task, error)
 
 		Range() Range
 		CanSplit(tasks.Key) bool
@@ -43,13 +44,14 @@ type (
 		Remaining() Iterator
 	}
 
-	PaginationFnProvider func(Range) collection.PaginationFn[tasks.Task]
+	PaginationFnProvider func(int32, Range) collection.PaginationFn[tasks.Task]
 
 	IteratorImpl struct {
 		paginationFnProvider PaginationFnProvider
 		remainingRange       Range
 
-		pagingIterator collection.Iterator[tasks.Task]
+		iteratorReaderID int32
+		pagingIterator   collection.Iterator[tasks.Task]
 	}
 )
 
@@ -62,20 +64,22 @@ func NewIterator(
 		remainingRange:       r,
 
 		// lazy initialized to prevent task pre-fetching on creating the iterator
-		pagingIterator: nil,
+		iteratorReaderID: 0,
+		pagingIterator:   nil,
 	}
 }
 
-func (i *IteratorImpl) HasNext() bool {
-	if i.pagingIterator == nil {
-		i.pagingIterator = collection.NewPagingIterator(i.paginationFnProvider(i.remainingRange))
+func (i *IteratorImpl) HasNext(readerID int32) bool {
+	if i.pagingIterator == nil || i.iteratorReaderID != readerID {
+		i.pagingIterator = collection.NewPagingIterator(i.paginationFnProvider(readerID, i.remainingRange))
+		i.iteratorReaderID = readerID
 	}
 
 	return i.pagingIterator.HasNext()
 }
 
-func (i *IteratorImpl) Next() (tasks.Task, error) {
-	if !i.HasNext() {
+func (i *IteratorImpl) Next(readerID int32) (tasks.Task, error) {
+	if !i.HasNext(readerID) {
 		panic("Iterator encountered Next call when there is no next item")
 	}
 

service/history/queues/iterator_test.go

@@ -71,7 +71,7 @@ func (s *iteratorSuite) TestNext_IncreaseTaskKey() {
 	taskKey := NewRandomKeyInRange(r)
 	mockTask := tasks.NewMockTask(s.controller)
 	mockTask.EXPECT().GetKey().Return(taskKey).Times(1)
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, paginationRange Range) collection.PaginationFn[tasks.Task] {
 		s.Equal(r, paginationRange)
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			return []tasks.Task{mockTask}, nil, nil
@@ -81,14 +81,14 @@ func (s *iteratorSuite) TestNext_IncreaseTaskKey() {
 	iterator := NewIterator(paginationFnProvider, r)
 	s.Equal(r, iterator.Range())
 
-	s.True(iterator.HasNext())
-	task, err := iterator.Next()
+	s.True(iterator.HasNext(DefaultReaderId))
+	task, err := iterator.Next(DefaultReaderId)
 	s.NoError(err)
 	s.Equal(mockTask, task)
 
 	s.Equal(NewRange(taskKey.Next(), r.ExclusiveMax), iterator.Range())
 
-	s.False(iterator.HasNext())
+	s.False(iterator.HasNext(DefaultReaderId))
 }
 
 func (s *iteratorSuite) TestCanSplit() {
@@ -113,7 +113,7 @@ func (s *iteratorSuite) TestCanSplit() {
 
 func (s *iteratorSuite) TestSplit() {
 	r := NewRandomRange()
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, _ Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			return []tasks.Task{}, nil, nil
 		}
@@ -127,8 +127,8 @@ func (s *iteratorSuite) TestSplit() {
 	leftIterator, rightIterator := iterator.Split(splitKey)
 	s.Equal(NewRange(r.InclusiveMin, splitKey), leftIterator.Range())
 	s.Equal(NewRange(splitKey, r.ExclusiveMax), rightIterator.Range())
-	s.False(leftIterator.HasNext())
-	s.False(leftIterator.HasNext())
+	s.False(leftIterator.HasNext(DefaultReaderId))
+	s.False(leftIterator.HasNext(DefaultReaderId))
 }
 
 func (s *iteratorSuite) TestCanMerge() {
@@ -170,60 +170,60 @@ func (s *iteratorSuite) TestMerge() {
 	r := NewRandomRange()
 
 	numLoad := 0
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, _ Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			numLoad++
 			return []tasks.Task{}, nil, nil
 		}
 	}
 
 	iterator := NewIterator(paginationFnProvider, r)
-	s.False(iterator.HasNext())
+	s.False(iterator.HasNext(DefaultReaderId))
 
 	incomingIterator := NewIterator(paginationFnProvider, r)
 	mergedIterator := iterator.Merge(incomingIterator)
 	s.Equal(r, mergedIterator.Range())
-	s.False(mergedIterator.HasNext())
+	s.False(mergedIterator.HasNext(DefaultReaderId))
 
 	incomingIterator = NewIterator(
 		paginationFnProvider,
 		NewRange(tasks.MinimumKey, r.InclusiveMin),
 	)
 	mergedIterator = iterator.Merge(incomingIterator)
 	s.Equal(NewRange(tasks.MinimumKey, r.ExclusiveMax), mergedIterator.Range())
-	s.False(mergedIterator.HasNext())
+	s.False(mergedIterator.HasNext(DefaultReaderId))
 
 	incomingIterator = NewIterator(
 		paginationFnProvider,
 		NewRange(r.ExclusiveMax, tasks.MaximumKey),
 	)
 	mergedIterator = iterator.Merge(incomingIterator)
 	s.Equal(NewRange(r.InclusiveMin, tasks.MaximumKey), mergedIterator.Range())
-	s.False(mergedIterator.HasNext())
+	s.False(mergedIterator.HasNext(DefaultReaderId))
 
 	incomingIterator = NewIterator(
 		paginationFnProvider,
 		NewRange(tasks.MinimumKey, NewRandomKeyInRange(r)),
 	)
 	mergedIterator = iterator.Merge(incomingIterator)
 	s.Equal(NewRange(tasks.MinimumKey, r.ExclusiveMax), mergedIterator.Range())
-	s.False(mergedIterator.HasNext())
+	s.False(mergedIterator.HasNext(DefaultReaderId))
 
 	incomingIterator = NewIterator(
 		paginationFnProvider,
 		NewRange(NewRandomKeyInRange(r), tasks.MaximumKey),
 	)
 	mergedIterator = iterator.Merge(incomingIterator)
 	s.Equal(NewRange(r.InclusiveMin, tasks.MaximumKey), mergedIterator.Range())
-	s.False(mergedIterator.HasNext())
+	s.False(mergedIterator.HasNext(DefaultReaderId))
 
 	incomingIterator = NewIterator(
 		paginationFnProvider,
 		NewRange(tasks.MinimumKey, tasks.MaximumKey),
 	)
 	mergedIterator = iterator.Merge(incomingIterator)
 	s.Equal(NewRange(tasks.MinimumKey, tasks.MaximumKey), mergedIterator.Range())
-	s.False(mergedIterator.HasNext())
+	s.False(mergedIterator.HasNext(DefaultReaderId))
 
 	// test if Merge returns a new iterator
 	s.Equal(7, numLoad)
@@ -238,7 +238,7 @@ func (s *iteratorSuite) TestRemaining() {
 	taskKey := NewRandomKeyInRange(r)
 	mockTask := tasks.NewMockTask(s.controller)
 	mockTask.EXPECT().GetKey().Return(taskKey).Times(1)
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, paginationRange Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			numLoad++
 			if paginationRange.ContainsKey(taskKey) {
@@ -249,14 +249,16 @@ func (s *iteratorSuite) TestRemaining() {
 	}
 
 	iterator := NewIterator(paginationFnProvider, r)
-	_, err := iterator.Next()
+	_, err := iterator.Next(DefaultReaderId)
 	s.NoError(err)
-	s.False(iterator.HasNext())
+	s.False(iterator.HasNext(DefaultReaderId))
 
 	remaining := iterator.Remaining()
 	s.Equal(iterator.Range(), remaining.Range())
-	s.False(remaining.HasNext())
+	s.False(remaining.HasNext(DefaultReaderId))
 
 	// test if Remaining returns a new iterator
 	s.Equal(2, numLoad)
 }
+
+// TODO: add a test on readerID change

service/history/queues/queue_base.go

@@ -47,7 +47,7 @@ import (
 )
 
 const (
-	DefaultReaderId = 0
+	DefaultReaderId = common.DefaultQueueReaderID
 
 	// Non-default readers will use critical pending task count * this coefficient
 	// as its max pending task count so that their loading will never trigger pending

service/history/queues/queue_base_test.go

@@ -257,7 +257,7 @@ func (s *queueBaseSuite) TestStartStop() {
 	)
 	mockShard.Resource.ClusterMetadata.EXPECT().GetCurrentClusterName().Return(cluster.TestCurrentClusterName).AnyTimes()
 
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, paginationRange Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			mockTask := tasks.NewMockTask(s.controller)
 			key := NewRandomKeyInRange(paginationRange)

service/history/queues/queue_immediate.go

@@ -62,14 +62,15 @@ func NewImmediateQueue(
 	logger log.Logger,
 	metricsHandler metrics.Handler,
 ) *immediateQueue {
-	paginationFnProvider := func(r Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(readerID int32, r Range) collection.PaginationFn[tasks.Task] { // readerID
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			ctx, cancel := newQueueIOContext()
 			defer cancel()
 
 			request := &persistence.GetHistoryTasksRequest{
 				ShardID:             shard.GetShardID(),
 				TaskCategory:        category,
+				ReaderID:            readerID,
 				InclusiveMinTaskKey: r.InclusiveMin,
 				ExclusiveMaxTaskKey: r.ExclusiveMax,
 				BatchSize:           options.BatchSize(),

service/history/queues/queue_scheduled.go

@@ -74,14 +74,15 @@ func NewScheduledQueue(
 	logger log.Logger,
 	metricsHandler metrics.Handler,
 ) *scheduledQueue {
-	paginationFnProvider := func(r Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(readerID int32, r Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			ctx, cancel := newQueueIOContext()
 			defer cancel()
 
 			request := &persistence.GetHistoryTasksRequest{
 				ShardID:             shard.GetShardID(),
 				TaskCategory:        category,
+				ReaderID:            readerID,
 				InclusiveMinTaskKey: tasks.NewKey(r.InclusiveMin.FireTime, 0),
 				ExclusiveMaxTaskKey: tasks.NewKey(r.ExclusiveMax.FireTime.Add(persistence.ScheduledTaskMinPrecision), 0),
 				BatchSize:           options.BatchSize(),
@@ -264,6 +265,7 @@ func (p *scheduledQueue) lookAheadTask() {
 	request := &persistence.GetHistoryTasksRequest{
 		ShardID:             p.shard.GetShardID(),
 		TaskCategory:        p.category,
+		ReaderID:            DefaultReaderId,
 		InclusiveMinTaskKey: tasks.NewKey(lookAheadMinTime, 0),
 		ExclusiveMaxTaskKey: tasks.NewKey(lookAheadMaxTime, 0),
 		BatchSize:           1,

service/history/queues/queue_scheduled_test.go

@@ -168,6 +168,7 @@ func (s *scheduledQueueSuite) TestPaginationFnProvider() {
 	s.mockExecutionManager.EXPECT().GetHistoryTasks(gomock.Any(), &persistence.GetHistoryTasksRequest{
 		ShardID:             s.mockShard.GetShardID(),
 		TaskCategory:        tasks.CategoryTimer,
+		ReaderID:            DefaultReaderId,
 		InclusiveMinTaskKey: tasks.NewKey(r.InclusiveMin.FireTime, 0),
 		ExclusiveMaxTaskKey: tasks.NewKey(r.ExclusiveMax.FireTime.Add(persistence.ScheduledTaskMinPrecision), 0),
 		BatchSize:           testQueueOptions.BatchSize(),
@@ -177,7 +178,7 @@ func (s *scheduledQueueSuite) TestPaginationFnProvider() {
 		NextPageToken: nextPageToken,
 	}, nil).Times(1)
 
-	paginationFn := paginationFnProvider(r)
+	paginationFn := paginationFnProvider(DefaultReaderId, r)
 	loadedTasks, actualNextPageToken, err := paginationFn(currentPageToken)
 	s.NoError(err)
 	for _, task := range loadedTasks {
@@ -278,6 +279,7 @@ func (s *scheduledQueueSuite) setupLookAheadMock(
 	s.mockExecutionManager.EXPECT().GetHistoryTasks(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, request *persistence.GetHistoryTasksRequest) (*persistence.GetHistoryTasksResponse, error) {
 		s.Equal(s.mockShard.GetShardID(), request.ShardID)
 		s.Equal(tasks.CategoryTimer, request.TaskCategory)
+		s.Equal(int32(DefaultReaderId), request.ReaderID)
 		s.Equal(lookAheadRange.InclusiveMin, request.InclusiveMinTaskKey)
 		s.Equal(1, request.BatchSize)
 		s.Nil(request.NextPageToken)

service/history/queues/reader_test.go

@@ -94,7 +94,7 @@ func (s *readerSuite) TestStartLoadStop() {
 	r := NewRandomRange()
 	scopes := []Scope{NewScope(r, predicates.Universal[tasks.Task]())}
 
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, paginationRange Range) collection.PaginationFn[tasks.Task] {
 		s.Equal(r, paginationRange)
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			mockTask := tasks.NewMockTask(s.controller)
@@ -256,7 +256,7 @@ func (s *readerSuite) TestShrinkSlices() {
 func (s *readerSuite) TestThrottle() {
 	scopes := NewRandomScopes(1)
 
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, _ Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			mockTask := tasks.NewMockTask(s.controller)
 			mockTask.EXPECT().GetKey().Return(NewRandomKeyInRange(scopes[0].Range)).AnyTimes()
@@ -319,7 +319,7 @@ func (s *readerSuite) TestLoadAndSubmitTasks_TooManyPendingTasks() {
 func (s *readerSuite) TestLoadAndSubmitTasks_MoreTasks() {
 	scopes := NewRandomScopes(1)
 
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, _ Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			result := make([]tasks.Task, 0, 100)
 			for i := 0; i != 100; i++ {
@@ -354,7 +354,7 @@ func (s *readerSuite) TestLoadAndSubmitTasks_MoreTasks() {
 func (s *readerSuite) TestLoadAndSubmitTasks_NoMoreTasks_HasNextSlice() {
 	scopes := NewRandomScopes(2)
 
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, _ Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			mockTask := tasks.NewMockTask(s.controller)
 			mockTask.EXPECT().GetKey().Return(NewRandomKeyInRange(scopes[0].Range)).AnyTimes()
@@ -384,7 +384,7 @@ func (s *readerSuite) TestLoadAndSubmitTasks_NoMoreTasks_HasNextSlice() {
 func (s *readerSuite) TestLoadAndSubmitTasks_NoMoreTasks_NoNextSlice() {
 	scopes := NewRandomScopes(1)
 
-	paginationFnProvider := func(paginationRange Range) collection.PaginationFn[tasks.Task] {
+	paginationFnProvider := func(_ int32, _ Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
 			mockTask := tasks.NewMockTask(s.controller)
 			mockTask.EXPECT().GetKey().Return(NewRandomKeyInRange(scopes[0].Range)).AnyTimes()

service/history/queues/slice.go

@@ -372,8 +372,8 @@ func (s *SliceImpl) SelectTasks(readerID int32, batchSize int) ([]Executable, er
 
 	executables := make([]Executable, 0, batchSize)
 	for len(executables) < batchSize && len(s.iterators) != 0 {
-		if s.iterators[0].HasNext() {
-			task, err := s.iterators[0].Next()
+		if s.iterators[0].HasNext(readerID) {
+			task, err := s.iterators[0].Next(readerID)
 			if err != nil {
 				s.iterators[0] = s.iterators[0].Remaining()
 				if len(executables) != 0 {