[common][rpc/membership] Fx integration fixes (#6955)

* [common][rpc/membership] Fx integration fixes

Author: 3vilhamster

common/membership/membershipfx/membershipfx.go

@@ -31,6 +31,7 @@ import (
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/membership"
 	"github.com/uber/cadence/common/metrics"
+	"github.com/uber/cadence/common/rpc"
 	"github.com/uber/cadence/common/service"
 )
 
@@ -41,6 +42,7 @@ type buildMembershipParams struct {
 	fx.In
 
 	Clock         clock.TimeSource
+	RPCFactory    rpc.Factory
 	PeerProvider  membership.PeerProvider
 	Logger        log.Logger
 	MetricsClient metrics.Client
@@ -70,10 +72,25 @@ func buildMembership(params buildMembershipParams) (buildMembershipResult, error
 		return buildMembershipResult{}, fmt.Errorf("create resolver: %w", err)
 	}
 
-	params.Lifecycle.Append(fx.StartStopHook(resolver.Start, resolver.Stop))
+	params.Lifecycle.Append(fx.StartStopHook(startResolver(resolver, params.RPCFactory), resolver.Stop))
 
 	return buildMembershipResult{
 		Rings:    rings,
 		Resolver: resolver,
 	}, nil
 }
+
+func startResolver(resolver membership.Resolver, rpcFactory rpc.Factory) func() error {
+	return func() error {
+		err := rpcFactory.Start(resolver)
+		if err != nil {
+			return fmt.Errorf("start rpc factory: %w", err)
+		}
+		err = rpcFactory.GetDispatcher().Start()
+		if err != nil {
+			return fmt.Errorf("start rpc factory dispatcher: %w", err)
+		}
+		resolver.Start()
+		return nil
+	}
+}

common/membership/membershipfx/memebershipfx_test.go

@@ -28,25 +28,42 @@ import (
 	"go.uber.org/fx"
 	"go.uber.org/fx/fxtest"
 	"go.uber.org/mock/gomock"
+	"go.uber.org/yarpc"
 
 	"github.com/uber/cadence/common/clock"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/log/testlogger"
 	"github.com/uber/cadence/common/membership"
 	"github.com/uber/cadence/common/metrics"
+	"github.com/uber/cadence/common/rpc"
+	"github.com/uber/cadence/common/service"
 )
 
 func TestFxStartStop(t *testing.T) {
 	app := fxtest.New(t, fx.Provide(func() appParams {
 		ctrl := gomock.NewController(t)
 		provider := membership.NewMockPeerProvider(ctrl)
+		provider.EXPECT().Start()
+		provider.EXPECT().Stop()
+		for _, s := range service.ListWithRing {
+			provider.EXPECT().Subscribe(s, gomock.Any()).Return(nil)
+			provider.EXPECT().GetMembers(s).Return([]membership.HostInfo{}, nil)
+			// this is also called by every ring, but could be called multiple times.
+			provider.EXPECT().Stop()
+		}
+		factory := rpc.NewMockFactory(ctrl)
+		factory.EXPECT().Start(gomock.Any()).Return(nil)
+		factory.EXPECT().GetDispatcher().Return(yarpc.NewDispatcher(yarpc.Config{
+			Name: "membership_test",
+		}))
 		return appParams{
 			Clock:         clock.NewMockedTimeSource(),
 			PeerProvider:  provider,
 			Logger:        testlogger.New(t),
 			MetricsClient: metrics.NewNoopMetricsClient(),
+			RPCFactory:    factory,
 		}
-	}), Module)
+	}), Module, fx.Invoke(func(resolver membership.Resolver) {}))
 	app.RequireStart().RequireStop()
 }
 
@@ -57,4 +74,5 @@ type appParams struct {
 	PeerProvider  membership.PeerProvider
 	Logger        log.Logger
 	MetricsClient metrics.Client
+	RPCFactory    rpc.Factory
 }

common/peerprovider/ringpopprovider/ringpopfx/ringpopfx.go

@@ -54,6 +54,5 @@ func buildRingpopProvider(params Params) (membership.PeerProvider, error) {
 	if err != nil {
 		return nil, err
 	}
-	params.Lifecycle.Append(fx.StartStopHook(provider.Start, provider.Stop))
 	return provider, nil
 }

common/peerprovider/ringpopprovider/ringpopfx/ringpopfx_test.go

@@ -25,13 +25,17 @@ package ringpopfx
 import (
 	"testing"
 
+	"github.com/stretchr/testify/require"
+	"github.com/uber/tchannel-go"
 	"go.uber.org/fx"
 	"go.uber.org/fx/fxtest"
 	"go.uber.org/mock/gomock"
 
 	"github.com/uber/cadence/common/config"
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/log/testlogger"
+	"github.com/uber/cadence/common/membership"
+	"github.com/uber/cadence/common/peerprovider/ringpopprovider"
 	"github.com/uber/cadence/common/rpc"
 )
 
@@ -41,23 +45,32 @@ func TestFxApp(t *testing.T) {
 			func() testSetupParams {
 				ctrl := gomock.NewController(t)
 				factory := rpc.NewMockFactory(ctrl)
-				factory.EXPECT().GetTChannel().Return(nil)
+				tch, err := tchannel.NewChannel("test-ringpop", nil)
+				require.NoError(t, err)
+				factory.EXPECT().GetTChannel().Return(tch)
 
 				return testSetupParams{
 					Service:    "test",
 					Logger:     testlogger.New(t),
 					RPCFactory: factory,
+					Config: config.Config{
+						Ringpop: ringpopprovider.Config{
+							Name:           "test-ringpop",
+							BootstrapMode:  ringpopprovider.BootstrapModeHosts,
+							BootstrapHosts: []string{"127.0.0.1:7933", "127.0.0.1:7934", "127.0.0.1:7935"},
+						},
+					},
 				}
 			}),
-		Module,
+		Module, fx.Invoke(func(provider membership.PeerProvider) {}),
 	)
 	app.RequireStart().RequireStop()
 }
 
 type testSetupParams struct {
 	fx.Out
 
-	Service       string `name:"service"`
+	Service       string `name:"service-full-name"`
 	Config        config.Config
 	ServiceConfig config.Service
 	Logger        log.Logger

common/rpc/factory.go

@@ -52,6 +52,8 @@ var (
 
 // Factory is an implementation of rpc.Factory interface
 type FactoryImpl struct {
+	startOnce      sync.Once
+	stopOnce       sync.Once
 	maxMessageSize int
 	channel        tchannel.Channel
 	dispatcher     *yarpc.Dispatcher
@@ -182,33 +184,38 @@ func (d *FactoryImpl) GetMaxMessageSize() int {
 }
 
 func (d *FactoryImpl) Start(peerLister PeerLister) error {
-	d.peerLister = peerLister
-	// subscribe to membership changes for history and matching. This is needed to update the peers for rpc
-	for _, svc := range servicesToTalkP2P {
-		ch := make(chan *membership.ChangedEvent, 1)
-		if err := d.peerLister.Subscribe(svc, factoryComponentName, ch); err != nil {
-			return fmt.Errorf("rpc factory failed to subscribe to membership updates for svc: %v, err: %v", svc, err)
+	var err error
+	d.startOnce.Do(func() {
+		d.peerLister = peerLister
+		// subscribe to membership changes for history and matching. This is needed to update the peers for rpc
+		for _, svc := range servicesToTalkP2P {
+			ch := make(chan *membership.ChangedEvent, 1)
+			if err = d.peerLister.Subscribe(svc, factoryComponentName, ch); err != nil {
+				err = fmt.Errorf("rpc factory failed to subscribe to membership updates for svc: %v, err: %w", svc, err)
+				return
+			}
+			d.wg.Add(1)
+			go d.listenMembershipChanges(svc, ch)
 		}
-		d.wg.Add(1)
-		go d.listenMembershipChanges(svc, ch)
-	}
-
-	return nil
+	})
+	return err
 }
 
 func (d *FactoryImpl) Stop() error {
-	d.logger.Info("stopping rpc factory")
+	d.stopOnce.Do(func() {
+		d.logger.Info("stopping rpc factory")
 
-	for _, svc := range servicesToTalkP2P {
-		if err := d.peerLister.Unsubscribe(svc, factoryComponentName); err != nil {
-			d.logger.Error("rpc factory failed to unsubscribe from membership updates", tag.Error(err), tag.Service(svc))
+		for _, svc := range servicesToTalkP2P {
+			if err := d.peerLister.Unsubscribe(svc, factoryComponentName); err != nil {
+				d.logger.Error("rpc factory failed to unsubscribe from membership updates", tag.Error(err), tag.Service(svc))
+			}
 		}
-	}
 
-	d.cancelFn()
-	d.wg.Wait()
+		d.cancelFn()
+		d.wg.Wait()
 
-	d.logger.Info("stopped rpc factory")
+		d.logger.Info("stopped rpc factory")
+	})
 	return nil
 }
 

common/rpc/rpcfx/rpcfx.go

@@ -69,6 +69,20 @@ type factoryParams struct {
 
 func buildFactory(p factoryParams) rpc.Factory {
 	res := rpc.NewFactory(p.Logger, p.RPCParams)
-	p.Lifecycle.Append(fx.StartStopHook(res.GetDispatcher().Start, res.GetDispatcher().Stop))
+	p.Lifecycle.Append(fx.StopHook(rpcStopper(res)))
 	return res
 }
+
+func rpcStopper(factory rpc.Factory) func() error {
+	return func() error {
+		err := factory.GetDispatcher().Stop()
+		if err != nil {
+			return fmt.Errorf("dispatcher stop: %w", err)
+		}
+		err = factory.Stop()
+		if err != nil {
+			return fmt.Errorf("factory stop: %w", err)
+		}
+		return nil
+	}
+}