temporalio · wxing1292 · Mar 29, 2022 · Mar 28, 2022 · Mar 28, 2022
diff --git a/service/history/historyEngine.go b/service/history/historyEngine.go
@@ -2504,7 +2504,7 @@ func (e *historyEngineImpl) ResetWorkflowExecution(
 	}
 
 	// also load the current run of the workflow, it can be different from the base runID
-	resp, err := e.executionManager.GetCurrentExecution(ctx, &persistence.GetCurrentExecutionRequest{
+	resp, err := e.shard.GetCurrentExecution(ctx, &persistence.GetCurrentExecutionRequest{
 		ShardID:     e.shard.GetShardID(),
 		NamespaceID: namespaceID.St
10000
ring(),
 		WorkflowID:  request.WorkflowExecution.GetWorkflowId(),
@@ -3596,7 +3596,7 @@ func (e *historyEngineImpl) loadWorkflow(
 		}
 
 		// workflow not running, need to check current record
-		resp, err := e.shard.GetExecutionManager().GetCurrentExecution(
+		resp, err := e.shard.GetCurrentExecution(
 			ctx,
 			&persistence.GetCurrentExecutionRequest{
 				ShardID:     e.shard.GetShardID(),

diff --git a/service/history/nDCTransactionMgr.go b/service/history/nDCTransactionMgr.go
@@ -398,7 +398,7 @@ func (r *nDCTransactionMgrImpl) checkWorkflowExists(
 	runID string,
 ) (bool, error) {
 
-	_, err := r.shard.GetExecutionManager().GetWorkflowExecution(
+	_, err := r.shard.GetWorkflowExecution(
 		ctx,
 		&persistence.GetWorkflowExecutionRequest{
 			ShardID:     r.shard.GetShardID(),
@@ -424,7 +424,7 @@ func (r *nDCTransactionMgrImpl) getCurrentWorkflowRunID(
 	workflowID string,
 ) (string, error) {
 
-	resp, err := r.shard.GetExecutionManager().GetCurrentExecution(
+	resp, err := r.shard.GetCurrentExecution(
 		ctx,
 		&persistence.GetCurrentExecutionRequest{
 			ShardID:     r.shard.GetShardID(),

diff --git a/service/history/shard/context.go b/service/history/shard/context.go
@@ -102,6 +102,8 @@ type (
 		UpdateWorkflowExecution(ctx context.Context, request *persistence.UpdateWorkflowExecutionRequest) (*persistence.UpdateWorkflowExecutionResponse, error)
 		ConflictResolveWorkflowExecution(ctx context.Context, request *persistence.ConflictResolveWorkflowExecutionRequest) (*persistence.ConflictResolveWorkflowExecutionResponse, error)
 		SetWorkflowExecution(ctx context.Context, request *persistence.SetWorkflowExecutionRequest) (*persistence.SetWorkflowExecutionResponse, error)
+		GetCurrentExecution(ctx context.Context, request *persistence.GetCurrentExecutionRequest) (*persistence.GetCurrentExecutionResponse, error)
+		GetWorkflowExecution(ctx context.Context, request *persistence.GetWorkflowExecutionRequest) (*persistence.GetWorkflowExecutionResponse, error)
 		// DeleteWorkflowExecution deletes workflow execution, current workflow execution, and add task to delete visibility.
 		// If branchToken != nil, then delete history also, otherwise leave history.
 		DeleteWorkflowExecution(ctx context.Context, workflowKey definition.WorkflowKey, branchToken []byte, version int64, startTime *time.Time, closeTime *time.Time) error

diff --git a/service/history/shard/context_impl.go b/service/history/shard/context_impl.go
@@ -34,14 +34,14 @@ import (
 	"go.temporal.io/api/enums/v1"
 	"go.temporal.io/api/serviceerror"
 
-	"go.temporal.io/server/api/adminservice/v1"
-	"go.temporal.io/server/api/historyservice/v1"
 	"go.temporal.io/server/client"
 	"go.temporal.io/server/common/archiver"
 	"go.temporal.io/server/common/cluster"
 	"go.temporal.io/server/common/persistence/serialization"
 	"go.temporal.io/server/common/searchattribute"
 
+	"go.temporal.io/server/api/adminservice/v1"
+	"go.temporal.io/server/api/historyservice/v1"
 	persistencespb "go.temporal.io/server/api/persistence/v1"
 	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/backoff"
@@ -599,7 +599,7 @@ func (s *ContextImpl) CreateWorkflowExecution(
 	currentRangeID := s.getRangeIDLocked()
 	request.RangeID = currentRangeID
 	resp, err := s.executionManager.CreateWorkflowExecution(ctx, request)
-	if err = s.handleErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
+	if err = s.handleWriteErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
 		return nil, err
 	}
 	return resp, nil
@@ -647,7 +647,7 @@ func (s *ContextImpl) UpdateWorkflowExecution(
 	currentRangeID := s.getRangeIDLocked()
 	request.RangeID = currentRangeID
 	resp, err := s.executionManager.UpdateWorkflowExecution(ctx, request)
-	if err = s.handleErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
+	if err = s.handleWriteErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
 		return nil, err
 	}
 	return resp, nil
@@ -705,7 +705,7 @@ func (s *ContextImpl) ConflictResolveWorkflowExecution(
 	currentRangeID := s.getRangeIDLocked()
 	request.RangeID = currentRangeID
 	resp, err := s.executionManager.ConflictResolveWorkflowExecution(ctx, request)
-	if err = s.handleErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
+	if err = s.handleWriteErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
 		return nil, err
 	}
 	return resp, nil
@@ -743,7 +743,29 @@ func (s *ContextImpl) SetWorkflowExecution(
 	currentRangeID := s.getRangeIDLocked()
 	request.RangeID = currentRangeID
 	resp, err := s.executionManager.SetWorkflowExecution(ctx, request)
-	if err = s.handleErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
+	if err = s.handleWriteErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
+		return nil, err
+	}
+	return resp, nil
+}
+
+func (s *ContextImpl) GetCurrentExecution(
+	ctx context.Context,
+	request *persistence.GetCurrentExecutionRequest,
+) (*persistence.GetCurrentExecutionResponse, error) {
+	resp, err := s.executionManager.GetCurrentExecution(ctx, request)
+	if err = s.handleReadError(err); err != nil {
+		return nil, err
+    	}
+	return resp, nil
+}
+
+func (s *ContextImpl) GetWorkflowExecution(
+	ctx context.Context,
+	request *persistence.GetWorkflowExecutionRequest,
+) (*persistence.GetWorkflowExecutionResponse, error) {
+	resp, err := s.executionManager.GetWorkflowExecution(ctx, request)
+	if err = s.handleReadError(err); err != nil {
 		return nil, err
 	}
 	return resp, nil
@@ -766,7 +788,7 @@ func (s *ContextImpl) addTasksLocked(
 
 	request.RangeID = s.getRangeIDLocked()
 	err := s.executionManager.AddHistoryTasks(ctx, request)
-	if err = s.handleErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
+	if err = s.handleWriteErrorAndUpdateMaxReadLevelLocked(err, transferMaxReadLevel); err != nil {
 		return err
 	}
 	s.engine.NotifyNewTasks(namespaceEntry.ActiveClusterName(), request.Tasks)
@@ -1006,7 +1028,7 @@ func (s *ContextImpl) renewRangeLocked(isStealing bool) error {
 			tag.ShardRangeID(updatedShardInfo.GetRangeId()),
 			tag.PreviousShardRangeID(s.shardInfo.GetRangeId()),
 		)
-		return s.handleErrorLocked(err)
+		return s.handleWriteErrorLocked(err)
 	}
 
 	// Range is successfully updated in cassandra now update shard context to reflect new range
@@ -1050,7 +1072,7 @@ func (s *ContextImpl) updateShardInfoLocked() error {
 		PreviousRangeID: s.shardInfo.GetRangeId(),
 	})
 	if err != nil {
-		return s.handleErrorLocked(err)
+		return s.handleWriteErrorLocked(err)
 	}
 
 	s.lastUpdated = now
@@ -1212,12 +1234,28 @@ func (s *ContextImpl) GetLastUpdatedTime() time.Time {
 	return s.lastUpdated
 }
 
-func (s *ContextImpl) handleErrorLocked(err error) error {
+func (s *ContextImpl) handleReadError(err error) error {
+	switch err.(type) {
+	case nil:
+		return nil
+
+	case *persistence.ShardOwnershipLostError:
+		// Shard is stolen, trigger shutdown of history engine.
+		// Handling of max read level doesn't matter here.
+		s.Unload()
+		return err
+
+	default:
+		return err
+	}
+}
+
+func (s *ContextImpl) handleWriteErrorLocked(err error) error {
 	// We can use 0 here since updateMaxReadLevelLocked ensures that the read level never goes backwards.
-	return s.handleErrorAndUpdateMaxReadLevelLocked(err, 0)
+	return s.handleWriteErrorAndUpdateMaxReadLevelLocked(err, 0)
 }
 
-func (s *ContextImpl) handleErrorAndUpdateMaxReadLevelLocked(err error, newMaxReadLevel int64) error {
+func (s *ContextImpl) handleWriteErrorAndUpdateMaxReadLevelLocked(err error, newMaxReadLevel int64) error {
 	switch err.(type) {
 	case nil:
 		// Persistence success: update max read level

diff --git a/service/history/shard/context_mock.go b/service/history/shard/context_mock.go
diff --git a/service/history/transferQueueActiveTaskExecutor.go b/service/history/transferQueueActiveTaskExecutor.go
@@ -740,7 +740,7 @@ func (t *transferQueueActiveTaskExecutor) processResetWorkflow(
 	if !currentMutableState.IsWorkflowExecutionRunning() {
 		// it means this this might not be current anymore, we need to check
 		var resp *persistence.GetCurrentExecutionResponse
-		resp, err = t.shard.GetExecutionManager().GetCurrentExecution(ctx, &persistence.GetCurrentExecutionRequest{
+		resp, err = t.shard.GetCurrentExecution(ctx, &persistence.GetCurrentExecutionRequest{
 			ShardID:     t.shard.GetShardID(),
 			NamespaceID: task.NamespaceID,
 			WorkflowID:  task.WorkflowID,

diff --git a/service/history/workflow/cache.go b/service/history/workflow/cache.go
@@ -69,11 +69,10 @@ type (
 
 	CacheImpl struct {
 		cache.Cache
-		shard            shard.Context
-		executionManager persistence.ExecutionManager
-		logger           log.Logger
-		metricsClient    metrics.Client
-		config           *configs.Config
+		shard         shard.Context
+		logger        log.Logger
+		metricsClient metrics.Client
+		config        *configs.Config
 	}
 
 	NewCacheFn func(shard shard.Context) Cache
@@ -94,12 +93,11 @@ func NewCache(shard shard.Context) Cache {
 	opts.Pin = true
 
 	return &CacheImpl{
-		Cache:            cache.New(config.HistoryCacheMaxSize(), opts),
-		shard:            shard,
-		executionManager: shard.GetExecutionManager(),
-		logger:           log.With(shard.GetLogger(), tag.ComponentHistoryCache),
-		metricsClient:    shard.GetMetricsClient(),
-		config:           config,
+		Cache:         cache.New(config.HistoryCacheMaxSize(), opts),
+		shard:         shard,
+		logger:        log.With(shard.GetLogger(), tag.ComponentHistoryCache),
+		metricsClient: shard.GetMetricsClient(),
+		config:        config,
 	}
 }
 
@@ -268,7 +266,7 @@ func (c *CacheImpl) getCurrentExecutionWithRetry(
 	var response *persistence.GetCurrentExecutionResponse
 	op := func() error {
 		var err error
-		response, err = c.executionManager.GetCurrentExecution(ctx, request)
+		response, err = c.shard.GetCurrentExecution(ctx, request)
 
 		return err
 	}

diff --git a/service/history/workflow/transaction_impl.go b/service/history/workflow/transaction_impl.go
@@ -455,7 +455,7 @@ func getWorkflowExecutionWithRetry(
 	var resp *persistence.GetWorkflowExecutionResponse
 	op := func() error {
 		var err error
-		resp, err = shard.GetExecutionManager().GetWorkflowExecution(context.TODO(), request)
+		resp, err = shard.GetWorkflowExecution(context.TODO(), request)
 
 		return err
 	}

diff --git a/service/history/workflowRebuilder.go b/service/history/workflowRebuilder.go
@@ -56,7 +56,6 @@ type (
 	workflowRebuilderImpl struct {
 		shard             shard.Context
 		workflowCache     workflow.Cache
-		executionMgr      persistence.ExecutionManager
 		newStateRebuilder nDCStateRebuilderProvider
 		transaction       workflow.Transaction
 		logger            log.Logger
@@ -73,7 +72,6 @@ func NewWorkflowRebuilder(
 	return &workflowRebuilderImpl{
 		shard:         shard,
 		workflowCache: workflowCache,
-		executionMgr:  shard.GetExecutionManager(),
 		newStateRebuilder: func() nDCStateRebuilder {
 			return newNDCStateRebuilder(shard, logger)
 		},
@@ -195,7 +193,7 @@ func (r *workflowRebuilderImpl) getMutableState(
 	ctx context.Context,
 	workflowKey definition.WorkflowKey,
 ) (*persistencespb.WorkflowMutableState, int64, error) {
-	record, err := r.executionMgr.GetWorkflowExecution(ctx, &persistence.GetWorkflowExecutionRequest{
+	record, err := r.shard.GetWorkflowExecution(ctx, &persistence.GetWorkflowExecutionRequest{
 		ShardID:     r.shard.GetShardID(),
 		NamespaceID: workflowKey.NamespaceID,
 		WorkflowID:  workflowKey.WorkflowID,