Remove special CHASM fx graph handling (#10582)

## What changed?

Access CHASM from functional tests via test hook, instead of
`fx.Populate`.

## Why?

One step closer to the goal of deleting the separate onebox.go fx graph
all-together and making use of the production fx graph from the temporal
package instead.

chasm/lib/tests/fx.go

@@ -1,3 +1,5 @@
+//go:build test_dep
+
 package tests
 
 import (

chasm/lib/tests/fx_noop.go

@@ -0,0 +1,7 @@
+//go:build !test_dep
+
+package tests
+
+import "go.uber.org/fx"
+
+var Module = fx.Options()

common/testing/testhooks/hooks.go

@@ -7,6 +7,7 @@ import (
 	"go.temporal.io/server/api/historyservice/v1"
 	persistencespb "go.temporal.io/server/api/persistence/v1"
 	replicationspb "go.temporal.io/server/api/replication/v1"
+	"go.temporal.io/server/chasm"
 	"go.temporal.io/server/common/namespace"
 	historytasks "go.temporal.io/server/service/history/tasks"
 )
@@ -25,6 +26,7 @@ var (
 	MatchingForwardTaskDelay                 = newKey[time.Duration, namespace.ID]()
 	HistoryReplicationTaskInterceptor        = newKey[func(*replicationspb.ReplicationTask, func() error) error, global]()
 	HistoryReplicationDLQWriteInterceptor    = newKey[func(*persistencespb.ReplicationTaskInfo, func() error) error, global]()
+	HistoryChasmRuntimeProvider              = newKey[func(chasm.Engine, chasm.VisibilityManager, *chasm.Registry), global]()
 	HistoryTransferTaskInterceptor           = newKey[func(historytasks.Task, func()), namespace.ID]()
 	HistoryDLQTaskDeleteInterceptor          = newKey[func(context.Context, *historyservice.DeleteDLQTasksRequest, func(context.Context, *historyservice.DeleteDLQTasksRequest) (*historyservice.DeleteDLQTasksResponse, error)) (*historyservice.DeleteDLQTasksResponse, error), global]()
 	NamespaceReplicationTaskInterceptor      = newKey[func(context.Context, *replicationspb.NamespaceTaskAttributes, func() error) error, namespace.Name]()

service/frontend/fx.go

@@ -13,6 +13,7 @@ import (
 	nexusoperationpb "go.temporal.io/server/chasm/lib/nexusoperation/gen/nexusoperationpb/v1"
 	chasmscheduler "go.temporal.io/server/chasm/lib/scheduler"
 	"go.temporal.io/server/chasm/lib/scheduler/gen/schedulerpb/v1"
+	chasmtests "go.temporal.io/server/chasm/lib/tests"
 	chasmworkflow "go.temporal.io/server/chasm/lib/workflow"
 	"go.temporal.io/server/client"
 	"go.temporal.io/server/common"
@@ -67,6 +68,7 @@ type (
 
 var Module = fx.Options(
 	resource.Module,
+	chasmtests.Module,
 	scheduler.Module,
 	workerdeployment.Module,
 	// Note that with this approach routes may be registered in arbitrary order.

service/history/fx.go

@@ -11,6 +11,7 @@ import (
 	"go.temporal.io/server/chasm/lib/callback"
 	chasmnexus "go.temporal.io/server/chasm/lib/nexusoperation"
 	"go.temporal.io/server/chasm/lib/scheduler"
+	chasmtests "go.temporal.io/server/chasm/lib/tests"
 	chasmworkflow "go.temporal.io/server/chasm/lib/workflow"
 	"go.temporal.io/server/common"
 	commoncache "go.temporal.io/server/common/cache"
@@ -34,6 +35,7 @@ import (
 	"go.temporal.io/server/common/rpc/interceptor"
 	"go.temporal.io/server/common/searchattribute"
 	"go.temporal.io/server/common/tasktoken"
+	"go.temporal.io/server/common/testing/testhooks"
 	"go.temporal.io/server/common/worker_versioning"
 	"go.temporal.io/server/components/callbacks"
 	hsmnexusoperations "go.temporal.io/server/components/nexusoperations"
@@ -66,6 +68,7 @@ var Module = fx.Options(
 	cache.Module,
 	archival.Module,
 	ChasmEngineModule,
+	chasmtests.Module,
 	fx.Provide(ConfigProvider), // might be worth just using provider for configs.Config directly
 	fx.Provide(workflow.NewCommandHandlerRegistry),
 	fx.Provide(ServiceErrorInterceptorProvider),
@@ -99,6 +102,20 @@ var Module = fx.Options(
 	workerdeployment.ClientModule,
 	fx.Provide(RoutingInfoCacheProvider),
 	fx.Invoke(ServiceLifetimeHooks),
+	fx.Invoke(func(
+		chasmEngine chasm.Engine,
+		chasmVisibilityManager chasm.VisibilityManager,
+		chasmRegistry *chasm.Registry,
+		testHooks testhooks.TestHooks,
+	) {
+		if hook, ok := testhooks.Get(
+			testHooks,
+			testhooks.HistoryChasmRuntimeProvider,
+			testhooks.GlobalScope,
+		); ok {
+			hook(chasmEngine, chasmVisibilityManager, chasmRegistry)
+		}
+	}),
 
 	callbacks.Module,
 	hsmnexusoperations.Module,

tests/chasm_test.go

@@ -69,15 +69,8 @@ func newChasmTestEnv(t *testing.T, unified bool) chasmTestEnv {
 		testcore.WithDynamicConfig(dynamicconfig.DeleteNamespaceUseChasmDeleteExecution, true),
 	)
 
-	chasmEngine, err := env.GetTestCluster().Host().ChasmEngine()
+	chasmCtx, err := env.GetTestCluster().Host().ChasmContext(env.Context())
 	require.NoError(t, err)
-	require.NotNil(t, chasmEngine)
-
-	chasmVisibilityMgr := env.GetTestCluster().Host().ChasmVisibilityManager()
-	require.NotNil(t, chasmVisibilityMgr)
-
-	chasmCtx := chasm.NewEngineContext(env.Context(), chasmEngine)
-	chasmCtx = chasm.NewVisibilityManagerContext(chasmCtx, chasmVisibilityMgr)
 
 	return chasmTestEnv{TestEnv: env, chasmCtx: chasmCtx}
 }

tests/testcore/onebox.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"maps"
 	"math/rand"
@@ -18,7 +19,6 @@ import (
 	"go.temporal.io/server/api/adminservice/v1"
 	"go.temporal.io/server/chasm"
 	chasmnexus "go.temporal.io/server/chasm/lib/nexusoperation"
-	chasmtests "go.temporal.io/server/chasm/lib/tests"
 	"go.temporal.io/server/client"
 	"go.temporal.io/server/common"
 	carchiver "go.temporal.io/server/common/archiver"
@@ -71,8 +71,6 @@ type (
 		namespaceRegistries []namespace.Registry
 		// Address for SDK to connect to, using membership grpc resolver.
 		frontendMembershipAddress string
-		chasmEngine               chasm.Engine
-		chasmVisibilityMgr        chasm.VisibilityManager
 
 		dcClient                         *dynamicconfig.MemoryClient
 		testHooks                        testhooks.TestHooks
@@ -107,7 +105,8 @@ type (
 		callbackLock              sync.RWMutex // Must be used for above callbacks
 		serviceFxOptions          map[primitives.ServiceName][]fx.Option
 		taskCategoryRegistry      tasks.TaskCategoryRegistry
-		chasmRegistry             *chasm.Registry
+		chasmEngine               chasm.Engine
+		chasmVisibilityMgr        chasm.VisibilityManager
 		replicationStreamRecorder *ReplicationStreamRecorder
 		taskQueueRecorder         *TaskQueueRecorder
 		spanExporters             map[telemetry.SpanExporterType]sdktrace.SpanExporter
@@ -303,15 +302,16 @@ func (c *TemporalImpl) NamespaceRegistries() []namespace.Registry {
 	return c.namespaceRegistries
 }
 
-func (c *TemporalImpl) ChasmEngine() (chasm.Engine, error) {
+func (c *TemporalImpl) ChasmContext(ctx context.Context) (context.Context, error) {
 	if numHistoryHosts := len(c.hostsByProtocolByService[grpcProtocol][primitives.HistoryService].All); numHistoryHosts != 1 {
-		return nil, fmt.Errorf("expected exactly one host for chasm engine, got %d", numHistoryHosts)
+		return nil, fmt.Errorf("expected exactly one history host for chasm context, got %d", numHistoryHosts)
 	}
-	return c.chasmEngine, nil
-}
-
-func (c *TemporalImpl) ChasmVisibilityManager() chasm.VisibilityManager {
-	return c.chasmVisibilityMgr
+	if c.chasmEngine == nil || c.chasmVisibilityMgr == nil {
+		return nil, errors.New("chasm context is not available")
+	}
+	ctx = chasm.NewEngineContext(ctx, c.chasmEngine)
+	ctx = chasm.NewVisibilityManagerContext(ctx, c.chasmVisibilityMgr)
+	return ctx, nil
 }
 
 func (c *TemporalImpl) copyPersistenceConfig() config.Persistence {
@@ -395,7 +395,6 @@ func (c *TemporalImpl) startFrontend() {
 			temporal.FxLogAdapter,
 			c.getFxOptionsForService(primitives.FrontendService),
 			chasm.Module,
-			chasmtests.Module,
 		)
 		err := app.Err()
 		if err != nil {
@@ -419,6 +418,15 @@ func (c *TemporalImpl) startFrontend() {
 func (c *TemporalImpl) startHistory() {
 	serviceName := primitives.HistoryService
 
+	testhooks.NewHook(testhooks.HistoryChasmRuntimeProvider, func(
+		chasmEngine chasm.Engine,
+		chasmVisibilityManager chasm.VisibilityManager,
+		_ *chasm.Registry,
+	) {
+		c.chasmEngine = chasmEngine
+		c.chasmVisibilityMgr = chasmVisibilityManager
+	}).Apply(c.testHooks, testhooks.GlobalScope)
+
 	for _, host := range c.hostsByProtocolByService[grpcProtocol][serviceName].All {
 		var namespaceRegistry namespace.Registry
 		logger := log.With(c.logger, tag.Host(host))
@@ -483,11 +491,7 @@ func (c *TemporalImpl) startHistory() {
 			temporal.FxLogAdapter,
 			c.getFxOptionsForService(primitives.HistoryService),
 			chasm.Module,
-			chasmtests.Module,
 			fx.Populate(&namespaceRegistry),
-			fx.Populate(&c.chasmEngine),
-			fx.Populate(&c.chasmVisibilityMgr),
-			fx.Populate(&c.chasmRegistry),
 		)
 		err := app.Err()
 		if err != nil {
@@ -543,7 +547,6 @@ func (c *TemporalImpl) startMatching() {
 			temporal.FxLogAdapter,
 			c.getFxOptionsForService(primitives.MatchingService),
 			chasm.Module,
-			chasmtests.Module,
 			fx.Populate(&namespaceRegistry),
 		)
 		err := app.Err()
@@ -610,7 +613,6 @@ func (c *TemporalImpl) startWorker() {
 			temporal.FxLogAdapter,
 			c.getFxOptionsForService(primitives.WorkerService),
 			chasm.Module,
-			chasmtests.Module,
 			fx.Populate(&namespaceRegistry),
 		)
 		err := app.Err()
@@ -666,10 +668,6 @@ func (c *TemporalImpl) GetTaskCategoryRegistry() tasks.TaskCategoryRegistry {
 	return c.taskCategoryRegistry
 }
 
-func (c *TemporalImpl) GetCHASMRegistry() *chasm.Registry {
-	return c.chasmRegistry
-}
-
 func (c *TemporalImpl) TlsConfigProvider() *encryption.FixedTLSConfigProvider {
 	return c.tlsConfigProvider
 }

tests/xdc/chasm_test.go

@@ -59,15 +59,9 @@ func (s *ChasmSuite) SetupSuite() {
 func (s *ChasmSuite) SetupTest() {
 	s.setupTest()
 
-	chasmEngine, err := s.clusters[0].Host().ChasmEngine()
+	chasmContext, err := s.clusters[0].Host().ChasmContext(context.Background())
 	s.Require().NoError(err)
-	s.Require().NotNil(chasmEngine)
-
-	chasmVisibilityMgr := s.clusters[0].Host().ChasmVisibilityManager()
-	s.Require().NotNil(chasmVisibilityMgr)
-
-	s.chasmContext = chasm.NewEngineContext(context.Background(), chasmEngine)
-	s.chasmContext = chasm.NewVisibilityManagerContext(s.chasmContext, chasmVisibilityMgr)
+	s.chasmContext = chasmContext
 }
 
 func (s *ChasmSuite) TearDownSuite() {
@@ -82,7 +76,6 @@ func (s *ChasmSuite) TestDeleteExecution_RunningExecution() {
 	})
 	s.NoError(err)
 	nsID := nsResp.NamespaceInfo.GetId()
-
 	tv := testvars.New(s.T())
 	storeID := tv.Any().String()
 
@@ -100,18 +93,12 @@ func (s *ChasmSuite) TestDeleteExecution_RunningExecution() {
 	)
 	s.NoError(err)
 
-	chasmRegistry := s.clusters[0].Host().GetCHASMRegistry()
-	archetypeID, ok := chasmRegistry.ComponentIDFor(&tests.PayloadStore{})
-	s.True(ok)
-	archetype, ok := chasmRegistry.ComponentFqnByID(archetypeID)
-	s.True(ok)
-
 	describeExecutionRequest := &adminservice.DescribeMutableStateRequest{
 		Namespace: nsName,
 		Execution: &commonpb.WorkflowExecution{
 			WorkflowId: storeID,
 		},
-		Archetype: archetype,
+		ArchetypeId: uint32(tests.ArchetypeID),
 	}
 	_, err = s.clusters[0].AdminClient().DescribeMutableState(testcore.NewContext(), describeExecutionRequest)
 	s.NoError(err)
@@ -177,18 +164,12 @@ func (s *ChasmSuite) TestRetentionTimer() {
 	)
 	s.NoError(err)
 
-	chasmRegistry := s.clusters[0].Host().GetCHASMRegistry()
-	archetypeID, ok := chasmRegistry.ComponentIDFor(&tests.PayloadStore{})
-	s.True(ok)
-	archetype, ok := chasmRegistry.ComponentFqnByID(archetypeID)
-	s.True(ok)
-
 	describeExecutionRequest := &adminservice.DescribeMutableStateRequest{
 		Namespace: nsName,
 		Execution: &commonpb.WorkflowExecution{
 			WorkflowId: storeID,
 		},
-		Archetype: archetype,
+		ArchetypeId: uint32(tests.ArchetypeID),
 	}
 	_, err = s.clusters[0].AdminClient().DescribeMutableState(testcore.NewContext(), describeExecutionRequest)
 	s.NoError(err)
@@ -282,14 +263,14 @@ func (s *ChasmSuite) TestActivityDispatchTaskStandbySpillover() {
 	s.NoError(err)
 	s.NotEmpty(startResp.GetRunId())
 
-	// Wait for replication to cluster 1 (standby).
 	describeExecutionRequest := &adminservice.DescribeMutableStateRequest{
 		Namespace: nsName,
 		Execution: &commonpb.WorkflowExecution{
 			WorkflowId: activityID,
 		},
-		Archetype: activity.Archetype,
+		ArchetypeId: uint32(activity.ArchetypeID),
 	}
+	// Wait for replication to cluster 1 (standby).
 	s.Eventually(func() bool {
 		_, err = s.clusters[1].AdminClient().DescribeMutableState(testcore.NewContext(), describeExecutionRequest)
 		return err == nil
@@ -406,18 +387,12 @@ func (s *ChasmSuite) TestDeleteExecution_ReplicatedToStandby() {
 	)
 	s.NoError(err)
 
-	chasmRegistry := s.clusters[0].Host().GetCHASMRegistry()
-	archetypeID, ok := chasmRegistry.ComponentIDFor(&tests.PayloadStore{})
-	s.True(ok)
-	archetype, ok := chasmRegistry.ComponentFqnByID(archetypeID)
-	s.True(ok)
-
 	describeExecutionRequest := &adminservice.DescribeMutableStateRequest{
 		Namespace: nsName,
 		Execution: &commonpb.WorkflowExecution{
 			WorkflowId: storeID,
 		},
-		Archetype: archetype,
+		ArchetypeId: uint32(tests.ArchetypeID),
 	}
 
 	s.Eventually(func() bool {