Keep CreateThriftServiceFunc as before

examples/keyvalue/server/handlers.go

@@ -27,7 +27,6 @@ import (
 	"go.uber.org/fx/examples/keyvalue/kv"
 	kvs "go.uber.org/fx/examples/keyvalue/kv/keyvalueserver"
 	"go.uber.org/fx/service"
-	"go.uber.org/yarpc"
 	"go.uber.org/yarpc/api/transport"
 )
 
@@ -37,7 +36,7 @@ type YARPCHandler struct {
 	items map[string]string
 }
 
-func NewYARPCThriftHandler(_ service.Host, _ *yarpc.Dispatcher) ([]transport.Procedure, error) {
+func NewYARPCThriftHandler(_ service.Host) ([]transport.Procedure, error) {
 	return kvs.New(&YARPCHandler{items: map[string]string{}}), nil
 }
 

modules/rpc/thrift.go

@@ -25,12 +25,11 @@ import (
 	"go.uber.org/fx/service"
 
 	"github.com/pkg/errors"
-	"go.uber.org/yarpc"
 	"go.uber.org/yarpc/api/transport"
 )
 
 // CreateThriftServiceFunc creates a Thrift service from a service host
-type CreateThriftServiceFunc func(svc service.Host, dispatcher *yarpc.Dispatcher) ([]transport.Procedure, error)
+type CreateThriftServiceFunc func(svc service.Host) ([]transport.Procedure, error)
 
 // ThriftModule creates a Thrift Module from a service func
 func ThriftModule(hookup CreateThriftServiceFunc, options ...modules.Option) service.ModuleCreateFunc {

modules/rpc/thrift_test.go

@@ -22,7 +22,6 @@ package rpc
 
 import (
 	"errors"
-	"sync"
 	"testing"
 
 	"go.uber.org/fx/config"
@@ -49,15 +48,7 @@ func (h testHost) Config() config.Provider {
 func TestThriftModule_OK(t *testing.T) {
 	dig.Reset()
 
-	wg := sync.WaitGroup{}
-	wg.Add(1)
-	testInbounds := func(_ service.Host, dispatcher *yarpc.Dispatcher) ([]transport.Procedure, error) {
-		require.Equal(t, 2, len(dispatcher.Inbounds()))
-		wg.Done()
-		return nil, nil
-	}
-
-	chip := ThriftModule(testInbounds, modules.WithRoles("rescue"))
+	chip := ThriftModule(okCreate, modules.WithRoles("rescue"))
 	dale := ThriftModule(okCreate, modules.WithRoles("ranges"))
 	cfg := []byte(`
 modules:
@@ -87,7 +78,11 @@ modules:
 
 	testInitRunModule(t, goofy[0], mci)
 	testInitRunModule(t, gopher[0], mci)
-	wg.Wait()
+
+	// Dispatcher must be resolved in the default graph
+	var dispatcher *yarpc.Dispatcher
+	assert.NoError(t, dig.Resolve(&dispatcher))
+	assert.Equal(t, 2, len(dispatcher.Inbounds()))
 }
 
 func TestThriftModule_BadOptions(t *testing.T) {
@@ -134,12 +129,13 @@ func errorOption(_ *service.ModuleCreateInfo) error {
 	return errors.New("bad option")
 }
 
-func okCreate(_ service.Host, dispatcher *yarpc.Dispatcher) ([]transport.Procedure, error) {
-	return thrift.BuildProcedures(thrift.Service{
+func okCreate(_ service.Host) ([]transport.Procedure, error) {
+	reg := thrift.BuildProcedures(thrift.Service{
 		Name: "foo",
-	}), nil
+	})
+	return reg, nil
 }
 
-func badCreateService(service.Host, *yarpc.Dispatcher) ([]transport.Procedure, error) {
+func badCreateService(_ service.Host) ([]transport.Procedure, error) {
 	return nil, errors.New("can't create service")
 }

modules/rpc/yarpc.go

@@ -128,7 +128,7 @@ type dispatcherController struct {
 
 	configs    []*yarpcConfig
 	handlers   []handlerWithDispatcher
-	dispatcher *yarpc.Dispatcher
+	dispatcher yarpc.Dispatcher
 }
 
 // Adds the config to the controller.
@@ -150,7 +150,7 @@ func (c *dispatcherController) appendHandler(handler handlerWithDispatcher) {
 // Apply handlers to the dispatcher.
 func (c *dispatcherController) applyHandlers() error {
 	for _, h := range c.handlers {
-		if err := h(c.dispatcher); err != nil {
+		if err := h(&c.dispatcher); err != nil {
 			return err
 		}
 	}
@@ -197,17 +197,19 @@ func (c *dispatcherController) Start(host service.Host) error {
 		_dispatcherMu.Lock()
 		defer _dispatcherMu.Unlock()
 
-		if c.dispatcher, err = _dispatcherFn(host, cfg); err != nil {
+		var d *yarpc.Dispatcher
+		if d, err = _dispatcherFn(host, cfg); err != nil {
 			c.startError = err
 			return
 		}
 
+		c.dispatcher = *d
 		if err := c.applyHandlers(); err != nil {
 			c.startError = err
 			return
 		}
 
-		c.startError = _starterFn(c.dispatcher)
+		c.startError = _starterFn(&c.dispatcher)
 	})
 
 	return c.startError
@@ -305,11 +307,16 @@ func newYARPCModule(
 		if errCr := dig.Register(module.controller); errCr != nil {
 			return nil, errs.Wrap(errCr, "can't register a dispatcher controller")
 		}
+
+		// Register dispatcher
+		if err := dig.Register(&module.controller.dispatcher); err != nil {
+			return nil, errs.Wrap(err, "unable to register the dispatcher")
+		}
 	}
 
 	module.controller.addConfig(module.config)
 	module.controller.appendHandler(func(dispatcher *yarpc.Dispatcher) error {
-		t, err := reg(mi.Host, dispatcher)
+		t, err := reg(mi.Host)
 		if err != nil {
 			return err
 		}