Files
tasks/service/history/replication/ack_manager_test.go
nikki-dagandGitHub a645917055 Pass context to mutable state transaction close methods (#8822)
## What changed?
The changes consistently add a context.Context parameter to:
  - CloseTransactionAsSnapshot() - the main transaction close method
  - TestCloneToProto() - test helper method
  - prepareNewWorkflow() - workflow preparation method

All call sites have been updated to pass the context through, with test
code using context.Background() where needed. This is a clean
refactoring to enable context propagation through the transaction close
path.


## Why?
We want to support accessing contextual information before closing the
mutable state transaction.

## How did you test it?
- [x] built
- [x] run locally and tested manually
- [x] covered by existing tests
- [ ] added new unit test(s)
- [ ] added new functional test(s)



---

> [!NOTE]
> Add context propagation to mutable state transaction closing by
updating method signatures and callers across history, ndc, workflow,
and tests.
> 
> - **API changes**:
> - Add `context.Context` parameter to
`MutableState.CloseTransactionAsSnapshot` and
`CloseTransactionAsMutation` in `history/interfaces` (and mocks).
>   - Update `workflow.TestCloneToProto(ctx, ms)` to accept/pass `ctx`.
> - **Implementations**:
> - Propagate `ctx` through `workflow/mutable_state_impl` (including
internal `closeTransaction`), `workflow/context`, and CHASM paths.
> - **Call sites updated**:
> - Replace calls to `CloseTransactionAsSnapshot/Mutation(...)` with
`(..., ctx, ...)` across `history_engine*`, `ndc/*`, `api/*`, queue
executors, reset/import/rebuilder, and replication.
> - Tests pass `context.Background()` where needed; numerous
`TestCloneToProto` usages updated accordingly.
> - **Related adjustments**:
> - `Starter.prepareNewWorkflow` and other helpers now take `ctx` and
forward it to snapshot closing.
> 
> <sup>Written by [Cursor
Bugbot]() for commit
ec8a15f5aff6a7c9e2a980e6569b4926425c0c7c. This will update automatically
on new commits. Configure
[here]().</sup>
2025-12-17 16:58:11 -06:00

519 lines
21 KiB
Go

package replication
import (
"context"
"strconv"
"testing"
"time"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
commonpb "go.temporal.io/api/common/v1"
"go.temporal.io/api/serviceerror"
historyspb "go.temporal.io/server/api/history/v1"
persistencespb "go.temporal.io/server/api/persistence/v1"
"go.temporal.io/server/chasm"
"go.temporal.io/server/common/cluster"
"go.temporal.io/server/common/definition"
"go.temporal.io/server/common/log"
"go.temporal.io/server/common/metrics"
"go.temporal.io/server/common/namespace"
"go.temporal.io/server/common/persistence"
"go.temporal.io/server/common/util"
"go.temporal.io/server/service/history/events"
"go.temporal.io/server/service/history/hsm"
historyi "go.temporal.io/server/service/history/interfaces"
"go.temporal.io/server/service/history/shard"
"go.temporal.io/server/service/history/tasks"
"go.temporal.io/server/service/history/tests"
"go.temporal.io/server/service/history/workflow"
wcache "go.temporal.io/server/service/history/workflow/cache"
"go.uber.org/mock/gomock"
)
type (
ackManagerSuite struct {
suite.Suite
*require.Assertions
controller *gomock.Controller
mockShard *shard.ContextTest
mockNamespaceRegistry *namespace.MockRegistry
mockMutableState *historyi.MockMutableState
mockClusterMetadata *cluster.MockMetadata
syncStateRetriever *MockSyncStateRetriever
mockExecutionMgr *persistence.MockExecutionManager
logger log.Logger
replicationAckManager *ackMgrImpl
}
)
func TestAckManagerSuite(t *testing.T) {
s := new(ackManagerSuite)
suite.Run(t, s)
}
func (s *ackManagerSuite) SetupSuite() {
}
func (s *ackManagerSuite) TearDownSuite() {
}
func (s *ackManagerSuite) SetupTest() {
s.Assertions = require.New(s.T())
s.controller = gomock.NewController(s.T())
s.mockMutableState = historyi.NewMockMutableState(s.controller)
s.mockShard = shard.NewTestContext(
s.controller,
&persistencespb.ShardInfo{
ShardId: 0,
RangeId: 1,
Owner: "test-shard-owner",
},
tests.NewDynamicConfig(),
)
reg := hsm.NewRegistry()
err := workflow.RegisterStateMachine(reg)
s.NoError(err)
s.mockShard.SetStateMachineRegistry(reg)
s.mockNamespaceRegistry = s.mockShard.Resource.NamespaceCache
s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(tests.NamespaceID).Return(tests.GlobalNamespaceEntry, nil).AnyTimes()
s.mockExecutionMgr = s.mockShard.Resource.ExecutionMgr
s.mockClusterMetadata = s.mockShard.Resource.ClusterMetadata
s.mockClusterMetadata.EXPECT().GetCurrentClusterName().Return(cluster.TestCurrentClusterName).AnyTimes()
s.mockClusterMetadata.EXPECT().GetClusterID().Return(int64(1)).AnyTimes()
s.mockClusterMetadata.EXPECT().IsGlobalNamespaceEnabled().Return(true).AnyTimes()
s.mockClusterMetadata.EXPECT().ClusterNameForFailoverVersion(true, gomock.Any()).Return(cluster.TestCurrentClusterName).AnyTimes()
s.logger = s.mockShard.GetLogger()
workflowCache := wcache.NewHostLevelCache(s.mockShard.GetConfig(), s.mockShard.GetLogger(), metrics.NoopMetricsHandler)
replicationProgressCache := NewProgressCache(s.mockShard.GetConfig(), s.mockShard.GetLogger(), metrics.NoopMetricsHandler)
s.syncStateRetriever = NewMockSyncStateRetriever(s.controller)
s.replicationAckManager = NewAckManager(
s.mockShard, workflowCache, nil, replicationProgressCache, s.mockExecutionMgr, s.syncStateRetriever, s.logger,
).(*ackMgrImpl)
}
func (s *ackManagerSuite) TearDownTest() {
s.controller.Finish()
s.mockShard.StopForTest()
}
func (s *ackManagerSuite) TestNotifyNewTasks_NotInitialized() {
s.replicationAckManager.maxTaskID = nil
s.replicationAckManager.NotifyNewTasks([]tasks.Task{
&tasks.HistoryReplicationTask{TaskID: 456},
&tasks.HistoryReplicationTask{TaskID: 123},
})
s.Equal(*s.replicationAckManager.maxTaskID, int64(456))
}
func (s *ackManagerSuite) TestNotifyNewTasks_Initialized() {
s.replicationAckManager.maxTaskID = util.Ptr(int64(123))
s.replicationAckManager.NotifyNewTasks([]tasks.Task{
&tasks.HistoryReplicationTask{TaskID: 100},
})
s.Equal(*s.replicationAckManager.maxTaskID, int64(123))
s.replicationAckManager.NotifyNewTasks([]tasks.Task{
&tasks.HistoryReplicationTask{TaskID: 234},
})
s.Equal(*s.replicationAckManager.maxTaskID, int64(234))
}
func (s *ackManagerSuite) TestTaskIDRange_NotInitialized() {
s.replicationAckManager.sanityCheckTime = time.Time{}
expectMaxTaskID := s.mockShard.GetQueueExclusiveHighReadWatermark(tasks.CategoryReplication).Prev().TaskID
expectMinTaskID := expectMaxTaskID - 100
s.replicationAckManager.maxTaskID = util.Ptr(expectMinTaskID - 100)
minTaskID, maxTaskID := s.replicationAckManager.taskIDsRange(expectMinTaskID)
s.Equal(expectMinTaskID, minTaskID)
s.Equal(expectMaxTaskID, maxTaskID)
s.NotEqual(time.Time{}, s.replicationAckManager.sanityCheckTime)
s.Equal(expectMaxTaskID, *s.replicationAckManager.maxTaskID)
}
func (s *ackManagerSuite) TestTaskIDRange_Initialized_UseHighestReplicationTaskID() {
now := time.Now().UTC()
sanityCheckTime := now.Add(2 * time.Minute)
s.replicationAckManager.sanityCheckTime = sanityCheckTime
expectMinTaskID := s.mockShard.GetQueueExclusiveHighReadWatermark(tasks.CategoryReplication).TaskID - 100
expectMaxTaskID := s.mockShard.GetQueueExclusiveHighReadWatermark(tasks.CategoryReplication).TaskID - 50
s.replicationAckManager.maxTaskID = util.Ptr(expectMaxTaskID)
minTaskID, maxTaskID := s.replicationAckManager.taskIDsRange(expectMinTaskID)
s.Equal(expectMinTaskID, minTaskID)
s.Equal(expectMaxTaskID, maxTaskID)
s.Equal(sanityCheckTime, s.replicationAckManager.sanityCheckTime)
s.Equal(expectMaxTaskID, *s.replicationAckManager.maxTaskID)
}
func (s *ackManagerSuite) TestTaskIDRange_Initialized_NoHighestReplicationTaskID() {
now := time.Now().UTC()
sanityCheckTime := now.Add(2 * time.Minute)
s.replicationAckManager.sanityCheckTime = sanityCheckTime
expectMinTaskID := s.mockShard.GetQueueExclusiveHighReadWatermark(tasks.CategoryReplication).Prev().TaskID - 100
expectMaxTaskID := s.mockShard.GetQueueExclusiveHighReadWatermark(tasks.CategoryReplication).Prev().TaskID
s.replicationAckManager.maxTaskID = nil
minTaskID, maxTaskID := s.replicationAckManager.taskIDsRange(expectMinTaskID)
s.Equal(expectMinTaskID, minTaskID)
s.Equal(expectMaxTaskID, maxTaskID)
s.Equal(sanityCheckTime, s.replicationAckManager.sanityCheckTime)
s.Equal(expectMaxTaskID, *s.replicationAckManager.maxTaskID)
}
func (s *ackManagerSuite) TestTaskIDRange_Initialized_UseHighestTransferTaskID() {
now := time.Now().UTC()
sanityCheckTime := now.Add(-2 * time.Minute)
s.replicationAckManager.sanityCheckTime = sanityCheckTime
expectMinTaskID := s.mockShard.GetQueueExclusiveHighReadWatermark(tasks.CategoryReplication).Prev().TaskID - 100
expectMaxTaskID := s.mockShard.GetQueueExclusiveHighReadWatermark(tasks.CategoryReplication).Prev().TaskID
s.replicationAckManager.maxTaskID = util.Ptr(s.mockShard.GetQueueExclusiveHighReadWatermark(tasks.CategoryReplication).TaskID - 50)
minTaskID, maxTaskID := s.replicationAckManager.taskIDsRange(expectMinTaskID)
s.Equal(expectMinTaskID, minTaskID)
s.Equal(expectMaxTaskID, maxTaskID)
s.NotEqual(sanityCheckTime, s.replicationAckManager.sanityCheckTime)
s.Equal(expectMaxTaskID, *s.replicationAckManager.maxTaskID)
}
func (s *ackManagerSuite) Test_GetMaxTaskInfo() {
now := time.Now()
taskSet := []tasks.Task{
&tasks.HistoryReplicationTask{
TaskID: 1,
VisibilityTimestamp: now,
},
&tasks.HistoryReplicationTask{
TaskID: 6,
VisibilityTimestamp: now.Add(time.Second),
},
&tasks.HistoryReplicationTask{
TaskID: 3,
VisibilityTimestamp: now.Add(time.Hour),
},
}
s.replicationAckManager.NotifyNewTasks(taskSet)
maxTaskID, maxVisibilityTimestamp := s.replicationAckManager.GetMaxTaskInfo()
s.Equal(int64(6), maxTaskID)
s.Equal(now.Add(time.Hour), maxVisibilityTimestamp)
}
func (s *ackManagerSuite) TestGetTasks_NoTasksInDB() {
ctx := context.Background()
minTaskID := int64(220878)
maxTaskID := minTaskID + 100
s.mockExecutionMgr.EXPECT().GetHistoryTasks(ctx, &persistence.GetHistoryTasksRequest{
ShardID: s.mockShard.GetShardID(),
TaskCategory: tasks.CategoryReplication,
InclusiveMinTaskKey: tasks.NewImmediateKey(minTaskID + 1),
ExclusiveMaxTaskKey: tasks.NewImmediateKey(maxTaskID + 1),
BatchSize: s.replicationAckManager.pageSize(),
NextPageToken: nil,
}).Return(s.getHistoryTasksResponse(0), nil)
replicationTasks, lastTaskID, err := s.replicationAckManager.getTasks(ctx, cluster.TestCurrentClusterName, minTaskID, maxTaskID)
s.NoError(err)
s.Empty(replicationTasks)
s.Equal(maxTaskID, lastTaskID)
}
func (s *ackManagerSuite) TestGetTasks_FirstPersistenceErrorReturnsErrorAndEmptyResult() {
ctx := context.Background()
minTaskID := int64(220878)
maxTaskID := minTaskID + 100
tasksResponse := s.getHistoryTasksResponse(2)
s.mockExecutionMgr.EXPECT().GetHistoryTasks(ctx, &persistence.GetHistoryTasksRequest{
ShardID: s.mockShard.GetShardID(),
TaskCategory: tasks.CategoryReplication,
InclusiveMinTaskKey: tasks.NewImmediateKey(minTaskID + 1),
ExclusiveMaxTaskKey: tasks.NewImmediateKey(maxTaskID + 1),
BatchSize: s.replicationAckManager.pageSize(),
NextPageToken: nil,
}).Return(tasksResponse, nil)
gweErr := serviceerror.NewUnavailable("random error")
s.mockExecutionMgr.EXPECT().GetWorkflowExecution(gomock.Any(), &persistence.GetWorkflowExecutionRequest{
ShardID: s.mockShard.GetShardID(),
NamespaceID: tasksResponse.Tasks[0].GetNamespaceID(),
WorkflowID: tasksResponse.Tasks[0].GetWorkflowID(),
RunID: tasksResponse.Tasks[0].GetRunID(),
ArchetypeID: chasm.WorkflowArchetypeID,
}).Return(nil, gweErr)
replicationTasks, lastTaskID, err := s.replicationAckManager.getTasks(ctx, cluster.TestCurrentClusterName, minTaskID, maxTaskID)
s.Error(err)
s.ErrorIs(err, gweErr)
s.Empty(replicationTasks)
s.EqualValues(0, lastTaskID)
}
func (s *ackManagerSuite) TestGetTasks_SecondPersistenceErrorReturnsPartialResult() {
ctx := context.Background()
minTaskID := int64(220878)
maxTaskID := minTaskID + 100
tasksResponse := s.getHistoryTasksResponse(2)
s.mockExecutionMgr.EXPECT().GetHistoryTasks(ctx, &persistence.GetHistoryTasksRequest{
ShardID: s.mockShard.GetShardID(),
TaskCategory: tasks.CategoryReplication,
InclusiveMinTaskKey: tasks.NewImmediateKey(minTaskID + 1),
ExclusiveMaxTaskKey: tasks.NewImmediateKey(maxTaskID + 1),
BatchSize: s.replicationAckManager.pageSize(),
NextPageToken: nil,
}).Return(tasksResponse, nil)
eventsCache := events.NewHostLevelEventsCache(
s.mockShard.GetExecutionManager(),
s.mockShard.GetConfig(),
s.mockShard.GetMetricsHandler(),
s.mockShard.GetLogger(),
false,
)
ms := workflow.TestLocalMutableState(s.mockShard, eventsCache, tests.GlobalNamespaceEntry, tests.WorkflowID, tests.RunID, log.NewTestLogger())
ei := ms.GetExecutionInfo()
ei.NamespaceId = tests.NamespaceID.String()
ei.VersionHistories = &historyspb.VersionHistories{
Histories: []*historyspb.VersionHistory{
{
Items: []*historyspb.VersionHistoryItem{
{
EventId: 1,
Version: 1,
},
},
},
},
}
s.mockExecutionMgr.EXPECT().GetWorkflowExecution(gomock.Any(), gomock.Any()).Return(&persistence.GetWorkflowExecutionResponse{
State: workflow.TestCloneToProto(context.Background(), ms)}, nil)
s.mockExecutionMgr.EXPECT().GetWorkflowExecution(gomock.Any(), gomock.Any()).Return(nil, serviceerror.NewUnavailable("some random error"))
s.mockExecutionMgr.EXPECT().ReadRawHistoryBranch(gomock.Any(), gomock.Any()).Return(&persistence.ReadRawHistoryBranchResponse{
HistoryEventBlobs: []*commonpb.DataBlob{{}}}, nil)
replicationTasks, lastTaskID, err := s.replicationAckManager.getTasks(ctx, cluster.TestCurrentClusterName, minTaskID, maxTaskID)
s.NoError(err)
s.Equal(1, len(replicationTasks))
s.Equal(tasksResponse.Tasks[0].GetTaskID(), lastTaskID)
}
func (s *ackManagerSuite) TestGetTasks_FullPage() {
tasksResponse := s.getHistoryTasksResponse(s.replicationAckManager.pageSize())
tasksResponse.NextPageToken = []byte{22, 3, 83} // There is more in DB.
minTaskID, maxTaskID := s.replicationAckManager.taskIDsRange(22)
s.mockExecutionMgr.EXPECT().GetHistoryTasks(gomock.Any(), &persistence.GetHistoryTasksRequest{
ShardID: s.mockShard.GetShardID(),
TaskCategory: tasks.CategoryReplication,
InclusiveMinTaskKey: tasks.NewImmediateKey(minTaskID + 1),
ExclusiveMaxTaskKey: tasks.NewImmediateKey(maxTaskID + 1),
BatchSize: s.replicationAckManager.pageSize(),
NextPageToken: nil,
}).Return(tasksResponse, nil)
eventsCache := events.NewHostLevelEventsCache(
s.mockShard.GetExecutionManager(),
s.mockShard.GetConfig(),
s.mockShard.GetMetricsHandler(),
s.mockShard.GetLogger(),
false,
)
ms := workflow.TestLocalMutableState(s.mockShard, eventsCache, tests.GlobalNamespaceEntry, tests.WorkflowID, tests.RunID, log.NewTestLogger())
ei := ms.GetExecutionInfo()
ei.NamespaceId = tests.NamespaceID.String()
ei.VersionHistories = &historyspb.VersionHistories{
Histories: []*historyspb.VersionHistory{
{
Items: []*historyspb.VersionHistoryItem{
{
EventId: 1,
Version: 1,
},
},
},
},
}
s.mockExecutionMgr.EXPECT().GetWorkflowExecution(gomock.Any(), gomock.Any()).Return(&persistence.GetWorkflowExecutionResponse{
State: workflow.TestCloneToProto(context.Background(), ms)}, nil).Times(s.replicationAckManager.pageSize())
s.mockExecutionMgr.EXPECT().ReadRawHistoryBranch(gomock.Any(), gomock.Any()).Return(&persistence.ReadRawHistoryBranchResponse{
HistoryEventBlobs: []*commonpb.DataBlob{{}}}, nil).Times(s.replicationAckManager.pageSize())
replicationMessages, err := s.replicationAckManager.GetTasks(context.Background(), cluster.TestCurrentClusterName, 22)
s.NoError(err)
s.NotNil(replicationMessages)
s.Len(replicationMessages.ReplicationTasks, s.replicationAckManager.pageSize())
s.Equal(tasksResponse.Tasks[len(tasksResponse.Tasks)-1].GetTaskID(), replicationMessages.LastRetrievedMessageId)
}
func (s *ackManagerSuite) TestGetTasks_PartialPage() {
numTasks := s.replicationAckManager.pageSize() / 2
tasksResponse := s.getHistoryTasksResponse(numTasks)
minTaskID, maxTaskID := s.replicationAckManager.taskIDsRange(22)
s.mockExecutionMgr.EXPECT().GetHistoryTasks(gomock.Any(), &persistence.GetHistoryTasksRequest{
ShardID: s.mockShard.GetShardID(),
TaskCategory: tasks.CategoryReplication,
InclusiveMinTaskKey: tasks.NewImmediateKey(minTaskID + 1),
ExclusiveMaxTaskKey: tasks.NewImmediateKey(maxTaskID + 1),
BatchSize: s.replicationAckManager.pageSize(),
NextPageToken: nil,
}).Return(tasksResponse, nil)
eventsCache := events.NewHostLevelEventsCache(
s.mockShard.GetExecutionManager(),
s.mockShard.GetConfig(),
s.mockShard.GetMetricsHandler(),
s.mockShard.GetLogger(),
false,
)
ms := workflow.TestLocalMutableState(s.mockShard, eventsCache, tests.GlobalNamespaceEntry, tests.WorkflowID, tests.RunID, log.NewTestLogger())
ei := ms.GetExecutionInfo()
ei.NamespaceId = tests.NamespaceID.String()
ei.VersionHistories = &historyspb.VersionHistories{
Histories: []*historyspb.VersionHistory{
{
Items: []*historyspb.VersionHistoryItem{
{
EventId: 1,
Version: 1,
},
},
},
},
}
s.mockExecutionMgr.EXPECT().GetWorkflowExecution(gomock.Any(), gomock.Any()).Return(&persistence.GetWorkflowExecutionResponse{
State: workflow.TestCloneToProto(context.Background(), ms)}, nil).Times(numTasks)
s.mockExecutionMgr.EXPECT().ReadRawHistoryBranch(gomock.Any(), gomock.Any()).Return(&persistence.ReadRawHistoryBranchResponse{
HistoryEventBlobs: []*commonpb.DataBlob{{}}}, nil).Times(numTasks)
replicationMessages, err := s.replicationAckManager.GetTasks(context.Background(), cluster.TestCurrentClusterName, 22)
s.NoError(err)
s.NotNil(replicationMessages)
s.Len(replicationMessages.ReplicationTasks, numTasks)
s.Equal(tasksResponse.Tasks[len(tasksResponse.Tasks)-1].GetTaskID(), replicationMessages.LastRetrievedMessageId)
}
func (s *ackManagerSuite) TestGetTasks_FilterNamespace() {
notExistOnTestClusterNamespaceID := namespace.ID("not-exist-on-" + cluster.TestCurrentClusterName)
notExistOnTestClusterNamespaceEntry := namespace.NewLocalNamespaceForTest(
&persistencespb.NamespaceInfo{},
&persistencespb.NamespaceConfig{},
"not-a-"+cluster.TestCurrentClusterName,
)
s.mockNamespaceRegistry.EXPECT().GetNamespaceByID(notExistOnTestClusterNamespaceID).Return(notExistOnTestClusterNamespaceEntry, nil).AnyTimes()
minTaskID, maxTaskID := s.replicationAckManager.taskIDsRange(22)
tasksResponse1 := s.getHistoryTasksResponse(s.replicationAckManager.pageSize())
// 2 of 25 tasks are for namespace that doesn't exist on poll cluster.
tasksResponse1.Tasks[1].(*tasks.HistoryReplicationTask).NamespaceID = notExistOnTestClusterNamespaceID.String()
tasksResponse1.Tasks[3].(*tasks.HistoryReplicationTask).NamespaceID = notExistOnTestClusterNamespaceID.String()
tasksResponse1.NextPageToken = []byte{22, 3, 83} // There is more in DB.
s.mockExecutionMgr.EXPECT().GetHistoryTasks(gomock.Any(), &persistence.GetHistoryTasksRequest{
ShardID: s.mockShard.GetShardID(),
TaskCategory: tasks.CategoryReplication,
InclusiveMinTaskKey: tasks.NewImmediateKey(minTaskID + 1),
ExclusiveMaxTaskKey: tasks.NewImmediateKey(maxTaskID + 1),
BatchSize: s.replicationAckManager.pageSize(),
NextPageToken: nil,
}).Return(tasksResponse1, nil)
tasksResponse2 := s.getHistoryTasksResponse(2)
// 1 of 2 task is for namespace that doesn't exist on poll cluster.
tasksResponse2.Tasks[1].(*tasks.HistoryReplicationTask).NamespaceID = notExistOnTestClusterNamespaceID.String()
tasksResponse2.NextPageToken = []byte{22, 8, 78} // There is more in DB.
s.mockExecutionMgr.EXPECT().GetHistoryTasks(gomock.Any(), &persistence.GetHistoryTasksRequest{
ShardID: s.mockShard.GetShardID(),
TaskCategory: tasks.CategoryReplication,
InclusiveMinTaskKey: tasks.NewImmediateKey(minTaskID + 1),
ExclusiveMaxTaskKey: tasks.NewImmediateKey(maxTaskID + 1),
BatchSize: s.replicationAckManager.pageSize(),
NextPageToken: []byte{22, 3, 83}, // previous token
}).Return(tasksResponse2, nil)
tasksResponse3 := s.getHistoryTasksResponse(1)
s.mockExecutionMgr.EXPECT().GetHistoryTasks(gomock.Any(), &persistence.GetHistoryTasksRequest{
ShardID: s.mockShard.GetShardID(),
TaskCategory: tasks.CategoryReplication,
InclusiveMinTaskKey: tasks.NewImmediateKey(minTaskID + 1),
ExclusiveMaxTaskKey: tasks.NewImmediateKey(maxTaskID + 1),
BatchSize: s.replicationAckManager.pageSize(),
NextPageToken: []byte{22, 8, 78}, // previous token
}).Return(tasksResponse3, nil)
eventsCache := events.NewHostLevelEventsCache(
s.mockShard.GetExecutionManager(),
s.mockShard.GetConfig(),
s.mockShard.GetMetricsHandler(),
s.mockShard.GetLogger(),
false,
)
ms := workflow.TestLocalMutableState(s.mockShard, eventsCache, tests.GlobalNamespaceEntry, tests.WorkflowID, tests.RunID, log.NewTestLogger())
ei := ms.GetExecutionInfo()
ei.NamespaceId = tests.NamespaceID.String()
ei.VersionHistories = &historyspb.VersionHistories{
Histories: []*historyspb.VersionHistory{
{
Items: []*historyspb.VersionHistoryItem{
{
EventId: 1,
Version: 1,
},
},
},
},
}
s.mockExecutionMgr.EXPECT().GetWorkflowExecution(gomock.Any(), gomock.Any()).Return(&persistence.GetWorkflowExecutionResponse{
State: workflow.TestCloneToProto(context.Background(), ms)}, nil).Times(s.replicationAckManager.pageSize())
s.mockExecutionMgr.EXPECT().ReadRawHistoryBranch(gomock.Any(), gomock.Any()).Return(&persistence.ReadRawHistoryBranchResponse{
HistoryEventBlobs: []*commonpb.DataBlob{{}}}, nil).Times(s.replicationAckManager.pageSize())
replicationMessages, err := s.replicationAckManager.GetTasks(context.Background(), cluster.TestCurrentClusterName, 22)
s.NoError(err)
s.NotNil(replicationMessages)
s.Len(replicationMessages.ReplicationTasks, s.replicationAckManager.pageSize())
s.Equal(tasksResponse3.Tasks[len(tasksResponse3.Tasks)-1].GetTaskID(), replicationMessages.LastRetrievedMessageId)
}
func (s *ackManagerSuite) getHistoryTasksResponse(size int) *persistence.GetHistoryTasksResponse {
result := &persistence.GetHistoryTasksResponse{}
for i := 1; i <= size; i++ {
result.Tasks = append(result.Tasks, &tasks.HistoryReplicationTask{
WorkflowKey: definition.WorkflowKey{
NamespaceID: tests.NamespaceID.String(),
WorkflowID: tests.WorkflowID + strconv.Itoa(i),
RunID: uuid.NewString(),
},
TaskID: int64(i),
FirstEventID: 1,
NextEventID: 1,
Version: 1,
},
)
}
return result
}