Fix membermship resolver used in integration tests (uber#6210)

aditigautam78 · Aug 7, 2024 · f2d0afb · f2d0afb
1 parent c77db73
commit f2d0afb
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/host/membership_resolver.go b/host/membership_resolver.go
@@ -33,13 +33,13 @@ type simpleResolver struct {
 }
 
 // NewSimpleResolver returns a membership resolver interface
-func NewSimpleResolver(serviceName string, hosts map[string][]membership.HostInfo) membership.Resolver {
+func NewSimpleResolver(serviceName string, hosts map[string][]membership.HostInfo, currentHost membership.HostInfo) membership.Resolver {
 	resolvers := make(map[string]*simpleHashring, len(hosts))
 	for service, hostList := range hosts {
 		resolvers[service] = newSimpleHashring(hostList)
 	}
 	return &simpleResolver{
-		hostInfo:  hosts[serviceName][0],
+		hostInfo:  currentHost,
 		resolvers: resolvers,
 	}
 }

diff --git a/host/onebox.go b/host/onebox.go
@@ -527,7 +527,7 @@ func (c *cadenceImpl) startFrontend(hosts map[string][]membership.HostInfo, star
 	params.PProfInitializer = newPProfInitializerImpl(c.logger, c.FrontendPProfPort())
 	params.RPCFactory = c.newRPCFactory(service.Frontend, c.FrontendHost())
 	params.MetricScope = tally.NewTestScope(service.Frontend, make(map[string]string))
-	params.MembershipResolver = newMembershipResolver(params.Name, hosts)
+	params.MembershipResolver = newMembershipResolver(params.Name, hosts, c.FrontendHost())
 	params.ClusterMetadata = c.clusterMetadata
 	params.MessagingClient = c.messagingClient
 	params.MetricsClient = metrics.NewClient(params.MetricScope, service.GetMetricsServiceIdx(params.Name, c.logger))
@@ -609,7 +609,7 @@ func (c *cadenceImpl) startHistory(hosts map[string][]membership.HostInfo, start
 		params.PProfInitializer = newPProfInitializerImpl(c.logger, pprofPorts[i])
 		params.RPCFactory = c.newRPCFactory(service.History, hostport)
 		params.MetricScope = tally.NewTestScope(service.History, make(map[string]string))
-		params.MembershipResolver = newMembershipResolver(params.Name, hosts)
+		params.MembershipResolver = newMembershipResolver(params.Name, hosts, hostport)
 		params.ClusterMetadata = c.clusterMetadata
 		params.MessagingClient = c.messagingClient
 		params.MetricsClient = metrics.NewClient(params.MetricScope, service.GetMetricsServiceIdx(params.Name, c.logger))
@@ -688,7 +688,7 @@ func (c *cadenceImpl) startMatching(hosts map[string][]membership.HostInfo, star
 		params.PProfInitializer = newPProfInitializerImpl(c.logger, pprofPorts[i])
 		params.RPCFactory = c.newRPCFactory(service.Matching, hostport)
 		params.MetricScope = tally.NewTestScope(service.Matching, map[string]string{"matching-host": matchingHost})
-		params.MembershipResolver = newMembershipResolver(params.Name, hosts)
+		params.MembershipResolver = newMembershipResolver(params.Name, hosts, hostport)
 		params.ClusterMetadata = c.clusterMetadata
 		params.MetricsClient = metrics.NewClient(params.MetricScope, service.GetMetricsServiceIdx(params.Name, c.logger))
 		params.DynamicConfig = newIntegrationConfigClient(dynamicconfig.NewNopClient(), c.matchingDynCfgOverrides)
@@ -747,7 +747,7 @@ func (c *cadenceImpl) startWorker(hosts map[string][]membership.HostInfo, startW
 	params.PProfInitializer = newPProfInitializerImpl(c.logger, c.WorkerPProfPort())
 	params.RPCFactory = c.newRPCFactory(service.Worker, c.WorkerServiceHost())
 	params.MetricScope = tally.NewTestScope(service.Worker, make(map[string]string))
-	params.MembershipResolver = newMembershipResolver(params.Name, hosts)
+	params.MembershipResolver = newMembershipResolver(params.Name, hosts, c.WorkerServiceHost())
 	params.ClusterMetadata = c.clusterMetadata
 	params.MetricsClient = metrics.NewClient(params.MetricScope, service.GetMetricsServiceIdx(params.Name, c.logger))
 	params.DynamicConfig = newIntegrationConfigClient(dynamicconfig.NewNopClient(), c.workerDynCfgOverrides)
@@ -965,8 +965,8 @@ func copyPersistenceConfig(pConfig config.Persistence) (config.Persistence, erro
 	return pConfig, nil
 }
 
-func newMembershipResolver(serviceName string, hosts map[string][]membership.HostInfo) membership.Resolver {
-	return NewSimpleResolver(serviceName, hosts)
+func newMembershipResolver(serviceName string, hosts map[string][]membership.HostInfo, currentHost membership.HostInfo) membership.Resolver {
+	return NewSimpleResolver(serviceName, hosts, currentHost)
 }
 
 func newPProfInitializerImpl(logger log.Logger, port int) common.PProfInitializer {