Add a hook point to the Matching service once a task has been matched (or not)

This adds a hook point to the matching service once it has matched a task or decided to spool it. This can be leveraged to react to task sync or no-sync match events, while keeping matching service decoupled from the event handling logic.

Author: 02strich

service/matching/handler.go

@@ -23,6 +23,7 @@ import (
 	"go.temporal.io/server/common/searchattribute"
 	"go.temporal.io/server/common/testing/testhooks"
 	"go.temporal.io/server/common/tqid"
+	"go.temporal.io/server/service/matching/hooks"
 	"go.temporal.io/server/service/matching/workers"
 	"go.temporal.io/server/service/worker/workerdeployment"
 	"go.uber.org/fx"
@@ -70,6 +71,7 @@ type (
 		RateLimiter                   TaskDispatchRateLimiter `optional:"true"`
 		WorkersRegistry               workers.Registry
 		Serializer                    serialization.Serializer
+		TaskMatchHooks                []hooks.TaskMatchHook `group:"TaskMatchHooks"`
 	}
 )
 
@@ -112,6 +114,7 @@ func NewHandler(
 			params.SearchAttributeMapperProvider,
 			params.RateLimiter,
 			params.Serializer,
+			params.TaskMatchHooks,
 		),
 		namespaceRegistry: params.NamespaceRegistry,
 		workersRegistry:   params.WorkersRegistry,

service/matching/hooks/matching_hooks.go

@@ -0,0 +1,24 @@
+package hooks
+
+import (
+	"context"
+
+	deploymentpb "go.temporal.io/api/deployment/v1"
+	enumspb "go.temporal.io/api/enums/v1"
+	taskqueuepb "go.temporal.io/api/taskqueue/v1"
+	"go.temporal.io/server/common/namespace"
+)
+
+type (
+	TaskMatchHookDetails struct {
+		Namespace               *namespace.Namespace
+		TaskQueueName           string
+		TaskQueueType           enumspb.TaskQueueType
+		DeploymentVersion       *deploymentpb.WorkerDeploymentVersion
+		IsSyncMatch             bool
+		TaskQueueStatsRetriever func() map[int32]*taskqueuepb.TaskQueueStats
+	}
+	TaskMatchHook interface {
+		ProcessTaskMatch(ctx context.Context, event *TaskMatchHookDetails)
+	}
+)

service/matching/matching_engine.go

@@ -60,6 +60,7 @@ import (
 	"go.temporal.io/server/common/util"
 	"go.temporal.io/server/common/worker_versioning"
 	"go.temporal.io/server/service/history/api"
+	"go.temporal.io/server/service/matching/hooks"
 	"go.temporal.io/server/service/worker/workerdeployment"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -175,6 +176,8 @@ type (
 		reachabilityCache reachabilityCache
 		// Rate limiter to limit the task dispatch
 		rateLimiter TaskDispatchRateLimiter
+
+		taskMatchHooks []hooks.TaskMatchHook
 	}
 )
 
@@ -254,6 +257,7 @@ func NewEngine(
 	saMapperProvider searchattribute.MapperProvider,
 	rateLimiter TaskDispatchRateLimiter,
 	historySerializer serialization.Serializer,
+	taskMatchHooks []hooks.TaskMatchHook,
 ) Engine {
 	scopedMetricsHandler := metricsHandler.WithTags(metrics.OperationTag(metrics.MatchingEngineScope))
 	e := &matchingEngineImpl{
@@ -296,6 +300,7 @@ func NewEngine(
 		namespaceReplicationQueue: namespaceReplicationQueue,
 		userDataUpdateBatchers:    collection.NewSyncMap[namespace.ID, *stream_batcher.Batcher[*userDataUpdate, error]](),
 		rateLimiter:               rateLimiter,
+		taskMatchHooks:            taskMatchHooks,
 	}
 	e.reachabilityCache = newReachabilityCache(
 		metrics.NoopMetricsHandler,
@@ -497,6 +502,7 @@ func (e *matchingEngineImpl) getTaskQueuePartitionManager(
 		throttledLogger,
 		metricsHandler,
 		userDataManager,
+		e.taskMatchHooks,
 	)
 	if err != nil {
 		return nil, false, err

service/matching/matching_engine_test.go

@@ -272,7 +272,7 @@ func newMatchingEngine(
 func (s *matchingEngineSuite) newPartitionManager(prtn tqid.Partition, config *Config) taskQueuePartitionManager {
 	tqConfig := newTaskQueueConfig(prtn.TaskQueue(), config, matchingTestNamespace)
 	logger, _, metricsHandler := s.matchingEngine.loggerAndMetricsForPartition(s.ns, prtn, tqConfig)
-	pm, err := newTaskQueuePartitionManager(s.matchingEngine, s.ns, prtn, tqConfig, logger, logger, metricsHandler, &mockUserDataManager{})
+	pm, err := newTaskQueuePartitionManager(s.matchingEngine, s.ns, prtn, tqConfig, logger, logger, metricsHandler, &mockUserDataManager{}, nil)
 	s.Require().NoError(err)
 	return pm
 }

service/matching/physical_task_queue_manager_test.go

@@ -82,7 +82,7 @@ func (s *PhysicalTaskQueueManagerTestSuite) SetupTest() {
 	onFatalErr := func(unloadCause) { s.T().Fatal("user data manager called onFatalErr") }
 	udMgr := newUserDataManager(engine.taskManager, engine.matchingRawClient, onFatalErr, nil, nil, prtn, tqConfig, engine.logger, engine.namespaceRegistry)
 
-	prtnMgr, err := newTaskQueuePartitionManager(engine, ns, prtn, tqConfig, engine.logger, nil, metrics.NoopMetricsHandler, udMgr)
+	prtnMgr, err := newTaskQueuePartitionManager(engine, ns, prtn, tqConfig, engine.logger, nil, metrics.NoopMetricsHandler, udMgr, nil)
 	s.NoError(err)
 	engine.partitions[prtn.Key()] = prtnMgr
 
@@ -533,7 +533,7 @@ func TestDrainCompletionNoReloadDraining(t *testing.T) {
 	onFatalErr := func(unloadCause) { t.Fatal("user data manager called onFatalErr") }
 	udMgr := newUserDataManager(engine.taskManager, engine.matchingRawClient, onFatalErr, nil, nil, prtn, tqConfig, engine.logger, engine.namespaceRegistry)
 
-	prtnMgr, err := newTaskQueuePartitionManager(engine, ns, prtn, tqConfig, engine.logger, nil, metrics.NoopMetricsHandler, udMgr)
+	prtnMgr, err := newTaskQueuePartitionManager(engine, ns, prtn, tqConfig, engine.logger, nil, metrics.NoopMetricsHandler, udMgr, nil)
 	require.NoError(t, err)
 	engine.partitions[prtn.Key()] = prtnMgr
 
@@ -578,7 +578,7 @@ func TestDrainCompletionNoReloadDraining(t *testing.T) {
 	prevPriStats, prevFairStats := priQueueData.persistenceStats(), fairQueueData.persistenceStats()
 
 	// create a new manager (reload)
-	prtnMgr2, err := newTaskQueuePartitionManager(engine, ns, prtn, tqConfig, engine.logger, nil, metrics.NoopMetricsHandler, udMgr)
+	prtnMgr2, err := newTaskQueuePartitionManager(engine, ns, prtn, tqConfig, engine.logger, nil, metrics.NoopMetricsHandler, udMgr, nil)
 	require.NoError(t, err)
 	engine.partitions[prtn.Key()] = prtnMgr2
 

service/matching/task_queue_partition_manager.go

@@ -36,6 +36,7 @@ import (
 	"go.temporal.io/server/common/tqid"
 	"go.temporal.io/server/common/util"
 	"go.temporal.io/server/common/worker_versioning"
+	"go.temporal.io/server/service/matching/hooks"
 	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -74,6 +75,8 @@ type (
 		// TODO(stephanos): move cache out of partition manager
 		cache cache.Cache // non-nil for root-partition
 
+		taskMatchHooks []hooks.TaskMatchHook
+
 		goroGroup goro.Group
 
 		autoEnableRateLimiter quotas.RateLimiter
@@ -113,6 +116,7 @@ func newTaskQueuePartitionManager(
 	throttledLogger log.Logger,
 	metricsHandler metrics.Handler,
 	userDataManager userDataManager,
+	taskMatchHooks []hooks.TaskMatchHook,
 ) (*taskQueuePartitionManagerImpl, error) {
 	rateLimitManager := newRateLimitManager(
 		userDataManager,
@@ -132,6 +136,7 @@ func newTaskQueuePartitionManager(
 		rateLimitManager:      rateLimitManager,
 		defaultQueueFuture:    future.NewFuture[physicalTaskQueueManager](),
 		autoEnableRateLimiter: quotas.NewRateLimiter(1.0/60, 1),
+		taskMatchHooks:        taskMatchHooks,
 	}
 	pm.initCtx, pm.initCancel = context.WithCancel(context.Background())
 
@@ -357,6 +362,20 @@ reredirectTask:
 	if isActive {
 		syncMatched, err = syncMatchQueue.TrySyncMatch(ctx, syncMatchTask)
 		if syncMatched && !pm.shouldBacklogSyncMatchTaskOnError(err) {
+			for _, l := range pm.taskMatchHooks {
+				var deploymentVersion *deploymentpb.WorkerDeploymentVersion
+				if targetVersion != nil {
+					deploymentVersion = &deploymentpb.WorkerDeploymentVersion{DeploymentName: targetVersion.DeploymentName, BuildId: targetVersion.BuildId}
+				}
+				l.ProcessTaskMatch(ctx, &hooks.TaskMatchHookDetails{
+					Namespace:               pm.ns,
+					TaskQueueName:           pm.partition.TaskQueue().Name(),
+					TaskQueueType:           pm.partition.TaskType(),
+					DeploymentVersion:       deploymentVersion,
+					IsSyncMatch:             syncMatched,
+					TaskQueueStatsRetriever: func() map[int32]*taskqueuepb.TaskQueueStats { return dbq.GetStatsByPriority(true) },
+				})
+			}
 
 			// Build ID is not returned for sync match. The returned build ID is used by History to update
 			// mutable state (and visibility) when the first workflow task is spooled.
@@ -382,6 +401,21 @@ reredirectTask:
 		assignedBuildId = spoolQueue.QueueKey().Version().BuildId()
 	}
 
+	for _, l := range pm.taskMatchHooks {
+		var deploymentVersion *deploymentpb.WorkerDeploymentVersion
+		if targetVersion != nil {
+			deploymentVersion = &deploymentpb.WorkerDeploymentVersion{DeploymentName: targetVersion.DeploymentName, BuildId: targetVersion.BuildId}
+		}
+		l.ProcessTaskMatch(ctx, &hooks.TaskMatchHookDetails{
+			Namespace:               pm.ns,
+			TaskQueueName:           pm.partition.TaskQueue().Name(),
+			TaskQueueType:           pm.partition.TaskType(),
+			DeploymentVersion:       deploymentVersion,
+			IsSyncMatch:             false,
+			TaskQueueStatsRetriever: func() map[int32]*taskqueuepb.TaskQueueStats { return dbq.GetStatsByPriority(true) },
+		})
+	}
+
 	return assignedBuildId, false, spoolQueue.SpoolTask(params.taskInfo)
 }
 
@@ -960,7 +994,7 @@ func (pm *taskQueuePartitionManagerImpl) Describe(
 			unversionedCurrentShareByPriority, unversionedRampingShareByPriority =
 				splitStatsByPriorityByRampPercentage(unversionedStatsByPriority, rampPercentage)
 		} else if currentExists {
-			// If there exist no ramping version, weattribute the entire unversioned backlog to the current version.
+			// If there exist no ramping version, we attribute the entire unversioned backlog to the current version.
 			unversionedCurrentShareByPriority = cloneStatsByPriority(unversionedStatsByPriority)
 		}
 	}