Enabling service-discovery driven shutdown of matching engine

common/dynamicconfig/constants.go

@@ -1672,6 +1672,13 @@ const (
 	// Default value: false
 	// Allowed filters: DomainID
 	MatchingEnableTaskInfoLogByDomainID
+	// MatchingEnableTasklistGuardAgainstOwnershipShardLoss
+	// enables guards to prevent tasklists from processing if there is any detection that the host
+	// no longer is active or owns the shard
+	// KeyName: matching.enableTasklistGuardAgainstOwnershipLoss
+	// Value type: Bool
+	// Default value: false
+	MatchingEnableTasklistGuardAgainstOwnershipShardLoss
 
 	// key for history
 
@@ -4102,6 +4109,11 @@ var BoolKeys = map[BoolKey]DynamicBool{
 		Description:  "MatchingEnableTaskInfoLogByDomainID is enables info level logs for decision/activity task based on the request domainID",
 		DefaultValue: false,
 	},
+	MatchingEnableTasklistGuardAgainstOwnershipShardLoss: {
+		KeyName:      "matching.enableTasklistGuardAgainstOwnershipLoss",
+		Description:  "allows guards to ensure that tasklists don't continue processing if there's signal that they've lost ownership",
+		DefaultValue: true,
+	},
 	EventsCacheGlobalEnable: {
 		KeyName:      "history.eventsCacheGlobalEnable",
 		Description:  "EventsCacheGlobalEnable is enables global cache over all history shards",

common/errors/taskListNotOwnedByHostError.go

@@ -24,21 +24,21 @@ package errors
 
 import "fmt"
 
-var _ error = &TaskListNotOwnnedByHostError{}
+var _ error = &TaskListNotOwnedByHostError{}
 
-type TaskListNotOwnnedByHostError struct {
+type TaskListNotOwnedByHostError struct {
 	OwnedByIdentity string
 	MyIdentity      string
 	TasklistName    string
 }
 
-func (m *TaskListNotOwnnedByHostError) Error() string {
+func (m *TaskListNotOwnedByHostError) Error() string {
 	return fmt.Sprintf("task list is not owned by this host: OwnedBy: %s, Me: %s, Tasklist: %s",
 		m.OwnedByIdentity, m.MyIdentity, m.TasklistName)
 }
 
-func NewTaskListNotOwnnedByHostError(ownedByIdentity string, myIdentity string, tasklistName string) *TaskListNotOwnnedByHostError {
-	return &TaskListNotOwnnedByHostError{
+func NewTaskListNotOwnedByHostError(ownedByIdentity string, myIdentity string, tasklistName string) *TaskListNotOwnedByHostError {
+	return &TaskListNotOwnedByHostError{
 		OwnedByIdentity: ownedByIdentity,
 		MyIdentity:      myIdentity,
 		TasklistName:    tasklistName,

common/log/tag/tags.go

@@ -956,6 +956,12 @@ func VisibilityQuery(query string) Tag {
 	return newStringTag("visibility-query", query)
 }
 
+// MembershipChangeEvent is a predefined tag for when logging hashring change events,
+// expected to be of type membership.ChangeEvent
+func MembershipChangeEvent(event interface{}) Tag {
+	return newPredefinedDynamicTag("membership-change-event", event)
+}
+
 // Dynamic Uses reflection based logging for arbitrary values
 // for not very performant logging
 func Dynamic(key string, v interface{}) Tag {

service/matching/config/config.go

@@ -79,6 +79,8 @@ type (
 		TaskDispatchRPSTTL time.Duration
 		// task gc configuration
 		MaxTimeBetweenTaskDeletes time.Duration
+
+		EnableTasklistOwnershipGuard dynamicconfig.BoolPropertyFn
 	}
 
 	ForwarderConfig struct {
@@ -154,6 +156,7 @@ func NewConfig(dc *dynamicconfig.Collection, hostName string) *Config {
 		EnableTasklistIsolation:         dc.GetBoolPropertyFilteredByDomain(dynamicconfig.EnableTasklistIsolation),
 		AllIsolationGroups:              mapIGs(dc.GetListProperty(dynamicconfig.AllIsolationGroups)()),
 		AsyncTaskDispatchTimeout:        dc.GetDurationPropertyFilteredByTaskListInfo(dynamicconfig.AsyncTaskDispatchTimeout),
+		EnableTasklistOwnershipGuard:    dc.GetBoolProperty(dynamicconfig.MatchingEnableTasklistGuardAgainstOwnershipShardLoss),
 		HostName:                        hostName,
 		TaskDispatchRPS:                 100000.0,
 		TaskDispatchRPSTTL:              time.Minute,

service/matching/handler/engine.go

@@ -78,6 +78,7 @@ type (
 	}
 
 	matchingEngineImpl struct {
+		shutdown             chan struct{}
 		taskManager          persistence.TaskManager
 		clusterMetadata      cluster.Metadata
 		historyService       history.Client
@@ -120,7 +121,8 @@ var (
 var _ Engine = (*matchingEngineImpl)(nil) // Asserts that interface is indeed implemented
 
 // NewEngine creates an instance of matching engine
-func NewEngine(taskManager persistence.TaskManager,
+func NewEngine(
+	taskManager persistence.TaskManager,
 	clusterMetadata cluster.Metadata,
 	historyService history.Client,
 	matchingClient matching.Client,
@@ -132,7 +134,9 @@ func NewEngine(taskManager persistence.TaskManager,
 	partitioner partition.Partitioner,
 	timeSource clock.TimeSource,
 ) Engine {
+
 	e := &matchingEngineImpl{
+		shutdown:             make(chan struct{}),
 		taskManager:          taskManager,
 		clusterMetadata:      clusterMetadata,
 		historyService:       historyService,
@@ -149,15 +153,20 @@ func NewEngine(taskManager persistence.TaskManager,
 		partitioner:          partitioner,
 		timeSource:           timeSource,
 	}
+
+	go e.subscribeToMembershipChanges()
+
 	e.waitForQueryResultFn = e.waitForQueryResult
 	return e
 }
 
 func (e *matchingEngineImpl) Start() {
-	// As task lists are initialized lazily nothing is done on startup at this point.
+	// reset the shutdown channel if there's any listeners
+	e.shutdown = make(chan struct{})
 }
 
 func (e *matchingEngineImpl) Stop() {
+	close(e.shutdown)
 	// Executes Stop() on each task list outside of lock
 	for _, l := range e.getTaskLists(math.MaxInt32) {
 		l.Stop()
@@ -200,26 +209,9 @@ func (e *matchingEngineImpl) getTaskListManager(taskList *tasklist.Identifier, t
 	}
 	e.taskListsLock.RUnlock()
 
-	// Defensive check to make sure we actually own the task list
-	//   If we try to create a task list manager for a task list that is not owned by us, return an error
-	//   The new task list manager will steal the task list from the current owner, which should only happen if
-	//   the task list is owned by the current host.
-	taskListOwner, err := e.membershipResolver.Lookup(service.Matching, taskList.GetName())
-	if err != nil {
-		return nil, fmt.Errorf("failed to lookup task list owner: %w", err)
-	}
-
-	self, err := e.membershipResolver.WhoAmI()
+	err := e.guardAgainstShardLoss(taskList)
 	if err != nil {
-		return nil, fmt.Errorf("failed to lookup self im membership: %w", err)
-	}
-
-	if taskListOwner.Identity() != self.Identity() {
-		return nil, cadence_errors.NewTaskListNotOwnnedByHostError(
-			taskListOwner.Identity(),
-			self.Identity(),
-			taskList.GetName(),
-		)
+		return nil, err
 	}
 
 	// If it gets here, write lock and check again in case a task list is created between the two locks
@@ -1202,6 +1194,132 @@ func (e *matchingEngineImpl) emitInfoOrDebugLog(
 	}
 }
 
+func (e *matchingEngineImpl) shutDownNonOwnedTasklists() error {
+	if !e.config.EnableTasklistOwnershipGuard() {
+		return nil
+	}
+	noLongerOwned, err := e.getNonOwnedTasklistsLocked()
+	if err != nil {
+		return err
+	}
+
+	for _, tl := range noLongerOwned {
+		// for each of the tasklists that are no longer owned, kick off the
+		// process of stopping them. The stopping process is IO heavy and
+		// can take a while, so do them in parallel so as to not hold the
+		// lock too long.
+		go func(tl tasklist.Manager) {
+
+			defer func() {
+				if r := recover(); r != nil {
+					e.logger.Error("panic occurred while trying to shut down tasklist", tag.Dynamic("recovered-panic", r))
+				}
+			}()
+
+			e.removeTaskListManager(tl)
+
+			e.logger.Info("shutting down tasklist preemptively because they are no longer owned by this host",
+				tag.WorkflowTaskListName(tl.TaskListID().GetName()),
+				tag.WorkflowDomainID(tl.TaskListID().GetDomainID()),
+				tag.Dynamic("tasklist-debug-info", tl.String()),
+			)
+
+			tl.Stop()
+		}(tl)
+	}
+	return nil
+}
+
+func (e *matchingEngineImpl) getNonOwnedTasklistsLocked() ([]tasklist.Manager, error) {
+	if !e.config.EnableTasklistOwnershipGuard() {
+		return nil, nil
+	}
+
+	var toShutDown []tasklist.Manager
+
+	e.taskListsLock.RLock()
+	defer e.taskListsLock.RUnlock()
+
+	self, err := e.membershipResolver.WhoAmI()
+	if err != nil {
+		return nil, fmt.Errorf("failed to lookup self im membership: %w", err)
+	}
+
+	for tl, manager := range e.taskLists {
+		taskListOwner, err := e.membershipResolver.Lookup(service.Matching, tl.GetName())
+		if err != nil {
+			return nil, fmt.Errorf("failed to lookup task list owner: %w", err)
+		}
+
+		if taskListOwner.Identity() != self.Identity() {
+			toShutDown = append(toShutDown, manager)
+		}
+	}
+
+	e.logger.Info("Got list of non-owned-tasklists",
+		tag.Dynamic("tasklist-debug-info", toShutDown),
+	)
+	return toShutDown, nil
+}
+
+func (e *matchingEngineImpl) guardAgainstShardLoss(taskList *tasklist.Identifier) error {
+	if !e.config.EnableTasklistOwnershipGuard() {
+		return nil
+	}
+
+	self, err := e.membershipResolver.WhoAmI()
+	if err != nil {
+		return fmt.Errorf("failed to lookup self im membership: %w", err)
+	}
+
+	if e.isShuttingDown() {
+		e.logger.Warn("request to get tasklist is being rejected be cause engine is Host is shut down",
+			tag.WorkflowDomainID(taskList.GetDomainID()),
+			tag.WorkflowTaskListType(taskList.GetType()),
+			tag.WorkflowTaskListName(taskList.GetName()),
+		)
+
+		return cadence_errors.NewTaskListNotOwnedByHostError(
+			"not known",
+			self.Identity(),
+			taskList.GetName(),
+		)
+	}
+
+	// Defensive check to make sure we actually own the task list
+	//   If we try to create a task list manager for a task list that is not owned by us, return an error
+	//   The new task list manager will steal the task list from the current owner, which should only happen if
+	//   the task list is owned by the current host.
+	taskListOwner, err := e.membershipResolver.Lookup(service.Matching, taskList.GetName())
+	if err != nil {
+		return fmt.Errorf("failed to lookup task list owner: %w", err)
+	}
+
+	if taskListOwner.Identity() != self.Identity() {
+		e.logger.Warn("Request to get tasklist is being rejected be cause engine does not own this shard",
+			tag.WorkflowDomainID(taskList.GetDomainID()),
+			tag.WorkflowTaskListType(taskList.GetType()),
+			tag.WorkflowTaskListName(taskList.GetName()),
+		)
+		return cadence_errors.NewTaskListNotOwnedByHostError(
+			taskListOwner.Identity(),
+			self.Identity(),
+			taskList.GetName(),
+		)
+	}
+
+	return nil
+}
+
+func (e *matchingEngineImpl) isShuttingDown() bool {
+	select {
+	case <-e.shutdown:
+		return true
+	default:
+		return false
+	}
+}
+
 func (m *lockableQueryTaskMap) put(key string, value chan *queryResult) {
 	m.Lock()
 	defer m.Unlock()

service/matching/handler/engine_integration_test.go

@@ -1317,7 +1317,7 @@ func (s *matchingEngineSuite) TestGetTaskListManager_OwnerShip() {
 			name:          "Not owned by current host",
 			lookUpResult:  "A",
 			whoAmIResult:  "B",
-			expectedError: new(cadence_errors.TaskListNotOwnnedByHostError),
+			expectedError: new(cadence_errors.TaskListNotOwnedByHostError),
 		},
 		{
 			name:          "LookupError",