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_id 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_id = 1;
     // blob is the serialized task.
     temporal.api.common.v1.DataBlob blob = 2;
   }

service/frontend/admin_handler.go

@@ -122,6 +122,8 @@ type (
 
 		// DEPRECATED
 		persistenceExecutionManager persistence.ExecutionManager
+
+		taskCategoryRegistry tasks.TaskCategoryRegistry
 	}
 
 	NewAdminHandlerArgs struct {
@@ -152,6 +154,8 @@ type (
 
 		// DEPRECATED
 		PersistenceExecutionManager persistence.ExecutionManager
+
+		CategoryRegistry tasks.TaskCategoryRegistry
 	}
 )
 
@@ -201,6 +205,7 @@ func NewAdminHandler(
 		saManager:                   args.SaManager,
 		clusterMetadata:             args.ClusterMetadata,
 		healthServer:                args.HealthServer,
+		taskCategoryRegistry:        args.CategoryRegistry,
 	}
 }
 
@@ -793,7 +798,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.CategoryId), 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
-}

service/history/api/addtasks/api_test.go

@@ -127,8 +127,8 @@ func TestInvoke(t *testing.T) {
 						blob, err := serializer.SerializeTask(task)
 						require.NoError(t, err)
 						params.req.Tasks = append(params.req.Tasks, &historyservice.AddTasksRequest_Task{
-							Category: task.GetCategory().Name(),
-							Blob:     &blob,
+							CategoryId: int32(task.GetCategory().ID()),
+							Blob:       &blob,
 						})
 					}
 				}
@@ -185,11 +185,11 @@ func TestInvoke(t *testing.T) {
 		{
 			name: "invalid task category",
 			configure: func(t *testing.T, params *testParams) {
-				params.req.Tasks[0].Category = "my-invalid-task-category"
+				params.req.Tasks[0].CategoryId = -1
 				params.expectation = func(resp *historyservice.AddTasksResponse, err error) {
 					require.ErrorAs(t, err, new(*serviceerror.InvalidArgument))
 					assert.ErrorContains(t, err, "Invalid task category")
-					assert.ErrorContains(t, err, "my-invalid-task-category")
+					assert.ErrorContains(t, err, "-1")
 				}
 			},
 		},
@@ -251,6 +251,7 @@ func TestInvoke(t *testing.T) {
 				params.deserializer,
 				params.numShards,
 				params.req,
+				tasks.NewDefaultTaskCategoryRegistry(),
 			)
 			params.expectation(resp, err)
 		})
@@ -276,8 +277,8 @@ func getDefaultTestParams(t *testing.T) *testParams {
 			ShardId: 1,
 			Tasks: []*historyservice.AddTasksRequest_Task{
 				{
-					Category: tasks.CategoryTransfer.Name(),
-					Blob:     &blob,
+					CategoryId: int32(tasks.CategoryTransfer.ID()),
+					Blob:       &blob,
 				},
 			},
 		},

service/history/api/deletedlqtasks/api.go

@@ -32,14 +32,16 @@ import (
 	"go.temporal.io/server/api/historyservice/v1"
 	"go.temporal.io/server/common/persistence"
 	"go.temporal.io/server/service/history/api"
+	"go.temporal.io/server/service/history/tasks"
 )
 
 func Invoke(
 	ctx context.Context,
 	historyTaskQueueManager persistence.HistoryTaskQueueManager,
 	req *historyservice.DeleteDLQTasksRequest,
+	registry tasks.TaskCategoryRegistry,
 ) (*historyservice.DeleteDLQTasksResponse, error) {
-	category, err := api.GetTaskCategory(int(req.DlqKey.TaskCategory))
+	category, err := api.GetTaskCategory(int(req.DlqKey.TaskCategory), registry)
 	if err != nil {
 		return nil, err
 	}

service/history/api/deletedlqtasks/api_test.go

@@ -38,6 +38,7 @@ import (
 	"go.temporal.io/server/common/persistence"
 	"go.temporal.io/server/common/persistence/persistencetest"
 	"go.temporal.io/server/service/history/api/deletedlqtasks"
+	"go.temporal.io/server/service/history/tasks"
 )
 
 func TestInvoke_InvalidCategory(t *testing.T) {
@@ -50,7 +51,7 @@ func TestInvoke_InvalidCategory(t *testing.T) {
 			SourceCluster: queueKey.SourceCluster,
 			TargetCluster: queueKey.TargetCluster,
 		},
-	})
+	}, tasks.NewDefaultTaskCategoryRegistry())
 	require.Error(t, err)
 	assert.Equal(t, codes.InvalidArgument, serviceerror.ToStatus(err).Code())
 	assert.ErrorContains(t, err, "-1")
@@ -62,11 +63,11 @@ func TestInvoke_ErrDeleteMissingMessageIDUpperBound(t *testing.T) {
 	queueKey := persistencetest.GetQueueKey(t, persistencetest.WithQueueType(persistence.QueueTypeHistoryDLQ))
 	_, err := deletedlqtasks.Invoke(context.Background(), nil, &historyservice.DeleteDLQTasksRequest{
 		DlqKey: &commonspb.HistoryDLQKey{
-			TaskCategory:  queueKey.Category.ID(),
+			TaskCategory:  int32(queueKey.Category.ID()),
 			SourceCluster: queueKey.SourceCluster,
 			TargetCluster: queueKey.TargetCluster,
 		},
-	})
+	}, tasks.NewDefaultTaskCategoryRegistry())
 	require.Error(t, err)
 	assert.Equal(t, codes.InvalidArgument, serviceerror.ToStatus(err).Code())
 	assert.ErrorContains(t, err, "inclusive_max_task_metadata")