TaskGroups for chasm tasks

chasm/lib/nexusoperation/library.go

@@ -51,12 +51,20 @@ func (l *Library) Components() []*chasm.RegistrableComponent {
 
 func (l *Library) Tasks() []*chasm.RegistrableTask {
 	return []*chasm.RegistrableTask{
-		chasm.NewRegistrableSideEffectTask("invocation", l.operationInvocationTaskHandler),
+		chasm.NewRegistrableSideEffectTask(
+			"invocation",
+			l.operationInvocationTaskHandler,
+			chasm.WithSideEffectTaskGroup(TaskGroupName),
+		),
 		chasm.NewRegistrablePureTask("invocationBackoff", l.operationBackoffTaskHandler),
 		chasm.NewRegistrablePureTask("scheduleToStartTimeout", l.operationScheduleToStartTimeoutTaskHandler),
 		chasm.NewRegistrablePureTask("startToCloseTimeout", l.operationStartToCloseTimeoutTaskHandler),
 		chasm.NewRegistrablePureTask("scheduleToCloseTimeout", l.operationScheduleToCloseTimeoutTaskHandler),
-		chasm.NewRegistrableSideEffectTask("cancellation", l.cancellationInvocationTaskHandler),
+		chasm.NewRegistrableSideEffectTask(
+			"cancellation",
+			l.cancellationInvocationTaskHandler,
+			chasm.WithSideEffectTaskGroup(TaskGroupName),
+		),
 		chasm.NewRegistrablePureTask("cancellationBackoff", l.cancellationBackoffTaskHandler),
 	}
 }

chasm/lib/nexusoperation/operation.go

@@ -32,6 +32,9 @@ type InvocationData struct {
 	NexusLink nexus.Link
 }
 
+// TaskGroupName groups invocation and cancellation together for the outbound queue
+const TaskGroupName = "nexus"
+
 // OperationStore defines the interface that must be implemented by any parent component that wants to manage Nexus operations.
 // It's the responsibility of the parrent component to apply the appropriate state transitions to the operation.
 type OperationStore interface {

chasm/registrable_task.go

@@ -16,6 +16,7 @@ type (
 		sideEffectTaskExecuteFn sideEffectTaskExecuteFn
 		sideEffectTaskDiscardFn sideEffectTaskDiscardFn
 		isPureTask              bool
+		outboundTaskGroup       string // For grouping on the outbound queue. See [WithSideEffectTaskGroup] for details.
 
 		// Those two fields are initialized when the component is registered to a library.
 		library    namer
@@ -155,9 +156,19 @@ func (rt *RegistrableTask) registerToLibrary(
 
 	fqn := rt.fqType()
 	rt.taskTypeID = GenerateTypeID(fqn)
+	// If outboundTaskGroup wasn't set on creation default it here,
+	// since this is the first place we will have the fqn.
+	if rt.outboundTaskGroup == "" {
+		rt.outboundTaskGroup = fqn
+	}
 	return fqn, rt.taskTypeID, nil
 }
 
+// TaskGroup returns the side-effect task group for the task.
+func (rt *RegistrableTask) TaskGroup() string {
+	return rt.outboundTaskGroup
+}
+
 // GoType returns the reflect.Type of the task's Go struct.
 func (rt *RegistrableTask) GoType() reflect.Type {
 	return rt.goType
@@ -173,3 +184,13 @@ func (rt *RegistrableTask) fqType() string {
 	}
 	return FullyQualifiedName(rt.library.Name(), rt.taskType)
 }
+
+// WithSideEffectTaskGroup sets the task group for the task. The task group is used when
+// the side effect's destination is specified for grouping semantics on the outbound queue,
+// affects multi-cursor and the circuit breaker.
+// If task group isn't provided, the task group will default to the fully qualified name at library registration.
+func WithSideEffectTaskGroup(taskgroup string) RegistrableTaskOption {
+	return func(rt *RegistrableTask) {
+		rt.outboundTaskGroup = taskgroup
+	}
+}

chasm/registry_test.go

@@ -120,6 +120,7 @@ func (s *RegistryTestSuite) TestRegistry_RegisterTasks_Success() {
 		chasm.NewRegistrableSideEffectTask(
 			"Task1",
 			chasm.NewMockSideEffectTaskHandler[*chasm.MockComponent, testTask1](ctrl),
+			chasm.WithSideEffectTaskGroup("test-task-group"),
 		),
 		chasm.NewRegistrablePureTask(
 			"Task2",
@@ -133,6 +134,7 @@ func (s *RegistryTestSuite) TestRegistry_RegisterTasks_Success() {
 	rt1, ok := r.Task("TestLibrary.Task1")
 	require.True(s.T(), ok)
 	require.Equal(s.T(), "TestLibrary.Task1", rt1.FqType())
+	s.Require().Equal("test-task-group", rt1.TaskGroup())
 
 	missingRT, ok := r.Task("TestLibrary.TaskMissing")
 	require.False(s.T(), ok)
@@ -142,6 +144,7 @@ func (s *RegistryTestSuite) TestRegistry_RegisterTasks_Success() {
 	rt2, ok := r.TaskFor(tInstance1)
 	require.True(s.T(), ok)
 	require.Equal(s.T(), "TestLibrary.Task2", rt2.FqType())
+	s.Require().Equal(rt2.FqType(), rt2.TaskGroup())
 
 	rt2, ok = r.TaskOf(reflect.TypeOf(tInstance1))
 	require.True(s.T(), ok)

service/history/outbound_queue_factory.go

@@ -3,6 +3,7 @@ package history
 import (
 	"fmt"
 
+	"go.temporal.io/server/chasm"
 	"go.temporal.io/server/client"
 	"go.temporal.io/server/common/collection"
 	"go.temporal.io/server/common/dynamicconfig"
@@ -307,6 +308,7 @@ func (f *outboundQueueFactory) CreateQueue(
 		logger,
 		metricsHandler,
 		factory,
+		outboundTaskGroupPostProcessor(f.ChasmRegistry),
 	)
 }
 
@@ -359,3 +361,18 @@ func getNamespaceNameOrDefault(
 	}
 	return nsName.String()
 }
+
+func outboundTaskGroupPostProcessor(registry *chasm.Registry) func([]tasks.Task) {
+	if registry == nil {
+		return nil
+	}
+	return func(taskSlice []tasks.Task) {
+		for _, t := range taskSlice {
+			if ct, ok := t.(*tasks.ChasmTask); ok {
+				if rt, ok := registry.TaskByID(ct.Info.GetTypeId()); ok {
+					ct.SetOutboundTaskGroup(rt.TaskGroup())
+				}
+			}
+		}
+	}
+}

service/history/queues/executable.go

@@ -744,10 +744,10 @@ func (e *executableImpl) GetDestination() string {
 	return ""
 }
 
-// StateMachineTaskType returns the embedded task's state machine task type if it exists. Defaults to 0.
-func (e *executableImpl) StateMachineTaskType() string {
-	if t, ok := e.Task.(tasks.HasStateMachineTaskType); ok {
-		return t.StateMachineTaskType()
+// OutboundTaskGroup returns the embedded task's outbound task group. Defaults to an empty string.
+func (e *executableImpl) OutboundTaskGroup() string {
+	if t, ok := e.Task.(tasks.HasOutboundTaskGroup); ok {
+		return t.OutboundTaskGroup()
 	}
 	return ""
 }

service/history/queues/grouper.go

@@ -44,10 +44,10 @@ func (GrouperStateMachineNamespaceIDAndDestination) KeyTyped(task tasks.Task) (k
 	if ok {
 		dest = destGetter.GetDestination()
 	}
-	smtGetter, ok := task.(tasks.HasStateMachineTaskType)
+	smtGetter, ok := task.(tasks.HasOutboundTaskGroup)
 	var smt string
 	if ok {
-		smt = smtGetter.StateMachineTaskType()
+		smt = smtGetter.OutboundTaskGroup()
 	}
 	return tasks.TaskGroupNamespaceIDAndDestination{
 		TaskGroup:   smt,

service/history/queues/metrics.go

@@ -193,7 +193,7 @@ func GetOutboundTaskTypeTagValue(
 
 	switch task := task.(type) {
 	case *tasks.StateMachineOutboundTask:
-		return prefix + "." + task.StateMachineTaskType()
+		return prefix + "." + task.OutboundTaskGroup()
 	case *tasks.ChasmTask:
 		return prefix + "." + getCHASMTaskTypeTagValue(task, chasmRegistry)
 	case *tasks.WorkerCommandsTask:

service/history/queues/queue_immediate.go

@@ -24,6 +24,8 @@ type (
 
 		notifyCh chan struct{}
 	}
+
+	taskPostProcessorFn func([]tasks.Task)
 )
 
 func NewImmediateQueue(
@@ -37,6 +39,7 @@ func NewImmediateQueue(
 	logger log.Logger,
 	metricsHandler metrics.Handler,
 	factory ExecutableFactory,
+	taskPostProcessor taskPostProcessorFn,
 ) *immediateQueue {
 	paginationFnProvider := func(r Range) collection.PaginationFn[tasks.Task] {
 		return func(paginationToken []byte) ([]tasks.Task, []byte, error) {
@@ -57,6 +60,10 @@ func NewImmediateQueue(
 				return nil, nil, err
 			}
 
+			if taskPostProcessor != nil {
+				taskPostProcessor(resp.Tasks)
+			}
+
 			return resp.Tasks, resp.NextPageToken, nil
 		}
 	}

service/history/queues/queue_immediate_test.go

@@ -64,6 +64,7 @@ func (s *immediateQueueSuite) SetupTest() {
 		log.NewTestLogger(),
 		metrics.NoopMetricsHandler,
 		nil, // execuable factory
+		nil, // taskPostProcessor
 	)
 }
 

service/history/tasks/chasm_task.go

@@ -68,11 +68,14 @@ type ChasmTask struct {
 	Info                   *persistencespb.ChasmTaskInfo
 
 	// In-memory only
-	DeserializedTask reflect.Value
+	outboundTaskGroup string // set to the registered task's taskgroup after deserialization for outbound tasks
+	DeserializedTask  reflect.Value
 }
 
 var _ Task = &ChasmTask{}
 var _ HasArchetypeID = &ChasmTask{}
+var _ HasOutboundTaskGroup = &ChasmTask{}
+var _ HasDestination = &ChasmTask{}
 
 func (t *ChasmTask) GetCategory() Category {
 	return t.Category
@@ -105,6 +108,19 @@ func (t *ChasmTask) GetVisibilityTime() time.Time {
 func (t *ChasmTask) SetVisibilityTime(timestamp time.Time) {
 	t.VisibilityTimestamp = timestamp
 }
+
 func (t *ChasmTask) GetArchetypeID() uint32 {
 	return t.Info.GetArchetypeId()
 }
+
+func (t *ChasmTask) GetDestination() string {
+	return t.Destination
+}
+
+func (t *ChasmTask) OutboundTaskGroup() string {
+	return t.outboundTaskGroup
+}
+
+func (t *ChasmTask) SetOutboundTaskGroup(taskgroup string) {
+	t.outboundTaskGroup = taskgroup
+}

service/history/tasks/predicates.go

@@ -122,11 +122,11 @@ func NewOutboundTaskGroupPredicate(
 }
 
 func (n *OutboundTaskGroupPredicate) Test(task Task) bool {
-	smTask, ok := task.(HasStateMachineTaskType)
+	smTask, ok := task.(HasOutboundTaskGroup)
 	if !ok {
 		return false
 	}
-	_, ok = n.Groups[smTask.StateMachineTaskType()]
+	_, ok = n.Groups[smTask.OutboundTaskGroup()]
 	return ok
 }
 
@@ -204,8 +204,8 @@ func (t *OutboundTaskPredicate) Test(task Task) bool {
 	group := TaskGroupNamespaceIDAndDestination{
 		NamespaceID: task.GetNamespaceID(),
 	}
-	if smTask, ok := task.(HasStateMachineTaskType); ok {
-		group.TaskGroup = smTask.StateMachineTaskType()
+	if smTask, ok := task.(HasOutboundTaskGroup); ok {
+		group.TaskGroup = smTask.OutboundTaskGroup()
 	}
 	if dTask, ok := task.(HasDestination); ok {
 		group.Destination = dTask.GetDestination()

service/history/tasks/state_machine_task.go

@@ -16,7 +16,7 @@ type StateMachineTask struct {
 	Info                *persistencespb.StateMachineTaskInfo
 }
 
-var _ HasStateMachineTaskType = &StateMachineTask{}
+var _ HasOutboundTaskGroup = &StateMachineTask{}
 
 func (t *StateMachineTask) GetTaskID() int64 {
 	return t.TaskID
@@ -34,7 +34,7 @@ func (t *StateMachineTask) SetVisibilityTime(timestamp time.Time) {
 	t.VisibilityTimestamp = timestamp
 }
 
-func (t *StateMachineTask) StateMachineTaskType() string {
+func (t *StateMachineTask) OutboundTaskGroup() string {
 	return t.Info.Type
 }
 

service/history/tasks/task.go

@@ -36,10 +36,11 @@ type (
 		GetVersion() int64
 	}
 
-	// HasStateMachineTaskType must be implemented by all HSM state machine tasks.
-	HasStateMachineTaskType interface {
-		StateMachineTaskType() string
+	// HasOutboundTaskGroup must be implemented by all tasks that participate in outbound task grouping.
+	HasOutboundTaskGroup interface {
+		OutboundTaskGroup() string
 	}
+
 	// HasDestination must be implemented by all tasks used in the outbound queue.
 	HasDestination interface {
 		GetDestination() string

service/history/transfer_queue_factory.go

@@ -202,5 +202,6 @@ func (f *transferQueueFactory) CreateQueue(
 		logger,
 		metricsHandler,
 		factory,
+		nil, // taskPostProcessor
 	)
 }

service/history/visibility_queue_factory.go

@@ -168,5 +168,6 @@ func (f *visibilityQueueFactory) CreateQueue(
 		logger,
 		metricsHandler,
 		factory,
+		nil, // taskPostProcessor
 	)
 }