From 44fb6c980a29f4410bc0c8a28a3052f606ed2378 Mon Sep 17 00:00:00 2001
From: David Reiss <david@temporal.io>
Date: Tue, 30 May 2023 20:46:19 -0400
Subject: [PATCH] Make server fx init more extendable

---
 temporal/fx.go     | 60 +++++++++++++++++++---------------------------
 temporal/server.go |  2 +-
 2 files changed, 25 insertions(+), 37 deletions(-)

diff --git a/temporal/fx.go b/temporal/fx.go
index 4827643d3c0..a9f820f7b33 100644
--- a/temporal/fx.go
+++ b/temporal/fx.go
@@ -129,12 +129,10 @@ type (
 	}
 )
 
-func NewServerFx(opts ...ServerOption) (*ServerFx, error) {
-	var s ServerFx
-	s.app = fx.New(
+var (
+	TopLevelModule = fx.Options(
 		pprof.Module,
 		fx.Provide(NewServerFxImpl),
-		fx.Supply(opts),
 		fx.Provide(ServerOptionsProvider),
 		TraceExportModule,
 
@@ -148,7 +146,14 @@ func NewServerFx(opts ...ServerOption) (*ServerFx, error) {
 		fx.Provide(ApplyClusterMetadataConfigProvider),
 		fx.Invoke(ServerLifetimeHooks),
 		FxLogAdapter,
+	)
+)
 
+func NewServerFx(topLevelModule fx.Option, opts ...ServerOption) (*ServerFx, error) {
+	var s ServerFx
+	s.app = fx.New(
+		topLevelModule,
+		fx.Supply(opts),
 		fx.Populate(&s.startupSynchronizationMode),
 		fx.Populate(&s.logger),
 	)
@@ -346,6 +351,17 @@ type (
 	}
 )
 
+func NewService(app *fx.App, serviceName primitives.ServiceName, logger log.Logger, stopChan chan struct{}) ServicesGroupOut {
+	return ServicesGroupOut{
+		Services: &ServicesMetadata{
+			app:         app,
+			serviceName: serviceName,
+			logger:      logger,
+			stopChan:    stopChan,
+		},
+	}
+}
+
 func HistoryServiceProvider(
 	params ServiceProviderParamsCommon,
 ) (ServicesGroupOut, error) {
@@ -393,14 +409,7 @@ func HistoryServiceProvider(
 		FxLogAdapter,
 	)
 
-	return ServicesGroupOut{
-		Services: &ServicesMetadata{
-			app:         app,
-			serviceName: serviceName,
-			logger:      params.Logger,
-			stopChan:    stopChan,
-		},
-	}, app.Err()
+	return NewService(app, serviceName, params.Logger, stopChan), app.Err()
 }
 
 func MatchingServiceProvider(
@@ -447,14 +456,7 @@ func MatchingServiceProvider(
 		FxLogAdapter,
 	)
 
-	return ServicesGroupOut{
-		Services: &ServicesMetadata{
-			app:         app,
-			serviceName: serviceName,
-			logger:      params.Logger,
-			stopChan:    stopChan,
-		},
-	}, app.Err()
+	return NewService(app, serviceName, params.Logger, stopChan), app.Err()
 }
 
 func FrontendServiceProvider(
@@ -531,14 +533,7 @@ func genericFrontendServiceProvider(
 		FxLogAdapter,
 	)
 
-	return ServicesGroupOut{
-		Services: &ServicesMetadata{
-			app:         app,
-			serviceName: serviceName,
-			logger:      params.Logger,
-			stopChan:    stopChan,
-		},
-	}, app.Err()
+	return NewService(app, serviceName, params.Logger, stopChan), app.Err()
 }
 
 func WorkerServiceProvider(
@@ -585,14 +580,7 @@ func WorkerServiceProvider(
 		FxLogAdapter,
 	)
 
-	return ServicesGroupOut{
-		Services: &ServicesMetadata{
-			app:         app,
-			serviceName: serviceName,
-			logger:      params.Logger,
-			stopChan:    stopChan,
-		},
-	}, app.Err()
+	return NewService(app, serviceName, params.Logger, stopChan), app.Err()
 }
 
 // ApplyClusterMetadataConfigProvider performs a config check against the configured persistence store for cluster metadata.
diff --git a/temporal/server.go b/temporal/server.go
index 59f43384840..58fc125d7df 100644
--- a/temporal/server.go
+++ b/temporal/server.go
@@ -66,5 +66,5 @@ var (
 
 // NewServer returns a new instance of server that serves one or many services.
 func NewServer(opts ...ServerOption) (Server, error) {
-	return NewServerFx(opts...)
+	return NewServerFx(TopLevelModule, opts...)
 }