Initialize shard controller with fx (#2319)

service/history/fx.go

@@ -58,12 +58,14 @@ import (
 	"go.temporal.io/server/common/searchattribute"
 	"go.temporal.io/server/service"
 	"go.temporal.io/server/service/history/configs"
+	"go.temporal.io/server/service/history/shard"
 	"go.temporal.io/server/service/history/workflow"
 )
 
 var Module = fx.Options(
 	resource.Module,
 	workflow.Module,
+	shard.Module,
 	fx.Provide(ParamsExpandProvider), // BootstrapParams should be deprecated
 	fx.Provide(dynamicconfig.NewCollection),
 	fx.Provide(ConfigProvider), // might be worth just using provider for configs.Config directly
@@ -78,6 +80,7 @@ var Module = fx.Options(
 	fx.Provide(HandlerProvider),
 	fx.Provide(ServiceProvider),
 	fx.Invoke(ServiceLifetimeHooks),
+	fx.Invoke(func(h *Handler, c *shard.ControllerImpl) { c.SetEngineFactory(h) }), // create cycle
 )
 
 func ServiceProvider(
@@ -132,6 +135,7 @@ func HandlerProvider(
 	archivalMetadata archiver.ArchivalMetadata,
 	hostInfoProvider resource.HostInfoProvider,
 	archiverProvider provider.ArchiverProvider,
+	shardController *shard.ControllerImpl,
 ) *Handler {
 	args := NewHandlerArgs{
 		config,
@@ -157,6 +161,7 @@ func HandlerProvider(
 		archivalMetadata,
 		hostInfoProvider,
 		archiverProvider,
+		shardController,
 	}
 	return NewHandler(args)
 }
@@ -273,5 +278,4 @@ func ServiceLifetimeHooks(
 			},
 		},
 	)
-
 }

service/history/handler.go

@@ -76,7 +76,6 @@ type (
 	Handler struct {
 		status int32
 
-		controller                  *shard.ControllerImpl
 		tokenSerializer             common.TaskTokenSerializer
 		startWG                     sync.WaitGroup
 		config                      *configs.Config
@@ -104,6 +103,7 @@ type (
 		clusterMetadata             cluster.Metadata
 		archivalMetadata            archiver.ArchivalMetadata
 		hostInfoProvider            resource.HostInfoProvider
+		controller                  *shard.ControllerImpl
 	}
 
 	NewHandlerArgs struct {
@@ -130,6 +130,7 @@ type (
 		ArchivalMetadata            archiver.ArchivalMetadata
 		HostInfoProvider            resource.HostInfoProvider
 		ArchiverProvider            provider.ArchiverProvider
+		ShardController             *shard.ControllerImpl
 	}
 )
 
@@ -184,6 +185,7 @@ func NewHandler(args NewHandlerArgs) *Handler {
 		archivalMetadata:            args.ArchivalMetadata,
 		hostInfoProvider:            args.HostInfoProvider,
 		archiverProvider:            args.ArchiverProvider,
+		controller:                  args.ShardController,
 	}
 
 	// prevent us from trying to serve requests before shard controller is started and ready
@@ -210,26 +212,6 @@ func (h *Handler) Start() {
 
 	h.replicationTaskFetchers.Start()
 
-	h.controller = shard.NewController(
-		h,
-		h.config,
-		h.logger,
-		h.throttledLogger,
-		h.persistenceExecutionManager,
-		h.persistenceShardManager,
-		h.clientBean,
-		h.historyClient,
-		h.historyServiceResolver,
-		h.metricsClient,
-		h.payloadSerializer,
-		h.timeSource,
-		h.namespaceRegistry,
-		h.saProvider,
-		h.saMapper,
-		h.clusterMetadata,
-		h.archivalMetadata,
-		h.hostInfoProvider,
-	)
 	h.eventNotifier = events.NewNotifier(h.timeSource, h.metricsClient, h.config.GetShardID)
 	// events notifier must starts before controller
 	h.eventNotifier.Start()

service/history/service.go

@@ -30,7 +30,6 @@ import (
 	"sync/atomic"
 	"time"
 
-	"go.uber.org/fx"
 	"google.golang.org/grpc"
 	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 
@@ -48,8 +47,6 @@ import (
 // Service represents the history service
 type (
 	Service struct {
-		self *fx.App
-
 		status            int32
 		handler           *Handler
 		visibilityManager manager.VisibilityManager

service/history/shard/controller_impl.go

@@ -88,53 +88,8 @@ type (
 	}
 )
 
-func NewController(
-	factory EngineFactory,
-	config *configs.Config,
-	logger log.Logger,
-	throttledLogger log.Logger,
-	persistenceExecutionManager persistence.ExecutionManager,
-	persistenceShardManager persistence.ShardManager,
-	clientBean client.Bean,
-	historyClient historyservice.HistoryServiceClient,
-	historyServiceResolver membership.ServiceResolver,
-	metricsClient metrics.Client,
-	payloadSerializer serialization.Serializer,
-	timeSource clock.TimeSource,
-	namespaceRegistry namespace.Registry,
-	saProvider searchattribute.Provider,
-	saMapper searchattribute.Mapper,
-	clusterMetadata cluster.Metadata,
-	archivalMetadata archiver.ArchivalMetadata,
-	hostInfoProvider resource.HostInfoProvider,
-) *ControllerImpl {
-	hostIdentity := hostInfoProvider.HostInfo().Identity()
-	return &ControllerImpl{
-		status:                      common.DaemonStatusInitialized,
-		membershipUpdateCh:          make(chan *membership.ChangedEvent, 10),
-		engineFactory:               factory,
-		historyShards:               make(map[int32]*ContextImpl),
-		shutdownCh:                  make(chan struct{}),
-		logger:                      logger,
-		contextTaggedLogger:         log.With(logger, tag.ComponentShardController, tag.Address(hostIdentity)),
-		throttledLogger:             log.With(throttledLogger, tag.ComponentShardController, tag.Address(hostIdentity)),
-		config:                      config,
-		metricsScope:                metricsClient.Scope(metrics.HistoryShardControllerScope),
-		persistenceExecutionManager: persistenceExecutionManager,
-		persistenceShardManager:     persistenceShardManager,
-		clientBean:                  clientBean,
-		historyClient:               historyClient,
-		historyServiceResolver:      historyServiceResolver,
-		metricsClient:               metricsClient,
-		payloadSerializer:           payloadSerializer,
-		timeSource:                  timeSource,
-		namespaceRegistry:           namespaceRegistry,
-		saProvider:                  saProvider,
-		saMapper:                    saMapper,
-		clusterMetadata:             clusterMetadata,
-		archivalMetadata:            archivalMetadata,
-		hostInfoProvider:            hostInfoProvider,
-	}
+func (c *ControllerImpl) SetEngineFactory(engineFactory EngineFactory) {
+	c.engineFactory = engineFactory
 }
 
 func (c *ControllerImpl) Start() {
@@ -146,6 +101,14 @@ func (c *ControllerImpl) Start() {
 		return
 	}
 
+	if c.engineFactory == nil {
+		panic("engineFactory was not injected")
+	}
+
+	hostIdentity := c.hostInfoProvider.HostInfo().Identity()
+	c.contextTaggedLogger = log.With(c.logger, tag.ComponentShardController, tag.Address(hostIdentity))
+	c.throttledLogger = log.With(c.throttledLogger, tag.ComponentShardController, tag.Address(hostIdentity))
+
 	c.acquireShards()
 	c.shutdownWG.Add(1)
 	go c.shardManagementPump()

service/history/shard/controller_test.go

@@ -43,6 +43,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 
+	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/cluster"
 	"go.temporal.io/server/common/dynamicconfig"
 	"go.temporal.io/server/common/log"
@@ -80,6 +81,41 @@ type (
 	}
 )
 
+func NewTestController(
+	engineFactory *MockEngineFactory,
+	config *configs.Config,
+	resource *resource.Test,
+	hostInfoProvider *resource.MockHostInfoProvider,
+) *ControllerImpl {
+	return &ControllerImpl{
+		config:                      config,
+		logger:                      resource.GetLogger(),
+		throttledLogger:             resource.GetThrottledLogger(),
+		contextTaggedLogger:         log.With(resource.GetLogger(), tag.ComponentShardController, tag.Address(resource.GetHostInfo().Identity())),
+		persistenceExecutionManager: resource.GetExecutionManager(),
+		persistenceShardManager:     resource.GetShardManager(),
+		clientBean:                  resource.GetClientBean(),
+		historyClient:               resource.GetHistoryClient(),
+		historyServiceResolver:      resource.GetHistoryServiceResolver(),
+		metricsClient:               resource.GetMetricsClient(),
+		payloadSerializer:           resource.GetPayloadSerializer(),
+		timeSource:                  resource.GetTimeSource(),
+		namespaceRegistry:           resource.GetNamespaceRegistry(),
+		saProvider:                  resource.GetSearchAttributesProvider(),
+		saMapper:                    resource.GetSearchAttributesMapper(),
+		clusterMetadata:             resource.GetClusterMetadata(),
+		archivalMetadata:            resource.GetArchivalMetadata(),
+		hostInfoProvider:            hostInfoProvider,
+
+		status:             common.DaemonStatusInitialized,
+		membershipUpdateCh: make(chan *membership.ChangedEvent, 10),
+		engineFactory:      engineFactory,
+		shutdownCh:         make(chan struct{}),
+		metricsScope:       resource.GetMetricsClient().Scope(metrics.HistoryShardControllerScope),
+		historyShards:      make(map[int32]*ContextImpl),
+	}
+}
+
 func TestShardControllerSuite(t *testing.T) {
 	s := new(controllerSuite)
 	suite.Run(t, s)
@@ -103,24 +139,10 @@ func (s *controllerSuite) SetupTest() {
 	s.logger = s.mockResource.Logger
 	s.config = tests.NewDynamicConfig()
 
-	s.shardController = NewController(
+	s.shardController = NewTestController(
 		s.mockEngineFactory,
 		s.config,
-		s.mockResource.Logger,
-		s.mockResource.GetThrottledLogger(),
-		s.mockResource.GetExecutionManager(),
-		s.mockResource.GetShardManager(),
-		s.mockResource.GetClientBean(),
-		s.mockResource.GetHistoryClient(),
-		s.mockResource.GetHistoryServiceResolver(),
-		s.mockResource.GetMetricsClient(),
-		s.mockResource.GetPayloadSerializer(),
-		s.mockResource.GetTimeSource(),
-		s.mockResource.GetNamespaceRegistry(),
-		s.mockResource.GetSearchAttributesProvider(),
-		s.mockResource.GetSearchAttributesMapper(),
-		s.mockResource.GetClusterMetadata(),
-		s.mockResource.GetArchivalMetadata(),
+		s.mockResource,
 		s.mockHostInfoProvider,
 	)
 }
@@ -476,24 +498,10 @@ func (s *controllerSuite) TestAcquireShardRenewLookupFailed() {
 func (s *controllerSuite) TestHistoryEngineClosed() {
 	numShards := int32(4)
 	s.config.NumberOfShards = numShards
-	s.shardController = NewController(
+	s.shardController = NewTestController(
 		s.mockEngineFactory,
 		s.config,
-		s.mockResource.Logger,
-		s.mockResource.GetThrottledLogger(),
-		s.mockResource.GetExecutionManager(),
-		s.mockResource.GetShardManager(),
-		s.mockResource.GetClientBean(),
-		s.mockResource.GetHistoryClient(),
-		s.mockResource.GetHistoryServiceResolver(),
-		s.mockResource.GetMetricsClient(),
-		s.mockResource.GetPayloadSerializer(),
-		s.mockResource.GetTimeSource(),
-		s.mockResource.GetNamespaceRegistry(),
-		s.mockResource.GetSearchAttributesProvider(),
-		s.mockResource.GetSearchAttributesMapper(),
-		s.mockResource.GetClusterMetadata(),
-		s.mockResource.GetArchivalMetadata(),
+		s.mockResource,
 		s.mockHostInfoProvider,
 	)
 	historyEngines := make(map[int32]*MockEngine)
@@ -586,24 +594,10 @@ func (s *controllerSuite) TestHistoryEngineClosed() {
 func (s *controllerSuite) TestShardControllerClosed() {
 	numShards := int32(4)
 	s.config.NumberOfShards = numShards
-	s.shardController = NewController(
+	s.shardController = NewTestController(
 		s.mockEngineFactory,
 		s.config,
-		s.mockResource.Logger,
-		s.mockResource.GetThrottledLogger(),
-		s.mockResource.GetExecutionManager(),
-		s.mockResource.GetShardManager(),
-		s.mockResource.GetClientBean(),
-		s.mockResource.GetHistoryClient(),
-		s.mockResource.GetHistoryServiceResolver(),
-		s.mockResource.GetMetricsClient(),
-		s.mockResource.GetPayloadSerializer(),
-		s.mockResource.GetTimeSource(),
-		s.mockResource.GetNamespaceRegistry(),
-		s.mockResource.GetSearchAttributesProvider(),
-		s.mockResource.GetSearchAttributesMapper(),
-		s.mockResource.GetClusterMetadata(),
-		s.mockResource.GetArchivalMetadata(),
+		s.mockResource,
 		s.mockHostInfoProvider,
 	)