Create a task category registry object

client/history/historytest/clienttest.go

@@ -97,7 +97,7 @@ func TestClient(t *testing.T, historyTaskQueueManager persistence.HistoryTaskQue
 		require.NoError(t, err)
 		enqueueTasks(t, historyTaskQueueManager, 2, queueKey.SourceCluster, queueKey.TargetCluster)
 		dlqKey := &commonspb.HistoryDLQKey{
-			TaskCategory:  tasks.CategoryTransfer.ID(),
+			TaskCategory:  int32(tasks.CategoryTransfer.ID()),
 			SourceCluster: queueKey.SourceCluster,
 			TargetCluster: queueKey.TargetCluster,
 		}
@@ -139,7 +139,7 @@ func readTasks(
 	for i := 0; i < numTasks; i++ {
 		res, err := client.GetDLQTasks(context.Background(), &historyservice.GetDLQTasksRequest{
 			DlqKey: &commonspb.HistoryDLQKey{
-				TaskCategory:  tasks.CategoryTransfer.ID(),
+				TaskCategory:  int32(tasks.CategoryTransfer.ID()),
 				SourceCluster: sourceCluster,
 				TargetCluster: targetCluster,
 			},
@@ -156,8 +156,9 @@ func readTasks(
 func createServer(historyTaskQueueManager persistence.HistoryTaskQueueManager) *grpc.Server {
 	// TODO: find a better way to create a history handler
 	historyHandler := historyserver.HandlerProvider(historyserver.NewHandlerArgs{
-		TaskQueueManager: historyTaskQueueManager,
-		TracerProvider:   fakeTracerProvider{},
+		TaskQueueManager:     historyTaskQueueManager,
+		TracerProvider:       fakeTracerProvider{},
+		TaskCategoryRegistry: tasks.NewDefaultTaskCategoryRegistry(),
 	})
 	grpcServer := grpc.NewServer()
 	historyservice.RegisterHistoryServiceServer(grpcServer, historyHandler)

cmd/tools/tdbg/main.go

@@ -27,10 +27,11 @@ package main
 import (
 	"os"
 
+	"go.temporal.io/server/service/history/tasks"
 	"go.temporal.io/server/tools/tdbg"
 )
 
 func main() {
-	app := tdbg.NewCliApp(tdbg.NewClientFactory())
+	app := tdbg.NewCliApp(tdbg.NewClientFactory(), tasks.NewDefaultTaskCategoryRegistry())
 	_ = app.Run(os.Args)
 }

common/persistence/sql/execution_tasks.go

@@ -146,7 +146,7 @@ func (m *sqlExecutionStore) getHistoryImmediateTasks(
 
 	rows, err := m.Db.RangeSelectFromHistoryImmediateTasks(ctx, sqlplugin.HistoryImmediateTasksRangeFilter{
 		ShardID:            request.ShardID,
-		CategoryID:         categoryID,
+		CategoryID:         int32(categoryID),
 		InclusiveMinTaskID: inclusiveMinTaskID,
 		ExclusiveMaxTaskID: exclusiveMaxTaskID,
 		PageSize:           request.BatchSize,
@@ -200,7 +200,7 @@ func (m *sqlExecutionStore) completeHistoryImmediateTask(
 
 	if _, err := m.Db.DeleteFromHistoryImmediateTasks(ctx, sqlplugin.HistoryImmediateTasksFilter{
 		ShardID:    request.ShardID,
-		CategoryID: categoryID,
+		CategoryID: int32(categoryID),
 		TaskID:     request.TaskKey.TaskID,
 	}); err != nil {
 		return serviceerror.NewUnavailable(
@@ -229,7 +229,7 @@ func (m *sqlExecutionStore) rangeCompleteHistoryImmediateTasks(
 
 	if _, err := m.Db.RangeDeleteFromHistoryImmediateTasks(ctx, sqlplugin.HistoryImmediateTasksRangeFilter{
 		ShardID:            request.ShardID,
-		CategoryID:         categoryID,
+		CategoryID:         int32(categoryID),
 		InclusiveMinTaskID: request.InclusiveMinTaskKey.TaskID,
 		ExclusiveMaxTaskID: request.ExclusiveMaxTaskKey.TaskID,
 	}); err != nil {
@@ -263,7 +263,7 @@ func (m *sqlExecutionStore) getHistoryScheduledTasks(
 
 	rows, err := m.Db.RangeSelectFromHistoryScheduledTasks(ctx, sqlplugin.HistoryScheduledTasksRangeFilter{
 		ShardID:                         request.ShardID,
-		CategoryID:                      categoryID,
+		CategoryID:                      int32(categoryID),
 		InclusiveMinVisibilityTimestamp: pageToken.Timestamp,
 		InclusiveMinTaskID:              pageToken.TaskID,
 		ExclusiveMaxVisibilityTimestamp: request.ExclusiveMaxTaskKey.FireTime,
@@ -313,7 +313,7 @@ func (m *sqlExecutionStore) completeHistoryScheduledTask(
 
 	if _, err := m.Db.DeleteFromHistoryScheduledTasks(ctx, sqlplugin.HistoryScheduledTasksFilter{
 		ShardID:             request.ShardID,
-		CategoryID:          categoryID,
+		CategoryID:          int32(categoryID),
 		VisibilityTimestamp: request.TaskKey.FireTime,
 		TaskID:              request.TaskKey.TaskID,
 	}); err != nil {
@@ -338,7 +338,7 @@ func (m *sqlExecutionStore) rangeCompleteHistoryScheduledTasks(
 	end := request.ExclusiveMaxTaskKey.FireTime
 	if _, err := m.Db.RangeDeleteFromHistoryScheduledTasks(ctx, sqlplugin.HistoryScheduledTasksRangeFilter{
 		ShardID:                         request.ShardID,
-		CategoryID:                      categoryID,
+		CategoryID:                      int32(categoryID),
 		InclusiveMinVisibilityTimestamp: start,
 		ExclusiveMaxVisibilityTimestamp: end,
 	}); err != nil {

common/persistence/sql/execution_util.go

@@ -698,7 +698,7 @@ func createImmediateTasks(
 	ctx context.Context,
 	tx sqlplugin.Tx,
 	shardID int32,
-	categoryID int32,
+	categoryID int,
 	immedidateTasks []p.InternalHistoryTask,
 ) error {
 	// This is for backward compatiblity.
@@ -721,7 +721,7 @@ func createImmediateTasks(
 	for _, task := range immedidateTasks {
 		immediateTasksRows = append(immediateTasksRows, sqlplugin.HistoryImmediateTasksRow{
 			ShardID:      shardID,
-			CategoryID:   categoryID,
+			CategoryID:   int32(categoryID),
 			TaskID:       task.Key.TaskID,
 			Data:         task.Blob.Data,
 			DataEncoding: task.Blob.EncodingType.String(),
@@ -748,7 +748,7 @@ func createScheduledTasks(
 	ctx context.Context,
 	tx sqlplugin.Tx,
 	shardID int32,
-	categoryID int32,
+	categoryID int,
 	scheduledTasks []p.InternalHistoryTask,
 ) error {
 	// This is for backward compatiblity.
@@ -766,7 +766,7 @@ func createScheduledTasks(
 	for _, task := range scheduledTasks {
 		scheduledTasksRows = append(scheduledTasksRows, sqlplugin.HistoryScheduledTasksRow{
 			ShardID:             shardID,
-			CategoryID:          categoryID,
+			CategoryID:          int32(categoryID),
 			VisibilityTimestamp: task.Key.FireTime,
 			TaskID:              task.Key.TaskID,
 			Data:                task.Blob.Data,

proto/internal/temporal/server/api/historyservice/v1/request_response.proto

@@ -778,9 +778,10 @@ message AddTasksRequest {
   int32 shard_id = 1;
 
   message Task {
-    // category is needed to deserialize the tasks. Examples include "transfer", "timer", etc. See the history/tasks
-    // package for a definitive list. Warning: this is not the same as the stringified value of a TaskCategory enum.
-    string category = 1;
+    // category is needed to deserialize the tasks. See TaskCategory for a list of options here. However, keep in mind
+    // that the list of valid options is registered dynamically with the server in the history/tasks package, so that
+    // enum is not comprehensive.
+    int32 category = 1;
     // blob is the serialized task.
     temporal.api.common.v1.DataBlob blob = 2;
   }

service/frontend/admin_handler.go

@@ -123,6 +123,8 @@ type (
 
 		// DEPRECATED
 		persistenceExecutionManager persistence.ExecutionManager
+
+		taskCategoryRegistry tasks.TaskCategoryRegistry
 	}
 
 	NewAdminHandlerArgs struct {
@@ -153,6 +155,8 @@ type (
 
 		// DEPRECATED
 		PersistenceExecutionManager persistence.ExecutionManager
+
+		CategoryRegistry tasks.TaskCategoryRegistry
 	}
 )
 
@@ -202,6 +206,7 @@ func NewAdminHandler(
 		saManager:                   args.SaManager,
 		clusterMetadata:             args.ClusterMetadata,
 		healthServer:                args.HealthServer,
+		taskCategoryRegistry:        args.CategoryRegistry,
 	}
 }
 
@@ -794,7 +799,7 @@ func (adh *AdminHandler) ListHistoryTasks(
 		return nil, errTaskRangeNotSet
 	}
 
-	taskCategory, ok := tasks.GetCategoryByID(int32(request.Category))
+	taskCategory, ok := adh.taskCategoryRegistry.GetCategoryByID(int(request.Category))
 	if !ok {
 		return nil, &serviceerror.InvalidArgument{
 			Message: fmt.Sprintf("unknown task category: %v", request.Category),

service/frontend/admin_handler_test.go

@@ -170,6 +170,7 @@ func (s *adminHandlerSuite) SetupTest() {
 		serialization.NewSerializer(),
 		clock.NewRealTimeSource(),
 		s.mockResource.GetExecutionManager(),
+		tasks.NewDefaultTaskCategoryRegistry(),
 	}
 	s.mockMetadata.EXPECT().GetCurrentClusterName().Return(uuid.New()).AnyTimes()
 	s.handler = NewAdminHandler(args)
@@ -1160,7 +1161,7 @@ func (s *adminHandlerSuite) TestGetDLQTasks() {
 			blob := &commonpb.DataBlob{}
 			expectation := s.mockHistoryClient.EXPECT().GetDLQTasks(gomock.Any(), &historyservice.GetDLQTasksRequest{
 				DlqKey: &commonspb.HistoryDLQKey{
-					TaskCategory:  tasks.CategoryTransfer.ID(),
+					TaskCategory:  int32(tasks.CategoryTransfer.ID()),
 					SourceCluster: "test-source-cluster",
 					TargetCluster: "test-target-cluster",
 				},
@@ -1187,7 +1188,7 @@ func (s *adminHandlerSuite) TestGetDLQTasks() {
 			}
 			response, err := s.handler.GetDLQTasks(context.Background(), &adminservice.GetDLQTasksRequest{
 				DlqKey: &commonspb.HistoryDLQKey{
-					TaskCategory:  tasks.CategoryTransfer.ID(),
+					TaskCategory:  int32(tasks.CategoryTransfer.ID()),
 					SourceCluster: "test-source-cluster",
 					TargetCluster: "test-target-cluster",
 				},

service/frontend/fx.go

@@ -65,6 +65,7 @@ import (
 	"go.temporal.io/server/common/telemetry"
 	"go.temporal.io/server/service"
 	"go.temporal.io/server/service/frontend/configs"
+	"go.temporal.io/server/service/history/tasks"
 )
 
 type FEReplicatorNamespaceReplicationQueue persistence.NamespaceReplicationQueue
@@ -503,6 +504,7 @@ func AdminHandlerProvider(
 	healthServer *health.Server,
 	eventSerializer serialization.Serializer,
 	timeSource clock.TimeSource,
+	taskCategoryRegistry tasks.TaskCategoryRegistry,
 ) *AdminHandler {
 	args := NewAdminHandlerArgs{
 		persistenceConfig,
@@ -530,6 +532,7 @@ func AdminHandlerProvider(
 		eventSerializer,
 		timeSource,
 		persistenceExecutionManager,
+		taskCategoryRegistry,
 	}
 	return NewAdminHandler(args)
 }

service/history/api/addtasks/api.go

@@ -34,6 +34,7 @@ import (
 	"go.temporal.io/server/api/historyservice/v1"
 	"go.temporal.io/server/common/definition"
 	"go.temporal.io/server/common/persistence"
+	"go.temporal.io/server/service/history/api"
 	"go.temporal.io/server/service/history/shard"
 	"go.temporal.io/server/service/history/tasks"
 )
@@ -63,6 +64,7 @@ func Invoke(
 	deserializer TaskDeserializer,
 	numShards int,
 	req *historyservice.AddTasksRequest,
+	taskRegistry tasks.TaskCategoryRegistry,
 ) (*historyservice.AddTasksResponse, error) {
 	if len(req.Tasks) > maxTasksPerRequest {
 		return nil, serviceerror.NewInvalidArgument(fmt.Sprintf(
@@ -83,12 +85,9 @@ func Invoke(
 			return nil, serviceerror.NewInvalidArgument(fmt.Sprintf("Nil task at index: %d", i))
 		}
 
-		category, ok := getCategoryByName(task.Category)
-		if !ok {
-			return nil, serviceerror.NewInvalidArgument(fmt.Sprintf(
-				"Invalid task category: %s",
-				task.Category,
-			))
+		category, err := api.GetTaskCategory(int(task.GetCategory()), taskRegistry)
+		if err != nil {
+			return nil, err
 		}
 
 		if task.Blob == nil {
@@ -139,14 +138,3 @@ func Invoke(
 
 	return &historyservice.AddTasksResponse{}, nil
 }
-
-func getCategoryByName(categoryName string) (tasks.Category, bool) {
-	categories := tasks.GetCategories()
-	for _, category := range categories {
-		if category.Name() == categoryName {
-			return category, true
-		}
-	}
-
-	return tasks.Category{}, false
-}