Conditionally register archival category (#3867)

service/history/queueFactoryBase.go

@@ -42,6 +42,7 @@ import (
 	"go.temporal.io/server/service/history/configs"
 	"go.temporal.io/server/service/history/queues"
 	"go.temporal.io/server/service/history/shard"
+	"go.temporal.io/server/service/history/tasks"
 	wcache "go.temporal.io/server/service/history/workflow/cache"
 )
 
@@ -136,8 +137,13 @@ func getQueueFactories(
 		queueFactorySet.TimerQueueFactory,
 		queueFactorySet.VisibilityQueueFactory,
 	}
+	c := tasks.CategoryArchival
+	// this will only affect tests because this method is only called once in production,
+	// but it may be called many times across test runs, which would leave the archival queue as a dangling category
+	tasks.RemoveCategory(c.ID())
 	if archivalMetadata.GetHistoryConfig().StaticClusterState() == archiver.ArchivalEnabled || archivalMetadata.GetVisibilityConfig().StaticClusterState() == archiver.ArchivalEnabled {
 		factories = append(factories, queueFactorySet.ArchivalQueueFactory)
+		tasks.NewCategory(c.ID(), c.Type(), c.Name())
 	}
 	return factories
 }

service/history/queue_factory_base_test.go

@@ -28,6 +28,7 @@ import (
 	"testing"
 
 	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/fx"
 
@@ -45,6 +46,7 @@ import (
 	"go.temporal.io/server/common/sdk"
 	"go.temporal.io/server/service/history/archival"
 	"go.temporal.io/server/service/history/configs"
+	"go.temporal.io/server/service/history/tasks"
 	"go.temporal.io/server/service/history/workflow"
 	"go.temporal.io/server/service/worker/archiver"
 )
@@ -93,7 +95,7 @@ type moduleTestCase struct {
 
 // Run runs the test case.
 func (c *moduleTestCase) Run(t *testing.T) {
-	t.Parallel()
+
 	controller := gomock.NewController(t)
 	dependencies := getModuleDependencies(controller, c)
 	var factories []QueueFactory
@@ -135,8 +137,10 @@ func (c *moduleTestCase) Run(t *testing.T) {
 	require.NotNil(t, viq)
 	if c.ExpectArchivalQueue {
 		require.NotNil(t, aq)
+		assert.Contains(t, tasks.GetCategories(), tasks.CategoryIDArchival)
 	} else {
 		require.Nil(t, aq)
+		assert.NotContains(t, tasks.GetCategories(), tasks.CategoryIDArchival)
 	}
 }
 

service/history/tasks/category.go

@@ -145,6 +145,14 @@ func GetCategories() map[int32]Category {
 	return maps.Clone(categories.m)
 }
 
+// RemoveCategory removes a registered Category.
+// This should only be used for testing.
+func RemoveCategory(id int32) {
+	categories.Lock()
+	defer categories.Unlock()
+	delete(categories.m, id)
+}
+
 // GetCategoryByID returns a registered Category with the same ID
 func GetCategoryByID(id int32) (Category, bool) {
 	categories.RLock()